Convolutional neural networks for characterisation and segmentation of anomalous diffusion trajectories
This repository contains code for the analysis of single particle tracking data for characterisation of anomalous diffusion, using convolutional neural networks (CNNs). Submitted to the Anomalous Diffusion (AnDi) Challenge.
The models in this repository perform three tasks on 1D and 2D tracks, as set out in the challenge:
- Task 1: inference of anomalous exponent
- Task 2: classification of diffusion model
- Task 3: segmentation of trajectories
The CNN architecture used here is adapted from that used in Granik et al., 2019, which is itself based on a CNN architecture described in Bai et al., 2018. As well as some minor changes to hyperparameters and model training (as outlined in the code), I've made some major modifications to the original methods which improve performance and applicability to real data:
- Consideration of tracks of all lengths (5-1000+ steps) with a single model
- Training models on tracks with variable signal to noise ratio (permitting inference without prior information about SNR)
- Segmentation of trajectories switching between diffusion modes (currently limited to single length tracks with one switch). NB. this uses a slightly different CNN architecture (regular convolutions instead of causal convolutions)
Please see the AnDi Challenge follow-up paper (available here) for a full description of the method and a thorough assessment of its performance (look out for team J!).
Note: the code in this repository has been modified and improved somewhat since the original challenge (mostly just to improve usability). For the exact code and models used in the original challenge, please see this fork.
Clone the repository:
git clone https://github.com/tsmbland/andi_challenge.git
cd andi_challenge
Create conda environment:
conda env create -f environment.yml
Activate conda environment:
conda activate andi_challenge
Validation and test data can be generated by running Test/Generate.py and Validation/Generate.py respectively. By default, this will generate 10,000 tracks per task per dimension (1D and 2D) for test and validation, using the andi-datasets package.
Training data is generated on-the-fly during training, and does not need to be generated beforehand.
To train models, run the scripts in the Train folder for the respective task (E.g. Task1_Exponent/Train/1D.py for task 1 in 1D). Pre-trained models are included as .h5 files in the Models folder for each task.
The repository includes three scripts for analysing challenge data in the ‘Challenge’ folder, named Task1.py, Task2.py and Task3.py for the three tasks. To perform analysis, run these scripts, specifying the path to the data folder at the top of the file.
Non-challenge data (e.g. real data) can be analysed in a similar way.
There are four notebooks in the Notebooks folder that lay out some of the features of this repository and provide further instructions on performing analysis:
- Notebook 1
- Generating and preprocessing trajectories
- Notebook 2
- Notebook 3
- Notebook 4
Bai, S., Kolter, J.Z., and Koltun, V. (2018). An empirical evaluation of generic convolutional and recurrent networks for sequence modeling. ArXiv.
Granik, N., Weiss, L.E., Nehme, E., Levin, M., Chein, M., Perlson, E., Roichman, Y., and Shechtman, Y. (2019). Single-Particle Diffusion Characterization by Deep Learning. Biophys. J. 117, 185–192
Muñoz-Gil, G., Volpe, G., García-March, M.A., Metzler, R., Lewenstein, M., and Manzo, C. (2020). The anomalous diffusion challenge: single trajectory characterisation as a competition. ArXiv
Muñoz-Gil, G., Volpe, G., Garcia-March, M.A. et al. Objective comparison of methods to decode anomalous diffusion. Nat Commun 12, 6253 (2021).
This work is licensed under a Creative Commons Attribution 4.0 International License.
