Computational Synthetic Biology Enabled through JAX: A Showcase

Mathematical modeling is indispensable in synthetic biology but remains underutilized. Tackling problems, from optimizing gene networks to simulating intracellular dynamics, can be facilitated by the ever-growing body of modeling approaches, be they mechanistic, stochastic, data-driven, or AI-enabled. Thanks to progress in the AI community, robust frameworks have emerged to enable researchers to access complex computational hardware and compilation. Previously, these frameworks focused solely on deep learning, but they have been developed to the point where running different forms of computation is relatively simple, as made possible, notably, by the JAX library. Running simulations at scale on GPUs speeds up research, which compounds enable larger-scale experiments and greater usability of code. As JAX remains underexplored in computational biology, we demonstrate its utility in three example projects ranging from synthetic biology to directed evolution, each with an accompanying demonstrative Jupyter notebook. We hope that these tutorials serve to democratize the flexible scaling, faster run-times, easy GPU portability, and mathematical enhancements (such as automatic differentiation) that JAX brings, all with only minor restructuring of code.

In order to optimise for adaptation to a perturbation, a step response is added as a signal input and the system's response can be assessed in the objective function.Using the above equation, each optimisation loop is initialised with x = 0.The step input is then introduced by adding a fixed amount 100 at t = 0 to a species x sig that we define as the signal species, and the system is simulated for time interval t 1 .The robustness of the circuit can be evaluated using the metrics of sensitivity and precision, which have been defined in Ma and Tang et al. 2009 2 as a measure of how responsive a system is to a perturbation (sensitivity) and how closely it can return to its pre-perturbed state (precision), defined as following: where S is sensitivity, P is precision, I is the input, and O is the output.I i and I f represent the initial and final inputs, while O i and O f represent the initial and final outputs.r The cost function for the optimisation can thus be defined as a weighted sum of sensitivity and precision.To prevent the trivial case arising where the circuit parameters tend towards zero, the overshoot is also added to the loss (scaled and negative).The cost is thus defined through the following functions: where r is robustness, o is the overshoot, w s and w p are the weights for sensitivity and robustness respectively, and l is the cost.For optimising the system, we backpropagate from the cost function to find optimal parameters.While an in-depth explanation of this model's assumptions, equations and properties can be found in the cited work, a basic discussion of its definition is reproduced here.

Supplementary note for showcase
A hybrid approach is used to simulate the behaviour of an E. coli cell hosting a synthetic gene circuit.The expression of the cell's native genes is captured by just a handful of lumped variables, each representing the average dynamics of a large set of abundances of species with similar behaviour and functions.Due to this coarse-graining, stochastic fluctuations of individual variables are averaged out, thereby simulating the host cell's state deterministically.
Conversely, each gene in the synthetic gene circuit is treated individually, so the stochasticity of gene expression is expected to significantly affect the dynamics of the circuitry-associated variables. 4dynamics of the host cell X are thus described by the following system of ordinary differential equations (ODEs): Here, the cell's genome is partitioned into three classes: ribosomal (r), metabolic (a) and housekeeping (q).Since the mass fraction of housekeeping genes is always ϕ q = 0.59, we can avoid modelling it explicitly.For the remaining two classes, we consider their mRNA concentrations m r and m a and protein concentrations p r and p a .By analogy, the mRNA and protein abundances for any synthetic gene x l in the set X = {x 1 , x 2 , . . .
For a synthetic circuit gene x l ∈ X with mRNA and protein counts m x l and p x l , we list the considered stochastic reactions below.In these expressions, F i is the gene-specific transcription regulation function and all other parameters and functions have the same meaning and form as for the host cell model ODEs.
• Transcription at rate • mRNA degradation at rate • mRNA dilution at rate λm

Gene circuit description
The model and algorithm described above allow modelling of the behaviour of an arbitrary gene circuit hosted by a bacterial cell.Here we describe the particular circuit simulated in the showcase provided in this paper, which has been proposed and analysed in Sechkar and Steel et al. 2024 3 and also described the Matlab implementation of the simulation script used as the benchmark for our JAX version.
The circuit is an antithetic integral feedback controller aiming to maintain a constant level of competition for ribosomes in the cell, which is comprised of four genes: the sensor (sens), the annihilator (anti), the actuator (act), and the amplifier (amp).The concentration of the constitutively expressed sensor gene protein depends on the number of ribosomes available for translation, and regulates the expression of annihilator RNAs.By binding and mutually annihilating with the constitutively expressed mRNAs of the actuator gene, the annihilator thus implements integral feedback with ribosome availability as the controlled variable.The actuator protein in turn regulates the amplifier gene's expression.By adjusting the abundance of amplifier gene transcripts, the controller is able to manipulate the extent of mRNA competition for the cell's ribosomes, keeping it constant despite disturbances.As an example of such disturbance, we consider the synthetic gene dist which becomes active at a time t dist .
The stochastic behaviour of this circuit is therefore simulated according to the previous section, with the set of synthetic gene in the cell being X = {sens, anti, act, amp, dist} and the mutual RNA annihilation reaction possible for x l,1 = act and x l,2 = anti.The transcription regulation functions are thus for the constitutive sensor and actuator gene, for the regulated annihilator and actuator genes, and for the time-dependent disturbance gene.All synthetic gene parameters are given in Table S3, copied with minimal changes from the Supplementary Table S9 in Sechkar and Steel et al.

Simulation details
In all versions of the code, the tau-leaping step used in the simulation was ∆t = 10 −6 h.As for ODE integration, in the JAX-enabled program it was performed using the Euler method with a time step of 10 −7 h.In the JAX-free Python, we used the scipy.integrate.solveivp function with the RK45 method and relative and absolute tolerances of 10 −6 and 10 −9 .In the Matlab version, we employed the ode15s solver with the same integration tolerances.
Supplementary Table S1: Parameters appearing in the cell model ODEs and Supplementary Table S2.
‡ Identical across all native genes i ∈ {a, r}.

2 :
Stochastic cell model simulations Cell model definition and simulation algorithm To demonstrate the use of JAX for stochastic model simulations, we considered the example of a coarse-grained resource-aware cell model proposed in Sechkar and Steel et al. 2024.3 x l } expressed by the cell are denoted as m x l and p x l .Meanwhile, t c and t u represent the concentrations of aminoacylated and free tRNA concentrations in the cell, respectively.The extent of competition for ribosomes in the cell is captured by the "resource competition denominator" D (Equation (S13)).The meaning, definitions and values of all other parameters and functions appearing in the ODEs are displayed in Tables S1 and Tables S2 copied with minimal changes from Supplementary TablesS1 and S2inSechkar and Steel et al. 2024. 3 Decreases m x l by Increases p x l by 1, decreases t c by n x l (charged tRNAs being used up), increases t u by n x l (uncharged tRNAs produced) This definition describing the continuous and stochastic parts of the hybrid model allows its simulation on the time interval t ∈ (t start , t f inish ) using a tau-leaping Algorithm S1, wherein all model variables are gathered into the vector v = (m a m r p a R t c t u m x 1 . . .m x L p x 1 . . .p x L ) ⊤ .Determine the number of times each stochastic reaction occurred over (t, t + ∆t) by sampling a Poisson distribution with mean equal to reaction rate × ∆t Sum the total changes in variables caused by stochastic each reaction over (t, t + ∆t) to obtain ∆ stoch v Set v = v + ∆ det v + ∆ stoch v for all entries v j in v do if v j < 0 then • Protein dilution at rate λp x l .Decreases p x l by 1 Additionally, if RNA transcripts of gene x l,1 and gene x l,2 can bind to each other and S5 thereby mutually annihilate, an additional stochastic annihilation reaction is introduced, where θ is the RNAs' mutual annihilation rate:• mRNA annihilation at rate θm l,1 m l,2 .Decreases m l,1 by Algorithm S1 Hybrid tau-leaping cell model simulation algorithm.Set t = t start , v = v start while t < t f inish do Find the deterministic change in variables ∆ det v by integrating ODEs (S7)-(S12) over the time interval (t, t + ∆t)