Nexus: A Multi-Scale Simulator for Biological Control and Causal Discovery

Biological systems exhibit complex multi-scale dynamics spanning subcellular gene regulatory processes to tissue-level spatial organization. However, causal analysis and control of biological systems face fundamental challenges due to the predominantly static nature of experimental measurements and the difficulty of collecting large-scale interventional datasets. We present Nexus, a hierarchical biological simulator that addresses these limitations by generating realistic, multi-modal biological data across cellular and subcellular scales. Nexus integrates gene regulatory network dynamics for RNA and protein simulation with spatial models of cell movement, chemical diffusion, and reactions. Built on the JAX framework, Nexus leverages automatic vectorization and just-in-time compilation to achieve significant speedups over existing simulators, particularly for large cell populations. The simulator supports extensive customization, backpropagation for parameter learning, and generation of interventional data, enabling large-scale dataset creation for causal discovery and analysis. Critically, Nexus provides a testbed for advancing control strategies in biological systems, particularly hierarchical optimal control where interventions are applied simultaneously at multiple levels to steer the system towards desired states. We demonstrate Nexus's accuracy through validation against established simulators and real-world biological data, showing its potential to advance both causal reasoning and control strategies in biological systems.