Clean-up Experiment dependencies inside Stimulus

stytra/stimulation/__init__.py

@@ -1,14 +1,16 @@
+from dataclasses import dataclass
 import datetime
 from copy import deepcopy
+import warnings
 
 from PyQt5.QtCore import pyqtSignal, QTimer, QObject
-from stytra.stimulation.stimuli import Pause, DynamicStimulus
+from stytra.stimulation.stimuli import Pause, DynamicStimulus, EnvironmentState
 from stytra.collectors.accumulators import DynamicLog, FramerateAccumulator
 from stytra.utilities import FramerateRecorder
 from lightparam.param_qt import ParametrizedQt, Param
 
 import logging
-
+        
 
 class ProtocolRunner(QObject):
     """Class for managing and running stimulation Protocols.
@@ -91,7 +93,9 @@ def __init__(self, experiment=None, target_dt=0, log_print=True):
         self.current_stimulus = None  # current stimulus object
         self.past_stimuli_elapsed = None  # time elapsed in previous stimuli
         self.dynamic_log = None  # dynamic log for stimuli
-
+
+        self.environment_state = EnvironmentState(calibrator = self.experiment.calibrator,)
+
         self.update_protocol()
         self.protocol.sig_param_changed.connect(self.update_protocol)
 
@@ -109,10 +113,19 @@ def update_protocol(self):
         self.stimuli = self.protocol._get_stimulus_list()
 
         self.current_stimulus = self.stimuli[0]
+
+        if hasattr(self.experiment, 'estimator'):
+            self.environment_state.estimator = self.experiment.estimator
 
         # pass experiment to stimuli for calibrator and asset folders:
         for stimulus in self.stimuli:
-            stimulus.initialise_external(self.experiment)
+            try:
+                stimulus.initialise_external(self.experiment, self.environment_state,)
+            except TypeError as e:
+                print("Error: {}".format(e))
+                stimulus.initialise_external(self.experiment)
+                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator and the estimator object!", FutureWarning)
+                warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator and the estimator object!", DeprecationWarning)
 
         if self.dynamic_log is None:
             self.dynamic_log = DynamicLog(self.stimuli, experiment=self.experiment)

stytra/stimulation/stimuli/closed_loop.py

@@ -69,7 +69,7 @@ def get_fish_vel(self):
         """ Function that update estimated fish velocty. Change to add lag or
         shunting.
         """
-        self.fish_vel = self._experiment.estimator.get_velocity()
+        self.fish_vel = self._environment_state.estimator.get_velocity()
 
     def bout_started(self):
         """ Function called on bout start.
@@ -90,7 +90,7 @@ def update(self):
                 self._experiment.logger.info(
                     "Experiment aborted! {} seconds without bouts".format(
                         self._elapsed - self.prev_bout_t
-                    )
+                    ) #! TOFIX: Remove
                 )
                 self.abort_experiment()
 
@@ -130,7 +130,7 @@ def calculate_final_vel(self):
         self.vel = self.base_vel - self.fish_vel * int(self.fish_swimming)
 
     def abort_experiment(self):
-        self._experiment.protocol_runner.stop()
+        self._experiment.protocol_runner.stop() #! TOFIX: Remove?
 
 
 class CalibratingClosedLoop1D(Basic_CL_1D):
@@ -177,7 +177,7 @@ def __init__(
 
     def bout_started(self):
         super().bout_started()
-        self.est_gain = self._experiment.estimator.base_gain
+        self.est_gain = self._environment_state.estimator.base_gain
 
     def bout_occurring(self):
         self.bout_vig.append(self.fish_vel / self.est_gain)
@@ -196,7 +196,7 @@ def bout_ended(self):
                 self.median_calib = self.median_vig * self.est_gain
                 self.est_gain = self.target_avg_fish_vel / self.median_vig
 
-                self._experiment.estimator.base_gain = self.est_gain
+                self._environment_state.estimator.base_gain = self.est_gain
 
         self.bout_vel = []
 
@@ -212,14 +212,14 @@ def stop(self):
                 "Experiment aborted! N bouts: {}; gain: {}".format(
                     len(self.bouts_vig_list), self.est_gain
                 )
-            )
+            ) #! TOFIX: Remove
 
         if len(self.bouts_vig_list) > self.calibrate_after:
             self._experiment.logger.info(
                 "Calibrated! Calculated gain  {} with {} bouts".format(
                     self.est_gain, len(self.bouts_vig_list)
                 )
-            )
+            ) #! TOFIX: Remove
 
 
 class GainChangerStimulus(Stimulus):
@@ -246,7 +246,7 @@ def __init__(self, newgain=1):
         self.newgain = newgain
 
     def start(self):
-        self._experiment.estimator.base_gain = self.newgain
+        self._environment_state.estimator.base_gain = self.newgain
 
 
 class GainLagClosedLoop1D(Basic_CL_1D):
@@ -277,8 +277,7 @@ def get_fish_vel(self):
         shunting.
         """
         super(GainLagClosedLoop1D, self).get_fish_vel()
-        self.lag_vel = self._experiment.estimator.get_velocity(self.lag)
-
+        self.lag_vel = self._environment_state.estimator.get_velocity(self.lag)
     def calculate_final_vel(self):
         subtract_to_base = self.gain * self.lag_vel
 
@@ -329,7 +328,7 @@ def bout_started(self):
             # print("set: {} gain and {} lag".format(self.gain, self.lag))
 
         # refresh lag if it was changed:
-        self.lag_vel = self._experiment.estimator.get_velocity(self.lag)
+        self.lag_vel = self._environment_state.estimator.get_velocity(self.lag)
 
 
 class PerpendicularMotion(BackgroundStimulus, InterpolatedStimulus):
@@ -338,7 +337,7 @@ class PerpendicularMotion(BackgroundStimulus, InterpolatedStimulus):
     """
 
     def update(self):
-        y, x, theta = self._experiment.estimator.get_position()
+        y, x, theta = self._environment_state.estimator.get_position()
         if np.isfinite(theta):
             self.theta = theta
         super().update()
@@ -352,7 +351,7 @@ def __init__(self, *args, **kwargs):
 
     def update(self):
         if self.is_tracking:
-            y, x, theta = self._experiment.estimator.get_position()
+            y, x, theta = self._environment_state.estimator.get_position()
             if np.isfinite(theta):
                 self.x = x
                 self.y = y
@@ -362,7 +361,7 @@ def update(self):
 
 class FishRelativeStimulus(BackgroundStimulus):
     def get_transform(self, w, h, x, y):
-        y_fish, x_fish, theta_fish = self._experiment.estimator.get_position()
+        y_fish, x_fish, theta_fish = self._environment_state.estimator.get_position()
         if np.isnan(y_fish):
             return super().get_transform(w, h, x, y)
         rot_fish = (theta_fish - np.pi / 2) * 180 / np.pi

stytra/stimulation/stimuli/conditional.py

@@ -36,9 +36,9 @@ def get_dynamic_state(self):
             state.update(self.active.get_dynamic_state())
         return state
 
-    def initialise_external(self, experiment):
-        super().initialise_external(experiment)
-        self.active.initialise_external(experiment)
+    def initialise_external(self, experiment, environment_state):
+        super().initialise_external(experiment, environment_state)
+        self.active.initialise_external(experiment, environment_state)
 
     def get_state(self):
         state = super().get_state()
@@ -50,7 +50,7 @@ def start(self):
         self.active.start()
 
     def check_condition(self):
-        y, x, theta = self._experiment.estimator.get_position()
+        y, x, theta = self._environment_state.estimator.get_position()
         return not np.isnan(y)
 
     def update(self):
@@ -157,10 +157,10 @@ def get_dynamic_state(self):
             state.update(self._stim_on.get_dynamic_state())
         return state
 
-    def initialise_external(self, experiment):
-        super().initialise_external(experiment)
-        self._stim_on.initialise_external(experiment)
-        self._stim_off.initialise_external(experiment)
+    def initialise_external(self, experiment, environment_state):
+        super().initialise_external(experiment, environment_state)
+        self._stim_on.initialise_external(experiment, environment_state)
+        self._stim_off.initialise_external(experiment, environment_state)
 
     def get_state(self):
         state = super().get_state()
@@ -270,8 +270,8 @@ def __init__(self, stimulus, *args, centering_stimulus=None, margin=45, **kwargs
         self.yc = 240
 
     def check_condition_on(self):
-        y, x, theta = self._experiment.estimator.get_position()
-        scale = self._experiment.calibrator.mm_px ** 2
+        y, x, theta = self._environment_state.estimator.get_position()
+        scale = self._environment_state.calibrator.mm_px ** 2 
         return (
             x > 0 and ((x - self.xc) ** 2 + (y - self.yc) ** 2) <= self.margin / scale
         )
@@ -323,15 +323,15 @@ def __init__(
         self.yc = 240
 
     def check_condition_on(self):
-        y, x, theta = self._experiment.estimator.get_position()
-        scale = self._experiment.calibrator.mm_px ** 2
+        y, x, theta = self._environment_state.estimator.get_position()
+        scale = self._environment_state.calibrator.mm_px ** 2
         return (not np.isnan(x)) and (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 <= self.margin_in / scale
         )
 
     def check_condition_off(self):
-        y, x, theta = self._experiment.estimator.get_position()
-        scale = self._experiment.calibrator.mm_px ** 2
+        y, x, theta = self._environment_state.estimator.get_position()
+        scale = self._environment_state.calibrator.mm_px ** 2
         return np.isnan(x) or (
             (x - self.xc) ** 2 + (y - self.yc) ** 2 > self.margin_out / scale
         )

stytra/stimulation/stimuli/external.py

@@ -42,7 +42,7 @@ def __init__(
         pulse_dur_str = str(pulse_dur_ms).zfill(3)
         self.mex = str("shock" + amp_dac + pulse_dur_str)
 
-    def initialise_external(self, experiment):
+    def initialise_external(self, experiment, environment_state):
         """
 
         Parameters

stytra/stimulation/stimuli/generic_stimuli.py

@@ -1,6 +1,18 @@
 import numpy as np
 import datetime
+import warnings
+from dataclasses import dataclass
 
+@dataclass
+class EnvironmentState:
+    def __init__(self, calibrator = None, estimator = None, height:int = 600, width:int = 800):
+        """
+        Holds Environment variables to pass from the protocol runner to the stimulus
+        """
+        self.calibrator = calibrator
+        self.estimator = estimator
+        self.height = height
+        self.width = width
 
 class Stimulus:
     """ Abstract class for a Stimulus.
@@ -66,6 +78,7 @@ def __init__(self, duration=0.0):
         self._elapsed = 0.0  # time from the beginning of the stimulus
         self.name = "undefined"
         self._experiment = None
+        self._environment_state = None
         self.real_time_start = None
         self.real_time_stop = None
 
@@ -111,7 +124,7 @@ def stop(self):
         """
         pass
 
-    def initialise_external(self, experiment):
+    def initialise_external(self, experiment, environment_state: EnvironmentState = None):
         """ Make a reference to the Experiment class inside the Stimulus.
         This is required to access from inside the Stimulus class to the
         Calibrator, the Pyboard, the asset directories with movies or the motor
@@ -130,7 +143,18 @@ def initialise_external(self, experiment):
             None
 
         """
+
+        if isinstance(environment_state, EnvironmentState):
+            self._environment_state = environment_state
+        else:
+            self._environment_state = experiment
+            warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", FutureWarning)
+            warnings.warn("Warning: 'initialise_external' will use the environment_state variable which holds the calibrator object!", DeprecationWarning)
+
+
         self._experiment = experiment
+
+
 
 
 class DynamicStimulus(Stimulus):
@@ -289,10 +313,10 @@ def update(self):
             s.update()
             s._elapsed = self._elapsed
 
-    def initialise_external(self, experiment):
-        super().initialise_external(experiment)
+    def initialise_external(self, experiment, environment_state):
+        super().initialise_external(experiment, environment_state)
         for s in self._stim_list:
-            s.initialise_external(experiment)
+            s.initialise_external(experiment, environment_state)
 
     @property
     def dynamic_parameter_names(self):

stytra/stimulation/stimuli/kinematograms.py

@@ -75,8 +75,8 @@ def get_dimensions(self):
         -------
         number of dots to display and the displacement amount in pixel coordinates
         """
-        if self._experiment.calibrator is not None:
-            mm_px = self._experiment.calibrator.mm_px
+        if self._environment_state.calibrator is not None:
+            mm_px = self._environment_state.calibrator.mm_px  
         else:
             mm_px = 1