<commit_before># -*- coding: utf-8 -*-
import glob
import os
import pickle
import numpy as np
import numpy.testing as npt
import scipy.signal

import matplotlib.pyplot as plt
import matplotlib.animation
import seaborn as sns

from antlia.record import load_file
from antlia.util import reduce_runs
from antlia.trial import Trial


LIDAR_NUM_ANGLES = 1521
LIDAR_FOV_DEG = 190
LIDAR_SAMPLE_RATE = 20
LIDAR_ANGLES = np.linspace( # in radians
    (90 - LIDAR_FOV_DEG/2)*np.pi/180,
    (90 + LIDAR_FOV_DEG/2)*np.pi/180,
    LIDAR_NUM_ANGLES
)

"""LIDAR datatype format is:
    (
        timestamp (long),
        flag (bool saved as int),
        accelerometer[3] (double),
        gps[3] (double),
        distance[LIDAR_NUM_ANGLES] (long),
    )

    'int' and 'long' are the same size on the raspberry pi (32 bits).
"""
LIDAR_RECORD_DTYPE = np.dtype(','.join(
    ['i4'] + ['i4'] + 3*['f8'] + 3*['f8'] + LIDAR_NUM_ANGLES*['i4']
    )
)
LIDAR_CONVERTED_DTYPE = np.dtype([
    ('time', 'f8'),
    ('sync', 'f8'),
    ('accelerometer x', 'f8'), # body x-axis may not be aligned with inertial!
    ('accelerometer y', 'f8'), # body y-axis may not be aligned with inertial!
    ('accelerometer z', 'f8'), # body z-axis may not be aligned with inertial!
    ('gps', 'f8', (3,)),
    ('distance', 'f8', (LIDAR_NUM_ANGLES,)),
])


def __get_record_files(extension):
    record_path = os.path.join(os.path.dirname(__file__),
                               r'../../data/lidar/')
    return sorted(glob.glob('{}*{}'.format(record_path, extension)))


LIDAR_RECORD_FILES =  __get_record_files('.bin')
BICYCLE_RECORD_FILES = __get_record_files('.csv')
assert LIDAR_RECORD_FILES, "No LIDAR records found!"
assert BICYCLE_RECORD_FILES, "No bicycle records found!"


def _get_lidar_records(convert_dtype=True):
    """Returns a list of LIDAR records.

    Parameters:
    convert_dtype: bool, Datatype of LIDAR records.
                   True    - dtype = LIDAR_CONVERTED_DTYPE
                   False   - dtype = LIDAR_RECORD_DTYPE
    Returns:
    records: list of array_like
    """
    float_dtype = np.dtype(','.join(len(LIDAR_RECORD_DTYPE) * ['f8']))

    records = []
    for filename in LIDAR_RECORD_FILES:
        with open(filename, 'rb') as f:
            x = np.fromfile(f, LIDAR_RECORD_DTYPE)
            if convert_dtype:
                # copy and convert data to new dtype
                y = x.astype(float_dtype)
                # create view which shares the same underlying memory
                x = y.view(LIDAR_CONVERTED_DTYPE).view(np.recarray)
                # convert timestamp to time and start at zero
                x.time /= 1000
                x.time -= x.time[0]
                # flip sync value to be active high
                x.sync = np.invert(x.sync.astype(bool)).astype(x.sync.dtype)
                # TODO: convert accelerometer
                # TODO: convert gps
                # convert distance from millimeters to meters
                x.distance /= 1000
            records.append(x)
    return records


def _get_bicycle_records():
    # load calibration data
    calibration_path = os.path.join(os.path.dirname(__file__),
                                    '..', r'config.p')
    with open(calibration_path, 'rb') as f:
        calibration = pickle.load(f)

    records = []
    for filename in BICYCLE_RECORD_FILES:
        r = load_file(filename, calibration['convbike'])
        records.append(r)
    return records

class LidarRecord(np.recarray):
    def __new__(cls, input_array):
        obj = np.asarray(input_array).view(cls)
        return obj

    def __array_finalize__(self, obj):
        if obj is None:
            return

    def frame_index(self, time):
        """Get lidar frame(s) at time.

        Parameters:
        time: float or indexing function
        """
        if not callable(time):
            # return a one element LidarRecord
            return np.where(self.time >= time)[0][:1]

        return np.where(time(self.time))[0]

    def frame(self, time):
        index = self.frame_index(time)
        return self[index]

    def cartesian(self, xlim=None, ylim=None):
        rho = self.distance
        x = rho*np.cos(LIDAR_ANGLES)
        y = rho*np.sin(LIDAR_ANGLES)

        index = np.zeros(x.shape, dtype=bool)
        if xlim is not None:
            xmin, xmax = xlim
            index = index | (x < xmin) | (x > xmax)
        if ylim is not None:
            ymin, ymax = ylim
            index = index | (y < ymin) | (y > ymax)
        return (np.ma.masked_array(x, index),
                np.ma.masked_array(y, index))

    def animate(self, xlim=None, ylim=None, speedup=1, color=None, **kwargs):
        fig, ax = plt.subplots(**kwargs)

        line, = ax.plot([], [], 'X', animated=True)
        time_template = 'time = {:0.03f} s'
        if speedup != 1:
            time_template += ' (x{:0.1f})'.format(speedup)
        if speedup <= 0:
            raise ValueError

        time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
        ax.set_xlabel('x position [meters]')
        ax.set_ylabel('y position [meters]')

        def init():
            if xlim is not None:
                ax.set_xlim(xlim)
            if ylim is not None:
                ax.set_ylim(ylim)
            line.set_data([], [])
            time_text.set_text('')
            return line, time_text<commit_msg>Add argument to change plot args in lidar animation<commit_after># -*- coding: utf-8 -*-
import glob
import os
import pickle
import numpy as np
import numpy.testing as npt
import scipy.signal

import matplotlib.pyplot as plt
import matplotlib.animation
import seaborn as sns

from antlia.record import load_file
from antlia.util import reduce_runs
from antlia.trial import Trial


LIDAR_NUM_ANGLES = 1521
LIDAR_FOV_DEG = 190
LIDAR_SAMPLE_RATE = 20
LIDAR_ANGLES = np.linspace( # in radians
    (90 - LIDAR_FOV_DEG/2)*np.pi/180,
    (90 + LIDAR_FOV_DEG/2)*np.pi/180,
    LIDAR_NUM_ANGLES
)

"""LIDAR datatype format is:
    (
        timestamp (long),
        flag (bool saved as int),
        accelerometer[3] (double),
        gps[3] (double),
        distance[LIDAR_NUM_ANGLES] (long),
    )

    'int' and 'long' are the same size on the raspberry pi (32 bits).
"""
LIDAR_RECORD_DTYPE = np.dtype(','.join(
    ['i4'] + ['i4'] + 3*['f8'] + 3*['f8'] + LIDAR_NUM_ANGLES*['i4']
    )
)
LIDAR_CONVERTED_DTYPE = np.dtype([
    ('time', 'f8'),
    ('sync', 'f8'),
    ('accelerometer x', 'f8'), # body x-axis may not be aligned with inertial!
    ('accelerometer y', 'f8'), # body y-axis may not be aligned with inertial!
    ('accelerometer z', 'f8'), # body z-axis may not be aligned with inertial!
    ('gps', 'f8', (3,)),
    ('distance', 'f8', (LIDAR_NUM_ANGLES,)),
])


def __get_record_files(extension):
    record_path = os.path.join(os.path.dirname(__file__),
                               r'../../data/lidar/')
    return sorted(glob.glob('{}*{}'.format(record_path, extension)))


LIDAR_RECORD_FILES =  __get_record_files('.bin')
BICYCLE_RECORD_FILES = __get_record_files('.csv')
assert LIDAR_RECORD_FILES, "No LIDAR records found!"
assert BICYCLE_RECORD_FILES, "No bicycle records found!"


def _get_lidar_records(convert_dtype=True):
    """Returns a list of LIDAR records.

    Parameters:
    convert_dtype: bool, Datatype of LIDAR records.
                   True    - dtype = LIDAR_CONVERTED_DTYPE
                   False   - dtype = LIDAR_RECORD_DTYPE
    Returns:
    records: list of array_like
    """
    float_dtype = np.dtype(','.join(len(LIDAR_RECORD_DTYPE) * ['f8']))

    records = []
    for filename in LIDAR_RECORD_FILES:
        with open(filename, 'rb') as f:
            x = np.fromfile(f, LIDAR_RECORD_DTYPE)
            if convert_dtype:
                # copy and convert data to new dtype
                y = x.astype(float_dtype)
                # create view which shares the same underlying memory
                x = y.view(LIDAR_CONVERTED_DTYPE).view(np.recarray)
                # convert timestamp to time and start at zero
                x.time /= 1000
                x.time -= x.time[0]
                # flip sync value to be active high
                x.sync = np.invert(x.sync.astype(bool)).astype(x.sync.dtype)
                # TODO: convert accelerometer
                # TODO: convert gps
                # convert distance from millimeters to meters
                x.distance /= 1000
            records.append(x)
    return records


def _get_bicycle_records():
    # load calibration data
    calibration_path = os.path.join(os.path.dirname(__file__),
                                    '..', r'config.p')
    with open(calibration_path, 'rb') as f:
        calibration = pickle.load(f)

    records = []
    for filename in BICYCLE_RECORD_FILES:
        r = load_file(filename, calibration['convbike'])
        records.append(r)
    return records

class LidarRecord(np.recarray):
    def __new__(cls, input_array):
        obj = np.asarray(input_array).view(cls)
        return obj

    def __array_finalize__(self, obj):
        if obj is None:
            return

    def frame_index(self, time):
        """Get lidar frame(s) at time.

        Parameters:
        time: float or indexing function
        """
        if not callable(time):
            # return a one element LidarRecord
            return np.where(self.time >= time)[0][:1]

        return np.where(time(self.time))[0]

    def frame(self, time):
        index = self.frame_index(time)
        return self[index]

    def cartesian(self, xlim=None, ylim=None):
        rho = self.distance
        x = rho*np.cos(LIDAR_ANGLES)
        y = rho*np.sin(LIDAR_ANGLES)

        index = np.zeros(x.shape, dtype=bool)
        if xlim is not None:
            xmin, xmax = xlim
            index = index | (x < xmin) | (x > xmax)
        if ylim is not None:
            ymin, ymax = ylim
            index = index | (y < ymin) | (y > ymax)
        return (np.ma.masked_array(x, index),
                np.ma.masked_array(y, index))

    def animate(self, xlim=None, ylim=None,
                speedup=1, color=None, plot_kwargs={}, **kwargs):
        fig, ax = plt.subplots(**kwargs)

        plot_kwargs.setdefault('linestyle', ' ')
        plot_kwargs.setdefault('marker', '.')

        line, = ax.plot([], [], animated=True, **plot_kwargs)
        time_template = 'time = {:0.03f} s'
        if speedup != 1:
            time_template += ' (x{:0.1f})'.format(speedup)
        if speedup <= 0:
            raise ValueError

        time_text = ax.text(0.05, 0.9, '', transform=ax.transAxes)
        ax.set_xlabel('x position [meters]')
        ax.set_ylabel('y position [meters]')

        def init():
            if xlim is not None:
                ax.set_xlim(xlim)
            if ylim is not None:
                ax.set_ylim(ylim)
            line.set_data([], [])
            time_text.set_text('')
            return line, time_text
{'fraction_non_alphanumeric': 0.12098246895544193, 'fraction_numerical': 0.012235208181154127, 'mean_word_length': 2.1960898745258244, 'pattern_counts': {'":': 0, '<': 9, '<?xml version=': 0, '>': 9, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 2, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '3270094', 'n_tokens_mistral': 3804, 'n_tokens_neox': 3650, 'n_words': 896}