From 4fdc0b89e699d04f1ad62de956df73ec0b6fce3f Mon Sep 17 00:00:00 2001
From: katherineperks <kperks@uw.edu>
Date: Thu, 23 Nov 2023 16:21:45 -0800
Subject: [PATCH 01/14] init; to-do update lindecoder.py

---
 features/bmi_task_features.py | 24 ++++++++++++++++++++++++
 1 file changed, 24 insertions(+)

diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index f3557c512..c28e3275e 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -11,6 +11,30 @@
 ########################################################################################################
 # Decoder/BMISystem add-ons
 ########################################################################################################
+class UnitDropout(traits.HasTraits):
+    '''
+    Randomly removes units from the decoder. Compatible with lindecoder.py only.
+    '''
+
+    def init(self):
+        super().init()
+        self.decoder_unit_mask = np.ones((len(self.decoder.units),), dtype='bool')
+        new_dtype = np.dtype(self.trial_dtype.descr + [('decoder_unit_mask', '?', self.decoder_unit_mask.shape)])
+        self.trial_dtype = new_dtype
+        if not hasattr(self.decoder.filt, "update_mask"):
+            raise ValueError(f"FeatureDropout feature not supported by decoder {repr(self.decoder)}")
+
+    def _start_wait(self):
+        '''
+        Override the decoder to drop a random unit. Keep a record of what's going on in the `trial` data.
+        '''
+        self.decoder_unit_mask[:] = True
+        self.decoder_unit_mask[np.random.choice(len(self.decoder_unit_mask), 1)] = False
+        self.decoder.filt.update_mask(self.decoder_unit_mask)
+        self.trial_record['decoder_unit_mask'] = self.decoder_unit_mask
+        super()._start_wait()
+
+
 class NormFiringRates(traits.HasTraits):
     ''' Docstring '''
     

From 00a81b68be3cb3388012d36b4249ec481205821c Mon Sep 17 00:00:00 2001
From: Leo Scholl <leo.scholl@gmail.com>
Date: Wed, 13 Mar 2024 15:40:10 -0700
Subject: [PATCH 02/14] make work

---
 features/__init__.py          |  3 ++-
 features/bmi_task_features.py | 24 ++++++++++++++++--------
 2 files changed, 18 insertions(+), 9 deletions(-)

diff --git a/features/__init__.py b/features/__init__.py
index 8793a8b8c..b8d81f728 100644
--- a/features/__init__.py
+++ b/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
@@ -78,6 +78,7 @@
     eye_calibration=EyeCalibration, 
     force_sensor=ForceControl,
     show_fixation_progress=Progressbar_fixation,
+    random_unit_dropout=RandomUnitDropout,
 )
 
 # >>> features.built_in_features['autostart'].__module__
diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index c28e3275e..33878322d 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -11,9 +11,10 @@
 ########################################################################################################
 # Decoder/BMISystem add-ons
 ########################################################################################################
-class UnitDropout(traits.HasTraits):
+class RandomUnitDropout(traits.HasTraits):
     '''
-    Randomly removes units from the decoder. Compatible with lindecoder.py only.
+    Randomly removes units from the decoder. Requires the decoder have a `unit_drop_prob` attribute
+    describing the probability that each unit will be dropped on a given trial.
     '''
 
     def init(self):
@@ -21,16 +22,23 @@ def init(self):
         self.decoder_unit_mask = np.ones((len(self.decoder.units),), dtype='bool')
         new_dtype = np.dtype(self.trial_dtype.descr + [('decoder_unit_mask', '?', self.decoder_unit_mask.shape)])
         self.trial_dtype = new_dtype
-        if not hasattr(self.decoder.filt, "update_mask"):
-            raise ValueError(f"FeatureDropout feature not supported by decoder {repr(self.decoder)}")
+        if not hasattr(self.decoder, "unit_drop_prob"):
+            self.decoder.unit_drop_prob = np.zeros((len(self.decoder.units),), dtype='float')
+        
+        # Save a copy of the mFR and sdFR from the decoder
+        self.decoder_mFR = self.decoder.mFR.copy()
+        self.decoder_sdFR = self.decoder.sdFR.copy()
 
     def _start_wait(self):
         '''
-        Override the decoder to drop a random unit. Keep a record of what's going on in the `trial` data.
+        Override the decoder to drop random units. Keep a record of what's going on in the `trial` data.
         '''
-        self.decoder_unit_mask[:] = True
-        self.decoder_unit_mask[np.random.choice(len(self.decoder_unit_mask), 1)] = False
-        self.decoder.filt.update_mask(self.decoder_unit_mask)
+        self.decoder_unit_mask = np.random.rand(len(self.decoder.units)) < self.decoder.unit_drop_prob
+        mFR_curr = self.decoder_mFR.copy()
+        sdFR_curr = self.decoder_sdFR.copy()
+        mFR_curr[self.decoder_unit_mask] = 0
+        sdFR_curr[self.decoder_unit_mask] = 0
+        self.decoder.init_zscore(self, mFR_curr, sdFR_curr)
         self.trial_record['decoder_unit_mask'] = self.decoder_unit_mask
         super()._start_wait()
 

From 180f8a61517a78487c248616c73648d41aba69e8 Mon Sep 17 00:00:00 2001
From: Leo Scholl <leo.scholl@gmail.com>
Date: Wed, 13 Mar 2024 15:42:57 -0700
Subject: [PATCH 03/14] add default prob

---
 riglib/bmi/bmi.py | 3 +++
 1 file changed, 3 insertions(+)

diff --git a/riglib/bmi/bmi.py b/riglib/bmi/bmi.py
index ec461a5b8..aac0189f2 100644
--- a/riglib/bmi/bmi.py
+++ b/riglib/bmi/bmi.py
@@ -512,6 +512,9 @@ def __init__(self, filt, units, ssm, binlen=0.1, n_subbins=1, tslice=[-1,-1], ca
         self.bminum = int(self.binlen/(1/call_rate))
         self.spike_counts = np.zeros([len(units), 1])
 
+        # For the unit drop feature - set some/all units to nonzero probability
+        self.unit_drop_prob = np.zeros((len(self.units),), dtype='float')
+
         self.set_call_rate(call_rate)
 
         self._pickle_init()

From 497c220d54cbb1ec35c1078ece633c5d5d8eecb0 Mon Sep 17 00:00:00 2001
From: Leo Scholl <leo.scholl@gmail.com>
Date: Wed, 13 Mar 2024 15:48:35 -0700
Subject: [PATCH 04/14] remove redundant variables

---
 features/bmi_task_features.py | 20 +++++++++-----------
 1 file changed, 9 insertions(+), 11 deletions(-)

diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index 33878322d..9aff0030f 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -19,27 +19,25 @@ class RandomUnitDropout(traits.HasTraits):
 
     def init(self):
         super().init()
-        self.decoder_unit_mask = np.ones((len(self.decoder.units),), dtype='bool')
-        new_dtype = np.dtype(self.trial_dtype.descr + [('decoder_unit_mask', '?', self.decoder_unit_mask.shape)])
+        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
         if not hasattr(self.decoder, "unit_drop_prob"):
+            print("WARNING: RandomUnitDropout feature requires the decoder to have a unit_drop_prob attribute")
             self.decoder.unit_drop_prob = np.zeros((len(self.decoder.units),), dtype='float')
         
-        # Save a copy of the mFR and sdFR from the decoder
+        # Save a copy of the mFR from the decoder
         self.decoder_mFR = self.decoder.mFR.copy()
-        self.decoder_sdFR = self.decoder.sdFR.copy()
 
     def _start_wait(self):
         '''
         Override the decoder to drop random units. Keep a record of what's going on in the `trial` data.
         '''
-        self.decoder_unit_mask = np.random.rand(len(self.decoder.units)) < self.decoder.unit_drop_prob
-        mFR_curr = self.decoder_mFR.copy()
-        sdFR_curr = self.decoder_sdFR.copy()
-        mFR_curr[self.decoder_unit_mask] = 0
-        sdFR_curr[self.decoder_unit_mask] = 0
-        self.decoder.init_zscore(self, mFR_curr, sdFR_curr)
-        self.trial_record['decoder_unit_mask'] = self.decoder_unit_mask
+        self.decoder_units_dropped = np.random.rand(len(self.decoder.units)) < self.decoder.unit_drop_prob
+        mFR_drop = self.decoder_mFR.copy()
+        mFR_drop[self.decoder_units_dropped] = 0
+        self.decoder.init_zscore(self, mFR_drop, self.decoder.sdFR)
+        self.trial_record['decoder_units_dropped'] = self.decoder_units_dropped
         super()._start_wait()
 
 

From faeeebf8e783d839159e6c880588a8bec11c5800 Mon Sep 17 00:00:00 2001
From: katherineperks <leo.scholl@gmail.com>
Date: Fri, 22 Mar 2024 14:48:12 -0700
Subject: [PATCH 05/14] not working

---
 built_in_tasks/bmimultitasks.py |  2 +-
 features/bmi_task_features.py   |  8 ++++++++
 features/clda_features.py       |  1 +
 riglib/bmi/clda.py              | 13 +++++++++++++
 4 files changed, 23 insertions(+), 1 deletion(-)

diff --git a/built_in_tasks/bmimultitasks.py b/built_in_tasks/bmimultitasks.py
index 9ea543172..6bb42ed6d 100644
--- a/built_in_tasks/bmimultitasks.py
+++ b/built_in_tasks/bmimultitasks.py
@@ -303,7 +303,7 @@ def _start_None(self):
         # Optionally save a new decoder zscored from this task
         if (not self.save_zscore) or (self.saveid is None):
             return
-        
+
         if not (np.all(self.decoder.mFR == 0) and np.all(self.decoder.sdFR) == 1):
             return
 
diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index 9aff0030f..7f355dd2d 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -4,6 +4,7 @@
 import time
 import numpy as np
 from riglib.experiment import traits, experiment
+from riglib.bmi import clda
 
 ###### CONSTANTS
 sec_per_min = 60
@@ -40,6 +41,13 @@ def _start_wait(self):
         self.trial_record['decoder_units_dropped'] = self.decoder_units_dropped
         super()._start_wait()
 
+    def create_learner(self):
+        '''Always be ready to implement neuron dropout'''        
+        self.learner = clda.EagerLearner()
+
+    def create_decoder(self):
+        self.updater = clda.UnitDropout(self.decoder_units_dropped)
+
 
 class NormFiringRates(traits.HasTraits):
     ''' Docstring '''
diff --git a/features/clda_features.py b/features/clda_features.py
index 6236d44e1..91dd95336 100644
--- a/features/clda_features.py
+++ b/features/clda_features.py
@@ -12,6 +12,7 @@ class CLDA_KFRML_IntendedVelocity(traits.HasTraits):
     clda_update_batch_time = traits.Float(60, desc="How frequently to update weights [s]")
     clda_update_half_life = traits.Float(120, desc="Half-life for exponential decay [s] to combine with previous weights.") #[s]
     clda_learner_batch_time = traits.Float(60, desc="How much data to update the learner with [s]") # Samples to update intended kinematics with
+    
     def create_learner(self):
         '''
         The "learner" uses knowledge of the task goals to determine the "intended"
diff --git a/riglib/bmi/clda.py b/riglib/bmi/clda.py
index ed55c9ce3..a6604af3f 100644
--- a/riglib/bmi/clda.py
+++ b/riglib/bmi/clda.py
@@ -210,6 +210,14 @@ def get_batch(self):
         '''DumbLearner never has any 'batch' data to retrieve'''
         raise NotImplementedError
 
+class EagerLearner(DumbLearner):
+    '''
+    Eager learner is always ready to learn but never learns anything. Used to make updates to the decoder that are determined externally.
+    '''
+    def is_ready(self):
+        '''Eager learner is always ready to change'''
+        return True
+
 class FeedbackControllerLearner(Learner):
     '''
     An intention estimator where the subject is assumed to operate like a state feedback controller
@@ -1018,6 +1026,11 @@ def calc(self, intended_kin=None, spike_counts=None, decoder=None, half_life=120
 
         return {'filt.C': self.C_est}
 
+class UnitDropout(Updater):
+    def calc(self):
+        pass
+
+
 
 ###############################
 ##### Deprecated updaters #####

From 8b3459c55e5043b85797f6c92a1bf4da976e09aa Mon Sep 17 00:00:00 2001
From: katherineperks <touchi@uw.edu>
Date: Wed, 3 Apr 2024 15:58:13 -0700
Subject: [PATCH 06/14] test

---
 features/clda_features.py | 14 ++++++++++-
 riglib/bmi/clda.py        | 49 +++++++++++++++++++++++++++++++--------
 2 files changed, 52 insertions(+), 11 deletions(-)

diff --git a/features/clda_features.py b/features/clda_features.py
index 91dd95336..04cf7956b 100644
--- a/features/clda_features.py
+++ b/features/clda_features.py
@@ -42,4 +42,16 @@ def create_updater(self):
         '''
         update_batch_size = int(self.clda_update_batch_time/self.decoder.binlen)
         self.updater = clda.KFRML(update_batch_size, self.clda_update_half_life)
-        self.updater.init(self.decoder)
\ No newline at end of file
+        self.updater.init(self.decoder)
+
+
+# class UnitDropout(traits.HasTraits):
+#     units_to_drop = traits.Float([], desc="Which units should have a chance to be dropped")
+#     unit_drop_prob = traits.Float(0, desc="Probability of dropping each unit on a given trial")
+#     targets_to_drop_units = traits.Float(None, desc="Which targets should unit dropping be applied on")
+
+#     def create_learner(self):
+#         self.learner  = clda.TargetSpecificLearner(self.targets_to_drop_units)
+
+#     def create_updater(self):
+#         self.updater = clda.UnitDropout(self.units_to_drop, self.unit_drop_prob)
\ No newline at end of file
diff --git a/riglib/bmi/clda.py b/riglib/bmi/clda.py
index a6604af3f..de9d9355b 100644
--- a/riglib/bmi/clda.py
+++ b/riglib/bmi/clda.py
@@ -85,6 +85,7 @@ def __init__(self, batch_size, *args, **kwargs):
         self.input_state_index = 0
         self.reset()
 
+
     def disable(self):
         '''Set a flag to disable forming intention estimates from new incoming data'''
         self.enabled = False
@@ -210,13 +211,20 @@ def get_batch(self):
         '''DumbLearner never has any 'batch' data to retrieve'''
         raise NotImplementedError
 
-class EagerLearner(DumbLearner):
-    '''
-    Eager learner is always ready to learn but never learns anything. Used to make updates to the decoder that are determined externally.
-    '''
-    def is_ready(self):
-        '''Eager learner is always ready to change'''
-        return True
+# class TargetSpecificLearner(Learner):
+#     '''
+#     Ready to learn only to specific targets
+#     '''
+#     def __init__(self, update_target_idx):
+#         self.update_target_idx = update_target_idx
+
+
+#     def is_ready(self):
+#         ''''''
+#         _is_ready = self.gen in self.update_target_idx
+
+
+#         return _is_ready
 
 class FeedbackControllerLearner(Learner):
     '''
@@ -1026,9 +1034,30 @@ def calc(self, intended_kin=None, spike_counts=None, decoder=None, half_life=120
 
         return {'filt.C': self.C_est}
 
-class UnitDropout(Updater):
-    def calc(self):
-        pass
+# class UnitDropout(Updater):
+#     def __init__(self, units_dropped=None, unit_drop_prob=0): 
+#         self.units_dropped = units_dropped
+#         self.unit_drop_prob = unit_drop_prob
+
+#     def calc(self, units_dropped=None, unit_drop_prob=0):
+#         '''
+#         Args:
+#             units_dropped (): If none, randomly drop any unit with the max(||C||_2). Else specific arry input
+#             unit_drop_prob (float): Probability of dropping any given unit
+#             targets (): If none, apply on any target, else only on trials to specific target(s)
+#         '''
+        
+#         # C_new = self.S * R_inv
+#         # C = copy.deepcopy(decoder.filt.C)
+#         # C[np.ix_(self.adapting_inds, self.state_adapting_inds)] = C_new[np.ix_(self.adapting_inds, self.state_adapting_inds)]
+        
+#         # Q = (1./self.ESS) * (self.T - self.S*C.T)
+#         # if hasattr(self, 'stable_inds_mesh'):
+#         #     if len(self.stable_inds) > 0:
+#         #         print('stable inds mesh: ', self.stable_inds, self.stable_inds_mesh)
+#         #         Q_old = decoder.filt.Q[self.stable_inds_mesh].copy()
+#         #         Q[self.stable_inds_mesh] = Q_old
+#         pass
 
 
 

From b3e28e3d8c5831b6145185ed7f1cc26a85cbc6d0 Mon Sep 17 00:00:00 2001
From: Katherine <leo.scholl@gmail.com>
Date: Thu, 4 Apr 2024 13:48:06 -0700
Subject: [PATCH 07/14] revert

---
 riglib/bmi/bmi.py  |  3 ---
 riglib/bmi/clda.py | 42 +-----------------------------------------
 2 files changed, 1 insertion(+), 44 deletions(-)

diff --git a/riglib/bmi/bmi.py b/riglib/bmi/bmi.py
index 5e8a534c4..df7bef55d 100644
--- a/riglib/bmi/bmi.py
+++ b/riglib/bmi/bmi.py
@@ -512,9 +512,6 @@ def __init__(self, filt, units, ssm, binlen=0.1, n_subbins=1, tslice=[-1,-1], ca
         self.bminum = int(self.binlen/(1/call_rate))
         self.spike_counts = np.zeros([len(units), 1])
 
-        # For the unit drop feature - set some/all units to nonzero probability
-        self.unit_drop_prob = np.zeros((len(self.units),), dtype='float')
-
         self.set_call_rate(call_rate)
 
         self._pickle_init()
diff --git a/riglib/bmi/clda.py b/riglib/bmi/clda.py
index 3b71c7e4a..9d2fe5dee 100644
--- a/riglib/bmi/clda.py
+++ b/riglib/bmi/clda.py
@@ -211,20 +211,6 @@ def get_batch(self):
         '''DumbLearner never has any 'batch' data to retrieve'''
         raise NotImplementedError
 
-# class TargetSpecificLearner(Learner):
-#     '''
-#     Ready to learn only to specific targets
-#     '''
-#     def __init__(self, update_target_idx):
-#         self.update_target_idx = update_target_idx
-
-
-#     def is_ready(self):
-#         ''''''
-#         _is_ready = self.gen in self.update_target_idx
-
-
-#         return _is_ready
 
 class FeedbackControllerLearner(Learner):
     '''
@@ -1034,32 +1020,6 @@ def calc(self, intended_kin=None, spike_counts=None, decoder=None, half_life=120
 
         return {'filt.C': self.C_est}
 
-# class UnitDropout(Updater):
-#     def __init__(self, units_dropped=None, unit_drop_prob=0): 
-#         self.units_dropped = units_dropped
-#         self.unit_drop_prob = unit_drop_prob
-
-#     def calc(self, units_dropped=None, unit_drop_prob=0):
-#         '''
-#         Args:
-#             units_dropped (): If none, randomly drop any unit with the max(||C||_2). Else specific arry input
-#             unit_drop_prob (float): Probability of dropping any given unit
-#             targets (): If none, apply on any target, else only on trials to specific target(s)
-#         '''
-        
-#         # C_new = self.S * R_inv
-#         # C = copy.deepcopy(decoder.filt.C)
-#         # C[np.ix_(self.adapting_inds, self.state_adapting_inds)] = C_new[np.ix_(self.adapting_inds, self.state_adapting_inds)]
-        
-#         # Q = (1./self.ESS) * (self.T - self.S*C.T)
-#         # if hasattr(self, 'stable_inds_mesh'):
-#         #     if len(self.stable_inds) > 0:
-#         #         print('stable inds mesh: ', self.stable_inds, self.stable_inds_mesh)
-#         #         Q_old = decoder.filt.Q[self.stable_inds_mesh].copy()
-#         #         Q[self.stable_inds_mesh] = Q_old
-#         pass
-
-
 
 ###############################
 ##### Deprecated updaters #####
@@ -1074,7 +1034,7 @@ class KFSmoothbatch(Updater):
     def __init__(self, batch_time, half_life):
         '''
         Constructor for KFSmoothbatch
-
+F
         Parameters
         ----------
         batch_time : float

From c374c7608a15448430bf6c8fbb3c1c038873ad84 Mon Sep 17 00:00:00 2001
From: Leo Scholl <leo.scholl@gmail.com>
Date: Thu, 4 Apr 2024 14:35:40 -0700
Subject: [PATCH 08/14] something like this

---
 features/bmi_task_features.py | 62 ++++++++++++++++++++++-------------
 setup.py                      |  2 +-
 2 files changed, 41 insertions(+), 23 deletions(-)

diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index 7f355dd2d..53637fe71 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -14,40 +14,58 @@
 ########################################################################################################
 class RandomUnitDropout(traits.HasTraits):
     '''
-    Randomly removes units from the decoder. Requires the decoder have a `unit_drop_prob` attribute
-    describing the probability that each unit will be dropped on a given trial.
+    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
-        if not hasattr(self.decoder, "unit_drop_prob"):
-            print("WARNING: RandomUnitDropout feature requires the decoder to have a unit_drop_prob attribute")
-            self.decoder.unit_drop_prob = np.zeros((len(self.decoder.units),), dtype='float')
-        
-        # Save a copy of the mFR from the decoder
-        self.decoder_mFR = self.decoder.mFR.copy()
+        self.unit_drop_group_idx = 0
+
+        # Save a copy of the decoder
+        self.decoder_orig = self.decoder.copy()
 
     def _start_wait(self):
-        '''
-        Override the decoder to drop random units. Keep a record of what's going on in the `trial` data.
-        '''
-        self.decoder_units_dropped = np.random.rand(len(self.decoder.units)) < self.decoder.unit_drop_prob
-        mFR_drop = self.decoder_mFR.copy()
-        mFR_drop[self.decoder_units_dropped] = 0
-        self.decoder.init_zscore(self, mFR_drop, self.decoder.sdFR)
-        self.trial_record['decoder_units_dropped'] = self.decoder_units_dropped
         super()._start_wait()
 
-    def create_learner(self):
-        '''Always be ready to implement neuron dropout'''        
-        self.learner = clda.EagerLearner()
-
-    def create_decoder(self):
-        self.updater = clda.UnitDropout(self.decoder_units_dropped)
+        # Decide which units to drop in this trial but don't actually drop them yet
+        if np.random.rand() < self.unit_drop_prob:
+            self.unit_drop_group_idx = (self.unit_drop_group_idx + 1) % len(self.unit_drop_groups)
+            self.decoder_units_dropped = np.isin(range(len(self.decoder.units)), self.unit_drop_groups[self.unit_drop_group_idx])
+            self.trial_record['decoder_units_dropped'] = self.decoder_units_dropped
 
+    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 '''
diff --git a/setup.py b/setup.py
index 0a3272a35..09aa4f970 100644
--- a/setup.py
+++ b/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(),

From 675c1838fb270211b921be73261ebcf9e3e4325c Mon Sep 17 00:00:00 2001
From: Leo Scholl <leo.scholl@gmail.com>
Date: Thu, 4 Apr 2024 15:18:59 -0700
Subject: [PATCH 09/14] more working but not tested

---
 features/bmi_task_features.py | 15 +++++++++++----
 riglib/bmi/clda.py            |  4 +---
 2 files changed, 12 insertions(+), 7 deletions(-)

diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index 53637fe71..5e7e63ea7 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -36,13 +36,20 @@ def init(self):
         self.decoder_orig = self.decoder.copy()
 
     def _start_wait(self):
-        super()._start_wait()
 
         # Decide which units to drop in this trial but don't actually drop them yet
-        if np.random.rand() < self.unit_drop_prob:
-            self.unit_drop_group_idx = (self.unit_drop_group_idx + 1) % len(self.unit_drop_groups)
+        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])
-            self.trial_record['decoder_units_dropped'] = self.decoder_units_dropped
+
+            # 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):
         '''
diff --git a/riglib/bmi/clda.py b/riglib/bmi/clda.py
index 9d2fe5dee..bed019fb9 100644
--- a/riglib/bmi/clda.py
+++ b/riglib/bmi/clda.py
@@ -85,7 +85,6 @@ def __init__(self, batch_size, *args, **kwargs):
         self.input_state_index = 0
         self.reset()
 
-
     def disable(self):
         '''Set a flag to disable forming intention estimates from new incoming data'''
         self.enabled = False
@@ -211,7 +210,6 @@ def get_batch(self):
         '''DumbLearner never has any 'batch' data to retrieve'''
         raise NotImplementedError
 
-
 class FeedbackControllerLearner(Learner):
     '''
     An intention estimator where the subject is assumed to operate like a state feedback controller
@@ -1034,7 +1032,7 @@ class KFSmoothbatch(Updater):
     def __init__(self, batch_time, half_life):
         '''
         Constructor for KFSmoothbatch
-F
+
         Parameters
         ----------
         batch_time : float

From db97b49ddd7c873c4b91a167a11a6724c1faf13a Mon Sep 17 00:00:00 2001
From: Katherine <leo.scholl@gmail.com>
Date: Fri, 5 Apr 2024 09:36:31 -0700
Subject: [PATCH 10/14] ryan comments

---
 features/bmi_task_features.py | 8 ++++----
 1 file changed, 4 insertions(+), 4 deletions(-)

diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index 5e7e63ea7..0cc705792 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -22,8 +22,8 @@ class RandomUnitDropout(traits.HasTraits):
     '''
 
     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")
+    unit_drop_groups = traits.Array(value=[[1, 256], [2]], desc="Groups of channels to drop from the decoder one at a time")
+    unit_drop_targets = traits.Array(value=[1, [7, 8]], desc="Target indices on which to drop groups of units from the decoder")
 
     def init(self):
         super().init()
@@ -38,9 +38,9 @@ def init(self):
     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 
+        if (self.gen_indices[self.target_index] in 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])
+            self.decoder_units_dropped = np.isin(self.decoder.channels, 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)

From ca21277e324243b726e1f71326ca0bb6fe296245 Mon Sep 17 00:00:00 2001
From: Katherine <leo.scholl@gmail.com>
Date: Fri, 5 Apr 2024 10:41:10 -0700
Subject: [PATCH 11/14] good enough

---
 features/bmi_task_features.py | 55 ++++++++++++++++++++---------------
 1 file changed, 32 insertions(+), 23 deletions(-)

diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index 0cc705792..b2ca92e54 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -4,7 +4,7 @@
 import time
 import numpy as np
 from riglib.experiment import traits, experiment
-from riglib.bmi import clda
+import copy
 
 ###### CONSTANTS
 sec_per_min = 60
@@ -14,31 +14,33 @@
 ########################################################################################################
 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.
+    Randomly removes units from the decoder. 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. Does not work with CLDA turned on. 
     '''
 
     unit_drop_prob = traits.Float(0, desc="Probability of dropping a group of units from the decoder")
-    unit_drop_groups = traits.Array(value=[[1, 256], [2]], desc="Groups of channels to drop from the decoder one at a time")
-    unit_drop_targets = traits.Array(value=[1, [7, 8]], desc="Target indices on which to drop groups of units from the decoder")
+    unit_drop_groups = traits.Array(value=[[1, 256]], desc="Groups of channels to drop from the decoder one at a time")
+    unit_drop_targets = traits.Array(value=[1], 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.add_dtype('decoder_units_dropped', '?', self.decoder_units_dropped.shape)
         self.unit_drop_group_idx = 0
+        super().init()
 
         # Save a copy of the decoder
-        self.decoder_orig = self.decoder.copy()
+        self.decoder_orig = copy.deepcopy(self.decoder)
+        self.reportstats['Units dropped'] = '' # Runtime stat displayed on the UI
 
-    def _start_wait(self):
 
+    def _start_wait(self):
+        super()._start_wait()
+        
         # Decide which units to drop in this trial but don't actually drop them yet
-        if (self.gen_indices[self.target_index] in self.unit_drop_targets[self.unit_drop_group_idx] and 
+        if (self.gen_indices[self.target_index] in np.array(self.unit_drop_targets[self.unit_drop_group_idx]) and 
             np.random.rand() < self.unit_drop_prob):
             self.decoder_units_dropped = np.isin(self.decoder.channels, self.unit_drop_groups[self.unit_drop_group_idx])
 
@@ -47,31 +49,38 @@ def _start_wait(self):
         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 self.target_index == -1:
+
+            # Came from a penalty state
+            pass
+        elif 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
+            self.task_data['decoder_units_dropped'] = self.decoder_units_dropped
+            self.reportstats['Units dropped'] = str(self.decoder.channels[self.decoder_units_dropped])
 
-    def _reset_decoder(self):
-        self.decoder = self.decoder_orig.copy()
+    def _reset_decoder_units(self):
+        if hasattr(self.decoder.filt, 'C'):
+            self.decoder.filt.C = self.decoder_orig.filt.C
+        elif hasattr(self.decoder.filt, 'unit_to_state'):
+            self.decoder.filt.unit_to_state = self.decoder_orig.filt.unit_to_state
+        self.task_data['decoder_units_dropped'] = np.zeros((len(self.decoder.units),), dtype='bool')
+        self.reportstats['Units dropped'] = '[]'
 
     def _increment_tries(self):
         super()._increment_tries()
-        self._reset_decoder()
+        self._reset_decoder_units()
 
     def _start_reward(self):
         super()._start_reward()
-        self._reset_decoder()
+        self._reset_decoder_units()
     
 
 class NormFiringRates(traits.HasTraits):

From abc25e39c3e317196c7533d207dbfb0bb6d1ebfb Mon Sep 17 00:00:00 2001
From: Katherine <leo.scholl@gmail.com>
Date: Fri, 5 Apr 2024 11:10:33 -0700
Subject: [PATCH 12/14] change defaults

---
 features/bmi_task_features.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index b2ca92e54..f424e4cff 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -22,8 +22,8 @@ class RandomUnitDropout(traits.HasTraits):
     '''
 
     unit_drop_prob = traits.Float(0, desc="Probability of dropping a group of units from the decoder")
-    unit_drop_groups = traits.Array(value=[[1, 256]], desc="Groups of channels to drop from the decoder one at a time")
-    unit_drop_targets = traits.Array(value=[1], desc="Target indices on which to drop groups of units from the decoder")
+    unit_drop_groups = traits.List(value=[], desc="Groups of channels to drop from the decoder one at a time")
+    unit_drop_targets = traits.List(value=[], desc="Target indices on which to drop groups of units from the decoder")
 
     def init(self):
         self.decoder_units_dropped = np.ones((len(self.decoder.units),), dtype='bool')

From 1bf6524208c22870b5b4980cd48f43c50fda6735 Mon Sep 17 00:00:00 2001
From: Katherine <leo.scholl@gmail.com>
Date: Fri, 5 Apr 2024 11:44:27 -0700
Subject: [PATCH 13/14] fix array

---
 db/tracker/json_param.py      |  2 ++
 features/bmi_task_features.py | 14 ++++++++++++--
 2 files changed, 14 insertions(+), 2 deletions(-)

diff --git a/db/tracker/json_param.py b/db/tracker/json_param.py
index a048d4424..b8befe36c 100644
--- a/db/tracker/json_param.py
+++ b/db/tracker/json_param.py
@@ -83,6 +83,8 @@ def norm_trait(trait, value):
             record = Model.objects.get(pk=value)
             value = record.get()
         # Otherwise, let's hope it's already an instance
+    elif ttype == 'Array':
+        value = np.array(value)
     elif ttype == 'DataFile':
         # Similar to Instance traits, except we always know to use models.DataFile as the database table to look up the primary key
         from . import models
diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index f424e4cff..d23bf2a21 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -1,6 +1,7 @@
 '''
 BMI task features
 '''
+import ast
 import time
 import numpy as np
 from riglib.experiment import traits, experiment
@@ -22,8 +23,8 @@ class RandomUnitDropout(traits.HasTraits):
     '''
 
     unit_drop_prob = traits.Float(0, desc="Probability of dropping a group of units from the decoder")
-    unit_drop_groups = traits.List(value=[], desc="Groups of channels to drop from the decoder one at a time")
-    unit_drop_targets = traits.List(value=[], desc="Target indices on which to drop groups of units from the decoder")
+    unit_drop_groups = traits.Array(value=[[1],[2],[3],[4]], desc="Groups of channels to drop from the decoder one at a time")
+    unit_drop_targets = traits.Array(value=[[1,2], [3,4], [5,6], [7,8]], desc="Target indices on which to drop groups of units from the decoder")
 
     def init(self):
         self.decoder_units_dropped = np.ones((len(self.decoder.units),), dtype='bool')
@@ -35,6 +36,15 @@ def init(self):
         self.decoder_orig = copy.deepcopy(self.decoder)
         self.reportstats['Units dropped'] = '' # Runtime stat displayed on the UI
 
+        print('--------------------\nUnit dropping settings:')
+        if len(self.unit_drop_groups) == len(self.unit_drop_targets):
+            for i in range(len(self.unit_drop_groups)):
+                print(f'group {i} channels: {self.unit_drop_groups[i]}, targets: {self.unit_drop_targets[i]}')
+        else:
+            print('Unit dropping settings invalid! Please reformat into groups [[ch1, ch2], ...] and targets [[t1, t2, ...], ...]')
+            print('Current groups:', self.unit_drop_groups)
+            print('Current targets:', self.unit_drop_targets)
+        print('--------------------')
 
     def _start_wait(self):
         super()._start_wait()

From 61e328317c1814a87c2d6b779c9c3909eae68077 Mon Sep 17 00:00:00 2001
From: Katherine <leo.scholl@gmail.com>
Date: Mon, 8 Apr 2024 09:55:34 -0700
Subject: [PATCH 14/14] sort of works

---
 features/bmi_task_features.py | 16 ++++++++++++++--
 1 file changed, 14 insertions(+), 2 deletions(-)

diff --git a/features/bmi_task_features.py b/features/bmi_task_features.py
index d23bf2a21..1c59b08da 100644
--- a/features/bmi_task_features.py
+++ b/features/bmi_task_features.py
@@ -23,12 +23,24 @@ class RandomUnitDropout(traits.HasTraits):
     '''
 
     unit_drop_prob = traits.Float(0, desc="Probability of dropping a group of units from the decoder")
-    unit_drop_groups = traits.Array(value=[[1],[2],[3],[4]], desc="Groups of channels to drop from the decoder one at a time")
-    unit_drop_targets = traits.Array(value=[[1,2], [3,4], [5,6], [7,8]], desc="Target indices on which to drop groups of units from the decoder")
+    unit_drop_group1_channels = traits.List(value=[], desc="Channels to drop from the decoder one at a time")
+    unit_drop_group1_targets = traits.List(value=[], desc="Target indices on which to drop groups of units from the decoder")
+    unit_drop_group2_channels = traits.List(value=[], desc="Channels to drop from the decoder one at a time")
+    unit_drop_group2_targets = traits.List(value=[], desc="Target indices on which to drop groups of units from the decoder")
+    unit_drop_group3_channels = traits.List(value=[], desc="Channels to drop from the decoder one at a time")
+    unit_drop_group3_targets = traits.List(value=[], desc="Target indices on which to drop groups of units from the decoder")
+    unit_drop_group4_channels = traits.List(value=[], desc="Channels to drop from the decoder one at a time")
+    unit_drop_group4_targets = traits.List(value=[], desc="Target indices on which to drop groups of units from the decoder")
 
     def init(self):
         self.decoder_units_dropped = np.ones((len(self.decoder.units),), dtype='bool')
         self.add_dtype('decoder_units_dropped', '?', self.decoder_units_dropped.shape)
+        self.unit_drop_groups = [self.unit_drop_group1_channels, self.unit_drop_group2_channels,
+                                 self.unit_drop_group3_channels, self.unit_drop_group4_channels]
+        self.unit_drop_groups = np.array([g for g in self.unit_drop_groups if not g == ['']])
+        self.unit_drop_targets = [self.unit_drop_group1_targets, self.unit_drop_group2_targets,
+                                 self.unit_drop_group3_targets, self.unit_drop_group4_targets]
+        self.unit_drop_targets = np.array([t for t in self.unit_drop_targets if not t == ['']])
         self.unit_drop_group_idx = 0
         super().init()