Neuron dropout

features/__init__.py

@@ -15,7 +15,7 @@
 from .plexon_features import PlexonBMI, RelayPlexon, RelayPlexByte
 from .hdf_features import SaveHDF
 from .video_recording_features import SingleChannelVideo, E3Video
-from .bmi_task_features import NormFiringRates
+from .bmi_task_features import NormFiringRates, RandomUnitDropout
 from .arduino_features import PlexonSerialDIORowByte
 from .blackrock_features import BlackrockBMI
 from .blackrock_features import RelayBlackrockByte
@@ -79,6 +79,7 @@
     eye_calibration=EyeCalibration, 
     force_sensor=ForceControl,
     show_fixation_progress=Progressbar_fixation,
+    random_unit_dropout=RandomUnitDropout,
     clda_kfrml=CLDA_KFRML_IntendedVelocity
 )
 

features/bmi_task_features.py

@@ -4,13 +4,76 @@
 import time
 import numpy as np
 from riglib.experiment import traits, experiment
+from riglib.bmi import clda
 
 ###### CONSTANTS
 sec_per_min = 60
 
 ########################################################################################################
 # Decoder/BMISystem add-ons
 ########################################################################################################
+class RandomUnitDropout(traits.HasTraits):
+    '''
+    Randomly removes units from the decoder. Does not work with CLDA turned on. Units are removed at the
+    end of the delay period on each trial and replaced when the trial ends (either in reward or penalty).
+    The same units will be dropped on repeated trials. The units to drop are specified in the 
+    `unit_drop_groups` attribute by a list of lists of unit indices. The `unit_drop_targets` attribute
+    specifies the target indices on which to drop each group of units.
+    '''
+
+    unit_drop_prob = traits.Float(0, desc="Probability of dropping a group of units from the decoder")
+    unit_drop_groups = traits.Array(value=[[0, 1], [2]], desc="Groups of unit indices to drop from the decoder one at a time")
+    unit_drop_targets = traits.Array(value=[1, 2], desc="Target indices on which to drop groups of units from the decoder")
+
+    def init(self):
+        super().init()
+        self.decoder_units_dropped = np.ones((len(self.decoder.units),), dtype='bool')
+        new_dtype = np.dtype(self.trial_dtype.descr + [('decoder_units_dropped', '?', self.decoder_units_dropped.shape)])
+        self.trial_dtype = new_dtype
+        self.unit_drop_group_idx = 0
+
+        # Save a copy of the decoder
+        self.decoder_orig = self.decoder.copy()
+
+    def _start_wait(self):
+
+        # Decide which units to drop in this trial but don't actually drop them yet
+        if (self.gen_indices[self.target_index] == self.unit_drop_targets[self.unit_drop_group_idx] and 
+            np.random.rand() < self.unit_drop_prob):
+            self.decoder_units_dropped = np.isin(range(len(self.decoder.units)), self.unit_drop_groups[self.unit_drop_group_idx])
+
+            # Update the group for next trial
+            self.unit_drop_group_idx = (self.unit_drop_group_idx + 1) % len(self.unit_drop_groups)
+        else:
+            self.decoder_units_dropped = np.zeros((len(self.decoder.units),), dtype='bool')
+
+        # Update the trial record
+        self.trial_record['decoder_units_dropped'] = self.decoder_units_dropped
+        super()._start_wait()
+
+    def _start_targ_transition(self):
+        '''
+        Override the decoder to drop random units. Keep a record of what's going on in the `trial` data.
+        '''
+        super()._start_targ_transition()
+        if self.target_index + 1 < self.chain_length and np.any(self.decoder_units_dropped):
+            if hasattr(self.decoder.filt, 'C'):
+                self.decoder.filt.C[self.decoder_units_dropped, :] = 0
+            elif hasattr(self.decoder.filt, 'unit_to_state'):
+                self.decoder.filt.unit_to_state[:, self.decoder_units_dropped] = 0
+
+    def _reset_decoder(self):
+        self.decoder = self.decoder_orig.copy()
+
+    def _increment_tries(self):
+        super()._increment_tries()
+        self._reset_decoder()
+
+    def _start_reward(self):
+        super()._start_reward()
+        self._reset_decoder()
+
+
 class NormFiringRates(traits.HasTraits):
     ''' Docstring '''
 

features/clda_features.py

@@ -41,4 +41,4 @@ def create_updater(self):
         alter the decoder parameters to better match the intention estimates.
         '''
         self.updater = clda.KFRML(self.clda_batch_time, self.clda_update_half_life)
-        self.updater.init(self.decoder)
+        self.updater.init(self.decoder)

setup.py

@@ -2,7 +2,7 @@
 
 setuptools.setup(
     name="aolab-bmi3d",
-    version="1.0.5",
+    version="1.0.6",
     author="Lots of people",
     description="electrophysiology experimental rig library",
     packages=setuptools.find_packages(),