lactationcurve.fitting

133def ali_schaeffer_model(t, a, b, c, d, k):
134    """Ali & Schaeffer lactation curve model.
135
136    Args:
137        t: Time since calving in days (DIM). Use `t >= 1` to avoid `log(0)`.
138        a: Intercept-like parameter (numerical).
139        b: Linear coefficient on scaled time `t/305` (numerical).
140        c: Quadratic coefficient on scaled time `t/305` (numerical).
141        d: Coefficient on `log(305/t)` (numerical).
142        k: Coefficient on `[log(305/t)]^2` (numerical).
143
144    Returns:
145        Predicted milk yield at `t`.
146
147    Notes:
148        Uses `t_scaled = t / 305` and `log_term = ln(305 / t)`.
149    """
150    t_scaled = t / 305
151    log_term = np.log(305 / t)
152    return a + b * t_scaled + c * (t_scaled**2) + d * log_term + k * (log_term**2)

711def bayesian_fit_milkbot_single_lactation(
712    dim, milkrecordings, key: str, parity=3, breed="H", continent="USA"
713) -> dict:
714    """
715    Fit a single lactation using the MilkBot API.
716
717    Args:
718        dim: List/array of DIM values.
719        milkrecordings: List/array of milk recordings (kg).
720        key: API key for MilkBot.
721        parity: Lactation number; values >= 3 are treated as one group in priors.
722        breed: "H" (Holstein) or "J" (Jersey).
723        continent: Prior source:
724            - "USA"   → MilkBot USA priors
725            - "EU"    → MilkBot EU priors
726            - "CHEN"  → Chen et al. published priors
727
728    Returns:
729        Dictionary with fitted parameters and metadata:
730            {
731                "scale": float,
732                "ramp": float,
733                "decay": float,
734                "offset": float,
735                "nPoints": int
736            }
737
738    Raises:
739        requests.HTTPError: For unsuccessful HTTP response codes.
740        RuntimeError: If the response format is unexpected.
741
742    Notes:
743        - When `continent == "CHEN"`, Chen et al. priors are included in the request payload.
744        - EU calls use the GCP EU endpoint; others use `milkbot.com`.
745    """
746    # check and prepare input
747    inputs = validate_and_prepare_inputs(
748        dim, milkrecordings, breed=breed, parity=parity, continent=continent
749    )
750
751    dim = inputs.dim
752    milkrecordings = inputs.milkrecordings
753    breed = inputs.breed
754    parity = inputs.parity
755    continent = inputs.continent
756
757    # -----------------------------
758    # Select server (USA vs EU)
759    # -----------------------------
760    if continent == "EU":
761        base_url = "https://europe-west1-numeric-analogy-337601.cloudfunctions.net/milkBot-fitter"
762    else:
763        base_url = "https://milkbot.com"
764
765    # -----------------------------
766    # Prepare headers
767    # -----------------------------
768    headers = {"Content-Type": "application/json", "X-API-KEY": key}
769
770    # -----------------------------
771    # Prepare milk points
772    # -----------------------------
773    points = sorted(
774        ({"dim": int(d), "milk": float(m)} for d, m in zip(dim, milkrecordings)),
775        key=lambda p: p["dim"],
776    )
777
778    # -----------------------------
779    # Lactation metadata
780    # -----------------------------
781    payload = {
782        "lactation": {
783            "lacKey": "single_lactation_fit",
784            "breed": breed,
785            "parity": parity,
786            "points": points,
787        },
788        "options": {
789            "returnInputData": False,
790            "returnPath": False,
791            "returnDiscriminatorPath": False,
792            # "fitEngine": "AnnealingFitter@2.0", #comment out to use the default fitter
793            # "fitObjective": "MB2@2.0",
794            "preferredMilkUnit": "kg",
795        },
796    }
797
798    # -----------------------------
799    # Add priors only if Chen
800    # -----------------------------
801    if continent == "CHEN":
802        assert parity is not None, "parity is required for Chen priors"
803        payload["priors"] = get_chen_priors(parity)
804
805    # -----------------------------
806    # Call API
807    # -----------------------------
808    response = requests.post(f"{base_url}/fitLactation", headers=headers, json=payload)
809    response.raise_for_status()
810    res = response.json()
811
812    # -----------------------------
813    # Normalize response (USA vs EU)
814    # -----------------------------
815    if "fittedParams" in res:
816        # USA-style response
817        fitted = res["fittedParams"]
818    elif "params" in res:
819        fitted = res["params"]
820    else:
821        raise RuntimeError(f"Unexpected MilkBot response format: {res}")
822
823    # -----------------------------
824    # Parse result
825    # -----------------------------
826
827    return {
828        "scale": fitted["scale"],
829        "ramp": fitted["ramp"],
830        "decay": fitted["decay"],
831        "offset": fitted["offset"],
832        "nPoints": len(points),
833    }

170def brody_model(t, a, k):
171    """Brody lactation curve model.
172
173    Args:
174        t: Time since calving in days (DIM).
175        a: Scale parameter (numerical).
176        k: Decay parameter (numerical).
177
178    Returns:
179        Predicted milk yield at `t`.
180    """
181    return a * np.exp(-k * t)

217def dhanoa_model(t, a, b, c):
218    """Dhanoa lactation curve model.
219
220    Args:
221        t: Time since calving in days (DIM).
222        a: Scale parameter (numerical).
223        b: Shape parameter (numerical).
224        c: Decay parameter (numerical).
225
226    Returns:
227        Predicted milk yield at `t`.
228
229    Notes:
230        Formula: `y(t) = a * t ** (b * c) * exp(-c * t)`.
231    """
232    return a * t ** (b * c) * np.exp(-c * t)

292def dijkstra_model(t, a, b, c, d):
293    """Dijkstra lactation curve model.
294
295    Args:
296        t: Time since calving in days (DIM).
297        a: Scale parameter (numerical).
298        b: Growth parameter (numerical).
299        c: Saturation rate parameter (numerical).
300        d: Decay parameter (numerical).
301
302    Returns:
303        Predicted milk yield at `t`.
304
305    Notes:
306        Formula: `y(t) = a * exp((b * (1 - exp(-c * t)) / c) - d * t)`.
307    """
308    return a * np.exp((b * (1 - np.exp(-c * t)) / c) - d * t)

235def emmans_model(t, a, b, c, d):
236    """Emmans lactation curve model.
237
238    Args:
239        t: Time since calving in days (DIM).
240        a: Scale parameter (numerical).
241        b: Growth parameter (numerical).
242        c: Decay parameter (numerical).
243        d: Location parameter in nested exponential (numerical).
244
245    Returns:
246        Predicted milk yield at `t`.
247
248    Notes:
249        Formula: `y(t) = a * exp(-exp(d - b*t)) * exp(-c*t)`.
250    """
251    return a * np.exp(-np.exp(d - b * t)) * np.exp(-c * t)

155def fischer_model(t, a, b, c):
156    """Fischer lactation curve model.
157
158    Args:
159        t: Time since calving in days (DIM).
160        a: Scale parameter (numerical).
161        b: Linear decline parameter (numerical).
162        c: Exponential decay parameter (numerical).
163
164    Returns:
165        Predicted milk yield at `t`.
166    """
167    return a - b * t - a * np.exp(-c * t)

373def fit_lactation_curve(
374    dim,
375    milkrecordings,
376    model="wood",
377    fitting="frequentist",
378    breed="H",
379    parity=3,
380    continent="USA",
381    key=None,
382):
383    """Fit lactation data to a lactation curve model and return predictions.
384
385    Depending on `fitting`:
386    - **frequentist**: Fits parameters using `minimize` and/or `curve_fit`
387      for the specified `model`, then predicts over DIM 1–305 (or up to `max(dim)` if greater).
388    - **bayesian**: (MilkBot only) Calls the MilkBot Bayesian fitting API and
389      returns predictions using the fitted parameters.
390
391    Args:
392        dim (Int): List/array of days in milk (DIM).
393        milkrecordings (Float): List/array of milk recordings (kg).
394        model (Str): Model name (lowercase), default "wood".
395            Supported for frequentist: "wood", "wilmink", "ali_schaeffer", "fischer", "milkbot".
396        fitting (Str): "frequentist" (default) or "bayesian".
397            Bayesian fitting is currently implemented only for "milkbot".
398        breed (Str): "H" (Holstein, default) or "J" (Jersey). Only used for Bayesian.
399        parity (Int): Lactation number; all parities >= 3 considered one group in priors (Bayesian).
400        continent (Str): Prior source for Bayesian, "USA" (default), "EU", or "CHEN".
401        key (Str | None): API key for MilkBot (required when `fitting == "bayesian"`).
402
403    Returns:
404        List/array of predicted milk yield for DIM 1–305 (or up to the maximum DIM if > 305).
405
406    Raises:
407        Exception: If an unknown model is requested (frequentist),
408            or Bayesian is requested for a non-MilkBot model,
409            or `key` is missing when Bayesian fitting is requested.
410
411    Notes:
412        Uses `validate_and_prepare_inputs` for input checking and normalization.
413    """
414    # check and prepare input
415    inputs = validate_and_prepare_inputs(
416        dim,
417        milkrecordings,
418        model=model,
419        fitting=fitting,
420        breed=breed,
421        parity=parity,
422        continent=continent,
423    )
424
425    dim = inputs.dim
426    milkrecordings = inputs.milkrecordings
427    model = inputs.model
428    fitting = inputs.fitting
429    breed = inputs.breed
430    parity = inputs.parity
431    continent = inputs.continent
432
433    if fitting == "frequentist":
434        if model == "wood":
435            params = get_lc_parameters(dim, milkrecordings, model)
436            assert params is not None, "Failed to fit Wood model parameters"
437            a_w, b_w, c_w = params[0], params[1], params[2]
438            if max(dim) > 305:
439                t_range = np.arange(1, (max(dim) + 1))
440                y_w = wood_model(t_range, a_w, b_w, c_w)
441            else:
442                t_range = np.arange(1, 306)
443                y_w = wood_model(t_range, a_w, b_w, c_w)
444            return y_w
445
446        elif model == "wilmink":
447            params = get_lc_parameters(dim, milkrecordings, model)
448            assert params is not None, "Failed to fit Wilmink model parameters"
449            a_wil, b_wil, c_wil, k_wil = (params[0], params[1], params[2], params[3])
450            if max(dim) > 305:
451                t_range = np.arange(1, (max(dim) + 1))
452                y_wil = wilmink_model(t_range, a_wil, b_wil, c_wil, k_wil)
453            else:
454                t_range = np.arange(1, 306)
455                y_wil = wilmink_model(t_range, a_wil, b_wil, c_wil, k_wil)
456            return y_wil
457
458        elif model == "ali_schaeffer":
459            params = get_lc_parameters(dim, milkrecordings, model)
460            assert params is not None, "Failed to fit Ali & Schaeffer model parameters"
461            a_as, b_as, c_as, d_as, k_as = (params[0], params[1], params[2], params[3], params[4])
462            if max(dim) > 305:
463                t_range = np.arange(1, (max(dim) + 1))
464                y_as = ali_schaeffer_model(t_range, a_as, b_as, c_as, d_as, k_as)
465            else:
466                t_range = np.arange(1, 306)
467                y_as = ali_schaeffer_model(t_range, a_as, b_as, c_as, d_as, k_as)
468            return y_as
469
470        elif model == "fischer":
471            params = get_lc_parameters(dim, milkrecordings, model)
472            assert params is not None, "Failed to fit Fischer model parameters"
473            a_f, b_f, c_f = params[0], params[1], params[2]
474            if max(dim) > 305:
475                t_range = np.arange(1, (max(dim) + 1))
476                y_f = fischer_model(t_range, a_f, b_f, c_f)
477            else:
478                t_range = np.arange(1, 306)
479                y_f = fischer_model(t_range, a_f, b_f, c_f)
480            return y_f
481
482        elif model == "milkbot":
483            params = get_lc_parameters(dim, milkrecordings, model)
484            assert params is not None, "Failed to fit MilkBot model parameters"
485            a_mb, b_mb, c_mb, d_mb = (params[0], params[1], params[2], params[3])
486            if max(dim) > 305:
487                t_range = np.arange(1, (max(dim) + 1))
488            else:
489                t_range = np.arange(1, 306)
490
491            y_mb = milkbot_model(t_range, a_mb, b_mb, c_mb, d_mb)
492            return y_mb
493
494        else:
495            raise Exception("Unknown model")
496    else:
497        if model == "milkbot":
498            if key is None:
499                raise Exception("Key needed to use Bayesian fitting engine milkbot")
500            else:
501                assert parity is not None, "parity is required for Bayesian fitting"
502                assert breed is not None, "breed is required for Bayesian fitting"
503                assert continent is not None, "continent is required for Bayesian fitting"
504                parameters = bayesian_fit_milkbot_single_lactation(
505                    dim, milkrecordings, key, parity, breed, continent
506                )
507                if max(dim) > 305:
508                    t_range = np.arange(1, (max(dim) + 1))
509                    y_mb_bay = milkbot_model(
510                        t_range,
511                        parameters["scale"],
512                        parameters["ramp"],
513                        parameters["offset"],
514                        parameters["decay"],
515                    )
516                else:
517                    t_range = np.arange(1, 306)
518                    y_mb_bay = milkbot_model(
519                        t_range,
520                        parameters["scale"],
521                        parameters["ramp"],
522                        parameters["offset"],
523                        parameters["decay"],
524                    )
525                return y_mb_bay
526        else:
527            raise Exception("Bayesian fitting is currently only implemented for milkbot models")

664def get_chen_priors(parity: int) -> dict:
665    """
666    Return Chen et al. priors in MilkBot format.
667
668    Args:
669        parity: Lactation number (1, 2, or >= 3).
670
671    Returns:
672        Dictionary with parameter priors:
673        - "scale": {"mean", "sd"}
674        - "ramp": {"mean", "sd"}
675        - "decay": {"mean", "sd"}
676        - "offset": {"mean", "sd"}
677        - "seMilk": Standard error of milk measurement.
678        - "milkUnit": Unit string (e.g., "kg").
679    """
680    if parity == 1:
681        return {
682            "scale": {"mean": 34.11, "sd": 7},
683            "ramp": {"mean": 29.96, "sd": 3},
684            "decay": {"mean": 0.001835, "sd": 0.000738},
685            "offset": {"mean": -0.5, "sd": 0.02},
686            "seMilk": 4,
687            "milkUnit": "kg",
688        }
689
690    if parity == 2:
691        return {
692            "scale": {"mean": 44.26, "sd": 9.57},
693            "ramp": {"mean": 22.52, "sd": 3},
694            "decay": {"mean": 0.002745, "sd": 0.000979},
695            "offset": {"mean": -0.78, "sd": 0.07},
696            "seMilk": 4,
697            "milkUnit": "kg",
698        }
699
700    # parity >= 3
701    return {
702        "scale": {"mean": 48.41, "sd": 10.66},
703        "ramp": {"mean": 22.54, "sd": 8.724},
704        "decay": {"mean": 0.002997, "sd": 0.000972},
705        "offset": {"mean": 0.0, "sd": 0.03},
706        "seMilk": 4,
707        "milkUnit": "kg",
708    }

591def get_lc_parameters(dim, milkrecordings, model="wood") -> tuple[float, ...]:
592    """Fit lactation data to a model and return fitted parameters (frequentist).
593
594    Depending on `model`, this uses `scipy.optimize.minimize` and/or
595    `scipy.optimize.curve_fit` with model-specific starting values and bounds.
596
597    Args:
598        dim (int): List/array of DIM values.
599        milkrecordings (float): List/array of milk recordings (kg).
600        model (str): One of "wood", "wilmink", "ali_schaeffer", "fischer", "milkbot".
601
602    Returns:
603        Fitted parameters as floats, in alphabetical order by parameter name:
604            - wood: (a, b, c)
605            - wilmink: (a, b, c, k) with k fixed at -0.05
606            - ali_schaeffer: (a, b, c, d, k)
607            - fischer: (a, b, c)
608            - milkbot: (a, b, c, d)
609    """
610    # check and prepare input
611    inputs = validate_and_prepare_inputs(dim, milkrecordings, model=model)
612
613    dim = inputs.dim
614    milkrecordings = inputs.milkrecordings
615    model = inputs.model
616
617    if model == "wood":
618        wood_guess = [30, 0.2, 0.01]
619        wood_bounds = [(1, 100), (0.01, 1.5), (0.0001, 0.1)]
620        wood_res = minimize(
621            wood_objective, wood_guess, args=(dim, milkrecordings), bounds=wood_bounds
622        )
623        a_w, b_w, c_w = wood_res.x
624        return a_w, b_w, c_w
625
626    elif model == "wilmink":
627        wil_guess = [10, 0.1, 30]
628        wil_params, _ = curve_fit(wilmink_model, dim, milkrecordings, p0=wil_guess)
629        a_wil, b_wil, c_wil = wil_params
630        k_wil = -0.05  # set fixed
631        return a_wil, b_wil, c_wil, k_wil
632
633    elif model == "ali_schaeffer":
634        ali_schaeffer_guess = [10, 10, -5, 1, 1]
635        ali_schaeffer_params, _ = curve_fit(
636            ali_schaeffer_model, dim, milkrecordings, p0=ali_schaeffer_guess
637        )
638        a_as, b_as, c_as, d_as, k_as = ali_schaeffer_params
639        return a_as, b_as, c_as, d_as, k_as
640
641    elif model == "fischer":
642        fischer_guess = [max(milkrecordings), 0.01, 0.01]
643        fischer_bounds = [(0, 100), (0, 1), (0.0001, 1)]
644        fischer_params, _ = curve_fit(
645            fischer_model,
646            dim,
647            milkrecordings,
648            p0=fischer_guess,
649            bounds=np.transpose(fischer_bounds),
650        )
651        a_f, b_f, c_f = fischer_params
652        return a_f, b_f, c_f
653
654    elif model == "milkbot":
655        mb_guess = [max(milkrecordings), 20.0, -0.7, 0.022]
656        mb_bounds = [(1, 100), (1, 100), (-600, 300), (0.0001, 0.1)]
657        mb_res = minimize(milkbot_objective, mb_guess, args=(dim, milkrecordings), bounds=mb_bounds)
658        a_mb, b_mb, c_mb, d_mb = mb_res.x
659        return a_mb, b_mb, c_mb, d_mb
660
661    raise ValueError(f"Unknown model: {model}")

530def get_lc_parameters_least_squares(dim, milkrecordings, model="milkbot"):
531    """Fit lactation data and return model parameters (least squares; frequentist).
532
533    This helper uses `scipy.optimize.least_squares` to fit the MilkBot model with bounds,
534    and returns the fitted parameters.
535    Currently implemented only for the MilkBot model, as it is
536    more complex and benefits from the robust optimization approach.
537    Other models can be fitted using `get_lc_parameters` with
538    numerical optimisation, which is generally faster for simpler
539    models.
540
541    Args:
542        dim (int): List/array of DIM values.
543        milkrecordings (float): List/array of milk recordings (kg).
544        model (str): Pre-defined model name; currently used with "milkbot".
545
546    Returns:
547        Parameters `(a, b, c, d)` as `np.float` in alphabetic order.
548
549    """
550    # check and prepare input
551    inputs = validate_and_prepare_inputs(dim, milkrecordings, model=model)
552
553    dim = inputs.dim
554    milkrecordings = inputs.milkrecordings
555    model = inputs.model
556
557    # ------------------------------
558    # Initial guess
559    # ------------------------------
560    a0 = np.max(milkrecordings)
561    b0 = 50.0
562    c0 = 30.0
563    d0 = 0.01
564    p0 = [a0, b0, c0, d0]
565
566    # ------------------------------
567    # Parameter bounds
568    # ------------------------------
569    lower = [np.max(milkrecordings) * 0.5, 1.0, -300.0, 1e-6]
570    upper = [np.max(milkrecordings) * 8.0, 400.0, 300.0, 1.0]
571
572    # ------------------------------
573    # Fit using least-squares
574    # ------------------------------
575    res = least_squares(
576        residuals_milkbot,
577        p0,
578        args=(dim, milkrecordings),
579        bounds=(lower, upper),
580        method="trf",  # trust region reflective, works well with bounds
581    )
582
583    # ------------------------------
584    # Extract parameters
585    # ------------------------------
586    a_mb, b_mb, c_mb, d_mb = res.x
587
588    return a_mb, b_mb, c_mb, d_mb

254def hayashi_model(t, a, b, c, d):
255    """Hayashi lactation curve model.
256
257    Args:
258        t: Time since calving in days (DIM).
259        a: Ratio parameter (> 0) (numerical).
260        b: Scale parameter (numerical).
261        c: Time constant for the first exponential term (numerical).
262        d: Parameter retained for compatibility with literature (unused in this expression).
263
264    Returns:
265        Predicted milk yield at `t`.
266
267    Notes:
268        Formula: `y(t) = b * (exp(-t / c) - exp(-t / (a * c)))`.
269    """
270    return b * (np.exp(-t / c) - np.exp(-t / (a * c)))

76def milkbot_model(t, a, b, c, d):
77    """MilkBot lactation curve model.
78
79    Args:
80        t: Time since calving in days (DIM), scalar or array-like.
81        a: Scale; overall level of milk production.
82        b: Ramp; governs the rate of rise in early lactation.
83        c: Offset; small (usually minor) correction around the theoretical start of lactation.
84        d: Decay; exponential decline rate, evident in late lactation.
85
86    Returns:
87        Predicted milk yield at `t` (same shape as `t`).
88
89    Notes:
90        Formula: `y(t) = a * (1 - exp((c - t) / b) / 2) * exp(-d * t)`.
91    """
92    return a * (1 - np.exp((c - t) / b) / 2) * np.exp(-d * t)

199def nelder_model(t, a, b, c):
200    """Nelder lactation curve model.
201
202    Args:
203        t: Time since calving in days (DIM).
204        a: Denominator intercept (numerical).
205        b: Denominator linear coefficient (numerical).
206        c: Denominator quadratic coefficient (numerical).
207
208    Returns:
209        Predicted milk yield at `t`.
210
211    Notes:
212        Formula: `y(t) = t / (a + b*t + c*t**2)`.
213    """
214    return t / (a + b * t + c * t**2)

311def prasad_model(t, a, b, c, d):
312    """Prasad lactation curve model.
313
314    Args:
315        t: Time since calving in days (DIM).
316        a: Intercept-like parameter (numerical).
317        b: Linear coefficient (numerical).
318        c: Quadratic coefficient (numerical).
319        d: Inverse-time coefficient (numerical).
320
321    Returns:
322        Predicted milk yield at `t`.
323
324    Notes:
325        Formula: `y(t) = a + b*t + c*t**2 + d/t`.
326    """
327    return a + b * t + c * t**2 + d / t

273def rook_model(t, a, b, c, d):
274    """Rook lactation curve model.
275
276    Args:
277        t: Time since calving in days (DIM).
278        a: Scale parameter (numerical).
279        b: Shape parameter in rational term (numerical).
280        c: Offset parameter in rational term (numerical).
281        d: Exponential decay parameter (numerical).
282
283    Returns:
284        Predicted milk yield at `t`.
285
286    Notes:
287        Formula: `y(t) = a * (1 / (1 + b / (c + t))) * exp(-d * t)`.
288    """
289    return a * (1 / (1 + b / (c + t))) * np.exp(-d * t)

184def sikka_model(t, a, b, c):
185    """Sikka lactation curve model.
186
187    Args:
188        t: Time since calving in days (DIM).
189        a: Scale parameter (numerical).
190        b: Growth parameter (numerical).
191        c: Quadratic decay parameter (numerical).
192
193    Returns:
194        Predicted milk yield at `t`.
195    """
196    return a * np.exp(b * t - c * t**2)

113def wilmink_model(t, a, b, c, k=-0.05):
114    """Wilmink lactation curve model.
115
116    Args:
117        t: Time since calving in days (DIM), scalar or array-like.
118        a: Intercept-like parameter (numerical).
119        b: Linear trend coefficient (numerical).
120        c: Exponential-term scale (numerical).
121        k: Fixed exponential rate (numerical), default -0.05.
122
123    Returns:
124        Predicted milk yield at `t`.
125
126    Notes:
127        Formula: `y(t) = a + b * t + c * exp(k * t)`.
128    """
129    t = np.asarray(t)
130    return a + b * t + c * np.exp(k * t)

 95def wood_model(t, a, b, c):
 96    """Wood lactation curve model.
 97
 98    Args:
 99        t: Time since calving in days (DIM), scalar or array-like.
100        a: Scale parameter (numerical).
101        b: Shape parameter controlling rise (numerical).
102        c: Decay parameter (numerical).
103
104    Returns:
105        Predicted milk yield at `t`.
106
107    Notes:
108        Formula: `y(t) = a * t**b * exp(-c * t)`.
109    """
110    return a * (t**b) * np.exp(-c * t)