Activation maximisation

Bases: Module

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.

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.
device (torch.device, optional) – Device for computation.

activation_function(x: Tensor, x_previous=typing.Optional[torch.Tensor]) → 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.
Returns:: The output tensor after applying the activation function.
Return type:: torch.Tensor

property biases

Get the biases of the neurons, applying absolute value if raw_biases are provided. If raw_biases is None, returns None.

Returns:: The biases of the neurons, or None if raw_biases is not set.
Return type:: torch.Tensor or None

divnorm(slopes: float | Tensor, x_previous: Tensor) → Tensor[source]

Change the slopes based on divisive normalization. The new slope bewteen pre and post is: original_slope * (1 - pre_post_weight * pre_activation * divisive_strength).

Parameters:

slopes (Union[float, torch.Tensor]) – The slopes to be normalized. If a float, it is assumed to be the same for all neurons.
x_previous (torch.Tensor) – The previous activations of the neurons.

Returns:

The output tensor after applying divisive normalization.

Return type:

torch.Tensor

forward(inputs: Tensor)[source]

Process inputs through the network. Use matrix-matrix multiplication for batched processing.

Parameters:: inputs (torch.Tensor) – Either a 2D tensor (num_input, time_steps) or a 3D tensor (batch_size, num_input, time_steps)
Returns:: The activations of all neurons across time steps (in batches).
Return type:: torch.Tensor

set_param_grads(slopes=False, raw_biases=False, divisive_strength=False)[source]: Set requires_grad for specific parameters.

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) → 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-6.
n_runs (int, optional) – The number of runs to consider for early stopping. Defaults to 10.
use_tqdm (bool, optional) – Whether to use tqdm progress bars to track optimization progress. Defaults to True.
print_output (bool, optional) – Whether to print loss information during optimization. Defaults to True.
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.

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. 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

Get selected input activation.

Parameters:

model_in (numpy.ndarray | torch.Tensor) – The input to the model. Shape should be (num_neurons, num_timepoints).
sensory_indices (arrayable) – All the indices of sensory neurons. Length should be equal to the first dimension of model_in.
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 activation. Neurons with activations below this threshold are not considered. Defaults to 0.

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, with the first: columns being group and sensory_indices. The rest are the timesteps.

Return type:

pd.DataFrame

Get the activations for specified indices across timepoints, include batch name and group information when available.

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

connectome_interpreter.activation_maximisation.input_from_df(df: DataFrame, sensory_indices: list, idx_to_group: dict, num_layers: int, timepoints: int | List[int] | ndarray | set = 0) → 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.

Returns:

The input for the model.

Return type:

np.ndarray

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]}
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)[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 average absolute difference 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.

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