+
+
+
+
+## Install dependencies
+```
+conda env create -f environment.yml
+```
+
+## ⏬ Dataset download
+
+Put dataset into data file:
+
+- [Norman](https://figshare.com/articles/dataset/Norman_et_al_2019_Science_labeled_Perturb-seq_data/24688110)
+- [Combosciplex subset of sciplex v3](https://figshare.com/articles/dataset/combosciplex/25062230?file=44229635)
+### Alternative Data Access
+
+We also provide the datasets via [Google Drive](https://drive.google.com/drive/folders/1cNpYAt9jVWZN82miNZtkP10YeSo7hufL?usp=sharing). This folder contains:
+- The **Norman** dataset and its corresponding data splits.
+- The **ComboSciPlex** dataset.
+
+Example directory layout after download (relative to repo root):
+```
+scDFM/
+├─ data/
+│ ├─ norman.h5ad
+│ └─ combosciplex.h5ad
+├─ src/
+│ └─ ...
+└─ run.sh
+```
+
+
+
+## 📥 Training
+
+An example on additive task.
+```bash
+bash run.sh
+```
+
+## 🫡 Citation
+
+If you find our work and/or our code useful, please cite us via:
+
+```bibtex
+@inproceedings{yu2026scdfm,
+ title={sc{DFM}: Distributional Flow Matching Model for Robust Single-Cell Perturbation Prediction},
+ author={Chenglei Yu and Chuanrui Wang and Bangyan Liao and Tailin Wu},
+ booktitle={The Fourteenth International Conference on Learning Representations},
+ year={2026},
+ url={https://openreview.net/forum?id=QSGanMEcUV}
+}
+```
+
+## 📚 Related Resources
+
+- AI for Scientific Simulation and Discovery Lab: https://github.com/AI4Science-WestlakeU
diff --git a/scDFM/assets/fig1.png b/scDFM/assets/fig1.png
new file mode 100644
index 0000000000000000000000000000000000000000..524547e8735dade267ec824cba6f4cb9e79870d2
--- /dev/null
+++ b/scDFM/assets/fig1.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:56e7b45649aa2268576eed5923db5f84a2782c7288fd538a0dc37d8e87bb587b
+size 931506
diff --git a/scDFM/assets/logo.png b/scDFM/assets/logo.png
new file mode 100644
index 0000000000000000000000000000000000000000..c9a04c36372c1403af3e0ddb369498feb673a697
--- /dev/null
+++ b/scDFM/assets/logo.png
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:8f0851a69c4456af600fd6044c2b2e37147ab532d6b8dab709a50c117a15b9cd
+size 176734
diff --git a/scDFM/assets/scdfm_PPT.pdf b/scDFM/assets/scdfm_PPT.pdf
new file mode 100644
index 0000000000000000000000000000000000000000..865a4c771e7154640422b2a6bafcaa45224a6400
--- /dev/null
+++ b/scDFM/assets/scdfm_PPT.pdf
@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:2d0fcbceb443faf49bdce0a3f6685a388649a920510c4c133e777164b7eab3da
+size 4959141
diff --git a/scDFM/config/config_flow.py b/scDFM/config/config_flow.py
new file mode 100644
index 0000000000000000000000000000000000000000..28a5bd567f062dac5d935769d93c7b5360b4f284
--- /dev/null
+++ b/scDFM/config/config_flow.py
@@ -0,0 +1,64 @@
+
+from dataclasses import dataclass
+import os
+@dataclass
+class FlowConfig:
+ # Flow model type
+ model_type: str = 'hierarchical'
+
+ # Flow Matching specific parameters
+ batch_size: int = 32
+ ntoken: int = 512
+ d_model: int = 512
+ lr: float = 1e-5
+ steps: int = 5000
+ eta_min: float = 1e-7
+ devices: str = "1"
+ test_only: bool = False
+ # Perturbation related parameters
+ data_name: str = "combosciplex"
+ perturbation_function: str = 'crisper'
+ noise_type: str = "Gaussian"
+ poisson_alpha: float = 0.8
+ poisson_target_sum: int = -1
+
+ print_every: int = 5000
+ mode: str = 'predict_y' # predict_y, predict_p
+ result_path: str = './result'
+ perturbation_fusion_method: str = 'sum' # mlp, sum
+ fusion_method: str = 'cross' # cross , concat, add
+ infer_top_gene: int = 1000
+ n_top_genes: int = 5000
+ checkpoint_path: str = ''
+ gamma: float = 0.0
+ split_method: str = 'additive'
+ use_mmd_loss: bool = False
+ fold: int = 0
+ use_negative_edge: bool = False
+ topk: int = 15
+
+ def __post_init__(self):
+ if self.data_name == 'norman_umi_go_filtered':
+ self.n_top_genes = 5054
+ if self.data_name == 'norman':
+ self.n_top_genes = 5000
+ path = self.make_path()
+
+ def make_path(self):
+ exp_name = '-'.join(['flow',
+ f'fusion_{self.fusion_method}',
+ f'{self.data_name}',
+ self.model_type,
+ self.mode,
+ f'gamma_{self.gamma}',
+ f'perturbation_function_{self.perturbation_function}',
+ f'lr_{self.lr}',
+ f'dim_model_{self.d_model}',
+ f'infer_top_gene_{self.infer_top_gene}',
+ f'split_method_{self.split_method}',
+ f'use_mmd_loss_{self.use_mmd_loss}',
+ f'fold_{self.fold}',
+ f'use_negative_edge_{self.use_negative_edge}',
+ f'topk_{self.topk}',
+ ])
+ return os.path.join(self.result_path, exp_name)
\ No newline at end of file
diff --git a/scDFM/environment.yml b/scDFM/environment.yml
new file mode 100644
index 0000000000000000000000000000000000000000..88802032db069b901c887ad0750e4336d8bb34df
--- /dev/null
+++ b/scDFM/environment.yml
@@ -0,0 +1,318 @@
+name: sc
+channels:
+ - conda-forge
+dependencies:
+ - _libgcc_mutex=0.1
+ - _openmp_mutex=4.5
+ - asttokens=3.0.0
+ - bzip2=1.0.8
+ - ca-certificates=2025.8.3
+ - comm=0.2.3
+ - decorator=5.2.1
+ - exceptiongroup=1.3.0
+ - executing=2.2.0
+ - importlib-metadata=8.7.0
+ - ipykernel=6.30.1
+ - ipython=9.4.0
+ - ipython_pygments_lexers=1.1.1
+ - jedi=0.19.2
+ - jupyter_client=8.6.3
+ - jupyter_core=5.8.1
+ - keyutils=1.6.1
+ - krb5=1.21.3
+ - ld_impl_linux-64=2.44
+ - libedit=3.1.20250104
+ - libexpat=2.7.1
+ - libffi=3.4.6
+ - libgcc=15.1.0
+ - libgcc-ng=15.1.0
+ - libgomp=15.1.0
+ - liblzma=5.8.1
+ - libnsl=2.0.1
+ - libsodium=1.0.20
+ - libsqlite=3.50.4
+ - libstdcxx=15.1.0
+ - libstdcxx-ng=15.1.0
+ - libuuid=2.38.1
+ - libxcrypt=4.4.36
+ - libzlib=1.3.1
+ - matplotlib-inline=0.1.7
+ - ncurses=6.5
+ - nest-asyncio=1.6.0
+ - openssl=3.5.2
+ - packaging=25.0
+ - parso=0.8.4
+ - pexpect=4.9.0
+ - pickleshare=0.7.5
+ - pip=25.2
+ - platformdirs=4.3.8
+ - prompt-toolkit=3.0.51
+ - ptyprocess=0.7.0
+ - pure_eval=0.2.3
+ - pygments=2.19.2
+ - python=3.11.13
+ - python-dateutil=2.9.0.post0
+ - python_abi=3.11
+ - readline=8.2
+ - setuptools=80.9.0
+ - six=1.17.0
+ - stack_data=0.6.3
+ - tk=8.6.13
+ - traitlets=5.14.3
+ - typing_extensions=4.14.1
+ - wcwidth=0.2.13
+ - wheel=0.45.1
+ - zeromq=4.3.5
+ - zipp=3.23.0
+ - pip:
+ - absl-py==2.3.1
+ - accelerate==1.9.0
+ - adjusttext==1.3.0
+ - adpbulk==0.1.4
+ - aiofiles==24.1.0
+ - aiohappyeyeballs==2.6.1
+ - aiohttp==3.12.15
+ - aiosignal==1.4.0
+ - anndata==0.12.1
+ - antlr4-python3-runtime==4.9.3
+ - anyio==4.10.0
+ - argon2-cffi==25.1.0
+ - argon2-cffi-bindings==25.1.0
+ - array-api-compat==1.12.0
+ - arrow==1.3.0
+ - arviz==0.22.0
+ - async-lru==2.0.5
+ - attrs==25.3.0
+ - babel==2.17.0
+ - beautifulsoup4==4.13.4
+ - bleach==6.2.0
+ - blitzgsea==1.3.54
+ - cached-property==2.0.1
+ - cell-eval==0.5.42
+ - cell-gears==0.1.2
+ - certifi==2025.8.3
+ - cffi==1.17.1
+ - charset-normalizer==3.4.2
+ - chex==0.1.90
+ - click==8.2.1
+ - contourpy==1.3.3
+ - cpa==0.4.1
+ - crc32c==2.7.1
+ - cycler==0.12.1
+ - datasets==2.21.0
+ - dcor==0.6
+ - debugpy==1.8.15
+ - defusedxml==0.7.1
+ - deprecated==1.2.18
+ - dill==0.3.8
+ - docrep==0.3.2
+ - docstring-parser==0.17.0
+ - donfig==0.8.1.post1
+ - einops==0.8.1
+ - equinox==0.13.0
+ - et-xmlfile==2.0.0
+ - etils==1.13.0
+ - fast-array-utils==1.2.1
+ - fastjsonschema==2.21.1
+ - filelock==3.18.0
+ - flax==0.11.0
+ - fonttools==4.59.0
+ - formulaic==1.2.0
+ - formulaic-contrasts==1.0.0
+ - fqdn==1.5.1
+ - frozenlist==1.7.0
+ - fsspec==2024.6.1
+ - grpcio==1.74.0
+ - h11==0.16.0
+ - h5netcdf==1.6.3
+ - h5py==3.14.0
+ - hf-xet==1.1.5
+ - httpcore==1.0.9
+ - httpx==0.28.1
+ - huggingface-hub==0.34.3
+ - humanize==4.12.3
+ - idna==3.10
+ - igraph==0.11.9
+ - importlib-resources==6.5.2
+ - iniconfig==2.1.0
+ - interface-meta==1.3.0
+ - ipywidgets==8.1.7
+ - isoduration==20.11.0
+ - jax==0.6.2
+ - jaxlib==0.6.2
+ - jaxopt==0.8.5
+ - jaxtyping==0.3.2
+ - jinja2==3.1.6
+ - joblib==1.5.1
+ - json5==0.12.0
+ - jsonpointer==3.0.0
+ - jsonschema==4.25.0
+ - jsonschema-specifications==2025.4.1
+ - jupyter==1.1.1
+ - jupyter-console==6.6.3
+ - jupyter-events==0.12.0
+ - jupyter-lsp==2.2.6
+ - jupyter-server==2.16.0
+ - jupyter-server-terminals==0.5.3
+ - jupyterlab==4.4.5
+ - jupyterlab-pygments==0.3.0
+ - jupyterlab-server==2.27.3
+ - jupyterlab-widgets==3.0.15
+ - kiwisolver==1.4.8
+ - lamin-utils==0.15.0
+ - lark==1.2.2
+ - legacy-api-wrap==1.4.1
+ - leidenalg==0.10.2
+ - lightning==2.5.2
+ - lightning-utilities==0.15.0
+ - lineax==0.0.8
+ - llvmlite==0.44.0
+ - markdown==3.8.2
+ - markdown-it-py==3.0.0
+ - markupsafe==3.0.2
+ - matplotlib==3.10.5
+ - mdurl==0.1.2
+ - mistune==3.1.3
+ - ml-collections==1.1.0
+ - ml-dtypes==0.5.3
+ - mpmath==1.3.0
+ - msgpack==1.1.1
+ - mudata==0.3.2
+ - multidict==6.6.3
+ - multipledispatch==1.0.0
+ - multiprocess==0.70.16
+ - narwhals==2.0.1
+ - natsort==8.4.0
+ - nbclient==0.10.2
+ - nbconvert==7.16.6
+ - nbformat==5.10.4
+ - networkx==3.5
+ - notebook==7.4.5
+ - notebook-shim==0.2.4
+ - numba==0.61.2
+ - numcodecs==0.16.1
+ - numpy==2.2.6
+ - numpyro==0.18.0
+ - nvidia-cublas-cu12==12.6.4.1
+ - nvidia-cuda-cupti-cu12==12.6.80
+ - nvidia-cuda-nvrtc-cu12==12.6.77
+ - nvidia-cuda-runtime-cu12==12.6.77
+ - nvidia-cudnn-cu12==9.5.1.17
+ - nvidia-cufft-cu12==11.3.0.4
+ - nvidia-cufile-cu12==1.11.1.6
+ - nvidia-curand-cu12==10.3.7.77
+ - nvidia-cusolver-cu12==11.7.1.2
+ - nvidia-cusparse-cu12==12.5.4.2
+ - nvidia-cusparselt-cu12==0.6.3
+ - nvidia-nccl-cu12==2.26.2
+ - nvidia-nvjitlink-cu12==12.6.85
+ - nvidia-nvtx-cu12==12.6.77
+ - omegaconf==2.3.0
+ - openpyxl==3.1.5
+ - opt-einsum==3.4.0
+ - optax==0.2.5
+ - orbax==0.1.7
+ - orbax-checkpoint==0.11.20
+ - ott-jax==0.5.1
+ - overrides==7.7.0
+ - pandas==2.3.1
+ - pandocfilters==1.5.1
+ - patsy==1.0.1
+ - pdex==0.1.21
+ - pertpy==1.0.1
+ - pillow==11.3.0
+ - pluggy==1.6.0
+ - ply==3.11
+ - polars==1.32.0
+ - pot==0.9.5
+ - prometheus-client==0.22.1
+ - propcache==0.3.2
+ - protobuf==6.31.1
+ - psutil==7.0.0
+ - pubchempy==1.0.4
+ - pyarrow==21.0.0
+ - pycparser==2.22
+ - pydeseq2==0.5.2
+ - pydot==4.0.1
+ - pynndescent==0.5.13
+ - pyomo==6.9.2
+ - pyparsing==3.2.3
+ - pyro-api==0.1.2
+ - pyro-ppl==1.9.1
+ - pytest==8.4.1
+ - python-json-logger==3.3.0
+ - pytorch-lightning==1.9.5
+ - pytz==2025.2
+ - pyyaml==6.0.2
+ - pyzmq==27.0.1
+ - rdkit==2025.3.5
+ - referencing==0.36.2
+ - requests==2.32.4
+ - rfc3339-validator==0.1.4
+ - rfc3986-validator==0.1.1
+ - rfc3987-syntax==1.1.0
+ - rich==14.1.0
+ - rpds-py==0.26.0
+ - safetensors==0.5.3
+ - scanpy==1.11.4
+ - scgpt==0.2.4
+ - scib==1.1.7
+ - scikit-learn==1.7.1
+ - scikit-misc==0.5.1
+ - scipy==1.16.1
+ - scvi-tools==0.20.3
+ - seaborn==0.13.2
+ - send2trash==1.8.3
+ - session-info==1.0.1
+ - session-info2==0.2
+ - shtab==1.7.2
+ - simplejson==3.20.1
+ - sniffio==1.3.1
+ - soupsieve==2.7
+ - sparse==0.17.0
+ - sparsecca==0.3.1
+ - statsmodels==0.14.5
+ - stdlib-list==0.11.1
+ - sympy==1.14.0
+ - tensorboard==2.20.0
+ - tensorboard-data-server==0.7.2
+ - tensorstore==0.1.76
+ - terminado==0.18.1
+ - texttable==1.7.0
+ - threadpoolctl==3.6.0
+ - timm==1.0.19
+ - tinycss2==1.4.0
+ - tomlkit==0.13.3
+ - toolz==1.0.0
+ - torch==2.7.1
+ - torch-geometric==2.6.1
+ - torchaudio==2.7.1
+ - torchdiffeq==0.2.5
+ - torchinfo==1.8.0
+ - torchmetrics==1.8.0
+ - torchtext==0.18.0
+ - torchvision==0.22.1
+ - tornado==6.5.1
+ - tqdm==4.67.1
+ - treescope==0.1.9
+ - triton==3.3.1
+ - typeguard==4.4.4
+ - types-python-dateutil==2.9.0.20250708
+ - tyro==0.9.27
+ - tzdata==2025.2
+ - umap-learn==0.5.9.post2
+ - uri-template==1.3.0
+ - urllib3==2.5.0
+ - wadler-lindig==0.1.7
+ - webcolors==24.11.1
+ - webencodings==0.5.1
+ - websocket-client==1.8.0
+ - werkzeug==3.1.3
+ - widgetsnbextension==4.0.14
+ - wrapt==1.17.2
+ - xarray==2025.7.1
+ - xarray-einstats==0.9.1
+ - xxhash==3.5.0
+ - yarl==1.20.1
+ - zarr==3.1.1
+prefix: /opt/conda/envs/sc
diff --git a/scDFM/run.sh b/scDFM/run.sh
new file mode 100644
index 0000000000000000000000000000000000000000..c7a0fbf33509be7a0c171ad75b4dff9024683d91
--- /dev/null
+++ b/scDFM/run.sh
@@ -0,0 +1,24 @@
+export PYTHONPATH=./
+
+python src/script/run.py \
+--batch_size=48 \
+--devices='0' \
+--model_type=origin \
+--lr=5e-5 \
+--steps=200000 \
+--data_name=norman \
+--d_model=128 \
+--eta_min=1e-6 \
+--fusion_method=differential_perceiver \
+--infer_top_gene=1000 \
+--n_top_genes=5000 \
+--result_path=./result/additive \
+--perturbation_function=crisper \
+--noise_type=Gaussian \
+--mode=predict_y \
+--gamma=0.5 \
+--split_method=additive \
+--use_mmd_loss \
+--fold=1 \
+--topk=30 \
+--use_negative_edge \
diff --git a/scDFM/src/data_process/_data.py b/scDFM/src/data_process/_data.py
new file mode 100644
index 0000000000000000000000000000000000000000..21ff590fd61e2583f58b9a128fa91dfde4c15a52
--- /dev/null
+++ b/scDFM/src/data_process/_data.py
@@ -0,0 +1,218 @@
+from collections.abc import Callable
+from dataclasses import dataclass
+from typing import Any
+
+import jax
+import jax.numpy as jnp
+
+from src.utils._types import ArrayLike
+
+__all__ = [
+ "BaseDataMixin",
+ "ConditionData",
+ "PredictionData",
+ "TrainingData",
+ "ValidationData",
+]
+
+
+@dataclass
+class ReturnData: # TODO: this should rather be a NamedTuple
+ split_covariates_mask: jnp.ndarray | None
+ split_idx_to_covariates: dict[int, tuple[Any, ...]]
+ perturbation_covariates_mask: jnp.ndarray | None
+ perturbation_idx_to_covariates: dict[int, tuple[Any, ...]]
+ perturbation_idx_to_id: dict[int, Any]
+ condition_data: dict[str, ArrayLike]
+ control_to_perturbation: dict[int, ArrayLike]
+ max_combination_length: int
+ condition_data_id: jnp.ndarray # mapping from perturbation_idx to condition_id
+
+
+class BaseDataMixin:
+ """Base class for data containers."""
+
+ @property
+ def n_controls(self) -> int:
+ """Returns the number of control covariate values."""
+ return len(self.split_idx_to_covariates) # type: ignore[attr-defined]
+
+ @property
+ def n_perturbations(self) -> int:
+ """Returns the number of perturbation covariate combinations."""
+ return len(self.perturbation_idx_to_covariates) # type: ignore[attr-defined]
+
+ @property
+ def n_perturbation_covariates(self) -> int:
+ """Returns the number of perturbation covariates."""
+ return len(self.condition_data) # type: ignore[attr-defined]
+
+ def _format_params(self, fmt: Callable[[Any], str]) -> str:
+ params = {
+ "n_controls": self.n_controls,
+ "n_perturbations": self.n_perturbations,
+ }
+ return ", ".join(f"{name}={fmt(val)}" for name, val in params.items())
+
+ def __repr__(self) -> str:
+ return f"{self.__class__.__name__}[{self._format_params(repr)}]"
+
+
+@dataclass
+class ConditionData(BaseDataMixin):
+ """Data container containing condition embeddings.
+
+ Parameters
+ ----------
+ condition_data
+ Dictionary with embeddings for conditions.
+ max_combination_length
+ Maximum number of covariates in a combination.
+ null_value
+ Token to use for masking `null_value`.
+ data_manager
+ Data manager used to generate the data.
+ condition_data_id
+ Mapping from perturbation index to condition ID.
+ """
+
+ condition_data: dict[str, ArrayLike]
+ max_combination_length: int
+ perturbation_idx_to_covariates: dict[int, tuple[str, ...]]
+ perturbation_idx_to_id: dict[int, Any]
+ null_value: Any
+ data_manager: Any
+ condition_data_id: jnp.ndarray
+
+@dataclass
+class TrainingData(BaseDataMixin):
+ """Training data.
+
+ Parameters
+ ----------
+ cell_data
+ The representation of cell data, e.g. PCA of gene expression data.
+ split_covariates_mask
+ Mask of the split covariates.
+ split_idx_to_covariates
+ Dictionary explaining values in ``split_covariates_mask``.
+ perturbation_covariates_mask
+ Mask of the perturbation covariates.
+ perturbation_idx_to_covariates
+ Dictionary explaining values in ``perturbation_covariates_mask``.
+ condition_data
+ Dictionary with embeddings for conditions.
+ control_to_perturbation
+ Mapping from control index to target distribution indices.
+ max_combination_length
+ Maximum number of covariates in a combination.
+ data_manager
+ The data manager
+ condition_data_id
+ Mapping from perturbation index to condition ID.
+ """
+
+ cell_data: jax.Array # (n_cells, n_features)
+ split_covariates_mask: jax.Array # (n_cells,), which cell assigned to which source distribution
+ split_idx_to_covariates: dict[int, tuple[Any, ...]] # (n_sources,) dictionary explaining split_covariates_mask
+ perturbation_covariates_mask: jax.Array # (n_cells,), which cell assigned to which target distribution
+ perturbation_idx_to_covariates: dict[
+ int, tuple[str, ...]
+ ] # (n_targets,), dictionary explaining perturbation_covariates_mask
+ perturbation_idx_to_id: dict[int, Any]
+ condition_data: dict[str, ArrayLike] # (n_targets,) all embeddings for conditions
+ control_to_perturbation: dict[int, ArrayLike] # mapping from control idx to target distribution idcs
+ max_combination_length: int
+ null_value: Any
+ data_manager: Any
+ condition_data_id: jnp.ndarray
+ cell_data_id: Any
+
+
+@dataclass
+class ValidationData(BaseDataMixin):
+ """Data container for the validation data.
+
+ Parameters
+ ----------
+ cell_data
+ The representation of cell data, e.g. PCA of gene expression data.
+ split_covariates_mask
+ Mask of the split covariates.
+ split_idx_to_covariates
+ Dictionary explaining values in ``split_covariates_mask``.
+ perturbation_covariates_mask
+ Mask of the perturbation covariates.
+ perturbation_idx_to_covariates
+ Dictionary explaining values in ``perturbation_covariates_mask``.
+ condition_data
+ Dictionary with embeddings for conditions.
+ control_to_perturbation
+ Mapping from control index to target distribution indices.
+ max_combination_length
+ Maximum number of covariates in a combination.
+ data_manager
+ The data manager
+ condition_data_id
+ Mapping from perturbation index to condition ID.
+ n_conditions_on_log_iteration
+ Number of conditions to use for computation callbacks at each logged iteration.
+ If :obj:`None`, use all conditions.
+ n_conditions_on_train_end
+ Number of conditions to use for computation callbacks at the end of training.
+ If :obj:`None`, use all conditions.
+ """
+
+ cell_data: jax.Array # (n_cells, n_features)
+ split_covariates_mask: jax.Array # (n_cells,), which cell assigned to which source distribution
+ split_idx_to_covariates: dict[int, tuple[Any, ...]] # (n_sources,) dictionary explaining split_covariates_mask
+ perturbation_covariates_mask: jax.Array # (n_cells,), which cell assigned to which target distribution
+ perturbation_idx_to_covariates: dict[
+ int, tuple[str, ...]
+ ] # (n_targets,), dictionary explaining perturbation_covariates_mask
+ perturbation_idx_to_id: dict[int, Any]
+ condition_data: dict[str, ArrayLike] # (n_targets,) all embeddings for conditions
+ control_to_perturbation: dict[int, jax.Array] # mapping from control idx to target distribution idcs
+ max_combination_length: int
+ null_value: Any
+ data_manager: Any
+ condition_data_id: jnp.ndarray
+ n_conditions_on_log_iteration: int | None = None
+ n_conditions_on_train_end: int | None = None
+
+
+@dataclass
+class PredictionData(BaseDataMixin):
+ """Data container to perform prediction.
+
+ Parameters
+ ----------
+ src_data
+ Dictionary with data for source cells.
+ condition_data
+ Dictionary with embeddings for conditions.
+ control_to_perturbation
+ Mapping from control index to target distribution indices.
+ covariate_encoder
+ Encoder for the primary covariate.
+ max_combination_length
+ Maximum number of covariates in a combination.
+ null_value
+ Token to use for masking ``null_value``.
+ condition_data_id
+ Mapping from condition name to condition ID embedding.
+ """
+
+ cell_data: jax.Array # (n_cells, n_features)
+ split_covariates_mask: jax.Array # (n_cells,), which cell assigned to which source distribution
+ split_idx_to_covariates: dict[int, tuple[Any, ...]] # (n_sources,) dictionary explaining split_covariates_mask
+ perturbation_idx_to_covariates: dict[
+ int, tuple[str, ...]
+ ] # (n_targets,), dictionary explaining perturbation_covariates_mask
+ perturbation_idx_to_id: dict[int, Any]
+ condition_data: dict[str, ArrayLike] # (n_targets,) all embeddings for conditions
+ control_to_perturbation: dict[int, ArrayLike]
+ max_combination_length: int
+ null_value: Any
+ data_manager: Any
+ condition_data_id: jnp.ndarray
diff --git a/scDFM/src/data_process/_dataloader.py b/scDFM/src/data_process/_dataloader.py
new file mode 100644
index 0000000000000000000000000000000000000000..be30c54b45b3b8c6367e12a4e474df6bcdcf5d4b
--- /dev/null
+++ b/scDFM/src/data_process/_dataloader.py
@@ -0,0 +1,235 @@
+import abc
+from typing import Any, Literal
+
+import jax
+import jax.numpy as jnp
+import numpy as np
+
+from src.data_process._data import PredictionData, TrainingData, ValidationData
+import pdb
+import torch
+
+__all__ = ["TrainSampler", "ValidationSampler", "PredictionSampler"]
+
+
+class TrainSampler:
+ """Data sampler for :class:`~cellflow.data.TrainingData`.
+
+ Parameters
+ ----------
+ data
+ The training data.
+ batch_size
+ The batch size.
+
+ """
+
+ def __init__(self, data: TrainingData, batch_size: int = 1024):
+ self._data = data
+ self._data_idcs = jnp.arange(data.cell_data.shape[0])
+ self.batch_size = batch_size
+ self.n_source_dists = data.n_controls
+ self.n_target_dists = data.n_perturbations
+ self.conditional_samplings = [
+ lambda key, i=i: jax.random.choice(key, self._data.control_to_perturbation[i])
+ for i in range(self.n_source_dists)
+ ]
+ self.get_embeddings = lambda idx: {
+ pert_cov: jnp.expand_dims(arr[idx], 0) for pert_cov, arr in self._data.condition_data.items()
+ }
+ # @jax.jit
+ def _sample(rng: jax.Array) -> Any:
+
+ rng_1, rng_2, rng_3, rng_4 = jax.random.split(rng, 4)
+ source_dist_idx = jax.random.choice(rng_1, self.n_source_dists)
+ source_cells_mask = self._data.split_covariates_mask == source_dist_idx
+ src_cond_p = source_cells_mask / jnp.count_nonzero(source_cells_mask)
+ source_batch_idcs = jax.random.choice(rng_2, self._data_idcs, [self.batch_size], replace=True, p=src_cond_p)
+ source_batch = self._data.cell_data[source_batch_idcs]
+
+ target_dist_idx = jax.lax.switch(source_dist_idx, self.conditional_samplings, rng_3)
+ target_cells_mask = self._data.perturbation_covariates_mask == target_dist_idx
+ tgt_cond_p = target_cells_mask / jnp.count_nonzero(target_cells_mask)
+ target_batch_idcs = jax.random.choice(
+ rng_4,
+ self._data_idcs,
+ [self.batch_size],
+ replace=True,
+ p=tgt_cond_p,
+ )
+ target_batch = self._data.cell_data[target_batch_idcs]
+ if self._data.condition_data is None:
+ return {"src_cell_data": source_batch, "tgt_cell_data": target_batch}
+
+ condition_batch = self.get_embeddings(target_dist_idx)
+
+ return {
+ "src_cell_data": source_batch,
+ "tgt_cell_data": target_batch,
+ "condition": condition_batch,
+ "condition_id": self._data.condition_data_id[target_dist_idx],
+ "src_cell_id": self._data.cell_data_id[source_batch_idcs],
+ "tgt_cell_id": self._data.cell_data_id[target_batch_idcs]
+ }
+
+ self.sample = _sample
+
+ @property
+ def data(self) -> TrainingData:
+ """The training data."""
+ return self._data
+
+ def precompute_latent_embedding(self, model,):
+ """
+ compute latent representation
+
+ return
+ ----
+ latent_embedding : torch.Tensor
+ """
+
+ cell_data = self._data.cell_data
+ with torch.no_grad():
+ self._data.cell_embedding = model(cell_data)
+
+
+class BaseValidSampler(abc.ABC):
+ @abc.abstractmethod
+ def sample(*args, **kwargs):
+ pass
+
+ def _get_key(self, cond_idx: int) -> tuple[str, ...]:
+ if len(self._data.perturbation_idx_to_id): # type: ignore[attr-defined]
+ return self._data.perturbation_idx_to_id[cond_idx] # type: ignore[attr-defined]
+ cov_combination = self._data.perturbation_idx_to_covariates[cond_idx] # type: ignore[attr-defined]
+ return tuple(cov_combination[i] for i in range(len(cov_combination)))
+
+ def _get_perturbation_to_control(self, data: ValidationData | PredictionData) -> dict[int, int]:
+ d = {}
+ for k, v in data.control_to_perturbation.items():
+ for el in v:
+ d[el] = k
+ return d
+
+ def _get_condition_data(self, cond_idx: int) -> jnp.ndarray:
+ return {k: v[[cond_idx], ...] for k, v in self._data.condition_data.items()} # type: ignore[attr-defined]
+
+
+class ValidationSampler(BaseValidSampler):
+ """Data sampler for :class:`~cellflow.data.ValidationData`.
+
+ Parameters
+ ----------
+ val_data
+ The validation data.
+ seed
+ Random seed.
+ """
+
+ def __init__(self, val_data: ValidationData, seed: int = 0) -> None:
+ self._data = val_data
+ self.perturbation_to_control = self._get_perturbation_to_control(val_data)
+ self.n_conditions_on_log_iteration = (
+ val_data.n_conditions_on_log_iteration
+ if val_data.n_conditions_on_log_iteration is not None
+ else val_data.n_perturbations
+ )
+ self.n_conditions_on_train_end = (
+ val_data.n_conditions_on_train_end
+ if val_data.n_conditions_on_train_end is not None
+ else val_data.n_perturbations
+ )
+ self.rng = np.random.default_rng(seed)
+ if self._data.condition_data is None:
+ raise NotImplementedError("Validation data must have condition data.")
+
+ def sample(self, mode: Literal["on_log_iteration", "on_train_end"]) -> Any:
+ """Sample data for validation.
+
+ Parameters
+ ----------
+ mode
+ Sampling mode. Either ``"on_log_iteration"`` or ``"on_train_end"``.
+
+ Returns
+ -------
+ Dictionary with source, condition, and target data from the validation data.
+ """
+ import pdb; pdb.set_trace()
+ size = self.n_conditions_on_log_iteration if mode == "on_log_iteration" else self.n_conditions_on_train_end
+ condition_idcs = self.rng.choice(self._data.n_perturbations, size=(size,), replace=False)
+
+ source_idcs = [self.perturbation_to_control[cond_idx] for cond_idx in condition_idcs]
+ source_cells_mask = [self._data.split_covariates_mask == source_idx for source_idx in source_idcs]
+ source_cells = [self._data.cell_data[mask] for mask in source_cells_mask]
+ target_cells_mask = [cond_idx == self._data.perturbation_covariates_mask for cond_idx in condition_idcs]
+ target_cells = [self._data.cell_data[mask] for mask in target_cells_mask]
+ conditions = [self._get_condition_data(cond_idx) for cond_idx in condition_idcs]
+ cell_rep_dict = {}
+ cond_dict = {}
+ true_dict = {}
+ for i in range(len(condition_idcs)):
+ k = self._get_key(condition_idcs[i])
+ cell_rep_dict[k] = source_cells[i]
+ cond_dict[k] = conditions[i]
+ true_dict[k] = target_cells[i]
+
+ return {"source": cell_rep_dict, "condition": cond_dict, "target": true_dict}
+
+ @property
+ def data(self) -> ValidationData:
+ """The validation data."""
+ return self._data
+
+
+class PredictionSampler(BaseValidSampler):
+ """Data sampler for :class:`~cellflow.data.PredictionData`.
+
+ Parameters
+ ----------
+ pred_data
+ The prediction data.
+
+ """
+
+ def __init__(self, pred_data: PredictionData) -> None:
+ self._data = pred_data
+ self.perturbation_to_control = self._get_perturbation_to_control(pred_data)
+ if self._data.condition_data is None:
+ raise NotImplementedError("Validation data must have condition data.")
+
+ def sample(self) -> Any:
+ """Sample data for prediction.
+
+ Returns
+ -------
+ Dictionary with source and condition data from the prediction data.
+ """
+ condition_idcs = range(self._data.n_perturbations)
+
+ source_idcs = [self.perturbation_to_control[cond_idx] for cond_idx in condition_idcs]
+ source_cells_mask = [self._data.split_covariates_mask == source_idx for source_idx in source_idcs]
+ source_cells = [self._data.cell_data[mask] for mask in source_cells_mask]
+ target_cells_mask = [cond_idx == self._data.perturbation_covariates_mask for cond_idx in condition_idcs]
+ target_cells = [self._data.cell_data[mask] for mask in target_cells_mask]
+
+ conditions = [self._get_condition_data(cond_idx) for cond_idx in condition_idcs]
+ cell_rep_dict = {}
+ cond_dict = {}
+ true_dict = {}
+ for i in range(len(condition_idcs)):
+ k = self._get_key(condition_idcs[i])
+ cell_rep_dict[k] = source_cells[i]
+ cond_dict[k] = conditions[i]
+ true_dict[k] = target_cells[i]
+
+ return {
+ "source": cell_rep_dict,
+ "target": true_dict,
+ "condition": cond_dict,
+ }
+
+ @property
+ def data(self) -> PredictionData:
+ """The training data."""
+ return self._data
diff --git a/scDFM/src/data_process/_datamanager.py b/scDFM/src/data_process/_datamanager.py
new file mode 100644
index 0000000000000000000000000000000000000000..9d1000417ae97bbdc35a1556e7c5bb4d288883ad
--- /dev/null
+++ b/scDFM/src/data_process/_datamanager.py
@@ -0,0 +1,988 @@
+from collections.abc import Sequence
+from typing import Any
+
+import anndata
+import jax
+import jax.numpy as jnp
+import numpy as np
+import pandas as pd
+import scipy.sparse as sp
+import sklearn.preprocessing as preprocessing
+from pandas.api.types import is_numeric_dtype
+from tqdm import tqdm
+
+from src.utils._logging import logger
+from src.utils._types import ArrayLike
+from src.data_process._data import ConditionData, PredictionData, ReturnData, TrainingData, ValidationData
+import pdb
+from src.utils._utils import _flatten_list, _to_list
+
+__all__ = ["DataManager"]
+
+
+class DataManager:
+ """Data manager for handling perturbation data.
+
+ Parameters
+ ----------
+ adata
+ An :class:`~anndata.AnnData` object.
+ covariate_encoder
+ Encoder for the primary covariate.
+ categorical
+ Whether the primary covariate is categorical.
+ max_combination_length
+ Maximum number of combinations of primary ``perturbation_covariates``.
+ sample_rep
+ Key in :attr:`~anndata.AnnData.obsm` where the sample representation is stored or ``'X'``
+ to use :attr:`~anndata.AnnData.X`.
+ covariate_data
+ Dataframe with covariates. If :obj:`None`, :attr:`~anndata.AnnData.obs` is used.
+ condition_id_key
+ Key in :attr:`~anndata.AnnData.obs` that defines the condition id.
+ perturbation_covariates
+ A dictionary where the keys indicate the name of the covariate group
+ and the values are keys in :attr:`~anndata.AnnData.uns`. The corresponding
+ columns should be either boolean (presence/abscence of the perturbation) or
+ numeric (concentration or magnitude of the perturbation). If multiple groups
+ are provided, the first is interpreted as the primary perturbation and the
+ others as covariates corresponding to these perturbations, e.g.
+ ``{"drug":("drugA", "drugB"), "time":("drugA_time", "drugB_time")}``.
+ perturbation_covariate_reps
+ A dictionary where the keys indicate the name of the covariate group and the
+ values are keys in :attr:`~anndata.AnnData.uns` storing a dictionary with
+ the representation of the covariates. E.g. ``{"drug":"drug_embeddings"}``
+ with ``adata.uns["drug_embeddings"] = {"drugA": np.array, "drugB": np.array}``.
+ sample_covariates
+ Keys in :attr:`~anndata.AnnData.obs` indicating sample covatiates to be taken
+ into account for training and prediction, e.g. ``["age", "cell_type"]``.
+ sample_covariate_reps
+ A dictionary where the keys indicate the name of the covariate group and the
+ values are keys in :attr:`~anndata.AnnData.uns` storing a dictionary with the
+ representation of the covariates. E.g. ``{"cell_type": "cell_type_embeddings"}`` with
+ ``adata.uns["cell_type_embeddings"] = {"cell_typeA": np.array, "cell_typeB": np.array}``.
+ split_covariates
+ Covariates in :attr:`~anndata.AnnData.obs` to split all control cells into
+ different control populations. The perturbed cells are also split according to these
+ columns, but if these covariates should also be encoded in the model, the corresponding
+ column should also be used in ``perturbation_covariates`` or ``sample_covariates``.
+ null_value
+ Value to use for padding to ``max_combination_length``.
+ """
+
+ def __init__(
+ self,
+ adata: anndata.AnnData,
+ sample_rep: str | dict[str, str],
+ control_key: str,
+ perturbation_covariates: dict[str, Sequence[str]] | None = None,
+ perturbation_covariate_reps: dict[str, str] | None = None,
+ sample_covariates: Sequence[str] | None = None,
+ sample_covariate_reps: dict[str, str] | None = None,
+ split_covariates: Sequence[str] | None = None,
+ max_combination_length: int | None = None,
+ null_value: float = 0.0,
+ ):
+ self._adata = adata
+ self._sample_rep = self._verify_sample_rep(sample_rep)
+ self._control_key = control_key
+ self._perturbation_covariates = self._verify_perturbation_covariates(perturbation_covariates)
+ self._perturbation_covariate_reps = self._verify_perturbation_covariate_reps(
+ adata,
+ perturbation_covariate_reps,
+ self._perturbation_covariates,
+ )
+ self._sample_covariates = self._verify_sample_covariates(sample_covariates)
+ self._sample_covariate_reps = self._verify_sample_covariate_reps(
+ adata, sample_covariate_reps, self._sample_covariates
+ )
+ self._split_covariates = self._verify_split_covariates(adata, split_covariates, control_key)
+ self._max_combination_length = self._get_max_combination_length(
+ self._perturbation_covariates, max_combination_length
+ )
+ self._null_value = null_value
+ self._primary_one_hot_encoder, self._is_categorical = self._get_primary_covar_encoder(
+ self._adata,
+ self._perturbation_covariates,
+ self._perturbation_covariate_reps,
+ )
+ self._linked_perturb_covars = self._get_linked_perturbation_covariates(self._perturbation_covariates)
+ sample_cov_groups = {covar: _to_list(covar) for covar in self._sample_covariates}
+ covariate_groups = self._perturbation_covariates | sample_cov_groups
+ self._covariate_reps = (self._perturbation_covariate_reps or {}) | (self._sample_covariate_reps or {})
+
+ self._covar_to_idx = self._get_covar_to_idx(covariate_groups) # type: ignore[arg-type]
+ perturb_covar_keys = _flatten_list(self._perturbation_covariates.values()) + list(self._sample_covariates)
+ perturb_covar_keys += [col for col in self._split_covariates if col not in perturb_covar_keys]
+ self._perturb_covar_keys = [k for k in perturb_covar_keys if k is not None]
+
+ # get the unique values of the split covariates
+ condition = np.unique(list(self.adata.obs['condition']))
+ unique_drug = []
+ np.array([unique_drug.extend(perturbation.split('+')) for perturbation in condition])
+ unique_drug = np.unique(unique_drug)
+ unique_drug.sort()
+ self.unique_drug = unique_drug
+ self.drug_dict = {drug: i for i, drug in enumerate(unique_drug)}
+
+ def get_train_data(self, adata: anndata.AnnData, one_hot: bool = False) -> Any:
+ """Get training data for the model.
+
+ Parameters
+ ----------
+ adata
+ An :class:`~anndata.AnnData` object.
+
+ Returns
+ -------
+ Training data for the model.
+ """
+ # create a
+ split_cov_combs = self._get_split_cov_combs(adata.obs)
+ cond_data = self._get_condition_data(split_cov_combs=split_cov_combs, adata=adata, one_hot=one_hot)
+ cell_data = self._get_cell_data(adata)
+ return TrainingData(
+ cell_data=cell_data,
+ split_covariates_mask=cond_data.split_covariates_mask,
+ split_idx_to_covariates=cond_data.split_idx_to_covariates,
+ perturbation_covariates_mask=cond_data.perturbation_covariates_mask,
+ perturbation_idx_to_covariates=cond_data.perturbation_idx_to_covariates,
+ perturbation_idx_to_id=cond_data.perturbation_idx_to_id,
+ condition_data=cond_data.condition_data,
+ control_to_perturbation=cond_data.control_to_perturbation,
+ max_combination_length=cond_data.max_combination_length,
+ null_value=self._null_value,
+ data_manager=self,
+ condition_data_id=cond_data.condition_data_id,
+ cell_data_id=adata.obs.index,
+ )
+
+ def get_validation_data(
+ self,
+ adata: anndata.AnnData,
+ n_conditions_on_log_iteration: int | None = None,
+ n_conditions_on_train_end: int | None = None,
+ one_hot: bool = False,
+ ) -> ValidationData:
+ """Get validation data for the model.
+
+ Parameters
+ ----------
+ adata
+ An :class:`~anndata.AnnData` object.
+ n_conditions_on_log_iteration
+ Number of conditions to validate on during logging.
+ n_conditions_on_train_end
+ Number of conditions to validate on at the end of training.
+
+ Returns
+ -------
+ Validation data for the model.
+ """
+ split_cov_combs = self._get_split_cov_combs(adata.obs)
+ cond_data = self._get_condition_data(split_cov_combs=split_cov_combs, adata=adata, one_hot=one_hot)
+ cell_data = self._get_cell_data(adata)
+ return ValidationData(
+ cell_data=cell_data,
+ split_covariates_mask=cond_data.split_covariates_mask,
+ split_idx_to_covariates=cond_data.split_idx_to_covariates,
+ perturbation_covariates_mask=cond_data.perturbation_covariates_mask,
+ perturbation_idx_to_covariates=cond_data.perturbation_idx_to_covariates,
+ perturbation_idx_to_id=cond_data.perturbation_idx_to_id,
+ condition_data=cond_data.condition_data,
+ control_to_perturbation=cond_data.control_to_perturbation,
+ max_combination_length=cond_data.max_combination_length,
+ null_value=self._null_value,
+ data_manager=self,
+ condition_data_id=cond_data.condition_data_id,
+ n_conditions_on_log_iteration=n_conditions_on_log_iteration,
+ n_conditions_on_train_end=n_conditions_on_train_end,
+ )
+
+ def get_prediction_data(
+ self,
+ adata: anndata.AnnData,
+ sample_rep: str,
+ covariate_data: pd.DataFrame,
+ rep_dict: dict[str, Any] | None = None,
+ condition_id_key: str | None = None,
+ one_hot: bool = False,
+ ) -> Any:
+ """Get predictions for control cells.
+
+ Extracts source distributions from ``'adata'`` and simulates cells perturbed
+ with covariates defined in ``'covariate_data'``.
+
+ Parameters
+ ----------
+ adata
+ An :class:`~anndata.AnnData` object to extract control cells from.
+ sample_rep
+ Key in :attr:`~anndata.AnnData.obsm` where the sample representation of the control
+ is stored or ``'X'`` to use :attr:`~anndata.AnnData.X`.
+ covariate_data
+ A :class:`~pandas.DataFrame` with columns defining the covariates as
+ in :meth:`cellflow.model.CellFlow.prepare_data` and stored in
+ :attr:`cellflow.model.CellFlow.data_manager`.
+ rep_dict
+ Dictionary with representations of the covariates.
+ If not provided, :attr:`~anndata.AnnData.uns` is used.
+ condition_id_key
+ Key in :class:`~pandas.DataFrame` that defines the condition names.
+
+ Returns
+ -------
+ Training data for the model.
+ """
+ self._verify_prediction_data(adata)
+ split_cov_combs = self._get_split_cov_combs(covariate_data=covariate_data)
+
+ # adata is None since we don't extract cell masks for predicted covariates
+ cond_data = self._get_condition_data(
+ split_cov_combs=split_cov_combs,
+ adata=None,
+ covariate_data=covariate_data,
+ rep_dict=adata.uns if rep_dict is None else rep_dict,
+ condition_id_key=condition_id_key,
+ one_hot=one_hot,
+ )
+
+ cell_data = self._get_cell_data(adata, sample_rep)
+ split_covariates_mask, split_idx_to_covariates = self._get_split_covariates_mask(
+ adata=adata, split_cov_combs=split_cov_combs
+ )
+
+ return PredictionData(
+ cell_data=cell_data,
+ split_covariates_mask=split_covariates_mask,
+ split_idx_to_covariates=split_idx_to_covariates,
+ condition_data=cond_data.condition_data,
+ control_to_perturbation=cond_data.control_to_perturbation,
+ perturbation_idx_to_covariates=cond_data.perturbation_idx_to_covariates,
+ perturbation_idx_to_id=cond_data.perturbation_idx_to_id,
+ max_combination_length=cond_data.max_combination_length,
+ null_value=self._null_value,
+ data_manager=self,
+ condition_data_id=cond_data.condition_data_id,
+ )
+
+ def get_condition_data(
+ self,
+ covariate_data: pd.DataFrame,
+ rep_dict: dict[str, Any] | None = None,
+ condition_id_key: str | None = None,
+ one_hot: bool = False,
+ ) -> ConditionData:
+ """Get condition data for the model.
+
+ Parameters
+ ----------
+ covariate_data
+ Dataframe with covariates.
+ condition_id_key
+ Key in ``covariate_data`` that defines the condition id.
+ rep_dict
+
+ Returns
+ -------
+ Condition data for the model.
+ """
+ self._verify_covariate_data(covariate_data, self._perturb_covar_keys)
+ split_cov_combs = self._get_split_cov_combs(covariate_data)
+ cond_data = self._get_condition_data(
+ split_cov_combs=split_cov_combs,
+ adata=None,
+ covariate_data=covariate_data,
+ rep_dict=rep_dict,
+ condition_id_key=condition_id_key,
+ one_hot=one_hot,
+ )
+ return ConditionData(
+ condition_data=cond_data.condition_data,
+ max_combination_length=cond_data.max_combination_length,
+ perturbation_idx_to_covariates=cond_data.perturbation_idx_to_covariates,
+ perturbation_idx_to_id=cond_data.perturbation_idx_to_id,
+ null_value=self._null_value,
+ data_manager=self,
+ condition_data_id=cond_data.condition_data_id,
+ )
+
+ def _get_split_cov_combs(self, covariate_data: pd.DataFrame) -> np.ndarray | list[list[Any]]: # type: ignore[type-arg]
+ if len(self._split_covariates) > 0:
+ return covariate_data[self._split_covariates].drop_duplicates().values
+ else:
+ return [[]]
+
+ def _get_condition_data(
+ self,
+ split_cov_combs: np.ndarray | list[list[Any]], # type: ignore[type-arg]
+ adata: anndata.AnnData | None,
+ covariate_data: pd.DataFrame | None = None,
+ rep_dict: dict[str, Any] | None = None,
+ condition_id_key: str | None = None,
+ one_hot: bool = False,
+ ) -> ReturnData:
+ # for prediction: adata is None, covariate_data is provided
+ # for training/validation: adata is provided and used to get cell masks, covariate_data is None
+ if adata is None and covariate_data is None:
+ raise ValueError("Either `adata` or `covariate_data` must be provided.")
+ covariate_data = covariate_data if covariate_data is not None else adata.obs # type: ignore[union-attr]
+ if rep_dict is None:
+ rep_dict = adata.uns if adata is not None else {}
+ # check if all perturbation/split covariates and control cells are present in the input
+ self._verify_covariate_data(
+ covariate_data,
+ {covar: _to_list(covar) for covar in self._sample_covariates},
+ )
+ self._verify_control_data(adata)
+ self._verify_covariate_data(covariate_data, _to_list(self._split_covariates))
+
+ # extract unique combinations of perturbation covariates
+ if condition_id_key is not None:
+ self._verify_condition_id_key(covariate_data, condition_id_key)
+ select_keys = self._perturb_covar_keys + [condition_id_key]
+ else:
+ select_keys = self._perturb_covar_keys
+ perturb_covar_df = covariate_data[select_keys].drop_duplicates()
+ if condition_id_key is not None:
+ perturb_covar_df = perturb_covar_df.set_index(condition_id_key)
+ else:
+ perturb_covar_df = perturb_covar_df.reset_index()
+
+ # get indices of cells belonging to each unique condition
+ _perturb_covar_df, _covariate_data = (
+ perturb_covar_df[self._perturb_covar_keys],
+ covariate_data[self._perturb_covar_keys],
+ )
+ _perturb_covar_df["row_id"] = range(len(perturb_covar_df))
+ _covariate_data["cell_index"] = _covariate_data.index
+ _perturb_covar_merged = _perturb_covar_df.merge(_covariate_data, on=self._perturb_covar_keys, how="inner")
+ perturb_covar_to_cells = _perturb_covar_merged.groupby("row_id")["cell_index"].apply(list).to_list()
+
+ # intialize data containers
+ if adata is not None:
+ split_covariates_mask = np.full((len(adata),), -1, dtype=jnp.int32)
+ perturbation_covariates_mask = np.full((len(adata),), -1, dtype=jnp.int32)
+ control_mask = covariate_data[self._control_key]
+ else:
+ split_covariates_mask = None
+ perturbation_covariates_mask = None
+ control_mask = jnp.ones((len(covariate_data),))
+
+ condition_data: dict[str, list[jnp.ndarray]] = {i: [] for i in self._covar_to_idx.keys()}
+
+ control_to_perturbation: dict[int, ArrayLike] = {}
+ split_idx_to_covariates: dict[int, tuple[Any]] = {}
+ perturbation_idx_to_covariates: dict[int, tuple[Any]] = {}
+ perturbation_idx_to_id: dict[int, Any] = {}
+ condition_data_id_list: list[jnp.ndarray] = []
+
+ src_counter = 0
+ tgt_counter = 0
+
+ # iterate over unique split covariate combinations
+ for split_combination in split_cov_combs:
+ # get masks for split covariates; for prediction, it's done outside this method
+ if adata is not None:
+ split_covariates_mask, split_idx_to_covariates, split_cov_mask = self._get_split_combination_mask(
+ covariate_data=adata.obs,
+ split_covariates_mask=split_covariates_mask, # type: ignore[arg-type]
+ split_combination=split_combination, # type: ignore[arg-type]
+ split_idx_to_covariates=split_idx_to_covariates,
+ control_mask=control_mask,
+ src_counter=src_counter,
+ )
+ conditional_distributions = []
+
+ # iterate over target conditions
+ filter_dict = dict(zip(self.split_covariates, split_combination, strict=False))
+ pc_df = perturb_covar_df[
+ (perturb_covar_df[list(filter_dict.keys())] == list(filter_dict.values())).all(axis=1)
+ ]
+ pbar = tqdm(pc_df.iterrows(), total=pc_df.shape[0])
+
+ for i, tgt_cond in pbar:
+ tgt_cond = tgt_cond[self._perturb_covar_keys]
+ # for train/validation, only extract covariate combinations that are present in adata
+ if adata is not None:
+ mask = covariate_data.index.isin(perturb_covar_to_cells[i])
+ mask *= (1 - control_mask) * split_cov_mask
+ mask = np.array(mask == 1)
+ if mask.sum() == 0:
+ continue
+ # map unique condition id to target id
+ perturbation_covariates_mask[mask] = tgt_counter # type: ignore[index]
+
+ # map target id to unique conditions and their ids
+ conditional_distributions.append(tgt_counter)
+ perturbation_idx_to_covariates[tgt_counter] = tgt_cond.values
+ if condition_id_key is not None:
+ perturbation_idx_to_id[tgt_counter] = i
+
+ # get embeddings for conditions
+ embedding = self._get_perturbation_covariates(
+ condition_data=tgt_cond,
+ rep_dict=rep_dict,
+ perturb_covariates={k: _to_list(v) for k, v in self._perturbation_covariates.items()},
+ )
+ for pert_cov, emb in embedding.items():
+ condition_data[pert_cov].append(emb)
+
+ condition_name, condition_id_emb = self._generate_condition_name_and_embedding(tgt_cond, tgt_counter, embedding, one_hot=one_hot)
+ condition_data_id_list.append(condition_id_emb)
+ tgt_counter += 1
+
+ # map source (control) to target condition ids
+ control_to_perturbation[src_counter] = np.array(conditional_distributions)
+ src_counter += 1
+ # convert outputs to jax arrays
+ for pert_cov, emb in condition_data.items():
+ condition_data[pert_cov] = jnp.array(emb)
+ condition_data_id = jnp.stack(condition_data_id_list, axis=0) # shape: (N, D)
+ split_covariates_mask = jnp.asarray(split_covariates_mask) if split_covariates_mask is not None else None
+ perturbation_covariates_mask = (
+ jnp.asarray(perturbation_covariates_mask) if perturbation_covariates_mask is not None else None
+ )
+ return ReturnData(
+ split_covariates_mask=split_covariates_mask,
+ split_idx_to_covariates=split_idx_to_covariates,
+ perturbation_covariates_mask=perturbation_covariates_mask,
+ perturbation_idx_to_covariates=perturbation_idx_to_covariates,
+ perturbation_idx_to_id=perturbation_idx_to_id,
+ condition_data=condition_data, # type: ignore[arg-type]
+ control_to_perturbation=control_to_perturbation,
+ max_combination_length=self._max_combination_length,
+ condition_data_id=condition_data_id,
+ )
+
+ def _generate_condition_name_and_embedding(
+ self,
+ tgt_cond: pd.Series,
+ tgt_counter: int,
+ embedding: dict[str, jnp.ndarray],
+ one_hot: bool = False,
+ ) -> tuple[str, jnp.ndarray]:
+ """
+ 基于 self.drug_dict 映射 tgt_cond 中的药物组合，生成唯一名称和 embedding。
+ """
+ # 展开所有 perturbation covariates（如 Drug1, Drug2）
+ drug_names = []
+ for covar_group in self._perturbation_covariates.values():
+ for col in covar_group:
+ val = tgt_cond[col]
+ if isinstance(val, str):
+ drug_names.extend(val.split('+'))
+ else:
+ drug_names.append(str(val))
+ drug_names = sorted(set(drug_names)) # 去重 + 排序
+ drug_ids = [self.drug_dict[name] for name in drug_names]
+
+ condition_name = 'drug_' + '_'.join(str(i) for i in drug_ids)
+ if one_hot:
+ one_hots = []
+ for idx in drug_ids:
+ vec = jnp.zeros((num_drugs,), dtype=jnp.float32).at[idx].set(1.0)
+ one_hots.append(vec)
+ condition_id_emb = jnp.concatenate(one_hots, axis=-1)
+ else:
+ condition_id_emb = jnp.array(drug_ids, dtype=jnp.int32)
+ return condition_name, condition_id_emb
+
+ @staticmethod
+ def _verify_condition_id_key(covariate_data: pd.DataFrame, condition_id_key: str | None) -> None:
+ if condition_id_key is not None and condition_id_key not in covariate_data.columns:
+ raise ValueError(f"The `condition_id_key` column ('{condition_id_key}') was not found in `adata.obs`.")
+ if not len(covariate_data[condition_id_key].unique()) == len(covariate_data):
+ raise ValueError(f"The `condition_id_key` column ('{condition_id_key}') must contain unique values.")
+
+ @staticmethod
+ def _verify_sample_rep(sample_rep: str | dict[str, str]) -> str | dict[str, str]:
+ if not (isinstance(sample_rep, str) or isinstance(sample_rep, dict)):
+ raise ValueError(
+ f"`sample_rep` should be of type `str` or `dict`, found {sample_rep} to be of type {type(sample_rep)}."
+ )
+ return sample_rep
+
+ def _get_cell_data(
+ self,
+ adata: anndata.AnnData,
+ sample_rep: str | None = None,
+ ) -> jax.Array:
+ sample_rep = self._sample_rep if sample_rep is None else sample_rep
+ if sample_rep == "X":
+ sample_rep = adata.X
+ if isinstance(sample_rep, sp.csr_matrix):
+ return jnp.asarray(sample_rep.toarray())
+ else:
+ return jnp.asarray(sample_rep)
+ if isinstance(self._sample_rep, str):
+ if self._sample_rep not in adata.obsm:
+ raise KeyError(f"Sample representation '{self._sample_rep}' not found in `adata.obsm`.")
+ return jnp.asarray(adata.obsm[self._sample_rep])
+ attr, key = next(iter(sample_rep.items())) # type: ignore[union-attr]
+ return jnp.asarray(getattr(adata, attr)[key])
+
+ def _verify_control_data(self, adata: anndata.AnnData | None) -> None:
+ if adata is None:
+ return None
+ if self._control_key not in adata.obs:
+ raise ValueError(f"Control column '{self._control_key}' not found in adata.obs.")
+ if not isinstance(adata.obs[self._control_key].dtype, pd.BooleanDtype):
+ try:
+ adata.obs[self._control_key] = adata.obs[self._control_key].astype("boolean")
+ except TypeError as e:
+ raise ValueError(f"Control column '{self._control_key}' could not be converted to boolean.") from e
+ if adata.obs[self._control_key].sum() == 0:
+ raise ValueError("No control cells found in adata.")
+
+ def _verify_prediction_data(self, adata: anndata.AnnData) -> None:
+ if self._control_key not in adata.obs:
+ raise ValueError(f"Control column '{self._control_key}' not found in adata.obs.")
+ if not isinstance(adata.obs[self._control_key].dtype, pd.BooleanDtype):
+ try:
+ adata.obs[self._control_key] = adata.obs[self._control_key].astype("boolean")
+ except ValueError as e:
+ raise ValueError(f"Control column '{self._control_key}' could not be converted to boolean.") from e
+ if not adata.obs[self._control_key].all():
+ raise ValueError(
+ f"For prediction, all cells in `adata` should be from control condition. Ensure that '{self._control_key}' is `True` for all cells, even if you're setting `.obs` to predicted condition."
+ )
+
+ def _get_split_combination_mask(
+ self,
+ covariate_data: pd.DataFrame,
+ split_covariates_mask: ArrayLike,
+ split_combination: ArrayLike,
+ split_idx_to_covariates: dict[int, tuple[Any]],
+ control_mask: ArrayLike,
+ src_counter: int,
+ ) -> tuple[ArrayLike, dict[int, tuple[Any]], ArrayLike]:
+ filter_dict = dict(zip(self.split_covariates, split_combination, strict=False))
+ split_cov_mask = (covariate_data[list(filter_dict.keys())] == list(filter_dict.values())).all(axis=1)
+ mask = jnp.array(control_mask * split_cov_mask).astype(bool)
+ split_covariates_mask[mask] = src_counter
+ split_idx_to_covariates[src_counter] = tuple(split_combination)
+ return split_covariates_mask, split_idx_to_covariates, split_cov_mask
+
+ def _get_split_covariates_mask(
+ self,
+ adata: anndata.AnnData,
+ split_cov_combs: np.ndarray | list[list[Any]], # type: ignore[type-arg]
+ ) -> tuple[ArrayLike, dict[int, tuple[Any]]]:
+ # here we assume that adata only contains source cells
+ if len(self.split_covariates) == 0:
+ return jnp.full((len(adata),), 0, dtype=jnp.int32), {}
+ split_covariates_mask = np.full((len(adata),), -1, dtype=jnp.int32)
+ split_idx_to_covariates: dict[int, Any] = {}
+ src_counter = 0
+ for split_combination in split_cov_combs:
+ split_covariates_mask_previous = split_covariates_mask.copy()
+ split_covariates_mask, split_idx_to_covariates, _ = self._get_split_combination_mask(
+ covariate_data=adata.obs,
+ split_covariates_mask=split_covariates_mask,
+ split_combination=split_combination, # type: ignore[arg-type]
+ split_idx_to_covariates=split_idx_to_covariates,
+ control_mask=jnp.ones((adata.n_obs,)),
+ src_counter=src_counter,
+ )
+ if (split_covariates_mask == split_covariates_mask_previous).all():
+ raise ValueError(f"No cells found in `adata` for split covariates {split_combination}.")
+ src_counter += 1
+ return jnp.asarray(split_covariates_mask), split_idx_to_covariates
+
+ @staticmethod
+ def _verify_perturbation_covariates(data: dict[str, Sequence[str]] | None) -> dict[str, list[str]]:
+ if data is None:
+ return {}
+ if not isinstance(data, dict):
+ raise ValueError(
+ f"`perturbation_covariates` should be a dictionary, found {data} to be of type {type(data)}."
+ )
+ if len(data) == 0:
+ raise ValueError("No perturbation covariates provided.")
+ for key, covars in data.items():
+ if not isinstance(key, str):
+ raise ValueError(f"Key should be a string, found {key} to be of type {type(key)}.")
+ if not isinstance(covars, tuple | list):
+ raise ValueError(f"Value should be a tuple, found {covars} to be of type {type(covars)}.")
+ if len(covars) == 0:
+ raise ValueError(f"No covariates provided for perturbation group {key}.")
+ lengths = [len(covs) for covs in data.values()]
+ if len(set(lengths)) != 1:
+ raise ValueError(f"Length of perturbation covariate groups must match, found lengths {lengths}.")
+ return {k: list(el) for k, el in data.items()}
+
+ @staticmethod
+ def _verify_sample_covariates(
+ sample_covariates: Sequence[str] | None,
+ ) -> list[str]:
+ if sample_covariates is None:
+ return []
+ if not isinstance(sample_covariates, tuple | list):
+ raise ValueError(
+ f"`sample_covariates` should be a tuple or list, found {sample_covariates} to be of type {type(sample_covariates)}."
+ )
+ for covar in sample_covariates:
+ if not isinstance(covar, str):
+ raise ValueError(f"Key should be a string, found {covar} to be of type {type(covar)}.")
+ return list(sample_covariates)
+
+ @staticmethod
+ def _verify_split_covariates(
+ adata: anndata.AnnData,
+ data: Sequence[str] | None,
+ control_key: str,
+ ) -> Sequence[str]:
+ if data is None:
+ return []
+ if not isinstance(data, tuple | list):
+ raise ValueError(f"`split_covariates` should be a tuple or list, found {data} to be of type {type(data)}.")
+ for covar in data:
+ if not isinstance(covar, str):
+ raise ValueError(f"Key should be a string, found {covar} to be of type {type(covar)}.")
+ source_splits = adata.obs[adata.obs[control_key]][data].drop_duplicates()
+ source_splits = map(tuple, source_splits.values)
+ target_splits = adata.obs[~adata.obs[control_key]][data].drop_duplicates()
+ target_splits = map(tuple, target_splits.values)
+ source_without_targets = set(source_splits) - set(target_splits)
+ if len(source_without_targets) > 0:
+ raise ValueError(
+ f"Source distribution with split covariate values {source_without_targets} do not have a corresponding target distribution."
+ )
+ return data
+
+ @staticmethod
+ def _verify_covariate_data(covariate_data: pd.DataFrame, covars) -> None:
+ for covariate in covars:
+ if covariate is not None and covariate not in covariate_data:
+ raise ValueError(f"Covariate {covariate} not found in adata.obs or covariate_data.")
+
+ @staticmethod
+ def _get_linked_perturbation_covariates(perturb_covariates: dict[str, list[str]]) -> dict[str, dict[Any, Any]]:
+ primary_group, primary_covars = next(iter(perturb_covariates.items()))
+ linked_perturb_covars: dict[str, dict[Any, Any]] = {k: {} for k in primary_covars}
+ for cov_group, covars in list(perturb_covariates.items())[1:]:
+ for primary_cov, linked_cov in zip(primary_covars, covars, strict=False):
+ linked_perturb_covars[primary_cov][cov_group] = linked_cov
+
+ return linked_perturb_covars
+
+ @staticmethod
+ def _verify_perturbation_covariate_reps(
+ adata: anndata.AnnData,
+ perturbation_covariate_reps: dict[str, str] | None,
+ perturbation_covariates: dict[str, list[str]],
+ ) -> dict[str, str]:
+ if perturbation_covariate_reps is None:
+ return {}
+ for key, value in perturbation_covariate_reps.items():
+ if key not in perturbation_covariates:
+ raise ValueError(f"Key '{key}' not found in covariates.")
+ if value not in adata.uns:
+ raise ValueError(f"Perturbation covariate representation '{value}' not found in `adata.uns`.")
+ if not isinstance(adata.uns[value], dict):
+ raise ValueError(
+ f"Perturbation covariate representation '{value}' in `adata.uns` should be of type `dict`, found {type(adata.uns[value])}."
+ )
+ return perturbation_covariate_reps
+
+ @staticmethod
+ def _verify_sample_covariate_reps(
+ adata: anndata.AnnData,
+ sample_covariate_reps: dict[str, str] | None,
+ covariates: list[str],
+ ) -> dict[str, str]:
+ if sample_covariate_reps is None:
+ return {}
+ for key, value in sample_covariate_reps.items():
+ if key not in covariates:
+ raise ValueError(f"Key '{key}' not found in covariates.")
+ if value not in adata.uns:
+ raise ValueError(f"Sample covariate representation '{value}' not found in `adata.uns`.")
+ if not isinstance(adata.uns[value], dict):
+ raise ValueError(
+ f"Sample covariate representation '{value}' in `adata.uns` should be of type `dict`, found {type(adata.uns[value])}."
+ )
+ return sample_covariate_reps
+
+ @staticmethod
+ def _get_max_combination_length(
+ perturbation_covariates: dict[str, list[str]],
+ max_combination_length: int | None,
+ ) -> int:
+ obs_max_combination_length = max(len(comb) for comb in perturbation_covariates.values())
+ if max_combination_length is None:
+ return obs_max_combination_length
+ elif max_combination_length < obs_max_combination_length:
+ logger.warning(
+ f"Provided `max_combination_length` is smaller than the observed maximum combination length of the perturbation covariates. Setting maximum combination length to {obs_max_combination_length}.",
+ stacklevel=2,
+ )
+ return obs_max_combination_length
+ else:
+ return max_combination_length
+
+ def _get_primary_covar_encoder(
+ self,
+ adata: anndata.AnnData,
+ perturbation_covariates: dict[str, list[str]],
+ perturbation_covariate_reps: dict[str, str],
+ ) -> tuple[preprocessing.OneHotEncoder | None, bool]:
+ primary_group, primary_covars = next(iter(perturbation_covariates.items()))
+ is_categorical = self._check_covariate_type(adata, primary_covars)
+ if perturbation_covariate_reps and primary_group in perturbation_covariate_reps:
+ return None, is_categorical
+ if is_categorical:
+ encoder = preprocessing.OneHotEncoder(sparse_output=False)
+ all_values = np.unique(adata.obs[primary_covars].values.flatten())
+ encoder.fit(all_values.reshape(-1, 1))
+ return encoder, is_categorical
+ encoder = preprocessing.OneHotEncoder(sparse_output=False)
+ encoder.fit(np.array(primary_covars).reshape(-1, 1))
+ return encoder, is_categorical
+
+ @staticmethod
+ def _check_covariate_type(adata: anndata.AnnData, covars: Sequence[str]) -> bool:
+ col_is_cat = []
+ for covariate in covars:
+ if is_numeric_dtype(adata.obs[covariate]):
+ col_is_cat.append(False)
+ continue
+ if adata.obs[covariate].isin(["True", "False", True, False]).all():
+ adata.obs[covariate] = adata.obs[covariate].astype(int)
+ col_is_cat.append(False)
+ continue
+ try:
+ adata.obs[covariate] = adata.obs[covariate].astype("category")
+ col_is_cat.append(True)
+ except ValueError as e:
+ raise ValueError(
+ f"Perturbation covariates `{covariate}` should be either numeric/boolean or categorical."
+ ) from e
+
+ if max(col_is_cat) != min(col_is_cat):
+ raise ValueError(
+ f"Groups of perturbation covariates `{covariate}` should be either all numeric/boolean or all categorical."
+ )
+
+ return max(col_is_cat)
+
+ @staticmethod
+ def _verify_covariate_type(covariate_data: pd.DataFrame, covars: Sequence[str], categorical: bool) -> None:
+ for covariate in covars:
+ if is_numeric_dtype(covariate_data[covariate]):
+ if categorical:
+ raise ValueError(f"Perturbation covariates `{covariate}` should be categorical, found numeric.")
+ continue
+ if covariate_data[covariate].isin(["True", "False", True, False]).all():
+ if categorical:
+ raise ValueError(f"Perturbation covariates `{covariate}` should be categorical, found boolean.")
+ continue
+ try:
+ covariate_data[covariate] = covariate_data[covariate].astype("category")
+ except ValueError as e:
+ raise ValueError(
+ f"Perturbation covariates `{covariate}` should be either numeric/boolean or categorical."
+ ) from e
+ else:
+ if not categorical:
+ raise ValueError(
+ f"Perturbation covariates `{covariate}` should be numeric/boolean, found categorical."
+ )
+
+ @staticmethod
+ def _check_shape(arr: float | ArrayLike) -> ArrayLike:
+ if not hasattr(arr, "shape") or len(arr.shape) == 0:
+ return jnp.ones((1, 1)) * arr
+ if arr.ndim == 1: # type: ignore[union-attr]
+ return jnp.expand_dims(arr, 0)
+ elif arr.ndim == 2: # type: ignore[union-attr]
+ if arr.shape[0] == 1:
+ return arr # type: ignore[return-value]
+ if arr.shape[1] == 1:
+ return jnp.transpose(arr)
+ raise ValueError(
+ "Condition representation has an unexpected shape. Should be (1, n_features) or (n_features, )."
+ )
+ elif arr.ndim > 2: # type: ignore[union-attr]
+ raise ValueError("Condition representation has too many dimensions. Should be 1 or 2.")
+
+ raise ValueError(
+ "Condition representation as an unexpected format. Expected an array of shape (1, n_features) or (n_features, )."
+ )
+
+ @staticmethod
+ def _get_covar_to_idx(covariate_groups: dict[str, Sequence[str]]) -> dict[str, int]:
+ idx_to_covar = {}
+ for idx, cov_group in enumerate(covariate_groups):
+ idx_to_covar[idx] = cov_group
+ covar_to_idx = {v: k for k, v in idx_to_covar.items()}
+ return covar_to_idx
+
+ @staticmethod
+ def _pad_to_max_length(arr: jax.Array, max_combination_length: int, null_value: Any) -> jax.Array:
+ if arr.shape[0] < max_combination_length:
+ null_arr = jnp.full((max_combination_length - arr.shape[0], arr.shape[1]), null_value)
+ arr = jnp.concatenate([arr, null_arr], axis=0)
+ return arr
+
+ def _get_perturbation_covariates(
+ self,
+ condition_data: pd.DataFrame,
+ rep_dict: dict[str, dict[str, ArrayLike]],
+ perturb_covariates: Any, # TODO: check if we can save as attribtue
+ ) -> dict[str, jax.Array]:
+ primary_group, primary_covars = next(iter(perturb_covariates.items()))
+
+ perturb_covar_emb: dict[str, list[jax.Array]] = {group: [] for group in perturb_covariates}
+ for primary_cov in primary_covars:
+ value = condition_data[primary_cov]
+ cov_name = value if self.is_categorical else primary_cov
+ if primary_group in self._covariate_reps:
+ rep_key = self._covariate_reps[primary_group]
+ try:
+ prim_arr = rep_dict[rep_key][cov_name]
+ except KeyError as err:
+ raise ValueError(
+ f"Representation for '{cov_name}' not found in `adata.uns['{rep_key}']` or `rep_dict`."
+ ) from err
+ prim_arr = jnp.asarray(prim_arr)
+ else:
+ prim_arr = jnp.asarray(
+ self.primary_one_hot_encoder.transform( # type: ignore[union-attr]
+ np.array(cov_name).reshape(-1, 1)
+ )
+ )
+
+ if not self.is_categorical:
+ prim_arr *= value
+
+ prim_arr = self._check_shape(prim_arr)
+ perturb_covar_emb[primary_group].append(prim_arr)
+
+ for linked_covar in self._linked_perturb_covars[primary_cov].items():
+ linked_group, linked_cov = list(linked_covar)
+
+ if linked_cov is None:
+ linked_arr = jnp.full((1, 1), self._null_value)
+ linked_arr = self._check_shape(linked_arr)
+ perturb_covar_emb[linked_group].append(linked_arr)
+ continue
+
+ cov_name = condition_data[linked_cov]
+
+ if linked_group in self._covariate_reps:
+ rep_key = self._covariate_reps[linked_group]
+ if cov_name not in rep_dict[rep_key]:
+ raise ValueError(f"Representation for '{cov_name}' not found in `adata.uns['{rep_key}']`.")
+ linked_arr = jnp.asarray(rep_dict[rep_key][cov_name])
+ else:
+ linked_arr = jnp.asarray(condition_data[linked_cov])
+
+ linked_arr = self._check_shape(linked_arr)
+ perturb_covar_emb[linked_group].append(linked_arr)
+
+ perturb_covar_emb = {
+ k: self._pad_to_max_length(
+ jnp.concatenate(v, axis=0),
+ self._max_combination_length,
+ self._null_value,
+ )
+ for k, v in perturb_covar_emb.items()
+ }
+
+ sample_covar_emb: dict[str, jax.Array] = {}
+ for sample_cov in self._sample_covariates:
+ value = condition_data[sample_cov]
+ if sample_cov in self._covariate_reps:
+ rep_key = self._covariate_reps[sample_cov]
+ if value not in rep_dict[rep_key]:
+ raise ValueError(f"Representation for '{value}' not found in `adata.uns['{sample_cov}']`.")
+ cov_arr = jnp.asarray(rep_dict[rep_key][value])
+ else:
+ cov_arr = jnp.asarray(value)
+
+ cov_arr = self._check_shape(cov_arr)
+ sample_covar_emb[sample_cov] = jnp.tile(cov_arr, (self._max_combination_length, 1))
+
+ return perturb_covar_emb | sample_covar_emb
+
+ @property
+ def is_categorical(self) -> bool:
+ """Whether the primary covariate is categorical."""
+ return self._is_categorical
+
+ @property
+ def adata(self) -> anndata.AnnData:
+ """An :class:`~anndata.AnnData` object used for instantiating the DataManager."""
+ return self._adata
+
+ @property
+ def control_key(self) -> str:
+ """Boolean key in :attr:`~anndata.AnnData.obs` indicating whether belongs to control group."""
+ return self._control_key
+
+ @property
+ def perturbation_covariates(self) -> dict[str, list[str]]:
+ """Dictionary with keys indicating the name of the covariate group and values are keys in :attr:`~anndata.AnnData.obs` which together define the covariates."""
+ return self._perturbation_covariates
+
+ @property
+ def perturbation_covariate_reps(self) -> dict[str, str]:
+ """Dictionary with keys indicating the name of the covariate group and values are keys in :attr:`~anndata.AnnData.uns` storing a dictionary with the representation of the covariates."""
+ return self._perturbation_covariate_reps
+
+ @property
+ def sample_covariates(self) -> Sequence[str]:
+ """Keys in :attr:`~anndata.AnnData.obs` indicating which sample the cell belongs to (e.g. cell line)."""
+ return self._sample_covariates
+
+ @property
+ def sample_covariate_reps(self) -> dict[str, str]:
+ """Dictionary with keys indicating the name of the sample covariate group and values are keys in :attr:`~anndata.AnnData.uns` storing a dictionary with the representation of the sample covariates."""
+ return self._sample_covariate_reps
+
+ @property
+ def split_covariates(self) -> Sequence[str]:
+ """Covariates in :attr:`~anndata.AnnData.obs` to split all control cells into different control populations."""
+ return self._split_covariates
+
+ @property
+ def max_combination_length(self) -> int:
+ """Maximum combination length of perturbation covariates."""
+ return self._max_combination_length
+
+ @property
+ def null_value(self) -> float:
+ """Value to use for padding to :attr:`~max_combination_length`."""
+ return self._null_value
+
+ @property
+ def primary_one_hot_encoder(self) -> preprocessing.OneHotEncoder | None:
+ """One-hot encoder for the primary covariate."""
+ return self._primary_one_hot_encoder
+
+ @property
+ def linked_perturb_covars(self) -> dict[str, dict[Any, Any]]:
+ """Dictionary with keys indicating the name of the primary covariate and values are dictionaries with keys indicating the name of the linked covariate group and values are the linked covariates."""
+ return self._linked_perturb_covars
+
+ @property
+ def covariate_reps(self) -> dict[str, str]:
+ """Dictionary which stores representation of covariates, i.e. the union of ``sample_covariate_reps`` and ``perturbation_covariate_reps``."""
+ return self._covariate_reps
+
+ @property
+ def covar_to_idx(self) -> dict[str, int]:
+ """TODO: add description"""
+ return self._covar_to_idx
+
+ @property
+ def perturb_covar_keys(self) -> list[str]:
+ """List of all perturbation covariates."""
+ return self._perturb_covar_keys
+
+ @property
+ def sample_rep(self) -> str | dict[str, str]:
+ """Key of the sample representation."""
+ return self._sample_rep
diff --git a/scDFM/src/data_process/data.py b/scDFM/src/data_process/data.py
new file mode 100644
index 0000000000000000000000000000000000000000..7851cdcfc8d6b994b2041c5b694153d67167a844
--- /dev/null
+++ b/scDFM/src/data_process/data.py
@@ -0,0 +1,592 @@
+import scanpy as sc
+import pandas as pd
+import numpy as np
+import json
+import torch
+import pickle
+from typing import Union, Optional
+from pathlib import Path
+import os
+from src.utils._preprocessing import annotate_compounds, get_molecular_fingerprints
+from src.data_process._datamanager import DataManager
+import jax
+import torch.nn as nn
+from torch.utils.data import Dataset, DataLoader
+import pdb
+import tqdm
+from random import shuffle
+from src.utils.utils import build_gene_coexpression_graph,sorted_pad_mask
+# combosciplex url: https://figshare.com/articles/dataset/combosciplex/25062230?file=44229635
+# 'norman' url = 'https://dataverse.harvard.edu/api/access/datafile/6154020'
+
+class Data:
+ def __init__(self, data_path='../../data'):
+ self.data_path = data_path
+ if not os.path.exists(data_path):
+ raise ValueError(data_path + ' does not exist')
+ # os.makedirs(data_path)
+
+
+ def load_data(self, data_name = None, data_path = None):
+ self.data_name = data_name
+ if data_name in ['norman', 'norman_umi_go_filtered',]:
+ self.adata = sc.read_h5ad(os.path.join(self.data_path, data_name + '.h5ad'))
+ elif data_name in ['combosciplex', ]:
+ self.adata = sc.read_h5ad(os.path.join(self.data_path, data_name + '.h5ad'))
+ else:
+ raise ValueError(data_name + ' is not a valid data name')
+
+ def process_data(self, n_top_genes = 2000,infer_top_gene=1000,split_method='additive',
+ use_negative_edge=True, k=30,
+ **kwargs):
+ os.makedirs(os.path.join(self.data_path, self.data_name), exist_ok=True)
+ if self.data_name == 'combosciplex':
+
+ if os.path.exists(os.path.join(self.data_path, self.data_name, 'processed.h5ad')):
+ self.adata = sc.read_h5ad(os.path.join(self.data_path, self.data_name, 'processed.h5ad'))
+ else:
+
+ self.adata.obs["condition"] = self.adata.obs.apply(
+ lambda x: "control" if x["condition"] == "control+control" else x["condition"], axis=1
+ )
+
+ self.adata.obs["is_control"] = self.adata.obs.apply(
+ lambda x: True if x["condition"] == "control" else False, axis=1
+ )
+
+ annotate_compounds(self.adata, compound_keys=["Drug1", "Drug2"])
+ get_molecular_fingerprints(self.adata, compound_keys=["Drug1", "Drug2"])
+ self.adata.uns["fingerprints"]["control"] = np.zeros(1024)
+
+ self.adata.write(os.path.join(self.data_path, self.data_name, 'processed.h5ad'))
+
+ self.adata.X = self.adata.layers["counts"].copy()
+ sc.pp.normalize_total(self.adata)
+ sc.pp.log1p(self.adata)
+ sc.pp.highly_variable_genes(self.adata, inplace=True, n_top_genes=n_top_genes)
+
+ if 'test_conditions' in kwargs.keys():
+ test_conditions = kwargs['test_conditions']
+ else:
+ test_conditions = ['Panobinostat+Crizotinib',
+ 'Panobinostat+Curcumin',
+ 'Panobinostat+SRT1720',
+ 'Panobinostat+Sorafenib',
+ 'SRT2104+Alvespimycin',
+ 'control+Alvespimycin',
+ 'control+Dacinostat']
+
+ self.adata = self.adata[:,self.adata.var['highly_variable']] # filter out low variable genes
+
+ self.adata.obs["mode"] = self.adata.obs.apply(lambda x: "test" if x["condition"] in test_conditions else "train", axis=1)
+ self.adata_train = self.adata[self.adata.obs["mode"] == "train"]
+ self.adata_test = self.adata[(self.adata.obs["mode"] == "test") | (self.adata.obs["condition"]=="control")]
+
+ sc.pp.highly_variable_genes(self.adata_test, inplace=True, n_top_genes=infer_top_gene)
+ self.adata_test = self.adata_test[:,self.adata_test.var['highly_variable']]
+
+ condition = np.unique(list(self.adata.obs['condition']))
+ unique_perturbation = []
+ np.array([unique_perturbation.extend(perturbation.split('+')) for perturbation in condition])
+ unique_perturbation = np.unique(unique_perturbation)
+ unique_perturbation.sort()
+ self.unique_perturbation = unique_perturbation
+ self.perturbation_dict = {perturbation: i for i, perturbation in enumerate(unique_perturbation)}
+ # self._val_manager =
+ elif self.data_name == 'norman' or self.data_name == 'norman_umi_go_filtered':
+
+ sc.pp.highly_variable_genes(self.adata, inplace=True, n_top_genes=n_top_genes)
+ unique_perturbation = []
+ [unique_perturbation.extend(perturbation.split('+')) for perturbation in self.adata.obs['condition'].unique()]
+ unique_perturbation = np.unique(unique_perturbation)
+
+ if self.data_name == 'norman':
+ for perturbation in unique_perturbation:
+ if perturbation in self.adata.var_names:
+ self.adata.var.loc[perturbation, 'highly_variable'] = True
+ else:
+ print(f"Warning: {perturbation} is not in the gene names")
+ self.adata = self.adata[:,self.adata.var['highly_variable']]
+ elif self.data_name == 'norman_umi_go_filtered':
+ all_gene_names = list(self.adata.var['gene_name']) + ['ctrl']
+ for perturbation in unique_perturbation:
+ if perturbation not in all_gene_names:
+ print(f"Warning: {perturbation} is not in the gene names")
+ self.adata.var['highly_variable'] = True
+
+
+ #### for split five times
+ if split_method == 'additive' or 'combinations':
+ split_file = os.path.join(self.data_path, self.data_name, 'split_results.pkl')
+ if os.path.exists(split_file):
+ with open(split_file, 'rb') as f:
+ self.split_results = pickle.load(f)
+ else:
+ perturbations = np.unique(self.adata.obs['condition'])
+ double_perturbation = [p for p in perturbations if 'ctrl' not in p]
+ double_perturbation = np.array(double_perturbation)
+
+ self.split_results = []
+
+ for i in range(5):
+ np.random.seed(42 + i)
+ shuffled = double_perturbation.copy()
+ np.random.shuffle(shuffled)
+
+ split_idx = int(len(shuffled) * 0.3)
+ test_double = shuffled[:split_idx]
+ train_double = shuffled[split_idx:]
+ self.split_results.append({
+ 'train': train_double.tolist(),
+ 'test': test_double.tolist()
+ })
+
+ with open(split_file, 'wb') as f:
+ pickle.dump(self.split_results, f)
+ print('split results saved')
+
+ elif split_method == 'unseen':
+ split_file = os.path.join(self.data_path, self.data_name, 'split_results_unseen.pkl')
+ if os.path.exists(split_file):
+ with open(split_file, 'rb') as f:
+ self.split_results = pickle.load(f)
+ else:
+ self.split_results = []
+ for i in range(5):
+ perturbations = np.unique(self.adata.obs['condition'])
+ double_perturbation = [p for p in perturbations if 'ctrl' not in p]
+ single = []
+ [single.extend(p.split('+')) for p in double_perturbation]
+ single = list(set(single))
+
+ shuffle(single)
+ remove_genes = single[:12]
+ p_count = {}
+ for p in double_perturbation:
+ ps = p.split('+')
+ count = int(ps[0] in remove_genes) + int(ps[1] in remove_genes)
+ p_count[p] = count
+ double_perturbation = [p for p, count in p_count.items() if count > 0]
+ double_perturbation = list(double_perturbation)
+ remove_genes_condition = [p+'+control' for p in remove_genes]
+ double_perturbation.extend(remove_genes_condition)
+ self.split_results.append({
+ 'p_count': p_count,
+ 'test': double_perturbation
+ })
+ with open(split_file, 'wb') as f:
+ pickle.dump(self.split_results, f)
+ print('split results unseen saved')
+
+ if 'fold' in kwargs.keys():
+ fold = kwargs['fold']
+ else:
+ fold = 0
+ self.adata.obs['condition'] = self.adata.obs['condition'].str.replace('ctrl', 'control')
+ self.adata.obs['Drug1'] = self.adata.obs['condition'].str.split('+').apply(lambda x: x[0])
+ self.adata.obs['Drug2'] = self.adata.obs['condition'].str.split('+').apply(lambda x: x[-1])
+ self.adata.obs['is_control'] = False
+ self.adata.obs.loc[self.adata.obs['control'] == 1, 'is_control'] = True
+ self.adata.obs['mode'] = 'train'
+
+
+ if split_method == 'combinations':
+ self.split_results[fold]['test'] = self.split_results[fold]['test'][:15]
+ remove_genes = []
+ [remove_genes.extend(p.split('+')) for p in self.split_results[fold]['test']]
+ remove_genes = set(remove_genes)
+ remove_genes_condition = [p+'+control' for p in remove_genes]
+
+ self.split_results[fold]['test'].extend(remove_genes_condition)
+
+
+ self.adata.obs.loc[self.adata.obs['condition'].isin(self.split_results[fold]['test']), 'mode'] = 'test'
+
+ self.adata_train = self.adata[self.adata.obs['mode'] == 'train']
+ self.adata_test = self.adata[(self.adata.obs['mode'] == 'test') | (self.adata.obs['control'] == 1)]
+
+
+
+ sc.pp.highly_variable_genes(self.adata_test, inplace=True, n_top_genes=infer_top_gene)
+ self.adata_test = self.adata_test[:,self.adata_test.var['highly_variable']]
+
+ condition = np.unique(list(self.adata.obs['condition']))
+ unique_perturbation = []
+ np.array([unique_perturbation.extend(perturbation.split('+')) for perturbation in condition])
+ unique_perturbation = np.unique(unique_perturbation)
+ unique_perturbation.sort()
+ self.unique_perturbation = unique_perturbation
+ self.perturbation_dict = {perturbation: i for i, perturbation in enumerate(unique_perturbation)}
+
+ else:
+ raise ValueError(self.data_name + ' is not a valid data name')
+
+ if 'fold' in kwargs.keys():
+ fold = kwargs['fold']
+ else:
+ fold = 0
+ if use_negative_edge:
+ mask_path = os.path.join(self.data_path, self.data_name,'mask_fold_'+str(fold)+'topk_'+str(k)+split_method+'_negative_edge'+'.pt')
+ else:
+ mask_path = os.path.join(self.data_path, self.data_name,'mask_fold_'+str(fold)+'topk_'+str(k)+split_method+'.pt')
+ if os.path.exists(mask_path):
+ self.mask = torch.load(mask_path)
+ else:
+ X = self.adata_train.X.toarray()
+ mask = build_gene_coexpression_graph(X,
+ method="pearson",
+ wgcna_beta=None,
+ sparsify="topk",
+ k=k,
+ use_negative_edge=use_negative_edge)
+ mask = sorted_pad_mask(mask, pad_size=4, gene_names=list(self.adata_train.var_names))
+ torch.save(mask, mask_path)
+ print('mask saved')
+
+
+ def load_flow_data(self, batch_size = 128):
+ if self.data_name == 'combosciplex':
+ train_sampler = TrainSampler(self.data_name, self.adata_train, ["Drug1", "Drug2"], self.perturbation_dict)
+ test_sampler = TestDataset(self.data_name, self.adata_test, ["Drug1", "Drug2"], self.perturbation_dict)
+
+ return train_sampler , test_sampler, []
+ elif self.data_name == 'norman' or self.data_name == 'norman_umi_go_filtered':
+ train_sampler = TrainSampler(self.data_name, self.adata_train, ["Drug1", "Drug2"], self.perturbation_dict)
+ test_sampler = TestDataset(self.data_name, self.adata_test, ["Drug1", "Drug2"], self.perturbation_dict)
+ return train_sampler , test_sampler, []
+ else:
+ raise ValueError(self.data_name + ' is not a valid data name')
+
+
+ def pretrain_data(self, batch_size = 128):
+ if self.data_name == 'combosciplex':
+
+ self.pretrain_train_data = PretrainData(self.adata_train, self.perturbation_dict)
+ self.pretrain_train_data_loader = DataLoader(self.pretrain_train_data, batch_size=batch_size, shuffle=True, pin_memory=True, num_workers=4)
+ self.pretrain_test_data = PretrainData(self.adata_test, self.perturbation_dict)
+ self.pretrain_test_data_loader = DataLoader(self.pretrain_test_data, batch_size=batch_size, shuffle=False, pin_memory=True, num_workers=4)
+ return self.pretrain_train_data_loader, self.pretrain_test_data_loader
+ else:
+ raise ValueError(self.data_name + ' is not a valid data name')
+
+class TrainSampler:
+ def __init__(self, data_name, adata: sc.AnnData, perturbation_covariates: list[str], perturbation_dict: dict,):
+ self.data_name = data_name
+ self.adata = adata
+ self.perturbation_covariates = perturbation_covariates
+ self.adata.obs['perturbation_covariates'] = self.adata.obs[perturbation_covariates].apply(lambda x: '+'.join(x), axis=1)
+ self._perturbation_covariates = adata.obs['perturbation_covariates'].unique()
+
+ self._perturbation_covariates = self._perturbation_covariates[self._perturbation_covariates != 'control+control']
+
+ self._perturbation_covariates.sort()
+ self.perturbation_covariates_dict = {perturbation: i for i, perturbation in enumerate(self._perturbation_covariates)}
+
+ perturbation_covariates_id = [adata.obs[perturbation_covariates[i]].apply(lambda x: perturbation_dict[x])
+ for i in range(len(perturbation_covariates))]
+ self.perturbation_covariates_id = np.array(perturbation_covariates_id).T
+
+
+ self.cells_name = self.adata.obs_names
+
+
+ def get_batch(self, batch_size: int, same_perturbation: bool = True):
+ if same_perturbation:
+ # random sample a perturbation from self._perturbation_covariates.
+ # the last one is control
+ perturbation_idx = np.random.choice(len(self._perturbation_covariates), 1)[0]
+
+ perturbation_id = self._perturbation_covariates[perturbation_idx]
+
+ # get the target data
+ tgt_idx = (self.adata.obs['perturbation_covariates'] == perturbation_id).to_numpy().nonzero()[0]
+ tgt_batch_idx = np.random.choice(tgt_idx, batch_size)
+
+ tgt_batch = torch.from_numpy(self.adata.X[tgt_batch_idx].toarray())
+
+ # get data from control
+ src_idx = (self.adata.obs['perturbation_covariates'] == 'control+control').to_numpy().nonzero()[0]
+ src_batch_idx = np.random.choice(src_idx, batch_size)
+
+ src_batch = torch.from_numpy(self.adata.X[src_batch_idx].toarray())
+
+ return {
+ 'src_cell_data': src_batch,
+ 'tgt_cell_data': tgt_batch,
+ 'src_cell_id': self.cells_name[src_batch_idx],
+ 'tgt_cell_id': self.cells_name[tgt_batch_idx],
+ 'condition_id': self.perturbation_covariates_id[tgt_batch_idx],
+ }
+
+ else:
+ raise ValueError('same_perturbation must be True')
+
+class TestDataset:
+ def __init__(self, data_name,adata: sc.AnnData, perturbation_covariates: list[str], perturbation_dict: dict,):
+ self.data_name = data_name
+ self.adata = adata
+ self.perturbation_covariates = perturbation_covariates
+ self.adata.obs['perturbation_covariates'] = self.adata.obs[perturbation_covariates].apply(lambda x: '+'.join(x), axis=1)
+ self._perturbation_covariates = adata.obs['perturbation_covariates'].unique()
+
+ self._perturbation_covariates = self._perturbation_covariates[self._perturbation_covariates != 'control+control']
+
+ self._perturbation_covariates.sort()
+ self.perturbation_covariates_dict = {perturbation: i for i, perturbation in enumerate(self._perturbation_covariates)}
+
+ perturbation_covariates_id = [adata.obs[perturbation_covariates[i]].apply(lambda x: perturbation_dict[x])
+ for i in range(len(perturbation_covariates))]
+ self.perturbation_covariates_id = np.array(perturbation_covariates_id).T
+
+
+ self.cells_name = self.adata.obs_names
+
+ def get_control_data(self,):
+ control_data = self.adata[self.adata.obs['is_control']]
+ return {
+ 'src_cell_data': torch.from_numpy(control_data.X.toarray()),
+ 'src_cell_id': control_data.obs_names,
+ 'condition_id': torch.tensor(self.perturbation_covariates_id[self.adata.obs['is_control']]),
+ }
+
+ def get_perturbation_data(self, perturbation: str):
+ perturbation_data = self.adata[self.adata.obs['perturbation_covariates'] == perturbation]
+ return {
+ 'tgt_cell_data': torch.from_numpy(perturbation_data.X.toarray()),
+ 'tgt_cell_id': perturbation_data.obs_names,
+ 'condition_id': torch.tensor(self.perturbation_covariates_id[self.adata.obs['perturbation_covariates'] == perturbation]),
+ }
+
+
+
+class PerturbationDataset(Dataset):
+ def __init__(self, sampler: TrainSampler, batch_size: int):
+ self.sampler = sampler
+ self.batch_size = batch_size
+ self.perturbations = sampler._perturbation_covariates
+
+ self.control_idx = (sampler.adata.obs['perturbation_covariates'] == 'control+control').to_numpy().nonzero()[0]
+
+ def __len__(self):
+
+ return len(self.perturbations) * 1000
+
+ def __getitem__(self, idx):
+ # 随机选一个 perturbation
+ perturbation_idx = np.random.choice(len(self.perturbations), 1)[0]
+ perturbation_id = self.perturbations[perturbation_idx]
+
+ # target batch
+ tgt_idx = (self.sampler.adata.obs['perturbation_covariates'] == perturbation_id).to_numpy().nonzero()[0]
+ tgt_batch_idx = np.random.choice(tgt_idx, self.batch_size)
+
+ # source (control) batch
+ src_batch_idx = np.random.choice(self.control_idx, self.batch_size)
+ if hasattr(self.sampler.adata.X[src_batch_idx], "toarray"):
+ src_batch = torch.from_numpy(self.sampler.adata.X[src_batch_idx].toarray())
+ tgt_batch = torch.from_numpy(self.sampler.adata.X[tgt_batch_idx].toarray())
+ else:
+ src_batch = torch.from_numpy(self.sampler.adata.X[src_batch_idx])
+ tgt_batch = torch.from_numpy(self.sampler.adata.X[tgt_batch_idx])
+
+ return {
+ 'src_cell_data': src_batch,
+ 'tgt_cell_data': tgt_batch,
+ 'src_cell_id': list(self.sampler.cells_name[src_batch_idx]),
+ 'tgt_cell_id': list(self.sampler.cells_name[tgt_batch_idx]),
+ 'condition_id': torch.tensor(self.sampler.perturbation_covariates_id[tgt_batch_idx]),
+ }
+class BinDiscretizer:
+ """
+ data = np.random.exponential(scale=2.0, size=1000)
+
+ bd = BinDiscretizer(n_bins=200)
+ bd.fit(data)
+
+ bd.save_edges('./data/combosciplex/bin_discretizer_edges.pkl')
+
+ new_bd = BinDiscretizer(n_bins=200)
+ new_bd.load_edges('./data/combosciplex/bin_discretizer_edges.pkl')
+
+ binned = bd.transform(data)
+
+ recon = bd.inverse_transform(binned, random=False)
+
+ Note: 0 is treated as a separate class (class 0), and non-zero values are discretized into classes 1 to n_bins.
+ """
+ def __init__(self, n_bins: int, strategy: str = "quantile"):
+ self.n_bins = n_bins
+ self.strategy = strategy
+ self.edges = None # will be (n_bins + 1, ) array
+
+ def fit(self, data: Union[np.ndarray, torch.Tensor]):
+ if isinstance(data, torch.Tensor):
+ data = data.detach().cpu().numpy()
+ data = data.flatten()
+ data = data[data > 0] # exclude zeros from fitting
+
+ if len(data) == 0:
+ raise ValueError("No non-zero entries in data to fit.")
+
+ if self.strategy == "quantile":
+ self.edges = np.quantile(data, np.linspace(0, 1, self.n_bins + 1))
+ elif self.strategy == "uniform":
+ self.edges = np.linspace(data.min(), data.max(), self.n_bins + 1)
+ else:
+ raise ValueError(f"Unknown strategy {self.strategy}")
+
+ def transform(self, data: Union[np.ndarray, torch.Tensor]) -> Union[np.ndarray, torch.Tensor]:
+ if self.edges is None:
+ raise RuntimeError("Call fit() before transform().")
+
+ is_torch = isinstance(data, torch.Tensor)
+ if is_torch:
+ orig_dtype = data.dtype
+ data = data.detach().cpu().numpy()
+
+ out = np.zeros_like(data, dtype=np.int64)
+ mask = data > 0
+
+ # Digitize non-zero entries: 0 remains 0, non-zero values get classes 1 to n_bins
+ if np.any(mask):
+ # np.digitize returns 0-based indices for the bins
+ # We want to map these to 1-based class indices
+ digitized = np.digitize(data[mask], self.edges[1:-1])
+ # Convert 0-based bin indices to 1-based class indices
+ # digitized=0 means it's in the first bin, which should be class 1
+ # digitized=1 means it's in the second bin, which should be class 2, etc.
+ out[mask] = digitized + 1
+
+ if is_torch:
+ return torch.from_numpy(out).to(dtype=torch.int64)
+ return out
+
+ def inverse_transform(self, digitized: Union[np.ndarray, torch.Tensor], random: bool = False) -> Union[np.ndarray, torch.Tensor]:
+ if self.edges is None:
+ raise RuntimeError("Call fit() before inverse_transform().")
+
+ is_torch = isinstance(digitized, torch.Tensor)
+ if is_torch:
+ orig_dtype = digitized.dtype
+ digitized = digitized.detach().cpu().numpy()
+
+ out = np.zeros_like(digitized, dtype=np.float64)
+ mask = digitized > 0
+
+ if np.any(mask):
+ ids = digitized[mask]
+ # Ensure ids are within valid range (1 to n_bins)
+ ids = np.clip(ids, 1, self.n_bins)
+ # Convert 1-based class indices back to 0-based bin indices
+ bin_ids = ids - 1
+ lefts = self.edges[bin_ids]
+ rights = self.edges[bin_ids + 1]
+
+ if random:
+ out[mask] = np.random.uniform(lefts, rights)
+ else:
+ out[mask] = (lefts + rights) / 2
+
+ if is_torch:
+ return torch.from_numpy(out).to(dtype=torch.float64)
+ return out
+
+ def save_edges(self, filepath: Union[str, Path]):
+ """Save edges to a file"""
+ if self.edges is None:
+ raise RuntimeError("No edges to save. Call fit() first.")
+
+ filepath = Path(filepath)
+ with open(filepath, 'wb') as f:
+ pickle.dump({'edges': self.edges, 'n_bins': self.n_bins}, f)
+
+ def load_edges(self, filepath: Union[str, Path]):
+ """Load edges from a file"""
+ filepath = Path(filepath)
+ if not filepath.exists():
+ raise FileNotFoundError(f"File {filepath} not found")
+
+ with open(filepath, 'rb') as f:
+ data = pickle.load(f)
+ loaded_n_bins = data['n_bins']
+ if loaded_n_bins != self.n_bins:
+ raise ValueError(f"Loaded n_bins ({loaded_n_bins}) does not match initialized n_bins ({self.n_bins})")
+ self.edges = data['edges']
+
+class PretrainData(Dataset):
+ def __init__(self, adata: sc.AnnData, drug_dict: dict):
+ self.adata = adata
+ self.drug_dict = drug_dict
+ self.X = torch.from_numpy(adata.X.toarray())
+ self.cell_id = adata.obs_names
+ # self.drug1 = torch.tensor(np.array(adata.obs['Drug1'].apply(lambda x: drug_dict[x])))
+ # self.drug2 = torch.tensor(np.array(adata.obs['Drug2'].apply(lambda x: drug_dict[x])))
+
+ def __len__(self):
+ return len(self.adata)
+
+ def __getitem__(self, idx):
+ return {
+ 'values' : self.X[idx],
+ 'cell_id': self.cell_id[idx],
+ }
+
+
+class FlowMatchingDataset(Dataset):
+ """PyTorch Dataset for flow matching training data"""
+
+ def __init__(self, jax_sampler, num_samples=10000, seed=42):
+ """
+ Args:
+ jax_sampler: JAX-based TrainSampler
+ num_samples: Number of samples to generate per epoch
+ seed: Random seed
+ """
+ self.jax_sampler = jax_sampler
+ self.num_samples = num_samples
+ self.rng = jax.random.PRNGKey(seed)
+
+ def __len__(self):
+ return self.num_samples
+
+ def __getitem__(self, idx):
+ """Sample a batch from the JAX sampler"""
+ # Generate new random key for each sample
+ self.rng, sample_key = jax.random.split(self.rng)
+
+ # Sample from JAX sampler
+ sample = self.jax_sampler.sample(sample_key)
+
+ # Convert JAX arrays to PyTorch tensors
+ src_cell_data = torch.from_numpy(np.array(sample['src_cell_data'])).float()
+ tgt_cell_data = torch.from_numpy(np.array(sample['tgt_cell_data'])).float()
+
+ sample['src_cell_id']
+ sample['tgt_cell_id']
+ # Convert condition embedding if available
+ condition_data = None
+ if 'condition' in sample:
+ condition_data = {
+ key: torch.from_numpy(np.array(val)).float()
+ for key, val in sample['condition'].items()
+ }
+
+ # Convert condition_id embedding if available
+ condition_id = None
+ if 'condition_id' in sample:
+ condition_id = torch.from_numpy(np.array(sample['condition_id'])).long()
+
+ return {
+ 'src_cell_data': src_cell_data,
+ 'tgt_cell_data': tgt_cell_data,
+ 'condition': condition_data,
+ 'condition_id': condition_id,
+ }
+
+if __name__ == "__main__":
+ data = Data(data_path='./data')
+ data.load_data(data_name='combosciplex')
+ data.process_data()
+
+
+
\ No newline at end of file
diff --git a/scDFM/src/data_process/scgpt_preprocess.py b/scDFM/src/data_process/scgpt_preprocess.py
new file mode 100644
index 0000000000000000000000000000000000000000..c4b1f8b9aa422f0dac7d9efc514545aecf1f4c8d
--- /dev/null
+++ b/scDFM/src/data_process/scgpt_preprocess.py
@@ -0,0 +1,303 @@
+from typing import Dict, Optional, Union
+
+import numpy as np
+import torch
+from scipy.sparse import issparse
+import scanpy as sc
+from scanpy.get import _get_obs_rep, _set_obs_rep
+from anndata import AnnData
+
+# from scgpt import logger
+
+
+class Preprocessor:
+ """
+ Prepare data into training, valid and test split. Normalize raw expression
+ values, binning or using other transform into the preset model input format.
+ """
+
+ def __init__(
+ self,
+ use_key: Optional[str] = None,
+ filter_gene_by_counts: Union[int, bool] = False,
+ filter_cell_by_counts: Union[int, bool] = False,
+ normalize_total: Union[float, bool] = 1e4,
+ result_normed_key: Optional[str] = "X_normed",
+ log1p: bool = False,
+ result_log1p_key: str = "X_log1p",
+ subset_hvg: Union[int, bool] = False,
+ hvg_use_key: Optional[str] = None,
+ hvg_flavor: str = "seurat_v3",
+ binning: Optional[int] = None,
+ result_binned_key: str = "X_binned",
+ ):
+ r"""
+ Set up the preprocessor, use the args to config the workflow steps.
+
+ Args:
+
+ use_key (:class:`str`, optional):
+ The key of :class:`~anndata.AnnData` to use for preprocessing.
+ filter_gene_by_counts (:class:`int` or :class:`bool`, default: ``False``):
+ Whther to filter genes by counts, if :class:`int`, filter genes with counts
+ filter_cell_by_counts (:class:`int` or :class:`bool`, default: ``False``):
+ Whther to filter cells by counts, if :class:`int`, filter cells with counts
+ normalize_total (:class:`float` or :class:`bool`, default: ``1e4``):
+ Whether to normalize the total counts of each cell to a specific value.
+ result_normed_key (:class:`str`, default: ``"X_normed"``):
+ The key of :class:`~anndata.AnnData` to store the normalized data. If
+ :class:`None`, will use normed data to replce the :attr:`use_key`.
+ log1p (:class:`bool`, default: ``True``):
+ Whether to apply log1p transform to the normalized data.
+ result_log1p_key (:class:`str`, default: ``"X_log1p"``):
+ The key of :class:`~anndata.AnnData` to store the log1p transformed data.
+ subset_hvg (:class:`int` or :class:`bool`, default: ``False``):
+ Whether to subset highly variable genes.
+ hvg_use_key (:class:`str`, optional):
+ The key of :class:`~anndata.AnnData` to use for calculating highly variable
+ genes. If :class:`None`, will use :attr:`adata.X`.
+ hvg_flavor (:class:`str`, default: ``"seurat_v3"``):
+ The flavor of highly variable genes selection. See
+ :func:`scanpy.pp.highly_variable_genes` for more details.
+ binning (:class:`int`, optional):
+ Whether to bin the data into discrete values of number of bins provided.
+ result_binned_key (:class:`str`, default: ``"X_binned"``):
+ The key of :class:`~anndata.AnnData` to store the binned data.
+ """
+ self.use_key = use_key
+ self.filter_gene_by_counts = filter_gene_by_counts
+ self.filter_cell_by_counts = filter_cell_by_counts
+ self.normalize_total = normalize_total
+ self.result_normed_key = result_normed_key
+ self.log1p = log1p
+ self.result_log1p_key = result_log1p_key
+ self.subset_hvg = subset_hvg
+ self.hvg_use_key = hvg_use_key
+ self.hvg_flavor = hvg_flavor
+ self.binning = binning
+ self.result_binned_key = result_binned_key
+
+ def __call__(self, adata: AnnData, batch_key: Optional[str] = None) -> Dict:
+ """
+ format controls the different input value wrapping, including categorical
+ binned style, fixed-sum normalized counts, log1p fixed-sum normalized counts, etc.
+
+ Args:
+
+ adata (:class:`AnnData`):
+ The :class:`AnnData` object to preprocess.
+ batch_key (:class:`str`, optional):
+ The key of :class:`AnnData.obs` to use for batch information. This arg
+ is used in the highly variable gene selection step.
+ """
+ key_to_process = self.use_key
+ # preliminary checks, will use later
+ if key_to_process == "X":
+ key_to_process = None # the following scanpy apis use arg None to use X
+ is_logged = self.check_logged(adata, obs_key=key_to_process)
+
+ # step 1: filter genes
+ if self.filter_gene_by_counts:
+ print("Filtering genes by counts ...")
+ sc.pp.filter_genes(
+ adata,
+ min_counts=self.filter_gene_by_counts
+ if isinstance(self.filter_gene_by_counts, int)
+ else None,
+ )
+
+ # step 2: filter cells
+ if (
+ isinstance(self.filter_cell_by_counts, int)
+ and self.filter_cell_by_counts > 0
+ ):
+ print("Filtering cells by counts ...")
+ sc.pp.filter_cells(
+ adata,
+ min_counts=self.filter_cell_by_counts
+ if isinstance(self.filter_cell_by_counts, int)
+ else None,
+ )
+
+ # step 3: normalize total
+ if self.normalize_total:
+ print("Normalizing total counts ...")
+ normed_ = sc.pp.normalize_total(
+ adata,
+ target_sum=self.normalize_total
+ if isinstance(self.normalize_total, float)
+ else None,
+ layer=key_to_process,
+ inplace=False,
+ )["X"]
+ key_to_process = self.result_normed_key or key_to_process
+ _set_obs_rep(adata, normed_, layer=key_to_process)
+
+ # step 4: log1p
+ if self.log1p:
+ print("Log1p transforming ...")
+ if is_logged:
+ print(
+ "The input data seems to be already log1p transformed. "
+ "Set `log1p=False` to avoid double log1p transform."
+ )
+ if self.result_log1p_key:
+ _set_obs_rep(
+ adata,
+ _get_obs_rep(adata, layer=key_to_process),
+ layer=self.result_log1p_key,
+ )
+ key_to_process = self.result_log1p_key
+ sc.pp.log1p(adata, layer=key_to_process)
+
+ # step 5: subset hvg
+ if self.subset_hvg:
+ print("Subsetting highly variable genes ...")
+ if batch_key is None:
+ print(
+ "No batch_key is provided, will use all cells for HVG selection."
+ )
+ sc.pp.highly_variable_genes(
+ adata,
+ layer=self.hvg_use_key,
+ n_top_genes=self.subset_hvg
+ if isinstance(self.subset_hvg, int)
+ else None,
+ batch_key=batch_key,
+ flavor=self.hvg_flavor,
+ subset=True,
+ )
+
+ # step 6: binning
+ if self.binning:
+ print("Binning data ...")
+ if not isinstance(self.binning, int):
+ raise ValueError(
+ "Binning arg must be an integer, but got {}.".format(self.binning)
+ )
+ n_bins = self.binning # NOTE: the first bin is always a spectial for zero
+ binned_rows = []
+ bin_edges = []
+ layer_data = _get_obs_rep(adata, layer=key_to_process)
+ layer_data = layer_data.toarray() if issparse(layer_data) else layer_data
+ if layer_data.min() < 0:
+ raise ValueError(
+ f"Assuming non-negative data, but got min value {layer_data.min()}."
+ )
+ for row in layer_data:
+ if row.max() == 0:
+ print(
+ "The input data contains all zero rows. Please make sure "
+ "this is expected. You can use the `filter_cell_by_counts` "
+ "arg to filter out all zero rows."
+ )
+ binned_rows.append(np.zeros_like(row, dtype=np.int64))
+ bin_edges.append(np.array([0] * n_bins))
+ continue
+ non_zero_ids = row.nonzero()
+ non_zero_row = row[non_zero_ids]
+ bins = np.quantile(non_zero_row, np.linspace(0, 1, n_bins - 1))
+ # bins = np.sort(np.unique(bins))
+ # NOTE: comment this line for now, since this will make the each category
+ # has different relative meaning across datasets
+ non_zero_digits = _digitize(non_zero_row, bins)
+ assert non_zero_digits.min() >= 1
+ assert non_zero_digits.max() <= n_bins - 1
+ binned_row = np.zeros_like(row, dtype=np.int64)
+ binned_row[non_zero_ids] = non_zero_digits
+ binned_rows.append(binned_row)
+ bin_edges.append(np.concatenate([[0], bins]))
+ adata.layers[self.result_binned_key] = np.stack(binned_rows)
+ adata.obsm["bin_edges"] = np.stack(bin_edges)
+
+ def check_logged(self, adata: AnnData, obs_key: Optional[str] = None) -> bool:
+ """
+ Check if the data is already log1p transformed.
+
+ Args:
+
+ adata (:class:`AnnData`):
+ The :class:`AnnData` object to preprocess.
+ obs_key (:class:`str`, optional):
+ The key of :class:`AnnData.obs` to use for batch information. This arg
+ is used in the highly variable gene selection step.
+ """
+ data = _get_obs_rep(adata, layer=obs_key)
+ max_, min_ = data.max(), data.min()
+ if max_ > 30:
+ return False
+ if min_ < 0:
+ return False
+
+ non_zero_min = data[data > 0].min()
+ if non_zero_min >= 1:
+ return False
+
+ return True
+
+
+def _digitize(x: np.ndarray, bins: np.ndarray, side="both") -> np.ndarray:
+ """
+ Digitize the data into bins. This method spreads data uniformly when bins
+ have same values.
+
+ Args:
+
+ x (:class:`np.ndarray`):
+ The data to digitize.
+ bins (:class:`np.ndarray`):
+ The bins to use for digitization, in increasing order.
+ side (:class:`str`, optional):
+ The side to use for digitization. If "one", the left side is used. If
+ "both", the left and right side are used. Default to "one".
+
+ Returns:
+
+ :class:`np.ndarray`:
+ The digitized data.
+ """
+ assert x.ndim == 1 and bins.ndim == 1
+
+ left_digits = np.digitize(x, bins)
+ if side == "one":
+ return left_digits
+
+ right_difits = np.digitize(x, bins, right=True)
+
+ rands = np.random.rand(len(x)) # uniform random numbers
+
+ digits = rands * (right_difits - left_digits) + left_digits
+ digits = np.ceil(digits).astype(np.int64)
+ return digits
+
+
+def binning(
+ row: Union[np.ndarray, torch.Tensor], n_bins: int
+) -> Union[np.ndarray, torch.Tensor]:
+ """Binning the row into n_bins."""
+ dtype = row.dtype
+ return_np = False if isinstance(row, torch.Tensor) else True
+ row = row.cpu().numpy() if isinstance(row, torch.Tensor) else row
+ # TODO: use torch.quantile and torch.bucketize
+
+ if row.max() == 0:
+ print(
+ "The input data contains row of zeros. Please make sure this is expected."
+ )
+ return (
+ np.zeros_like(row, dtype=dtype)
+ if return_np
+ else torch.zeros_like(row, dtype=dtype)
+ )
+
+ if row.min() <= 0:
+ non_zero_ids = row.nonzero()
+ non_zero_row = row[non_zero_ids]
+ bins = np.quantile(non_zero_row, np.linspace(0, 1, n_bins - 1))
+ non_zero_digits = _digitize(non_zero_row, bins)
+ binned_row = np.zeros_like(row, dtype=np.int64)
+ binned_row[non_zero_ids] = non_zero_digits
+ else:
+ bins = np.quantile(row, np.linspace(0, 1, n_bins - 1))
+ binned_row = _digitize(row, bins)
+ return torch.from_numpy(binned_row) if not return_np else binned_row.astype(dtype)
\ No newline at end of file
diff --git a/scDFM/src/data_process/utils.py b/scDFM/src/data_process/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..71d7de92f507e726d731cb07ebf669e3de33e8bc
--- /dev/null
+++ b/scDFM/src/data_process/utils.py
@@ -0,0 +1,35 @@
+import anndata as ad
+import pandas as pd
+import numpy as np
+
+def build_generated_anndata(
+ embeddings: np.ndarray, # shape: (N, D)
+ condition_id: np.ndarray, # shape: (N, 2)
+ pert_col: str = "perturbation", # first perturbation column name
+ pert_col2: str = "perturbation_2", # second perturbation (if combo)
+ combo_col: str = "pert_combo", # combined label
+) -> ad.AnnData:
+ """
+ Construct AnnData from model-generated embeddings and condition IDs.
+ """
+ N = embeddings.shape[0]
+ if condition_id.shape[0] != N:
+ raise ValueError("condition_id and embeddings must have same first dimension.")
+
+ # obs: per-cell annotation
+ obs = pd.DataFrame({
+ pert_col: condition_id[:, 0],
+ pert_col2: condition_id[:, 1],
+ combo_col: [f"{a}_{b}" for a, b in condition_id],
+ }, index=[f"cell_{i}" for i in range(N)])
+
+ # construct AnnData
+ adata = ad.AnnData(
+ X=embeddings, # could also store in obsm
+ obs=obs
+ )
+
+ # also store embedding explicitly for plotting
+ adata.obsm["embedding"] = embeddings
+
+ return adata
\ No newline at end of file
diff --git a/scDFM/src/flow_matching/__init__.py b/scDFM/src/flow_matching/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..719ae055c6a13d749beca75c7a496f3cd7cd8948
--- /dev/null
+++ b/scDFM/src/flow_matching/__init__.py
@@ -0,0 +1,7 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+__version__ = "1.0.10"
diff --git a/scDFM/src/flow_matching/ot/__init__.py b/scDFM/src/flow_matching/ot/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..4a7eee1fa298559ee8d944804ac860cfb4767a25
--- /dev/null
+++ b/scDFM/src/flow_matching/ot/__init__.py
@@ -0,0 +1,5 @@
+from .optimal_transport import OTPlanSampler
+
+__all__ = [
+ "OTPlanSampler",
+]
\ No newline at end of file
diff --git a/scDFM/src/flow_matching/ot/optimal_transport.py b/scDFM/src/flow_matching/ot/optimal_transport.py
new file mode 100644
index 0000000000000000000000000000000000000000..e0cd855a572274e82de4f1c563f6be592a5cc2a8
--- /dev/null
+++ b/scDFM/src/flow_matching/ot/optimal_transport.py
@@ -0,0 +1,291 @@
+import math
+import warnings
+from functools import partial
+from typing import Optional, Union
+
+import numpy as np
+import ot as pot
+import torch
+
+
+class OTPlanSampler:
+ """OTPlanSampler implements sampling coordinates according to an OT plan (wrt squared Euclidean
+ cost) with different implementations of the plan calculation."""
+
+ def __init__(
+ self,
+ method: str,
+ reg: float = 0.05,
+ reg_m: float = 1.0,
+ normalize_cost: bool = False,
+ num_threads: Union[int, str] = 1,
+ warn: bool = True,
+ ) -> None:
+ """Initialize the OTPlanSampler class.
+
+ Parameters
+ ----------
+ method: str
+ choose which optimal transport solver you would like to use.
+ Currently supported are ["exact", "sinkhorn", "unbalanced",
+ "partial"] OT solvers.
+ reg: float, optional
+ regularization parameter to use for Sinkhorn-based iterative solvers.
+ reg_m: float, optional
+ regularization weight for unbalanced Sinkhorn-knopp solver.
+ normalize_cost: bool, optional
+ normalizes the cost matrix so that the maximum cost is 1. Helps
+ stabilize Sinkhorn-based solvers. Should not be used in the vast
+ majority of cases.
+ num_threads: int or str, optional
+ number of threads to use for the "exact" OT solver. If "max", uses
+ the maximum number of threads.
+ warn: bool, optional
+ if True, raises a warning if the algorithm does not converge
+ """
+ # ot_fn should take (a, b, M) as arguments where a, b are marginals and
+ # M is a cost matrix
+ if method == "exact":
+ self.ot_fn = partial(pot.emd, numThreads=num_threads)
+ elif method == "sinkhorn":
+ self.ot_fn = partial(pot.sinkhorn, reg=reg)
+ elif method == "unbalanced":
+ self.ot_fn = partial(pot.unbalanced.sinkhorn_knopp_unbalanced, reg=reg, reg_m=reg_m)
+ elif method == "partial":
+ self.ot_fn = partial(pot.partial.entropic_partial_wasserstein, reg=reg)
+ else:
+ raise ValueError(f"Unknown method: {method}")
+ self.reg = reg
+ self.reg_m = reg_m
+ self.normalize_cost = normalize_cost
+ self.warn = warn
+
+ def get_map(self, x0, x1):
+ """Compute the OT plan (wrt squared Euclidean cost) between a source and a target
+ minibatch.
+
+ Parameters
+ ----------
+ x0 : Tensor, shape (bs, *dim)
+ represents the source minibatch
+ x1 : Tensor, shape (bs, *dim)
+ represents the source minibatch
+
+ Returns
+ -------
+ p : numpy array, shape (bs, bs)
+ represents the OT plan between minibatches
+ """
+ a, b = pot.unif(x0.shape[0]), pot.unif(x1.shape[0])
+ if x0.dim() > 2:
+ x0 = x0.reshape(x0.shape[0], -1)
+ if x1.dim() > 2:
+ x1 = x1.reshape(x1.shape[0], -1)
+ M = torch.cdist(x0, x1) ** 2
+ if self.normalize_cost:
+ M = M / M.max() # should not be normalized when using minibatches
+ p = self.ot_fn(a, b, M.detach().cpu().numpy())
+ if not np.all(np.isfinite(p)):
+ print("ERROR: p is not finite")
+ print(p)
+ print("Cost mean, max", M.mean(), M.max())
+ print(x0, x1)
+ if np.abs(p.sum()) < 1e-8:
+ if self.warn:
+ warnings.warn("Numerical errors in OT plan, reverting to uniform plan.")
+ p = np.ones_like(p) / p.size
+ return p
+
+ def sample_map(self, pi, batch_size, replace=True):
+ r"""Draw source and target samples from pi $(x,z) \sim \pi$
+
+ Parameters
+ ----------
+ pi : numpy array, shape (bs, bs)
+ represents the source minibatch
+ batch_size : int
+ represents the OT plan between minibatches
+ replace : bool
+ represents sampling or without replacement from the OT plan
+
+ Returns
+ -------
+ (i_s, i_j) : tuple of numpy arrays, shape (bs, bs)
+ represents the indices of source and target data samples from $\pi$
+ """
+ p = pi.flatten()
+ p = p / p.sum()
+ choices = np.random.choice(
+ pi.shape[0] * pi.shape[1], p=p, size=batch_size, replace=replace
+ )
+ return np.divmod(choices, pi.shape[1])
+
+ def sample_plan(self, x0, x1, replace=True):
+ r"""Compute the OT plan $\pi$ (wrt squared Euclidean cost) between a source and a target
+ minibatch and draw source and target samples from pi $(x,z) \sim \pi$
+
+ Parameters
+ ----------
+ x0 : Tensor, shape (bs, *dim)
+ represents the source minibatch
+ x1 : Tensor, shape (bs, *dim)
+ represents the source minibatch
+ replace : bool
+ represents sampling or without replacement from the OT plan
+
+ Returns
+ -------
+ x0[i] : Tensor, shape (bs, *dim)
+ represents the source minibatch drawn from $\pi$
+ x1[j] : Tensor, shape (bs, *dim)
+ represents the source minibatch drawn from $\pi$
+ """
+ pi = self.get_map(x0, x1)
+ i, j = self.sample_map(pi, x0.shape[0], replace=replace)
+ return x0[i], x1[j]
+
+ def sample_plan_fix_x0(self, x0, x1, replace=True):
+ """
+ Sample matched samples from x1 according to the OT plan, keeping the order of x0 unchanged.
+
+ Args:
+ x0 : Tensor, shape (bs, *dim) - kept unchanged
+ x1 : Tensor, shape (bs, *dim) - to be sampled
+ replace : bool - whether to sample with replacement
+
+ Returns:
+ x0 : original x0, unchanged
+ x1_matched : matched samples sampled from x1
+ """
+ pi = self.get_map(x0, x1)
+ matched_indices = []
+
+ for i in range(pi.shape[0]):
+ prob = pi[i] / pi[i].sum()
+ j = np.random.choice(pi.shape[1], p=prob)
+ matched_indices.append(j)
+
+ matched_indices = torch.tensor(matched_indices, dtype=torch.long, device=x1.device)
+ x1_matched = x1[matched_indices]
+ return x0, x1_matched
+
+ def sample_plan_with_labels(self, x0, x1, y0=None, y1=None, replace=True):
+ r"""Compute the OT plan $\pi$ (wrt squared Euclidean cost) between a source and a target
+ minibatch and draw source and target labeled samples from pi $(x,z) \sim \pi$
+
+ Parameters
+ ----------
+ x0 : Tensor, shape (bs, *dim)
+ represents the source minibatch
+ x1 : Tensor, shape (bs, *dim)
+ represents the target minibatch
+ y0 : Tensor, shape (bs)
+ represents the source label minibatch
+ y1 : Tensor, shape (bs)
+ represents the target label minibatch
+ replace : bool
+ represents sampling or without replacement from the OT plan
+
+ Returns
+ -------
+ x0[i] : Tensor, shape (bs, *dim)
+ represents the source minibatch drawn from $\pi$
+ x1[j] : Tensor, shape (bs, *dim)
+ represents the target minibatch drawn from $\pi$
+ y0[i] : Tensor, shape (bs, *dim)
+ represents the source label minibatch drawn from $\pi$
+ y1[j] : Tensor, shape (bs, *dim)
+ represents the target label minibatch drawn from $\pi$
+ """
+ pi = self.get_map(x0, x1)
+ i, j = self.sample_map(pi, x0.shape[0], replace=replace)
+ return (
+ x0[i],
+ x1[j],
+ y0[i] if y0 is not None else None,
+ y1[j] if y1 is not None else None,
+ )
+
+ def sample_trajectory(self, X):
+ """Compute the OT trajectories between different sample populations moving from the source
+ to the target distribution.
+
+ Parameters
+ ----------
+ X : Tensor, (bs, times, *dim)
+ different populations of samples moving from the source to the target distribution.
+
+ Returns
+ -------
+ to_return : Tensor, (bs, times, *dim)
+ represents the OT sampled trajectories over time.
+ """
+ times = X.shape[1]
+ pis = []
+ for t in range(times - 1):
+ pis.append(self.get_map(X[:, t], X[:, t + 1]))
+
+ indices = [np.arange(X.shape[0])]
+ for pi in pis:
+ j = []
+ for i in indices[-1]:
+ j.append(np.random.choice(pi.shape[1], p=pi[i] / pi[i].sum()))
+ indices.append(np.array(j))
+
+ to_return = []
+ for t in range(times):
+ to_return.append(X[:, t][indices[t]])
+ to_return = np.stack(to_return, axis=1)
+ return to_return
+
+
+def wasserstein(
+ x0: torch.Tensor,
+ x1: torch.Tensor,
+ method: Optional[str] = None,
+ reg: float = 0.05,
+ power: int = 2,
+ **kwargs,
+) -> float:
+ """Compute the Wasserstein (1 or 2) distance (wrt Euclidean cost) between a source and a target
+ distributions.
+
+ Parameters
+ ----------
+ x0 : Tensor, shape (bs, *dim)
+ represents the source minibatch
+ x1 : Tensor, shape (bs, *dim)
+ represents the source minibatch
+ method : str (default : None)
+ Use exact Wasserstein or an entropic regularization
+ reg : float (default : 0.05)
+ Entropic regularization coefficients
+ power : int (default : 2)
+ power of the Wasserstein distance (1 or 2)
+ Returns
+ -------
+ ret : float
+ Wasserstein distance
+ """
+ assert power == 1 or power == 2
+ # ot_fn should take (a, b, M) as arguments where a, b are marginals and
+ # M is a cost matrix
+ if method == "exact" or method is None:
+ ot_fn = pot.emd2
+ elif method == "sinkhorn":
+ ot_fn = partial(pot.sinkhorn2, reg=reg)
+ else:
+ raise ValueError(f"Unknown method: {method}")
+
+ a, b = pot.unif(x0.shape[0]), pot.unif(x1.shape[0])
+ if x0.dim() > 2:
+ x0 = x0.reshape(x0.shape[0], -1)
+ if x1.dim() > 2:
+ x1 = x1.reshape(x1.shape[0], -1)
+ M = torch.cdist(x0, x1)
+ if power == 2:
+ M = M**2
+ ret = ot_fn(a, b, M.detach().cpu().numpy(), numItermax=int(1e7))
+ if power == 2:
+ ret = math.sqrt(ret)
+ return ret
diff --git a/scDFM/src/flow_matching/path/__init__.py b/scDFM/src/flow_matching/path/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..88d29a2575c1be14d2ec10bd9eda799e7599dd1d
--- /dev/null
+++ b/scDFM/src/flow_matching/path/__init__.py
@@ -0,0 +1,22 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .affine import AffineProbPath, CondOTProbPath
+from .geodesic import GeodesicProbPath
+from .mixture import MixtureDiscreteProbPath
+from .path import ProbPath
+from .path_sample import DiscretePathSample, PathSample
+
+
+__all__ = [
+ "ProbPath",
+ "AffineProbPath",
+ "CondOTProbPath",
+ "MixtureDiscreteProbPath",
+ "GeodesicProbPath",
+ "PathSample",
+ "DiscretePathSample",
+]
diff --git a/scDFM/src/flow_matching/path/affine.py b/scDFM/src/flow_matching/path/affine.py
new file mode 100644
index 0000000000000000000000000000000000000000..90bb06a22f84055fe009c0f6257eb61493ed6f1c
--- /dev/null
+++ b/scDFM/src/flow_matching/path/affine.py
@@ -0,0 +1,260 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch import Tensor
+
+from src.flow_matching.path.path import ProbPath
+from src.flow_matching.path.path_sample import PathSample
+from src.flow_matching.path.scheduler.scheduler import CondOTScheduler, Scheduler
+from src.flow_matching.utils import expand_tensor_like
+
+
+class AffineProbPath(ProbPath):
+ r"""The ``AffineProbPath`` class represents a specific type of probability path where the transformation between distributions is affine.
+ An affine transformation can be represented as:
+
+ .. math::
+
+ X_t = \alpha_t X_1 + \sigma_t X_0,
+
+ where :math:`X_t` is the transformed data point at time `t`. :math:`X_0` and :math:`X_1` are the source and target data points, respectively. :math:`\alpha_t` and :math:`\sigma_t` are the parameters of the affine transformation at time `t`.
+
+ The scheduler is responsible for providing the time-dependent parameters :math:`\alpha_t` and :math:`\sigma_t`, as well as their derivatives, which define the affine transformation at any given time `t`.
+
+ Using ``AffineProbPath`` in the flow matching framework:
+
+ .. code-block:: python
+
+ # Instantiates a probability path
+ my_path = AffineProbPath(...)
+ mse_loss = torch.nn.MSELoss()
+
+ for x_1 in dataset:
+ # Sets x_0 to random noise
+ x_0 = torch.randn()
+
+ # Sets t to a random value in [0,1]
+ t = torch.rand()
+
+ # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
+ path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
+
+ # Computes the MSE loss w.r.t. the velocity
+ loss = mse_loss(path_sample.dx_t, my_model(x_t, t))
+ loss.backward()
+
+ Args:
+ scheduler (Scheduler): An instance of a scheduler that provides the parameters :math:`\alpha_t`, :math:`\sigma_t`, and their derivatives over time.
+
+ """
+
+ def __init__(self, scheduler: Scheduler):
+ self.scheduler = scheduler
+
+ def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample:
+ r"""Sample from the affine probability path:
+
+ | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`.
+ | return :math:`X_0, X_1, X_t = \alpha_t X_1 + \sigma_t X_0`, and the conditional velocity at :math:`X_t, \dot{X}_t = \dot{\alpha}_t X_1 + \dot{\sigma}_t X_0`.
+
+ Args:
+ x_0 (Tensor): source data point, shape (batch_size, ...).
+ x_1 (Tensor): target data point, shape (batch_size, ...).
+ t (Tensor): times in [0,1], shape (batch_size).
+
+ Returns:
+ PathSample: a conditional sample at :math:`X_t \sim p_t`.
+ """
+ self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t)
+
+ scheduler_output = self.scheduler(t)
+
+ alpha_t = expand_tensor_like(
+ input_tensor=scheduler_output.alpha_t, expand_to=x_1
+ )
+ sigma_t = expand_tensor_like(
+ input_tensor=scheduler_output.sigma_t, expand_to=x_1
+ )
+ d_alpha_t = expand_tensor_like(
+ input_tensor=scheduler_output.d_alpha_t, expand_to=x_1
+ )
+ d_sigma_t = expand_tensor_like(
+ input_tensor=scheduler_output.d_sigma_t, expand_to=x_1
+ )
+
+ # construct xt ~ p_t(x|x1).
+ x_t = sigma_t * x_0 + alpha_t * x_1
+ dx_t = d_sigma_t * x_0 + d_alpha_t * x_1
+
+ return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)
+
+ def target_to_velocity(self, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+ r"""Convert from x_1 representation to velocity.
+
+ | given :math:`X_1`.
+ | return :math:`\dot{X}_t`.
+
+ Args:
+ x_1 (Tensor): target data point.
+ x_t (Tensor): path sample at time t.
+ t (Tensor): time in [0,1].
+
+ Returns:
+ Tensor: velocity.
+ """
+ scheduler_output = self.scheduler(t)
+
+ alpha_t = scheduler_output.alpha_t
+ d_alpha_t = scheduler_output.d_alpha_t
+ sigma_t = scheduler_output.sigma_t
+ d_sigma_t = scheduler_output.d_sigma_t
+
+ a_t = d_sigma_t / sigma_t
+ b_t = (d_alpha_t * sigma_t - d_sigma_t * alpha_t) / sigma_t
+
+ return a_t * x_t + b_t * x_1
+
+ def epsilon_to_velocity(self, epsilon: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+ r"""Convert from epsilon representation to velocity.
+
+ | given :math:`\epsilon`.
+ | return :math:`\dot{X}_t`.
+
+ Args:
+ epsilon (Tensor): noise in the path sample.
+ x_t (Tensor): path sample at time t.
+ t (Tensor): time in [0,1].
+
+ Returns:
+ Tensor: velocity.
+ """
+ scheduler_output = self.scheduler(t)
+
+ alpha_t = scheduler_output.alpha_t
+ d_alpha_t = scheduler_output.d_alpha_t
+ sigma_t = scheduler_output.sigma_t
+ d_sigma_t = scheduler_output.d_sigma_t
+
+ a_t = d_alpha_t / alpha_t
+ b_t = (d_sigma_t * alpha_t - d_alpha_t * sigma_t) / alpha_t
+
+ return a_t * x_t + b_t * epsilon
+
+ def velocity_to_target(self, velocity: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+ r"""Convert from velocity to x_1 representation.
+
+ | given :math:`\dot{X}_t`.
+ | return :math:`X_1`.
+
+ Args:
+ velocity (Tensor): velocity at the path sample.
+ x_t (Tensor): path sample at time t.
+ t (Tensor): time in [0,1].
+
+ Returns:
+ Tensor: target data point.
+ """
+ scheduler_output = self.scheduler(t)
+
+ alpha_t = scheduler_output.alpha_t
+ d_alpha_t = scheduler_output.d_alpha_t
+ sigma_t = scheduler_output.sigma_t
+ d_sigma_t = scheduler_output.d_sigma_t
+
+ a_t = -d_sigma_t / (d_alpha_t * sigma_t - d_sigma_t * alpha_t)
+ b_t = sigma_t / (d_alpha_t * sigma_t - d_sigma_t * alpha_t)
+
+ return a_t * x_t + b_t * velocity
+
+ def epsilon_to_target(self, epsilon: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+ r"""Convert from epsilon representation to x_1 representation.
+
+ | given :math:`\epsilon`.
+ | return :math:`X_1`.
+
+ Args:
+ epsilon (Tensor): noise in the path sample.
+ x_t (Tensor): path sample at time t.
+ t (Tensor): time in [0,1].
+
+ Returns:
+ Tensor: target data point.
+ """
+ scheduler_output = self.scheduler(t)
+
+ alpha_t = scheduler_output.alpha_t
+ sigma_t = scheduler_output.sigma_t
+
+ a_t = 1 / alpha_t
+ b_t = -sigma_t / alpha_t
+
+ return a_t * x_t + b_t * epsilon
+
+ def velocity_to_epsilon(self, velocity: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+ r"""Convert from velocity to noise representation.
+
+ | given :math:`\dot{X}_t`.
+ | return :math:`\epsilon`.
+
+ Args:
+ velocity (Tensor): velocity at the path sample.
+ x_t (Tensor): path sample at time t.
+ t (Tensor): time in [0,1].
+
+ Returns:
+ Tensor: noise in the path sample.
+ """
+ scheduler_output = self.scheduler(t)
+
+ alpha_t = scheduler_output.alpha_t
+ d_alpha_t = scheduler_output.d_alpha_t
+ sigma_t = scheduler_output.sigma_t
+ d_sigma_t = scheduler_output.d_sigma_t
+
+ a_t = -d_alpha_t / (d_sigma_t * alpha_t - d_alpha_t * sigma_t)
+ b_t = alpha_t / (d_sigma_t * alpha_t - d_alpha_t * sigma_t)
+
+ return a_t * x_t + b_t * velocity
+
+ def target_to_epsilon(self, x_1: Tensor, x_t: Tensor, t: Tensor) -> Tensor:
+ r"""Convert from x_1 representation to velocity.
+
+ | given :math:`X_1`.
+ | return :math:`\epsilon`.
+
+ Args:
+ x_1 (Tensor): target data point.
+ x_t (Tensor): path sample at time t.
+ t (Tensor): time in [0,1].
+
+ Returns:
+ Tensor: noise in the path sample.
+ """
+ scheduler_output = self.scheduler(t)
+
+ alpha_t = scheduler_output.alpha_t
+ sigma_t = scheduler_output.sigma_t
+
+ a_t = 1 / sigma_t
+ b_t = -alpha_t / sigma_t
+
+ return a_t * x_t + b_t * x_1
+
+
+class CondOTProbPath(AffineProbPath):
+ r"""The ``CondOTProbPath`` class represents a conditional optimal transport probability path.
+
+ This class is a specialized version of the ``AffineProbPath`` that uses a conditional optimal transport scheduler to determine the parameters of the affine transformation.
+
+ The parameters :math:`\alpha_t` and :math:`\sigma_t` for the conditional optimal transport path are defined as:
+
+ .. math::
+
+ \alpha_t = t \quad \text{and} \quad \sigma_t = 1 - t.
+ """
+
+ def __init__(self):
+ self.scheduler = CondOTScheduler()
diff --git a/scDFM/src/flow_matching/path/geodesic.py b/scDFM/src/flow_matching/path/geodesic.py
new file mode 100644
index 0000000000000000000000000000000000000000..23144cc842e72caa79797044dbac3738b8245977
--- /dev/null
+++ b/scDFM/src/flow_matching/path/geodesic.py
@@ -0,0 +1,100 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+
+from torch import Tensor
+from torch.func import jvp, vmap
+
+from src.flow_matching.path.path import ProbPath
+
+from src.flow_matching.path.path_sample import PathSample
+from src.flow_matching.path.scheduler import ConvexScheduler
+from src.flow_matching.utils import expand_tensor_like
+
+from src.flow_matching.utils.manifolds import geodesic, Manifold
+
+
+class GeodesicProbPath(ProbPath):
+ r"""The ``GeodesicProbPath`` class represents a specific type of probability path where the transformation between distributions is defined through the geodesic path.
+ Mathematically, a geodesic path can be represented as:
+
+ .. math::
+
+ X_t = \psi_t(X_0 | X_1) = \exp_{X_1}(\kappa_t \log_{X_1}(X_0)),
+
+ where :math:`X_t` is the transformed data point at time `t`, :math:`X_0` and :math:`X_1` are the source and target data points, respectively, and :math:`\kappa_t` is a scheduler.
+
+ The scheduler is responsible for providing the time-dependent :math:`\kappa_t` and must be differentiable.
+
+ Using ``GeodesicProbPath`` in the flow matching framework:
+
+ .. code-block:: python
+ # Instantiates a manifold
+ manifold = FlatTorus()
+
+ # Instantiates a scheduler
+ scheduler = CondOTScheduler()
+
+ # Instantiates a probability path
+ my_path = GeodesicProbPath(scheduler, manifold)
+ mse_loss = torch.nn.MSELoss()
+
+ for x_1 in dataset:
+ # Sets x_0 to random noise
+ x_0 = torch.randn()
+
+ # Sets t to a random value in [0,1]
+ t = torch.rand()
+
+ # Samples the conditional path :math:`X_t \sim p_t(X_t|X_0,X_1)`
+ path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
+
+ # Computes the MSE loss w.r.t. the velocity
+ loss = mse_loss(path_sample.dx_t, my_model(x_t, t))
+ loss.backward()
+
+ Args:
+ scheduler (ConvexScheduler): The scheduler that provides :math:`\kappa_t`.
+ manifold (Manifold): The manifold on which the probability path is defined.
+
+ """
+
+ def __init__(self, scheduler: ConvexScheduler, manifold: Manifold):
+ self.scheduler = scheduler
+ self.manifold = manifold
+
+ def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample:
+ r"""Sample from the Riemannian probability path with geodesic interpolation:
+
+ | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`\kappa_t`.
+ | return :math:`X_0, X_1, X_t = \exp_{X_1}(\kappa_t \log_{X_1}(X_0))`, and the conditional velocity at :math:`X_t, \dot{X}_t`.
+
+ Args:
+ x_0 (Tensor): source data point, shape (batch_size, ...).
+ x_1 (Tensor): target data point, shape (batch_size, ...).
+ t (Tensor): times in [0,1], shape (batch_size).
+
+ Returns:
+ PathSample: A conditional sample at :math:`X_t \sim p_t`.
+ """
+ self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t)
+ t = expand_tensor_like(input_tensor=t, expand_to=x_1[..., 0:1]).clone()
+
+ def cond_u(x_0, x_1, t):
+ path = geodesic(self.manifold, x_0, x_1)
+ x_t, dx_t = jvp(
+ lambda t: path(self.scheduler(t).alpha_t),
+ (t,),
+ (torch.ones_like(t).to(t),),
+ )
+ return x_t, dx_t
+
+ x_t, dx_t = vmap(cond_u)(x_0, x_1, t)
+ x_t = x_t.reshape_as(x_1)
+ dx_t = dx_t.reshape_as(x_1)
+
+ return PathSample(x_t=x_t, dx_t=dx_t, x_1=x_1, x_0=x_0, t=t)
diff --git a/scDFM/src/flow_matching/path/mixture.py b/scDFM/src/flow_matching/path/mixture.py
new file mode 100644
index 0000000000000000000000000000000000000000..3c5b5f841be04f5f213cb21889df7df486770545
--- /dev/null
+++ b/scDFM/src/flow_matching/path/mixture.py
@@ -0,0 +1,117 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+import torch.nn.functional as F
+
+from torch import Tensor
+
+from src.flow_matching.path.path import ProbPath
+
+from src.flow_matching.path.path_sample import DiscretePathSample
+from src.flow_matching.path.scheduler import ConvexScheduler
+from src.flow_matching.utils import expand_tensor_like, unsqueeze_to_match
+
+
+class MixtureDiscreteProbPath(ProbPath):
+ r"""The ``MixtureDiscreteProbPath`` class defines a factorized discrete probability path.
+
+ This path remains constant at the source data point :math:`X_0` until a random time, determined by the scheduler, when it flips to the target data point :math:`X_1`.
+ The scheduler determines the flip probability using the parameter :math:`\sigma_t`, which is a function of time `t`. Specifically, :math:`\sigma_t` represents the probability of remaining at :math:`X_0`, while :math:`1 - \sigma_t` is the probability of flipping to :math:`X_1`:
+
+ .. math::
+
+ P(X_t = X_0) = \sigma_t \quad \text{and} \quad P(X_t = X_1) = 1 - \sigma_t,
+
+ where :math:`\sigma_t` is provided by the scheduler.
+
+ Example:
+
+ .. code-block:: python
+
+ >>> x_0 = torch.zeros((1, 3, 3))
+ >>> x_1 = torch.ones((1, 3, 3))
+
+ >>> path = MixtureDiscreteProbPath(PolynomialConvexScheduler(n=1.0))
+ >>> result = path.sample(x_0, x_1, t=torch.tensor([0.1])).x_t
+ >>> result
+ tensor([[[0.0, 0.0, 0.0],
+ [0.0, 0.0, 1.0],
+ [0.0, 0.0, 0.0]]])
+
+ >>> result = path.sample(x_0, x_1, t=torch.tensor([0.5])).x_t
+ >>> result
+ tensor([[[1.0, 0.0, 1.0],
+ [0.0, 1.0, 0.0],
+ [0.0, 1.0, 0.0]]])
+
+ >>> result = path.sample(x_0, x_1, t=torch.tensor([1.0])).x_t
+ >>> result
+ tensor([[[1.0, 1.0, 1.0],
+ [1.0, 1.0, 1.0],
+ [1.0, 1.0, 1.0]]])
+
+ Args:
+ scheduler (ConvexScheduler): The scheduler that provides :math:`\sigma_t`.
+ """
+
+ def __init__(self, scheduler: ConvexScheduler):
+ assert isinstance(
+ scheduler, ConvexScheduler
+ ), "Scheduler for ConvexProbPath must be a ConvexScheduler."
+
+ self.scheduler = scheduler
+
+ def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> DiscretePathSample:
+ r"""Sample from the affine probability path:
+ | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)` and a scheduler :math:`(\alpha_t,\sigma_t)`.
+ | return :math:`X_0, X_1, t`, and :math:`X_t \sim p_t`.
+ Args:
+ x_0 (Tensor): source data point, shape (batch_size, ...).
+ x_1 (Tensor): target data point, shape (batch_size, ...).
+ t (Tensor): times in [0,1], shape (batch_size).
+
+ Returns:
+ DiscretePathSample: a conditional sample at :math:`X_t ~ p_t`.
+ """
+ self.assert_sample_shape(x_0=x_0, x_1=x_1, t=t)
+
+ sigma_t = self.scheduler(t).sigma_t
+
+ sigma_t = expand_tensor_like(input_tensor=sigma_t, expand_to=x_1)
+
+ source_indices = torch.rand(size=x_1.shape, device=x_1.device) < sigma_t
+ x_t = torch.where(condition=source_indices, input=x_0, other=x_1)
+
+ return DiscretePathSample(x_t=x_t, x_1=x_1, x_0=x_0, t=t)
+
+ def posterior_to_velocity(
+ self, posterior_logits: Tensor, x_t: Tensor, t: Tensor
+ ) -> Tensor:
+ r"""Convert the factorized posterior to velocity.
+
+ | given :math:`p(X_1|X_t)`. In the factorized case: :math:`\prod_i p(X_1^i | X_t)`.
+ | return :math:`u_t`.
+
+ Args:
+ posterior_logits (Tensor): logits of the x_1 posterior conditional on x_t, shape (..., vocab size).
+ x_t (Tensor): path sample at time t, shape (...).
+ t (Tensor): time in [0,1].
+
+ Returns:
+ Tensor: velocity.
+ """
+ posterior = torch.softmax(posterior_logits, dim=-1)
+ vocabulary_size = posterior.shape[-1]
+ x_t = F.one_hot(x_t, num_classes=vocabulary_size)
+ t = unsqueeze_to_match(source=t, target=x_t)
+
+ scheduler_output = self.scheduler(t)
+
+ kappa_t = scheduler_output.alpha_t
+ d_kappa_t = scheduler_output.d_alpha_t
+
+ return (d_kappa_t / (1 - kappa_t)) * (posterior - x_t)
diff --git a/scDFM/src/flow_matching/path/path.py b/scDFM/src/flow_matching/path/path.py
new file mode 100644
index 0000000000000000000000000000000000000000..0dfcd477eb0d4986c0022cee14b176aaa6e76e6d
--- /dev/null
+++ b/scDFM/src/flow_matching/path/path.py
@@ -0,0 +1,61 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from abc import ABC, abstractmethod
+
+from torch import Tensor
+
+from src.flow_matching.path.path_sample import PathSample
+
+
+class ProbPath(ABC):
+ r"""Abstract class, representing a probability path.
+
+ A probability path transforms the distribution :math:`p(X_0)` into :math:`p(X_1)` over :math:`t=0\rightarrow 1`.
+
+ The ``ProbPath`` class is designed to support model training in the flow matching framework. It supports two key functionalities: (1) sampling the conditional probability path and (2) conversion between various training objectives.
+ Here is a high-level example
+
+ .. code-block:: python
+
+ # Instantiate a probability path
+ my_path = ProbPath(...)
+
+ for x_0, x_1 in dataset:
+ # Sets t to a random value in [0,1]
+ t = torch.rand()
+
+ # Samples the conditional path X_t ~ p_t(X_t|X_0,X_1)
+ path_sample = my_path.sample(x_0=x_0, x_1=x_1, t=t)
+
+ # Optimizes the model. The loss function varies, depending on model and path.
+ loss(path_sample, my_model(x_t, t)).backward()
+
+ """
+
+ @abstractmethod
+ def sample(self, x_0: Tensor, x_1: Tensor, t: Tensor) -> PathSample:
+ r"""Sample from an abstract probability path:
+
+ | given :math:`(X_0,X_1) \sim \pi(X_0,X_1)`.
+ | returns :math:`X_0, X_1, X_t \sim p_t(X_t)`, and a conditional target :math:`Y`, all objects are under ``PathSample``.
+
+ Args:
+ x_0 (Tensor): source data point, shape (batch_size, ...).
+ x_1 (Tensor): target data point, shape (batch_size, ...).
+ t (Tensor): times in [0,1], shape (batch_size).
+
+ Returns:
+ PathSample: a conditional sample.
+ """
+
+ def assert_sample_shape(self, x_0: Tensor, x_1: Tensor, t: Tensor):
+ assert (
+ t.ndim == 1
+ ), f"The time vector t must have shape [batch_size]. Got {t.shape}."
+ assert (
+ t.shape[0] == x_0.shape[0] == x_1.shape[0]
+ ), f"Time t dimension must match the batch size [{x_1.shape[0]}]. Got {t.shape}"
diff --git a/scDFM/src/flow_matching/path/path_sample.py b/scDFM/src/flow_matching/path/path_sample.py
new file mode 100644
index 0000000000000000000000000000000000000000..867032e3dfe06fc5926316c9c4f133f57a025082
--- /dev/null
+++ b/scDFM/src/flow_matching/path/path_sample.py
@@ -0,0 +1,53 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from dataclasses import dataclass, field
+
+from torch import Tensor
+
+
+@dataclass
+class PathSample:
+ r"""Represents a sample of a conditional-flow generated probability path.
+
+ Attributes:
+ x_1 (Tensor): the target sample :math:`X_1`.
+ x_0 (Tensor): the source sample :math:`X_0`.
+ t (Tensor): the time sample :math:`t`.
+ x_t (Tensor): samples :math:`X_t \sim p_t(X_t)`, shape (batch_size, ...).
+ dx_t (Tensor): conditional target :math:`\frac{\partial X}{\partial t}`, shape: (batch_size, ...).
+
+ """
+
+ x_1: Tensor = field(metadata={"help": "target samples X_1 (batch_size, ...)."})
+ x_0: Tensor = field(metadata={"help": "source samples X_0 (batch_size, ...)."})
+ t: Tensor = field(metadata={"help": "time samples t (batch_size, ...)."})
+ x_t: Tensor = field(
+ metadata={"help": "samples x_t ~ p_t(X_t), shape (batch_size, ...)."}
+ )
+ dx_t: Tensor = field(
+ metadata={"help": "conditional target dX_t, shape: (batch_size, ...)."}
+ )
+
+
+@dataclass
+class DiscretePathSample:
+ """
+ Represents a sample of a conditional-flow generated discrete probability path.
+
+ Attributes:
+ x_1 (Tensor): the target sample :math:`X_1`.
+ x_0 (Tensor): the source sample :math:`X_0`.
+ t (Tensor): the time sample :math:`t`.
+ x_t (Tensor): the sample along the path :math:`X_t \sim p_t`.
+ """
+
+ x_1: Tensor = field(metadata={"help": "target samples X_1 (batch_size, ...)."})
+ x_0: Tensor = field(metadata={"help": "source samples X_0 (batch_size, ...)."})
+ t: Tensor = field(metadata={"help": "time samples t (batch_size, ...)."})
+ x_t: Tensor = field(
+ metadata={"help": "samples X_t ~ p_t(X_t), shape (batch_size, ...)."}
+ )
diff --git a/scDFM/src/flow_matching/path/scheduler/__init__.py b/scDFM/src/flow_matching/path/scheduler/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..f3b1a43f08b3c85a9c69ebe87b4fdeaf9eff29b1
--- /dev/null
+++ b/scDFM/src/flow_matching/path/scheduler/__init__.py
@@ -0,0 +1,29 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .schedule_transform import ScheduleTransformedModel
+from .scheduler import (
+ CondOTScheduler,
+ ConvexScheduler,
+ CosineScheduler,
+ LinearVPScheduler,
+ PolynomialConvexScheduler,
+ Scheduler,
+ SchedulerOutput,
+ VPScheduler,
+)
+
+__all__ = [
+ "CondOTScheduler",
+ "CosineScheduler",
+ "ConvexScheduler",
+ "PolynomialConvexScheduler",
+ "ScheduleTransformedModel",
+ "Scheduler",
+ "VPScheduler",
+ "LinearVPScheduler",
+ "SchedulerOutput",
+]
diff --git a/scDFM/src/flow_matching/path/scheduler/schedule_transform.py b/scDFM/src/flow_matching/path/scheduler/schedule_transform.py
new file mode 100644
index 0000000000000000000000000000000000000000..8e2e436ee959c30d52500a422dd60b41b06f065e
--- /dev/null
+++ b/scDFM/src/flow_matching/path/scheduler/schedule_transform.py
@@ -0,0 +1,148 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from torch import Tensor
+
+from src.flow_matching.path.scheduler.scheduler import Scheduler
+from src.flow_matching.utils import ModelWrapper
+
+
+class ScheduleTransformedModel(ModelWrapper):
+ """
+ Change of scheduler for a velocity model.
+
+ This class wraps a given velocity model and transforms its scheduling
+ to a new scheduler function. It modifies the time
+ dynamics of the model according to the new scheduler while maintaining
+ the original model's behavior.
+
+ Example:
+
+ .. code-block:: python
+
+ import torch
+ from flow_matching.path.scheduler import CondOTScheduler, CosineScheduler, ScheduleTransformedModel
+ from flow_matching.solver import ODESolver
+
+ # Initialize the model and schedulers
+ model = ...
+
+ original_scheduler = CondOTScheduler()
+ new_scheduler = CosineScheduler()
+
+ # Create the transformed model
+ transformed_model = ScheduleTransformedModel(
+ velocity_model=model,
+ original_scheduler=original_scheduler,
+ new_scheduler=new_scheduler
+ )
+
+ # Set up the solver
+ solver = ODESolver(velocity_model=transformed_model)
+
+ x_0 = torch.randn([10, 2]) # Example initial condition
+
+ x_1 = solver.sample(
+ time_steps=torch.tensor([0.0, 1.0]),
+ x_init=x_0,
+ step_size=1/1000
+ )[1]
+
+ Args:
+ velocity_model (ModelWrapper): The original velocity model to be transformed.
+ original_scheduler (Scheduler): The scheduler used by the original model. Must implement the snr_inverse function.
+ new_scheduler (Scheduler): The new scheduler to be applied to the model.
+ """
+
+ def __init__(
+ self,
+ velocity_model: ModelWrapper,
+ original_scheduler: Scheduler,
+ new_scheduler: Scheduler,
+ ):
+ super().__init__(model=velocity_model)
+ self.original_scheduler = original_scheduler
+ self.new_scheduler = new_scheduler
+
+ assert hasattr(self.original_scheduler, "snr_inverse") and callable(
+ getattr(self.original_scheduler, "snr_inverse")
+ ), "The original scheduler must have a callable 'snr_inverse' method."
+
+ def forward(self, x: Tensor, t: Tensor, **extras) -> Tensor:
+ r"""
+ Compute the transformed marginal velocity field for a new scheduler.
+ This method implements a post-training velocity scheduler change for
+ affine conditional flows. It transforms a generating marginal velocity
+ field :math:`u_t(x)` based on an original scheduler to a new marginal velocity
+ field :math:`\bar{u}_r(x)` based on a different scheduler, while maintaining
+ the same data coupling.
+ The transformation is based on the scale-time (ST) transformation
+ between the two conditional flows, defined as:
+
+ .. math::
+
+ \bar{X}_r = s_r X_{t_r},
+
+ where :math:`X_t` and :math:`\bar{X}_r` are defined by their respective schedulers.
+ The ST transformation is computed as:
+
+ .. math::
+
+ t_r = \rho^{-1}(\bar{\rho}(r)) \quad \text{and} \quad s_r = \frac{\bar{\sigma}_r}{\sigma_{t_r}}.
+
+ Here, :math:`\rho(t)` is the signal-to-noise ratio (SNR) defined as:
+
+ .. math::
+
+ \rho(t) = \frac{\alpha_t}{\sigma_t}.
+
+ :math:`\bar{\rho}(r)` is similarly defined for the new scheduler.
+ The marginal velocity for the new scheduler is then given by:
+
+ .. math::
+
+ \bar{u}_r(x) = \left(\frac{\dot{s}_r}{s_r}\right) x + s_r \dot{t}_r u_{t_r}\left(\frac{x}{s_r}\right).
+
+ Args:
+ x (Tensor): :math:`x_t`, the input tensor.
+ t (Tensor): The time tensor (denoted as :math:`r` above).
+ **extras: Additional arguments for the model.
+ Returns:
+ Tensor: The transformed velocity.
+ """
+ r = t
+
+ r_scheduler_output = self.new_scheduler(t=r)
+
+ alpha_r = r_scheduler_output.alpha_t
+ sigma_r = r_scheduler_output.sigma_t
+ d_alpha_r = r_scheduler_output.d_alpha_t
+ d_sigma_r = r_scheduler_output.d_sigma_t
+
+ t = self.original_scheduler.snr_inverse(alpha_r / sigma_r)
+
+ t_scheduler_output = self.original_scheduler(t=t)
+
+ alpha_t = t_scheduler_output.alpha_t
+ sigma_t = t_scheduler_output.sigma_t
+ d_alpha_t = t_scheduler_output.d_alpha_t
+ d_sigma_t = t_scheduler_output.d_sigma_t
+
+ s_r = sigma_r / sigma_t
+
+ dt_r = (
+ sigma_t
+ * sigma_t
+ * (sigma_r * d_alpha_r - alpha_r * d_sigma_r)
+ / (sigma_r * sigma_r * (sigma_t * d_alpha_t - alpha_t * d_sigma_t))
+ )
+
+ ds_r = (sigma_t * d_sigma_r - sigma_r * d_sigma_t * dt_r) / (sigma_t * sigma_t)
+
+ u_t = self.model(x=x / s_r, t=t, **extras)
+ u_r = ds_r * x / s_r + dt_r * s_r * u_t
+
+ return u_r
diff --git a/scDFM/src/flow_matching/path/scheduler/scheduler.py b/scDFM/src/flow_matching/path/scheduler/scheduler.py
new file mode 100644
index 0000000000000000000000000000000000000000..422618acd255bcc50779f1df60ce4f2666cea3f6
--- /dev/null
+++ b/scDFM/src/flow_matching/path/scheduler/scheduler.py
@@ -0,0 +1,199 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+
+from typing import Union
+
+import torch
+
+from torch import Tensor
+
+
+@dataclass
+class SchedulerOutput:
+ r"""Represents a sample of a conditional-flow generated probability path.
+
+ Attributes:
+ alpha_t (Tensor): :math:`\alpha_t`, shape (...).
+ sigma_t (Tensor): :math:`\sigma_t`, shape (...).
+ d_alpha_t (Tensor): :math:`\frac{\partial}{\partial t}\alpha_t`, shape (...).
+ d_sigma_t (Tensor): :math:`\frac{\partial}{\partial t}\sigma_t`, shape (...).
+
+ """
+
+ alpha_t: Tensor = field(metadata={"help": "alpha_t"})
+ sigma_t: Tensor = field(metadata={"help": "sigma_t"})
+ d_alpha_t: Tensor = field(metadata={"help": "Derivative of alpha_t."})
+ d_sigma_t: Tensor = field(metadata={"help": "Derivative of sigma_t."})
+
+
+class Scheduler(ABC):
+ """Base Scheduler class."""
+
+ @abstractmethod
+ def __call__(self, t: Tensor) -> SchedulerOutput:
+ r"""
+ Args:
+ t (Tensor): times in [0,1], shape (...).
+
+ Returns:
+ SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t`
+ """
+ ...
+
+ @abstractmethod
+ def snr_inverse(self, snr: Tensor) -> Tensor:
+ r"""
+ Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`.
+
+ Args:
+ snr (Tensor): The signal-to-noise, shape (...)
+
+ Returns:
+ Tensor: t, shape (...)
+ """
+ ...
+
+
+class ConvexScheduler(Scheduler):
+ @abstractmethod
+ def __call__(self, t: Tensor) -> SchedulerOutput:
+ """Scheduler for convex paths.
+
+ Args:
+ t (Tensor): times in [0,1], shape (...).
+
+ Returns:
+ SchedulerOutput: :math:`\alpha_t,\sigma_t,\frac{\partial}{\partial t}\alpha_t,\frac{\partial}{\partial t}\sigma_t`
+ """
+ ...
+
+ @abstractmethod
+ def kappa_inverse(self, kappa: Tensor) -> Tensor:
+ """
+ Computes :math:`t` from :math:`\kappa_t`.
+
+ Args:
+ kappa (Tensor): :math:`\kappa`, shape (...)
+
+ Returns:
+ Tensor: t, shape (...)
+ """
+ ...
+
+ def snr_inverse(self, snr: Tensor) -> Tensor:
+ r"""
+ Computes :math:`t` from the signal-to-noise ratio :math:`\frac{\alpha_t}{\sigma_t}`.
+
+ Args:
+ snr (Tensor): The signal-to-noise, shape (...)
+
+ Returns:
+ Tensor: t, shape (...)
+ """
+ kappa_t = snr / (1.0 + snr)
+
+ return self.kappa_inverse(kappa=kappa_t)
+
+
+class CondOTScheduler(ConvexScheduler):
+ """CondOT Scheduler."""
+
+ def __call__(self, t: Tensor) -> SchedulerOutput:
+ return SchedulerOutput(
+ alpha_t=t,
+ sigma_t=1 - t,
+ d_alpha_t=torch.ones_like(t),
+ d_sigma_t=-torch.ones_like(t),
+ )
+
+ def kappa_inverse(self, kappa: Tensor) -> Tensor:
+ return kappa
+
+
+class PolynomialConvexScheduler(ConvexScheduler):
+ """Polynomial Scheduler."""
+
+ def __init__(self, n: Union[float, int]) -> None:
+ assert isinstance(
+ n, (float, int)
+ ), f"`n` must be a float or int. Got {type(n)=}."
+ assert n > 0, f"`n` must be positive. Got {n=}."
+
+ self.n = n
+
+ def __call__(self, t: Tensor) -> SchedulerOutput:
+ return SchedulerOutput(
+ alpha_t=t**self.n,
+ sigma_t=1 - t**self.n,
+ d_alpha_t=self.n * (t ** (self.n - 1)),
+ d_sigma_t=-self.n * (t ** (self.n - 1)),
+ )
+
+ def kappa_inverse(self, kappa: Tensor) -> Tensor:
+ return torch.pow(kappa, 1.0 / self.n)
+
+
+class VPScheduler(Scheduler):
+ """Variance Preserving Scheduler."""
+
+ def __init__(self, beta_min: float = 0.1, beta_max: float = 20.0) -> None:
+ self.beta_min = beta_min
+ self.beta_max = beta_max
+ super().__init__()
+
+ def __call__(self, t: Tensor) -> SchedulerOutput:
+ b = self.beta_min
+ B = self.beta_max
+ T = 0.5 * (1 - t) ** 2 * (B - b) + (1 - t) * b
+ dT = -(1 - t) * (B - b) - b
+
+ return SchedulerOutput(
+ alpha_t=torch.exp(-0.5 * T),
+ sigma_t=torch.sqrt(1 - torch.exp(-T)),
+ d_alpha_t=-0.5 * dT * torch.exp(-0.5 * T),
+ d_sigma_t=0.5 * dT * torch.exp(-T) / torch.sqrt(1 - torch.exp(-T)),
+ )
+
+ def snr_inverse(self, snr: Tensor) -> Tensor:
+ T = -torch.log(snr**2 / (snr**2 + 1))
+ b = self.beta_min
+ B = self.beta_max
+ t = 1 - ((-b + torch.sqrt(b**2 + 2 * (B - b) * T)) / (B - b))
+ return t
+
+
+class LinearVPScheduler(Scheduler):
+ """Linear Variance Preserving Scheduler."""
+
+ def __call__(self, t: Tensor) -> SchedulerOutput:
+ return SchedulerOutput(
+ alpha_t=t,
+ sigma_t=(1 - t**2) ** 0.5,
+ d_alpha_t=torch.ones_like(t),
+ d_sigma_t=-t / (1 - t**2) ** 0.5,
+ )
+
+ def snr_inverse(self, snr: Tensor) -> Tensor:
+ return torch.sqrt(snr**2 / (1 + snr**2))
+
+
+class CosineScheduler(Scheduler):
+ """Cosine Scheduler."""
+
+ def __call__(self, t: Tensor) -> SchedulerOutput:
+ pi = torch.pi
+ return SchedulerOutput(
+ alpha_t=torch.sin(pi / 2 * t),
+ sigma_t=torch.cos(pi / 2 * t),
+ d_alpha_t=pi / 2 * torch.cos(pi / 2 * t),
+ d_sigma_t=-pi / 2 * torch.sin(pi / 2 * t),
+ )
+
+ def snr_inverse(self, snr: Tensor) -> Tensor:
+ return 2.0 * torch.atan(snr) / torch.pi
diff --git a/scDFM/src/flow_matching/solver/__init__.py b/scDFM/src/flow_matching/solver/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..6bd7b01c0992b5469d35885961f81da1bc9fe745
--- /dev/null
+++ b/scDFM/src/flow_matching/solver/__init__.py
@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .discrete_solver import MixtureDiscreteEulerSolver
+from .ode_solver import ODESolver
+from .riemannian_ode_solver import RiemannianODESolver
+from .solver import Solver
+
+__all__ = [
+ "ODESolver",
+ "Solver",
+ "ModelWrapper",
+ "MixtureDiscreteEulerSolver",
+ "RiemannianODESolver",
+]
diff --git a/scDFM/src/flow_matching/solver/discrete_solver.py b/scDFM/src/flow_matching/solver/discrete_solver.py
new file mode 100644
index 0000000000000000000000000000000000000000..2d0b04fbf5cb3622d809226c111d4cd38980ac5a
--- /dev/null
+++ b/scDFM/src/flow_matching/solver/discrete_solver.py
@@ -0,0 +1,260 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from contextlib import nullcontext
+from math import ceil
+from typing import Callable, Optional, Union
+
+import torch
+from torch import Tensor
+
+from torch.nn import functional as F
+
+from src.flow_matching.path import MixtureDiscreteProbPath
+
+from src.flow_matching.solver.solver import Solver
+from src.flow_matching.utils import categorical, ModelWrapper
+from .utils import get_nearest_times
+
+try:
+ from tqdm import tqdm
+
+ TQDM_AVAILABLE = True
+except ImportError:
+ TQDM_AVAILABLE = False
+
+
+class MixtureDiscreteEulerSolver(Solver):
+ r"""Solver that simulates the CTMC process :math:`(X_t)_{t_{\text{init}}\leq t\leq t_{\text{final}}}` defined by :math:`p_t` the marginal probability path of ``path``.
+ Given :math:`X_t \sim p_t`, the algorithm of solver step from :math:`t` to :math:`t+h` for the i-th coordinate is:
+
+ .. math::
+
+ \begin{align*}
+ & X_1^i \sim p_{1|t}^i(\cdot|X_t)\\
+ & \lambda^i \gets \sum_{x^i\ne X_t^i} u_t^i(x^i, X_t^i|X_1^i)\\
+ & Z^i_{\text{change}} \sim U[0,1]\\
+ & X_{t+h}^i \sim \begin{cases}
+ \frac{u_t^i(\cdot, X_t^i|X_1^i)}{\lambda^i}(1-\delta_{X_t^i}(\cdot)) \text{ if $Z^i_{\text{change}}\le 1-e^{-h\lambda^i}$}\\
+ \delta_{X_t^i}(\cdot) \text{ else }
+ \end{cases}
+ \end{align*}
+
+ Where :math:`p_{1|t}(\cdot|X_t)` is the output of ``model``, and the conditional probability velocity is of the mixture probability path is:
+
+ .. math::
+
+ u_t^i(x^i, y^i|x_1^i) = \hat{u}_t^i(x^i, y^i|x_1^i) + c_{\text{div\_free}}\left[\hat{u}_t^i(x^i, y^i|x_1^i) - \check{u}_t^i(x^i, y^i|x_1^i) \right],
+
+ where
+
+ .. math::
+ \hat{u}_t^i(x^i, y^i|x_1^i) = \frac{\dot{\kappa}_t}{1-\kappa_t} \left[ \delta_{x_1^i}(x^i) - \delta_{y^i}(x^i) \right],
+
+ and
+
+ .. math::
+
+ \check{u}_t^i(x^i, y^i|x_1^i) = \frac{\dot{\kappa}_t}{\kappa_t}\left[ \delta_{y^i}(x^i) - p(x^i) \right].
+
+ The source distribution :math:`p(x^i)` is given by ``p``.
+
+ Args:
+ model (ModelWrapper): trained with x-prediction, outputting posterior probabilities (in the range :math:`[0,1]`), output must be [..., vocabulary_size].
+ path (MixtureDiscreteProbPath): Probability path used for x-prediction training.
+ vocabulary_size (int): size of the discrete vocabulary.
+ source_distribution_p (Optional[Tensor], optional): Source distribution, must be of shape [vocabulary_size]. Required only when divergence-free term for the probability velocity is non-zero. Defaults to None.
+ """
+
+ def __init__(
+ self,
+ model: ModelWrapper,
+ path: MixtureDiscreteProbPath,
+ vocabulary_size: int,
+ source_distribution_p: Optional[Tensor] = None,
+ ):
+ super().__init__()
+ self.model = model
+ self.path = path
+ self.vocabulary_size = vocabulary_size
+
+ if source_distribution_p is not None:
+ assert source_distribution_p.shape == torch.Size(
+ [vocabulary_size]
+ ), f"Source distribution p dimension must match the vocabulary size {vocabulary_size}. Got {source_distribution_p.shape}."
+
+ self.source_distribution_p = source_distribution_p
+
+ @torch.no_grad()
+ def sample(
+ self,
+ x_init: Tensor,
+ step_size: Optional[float],
+ div_free: Union[float, Callable[[float], float]] = 0.0,
+ dtype_categorical: torch.dtype = torch.float32,
+ time_grid: Tensor = torch.tensor([0.0, 1.0]),
+ return_intermediates: bool = False,
+ verbose: bool = False,
+ **model_extras,
+ ) -> Tensor:
+ """
+ Sample a sequence of discrete values from the given model.
+
+ .. code-block:: python
+
+ import torch
+ from flow_matching.utils import ModelWrapper
+ from flow_matching.solver import MixtureDiscreteEulerSolver
+
+ class DummyModel(ModelWrapper):
+ def __init__(self):
+ super().__init__(None)
+ def forward(self, x: torch.Tensor, t: torch.Tensor, **extras) -> torch.Tensor:
+ return ...
+
+ model = DummyModel()
+ solver = MixtureDiscreteEulerSolver(model=model)
+
+ x_init = torch.LongTensor([122, 725])
+ step_size = 0.001
+ time_grid = torch.tensor([0.0, 1.0])
+
+ result = solver.sample(x_init=x_init, step_size=step_size, time_grid=time_grid)
+
+ Args:
+ x_init (Tensor): The initial state.
+ step_size (Optional[float]): If float then time discretization is uniform with the given step size. If None then time discretization is set to be time_grid.
+ div_free (Union[float, Callable[[float], float]]): The coefficient of the divergence-free term in the probability velocity. Can be either a float or a time dependent function. Defaults to 0.0.
+ dtype_categorical (torch.dtype): Precision to use for categorical sampler. Defaults to torch.float32.
+ time_grid (Tensor): The CTMC process is solved in the interval [time_grid[0], time_grid[-1]] and if step_size is None then time discretization is set by the time grid. Defaults to torch.tensor([0.0,1.0]).
+ return_intermediates (bool): If True then return intermediate time steps according to time_grid. Defaults to False.
+ verbose (bool): Whether to print progress bars. Defaults to False.
+ **model_extras: Additional input for the model.
+
+ Returns:
+ Tensor: The sampled sequence of discrete values.
+
+ Raises:
+ ImportError: To run in verbose mode, tqdm must be installed.
+ """
+ if not div_free == 0.0:
+ assert (
+ self.source_distribution_p is not None
+ ), "Source distribution p must be specified in order to add a divergence-free term to the probability velocity."
+
+ # Initialize the current state `x_t` with the initial state `X_0`.
+ time_grid = time_grid.to(device=x_init.device)
+
+ if step_size is None:
+ # If step_size is None then set the t discretization to time_grid.
+ t_discretization = time_grid
+ n_steps = len(time_grid) - 1
+ else:
+ # If step_size is float then t discretization is uniform with step size set by step_size.
+ t_init = time_grid[0].item()
+ t_final = time_grid[-1].item()
+ assert (
+ t_final - t_init
+ ) > step_size, f"Time interval [time_grid[0], time_grid[-1]] must be larger than step_size. Got a time interval [{t_init}, {t_final}] and step_size {step_size}."
+
+ n_steps = ceil((t_final - t_init) / step_size)
+ t_discretization = torch.tensor(
+ [t_init + step_size * i for i in range(n_steps)] + [t_final],
+ device=x_init.device,
+ )
+
+ if return_intermediates:
+ # get order of intermediate steps:
+ order = torch.argsort(time_grid)
+ # Compute intermediate steps to return via nearest points in t_discretization to time_grid.
+ time_grid = get_nearest_times(
+ time_grid=time_grid, t_discretization=t_discretization
+ )
+
+ x_t = x_init.clone()
+ steps_counter = 0
+ res = []
+
+ if return_intermediates:
+ res = [x_init.clone()]
+
+ if verbose:
+ if not TQDM_AVAILABLE:
+ raise ImportError(
+ "tqdm is required for verbose mode. Please install it."
+ )
+ ctx = tqdm(total=t_final, desc=f"NFE: {steps_counter}")
+ else:
+ ctx = nullcontext()
+
+ with ctx:
+ for i in range(n_steps):
+ t = t_discretization[i : i + 1]
+ h = t_discretization[i + 1 : i + 2] - t_discretization[i : i + 1]
+
+ # Sample x_1 ~ p_1|t( \cdot |x_t)
+ p_1t = self.model(x=x_t, t=t.repeat(x_t.shape[0]), **model_extras)
+ x_1 = categorical(p_1t.to(dtype=dtype_categorical))
+
+ # Checks if final step
+ if i == n_steps - 1:
+ x_t = x_1
+ else:
+ # Compute u_t(x|x_t,x_1)
+ scheduler_output = self.path.scheduler(t=t)
+
+ k_t = scheduler_output.alpha_t
+ d_k_t = scheduler_output.d_alpha_t
+
+ delta_1 = F.one_hot(x_1, num_classes=self.vocabulary_size).to(
+ k_t.dtype
+ )
+ u = d_k_t / (1 - k_t) * delta_1
+
+ # Add divergence-free part
+ div_free_t = div_free(t) if callable(div_free) else div_free
+
+ if div_free_t > 0:
+ p_0 = self.source_distribution_p[(None,) * x_t.dim()]
+ u = u + div_free_t * d_k_t / (k_t * (1 - k_t)) * (
+ (1 - k_t) * p_0 + k_t * delta_1
+ )
+
+ # Set u_t(x_t|x_t,x_1) = 0
+ delta_t = F.one_hot(x_t, num_classes=self.vocabulary_size)
+ u = torch.where(
+ delta_t.to(dtype=torch.bool), torch.zeros_like(u), u
+ )
+
+ # Sample x_t ~ u_t( \cdot |x_t,x_1)
+ intensity = u.sum(dim=-1) # Assuming u_t(xt|xt,x1) := 0
+ mask_jump = torch.rand(
+ size=x_t.shape, device=x_t.device
+ ) < 1 - torch.exp(-h * intensity)
+
+ if mask_jump.sum() > 0:
+ x_t[mask_jump] = categorical(
+ u[mask_jump].to(dtype=dtype_categorical)
+ )
+
+ steps_counter += 1
+ t = t + h
+
+ if return_intermediates and (t in time_grid):
+ res.append(x_t.clone())
+
+ if verbose:
+ ctx.n = t.item()
+ ctx.refresh()
+ ctx.set_description(f"NFE: {steps_counter}")
+
+ if return_intermediates:
+ if step_size is None:
+ return torch.stack(res, dim=0)
+ else:
+ return torch.stack(res, dim=0)[order]
+ else:
+ return x_t
diff --git a/scDFM/src/flow_matching/solver/ode_solver.py b/scDFM/src/flow_matching/solver/ode_solver.py
new file mode 100644
index 0000000000000000000000000000000000000000..60cf30268735368e91e1f4fa746f27e4b672bf6a
--- /dev/null
+++ b/scDFM/src/flow_matching/solver/ode_solver.py
@@ -0,0 +1,197 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable, Optional, Sequence, Tuple, Union
+
+import torch
+from torch import Tensor
+from torchdiffeq import odeint
+
+from src.flow_matching.solver.solver import Solver
+from src.flow_matching.utils import gradient, ModelWrapper
+
+
+class ODESolver(Solver):
+ """A class to solve ordinary differential equations (ODEs) using a specified velocity model.
+
+ This class utilizes a velocity field model to solve ODEs over a given time grid using numerical ode solvers.
+
+ Args:
+ velocity_model (Union[ModelWrapper, Callable]): a velocity field model receiving :math:`(x,t)` and returning :math:`u_t(x)`
+ """
+
+ def __init__(self, velocity_model: Union[ModelWrapper, Callable]):
+ super().__init__()
+ self.velocity_model = velocity_model
+
+ def sample(
+ self,
+ x_init: Tensor,
+ step_size: Optional[float],
+ method: str = "euler",
+ atol: float = 1e-5,
+ rtol: float = 1e-5,
+ time_grid: Tensor = torch.tensor([0.0, 1.0]),
+ return_intermediates: bool = False,
+ enable_grad: bool = False,
+ **model_extras,
+ ) -> Union[Tensor, Sequence[Tensor]]:
+ r"""Solve the ODE with the velocity field.
+
+ Example:
+
+ .. code-block:: python
+
+ import torch
+ from flow_matching.utils import ModelWrapper
+ from flow_matching.solver import ODESolver
+
+ class DummyModel(ModelWrapper):
+ def __init__(self):
+ super().__init__(None)
+
+ def forward(self, x: torch.Tensor, t: torch.Tensor, **extras) -> torch.Tensor:
+ return torch.ones_like(x) * 3.0 * t**2
+
+ velocity_model = DummyModel()
+ solver = ODESolver(velocity_model=velocity_model)
+ x_init = torch.tensor([0.0, 0.0])
+ step_size = 0.001
+ time_grid = torch.tensor([0.0, 1.0])
+
+ result = solver.sample(x_init=x_init, step_size=step_size, time_grid=time_grid)
+
+ Args:
+ x_init (Tensor): initial conditions (e.g., source samples :math:`X_0 \sim p`). Shape: [batch_size, ...].
+ step_size (Optional[float]): The step size. Must be None for adaptive step solvers.
+ method (str): A method supported by torchdiffeq. Defaults to "euler". Other commonly used solvers are "dopri5", "midpoint" and "heun3". For a complete list, see torchdiffeq.
+ atol (float): Absolute tolerance, used for adaptive step solvers.
+ rtol (float): Relative tolerance, used for adaptive step solvers.
+ time_grid (Tensor): The process is solved in the interval [min(time_grid, max(time_grid)] and if step_size is None then time discretization is set by the time grid. May specify a descending time_grid to solve in the reverse direction. Defaults to torch.tensor([0.0, 1.0]).
+ return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Defaults to False.
+ enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False.
+ **model_extras: Additional input for the model.
+
+ Returns:
+ Union[Tensor, Sequence[Tensor]]: The last timestep when return_intermediates=False, otherwise all values specified in time_grid.
+ """
+
+ time_grid = time_grid.to(x_init.device)
+
+ def ode_func(t, x):
+ return self.velocity_model(x=x, t=t, **model_extras)
+
+ ode_opts = {"step_size": step_size} if step_size is not None else {}
+
+ with torch.set_grad_enabled(enable_grad):
+ # Approximate ODE solution with numerical ODE solver
+ sol = odeint(
+ ode_func,
+ x_init,
+ time_grid,
+ method=method,
+ options=ode_opts,
+ atol=atol,
+ rtol=rtol,
+ )
+
+ if return_intermediates:
+ return sol
+ else:
+ return sol[-1]
+
+ def compute_likelihood(
+ self,
+ x_1: Tensor,
+ log_p0: Callable[[Tensor], Tensor],
+ step_size: Optional[float],
+ method: str = "euler",
+ atol: float = 1e-5,
+ rtol: float = 1e-5,
+ time_grid: Tensor = torch.tensor([1.0, 0.0]),
+ return_intermediates: bool = False,
+ exact_divergence: bool = False,
+ enable_grad: bool = False,
+ **model_extras,
+ ) -> Union[Tuple[Tensor, Tensor], Tuple[Sequence[Tensor], Tensor]]:
+ r"""Solve for log likelihood given a target sample at :math:`t=0`.
+
+ Works similarly to sample, but solves the ODE in reverse to compute the log-likelihood. The velocity model must be differentiable with respect to x.
+ The function assumes log_p0 is the log probability of the source distribution at :math:`t=0`.
+
+ Args:
+ x_1 (Tensor): target sample (e.g., samples :math:`X_1 \sim p_1`).
+ log_p0 (Callable[[Tensor], Tensor]): Log probability function of the source distribution.
+ step_size (Optional[float]): The step size. Must be None for adaptive step solvers.
+ method (str): A method supported by torchdiffeq. Defaults to "euler". Other commonly used solvers are "dopri5", "midpoint" and "heun3". For a complete list, see torchdiffeq.
+ atol (float): Absolute tolerance, used for adaptive step solvers.
+ rtol (float): Relative tolerance, used for adaptive step solvers.
+ time_grid (Tensor): If step_size is None then time discretization is set by the time grid. Must start at 1.0 and end at 0.0, otherwise the likelihood computation is not valid. Defaults to torch.tensor([1.0, 0.0]).
+ return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Otherwise only return the final sample. Defaults to False.
+ exact_divergence (bool): Whether to compute the exact divergence or use the Hutchinson estimator.
+ enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False.
+ **model_extras: Additional input for the model.
+
+ Returns:
+ Union[Tuple[Tensor, Tensor], Tuple[Sequence[Tensor], Tensor]]: Samples at time_grid and log likelihood values of given x_1.
+ """
+ assert (
+ time_grid[0] == 1.0 and time_grid[-1] == 0.0
+ ), f"Time grid must start at 1.0 and end at 0.0. Got {time_grid}"
+
+ # Fix the random projection for the Hutchinson divergence estimator
+ if not exact_divergence:
+ z = (torch.randn_like(x_1).to(x_1.device) < 0) * 2.0 - 1.0
+
+ def ode_func(x, t):
+ return self.velocity_model(x=x, t=t, **model_extras)
+
+ def dynamics_func(t, states):
+ xt = states[0]
+ with torch.set_grad_enabled(True):
+ xt.requires_grad_()
+ ut = ode_func(xt, t)
+
+ if exact_divergence:
+ # Compute exact divergence
+ div = 0
+ for i in range(ut.flatten(1).shape[1]):
+ div += gradient(ut[:, i], xt, create_graph=True)[:, i]
+ else:
+ # Compute Hutchinson divergence estimator E[z^T D_x(ut) z]
+ ut_dot_z = torch.einsum(
+ "ij,ij->i", ut.flatten(start_dim=1), z.flatten(start_dim=1)
+ )
+ grad_ut_dot_z = gradient(ut_dot_z, xt)
+ div = torch.einsum(
+ "ij,ij->i",
+ grad_ut_dot_z.flatten(start_dim=1),
+ z.flatten(start_dim=1),
+ )
+
+ return ut.detach(), div.detach()
+
+ y_init = (x_1, torch.zeros(x_1.shape[0], device=x_1.device))
+ ode_opts = {"step_size": step_size} if step_size is not None else {}
+
+ with torch.set_grad_enabled(enable_grad):
+ sol, log_det = odeint(
+ dynamics_func,
+ y_init,
+ time_grid,
+ method=method,
+ options=ode_opts,
+ atol=atol,
+ rtol=rtol,
+ )
+
+ x_source = sol[-1]
+ source_log_p = log_p0(x_source)
+
+ if return_intermediates:
+ return sol, source_log_p + log_det[-1]
+ else:
+ return sol[-1], source_log_p + log_det[-1]
diff --git a/scDFM/src/flow_matching/solver/riemannian_ode_solver.py b/scDFM/src/flow_matching/solver/riemannian_ode_solver.py
new file mode 100644
index 0000000000000000000000000000000000000000..eaeb682e9bab49fd2625d77f6b71a19250a24d71
--- /dev/null
+++ b/scDFM/src/flow_matching/solver/riemannian_ode_solver.py
@@ -0,0 +1,261 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+from typing import Callable
+
+import torch
+from torch import Tensor
+
+from src.flow_matching.solver.solver import Solver
+from src.flow_matching.utils import ModelWrapper
+from src.flow_matching.utils.manifolds import geodesic, Manifold
+
+try:
+ from tqdm import tqdm
+
+ TQDM_AVAILABLE = True
+except ImportError:
+ TQDM_AVAILABLE = False
+
+
+class RiemannianODESolver(Solver):
+ r"""Riemannian ODE solver
+ Initialize the ``RiemannianODESolver``.
+
+ Args:
+ manifold (Manifold): the manifold to solve on.
+ velocity_model (ModelWrapper): a velocity field model receiving :math:`(x,t)`
+ and returning :math:`u_t(x)` which is assumed to lie on the tangent plane at `x`.
+ """
+
+ def __init__(self, manifold: Manifold, velocity_model: ModelWrapper):
+ super().__init__()
+ self.manifold = manifold
+ self.velocity_model = velocity_model
+
+ def sample(
+ self,
+ x_init: Tensor,
+ step_size: float,
+ projx: bool = True,
+ proju: bool = True,
+ method: str = "euler",
+ time_grid: Tensor = torch.tensor([0.0, 1.0]),
+ return_intermediates: bool = False,
+ verbose: bool = False,
+ enable_grad: bool = False,
+ **model_extras,
+ ) -> Tensor:
+ r"""Solve the ODE with the `velocity_field` on the manifold.
+
+ Args:
+ x_init (Tensor): initial conditions (e.g., source samples :math:`X_0 \sim p`).
+ step_size (float): The step size.
+ projx (bool): Whether to project the point onto the manifold at each step. Defaults to True.
+ proju (bool): Whether to project the vector field onto the tangent plane at each step. Defaults to True.
+ method (str): One of ["euler", "midpoint", "rk4"]. Defaults to "euler".
+ time_grid (Tensor, optional): The process is solved in the interval [min(time_grid, max(time_grid)] and if step_size is None then time discretization is set by the time grid. Defaults to torch.tensor([0.0,1.0]).
+ return_intermediates (bool, optional): If True then return intermediate time steps according to time_grid. Defaults to False.
+ verbose (bool, optional): Whether to print progress bars. Defaults to False.
+ enable_grad (bool, optional): Whether to compute gradients during sampling. Defaults to False.
+ **model_extras: Additional input for the model.
+
+ Returns:
+ Tensor: The sampled sequence. Defaults to returning samples at :math:`t=1`.
+
+ Raises:
+ ImportError: To run in verbose mode, tqdm must be installed.
+ """
+ step_fns = {
+ "euler": _euler_step,
+ "midpoint": _midpoint_step,
+ "rk4": _rk4_step,
+ }
+ assert method in step_fns.keys(), f"Unknown method {method}"
+ step_fn = step_fns[method]
+
+ def velocity_func(x, t):
+ return self.velocity_model(x=x, t=t, **model_extras)
+
+ # --- Factor this out.
+ time_grid = torch.sort(time_grid.to(device=x_init.device)).values
+
+ if step_size is None:
+ # If step_size is None then set the t discretization to time_grid.
+ t_discretization = time_grid
+ n_steps = len(time_grid) - 1
+ else:
+ # If step_size is float then t discretization is uniform with step size set by step_size.
+ t_init = time_grid[0].item()
+ t_final = time_grid[-1].item()
+ assert (
+ t_final - t_init
+ ) > step_size, f"Time interval [min(time_grid), max(time_grid)] must be larger than step_size. Got a time interval [{t_init}, {t_final}] and step_size {step_size}."
+
+ n_steps = math.ceil((t_final - t_init) / step_size)
+ t_discretization = torch.tensor(
+ [step_size * i for i in range(n_steps)] + [t_final],
+ device=x_init.device,
+ )
+ # ---
+ t0s = t_discretization[:-1]
+
+ if verbose:
+ if not TQDM_AVAILABLE:
+ raise ImportError(
+ "tqdm is required for verbose mode. Please install it."
+ )
+ t0s = tqdm(t0s)
+
+ if return_intermediates:
+ xts = []
+ i_ret = 0
+
+ with torch.set_grad_enabled(enable_grad):
+ xt = x_init
+ for t0, t1 in zip(t0s, t_discretization[1:]):
+ dt = t1 - t0
+ xt_next = step_fn(
+ velocity_func,
+ xt,
+ t0,
+ dt,
+ manifold=self.manifold,
+ projx=projx,
+ proju=proju,
+ )
+ if return_intermediates:
+ while (
+ i_ret < len(time_grid)
+ and t0 <= time_grid[i_ret]
+ and time_grid[i_ret] <= t1
+ ):
+ xts.append(
+ interp(self.manifold, xt, xt_next, t0, t1, time_grid[i_ret])
+ )
+ i_ret += 1
+ xt = xt_next
+
+ if return_intermediates:
+ return torch.stack(xts, dim=0)
+ else:
+ return xt
+
+
+def interp(manifold, xt, xt_next, t, t_next, t_ret):
+ return geodesic(manifold, xt, xt_next)(
+ (t_ret - t) / (t_next - t).reshape(1)
+ ).reshape_as(xt)
+
+
+def _euler_step(
+ velocity_model: Callable,
+ xt: Tensor,
+ t0: Tensor,
+ dt: Tensor,
+ manifold: Manifold,
+ projx: bool = True,
+ proju: bool = True,
+) -> Tensor:
+ r"""Perform an Euler step on a manifold.
+
+ Args:
+ velocity_model (Callable): the velocity model
+ xt (Tensor): tensor containing the state at time t0
+ t0 (Tensor): the time at which this step is taken
+ dt (Tensor): the step size
+ manifold (Manifold): a manifold object
+ projx (bool, optional): whether to project the state onto the manifold. Defaults to True.
+ proju (bool, optional): whether to project the velocity onto the tangent plane. Defaults to True.
+
+ Returns:
+ Tensor: tensor containing the state after the step
+ """
+ velocity_fn = lambda x, t: (
+ manifold.proju(x, velocity_model(x, t)) if proju else velocity_model(x, t)
+ )
+ projx_fn = lambda x: manifold.projx(x) if projx else x
+
+ vt = velocity_fn(xt, t0)
+
+ xt = xt + dt * vt
+
+ return projx_fn(xt)
+
+
+def _midpoint_step(
+ velocity_model: Callable,
+ xt: Tensor,
+ t0: Tensor,
+ dt: Tensor,
+ manifold: Manifold,
+ projx: bool = True,
+ proju: bool = True,
+) -> Tensor:
+ r"""Perform a midpoint step on a manifold.
+
+ Args:
+ velocity_model (Callable): the velocity model
+ xt (Tensor): tensor containing the state at time t0
+ t0 (Tensor): the time at which this step is taken
+ dt (Tensor): the step size
+ manifold (Manifold): a manifold object
+ projx (bool, optional): whether to project the state onto the manifold. Defaults to True.
+ proju (bool, optional): whether to project the velocity onto the tangent plane. Defaults to True.
+
+ Returns:
+ Tensor: tensor containing the state after the step
+ """
+ velocity_fn = lambda x, t: (
+ manifold.proju(x, velocity_model(x, t)) if proju else velocity_model(x, t)
+ )
+ projx_fn = lambda x: manifold.projx(x) if projx else x
+
+ half_dt = 0.5 * dt
+ vt = velocity_fn(xt, t0)
+ x_mid = xt + half_dt * vt
+ x_mid = projx_fn(x_mid)
+
+ xt = xt + dt * velocity_fn(x_mid, t0 + half_dt)
+
+ return projx_fn(xt)
+
+
+def _rk4_step(
+ velocity_model: Callable,
+ xt: Tensor,
+ t0: Tensor,
+ dt: Tensor,
+ manifold: Manifold,
+ projx: bool = True,
+ proju: bool = True,
+) -> Tensor:
+ r"""Perform an RK4 step on a manifold.
+
+ Args:
+ velocity_model (Callable): the velocity model
+ xt (Tensor): tensor containing the state at time t0
+ t0 (Tensor): the time at which this step is taken
+ dt (Tensor): the step size
+ manifold (Manifold): a manifold object
+ projx (bool, optional): whether to project the state onto the manifold. Defaults to True.
+ proju (bool, optional): whether to project the velocity onto the tangent plane. Defaults to True.
+
+ Returns:
+ Tensor: tensor containing the state after the step
+ """
+ velocity_fn = lambda x, t: (
+ manifold.proju(x, velocity_model(x, t)) if proju else velocity_model(x, t)
+ )
+ projx_fn = lambda x: manifold.projx(x) if projx else x
+
+ k1 = velocity_fn(xt, t0)
+ k2 = velocity_fn(projx_fn(xt + dt * k1 / 3), t0 + dt / 3)
+ k3 = velocity_fn(projx_fn(xt + dt * (k2 - k1 / 3)), t0 + dt * 2 / 3)
+ k4 = velocity_fn(projx_fn(xt + dt * (k1 - k2 + k3)), t0 + dt)
+
+ return projx_fn(xt + (k1 + 3 * (k2 + k3) + k4) * dt * 0.125)
diff --git a/scDFM/src/flow_matching/solver/solver.py b/scDFM/src/flow_matching/solver/solver.py
new file mode 100644
index 0000000000000000000000000000000000000000..4819e1c0ff148c411648c6e248380f471efd0a6c
--- /dev/null
+++ b/scDFM/src/flow_matching/solver/solver.py
@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from abc import ABC, abstractmethod
+
+from torch import nn, Tensor
+
+
+class Solver(ABC, nn.Module):
+ """Abstract base class for solvers."""
+
+ @abstractmethod
+ def sample(self, x_0: Tensor = None) -> Tensor:
+ ...
diff --git a/scDFM/src/flow_matching/solver/utils.py b/scDFM/src/flow_matching/solver/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..f3a34ee4b0ca4c3ef64dfa1b55cb13da1187b965
--- /dev/null
+++ b/scDFM/src/flow_matching/solver/utils.py
@@ -0,0 +1,19 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch import Tensor
+
+
+def get_nearest_times(time_grid: Tensor, t_discretization: Tensor) -> Tensor:
+ distances = torch.cdist(
+ time_grid.unsqueeze(1),
+ t_discretization.unsqueeze(1),
+ compute_mode="donot_use_mm_for_euclid_dist",
+ )
+ nearest_indices = distances.argmin(dim=1)
+
+ return t_discretization[nearest_indices]
diff --git a/scDFM/src/flow_matching/utils/__init__.py b/scDFM/src/flow_matching/utils/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..0085c44c5e044d5a7dd0c2bde9755e9d3c18c0b1
--- /dev/null
+++ b/scDFM/src/flow_matching/utils/__init__.py
@@ -0,0 +1,17 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .categorical_sampler import categorical
+from .model_wrapper import ModelWrapper
+from .utils import expand_tensor_like, gradient, unsqueeze_to_match
+
+__all__ = [
+ "unsqueeze_to_match",
+ "expand_tensor_like",
+ "gradient",
+ "categorical",
+ "ModelWrapper",
+]
diff --git a/scDFM/src/flow_matching/utils/categorical_sampler.py b/scDFM/src/flow_matching/utils/categorical_sampler.py
new file mode 100644
index 0000000000000000000000000000000000000000..70937af5af1da4c3c2b18a85745197e240724a96
--- /dev/null
+++ b/scDFM/src/flow_matching/utils/categorical_sampler.py
@@ -0,0 +1,23 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch import Tensor
+
+
+def categorical(probs: Tensor) -> Tensor:
+ r"""Categorical sampler according to weights in the last dimension of ``probs`` using :func:`torch.multinomial`.
+
+ Args:
+ probs (Tensor): probabilities.
+
+ Returns:
+ Tensor: Samples.
+ """
+
+ return torch.multinomial(probs.flatten(0, -2), 1, replacement=True).view(
+ *probs.shape[:-1]
+ )
diff --git a/scDFM/src/flow_matching/utils/manifolds/__init__.py b/scDFM/src/flow_matching/utils/manifolds/__init__.py
new file mode 100644
index 0000000000000000000000000000000000000000..1148872f29c5e42e2578eab31751abaf6def51da
--- /dev/null
+++ b/scDFM/src/flow_matching/utils/manifolds/__init__.py
@@ -0,0 +1,18 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from .manifold import Euclidean, Manifold
+from .sphere import Sphere
+from .torus import FlatTorus
+from .utils import geodesic
+
+__all__ = [
+ "Euclidean",
+ "Manifold",
+ "Sphere",
+ "FlatTorus",
+ "geodesic",
+]
diff --git a/scDFM/src/flow_matching/utils/manifolds/manifold.py b/scDFM/src/flow_matching/utils/manifolds/manifold.py
new file mode 100644
index 0000000000000000000000000000000000000000..52a6a1bc6ea7d83e3aeefd335edb19a842aeca21
--- /dev/null
+++ b/scDFM/src/flow_matching/utils/manifolds/manifold.py
@@ -0,0 +1,93 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import abc
+
+import torch.nn as nn
+from torch import Tensor
+
+
+class Manifold(nn.Module, metaclass=abc.ABCMeta):
+ """A manifold class that contains projection operations and logarithm and exponential maps."""
+
+ @abc.abstractmethod
+ def expmap(self, x: Tensor, u: Tensor) -> Tensor:
+ r"""Computes exponential map :math:`\exp_x(u)`.
+
+ Args:
+ x (Tensor): point on the manifold
+ u (Tensor): tangent vector at point :math:`x`
+
+ Raises:
+ NotImplementedError: if not implemented
+
+ Returns:
+ Tensor: transported point
+ """
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def logmap(self, x: Tensor, y: Tensor) -> Tensor:
+ r"""Computes logarithmic map :math:`\log_x(y)`.
+
+ Args:
+ x (Tensor): point on the manifold
+ y (Tensor): point on the manifold
+
+ Raises:
+ NotImplementedError: if not implemented
+
+ Returns:
+ Tensor: tangent vector at point :math:`x`
+ """
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def projx(self, x: Tensor) -> Tensor:
+ """Project point :math:`x` on the manifold.
+
+ Args:
+ x (Tensor): point to be projected
+
+ Raises:
+ NotImplementedError: if not implemented
+
+ Returns:
+ Tensor: projected point on the manifold
+ """
+ raise NotImplementedError
+
+ @abc.abstractmethod
+ def proju(self, x: Tensor, u: Tensor) -> Tensor:
+ """Project vector :math:`u` on a tangent space for :math:`x`.
+
+ Args:
+ x (Tensor): point on the manifold
+ u (Tensor): vector to be projected
+
+ Raises:
+ NotImplementedError: if not implemented
+
+ Returns:
+ Tensor: projected tangent vector
+ """
+ raise NotImplementedError
+
+
+class Euclidean(Manifold):
+ """The Euclidean manifold."""
+
+ def expmap(self, x: Tensor, u: Tensor) -> Tensor:
+ return x + u
+
+ def logmap(self, x: Tensor, y: Tensor) -> Tensor:
+ return y - x
+
+ def projx(self, x: Tensor) -> Tensor:
+ return x
+
+ def proju(self, x: Tensor, u: Tensor) -> Tensor:
+ return u
diff --git a/scDFM/src/flow_matching/utils/manifolds/sphere.py b/scDFM/src/flow_matching/utils/manifolds/sphere.py
new file mode 100644
index 0000000000000000000000000000000000000000..23122926530add5b72420f31011476c5e0c5e7ad
--- /dev/null
+++ b/scDFM/src/flow_matching/utils/manifolds/sphere.py
@@ -0,0 +1,45 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import torch
+from torch import Tensor
+
+from src.flow_matching.utils.manifolds import Manifold
+
+
+class Sphere(Manifold):
+ """Represents a hyperpshere in :math:`R^D`. Isometric to the product of 1-D spheres."""
+
+ EPS = {torch.float32: 1e-4, torch.float64: 1e-7}
+
+ def expmap(self, x: Tensor, u: Tensor) -> Tensor:
+ norm_u = u.norm(dim=-1, keepdim=True)
+ exp = x * torch.cos(norm_u) + u * torch.sin(norm_u) / norm_u
+ retr = self.projx(x + u)
+ cond = norm_u > self.EPS[norm_u.dtype]
+
+ return torch.where(cond, exp, retr)
+
+ def logmap(self, x: Tensor, y: Tensor) -> Tensor:
+ u = self.proju(x, y - x)
+ dist = self.dist(x, y, keepdim=True)
+ cond = dist.gt(self.EPS[x.dtype])
+ result = torch.where(
+ cond,
+ u * dist / u.norm(dim=-1, keepdim=True).clamp_min(self.EPS[x.dtype]),
+ u,
+ )
+ return result
+
+ def projx(self, x: Tensor) -> Tensor:
+ return x / x.norm(dim=-1, keepdim=True)
+
+ def proju(self, x: Tensor, u: Tensor) -> Tensor:
+ return u - (x * u).sum(dim=-1, keepdim=True) * x
+
+ def dist(self, x: Tensor, y: Tensor, *, keepdim=False) -> Tensor:
+ inner = (x * y).sum(-1, keepdim=keepdim)
+ return torch.acos(inner)
diff --git a/scDFM/src/flow_matching/utils/manifolds/torus.py b/scDFM/src/flow_matching/utils/manifolds/torus.py
new file mode 100644
index 0000000000000000000000000000000000000000..8c97703b6b565e5d106b0d61def6470acb2f60de
--- /dev/null
+++ b/scDFM/src/flow_matching/utils/manifolds/torus.py
@@ -0,0 +1,28 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+import math
+
+import torch
+from torch import Tensor
+
+from src.flow_matching.utils.manifolds import Manifold
+
+
+class FlatTorus(Manifold):
+ r"""Represents a flat torus on the :math:`[0, 2\pi]^D` subspace. Isometric to the product of 1-D spheres."""
+
+ def expmap(self, x: Tensor, u: Tensor) -> Tensor:
+ return (x + u) % (2 * math.pi)
+
+ def logmap(self, x: Tensor, y: Tensor) -> Tensor:
+ return torch.atan2(torch.sin(y - x), torch.cos(y - x))
+
+ def projx(self, x: Tensor) -> Tensor:
+ return x % (2 * math.pi)
+
+ def proju(self, x: Tensor, u: Tensor) -> Tensor:
+ return u
diff --git a/scDFM/src/flow_matching/utils/manifolds/utils.py b/scDFM/src/flow_matching/utils/manifolds/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..af19591931b45831b70b71af02090aded937c422
--- /dev/null
+++ b/scDFM/src/flow_matching/utils/manifolds/utils.py
@@ -0,0 +1,45 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Callable
+
+import torch
+from torch import Tensor
+
+from src.flow_matching.utils.manifolds import Manifold
+
+
+def geodesic(
+ manifold: Manifold, start_point: Tensor, end_point: Tensor
+) -> Callable[[Tensor], Tensor]:
+ """Generate parameterized function for geodesic curve.
+
+ Args:
+ manifold (Manifold): the manifold to compute geodesic on.
+ start_point (Tensor): point on the manifold at :math:`t=0`.
+ end_point (Tensor): point on the manifold at :math:`t=1`.
+
+ Returns:
+ Callable[[Tensor], Tensor]: a function that takes in :math:`t` and outputs the geodesic at time :math:`t`.
+ """
+
+ shooting_tangent_vec = manifold.logmap(start_point, end_point)
+
+ def path(t: Tensor) -> Tensor:
+ """Generate parameterized function for geodesic curve.
+
+ Args:
+ t (Tensor): Times at which to compute points of the geodesics.
+
+ Returns:
+ Tensor: geodesic path evaluated at time t.
+ """
+ tangent_vecs = torch.einsum("i,...k->...ik", t, shooting_tangent_vec)
+ points_at_time_t = manifold.expmap(start_point.unsqueeze(-2), tangent_vecs)
+
+ return points_at_time_t
+
+ return path
diff --git a/scDFM/src/flow_matching/utils/model_wrapper.py b/scDFM/src/flow_matching/utils/model_wrapper.py
new file mode 100644
index 0000000000000000000000000000000000000000..ac7d932c851b3d2194bd63ce597ef924dc819a7f
--- /dev/null
+++ b/scDFM/src/flow_matching/utils/model_wrapper.py
@@ -0,0 +1,43 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from abc import ABC
+
+from torch import nn, Tensor
+
+
+class ModelWrapper(ABC, nn.Module):
+ """
+ This class is used to wrap around another model, adding custom forward pass logic.
+ """
+
+ def __init__(self, model: nn.Module):
+ super().__init__()
+ self.model = model
+
+ def forward(self, x: Tensor, t: Tensor, **extras) -> Tensor:
+ r"""
+ This method defines how inputs should be passed through the wrapped model.
+ Here, we're assuming that the wrapped model takes both :math:`x` and :math:`t` as input,
+ along with any additional keyword arguments.
+
+ Optional things to do here:
+ - check that t is in the dimensions that the model is expecting.
+ - add a custom forward pass logic.
+ - call the wrapped model.
+
+ | given x, t
+ | returns the model output for input x at time t, with extra information `extra`.
+
+ Args:
+ x (Tensor): input data to the model (batch_size, ...).
+ t (Tensor): time (batch_size).
+ **extras: additional information forwarded to the model, e.g., text condition.
+
+ Returns:
+ Tensor: model output.
+ """
+ return self.model(x=x, t=t, **extras)
diff --git a/scDFM/src/flow_matching/utils/utils.py b/scDFM/src/flow_matching/utils/utils.py
new file mode 100644
index 0000000000000000000000000000000000000000..beb31ff27343c40955ec52513c44f75a6ab5d665
--- /dev/null
+++ b/scDFM/src/flow_matching/utils/utils.py
@@ -0,0 +1,90 @@
+# Copyright (c) Meta Platforms, Inc. and affiliates.
+# All rights reserved.
+#
+# This source code is licensed under the CC-by-NC license found in the
+# LICENSE file in the root directory of this source tree.
+
+from typing import Optional
+
+import torch
+from torch import Tensor
+
+
+def unsqueeze_to_match(source: Tensor, target: Tensor, how: str = "suffix") -> Tensor:
+ """
+ Unsqueeze the source tensor to match the dimensionality of the target tensor.
+
+ Args:
+ source (Tensor): The source tensor to be unsqueezed.
+ target (Tensor): The target tensor to match the dimensionality of.
+ how (str, optional): Whether to unsqueeze the source tensor at the beginning
+ ("prefix") or end ("suffix"). Defaults to "suffix".
+
+ Returns:
+ Tensor: The unsqueezed source tensor.
+ """
+ assert (
+ how == "prefix" or how == "suffix"
+ ), f"{how} is not supported, only 'prefix' and 'suffix' are supported."
+
+ dim_diff = target.dim() - source.dim()
+
+ for _ in range(dim_diff):
+ if how == "prefix":
+ source = source.unsqueeze(0)
+ elif how == "suffix":
+ source = source.unsqueeze(-1)
+
+ return source
+
+
+def expand_tensor_like(input_tensor: Tensor, expand_to: Tensor) -> Tensor:
+ """`input_tensor` is a 1d vector of length equal to the batch size of `expand_to`,
+ expand `input_tensor` to have the same shape as `expand_to` along all remaining dimensions.
+
+ Args:
+ input_tensor (Tensor): (batch_size,).
+ expand_to (Tensor): (batch_size, ...).
+
+ Returns:
+ Tensor: (batch_size, ...).
+ """
+ assert input_tensor.ndim == 1, "Input tensor must be a 1d vector."
+ assert (
+ input_tensor.shape[0] == expand_to.shape[0]
+ ), f"The first (batch_size) dimension must match. Got shape {input_tensor.shape} and {expand_to.shape}."
+
+ dim_diff = expand_to.ndim - input_tensor.ndim
+
+ t_expanded = input_tensor.clone()
+ t_expanded = t_expanded.reshape(-1, *([1] * dim_diff))
+
+ return t_expanded.expand_as(expand_to)
+
+
+def gradient(
+ output: Tensor,
+ x: Tensor,
+ grad_outputs: Optional[Tensor] = None,
+ create_graph: bool = False,
+) -> Tensor:
+ """
+ Compute the gradient of the inner product of output and grad_outputs w.r.t :math:`x`.
+
+ Args:
+ output (Tensor): [N, D] Output of the function.
+ x (Tensor): [N, d_1, d_2, ... ] input
+ grad_outputs (Optional[Tensor]): [N, D] Gradient of outputs, if `None`,
+ then will use a tensor of ones
+ create_graph (bool): If True, graph of the derivative will be constructed, allowing
+ to compute higher order derivative products. Defaults to False.
+ Returns:
+ Tensor: [N, d_1, d_2, ... ]. the gradient w.r.t x.
+ """
+
+ if grad_outputs is None:
+ grad_outputs = torch.ones_like(output).detach()
+ grad = torch.autograd.grad(
+ output, x, grad_outputs=grad_outputs, create_graph=create_graph
+ )[0]
+ return grad
diff --git a/scDFM/src/loss/metrics.py b/scDFM/src/loss/metrics.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0ba0e2e24560286bffd75ff225c9080196f1d6c
--- /dev/null
+++ b/scDFM/src/loss/metrics.py
@@ -0,0 +1,328 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import numpy as np
+from scipy import stats
+from scipy.spatial.distance import jensenshannon
+from sklearn.metrics import r2_score
+from sklearn.decomposition import PCA
+from sklearn.neighbors import NearestNeighbors
+import warnings
+warnings.filterwarnings('ignore')
+
+
+# TODO: To be implemented into flow generation
+class FlowMatchingMetrics:
+
+ def __init__(self, device='cuda'):
+ self.device = device
+
+ def reconstruction_metrics(self, generated, target):
+
+ if isinstance(generated, torch.Tensor):
+ generated = generated.detach().cpu().numpy()
+ if isinstance(target, torch.Tensor):
+ target = target.detach().cpu().numpy()
+
+ mae = np.mean(np.abs(generated - target))
+ mse = np.mean((generated - target) ** 2)
+ rmse = np.sqrt(mse)
+
+
+ pos_mask = target > 0
+ zero_mask = target == 0
+
+ pos_mae = np.mean(np.abs(generated[pos_mask] - target[pos_mask])) if pos_mask.any() else 0.0
+ zero_mae = np.mean(np.abs(generated[zero_mask] - target[zero_mask])) if zero_mask.any() else 0.0
+
+
+ r2 = r2_score(target.flatten(), generated.flatten())
+
+ rel_error = np.mean(np.abs((generated - target) / (target + 1e-8)))
+
+ return {
+ 'mae': mae,
+ 'mse': mse,
+ 'rmse': rmse,
+ 'pos_mae': pos_mae,
+ 'zero_mae': zero_mae,
+ 'r2_score': r2,
+ 'relative_error': rel_error
+ }
+
+ def distribution_metrics(self, generated, target):
+ if isinstance(generated, torch.Tensor):
+ generated = generated.detach().cpu().numpy()
+ if isinstance(target, torch.Tensor):
+ target = target.detach().cpu().numpy()
+
+ metrics = {}
+
+ mean_diff = np.mean(np.abs(np.mean(generated, axis=0) - np.mean(target, axis=0)))
+ var_diff = np.mean(np.abs(np.var(generated, axis=0) - np.var(target, axis=0)))
+
+ metrics['mean_difference'] = mean_diff
+ metrics['variance_difference'] = var_diff
+
+
+ correlations = []
+ for i in range(generated.shape[1]):
+ if np.var(generated[:, i]) > 1e-8 and np.var(target[:, i]) > 1e-8:
+ corr, _ = stats.pearsonr(generated[:, i], target[:, i])
+ if not np.isnan(corr):
+ correlations.append(corr)
+
+ metrics['mean_correlation'] = np.mean(correlations) if correlations else 0.0
+
+
+ kl_divs = []
+ for i in range(min(generated.shape[1], 100)):
+ try:
+ bins = np.linspace(
+ min(generated[:, i].min(), target[:, i].min()),
+ max(generated[:, i].max(), target[:, i].max()),
+ 50
+ )
+ hist_gen, _ = np.histogram(generated[:, i], bins=bins, density=True)
+ hist_target, _ = np.histogram(target[:, i], bins=bins, density=True)
+
+
+ hist_gen = hist_gen + 1e-8
+ hist_target = hist_target + 1e-8
+
+
+ hist_gen = hist_gen / np.sum(hist_gen)
+ hist_target = hist_target / np.sum(hist_target)
+
+ kl_div = stats.entropy(hist_gen, hist_target)
+ if not np.isnan(kl_div) and not np.isinf(kl_div):
+ kl_divs.append(kl_div)
+ except:
+ continue
+
+ metrics['mean_kl_divergence'] = np.mean(kl_divs) if kl_divs else float('inf')
+
+ js_divs = []
+ for i in range(min(generated.shape[1], 100)):
+ try:
+ bins = np.linspace(
+ min(generated[:, i].min(), target[:, i].min()),
+ max(generated[:, i].max(), target[:, i].max()),
+ 50
+ )
+ hist_gen, _ = np.histogram(generated[:, i], bins=bins, density=True)
+ hist_target, _ = np.histogram(target[:, i], bins=bins, density=True)
+
+ hist_gen = hist_gen + 1e-8
+ hist_target = hist_target + 1e-8
+ hist_gen = hist_gen / np.sum(hist_gen)
+ hist_target = hist_target / np.sum(hist_target)
+
+ js_div = jensenshannon(hist_gen, hist_target)
+ if not np.isnan(js_div):
+ js_divs.append(js_div)
+ except:
+ continue
+
+ metrics['mean_js_divergence'] = np.mean(js_divs) if js_divs else 1.0
+
+ return metrics
+
+ def embedding_metrics(self, generated, target):
+
+ if isinstance(generated, torch.Tensor):
+ generated = generated.detach().cpu().numpy()
+ if isinstance(target, torch.Tensor):
+ target = target.detach().cpu().numpy()
+
+ metrics = {}
+
+
+ cosine_sims = []
+ for i in range(generated.shape[0]):
+ sim = np.dot(generated[i], target[i]) / (
+ np.linalg.norm(generated[i]) * np.linalg.norm(target[i]) + 1e-8
+ )
+ cosine_sims.append(sim)
+
+ metrics['mean_cosine_similarity'] = np.mean(cosine_sims)
+
+
+ euclidean_dists = np.linalg.norm(generated - target, axis=1)
+ metrics['mean_euclidean_distance'] = np.mean(euclidean_dists)
+
+
+ if generated.shape[0] > 10:
+ k = min(5, generated.shape[0] // 2)
+
+ nn_target = NearestNeighbors(n_neighbors=k)
+ nn_target.fit(target)
+ _, target_indices = nn_target.kneighbors(target)
+
+ nn_generated = NearestNeighbors(n_neighbors=k)
+ nn_generated.fit(generated)
+ _, generated_indices = nn_generated.kneighbors(generated)
+
+ preservation = 0
+ for i in range(generated.shape[0]):
+ intersection = len(set(target_indices[i]) & set(generated_indices[i]))
+ preservation += intersection / k
+
+ metrics['neighbor_preservation'] = preservation / generated.shape[0]
+ else:
+ metrics['neighbor_preservation'] = 0.0
+
+ return metrics
+
+ def flow_specific_metrics(self, generated, target, flow_model=None):
+
+ metrics = {}
+
+ if isinstance(generated, torch.Tensor):
+ generated = generated.detach().cpu().numpy()
+ if isinstance(target, torch.Tensor):
+ target = target.detach().cpu().numpy()
+
+
+ pairwise_dists = []
+ for i in range(min(generated.shape[0], 100)):
+ for j in range(i + 1, min(generated.shape[0], 100)):
+ dist = np.linalg.norm(generated[i] - generated[j])
+ pairwise_dists.append(dist)
+
+ metrics['generation_diversity'] = np.mean(pairwise_dists) if pairwise_dists else 0.0
+
+ if generated.shape[0] > 10 and target.shape[0] > 10:
+
+ combined = np.vstack([generated, target])
+ pca = PCA(n_components=min(10, combined.shape[1]))
+ combined_pca = pca.fit_transform(combined)
+
+ generated_pca = combined_pca[:generated.shape[0]]
+ target_pca = combined_pca[generated.shape[0]:]
+
+
+ coverage = 0
+ for dim in range(generated_pca.shape[1]):
+ gen_min, gen_max = generated_pca[:, dim].min(), generated_pca[:, dim].max()
+ target_min, target_max = target_pca[:, dim].min(), target_pca[:, dim].max()
+
+ if target_max > target_min:
+ overlap = max(0, min(gen_max, target_max) - max(gen_min, target_min))
+ coverage += overlap / (target_max - target_min)
+
+ metrics['mode_coverage'] = coverage / generated_pca.shape[1]
+ else:
+ metrics['mode_coverage'] = 0.0
+
+ return metrics
+
+ def comprehensive_evaluation(self, generated, target, embeddings_generated=None, embeddings_target=None):
+
+ all_metrics = {}
+
+ recon_metrics = self.reconstruction_metrics(generated, target)
+ all_metrics.update({f'recon_{k}': v for k, v in recon_metrics.items()})
+
+
+ dist_metrics = self.distribution_metrics(generated, target)
+ all_metrics.update({f'dist_{k}': v for k, v in dist_metrics.items()})
+
+ flow_metrics = self.flow_specific_metrics(generated, target)
+ all_metrics.update({f'flow_{k}': v for k, v in flow_metrics.items()})
+
+ if embeddings_generated is not None and embeddings_target is not None:
+ emb_metrics = self.embedding_metrics(embeddings_generated, embeddings_target)
+ all_metrics.update({f'emb_{k}': v for k, v in emb_metrics.items()})
+
+ return all_metrics
+
+def evaluate_flow_generation(generated_data, target_data, generated_embeddings=None, target_embeddings=None, device='cuda'):
+
+ metrics = FlowMatchingMetrics(device=device)
+ return metrics.comprehensive_evaluation(
+ generated_data, target_data,
+ generated_embeddings, target_embeddings
+ )
+
+def print_metrics(metrics_dict, title="Metrics"):
+
+ print(f"\n{'='*50}")
+ print(f"{title:^50}")
+ print(f"{'='*50}")
+
+
+ categories = {
+ 'Reconstruction': [k for k in metrics_dict.keys() if k.startswith('recon_')],
+ 'Distribution': [k for k in metrics_dict.keys() if k.startswith('dist_')],
+ 'Flow-specific': [k for k in metrics_dict.keys() if k.startswith('flow_')],
+ 'Embedding': [k for k in metrics_dict.keys() if k.startswith('emb_')]
+ }
+
+ for category, keys in categories.items():
+ if keys:
+ print(f"\n{category}:")
+ print("-" * 20)
+ for key in keys:
+ clean_key = key.replace(f"{category.lower().replace('-', '_')}_", "")
+ value = metrics_dict[key]
+ if isinstance(value, float):
+ print(f" {clean_key:<20}: {value:.6f}")
+ else:
+ print(f" {clean_key:<20}: {value}")
+
+ print(f"{'='*50}\n")
+
+
+def test_metrics():
+
+ batch_size, n_features = 100, 2000
+ generated = torch.randn(batch_size, n_features)
+ target = torch.randn(batch_size, n_features)
+
+ metrics = evaluate_flow_generation(generated, target)
+
+ print_metrics(metrics, "Flow Matching Evaluation")
+
+ return metrics
+
+
+def pairwise_sq_dists(X, Y):
+ # X:[m,d], Y:[n,d] -> [m,n]
+ return torch.cdist(X, Y, p=2)**2
+
+@torch.no_grad()
+def median_sigmas(X, scales=(0.5, 1.0, 2.0, 4.0)):
+ Z = X
+ D2 = pairwise_sq_dists(Z, Z)
+ tri = D2[~torch.eye(D2.size(0), dtype=bool, device=D2.device)]
+ m = torch.median(tri).clamp_min(1e-12)
+ s2 = torch.tensor(scales, device=Z.device) * m
+ sigmas = torch.sqrt(s2)
+ return [float(s.item()) for s in sigmas]
+
+def mmd2_unbiased_multi_sigma(X, Y, sigmas):
+ """
+ """
+ m, n = X.size(0), Y.size(0)
+ Dxx = pairwise_sq_dists(X, X) # [m,m]
+ Dyy = pairwise_sq_dists(Y, Y) # [n,n]
+ Dxy = pairwise_sq_dists(X, Y) # [m,n]
+
+ vals = []
+ for sigma in sigmas:
+ beta = 1.0 / (2.0 * (sigma ** 2) + 1e-12)
+ Kxx = torch.exp(-beta * Dxx)
+ Kyy = torch.exp(-beta * Dyy)
+ Kxy = torch.exp(-beta * Dxy)
+
+
+ term_xx = (Kxx.sum() - Kxx.diag().sum()) / (m * (m - 1) + 1e-12)
+ term_yy = (Kyy.sum() - Kyy.diag().sum()) / (n * (n - 1) + 1e-12)
+ term_xy = Kxy.mean() # / (m*n)
+ vals.append(term_xx + term_yy - 2.0 * term_xy)
+
+ return torch.stack(vals).mean()
+
+if __name__ == "__main__":
+ test_metrics()
diff --git a/scDFM/src/models/instantiate_model.py b/scDFM/src/models/instantiate_model.py
new file mode 100644
index 0000000000000000000000000000000000000000..f651aae1e82a71017143f24d23f9488959207c32
--- /dev/null
+++ b/scDFM/src/models/instantiate_model.py
@@ -0,0 +1,25 @@
+from src.models.scGPT.model import TransformerModel
+from src.models.perturbation.model import Model as FlowModel
+from src.models.perturbation.model import TimedTransformer
+from src.models.origin.model import model as OriginModel
+import torch
+
+def instantiate_model(model_type: str, **kwargs):
+
+ if model_type == 'origin':
+ if kwargs['fusion_method'] == 'differential_transformer':
+ layers = 8
+ elif kwargs['fusion_method'] == 'differential_perceiver':
+ layers = 4
+ else:
+ layers = 8
+ return OriginModel(fusion_method=kwargs['fusion_method'], nlayers=layers, perturbation_function=kwargs['perturbation_function'],mask_path=kwargs['mask_path'])
+ else:
+ raise ValueError(f"Invalid model type: {model_type}")
+
+if __name__ == "__main__":
+ model = instantiate_model("punet128")
+ x = torch.randn(32, 128, 128)
+ t = torch.randn(32)
+ out = model( x,t)
+ print(out.shape)
\ No newline at end of file
diff --git a/scDFM/src/models/origin/blocks.py b/scDFM/src/models/origin/blocks.py
new file mode 100644
index 0000000000000000000000000000000000000000..3fab34c6e0656af4287463607ad976e36871b60c
--- /dev/null
+++ b/scDFM/src/models/origin/blocks.py
@@ -0,0 +1,355 @@
+import pdb
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Optional, Dict, Tuple
+import math
+from .rotary import apply_rotary_emb
+
+def lambda_init_fn(depth):
+ return 0.8 - 0.6 * math.exp(-0.3 * depth)
+
+def repeat_kv(x: torch.Tensor, n_rep: int) -> torch.Tensor:
+ """torch.repeat_interleave(x, dim=1, repeats=n_rep)"""
+ bs, n_kv_heads, slen, head_dim = x.shape
+ if n_rep == 1:
+ return x
+ return (
+ x[:, :, None, :, :]
+ .expand(bs, n_kv_heads, n_rep, slen, head_dim)
+ .reshape(bs, n_kv_heads * n_rep, slen, head_dim)
+ )
+
+def modulate(x, shift, scale):
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+class RMSNorm(nn.Module):
+ def __init__(self, dim: int, eps: float = 1e-6, elementwise_affine=True, memory_efficient=False):
+ super().__init__()
+ self.dim = dim
+ self.eps = eps
+ self.elementwise_affine = elementwise_affine
+ if self.elementwise_affine:
+ self.weight = nn.Parameter(torch.ones(dim))
+ else:
+ self.register_parameter('weight', None)
+
+ def _norm(self, x):
+ return x * torch.rsqrt(x.pow(2).mean(-1, keepdim=True) + self.eps)
+
+ def forward(self, x):
+ output = self._norm(x.float()).type_as(x)
+ if self.weight is not None:
+ output = output * self.weight
+ return output
+
+ def extra_repr(self) -> str:
+ return f'dim={self.dim}, eps={self.eps}, elementwise_affine={self.elementwise_affine}'
+
+
+class MultiheadDiffAttn_Origin(nn.Module):
+ def __init__(
+ self,
+ embed_dim,
+ depth, # current layer index
+ num_heads,
+ num_kv_heads=None,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+
+ # arg num_heads set to half of baseline Transformer's num_heads
+ # for e.g., to compare with a baseline Transformer with 16 heads, pass in num_heads=8 for DIFF Transformer
+ self.num_heads = num_heads
+
+ # arg num_kv_heads set to half of baseline Transformer's num_kv_heads if use GQA
+ # for e.g., to compare with a baseline Transformer with 16 heads and 8 kv_heads,
+ # pass in num_heads=8, num_kv_heads=4 for DIFF Transformer
+ # if use MHA, pass in num_kv_heads=None
+ self.num_kv_heads = num_kv_heads if num_kv_heads is not None else num_heads
+ self.n_rep = self.num_heads // self.num_kv_heads
+
+ self.head_dim = embed_dim // num_heads // 2
+ self.scaling = self.head_dim ** -0.5
+
+ self.q_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.k_proj = nn.Linear(embed_dim, embed_dim // self.n_rep, bias=False)
+ self.v_proj = nn.Linear(embed_dim, embed_dim // self.n_rep, bias=False)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+
+ # depth means current layer index
+ self.lambda_init = lambda_init_fn(depth)
+ self.lambda_q1 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
+ self.lambda_k1 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
+ self.lambda_q2 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
+ self.lambda_k2 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
+
+ self.subln = RMSNorm(2 * self.head_dim, eps=1e-5, elementwise_affine=True)
+
+ def forward(
+ self,
+ x,
+ rel_pos,
+ attn_mask=None,
+ ):
+ bsz, tgt_len, embed_dim = x.size()
+ src_len = tgt_len
+
+ q = self.q_proj(x)
+ k = self.k_proj(x)
+ v = self.v_proj(x)
+
+ q = q.view(bsz, tgt_len, 2 * self.num_heads, self.head_dim)
+ k = k.view(bsz, src_len, 2 * self.num_kv_heads, self.head_dim)
+ v = v.view(bsz, src_len, self.num_kv_heads, 2 * self.head_dim)
+
+ q = apply_rotary_emb(q, *rel_pos, interleaved=True)
+ k = apply_rotary_emb(k, *rel_pos, interleaved=True)
+
+ offset = src_len - tgt_len
+ q = q.transpose(1, 2)
+ k = repeat_kv(k.transpose(1, 2), self.n_rep)
+ v = repeat_kv(v.transpose(1, 2), self.n_rep)
+ q *= self.scaling
+ attn_weights = torch.matmul(q, k.transpose(-1, -2))
+ attn_weights = torch.nan_to_num(attn_weights)
+ attn_weights = torch.nn.functional.softmax(attn_weights, dim=-1, dtype=torch.float32).type_as(
+ attn_weights
+ )
+
+ lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1).float()).type_as(q)
+ lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1).float()).type_as(q)
+ lambda_full = lambda_1 - lambda_2 + self.lambda_init
+ attn_weights = attn_weights.view(bsz, self.num_heads, 2, tgt_len, src_len)
+ attn_weights = attn_weights[:, :, 0] - lambda_full * attn_weights[:, :, 1]
+
+ attn = torch.matmul(attn_weights, v)
+ attn = self.subln(attn)
+ attn = attn * (1 - self.lambda_init)
+ attn = attn.transpose(1, 2).reshape(bsz, tgt_len, self.num_heads * 2 * self.head_dim)
+
+ attn = self.out_proj(attn)
+ return attn
+
+class MultiheadDiffAttn(nn.Module):
+ def __init__(
+ self,
+ embed_dim,
+ num_heads,
+ depth, # current layer index
+ cross=False,
+ ):
+ super().__init__()
+ self.embed_dim = embed_dim
+ self.cross = cross
+ # The code `pdb.set_trace()` is setting a breakpoint in the code using the Python debugger
+ # (pdb). When the code is executed, it will pause the program's execution at that point and
+ # allow you to interactively debug the program.
+ # pdb.set_trace()
+ self.num_heads = num_heads
+ self.head_dim = embed_dim // num_heads
+ self.scaling = self.head_dim ** -0.5
+
+ self.q_proj_1 = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.k_proj_1 = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.q_proj_2 = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.k_proj_2 = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.v_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+ self.out_proj = nn.Linear(embed_dim, embed_dim, bias=False)
+
+ # depth means current layer index
+ self.lambda_init = lambda_init_fn(depth)
+ self.lambda_q1 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
+ self.lambda_k1 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
+ self.lambda_q2 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
+ self.lambda_k2 = nn.Parameter(torch.zeros(self.head_dim, dtype=torch.float32).normal_(mean=0,std=0.1))
+
+ self.subln = RMSNorm(self.head_dim, eps=1e-5, elementwise_affine=True)
+
+ def forward(
+ self,
+ noisy_y,
+ x,
+ ):
+ bsz, tgt_len, embed_dim = x.size()
+ src_len = tgt_len
+
+ if self.cross:
+ q_1 = self.q_proj_1(noisy_y)
+ k_1 = self.k_proj_1(x)
+ q_2 = self.q_proj_2(noisy_y)
+ k_2 = self.k_proj_2(x)
+ else:
+ q_1 = self.q_proj_1(noisy_y)
+ k_1 = self.k_proj_1(noisy_y)
+ q_2 = self.q_proj_2(x)
+ k_2 = self.k_proj_2(x)
+ v = self.v_proj(noisy_y)
+
+ q_1 = q_1.view(bsz, tgt_len, self.num_heads, self.head_dim)
+ k_1 = k_1.view(bsz, src_len, self.num_heads, self.head_dim)
+ q_2 = q_2.view(bsz, tgt_len, self.num_heads, self.head_dim)
+ k_2 = k_2.view(bsz, src_len, self.num_heads, self.head_dim)
+
+ v = v.view(bsz, src_len, self.num_heads, self.head_dim)
+
+
+ q_1 = q_1.transpose(1, 2)
+ q_2 = q_2.transpose(1, 2)
+ k_1 = k_1.transpose(1, 2)
+ k_2 = k_2.transpose(1, 2)
+ q_1 *= self.scaling
+ q_2 *= self.scaling
+
+ attn_weights_1 = torch.matmul(q_1, k_1.transpose(-1, -2))
+ attn_weights_2 = torch.matmul(q_2, k_2.transpose(-1, -2))
+
+ attn_weights_1 = torch.nn.functional.softmax(attn_weights_1, dim=-1, dtype=torch.float32).type_as(
+ attn_weights_1
+ )
+ attn_weights_2 = torch.nn.functional.softmax(attn_weights_2, dim=-1, dtype=torch.float32).type_as(
+ attn_weights_2
+ )
+
+ lambda_1 = torch.exp(torch.sum(self.lambda_q1 * self.lambda_k1, dim=-1).float()).type_as(q_1)
+ lambda_2 = torch.exp(torch.sum(self.lambda_q2 * self.lambda_k2, dim=-1).float()).type_as(q_1)
+
+ lambda_full = lambda_1 - lambda_2 + self.lambda_init
+ attn_weights = attn_weights_1 - lambda_full * attn_weights_2
+
+
+ attn = torch.matmul(attn_weights, v.transpose(1,2))
+ # attn: (bsz, num_heads, tgt_len, head_dim)
+ attn = self.subln(attn)
+ attn = attn * (1 - self.lambda_init)
+ attn = attn.transpose(1, 2).reshape(bsz, tgt_len, self.num_heads * self.head_dim)
+
+ attn = self.out_proj(attn)
+ return attn
+
+
+
+class CrossAttentionTransformerLayer(nn.Module):
+ """
+ One-layer Transformer block with ONLY cross-attention:
+ - Pre-norm -> Multihead cross-attention (Q=tgt, K/V=memory)
+ - Residual
+ - Pre-norm -> FeedForward (MLP)
+ - Residual
+
+ Shapes (batch_first=True):
+ tgt: (B, T_tgt, d_model)
+ memory: (B, T_mem, d_model)
+ """
+ def __init__(
+ self,
+ d_model: int,
+ nhead: int,
+ mlp_ratio: int = 4,
+ dropout: float = 0.1,
+ activation: str = "gelu",
+ batch_first: bool = True
+ ):
+ super().__init__()
+ self.batch_first = batch_first
+
+ # Cross attention: Q from tgt, K/V from memory
+ self.cross_attn = nn.MultiheadAttention(
+ embed_dim=d_model, num_heads=nhead, dropout=dropout, batch_first=batch_first
+ )
+ self.cross_attn_dropout = nn.Dropout(dropout)
+ self.norm1 = nn.LayerNorm(d_model)
+
+ # Feedforward
+ self.linear1 = nn.Linear(d_model, int(mlp_ratio * d_model))
+ self.linear2 = nn.Linear(int(mlp_ratio * d_model), d_model)
+ self.ffn_dropout = nn.Dropout(dropout)
+ self.norm2 = nn.LayerNorm(d_model)
+
+ if activation == "relu":
+ self.activation = nn.ReLU()
+ elif activation == "gelu":
+ self.activation = nn.GELU()
+ else:
+ raise ValueError(f"Unsupported activation: {activation}")
+
+ def forward(
+ self,
+ tgt: torch.Tensor,
+ memory: torch.Tensor,
+ attn_mask: Optional[torch.Tensor] = None,
+ memory_key_padding_mask: Optional[torch.Tensor] = None,
+ need_weights: bool = False,
+ average_attn_weights: bool = True,
+ ) -> Tuple[torch.Tensor, Optional[torch.Tensor]]:
+ """
+ Args:
+ tgt: (B, T_tgt, d_model)
+ memory: (B, T_mem, d_model)
+ attn_mask: optional mask applied to (tgt_len, mem_len) or (B*nhead, tgt_len, mem_len)
+ memory_key_padding_mask: (B, T_mem) with True for positions to mask
+ need_weights: return attention weights
+ average_attn_weights: average heads if True
+
+ Returns:
+ out: (B, T_tgt, d_model)
+ attn: (B, T_tgt, T_mem) or (B, nhead, T_tgt, T_mem) if need_weights else None
+ """
+ # --- Cross Attention (pre-norm) ---
+ x = self.norm1(tgt)
+ attn_out, attn_w = self.cross_attn(
+ query=x,
+ key=memory,
+ value=memory,
+ attn_mask=attn_mask,
+ key_padding_mask=memory_key_padding_mask,
+ need_weights=need_weights,
+ average_attn_weights=average_attn_weights,
+ )
+ x = tgt + self.cross_attn_dropout(attn_out)
+
+ # --- Feedforward (pre-norm) ---
+ y = self.norm2(x)
+ y = self.linear2(self.ffn_dropout(self.activation(self.linear1(y))))
+ out = x + self.ffn_dropout(y)
+
+ return out, attn_w if need_weights else None
+# class PerceiverDiffTransBlock(nn.Module):
+# def __init__(self, d_in, d_latent, heads=8, mlp_ratio=4):
+# super().__init__()
+
+# self.ln_z1 = nn.LayerNorm(d_latent)
+# self.q = nn.Linear(d_latent, d_latent)
+# self.k = nn.Linear(d_in, d_latent)
+# self.v = nn.Linear(d_in, d_latent)
+
+# self.q2 = nn.Linear(d_latent, d_latent)
+# self.k2 = nn.Linear(d_latent, d_latent)
+# self.v2 = nn.Linear(d_latent, d_latent)
+# self.cross = nn.MultiheadAttention(d_latent, heads, dropout=0.1, batch_first=True)
+
+# self.ln_z2 = nn.LayerNorm(d_latent)
+# self.self_attn = nn.MultiheadAttention(d_latent, heads, dropout=0.1, batch_first=True)
+# self.ln_z3 = nn.LayerNorm(d_latent)
+# self.mlp = nn.Sequential(
+# nn.Linear(d_latent, int(mlp_ratio * d_latent)), nn.GELU(),
+# nn.Dropout(0.1),
+# nn.Linear(int(mlp_ratio * d_latent), d_latent),
+# nn.Dropout(0.1),
+# )
+# self.adaLN_modulation = nn.Sequential(
+# nn.SiLU(),
+# nn.Linear(d_latent, 6 * d_latent, bias=True)
+# )
+
+
+# def forward(self, z, x, t):
+# shift_self, scale_self, gate_self, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t).chunk(6, dim=1)
+# z = z + self.cross(self.q(self.ln_z1(z)),
+# self.k(x), self.v(x))[0]
+
+# z = modulate(self.ln_z2(z), shift_self, scale_self)
+# z = z + gate_self.unsqueeze(1) * self.self_attn(self.q2(z), self.k2(z), self.v2(z))[0]
+
+# z = z + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.ln_z3(z), shift_mlp, scale_mlp))
+# return z
\ No newline at end of file
diff --git a/scDFM/src/models/origin/layers.py b/scDFM/src/models/origin/layers.py
new file mode 100644
index 0000000000000000000000000000000000000000..23af1b23d78d2ae17a5919dca93e6d8408feab03
--- /dev/null
+++ b/scDFM/src/models/origin/layers.py
@@ -0,0 +1,194 @@
+import pdb
+from re import X
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Optional, Dict
+import math
+from .blocks import CrossAttentionTransformerLayer
+import copy
+class GeneadaLN(nn.Module):
+ def __init__(self, hidden_size: int, dropout: float = 0.1):
+ super().__init__()
+
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(hidden_size, 3 * hidden_size, bias=True)
+ )
+ self.norm = nn.LayerNorm(hidden_size)
+
+ def forward(self, gene_emb: Tensor, value_emb: Tensor) -> Tensor:
+ shift, gate, scale = self.adaLN_modulation(gene_emb).chunk(3, dim=-1)
+
+ x = value_emb + gate * (self.norm(value_emb) * scale + shift )
+ return x
+
+class ContinuousValueEncoder(nn.Module):
+ """
+ Encode real number values to a vector using neural nets projection.
+ """
+
+ def __init__(self, d_model: int, dropout: float = 0.1, max_value: int = 512):
+ super().__init__()
+ self.dropout = nn.Dropout(p=dropout)
+ self.linear1 = nn.Linear(1, d_model)
+ self.activation = nn.ReLU()
+ self.linear2 = nn.Linear(d_model, d_model)
+ self.norm = nn.LayerNorm(d_model)
+ self.max_value = max_value
+
+ def forward(self, x: Tensor) -> Tensor:
+ """
+ Args:
+ x: Tensor, shape [batch_size, seq_len]
+ """
+ # TODO: test using actual embedding layer if input is categorical
+ # expand last dimension
+ x = x.unsqueeze(-1)
+ # clip x to [-inf, max_value]
+ x = torch.clamp(x, max=self.max_value)
+ x = self.activation(self.linear1(x))
+ x = self.linear2(x)
+ x = self.norm(x)
+ return self.dropout(x)
+
+class GeneEncoder(nn.Module):
+ def __init__(
+ self,
+ num_embeddings: int,
+ embedding_dim: int,
+ padding_idx: Optional[int] = None,
+ nhead:int = 8 ,
+ use_perturbation_interaction: bool = False,
+ dropout:float = 0.1,
+ mask_path: str = None,
+
+ ):
+ super().__init__()
+ self.embedding = nn.Embedding(
+ num_embeddings, embedding_dim, padding_idx=padding_idx
+ )
+ self.enc_norm = nn.LayerNorm(embedding_dim)
+ self.use_perturbation_interaction = use_perturbation_interaction
+ if use_perturbation_interaction:
+ self.data_name = mask_path.split('/')[-2]
+ self.perturbation_interaction = CrossAttentionTransformerLayer(embedding_dim, nhead, mlp_ratio=4.0, dropout=dropout)
+ self.mask_padded = torch.load(mask_path)
+ self.mask_num = self.mask_padded.shape[0]
+ def forward(self, x: Tensor) -> Tensor:
+ if self.use_perturbation_interaction:
+ # NOTE using the same perturbation and gene names
+ if self.mask_padded.device != x.device:
+
+ self.mask_padded = self.mask_padded.to(x.device)
+
+ mask = self.mask_padded[x[0]]
+
+ x = self.embedding(x) # (batch, seq_len, embsize)
+
+ x = self.enc_norm(x)
+ if self.use_perturbation_interaction:
+
+ memory_id = torch.arange(self.mask_num,device=x.device)
+ memory_emb = self.embedding(memory_id)
+ memory_emb = self.enc_norm(memory_emb).expand(x.shape[0],-1,-1)
+ x = self.perturbation_interaction(x, memory_emb,mask)[0]
+ return x
+
+class BatchLabelEncoder(nn.Module):
+ def __init__(
+ self,
+ num_embeddings: int,
+ embedding_dim: int,
+ padding_idx: Optional[int] = None,
+ ):
+ super().__init__()
+ self.embedding = nn.Embedding(
+ num_embeddings, embedding_dim, padding_idx=padding_idx
+ )
+ self.enc_norm = nn.LayerNorm(embedding_dim)
+
+ def forward(self, x: Tensor) -> Tensor:
+ x = self.embedding(x) # (batch, embsize)
+ x = self.enc_norm(x)
+ return x
+
+class ExprDecoder(nn.Module):
+ def __init__(
+ self,
+ d_model: int,
+ explicit_zero_prob: bool = False,
+ use_batch_labels: bool = False,
+ ):
+ super().__init__()
+ d_in = d_model * 2 if use_batch_labels else d_model
+ self.fc = nn.Sequential(
+ nn.Linear(d_in, d_model),
+ nn.LeakyReLU(),
+ nn.Linear(d_model, d_model),
+ nn.LeakyReLU(),
+ nn.Linear(d_model, 1),
+ )
+ self.explicit_zero_prob = explicit_zero_prob
+ if explicit_zero_prob:
+ self.zero_logit = nn.Sequential(
+ nn.Linear(d_in, d_model),
+ nn.LeakyReLU(),
+ nn.Linear(d_model, d_model),
+ nn.LeakyReLU(),
+ nn.Linear(d_model, 1),
+ )
+
+ def forward(self, x: Tensor) -> Dict[str, Tensor]:
+ """x is the output of the transformer, (batch, seq_len, d_model)"""
+ pred_value = self.fc(x).squeeze(-1) # (batch, seq_len)
+
+ if not self.explicit_zero_prob:
+ return dict(pred=pred_value)
+ zero_logits = self.zero_logit(x).squeeze(-1) # (batch, seq_len)
+ zero_probs = torch.sigmoid(zero_logits)
+ return dict(pred=pred_value, zero_probs=zero_probs)
+ # TODO: note that the return currently is only for training. Since decoder
+ # is not used in the test setting for the integration task, the eval/inference
+ # logic is not implemented yet. However, remember to implement it when
+ # the decoder is used in any test setting. The inference logic will need
+ # to sample from the bernoulli distribution with the zero_probs.
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+ t_emb = self.mlp(t_freq)
+ return t_emb
\ No newline at end of file
diff --git a/scDFM/src/models/origin/model.py b/scDFM/src/models/origin/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..a51d87851f25f6cfd424e4365a12dca522c9e046
--- /dev/null
+++ b/scDFM/src/models/origin/model.py
@@ -0,0 +1,255 @@
+import torch
+import torch.nn as nn
+from torch import Tensor
+from typing import Optional, Dict
+import math
+from timm.models.vision_transformer import PatchEmbed, Attention, Mlp
+
+import pdb
+
+from .layers import GeneadaLN, ContinuousValueEncoder, GeneEncoder, BatchLabelEncoder, TimestepEmbedder, ExprDecoder
+from .blocks import MultiheadDiffAttn, modulate , CrossAttentionTransformerLayer
+class Block(nn.Module):
+ """
+ A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+ """
+ def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
+ super().__init__()
+ self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
+ self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ mlp_hidden_dim = int(hidden_size * mlp_ratio)
+ approx_gelu = lambda: nn.GELU(approximate="tanh")
+ self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+ )
+
+ def forward(self, x, c):
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+ x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+ x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+ return x
+
+class DifferentialTransformerBlock(nn.Module):
+ def __init__(self, hidden_size, num_heads, depth, mlp_ratio=4.0, cross=False, **block_kwargs):
+ super().__init__()
+ self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.attn = MultiheadDiffAttn(hidden_size, num_heads, depth, cross=cross)
+ self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ mlp_hidden_dim = int(hidden_size * mlp_ratio)
+ approx_gelu = lambda: nn.GELU(approximate="tanh")
+ self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+ )
+
+ def forward(self, y, x, c):
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+ y = y + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(y), shift_msa, scale_msa), x)
+ y = y + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(y), shift_mlp, scale_mlp))
+ return y
+
+class PerceiverBlock(nn.Module):
+ def __init__(self, d_in, d_latent, heads=8, mlp_ratio=4, dropout=0.0):
+ super().__init__()
+ self.ln_z1 = nn.LayerNorm(d_latent)
+ self.q = nn.Linear(d_latent, d_latent)
+ self.k = nn.Linear(d_in, d_latent)
+ self.v = nn.Linear(d_in, d_latent)
+
+ self.q2 = nn.Linear(d_latent, d_latent)
+ self.k2 = nn.Linear(d_latent, d_latent)
+ self.v2 = nn.Linear(d_latent, d_latent)
+ self.cross = nn.MultiheadAttention(d_latent, heads, dropout=dropout, batch_first=True)
+
+ self.ln_z2 = nn.LayerNorm(d_latent)
+ self.self_attn = nn.MultiheadAttention(d_latent, heads, dropout=dropout, batch_first=True)
+ self.ln_z3 = nn.LayerNorm(d_latent)
+ self.mlp = nn.Sequential(
+ nn.Linear(d_latent, int(mlp_ratio * d_latent)), nn.GELU(),
+ nn.Linear(int(mlp_ratio * d_latent), d_latent)
+ )
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(d_latent, 6 * d_latent, bias=True)
+ )
+
+
+ def forward(self, z, x, t):
+ shift_self, scale_self, gate_self, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(t).chunk(6, dim=1)
+ z = z + self.cross(self.q(self.ln_z1(z)),
+ self.k(x), self.v(x))[0]
+
+ z = modulate(self.ln_z2(z), shift_self, scale_self)
+ z = z + gate_self.unsqueeze(1) * self.self_attn(self.q2(z), self.k2(z), self.v2(z))[0]
+
+ z = z + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.ln_z3(z), shift_mlp, scale_mlp))
+ return z
+
+class DiffPerceiverBlock(nn.Module):
+ def __init__(self,hidden_size, num_heads, depth, mlp_ratio=4.0):
+ super().__init__()
+ self.diff_self_attn = DifferentialTransformerBlock(hidden_size, num_heads, depth, mlp_ratio=mlp_ratio,cross=True)
+ self.diff_cross_attn = DifferentialTransformerBlock(hidden_size, num_heads, depth, mlp_ratio=mlp_ratio,cross=False)
+
+ def forward(self, y, x, c):
+ y = self.diff_self_attn(y, y, c)
+ y = self.diff_cross_attn(y, x, c)
+ return y
+
+class model(nn.Module):
+ def __init__(self,
+ ntoken: int = 6000,
+ d_model: int = 512,
+ nhead: int = 8,
+ d_hid: int = 2048,
+ nlayers: int = 8, # 8
+ dropout: float = 0.1,
+ fusion_method: str = 'cross', # concat, add, cross
+ perturbation_function: str = 'crisper',
+ use_perturbation_interaction: bool = True,
+ mask_path: str = None,
+ ):
+ super().__init__()
+ self.t_embedder = TimestepEmbedder(d_model)
+ self.perturbation_embedder = BatchLabelEncoder(ntoken, d_model,)
+ self.fusion_method = fusion_method
+ self.perturbation_function = perturbation_function
+ self.fusion_layer = nn.Sequential(nn.Linear(2*d_model, d_model),
+ nn.GELU(),
+ nn.Linear(d_model, d_model),
+ nn.LayerNorm(d_model))
+ self.value_encoder_1 = ContinuousValueEncoder(d_model, dropout)
+ self.value_encoder_2 = ContinuousValueEncoder(d_model, dropout)
+ self.encoder = GeneEncoder(ntoken, d_model,use_perturbation_interaction=use_perturbation_interaction,mask_path=mask_path)
+ self.use_perturbation_interaction = use_perturbation_interaction
+ # if use_perturbation_interaction:
+ # self.perturbation_interaction = CrossAttentionTransformerLayer(d_model, nhead, mlp_ratio=4.0, dropout=dropout)
+
+ if self.fusion_method == 'differential_transformer':
+ self.blocks = nn.ModuleList([
+ DifferentialTransformerBlock(d_model, nhead, i, mlp_ratio=4.0) for i in range(nlayers)
+ ])
+ elif self.fusion_method == 'differential_perceiver':
+ self.blocks = nn.ModuleList([
+ DiffPerceiverBlock(d_model, nhead, i, mlp_ratio=4.0) for i in range(nlayers)
+ ])
+ elif self.fusion_method == 'perceiver':
+ self.blocks = nn.ModuleList([
+ PerceiverBlock(d_model, d_model, heads=nhead, mlp_ratio=4.0, dropout=0.1) for _ in range(nlayers)
+ ])
+ else:
+ raise ValueError(f"Invalid fusion method: {self.fusion_method}")
+
+
+ self.gene_adaLN = nn.ModuleList([
+ GeneadaLN(d_model, dropout) for _ in range(nlayers)
+ ])
+ self.adapter_layer = nn.ModuleList([
+ nn.Sequential(
+ nn.Linear(2*d_model, d_model),
+ nn.LeakyReLU(),
+ nn.Dropout(dropout),
+ nn.Linear(d_model, d_model),
+ nn.LeakyReLU()) for _ in range(nlayers)
+ ])
+
+ # predict_p task with embedding prediction
+ self.p_mask_embed = nn.Parameter(torch.randn(d_model)) # (d_model,)
+ self.p_head = nn.Sequential(
+ nn.LayerNorm(d_model),
+ nn.Linear(d_model, d_model)
+ )
+
+ self.final_layer = ExprDecoder(d_model, explicit_zero_prob=False, use_batch_labels=True)
+ self.initialize_weights()
+
+ def initialize_weights(self):
+ # Initialize transformer layers:
+ def _basic_init(module):
+ if isinstance(module, nn.Linear):
+ torch.nn.init.xavier_uniform_(module.weight)
+ if module.bias is not None:
+ nn.init.constant_(module.bias, 0)
+ self.apply(_basic_init)
+
+
+ def get_perturbation_emb(self, perturbation_id=None, perturbation_emb=None,
+ cell_1=None, use_mask: bool=False):
+ if use_mask:
+ B = cell_1.size(0)
+ return self.p_mask_embed[None, :].expand(B, -1).to(cell_1.device, dtype=cell_1.dtype)
+
+ assert perturbation_emb is None or perturbation_id is None
+ if perturbation_id is not None:
+ if self.perturbation_function == 'crisper':
+ perturbation_emb = self.encoder(perturbation_id)
+
+ else:
+ perturbation_emb = self.perturbation_embedder(perturbation_id)
+ perturbation_emb = perturbation_emb.mean(1) # (B,d)
+ elif perturbation_emb is not None:
+ perturbation_emb = perturbation_emb.to(cell_1.device, dtype=cell_1.dtype)
+ if perturbation_emb.dim() == 1:
+ perturbation_emb = perturbation_emb.unsqueeze(0)
+ if perturbation_emb.size(0) == 1:
+ perturbation_emb = perturbation_emb.expand(cell_1.shape[0], -1).contiguous()
+ perturbation_emb = self.perturbation_embedder.enc_norm(perturbation_emb)
+
+ return perturbation_emb
+
+ def forward(self,gene_id, cell_1, t, cell_2, perturbation_id=None, gene_id_all=None, perturbation_emb=None, mode="predict_y"):
+ if t.dim() == 0:
+ t = t.repeat(cell_1.size(0))
+
+ gene_emb = self.encoder(gene_id)
+ # gene_emb_all = self.encoder(gene_id_all)
+ gene_emb_all = gene_emb
+ value_emb_1 = self.value_encoder_1(cell_1)
+ value_emb_2 = self.value_encoder_2(cell_2)
+
+ value_emb_1 = value_emb_1 + gene_emb
+ value_emb_2 = value_emb_2 + gene_emb_all
+
+ value_emb = torch.cat([value_emb_1, value_emb_2], dim=-1)
+ value_emb = self.fusion_layer(value_emb)
+
+ t_emb = self.t_embedder(t)
+
+ x = value_emb
+
+ perturbation_emb = self.get_perturbation_emb(perturbation_id, perturbation_emb, cell_1)
+
+ for i,block in enumerate(self.blocks):
+ x = self.gene_adaLN[i](gene_emb, x)
+ perturbation_exp = perturbation_emb[:, None, :].expand(-1, x.size(1), -1) # (B, T, emb)
+ x = torch.cat([x, perturbation_exp], dim=-1)
+ x = self.adapter_layer[i](x)
+ x = block(x, value_emb_2, t_emb)
+
+
+ if mode=="predict_p":
+ x_pooling = x.mean(dim=1)
+ return self.p_head(x_pooling)
+
+ x = torch.cat([x, perturbation_emb[:, None, :].expand(-1, x.size(1), -1)], dim=-1)
+ x = self.final_layer(x)
+
+ return x['pred']
+
+
+
+if __name__ == "__main__":
+ model = model(ntoken=100, d_model=128, nhead=8, d_hid=128, nlayers=12)
+ batch_size = 10
+ gene_id = torch.randint(0, 100, (batch_size, 10))
+ x = torch.randn(batch_size, 10)
+ t = torch.randint(0, 100, (batch_size,))
+ y = torch.randn(batch_size, 10)
+ perturbation = torch.randint(0, 100, (batch_size,2))
+
+ output = model(gene_id, x, t, y, perturbation)
+ pdb.set_trace()
\ No newline at end of file
diff --git a/scDFM/src/models/origin/rotary.py b/scDFM/src/models/origin/rotary.py
new file mode 100644
index 0000000000000000000000000000000000000000..1f4e6fb77dfb37bc431c8de63798e42cf2510b9b
--- /dev/null
+++ b/scDFM/src/models/origin/rotary.py
@@ -0,0 +1,333 @@
+# Copyright (c) 2023, Tri Dao.
+# Adapted from Diff-Transformer codebase
+
+from typing import Optional, Union
+
+import torch
+
+import triton
+import triton.language as tl
+
+
+# @triton.autotune(
+# configs=[
+# triton.Config({"BLOCK_M": 2}),
+# triton.Config({"BLOCK_M": 4}),
+# triton.Config({"BLOCK_M": 8}),
+# triton.Config({"BLOCK_M": 16}),
+# ],
+# key=["CACHE_KEY_SEQLEN", "BLOCK_K", "INTERLEAVED"],
+# )
+@triton.jit
+def rotary_kernel(
+ OUT, # Pointers to matrices
+ X,
+ COS,
+ SIN,
+ CU_SEQLENS,
+ SEQLEN_OFFSETS, # this could be int or a pointer
+ # Matrix dimensions
+ seqlen,
+ nheads,
+ rotary_dim,
+ seqlen_ro,
+ CACHE_KEY_SEQLEN,
+ # strides
+ stride_out_batch,
+ stride_out_seqlen,
+ stride_out_nheads,
+ stride_out_headdim,
+ stride_x_batch,
+ stride_x_seqlen,
+ stride_x_nheads,
+ stride_x_headdim,
+ # Meta-parameters
+ BLOCK_K: tl.constexpr,
+ IS_SEQLEN_OFFSETS_TENSOR: tl.constexpr,
+ IS_VARLEN: tl.constexpr,
+ INTERLEAVED: tl.constexpr,
+ CONJUGATE: tl.constexpr,
+ BLOCK_M: tl.constexpr,
+):
+ pid_m = tl.program_id(axis=0)
+ pid_batch = tl.program_id(axis=1)
+ pid_head = tl.program_id(axis=2)
+ rotary_dim_half = rotary_dim // 2
+
+ if not IS_VARLEN:
+ X = X + pid_batch * stride_x_batch + pid_head * stride_x_nheads
+ OUT = OUT + pid_batch * stride_out_batch + pid_head * stride_out_nheads
+ else:
+ start_idx = tl.load(CU_SEQLENS + pid_batch)
+ seqlen = tl.load(CU_SEQLENS + pid_batch + 1) - start_idx
+ X = X + start_idx * stride_x_seqlen + pid_head * stride_x_nheads
+ OUT = OUT + start_idx * stride_out_seqlen + pid_head * stride_out_nheads
+
+ if pid_m * BLOCK_M >= seqlen:
+ return
+ rm = pid_m * BLOCK_M + tl.arange(0, BLOCK_M)
+ if not IS_SEQLEN_OFFSETS_TENSOR:
+ rm_cs = rm + SEQLEN_OFFSETS
+ else:
+ rm_cs = rm + tl.load(SEQLEN_OFFSETS + pid_batch)
+ rk = tl.arange(0, BLOCK_K)
+ rk_half = tl.arange(0, BLOCK_K // 2)
+
+ if not INTERLEAVED:
+ # Load the 1st and 2nd halves of X, do calculation, then store to 1st and 2nd halves of OUT
+ X = X + (rm[:, None] * stride_x_seqlen + rk_half[None, :] * stride_x_headdim)
+ COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :])
+ SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_half[None, :])
+ cos = tl.load(
+ COS, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=1.0
+ ).to(tl.float32)
+ sin = tl.load(
+ SIN, mask=(rm_cs[:, None] < seqlen_ro) & (rk_half[None, :] < rotary_dim_half), other=0.0
+ ).to(tl.float32)
+ x0 = tl.load(
+ X, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half), other=0.0
+ ).to(tl.float32)
+ x1 = tl.load(
+ X + rotary_dim_half * stride_x_headdim,
+ mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half),
+ other=0.0,
+ ).to(tl.float32)
+ if CONJUGATE:
+ sin = -sin
+ o0 = x0 * cos - x1 * sin
+ o1 = x0 * sin + x1 * cos
+ # write back result
+ OUT = OUT + (rm[:, None] * stride_out_seqlen + rk_half[None, :] * stride_out_headdim)
+ tl.store(OUT, o0, mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half))
+ tl.store(
+ OUT + rotary_dim_half * stride_out_headdim,
+ o1,
+ mask=(rm[:, None] < seqlen) & (rk_half[None, :] < rotary_dim_half),
+ )
+ else:
+ # We don't want to load X[0, 2, 4, ...] and X[1, 3, 5, ...] separately since both are slow.
+ # Instead, we load x0 = X[0, 1, 2, 3, ...] and x1 = X[1, 0, 3, 2, ...].
+ # Loading x0 will be fast but x1 will be slow.
+ # Then we load cos = COS[0, 0, 1, 1, ...] and sin = SIN[0, 0, 1, 1, ...].
+ # Then we do the calculation and use tl.where to pick put the right outputs for the even
+ # and for the odd indices.
+ rk_swap = rk + ((rk + 1) % 2) * 2 - 1 # 1, 0, 3, 2, 5, 4, ...
+ rk_repeat = tl.arange(0, BLOCK_K) // 2
+ X0 = X + (rm[:, None] * stride_x_seqlen + rk[None, :] * stride_x_headdim)
+ X1 = X + (rm[:, None] * stride_x_seqlen + rk_swap[None, :] * stride_x_headdim)
+ COS = COS + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :])
+ SIN = SIN + (rm_cs[:, None] * rotary_dim_half + rk_repeat[None, :])
+ cos = tl.load(
+ COS,
+ mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half),
+ other=1.0,
+ ).to(tl.float32)
+ sin = tl.load(
+ SIN,
+ mask=(rm_cs[:, None] < seqlen_ro) & (rk_repeat[None, :] < rotary_dim_half),
+ other=0.0,
+ ).to(tl.float32)
+ x0 = tl.load(X0, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim), other=0.0).to(
+ tl.float32
+ )
+ x1 = tl.load(
+ X1, mask=(rm[:, None] < seqlen) & (rk_swap[None, :] < rotary_dim), other=0.0
+ ).to(tl.float32)
+ if CONJUGATE:
+ sin = -sin
+ x0_cos = x0 * cos
+ x1_sin = x1 * sin
+ out = tl.where(rk[None, :] % 2 == 0, x0_cos - x1_sin, x0_cos + x1_sin)
+ OUT = OUT + (rm[:, None] * stride_out_seqlen + rk[None, :] * stride_out_headdim)
+ tl.store(OUT, out, mask=(rm[:, None] < seqlen) & (rk[None, :] < rotary_dim))
+
+
+def apply_rotary(
+ x: torch.Tensor,
+ cos: torch.Tensor,
+ sin: torch.Tensor,
+ seqlen_offsets: Union[int, torch.Tensor] = 0,
+ cu_seqlens: Optional[torch.Tensor] = None,
+ max_seqlen: Optional[int] = None,
+ interleaved=False,
+ inplace=False,
+ conjugate=False,
+) -> torch.Tensor:
+ """
+ Arguments:
+ x: (batch, seqlen, nheads, headdim) if cu_seqlens is None
+ else (total_seqlen, nheads, headdim).
+ cos: (seqlen_ro, rotary_dim / 2)
+ sin: (seqlen_ro, rotary_dim / 2)
+ seqlen_offsets: integer or integer tensor of size (batch,)
+ cu_seqlens: (batch + 1,) or None
+ max_seqlen: int
+ Returns:
+ y: (batch, seqlen, nheads, headdim)
+ """
+ is_varlen = cu_seqlens is not None
+ if not is_varlen:
+ batch, seqlen, nheads, headdim = x.shape
+ else:
+ assert max_seqlen is not None, "If cu_seqlens is passed in, then max_seqlen must be passed"
+ total_seqlen, nheads, headdim = x.shape
+ batch_p_1 = cu_seqlens.shape[0]
+ batch = batch_p_1 - 1
+ seqlen = max_seqlen
+ seqlen_ro, rotary_dim = cos.shape
+ assert sin.shape == cos.shape
+ rotary_dim *= 2
+ assert rotary_dim <= headdim, "rotary_dim must be <= headdim"
+ assert headdim <= 256, "Only support headdim <= 256"
+ assert seqlen_ro >= seqlen, "seqlen_ro must be >= seqlen"
+
+ assert (
+ cos.dtype == sin.dtype
+ ), f"cos and sin must have the same dtype, got {cos.dtype} and {sin.dtype}"
+ assert (
+ x.dtype == cos.dtype
+ ), f"Input and cos/sin must have the same dtype, got {x.dtype} and {cos.dtype}"
+
+ cos, sin = cos.contiguous(), sin.contiguous()
+ if isinstance(seqlen_offsets, torch.Tensor):
+ assert seqlen_offsets.shape == (batch,)
+ assert seqlen_offsets.dtype in [torch.int32, torch.int64]
+ seqlen_offsets = seqlen_offsets.contiguous()
+ else:
+ assert seqlen_offsets + seqlen <= seqlen_ro
+
+ output = torch.empty_like(x) if not inplace else x
+ if rotary_dim < headdim and not inplace:
+ output[..., rotary_dim:].copy_(x[..., rotary_dim:])
+
+ BLOCK_K = (
+ 32
+ if rotary_dim <= 32
+ else (64 if rotary_dim <= 64 else (128 if rotary_dim <= 128 else 256))
+ )
+ grid = lambda META: (triton.cdiv(seqlen, META["BLOCK_M"]), batch, nheads) # noqa
+ BLOCK_M = 4 if interleaved else (8 if rotary_dim <= 64 else 4)
+
+ # Need this, otherwise Triton tries to launch from cuda:0 and we get
+ # ValueError: Pointer argument (at 0) cannot be accessed from Triton (cpu tensor?)
+ with torch.cuda.device(x.device.index):
+ rotary_kernel[grid](
+ output, # data ptrs
+ x,
+ cos,
+ sin,
+ cu_seqlens,
+ seqlen_offsets,
+ seqlen, # shapes
+ nheads,
+ rotary_dim,
+ seqlen_ro,
+ seqlen // 128, # key for triton cache (limit number of compilations)
+ output.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0
+ output.stride(-3), # seqlen_stride or total_seqlen_stride
+ output.stride(-2), # nheads_stride
+ output.stride(-1), # headdim_stride
+ x.stride(0) if not is_varlen else 0, # batch_strides if not varlen else 0
+ x.stride(-3), # seqlen stride or total_seqlen_stride
+ x.stride(-2), # nheads stride
+ x.stride(-1), # headdim stride
+ BLOCK_K,
+ isinstance(seqlen_offsets, torch.Tensor),
+ is_varlen,
+ interleaved,
+ conjugate,
+ BLOCK_M,
+ )
+ return output
+
+class ApplyRotaryEmb(torch.autograd.Function):
+ @staticmethod
+ def forward(
+ ctx,
+ x,
+ cos,
+ sin,
+ interleaved=False,
+ inplace=False,
+ seqlen_offsets: Union[int, torch.Tensor] = 0,
+ cu_seqlens: Optional[torch.Tensor] = None,
+ max_seqlen: Optional[int] = None,
+ ):
+ out = apply_rotary(
+ x,
+ cos,
+ sin,
+ seqlen_offsets=seqlen_offsets,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=max_seqlen,
+ interleaved=interleaved,
+ inplace=inplace,
+ )
+ if isinstance(seqlen_offsets, int):
+ # Can't save int with save_for_backward
+ ctx.save_for_backward(cos, sin, cu_seqlens)
+ ctx.seqlen_offsets = seqlen_offsets
+ else:
+ ctx.save_for_backward(cos, sin, cu_seqlens, seqlen_offsets)
+ ctx.seqlen_offsets = None
+ ctx.interleaved = interleaved
+ ctx.inplace = inplace
+ ctx.max_seqlen = max_seqlen
+ return out if not inplace else x
+
+ @staticmethod
+ def backward(ctx, do):
+ seqlen_offsets = ctx.seqlen_offsets
+ if seqlen_offsets is None:
+ cos, sin, cu_seqlens, seqlen_offsets = ctx.saved_tensors
+ else:
+ cos, sin, cu_seqlens = ctx.saved_tensors
+ # TD [2023-09-02]: For some reason Triton (2.0.0.post1) errors with
+ # "[CUDA]: invalid device context", and cloning makes it work. Idk why. Triton 2.1.0 works.
+ if not ctx.interleaved and not ctx.inplace:
+ do = do.clone()
+ dx = apply_rotary(
+ do,
+ cos,
+ sin,
+ seqlen_offsets=seqlen_offsets,
+ cu_seqlens=cu_seqlens,
+ max_seqlen=ctx.max_seqlen,
+ interleaved=ctx.interleaved,
+ inplace=ctx.inplace,
+ conjugate=True,
+ )
+ return dx, None, None, None, None, None, None, None
+
+
+def apply_rotary_emb(
+ x,
+ cos,
+ sin,
+ interleaved=False,
+ inplace=False,
+ seqlen_offsets: Union[int, torch.Tensor] = 0,
+ cu_seqlens: Optional[torch.Tensor] = None,
+ max_seqlen: Optional[int] = None,
+):
+ """
+ Arguments:
+ x: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None
+ else (total_seqlen, nheads, headdim)
+ cos, sin: (seqlen_rotary, rotary_dim / 2)
+ interleaved: if True, rotate pairs of even and odd dimensions (GPT-J style) instead
+ of 1st half and 2nd half (GPT-NeoX style).
+ inplace: if True, apply rotary embedding in-place.
+ seqlen_offsets: (batch_size,) or int. Each sequence in x is shifted by this amount.
+ Most commonly used in inference when we have KV cache.
+ cu_seqlens: (batch + 1,) or None
+ max_seqlen: int
+ Return:
+ out: (batch_size, seqlen, nheads, headdim) if cu_seqlens is None
+ else (total_seqlen, nheads, headdim)
+ rotary_dim must be <= headdim
+ Apply rotary embedding to the first rotary_dim of x.
+ """
+ return ApplyRotaryEmb.apply(
+ x, cos, sin, interleaved, inplace, seqlen_offsets, cu_seqlens, max_seqlen
+ )
\ No newline at end of file
diff --git a/scDFM/src/models/perturbation/model.py b/scDFM/src/models/perturbation/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..d0e9c6fc236f7d9a6d08a3056da1402ae0c5ceed
--- /dev/null
+++ b/scDFM/src/models/perturbation/model.py
@@ -0,0 +1,393 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+from typing import Dict, Mapping, Optional, Tuple, Any, Union
+import pdb
+from timm.models.vision_transformer import PatchEmbed, Attention, Mlp
+from src.flow_matching.utils.model_wrapper import ModelWrapper
+
+def modulate(x, shift, scale):
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+#################################################################################
+# Embedding Layers for Timesteps and Class Labels #
+#################################################################################
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+class PerturbationEmbedder(nn.Module):
+ """
+ Embeds class labels into vector representations. Also handles label dropout for classifier-free guidance.
+ """
+ def __init__(self, num_classes, hidden_size, dropout_prob):
+ super().__init__()
+ use_cfg_embedding = dropout_prob > 0
+ self.embedding_table = nn.Embedding(num_classes + use_cfg_embedding, hidden_size)
+ self.num_classes = num_classes
+ self.dropout_prob = dropout_prob
+
+ def token_drop(self, labels, force_drop_ids=None):
+ """
+ Drops labels to enable classifier-free guidance.
+ """
+ if force_drop_ids is None:
+ drop_ids = torch.rand(labels.shape[0], device=labels.device) < self.dropout_prob
+ else:
+ drop_ids = force_drop_ids == 1
+ labels = torch.where(drop_ids, self.num_classes, labels)
+ return labels
+
+ def forward(self, labels, train=True, force_drop_ids=None):
+ use_dropout = self.dropout_prob > 0
+ if (train and use_dropout) or (force_drop_ids is not None):
+ labels = self.token_drop(labels, force_drop_ids)
+ embeddings = self.embedding_table(labels)
+ return embeddings
+
+class CrossAttentionModule(nn.Module):
+ def __init__(self, embed_dim, num_heads, dropout=0.0):
+ super().__init__()
+ self.attn = nn.MultiheadAttention(embed_dim=embed_dim,
+ num_heads=num_heads,
+ dropout=dropout,
+ batch_first=True) # using (batch, seq, dim)
+ self.norm = nn.LayerNorm(embed_dim)
+ self.ff = nn.Sequential(
+ nn.Linear(embed_dim, embed_dim * 4),
+ nn.GELU(),
+ nn.Linear(embed_dim * 4, embed_dim),
+ nn.Dropout(dropout),
+ )
+
+ def forward(self, x1, x2, attn_mask=None, key_padding_mask=None):
+ # x1: (b, l1, dim), x2: (b, l2, dim)
+ residual = x1
+ attn_out, attn_weights = self.attn(query=x1, key=x2, value=x2,
+ attn_mask=attn_mask,
+ key_padding_mask=key_padding_mask)
+ x = self.norm(attn_out + residual)
+
+ # add FFN
+ residual2 = x
+ x = self.ff(x)
+ x = self.norm(x + residual2)
+
+ return x, attn_weights
+
+
+#################################################################################
+# Core DiT Model #
+#################################################################################
+
+class DiTBlock(nn.Module):
+ """
+ A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+ """
+ def __init__(self, hidden_size, num_heads, mlp_ratio=4.0, **block_kwargs):
+ super().__init__()
+ self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.attn = Attention(hidden_size, num_heads=num_heads, qkv_bias=True, **block_kwargs)
+ self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ mlp_hidden_dim = int(hidden_size * mlp_ratio)
+ approx_gelu = lambda: nn.GELU(approximate="tanh")
+ self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=0)
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+ )
+
+ def forward(self, x, c):
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=1)
+ x = x + gate_msa.unsqueeze(1) * self.attn(modulate(self.norm1(x), shift_msa, scale_msa))
+ x = x + gate_mlp.unsqueeze(1) * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+ return x
+
+class HierarchicalBlock(nn.Module):
+ def __init__(self,d_model,nhead,dropout):
+ super().__init__()
+ self.cell1_block = DiTBlock(d_model,nhead,dropout)
+ self.cell2_block = DiTBlock(d_model,nhead,dropout)
+ self.cross_attention = CrossAttentionModule(d_model,nhead,dropout)
+
+ def forward(self,cell1,cell2,perturbation,t):
+ cell1 = self.cell1_block(cell1,perturbation+t)
+ cell2 = self.cell2_block(cell2,t)
+ cell1, _ = self.cross_attention(cell2,cell1)
+
+ return cell1,cell2,
+
+class SymmetricBlock(nn.Module):
+ def __init__(self,d_model,nhead,dropout):
+ super().__init__()
+ # self.hierarchical_block1 = HierarchicalBlock(d_model,nhead,dropout)
+ # self.hierarchical_block2 = HierarchicalBlock(d_model,nhead,dropout)
+ self.hierarchical_block = HierarchicalBlock(d_model,nhead,dropout)
+ self.norm = nn.LayerNorm(d_model)
+
+ def forward(self,cell1,cell2,perturbation,t):
+ # cell1_1,cell2_1 = self.hierarchical_block1(cell1,cell2,perturbation,t)
+ # cell2_2,cell1_2 = self.hierarchical_block2(cell2,cell1,-perturbation,t)
+ cell1_1,cell2_1 = self.hierarchical_block(cell1,cell2,perturbation,t)
+ cell2_2,cell1_2 = self.hierarchical_block(cell2,cell1,-perturbation,t)
+ cell1 = self.norm(cell1_1 + cell1_2)
+ cell2 = self.norm(cell2_1 + cell2_2)
+ return cell1,cell2
+
+class Model(nn.Module):
+ def __init__(
+ self,
+ model_type: str = 'hierarchical',
+ ntoken: int = 512,
+ d_model: int = 512,
+ nhead: int = 8,
+ nlayers: int = 6 ,
+ n_combination: int = 2,
+ nlayers_cls: int = 3,
+ d_perturbation: int = 512,
+ vocab: Any = None,
+ dropout: float = 0.5,
+ ):
+ super().__init__()
+ self.n_combination = n_combination
+ self.d_model = d_model
+ self.model_type = model_type
+ self.ntoken = ntoken
+ self.d_perturbation = d_perturbation
+ # self.p_embedder = PerturbationEmbedder(n_perturbation, d_model, dropout_prob=0)
+ self.p_embedder = nn.Sequential(
+ nn.Linear(d_perturbation,d_model),
+ nn.SiLU(),
+ nn.Linear(d_model,d_model),
+ )
+ self.t_embedder = TimestepEmbedder(d_model)
+ self.blocks = nn.ModuleList([
+ DiTBlock(d_model, nhead, dropout) for _ in range(nlayers)
+ ])
+
+ if model_type == 'hierarchical':
+ self.blocks = nn.ModuleList([
+ SymmetricBlock(d_model, nhead, dropout) for _ in range(nlayers)
+ ])
+ else:
+ self.blocks = nn.ModuleList([
+ DiTBlock(d_model, nhead, dropout) for _ in range(nlayers)
+ ])
+ self.perturbation_decoder = nn.Sequential(
+ nn.Linear(d_model,d_model*4),
+ nn.SiLU(),
+ nn.Linear(d_model*4,d_model),
+ )
+
+
+ def forward(self,cell1,cell2,t,perturbation=None):
+ """_summary_
+
+ Args:
+ cell1 (tensor): (B,ntoken,d_model) for control cell
+ cell2 (tensor): (B,ntoken,d_model) for perturbed cell
+ t (long): (B,) when inference is a number
+ perturbation (long): (B,n_combination)
+ """
+ t = self.t_embedder(t) # (N, D)
+ # pdb.set_trace()
+ if perturbation is not None:
+ perturbation_ = perturbation.reshape(-1,self.d_perturbation)
+ p = self.p_embedder(perturbation_) # (N, D)
+ # p = p.reshape(-1,self.n_combination,self.d_model)
+ # p = p.sum(dim=1)
+ c = t + p
+ else:
+ c = t
+ cells = torch.cat([cell1,cell2],dim=1)
+ for block in self.blocks:
+ if self.model_type == 'hierarchical':
+ cell1,cell2 = block(cell1,cell2,p,t)
+ else:
+ cells = block(cells,c)
+ if self.model_type != 'hierarchical':
+ cells = cells.reshape(-1,2,self.ntoken,self.d_model)
+ cell1 = cells[:,0,:,:]
+ cell2 = cells[:,1,:,:]
+ p = self.perturbation_decoder((cell1 + cell2).mean(dim=1))
+ # p = p.reshape(-1,self.n_combination)
+
+ return cell1, cell2, p
+
+
+
+class TimedTransformer(nn.Module):
+ def __init__(
+ self,
+ model_type: str = 'hierarchical',
+ ntoken: int = 512,
+ d_model: int = 512,
+ nhead: int = 8,
+ nlayers: int = 8 ,
+ n_combination: int = 2,
+ nlayers_cls: int = 3,
+ d_perturbation: int = 512,
+ vocab: Any = None,
+ dropout: float = 0.5,
+ ):
+ super().__init__()
+ self.n_combination = n_combination
+ self.d_model = d_model
+ self.model_type = model_type
+ self.ntoken = ntoken
+ # self.p_embedder = PerturbationEmbedder(n_perturbation, d_model, dropout_prob=0)
+ self.t_embedder = TimestepEmbedder(d_model)
+ self.blocks = nn.ModuleList([
+ DiTBlock(d_model, nhead, dropout) for _ in range(nlayers)
+ ])
+
+
+
+ def forward(self,cell1,t):
+ """_summary_
+
+ Args:
+ cell1 (tensor): (B,ntoken,d_model) for control cell
+ cell2 (tensor): (B,ntoken,d_model) for perturbed cell
+ t (long): (B,) when inference is a number
+ perturbation (long): (B,n_combination)
+ """
+ t = self.t_embedder(t) # (N, D)
+ c = t
+ cells = cell1
+ for block in self.blocks:
+ cells = block(cells,c)
+ return cells
+
+
+# class WrappedModel(ModelWrapper):
+# def __init__(self, model, mode="noisy_y"):
+# super().__init__(model)
+# assert mode in ["noisy_y", "noisy_p"]
+# self.mode = mode
+
+# def forward(self, x: torch.Tensor, t: torch.Tensor, x_t: torch.Tensor, condition_vec: torch.Tensor, mode=None):
+# """
+# x: x_0
+# x_t: path sample at time t
+# t: time step
+# condition_vec: condition vector
+# """
+# if mode is not None:
+# self.mode = mode
+# if self.mode == "noisy_y":
+# _, predicted_x_t, _ = self.model(x, x_t, t, condition_vec)
+# return (x, predicted_x_t, condition_vec)
+# elif self.mode == "noisy_p":
+# _, _, predicted_p_t = self.model(x, x_t, t, condition_vec)
+# return (x, x_t, predicted_p_t)
+# else:
+# raise ValueError(f"Unsupported mode: {self.mode}")
+
+class WrappedModel(ModelWrapper):
+ def __init__(self, model, mode="noisy_y"):
+ super().__init__(model)
+ assert mode in ["noisy_y", "noisy_p"]
+ self.mode = mode
+ self.model = model
+
+ def change_mode(self, mode):
+ self.mode = mode
+
+ def forward(self, x: torch.Tensor, t: torch.Tensor, **extras):
+ if t.ndim == 0:
+ t = t.expand(x.shape[0]).to(x.device)
+
+ if self.mode == "noisy_y":
+ condition_vec = extras["condition_vec"]
+ x_0 = extras["x_0"]
+ _, predicted_x_t, _ = self.model(x_0, x, t, condition_vec)
+ return predicted_x_t
+ elif self.mode == "noisy_p":
+ target_y = extras["target_y"]
+ x_0 = extras["x_0"]
+ _, _, predicted_p_t = self.model(x_0, target_y, t, x)
+ return predicted_p_t
+ else:
+ raise ValueError(f"Unsupported mode: {self.mode}")
+
+if __name__ == '__main__':
+ batch_size = 128
+ model = Model(model_type='hs',ntoken=512).cuda()
+ cell1 = torch.randn(batch_size, 1024,512).cuda()
+ cell2 = torch.randn(batch_size, 1024,512).cuda()
+ t = torch.randint(0, 1000, (batch_size,)).cuda()
+ # p = torch.randint(0,2,(5,2)).cuda()
+ p = torch.randn(batch_size,512).cuda()
+ model = nn.DataParallel(model,device_ids=[0,1,2,3,4,5,6,7])
+
+ from tqdm import tqdm
+ for i in tqdm(range(5000)):
+ out = model(cell1,cell2, t,p)
+ out[0].sum().backward()
+ pdb.set_trace()
+ print(out.shape)
+
+ # from accelerate import Accelerator
+ # import time
+
+ # accelerator = Accelerator()
+ # batch_size = 128
+ # model = Model(model_type='hs', ntoken=512)
+ # cell1 = torch.randn(batch_size, 1024, 512)
+ # cell2 = torch.randn(batch_size, 1024, 512)
+ # t = torch.randint(0, 1000, (batch_size,))
+ # p = torch.randn(batch_size, 512)
+
+ # # 用 accelerator.prepare 包装
+ # model, cell1, cell2, t, p = accelerator.prepare(model, cell1, cell2, t, p)
+
+ # from tqdm import tqdm
+ # start = time.time()
+ # for i in tqdm(range(5000)):
+ # out = model(cell1, cell2, t, p)
+ # loss = out[0].sum() + out[1].sum() + out[2].sum()
+ # accelerator.backward(loss)
+ # end = time.time()
+ # print(f"Total time: {end - start:.2f} seconds")
+ # import pdb; pdb.set_trace()
+ # print(out[0].shape)
\ No newline at end of file
diff --git a/scDFM/src/models/perturbation/moduls.py b/scDFM/src/models/perturbation/moduls.py
new file mode 100644
index 0000000000000000000000000000000000000000..660c8de1134cc22cf78cb67b0a74e98ee7d69670
--- /dev/null
+++ b/scDFM/src/models/perturbation/moduls.py
@@ -0,0 +1,136 @@
+import torch.nn as nn
+from typing import Optional
+from torch import nn, Tensor
+import pdb
+class CategoryValueEncoder(nn.Module):
+ def __init__(
+ self,
+ num_embeddings: int,
+ embedding_dim: int,
+ padding_idx: Optional[int] = None,
+ ):
+ super().__init__()
+ self.embedding = nn.Embedding(
+ num_embeddings, embedding_dim, padding_idx=padding_idx
+ )
+ self.enc_norm = nn.LayerNorm(embedding_dim)
+
+ def forward(self, x: Tensor) -> Tensor:
+ x = x.long()
+ x = self.embedding(x) # (batch, seq_len, embsize)
+ x = self.enc_norm(x)
+ return x
+
+class GeneEncoder(nn.Module):
+ def __init__(
+ self,
+ num_embeddings: int,
+ embedding_dim: int,
+ padding_idx: Optional[int] = None,
+ ):
+ super().__init__()
+ self.embedding = nn.Embedding(
+ num_embeddings, embedding_dim, padding_idx=padding_idx
+ )
+ self.enc_norm = nn.LayerNorm(embedding_dim)
+
+ def forward(self, x: Tensor) -> Tensor:
+ x = self.embedding(x) # (batch, seq_len, embsize)
+ x = self.enc_norm(x)
+ return x
+
+class PerturbationEmbedding(nn.Module):
+ def __init__(self, num_perturbations, emb_dim, max_comb_len=2, fusion_method='mlp', output_matrix=False):
+ super().__init__()
+ self.embedding = nn.Embedding(num_perturbations, emb_dim)
+ self.fusion_method = fusion_method
+ self.max_comb_len = max_comb_len
+ self.output_matrix = output_matrix
+ self.output_dim = emb_dim if not output_matrix else emb_dim * emb_dim
+
+ if fusion_method == 'mlp':
+ self.fusion = nn.Sequential(
+ nn.Linear(emb_dim * max_comb_len, emb_dim * 2),
+ nn.ReLU(),
+ nn.Linear(emb_dim * 2, self.output_dim)
+ )
+ elif fusion_method == 'sum':
+ self.fusion = None
+ else:
+ raise ValueError(f"Unsupported fusion method: {fusion_method}")
+
+ def forward(self, ids):
+ emb = self.embedding(ids) # [B, C, D]
+
+ if self.fusion_method == 'mlp':
+ emb = emb.view(emb.size(0), -1) # [B, C*D]
+ fused = self.fusion(emb) # [B, D] or [B, D*D]
+
+ if self.output_matrix:
+ B = fused.size(0)
+ D = int(self.output_dim ** 0.5)
+ return fused.view(B, D, D) # [B, D, D]
+ else:
+ return fused
+
+ elif self.fusion_method == 'sum':
+ out = emb.sum(dim=1) # [B, D]
+ if self.output_matrix:
+ B = out.size(0)
+ D = out.size(1)
+ return out.view(B, D, 1).expand(B, D, D) # dummy expansion
+ return out
+
+# class PerturbationEmbedding(nn.Module):
+# def __init__(self, num_perturbations, emb_dim, max_comb_len=2, fusion_method='mlp'):
+# """
+# Args:
+# num_perturbations: 词表大小
+# emb_dim: 嵌入维度
+# max_comb_len: 每个 condition 最多包含的 token 数量（如 drug1, drug2）
+# fusion_method: 'mlp' 或 'sum'
+# """
+# super().__init__()
+# self.embedding = nn.Embedding(num_perturbations, emb_dim)
+# self.fusion_method = fusion_method
+# self.max_comb_len = max_comb_len
+
+# if fusion_method == 'mlp':
+# self.fusion = nn.Sequential(
+# nn.Linear(emb_dim * max_comb_len, emb_dim),
+# nn.ReLU(),
+# nn.Linear(emb_dim, emb_dim)
+# )
+# elif fusion_method == 'sum':
+# self.fusion = None
+# else:
+# raise ValueError(f"Unsupported fusion method: {fusion_method}")
+# def init_weights(self, m):
+# if isinstance(m, nn.Linear):
+# nn.init.xavier_uniform_(m.weight)
+# nn.init.zeros_(m.bias)
+
+# def initialize_weights(self):
+# self.apply(self.init_weights)
+
+# def forward(self, ids):
+# """
+# Args:
+# ids: LongTensor of shape [B, max_comb_len]
+# Returns:
+# fused: Tensor of shape [B, emb_dim]
+# """
+
+# emb = self.embedding(ids) # [B, C, D]
+
+# if self.fusion_method == 'mlp':
+# emb = emb.view(emb.size(0), -1) # [B, C*D]
+# return self.fusion(emb) # [B, D]
+
+# elif self.fusion_method == 'sum':
+# if emb.dim() == 2:
+# return emb.sum(dim=0) # [B, D]
+# elif emb.dim() == 3:
+# return emb.sum(dim=1) # [B, C, D]
+# else:
+# raise ValueError(f"Unsupported dimension: {ids.dim()}")
\ No newline at end of file
diff --git a/scDFM/src/models/perturbation/timed_transformer.py b/scDFM/src/models/perturbation/timed_transformer.py
new file mode 100644
index 0000000000000000000000000000000000000000..5f43b4deee5291b99144d6538043268df5bb4e3e
--- /dev/null
+++ b/scDFM/src/models/perturbation/timed_transformer.py
@@ -0,0 +1,187 @@
+import torch
+import torch.nn as nn
+import numpy as np
+import math
+
+#### Adapted from Diffusion Transformer codebase https://github.com/facebookresearch/DiT ###
+
+def modulate(x, shift, scale):
+ return x * (1 + scale.unsqueeze(1)) + shift.unsqueeze(1)
+
+
+# TODO: How to embed perturbation?
+# Since they are mainly labels, we could easily embed them via a trained embedding layer
+
+class DiTAttention(nn.Module):
+ def __init__(self, cfg):
+ super().__init__()
+ self.t_embedder = TimestepEmbedder(cfg.hidden_size)
+ self.trans_conv = nn.ModuleList([
+ DiTBlock(cfg) for _ in range(cfg.num_layers)
+ ])
+
+ def forward(self, t, x, mask=None):
+ t = self.t_embedder(t)
+ for block in self.trans_conv:
+ x = block(x, t, mask)
+ return x
+
+
+
+class TimestepEmbedder(nn.Module):
+ """
+ Embeds scalar timesteps into vector representations.
+ """
+ def __init__(self, hidden_size, frequency_embedding_size=256):
+ super().__init__()
+ self.mlp = nn.Sequential(
+ nn.Linear(frequency_embedding_size, hidden_size, bias=True),
+ nn.SiLU(),
+ nn.Linear(hidden_size, hidden_size, bias=True),
+ )
+ self.frequency_embedding_size = frequency_embedding_size
+
+ @staticmethod
+ def timestep_embedding(t, dim, max_period=10000):
+ """
+ Create sinusoidal timestep embeddings.
+ :param t: a 1-D Tensor of N indices, one per batch element.
+ These may be fractional.
+ :param dim: the dimension of the output.
+ :param max_period: controls the minimum frequency of the embeddings.
+ :return: an (N, D) Tensor of positional embeddings.
+ """
+ # https://github.com/openai/glide-text2im/blob/main/glide_text2im/nn.py
+ half = dim // 2
+ freqs = torch.exp(
+ -math.log(max_period) * torch.arange(start=0, end=half, dtype=torch.float32) / half
+ ).to(device=t.device)
+ args = t[:, None].float() * freqs[None]
+ embedding = torch.cat([torch.cos(args), torch.sin(args)], dim=-1)
+ if dim % 2:
+ embedding = torch.cat([embedding, torch.zeros_like(embedding[:, :1])], dim=-1)
+ return embedding
+
+ def forward(self, t):
+ t_freq = self.timestep_embedding(t, self.frequency_embedding_size)
+ t_emb = self.mlp(t_freq)
+ return t_emb
+
+
+
+class Mlp(nn.Module):
+ """ MLP as used in Vision Transformer, MLP-Mixer and related networks
+
+ NOTE: When use_conv=True, expects 2D NCHW tensors, otherwise N*C expected.
+ """
+ def __init__(
+ self,
+ in_features,
+ hidden_features=None,
+ out_features=None,
+ act_layer=nn.GELU,
+ norm_layer=None,
+ bias=True,
+ drop=0.,
+ ):
+ super().__init__()
+ out_features = out_features or in_features
+ hidden_features = hidden_features or in_features
+ linear_layer = nn.Linear
+
+ self.fc1 = linear_layer(in_features, hidden_features, bias=bias)
+ self.act = act_layer()
+ self.drop1 = nn.Dropout(drop)
+ self.norm = norm_layer(hidden_features) if norm_layer is not None else nn.Identity()
+ self.fc2 = linear_layer(hidden_features, out_features, bias=bias)
+ self.drop2 = nn.Dropout(drop)
+
+ def forward(self, x):
+ x = self.fc1(x)
+ x = self.act(x)
+ x = self.drop1(x)
+ x = self.norm(x)
+ x = self.fc2(x)
+ x = self.drop2(x)
+ return x
+
+class DiTBlock(nn.Module):
+ """
+ A DiT block with adaptive layer norm zero (adaLN-Zero) conditioning.
+ """
+ def __init__(self, cfg):
+ super().__init__()
+ hidden_size = cfg.model.hidden_channels + cfg.model.edge_dim
+ num_heads = cfg.model.trans_num_heads
+ dropout = cfg.model.trans_dropout
+ mlp_ratio= 4.0
+
+ self.norm1 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ self.attn = BertSelfAttention(num_attention_heads=num_heads, hidden_size=hidden_size, dropout_prob=dropout)
+ self.norm2 = nn.LayerNorm(hidden_size, elementwise_affine=False, eps=1e-6)
+ mlp_hidden_dim = int(hidden_size * mlp_ratio)
+ approx_gelu = lambda: nn.GELU(approximate="tanh")
+ self.mlp = Mlp(in_features=hidden_size, hidden_features=mlp_hidden_dim, act_layer=approx_gelu, drop=dropout)
+ self.adaLN_modulation = nn.Sequential(
+ nn.SiLU(),
+ nn.Linear(hidden_size, 6 * hidden_size, bias=True)
+ )
+
+ def forward(self, x, c, mask):
+ shift_msa, scale_msa, gate_msa, shift_mlp, scale_mlp, gate_mlp = self.adaLN_modulation(c).chunk(6, dim=-1)
+ x = x + gate_msa * self.attn(modulate(self.norm1(x), shift_msa, scale_msa), mask)
+ x = x + gate_mlp * self.mlp(modulate(self.norm2(x), shift_mlp, scale_mlp))
+ return x
+
+
+class BertSelfAttention(nn.Module):
+ def __init__(self, num_attention_heads, hidden_size, dropout_prob):
+ super().__init__()
+ self.num_attention_heads = num_attention_heads
+ self.attention_head_size = int(hidden_size / num_attention_heads)
+ self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+ self.query = nn.Linear(hidden_size, hidden_size)
+ self.key = nn.Linear(hidden_size, hidden_size)
+ self.value = nn.Linear(hidden_size, hidden_size)
+ self.dropout = nn.Dropout(dropout_prob)
+
+ def transpose_for_scores(self, x: torch.Tensor) -> torch.Tensor:
+ new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+ x = x.view(new_x_shape)
+ return x.permute(0, 2, 1, 3)
+
+ def forward(
+ self,
+ hidden_states: torch.Tensor,
+ attention_mask: None,
+ ):
+ mixed_query_layer = self.query(hidden_states)
+
+ key_layer = self.transpose_for_scores(self.key(hidden_states))
+ value_layer = self.transpose_for_scores(self.value(hidden_states))
+ query_layer = self.transpose_for_scores(mixed_query_layer)
+
+ # Take the dot product between "query" and "key" to get the raw attention scores.
+ attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+ attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+
+ if attention_mask is not None:
+ # our attention mask is 1 for positions we want to attend and 0 for ones we want to ignore
+ # let the mask have -10000.0 where we want to ignore
+ attention_mask = torch.einsum('bi,bj->bij', attention_mask, attention_mask).unsqueeze(1)
+ attention_mask = attention_mask.masked_fill(attention_mask == 0, -10000.0)
+ attention_scores = attention_scores + attention_mask
+
+ # Normalize the attention scores to probabilities.
+ attention_probs = nn.functional.softmax(attention_scores, dim=-1)
+
+ # This is actually dropping out entire tokens to attend to, which might
+ # seem a bit unusual, but is taken from the original Transformer paper.
+ attention_probs = self.dropout(attention_probs)
+ context_layer = torch.matmul(attention_probs, value_layer)
+ context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+ new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+ context_layer = context_layer.view(new_context_layer_shape)
+
+ return context_layer
\ No newline at end of file
diff --git a/scDFM/src/models/scGPT/dsbn.py b/scDFM/src/models/scGPT/dsbn.py
new file mode 100644
index 0000000000000000000000000000000000000000..9b904f30bb317a46457f5d35c0974ebb2880163d
--- /dev/null
+++ b/scDFM/src/models/scGPT/dsbn.py
@@ -0,0 +1,82 @@
+from typing import Optional, Tuple
+
+import torch
+from torch import nn
+
+# The code is modified from https://github.com/wgchang/DSBN/blob/master/model/dsbn.py
+class _DomainSpecificBatchNorm(nn.Module):
+ _version = 2
+
+ def __init__(
+ self,
+ num_features: int,
+ num_domains: int,
+ eps: float = 1e-5,
+ momentum: float = 0.1,
+ affine: bool = True,
+ track_running_stats: bool = True,
+ ):
+ super(_DomainSpecificBatchNorm, self).__init__()
+ self._cur_domain = None
+ self.num_domains = num_domains
+ self.bns = nn.ModuleList(
+ [
+ self.bn_handle(num_features, eps, momentum, affine, track_running_stats)
+ for _ in range(num_domains)
+ ]
+ )
+
+ @property
+ def bn_handle(self) -> nn.Module:
+ raise NotImplementedError
+
+ @property
+ def cur_domain(self) -> Optional[int]:
+ return self._cur_domain
+
+ @cur_domain.setter
+ def cur_domain(self, domain_label: int):
+ self._cur_domain = domain_label
+
+ def reset_running_stats(self):
+ for bn in self.bns:
+ bn.reset_running_stats()
+
+ def reset_parameters(self):
+ for bn in self.bns:
+ bn.reset_parameters()
+
+ def _check_input_dim(self, input: torch.Tensor):
+ raise NotImplementedError
+
+ def forward(self, x: torch.Tensor, domain_label: int) -> torch.Tensor:
+ self._check_input_dim(x)
+ if domain_label >= self.num_domains:
+ raise ValueError(
+ f"Domain label {domain_label} exceeds the number of domains {self.num_domains}"
+ )
+ bn = self.bns[domain_label]
+ self.cur_domain = domain_label
+ return bn(x)
+
+
+class DomainSpecificBatchNorm1d(_DomainSpecificBatchNorm):
+ @property
+ def bn_handle(self) -> nn.Module:
+ return nn.BatchNorm1d
+
+ def _check_input_dim(self, input: torch.Tensor):
+ if input.dim() > 3:
+ raise ValueError(
+ "expected at most 3D input (got {}D input)".format(input.dim())
+ )
+
+
+class DomainSpecificBatchNorm2d(_DomainSpecificBatchNorm):
+ @property
+ def bn_handle(self) -> nn.Module:
+ return nn.BatchNorm2d
+
+ def _check_input_dim(self, input: torch.Tensor):
+ if input.dim() != 4:
+ raise ValueError("expected 4D input (got {}D input)".format(input.dim()))
\ No newline at end of file
diff --git a/scDFM/src/models/scGPT/model.py b/scDFM/src/models/scGPT/model.py
new file mode 100644
index 0000000000000000000000000000000000000000..e71c38dfe7881366ab1dc539d2c9921d5f17fcf1
--- /dev/null
+++ b/scDFM/src/models/scGPT/model.py
@@ -0,0 +1,1393 @@
+import gc
+import math
+from typing import Dict, Mapping, Optional, Tuple, Any, Union
+import pdb
+import torch
+import numpy as np
+from torch import nn, Tensor
+import torch.distributed as dist
+import torch.nn.functional as F
+from torch.nn import TransformerEncoder, TransformerEncoderLayer
+from torch.distributions import Bernoulli
+from tqdm import trange
+import torch
+from torch.autograd import Function
+from .dsbn import DomainSpecificBatchNorm1d
+
+
+class TransformerModel(nn.Module):
+ def __init__(
+ self,
+ ntoken: int,
+ d_model: int,
+ nhead: int,
+ d_hid: int,
+ nlayers: int,
+ nlayers_cls: int = 3,
+ n_cls: int = 1,
+ vocab: Any = None,
+ dropout: float = 0.5,
+ pad_token: str = "| \n", + " | file_name | \n", + "n_cells | \n", + "n_genes | \n", + "candidate_perturbation_columns | \n", + "candidate_count_columns | \n", + "candidate_gene_count_columns | \n", + "x_storage | \n", + "has_raw | \n", + "layers | \n", + "
|---|---|---|---|---|---|---|---|---|---|
| 0 | \n", + "AdamsonWeissman2016_GSM2406675_10X001.h5ad | \n", + "5768 | \n", + "35635 | \n", + "[perturbation] | \n", + "[ncounts, UMI count, read count] | \n", + "[ngenes] | \n", + "_CSCDataset | \n", + "False | \n", + "[] | \n", + "
| 1 | \n", + "AdamsonWeissman2016_GSM2406677_10X005.h5ad | \n", + "15006 | \n", + "32738 | \n", + "[perturbation] | \n", + "[ncounts, UMI count, read count] | \n", + "[ngenes] | \n", + "_CSCDataset | \n", + "False | \n", + "[] | \n", + "
| 2 | \n", + "AdamsonWeissman2016_GSM2406681_10X010.h5ad | \n", + "65337 | \n", + "32738 | \n", + "[perturbation] | \n", + "[ncounts, UMI count, read count] | \n", + "[ngenes] | \n", + "_CSCDataset | \n", + "False | \n", + "[] | \n", + "
| 3 | \n", + "ReplogleWeissman2022_K562_essential.h5ad | \n", + "310385 | \n", + "8563 | \n", + "[perturbation, gene, gene_id, gene_transcript,... | \n", + "[ncounts, UMI_count] | \n", + "[ngenes] | \n", + "Dataset | \n", + "False | \n", + "[] | \n", + "
| \n", + " | dataset_name | \n", + "file_name | \n", + "cell_barcode | \n", + "perturbation_label | \n", + "perturbation_gene | \n", + "perturbation_gene_id | \n", + "guide_id | \n", + "perturbation_type | \n", + "cell_line | \n", + "celltype | \n", + "... | \n", + "ncounts | \n", + "ngenes | \n", + "nperts | \n", + "is_control | \n", + "is_single_perturbation | \n", + "expression_location | \n", + "source_perturbation_col | \n", + "source_gene_col | \n", + "source_gene_id_col | \n", + "source_guide_col | \n", + "
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| 0 | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "AdamsonWeissman2016_GSM2406675_10X001.h5ad | \n", + "AAACATACACCGAT | \n", + "CREB1_pDS269 | \n", + "CREB1_pDS269 | \n", + "None | \n", + "None | \n", + "CRISPR | \n", + "K562 | \n", + "lymphoblasts | \n", + "... | \n", + "8138.0 | \n", + "2412 | \n", + "2 | \n", + "False | \n", + "False | \n", + "X only | \n", + "perturbation | \n", + "perturbation | \n", + "None | \n", + "None | \n", + "
| 1 | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "AdamsonWeissman2016_GSM2406675_10X001.h5ad | \n", + "AAACATACAGAGAT | \n", + "SNAI1_pDS266 | \n", + "SNAI1_pDS266 | \n", + "None | \n", + "None | \n", + "CRISPR | \n", + "K562 | \n", + "lymphoblasts | \n", + "... | \n", + "8980.0 | \n", + "2386 | \n", + "2 | \n", + "False | \n", + "False | \n", + "X only | \n", + "perturbation | \n", + "perturbation | \n", + "None | \n", + "None | \n", + "
| 2 | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "AdamsonWeissman2016_GSM2406675_10X001.h5ad | \n", + "AAACATACCAGAAA | \n", + "62(mod)_pBA581 | \n", + "62(mod)_pBA581 | \n", + "None | \n", + "None | \n", + "CRISPR | \n", + "K562 | \n", + "lymphoblasts | \n", + "... | \n", + "28610.0 | \n", + "4404 | \n", + "2 | \n", + "False | \n", + "False | \n", + "X only | \n", + "perturbation | \n", + "perturbation | \n", + "None | \n", + "None | \n", + "
| 3 | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "AdamsonWeissman2016_GSM2406675_10X001.h5ad | \n", + "AAACATACGTTGAC | \n", + "EP300_pDS268 | \n", + "EP300_pDS268 | \n", + "None | \n", + "None | \n", + "CRISPR | \n", + "K562 | \n", + "lymphoblasts | \n", + "... | \n", + "11346.0 | \n", + "2815 | \n", + "2 | \n", + "False | \n", + "False | \n", + "X only | \n", + "perturbation | \n", + "perturbation | \n", + "None | \n", + "None | \n", + "
| 4 | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "AdamsonWeissman2016_GSM2406675_10X001.h5ad | \n", + "AAACATACTGTTCT | \n", + "62(mod)_pBA581 | \n", + "62(mod)_pBA581 | \n", + "None | \n", + "None | \n", + "CRISPR | \n", + "K562 | \n", + "lymphoblasts | \n", + "... | \n", + "9864.0 | \n", + "2584 | \n", + "2 | \n", + "False | \n", + "False | \n", + "X only | \n", + "perturbation | \n", + "perturbation | \n", + "None | \n", + "None | \n", + "
5 rows × 22 columns
\n", + "| \n", + " | dataset_name | \n", + "n_cells | \n", + "n_single_perturbation | \n", + "n_control | \n", + "perturbation_type | \n", + "cell_line | \n", + "source_perturbation_col | \n", + "source_gene_col | \n", + "source_guide_col | \n", + "expression_location | \n", + "single_perturbation_ratio | \n", + "
|---|---|---|---|---|---|---|---|---|---|---|---|
| 0 | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "5768 | \n", + "6 | \n", + "16 | \n", + "[CRISPR] | \n", + "[K562] | \n", + "perturbation | \n", + "perturbation | \n", + "None | \n", + "X only | \n", + "0.001040 | \n", + "
| 1 | \n", + "AdamsonWeissman2016_GSM2406677_10X005 | \n", + "15006 | \n", + "13 | \n", + "3674 | \n", + "[CRISPR] | \n", + "[K562] | \n", + "perturbation | \n", + "perturbation | \n", + "None | \n", + "X only | \n", + "0.000866 | \n", + "
| 2 | \n", + "AdamsonWeissman2016_GSM2406681_10X010 | \n", + "65337 | \n", + "101 | \n", + "2714 | \n", + "[CRISPR] | \n", + "[K562] | \n", + "perturbation | \n", + "perturbation | \n", + "None | \n", + "X only | \n", + "0.001546 | \n", + "
| 3 | \n", + "ReplogleWeissman2022_K562_essential | \n", + "310385 | \n", + "299694 | \n", + "10816 | \n", + "[CRISPR] | \n", + "[K562] | \n", + "perturbation | \n", + "gene | \n", + "guide_id | \n", + "X only | \n", + "0.965556 | \n", + "
| \n", + " | dataset_name | \n", + "n_cells | \n", + "single_ratio | \n", + "control_ratio | \n", + "perturbation_type | \n", + "
|---|---|---|---|---|---|
| 3 | \n", + "ReplogleWeissman2022_K562_essential | \n", + "310385 | \n", + "0.965556 | \n", + "0.034847 | \n", + "[CRISPR] | \n", + "
| 2 | \n", + "AdamsonWeissman2016_GSM2406681_10X010 | \n", + "65337 | \n", + "0.001546 | \n", + "0.041538 | \n", + "[CRISPR] | \n", + "
| 1 | \n", + "AdamsonWeissman2016_GSM2406677_10X005 | \n", + "15006 | \n", + "0.000866 | \n", + "0.244835 | \n", + "[CRISPR] | \n", + "
| 0 | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "5768 | \n", + "0.001040 | \n", + "0.002774 | \n", + "[CRISPR] | \n", + "
| \n", + " | ensembl_id | \n", + "ncounts | \n", + "ncells | \n", + "reference_var_name | \n", + "
|---|---|---|---|---|
| ensembl_id | \n", + "\n", + " | \n", + " | \n", + " | \n", + " |
| ENSG00000237491 | \n", + "ENSG00000237491 | \n", + "18.0 | \n", + "18 | \n", + "RP11-206L10.9 | \n", + "
| ENSG00000188976 | \n", + "ENSG00000188976 | \n", + "6595.0 | \n", + "3772 | \n", + "NOC2L | \n", + "
| ENSG00000187961 | \n", + "ENSG00000187961 | \n", + "257.0 | \n", + "254 | \n", + "KLHL17 | \n", + "
| ENSG00000188290 | \n", + "ENSG00000188290 | \n", + "1815.0 | \n", + "1445 | \n", + "HES4 | \n", + "
| ENSG00000187608 | \n", + "ENSG00000187608 | \n", + "2946.0 | \n", + "2180 | \n", + "ISG15 | \n", + "
| \n", + " | dataset_name | \n", + "is_single_perturbation | \n", + "is_control | \n", + "perturbation_type | \n", + "svd_latent_1 | \n", + "svd_latent_2 | \n", + "scvi_latent_1 | \n", + "scvi_latent_2 | \n", + "svd_umap_1 | \n", + "svd_umap_2 | \n", + "scvi_umap_1 | \n", + "scvi_umap_2 | \n", + "is_single_perturbation_label | \n", + "is_control_label | \n", + "
|---|---|---|---|---|---|---|---|---|---|---|---|---|---|---|
| cell_key | \n", + "\n", + " | \n", + " | \n", + " | \n", + " | \n", + " | \n", + " | \n", + " | \n", + " | \n", + " | \n", + " | \n", + " | \n", + " | \n", + " | \n", + " |
| AdamsonWeissman2016_GSM2406675_10X001::GAAGTCTGGTGTTG | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "False | \n", + "False | \n", + "CRISPR | \n", + "35.384010 | \n", + "-6.692269 | \n", + "-0.310886 | \n", + "-0.143230 | \n", + "1.337706 | \n", + "-3.307590 | \n", + "10.162565 | \n", + "7.512173 | \n", + "non_single | \n", + "perturbed | \n", + "
| AdamsonWeissman2016_GSM2406675_10X001::AGGCAGGAGCGTTA | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "False | \n", + "False | \n", + "CRISPR | \n", + "37.330631 | \n", + "-8.365924 | \n", + "0.553059 | \n", + "-0.038306 | \n", + "2.900387 | \n", + "-2.000047 | \n", + "10.169436 | \n", + "9.600159 | \n", + "non_single | \n", + "perturbed | \n", + "
| AdamsonWeissman2016_GSM2406675_10X001::ACGTTGGAAACCAC | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "False | \n", + "False | \n", + "CRISPR | \n", + "37.511005 | \n", + "-8.362697 | \n", + "0.189952 | \n", + "0.852251 | \n", + "2.820119 | \n", + "-3.790564 | \n", + "11.277682 | \n", + "8.332374 | \n", + "non_single | \n", + "perturbed | \n", + "
| AdamsonWeissman2016_GSM2406675_10X001::TAACGTCTAGCATC | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "False | \n", + "False | \n", + "CRISPR | \n", + "38.322598 | \n", + "-7.435131 | \n", + "-0.409475 | \n", + "0.323751 | \n", + "3.220307 | \n", + "-3.741003 | \n", + "10.979252 | \n", + "8.406797 | \n", + "non_single | \n", + "perturbed | \n", + "
| AdamsonWeissman2016_GSM2406675_10X001::AGCACTGAAGATCC | \n", + "AdamsonWeissman2016_GSM2406675_10X001 | \n", + "False | \n", + "False | \n", + "CRISPR | \n", + "36.738476 | \n", + "-6.706728 | \n", + "0.605040 | \n", + "0.196592 | \n", + "0.960080 | \n", + "-3.801914 | \n", + "10.216475 | \n", + "7.173553 | \n", + "non_single | \n", + "perturbed | \n", + "