Autodiscjax is a python library built on top of jax and equinox to facilitate automated exploration and simulation of computational models of biological processes (such as gene, proteins or metabolites networks). It provides several already-implemented modules and pipelines to organize experimentation on these biological network pathways using curiosity-driven learning and exploration algorithms.
pip install autodiscjax
AutoDiscJax follows two main design principles:
- Everything is a module, where a module is simply a parametrized function that takes inputs and returns outputs (and log_data). All autodiscjax modules
adx.Moduleare implemented as equinox moduleseqx.Module, which essentially allows to represent the function as a callable PyTree (and hence to be compatible with jax transformations) while keeping an intuitive API for model building (python class with a _call_ method). The only add-on with respect to equinox is that when instantiating aadx.Module, the user must specify the module's outputs PyTree structure, shape and dtype. - An experiment pipeline defines (i) how modules interact sequentially and exchange information, and (ii) what information should be collected and saved in the experiment history.
AutoDiscJax provides a handful of already-implement modules and pipelines to
- Simulate biological networks while intervening on them according to our needs
- Automatically organize experimentation in those systems, by implementing a variety of exploration approaches such as random, optimization-driven and curiosity-driven search
- Analyze the discoveries of the exploration method, for instance by testing their robustness to various perturbations
Finally, AutoDiscJax takes advantage of JAX mains features (just-in-time compilation, automatic vectorization and automatic differentation) which are especially advantageous for parallel experimentation and computational speedups, as well as gradient-based optimization.
The project is licensed under the MIT license.
AutoDiscJax is inspired by:
- the auto_disc library purpose and structure (by the FLOWERS team)
- the equinox library module definition (by Patrick Kidger)
Library to: sbmltoodejax