Computational\nSynthetic Biology Enabled through JAX:\nA Showcase

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