Control

jax_morph.control.ODEController #

ODEController(
    input_specs, output_specs, hidden_size, *, tag
)

Abstract dynamic controller that integrates a per-cell ODE over one macro-step.

Subclasses declare their parameter fields and implement vector_field. A controller's evolving state is the concatenation of a latent hidden state and its output fields; hidden_size=0 yields a pure-output ODE without allocating a hidden-state field. Constructors take parameter values directly and do no random initialization - unset parameters default to zeros from the field specs.

Attributes:

input_specs –

Static tuple of fixed driver-field specifications.
output_specs –

Static tuple of observed evolving output-field specifications.
hidden_size –

Static latent-state width; hidden arrays have shape (capacity, hidden_size).
in_size –

Static flattened width of all input fields.
out_size –

Static flattened width of all output fields.
tag –

Static namespace for the optional latent state field.

Methods:

vector_field –

Evaluate the per-cell ODE derivative.
__call__ –

Integrate the ODE for one macro-step.

Parameters:

input_specs –

Per-cell fields held fixed while integrating the ODE.
output_specs –

Per-cell fields that form the observed evolving state.
hidden_size –

Width of the controller's latent evolving state.
tag –

Namespace for the latent state field.

Raises:

ValueError –

If hidden_size is negative.

vector_field #

vector_field(t, evolving, inputs)

Return the derivative of the evolving block with fixed per-cell inputs.

Parameters:

t –

Physical integration time within the macro-step.
evolving –

Batched concatenation of hidden and output state.
inputs –

Batched, fixed driver fields.

Returns:

–

Batched derivative with the same shape as evolving.

state_reads #

state_reads()

Return the driver fields read by the controller.

state_writes #

state_writes()

Return output fields and, when present, the heritable latent state field.

jax_morph.control.GeneNetworkConnectionist #

GeneNetworkConnectionist(
    input_specs,
    output_specs,
    hidden_size,
    *,
    W_gene=None,
    W_in=None,
    b=None,
    gamma=0.1,
    tag='gene',
)

Connectionist gene circuit: sigmoid-saturated production from a linear regulatory drive.

Each gene concentration \(g\) evolves as production minus linear degradation,

\[\dot{g} = \sigma\!\left(W_\mathrm{gene}\, g + W_\mathrm{in}\, u + b\right) - \gamma\, g,\]

where \(u\) are the driver inputs, \(\sigma\) is a saturating sigmoid, and \(\gamma\) is the per-gene degradation rate. Latent (hidden) genes, when present, evolve alongside the observed outputs and regulate every gene in the circuit.

Attributes:

W_gene –

Gene interaction matrix of shape (n_gene, n_gene).
W_in –

Input mixing matrix of shape (n_gene, n_in).
b –

Per-gene bias array of shape (n_gene,).
gamma –

Scalar or per-gene linear degradation rate.

Parameters:

input_specs –

Per-cell driver field specifications.
output_specs –

Per-cell observed gene field specifications.
hidden_size –

Number of latent regulator genes.
W_gene –

Gene interaction matrix of shape (n_gene, n_gene). Defaults to None (zeros).
W_in –

Input mixing matrix of shape (n_gene, n_in). Defaults to None (zeros).
b –

Per-gene bias of shape (n_gene,). Defaults to None (zeros).
gamma –

Per-gene linear degradation rate, scalar or shape (n_gene,). Defaults to 0.1.
tag –

Namespace for any latent gene field. Defaults to 'gene'.

Raises:

ValueError –

If hidden_size is negative or an explicit regulatory parameter has an invalid shape.

vector_field #

vector_field(t, evolving, inputs)

Return linear-drive production minus linear degradation for every cell.

Parameters:

t –

Physical integration time within the macro-step, unused by this autonomous field.
evolving –

Batched latent and observed gene concentrations.
inputs –

Batched driver concentrations.

Returns:

–

Batched gene concentration derivatives.

jax_morph.control.NeuralODE #

NeuralODE(
    input_specs,
    output_specs,
    hidden_size,
    *,
    mlp,
    tag='ode',
)

Generic neural ODE controller with an MLP per-cell vector field.

The evolving state \(y\) (hidden state concatenated with outputs) obeys \(\dot{y} = \mathrm{MLP}([u, y])\), with \(u\) the fixed driver inputs over the macro-step. Build a correctly sized MLP with make_mlp.

Attributes:

mlp –

Equinox MLP mapping flattened inputs plus evolving state to an evolving derivative.

Parameters:

input_specs –

Per-cell driver field specifications.
output_specs –

Per-cell observed evolving field specifications.
hidden_size –

Number of latent evolving values.
mlp –

An eqx.nn.MLP mapping in_size + hidden_size + out_size to hidden_size + out_size; build one with make_mlp.
tag –

Namespace for any latent state field. Defaults to 'ode'.

Raises:

ValueError –

If hidden_size is negative or mlp does not have the required input and output sizes.

make_mlp `classmethod` #

make_mlp(
    input_specs,
    output_specs,
    hidden_size,
    *,
    key,
    width=64,
    depth=2,
)

Return a correctly sized eqx.nn.MLP for a NeuralODE over these fields.

Parameters:

input_specs –

Per-cell driver field specifications.
output_specs –

Per-cell observed evolving field specifications.
hidden_size –

Number of latent evolving values.
key –

JAX random key used for MLP initialization.
width –

Width of every MLP hidden layer. Defaults to 64.
depth –

Number of MLP hidden layers. Defaults to 2.

Returns:

–

An MLP mapping in_size + hidden_size + out_size to hidden_size + out_size.

vector_field #

vector_field(t, evolving, inputs)

Return the MLP derivative for every cell's evolving controller state.

Parameters:

t –

Physical integration time within the macro-step, unused by this autonomous field.
evolving –

Batched latent and observed controller state.
inputs –

Batched driver fields.

Returns:

–

Batched derivatives of the controller state.

jax_morph.control.GeneNetworkMWC #

GeneNetworkMWC(
    input_specs,
    output_specs,
    *,
    hidden_size=0,
    log_rho=None,
    log_tau=None,
    F0=None,
    H_gene=None,
    log_K_gene=None,
    H_in=None,
    log_K_in=None,
    tag='gene',
)

Thermodynamic gene circuit with a log-occupancy (MWC-style) regulatory drive.

Each gene concentration \(g_i\) evolves as saturating production minus linear decay,

\[\dot{g}_i = \rho_i\,\sigma(F_i) - \frac{g_i}{\tau_i},\]

with the regulatory drive

\[F_i = F^0_i + \sum_j H^\mathrm{gene}_{ij}\,\ln\!\left(1 + \frac{g_j}{K^\mathrm{gene}_{ij}}\right) + \sum_k H^\mathrm{in}_{ik}\,\ln\!\left(1 + \frac{u_k}{K^\mathrm{in}_{ik}}\right),\]

where \(u\) are the driver inputs, \(\sigma\) is the sigmoid, \(\rho\) the maximum production rates, \(\tau\) the lifetimes, and \(K\) the concentration thresholds. The positive quantities \(\rho\), \(\tau\), and \(K\) are stored in log space and exponentiated within their dtype's finite positive range. Latent genes, when requested, evolve alongside the observed outputs and regulate every gene.

Attributes:

log_rho –

Log maximum production-rate array of shape (n_gene,).
log_tau –

Log lifetime array of shape (n_gene,).
F0 –

Per-gene activation-bias array of shape (n_gene,).
H_gene –

Gene-interaction array of shape (n_gene, n_gene).
log_K_gene –

Log gene-threshold array of shape (n_gene, n_gene).
H_in –

Input-mixing array of shape (n_gene, n_in).
log_K_in –

Log input-threshold array of shape (n_gene, n_in).

All parameters default to zeros, giving unit production rates and lifetimes, unit concentration thresholds, and an inert (zero-interaction) but well-posed circuit.

Parameters:

input_specs –

Per-cell driver field specifications.
output_specs –

Per-cell observed gene field specifications.
hidden_size –

Number of latent regulator genes. Defaults to 0.
log_rho –

Log maximum production rate of shape (n_gene,). Defaults to None (zeros).
log_tau –

Log lifetime of shape (n_gene,). Defaults to None (zeros).
F0 –

Per-gene activation bias of shape (n_gene,). Defaults to None (zeros).
H_gene –

Gene interaction strengths of shape (n_gene, n_gene). Defaults to None (zeros).
log_K_gene –

Log gene concentration thresholds of shape (n_gene, n_gene). Defaults to None (zeros).
H_in –

Input mixing strengths of shape (n_gene, n_in). Defaults to None (zeros).
log_K_in –

Log driver concentration thresholds of shape (n_gene, n_in). Defaults to None (zeros).
tag –

Namespace for any latent gene field. Defaults to 'gene'.

Raises:

ValueError –

If hidden_size is negative or an explicit thermodynamic parameter has an invalid shape.

vector_field #

vector_field(t, evolving, inputs)

Return thermodynamic production minus linear degradation for every cell.

Parameters:

t –

Physical integration time within the macro-step, unused by this autonomous field.
evolving –

Batched latent and observed gene concentrations.
inputs –

Batched driver concentrations.

Returns:

–

Batched gene concentration derivatives.

Control

jax_morph.control.ODEController #

vector_field #

state_reads #

state_writes #

jax_morph.control.GeneNetworkConnectionist #

vector_field #

jax_morph.control.NeuralODE #

make_mlp classmethod #

vector_field #

jax_morph.control.GeneNetworkMWC #

vector_field #

make_mlp `classmethod` #