lactationcurve.fitting

158def ali_schaeffer_model(t, a, b, c, d, k) -> np.floating | np.ndarray:
159    """Ali & Schaeffer lactation curve model.
160
161    Args:
162        t: Time since calving in days (DIM). Use `t >= 1` to avoid `log(0)`.
163        a: Fitted coefficient (numerical).
164        b: Fitted coefficient (numerical).
165        c: Fitted coefficient (numerical).
166        d: Fitted coefficient (numerical).
167        k: Fitted coefficient (numerical).
168
169    Returns:
170        Predicted milk yield at `t`.
171
172    Notes:
173        Formula: `t_scaled = t / 305` and `log_term = ln(305 / t)`
174        then `y(t) = a + b * t_scaled + c * t_scaled^2 + d * log_term + k * log_term^2`.
175    """
176    t_scaled = t / 305
177    log_term = np.log(305 / t)
178    return a + b * t_scaled + c * (t_scaled**2) + d * log_term + k * (log_term**2)

810def bayesian_fit_milkbot_single_lactation(
811    dim,
812    milkrecordings,
813    key: str,
814    parity=3,
815    breed="H",
816    custom_priors: MilkBotPriors | str | None = None,
817    continent="USA",
818    milk_unit="kg",
819) -> dict:
820    """
821    Fit a single lactation using the MilkBot API.
822
823    Args:
824        dim: List/array of DIM values.
825        milkrecordings: List/array of milk recordings (kg).
826        key: API key for MilkBot.
827        parity: Lactation number; values >= 3 are treated as one group in priors.
828        breed: "H" (Holstein) or "J" (Jersey).
829        custom_priors:
830            - "CHEN"  → Chen et al. published priors
831            - dict    → Custom priors in MilkBot format (overrides `continent`)
832        continent: priors used by MilkBot API for fitting; options:
833            - "USA"   → MilkBot USA priors
834            - "EU"    → MilkBot EU priors > estimates lower milk production
835
836
837    Returns:
838        Dictionary with fitted parameters and metadata:
839            {
840                "scale": float,
841                "ramp": float,
842                "decay": float,
843                "offset": float,
844                "nPoints": int
845            }
846
847    Raises:
848        requests.HTTPError: For unsuccessful HTTP response codes.
849        RuntimeError: If the response format is unexpected.
850
851    Notes:
852        - When `continent == "CHEN"`, Chen et al. priors are included in the request payload.
853        - EU calls use the GCP EU endpoint; others use `milkbot.com`.
854    """
855    # check and prepare input
856    inputs = validate_and_prepare_inputs(
857        dim,
858        milkrecordings,
859        breed=breed,
860        parity=parity,
861        custom_priors=custom_priors,
862        continent=continent,
863        milk_unit=milk_unit,
864    )
865
866    dim = inputs.dim
867    milkrecordings = inputs.milkrecordings
868    breed = inputs.breed
869    parity = inputs.parity
870    continent = inputs.continent
871    custom_priors = inputs.custom_priors
872    milk_unit = inputs.milk_unit
873
874    # -----------------------------
875    # Select server (USA vs EU)
876    # -----------------------------
877    if continent == "EU":
878        base_url = "https://europe-west1-numeric-analogy-337601.cloudfunctions.net/milkBot-fitter"
879    else:
880        base_url = "https://milkbot.com"
881
882    # -----------------------------
883    # Prepare headers
884    # -----------------------------
885    headers = {"Content-Type": "application/json", "X-API-KEY": key}
886
887    # -----------------------------
888    # Prepare milk points
889    # -----------------------------
890    points = sorted(
891        ({"dim": int(d), "milk": float(m)} for d, m in zip(dim, milkrecordings)),
892        key=lambda p: p["dim"],
893    )
894
895    # -----------------------------
896    # Lactation metadata
897    # -----------------------------
898    payload = {
899        "lactation": {
900            "lacKey": "single_lactation_fit",
901            "breed": breed,
902            "parity": parity,
903            "points": points,
904        },
905        "options": {
906            "returnInputData": False,
907            "returnPath": False,
908            "returnDiscriminatorPath": False,
909            # "fitEngine": "AnnealingFitter@2.0", #comment out to use the default fitter
910            # "fitObjective": "MB2@2.0",
911            "preferredMilkUnit": milk_unit,
912        },
913    }
914
915    # -----------------------------
916    # Add priors if provided or when using Chen et al. priors
917    # -----------------------------
918    if custom_priors == "CHEN":
919        assert parity is not None, "parity is required for Chen priors"
920        payload["priors"] = get_chen_priors(parity)
921
922    elif isinstance(custom_priors, dict):
923        payload["priors"] = dict(custom_priors)
924
925    # -----------------------------
926    # Call API
927    # -----------------------------
928    response = requests.post(f"{base_url}/fitLactation", headers=headers, json=payload)
929    response.raise_for_status()
930    res = response.json()
931
932    # -----------------------------
933    # Normalize response (USA vs EU)
934    # -----------------------------
935    if "fittedParams" in res:
936        # USA-style response
937        fitted = res["fittedParams"]
938    elif "params" in res:
939        fitted = res["params"]
940    else:
941        raise RuntimeError(f"Unexpected MilkBot response format: {res}")
942
943    # -----------------------------
944    # Parse result
945    # -----------------------------
946
947    return {
948        "scale": fitted["scale"],
949        "ramp": fitted["ramp"],
950        "decay": fitted["decay"],
951        "offset": fitted["offset"],
952        "nPoints": len(points),
953    }

199def brody_model(t, a, k) -> float:
200    """First Brody lactation curve model.
201
202    Args:
203        t: Time since calving in days (DIM).
204        a: Scale parameter (numerical).
205        k: Decay parameter (numerical).
206
207    Returns:
208        Predicted milk yield at `t`.
209
210    Notes:
211        Formula: `y(t) = a * exp(-k * t)`.
212
213        This was the first lactation curve model ever developed in 1923.
214    """
215    return a * np.exp(-k * t)

254def dhanoa_model(t, a, b, c) -> float:
255    """Dhanoa lactation curve model.
256
257    Args:
258        t: Time since calving in days (DIM).
259        a: Scale parameter (numerical).
260        b: Shape parameter (numerical).
261        c: Decay parameter (numerical).
262
263    Returns:
264        Predicted milk yield at `t`.
265
266    Notes:
267        Formula: `y(t) = a * t ** (b * c) * exp(-c * t)`.
268    """
269    return a * t ** (b * c) * np.exp(-c * t)

329def dijkstra_model(t, a, b, c, d) -> float:
330    """Dijkstra lactation curve model.
331
332    Args:
333        t: Time since calving in days (DIM).
334        a: Scale parameter (numerical).
335        b: Growth parameter (numerical).
336        c: Saturation rate parameter (numerical).
337        d: Decay parameter (numerical).
338
339    Returns:
340        Predicted milk yield at `t`.
341
342    Notes:
343        Formula: `y(t) = a * exp((b * (1 - exp(-c * t)) / c) - d * t)`.
344    """
345    return a * np.exp((b * (1 - np.exp(-c * t)) / c) - d * t)

272def emmans_model(t, a, b, c, d) -> float:
273    """Emmans lactation curve model.
274
275    Args:
276        t: Time since calving in days (DIM).
277        a: Scale parameter (numerical).
278        b: Growth parameter (numerical).
279        c: Decay parameter (numerical).
280        d: Location parameter in nested exponential (numerical).
281
282    Returns:
283        Predicted milk yield at `t`.
284
285    Notes:
286        Formula: `y(t) = a * exp(-exp(d - b*t)) * exp(-c*t)`.
287    """
288    return a * np.exp(-np.exp(d - b * t)) * np.exp(-c * t)

181def fischer_model(t, a, b, c) -> np.floating | np.ndarray:
182    """Fischer lactation curve model.
183
184    Args:
185        t: Time since calving in days (DIM).
186        a: Fitted coefficient (numerical).
187        b: Fitted coefficient (numerical).
188        c: Fitted coefficient (numerical).
189
190    Returns:
191        Predicted milk yield at `t`.
192
193    Notes:
194        Formula: `y(t) = a - b * t - a * exp(-c * t)`.
195    """
196    return a - b * t - a * np.exp(-c * t)

410def fit_lactation_curve(
411    dim,
412    milkrecordings,
413    model="wood",
414    fitting="frequentist",
415    breed="H",
416    parity=3,
417    continent="USA",
418    custom_priors=None,
419    key=None,
420    milk_unit="kg",
421    lactation_length=None,
422) -> np.ndarray:
423    """Fit lactation data to a lactation curve model and return predictions.
424
425    Depending on `fitting`:
426    - **frequentist**: Fits parameters using `minimize` and/or `curve_fit`
427      for the specified `model`, then predicts over DIM 1–305 (or up to `max(dim)` if greater).
428    - **bayesian**: (MilkBot only) Calls the MilkBot Bayesian fitting API and
429      returns predictions using the fitted parameters.
430
431    Args:
432        dim (Int): List/array of days in milk (DIM).
433        milkrecordings (Float): List/array of milk recordings (kg).
434        model (Str): Model name (lowercase), default "wood".
435            Supported for frequentist: "wood", "wilmink", "ali_schaeffer", "fischer", "milkbot".
436        fitting (Str): "frequentist" (default) or "bayesian".
437            Bayesian fitting is currently implemented only for "milkbot".
438        breed (Str): "H" (Holstein, default) or "J" (Jersey). Only used for Bayesian.
439        parity (Int): Lactation number; all parities >= 3 considered one group in priors.
440            Only used for Bayesian.
441        continent (Str): priors chosen by MilkBot API based on continent averages.
442            Only used for Bayesian, options: "USA" (default) and "EU".
443        custom_priors (Dict | str | None): Custom prior
444            distributions for Bayesian fitting.
445            If a dict is provided, it must be a dictionary
446            of prior distributions for each parameter
447            in the model.
448            Set the correct dictionary using the `build_prior` helper function.
449            If the string "CHEN" is provided, the default Chen et al. priors are used.
450            Only used for Bayesian.
451        key = Str: API key for MilkBot API (required for Bayesian fitting).
452            Only used for Bayesian.
453        milk_unit (Str): Unit of milk yield measurements. Must be either "kg" or "lbs".
454            Default is "kg".
455            Only used for Bayesian.
456        lactation_length (int): Length of lactation in days used for lactation curve fitting
457            (default 305 days or up to the maximum DIM if > 305).
458
459
460    Returns:
461        List/array of predicted milk yield for DIM 1–305 (or up to the maximum DIM if > 305).
462
463    Raises:
464        Exception: If an unknown model is requested (frequentist),
465            or Bayesian is requested for a non-MilkBot model,
466            or `key` is missing when Bayesian fitting is requested.
467
468    Notes:
469        Uses `validate_and_prepare_inputs` for input checking and normalization.
470    """
471    # check and prepare input
472    inputs = validate_and_prepare_inputs(
473        dim,
474        milkrecordings,
475        model=model,
476        fitting=fitting,
477        breed=breed,
478        parity=parity,
479        continent=continent,
480        custom_priors=custom_priors,
481        milk_unit=milk_unit,
482        lactation_length=lactation_length,
483    )
484
485    dim = inputs.dim
486    milkrecordings = inputs.milkrecordings
487    model = inputs.model
488    fitting = inputs.fitting
489    breed = inputs.breed
490    parity = inputs.parity
491    continent = inputs.continent
492    custom_priors = inputs.custom_priors
493    milk_unit = inputs.milk_unit
494    lactation_length = inputs.lactation_length
495
496    if lactation_length is not None and not (
497        isinstance(lactation_length, str) and lactation_length.lower() == "max"
498    ):
499        mask = dim <= lactation_length
500        dim = dim[mask]
501        milkrecordings = milkrecordings[mask]
502
503    if fitting == "frequentist":
504        if model == "wood":
505            params = get_lc_parameters(dim, milkrecordings, model)
506            assert params is not None, "Failed to fit Wood model parameters"
507            a_w, b_w, c_w = params[0], params[1], params[2]
508            if max(dim) > 305:
509                t_range = np.arange(1, (max(dim) + 1))
510                y_w = wood_model(t_range, a_w, b_w, c_w)
511            else:
512                t_range = np.arange(1, 306)
513                y_w = wood_model(t_range, a_w, b_w, c_w)
514            return np.asarray(y_w)
515
516        elif model == "wilmink":
517            params = get_lc_parameters(dim, milkrecordings, model)
518            assert params is not None, "Failed to fit Wilmink model parameters"
519            a_wil, b_wil, c_wil, k_wil = (params[0], params[1], params[2], params[3])
520            if max(dim) > 305:
521                t_range = np.arange(1, (max(dim) + 1))
522                y_wil = wilmink_model(t_range, a_wil, b_wil, c_wil, k_wil)
523            else:
524                t_range = np.arange(1, 306)
525                y_wil = wilmink_model(t_range, a_wil, b_wil, c_wil, k_wil)
526            return np.asarray(y_wil)
527
528        elif model == "ali_schaeffer":
529            params = get_lc_parameters(dim, milkrecordings, model)
530            assert params is not None, "Failed to fit Ali & Schaeffer model parameters"
531            a_as, b_as, c_as, d_as, k_as = (params[0], params[1], params[2], params[3], params[4])
532            if max(dim) > 305:
533                t_range = np.arange(1, (max(dim) + 1))
534                y_as = ali_schaeffer_model(t_range, a_as, b_as, c_as, d_as, k_as)
535            else:
536                t_range = np.arange(1, 306)
537                y_as = ali_schaeffer_model(t_range, a_as, b_as, c_as, d_as, k_as)
538            return np.asarray(y_as)
539
540        elif model == "fischer":
541            params = get_lc_parameters(dim, milkrecordings, model)
542            assert params is not None, "Failed to fit Fischer model parameters"
543            a_f, b_f, c_f = params[0], params[1], params[2]
544            if max(dim) > 305:
545                t_range = np.arange(1, (max(dim) + 1))
546                y_f = fischer_model(t_range, a_f, b_f, c_f)
547            else:
548                t_range = np.arange(1, 306)
549                y_f = fischer_model(t_range, a_f, b_f, c_f)
550            return np.asarray(y_f)
551
552        elif model == "milkbot":
553            params = get_lc_parameters(dim, milkrecordings, model)
554            assert params is not None, "Failed to fit MilkBot model parameters"
555            a_mb, b_mb, c_mb, d_mb = (params[0], params[1], params[2], params[3])
556            if max(dim) > 305:
557                t_range = np.arange(1, (max(dim) + 1))
558            else:
559                t_range = np.arange(1, 306)
560
561            y_mb = milkbot_model(t_range, a_mb, b_mb, c_mb, d_mb)
562            return np.asarray(y_mb)
563
564        else:
565            raise Exception("Unknown model")
566    else:
567        if model == "milkbot":
568            if key is None:
569                raise Exception("Key needed to use Bayesian fitting engine milkbot")
570            else:
571                assert parity is not None, "parity is required for Bayesian fitting"
572                assert breed is not None, "breed is required for Bayesian fitting"
573                assert continent is not None, "continent is required for Bayesian fitting"
574                parameters = bayesian_fit_milkbot_single_lactation(
575                    dim,
576                    milkrecordings,
577                    key,
578                    parity,
579                    breed,
580                    custom_priors,
581                    continent,
582                    milk_unit or "kg",
583                )
584                if max(dim) > 305:
585                    t_range = np.arange(1, (max(dim) + 1))
586                    y_mb_bay = milkbot_model(
587                        t_range,
588                        parameters["scale"],
589                        parameters["ramp"],
590                        parameters["offset"],
591                        parameters["decay"],
592                    )
593                else:
594                    t_range = np.arange(1, 306)
595                    y_mb_bay = milkbot_model(
596                        t_range,
597                        parameters["scale"],
598                        parameters["ramp"],
599                        parameters["offset"],
600                        parameters["decay"],
601                    )
602                return np.asarray(y_mb_bay)
603        else:
604            raise Exception("Bayesian fitting is currently only implemented for milkbot models")

743def get_chen_priors(parity: int) -> dict:
744    """
745    Return Chen et al. priors in MilkBot format.
746
747    Args:
748        parity: Lactation number (1, 2, or >= 3).
749
750    Returns:
751        Dictionary with parameter priors:
752        - "scale": {"mean", "sd"}
753        - "ramp": {"mean", "sd"}
754        - "decay": {"mean", "sd"}
755        - "offset": {"mean", "sd"}
756        - "seMilk": Standard error of milk measurement.
757        - "milkUnit": Unit string (e.g., "kg").
758    """
759    if parity == 1:
760        return {
761            "scale": {"mean": 34.11, "sd": 7},
762            "ramp": {"mean": 29.96, "sd": 3},
763            "decay": {"mean": 0.001835, "sd": 0.000738},
764            "offset": {"mean": -0.5, "sd": 0.02},
765            "seMilk": 4,
766            "milkUnit": "kg",
767        }
768
769    if parity == 2:
770        return {
771            "scale": {"mean": 44.26, "sd": 9.57},
772            "ramp": {"mean": 22.52, "sd": 3},
773            "decay": {"mean": 0.002745, "sd": 0.000979},
774            "offset": {"mean": -0.78, "sd": 0.07},
775            "seMilk": 4,
776            "milkUnit": "kg",
777        }
778
779    # parity >= 3
780    return {
781        "scale": {"mean": 48.41, "sd": 10.66},
782        "ramp": {"mean": 22.54, "sd": 8.724},
783        "decay": {"mean": 0.002997, "sd": 0.000972},
784        "offset": {"mean": 0.0, "sd": 0.03},
785        "seMilk": 4,
786        "milkUnit": "kg",
787    }

670def get_lc_parameters(dim, milkrecordings, model="wood") -> tuple[float, ...]:
671    """Fit lactation data to a model and return fitted parameters (frequentist).
672
673    Depending on `model`, this uses `scipy.optimize.minimize` and/or
674    `scipy.optimize.curve_fit` with model-specific starting values and bounds.
675
676    Args:
677        dim (int): List/array of DIM values.
678        milkrecordings (float): List/array of milk recordings (kg).
679        model (str): One of "wood", "wilmink", "ali_schaeffer", "fischer", "milkbot".
680
681    Returns:
682        Fitted parameters as floats, in alphabetical order by parameter name:
683            - wood: (a, b, c)
684            - wilmink: (a, b, c, k) with k fixed at -0.05
685            - ali_schaeffer: (a, b, c, d, k)
686            - fischer: (a, b, c)
687            - milkbot: (a, b, c, d)
688    """
689    # check and prepare input
690    inputs = validate_and_prepare_inputs(dim, milkrecordings, model=model)
691
692    dim = inputs.dim
693    milkrecordings = inputs.milkrecordings
694    model = inputs.model
695
696    if model == "wood":
697        wood_guess = [30, 0.2, 0.01]
698        wood_bounds = [(1, 100), (0.01, 1.5), (0.0001, 0.1)]
699        wood_res = minimize(
700            wood_objective, wood_guess, args=(dim, milkrecordings), bounds=wood_bounds
701        )
702        a_w, b_w, c_w = wood_res.x
703        return a_w, b_w, c_w
704
705    elif model == "wilmink":
706        wil_guess = [10, 0.1, 30]
707        wil_params, _ = curve_fit(wilmink_model, dim, milkrecordings, p0=wil_guess)
708        a_wil, b_wil, c_wil = wil_params
709        k_wil = -0.05  # set fixed
710        return a_wil, b_wil, c_wil, k_wil
711
712    elif model == "ali_schaeffer":
713        ali_schaeffer_guess = [10, 10, -5, 1, 1]
714        ali_schaeffer_params, _ = curve_fit(
715            ali_schaeffer_model, dim, milkrecordings, p0=ali_schaeffer_guess
716        )
717        a_as, b_as, c_as, d_as, k_as = ali_schaeffer_params
718        return a_as, b_as, c_as, d_as, k_as
719
720    elif model == "fischer":
721        fischer_guess = [max(milkrecordings), 0.01, 0.01]
722        fischer_bounds = [(0, 100), (0, 1), (0.0001, 1)]
723        fischer_params, _ = curve_fit(
724            fischer_model,
725            dim,
726            milkrecordings,
727            p0=fischer_guess,
728            bounds=np.transpose(fischer_bounds),
729        )
730        a_f, b_f, c_f = fischer_params
731        return a_f, b_f, c_f
732
733    elif model == "milkbot":
734        mb_guess = [max(milkrecordings), 20.0, -0.7, 0.022]
735        mb_bounds = [(1, 100), (1, 100), (-600, 300), (0.0001, 0.1)]
736        mb_res = minimize(milkbot_objective, mb_guess, args=(dim, milkrecordings), bounds=mb_bounds)
737        a_mb, b_mb, c_mb, d_mb = mb_res.x
738        return a_mb, b_mb, c_mb, d_mb
739
740    raise ValueError(f"Unknown model: {model}")

607def get_lc_parameters_least_squares(
608    dim, milkrecordings, model="milkbot"
609) -> tuple[float, float, float, float]:
610    """Fit lactation data and return model parameters (least squares; frequentist).
611
612    This helper uses `scipy.optimize.least_squares` to fit the MilkBot model with bounds,
613    and returns the fitted parameters.
614    Currently implemented only for the MilkBot model, as it is
615    more complex and benefits from the robust optimization approach.
616    Other models can be fitted using `get_lc_parameters` with
617    numerical optimisation, which is generally faster for simpler
618    models.
619
620    Args:
621        dim (int): List/array of DIM values.
622        milkrecordings (float): List/array of milk recordings (kg).
623        model (str): Pre-defined model name; currently used with "milkbot".
624
625    Returns:
626        Parameters `(a, b, c, d)` as `np.float` in alphabetic order.
627
628    """
629    # check and prepare input
630    inputs = validate_and_prepare_inputs(dim, milkrecordings, model=model)
631
632    dim = inputs.dim
633    milkrecordings = inputs.milkrecordings
634    model = inputs.model
635
636    # ------------------------------
637    # Initial guess
638    # ------------------------------
639    a0 = np.max(milkrecordings)
640    b0 = 50.0
641    c0 = 30.0
642    d0 = 0.01
643    p0 = [a0, b0, c0, d0]
644
645    # ------------------------------
646    # Parameter bounds
647    # ------------------------------
648    lower = [np.max(milkrecordings) * 0.5, 1.0, -300.0, 1e-6]
649    upper = [np.max(milkrecordings) * 8.0, 400.0, 300.0, 1.0]
650
651    # ------------------------------
652    # Fit using least-squares
653    # ------------------------------
654    res = least_squares(
655        residuals_milkbot,
656        p0,
657        args=(dim, milkrecordings),
658        bounds=(lower, upper),
659        method="trf",  # trust region reflective, works well with bounds
660    )
661
662    # ------------------------------
663    # Extract parameters
664    # ------------------------------
665    a_mb, b_mb, c_mb, d_mb = res.x
666
667    return a_mb, b_mb, c_mb, d_mb

291def hayashi_model(t, a, b, c, d) -> float:
292    """Hayashi lactation curve model.
293
294    Args:
295        t: Time since calving in days (DIM).
296        a: Ratio parameter (> 0) (numerical).
297        b: Scale parameter (numerical).
298        c: Time constant for the first exponential term (numerical).
299        d: Parameter retained for compatibility with literature (unused in this expression).
300
301    Returns:
302        Predicted milk yield at `t`.
303
304    Notes:
305        Formula: `y(t) = b * (exp(-t / c) - exp(-t / (a * c)))`.
306    """
307    return b * (np.exp(-t / c) - np.exp(-t / (a * c)))

101def milkbot_model(t, a, b, c, d) -> np.floating | np.ndarray:
102    """MilkBot lactation curve model.
103
104    Args:
105        t: Time since calving in days (DIM), scalar or array-like.
106        a: Scale; overall level of milk production.
107        b: Ramp; governs the rate of rise in early lactation.
108        c: Offset; small (usually minor) correction around the theoretical start of lactation.
109        d: Decay; exponential decline rate, evident in late lactation.
110
111    Returns:
112        Predicted milk yield at `t` (same shape as `t`).
113
114    Notes:
115        Formula: `y(t) = a * (1 - exp((c - t) / b) / 2) * exp(-d * t)`.
116    """
117    return a * (1 - np.exp((c - t) / b) / 2) * np.exp(-d * t)

236def nelder_model(t, a, b, c) -> float:
237    """Nelder lactation curve model.
238
239    Args:
240        t: Time since calving in days (DIM).
241        a: Denominator intercept (numerical).
242        b: Denominator linear coefficient (numerical).
243        c: Denominator quadratic coefficient (numerical).
244
245    Returns:
246        Predicted milk yield at `t`.
247
248    Notes:
249        Formula: `y(t) = t / (a + b*t + c*t^2)`.
250    """
251    return t / (a + b * t + c * t**2)

348def prasad_model(t, a, b, c, d) -> float:
349    """Prasad lactation curve model.
350
351    Args:
352        t: Time since calving in days (DIM).
353        a: Intercept-like parameter (numerical).
354        b: Linear coefficient (numerical).
355        c: Quadratic coefficient (numerical).
356        d: Inverse-time coefficient (numerical).
357
358    Returns:
359        Predicted milk yield at `t`.
360
361    Notes:
362        Formula: `y(t) = a + b*t + c*t^2 + d/t`.
363    """
364    return a + b * t + c * t**2 + d / t

310def rook_model(t, a, b, c, d) -> float:
311    """Rook lactation curve model.
312
313    Args:
314        t: Time since calving in days (DIM).
315        a: Scale parameter (numerical).
316        b: Shape parameter in rational term (numerical).
317        c: Offset parameter in rational term (numerical).
318        d: Exponential decay parameter (numerical).
319
320    Returns:
321        Predicted milk yield at `t`.
322
323    Notes:
324        Formula: `y(t) = a * (1 / (1 + b / (c + t))) * exp(-d * t)`.
325    """
326    return a * (1 / (1 + b / (c + t))) * np.exp(-d * t)

218def sikka_model(t, a, b, c) -> float:
219    """Sikka lactation curve model.
220
221    Args:
222        t: Time since calving in days (DIM).
223        a: Scale parameter (numerical).
224        b: Growth parameter (numerical).
225        c: Quadratic decay parameter (numerical).
226
227    Returns:
228        Predicted milk yield at `t`.
229
230    Notes:
231        Formula: `y(t) = a * exp(b * t - c * t^2)`.
232    """
233    return a * np.exp(b * t - c * t**2)

138def wilmink_model(t, a, b, c, k=-0.05) -> np.floating | np.ndarray:
139    """Wilmink lactation curve model.
140
141    Args:
142        t: Time since calving in days (DIM), scalar or array-like.
143        a: Fitted coefficient (numerical).
144        b: Fitted coefficient (numerical).
145        c: Fitted coefficient (numerical).
146        k: Fixed exponential coefficient (numerical), default -0.05.
147
148    Returns:
149        Predicted milk yield at `t`.
150
151    Notes:
152        Formula: `y(t) = a + b * t + c * exp(k * t)`.
153    """
154    t = np.asarray(t)
155    return a + b * t + c * np.exp(k * t)

120def wood_model(t, a, b, c) -> np.floating | np.ndarray:
121    """Wood lactation curve model.
122
123    Args:
124        t: Time since calving in days (DIM), scalar or array-like.
125        a: Fitted coefficient (numerical).
126        b: Fitted coefficient (numerical).
127        c: Fitted coefficient (numerical).
128
129    Returns:
130        Predicted milk yield at `t`.
131
132    Notes:
133        Formula: `y(t) = a * t^b * exp(-c * t)`.
134    """
135    return a * (t**b) * np.exp(-c * t)

790def build_prior(
791    scale_mean: float,
792    scale_sd: float,
793    ramp_mean: float,
794    ramp_sd: float,
795    decay_mean: float,
796    decay_sd: float,
797    offset_mean: float,
798    offset_sd: float,
799    se_milk: float = 4,
800) -> dict:
801    return {
802        "scale": {"mean": scale_mean, "sd": scale_sd},
803        "ramp": {"mean": ramp_mean, "sd": ramp_sd},
804        "decay": {"mean": decay_mean, "sd": decay_sd},
805        "offset": {"mean": offset_mean, "sd": offset_sd},
806        "seMilk": se_milk,
807    }