Open-Loop Controller

The OpenLoopNeuralController provides methods to learn time-dependent control inputs \(u(t)\).

Example

We first define a one-species, one-resource consumer-resource model (CRM) and assign its parameters.

from mgrowthctrl.models import CRModel, CRModelParams
from mgrowthctrl.models.base import ModelNames

params = CRModelParams.from_shapes(n=1, m=1)
names = ModelNames(["biomass"], ["substrate"])
model = CRModel(names=names, params=params, backend="torch")

model.r[0, 0] = 0.5
model.a[0, 0] = 1.0
model.k[0] = 0.01

The controller is trained to reach a target final biomass while penalizing control effort.

import numpy as np
import torch

torch.manual_seed(42)
np.random.seed(42)

def steering_loss(trajectory, controller, model):
    final_biomass = trajectory[-1, 0]
    target = torch.tensor(1e1, dtype=torch.float32)

    error = (torch.log(final_biomass + 1e-6) - torch.log(target + 1e-6)) ** 2

    steps = trajectory.shape[0]
    times = torch.linspace(0, controller.total_time, steps, dtype=torch.float32)
    t_norm = (times / controller.total_time).unsqueeze(1)

    U_tensor = controller.net(t_norm)
    cost = 1e-3 * U_tensor.pow(2).mean()

    return error + cost, U_tensor

Next, we initialize the open-loop neural controller.

from torch import nn
from mgrowthctrl.controllers import OpenLoopNeuralController

controller = OpenLoopNeuralController(
    total_time=100,
    n_species=model.n,
    n_resources=model.m,
    s_idx=[0],
    criterion=steering_loss,
    hidden_dims=[5, 5, 5],
    activation=nn.ELU(),
)

We define a logging function, set the initial condition, and train the controller.

def log_fn(model, epoch, loss, U_trajectory, sol_tensor):
    if epoch % 1 == 0:
        actual = sol_tensor[-1, 0].item()
        print(
            f"Epoch {epoch:03d} | Loss: {loss.item():.4f} "
            f"| Target: {1e1:.1e} | Actual: {actual:.1e}"
        )

y0 = torch.tensor([0.1, 10.0], dtype=torch.float32)
t_eval = np.linspace(0, 100, 100)
t_span = (0, 100)

controller.fit(model, y0, t_eval, t_span, epochs=100, lr=1e-2, log_fn=log_fn)

We visually compare the uncontrolled and controlled biomass trajectories.

import matplotlib.pyplot as plt

uncontrolled_sol = model.predict(y0.numpy(), t_eval)
controlled_sol = controller.simulate(model, y0, t_eval, t_span)

fig, ax = plt.subplots(figsize=(6, 3.2))
ax.plot(uncontrolled_sol.t, uncontrolled_sol.y[0], label="Uncontrolled")
ax.plot(controlled_sol.t, controlled_sol.y[0], label="Controlled")
ax.scatter([100], [1e1], color="red", s=150, marker="*", label="Target", zorder=10)

ax.set_xlabel("Time (h)")
ax.set_ylabel("Biomass")
ax.legend()
plt.tight_layout()
plt.savefig("basic_open_loop_result.png")
plt.show()

We observe that the controlled dynamics reaches the desired concentration:

Open-loop control predictions

Finally, we can inspect the learned control input over time.

raw_u = controller.get_input_history(t_eval, model).detach()

fig, ax = plt.subplots(figsize=(6, 3.2))
ax.plot(t_eval, raw_u[:, 0], label="Substrate input")
ax.set_xlabel("Time (h)")
ax.set_ylabel("Injection rate (mM/h)")
ax.legend()
plt.tight_layout()
plt.savefig("basic_open_loop_control_inputs.png")
plt.show()

Here’s what the result looks like:

Open-loop control inputs

For a detailed tutorial on open-loop control of a fitted consumer-resource model, please see the tutorial Advanced: Open-Loop Control of Microbiome Dynamics.

API Documentation

class mgrowthctrl.controllers.open_loop.OpenLoopNeuralController(
total_time: float,
n_species: int,
n_resources: int,
*,
hidden_dims: Sequence[int] | None = None,
activation: Module = Tanh(),
criterion: Callable[[Tensor, BaseController, BaseODEModel], Tuple[Tensor, Any]] | None = None,
injection_only: bool = True,
x_idx: Sequence[int] | None = None,
s_idx: Sequence[int] | None = None,
target_substrate_idx: int | None = None,
)

Bases: BaseController

Open-Loop Neural Controller (Time-Dependent Policy).

This controller learns a fixed schedule of interventions u(t) based solely on the current time t. It does not observe the current biological state (X, S) during execution.

This is analogous to a “blind archer” or a pre-calculated drug dosing schedule.

The controller outputs positive injection rates (via Softplus) which are added directly to the ODE derivatives:

  • dX/dt += u_x(t) (Probiotic injection)

  • dS/dt += u_s(t) (Nutrient/Prebiotic injection)

Parameters:

  • total_time (float) – The maximum time (T_final) of the experiment. Inputs are normalized to the range [0, 1] as t_norm = t / total_time.

  • n_species (int) – Number of species in the system (dimension of X).

  • n_resources (int) – Number of substrates in the system (dimension of S).

  • hidden_dims (Sequence[int], optional) – Architecture of the neural network (neurons per hidden layer). Default: [16, 16].

  • activation (nn.Module) – Activation function for hidden layers.

  • criterion (LossFunction, optional) – The objective function used for training. Signature: (trajectory, controller, model) -> (loss, aux_data).

  • injection_only (bool) – If True, applies Softplus to the network output so controls are nonnegative.

  • x_idx (Sequence[int], optional) – Indices of species (X) to control (add/inject). Default is None (control no species).

  • s_idx (Sequence[int], optional) –

    Indices of substrates (S) to control (add/inject).

    • If None: controls ALL substrates.

    • If []: controls NO substrates.

    • If [i, j]: controls specific substrates i and j.

  • target_substrate_idx (int, optional) – Convenience argument. If s_idx is None and this is set, the controller controls only this substrate index.

net

The neural network mapping time t to control signal u

forward(
model: BaseODEModel,
t: Tensor,
X: Tensor,
S: Tensor,
) → Tuple[Tensor, Tensor]

Calculates control signal at time t.

Parameters:

  • model (BaseODEModel) – The metadata wrapper for the system (used for dimensions, etc.).

  • t (torch.Tensor) – Current simulation time.

  • X (torch.Tensor) – Current biological state (Ignored by OpenLoopNeuralController, but required by API).

  • S (torch.Tensor) – Current biological state (Ignored by OpenLoopNeuralController, but required by API).

Returns:

  • dXc (torch.Tensor) – Additive change to biomass derivatives (same shape as X).

  • dSc (torch.Tensor) – Additive change to substrate derivatives (same shape as S).

get_input_history(trajectory: Tensor, model: BaseODEModel) → Tensor

Reconstructs the full control schedule u(t) for visualization.

Since this is an Open Loop controller, the schedule is fixed and independent of the actual bacterial state trajectory.

Parameters:

trajectory (torch.Tensor) – Shape (Steps, State_Dim). Used to determine the time grid resolution.

Returns:

U_tensor – Shape (Steps, Output_Dim). The control signal at each time step.

Return type:

torch.Tensor