Source code for pykinbiont.fit

"""Public fit() and preprocess() functions."""

from __future__ import annotations

from typing import Optional, Union
import numpy as np

from pykinbiont._core import get_jl
from pykinbiont._convert import (
    py_to_jl_growth_data,
    py_to_jl_irregular,
    py_to_jl_fit_options,
    py_to_jl_model_spec,
    jl_to_py_growth_fit_results,
)
from pykinbiont.types import (
    GrowthData,
    IrregularGrowthData,
    FitOptions,
    ModelSpec,
    GrowthFitResults,
)

_AnyData = Union[GrowthData, IrregularGrowthData]


def _to_jl_data(data: _AnyData):
    if isinstance(data, IrregularGrowthData):
        return py_to_jl_irregular(data)
    return py_to_jl_growth_data(data)




def fit(
    data: _AnyData,
    spec: ModelSpec,
    opts: Optional["FitOptions"] = None,
) -> GrowthFitResults:
    """Fit every curve in data to every model in spec and select best by AICc.

    Preprocessing (smoothing, blank subtraction, clustering, …) is applied
    according to opts before fitting.

    Parameters
    ----------
    data:
        GrowthData or IrregularGrowthData container.
    spec:
        Models to try and their initial parameters.
    opts:
        All preprocessing and fitting configuration. Defaults to ``FitOptions()``
        (no preprocessing, log-linear model selection).

    Returns
    -------
    GrowthFitResults
        One CurveFitResult per curve plus the post-preprocessing GrowthData
        (with cluster assignments if clustering was requested).
    """
    if opts is None:
        opts = FitOptions()
    jl = get_jl()
    jl_data = _to_jl_data(data)
    jl_opts = py_to_jl_fit_options(opts)
    jl_spec = py_to_jl_model_spec(spec)
    jl_res  = jl.Kinbiont.kinbiont_fit(jl_data, jl_spec, jl_opts)
    return jl_to_py_growth_fit_results(jl_res)





def preprocess(
    data: _AnyData,
    opts: FitOptions,
) -> Union[GrowthData, IrregularGrowthData]:
    """Apply the preprocessing pipeline to data and return a new GrowthData.

    When opts.cluster=True, the returned GrowthData has .clusters, .centroids,
    and .wcss populated.

    Parameters
    ----------
    data:
        GrowthData or IrregularGrowthData.
    opts:
        Preprocessing configuration.

    Returns
    -------
    GrowthData
        Preprocessed data (new instance, input is never modified).
    """
    jl = get_jl()
    jl_data = _to_jl_data(data)
    jl_opts = py_to_jl_fit_options(opts)
    jl_processed = jl.Kinbiont.preprocess(jl_data, jl_opts)

    clusters = (
        np.array(jl_processed.clusters, dtype=int)
        if jl_processed.clusters is not None else None
    )
    centroids = (
        np.array(jl_processed.centroids, dtype=np.float64)
        if jl_processed.centroids is not None else None
    )
    wcss = float(jl_processed.wcss) if jl_processed.wcss is not None else None

    return GrowthData(
        curves    = np.array(jl_processed.curves, dtype=np.float64),
        times     = np.array(jl_processed.times,  dtype=np.float64),
        labels    = [str(lb) for lb in jl_processed.labels],
        clusters  = clusters,
        centroids = centroids,
        wcss      = wcss,
    )