MicroPython driver for LIS2DH/LIS3DH with FIFO support

lis2dh.py

"""
LIS2DH accelerometer driver.

Should also work with LIS3DH

License: MIT
Copyright: Soundsensing AS
"""
from machine import I2C, Pin
import struct
import time
import math

# Register and bitfield defs
#
REG_WHO_AM_I = 0x0F
REG_CTRL_REG1 = 0x20
REG_CTRL_REG2 = 0x21
REG_CTRL_REG3 = 0x22
REG_CTRL_REG4 = 0x23
REG_CTRL_REG5 = 0x24
REG_CTRL_REG6 = 0x25
REG_STATUS = 0x27
REG_OUT_X_L = 0x28 # followed by high byte, then Y,Z

REG_FIFO_CTRL = 0x2E
REG_FIFO_SRC = 0x2F

DEV_ID = 0x33
MULTI_READ = 0x80

DATA_RATES = {
    0: 0 << 4, # power-down,
    100: 5 << 4,
    200: 6 << 4,
    400: 7 << 4,
    #1620: 8 << 4, # low-power mode only
    1344: 9 << 4, # also 5376 in low-power mode, not handled
}
SCALES = {
    2: 0 << 4,
    4: 1 << 4,
    8: 2 << 4,
    16: 3 << 4,
}


class LIS3DH:
    def __init__(self, i2c,
            addr=24, irq=3, odr=100, full_scale=2):
        self.i2c = i2c
        self.addr = addr
        self.odr = odr
        self.full_scale = full_scale

    def read_register(self, reg):
        """Read a register value (1 byte)"""
        data = self.i2c.readfrom_mem(self.addr, reg, 1)
        value = data[0]
        return value

    def write_register(self, reg, value):
        """Write a register value (1 byte)"""
        data = bytearray((value,))
        self.i2c.writeto_mem(self.addr, reg, data)

    def device_check(self):
        who = self.read_register(REG_WHO_AM_I)
        if who != DEV_ID:
            raise ValueError("Unknown device id")
        # TODO: also run self-check

    def get_fifo_count(self):
        """
        Return the number of samples ready in the FIFO
        """
        FF_FSS_MASK = 0x1F
        fifo_count = self.read_register(REG_FIFO_SRC) & FF_FSS_MASK
        return fifo_count

    def read_samples_into(self, buf):
        """
        Read accelerometer samples from the FIFO

        Data is returned little-endian, unless the BLE bit is set in CTRL_REG4
        ?? two’s complement left-justified 

        NOTE: caller is responsible for ensuring that enough samples are ready.
        Typically by calling get_fifo_count() first
        """
        n_bytes = len(buf)
        if (n_bytes % 6) != 0:
            raise ValueError("Buffer should be a multiple of 6")
        samples = n_bytes // 6
        if samples > 31:
            raise ValueError("Requested samples exceeds FIFO capacity")

        # MULTI_READ is a lis2dh specific marker that enables auto-increment
        self.i2c.readfrom_mem_into(self.addr, (REG_OUT_X_L | MULTI_READ), buf)

    def set_odr(self, odr : int):
        XEN = 1 << 0
        YEN = 1 << 1
        ZEN = 1 << 2

        # Enable each axis and set ODR mode
        data_rate = DATA_RATES[odr]
        value = self.read_register(REG_CTRL_REG1)
        value = data_rate | XEN | YEN | ZEN
        self.write_register(REG_CTRL_REG1, value)

    def set_mode(self):
        # XXX: only high-resolution supported, not low-power or normal
        HR = 1 << 3 # high-resolution mode
        value = self.read_register(REG_CTRL_REG4)
        value |= HR
        self.write_register(REG_CTRL_REG4, value)

    def set_scale(self, scale):
        value = self.read_register(REG_CTRL_REG4)
        value |= SCALES[scale]
        self.write_register(REG_CTRL_REG4, value)

    def set_highpass(self, enable):
        # High-pass cutoff frequencies. Not described in dataset :/
        # These definitions are from the portable driver by ST, lis2dh12_hpcf_t
        # https://github.com/STMicroelectronics/lis2dh12-pid/blob/master/lis2dh12_reg.h 
        HPCF_AGRESSIVE = 0
        HPCF_STRONG = 1
        HPCF_MEDIUM = 2
        HPCF_LIGHT = 3

        # High-pass mode. lis2dh12_hpm_t in ST driver
        HPM_NORMAL = 2

        high_pass_cutoff = HPCF_LIGHT
        high_pass_mode = HPM_NORMAL
        FDS_BIT = 3

        value = self.read_register(REG_CTRL_REG2)
        value &= (~0b11111000)
        value |= (high_pass_mode << 6)
        value |= (high_pass_cutoff << 4)
        if enable:
            # send filtered data to output/FIFO
            value |= (1 << FDS_BIT)
        else:
            # internal high-pass is bypassed
            pass
        print('set-highpass', hex(value))
        self.write_register(REG_CTRL_REG2, value)

    def fifo_enable(self, enable):
        BDU = 1 << 7 # block data update
        FIFO_EN = 1 << 6

        # FIFO modes
        FM_BYPASS = 0 << 6 # FIFO disabled. Also used to clear FIFO
        FM_FIFO = 1 << 6 # stops when FIFO full
        FM_STREAM = 2 << 6 # overwrites oldest when FIFO full

        if enable:
            value = self.read_register(REG_CTRL_REG4)
            value |= BDU
            self.write_register(REG_CTRL_REG4, value)

            value = self.read_register(REG_CTRL_REG5)
            value |= FIFO_EN
            self.write_register(REG_CTRL_REG5, value)

            value = self.read_register(REG_FIFO_CTRL)
            value |= FM_STREAM
            self.write_register(REG_FIFO_CTRL, value)

        else:
            value = FM_BYPASS
            self.write_register(REG_FIFO_CTRL, value)


    def start(self):
        print("lis2dh-start")
        self.set_mode()
        self.set_odr(self.odr)
        self.set_scale(self.full_scale)
        self.fifo_enable(True)

    def stop(self):
        print("lis2dh-stop")
        self.fifo_enable(False) # clear/disable FIFO
        self.set_odr(0) # put into power-down


def record_samples(accel, buf, chunk=25, overrun=31):

    start_time = time.ticks_ms()
    n_samples = len(buf) // 6
    n_chunks = math.ceil(n_samples / chunk)
    total_samples = 0    

    #print('record-samples-start', len(buf), n_samples, n_chunks)

    chunk_number = 0
    while chunk_number < n_chunks:

        chunk_samples = None
        count = accel.get_fifo_count()
        if chunk_number != 0 and count >= overrun:
            raise ValueError("FIFO overrun")
        if count > chunk:
            s = chunk_number*chunk*6
            e = min(s + chunk*6, buf)
            chunk_buf = memoryview(buf)[s:e]
            accel.read_samples_into(chunk_buf)
            chunk_samples = len(chunk_buf) // 6
            total_samples += chunk_samples
            chunk_number += 1

        yield chunk_samples
    
    assert total_samples == n_samples, (total_samples, n_samples)
    end_time = time.ticks_ms()
    dur = time.ticks_diff(end_time, start_time)


if __name__ == "__main__":

    i2c = I2C(0, sda=Pin(25), scl=Pin(0))

    print("i2c-scan")
    found = i2c.scan()
    for f in found:
        print(f)

    # address is normally 24 or 25
    samplerate = 400
    ACCELEROMETER_SAMPLE_INTERVAL = 0.0
    SLAVE_ADDR = 24

    accel = LIS3DH(i2c, addr=SLAVE_ADDR, odr=samplerate)

    SAMPLE_DURATION = 2.0
    n_samples = int(SAMPLE_DURATION*samplerate)
    samples = bytearray((0 for _ in range(6*n_samples)))

    last_sample = time.ticks_ms() / 1000.0

    def accel_interpret(buf):
        sarray = struct.unpack('hhh', buf)
        gs = [ (v/16.0)*1.0 for v in sarray ]
        return gs

    while True:

        if (time.ticks_ms()/1000.0) > last_sample + ACCELEROMETER_SAMPLE_INTERVAL:

            #print('accel-sample-start')
            accel.device_check()
            accel.set_highpass(True)
            accel.start()

            sampler = record_samples(accel, samples)
            for s in sampler:
                time.sleep_ms(1)

            print('raw-start', accel_interpret(samples[0:6]))
            print('raw-end', accel_interpret(samples[-6:]))

            accel.stop()