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

canns/analyzer/visualization/energy_plots.py

@@ -45,9 +45,10 @@ def _render_single_energy_1d_frame(
     options: _Energy1DRenderOptions,
 ) -> np.ndarray:
     """Render a single frame for 1D energy landscape animation (module-level for pickling)."""
+    from io import BytesIO
+
     import matplotlib.pyplot as plt
     import numpy as np
-    from io import BytesIO
 
     fig, ax = plt.subplots(figsize=options.figsize)
     sim_index = options.sim_indices_to_render[frame_index]
@@ -121,9 +122,10 @@ def _render_single_energy_2d_frame(
     options: _Energy2DRenderOptions,
 ) -> np.ndarray:
     """Render a single frame for 2D energy landscape animation (module-level for pickling)."""
+    from io import BytesIO
+
     import matplotlib.pyplot as plt
     import numpy as np
-    from io import BytesIO
 
     fig, ax = plt.subplots(figsize=options.figsize)
     sim_index = options.sim_indices_to_render[frame_index]
@@ -211,16 +213,10 @@ def energy_landscape_1d_static(
     show: bool = True,
     **kwargs: Any,
 ):
-    """Plot a 1D static energy landscape using Matplotlib.
-
-    This mirrors the long-form description from the pre-reorganisation module so
-    existing documentation references stay accurate. The function accepts a
-    dictionary of datasets, plotting each curve on the same set of axes while
-    honouring the ``PlotConfig`` defaults callers relied on previously.
+    """Plot a 1D static energy landscape.
 
     Args:
-        data_sets: Mapping of series labels to ``(x, y)`` tuples representing
-            the energy curve to draw.
+        data_sets: Mapping ``label -> (x, y)`` where ``x`` and ``y`` are 1D arrays.
         config: Optional :class:`PlotConfig` carrying shared styling.
         title: Plot title when no config override is supplied.
         xlabel: X-axis label when no config override is supplied.
@@ -234,6 +230,17 @@ def energy_landscape_1d_static(
 
     Returns:
         Tuple[plt.Figure, plt.Axes]: The created figure and axes handles.
+
+    Examples:
+        >>> import numpy as np
+        >>> from canns.analyzer.visualization import energy_landscape_1d_static, PlotConfigs
+        >>>
+        >>> x = np.linspace(0, 1, 5)
+        >>> data_sets = {"u": (x, np.sin(x)), "Iext": (x, np.cos(x))}
+        >>> config = PlotConfigs.energy_landscape_1d_static(show=False)
+        >>> fig, ax = energy_landscape_1d_static(data_sets, config=config)
+        >>> print(fig is not None)
+        True
     """
 
     config = _ensure_plot_config(
@@ -291,16 +298,10 @@ def energy_landscape_1d_animation(
 ) -> animation.FuncAnimation:
     """Create an animation of an evolving 1D energy landscape.
 
-    The docstring intentionally preserves the guidance from the previous
-    implementation so existing callers can rely on the same parameter
-    explanations.
-
     Args:
-        data_sets: Dictionary whose keys are legend labels and values are
-            ``(x_data, y_data)`` tuples where ``y_data`` is shaped as
-            ``(time, state)``.
-        time_steps_per_second: Number of simulation time steps per second of
-            wall-clock time (e.g., ``1/dt``).
+        data_sets: Mapping ``label -> (x, y_series)``, where ``y_series`` is
+            shaped ``(timesteps, npoints)``.
+        time_steps_per_second: Simulation steps per second (e.g., ``1/dt``).
         config: Optional :class:`PlotConfig` with shared styling overrides.
         fps: Frames per second to render in the resulting animation.
         title: Title used when ``config`` is not provided.
@@ -320,6 +321,22 @@ def energy_landscape_1d_animation(
 
     Returns:
         ``matplotlib.animation.FuncAnimation``: The constructed animation.
+
+    Examples:
+        >>> import numpy as np
+        >>> from canns.analyzer.visualization import energy_landscape_1d_animation, PlotConfigs
+        >>>
+        >>> x = np.linspace(0, 1, 5)
+        >>> y_series = np.stack([np.sin(x), np.cos(x)], axis=0)
+        >>> data_sets = {"u": (x, y_series), "Iext": (x, y_series)}
+        >>> config = PlotConfigs.energy_landscape_1d_animation(
+        ...     time_steps_per_second=10,
+        ...     fps=2,
+        ...     show=False,
+        ... )
+        >>> anim = energy_landscape_1d_animation(data_sets, config=config)
+        >>> print(anim is not None)
+        True
     """
 
     config = _ensure_plot_config(
@@ -442,8 +459,12 @@ def energy_landscape_1d_animation(
 
             if backend == "imageio":
                 # Use imageio backend with parallel rendering
-                workers = render_workers if render_workers is not None else get_optimal_worker_count()
-                ctx = get_multiprocessing_context(prefer_fork=(render_start_method == "fork"))
+                workers = (
+                    render_workers if render_workers is not None else get_optimal_worker_count()
+                )
+                ctx, start_method = get_multiprocessing_context(
+                    prefer_fork=(render_start_method == "fork")
+                )
 
                 # Create render options
                 render_options = _Energy1DRenderOptions(
@@ -466,9 +487,10 @@ def energy_landscape_1d_animation(
                 writer_kwargs, mode = get_imageio_writer_kwargs(config.save_path, config.fps)
 
                 try:
-                    import imageio
                     from functools import partial
 
+                    import imageio
+
                     # Create partial function with options
                     render_func = partial(_render_single_energy_1d_frame, options=render_options)
 
@@ -504,6 +526,7 @@ def energy_landscape_1d_animation(
 
                 except Exception as e:
                     import warnings
+
                     warnings.warn(
                         f"imageio rendering failed: {e}. Falling back to matplotlib.",
                         RuntimeWarning,
@@ -605,6 +628,16 @@ def energy_landscape_2d_static(
 
     Returns:
         Tuple[plt.Figure, plt.Axes]: The Matplotlib figure and axes objects.
+
+    Examples:
+        >>> import numpy as np
+        >>> from canns.analyzer.visualization import energy_landscape_2d_static, PlotConfigs
+        >>>
+        >>> z = np.random.rand(4, 4)
+        >>> config = PlotConfigs.energy_landscape_2d_static(show=False)
+        >>> fig, ax = energy_landscape_2d_static(z, config=config)
+        >>> print(fig is not None)
+        True
     """
 
     config = _ensure_plot_config(
@@ -673,9 +706,6 @@ def energy_landscape_2d_animation(
 ) -> animation.FuncAnimation:
     """Create an animation of an evolving 2D landscape.
 
-    The long-form description mirrors the previous implementation to maintain
-    backwards-compatible documentation for downstream users.
-
     Args:
         zs_data: Array of shape ``(timesteps, dim_y, dim_x)`` describing the
             landscape at each simulation step.
@@ -701,6 +731,20 @@ def energy_landscape_2d_animation(
 
     Returns:
         ``matplotlib.animation.FuncAnimation``: The constructed animation.
+
+    Examples:
+        >>> import numpy as np
+        >>> from canns.analyzer.visualization import energy_landscape_2d_animation, PlotConfigs
+        >>>
+        >>> zs = np.random.rand(3, 4, 4)
+        >>> config = PlotConfigs.energy_landscape_2d_animation(
+        ...     time_steps_per_second=10,
+        ...     fps=2,
+        ...     show=False,
+        ... )
+        >>> anim = energy_landscape_2d_animation(zs, config=config)
+        >>> print(anim is not None)
+        True
     """
 
     config = _ensure_plot_config(
@@ -817,8 +861,12 @@ def energy_landscape_2d_animation(
 
             if backend == "imageio":
                 # Use imageio backend with parallel rendering
-                workers = render_workers if render_workers is not None else get_optimal_worker_count()
-                ctx = get_multiprocessing_context(prefer_fork=(render_start_method == "fork"))
+                workers = (
+                    render_workers if render_workers is not None else get_optimal_worker_count()
+                )
+                ctx, start_method = get_multiprocessing_context(
+                    prefer_fork=(render_start_method == "fork")
+                )
 
                 # Create render options
                 render_options = _Energy2DRenderOptions(
@@ -841,9 +889,10 @@ def energy_landscape_2d_animation(
                 writer_kwargs, mode = get_imageio_writer_kwargs(config.save_path, config.fps)
 
                 try:
-                    import imageio
                     from functools import partial
 
+                    import imageio
+
                     # Create partial function with options
                     render_func = partial(_render_single_energy_2d_frame, options=render_options)
 
@@ -879,6 +928,7 @@ def energy_landscape_2d_animation(
 
                 except Exception as e:
                     import warnings
+
                     warnings.warn(
                         f"imageio rendering failed: {e}. Falling back to matplotlib.",
                         RuntimeWarning,

canns/analyzer/visualization/spatial_plots.py

@@ -49,9 +49,10 @@ def _render_single_grid_tracking_frame(
     options: _GridCellTrackingRenderOptions,
 ) -> np.ndarray:
     """Render a single frame for grid cell tracking animation (module-level for pickling)."""
+    from io import BytesIO
+
     import matplotlib.pyplot as plt
     import numpy as np
-    from io import BytesIO
 
     fig, (ax1, ax2, ax3) = plt.subplots(1, 3, figsize=options.figsize)
     sim_idx = options.sim_indices_to_render[frame_index]
@@ -84,12 +85,19 @@ def _render_single_grid_tracking_frame(
 
     # Panel 3: Rate map with position
     im = ax3.imshow(
-        options.rate_map.T, origin="lower", cmap="hot",
-        extent=[0, options.env_size, 0, options.env_size], aspect="auto"
+        options.rate_map.T,
+        origin="lower",
+        cmap="hot",
+        extent=[0, options.env_size, 0, options.env_size],
+        aspect="auto",
     )
     ax3.plot(
-        [options.position[sim_idx, 0]], [options.position[sim_idx, 1]], "c*",
-        markersize=15, markeredgecolor="white", markeredgewidth=1.5
+        [options.position[sim_idx, 0]],
+        [options.position[sim_idx, 1]],
+        "c*",
+        markersize=15,
+        markeredgecolor="white",
+        markeredgewidth=1.5,
     )
     ax3.set_xlabel("X Position (m)", fontsize=10)
     ax3.set_ylabel("Y Position (m)", fontsize=10)
@@ -97,7 +105,9 @@ def _render_single_grid_tracking_frame(
     plt.colorbar(im, ax=ax3, fraction=0.046, pad=0.04)
 
     # Overall title with time
-    fig.suptitle(f"{options.title}  |  Time: {current_time_s:.2f} s", fontsize=13, fontweight="bold")
+    fig.suptitle(
+        f"{options.title}  |  Time: {current_time_s:.2f} s", fontsize=13, fontweight="bold"
+    )
 
     fig.tight_layout()
 
@@ -179,15 +189,15 @@ def plot_firing_field_heatmap(
         tuple[plt.Figure, plt.Axes]: The figure and axis objects for further customization.
 
     Example:
-        >>> from canns.analyzer.metrics.spatial_metrics import compute_firing_field
+        >>> import numpy as np
         >>> from canns.analyzer.visualization import plot_firing_field_heatmap, PlotConfig
-        >>> # Compute firing field
-        >>> heatmaps = compute_firing_field(activity, positions, 5.0, 5.0, 50, 50)
-        >>> # Plot single neuron with PlotConfig
-        >>> config = PlotConfig(figsize=(6, 6), title='Neuron 0', save_path='neuron_0.png', show=False)
-        >>> fig, ax = plot_firing_field_heatmap(heatmaps[0], config=config)
-        >>> # Plot with legacy parameters
-        >>> fig, ax = plot_firing_field_heatmap(heatmaps[1], title='Neuron 1', cmap='viridis', save_path='neuron_1.png')
+        >>>
+        >>> # Dummy input heatmap (M x K)
+        >>> heatmap = np.random.rand(6, 6)
+        >>> config = PlotConfig(title="Neuron 0", show=False)
+        >>> fig, ax = plot_firing_field_heatmap(heatmap, config=config)
+        >>> print(fig is not None)
+        True
     """
     # Handle configuration
     if config is None:
@@ -272,14 +282,16 @@ def plot_autocorrelation(
         tuple[plt.Figure, plt.Axes]: Figure and axes objects.
 
     Example:
+        >>> import numpy as np
         >>> from canns.analyzer.metrics.spatial_metrics import compute_spatial_autocorrelation
         >>> from canns.analyzer.visualization import plot_autocorrelation, PlotConfigs
+        >>>
+        >>> rate_map = np.random.rand(10, 10)
         >>> autocorr = compute_spatial_autocorrelation(rate_map)
-        >>> # Modern approach
-        >>> config = PlotConfigs.grid_autocorrelation(save_path='autocorr.png')
+        >>> config = PlotConfigs.grid_autocorrelation(show=False)
         >>> fig, ax = plot_autocorrelation(autocorr, config=config)
-        >>> # Legacy approach
-        >>> fig, ax = plot_autocorrelation(autocorr, cmap='RdBu_r', save_path='autocorr.png')
+        >>> print(fig is not None)
+        True
 
     References:
         Sargolini et al. (2006). Conjunctive representation of position, direction,
@@ -369,12 +381,15 @@ def plot_grid_score(
         tuple[plt.Figure, plt.Axes]: Figure and axes objects.
 
     Example:
+        >>> import numpy as np
         >>> from canns.analyzer.metrics.spatial_metrics import compute_grid_score
         >>> from canns.analyzer.visualization import plot_grid_score
+        >>>
+        >>> autocorr = np.random.rand(10, 10)
         >>> grid_score, rotated_corrs = compute_grid_score(autocorr)
-        >>> fig, ax = plot_grid_score(rotated_corrs, grid_score)
-        >>> print(f"Grid score: {grid_score:.3f}")
-        Grid score: 0.456
+        >>> fig, ax = plot_grid_score(rotated_corrs, grid_score, show=False)
+        >>> print(isinstance(grid_score, float))
+        True
     """
     config = _ensure_plot_config(
         config,
@@ -483,11 +498,20 @@ def plot_grid_spacing_analysis(
         tuple[plt.Figure, plt.Axes]: Figure and axes objects.
 
     Example:
+        >>> import numpy as np
         >>> from canns.analyzer.metrics.spatial_metrics import find_grid_spacing
         >>> from canns.analyzer.visualization import plot_grid_spacing_analysis
-        >>> spacing_bins, spacing_m = find_grid_spacing(autocorr, bin_size=0.06)
-        >>> fig, ax = plot_grid_spacing_analysis(autocorr, spacing_bins, bin_size=0.06)
-        >>> print(f"Spacing: {spacing_m:.3f}m")
+        >>>
+        >>> autocorr = np.random.rand(12, 12)
+        >>> spacing_bins, spacing_m = find_grid_spacing(autocorr, bin_size=0.05)
+        >>> fig, ax = plot_grid_spacing_analysis(
+        ...     autocorr,
+        ...     spacing_bins,
+        ...     bin_size=0.05,
+        ...     show=False,
+        ... )
+        >>> print(spacing_m is not None)
+        True
     """
     config = _ensure_plot_config(
         config,
@@ -620,18 +644,29 @@ def create_grid_cell_tracking_animation(
         FuncAnimation | None: Animation object, or None if displayed in Jupyter.
 
     Example:
-        >>> from canns.analyzer.visualization import create_grid_cell_tracking_animation, PlotConfigs
-        >>> # Create animation
+        >>> import numpy as np
+        >>> from canns.analyzer.visualization import (
+        ...     create_grid_cell_tracking_animation,
+        ...     PlotConfigs,
+        ... )
+        >>>
+        >>> position = np.array([[0.0, 0.0], [0.1, 0.1], [0.2, 0.2]])
+        >>> activity = np.array([0.0, 0.5, 1.0])
+        >>> rate_map = np.random.rand(5, 5)
         >>> config = PlotConfigs.grid_cell_tracking_animation(
-        ...     time_steps_per_second=1000,  # dt=1.0ms
-        ...     fps=20,
-        ...     save_path="tracking.gif"
+        ...     time_steps_per_second=10,
+        ...     fps=2,
+        ...     show=False,
         ... )
         >>> anim = create_grid_cell_tracking_animation(
-        ...     position, activity, rate_map,
+        ...     position,
+        ...     activity,
+        ...     rate_map,
         ...     config=config,
-        ...     env_size=3.0
+        ...     env_size=1.0,
         ... )
+        >>> print(anim is not None)
+        True
     """
     config = _ensure_plot_config(
         config,
@@ -767,8 +802,12 @@ def create_grid_cell_tracking_animation(
 
             if backend == "imageio":
                 # Use imageio backend with parallel rendering
-                workers = render_workers if render_workers is not None else get_optimal_worker_count()
-                ctx = get_multiprocessing_context(prefer_fork=(render_start_method == "fork"))
+                workers = (
+                    render_workers if render_workers is not None else get_optimal_worker_count()
+                )
+                ctx, start_method = get_multiprocessing_context(
+                    prefer_fork=(render_start_method == "fork")
+                )
 
                 # Create render options
                 render_options = _GridCellTrackingRenderOptions(
@@ -788,11 +827,14 @@ def create_grid_cell_tracking_animation(
                 writer_kwargs, mode = get_imageio_writer_kwargs(config.save_path, config.fps)
 
                 try:
-                    import imageio
                     from functools import partial
 
+                    import imageio
+
                     # Create partial function with options
-                    render_func = partial(_render_single_grid_tracking_frame, options=render_options)
+                    render_func = partial(
+                        _render_single_grid_tracking_frame, options=render_options
+                    )
 
                     with imageio.get_writer(config.save_path, mode=mode, **writer_kwargs) as writer:
                         if workers > 1 and ctx is not None:
@@ -826,6 +868,7 @@ def create_grid_cell_tracking_animation(
 
                 except Exception as e:
                     import warnings
+
                     warnings.warn(
                         f"imageio rendering failed: {e}. Falling back to matplotlib.",
                         RuntimeWarning,
@@ -855,7 +898,9 @@ def create_grid_cell_tracking_animation(
                         pbar.update(1)
 
                     try:
-                        ani.save(config.save_path, writer=writer, progress_callback=progress_callback)
+                        ani.save(
+                            config.save_path, writer=writer, progress_callback=progress_callback
+                        )
                         print(f"Animation saved to: {config.save_path}")
                     finally:
                         pbar.close()

canns/analyzer/visualization/spike_plots.py

@@ -47,9 +47,6 @@ def raster_plot(
 ):
     """Generate a raster plot from a spike train matrix.
 
-    The explanatory text mirrors the former ``visualize`` module so callers see
-    the same guidance after the reorganisation.
-
     Args:
         spike_train: Boolean/integer array of shape ``(timesteps, neurons)``.
         config: Optional :class:`PlotConfig` with shared styling options.
@@ -62,6 +59,17 @@ def raster_plot(
         save_path: Optional path used to persist the plot.
         show: Whether to display the plot interactively.
         **kwargs: Additional keyword arguments passed through to Matplotlib.
+
+    Examples:
+        >>> import numpy as np
+        >>> from canns.analyzer.visualization import raster_plot, PlotConfigs
+        >>>
+        >>> spike_train = np.zeros((5, 3), dtype=int)
+        >>> spike_train[::2, 0] = 1
+        >>> config = PlotConfigs.raster_plot(show=False)
+        >>> fig, ax = raster_plot(spike_train, config=config)
+        >>> print(fig is not None)
+        True
     """
 
     config = _ensure_plot_config(
@@ -158,6 +166,16 @@ def average_firing_rate_plot(
         save_path: Optional path used to persist the plot.
         show: Whether to display the plot interactively.
         **kwargs: Additional keyword arguments forwarded to Matplotlib.
+
+    Examples:
+        >>> import numpy as np
+        >>> from canns.analyzer.visualization import average_firing_rate_plot, PlotConfigs
+        >>>
+        >>> spike_train = np.random.randint(0, 2, size=(10, 4))
+        >>> config = PlotConfigs.average_firing_rate_plot(mode="population", show=False)
+        >>> fig, ax = average_firing_rate_plot(spike_train, dt=0.1, config=config)
+        >>> print(fig is not None)
+        True
     """
 
     config = _ensure_plot_config(
@@ -267,10 +285,12 @@ def population_activity_heatmap(
 
     Example:
         >>> import numpy as np
-        >>> from canns.analyzer.visualization.spike_plots import population_activity_heatmap
-        >>> # Simulate some activity data
-        >>> activity = np.random.rand(1000, 100)  # 1000 timesteps, 100 neurons
-        >>> fig, ax = population_activity_heatmap(activity, dt=0.001)
+        >>> from canns.analyzer.visualization import population_activity_heatmap, PlotConfig
+        >>> activity = np.random.rand(10, 5)
+        >>> config = PlotConfig(show=False)
+        >>> fig, ax = population_activity_heatmap(activity, dt=0.1, config=config)
+        >>> print(fig is not None)
+        True
     """
     if config is None:
         config = PlotConfig(