bluecellulab 2.6.42__py3-none-any.whl → 2.6.43__py3-none-any.whl

bluecellulab/stimulus/factory.py

@@ -1,361 +1,25 @@
-from __future__ import annotations
-from abc import ABC, abstractmethod
-from typing import Optional
-import logging
-import matplotlib.pyplot as plt
-import numpy as np
-
-logger = logging.getLogger(__name__)
-
-
-class Stimulus(ABC):
-    def __init__(self, dt: float) -> None:
-        self.dt = dt
-
-    @property
-    @abstractmethod
-    def time(self) -> np.ndarray:
-        """Time values of the stimulus."""
-        ...
-
-    @property
-    @abstractmethod
-    def current(self) -> np.ndarray:
-        """Current values of the stimulus."""
-        ...
-
-    def __len__(self) -> int:
-        return len(self.time)
-
-    @property
-    def stimulus_time(self) -> float:
-        return len(self) * self.dt
-
-    def __repr__(self) -> str:
-        return f"{self.__class__.__name__}(dt={self.dt})"
-
-    def plot(self, ax=None, **kwargs):
-        if ax is None:
-            ax = plt.gca()
-        ax.plot(self.time, self.current, **kwargs)
-        ax.set_xlabel("Time (ms)")
-        ax.set_ylabel("Current (nA)")
-        ax.set_title(self.__class__.__name__)
-        return ax
-
-    def __add__(self, other: Stimulus) -> CombinedStimulus:
-        """Override + operator to concatenate Stimulus objects."""
-        if self.dt != other.dt:
-            raise ValueError(
-                "Stimulus objects must have the same dt to be concatenated"
-            )
-        if len(self.time) == 0:
-            return CombinedStimulus(other.dt, other.time, other.current)
-        elif len(other.time) == 0:
-            return CombinedStimulus(self.dt, self.time, self.current)
-        else:
-            # shift other time
-            other_time = other.time + self.time[-1] + self.dt
-            combined_time = np.concatenate([self.time, other_time])
-            # Concatenate the current arrays
-            combined_current = np.concatenate([self.current, other.current])
-            return CombinedStimulus(self.dt, combined_time, combined_current)
-
-    def __eq__(self, other: object) -> bool:
-        if not isinstance(other, Stimulus):
-            return NotImplemented
-        else:
-            return (
-                np.allclose(self.time, other.time)
-                and np.allclose(self.current, other.current)
-                and self.dt == other.dt
-            )
-
-
-class CombinedStimulus(Stimulus):
-    """Represents the Stimulus created by combining multiple stimuli."""
-
-    def __init__(self, dt: float, time: np.ndarray, current: np.ndarray) -> None:
-        super().__init__(dt)
-        self._time = time
-        self._current = current
-
-    @property
-    def time(self) -> np.ndarray:
-        return self._time
-
-    @property
-    def current(self) -> np.ndarray:
-        return self._current
-
-
-class Empty(Stimulus):
-    """Represents empty stimulus (all zeros) that has no impact on the cell.
-
-    This is required by some Stimuli that expect the cell to rest.
-    """
-
-    def __init__(self, dt: float, duration: float) -> None:
-        super().__init__(dt)
-        self.duration = duration
-
-    @property
-    def time(self) -> np.ndarray:
-        return np.arange(0.0, self.duration, self.dt)
-
-    @property
-    def current(self) -> np.ndarray:
-        return np.zeros_like(self.time)
-
-
-class Flat(Stimulus):
-    def __init__(self, dt: float, duration: float, amplitude: float) -> None:
-        super().__init__(dt)
-        self.duration = duration
-        self.amplitude = amplitude
-
-    @property
-    def time(self) -> np.ndarray:
-        return np.arange(0.0, self.duration, self.dt)
-
-    @property
-    def current(self) -> np.ndarray:
-        return np.full_like(self.time, self.amplitude)
-
-
-class Slope(Stimulus):
-    def __init__(
-        self, dt: float, duration: float, amplitude_start: float, amplitude_end: float
-    ) -> None:
-        super().__init__(dt)
-        self.duration = duration
-        self.amplitude_start = amplitude_start
-        self.amplitude_end = amplitude_end
-
-    @property
-    def time(self) -> np.ndarray:
-        return np.arange(0.0, self.duration, self.dt)
-
-    @property
-    def current(self) -> np.ndarray:
-        return np.linspace(self.amplitude_start, self.amplitude_end, len(self.time))
-
-
-class Zap(Stimulus):
-    def __init__(self, dt: float, duration: float, amplitude: float) -> None:
-        super().__init__(dt)
-        self.duration = duration
-        self.amplitude = amplitude
-
-    @property
-    def time(self) -> np.ndarray:
-        return np.arange(0.0, self.duration, self.dt)
-
-    @property
-    def current(self) -> np.ndarray:
-        return self.amplitude * np.sin(
-            2.0 * np.pi * (1.0 + (1.0 / (5.15 - (self.time - 0.1)))) * (self.time - 0.1)
-        )
-
-
-class Step(Stimulus):
-
-    def __init__(self):
-        raise NotImplementedError(
-            "This class cannot be instantiated directly. "
-            "Please use the class methods 'amplitude_based' "
-            "or 'threshold_based' to create objects."
-        )
-
-    @classmethod
-    def amplitude_based(
-        cls,
-        dt: float,
-        pre_delay: float,
-        duration: float,
-        post_delay: float,
-        amplitude: float,
-    ) -> CombinedStimulus:
-        """Create a Step stimulus from given time events and amplitude.
-
-        Args:
-            dt: The time step of the stimulus.
-            pre_delay: The delay before the start of the step.
-            duration: The duration of the step.
-            post_delay: The time to wait after the end of the step.
-            amplitude: The amplitude of the step.
-        """
-        return (
-            Empty(dt, duration=pre_delay)
-            + Flat(dt, duration=duration, amplitude=amplitude)
-            + Empty(dt, duration=post_delay)
-        )
-
-    @classmethod
-    def threshold_based(
-        cls,
-        dt: float,
-        pre_delay: float,
-        duration: float,
-        post_delay: float,
-        threshold_current: float,
-        threshold_percentage: float,
-    ) -> CombinedStimulus:
-        """Creates a Step stimulus with respect to the threshold current.
-
-        Args:
-
-            dt: The time step of the stimulus.
-            pre_delay: The delay before the start of the step.
-            duration: The duration of the step.
-            post_delay: The time to wait after the end of the step.
-            threshold_current: The threshold current of the Cell.
-            threshold_percentage: Percentage of desired threshold_current amplification.
-        """
-        amplitude = threshold_current * threshold_percentage / 100
-        res = cls.amplitude_based(
-            dt,
-            pre_delay=pre_delay,
-            duration=duration,
-            post_delay=post_delay,
-            amplitude=amplitude,
-        )
-        return res
+# Copyright 2023-2024 Blue Brain Project / EPFL
+# Copyright 2025 Open Brain Institute
 
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
 
-class Ramp(Stimulus):
+#     http://www.apache.org/licenses/LICENSE-2.0
 
-    def __init__(self):
-        raise NotImplementedError(
-            "This class cannot be instantiated directly. "
-            "Please use the class methods 'amplitude_based' "
-            "or 'threshold_based' to create objects."
-        )
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
 
-    @classmethod
-    def amplitude_based(
-        cls,
-        dt: float,
-        pre_delay: float,
-        duration: float,
-        post_delay: float,
-        amplitude: float,
-    ) -> CombinedStimulus:
-        """Create a Ramp stimulus from given time events and amplitudes.
-
-        Args:
-            dt: The time step of the stimulus.
-            pre_delay: The delay before the start of the ramp.
-            duration: The duration of the ramp.
-            post_delay: The time to wait after the end of the ramp.
-            amplitude: The final amplitude of the ramp.
-        """
-        return (
-            Empty(dt, duration=pre_delay)
-            + Slope(
-                dt,
-                duration=duration,
-                amplitude_start=0.0,
-                amplitude_end=amplitude,
-            )
-            + Empty(dt, duration=post_delay)
-        )
-
-    @classmethod
-    def threshold_based(
-        cls,
-        dt: float,
-        pre_delay: float,
-        duration: float,
-        post_delay: float,
-        threshold_current: float,
-        threshold_percentage: float,
-    ) -> CombinedStimulus:
-        """Creates a Ramp stimulus with respect to the threshold current.
-
-        Args:
-
-            dt: The time step of the stimulus.
-            pre_delay: The delay before the start of the ramp.
-            duration: The duration of the ramp.
-            post_delay: The time to wait after the end of the ramp.
-            threshold_current: The threshold current of the Cell.
-            threshold_percentage: Percentage of desired threshold_current amplification.
-        """
-        amplitude = threshold_current * threshold_percentage / 100
-        res = cls.amplitude_based(
-            dt,
-            pre_delay=pre_delay,
-            duration=duration,
-            post_delay=post_delay,
-            amplitude=amplitude,
-        )
-        return res
-
-
-class DelayedZap(Stimulus):
-
-    def __init__(self):
-        raise NotImplementedError(
-            "This class cannot be instantiated directly. "
-            "Please use the class methods 'amplitude_based' "
-            "or 'threshold_based' to create objects."
-        )
-
-    @classmethod
-    def amplitude_based(
-        cls,
-        dt: float,
-        pre_delay: float,
-        duration: float,
-        post_delay: float,
-        amplitude: float,
-    ) -> CombinedStimulus:
-        """Create a DelayedZap stimulus from given time events and amplitude.
-
-        Args:
-            dt: The time step of the stimulus.
-            pre_delay: The delay before the start of the step.
-            duration: The duration of the step.
-            post_delay: The time to wait after the end of the step.
-            amplitude: The amplitude of the step.
-        """
-        return (
-            Empty(dt, duration=pre_delay)
-            + Zap(dt, duration=duration, amplitude=amplitude)
-            + Empty(dt, duration=post_delay)
-        )
-
-    @classmethod
-    def threshold_based(
-        cls,
-        dt: float,
-        pre_delay: float,
-        duration: float,
-        post_delay: float,
-        threshold_current: float,
-        threshold_percentage: float,
-    ) -> CombinedStimulus:
-        """Creates a SineSpec stimulus with respect to the threshold current.
-
-        Args:
+from __future__ import annotations
+from typing import Optional
+import logging
+from bluecellulab.stimulus.stimulus import DelayedZap, Empty, Ramp, Slope, Step, StepNoise, Stimulus, OrnsteinUhlenbeck, ShotNoise
+from bluecellulab.stimulus.circuit_stimulus_definitions import Stimulus as CircuitStimulus
 
-            dt: The time step of the stimulus.
-            pre_delay: The delay before the start of the step.
-            duration: The duration of the step.
-            post_delay: The time to wait after the end of the step.
-            threshold_current: The threshold current of the Cell.
-            threshold_percentage: Percentage of desired threshold_current amplification.
-        """
-        amplitude = threshold_current * threshold_percentage / 100
-        res = cls.amplitude_based(
-            dt,
-            pre_delay=pre_delay,
-            duration=duration,
-            post_delay=post_delay,
-            amplitude=amplitude,
-        )
-        return res
+logger = logging.getLogger(__name__)
 
 
 class StimulusFactory:
@@ -374,15 +38,39 @@ class StimulusFactory:
         )
 
     def ramp(
-        self, pre_delay: float, duration: float, post_delay: float, amplitude: float
+        self,
+        pre_delay: float,
+        duration: float,
+        post_delay: float,
+        amplitude: Optional[float] = None,
+        threshold_current: Optional[float] = None,
+        threshold_percentage: Optional[float] = 220.0,
     ) -> Stimulus:
-        return Ramp.amplitude_based(
-            self.dt,
-            pre_delay=pre_delay,
-            duration=duration,
-            post_delay=post_delay,
-            amplitude=amplitude,
-        )
+        if amplitude is not None:
+            if threshold_current is not None and threshold_current != 0 and threshold_percentage is not None:
+                logger.info(
+                    "amplitude, threshold_current and threshold_percentage are all set in ramp."
+                    " Will only keep amplitude value."
+                )
+            return Ramp.amplitude_based(
+                self.dt,
+                pre_delay=pre_delay,
+                duration=duration,
+                post_delay=post_delay,
+                amplitude=amplitude,
+            )
+
+        if threshold_current is not None and threshold_current != 0 and threshold_percentage is not None:
+            return Ramp.threshold_based(
+                self.dt,
+                pre_delay=pre_delay,
+                duration=duration,
+                post_delay=post_delay,
+                threshold_current=threshold_current,
+                threshold_percentage=threshold_percentage,
+            )
+
+        raise TypeError("You have to give either threshold_current or amplitude")
 
     def ap_waveform(
         self,
@@ -687,3 +375,243 @@ class StimulusFactory:
             )
 
         raise TypeError("You have to give either threshold_current or amplitude")
+
+    def ornstein_uhlenbeck(
+        self,
+        pre_delay: float,
+        post_delay: float,
+        duration: float,
+        tau: float,
+        sigma: Optional[float] = None,
+        mean: Optional[float] = None,
+        mean_percent: Optional[float] = None,
+        sigma_percent: Optional[float] = None,
+        threshold_current: Optional[float] = None,
+        seed: Optional[int] = None
+    ) -> Stimulus:
+        """Creates an Ornstein-Uhlenbeck process stimulus (factory-compatible).
+
+        Args:
+            pre_delay: Delay before the noise starts (ms).
+            post_delay: Delay after the noise ends (ms).
+            duration: Duration of the stimulus (ms).
+            tau: Time constant of the noise process.
+            sigma: Standard deviation of the noise (used when mean is provided).
+            mean: Absolute mean current value (used if provided).
+            mean_percent: Mean current as a percentage of threshold current (used if mean is None).
+            sigma_percent: Standard deviation as a percentage of threshold current (used if sigma is None).
+            threshold_current: Reference threshold current for percentage-based calculation.
+            seed: Optional random seed for reproducibility.
+
+        Returns:
+            A `Stimulus` object (OrnsteinUhlenbeck) that can be plotted and injected.
+
+        Notes:
+            - If `mean` is provided, `mean_percent` is ignored.
+            - If `threshold_current` is not provided, threshold-based parameters cannot be used.
+        """
+        is_amplitude_based = mean is not None and sigma is not None
+        is_threshold_based = (
+            threshold_current is not None
+            and threshold_current != 0
+            and mean_percent is not None
+            and sigma_percent is not None
+        )
+
+        if is_amplitude_based:
+            if is_threshold_based:
+                logger.info(
+                    "mean, threshold_current, and mean_percent are all set in Ornstein-Uhlenbeck."
+                    " Using mean and ignoring threshold-based parameters."
+                )
+
+            return OrnsteinUhlenbeck.amplitude_based(
+                dt=self.dt,
+                pre_delay=pre_delay,
+                post_delay=post_delay,
+                duration=duration,
+                tau=tau,
+                sigma=sigma,  # type: ignore[arg-type]
+                mean=mean,  # type: ignore[arg-type]
+                seed=seed
+            )
+
+        if is_threshold_based:
+            return OrnsteinUhlenbeck.threshold_based(
+                dt=self.dt,
+                pre_delay=pre_delay,
+                post_delay=post_delay,
+                duration=duration,
+                mean_percent=mean_percent,  # type: ignore[arg-type]
+                sigma_percent=sigma_percent,  # type: ignore[arg-type]
+                threshold_current=threshold_current,  # type: ignore[arg-type]
+                tau=tau,
+                seed=seed
+            )
+
+        raise TypeError("You have to give either `mean` and `sigma` or `threshold_current` and `mean_percent` and `sigma_percent`.")
+
+    def shot_noise(
+        self,
+        pre_delay: float,
+        post_delay: float,
+        duration: float,
+        rate: float,
+        rise_time: float,
+        decay_time: float,
+        mean: Optional[float] = None,
+        sigma: Optional[float] = None,
+        mean_percent: Optional[float] = None,
+        sigma_percent: Optional[float] = None,
+        relative_skew: float = 0.5,
+        threshold_current: Optional[float] = None,
+        seed: Optional[int] = None
+    ) -> Stimulus:
+        """Creates a ShotNoise instance, either with an absolute amplitude or
+        relative to a threshold current.
+
+        Args:
+            pre_delay: Delay before the noise starts (ms).
+            post_delay: Delay after the noise ends (ms).
+            duration: Duration of the stimulus (ms).
+            rate: Mean rate of synaptic events (Hz).
+            mean: Mean amplitude of events (nA), used if provided.
+            sigma: Standard deviation of event amplitudes.
+            rise_time: Rise time of synaptic events (ms).
+            decay_time: Decay time of synaptic events (ms).
+            mean_percent: Mean current as a percentage of threshold current (used if mean is None).
+            sigma_percent: Standard deviation as a percentage of threshold current (used if sigma is None).
+            relative_skew: Skew factor for the shot noise process (default: 0.5).
+            threshold_current: Reference threshold current for percentage-based calculation.
+            seed: Optional random seed for reproducibility.
+
+        Returns:
+            A `Stimulus` object that can be plotted and injected.
+
+        Notes:
+            - If `mean` is provided, `mean_percent` is ignored.
+            - If `threshold_current` is not provided, threshold-based parameters cannot be used.
+        """
+        is_amplitude_based = mean is not None and sigma is not None
+        is_threshold_based = (
+            threshold_current is not None
+            and threshold_current != 0
+            and mean_percent is not None
+            and sigma_percent is not None
+        )
+
+        if is_amplitude_based:
+            if is_threshold_based:
+                logger.info(
+                    "mean, threshold_current, and mean_percent are all set in ShotNoise."
+                    " Using mean and ignoring threshold-based parameters."
+                )
+
+            return ShotNoise.amplitude_based(
+                dt=self.dt,
+                pre_delay=pre_delay,
+                post_delay=post_delay,
+                duration=duration,
+                rate=rate,
+                mean=mean,  # type: ignore[arg-type]
+                sigma=sigma,  # type: ignore[arg-type]
+                rise_time=rise_time,
+                decay_time=decay_time,
+                seed=seed
+            )
+
+        if is_threshold_based:
+            return ShotNoise.threshold_based(
+                dt=self.dt,
+                pre_delay=pre_delay,
+                post_delay=post_delay,
+                duration=duration,
+                rise_time=rise_time,
+                decay_time=decay_time,
+                mean_percent=mean_percent,  # type: ignore[arg-type]
+                sigma_percent=sigma_percent,  # type: ignore[arg-type]
+                threshold_current=threshold_current,  # type: ignore[arg-type]
+                relative_skew=relative_skew,
+                seed=seed
+            )
+
+        raise TypeError("You must provide either `mean` and `sigma`, or `threshold_current` and `mean_percent` and `sigma_percent` with percentage values.")
+
+    def step_noise(
+        self,
+        pre_delay: float,
+        post_delay: float,
+        duration: float,
+        step_interval: float,
+        mean: Optional[float] = None,
+        sigma: Optional[float] = None,
+        mean_percent: Optional[float] = None,
+        sigma_percent: Optional[float] = None,
+        threshold_current: Optional[float] = None,
+        seed: Optional[int] = None,
+    ) -> Stimulus:
+        """Creates a StepNoise instance, either with an absolute amplitude or
+        relative to a threshold current.
+
+        Args:
+            pre_delay: Delay before the step noise starts (ms).
+            post_delay: Delay after the step noise ends (ms).
+            duration: Duration of the stimulus (ms).
+            step_interval: Interval at which noise amplitude changes.
+            mean: Mean amplitude of step noise (nA), used if provided.
+            sigma: Standard deviation of step noise.
+            mean_percent: Mean current as a percentage of threshold current (used if mean is None).
+            sigma_percent: Standard deviation as a percentage of threshold current (used if sigma is None).
+            threshold_current: Reference threshold current for percentage-based calculation.
+            seed: Optional random seed for reproducibility.
+
+        Returns:
+            A `Stimulus` object that can be plotted and injected.
+
+        Notes:
+            - If `mean` is provided, `mean_percent` is ignored.
+            - If `threshold_current` is not provided, threshold-based parameters cannot be used.
+        """
+        is_amplitude_based = mean is not None and sigma is not None
+        is_threshold_based = (
+            threshold_current is not None
+            and threshold_current != 0
+            and mean_percent is not None
+            and sigma_percent is not None
+        )
+
+        if is_amplitude_based:
+            if is_threshold_based:
+                logger.info(
+                    "mean, threshold_current, and mean_percent are all set in StepNoise."
+                    " Using mean and ignoring threshold-based parameters."
+                )
+            return StepNoise.amplitude_based(
+                dt=self.dt,
+                pre_delay=pre_delay,
+                post_delay=post_delay,
+                duration=duration,
+                step_interval=step_interval,
+                mean=mean,  # type: ignore[arg-type]
+                sigma=sigma,  # type: ignore[arg-type]
+                seed=seed,
+            )
+
+        if is_threshold_based:
+            return StepNoise.threshold_based(
+                dt=self.dt,
+                pre_delay=pre_delay,
+                post_delay=post_delay,
+                duration=duration,
+                step_interval=step_interval,
+                mean_percent=mean_percent,  # type: ignore[arg-type]
+                sigma_percent=sigma_percent,  # type: ignore[arg-type]
+                threshold_current=threshold_current,  # type: ignore[arg-type]
+                seed=seed,
+            )
+
+        raise TypeError("You must provide either `mean` and `sigma`, or `threshold_current` and `mean_percent` and `sigma_percent`  with percentage values.")
+
+    def from_sonata(cls, circuit_stimulus: CircuitStimulus):
+        """Convert a SONATA stimulus into a factory-based stimulus."""
+        raise ValueError(f"Unsupported circuit stimulus type: {type(circuit_stimulus)}")