lactationcurve.fitting

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

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

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

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

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

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

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

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

640def get_chen_priors(parity: int) -> dict:
641    """
642    Return Chen et al. priors in MilkBot format.
643
644    Args:
645        parity: Lactation number (1, 2, or >= 3).
646
647    Returns:
648        Dictionary with parameter priors:
649        - "scale": {"mean", "sd"}
650        - "ramp": {"mean", "sd"}
651        - "decay": {"mean", "sd"}
652        - "offset": {"mean", "sd"}
653        - "seMilk": Standard error of milk measurement.
654        - "milkUnit": Unit string (e.g., "kg").
655    """
656    if parity == 1:
657        return {
658            "scale": {"mean": 34.11, "sd": 7},
659            "ramp": {"mean": 29.96, "sd": 3},
660            "decay": {"mean": 0.001835, "sd": 0.000738},
661            "offset": {"mean": -0.5, "sd": 0.02},
662            "seMilk": 4,
663            "milkUnit": "kg",
664        }
665
666    if parity == 2:
667        return {
668            "scale": {"mean": 44.26, "sd": 9.57},
669            "ramp": {"mean": 22.52, "sd": 3},
670            "decay": {"mean": 0.002745, "sd": 0.000979},
671            "offset": {"mean": -0.78, "sd": 0.07},
672            "seMilk": 4,
673            "milkUnit": "kg",
674        }
675
676    # parity >= 3
677    return {
678        "scale": {"mean": 48.41, "sd": 10.66},
679        "ramp": {"mean": 22.54, "sd": 8.724},
680        "decay": {"mean": 0.002997, "sd": 0.000972},
681        "offset": {"mean": 0.0, "sd": 0.03},
682        "seMilk": 4,
683        "milkUnit": "kg",
684    }

569def get_lc_parameters(dim, milkrecordings, model="wood"):
570    """Fit lactation data to a model and return fitted parameters (frequentist).
571
572    Depending on `model`, this uses `scipy.optimize.minimize` and/or
573    `scipy.optimize.curve_fit` with model-specific starting values and bounds.
574
575    Args:
576        dim (int): List/array of DIM values.
577        milkrecordings (float): List/array of milk recordings (kg).
578        model (str): One of "wood", "wilmink", "ali_schaeffer", "fischer", "milkbot".
579
580    Returns:
581        Fitted parameters as floats, in alphabetical order by parameter name:
582            - wood: (a, b, c)
583            - wilmink: (a, b, c, k) with k fixed at -0.05
584            - ali_schaeffer: (a, b, c, d, k)
585            - fischer: (a, b, c)
586            - milkbot: (a, b, c, d)
587    """
588    # check and prepare input
589    inputs = validate_and_prepare_inputs(dim, milkrecordings, model=model)
590
591    dim = inputs.dim
592    milkrecordings = inputs.milkrecordings
593    model = inputs.model
594
595    if model == "wood":
596        wood_guess = [30, 0.2, 0.01]
597        wood_bounds = [(1, 100), (0.01, 1.5), (0.0001, 0.1)]
598        wood_res = minimize(
599            wood_objective, wood_guess, args=(dim, milkrecordings), bounds=wood_bounds
600        )
601        a_w, b_w, c_w = wood_res.x
602        return a_w, b_w, c_w
603
604    elif model == "wilmink":
605        wil_guess = [10, 0.1, 30]
606        wil_params, _ = curve_fit(wilmink_model, dim, milkrecordings, p0=wil_guess)
607        a_wil, b_wil, c_wil = wil_params
608        k_wil = -0.05  # set fixed
609        return a_wil, b_wil, c_wil, k_wil
610
611    elif model == "ali_schaeffer":
612        ali_schaeffer_guess = [10, 10, -5, 1, 1]
613        ali_schaeffer_params, _ = curve_fit(
614            ali_schaeffer_model, dim, milkrecordings, p0=ali_schaeffer_guess
615        )
616        a_as, b_as, c_as, d_as, k_as = ali_schaeffer_params
617        return a_as, b_as, c_as, d_as, k_as
618
619    elif model == "fischer":
620        fischer_guess = [max(milkrecordings), 0.01, 0.01]
621        fischer_bounds = [(0, 100), (0, 1), (0.0001, 1)]
622        fischer_params, _ = curve_fit(
623            fischer_model,
624            dim,
625            milkrecordings,
626            p0=fischer_guess,
627            bounds=np.transpose(fischer_bounds),
628        )
629        a_f, b_f, c_f = fischer_params
630        return a_f, b_f, c_f
631
632    elif model == "milkbot":
633        mb_guess = [max(milkrecordings), 20.0, -0.7, 0.022]
634        mb_bounds = [(1, 100), (1, 100), (-600, 300), (0.0001, 0.1)]
635        mb_res = minimize(milkbot_objective, mb_guess, args=(dim, milkrecordings), bounds=mb_bounds)
636        a_mb, b_mb, c_mb, d_mb = mb_res.x
637        return a_mb, b_mb, c_mb, d_mb

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

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

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

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

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

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

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

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

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