Single-cell representation learning

README.md

@@ -102,6 +102,36 @@ The robust virtual staining models (i.e *VSCyto2D*, *VSCyto3D*, *VSNeuromast*),
 
 A full illustration of the virtual staining pipeline can be found [here](https://github.com/mehta-lab/VisCy/blob/dde3e27482e58a30f7c202e56d89378031180c75/docs/virtual_staining.md).
 
+## Image representation learning
+
+We are currently developing self-supervised representation learning to map cell state dynamics in response to perturbations,
+with focus on cell and organelle remodeling due to viral infection.
+
+See our recent work on temporally regularized contrastive sampling method
+for representation learning on [arXiv](https://arxiv.org/abs/2410.11281).
+
+<details>
+ <summary> Pradeep, Imran, Liu et al., 2024 </summary>
+
+  <pre><code>
+@misc{pradeep_contrastive_2024,
+      title={Contrastive learning of cell state dynamics in response to perturbations},
+      author={Soorya Pradeep and Alishba Imran and Ziwen Liu and Taylla Milena Theodoro and Eduardo Hirata-Miyasaki and Ivan Ivanov and Madhura Bhave and Sudip Khadka and Hunter Woosley and Carolina Arias and Shalin B. Mehta},
+      year={2024},
+      eprint={2410.11281},
+      archivePrefix={arXiv},
+      primaryClass={cs.CV},
+      url={https://arxiv.org/abs/2410.11281},
+}
+    </code></pre>
+  </details>
+
+### Workflow demo
+
+[Exploration of learned embeddings with napari-iohub](https://drive.google.com/file/d/16WSoTvXJ-siLb7iyOueOag_cKn9Iwckc/view?usp=drive_link)
+
+![DynaCLR](https://github.com/mehta-lab/VisCy/blob/9eaab7eca50d684d8a473ad9da089aeab0e8f6a0/docs/figures/dynaCLR_schematic.png?raw=true)
+
 ## Installation
 
 1. We recommend using a new Conda/virtual environment.
@@ -148,4 +178,3 @@ for reading and writing data in [OME-Zarr](https://www.nature.com/articles/s4159
 The full functionality is tested on Linux `x86_64` with NVIDIA Ampere GPUs (CUDA 12.4).
 Some features (e.g. mixed precision and distributed training) may not be available with other setups,
 see [PyTorch documentation](https://pytorch.org) for details.
-

applications/contrastive_phenotyping/contrastive_cli/fit_ctc_mps.yml

@@ -0,0 +1,117 @@
+# See help here on how to configure hyper-parameters with config files:
+# https://lightning.ai/docs/pytorch/stable/cli/lightning_cli_advanced.html
+seed_everything: 42
+trainer:
+  accelerator: gpu
+  strategy: auto
+  devices: 1
+  num_nodes: 1
+  precision: 32-true
+  logger:
+    class_path: lightning.pytorch.loggers.TensorBoardLogger
+    # Nesting the logger config like this is equivalent to
+    # supplying the following argument to `lightning.pytorch.Trainer`:
+    # logger=TensorBoardLogger(
+    #     "/hpc/projects/intracellular_dashboard/viral-sensor/infection_classification/models/contrastive_tune_augmentations",
+    #     log_graph=True,
+    #     version="vanilla",
+    # )
+    init_args:
+      save_dir: /Users/ziwen.liu/Projects/test-time
+      # this is the name of the experiment.
+      # The logs will be saved in `save_dir/lightning_logs/version`
+      version: time_interval_1
+      log_graph: True
+  callbacks:
+    - class_path: lightning.pytorch.callbacks.LearningRateMonitor
+      init_args:
+        logging_interval: step
+    - class_path: lightning.pytorch.callbacks.ModelCheckpoint
+      init_args:
+        monitor: loss/val
+        every_n_epochs: 1
+        save_top_k: 4
+        save_last: true
+  fast_dev_run: false
+  max_epochs: 100
+  log_every_n_steps: 10
+  enable_checkpointing: true
+  inference_mode: true
+  use_distributed_sampler: true
+  # synchronize batchnorm parameters across multiple GPUs.
+  # important for contrastive learning to normalize the tensors across the whole batch.
+  sync_batchnorm: true
+model:
+  class_path: viscy.representation.engine.ContrastiveModule
+  init_args:
+    encoder:
+      class_path: viscy.representation.contrastive.ContrastiveEncoder
+      init_args:
+        backbone: convnext_tiny
+        in_channels: 1
+        in_stack_depth: 1
+        stem_kernel_size: [1, 4, 4]
+        stem_stride: [1, 4, 4]
+        embedding_dim: 768
+        projection_dim: 32
+        drop_path_rate: 0.0
+    loss_function:
+      class_path: torch.nn.TripletMarginLoss
+      init_args:
+        margin: 0.5
+    lr: 0.0002
+    log_batches_per_epoch: 3
+    log_samples_per_batch: 2
+    example_input_array_shape: [1, 1, 1, 128, 128]
+data:
+  class_path: viscy.data.triplet.TripletDataModule
+  init_args:
+    data_path: /Users/ziwen.liu/Downloads/Hela_CTC.zarr
+    tracks_path: /Users/ziwen.liu/Downloads/Hela_CTC.zarr
+    source_channel:
+      - DIC
+    z_range: [0, 1]
+    batch_size: 16
+    num_workers: 4
+    initial_yx_patch_size: [256, 256]
+    final_yx_patch_size: [128, 128]
+    time_interval: 1
+    normalizations:
+      - class_path: viscy.transforms.NormalizeSampled
+        init_args:
+          keys: [DIC]
+          level: fov_statistics
+          subtrahend: mean
+          divisor: std
+    augmentations:
+      - class_path: viscy.transforms.RandAffined
+        init_args:
+          keys: [DIC]
+          prob: 0.8
+          scale_range: [0, 0.2, 0.2]
+          rotate_range: [3.14, 0.0, 0.0]
+          shear_range: [0.0, 0.01, 0.01]
+          padding_mode: zeros
+      - class_path: viscy.transforms.RandAdjustContrastd
+        init_args:
+          keys: [DIC]
+          prob: 0.5
+          gamma: [0.8, 1.2]
+      - class_path: viscy.transforms.RandScaleIntensityd
+        init_args:
+          keys: [DIC]
+          prob: 0.5
+          factors: 0.5
+      - class_path: viscy.transforms.RandGaussianSmoothd
+        init_args:
+          keys: [DIC]
+          prob: 0.5
+          sigma_x: [0.25, 0.75]
+          sigma_y: [0.25, 0.75]
+          sigma_z: [0.0, 0.0]
+      - class_path: viscy.transforms.RandGaussianNoised
+        init_args:
+          keys: [DIC]
+          prob: 0.5
+          mean: 0.0
+          std: 0.2

applications/contrastive_phenotyping/contrastive_cli/predict_ctc_mps.yml

@@ -0,0 +1,48 @@
+seed_everything: 42
+trainer:
+  accelerator: gpu
+  strategy: auto
+  devices: auto
+  num_nodes: 1
+  precision: 32-true
+  callbacks:
+    - class_path: viscy.representation.embedding_writer.EmbeddingWriter
+      init_args:
+        output_path: /Users/ziwen.liu/Projects/test-time/predict/time_interval_1.zarr
+  inference_mode: true
+model:
+  class_path: viscy.representation.engine.ContrastiveModule
+  init_args:
+    encoder:
+      class_path: viscy.representation.contrastive.ContrastiveEncoder
+      init_args:
+        backbone: convnext_tiny
+        in_channels: 1
+        in_stack_depth: 1
+        stem_kernel_size: [1, 4, 4]
+        stem_stride: [1, 4, 4]
+        embedding_dim: 768
+        projection_dim: 32
+        drop_path_rate: 0.0
+    example_input_array_shape: [1, 1, 1, 128, 128]
+data:
+  class_path: viscy.data.triplet.TripletDataModule
+  init_args:
+    data_path: /Users/ziwen.liu/Downloads/Hela_CTC.zarr
+    tracks_path: /Users/ziwen.liu/Downloads/Hela_CTC.zarr
+    source_channel: DIC
+    z_range: [0, 1]
+    batch_size: 16
+    num_workers: 4
+    initial_yx_patch_size: [128, 128]
+    final_yx_patch_size: [128, 128]
+    time_interval: 1
+    normalizations:
+      - class_path: viscy.transforms.NormalizeSampled
+        init_args:
+          keys: [DIC]
+          level: fov_statistics
+          subtrahend: mean
+          divisor: std
+return_predictions: false
+ckpt_path: /Users/ziwen.liu/Projects/test-time/lightning_logs/time_interval_1/checkpoints/last.ckpt