DuMM_bacteria_track

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

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