canns 0.12.6__py3-none-any.whl → 0.13.0__py3-none-any.whl

canns/task/open_loop_navigation.py

@@ -27,31 +27,63 @@ __all__ = [
 
 
 def map2pi(a):
+    """Wrap angles to the range [-pi, pi].
+
+    Workflow:
+        Setup -> Provide angles (scalar or array-like).
+        Execute -> Call ``map2pi``.
+        Result -> Angles wrapped into [-pi, pi].
+
+    Examples:
+        >>> import numpy as np
+        >>> import brainpy.math as bm
+        >>> from canns.task.open_loop_navigation import map2pi
+        >>>
+        >>> angles = bm.array([3.5, -4.0])
+        >>> wrapped = map2pi(angles)
+        >>> bool(((wrapped >= -np.pi) & (wrapped <= np.pi)).all())
+        True
+    """
     b = bm.where(a > np.pi, a - np.pi * 2, a)
     c = bm.where(b < -np.pi, b + np.pi * 2, b)
     return c
 
 
 class ActionPolicy(ABC):
-    """
-    Abstract base class for action policies that control agent movement.
-
-    Action policies compute parameters for agent.update() at each simulation step,
-    enabling reusable, testable, and composable control strategies.
-
-    Example:
-        ```python
-        class ConstantDriftPolicy(ActionPolicy):
-            def __init__(self, drift_direction):
-                self.drift = np.array(drift_direction)
-
-            def compute_action(self, step_idx, agent):
-                return {'drift_velocity': self.drift}
-
-        task = CustomOpenLoopNavigationTask(
-            duration=100, action_policy=ConstantDriftPolicy([0.1, 0.0])
-        )
-        ```
+    """Abstract base class for action policies that control agent movement.
+
+    Action policies compute parameters for ``agent.update()`` at each simulation
+    step, enabling reusable and testable control strategies.
+
+    Workflow:
+        Setup -> Implement ``compute_action``.
+        Execute -> Pass the policy into a task and call ``get_data()``.
+        Result -> Task data is generated using the policy-controlled actions.
+
+    Examples:
+        >>> import numpy as np
+        >>> import brainpy.math as bm
+        >>> from canns.task.open_loop_navigation import ActionPolicy, CustomOpenLoopNavigationTask
+        >>>
+        >>> class ConstantDriftPolicy(ActionPolicy):
+        ...     def __init__(self, drift_direction):
+        ...         self.drift = np.array(drift_direction, dtype=float)
+        ...
+        ...     def compute_action(self, step_idx, agent):
+        ...         return {"drift_velocity": self.drift, "drift_to_random_strength_ratio": 10.0}
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> task = CustomOpenLoopNavigationTask(
+        ...     duration=0.5,
+        ...     width=1.0,
+        ...     height=1.0,
+        ...     dt=bm.get_dt(),
+        ...     action_policy=ConstantDriftPolicy([0.1, 0.0]),
+        ...     progress_bar=False,
+        ... )
+        >>> task.get_data()
+        >>> task.data.position.shape[1]
+        2
     """
 
     @abstractmethod
@@ -75,15 +107,31 @@ class ActionPolicy(ABC):
 
 @dataclass
 class OpenLoopNavigationData:
-    """
-    Container for the inputs recorded during the open-loop navigation task.
-    It contains the position, velocity, speed, movement direction, head direction,
-    and rotational velocity of the agent.
-
-    Additional fields for theta sweep analysis:
-    - ang_velocity: Angular velocity calculated using unwrap method
-    - linear_speed_gains: Normalized linear speed gains [0,1]
-    - ang_speed_gains: Normalized angular speed gains [-1,1]
+    """Container for open-loop navigation trajectories and derived signals.
+
+    It stores position, velocity, speed, movement direction, head direction,
+    and rotational velocity. Optional fields are added for theta sweep analysis.
+
+    Workflow:
+        Setup -> Create an ``OpenLoopNavigationTask``.
+        Execute -> Call ``get_data()``.
+        Result -> Access trajectories from ``task.data``.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.task.open_loop_navigation import OpenLoopNavigationTask
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> task = OpenLoopNavigationTask(
+        ...     duration=1.0,
+        ...     width=1.0,
+        ...     height=1.0,
+        ...     dt=bm.get_dt(),
+        ...     progress_bar=False,
+        ... )
+        >>> task.get_data()
+        >>> task.data.position.shape[1]
+        2
     """
 
     position: np.ndarray
@@ -100,9 +148,32 @@ class OpenLoopNavigationData:
 
 
 class OpenLoopNavigationTask(BaseNavigationTask):
-    """
-    Open-loop spatial navigation task that synthesises trajectories without
-    incorporating real-time feedback from a controller.
+    """Open-loop navigation task that synthesizes trajectories.
+
+    The trajectory is generated without real-time feedback control. This is
+    useful for producing reproducible paths for model evaluation.
+
+    Workflow:
+        Setup -> Instantiate the task with environment and motion settings.
+        Execute -> Call ``get_data()`` to generate a trajectory.
+        Result -> Read ``task.data`` for positions, velocities, and speed.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.task.open_loop_navigation import OpenLoopNavigationTask
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> task = OpenLoopNavigationTask(
+        ...     duration=1.0,
+        ...     width=1.0,
+        ...     height=1.0,
+        ...     start_pos=(0.5, 0.5),
+        ...     dt=bm.get_dt(),
+        ...     progress_bar=False,
+        ... )
+        >>> task.get_data()
+        >>> task.data.position.shape[0] == task.total_steps
+        True
     """
 
     def __init__(
@@ -621,7 +692,9 @@ class OpenLoopNavigationTask(BaseNavigationTask):
 
         if not hasattr(self, "agent") or self.agent is None:
             self.agent = Agent(
-                environment=self.env, params=copy.deepcopy(self.agent_params), rng_seed=self.rng_seed
+                environment=self.env,
+                params=copy.deepcopy(self.agent_params),
+                rng_seed=self.rng_seed,
             )
 
         # Set initial position
@@ -696,29 +769,36 @@ class OpenLoopNavigationTask(BaseNavigationTask):
 
 
 class CustomOpenLoopNavigationTask(OpenLoopNavigationTask):
-    """
-    Template class for policy-based open-loop navigation tasks.
-
-    This class enables custom action policies by accepting an ActionPolicy object
-    that controls agent movement at each simulation step.
-
-    Args:
-        action_policy: ActionPolicy instance controlling agent movement
-        **kwargs: All other arguments passed to OpenLoopNavigationTask
-
-    Example:
-        ```python
-        # Define custom policy
-        class MyPolicy(ActionPolicy):
-            def compute_action(self, step_idx, agent):
-                return {'drift_velocity': np.array([0.1, 0.0])}
-
-        # Use it
-        task = CustomOpenLoopNavigationTask(
-            duration=100, action_policy=MyPolicy()
-        )
-        task.get_data()
-        ```
+    """Open-loop navigation task driven by a custom action policy.
+
+    Provide an ``ActionPolicy`` to control how the agent moves at each step.
+
+    Workflow:
+        Setup -> Implement a policy and create the task.
+        Execute -> Call ``get_data()``.
+        Result -> Trajectory data reflects the policy-driven actions.
+
+    Examples:
+        >>> import numpy as np
+        >>> import brainpy.math as bm
+        >>> from canns.task.open_loop_navigation import ActionPolicy, CustomOpenLoopNavigationTask
+        >>>
+        >>> class MyPolicy(ActionPolicy):
+        ...     def compute_action(self, step_idx, agent):
+        ...         return {"drift_velocity": np.array([0.05, 0.0])}
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> task = CustomOpenLoopNavigationTask(
+        ...     duration=0.5,
+        ...     width=1.0,
+        ...     height=1.0,
+        ...     dt=bm.get_dt(),
+        ...     action_policy=MyPolicy(),
+        ...     progress_bar=False,
+        ... )
+        >>> task.get_data()
+        >>> task.data.velocity.shape[1]
+        2
     """
 
     def __init__(self, *args, action_policy: ActionPolicy | None = None, **kwargs):
@@ -733,46 +813,38 @@ class CustomOpenLoopNavigationTask(OpenLoopNavigationTask):
 
 
 class StateAwareRasterScanPolicy(ActionPolicy):
-    """
-    State-aware raster scan policy with cyclic dual-mode exploration.
-
-    Scanning strategy (循环扫描):
-    1. Horizontal mode: Left-right sweeps moving downward
-       → When reaching bottom: switch to Vertical mode
-    2. Vertical mode: Up-down sweeps moving rightward
-       → When reaching right edge: switch back to Horizontal mode
-    3. Cycles continuously: H → V → H → V → ... (避免撞墙)
-
-    This cyclic dual-mode approach achieves comprehensive coverage by combining
-    orthogonal scanning patterns and avoiding wall collisions.
-
-    Tested performance (200s, 1.0m x 1.0m environment):
-        - Cyclic dual-mode: ~75-80%+ coverage (continuous cycling)
-        - Single horizontal: 54.1% coverage (29 rows)
-
-    Args:
-        width: Environment width in meters
-        height: Environment height in meters
-        margin: Distance from wall to trigger turn (default: 0.05)
-        step_size: Movement per turn in perpendicular direction (default: 0.03)
-        speed: Movement speed (default: 0.15)
-        drift_strength: Drift-to-random ratio for agent.update() (default: 15.0)
-
-    Example:
-        ```python
-        policy = StateAwareRasterScanPolicy(
-            width=1.0, height=1.0,
-            step_size=0.03,  # Dense scanning for high coverage
-            drift_strength=15.0
-        )
-        task = CustomOpenLoopNavigationTask(
-            duration=200,
-            action_policy=policy,
-            width=1.0,
-            height=1.0,
-            start_pos=(0.05, 0.95)  # Start at top-left
-        )
-        ```
+    """State-aware raster scan policy with cyclic dual-mode exploration.
+
+    Scanning strategy:
+        1) Horizontal mode: left-right sweeps moving downward
+        2) Vertical mode: up-down sweeps moving rightward
+        3) Cycles continuously to avoid walls and improve coverage
+
+    Workflow:
+        Setup -> Instantiate the policy with environment size.
+        Execute -> Use it in ``CustomOpenLoopNavigationTask.get_data()``.
+        Result -> The trajectory follows a raster-scan pattern.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.task.open_loop_navigation import (
+        ...     StateAwareRasterScanPolicy,
+        ...     CustomOpenLoopNavigationTask,
+        ... )
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> policy = StateAwareRasterScanPolicy(width=1.0, height=1.0)
+        >>> task = CustomOpenLoopNavigationTask(
+        ...     duration=0.5,
+        ...     width=1.0,
+        ...     height=1.0,
+        ...     dt=bm.get_dt(),
+        ...     action_policy=policy,
+        ...     progress_bar=False,
+        ... )
+        >>> task.get_data()
+        >>> task.data.position.shape[1]
+        2
     """
 
     def __init__(
@@ -905,47 +977,32 @@ class StateAwareRasterScanPolicy(ActionPolicy):
 
 
 class RasterScanNavigationTask(CustomOpenLoopNavigationTask):
-    """
-    Preset task for cyclic dual-mode state-aware raster scan exploration.
-
-    Systematically explores the environment using cyclic mode switching:
-    1. Horizontal phase: Left-right sweeps moving downward
-       → Switches to Vertical when reaching bottom
-    2. Vertical phase: Up-down sweeps moving rightward
-       → Switches back to Horizontal when reaching right edge
-    3. Cycles continuously: H → V → H → V → ...
-
-    This cyclic dual-mode strategy achieves superior coverage by combining
-    orthogonal scanning patterns and continuously adapting to avoid walls.
-
-    Performance (200s, 1.0m x 1.0m):
-        - Cyclic dual-mode: ~75-80% coverage (continuous cycling)
-        - Single horizontal: 54.1% coverage (29 rows)
-        - +20-30% improvement over random walk
-
-    Args:
-        duration: Simulation duration in seconds
-        width: Environment width (default: 1.0)
-        height: Environment height (default: 1.0)
-        step_size: Scan density - smaller = denser scanning (default: 0.03)
-        margin: Wall detection margin (default: 0.05)
-        speed: Movement speed in m/s (default: 0.15)
-        drift_strength: Drift control strength (default: 15.0)
-        **kwargs: Additional arguments passed to OpenLoopNavigationTask
-
-    Example:
-        ```python
-        # High coverage dual-mode exploration
-        task = RasterScanNavigationTask(
-            duration=200,
-            width=1.0,
-            height=1.0,
-            step_size=0.03,  # Dense scanning in both directions
-            speed=0.15  # Movement speed
-        )
-        task.get_data()
-        task.show_trajectory_analysis()
-        ```
+    """Preset open-loop task for cyclic dual-mode raster scan exploration.
+
+    The task alternates between horizontal and vertical sweep phases to cover
+    the environment while avoiding walls.
+
+    Workflow:
+        Setup -> Instantiate the task with scan parameters.
+        Execute -> Call ``get_data()``.
+        Result -> Access the generated trajectory in ``task.data``.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.task.open_loop_navigation import RasterScanNavigationTask
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> task = RasterScanNavigationTask(
+        ...     duration=0.5,
+        ...     width=1.0,
+        ...     height=1.0,
+        ...     step_size=0.05,
+        ...     dt=bm.get_dt(),
+        ...     progress_bar=False,
+        ... )
+        >>> task.get_data()
+        >>> task.data.position.shape[1]
+        2
     """
 
     def __init__(
@@ -983,11 +1040,24 @@ class RasterScanNavigationTask(CustomOpenLoopNavigationTask):
 
 
 class TMazeOpenLoopNavigationTask(OpenLoopNavigationTask):
-    """
-    Open-loop navigation task in a T-maze environment.
-
-    This subclass configures the environment with a T-maze boundary, which is useful
-    for studying decision-making and spatial navigation in a controlled setting.
+    """Open-loop navigation task in a T-maze environment.
+
+    The environment boundary is configured to a classic T-maze layout.
+
+    Workflow:
+        Setup -> Instantiate the task with maze geometry.
+        Execute -> Call ``get_data()``.
+        Result -> Use ``task.data.position`` as the trajectory.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.task.open_loop_navigation import TMazeOpenLoopNavigationTask
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> task = TMazeOpenLoopNavigationTask(duration=0.5, dt=bm.get_dt(), progress_bar=False)
+        >>> task.get_data()
+        >>> task.data.position.shape[1]
+        2
     """
 
     def __init__(
@@ -1038,12 +1108,25 @@ class TMazeOpenLoopNavigationTask(OpenLoopNavigationTask):
 
 
 class TMazeRecessOpenLoopNavigationTask(TMazeOpenLoopNavigationTask):
-    """
-    Open-loop navigation task in a T-maze environment with recesses at stem-arm junctions.
-
-    This variant adds small rectangular indentations at the T-junction, creating
-    additional spatial features that may be useful for studying spatial navigation
-    and decision-making.
+    """Open-loop navigation task in a T-maze with recesses at the junction.
+
+    Recesses add small indentations near the stem-arm junctions, providing
+    extra spatial structure.
+
+    Workflow:
+        Setup -> Instantiate the task with recess geometry.
+        Execute -> Call ``get_data()``.
+        Result -> Use ``task.data`` for downstream modeling.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.task.open_loop_navigation import TMazeRecessOpenLoopNavigationTask
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> task = TMazeRecessOpenLoopNavigationTask(duration=0.5, dt=bm.get_dt(), progress_bar=False)
+        >>> task.get_data()
+        >>> task.data.position.shape[1]
+        2
     """
 
     def __init__(

canns/task/tracking.py

@@ -429,10 +429,34 @@ class SmoothTracking(TrackingTask):
 
 
 class PopulationCoding1D(PopulationCoding):
-    """
-    Population coding task for 1D continuous attractor networks.
-    In this task, a stimulus is presented for a specific duration, preceded and followed by
-    periods of no stimulation, to test the network's ability to form and maintain a memory bump.
+    """Population coding task for 1D continuous attractor networks.
+
+    A stimulus is presented for a specific duration, preceded and followed by
+    periods of no stimulation.
+
+    Workflow:
+        Setup -> Create a 1D CANN and the task.
+        Execute -> Call ``get_data()``.
+        Result -> Use ``task.data`` as the input sequence.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.models.basic import CANN1D
+        >>> from canns.task.tracking import PopulationCoding1D
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> model = CANN1D(num=64)
+        >>> task = PopulationCoding1D(
+        ...     cann_instance=model,
+        ...     before_duration=1.0,
+        ...     after_duration=1.0,
+        ...     Iext=0.0,
+        ...     duration=2.0,
+        ...     time_step=bm.get_dt(),
+        ... )
+        >>> task.get_data()
+        >>> task.data.shape[0] == task.total_steps
+        True
     """
 
     def __init__(
@@ -472,10 +496,32 @@ class PopulationCoding1D(PopulationCoding):
 
 
 class TemplateMatching1D(TemplateMatching):
-    """
-    Template matching task for 1D continuous attractor networks.
-    This task presents a stimulus with added noise to test the network's ability to
-    denoise the input and settle on the correct underlying pattern (template).
+    """Template matching task for 1D continuous attractor networks.
+
+    A fixed stimulus template is presented with noise at each step, testing
+    the network's denoising dynamics.
+
+    Workflow:
+        Setup -> Create a 1D CANN and the task.
+        Execute -> Call ``get_data()``.
+        Result -> Use ``task.data`` as the noisy input sequence.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.models.basic import CANN1D
+        >>> from canns.task.tracking import TemplateMatching1D
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> model = CANN1D(num=64)
+        >>> task = TemplateMatching1D(
+        ...     cann_instance=model,
+        ...     Iext=0.0,
+        ...     duration=1.0,
+        ...     time_step=bm.get_dt(),
+        ... )
+        >>> task.get_data()
+        >>> task.data.shape[1] == model.shape[0]
+        True
     """
 
     def __init__(
@@ -504,10 +550,31 @@ class TemplateMatching1D(TemplateMatching):
 
 
 class SmoothTracking1D(SmoothTracking):
-    """
-    Smooth tracking task for 1D continuous attractor networks.
-    This task provides an external input that moves smoothly over time, testing the network's
-    ability to track a continuously changing stimulus.
+    """Smooth tracking task for 1D continuous attractor networks.
+
+    The external input moves smoothly between key positions.
+
+    Workflow:
+        Setup -> Create a 1D CANN and the task.
+        Execute -> Call ``get_data()``.
+        Result -> ``task.data`` contains the smoothly varying stimulus.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.models.basic import CANN1D
+        >>> from canns.task.tracking import SmoothTracking1D
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> model = CANN1D(num=64)
+        >>> task = SmoothTracking1D(
+        ...     cann_instance=model,
+        ...     Iext=(0.0, 1.0, 0.5),
+        ...     duration=(0.5, 0.5),
+        ...     time_step=bm.get_dt(),
+        ... )
+        >>> task.get_data()
+        >>> task.data.shape[0] == task.total_steps
+        True
     """
 
     def __init__(
@@ -565,10 +632,33 @@ class CustomTracking1D(TrackingTask):
 
 
 class PopulationCoding2D(PopulationCoding):
-    """
-    Population coding task for 2D continuous attractor networks.
-    In this task, a stimulus is presented for a specific duration, preceded and followed by
-    periods of no stimulation, to test the network's ability to form and maintain a memory bump.
+    """Population coding task for 2D continuous attractor networks.
+
+    A 2D stimulus is presented for a duration with pre- and post-silence.
+
+    Workflow:
+        Setup -> Create a 2D CANN and the task.
+        Execute -> Call ``get_data()``.
+        Result -> Use ``task.data`` as the input sequence.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.models.basic import CANN2D
+        >>> from canns.task.tracking import PopulationCoding2D
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> model = CANN2D(length=8)
+        >>> task = PopulationCoding2D(
+        ...     cann_instance=model,
+        ...     before_duration=1.0,
+        ...     after_duration=1.0,
+        ...     Iext=(0.0, 0.0),
+        ...     duration=1.0,
+        ...     time_step=bm.get_dt(),
+        ... )
+        >>> task.get_data()
+        >>> task.data.shape[1:] == model.shape
+        True
     """
 
     def __init__(
@@ -609,10 +699,31 @@ class PopulationCoding2D(PopulationCoding):
 
 
 class TemplateMatching2D(TemplateMatching):
-    """
-    Template matching task for 2D continuous attractor networks.
-    This task presents a stimulus with added noise to test the network's ability to
-    denoise the input and settle on the correct underlying pattern (template).
+    """Template matching task for 2D continuous attractor networks.
+
+    A 2D template is presented with noise at each step.
+
+    Workflow:
+        Setup -> Create a 2D CANN and the task.
+        Execute -> Call ``get_data()``.
+        Result -> ``task.data`` contains noisy 2D inputs.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.models.basic import CANN2D
+        >>> from canns.task.tracking import TemplateMatching2D
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> model = CANN2D(length=8)
+        >>> task = TemplateMatching2D(
+        ...     cann_instance=model,
+        ...     Iext=(0.0, 0.0),
+        ...     duration=1.0,
+        ...     time_step=bm.get_dt(),
+        ... )
+        >>> task.get_data()
+        >>> task.data.shape[1:] == model.shape
+        True
     """
 
     def __init__(
@@ -642,10 +753,31 @@ class TemplateMatching2D(TemplateMatching):
 
 
 class SmoothTracking2D(SmoothTracking):
-    """
-    Smooth tracking task for 2D continuous attractor networks.
-    This task provides an external input that moves smoothly over time, testing the network's
-    ability to track a continuously changing stimulus.
+    """Smooth tracking task for 2D continuous attractor networks.
+
+    The external 2D input moves smoothly between key positions.
+
+    Workflow:
+        Setup -> Create a 2D CANN and the task.
+        Execute -> Call ``get_data()``.
+        Result -> ``task.data`` contains smoothly varying 2D inputs.
+
+    Examples:
+        >>> import brainpy.math as bm
+        >>> from canns.models.basic import CANN2D
+        >>> from canns.task.tracking import SmoothTracking2D
+        >>>
+        >>> bm.set_dt(0.1)
+        >>> model = CANN2D(length=8)
+        >>> task = SmoothTracking2D(
+        ...     cann_instance=model,
+        ...     Iext=((0.0, 0.0), (1.0, 1.0), (0.5, 0.5)),
+        ...     duration=(0.5, 0.5),
+        ...     time_step=bm.get_dt(),
+        ... )
+        >>> task.get_data()
+        >>> task.data.shape[1:] == model.shape
+        True
     """
 
     def __init__(