Metadata-Version: 2.4
Name: cell_attn
Version: 0.1.0
Summary: A package for using deep attention networks to analyse cell movement data.
Project-URL: Homepage, https://github.com/Jamie-hb/cell_attn
Author: James Boyle
License-Expression: MIT
License-File: LICENCE
Classifier: Operating System :: OS Independent
Classifier: Programming Language :: Python :: 3
Requires-Python: >=3.9
Requires-Dist: matplotlib>=3.9.2
Requires-Dist: numpy>=1.26.2
Requires-Dist: pandas>=2.2.2
Requires-Dist: seaborn>=0.13.2
Requires-Dist: torch>=2.5.0
Description-Content-Type: text/markdown

<img src="images/logo.png" alt="markdown language" width="80" height="80">

# cell_attn

Deep attention networks are an exploratory tool for inferring how cell-cell interactions affect cell movement from cell trajectory data. cell_attn provides the necessary functions to easily train a deep attention network on a cell trajectory data set, and use the trained DAN to infer how cell-cell interactions affect cell movement in the data set.

Here we provide a brief tutorial on how to do this, using an example of publically available cell trajectory data from an assay of collectively migrating MDCK (kindey epithelium) cells. The model and inference tools are described in more detail in [1], and the data in [2].

## Installation

This package is installable via `pip` usinng the command `pip install cell_attn`. cell_attn has pytorch as a dependency, and you may wish to install this separately, see the [pytorch installation page](https://pytorch.org/get-started/locally/) for details.

### Dependencies

cell_attn has the following dependencies. cell_attn only uses standard aspects of these packages, so we expect that it should be compatible with similar versions.

- `numpy >= 1.26.4`
- `matplotlib >= 3.9.2`
- `seaborn >= 0.13.2`
- `pytorch >= 2.5.0`
- `pandas >= 2.2.2`

## What data do I need?

Traning a deep attention network requires co-ordinate data for each cell at evenly spaced time points, in the format of a $2 \times N \times M$ numpy array $X$, where $N$ is the number of cells, $M$ is the number of timepoints, and $X[0,i,t]$, $X[1,i,t]$ are the $x$ and $y$ co-ordinates of cell $i$ at time $t$ respectively.

$$ X[0,:,:] = 
   \begin{pmatrix}
   x_1(t_1) & x_1(t_2) & \cdots & x_1(t_M) \\
   x_2(t_1) & x_2(t_2) & \cdots & x_2(t_M) \\
   \vdots   & \vdots   & \ddots & \vdots \\
   x_N(t_1) & x_N(t_2) & \cdots & x_N(t_M) \\
   \end{pmatrix} 
   \hspace{20mm}
   X[1,:,:] = 
   \begin{pmatrix}
   y_1(t_1) & y_1(t_2) & \cdots & y_1(t_M) \\
   y_2(t_1) & y_2(t_2) & \cdots & y_2(t_M) \\
   \vdots   & \vdots   & \ddots & \vdots \\
   y_N(t_1) & y_N(t_2) & \cdots & y_N(t_M) \\
   \end{pmatrix}
   $$

Where $(x_i(t), y_i(t))$ are the co-ordinates of cell $i$ at time $t$. 

$X$ can contain missing values. If analysing data from multiple experiments, you will need one trajectory data matrix for each experiment.

Given trajectory data in this format, the procedure for training and using a deep attention network is as follows:
1) Use `create_dataset_from_trajectories()` to create the input and label data required to train a deep attention network.
2) Use `train_deep_attention_network()` to train a deep attention network.
3) Use various inference tools to infer from the trained deep attention network how cell-cell interactions affect cell movement.

In this tutorial we shall use cell trajectory data from an assay of collectively migrating kidney epithelial cells, namely Madin-Darby Canine Kidney (MDCK) cells. This assay is analysed using our deep attention network model in [1], and was originally published and analysed using an earlier deep attention network model in [2]. The data contains fifteen replicates of the same experiment, and in this tutorial, we shall suppose that all the trajectory matrices lie in one directory `mdck_data`, with `mdck_data/trajectories_i.npy` being the trajectory matrix $X$ for experiment $i$.

## 1) Input and label data creation

The first step is to create the input and label data required to train a deep attention network.

### Folder structure

For training, the input and label data must be in a single directory (let's call this `trajectory_data`), within which there must be two folders, `training_data` and `validation_data`. Within each of these there must be a folder for each experiment, containing the input and label data generated by the trajectory matrix from that experiment.

**Note**: There must be at least one experiment in each of the training and validation data directories, if possible we recommend a 70/30 split. If you only have data from one experiment, then we recommend splitting that data in two.

The MDCK data contains 15 experiments, which we shall call `experiment_1`, `experiment_2` etc', so we put 11 of these in the training data directory, and 4 in the validation data directory. The file structure we are aiming for is:

- `trajectory_data`
  - `training_data`
    - `experiment_1`
      - `input_data.npy`
      - `labels.npy`
    - `experiment_2`
      - `input_data.npy`
      - `labels.npy`
    - ...
    - `experiment_11`
      - `input_data.npy`
      - `labels.npy`
  - `validation_data`
    - `experiment_12`
      - `input_data.npy`
      - `labels.npy`
    - ...
    - `experiment_15`
      - `input_data.npy`
      - `labels.npy`

where `input_data.npy` and `labels.npy` are the input and label data, which we shall now explain how to create. For now, ensure you have the correct folder structure with empty folders.

**Note**: The overall directory, and folders for each experiment, can be called whatever you want, but the folders containing the training and validation data *must* be called `training_data` and `validation_data`.

### Input and Label Creation

Input and label data can be created from a trajectory matrix using the `create_data_from_trajectories()`. This function requires the following inputs:

- `X`: The trajectory matrix $X$
- `nn`: The number of nearest neighbours to use to predict cell movement
- `timestep`: Over what length of time to predict cell movement (measured in indices of $X$)
- `stdir`: A path specifying where to store the input and label data. 

This function will then create the input and label data, and deposit them in `stdir` under the names `input_data.npy` and `labels.npy`. It will also create a file called `input_data_parameters.csv` recording the different parameters used to generate the data from the trajectory matrix.

We recommend picking `nn` large enough such that a cell's $\textrm{nn}^{th}$ nearest neighbour is far enough away that you would expect them not to interact significantly, and `timestep` to be small (we usually take `timestep` = 1). 

For the MDCK data, we use 10 nearest neighbours, with a prediction interval of 20 minutes (following [2]). The time resolution of the data is 10 minutes, so this corresponds to `timestep=2`. We can process the data using the code


```python
import numpy as np
import cell_attn 

for i in range(1,12):
    X = np.load('mdck_data/trajectories_i.npy')
    cell_attn.create_data_from_trajectories(X=X, nn=10, timestep=2, stdir='trajectory_data/training_data/experiment_i')            

for i in range(12,16):
    X = np.load('mdck_data/trajectories_i.npy')
    cell_attn.create_data_from_trajectories(X=X, nn=10, timestep=2, stdir='trajectory_data/validation_data/experiment_i')
```

This processes the data then for experiments 1 - 11 and deposits the input and label data in the training data directory, and for experiments 12 - 15 deposits the input and label data in the validation data directory.

**Note**: If the trajectory matrix is very large but also contains lots of missing data (as may happen if cells often die or are lost by the tracking algorithm), we recommend passing `memory_saver=True`. This creates the input and label data in a way that uses as little memory as possible, but that may be slower for trajectory matrices with minimal missing data.

**Note**: To use only trajectory data from a limited time range, you can pass the start and end indices of this range to `t_min` and `t_max` respectively.

## 2) Model Training

Given input and label data in the correct directory structure, we can train a deep attention network using the `train_deep_attention_network()` function, which requires only two arguments:

- `datadir`: A path to the top level directory containing the input and label data.
- `stdir`: A path to the directory to store the trained model in.

**Note**: Further arguments can be passed to control the details of the training process, see the documentation. In particular if you are getting results which look like the model may have failed to train properly, you can pass `training_speed='long'` to have the training process run for more iterations.

For us `datadir='trajectory_data'`, and let's store our trained model in a directory called `mdck_dan` (if this directory doesn't already exist, `train_deep_attention_network()` will attempt to create it).

We thus run the code


```python
cell_attn.train_deep_attention_network(datadir='trajectory_data', stdir='mdck_dan')
```

**Note**: Under the hood, `train_deep_attention_network()` is training two feed forward neural networks. This can be a computationally expensive process, and we strongly recommend running this function on a GPU if possible (cell_attn will automatically search for a GPU, using the `torch.cuda.is_available()` function in `pytorch`, and will only run on CPU is this returns `False`, see the [pytorch documentation](https://docs.pytorch.org/docs/2.12/generated/torch.cuda.is_available.html) for details).

This creates the following files in `stdir`:

- `model.pth`: The trained deep attention network.
- `model_training_parameters.csv`: A log of the parameters used to train the network.
- `training_log.txt`: A text file tracking the model training, useful for debugging.
- `losses.csv`: Contains the train and validation loss after each batch.
- `metrics_per_epoch.csv`: Contains the loss after every epoch.

## 3) Inference

Now for the exciting bit! We can use the trained deep attention network to infer how cell-cell interactions affect cell movement in the MDCK cell data.

First let's load in the trained deep attention network using `load_model()`:


```python
model = cell_attn.load_model(stdir='mdck_dan')
```

### Attention heatmaps

We can produce an Attention Heatmap, showing how cell-cell interaction strength is affected by the displacement of the interacting cells, using `plot_attention_heatmap()`:


```python
plt.figure()
cell_attn.plot_attention_heatmap(model=model, datadir='trajectory_data/training_data')
plt.show()
```

<img src="images/MDCK_AHM.png" alt="markdown language" width="500" height="500">

We see that MDCK cells take movement cues predominantly from their frontward neighbours.

**Note**: Here we provide `plot_attention_heatmap()` with the path to the directory containing the input and label data. `plot_attention_heatmap()` then builds a pytorch dataloader which the deep attention network can read data from. We could instead create the dataloader separately:


```python
dataloader = cell_attn.load_data(datadir='trajectory_data/training_data')
```

and pass this directly to `plot_attention_heatmap()`:


```python
cell_attn.plot_attention_heatmap(model=model, dataloader=dataloader)
```

This option, to pass a dataloader rather than have the function construct the dataloader, is available for all inference functions, and can save time if you are using multiple of them in a single session, as the dataloader only needs to be constructed once.

### Attraction-Repulsion and Alignment heatmaps

The syntax for producing attraction-repulsion and alignment heatmaps is very similar.


```python
plt.figure()
cell_attn.plot_attraction_repulsion_heatmap(model=model, datadir='trajectory_data/training_data')
plt.show()
```

<img src="images/MDCK_ARHM.png" alt="markdown language" width="500" height="500">


```python
plt.figure()
cell_attn.plot_alignment_heatmap(model=model, datadir='trajectory_data/training_data')
plt.show()
```

<img src="images/MDCK_ALHM.png" alt="markdown language" width="500" height="500">

## 4) References

[1]: BioRxiv pre-print

[2]: J. LaChance, K. Suh, J. Clausen, D. J. Cohen. Learning the rules of collective cell migration using deep attention networks. PLoS computational biology (2022).
