neurospatial.environment._protocols

_protocols

Protocol definitions for Environment mixins.

This module defines Protocol classes that specify the interface mixins expect from the Environment class. Using Protocols lets mypy understand the mixin pattern without type: ignore comments or disabled error codes.

Docstrings here are intentionally one-line summaries. The full parameter and return documentation lives on the concrete Environment (and sibling) methods that users actually read; keeping this typing-only file to bare signatures avoids the docstring drift that duplication invites.

Index-dtype convention: point/region index arrays (bin_at, bins_in_region, boundary_bins, point_to_bin_index, _allocate_time_linear) are np.intp -- pointer-sized (64-bit on every platform, including Windows where np.int_ is 32-bit) and the dtype NumPy fancy-indexing returns. Trajectory sequence arrays (bin_sequence, BinSequenceWithRuns.bins) stay compact np.int32 on purpose, and the run_starts / run_lengths offsets are np.int64. The split is intentional -- portability for indices, memory for long sequences -- so do not collapse them to one dtype.

See: https://mypy.readthedocs.io/en/latest/more_types.html#mixin-classes

Classes

DiffusionGeometry

Bases: NamedTuple

Finite-volume geometry cached by _diffusion_geometry.

Built once and dropped wholesale on any _state_version bump; shared by diffuse's eigenbasis build and the smoothing consumers' W-component support gates.

EnvironmentProtocol

Bases: Protocol

Protocol defining the interface that Environment provides to mixins.

Mixins annotate self as EnvironmentProtocol (or the SelfEnv TypeVar) instead of Environment to avoid "erased type" errors in mypy.

Attributes

n_bins property

n_bins: int

Number of spatial bins in the environment.

n_dims property

n_dims: int

Number of spatial dimensions (1, 2, or 3).

is_linearized_track property

is_linearized_track: bool

Whether this is a 1D (linearized) environment.

bin_sizes property

bin_sizes: NDArray[float64]

Per-bin cell volume; the canonical mass M, shape (n_bins,).

layout_type property

layout_type: str | None

Type of layout engine used to create this environment, or None.

layout_parameters property

layout_parameters: dict[str, Any] | None

Parameters used to create the layout, or None.

boundary_bins property

boundary_bins: NDArray[intp]

Integer indices of bins at the environment boundary.

Functions

get_differential_operator

get_differential_operator() -> sparse.csc_matrix

Build (or fetch a cached) edge-oriented differential operator D.

Source code in src/neurospatial/environment/_protocols.py

def get_differential_operator(self) -> sparse.csc_matrix:
    """Build (or fetch a cached) edge-oriented differential operator ``D``."""
    ...

bin_at

bin_at(points_nd: NDArray[float64]) -> NDArray[np.intp]

Find bin indices for given points (-1 for points outside).

Source code in src/neurospatial/environment/_protocols.py

def bin_at(self, points_nd: NDArray[np.float64]) -> NDArray[np.intp]:
    """Find bin indices for given points (-1 for points outside)."""
    ...

bin_sequence

bin_sequence(times: NDArray[float64], positions: NDArray[float64], *, dedup: bool = True, outside_value: int | None = -1) -> NDArray[np.int32]

Convert a trajectory to a sequence of bin indices.

Source code in src/neurospatial/environment/_protocols.py

def bin_sequence(
    self,
    times: NDArray[np.float64],
    positions: NDArray[np.float64],
    *,
    dedup: bool = True,
    outside_value: int | None = -1,
) -> NDArray[np.int32]:
    """Convert a trajectory to a sequence of bin indices."""
    ...

bin_sequence_with_runs

bin_sequence_with_runs(times: NDArray[float64], positions: NDArray[float64], *, outside_value: int | None = -1) -> BinSequenceWithRuns

Convert a trajectory to a bin sequence plus per-run boundaries.

Source code in src/neurospatial/environment/_protocols.py

def bin_sequence_with_runs(
    self,
    times: NDArray[np.float64],
    positions: NDArray[np.float64],
    *,
    outside_value: int | None = -1,
) -> BinSequenceWithRuns:
    """Convert a trajectory to a bin sequence plus per-run boundaries."""
    ...

bins_in_region

bins_in_region(region_name: str) -> NDArray[np.intp]

Get integer bin indices within a named region.

Source code in src/neurospatial/environment/_protocols.py

def bins_in_region(self, region_name: str) -> NDArray[np.intp]:
    """Get integer bin indices within a named region."""
    ...

region_mask

region_mask(regions: str | list[str] | Region | Regions, *, include_boundary: bool = True) -> NDArray[np.bool_]

Boolean mask (shape (n_bins,)) for one or more regions.

Source code in src/neurospatial/environment/_protocols.py

def region_mask(
    self,
    regions: str | list[str] | Region | Regions,
    *,
    include_boundary: bool = True,
) -> NDArray[np.bool_]:
    """Boolean mask (shape ``(n_bins,)``) for one or more regions."""
    ...

region_membership

region_membership(regions: Regions | None = None, *, include_boundary: bool = True) -> NDArray[np.bool_]

Boolean membership mask for all regions, shape (n_bins, n_regions).

Source code in src/neurospatial/environment/_protocols.py

def region_membership(
    self,
    regions: Regions | None = None,
    *,
    include_boundary: bool = True,
) -> NDArray[np.bool_]:
    """Boolean membership mask for all regions, shape ``(n_bins, n_regions)``."""
    ...

compute_kernel

compute_kernel(bandwidth: float, *, mode: KernelMode = 'density', cache: bool = True) -> NDArray[np.float64]

Compute a smoothing kernel matrix, shape (n_bins, n_bins).

Source code in src/neurospatial/environment/_protocols.py

def compute_kernel(
    self,
    bandwidth: float,
    *,
    mode: KernelMode = "density",
    cache: bool = True,
) -> NDArray[np.float64]:
    """Compute a smoothing kernel matrix, shape ``(n_bins, n_bins)``."""
    ...

smooth

smooth(field: NDArray[float64], bandwidth: float, *, mode: KernelMode = 'density') -> NDArray[np.float64]

Apply graph-based smoothing to a spatial field.

Source code in src/neurospatial/environment/_protocols.py

def smooth(
    self,
    field: NDArray[np.float64],
    bandwidth: float,
    *,
    mode: KernelMode = "density",
) -> NDArray[np.float64]:
    """Apply graph-based smoothing to a spatial field."""
    ...

diffuse

diffuse(fields: Any, bandwidth: float, *, mode: KernelMode = 'density', backend: Literal['numpy', 'jax'] = 'numpy') -> NDArray[np.float64]

Matrix-free diffusion smoothing (no dense (n_bins, n_bins) kernel).

Source code in src/neurospatial/environment/_protocols.py

def diffuse(
    self,
    # ``Any`` is load-bearing: on ``backend="jax"`` diffuse accepts (and
    # returns) a JAX ``Array``, and jax is an optional dependency whose array
    # type is not importable here -- so ``NDArray | jax.Array`` cannot be
    # spelled. ``Any`` keeps the backend-polymorphic surface without breaking
    # jax-absent type checks.
    fields: Any,
    bandwidth: float,
    *,
    mode: KernelMode = "density",
    backend: Literal["numpy", "jax"] = "numpy",
) -> NDArray[np.float64]:
    """Matrix-free diffusion smoothing (no dense ``(n_bins, n_bins)`` kernel)."""
    ...

occupancy

occupancy(times: NDArray[float64], positions: NDArray[float64], *, speed: NDArray[float64] | None = None, min_speed: float | None = None, max_gap: float | None = 0.5, bandwidth: float | None = None, time_allocation: Literal['start', 'linear'] = 'start', return_seconds: bool = True) -> NDArray[np.float64]

Compute spatial occupancy (time per bin) from position data.

Source code in src/neurospatial/environment/_protocols.py

def occupancy(
    self,
    times: NDArray[np.float64],
    positions: NDArray[np.float64],
    *,
    speed: NDArray[np.float64] | None = None,
    min_speed: float | None = None,
    max_gap: float | None = 0.5,
    bandwidth: float | None = None,
    time_allocation: Literal["start", "linear"] = "start",
    return_seconds: bool = True,
) -> NDArray[np.float64]:
    """Compute spatial occupancy (time per bin) from position data."""
    ...

distance_between

distance_between(point1: NDArray[float64], point2: NDArray[float64], edge_weight: str = 'distance') -> float

Compute graph (shortest-path) distance between two points.

Source code in src/neurospatial/environment/_protocols.py

def distance_between(
    self,
    point1: NDArray[np.float64],
    point2: NDArray[np.float64],
    edge_weight: str = "distance",
) -> float:
    """Compute graph (shortest-path) distance between two points."""
    ...

to_linear

to_linear(nd_position: NDArray[float64]) -> NDArray[np.float64]

Convert N-D positions to 1D linear positions (1D environments only).

Source code in src/neurospatial/environment/_protocols.py

def to_linear(self, nd_position: NDArray[np.float64]) -> NDArray[np.float64]:
    """Convert N-D positions to 1D linear positions (1D environments only)."""
    ...

linear_to_nd

linear_to_nd(linear_position: NDArray[float64]) -> NDArray[np.float64]

Convert 1D linear positions to N-D coordinates (1D environments only).

Source code in src/neurospatial/environment/_protocols.py

def linear_to_nd(self, linear_position: NDArray[np.float64]) -> NDArray[np.float64]:
    """Convert 1D linear positions to N-D coordinates (1D environments only)."""
    ...

animate_fields

animate_fields(fields: Sequence[NDArray[float64]] | NDArray[float64], *, frame_times: NDArray[float64], backend: Literal['auto', 'napari', 'video', 'html', 'widget'] = 'auto', save_path: str | None = None, speed: float = 1.0, cmap: str = 'viridis', vmin: float | None = None, vmax: float | None = None, frame_labels: Sequence[str] | None = None, overlay_trajectory: NDArray[float64] | None = None, title: str = 'Spatial Field Animation', dpi: int = 100, codec: str = 'h264', bitrate: int = 5000, n_workers: int | None = None, dry_run: bool = False, image_format: Literal['png', 'jpeg'] = 'png', max_html_frames: int = 500, contrast_limits: tuple[float, float] | None = None, show_colorbar: bool = False, colorbar_label: str = '', overlays: list[OverlayProtocol] | None = None, show_regions: bool | list[str] = False, region_alpha: float = 0.3, scale_bar: bool | ScaleBarConfig = False, **kwargs: Any) -> Any

Animate spatial fields over time (backend-specific return value).

Source code in src/neurospatial/environment/_protocols.py

def animate_fields(
    self,
    fields: Sequence[NDArray[np.float64]] | NDArray[np.float64],
    *,
    frame_times: NDArray[np.float64],
    backend: Literal["auto", "napari", "video", "html", "widget"] = "auto",
    save_path: str | None = None,
    speed: float = 1.0,
    cmap: str = "viridis",
    vmin: float | None = None,
    vmax: float | None = None,
    frame_labels: Sequence[str] | None = None,
    overlay_trajectory: NDArray[np.float64] | None = None,
    title: str = "Spatial Field Animation",
    dpi: int = 100,
    codec: str = "h264",
    bitrate: int = 5000,
    n_workers: int | None = None,
    dry_run: bool = False,
    image_format: Literal["png", "jpeg"] = "png",
    max_html_frames: int = 500,
    contrast_limits: tuple[float, float] | None = None,
    show_colorbar: bool = False,
    colorbar_label: str = "",
    overlays: list[OverlayProtocol] | None = None,
    show_regions: bool | list[str] = False,
    region_alpha: float = 0.3,
    scale_bar: bool | ScaleBarConfig = False,
    **kwargs: Any,
) -> Any:
    """Animate spatial fields over time (backend-specific return value)."""
    ...