rtgym.utils package

Subpackages

• rtgym.utils.visualization package
  • Submodules
  • rtgym.utils.visualization.plot_ratemap module
    • plot_ratemaps()
  • Module contents

Submodules

rtgym.utils.common module

hash_seed(seed, sensory_key)

Generate a hashed seed from base seed and sensory key.

Ensures different sensory cells have different random number generators even when using the same base seed by incorporating the sensory key into the seed generation.

Parameters:

• seed (int) – Base seed value.

• sensory_key (str) – Unique key for the sensory modality.

Returns:

Hashed seed value.

Return type:

int

rtgym.utils.data_processing module

class RatemapAggregator(arena_map, device=None)

Bases: object

get_ratemap()

Returns the final normalized firing fields (n_cells, n_x, n_y). Unvisited points (visit_count=0) will be NaN.

init_counts()

reset()

Reset the aggregator (clears all partial sums and counts).

update(states, coords)

Accumulate partial sums and visit counts from new data.

Parameters:

• states (torch.Tensor or np.ndarray) – Shape=(n_batches, n_timesteps, n_cells) or (n_timesteps, n_batches, n_cells).

• coords (torch.Tensor or np.ndarray) – Shape=(n_batches, n_timesteps, 2) or (n_timesteps, n_batches, 2).

combine_trajectories(traj_list: list)

Merge a list of Trajectory objects into a single Trajectory object

get_gym_dimensions(coords, arena_map)

Get the dimensions of the gym background

restrain2ff(coords, states, arena_map)

states2ff(states, coords, arena_map)

Takes the states and coords and returns the firing fields.

Parameters:

• states – shape=(n_batches, n_timesteps, n_cells) or (n_timesteps, n_batches, n_cells)

• coords – shape=(n_batches, n_timesteps, 2) or (n_timesteps, n_batches, 2)

• gym_bg – shape=(n_x, n_y), 0 for free space, 1 for walls

Returns:

shape=(n_cells, n_x, n_y)

Return type:

field

states_to_ratemap(states, coords, arena_map)

Takes the states and coords and returns the firing fields using PyTorch.

Parameters:

• states – shape=(n_batches, n_timesteps, n_cells) or (n_timesteps, n_batches, n_cells)

• coords – shape=(n_batches, n_timesteps, 2) or (n_timesteps, n_batches, 2)

• arena_map – shape=(n_x, n_y), 0 for free space, 1 for walls (torch.Tensor)

Returns:

shape=(n_cells, n_x, n_y) (torch.Tensor)

Return type:

field

rtgym.utils.decode_response module

Utility functions for decoding sensory responses to spatial trajectories.

This module contains algorithms for converting high-dimensional sensory responses (e.g., from place cells, grid cells) back to spatial coordinates using various optimization techniques including nearest neighbor search, spatial interpolation, and approximate search methods.

Mathematical Background

Given sensory response vectors \(\mathbf{r} \in \mathbb{R}^d\) and response maps \(\mathbf{M} \in \mathbb{R}^{H \times W \times d}\), we seek to find coordinates \((y,x)\) such that:

\[(y,x) = \arg\min_{(i,j)} \|\mathbf{r} - \mathbf{M}[i,j,:]\|_2^2\]

where \(\|\cdot\|_2\) denotes the Euclidean norm.

Authors: RatatouGym Development Team

create_dataclass_result(pred_coords: numpy.ndarray, is_trajectory: bool) → Trajectory | AgentState

Create appropriate dataclass result from decoded coordinates.

Parameters:

• pred_coords (np.ndarray) – Decoded coordinate array

• is_trajectory (bool) – Whether to create Trajectory or AgentState object

Returns:

• Trajectory or AgentState – Dataclass object containing the coordinates

• r

decode_response_euclidean(response: numpy.ndarray, res_maps: numpy.ndarray) → Tuple[numpy.ndarray, bool]

Decode sensory responses using brute-force Euclidean distance computation.

This is the reference implementation that computes the exact nearest neighbor by evaluating the Euclidean distance between each query response and all template responses in the response maps.

Mathematical Formulation

For response vector \(\mathbf{r} \in \mathbb{R}^d\) and response map \(\mathbf{M} \in \mathbb{R}^{H \times W \times d}\), compute:

\[d^2(\mathbf{r}, \mathbf{M}[i,j]) = \sum_{k=1}^{d} (r_k - M[i,j,k])^2\]

Return:

\[\arg\min_{(i,j)} d^2(\mathbf{r}, \mathbf{M}[i,j])\]

param response:

Sensory response array of shape:

• (B, T, D) for trajectory decoding

• (B, D) for single state decoding

where B = batch size, T = time steps, D = feature dimensions

type response:

np.ndarray

param res_maps:

Response template maps of shape (D, H, W) where H = arena height, W = arena width

type res_maps:

np.ndarray

returns:

• tuple of (np.ndarray, bool) –

  • decoded_coordinates : Array of shape (B,T,2) or (B,2) containing spatial coordinates

  • is_trajectory : Boolean indicating if input was trajectory (3D) or state (2D)

• Complexity

• ———-

• - **Time** (\(O(B \cdot T \cdot H \cdot W \cdot D)\) for trajectory, \(O(B \cdot H \cdot W \cdot D)\) for states)

• - **Space** (\(O(H \cdot W \cdot D)\))

Note

This method provides exact results but can be computationally expensive for large arenas or high-dimensional feature spaces. Consider using approximate methods for better performance. r

decode_response_faiss(response: numpy.ndarray, res_maps: numpy.ndarray, n_clusters: int = 100) → Tuple[numpy.ndarray, bool]

Decode sensory responses using Facebook AI Similarity Search (FAISS).

FAISS provides state-of-the-art performance for large-scale nearest neighbor search through optimized indexing structures. This implementation uses Inverted File Index (IVF) with exact search within clusters.

Algorithm

1. Cluster response map features into n_clusters groups using k-means

2. Build inverted index mapping clusters to their members

3. For each query, search only relevant clusters for nearest neighbors

4. Return exact nearest neighbor within searched clusters

param response:

Sensory response array

type response:

np.ndarray

param res_maps:

Response template maps

type res_maps:

np.ndarray

param n_clusters:

Number of clusters for IVF index:

• More clusters: faster search, potentially lower recall

• Fewer clusters: slower search, higher recall

• Recommended: \(\sqrt{\text{number of points}}\) to \(\text{n_points}/39\)

type n_clusters:

int, default=100

returns:

• tuple of (np.ndarray, bool) – (decoded_coordinates, is_trajectory)

• Performance Characteristics

• —————————

• - **Index build time** (\(O(H \cdot W \cdot D \cdot \log(\text{n_clusters}))\))

• - **Query time** (\(O(B \cdot T \cdot \sqrt{H \cdot W})\) approximately)

• - **Memory usage** (~1.5x the size of original data)

Notes

• Requires FAISS library installation

• Optimized for very large datasets (>10k points)

• Can utilize GPU acceleration with appropriate FAISS build

• May return approximate results depending on nprobe parameter

r

decode_response_interpolation(response: numpy.ndarray, res_maps: numpy.ndarray, n_anchors: int = 1000, random_state: int = 42) → Tuple[numpy.ndarray, bool]

Decode sensory responses using spatial interpolation with anchor points.

This method exploits the spatial continuity assumption: nearby locations should have similar sensory responses. Instead of exact nearest neighbor search, it interpolates coordinates from multiple nearby anchor points.

Algorithm

1. Select n_anchors representative points using k-means clustering

2. For each query response, find k nearest anchor points

3. Compute inverse-distance weighted interpolation of anchor coordinates

4. Return interpolated coordinates (may be non-integer)

Mathematical Formulation

Given \(k\) nearest anchors with coordinates \(\mathbf{c}_1,\ldots,\mathbf{c}_k\) and distances \(d_1,\ldots,d_k\):

\[ \begin{align}\begin{aligned}w_i &= \frac{1}{d_i + \varepsilon} \quad \text{where } \varepsilon \text{ prevents division by zero}\\\tilde{w}_i &= \frac{w_i}{\sum_{j=1}^k w_j} \quad \text{(normalized weights)}\\\hat{\mathbf{c}} &= \sum_{i=1}^k \tilde{w}_i \cdot \mathbf{c}_i \quad \text{(interpolated coordinate)}\end{aligned}\end{align} \]

param response:

Sensory response array

type response:

np.ndarray

param res_maps:

Response template maps

type res_maps:

np.ndarray

param n_anchors:

Number of anchor points to select:

• More anchors: better spatial resolution, slower computation

• Fewer anchors: faster computation, lower spatial precision

type n_anchors:

int, default=1000

param random_state:

Random seed for reproducible anchor selection

type random_state:

int, default=42

returns:

• tuple of (np.ndarray, bool) – (decoded_coordinates, is_trajectory) Coordinates are rounded to integers and clipped to valid arena bounds

• Advantages

• ———-

• - Smooth spatial interpolation reduces noise

• - Faster than exact nearest neighbor for large datasets

• - Naturally handles uncertainty through weighted averaging

• Limitations

• ———–

• - May not find exact nearest neighbor

• • Requires tuning of n_anchors parameter

• - Assumes local spatial smoothness in response patterns

• Complexity

• ———-

• - **Anchor selection** (\(O(H \cdot W \cdot D \cdot \log(\text{n_anchors}))\))

• - **Query** (\(O(B \cdot T \cdot \log(\text{n_anchors}))\))

• r

decode_response_kdtree(response: numpy.ndarray, res_maps: numpy.ndarray) → Tuple[numpy.ndarray, bool]

Decode sensory responses using K-d tree for accelerated nearest neighbor search.

K-d trees provide significant speedup over brute force methods, especially for high-dimensional feature spaces. The tree is constructed once and then queried efficiently for all response vectors.

Algorithm

1. Build K-d tree on flattened response map features

2. Query tree for nearest neighbors of each response vector

3. Map indices back to spatial coordinates

param response:

Sensory response array (see decode_response_euclidean())

type response:

np.ndarray

param res_maps:

Response template maps (see decode_response_euclidean())

type res_maps:

np.ndarray

returns:

• tuple of (np.ndarray, bool) – Same format as decode_response_euclidean()

• Complexity

• ———-

• - **Tree construction** (\(O(H \cdot W \cdot D \cdot \log(H \cdot W))\))

• - **Query** (\(O(B \cdot T \cdot \log(H \cdot W))\) for trajectory)

• - **Space** (\(O(H \cdot W \cdot D)\))

• Advantages

• ———-

• - Significantly faster than brute force for large arenas

• - Exact nearest neighbor results

• - Memory efficient tree structure

• Limitations

• ———–

• - Performance degrades in very high dimensions (curse of dimensionality)

• - Tree construction overhead for small datasets

• r

decode_response_torch(response: numpy.ndarray | torch.Tensor, res_maps: numpy.ndarray | torch.Tensor, device: str | torch.device | None = None, chunk_size: int = 1024) → Tuple[numpy.ndarray, bool]

Decode sensory responses using PyTorch for GPU-accelerated computation.

This implementation leverages PyTorch’s optimized tensor operations and optional GPU acceleration for faster distance computation. Memory usage is controlled through chunked processing.

Mathematical Optimization

Instead of computing \(\|\mathbf{r} - \mathbf{M}\|^2\) directly, we use the identity:

\[\|\mathbf{r} - \mathbf{M}\|^2 = \|\mathbf{r}\|^2 - 2\langle\mathbf{r},\mathbf{M}\rangle + \|\mathbf{M}\|^2\]

Since \(\|\mathbf{r}\|^2\) is constant for \(\arg\min\), we compute:

\[\text{distance} \propto -2\langle\mathbf{r},\mathbf{M}\rangle + \|\mathbf{M}\|^2\]

param response:

Sensory response array (numpy or torch tensor)

type response:

np.ndarray or torch.Tensor

param res_maps:

Response template maps (numpy or torch tensor)

type res_maps:

np.ndarray or torch.Tensor

param device:

Computation device ('cpu', 'cuda', or torch.device object). If None, automatically selects CUDA if available

type device:

str, torch.device, or None

param chunk_size:

Number of responses to process simultaneously. Larger values use more memory but may be faster

type chunk_size:

int, default=1024

returns:

• tuple of (np.ndarray, bool) – (decoded_coordinates, is_trajectory) as numpy arrays

• Performance Notes

• —————–

• - GPU acceleration provides 10-100x speedup for large problems

• - Chunked processing prevents out-of-memory errors

• - Automatic mixed precision could be added for further optimization

• Memory Usage

• ————

• Peak memory ≈ chunk_size × H × W × 4 bytes (for float32)

• r

rtgym.utils.masking module

class Masking(m_max: float = 0.5, m_min: float = 0.0, sigma_t: float = 0, sigma_d: float = 0, t_warmup: int = 0, device: torch.device = torch.device)

Bases: object

Base class for all masking utilities.

apply_gaussian_blur(mask: torch.Tensor)

gaussian_kernel_1d(kernel_size: int, sigma: float)

Creates a 1D Gaussian kernel.

mask(x, mask_idx=None)

new_mask(x)

Generates a mask with different masking thresholds for each dimension. The thresholds are determined based on per-batch and per-dimension masking ratios.

Parameters:

x (torch.Tensor) – Input tensor of shape (batch_size, T, D)

Returns:

Mask tensor of the same shape as x

Return type:

torch.Tensor

set_m_max(m_max: float)

set_m_min(m_min: float)

to(device: torch.device)

to_tensor(x)

rtgym.utils.verbose module

compute_layout(fields, n_cols)

Computes the layout for the grid of plots.

generate_titles(fields, titles)

Generates titles for the plots.

plot_dual_field(ax_, field, cmap, title)

Example dual field plotting function.

plot_dual_fields(ax_, field_0, field_1, title)

Plots dual firing fields on the given axis.

plot_single_field(ax_, field, cmap, title)

Default function for plotting a single field.

prepare_fields(fields, mask=None, bad_color='black', cmap='jet')

Prepares the fields by applying a mask and setting up the colormap.

print_dict(params)

title_single_field(i, field)

Title generator for single fields.

visualize_fields(fields, titles=None, n_cols=10, cmap='jet', mask=None, bad_color='black', plot_func=None, title_func=None)

Visualizes fields of cells using a custom plotting and title generation function.

Parameters:

• fields (list or array) – Fields to be visualized.

• titles (list, optional) – Titles for each field. Default is None.

• n_cols (int) – Number of columns in the grid layout. Default is 5.

• cmap (str) – Colormap for visualization. Default is ‘jet’.

• mask (array, optional) – Mask to apply to fields, setting masked areas to NaN. Default is None.

• bad_color (str) – Color for masked regions. Default is ‘black’.

• plot_func (callable) – Function to handle the plotting for each field.

• title_func (callable) – Function to generate titles for each field.

Returns:

Matplotlib figure and axes objects.

Return type:

fig, ax

Module contents