---
license: mit
language:
- en
metrics:
- accuracy 0.9860
- roc_auc 0.9979
pipeline_tag: image-feature-extraction
tags:
- gnn
- link-prediction
- pytorch
- pytorch-geometric
- graph-neural-networks
- biology
- microscopy
- cell-tracking
- lineage-tracking

---
# DuMM Bacteria Tracker Model
## Overview

This repository contains the trained weights and documentation for the **DuMM Bacteria Tracker Model**, a **Graph Neural Network (GNN)** designed for **cell lineage link prediction** in time-lapse microscopy data.

The model uses a custom implementation of a **Parameter Decoupled Network (PDN)** variant within an **Edge-Propagation Message Passing Neural Network (EP-MPNN)** architecture, inspired by recent advancements in dynamic graph representation learning.

## Model Architecture and Implementation Details

* **Framework:** PyTorch and PyTorch Geometric (PyG).
* **Architecture:** `LineageLinkPredictionGNN` (custom `nn.Module`).
* **Core Layers:** Utilizes custom `EP_MPNN_Block` which incorporates **Distance & Similarity (DS)** features and a **Jumping Knowledge (JK)** network for aggregating features across multiple layers.
* **Task:** Binary classification (Link Prediction) on candidate edges between adjacent time frames (continuation or division links).
* **Loss Function:** `BCEWithLogitsLoss`.
* **Input Features:** Node features are scaled using a `StandardScalerTransform` and dynamically used to compute edge attributes (Absolute Difference + Cosine Similarity) via the `DS_block`.

### Input Features

The model uses a **10-dimensional feature vector** for each node (cell) derived from image analysis. These features capture morphological and intensity properties of the bacterial cells across two channels (Phase Contrast and Fluorescence).

| Feature Name | Type | Description |
| :--- | :--- | :--- |
| `area` | Morphological | Area of the cell segment. |
| `centroid_y` | Positional | Y-coordinate of the cell's centroid (critical for 1D growth systems). |
| `axis_major_length` | Morphological | Length of the cell's major axis. |
| `axis_minor_length` | Morphological | Length of the cell's minor axis. |
| `intensity_mean/max/min_phase` | Intensity | Mean, max, and min pixel intensity in the **Phase Contrast** channel. |
| `intensity_mean/max/min_fluor` | Intensity | Mean, max, and min pixel intensity in the **Fluorescence** channel. |


 The model was trained on microscopy images of the duplex mother machine developed by the Jun lab (https://jun.ucsd.edu/mother_machine.php)

### Data Preprocessing and Splitting

#### Splitting Strategy
To ensure the model generalizes to future, unseen data, a **time-based temporal split** was employed:
1.  **Training Set:** First 60% of unique time frames (`sorted_time_frames[:train_split_idx]`).
2.  **Validation Set:** Next 20% of unique time frames (`sorted_time_frames[train_split_idx:val_split_idx]`).
3.  **Test Set:** Final 20% of unique time frames (`sorted_time_frames[val_split_idx:]`).

#### Normalization
* **Method:** Node features were normalized using **Standard Scaling (`sklearn.preprocessing.StandardScaler`)**.
* **Fit:** The scaler was **fitted *only* on the training set features** (`all_train_node_features_df`).
* **Application:** The fitted scaler was then applied to transform the features in the Training, Validation, and Test sets via the `StandardScalerTransform`. This avoids data leakage.

#### Graph Creation (Candidate Generation)
Candidate edges (links between cells in adjacent time frames) were generated based on a custom set of geometric and morphological heuristics:
* **Distance Constraint:** Max distance between centroids is limited by `max_dist_link` (default 50.0).
* **Area Ratio Constraints:**
    * **Continuation (1-to-1):** `min_area_ratio_continuation` (0.8) to `max_area_ratio_continuation` (1.2).
    * **Division (1-to-2):** `min_area_ratio_division` (1.8) to `max_area_ratio_division` (2.2).

### Training Protocol

| Hyperparameter | Value | Description |
| :--- | :--- | :--- |
| **GNN Layers (`num_blocks`)** | 2 | Number of sequential EP-MPNN blocks. |
| **Hidden Channels** | 128 | Dimension for node and edge embeddings. |
| **Optimizer** | Adam | Standard optimization algorithm. |
| **Learning Rate** | 0.001 | Base learning rate. |
| **Weight Decay** | 0.0005 | L2 regularization applied to prevent overfitting. |
| **Batch Size** | 32 | Number of graphs processed per iteration. |
| **Evaluation Metric** | Validation Accuracy (`val_acc`) | Used for saving the `best_link_prediction_model.pt`. |
| **Early Stopping** | Yes | Monitors Validation Loss (`val_loss`) with a **patience of 10 epochs**. |
| **Max Epochs** | 500 | Maximum number of training epochs. |