Source code for pykinbiont.models

"""Growth model types and the built-in model registry."""

from __future__ import annotations

from dataclasses import dataclass, field
from typing import Callable, Optional




class AbstractGrowthModel:
    """Base class for all growth models. Do not instantiate directly."""





@dataclass
class NLModel(AbstractGrowthModel):
    """Non-linear (closed-form) growth model.

    Parameters
    ----------
    name:
        Unique identifier used in results.
    func:
        Python callable with signature ``(p, t) -> y`` where p is a 1-D array
        of parameters and t is a scalar or 1-D array of time points.
        None for built-in registry models (Julia has the function).
    param_names:
        Human-readable name for each parameter.
    """

    name:        str
    func:        Optional[Callable] = None
    param_names: list[str] = field(default_factory=list)





@dataclass
class ODEModel(AbstractGrowthModel):
    """ODE growth model (SciML in-place form ``f!(du, u, p, t)``).

    Parameters
    ----------
    name:
        Unique identifier used in results.
    func:
        Python callable with signature ``f(du, u, p, t)`` (in-place).
        None for built-in registry models.
    param_names:
        Human-readable name for each parameter.
    n_eq:
        Number of state equations.
    """

    name:        str
    func:        Optional[Callable] = None
    param_names: list[str] = field(default_factory=list)
    n_eq:        int = 1





@dataclass
class LogLinModel(AbstractGrowthModel):
    """Sentinel for log-linear (exponential phase) fitting. No parameters."""





@dataclass
class DDDEModel(AbstractGrowthModel):
    """Data-Driven Differential Equation discovery model (sparse regression).

    Parameters
    ----------
    max_degree:
        Maximum polynomial degree in the candidate basis.
    lambda_min:
        log₁₀ of the minimum STLSQ sparsity threshold.
    lambda_max:
        log₁₀ of the maximum STLSQ sparsity threshold.
    lambda_step:
        Step in log₁₀ space between threshold values.

    Note: requires DataDrivenDiffEq + DataDrivenSparse + ModelingToolkit in
    the Julia environment. Manage these with Pkg.add in Julia.
    """

    max_degree:  int   = 4
    lambda_min:  float = -5.0
    lambda_max:  float = -1.0
    lambda_step: float = 0.5



class _LazyRegistry(dict):
    """dict subclass that populates itself from Julia on first access."""

    def __init__(self):
        super().__init__()
        self._populated: bool = False

    def _populate(self) -> None:
        if self._populated:
            return
        from pykinbiont._core import get_jl
        jl = get_jl()
        # Keys with "NL_" prefix are closed-form NLModel; all others are ODEModel.
        # This mirrors the naming convention in KinBiont.jl's MODEL_REGISTRY.
        for name, jl_model in jl.Kinbiont.MODEL_REGISTRY.items():
            name_str = str(name)
            param_names = [str(p) for p in jl_model.param_names]
            if name_str.startswith("NL_"):
                self[name_str] = NLModel(name=name_str, param_names=param_names)
            else:
                n_eq = int(jl_model.n_eq) if hasattr(jl_model, "n_eq") else 1
                self[name_str] = ODEModel(name=name_str, param_names=param_names, n_eq=n_eq)
        self._populated = True

    def __getitem__(self, key):
        self._populate()
        return super().__getitem__(key)

    def __contains__(self, key):
        self._populate()
        return super().__contains__(key)

    def __len__(self):
        self._populate()
        return super().__len__()

    def __iter__(self):
        self._populate()
        return super().__iter__()

    def get(self, key, default=None):
        self._populate()
        return super().get(key, default)

    def keys(self):
        self._populate()
        return super().keys()

    def values(self):
        self._populate()
        return super().values()

    def items(self):
        self._populate()
        return super().items()


MODEL_REGISTRY: dict[str, NLModel | ODEModel] = _LazyRegistry()
"""Built-in model registry. Keys with 'NL_' prefix are NLModel; others are ODEModel.
Populated lazily on first access (triggers Julia startup if not already running)."""