Skip to content

Commit

Permalink
added mask module and deprecation warnings
Browse files Browse the repository at this point in the history
  • Loading branch information
@zhaozewang
zhaozewang committed Jan 8, 2024
1 parent a2069e6 commit 9be99f9
Show file tree
Hide file tree
Showing 23 changed files with 902 additions and 41 deletions.
20 changes: 14 additions & 6 deletions docs/change_logs/v1.1.0.md
View file
Original file line number Diff line number Diff line change
@@ -1,4 +1,5 @@
## Verson 1.1.0 Change Logs
### Changes for `CTRNN`
1. Merged `input_dim`, `hidden_size`, and `output_dim` into `dims`.
2. Removed `scaling`. This is now split into `weights` by manually scaling the weight matrices.
3. Removed `layer_masks`. This is now split into `plasticity_masks`, `ei_masks`, and `sparsity_masks` to increase flexibility and reduce confusion.
Expand All @@ -9,6 +10,13 @@
8. Rename and upgraded `learnable` to `plasticity_masks`. `plasticity_masks` inherits everything from `learnable`, but can also be masks that can be used to control the plasticity of every single weight in the network.
9. Add `auto_rescale()` to `CTRNN`. If you set a sparsity mask, calling this function will help you automatically rescale the weights to make sure the sparse matrix can still drive the network dynamics. (Consider if you make sparsity = 0.01, the output is unlikely to be able to drive the next layer).

### Changes for `structure`
1. The structure module won't be maintained in future versions. Everything is now inherited by `mask`.

### Changes for `mask`
1. Inherited everything from `structure`.
2. The `mask` module now recieve a `dims` parameter instead of `input_dim`, `hidden_size`, and `output_dim`.

## TODOs
- [x] Put initialization to structures, i.e., the distribution of weights, biases, etc.
- [x] Add custom controls to the update speeds of different parts of the network.
Expand All @@ -25,12 +33,12 @@
- [x] Remove `layer_masks`. This is now split into `plasticity_masks`, `ei_masks`, and `sparsity_masks` to increase flexibility and reduce confusion.
- [x] Remove `scaling`. This is now split into `weights` by manually scaling the weight matrices.
- [x] Merge `input_dim`, `hidden_size`, and `output_dim` into `dims`.
- [ ] Check the checker function dimension check feature.
- [ ] Add deprecation auto-fixing.
- [ ] Test `init_state`.
- [ ] Rewrite the dimension check for masks. Currently the masks are not intuitive. One way is to transpose everything in `mask` module, and transpose them back when received by `CTRNN`. Don't add auto-compatible feature for now as it might mess up the hidden layer since its a sqare matrix.
- [ ] Add deprecation autofix.
- [ ] Test `init_state` (how hidden states are initialized).
- [x] Add `auto_rescale()` to `CTRNN`.
- [ ] Check _balance_excitatory_inhibitory().
- [ ] Put `self_connections` into `sparsity_masks`.
- [ ] Rename `structure` module to `masks`.
- [x] Add function `adjust_gradients()` to accommodate the new `plasticity_masks`.
- [ ] Add warning for unnecessary call of `adjust_gradients()` when all plasticity masks are 1.
- [x] Rename `structure` module to `mask`.
- [x] Add function `apply_plasticity()` to accommodate the new `plasticity_masks`.
- [ ] ~~Add warning for unnecessary call of `apply_plasticity()` when all plasticity masks are 1.~~ (no need as it might be good to check gradients anyway)
107 changes: 107 additions & 0 deletions docs/mask/index.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,107 @@
# RNN mask(s) (`nn4n.mask`)
[Back to Home](https://github.com/zhaozewang/NN4Neurosci/blob/main/README.md)
## Table of Contents
- [Introduction](#introduction)
- [Masks](#masks)
- [BaseMask](#basemask)
- [Parameters](#basemask-parameters)
- [Methods](#basemask-methods)
- [Multi Area](#multiarea)
- [Parameters](#multiarea-parameters)
- [Forward Backward Specifications](#forward-backward-specifications)
- [Multi Area EI](#multiareaei)
- [Parameters](#multiareaei-parameters)
- [Inter-Area Connections Under EI Constraints](#inter-area-connections-under-ei-constraints)
- [Random Input](#randominput)
- [Parameters](#randominput-parameters)

## Introduction
This module defines masks for any RNN in the standard 3-layer architectures (as shown below). The masks of the hidden layer in this project are defined using masks. Therefore, classes in this module will generate input_mask, hidden_mask, and readout_mask that are used in the `model` module<br>

<p align="center"><img src="https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/images/basics/RNN_structure.png" width="400"></p>

Where yellow nodes are in the InputLayer, green nodes are in the HiddenLayer, and purple nodes are in the ReadoutLayer.

## Masks
### BaseMask
Base class for all masks. It defines the basic mask of a RNN. It serves as a boilerplate for other masks. It is not meant to be used directly.

#### BaseMask Parameters
| Parameter | Default | Required | Type | Description |
|:---------------------|:-------------:|:-------------:|:----------------:|:-------------------------------------------|
| dims | [1, 100, 1] | True |`list` | Dimensions of the RNN. It should be a list of three elements, where the first element is the input dimension, the second element is the hidden size, and the third element is the output dimension. |

#### BaseMask Methods
Methods that are shared by all masks. <br>
| Method | Description |
|:-----------------------------------------------------|:----------------------------------------------------|
| [`get_input_idx()`](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_input_idx.md) | Get indices of neurons that receive input. |
| [`get_readout_idx()`](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_readout_idx.md) | Get indices of neurons that readout from. |
| [`get_non_input_idx()`](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_non_input_idx.md) | Get indices of neurons that don't receive input. |
| [`visualize()`](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/visualize.md) | Visualize the generated masks. |
| [`masks()`](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/masks.md) | Return a list of np.ndarray masks. It will be of length 3, where the first element is the input mask, the second element is the hidden mask, and the third element is the readout mask. For those masks that do not have specification for a certain mask, it will be an all-one matrix. |
| [`get_areas()`](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_areas.md) | Get a list of areas names. |
| [`get_area_idx()`](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_area_idx.md) | Get indices of neurons in a specific area. The parameter `area` could be either a string from the `get_areas()` or a index of the area. |


### MultiArea
See [Examples](https://github.com/zhaozewang/NN4Neurosci/blob/main/examples/MultiArea.ipynb) <br>
This will generate a multi-area RNN without E/I constraints. Therefore, by default, the input/hidden/readout masks are binary masks. Use cautious when the `positivity_constraints` parameter of CTRNN is set to `True`, because it will make all neurons to be excitatory.
**NOTE:** This also implicitly covers single area case. If `n_area` is set to 1. All other parameters that conflict this setting will be ignored.
#### MultiArea Parameters
| Parameter | Default | Required | Type | Description |
|:---------------------|:-------------:|:-------------:|:-------------------------:|:-------------------------------------------|
| n_areas | 2 | False |`int` or `list` | Number of areas.<br>- If `n_areas` is an integer, `n_areas` must be a divisor of `hidden_size`. It will divide the HiddenLayer into three equal size regions.<br>- If `n_areas` is a list, it must sums up to `hidden_size`, where each element in the list denote the number of neurons in that area. |
| area_connectivities | [0.1, 0.1] | False |`list` or `np.ndarray` | Area-to-area connection connectivity. Entries must between `[0,1]`<br>- If its a list of two elements, the first element is the forward connectivity, and the second is the backward connectivity. The within-area connectivity will be 1.<br>- If its a list of three elements, the last element will be the within-area connectivity.<br>- If `area_connectivities` is an `np.ndarray`, it must be of shape (`n_areas`, `n_areas`). See [forward/backward specifications](#forward-backward-specifications)|
| input_areas | `None` | False |`list` or `None` | Areas that receive input. If set to `None`, all neurons will receive inputs. If set to a `list`, list elements should be the index of the areas that receive input. Set it to a list of one element if only one area receives input. |
| readout_areas | `None` | False |`list` or `None` | Areas that readout from. If set to `None`, all neurons will readout from. If set to a `list`, list elements should be the index of the areas that readout from. Set it to a list of one element if only one area readout from. |


| Attributes | Type | Description |
|:-------------------------|:--------------------------:|:-------------------------------------------|
| n_areas | `int` | Number of areas |
| node_assignment | `list` | Nodes area assignment |
| hidden_size | `int` | Number of nodes in the HiddenLayer |
| input_dim | `int` | Input dimension |
| output_dim | `int` | Output dimension |
| area_connectivities | `np.ndarray` | Area-to-area connectivity matrix. If it is a list in params, it will be transformed into a numpy matrix after initialization |

#### Forward Backward Specifications
RNNs can be implemented in various ways, in this library,
$$s W^T + b$$
is used in the HiddenLayer forward pass, where $W$ is the connectivity matrix of the HiddenLayer and $s$ is the current HiddenLayer state.<br>
$W$ may not matter if your connectivity matrix is symmetric. But if it's not, you might want to pay attention to the forward connections and backward connections. In the figure below, three networks (`n_areas` = 2, 3, 4) and their corresponding forward/backward connection matrix are provided. The blue regions are intra-area connectivity, the green regions are forward connections, and the red regions are backward connections.

<p align="center"><img src="https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/images/basics/Multi_Area.png" width="700"></p>

<!-- ![area_connectivities](../img/Multi_Area_Transpose.png) -->


### MultiAreaEI
[Examples](https://github.com/zhaozewang/NN4Neurosci/blob/main/examples/MultiArea.ipynb) <br>
This class is a child class of `MultiArea`. It will generate a multi-area RNN with E/I constraints. Therefore, by default, the input/hidden/readout masks are signed masks. Use cautious as it will change the sign of the weights.
#### MultiAreaEI Parameters
| Parameter | Default | Type | Description |
|:------------------------------|:-----------------------:|:--------------------------:|:-------------------------------------------|
| ext_pct | 0.8 | `float` | Percentage of excitatory neurons |
| inter_area_connections |[True, True, True, True] | `list` (of booleans) | Allows for what type of inter-area connections. `inter_area_connections` must be a `boolean` list of 4 elements, denoting whether 'exc-exc', 'exc-inh', 'inh-exc', and 'inh-inh' connections are allowed between areas. see [inter-area connections under EI constraints](#inter-area-connections-under-ei-constraints). |
| inh_readout | True | `boolean` | Whether to readout inhibitory neurons |


#### Inter-Area Connections Under EI Constraints
Depending on the specific problem you are investigating on, it is possible that you want to eliminate inhibitory connections between areas. Or, you might not want excitatory neurons to connect to inhibitory neurons in other areas. See figure below for different cases of inter-area connections under EI constraints.

<p align="center"><img src="https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/images/basics/Multi_Area_EI.png" width="550"></p>

To specify what kinds of inter-area connections you want to keep, simple pass a 4-element boolean list to `inter_area_connections`. The 4 elements denotes whether to keep inter-area 'exc-exc', 'exc-inh', 'inh-exc', and 'inh-inh' connections.

### RandomInput
Randomly inject input to the network. Neurons' dynamic receiving input will be heavily driven by the inputting signal. Injecting signal to only part of the neuron will result in more versatile and hierarchical dynamics. See [A Versatile Hub Model For Efficient Information Propagation And Feature Selection](https://arxiv.org/abs/2307.02398) <br>

#### RandomInput Parameters
| Parameter | Default | Type | Description |
|:------------------------------|:-----------------------:|:--------------------------:|:-------------------------------------------|
| input_spar | 1 | `float` | Input sparsity. Percentage of neurons that receive input. |
| readout_spars | 1 | `float` | Readout sparsity. Percentage of neurons that readout from. |
| hidden_spar | 1 | `float` | Hidden sparsity. Percentage of edges that are non-zero. |
| overlap | `True` | `boolean` | Whether to allow overlap between input and readout neurons. |
38 changes: 38 additions & 0 deletions docs/mask/methods/get_area_idx.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,38 @@
# nn4n.mask.get_area_idx()

[Back to mask](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/index.md) </br>

## Introduction
Get the node indices of a specific area.

## Parameters
> `area` (str) or (int): required.
>> The name of the area or the index of the area.
## Returns
> `indices` (np.ndarray):
>> The indices of the nodes in the specified area.
## Usage
```python
import torch
from nn4n.mask import MultiArea

params = {
"hidden_size": 100,
"input_dim": 2,
"output_dim": 2,
}
multiarea = MultiArea(**params)
multiarea.get_areas() # ['area_1', 'area_2']
multiarea.get_area_idx('area_1')
```

Output
```
array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33,
34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49])
```

see also: [get_areas()]('https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_areas.md')
31 changes: 31 additions & 0 deletions docs/mask/methods/get_areas.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,31 @@
# nn4n.mask.get_areas()

[Back to mask](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/index.md) </br>

## Introduction
Retrieve the area names

## Returns
> `area_names` (list):
>> A list of area names.
## Usage
```python
import torch
from nn4n.mask import MultiArea

params = {
"hidden_size": 100,
"input_dim": 2,
"output_dim": 2,
}
multiarea = MultiArea(**params)
multiarea.get_areas()
```

Output
```
['area_1', 'area_2']
```

see also: [get_area_idx()]('https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_area_idx.md')
36 changes: 36 additions & 0 deletions docs/mask/methods/get_input_idx.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
# nn4n.mask.get_non_input_idx()

[Back to mask](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/index.md) </br>

## Introduction
Get the indices of the neurons that receive inputs.

## Returns
> `input_idx` (np.ndarray):
>> The indices of the input neurons. Will be a 1D array of shape `(input_dim,)`.
## Usage
```python
import torch
from nn4n.mask import MultiArea

params = {
"hidden_size": 100,
"input_dim": 2,
"output_dim": 2,
}
multiarea = MultiArea(**params)
multiarea.get_input_idx()
```

Output
```
array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33,
34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50,
51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67,
68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84,
85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99])
```

see also: [get_readout_idx()](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_readout_idx.md) | [get_non_input_idx()](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_non_input_idx.md)
40 changes: 40 additions & 0 deletions docs/mask/methods/get_non_input_idx.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,40 @@
# nn4n.mask.get_non_input_idx()

[Back to mask](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/index.md) </br>

## Introduction
Get the indices of the neurons that receive inputs.

## Returns
> `non_input_idx` (np.ndarray):
>> The indices of neurons that don't receive inputs.
## Usage
The example below will create a `MultiArea` object with 3 areas, where the first area is the input area. The `get_non_input_idx()` method will return indices of neurons in area 2 and 3, as they don't receive inputs.

```python
import torch
from nn4n.mask import MultiArea

params = {
"n_areas": 3,
"area_connectivities": [0.1, 0.1],
"hidden_size": 90,
"input_dim": 1,
"output_dim": 1,
"input_areas": [0],
"readout_areas": [2],
}
multiarea = MultiArea(**params)
multiarea.get_non_input_idx()
```

Output
```
array([30, 31, 32, 33, 34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46,
47, 48, 49, 50, 51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63,
64, 65, 66, 67, 68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80,
81, 82, 83, 84, 85, 86, 87, 88, 89])
```

see also: [get_readout_idx()](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_readout_idx.md) | [get_input_idx()](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_input_idx.md)
36 changes: 36 additions & 0 deletions docs/mask/methods/get_readout_idx.md
View file
Original file line number Diff line number Diff line change
@@ -0,0 +1,36 @@
# nn4n.mask.get_readout_idx()

[Back to mask](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/index.md) </br>

## Introduction
Get the indices of the neurons that are being readout.

## Returns
> `readout_idx` (np.ndarray):
>> The indices of the neurons that are being readout.
## Usage
```python
import torch
from nn4n.mask import MultiArea

params = {
"hidden_size": 100,
"input_dim": 2,
"output_dim": 2,
}
multiarea = MultiArea(**params)
multiarea.get_readout_idx()
```

Output
```
array([ 0, 1, 2, 3, 4, 5, 6, 7, 8, 9, 10, 11, 12, 13, 14, 15, 16,
17, 18, 19, 20, 21, 22, 23, 24, 25, 26, 27, 28, 29, 30, 31, 32, 33,
34, 35, 36, 37, 38, 39, 40, 41, 42, 43, 44, 45, 46, 47, 48, 49, 50,
51, 52, 53, 54, 55, 56, 57, 58, 59, 60, 61, 62, 63, 64, 65, 66, 67,
68, 69, 70, 71, 72, 73, 74, 75, 76, 77, 78, 79, 80, 81, 82, 83, 84,
85, 86, 87, 88, 89, 90, 91, 92, 93, 94, 95, 96, 97, 98, 99])
```

see also: [get_input_idx()](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_input_idx.md) | [get_non_input_idx()](https://github.com/zhaozewang/NN4Neurosci/blob/main/docs/mask/methods/get_non_input_idx.md)
Loading

0 comments on commit 9be99f9

Please sign in to comment.