diffbio 0.1.0__tar.gz → 0.1.1__tar.gz

{diffbio-0.1.0 → diffbio-0.1.1}/.gitignore

@@ -217,3 +217,11 @@ CLAUDE.md
 # Benchmark results (generated)
 benchmarks/results/
 data/
+
+# Local tracer-bullet tickets and scratch work
+.scratch/
+
+# Private working notes and per-editor assistant configuration. Local only.
+.claude/
+.cursor/
+AGENTS.md

{diffbio-0.1.0 → diffbio-0.1.1}/PKG-INFO

@@ -1,6 +1,6 @@
-Metadata-Version: 2.4
+Metadata-Version: 2.5
 Name: diffbio
-Version: 0.1.0
+Version: 0.1.1
 Summary: DiffBio: End-to-end differentiable bioinformatics pipelines built on Datarax, Artifex, Opifex, and Calibrax
 Project-URL: Bug Tracker, https://github.com/avitai/DiffBio/issues
 Project-URL: Documentation, https://diffbio.readthedocs.io
@@ -35,7 +35,6 @@ Classifier: Intended Audience :: Education
 Classifier: Intended Audience :: Science/Research
 Classifier: License :: OSI Approved :: MIT License
 Classifier: Programming Language :: Python :: 3 :: Only
-Classifier: Programming Language :: Python :: 3.11
 Classifier: Programming Language :: Python :: 3.12
 Classifier: Topic :: Scientific/Engineering
 Classifier: Topic :: Scientific/Engineering :: Artificial Intelligence
@@ -43,32 +42,30 @@ Classifier: Topic :: Scientific/Engineering :: Bio-Informatics
 Classifier: Topic :: Software Development
 Classifier: Topic :: Software Development :: Libraries
 Classifier: Topic :: Software Development :: Libraries :: Python Modules
-Requires-Python: >=3.11
+Requires-Python: >=3.12
 Requires-Dist: anndata>=0.9.1
-Requires-Dist: avitai-artifex>=0.1.0
+Requires-Dist: avitai-artifex>=0.1.4
 Requires-Dist: beartype>=0.14.1
 Requires-Dist: biopython>=1.81
-Requires-Dist: calibrax>=0.1.1
+Requires-Dist: calibrax>=0.1.2
 Requires-Dist: chex>=0.1.7
-Requires-Dist: datarax>=0.1.3
-Requires-Dist: flax>=0.12.0
+Requires-Dist: datarax>=0.1.5
+Requires-Dist: flax>=0.12.1
 Requires-Dist: h5py>=3.7
 Requires-Dist: jax-md>=0.2.27
 Requires-Dist: jax>=0.6.1
 Requires-Dist: jaxtyping>=0.2.20
 Requires-Dist: numpy>=1.24
-Requires-Dist: opifex>=0.1.0
-Requires-Dist: optax>=0.1.4
+Requires-Dist: opifex>=0.2.1
+Requires-Dist: optax>=0.2.8
 Requires-Dist: orbax-checkpoint>=0.11.10
-Requires-Dist: pre-commit>=4.3.0
-Requires-Dist: rdkit>=2025.9.3
-Requires-Dist: ruff>=0.1.5
 Requires-Dist: scipy>=1.10
 Provides-Extra: all
 Requires-Dist: bandit[toml]>=1.8.6; extra == 'all'
 Requires-Dist: beartype>=0.14.1; extra == 'all'
 Requires-Dist: build>=1.0.3; extra == 'all'
 Requires-Dist: coverage>=7; extra == 'all'
+Requires-Dist: datasets>=2.14; extra == 'all'
 Requires-Dist: deepchem>=2.8.0; extra == 'all'
 Requires-Dist: flake8-functions-names>=0.4; extra == 'all'
 Requires-Dist: flake8>=7.0; extra == 'all'
@@ -77,7 +74,6 @@ Requires-Dist: import-linter>=2.5; extra == 'all'
 Requires-Dist: interrogate>=1.7.0; extra == 'all'
 Requires-Dist: ipykernel>=6.29.5; extra == 'all'
 Requires-Dist: jax[cuda12]>=0.6.1; extra == 'all'
-Requires-Dist: jaxlib>=0.6.1; extra == 'all'
 Requires-Dist: lineax>=0.0.8; extra == 'all'
 Requires-Dist: matplotlib>=3.7; extra == 'all'
 Requires-Dist: mkdocs-include-exclude-files>=0.1; extra == 'all'
@@ -105,6 +101,7 @@ Requires-Dist: pytest-xdist>=3.6; extra == 'all'
 Requires-Dist: pytest>=8.3.5; extra == 'all'
 Requires-Dist: python-dotenv>=1; extra == 'all'
 Requires-Dist: radon>=6.0.1; extra == 'all'
+Requires-Dist: rdkit>=2025.9.3; extra == 'all'
 Requires-Dist: ruff>=0.1.5; extra == 'all'
 Requires-Dist: scib-metrics>=0.5; extra == 'all'
 Requires-Dist: shellcheck-py>=0.10.0.1; extra == 'all'
@@ -113,6 +110,7 @@ Requires-Dist: tabulate>=0.9; extra == 'all'
 Requires-Dist: torch>=1.13.0; extra == 'all'
 Requires-Dist: wemake-python-styleguide>=1.0; extra == 'all'
 Provides-Extra: benchmark
+Requires-Dist: datasets>=2.14; extra == 'benchmark'
 Requires-Dist: deepchem>=2.8.0; extra == 'benchmark'
 Requires-Dist: matplotlib>=3.7; extra == 'benchmark'
 Requires-Dist: pandas>=2.0; extra == 'benchmark'
@@ -120,6 +118,8 @@ Requires-Dist: pynndescent>=0.5; extra == 'benchmark'
 Requires-Dist: scib-metrics>=0.5; extra == 'benchmark'
 Requires-Dist: squidpy>=1.4; extra == 'benchmark'
 Requires-Dist: tabulate>=0.9; extra == 'benchmark'
+Provides-Extra: chem
+Requires-Dist: rdkit>=2025.9.3; extra == 'chem'
 Provides-Extra: cuda-dev
 Requires-Dist: bandit[toml]>=1.8.6; extra == 'cuda-dev'
 Requires-Dist: build>=1.0.3; extra == 'cuda-dev'
@@ -130,7 +130,6 @@ Requires-Dist: import-linter>=2.5; extra == 'cuda-dev'
 Requires-Dist: interrogate>=1.7.0; extra == 'cuda-dev'
 Requires-Dist: ipykernel>=6.29.5; extra == 'cuda-dev'
 Requires-Dist: jax[cuda12]>=0.6.1; extra == 'cuda-dev'
-Requires-Dist: jaxlib>=0.6.1; extra == 'cuda-dev'
 Requires-Dist: pylint>=3.3.8; extra == 'cuda-dev'
 Requires-Dist: pyright>=1.1.336; extra == 'cuda-dev'
 Requires-Dist: pytest-asyncio>=0.23; extra == 'cuda-dev'
@@ -147,6 +146,8 @@ Requires-Dist: radon>=6.0.1; extra == 'cuda-dev'
 Requires-Dist: ruff>=0.1.5; extra == 'cuda-dev'
 Requires-Dist: shellcheck-py>=0.10.0.1; extra == 'cuda-dev'
 Requires-Dist: wemake-python-styleguide>=1.0; extra == 'cuda-dev'
+Provides-Extra: cuda12
+Requires-Dist: jax[cuda12]>=0.6.1; extra == 'cuda12'
 Provides-Extra: dev
 Requires-Dist: bandit[toml]>=1.8.6; extra == 'dev'
 Requires-Dist: build>=1.0.3; extra == 'dev'
@@ -183,9 +184,6 @@ Requires-Dist: pymdown-extensions>=10.14.3; extra == 'docs'
 Provides-Extra: genomics
 Requires-Dist: pyfaidx>=0.8.0; extra == 'genomics'
 Requires-Dist: pysam>=0.22.0; extra == 'genomics'
-Provides-Extra: gpu
-Requires-Dist: jax[cuda12]>=0.6.1; extra == 'gpu'
-Requires-Dist: jaxlib>=0.6.1; extra == 'gpu'
 Provides-Extra: soft-ops-advanced
 Requires-Dist: lineax>=0.0.8; extra == 'soft-ops-advanced'
 Requires-Dist: optimistix>=0.0.9; extra == 'soft-ops-advanced'
@@ -211,7 +209,7 @@ Description-Content-Type: text/markdown
 # DiffBio
 
 <p align="center">
-  <a href="https://www.python.org/downloads/"><img src="https://img.shields.io/badge/python-3.11+-blue.svg" alt="Python 3.11+"></a>
+  <a href="https://www.python.org/downloads/"><img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="Python 3.12+"></a>
   <a href="https://jax.readthedocs.io/"><img src="https://img.shields.io/badge/JAX-0.6.1+-green.svg" alt="JAX"></a>
   <a href="https://flax.readthedocs.io/"><img src="https://img.shields.io/badge/Flax-0.12+-orange.svg" alt="Flax"></a>
   <a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-blue.svg" alt="License"></a>
@@ -222,11 +220,28 @@ Description-Content-Type: text/markdown
 </p>
 
 <p align="center">
-  Built on <a href="https://github.com/avitai/datarax">Datarax</a>, <a href="https://github.com/avitai/artifex">Artifex</a>, <a href="https://github.com/avitai/Opifex">Opifex</a>, and <a href="https://github.com/avitai/calibrax">Calibrax</a> | Powered by <a href="https://jax.readthedocs.io/">JAX</a> & <a href="https://flax.readthedocs.io/">Flax NNX</a>
+  Built on <a href="https://github.com/avitai/datarax">Datarax</a>, <a href="https://github.com/avitai/artifex">Artifex</a>, <a href="https://github.com/avitai/opifex">Opifex</a>, and <a href="https://github.com/avitai/calibrax">Calibrax</a> | Powered by <a href="https://jax.readthedocs.io/">JAX</a> & <a href="https://flax.readthedocs.io/">Flax NNX</a>
 </p>
 
 ---
 
+> **Research preview.** DiffBio is under rapid iteration and the API will change while we iterate
+> toward v1.0. What that means concretely:
+>
+> | Area | Status | Impact |
+> |------|--------|--------|
+> | **API** | 🔄 Unstable | Breaking changes are expected. Public interfaces may change without deprecation warnings. Pin to specific commits if stability is required. |
+> | **Tests** | 🔄 In Flux | Test suite is being expanded. Some tests may fail or be skipped. Coverage metrics are improving but not yet full. |
+> | **Documentation** | 🔄 Evolving | Docs may not reflect current implementation. Code examples might be outdated. Refer to source code and tests for accurate usage. |
+>
+> Pin a version if you need stability, and do not put it in production yet. For research and
+> experimentation it is ready to use today, with the understanding that APIs will evolve.
+>
+> This is public this early on purpose. Issues, questions and pull requests genuinely steer
+> what gets built next, and a star tells us which layer to push on.
+
+---
+
 ## Overview
 
 DiffBio is a framework for building **end-to-end differentiable bioinformatics
@@ -258,7 +273,7 @@ This enables learning optimal pipeline parameters directly from data, rather tha
 - **Composable Architecture** built on the Datarax, Artifex, Opifex, and Calibrax stack
 - **Training Utilities** with gradient clipping, custom loss functions, and synthetic data generation
 
-For complete operator and pipeline listings, see the [Operators Overview](https://docs.avitai.bio/diffbio/user-guide/operators/overview/) and [Pipelines Overview](https://docs.avitai.bio/diffbio/user-guide/pipelines/overview/) in the documentation.
+For complete operator and pipeline listings, see the [Operators Overview](https://diffbio.readthedocs.io/en/latest/user-guide/operators/overview/) and [Pipelines Overview](https://diffbio.readthedocs.io/en/latest/user-guide/pipelines/overview/) in the documentation.
 
 ## Installation
 
@@ -280,24 +295,22 @@ import jax
 import jax.numpy as jnp
 from flax import nnx
 
-from diffbio.operators import DifferentiableQualityFilter
-from diffbio.operators.variant.pileup import DifferentiablePileup
-from diffbio.operators.alignment.smith_waterman import SmoothSmithWaterman
+from diffbio.operators import DifferentiableQualityFilter, QualityFilterConfig
 
-# Quality filtering with learnable threshold
+# Quality filtering with learnable threshold (default initial_threshold=20.0)
 quality_filter = DifferentiableQualityFilter(
-    threshold=20.0,
-    temperature=1.0,
+    QualityFilterConfig(initial_threshold=20.0),
     rngs=nnx.Rngs(0),
 )
 
-# Apply to reads
+# Apply to a one-hot encoded sequence with per-position quality scores
 quality_scores = jnp.array([35.0, 15.0, 28.0, 10.0])
-reads = jax.nn.one_hot(jnp.array([[0, 1, 2, 3]] * 4), 4)
-data = {"reads": reads, "quality": quality_scores}
+sequence = jax.nn.one_hot(jnp.array([0, 1, 2, 3]), 4)  # (length, alphabet=4)
+data = {"sequence": sequence, "quality_scores": quality_scores}
 
 filtered_data, _, _ = quality_filter.apply(data, {}, None)
-# filtered_data["weights"] contains soft weights for each read
+# filtered_data["sequence"]        - sequence with low-quality positions softly suppressed
+# filtered_data["quality_scores"]  - pass-through quality values
 ```
 
 ### Using the Variant Calling Pipeline
@@ -386,9 +399,12 @@ DiffBio sits on a layered ecosystem rather than standing alone:
 |---|---|---|
 | Execution contracts | [Datarax](https://github.com/avitai/datarax) | Operator, data-source, and pipeline contracts |
 | Modeling substrate | [Artifex](https://github.com/avitai/artifex) | Reusable transformer and generative-model components |
-| Scientific ML substrate | [Opifex](https://github.com/avitai/Opifex) | Scientific optimization, operator learning, and advanced training methods |
+| Scientific ML substrate | [Opifex](https://github.com/avitai/opifex) | Scientific optimization, operator learning, and advanced training methods |
 | Evaluation substrate | [Calibrax](https://github.com/avitai/calibrax) | Metrics, benchmarking, comparison, profiling, and regression checks |
-| Biology-specific layer | DiffBio | Differentiable biological operators and domain compositions |
+
+DiffBio itself sits on top of these as the biology-specific layer: differentiable
+biological operators and end-to-end pipeline compositions (alignment, variant
+calling, single-cell analysis, drug discovery, structural biology, multi-omics).
 
 Each DiffBio operator inherits from Datarax's `OperatorModule` and implements:
 
@@ -403,18 +419,31 @@ This enables:
 
 ### Operator Composition
 
-Operators are chained by threading the `(data, state, metadata)` triple
-returned by `apply()` into the next operator:
+`apply()` runs an operator on a single element (no batch dimension). Operators
+are chained by threading the `(data, state, metadata)` triple returned by
+`apply()` into the next operator:
 
 ```python
-data, state, metadata = quality_filter.apply(batch_data, {}, None)
+data, state, metadata = quality_filter.apply(element_data, {}, None)
 data, state, metadata = pileup.apply(data, state, metadata)
 data, state, metadata = classifier.apply(data, state, metadata)
 
-# `data` is a dict of JAX arrays — read out the per-position predictions
+# `data` is a dict of JAX arrays - read out the per-position predictions
 predictions = data["logits"]
 ```
 
+For batched data wrapped in a Datarax `Batch`, call the operator directly
+(or use `apply_batch()`); both delegate to the same code path:
+
+```python
+from datarax import Batch
+
+batch = Batch.from_parts(...)        # construct from a list of elements
+batch = quality_filter(batch)         # equivalent to quality_filter.apply_batch(batch)
+batch = pileup(batch)
+batch = classifier(batch)
+```
+
 ## Testing
 
 ```bash
@@ -435,30 +464,44 @@ uv run pytest tests/integration/ -vv
 ```
 DiffBio/
 ├── src/diffbio/
-│   ├── core/               # Base operators, graph utils, soft ops
-│   ├── operators/           # 35+ differentiable operators
+│   ├── core/                # Base operators, graph utils, soft ops, neural components
+│   ├── operators/           # 40+ differentiable operators
 │   │   ├── alignment/       # Smith-Waterman, profile HMM, soft MSA
-│   │   ├── variant/         # Pileup, classifiers, CNV segmentation
-│   │   ├── singlecell/      # Clustering, trajectory, velocity, GRN, ...
-│   │   ├── drug_discovery/  # Fingerprints, property prediction, ADMET
-│   │   ├── epigenomics/     # Peak calling, chromatin state
-│   │   ├── normalization/   # VAE normalizer, UMAP, PHATE
+│   │   ├── assembly/        # GNN assembly, metagenomic binning
+│   │   ├── crispr/          # Guide RNA scoring
+│   │   ├── drug_discovery/  # Fingerprints, ADMET, AttentiveFP, MACCS keys
+│   │   ├── epigenomics/     # Peak calling, chromatin state, contextual epigenomics
+│   │   ├── foundation_models/ # Geneformer/scGPT adapters, transformer encoders
+│   │   ├── mapping/         # Neural read mapping
+│   │   ├── metabolomics/    # Spectral similarity
+│   │   ├── molecular_dynamics/ # Force fields, MD integrators
+│   │   ├── multiomics/      # Hi-C, spatial deconvolution, multi-omics VAE
+│   │   ├── normalization/   # VAE normalizer, UMAP, PHATE, embeddings
+│   │   ├── population/      # Ancestry estimation
+│   │   ├── preprocessing/   # Adapter removal, duplicate weighting, error correction
+│   │   ├── protein/         # Secondary structure
+│   │   ├── rna_structure/   # RNA folding
+│   │   ├── rnaseq/          # Splicing PSI, motif discovery
+│   │   ├── singlecell/      # Clustering, trajectory, velocity, GRN, batch correction, ...
 │   │   ├── statistical/     # HMM, NB GLM, EM quantification
-│   │   ├── multiomics/      # Hi-C, spatial deconvolution
-│   │   └── ...              # preprocessing, protein, RNA, assembly, ...
-│   ├── pipelines/           # End-to-end pipelines
-│   ├── losses/              # Alignment, single-cell, statistical losses
-│   ├── sources/             # Data loaders (FASTA, BAM, MolNet, ...)
-│   ├── splitters/           # Dataset splitting strategies
-│   └── utils/               # Training utilities
+│   │   └── variant/         # Pileup, classifiers, CNV segmentation
+│   ├── pipelines/           # 6 end-to-end pipelines
+│   ├── losses/              # Alignment, biological-regularization, single-cell, statistical, metric
+│   ├── sources/             # Data loaders (FASTA, BAM, AnnData, MoleculeNet, indexed views)
+│   ├── splitters/           # Random, stratified, scaffold, Tanimoto, sequence-identity
+│   ├── samplers/            # Perturbation samplers
+│   ├── sequences/           # DNA / RNA encoding utilities
+│   ├── evaluation/          # Evaluation runner and graders
+│   └── utils/               # Training utilities, dependency-runtime checks
 ├── tests/                   # Unit, integration, and benchmark tests
 ├── benchmarks/              # Domain benchmarks with training + baselines
+├── examples/                # Runnable example scripts paired with notebooks
 └── docs/                    # MkDocs documentation
 ```
 
 ## Requirements
 
-- Python 3.11+
+- Python 3.12+
 - JAX 0.6.1+
 - Flax 0.12+
 - Optax 0.1.4+
@@ -472,9 +515,8 @@ MIT License. See [LICENSE](LICENSE) for details.
 ## Acknowledgments
 
 DiffBio builds on ideas from:
-- [SMURF](https://www.biorxiv.org/content/10.1101/2021.10.23.465204): Differentiable Smith-Waterman for end-to-end MSA learning
 - [Datarax](https://github.com/avitai/datarax): Composable data processing framework
 - [Artifex](https://github.com/avitai/artifex): Generative-model and transformer substrate
-- [Opifex](https://github.com/avitai/Opifex): Scientific ML and advanced optimization substrate
+- [Opifex](https://github.com/avitai/opifex): Scientific ML and advanced optimization substrate
 - [Calibrax](https://github.com/avitai/calibrax): Benchmarking, comparison, and regression substrate
 - [Flax NNX](https://flax.readthedocs.io/): Neural network library for JAX

{diffbio-0.1.0 → diffbio-0.1.1}/README.md

@@ -1,7 +1,7 @@
 # DiffBio
 
 <p align="center">
-  <a href="https://www.python.org/downloads/"><img src="https://img.shields.io/badge/python-3.11+-blue.svg" alt="Python 3.11+"></a>
+  <a href="https://www.python.org/downloads/"><img src="https://img.shields.io/badge/python-3.12+-blue.svg" alt="Python 3.12+"></a>
   <a href="https://jax.readthedocs.io/"><img src="https://img.shields.io/badge/JAX-0.6.1+-green.svg" alt="JAX"></a>
   <a href="https://flax.readthedocs.io/"><img src="https://img.shields.io/badge/Flax-0.12+-orange.svg" alt="Flax"></a>
   <a href="LICENSE"><img src="https://img.shields.io/badge/license-MIT-blue.svg" alt="License"></a>
@@ -12,11 +12,28 @@
 </p>
 
 <p align="center">
-  Built on <a href="https://github.com/avitai/datarax">Datarax</a>, <a href="https://github.com/avitai/artifex">Artifex</a>, <a href="https://github.com/avitai/Opifex">Opifex</a>, and <a href="https://github.com/avitai/calibrax">Calibrax</a> | Powered by <a href="https://jax.readthedocs.io/">JAX</a> & <a href="https://flax.readthedocs.io/">Flax NNX</a>
+  Built on <a href="https://github.com/avitai/datarax">Datarax</a>, <a href="https://github.com/avitai/artifex">Artifex</a>, <a href="https://github.com/avitai/opifex">Opifex</a>, and <a href="https://github.com/avitai/calibrax">Calibrax</a> | Powered by <a href="https://jax.readthedocs.io/">JAX</a> & <a href="https://flax.readthedocs.io/">Flax NNX</a>
 </p>
 
 ---
 
+> **Research preview.** DiffBio is under rapid iteration and the API will change while we iterate
+> toward v1.0. What that means concretely:
+>
+> | Area | Status | Impact |
+> |------|--------|--------|
+> | **API** | 🔄 Unstable | Breaking changes are expected. Public interfaces may change without deprecation warnings. Pin to specific commits if stability is required. |
+> | **Tests** | 🔄 In Flux | Test suite is being expanded. Some tests may fail or be skipped. Coverage metrics are improving but not yet full. |
+> | **Documentation** | 🔄 Evolving | Docs may not reflect current implementation. Code examples might be outdated. Refer to source code and tests for accurate usage. |
+>
+> Pin a version if you need stability, and do not put it in production yet. For research and
+> experimentation it is ready to use today, with the understanding that APIs will evolve.
+>
+> This is public this early on purpose. Issues, questions and pull requests genuinely steer
+> what gets built next, and a star tells us which layer to push on.
+
+---
+
 ## Overview
 
 DiffBio is a framework for building **end-to-end differentiable bioinformatics
@@ -48,7 +65,7 @@ This enables learning optimal pipeline parameters directly from data, rather tha
 - **Composable Architecture** built on the Datarax, Artifex, Opifex, and Calibrax stack
 - **Training Utilities** with gradient clipping, custom loss functions, and synthetic data generation
 
-For complete operator and pipeline listings, see the [Operators Overview](https://docs.avitai.bio/diffbio/user-guide/operators/overview/) and [Pipelines Overview](https://docs.avitai.bio/diffbio/user-guide/pipelines/overview/) in the documentation.
+For complete operator and pipeline listings, see the [Operators Overview](https://diffbio.readthedocs.io/en/latest/user-guide/operators/overview/) and [Pipelines Overview](https://diffbio.readthedocs.io/en/latest/user-guide/pipelines/overview/) in the documentation.
 
 ## Installation
 
@@ -70,24 +87,22 @@ import jax
 import jax.numpy as jnp
 from flax import nnx
 
-from diffbio.operators import DifferentiableQualityFilter
-from diffbio.operators.variant.pileup import DifferentiablePileup
-from diffbio.operators.alignment.smith_waterman import SmoothSmithWaterman
+from diffbio.operators import DifferentiableQualityFilter, QualityFilterConfig
 
-# Quality filtering with learnable threshold
+# Quality filtering with learnable threshold (default initial_threshold=20.0)
 quality_filter = DifferentiableQualityFilter(
-    threshold=20.0,
-    temperature=1.0,
+    QualityFilterConfig(initial_threshold=20.0),
     rngs=nnx.Rngs(0),
 )
 
-# Apply to reads
+# Apply to a one-hot encoded sequence with per-position quality scores
 quality_scores = jnp.array([35.0, 15.0, 28.0, 10.0])
-reads = jax.nn.one_hot(jnp.array([[0, 1, 2, 3]] * 4), 4)
-data = {"reads": reads, "quality": quality_scores}
+sequence = jax.nn.one_hot(jnp.array([0, 1, 2, 3]), 4)  # (length, alphabet=4)
+data = {"sequence": sequence, "quality_scores": quality_scores}
 
 filtered_data, _, _ = quality_filter.apply(data, {}, None)
-# filtered_data["weights"] contains soft weights for each read
+# filtered_data["sequence"]        - sequence with low-quality positions softly suppressed
+# filtered_data["quality_scores"]  - pass-through quality values
 ```
 
 ### Using the Variant Calling Pipeline
@@ -176,9 +191,12 @@ DiffBio sits on a layered ecosystem rather than standing alone:
 |---|---|---|
 | Execution contracts | [Datarax](https://github.com/avitai/datarax) | Operator, data-source, and pipeline contracts |
 | Modeling substrate | [Artifex](https://github.com/avitai/artifex) | Reusable transformer and generative-model components |
-| Scientific ML substrate | [Opifex](https://github.com/avitai/Opifex) | Scientific optimization, operator learning, and advanced training methods |
+| Scientific ML substrate | [Opifex](https://github.com/avitai/opifex) | Scientific optimization, operator learning, and advanced training methods |
 | Evaluation substrate | [Calibrax](https://github.com/avitai/calibrax) | Metrics, benchmarking, comparison, profiling, and regression checks |
-| Biology-specific layer | DiffBio | Differentiable biological operators and domain compositions |
+
+DiffBio itself sits on top of these as the biology-specific layer: differentiable
+biological operators and end-to-end pipeline compositions (alignment, variant
+calling, single-cell analysis, drug discovery, structural biology, multi-omics).
 
 Each DiffBio operator inherits from Datarax's `OperatorModule` and implements:
 
@@ -193,18 +211,31 @@ This enables:
 
 ### Operator Composition
 
-Operators are chained by threading the `(data, state, metadata)` triple
-returned by `apply()` into the next operator:
+`apply()` runs an operator on a single element (no batch dimension). Operators
+are chained by threading the `(data, state, metadata)` triple returned by
+`apply()` into the next operator:
 
 ```python
-data, state, metadata = quality_filter.apply(batch_data, {}, None)
+data, state, metadata = quality_filter.apply(element_data, {}, None)
 data, state, metadata = pileup.apply(data, state, metadata)
 data, state, metadata = classifier.apply(data, state, metadata)
 
-# `data` is a dict of JAX arrays — read out the per-position predictions
+# `data` is a dict of JAX arrays - read out the per-position predictions
 predictions = data["logits"]
 ```
 
+For batched data wrapped in a Datarax `Batch`, call the operator directly
+(or use `apply_batch()`); both delegate to the same code path:
+
+```python
+from datarax import Batch
+
+batch = Batch.from_parts(...)        # construct from a list of elements
+batch = quality_filter(batch)         # equivalent to quality_filter.apply_batch(batch)
+batch = pileup(batch)
+batch = classifier(batch)
+```
+
 ## Testing
 
 ```bash
@@ -225,30 +256,44 @@ uv run pytest tests/integration/ -vv
 ```
 DiffBio/
 ├── src/diffbio/
-│   ├── core/               # Base operators, graph utils, soft ops
-│   ├── operators/           # 35+ differentiable operators
+│   ├── core/                # Base operators, graph utils, soft ops, neural components
+│   ├── operators/           # 40+ differentiable operators
 │   │   ├── alignment/       # Smith-Waterman, profile HMM, soft MSA
-│   │   ├── variant/         # Pileup, classifiers, CNV segmentation
-│   │   ├── singlecell/      # Clustering, trajectory, velocity, GRN, ...
-│   │   ├── drug_discovery/  # Fingerprints, property prediction, ADMET
-│   │   ├── epigenomics/     # Peak calling, chromatin state
-│   │   ├── normalization/   # VAE normalizer, UMAP, PHATE
+│   │   ├── assembly/        # GNN assembly, metagenomic binning
+│   │   ├── crispr/          # Guide RNA scoring
+│   │   ├── drug_discovery/  # Fingerprints, ADMET, AttentiveFP, MACCS keys
+│   │   ├── epigenomics/     # Peak calling, chromatin state, contextual epigenomics
+│   │   ├── foundation_models/ # Geneformer/scGPT adapters, transformer encoders
+│   │   ├── mapping/         # Neural read mapping
+│   │   ├── metabolomics/    # Spectral similarity
+│   │   ├── molecular_dynamics/ # Force fields, MD integrators
+│   │   ├── multiomics/      # Hi-C, spatial deconvolution, multi-omics VAE
+│   │   ├── normalization/   # VAE normalizer, UMAP, PHATE, embeddings
+│   │   ├── population/      # Ancestry estimation
+│   │   ├── preprocessing/   # Adapter removal, duplicate weighting, error correction
+│   │   ├── protein/         # Secondary structure
+│   │   ├── rna_structure/   # RNA folding
+│   │   ├── rnaseq/          # Splicing PSI, motif discovery
+│   │   ├── singlecell/      # Clustering, trajectory, velocity, GRN, batch correction, ...
 │   │   ├── statistical/     # HMM, NB GLM, EM quantification
-│   │   ├── multiomics/      # Hi-C, spatial deconvolution
-│   │   └── ...              # preprocessing, protein, RNA, assembly, ...
-│   ├── pipelines/           # End-to-end pipelines
-│   ├── losses/              # Alignment, single-cell, statistical losses
-│   ├── sources/             # Data loaders (FASTA, BAM, MolNet, ...)
-│   ├── splitters/           # Dataset splitting strategies
-│   └── utils/               # Training utilities
+│   │   └── variant/         # Pileup, classifiers, CNV segmentation
+│   ├── pipelines/           # 6 end-to-end pipelines
+│   ├── losses/              # Alignment, biological-regularization, single-cell, statistical, metric
+│   ├── sources/             # Data loaders (FASTA, BAM, AnnData, MoleculeNet, indexed views)
+│   ├── splitters/           # Random, stratified, scaffold, Tanimoto, sequence-identity
+│   ├── samplers/            # Perturbation samplers
+│   ├── sequences/           # DNA / RNA encoding utilities
+│   ├── evaluation/          # Evaluation runner and graders
+│   └── utils/               # Training utilities, dependency-runtime checks
 ├── tests/                   # Unit, integration, and benchmark tests
 ├── benchmarks/              # Domain benchmarks with training + baselines
+├── examples/                # Runnable example scripts paired with notebooks
 └── docs/                    # MkDocs documentation
 ```
 
 ## Requirements
 
-- Python 3.11+
+- Python 3.12+
 - JAX 0.6.1+
 - Flax 0.12+
 - Optax 0.1.4+
@@ -262,9 +307,8 @@ MIT License. See [LICENSE](LICENSE) for details.
 ## Acknowledgments
 
 DiffBio builds on ideas from:
-- [SMURF](https://www.biorxiv.org/content/10.1101/2021.10.23.465204): Differentiable Smith-Waterman for end-to-end MSA learning
 - [Datarax](https://github.com/avitai/datarax): Composable data processing framework
 - [Artifex](https://github.com/avitai/artifex): Generative-model and transformer substrate
-- [Opifex](https://github.com/avitai/Opifex): Scientific ML and advanced optimization substrate
+- [Opifex](https://github.com/avitai/opifex): Scientific ML and advanced optimization substrate
 - [Calibrax](https://github.com/avitai/calibrax): Benchmarking, comparison, and regression substrate
 - [Flax NNX](https://flax.readthedocs.io/): Neural network library for JAX

diffbio-0.1.1/benchmarks/crossmodality/README.md

@@ -0,0 +1,76 @@
+# Cross-modality reproduction benchmarks
+
+Scripts that reproduce the experiments in the DiffBio manuscript
+(`paper/diffbio.tex`) beyond the single-cell annotation result (which lives in
+`benchmarks/singlecell/`). Every script uses **real data**, fits all preprocessing on
+the **training split only**, reports **held-out** metrics over 3 seeds, and writes its
+summary to `benchmarks/results/crossmodality/`.
+
+Each script reads its dataset path from an environment variable with a default under
+`/mnt/ssd2/Data`; override the variable to point at your own copy.
+
+## Data sources
+
+| Pipeline | Dataset | Source | Env var |
+|---|---|---|---|
+| DNA regulatory elements | Genomic Benchmarks `human_ocr_ensembl` | HuggingFace `katarinagresova/Genomic_Benchmarks_human_ocr_ensembl` (DOI 10.1186/s12863-023-01123-8) | `HF_HOME` |
+| scATAC annotation | CATLAS GSE184462 | GEO `GSE184462_RAW.tar` + `GSE184462_metadata.tsv.gz` (Zhang 2021, 10.1016/j.cell.2021.10.024) | `DIFFBIO_CATLAS_DIR` |
+| Perturb-seq identity | Virtual Cell Challenge Perturb-seq | challenge `adata_Training.h5ad` | `DIFFBIO_VCC_NPZ` |
+| Variant calling | GIAB HG002 chr20 | `HG002.chr20.bam` + GIAB v4.2.1 truth VCF + GRCh38 chr20 FASTA (Zook 2019, 10.1038/s41587-019-0074-6) | `DIFFBIO_GIAB_DIR` |
+
+## Data preparation
+
+Prep scripts materialize the model-ready arrays the benchmarks consume:
+
+```bash
+python benchmarks/crossmodality/prepare_scatac.py       # CATLAS fragments -> cell x bin npz
+python benchmarks/crossmodality/prepare_variant.py       # GIAB BAM+truth -> labelled pileup windows
+python benchmarks/crossmodality/prepare_perturbation.py  # VCC h5ad -> target-masked npz (chunked)
+```
+
+DNA needs no prep (the HuggingFace dataset is featurized on the fly by the k-mer operator).
+
+## Reproducing paper results
+
+```bash
+source ./activate.sh
+
+# Table 3 / Figure 4 -- the joint gain across modalities and machineries (gain vs k)
+python benchmarks/crossmodality/dna_classification.py         # DNA, +7.5pp @ k=5
+python benchmarks/crossmodality/scatac_annotation.py          # scATAC, +12.2pp @ k=5
+python benchmarks/crossmodality/perturbation_identity.py      # Perturb-seq, +8.3pp @ k=5
+python benchmarks/crossmodality/variant_read_pooling.py       # variant (set pooling), +16.9pp @ k=5
+
+# Figure 5a / section 4.4 depth -- scaling crossover and domain baselines
+python benchmarks/crossmodality/scatac_deepening.py           # scATAC scaling + LSI+LogReg/kNN baseline
+python benchmarks/crossmodality/dna_deepening.py              # DNA scaling + 6-mer+linear baseline
+
+# Section 4.5 -- knob attribution on Perturb-seq (the projection carries the gain)
+python benchmarks/crossmodality/perturbation_ablation.py
+
+# Section 4.6 -- trainable and orthonormal reduction bases (MatrixFreePCA, Stiefel projection)
+python benchmarks/crossmodality/trainable_basis.py
+
+# Rigor audits (data leakage / overfitting)
+python benchmarks/crossmodality/audit_leakfree_scatac.py      # held-out donors + train-only bins, train/test F1
+python benchmarks/crossmodality/audit_overfitting.py          # 3-way train/val/test, DNA + Perturb-seq
+python benchmarks/crossmodality/audit_regularization_scatac.py  # weight-decay / early-stop / Stiefel vs joint overfit
+```
+
+## Data splits (no fitting on held-out data)
+
+Every benchmark fits **all** preprocessing (PCA / LSI / normalization / linear
+baselines) on the **training split only**, then applies the fitted transform to the
+held-out data. Splits:
+
+| Pipeline | Split |
+|---|---|
+| scRNA annotation | 80/20 stratified by cell type |
+| DNA regulatory elements | official Genomic Benchmarks disjoint train/test sequences (+15% validation carved from train for `audit_overfitting.py`) |
+| Perturb-seq identity | held-out experimental **batches** (target-gene columns masked) |
+| scATAC annotation | stratified 70/30 (headline); **held-out donors + train-only bin selection** for `audit_leakfree_scatac.py` |
+| Variant calling | disjoint **genomic regions** (chr20:2–22 Mb train / 30–40 Mb test) |
+
+Each script prints a per-`k` (or per-arm) summary and writes a JSON to
+`benchmarks/results/crossmodality/`. Numbers are means over 3 seeds; small differences
+from the paper are expected from GPU non-determinism in the convolution/scan kernels.