Modelling

Bases: _NetworkBase

A PyTorch module representing a multilayered neural network model with trainable parameters and flexible activation functions.

This network architecture is designed to process temporal sequences of input data through multiple synaptic hops, with the initial layer handling only external inputs and subsequent layers processing external+internal input.

Excitabililty is implemented as the slope of the tanh activation function. Users can specify default excitability (tanh_steepness) for all neurons, as well as specific excitabilities for neuron groups. In training, the excitabilities are shared per group specified by idx_to_group.

Baseline activity is taken into account through biases. Users can specify default biases for all neurons, as well as specific biases for neuron groups. In training, the biases are shared per group specified by idx_to_group.

In divisive normalisation, the synaptic connections specified in divisive_normalization no longer participate in the normal subtractive inhibition. Rather, they are used to change the slope of the post-synaptic neuron: the new slope is: max(0, original_slope * (1 + pre_post_weight * pre_activation * divisive_strength)). divisive_strength is a parameter the user can set, to specity how strong the divisive normalization is. max(0, …) is to ensure that the slope does not become negative.

all_weights

The connectome. Input neurons are in the columns.

Type:: torch.nn.Parameter

sensory_indices

Indices indicating which rows/columns in the all_weights matrix correspond to sensory neurons.

Type:: list[int]

num_layers

The number of layers in the network.

Type:: int

threshold

The activation threshold for neurons in the network.

Type:: float

activations

An array storing the activations of all (batches of) neurons (rows) across time steps (columns).

Type:: numpy.ndarray

custom_activation_function

A custom activation function to use instead of the default. The function should take the model, the input tensor x, and the previous activations x_previous as arguments, and return the output tensor after applying the activation function.

Type:: Callable

default_bias

Default bias value for all neurons.

Type:: float

slope

Trainable parameter for the steepness of the tanh activation function, grouped by neuron type.

Type:: torch.nn.Parameter

raw_biases

Trainable parameter for the biases of neurons, grouped by neuron type.

Type:: torch.nn.Parameter

indices

Indices mapping each neuron to its group for parameter sharing.

Type:: torch.Tensor

divisive_strength

Trainable parameter for the strength of divisive normalization, if applicable. In training, the parameter is shared per group specified by idx_to_group. This parameter only exists for the pres in divisive_normalization.

Type:: torch.nn.Parameter

divisive_normalization

A dictionary where keys are pre-synaptic neuron groups and values are lists of post-synaptic neuron groups. These inhibitory connections are implemented divisively instead of subtractively.

Type:: Dict[str, List[str]]

Parameters:

all_weights (Union[torch.Tensor, scipy.sparse.spmatrix]) – The connectome. Input neurons are in the columns.
sensory_indices (list[int]) – A list indicating the indices of sensory neurons within the network.
num_layers (int, optional) – The number of temporal layers to unroll the network through. Defaults to 2.
threshold (float, optional) – The threshold for activation of neurons. Defaults to 0.01.
tanh_steepness (float, optional) – Default steepness for tanh activation. Defaults to 5.
idx_to_group (dict, optional) – Mapping from neuron indices to group names for parameter sharing (all neurons in the same group share the same parameter). Defaults to None (no parameter sharing).
default_bias (float, optional) – Default bias value for all groups. Defaults to 0.0.
bias_dict (dict, optional) – Dict mapping group names to bias values. Defaults to None (uses default_bias).
slope_dict (dict, optional) – Dict mapping group names to slope values. Uses tanh_steepness if None.
divisive_normalization (Dict[str, List[str]], optional) – A dictionary where keys are pre-synaptic neuron groups and values are lists of post-synaptic neuron groups. These inhibitory connections are implemented divisively instead of subtractively. Defaults to None.
divisive_strength (Union[float, int, dict], optional) – The strength of the divisive normalization. If a float or int, it applies to all connections. If a dict, it maps pre-synaptic neuron groups to their specific divisive strengths. Defaults to 1.
activation_function (Callable, optional) – Custom activation function. If None, uses default implementation.
tau (float, optional) – Time constant. Higher tau results in slower changes. Minimum 1, where the activation at the current time step is solely determined by the current input. Defaults to 10.
device (torch.device, optional) – Device for computation.

activation_function(x: Tensor, x_previous: Tensor | None = None, slopes_full: Tensor | None = None, biases_full: Tensor | None = None, taus_full=None) → Tensor[source]

Apply the activation function to the input tensor. Currently it is tanh(relu(slope * x + bias)). x = weights @ inputs (done outside the activation function).

Parameters:

x (torch.Tensor) – The input tensor to the activation function.
x_previous (Optional[torch.Tensor]) – The previous activations of the neurons.
slopes_full (Optional[torch.Tensor]) – Pre-expanded slopes, shape (num_neurons, 1). If None, computed from self.slope[self.indices].
biases_full (Optional[torch.Tensor]) – Pre-expanded biases, shape (num_neurons, 1). If None, computed from self.biases[self.indices].
taus_full – Pre-expanded taus, shape (num_neurons, 1), or scalar. If None, computed from self.effective_tau[self.indices].

Returns:

The output tensor after applying the activation function.

Return type:

torch.Tensor

forward(inputs: Tensor, manipulate: Dict[int, Dict[int | str, float]] | Dict[int, Dict[int, Dict[int | str, float]]] | None = None, checkpoint_steps: int = 0, return_layer_list: bool = False)[source]

Parameters:: return_layer_list (bool) – If True, return List[Tensor] of length num_layers, each shape (neurons, batch), on GPU with grad_fn preserved. Required for layer-wise gradient attribution via torch.autograd.grad — stacking/transposing makes per-layer tensors siblings of a shared parent, which breaks the graph path autograd.grad needs. When True, self.activations is not populated.

Bases: _NetworkBase

A PyTorch module representing a multilayered neural network model with trainable parameters and flexible activation functions.

This network architecture is designed to process temporal sequences of input data through multiple synaptic hops, with the initial layer handling only external inputs and subsequent layers processing external+internal input.

Excitabililty is implemented as the slope of the tanh activation function. Users can specify default excitability (tanh_steepness) for all neurons, as well as specific excitabilities for neuron groups. In training, the excitabilities are shared per group specified by idx_to_group.

Baseline activity is taken into account through biases. Users can specify default biases for all neurons, as well as specific biases for neuron groups. In training, the biases are shared per group specified by idx_to_group.

In divisive normalisation, the synaptic connections specified in divisive_normalization no longer participate in the normal subtractive inhibition. Rather, they are used to change the slope of the post-synaptic neuron: the new slope is: max(0, original_slope * (1 + pre_post_weight * pre_activation * divisive_strength)). divisive_strength is a parameter the user can set, to specity how strong the divisive normalization is. max(0, …) is to ensure that the slope does not become negative.

all_weights

The connectome. Input neurons are in the columns.

Type:: torch.nn.Parameter

sensory_indices

Indices indicating which rows/columns in the all_weights matrix correspond to sensory neurons.

Type:: list[int]

num_layers

The number of layers in the network.

Type:: int

threshold

The activation threshold for neurons in the network.

Type:: float

activations

An array storing the activations of all (batches of) neurons (rows) across time steps (columns).

Type:: numpy.ndarray

custom_activation_function

A custom activation function to use instead of the default. The function should take the model, the input tensor x, and the previous activations x_previous as arguments, and return the output tensor after applying the activation function.

Type:: Callable

default_bias

Default bias value for all neurons.

Type:: float

slope

Trainable parameter for the steepness of the tanh activation function, grouped by neuron type.

Type:: torch.nn.Parameter

raw_biases

Trainable parameter for the biases of neurons, grouped by neuron type.

Type:: torch.nn.Parameter

indices

Indices mapping each neuron to its group for parameter sharing.

Type:: torch.Tensor

divisive_strength

Trainable parameter for the strength of divisive normalization, if applicable. In training, the parameter is shared per group specified by idx_to_group. This parameter only exists for the pres in divisive_normalization.

Type:: torch.nn.Parameter

divisive_normalization

A dictionary where keys are pre-synaptic neuron groups and values are lists of post-synaptic neuron groups. These inhibitory connections are implemented divisively instead of subtractively.

Type:: Dict[str, List[str]]

Parameters:

all_weights (Union[torch.Tensor, scipy.sparse.spmatrix]) – The connectome. Input neurons are in the columns.
sensory_indices (list[int]) – A list indicating the indices of sensory neurons within the network.
num_layers (int, optional) – The number of temporal layers to unroll the network through. Defaults to 2.
threshold (float, optional) – The threshold for activation of neurons. Defaults to 0.01.
tanh_steepness (float, optional) – Default steepness for tanh activation. Defaults to 5.
idx_to_group (dict, optional) – Mapping from neuron indices to group names for parameter sharing (all neurons in the same group share the same parameter). Defaults to None (no parameter sharing).
default_bias (float, optional) – Default bias value for all groups. Defaults to 0.0.
bias_dict (dict, optional) – Dict mapping group names to bias values. Defaults to None (uses default_bias).
slope_dict (dict, optional) – Dict mapping group names to slope values. Uses tanh_steepness if None.
divisive_normalization (Dict[str, List[str]], optional) – A dictionary where keys are pre-synaptic neuron groups and values are lists of post-synaptic neuron groups. These inhibitory connections are implemented divisively instead of subtractively. Defaults to None.
divisive_strength (Union[float, int, dict], optional) – The strength of the divisive normalization. If a float or int, it applies to all connections. If a dict, it maps pre-synaptic neuron groups to their specific divisive strengths. Defaults to 1.
activation_function (Callable, optional) – Custom activation function. If None, uses default implementation.
tau (float, optional) – Time constant. Higher tau results in slower changes. Minimum 1, where the activation at the current time step is solely determined by the current input. Defaults to 10.
device (torch.device, optional) – Device for computation.

activation_function(x: Tensor, x_previous: Tensor | None = None, slopes_full: Tensor | None = None, biases_full: Tensor | None = None, taus_full=None) → Tensor[source]

Apply the activation function to the input tensor. Currently it is tanh(relu(slope * x + bias)). x = weights @ inputs (done outside the activation function).

Parameters:

x (torch.Tensor) – The input tensor to the activation function.
x_previous (Optional[torch.Tensor]) – The previous activations of the neurons.
slopes_full (Optional[torch.Tensor]) – Pre-expanded slopes, shape (num_neurons, 1). If None, computed from self.slope[self.indices].
biases_full (Optional[torch.Tensor]) – Pre-expanded biases, shape (num_neurons, 1). If None, computed from self.biases[self.indices].
taus_full – Pre-expanded taus, shape (num_neurons, 1), or scalar. If None, computed from self.effective_tau[self.indices].

Returns:

The output tensor after applying the activation function.

Return type:

torch.Tensor

forward(inputs: Tensor, manipulate: Dict[int, Dict[int | str, float]] | Dict[int, Dict[int, Dict[int | str, float]]] | None = None, checkpoint_steps: int = 0, return_layer_list: bool = False)[source]

Parameters:: return_layer_list (bool) – If True, return List[Tensor] of length num_layers, each shape (neurons, batch), on GPU with grad_fn preserved. Required for layer-wise gradient attribution via torch.autograd.grad — stacking/transposing makes per-layer tensors siblings of a shared parent, which breaks the graph path autograd.grad needs. When True, self.activations is not populated.

class connectome_interpreter.activation_maximisation.TargetActivation(targets: Dict[int, Dict[int, float]] | DataFrame, batch_size: int | None = None)[source]

Bases: object

Dataclass to handle target activations for activation maximisation.

The target activations can be specified as a dictionary or a DataFrame. The dictionary should have the following structure:

{layer: {neuron_index: target_activation_value}}

The DataFrame should have the following columns:

‘batch’: The batch index.
‘layer’: The layer index.
‘neuron’: The neuron index.
‘value’: The target activation value.

Parameters:

targets (Union[Dict[int, Dict[int, float]], pd.DataFrame]) – The target activations. If a dictionary, all batches will have the same target. If a DataFrame, each row represents a target activation for a specific batch.
batch_size (Optional[int], optional) – The number of batches. Defaults to None.

batch_size: int | None = None

get_batch_targets(batch_idx: int) → Dict[int, Dict[int, float]][source]

targets: Dict[int, Dict[int, float]] | DataFrame

connectome_interpreter.activation_maximisation.activated_path_for_ngl(path)[source]

Convert a path DataFrame [with ‘pre_activation’ and ‘post_activation’ columns] to a format suitable for Neuroglancer visualization (get_ngl_link(df_format=’long’)). Neurons are coloured by their activation.

Parameters:: path (pd.DataFrame) – A DataFrame containing the columns ‘pre’, ‘post’, ‘layer’, ‘pre_activation’, and ‘post_activation’ (standard output from function activations_to_df() and get_activations_for_path()).
Returns:: A DataFrame with columns ‘neuron_id’, ‘layer’, and ‘activation’, suitable for Neuroglancer visualization.
Return type:: pd.DataFrame

connectome_interpreter.activation_maximisation.activation_maximisation(model: MultilayeredNetwork, target_activations: TargetActivation, input_tensor: Tensor | None = None, num_iterations: int = 50, learning_rate: float = 0.1, in_reg_lambda: float = 0.01, out_reg_lambda: float = 0.01, custom_reg_functions: Dict[str, Callable[[Tensor], Tensor]] | None = None, early_stopping: bool = True, stopping_threshold: float = 1e-05, n_runs: int = 10, use_tqdm: bool = True, print_output: bool = True, device: device | None = None, wandb: bool = False, seed: int | None = None, normalize_gradients: bool = False, quiet: bool = False) → Tuple[ndarray, ndarray, List[float], List[float], List[float], List[ndarray]][source]

Performs activation maximisation on a given model to identify input patterns that result in the target activations.

This is done by adjusting the input tensor over num_iterations using gradient descent, while also regularising the overall input and output (to keep activated neurons sparse).

Parameters:

model – A PyTorch model with activations, sensory_indices, and threshold attributes.
target_activations (TargetActivation) – Target activations specification.
input_tensor (torch.Tensor, optional) – The initial tensor to optimize. If None, a random tensor is created. Defaults to None.
num_iterations (int, optional) – The number of iterations to run the optimization for. Defaults to 50.
learning_rate (float, optional) – The learning rate for the optimizer. Defaults to 0.1.
in_reg_lambda (float, optional) – The coefficient for input regularization. Defaults to 0.01.
out_reg_lambda (float, optional) – The coefficient for output regularization. Defaults to 0.01.
custom_reg_functions (Dict[str, Callable[[torch.Tensor]], optional) – A dictionary with keys ‘in’ and ‘out’ that map to functions that calculate the input and output regularization losses, respectively. If None, the default regularization function (L1 plus L2) is used. Defaults to None.
early_stopping (bool, optional) – Whether to stop the optimization early if the difference between the biggest and the smallest loss within the last n_runs falls below stopping_threshold. Defaults to True.
stopping_threshold (float, optional) – The threshold for early stopping. Defaults to 1e-5.
n_runs (int, optional) – The number of runs to consider for early stopping. Defaults to 10.
report_memory_usage (bool, optional) – Whether to report GPU memory usage during optimization. Defaults to False.
device – The device to run the optimization on. If None, automatically selects a device. Defaults to None.
wandb (bool, optional) – Whether to log optimization details to Weights & Biases (https://wandb.ai/site/). Defaults to True. Requires wandb to be installed.
seed (int, optional) – Random seed for reproducible input tensor initialization. Only used when input_tensor is None. Defaults to None (no seed set).
normalize_gradients (bool, optional) – Whether to normalize gradients per timepoint to mitigate vanishing gradients for longer path lengths. When True, gradients are normalized by their L2 norm for each timepoint independently. This helps ensure that inputs targeting neurons at deeper layers receive comparable gradient magnitudes to those at earlier layers. Defaults to False.

Returns:

A tuple containing:

numpy.ndarray: The optimized input as a numpy array.

numpy.ndarray: The output of the model after optimization as a numpy array.

list(float): A list of output activation losses over iterations.

list(float): A list of input activation regularization losses over iterations.

list(float): A list of output activation regularization losses over iterations.

list(numpy.ndarray): A list of input tensor snapshots taken during optimization.

Return type:

tuple

Examples

# Single target for multiple batches
targets_dict = {
    # layer 0: neuron 1 -> 0.5, neuron 2 -> 0.8
    0: {1: 0.5, 2: 0.8},
    # layer 1: neuron 0 -> 0.3
    1: {0: 0.3}
}
targets = TargetActivation(targets=targets_dict, batch_size=4)
inputs, outputs, *losses = activation_maximisation(model, targets)

# Different targets per batch using DataFrame
targets_df = pd.DataFrame([
    {'batch': 0, 'layer': 0, 'neuron': 1, 'value': 0.5},
    {'batch': 0, 'layer': 0, 'neuron': 2, 'value': 0.8},
    {'batch': 1, 'layer': 1, 'neuron': 0, 'value': 0.3}
])
# batch_size inferred
targets = TargetActivation(targets=targets_df)
results = activation_maximisation(model, targets)

# Custom regularization
def sparse_reg(x):
    return torch.sum(torch.abs(x))
custom_reg = {'in': sparse_reg, 'out': sparse_reg}
results = activation_maximisation(
    model, targets, custom_reg_functions=custom_reg
)

connectome_interpreter.activation_maximisation.activations_to_df(inprop, model_input: ndarray, out: ndarray, sensory_indices: List[int], inidx_mapping: dict | None = None, outidx_mapping: dict | None = None, activation_threshold: float = 0, connectivity_threshold: float = 0, high_ram: bool = True) → DataFrame[source]

Generates a dataframe representing the paths in a layered plot,filtering by activation and connectivity thresholds.

This function takes the direct connectivity matrix (inprop), input neuron activity, output neuron activity, indices for sensory neurons, and mapping between input and output indices to groups. It generates a dataframe that represents the paths through the network layers.

Parameters:

inprop (scipy.sparse matrix or numpy.ndarray) – Matrix representing the synaptic strengths between neurons, can be dense or sparse. Presynaptic is in the rows, postsynaptic in the columns.
model_input (numpy.ndarray) – A 2D array representing input to the network. Neurons are in the rows, timepoints in the columns. Only the first timepoint is used, since out is expected to have activity of all neurons, including input neurons.
out (numpy.ndarray) – A 2D array representing the output from the network. The second dimension represents timepoints.
sensory_indices (list of int) – A list of indices corresponding to sensory neurons in inprop.
inidx_mapping (dict, optional) – A dictionary mapping indices in inprop to new indices (e.g. cell type). If None, indices are not remapped. Defaults to None.
outidx_mapping (dict, optional) – A dictionary mapping indices in out to new indices. If None, inidx_mapping is used for mapping. Defaults to None.
activation_threshold (float, optional) – A threshold value for activation. Neurons with activations below this threshold are not considered. Defaults to 0.
connectivity_threshold (float, optional) – A threshold for filtering connections. Connections with weights below this threshold are ignored. Defaults to 0.
high_ram (bool, optional) – Whether to use a high RAM implementation (which is slightly faster). This implementation gets direct connections between all relevant neurons at once, instead of within each layer. Defaults to True.

Returns:

A dataframe representing the paths in the network. Each row is a connection, with columns for ‘pre’ and ‘post’ neuron indices, ‘layer’, and their respective activations (‘pre_activation’, ‘post_activation’).

Return type:

pandas.DataFrame

connectome_interpreter.activation_maximisation.activations_to_df_batched(inprop, opt_in: ndarray, out: ndarray, sensory_indices: List[int], inidx_mapping: dict | None = None, outidx_mapping: dict | None = None, activation_threshold: float = 0, connectivity_threshold: float = 0, high_ram: bool = True) → DataFrame[source]

Generates a dataframe representing the paths in a layered plot, filtering by activation and connectivity thresholds.

This function takes the direct connectivity matrix (inprop), optimal input neuron activity, output neuron activity, indices for sensory neurons, and mapping between input and output indices to groups. It generates a dataframe that represents the paths through the network layers.

Parameters:

inprop (scipy.sparse matrix or numpy.ndarray) – Matrix representing the synaptic strengths between neurons, can be dense or sparse. Presynaptic is in the rows, postsynaptic in the columns.
opt_in (numpy.ndarray) – A 3D array representing optimal input to the network. The first dimension represents the batch size, the second dimension represents input neurons, and the third dimension represents timepoints. Only the first timepoint is used, since out is expected to have activity of all neurons, including input neurons.
out (numpy.ndarray) – A 3D array representing the activation of all neurons. The first dimension represents the batch size, the second dimension represents all neurons, and the third dimension represents timepoints.
sensory_indices (list of int) – A list of indices corresponding to sensory neurons in inprop.
inidx_mapping (dict, optional) – A dictionary mapping indices in inprop to new indices (e.g. cell type). If None, indices are not remapped. Defaults to None.
outidx_mapping (dict, optional) – A dictionary mapping indices in out to new indices. If None, inidx_mapping is used for mapping. Defaults to None.
activation_threshold (float, optional) – A threshold value for activation. Neurons with activations below this threshold are not considered. Defaults to 0.
connectivity_threshold (float, optional) – A threshold for filtering connections. Connections with weights below this threshold are ignored. Defaults to 0.
high_ram (bool, optional) – Whether to use a high RAM implementation (which is slightly faster). This implementation gets direct connections between all relevant neurons at once, instead of within each layer. Defaults to True.

Returns:

A dataframe representing the paths in the network. Each row is a connection, with columns for ‘pre’ and ‘post’ neuron indices, ‘layer’, and their respective activations (‘pre_activation’, ‘post_activation’).

Return type:

pandas.DataFrame

Given the input arguments, return a dataframe with columns ‘normalised_column’, ‘cell_group’, ‘time_point’, ‘activation’, and ‘time_step’. This gives the raw data for plot_activity_by_column(). Note, if both model_input and sensory_indices are provided, the first timepoint of model_input is also included (time_step = 0). Other timesteps are already included in activity.

Parameters:

activity (torch.Tensor | numpy.ndarray) – The activity of the model. Shape should be (num_neurons, num_timepoints).
idx_to_group (dict) – A dictionary mapping indices from the model to the groups of interest (e.g. cell type). max(idx_to_group.keys()) should be equal to number of units in the model.
idx_to_column (dict) – A dictionary mapping indices from the model to the columns of interest (e.g. column in the central complex).
selected_group (arrayable) – The groups to select from the activity. This should be a list of groups that are present in idx_to_group.values().
model_input (numpy.ndarray | torch.Tensor, optional) – The input to the model. Shape should be (num_neurons, num_timepoints). If provided, the first timepoint of model_input is also included in the output dataframe (time_step = 0). Defaults to None.
sensory_indices (arrayable, optional) – The indices of sensory neurons. If provided, it should be a list of indices that are present in idx_to_group. If provided, it must also be provided with model_input. Defaults to None.

Returns:

A dataframe with columns ‘normalised_column’, ‘cell_group’, ‘time_point’, ‘activation’, and ‘time_step’. The first timepoint of model_input is also included if model_input is provided.

Return type:

pd.DataFrame

connectome_interpreter.activation_maximisation.add_sign(inprop: spmatrix, idx_to_sign: dict)[source]

Add sign to the inprop matrix based on the idx_to_sign dictionary, and transpose it such that the pre is in the columns (required by torch).

Parameters:

inprop (scipy.sparse matrix) – The matrix to which the sign will be added.
idx_to_sign (dict) – A dictionary mapping indices to their sign (-1 or 1).

Returns:

The matrix with the sign added.

Return type:

scipy.sparse matrix

Get the activations for the pre and post neurons in the path, based on the activations of the model and the input.

Parameters:

path (pd.DataFrame) – A dataframe representing the paths in the network. Each row is a connection, with columns for ‘pre’ and ‘post’ neuron indices, and ‘layer’.
activations (torch.Tensor | numpy.ndarray) – The activations of the model. Shape should be (num_neurons, num_layers).
model_in (torch.Tensor | numpy.ndarray) – The input to the model. Shape should be (num_neurons, something) - only the first column (num_neurons, 0), is used, when there is 1 in ‘layer’ in path. It is otherwise not used.
sensory_indices (arrayable) – The indices of sensory neurons.
idx_to_group (dict, optional) – A dictionary mapping indices from the model to the groups in path (e.g. cell type). Defaults to None.
activation_start (int, optional) – Which layer corresponds to the start layer of path. By default, activation_start = 0, the sensory-only layer in the model corresponds to path.pre[path.layer == 1]. If you want activations[:,1] (i.e. two timesteps forward) to correspond to path.pre[path.layer == 1], set activation_start = 2. If you want the last timepoint to correspond to the last layer in path, set activation_start = activations.shape[1] - path.layer.max().

Returns:

The activations for the pre and post neurons in the path.

Return type:

pd.DataFrame

Compute gradients of target neurons with regard to monitor neurons, i.e. the rate of change of target neuron activations with respect to monitor neurons.

Parameters:

model – MultilayeredNetwork instance.
input_tensor – Input tensor to the model. Shape: (batch_size, num_input_neurons, num_timesteps) or (num_input_neurons, num_timesteps)
monitor_neurons – List of neuron indices to monitor gradients for. If model’s idx_to_group attribute is not None (e.g. idx_to_type), these should be group names.
target_neurons – Dictionary mapping layer/timepoint indices to lists of neurons. If model’s idx_to_group attribute is not None (e.g. idx_to_type), these should be group names. e.g. {5: [‘DA1_lPN’]} means calculating gradients of layer 5 neurons in group ‘DA1_lPN’.
monitor_layers – List of layer indices to monitor gradients at. If None, monitor all layers.
method – Gradient computation method. “vanilla” or “input_x_gradient”.
batch_names – Optional list of batch names corresponding to the input tensor. If not provided, defaults to [“batch_0”, “batch_1”, …].
device – Optional torch device to run the computations on. If None, uses CUDA if available, else CPU.
checkpoint_steps (int) – Number of layers to include in each checkpoint when using gradient checkpointing. Adjust based on available memory and model size. Defaults to 50.
normalize_per_layer (bool) – If True, divide each layer’s gradient tensor by its own L2 norm before returning. This mitigates exploding/vanishing gradients across time (early layers can have gradients many orders of magnitude larger than late ones when the per-step Jacobian’s spectral radius is far from 1). With this enabled, cross-layer comparisons reflect the relative pattern of which inputs matter most at each layer, not the absolute magnitude. Defaults to False.

Returns:

DataFrame containing gradients with columns: ‘group’, ‘batch_name’, ‘time_0’, ‘time_1’, …, ‘time_N’.

Return type:

pd.DataFrame

Get selected input activation.

Parameters:

model_in (numpy.ndarray | torch.Tensor) – The input to the model. Shape should be (num_neurons, num_timepoints) for 2D or (batch_size, num_neurons, num_timepoints) for 3D.
sensory_indices (arrayable) – All the indices of sensory neurons. Length should be equal to the first dimension of model_in (2D) or second dimension (3D).
idx_to_group (dict) – A dictionary mapping indices from the model to the groups of interest (e.g. cell type). max(idx_to_group.keys()) should be equal to number of units in the model.
selected_indices (arrayable, optional) – The indices of the neurons to select. Defaults to None. If None, all neurons are selected.
activation_threshold (float, optional) – A threshold value for the absolute value of activation. Neurons with activations below this threshold are not considered. Defaults to 0.
batch_names (arrayable, optional) – The names of the batches. Defaults to None. If model_in.ndim == 3, then this should be supplied. If not, batch names will be e.g. ‘batch_0’, ‘batch_1’, etc.

Notes

If there are two neurons in a cell type, one passes the activation threshold and the other doesn’t, only the one that passes the threshold is selected.

Returns:: The input activation for the model. For 2D input, columns are ‘group’ and ‘time_0’, ‘time_1’, etc. For 3D input, columns are ‘batch_name’, ‘group’ and ‘time_0’, ‘time_1’, etc.
Return type:: pd.DataFrame

Get the activations for specified indices across timepoints, include batch name and group information when available. Memory efficiency provided by Claude.

Parameters:

activations (torch.Tensor | numpy.ndarray) – Output activation from the model. Shape should be (batch_size, num_neurons, num_timepoints) or (num_neurons, num_timepoints).
neuron_indices (arrayable) – The indices of the neurons to get activations for.
batch_names (arrayable, optional) – The names of the batches. Defaults to None. If activations.ndim == 3, then this should be supplied. If not, batch names will be e.g. ‘batch_0’, ‘batch_1’, etc.
idx_to_group (dict, optional) – A dictionary mapping indices to groups. Defaults to None.

Returns:

The activations for the neurons, with the first columns being batch_names, neuron_indices, and group. The rest are the timesteps.

Return type:

pd.DataFrame

Guesses optimal stimulus based on signed effective connectivity (no non-linearity). Since it calculates effective connectivity on the fly (behind the scene using signed_effective_conn_from_paths()), it would be good to have a small-ish path length (e.g. 5).

When the layer number in target_activations is > longest_plen, the function returns a stimulus of shape ((batch,) num_sensory_neurons, longest_plen), instead of shape ((batch,) num_sensory_neurons, layer), appropriate for the stimulus immediately before the large layer number in target_activations.

e.g. if target_activations has layer 10, and longest_plen = 5, the user should first make a random stimulus of shape ((batch,) num_sensory_neurons, 10), and then replace the last 5 timesteps with the output of this function.

Parameters:

inprop (sparse.spmatrix) – The inprop matrix of the model. Pre is in the rows, post in the columns. There are only positive values in this matrix, the sign information is added using idx_to_sign.
sensory_indices (arrayable) – The sensory indices of the model.
idx_to_sign (dict) – A dictionary mapping indices to their sign (-1 or 1).
target_activations (TargetActivation) – The target activations.
longest_plen (int, optional) – The longest path length to consider. Defaults to 5.

Returns:

The guessed optimal stimulus. Shape is ((batch,) num_sensory_neurons, min(longest_plen, max_layer_in_target_activations)).

Return type:

np.ndarray

connectome_interpreter.activation_maximisation.input_from_df(df: ~pandas.core.frame.DataFrame, sensory_indices: list, idx_to_group: dict, num_layers: int, timepoints: int | ~typing.List[int] | ~numpy.ndarray | set = 0, dtype: type = <class 'numpy.float32'>) → ndarray[source]

Make well-formatted input for the model, based on defined vectors of input neuron activation (df). The function returens a 3D tensor, with the first dimension being the batch size, the second dimension being the number of input neurons of the model, and the third dimension being the number of layers.

Parameters:

df – pd.DataFrame Rows correspond to the values in idx_to_group, and columns correspond to the batches (number of columns = number of batch). For instance, rows of df can be olfactory glomeruli, and columns can be odours. df is the reaction of each glomerulus to each odour.
sensory_indices – list The indices of sensory neurons. For instance, these could be the indices of individual olfactory receptor neurons.
idx_to_group – dict A dictionary that maps indices to group. For instance, this could map from indices of indivudal olfactory receptor neuron to the glomerulus they innervate (rows of df).
num_layers – int The number of layers in the model.
timepoints – int or list or np.ndarray or set, optional The timepoints to put the input in. For instance, if timepoints=[0, 1], then the same input will be put in timepoint 0 and 1. Defaults to 0.
dtype – type, optional The data type of the output array. Defaults to np.float32.

Returns:

The input for the model.

Return type:

np.ndarray

connectome_interpreter.activation_maximisation.legacy_activation_function(self, x: Tensor, x_previous: Tensor | None = None) → Tensor[source]

Apply the activation function to the input tensor. Currently it is tanh(relu(slope * x + bias)). x = weights @ inputs (done outside the activation function). Note if using this, then tau should be a number between 0 and 1, as it represents the extent to which the activation persists to the next timestep.

Parameters:

x (torch.Tensor) – The input tensor to the activation function.
x_previous (Optional[torch.Tensor]) – The previous activations of the neurons.

Returns:

The output tensor after applying the activation function.

Return type:

torch.Tensor

connectome_interpreter.activation_maximisation.legacy_activation_function_linear(self, x: Tensor, x_previous: Tensor | None = None) → Tensor[source]

Apply the activation function to the input tensor. Currently it is tanh(relu(slope * x + bias)). x = weights @ inputs (done outside the activation function). Note if using this, then tau should be a number between 0 and 1, as it represents the extent to which the activation persists to the next timestep.

Parameters:

x (torch.Tensor) – The input tensor to the activation function.
x_previous (Optional[torch.Tensor]) – The previous activations of the neurons.

Returns:

The output tensor after applying the activation function.

Return type:

torch.Tensor

Take output from activity_by_column() and plot the activity per neuron group, per time step, per column. The x axis is the normalised column, normalised within each neuron group.

Parameters:

activity (torch.Tensor | numpy.ndarray) – The activity of the model. Shape should be (num_neurons, num_timepoints).
idx_to_group (dict) – A dictionary mapping indices from the model to the groups of interest (e.g. cell type). max(idx_to_group.keys()) should be equal to number of units in the model.
idx_to_column (dict) – A dictionary mapping indices from the model to the columns of interest (e.g. column in the central complex).
selected_group (arrayable) – The groups to select from the activity. This should be a list of groups that are present in idx_to_group.values().
plot_type (str, optional) – The type of plot to create. Can be either ‘scatter’ or ‘line’. Defaults to ‘line’.
model_input (numpy.ndarray | torch.Tensor, optional) – The input to the model. Shape should be (num_neurons, num_timepoints). If provided, the first timepoint of model_input is also included in the output dataframe (time_step = 0). Defaults to None.
sensory_indices (arrayable, optional) – The indices of sensory neurons. If provided, it should be a list of indices that are present in idx_to_group. If provided, it must also be provided with model_input. Defaults to None.
figsize (tuple, optional) – The size of the figure in pixels. Defaults to (800, 600).
global_min (float, optional) – The minimum value for the y-axis. If None, the minimum value is set to the smaller of 0, and the minimum activation value across all groups and time steps. Defaults to None.
global_max (float, optional) – The maximum value for the y-axis. If None, the maximum value is set to the bigger of 1, and the maximum activation value across all groups and time steps. Defaults to None.

Returns:

A Plotly figure object.

Return type:

plotly.graph_objects.Figure

connectome_interpreter.activation_maximisation.plot_timeseries(df: DataFrame, style: dict | None = None, sizing: dict | None = None, x_label: str = 'Time', y_label: str = 'Activation', title: str | None = None, slider_dim: str | None = 'batch', ymin: float | None = None, ymax: float | None = None, scatter_mode: str = 'lines') → Figure[source]

Generate an interactive time-series plot of neural activations, based on a dataframe that’s like the output of get_neuron_activation().

If both ‘batch_name’ and ‘group’ columns exist, the slider_dim argument decides which dimension the slider animates over.

Parameters:

df (pd.DataFrame) – Must contain ‘group’ plus one or more ‘time_*’ columns (e.g. ‘time_0’, ‘time_1’ …). A ‘batch_name’ column is optional.
style (Optional[dict]) –
Plot styling; keys:
- ’font_type’: str, default=’Arial’
- ’linecolor’: str, default=’black’
- ’papercolor’: str, default=’rgba(255,255,255,255)’
sizing (Optional[dict]) –
Layout sizing; keys (px / pt):
- ’fig_width’: int, default=600
- ’fig_height’: int, default=400
- ’fig_margin’: int, default=0
- ’fsize_ticks_pt’: int, default=12
- ’fsize_title_pt’: int, default=16
- ’markersize’: int, default=8
- ’ticklen’: int, default=5
- ’tickwidth’: int, default=1
- ’axislinewidth’: int, default=1.5
- ’markerlinewidth’: int, default=1
x_label (str) – X-axis label.
y_label (str) – Y-axis label.
title (Optional[str]) – Figure title.
slider_dim (str | None) –
Dimension for slider if both are present.
- ’batch’ (default): slider toggles batch_name, traces coloured by group.
- ’group’: slider toggles group, traces coloured by batch_name.
- None: no slider (all traces in one panel).
ymin (Optional[float]) – Minimum y-axis value. If None, set to minimum activation.
ymax (Optional[float]) – Maximum y-axis value. If None, set to maximum activation.
scatter_mode (str) – Plot mode for traces, e.g. ‘lines’, ‘markers’, ‘lines+markers’.

Returns:

fig (go.Figure)

connectome_interpreter.activation_maximisation.saliency(model: MultilayeredNetwork, input_tensor: Tensor, neurons_of_interest: Dict[int, List[int]], method: str = 'vanilla', normalize: bool = False, device: device | None = None) → Tensor[source]

Computes saliency maps: given the current input, to what extent will each element’s change result in change in the activations of the neurons of interest?

Parameters:

model – A MultilayeredNetwork model
input_tensor – Input tensor to analyze
neurons_of_interest – Dictionary mapping layer indices to lists of neuron indices to analyze: {layer_idx: [neuron_indices]}. Layer_idx = 0 corresponds to the first layer after the input layer.
method – Saliency computation method (“vanilla” or “input_x_gradient”)
normalize – Whether to normalize saliency maps by their maximum value
device – Computation device

Returns:

Saliency maps with same shape as input tensor

Return type:

torch.Tensor

connectome_interpreter.activation_maximisation.train_model(model: MultilayeredNetwork, inputs: Tensor, targets: DataFrame, num_epochs: int = 100, learning_rate: float = 0.01, param_reg_lambda: float = 0.01, wandb: bool = False, wandb_project_name: str = 'connectome_interpreter', train_fraction: float = 0.8, seed: int = 42, train_slopes: bool = True, train_biases: bool = True, train_divisive_strength: bool = True, train_tau: bool = True, checkpoint_steps: int = 50, activation_loss_fn: str | Callable = 'mse')[source]

Train the model to approximate the targets, while keeping the model parameter change minumum (param_reg_lambda decdes how strongly). The loss is the mean squared error (default) between the model activations and the target activations.

Parameters:

model (MultilayeredNetwork) – The model to train.
inputs (torch.Tensor) – The input data. Shape: (batch_size, num_input_neurons, num_layers).
targets (pd.DataFrame) – The target activations. A DataFrame with columns: “batch”, “neuron_idx”, “layer” (optional), “value”. The “layer” column specifies the timestep (0-indexed). If ‘layer’ is not present, the average activation across all timesteps is used.
num_epochs (int, optional) – Number of training epochs. Defaults to 100.
learning_rate (float, optional) – Learning rate for the optimizer. Defaults to 0.01.
param_reg_lambda (float, optional) – Regularization parameter for the model parameters. Defaults to 0.01.
wandb (bool, optional) – Whether to use Weights & Biases logging.
wandb_project_name (str, optional) – Project name for wandb.
train_fraction (float, optional) – Fraction of data to use for training.
seed (int, optional) – Random seed for reproducibility.
train_slopes (bool, optional) – Whether to train slopes. Defaults to True.
train_biases (bool, optional) – Whether to train biases. Defaults to True. Note: biases use abs(raw_biases) internally to keep them positive. When enabled, any raw_biases sitting exactly at 0 are nudged to 1e-6 to break the abs() kink where gradients vanish.
train_divisive_strength (bool, optional) – Whether to train divisive strength. Defaults to True.
train_tau (bool, optional) – Whether to train tau. Defaults to True.
activation_loss_fn (str or callable, optional) – Loss function for activations. Either “mae”, “mse” (default), or a callable with signature fn(pred: torch.Tensor, target: torch.Tensor) -> torch.Tensor returning a scalar loss.

connectome_interpreter.activation_maximisation.training_mode(model: MultilayeredNetwork, train_slopes=True, train_biases=True, divisive_strength=True, train_tau=True)[source]: Context manager for training mode - enables gradients for slopes and biases.