Cellifex

---

---

⚠️ Early Development — API Unstable

Cellifex is in active alpha development. Public interfaces may change without deprecation warnings; gate tolerances continue to tighten as new parity tests land. Pin to a specific commit if reproducibility is required.

Area	Status	Current expectation
API surface	Unstable	Breaking changes are preferred over preserving bad interfaces; the layered-architecture contract (domain → service → infrastructure) is the stable boundary.
Differentiable cores	Active	FBA / dFBA / RBA / ec-FBA / GECKO / SKiMpy / regulation-CRN / signalling all hold tight SOTA-parity. The JAX-native QP / LP solver framework (jax_qp) is bit-exact with cobra HiGHS and jit/vmap/grad-composable end-to-end.
Learned gates	Active	All five genome-scale learned MLP gates (MEMOTE, OptKnock, FVA, MFA, DGPrime) have training-loop convergence tests. OptKnock, FVA, MFA and DGPrime have held-out parity tests against reference tools; on MEMOTE's own scores the MEMOTE gate does not yet generalise across hosts (see the held-out table).
Docs & tests	Maintained	Pre-commit blocks docstring coverage below 80%, layered-architecture violations and security findings. The test suite exercises every dynamics, gate and scenario module.

---

Overview

Cellifex treats the cell as a differentiable dynamical system. Every dynamics, constraint, loss, and surrogate module is JAX-native: fully jax.jit-traceable and jax.grad-differentiable. Classical constraint-based (FBA, dFBA, RBA) and kinetic (Michaelis–Menten, lin-log, COBRApy) simulators live strictly under cellifex.baselines as data-generation sources and SOTA comparison references — they never appear in the differentiable import chain, enforced by import-linter contracts.

The core thesis: traditional cell-state modelling splits into two non-interoperable worlds — rigid mechanistic simulators (accurate but non-differentiable) and surrogate ML models (fast but ungrounded). Cellifex unifies them. Every classical mechanism gets a jittable, differentiable replacement registered behind a six-gate module contract: I/O typing, accuracy spec, gradient cleanness, feasibility, UQ, and registered identity. The result is a single composable surface that supports gradient-based pathway design, end-to-end Bayesian inference, and sub-millisecond inference loops on genome-scale models.

Why Cellifex?

• Differentiable through the whole stack — From metabolite uptake to biomass to pathway score, every layer flows jax.grad. No callouts to scipy solvers in the critical path; jax_qp provides the LP/QP forward + OptNet KKT-IFT backward as pure JAX.
• SOTA-parity verified, not just executable — Each replacement module has held-out parity tests against the current leading open-source tool (cobra, straindesign, memote, mfapy, equilibrator-api, efmtool, equilibrator-pathway, tellurium, sklearn Lasso). The test contracts use empirical-floor thresholds with documented headroom margins.
• 8 host configurations out of the box — E. coli (core + iJO1366), S. cerevisiae (iMM904, standing in for Yeast9), B. subtilis, K. phaffii, Y. lipolytica, C. glutamicum, P. putida. Five GEMs are committed under models/; iJO1366 comes from cobra's bundled models, and iMM904 and iYO844 are downloaded from BiGG on first use.
• Six-gate module contract — Every public dynamics or gate module exposes (1) typed I/O, (2) accuracy spec, (3) clean jax.grad, (4) feasibility invariant, (5) UQ surface, (6) registered identity. The contract is type-checked at the protocol layer and grad-checked at runtime.
• Layered architecture enforced statically — domain → service → infrastructure import direction is enforced by import-linter; baselines live in a sealed compartment by a second import-linter contract. Bad-layer imports break CI.

Features

Differentiable cell-process modules

Module	Replaces
DifferentiableFBA	cobra.flux_balance_analysis
DifferentiableCRN (regulation)	bioscrape.py_simulate_model
DifferentiableRBAFull	scipy-HiGHS bilevel Scott-Hwa
DifferentiableGECKO	scipy-HiGHS ec-FBA
DifferentiableSkimpy	reversible-MM + Haldane closure
DifferentiableDFBA	COBRApy dFBA Mahadevan loop
DifferentiableECFBA	cobra solve_ec_fba
DifferentiableSignalling	tellurium / libRoadRunner BIOMD0000000010 MAPK cascade
KineticDynamics	tellurium / libRoadRunner symbolic MM rate-rule evaluation (also matches SKiMpy reversible-MM at zero product)

Learned MLP gates (sub-millisecond inference)

Gate	Replaces
MemoteLearnedGate	memote.suite.api.test_model
OptKnockLearnedGate	straindesign.compute_strain_designs (Burgard-Maranas bilevel MILP)
FVALearnedGate	cobra.flux_variability_analysis
MfaLearnedGate	mfapy.metabolicmodel.fitting_flux
DGPrimeLearnedGate	equilibrator_api.standard_dg_prime
RetrosynthRerankerGate	aizynthfinder StateScorer

All five genome-scale gates consume the typed 22-dim v2 GEM-feature surface from cellifex.data.gem_features.GEMFeaturesV2; jax.grad and nnx.grad flow cleanly through every gate.

Neural-operator + generative backbones

• Neural operators — FNO (1D-heat operator parity), DeepONet (antiderivative-operator parity), StandardPINN (1D-Poisson parity)
• Generative families — VAE, Flow (RealNVP), DDPM via artifex.generative_models.factory; wrapper-drift parity tests at atol=1e-6 against the raw sibling implementations
• L2O — LearnToOptimize via opifex.optimization.meta_optimization
• Neural-ODE / SDE / CDE — diffrax-based with nnx.Module wrapping
• SINDy — symbolic regression via opifex.discovery.sindy with numpy lstsq parity

JAX-native differentiable QP / LP solver

cellifex.dynamics.metabolism.jax_qp ships a fully JAX-compatible QP/LP solver framework. Forward: MPAX r²HPDHG (Lu-Yang 2023; bit-exact with cobra HiGHS at FBA scale) plus a pure-JAX Chambolle-Pock PDHG fallback. Backward: OptNet KKT-IFT via Murty iterative-refinement active-set polish + lineax.AutoLinearSolver adjoint solve. Reverse-mode jax.grad composes with jit/vmap end-to-end. No scipy callbacks — fully GPU-native.

Proteome-economics gates (analytical + learned)

MDF, ECM, EFM, cMCS, Dekel–Alon ridge — both analytical and learned variants. MDF gate matches equilibrator-pathway within 1e-3 kJ/mol on the Beber 2022 CCM fixture; EFM gate matches efmtool top-k ranking on a 4-metabolite × 6-reaction branched network.

Closed-loop DBTL

Real JBEI Flaviolin multi-cycle yield data (241 rows × 5 cycles, Zenodo 10.5281/zenodo.15093363) drives scenario I; scenario J adds human-in-the-loop gate review through the ai-for-science-hub bridge.

Parity & held-out verification

This section lists the tests that compare cellifex's implementations with external reference tools. Without such a comparison a differentiable replacement only shows that it runs; with it, the module is checked against the reference tool's output while staying composable through jax.grad / jit / vmap.

Each module below is compared with the current leading open-source tool for its task. The SOTA tools are not selected once and frozen — they are tracked: if the canonical reference for a task moves (e.g. OptKnock's open-source SOTA shifted from the COBRApy single-deletion proxy to straindesign's bilevel MILP), the cellifex reference moves with it. Tests are pytest.importorskip-gated on the reference package so the suite still runs when an extra is not installed. CI runs them in a separate reference-tools job that installs the inference extra and the mfapy group; bioscrape is not part of that extra, so its parity tests still skip there.

SOTA parity per module

Forward parity: the cellifex module's output is compared against the SOTA reference on a fixed input fixture and asserted within a documented tolerance.

Module	SOTA reference	Tolerance
DifferentiableFBA	cobra.flux_balance_analysis	flux MSE < 1%, biomass rel-err < 0.5% across 60 conditions
DifferentiableCRN	bioscrape.py_simulate_model	per-timepoint rel-err < 5% across all dynamic species
DifferentiableRBAFull	scipy-HiGHS bilevel Scott-Hwa	inner-LP biomass < 1%; μ* < 1e-3 h⁻¹
DifferentiableGECKO	scipy-HiGHS ec-FBA	biomass rel-err < 1% with real BRENDA k_cats
DifferentiableDFBA	COBRApy dFBA Mahadevan loop	biomass rel-L2 < 2%
DifferentiableECFBA	cobra solve_ec_fba	per-condition rel-err < 1% across 5 enzyme levels; envelope-theorem gradient non-zero in binding regime, zero in saturated
DifferentiableSignalling	tellurium / libRoadRunner	endpoint and trajectory rel-err < 5% on BIOMD0000000010
KineticDynamics	tellurium / libRoadRunner	per-reaction rate rel-err < 1e-12 on a 3-reaction MM cascade (A→B→C→D)
MemoteLearnedGate	MEMOTE snapshot-report score (memote.suite.reporting.SnapshotReport)	default-init gate score within 0.05 of the fixture's MEMOTE total
OptKnockLearnedGate	straindesign.compute_strain_designs	every SOTA design dominates wildtype target flux; gate finite + bounded on every design mask
FVALearnedGate	cobra.flux_variability_analysis	per-reaction bias-calibrated parity within 1e-2 mmol/gDW/h
MfaLearnedGate	synthetic stand-in fixture (pFBA mean, 10% sigma; mfapy cannot fit e_coli_core)	per-reaction (flux_mean, flux_sigma) bias-calibrated parity within 1e-3
DGPrimeLearnedGate	equilibrator_api.standard_dg_prime	per-reaction ΔG'° bias-calibrated parity within 1e-3 kJ/mol
mdf_gate	equilibrator_pathway.ThermodynamicModel.mdf_analysis	within 1e-3 kJ/mol on the Beber 2022 CCM fixture
efm_gate	efmtool.calculate_efms	top-k ranking matches enumerated EFMs on a 4-met × 6-rxn branched network
Surrogate backbone wrappers	raw artifex / opifex reference modules	per-seed wrapped(x) == direct(x) at atol=1e-6
FNO (neural_operators.py)	analytical 1D-heat operator (α=0.01, T=0.1)	held-out RMSE < 0.1 after 300 epochs on 64 analytical training pairs
DeepONet (neural_operators.py)	analytical antiderivative operator (Lu 2021 sin/cos family)	held-out RMSE < 0.05 after 1500 Adam steps
StandardPINN (pinns.py)	analytical 1D-Poisson solution	max-abs error < 0.05 after 1500 Adam steps

Held-out generalisation

Held-out parity goes deeper than forward parity: the learned gate is trained end-to-end on N − 1 conditions and verified on the N-th held-out condition, against the SOTA tool's output at that held-out condition. This is the test that distinguishes "gate matches the training data" from "gate generalises across the natural perturbation axis the gate is designed for". A gate can pass forward parity at a single calibration point yet fail to generalise — held-out tests catch that gap.

Gate	Perturbation axis	Held-out test	SOTA baseline	Tolerance
FVALearnedGate	biomass_fraction ∈ {0.5, 0.8, 0.95, 0.99}	leave-one-out at α = 0.8 (interior)	cobra.flux_variability_analysis per α	p95 per-reaction abs-err < 15 mmol/gDW/h
DGPrimeLearnedGate	(pH, ionic_strength, T, pMg) (4 tuples)	leave-one-out at physiological standard	equilibrator_api.standard_dg_prime per tuple	p95 per-reaction ΔG_mean < 5 kJ/mol; ΔG_sigma < 1 kJ/mol
OptKnockLearnedGate	multi-knockout design space	4-fold cross-validation on 39 SOTA bilevel-MILP designs	straindesign.compute_strain_designs succinate target	mean held-out abs-err < 1.5 mmol/gDW/h
MemoteLearnedGate	host organism (8-host sweep)	leave-one-host-out with L1-regularised training	sklearn.linear_model.Lasso (small-data sparse linear regression)	target: gate mean abs-err ≤ 110% of Lasso. Not met on memote 0.17.0 scores (gate 0.232, Lasso 0.156, training-mean predictor 0.216); the test is a strict expected failure

Training-loop convergence tests (10-step descent + 500-1500 step Adam fit) verify that the loss surface, gradient computation, and optimizer configuration work correctly end-to-end for each learned gate.

The shared tests/_gate_training.py::train_for_n_steps helper exposes an optional l1_strength kwarg for small-data sparse-regression training; both the L1 regularizer and the Adam loop run inside nnx.jit for fully JAX-native training. The SOTA baseline does not need to be JAX-compatible (e.g. sklearn.linear_model.Lasso for MEMOTE held-out); only the gate's own forward and training do.

Installation

Cellifex uses uv exclusively for dependency management. Conda, pip, and poetry are not supported.

# Linux / macOS CPU
uv pip install -e ".[all-cpu]"

# Linux GPU (CUDA 12 local)
uv pip install -e ".[all-gpu]"

# macOS Apple Silicon (Metal)
uv pip install -e ".[all-macos]"

# Inference extras (memote, straindesign, equilibrator-api, ...)
uv pip install -e ".[inference]"

# mfapy, the 13C-MFA reference, is published only on GitHub, so it is a uv
# dependency group rather than an extra
uv sync --extra inference --group mfapy

For first-time setup from a source checkout:

# One-time GPU-enabled venv setup
bash setup.sh
source activate.sh
uv pip install -e ".[dev,inference]"

The setup script:

• Creates a .venv with uv and a CUDA 12 backend if a GPU is detected
• Pins JAX to the matching backend automatically
• Writes a generated .cellifex.env and leaves .env for user-owned overrides

source activate.sh is required every shell session; it refreshes the managed backend state before applying user overrides.

Quickstart

Run the test suite

source activate.sh
uv run pytest -x

The tests exercise every dynamics module, learned gate, scenario, and held-out parity contract.

Solve a forward FBA

from flax import nnx
from cellifex.dynamics.metabolism.diff_fba import DifferentiableFBA, FBASolverConfig

# Load E. coli core and build a jittable FBA module.
module = DifferentiableFBA.from_host_id(
    "ecoli_core",
    rngs=nnx.Rngs(0),
    config=FBASolverConfig(strategy="relaxed_qp"),
)

# Forward solve — sub-millisecond on e_coli_core.
result = module()
print(f"biomass flux = {result['biomass']:.4f} h^-1")

Run a learned gate

import jax.numpy as jnp
from flax import nnx
from cellifex.gates.memote_learned import MemoteLearnedGate
from cellifex.data.gem_features import extract_gem_features_v2, gem_features_v2_to_mlp_input
from cellifex.hosts import HOSTS

# 22-dim v2 GEM features for e_coli_core.
model = HOSTS["ecoli_core"].load_model()
features = gem_features_v2_to_mlp_input(extract_gem_features_v2(model))

gate = MemoteLearnedGate(rngs=nnx.Rngs(0), reference_score=0.5)
score = float(gate(jnp.asarray(features)))
print(f"predicted MEMOTE total score = {score:.4f}")

Regenerate a benchmark report

The reports list the host and backbone roster of a version; every metric is marked not measured.

uv run python scripts/run_benchmark_sweep.py --backbones 14 --output docs/benchmarks/v2.0

Run the JBEI HITL review loop

uv run python scripts/run_scenario_j.py

Documentation

Start here

• Architecture — Three-tier layered design (domain / service / infrastructure), import contracts, six-gate module contract.
• Roadmap — Module catalogue, baseline-replacement trajectory, current status table.
• Index — Top-level feature tour with cards.

Domain references

• Host profiles — Eight host organism configs with GEM source provenance.
• Experiments — Nine reference scenarios (A–F, H–J) mapped to strain-design workflow steps.
• Strain-design alignment — DBTL-step ↔ cellifex-module mapping.
• Proteome economics — MDF / ECM / EFM / cMCS / Dekel–Alon gates.
• Coupling contracts — Module ↔ module typed interfaces.

API reference

• Index • Dynamics • Gates • Surrogates
• Hosts • Composers • Adapters • Schema • Baselines

Ecosystem

• Ecosystem — Sibling-repo dependency graph + contribution-back contract.
• Sibling pinning — Pin policy for cross-repo deps.
• Benchmarks — Host and backbone roster per version (metrics not measured yet).

Architecture

cellifex/
├── src/cellifex/
│   ├── dynamics/             # Differentiable cell processes (FBA, dFBA, RBA, GECKO, signalling, ...)
│   │   └── metabolism/jax_qp/   # JAX-native QP / LP solver (MPAX + OptNet KKT-IFT)
│   ├── gates/                # Analytical + learned gate modules (MEMOTE, OptKnock, FVA, MFA, ΔG'°, ...)
│   ├── surrogates/           # Backbones (Neural-ODE, FNO, DeepONet, PINN, SINDy, generative, L2O) + registry
│   ├── schema/               # Typed state, protocols, contracts (cross-cutting domain-layer)
│   ├── losses/               # Biology-informed + composite loss functions
│   ├── constraints/          # Mass-balance, thermodynamic, proteome-economic relaxations
│   ├── derivers/             # Domain-layer derived quantities
│   ├── hosts/                # Eight organism configs
│   ├── data/                 # Concentration contracts, GEM features (v2), provenance enum
│   ├── composers/            # Service-layer assemblies (ec_core_cell, full_cell)
│   ├── engines/              # Rollout + active-learning drivers
│   ├── metrics/              # Flux + trajectory metrics
│   ├── baselines/            # Sealed SynBio 1.0 + SOTA references (isolated via import-linter)
│   ├── adapters/             # Sibling-repo bridges (artifex, opifex, datarax, calibrax, hub)
│   ├── io/                   # Dataset readers / writers
│   └── monitoring.py         # Profiler wrapper
├── docs/                     # Public documentation
├── tests/                    # Unit, integration, scenario, parity tests
├── scripts/                  # Operational entry points (benchmark sweep, regen, audit)
├── notebooks/                # Demo + walkthrough notebooks
└── models/                   # Committed SBML + COBRA models

See Architecture for the layered import map and module contracts.

Sibling repositories

Cellifex is part of the Avitai Bio ecosystem and consumes stable primitives from eight sibling packages. Every cross-repo dependency is pinned in pyproject.toml and its role is documented in docs/ecosystem.md.

Sibling	Role in cellifex
datarax	Dataset + DAG pipeline primitives (Element, Batch, Source, Pipeline)
artifex	Generative-model factory (VAE / Flow / DDPM backbones, MLP, Trainer)
opifex	Neural operators, PINNs, SINDy, L2O, Neural-ODE adapters
calibrax	Metrics + benchmark composition (Metric, MetricSuite, Comparison)
DiffBio	Differentiable bioinformatics operators + soft_ops (soft argmax / quantile / sort)
ai-for-science-hub	HITL review queue + LLM-completion bridge for scenario J
playground	Monitoring profiler + neural-vitai species embedding service
avitai-knowledge	Knowledge-base lookups (kinetic priors, kcat tables)

Development

Pre-commit gates

19 hooks block any commit that fails. The non-trivial ones:

uv run pre-commit run --all-files

• ruff linter + ruff format formatter
• pyright type checker (zero errors enforced on src/)
• bandit security scanner (no medium/high findings)
• import-linter layered-architecture contracts (domain ≠ infrastructure; baselines isolated)
• interrogate docstring coverage (≥ 80%)
• shellcheck for bash / sh scripts
• A custom Ruff config blocking print statements, mutable defaults, bare except, and other SWE-bans

Test suite

uv run pytest                                    # full suite (about 17 min on a CI runner)
uv run pytest -x --ignore=tests/test_jax_qp.py   # fast path (~30 s)
uv run pytest tests/test_memote_learned.py       # single module

The tests cover:

• Per-module unit tests (forward shape, grad finiteness, registered identity)
• Per-module SOTA-parity tests (importorskip-gated when extras absent)
• Scenario regression tests (A–J)
• Strict jit+grad smoke tests across every public module
• Held-out parity tests against the current leading open-source SOTA tool per module

Engineering principles

Cellifex follows the avitai-ecosystem-wide standards:

1. Zero tech debt — refactor toward the better design now, never defer.
2. TDD mandatory — tests first; thresholds set from empirical floors; never loosen to make a failing test pass.
3. DRY essential — extract on the third occurrence; centralise constants; one source of truth per knowledge.
4. No backward compatibility constraints — breaking changes are preferred when they produce cleaner foundations.

See Architecture and Coupling contracts for the design framework these principles operate within.

Project status

Cellifex is in active alpha.

• Installation, quickstart, and developer-guide docs are built with mkdocs build --strict in CI.
• CI runs the test suite on every push to main and on pull requests; the tests that compare against reference tools, including the held-out parity tests, run in the reference-tools job.
• Public API surfaces evolve when a simpler or more truthful runtime design requires it. There is no LTS branch.

Use Architecture, Roadmap, and Ecosystem as the current source of truth for supported workflows.

Contributing

Cellifex accepts contributions through the standard repository workflow:

1. Clone the repository and run bash setup.sh.
2. Activate the environment: source activate.sh.
3. Create a feature branch.
4. Add or update tests first (TDD); document any new public surface.
5. Run uv run pytest and uv run pre-commit run --all-files.
6. Open a PR. CI runs the test suite, pre-commit, pyright, and import-linter.

New public modules document the literature formulation they implement, the reference tool their parity tests compare against, and their differentiability contract.

Citation

If you use Cellifex in research, please cite:

@software{cellifex_2026,
  title = {Cellifex: An AI/ML-Native Cell-State Dynamics Foundation Model},
  author = {Shafiei, Mahdi and contributors},
  year = {2026},
  url = {https://github.com/avitai/cellifex},
  version = {0.0.1}
}

License

This project is licensed under the MIT License — see LICENSE for details.

Acknowledgments

Cellifex builds on several strong open-source projects:

• JAX — Numerical computing and transformations
• Flax — Neural network modules with NNX support
• Optax — Optimization utilities
• Orbax — Checkpointing
• diffrax — ODE / SDE / CDE solvers in JAX
• lineax — JAX-native linear solves (adjoint backward in jax_qp)
• MPAX — JAX-native r²HPDHG LP solver (forward in jax_qp)
• equinox — Library backbones for jax_qp/_pdhg_solver.py
• COBRApy — Reference FBA / FVA / dFBA implementations
• memote — SOTA GEM quality scoring
• straindesign — SOTA bilevel-MILP OptKnock (Schneider 2022)
• mfapy — Canonical 13C-MFA SLSQP reference
• equilibrator-api — eQuilibrator-CC ΔG'° (Beber 2022)
• efmtool — Schuster 2000 EFM enumeration
• BlackJAX — Bayesian flux posterior inference
• tellurium — libRoadRunner SBML reference simulator

---