Bergmann_2022_DataRepository/04_analysis/trial_envelope.py

from collections import namedtuple
import numpy as np
import sliding1d as sliding

def interpolate(vec, min_size=10):
    valid = ~np.isnan(vec)
    if np.count_nonzero(valid) < min_size:
        return np.empty(vec.size) * np.nan
    t     = np.arange(vec.size)
    return np.interp(t, t[valid], vec[valid])

def whisker(trial, side="left", radius_sample=10, smooth=True):
    return Envelope.whisker(trial, side=side, radius_sample=radius_sample, smooth=smooth)

class Envelope(namedtuple("_Envelope", ("time", "raw", "bottom", "top"))):
    @classmethod
    def whisker(cls, trial, side="left", radius_sample=10, smooth=True):
        vec  = interpolate(trial.tracking[f"{side}_whisker_angle_deg"])
        vec  = (vec - vec.min()) / (vec.max() - vec.min())
        time = np.array(trial.tracking["time"])
        return cls.compute(time, vec, radius_sample=radius_sample, smooth=smooth)
    
    @classmethod
    def compute(cls, time, vec, radius_sample=10, smooth=True):
        bottom = sliding.nanmin(vec, radius_sample)
        top    = sliding.nanmax(vec, radius_sample)
        if smooth == True:
            bottom = sliding.nanmean(bottom, radius_sample)
            top    = sliding.nanmean(top, radius_sample)
        return cls(time, vec, bottom, top)
    
    @property
    def amplitude(self):
        return self.top - self.bottom
    
    def with_range(self, rng):
        return self.__class__(self.time[rng],
                              self.raw[rng],
                              self.bottom[rng],
                              self.top[rng])