**********************
Closed-Loop Controller
**********************

The :class:`~mgrowthctrl.controllers.closed_loop.ClosedLoopNeuralController` provides methods to learn state-dependent control inputs :math:`u(X, S)`.

Example
=======

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

.. code-block:: python

    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.

.. code-block:: python

    import numpy as np
    import torch

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

    def regulation_loss(trajectory, controller, model):
        target = torch.tensor(1e1, dtype=torch.float32)

        biomass = trajectory[:, 0]
        tracking_error = (
            torch.log(biomass + 1e-6) - torch.log(target + 1e-6)
        ).pow(2).mean()

        U_tensor = controller.net(torch.log(trajectory + 1e-6))
        effort = 1e-3 * U_tensor.pow(2).mean()

        return tracking_error + effort, U_tensor

We log-transform the input to `controller.net()` because bacterial counts and metabolite concentrations can span several orders of magnitude.

Next, we initialize the closed-loop neural controller. We assume that the controller has access to the full state (i.e., all bacterial counts and metabolite concentrations). Partial observability can be specified via the `obs_x_idx` and `obs_s_idx` arguments.

.. code-block:: python

    from torch import nn
    from mgrowthctrl.controllers import ClosedLoopNeuralController

    controller = ClosedLoopNeuralController(
        state_dim=model.n + model.m,
        n_species=model.n,
        n_resources=model.m,
        s_idx=[0],
        criterion=regulation_loss,
        hidden_dims=[5, 5, 5],
        activation=nn.ELU(),
    )

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

.. code-block:: python

    def log_fn(model, epoch, loss, U_trajectory, sol_tensor):
        if epoch % 1 == 0:
            final_biomass = sol_tensor[-1, 0].item()
            mean_input = float(U_trajectory.abs().max().item())
            print(
                f"Epoch {epoch:03d} | Loss: {loss.item():.4f} "
                f"| Final biomass: {final_biomass:.1e} "
                f"| Mean |u|: {mean_input:.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=3e-2, log_fn=log_fn)

We visually compare the uncontrolled and controlled biomass trajectories.

.. code-block:: python

    import matplotlib.pyplot as plt

    t_eval = np.linspace(0, 200, 200)
    t_span = (0, 200)
    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.hlines(1e1, 0, 200, ls="--", color="red", label="Target")
    ax.set_xlabel("Time (h)")
    ax.set_ylabel("Biomass")
    ax.legend()
    plt.tight_layout()
    plt.savefig("basic_closed_loop_result.png")
    plt.show()

We observe that the controlled dynamics reaches the desired concentration:

.. image:: ../_static/basic_closed_loop_result.png
   :alt: Closed-loop control predictions
   :align: center

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

.. code-block:: python

    raw_u = controller.get_input_history(torch.tensor(controlled_sol.y.T), model).detach()

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

Here's what the result looks like:

.. image:: ../_static/basic_closed_loop_control_inputs.png
   :alt: Closed-loop control inputs
   :align: center

For a detailed tutorial on closed-loop control of a fitted consumer-resource model, please see the tutorial :doc:`../tutorials/advanced_closed_loop_control_fitting`.

API Documentation
=================

.. automodule:: mgrowthctrl.controllers.closed_loop
   :members:
   :member-order: bysource
   :undoc-members:
   :show-inheritance: