lactationcurve.characteristics

220def ISLC(
221    df: pd.DataFrame,
222    days_in_milk_col: str | None = None,
223    milking_yield_col: str | None = None,
224    test_id_col: str | None = None,
225    default_test_id: int = 0,
226    standard_lc_305: pd.Series | npt.NDArray[np.float64] = STANDARD_CURVE,
227    correlation_matrix: pd.DataFrame = CORR_MATRIX,
228    std_per_grid_day: npt.NDArray[np.float64] = STDs,
229    max_dim: int | str = 305,
230    fit_standard_lc_from_data: bool = False,
231    reference_df: pd.DataFrame | None = None,
232) -> pd.DataFrame:
233    """Estimate cumulative lactation yield per TestId using ISLC.
234
235    This variant applies :func:`ISLC_method` to standardized lactation records.
236    By default it uses the package-provided standard curve, correlation matrix,
237    and per-grid day standard deviations. Optionally, these assets can be refit
238    from a reference dataset by providing a pandas dataframe at 'reference_df ='
239    when ``fit_standard_lc_from_data`` is True.
240    Alternative for customization is to set standard_lc_305,
241    correlation_matrix, and std_per_grid_day directly.
242
243    Args:
244        df: Input test-day records.
245        days_in_milk_col: Name of DIM column in ``df``.
246        milking_yield_col: Name of milk-yield column in ``df``.
247        test_id_col: Name of lactation identifier column in ``df``.
248        default_test_id: Fallback TestId value when ``test_id_col`` is absent.
249        standard_lc_305: Standard lactation curve values used by ISLC.
250        correlation_matrix: Correlation matrix aligned with ISLC grid.
251        std_per_grid_day: Standard deviations aligned with ISLC grid.
252        max_dim: Maximum DIM to include
253            (or ``"max"`` so highest test day DIM is considered the final lactation day).
254        fit_standard_lc_from_data: If True, fit representation from ``reference_df``.
255        reference_df: Reference records used when fitting curve and covariance.
256
257    Returns:
258        DataFrame with columns ``TestId`` and ``LactationMilkYield``.
259
260    Raises:
261        ValueError: If required columns are missing or fit mode has no reference data.
262    """
263    # Standardize columns and filter DIM <= 305
264    df = standardize_lactation_columns(
265        df,
266        days_in_milk_col=days_in_milk_col,
267        milking_yield_col=milking_yield_col,
268        test_id_col=test_id_col,
269        default_test_id=default_test_id,
270        max_dim=max_dim,
271    )
272
273    # Fit covariance/standard_lc if requested
274    if fit_standard_lc_from_data is True:
275        if reference_df is None:
276            raise ValueError("Provide reference_df to fit your own standard lactation curve.")
277        reference_df = standardize_lactation_columns(
278            reference_df,
279            days_in_milk_col=days_in_milk_col,
280            milking_yield_col=milking_yield_col,
281            test_id_col=test_id_col,
282            default_test_id=default_test_id,
283            max_dim=305,
284        )
285        dim = reference_df["DaysInMilk"]
286        milk_yield = reference_df["MilkingYield"]
287        standard_full_lc_305 = fit_lactation_curve(dim, milkrecordings=milk_yield, model="wilmink")
288        # convert to pandas series
289        standard_full_lc_305 = pd.Series(standard_full_lc_305)
290
291        # fit matrix
292        (
293            correlation_matrix,
294            std_per_grid_day,
295            standard_curve_fitted,
296        ) = create_standard_lc_representation(
297            df=reference_df,
298            standard_lactation_curve=pd.Series(standard_full_lc_305),
299            # reference_df is standardized above, so fixed names are required here
300            days_in_milk_col="DaysInMilk",
301            milking_yield_col="MilkingYield",
302            interpolation_method=interpolation_standard_lc,
303            small_grid=False,
304        )
305
306        # Align fitted curve with default internal convention (day 0 equals day 1).
307        standard_lc_305 = np.insert(
308            np.asarray(standard_curve_fitted, dtype=float),
309            0,
310            float(standard_curve_fitted.iloc[0]),
311        )
312
313    rows: list[dict[str, Any]] = []
314    for test_id, group in df.groupby("TestId", sort=False):
315        if group.empty:
316            rows.append({"TestId": test_id, "LactationMilkYield": np.nan})
317            continue
318
319        try:
320            cumulative_yield = ISLC_method(
321                df=group,
322                standard_lc=standard_lc_305,
323                correlation_matrix=correlation_matrix,
324                std_per_grid_day=std_per_grid_day,
325                days_in_milk_col="DaysInMilk",
326                milking_yield_col="MilkingYield",
327            )
328        except ValueError:
329            cumulative_yield = np.nan
330
331        rows.append({"TestId": test_id, "LactationMilkYield": cumulative_yield})
332
333    return pd.DataFrame(rows, columns=pd.Index(["TestId", "LactationMilkYield"]))

336def ISLC_method(
337    df: pd.DataFrame,
338    standard_lc: pd.Series | npt.NDArray[np.float64],
339    correlation_matrix: pd.DataFrame,
340    std_per_grid_day: npt.NDArray[np.float64],
341    days_in_milk_col: str = "DaysInMilk",
342    milking_yield_col: str = "MilkingYield",
343) -> float:
344    """Estimate total lactation yield by predicting missing grid days, then integrating.
345
346    The method fills missing grid-day yields with the ICAR-style relation:
347
348        yp = gp + b* (yi - gi)
349
350        where for each missing day p, the nearest measured day i is used and
351        b* = corr(i,p) * std(p) / std(i). The resulting complete grid profile is
352        integrated with the test-interval style weighting used elsewhere in this
353        package.
354
355    Notes:
356        The integration is performed over the days present after merging
357        measured and predicted points. If the input contains DIM beyond 305,
358        total yield can reflect that extended horizon.
359    """
360    # add zero and 305 to the grid
361    grid = [0] + GRID_DAYS + [305]
362    dim_col = days_in_milk_col
363    yield_col = milking_yield_col
364
365    standard_lc_series = _curve_to_series(standard_lc)
366
367    # Extract the measured days
368    measured_days = df[dim_col].to_numpy()
369    if measured_days.size == 0:
370        raise ValueError("Input df must contain at least one measured day.")
371
372    # interpolate the measured days to the grid using the standard lactation curve
373    interpolated_df = interpolation_standard_lc(
374        group=df,
375        days_in_milk_col=dim_col,
376        milking_yield_col=yield_col,
377        standard_lc=standard_lc_series,
378    )
379
380    # drop unnecessary columns
381    interpolated_df = interpolated_df.drop(columns=[dim_col, yield_col])
382
383    # rename grid colum
384    interpolated_df = interpolated_df.rename(
385        columns={"GridDay": dim_col, "MilkYieldInterp": yield_col}
386    )[[dim_col, yield_col]]
387
388    measured_subset = cast(pd.DataFrame, df.loc[:, [dim_col, yield_col]])
389    interpolated_subset = cast(pd.DataFrame, interpolated_df.loc[:, [dim_col, yield_col]])
390
391    # join the interpolated values with the original measurements, keeping all measurements
392    merged_df = cast(
393        pd.DataFrame,
394        pd.merge(
395            measured_subset,
396            interpolated_subset,
397            on=dim_col,
398            how="outer",
399            suffixes=("_meas", "_interp"),
400        ),
401    )
402    meas_series = cast(pd.Series, merged_df[f"{yield_col}_meas"])
403    interp_series = cast(pd.Series, merged_df[f"{yield_col}_interp"])
404    merged_df[yield_col] = cast(pd.Series, meas_series.combine_first(interp_series))
405    df = cast(
406        pd.DataFrame,
407        merged_df.loc[:, [dim_col, yield_col]].sort_values(by=dim_col),
408    )
409
410    # days missing from measurement → need prediction
411    days_to_predict = [day for day in grid if day not in df[dim_col].values]
412
413    predicted_rows: list[dict[str, float | int]] = []
414
415    for day in days_to_predict:
416        # identify closest measured day
417        dim_series = cast(pd.Series, interpolated_df[dim_col])
418        grid_vals = np.asarray(dim_series.to_numpy(dtype=float), dtype=float)
419        closest_idx = int(np.argmin(np.abs(grid_vals - float(day))))
420        closest_day = int(dim_series.iloc[closest_idx])
421
422        # extract mean milk yield from the standard lactation curve for the day to predict
423        expected_yield_grid_day = _curve_value(standard_lc_series, day)
424        expected_yield_measured_day = _curve_value(standard_lc_series, closest_day)
425
426        # calculate difference between measured and expected yield at closest measured day
427        diff = (
428            interpolated_df.loc[interpolated_df[dim_col] == closest_day, yield_col].iloc[0]
429            - expected_yield_measured_day
430        )
431
432        # calculate the correlation-based correction factor b*
433        # b* = corr(i,p) * std(p) / std(i)
434        pi = list(grid).index(closest_day)
435        qi = list(grid).index(day)
436        bpj = correlation_matrix.iloc[pi, qi]
437        std_q = std_per_grid_day[qi]
438        std_p = std_per_grid_day[pi]
439        b_star = 0.0 if std_p == 0 else float(bpj * std_q / std_p)
440
441        # predict the yield for the missing day
442        predicted_yield = expected_yield_grid_day + b_star * diff
443        predicted_rows.append({dim_col: int(day), yield_col: float(predicted_yield)})
444
445    if predicted_rows:
446        df = pd.concat([df, pd.DataFrame(predicted_rows)], ignore_index=True)
447
448    # apply the test interval method to calculate total milk yield over 305 days
449    # sort the dataframe by days in milk and keep one value per day
450    df = df.drop_duplicates(subset=[dim_col], keep="last")
451    df = df.sort_values(by=dim_col)
452    # Trapezoidal contributions
453    df["width"] = df[dim_col].diff().shift(-1)
454    df["avg_yield"] = (df[yield_col] + df[yield_col].shift(-1)) / 2
455    df["trapezoid_area"] = df["width"] * df["avg_yield"]
456
457    trapezoid_values = np.asarray(pd.to_numeric(df["trapezoid_area"], errors="coerce"), dtype=float)
458    total_yield = float(np.nansum(trapezoid_values))
459
460    return total_yield

576def ISLC_original(
577    df: pd.DataFrame,
578    days_in_milk_col: str | None = None,
579    milking_yield_col: str | None = None,
580    test_id_col: str | None = None,
581    default_test_id: int = 0,
582    standard_lc_305: pd.Series | npt.NDArray[np.float64] = STANDARD_CURVE,
583    correlation_matrix: pd.DataFrame = CORR_MATRIX,
584    std_per_grid_day: npt.NDArray[np.float64] = STDs,
585    fit_standard_lc_from_data: bool = False,
586    reference_df: pd.DataFrame | None = None,
587) -> pd.DataFrame:
588    """Compute estimated 305-day yields using the original grid-based ISLC variant
589    from the original paper by Dr. Wilmink in 1987.
590
591    The function groups standardized input by ``TestId``, interpolates to the
592    fixed GRID_DAYS grid, and applies :func:`ISLC_original_method` per lactation.
593
594    Args:
595       df: Input records with DIM and milk yield columns.
596       days_in_milk_col: Name of DIM column in ``df``.
597       milking_yield_col: Name of milk-yield column in ``df``.
598       test_id_col: Name of lactation identifier column in ``df``.
599       default_test_id: Fallback TestId when ``test_id_col`` is absent.
600       standard_lc_305: Standard lactation curve used for interpolation.
601       correlation_matrix: Correlation matrix over GRID_DAYS.
602       std_per_grid_day: Standard deviations per GRID_DAYS element.
603       fit_standard_lc_from_data: If True, refit standard curve and covariance
604           representation from ``reference_df`` before prediction.
605       reference_df: Reference population used when
606           ``fit_standard_lc_from_data=True``. The reference data is
607           standardized to ``DaysInMilk``/``MilkingYield`` and fit using
608           ``small_grid=True`` so correlation and standard-deviation assets are
609           aligned to ``GRID_DAYS``.
610
611    Returns:
612        DataFrame with columns ``LactationMilkYield`` and ``TestId``.
613
614    Raises:
615        ValueError: If required columns (DaysInMilk or MilkingYield) cannot be
616            found, or if ``fit_standard_lc_from_data=True`` and
617            ``reference_df`` is not provided.
618
619    Notes:
620        Records with DIM > 305 are dropped before interpolation and scoring.
621    """
622    # Standardize columns and filter DIM <= 305
623    df = standardize_lactation_columns(
624        df,
625        days_in_milk_col=days_in_milk_col,
626        milking_yield_col=milking_yield_col,
627        test_id_col=test_id_col,
628        default_test_id=default_test_id,
629        max_dim=305,
630    )
631
632    # Fit covariance/standard_lc if requested
633    if fit_standard_lc_from_data is True:
634        if reference_df is None:
635            raise ValueError("Provide reference_df to fit your own standard lactation curve.")
636        reference_df = standardize_lactation_columns(
637            reference_df,
638            days_in_milk_col=days_in_milk_col,
639            milking_yield_col=milking_yield_col,
640            test_id_col=test_id_col,
641            default_test_id=default_test_id,
642            max_dim=305,
643        )
644        dim = reference_df["DaysInMilk"]
645        milk_yield = reference_df["MilkingYield"]
646        standard_full_lc_305 = fit_lactation_curve(dim, milkrecordings=milk_yield, model="wilmink")
647        standard_full_lc_305 = pd.Series(standard_full_lc_305)
648
649        # fit matrix
650        (
651            correlation_matrix,
652            std_per_grid_day,
653            standard_curve_fitted,
654        ) = create_standard_lc_representation(
655            df=reference_df,
656            standard_lactation_curve=pd.Series(standard_full_lc_305),
657            # reference_df is standardized above, so fixed names are required here
658            days_in_milk_col="DaysInMilk",
659            milking_yield_col="MilkingYield",
660            interpolation_method=interpolation_standard_lc,
661            small_grid=True,
662        )
663
664        # Align fitted curve with default internal convention (day 0 equals day 1).
665        standard_lc_305 = np.insert(
666            np.asarray(standard_curve_fitted, dtype=float),
667            0,
668            float(standard_curve_fitted.iloc[0]),
669        )
670    rows: list[dict[str, Any]] = []
671    for test_id, group in df.groupby("TestId", sort=False):
672        lactation = interpolation_standard_lc(
673            group,
674            days_in_milk_col="DaysInMilk",
675            milking_yield_col="MilkingYield",
676            standard_lc=standard_lc_305,
677        )
678        if lactation.empty:
679            rows.append({"TestId": test_id, "305_milk_yield": np.nan})
680            continue
681
682        cumulative_yield = ISLC_original_method(
683            df=lactation,
684            standard_lc=standard_lc_305,
685            correlation_matrix=correlation_matrix,
686            std_per_grid_day=std_per_grid_day,
687        )
688        rows.append({"TestId": test_id, "LactationMilkYield": cumulative_yield})
689
690    return pd.DataFrame(rows, columns=pd.Index(["TestId", "LactationMilkYield"]))

463def ISLC_original_method(
464    df: pd.DataFrame,
465    standard_lc: pd.Series | npt.NDArray[np.float64],
466    correlation_matrix: pd.DataFrame,
467    std_per_grid_day: npt.NDArray[np.float64],
468) -> float:
469    """Estimate 305-day milk yield for a single interpolated lactation
470    from the original paper by Dr. Wilmink in 1987.
471
472    This function computes an estimated 305-day yield by combining three
473    components: known production from interpolated measurements, the
474    population mean for missing grid days (from ``standard_lc``), and a
475    correlation-based deviation correction derived from ``correlation_matrix``.
476
477    Args:
478        df: DataFrame containing at least the columns ``GridDay`` and
479            ``MilkYieldInterp`` (values already interpolated onto the grid).
480        standard_lc: Standard lactation curve values indexed by grid day.
481        correlation_matrix: Correlation coefficients between grid-day yields
482            (matrix-like, rows/columns ordered as the grid).
483        std_per_grid_day: 1-D array of standard deviations for each grid day.
484
485    Returns:
486        A float with the estimated 305-day milk yield.
487
488    Raises:
489        ValueError: if any measured day in ``df`` is not present in the
490            expected global ``grid``.
491
492    Notes:
493        - The routine assumes fixed recording intervals of 20 days, except
494          that the final interval (if day 290 is present) is 25 days.
495        - The implementation expects a global ``grid`` variable to be
496          defined and aligned with ``standard_lc`` and
497          ``std_per_grid_day``.
498        - The amount of estimated milk yields depends on the
499            availability of measured milk yields.
500        - Lactation yield is:
501            sum(known measurements (p) + estimated measurements (q)).
502        - In the implementation according to the CRV E2 document,
503            there should also be a correction factor for the difference between
504             the expected and actual complete lactation yield of the previous lactation.
505            However, this correction factor is not included in the current implementation
506            as it requires additional data (previous lactation yield) that is not always available.
507    """
508    grid = GRID_DAYS
509    standard_lc_series = _curve_to_series(standard_lc)
510
511    # --- Validate that all measurements lie on the expected grid -------
512    measured_days = df["GridDay"].values
513    if any(day not in grid for day in measured_days):
514        raise ValueError("At least one DIM in the input df is not part of the grid.")
515
516    # --- Prepare basic components -------------------------------------
517    df = df.sort_values("GridDay")
518
519    # days missing from measurement → need prediction
520    days_to_predict = [day for day in grid if day not in measured_days]
521    pred_idx = [list(grid).index(d) for d in days_to_predict]
522
523    # mean expected milk yield from standard curve
524    mean_lc = pd.Series([_curve_value(standard_lc_series, g) for g in grid], index=grid)
525
526    # --- Compute known production (sum of actual measurements) ---------
527    if 290 in measured_days:
528        # special last interval: 25 days
529        known_prod = df["MilkYieldInterp"][:-1].sum() * 20 + df["MilkYieldInterp"].iloc[-1] * 25
530    else:
531        known_prod = df["MilkYieldInterp"].sum() * 20
532
533    # --- Compute population mean for missing days ----------------------
534    pop_mean_missing = mean_lc.loc[days_to_predict]
535
536    if 290 in pop_mean_missing.index:
537        population_mean = pop_mean_missing.iloc[:-1].sum() * 20 + pop_mean_missing.iloc[-1] * 25
538    else:
539        population_mean = pop_mean_missing.sum() * 20
540
541    # --- Compute correlation-based correction --------------------------
542    days_to_predict_np = np.array(days_to_predict)
543    measured_days_np = np.array(measured_days)
544
545    # for each missing day: find closest measured day
546    closest_measured = [
547        measured_days_np[np.argmin(np.abs(measured_days_np - g))] for g in days_to_predict_np
548    ]
549    closest_idx = [list(grid).index(d) for d in closest_measured]
550
551    b_star_list = []
552    for pi, qi in zip(closest_idx, pred_idx):
553        bpj = correlation_matrix.iloc[pi, qi]
554        stdj = std_per_grid_day[qi]
555        stdp = std_per_grid_day[pi]
556        b_star_list.append(bpj * stdj / stdp)
557
558    # scale correction by interval lengths
559    if 290 in measured_days:
560        # last interval belongs to prediction group
561        correction_weighted = sum(b_star_list[:-1]) * 20
562    else:
563        correction_weighted = sum(b_star_list[:-1]) * 20 + b_star_list[-1] * 25
564
565    last_day = df["GridDay"].iloc[-1]
566    last_yield = df["MilkYieldInterp"].iloc[-1]
567    mean_last = mean_lc.loc[last_day]
568
569    correction = correction_weighted * (last_yield - mean_last)
570
571    # --- Final sum -----------------------------------------------------
572    final_milk_yield = known_prod + population_mean + correction
573    return float(final_milk_yield)

494def best_predict_method(
495    df: pd.DataFrame,
496    standard_lc: np.ndarray = STANDARD_CURVE,
497    days_in_milk_col: str | None = None,
498    milking_yield_col: str | None = None,
499    test_id_col: str | None = None,
500    default_test_id: int = 0,
501    covariance_matrix: np.ndarray | None = COV_MATRIX,
502    fit_standard_lc_from_data: bool = False,
503    reference_df: pd.DataFrame | None = None,
504) -> pd.DataFrame:
505    """Apply best prediction to one or more lactations.
506
507    By default this function uses the package-provided standard curve and covariance matrix.
508    But it is also possible to provide your own standard curve and covariance matrix,
509    for example when you want to fit these ingredients from your own reference population.
510    This can be done in two ways: either by fitting the covariance matrix and standard curve
511    directly from a reference dataset by providing a pandas dataframe at 'reference_df ='
512    when ``fit_standard_lc_from_data`` is True.
513    Alternative for customization is to set standard_lc_305 and covariance_matrix
514    directly in the function call.
515
516    Args:
517        df: Input observations. If ``TestId`` is missing, all rows are treated
518            as one lactation.
519        standard_lc: Expected daily milk yield lactation curve on days 1..305.
520            If not provided, the package's default curve is used.
521            Or fit your own standard curve from a reference dataset by
522            providing a pandas dataframe at 'reference_df ='
523            when ``fit_standard_lc_from_data`` is True.
524        days_in_milk_col: Optional input column name for days in milk. If
525            provided, it is mapped to ``DaysInMilk``.
526        milking_yield_col: Optional input column name for milk yield. If
527            provided, it is mapped to ``MilkingYield``.
528        test_id_col: Optional input column name for lactation/test identifier.
529            If provided, it is mapped to ``TestId``.
530        default_test_id: Fallback test id used when no test-id column is
531            available.
532        covariance_matrix: Optional prefit covariance matrix. If omitted,
533            the default matrix is used or
534             ``reference_df`` can be used to fit one for your own data.
535        fit_standard_lc_from_data: Whether to fit covariance information from
536            ``reference_df`` instead of using a provided covariance matrix.
537        reference_df: Reference dataframe used when ``covariance_matrix`` and
538            ``standard_lc`` are not provided and ``fit_standard_lc_from_data``
539            is True.
540
541    Returns:
542        Dataframe with columns ``TestId`` and ``LactationMilkYield``.
543
544    Raises:
545        ValueError: If neither ``covariance_matrix`` nor ``reference_df`` is
546            provided.
547    """
548    # Standardize columns and filter DIM <= 305
549    df = standardize_lactation_columns(
550        df,
551        days_in_milk_col=days_in_milk_col,
552        milking_yield_col=milking_yield_col,
553        test_id_col=test_id_col,
554        default_test_id=default_test_id,
555        max_dim=305,
556    )
557
558    # Fit covariance if not provided
559    if fit_standard_lc_from_data:
560        if reference_df is None:
561            raise ValueError("Provide reference_df to fit your own standard lactation curve.")
562        reference_df = standardize_lactation_columns(
563            reference_df,
564            days_in_milk_col=days_in_milk_col,
565            milking_yield_col=milking_yield_col,
566            test_id_col=test_id_col,
567            default_test_id=default_test_id,
568            max_dim=305,
569        )
570        covariance_matrix = cast(
571            np.ndarray, fit_autocorrelation_matrix(reference_df, standard_lc)["B_hat"]
572        )
573
574    covariance_matrix_array = cast(np.ndarray, covariance_matrix)
575
576    df = df.copy()
577
578    results = []
579
580    for test_id, lactation in df.groupby("TestId"):
581        pred = best_predict_method_single_lac(
582            lactation,
583            standard_lc,
584            covariance_matrix_array,
585        )
586        results.append({"TestId": test_id, "LactationMilkYield": pred})
587
588    return pd.DataFrame(results)

424def best_predict_method_single_lac(
425    lactation: pd.DataFrame,
426    standard_lc: np.ndarray = STANDARD_CURVE,
427    covariance_matrix: np.ndarray = COV_MATRIX,
428) -> float:
429    """Predict 305-day cumulative yield for one lactation.
430
431    Observed test-day deviations are projected over all 305 days using the
432    covariance structure and then added to the baseline cumulative standard
433    curve.
434
435    By default this function uses the package-provided standard curve and covariance matrix.
436    But it is also possible to provide your own standard curve and covariance matrix,
437    for example when you want to fit these ingredients from your own reference population.
438
439
440    Args:
441        lactation: Observed records for one lactation.
442        standard_lc: Population mean daily yield profile.
443        covariance_matrix: Day-to-day covariance matrix on the 305-day grid.
444
445    Returns:
446        Predicted cumulative 305-day milk yield.
447
448    Notes:
449        Duplicate day records are resolved with ``keep="last"`` before
450        prediction. If no valid observations remain in days 1..305, the method
451        returns the cumulative standard curve.
452    """
453    filtered_lactation = lactation.loc[
454        (lactation["DaysInMilk"] >= 1) & (lactation["DaysInMilk"] <= 305)
455    ].copy()
456    filtered_lactation = filtered_lactation.drop_duplicates(subset=["DaysInMilk"], keep="last")
457    filtered_lactation = filtered_lactation.sort_values("DaysInMilk")
458
459    corrected_series = preprocess_measured_data(
460        filtered_lactation,
461        standard_lc=standard_lc,
462    )
463
464    if filtered_lactation.empty:
465        return float(np.sum(standard_lc))
466
467    obs_idx_1based = filtered_lactation["DaysInMilk"].to_numpy(dtype=int)  # DaysInMilk: 1-305
468    obs_idx_0based = obs_idx_1based - 1  # Convert to 0-based matrix indices: 0-304
469    y_obs = corrected_series.loc[obs_idx_1based].to_numpy(
470        dtype=float
471    )  # corrected_series is indexed by DaysInMilk (1-305)
472
473    # Extract covariance blocks
474    B_oo = covariance_matrix[
475        np.ix_(obs_idx_0based, obs_idx_0based)
476    ]  # Use 0-based indices for matrix
477    B_mo = covariance_matrix[:, obs_idx_0based]  # Use 0-based indices for matrix
478
479    # solve
480    c, lower = cho_factor(B_oo)
481    alpha = cho_solve((c, lower), y_obs)
482
483    # Predict full deviation curve
484    y_estimate = B_mo @ alpha
485
486    # Total milk = baseline + deviation
487    deviation = np.sum(y_estimate)
488
489    total = np.sum(standard_lc) + deviation
490
491    return total

271def build_covariance_matrix(rho: float, size: int) -> np.ndarray:
272    """Construct a covariance matrix.
273
274    Cole et al. (2007) estimated correlations among test-day yields using a
275    simplified model with an identity matrix (I) for daily measurement error
276    and an autoregressive matrix (E) for biological change. E is defined as
277    ``Eij = r ** |i-j|`` where ``i`` and ``j`` are test-day DIM and
278    ``0 < r < 1``.
279
280    Element ``(i, j)`` is ``rho ** abs(i - j)``.
281
282    Args:
283        rho: AR(1) correlation parameter.
284        size: Matrix dimension.
285
286    Returns:
287        A ``(size, size)`` AR(1) correlation matrix.
288    """
289    idx = np.arange(size)
290    M = np.abs(idx[:, None] - idx[None, :])
291    return rho**M

323def calculate_characteristic(
324    dim,
325    milkrecordings,
326    model="wood",
327    characteristic="cumulative_milk_yield",
328    fitting="frequentist",
329    key=None,
330    parity=3,
331    breed="H",
332    continent="USA",
333    custom_priors=None,
334    milk_unit="kg",
335    persistency_method="derived",
336    lactation_length=305,
337) -> float | None:
338    """Evaluate a lactation curve characteristic from observed test-day data.
339
340    This function fits the requested model (frequentist or Bayesian via MilkBot),
341    retrieves model parameters, and evaluates the requested characteristic using the
342    symbolic expression (if available), falling back to numeric methods when needed.
343
344    Args:
345        dim (Int): Days in milk (DIM).
346        milkrecordings (Float): Milk recordings (kg or lbs) for each DIM.
347        model (str): Model name. Supported for this function:
348            'milkbot', 'wood', 'wilmink', 'ali_schaeffer', 'fischer'.
349        characteristic (str): One of 'time_to_peak', 'peak_yield',
350            'cumulative_milk_yield', 'persistency'.
351        fitting (str): 'frequentist' (default) or 'bayesian'.
352        key (str | None): API key for MilkBot Bayesian fitting.
353        parity (Int): Parity of the cow; values above 3 are considered as 3 (Bayesian).
354        breed (str): 'H' (Holstein) or 'J' (Jersey) (Bayesian).
355        custom_priors: provide your own priors for Bayesian fitting.
356            provide as dictionary using build_prior() function
357            to create the priors in the right format.
358            Alternative use priors from the literature provided by the string command 'CHEN'
359        milk_unit: Unit of milk recordings ('kg' or 'lb') for Bayesian fitting.
360        continent (str): 'USA' or 'EU' (Bayesian).
361        persistency_method (str): 'derived' (average slope after peak; default) or 'literature'
362            (only for Wood and MilkBot).
363        lactation_length (Int | str): Horizon for persistency calculation:
364            305 (default), 'max' (use max DIM in data), or an integer.
365
366    Returns:
367        float | None: The requested characteristic value.
368        Or None if value is negative (except persistency)
369        or cannot be computed symbolically.
370
371    Raises:
372        Exception: If inputs are invalid, the model/characteristic is unsupported,
373            an API key is missing for Bayesian fitting, or the characteristic cannot be computed.
374    """
375    # check and prepare input
376    inputs = validate_and_prepare_inputs(
377        dim,
378        milkrecordings,
379        model=model,
380        fitting=fitting,
381        breed=breed,
382        parity=parity,
383        continent=continent,
384        custom_priors=custom_priors,
385        milk_unit=milk_unit,
386        persistency_method=persistency_method,
387        lactation_length=lactation_length,
388    )
389
390    dim = inputs.dim
391    milkrecordings = inputs.milkrecordings
392    model = inputs.model
393    fitting = inputs.fitting
394    breed = inputs.breed
395    parity = inputs.parity
396    continent = inputs.continent
397    custom_priors = inputs.custom_priors
398    milk_unit = inputs.milk_unit
399    persistency_method = inputs.persistency_method
400    lactation_length = inputs.lactation_length
401
402    # After validation, these fields are guaranteed non-None for the paths below
403    assert model is not None
404    assert fitting is not None
405
406    if model not in ["milkbot", "wood", "wilmink", "ali_schaeffer", "fischer"]:
407        msg = (
408            "this function currently only works for"
409            " the milkbot, wood, wilmink,"
410            " ali_schaeffer and fischer models"
411        )
412        raise Exception(msg)
413
414    characteristic_options: list[str] = [
415        "time_to_peak",
416        "peak_yield",
417        "cumulative_milk_yield",
418        "persistency",
419    ]
420
421    if characteristic in characteristic_options:
422        if fitting == "frequentist":
423            # Get fitted parameters from your fitting function
424            fitted_params = get_lc_parameters(dim, milkrecordings, model)
425            assert fitted_params is not None
426
427            if characteristic != "persistency":
428                # Try symbolic formula first
429                assert isinstance(lactation_length, int)
430                try:
431                    expr, params, fn = lactation_curve_characteristic_function(
432                        model, characteristic, lactation_length
433                    )
434                    with np.errstate(
435                        divide="ignore", invalid="ignore"
436                    ):  # get rid of warnings for invalid operations
437                        value = fn(*fitted_params)
438                except Exception:
439                    value = None
440
441                # If symbolic formula fails or is invalid (use numeric approach)
442                if (
443                    value is None
444                    or not np.isfinite(value)
445                    or (characteristic == "time_to_peak" and value <= 0)
446                ):
447                    assert parity is not None
448                    assert breed is not None
449                    assert continent is not None
450                    if characteristic == "time_to_peak":
451                        value = numeric_time_to_peak(
452                            dim,
453                            milkrecordings,
454                            model,
455                            fitting=fitting,
456                            key=key,
457                            parity=parity,
458                            breed=breed,
459                            continent=continent,
460                        )
461                    elif characteristic == "peak_yield":
462                        value = numeric_peak_yield(
463                            dim,
464                            milkrecordings,
465                            model,
466                            fitting=fitting,
467                            key=key,
468                            parity=parity,
469                            breed=breed,
470                            continent=continent,
471                        )
472                    elif characteristic == "cumulative_milk_yield":
473                        value = numeric_cumulative_yield(
474                            dim,
475                            milkrecordings,
476                            model,
477                            fitting=fitting,
478                            lactation_length=lactation_length,
479                            key=key,
480                            parity=parity,
481                            breed=breed,
482                            continent=continent,
483                        )
484
485            else:
486                if persistency_method == "derived":
487                    # find lactation length from data
488                    if lactation_length == "max":
489                        lactation_length = max(dim)
490                    elif isinstance(lactation_length, int):
491                        lactation_length = lactation_length
492                    else:
493                        lactation_length = 305
494                    value = persistency_fitted_curve(
495                        dim,
496                        milkrecordings,
497                        model,
498                        fitting="frequentist",
499                        lactation_length=lactation_length,
500                    )
501
502                else:
503                    if model == "wood":
504                        value = persistency_wood(fitted_params[1], fitted_params[2])
505
506                    elif model == "milkbot":
507                        assert len(fitted_params) >= 4
508                        value = persistency_milkbot(fitted_params[3])
509
510                    else:
511                        msg = (
512                            "Currently only the Wood model"
513                            " and MilkBot model have a"
514                            " separate model function from"
515                            " the literature integrated for"
516                            " persistency. if"
517                            ' persistency="derived" is'
518                            " selected, persistency can be"
519                            " calculated for every model as"
520                            " the average slope of the"
521                            " lactation after the peak."
522                        )
523                        raise Exception(msg)
524            try:
525                if value is None:
526                    return None
527                if characteristic != "persistency" and value < 0:
528                    return None
529                return float(value)
530            except ValueError:
531                raise Exception("Could not compute characteristic")
532
533        else:
534            if model == "milkbot":
535                if key is None:
536                    raise Exception("Key needed to use Bayesian fitting engine MilkBot")
537                else:
538                    assert parity is not None
539                    assert breed is not None
540                    assert continent is not None
541                    assert isinstance(milk_unit, str)
542                    fitted_params_bayes = bayesian_fit_milkbot_single_lactation(
543                        dim,
544                        milkrecordings,
545                        key=key,
546                        parity=parity,
547                        breed=breed,
548                        custom_priors=custom_priors,
549                        continent=continent,
550                        milk_unit=milk_unit,
551                    )
552                    fitted_params_bayes = (
553                        fitted_params_bayes["scale"],
554                        fitted_params_bayes["ramp"],
555                        fitted_params_bayes["offset"],
556                        fitted_params_bayes["decay"],
557                    )
558
559                    if characteristic != "persistency":
560                        # Get the symbolic expression and model parameters
561                        try:
562                            expr, param_symbols, fn = lactation_curve_characteristic_function(
563                                model, characteristic
564                            )
565                            with np.errstate(divide="ignore", invalid="ignore"):
566                                value = fn(*fitted_params_bayes)
567                        except Exception:
568                            value = None
569
570                        if (
571                            value is None
572                            or not np.isfinite(value)
573                            or (characteristic == "time_to_peak" and value <= 0)
574                        ):
575                            if characteristic == "time_to_peak":
576                                value = numeric_time_to_peak(
577                                    dim,
578                                    milkrecordings,
579                                    model,
580                                    fitting=fitting,
581                                    key=key,
582                                    parity=parity,
583                                    breed=breed,
584                                    custom_priors=custom_priors,
585                                    milk_unit=milk_unit,
586                                    continent=continent,
587                                )
588                            elif characteristic == "peak_yield":
589                                value = numeric_peak_yield(
590                                    dim,
591                                    milkrecordings,
592                                    model,
593                                    fitting=fitting,
594                                    key=key,
595                                    parity=parity,
596                                    breed=breed,
597                                    custom_priors=custom_priors,
598                                    milk_unit=milk_unit,
599                                    continent=continent,
600                                )
601                            elif characteristic == "cumulative_milk_yield":
602                                if lactation_length == "max":
603                                    numeric_lactation_length = max(dim)
604                                elif isinstance(lactation_length, int):
605                                    numeric_lactation_length = lactation_length
606                                else:
607                                    numeric_lactation_length = 305
608                                value = numeric_cumulative_yield(
609                                    dim,
610                                    milkrecordings,
611                                    model,
612                                    fitting=fitting,
613                                    lactation_length=numeric_lactation_length,
614                                    key=key,
615                                    parity=parity,
616                                    breed=breed,
617                                    custom_priors=custom_priors,
618                                    milk_unit=milk_unit,
619                                    continent=continent,
620                                )
621
622                    else:
623                        if persistency_method == "derived":
624                            # find lactation length from data
625                            if lactation_length == "max":
626                                lactation_length = max(dim)
627                            elif isinstance(lactation_length, int):
628                                lactation_length = lactation_length
629                            else:
630                                lactation_length = 305
631                            assert isinstance(milk_unit, str)
632                            value = persistency_fitted_curve(
633                                dim,
634                                milkrecordings,
635                                model,
636                                fitting="bayesian",
637                                key=key,
638                                parity=parity,
639                                breed=breed,
640                                custom_priors=custom_priors,
641                                milk_unit=milk_unit,
642                                continent=continent,
643                                lactation_length=lactation_length,
644                            )
645
646                        else:
647                            value = persistency_milkbot(fitted_params_bayes[3])
648
649                try:
650                    if value is None:
651                        return None
652                    if characteristic != "persistency" and value < 0:
653                        return None
654                    return float(value)
655                except ValueError:
656                    raise Exception("Could not compute characteristic")
657            else:
658                raise Exception("Bayesian fitting is currently only implemented for MilkBot models")
659
660    else:
661        raise Exception("Unknown characteristic")

242def center_lactation_data(
243    milk_matrix: np.ndarray,
244    standard_lc: np.ndarray,
245    day_mean_method: str = "standard_lc",
246) -> np.ndarray:
247    """Center lactation yields before covariance estimation.
248
249    Args:
250        milk_matrix: Yield matrix with lactations in rows and days in columns.
251        standard_lc: Expected day-wise milk yield profile.
252        day_mean_method: Mean-centering strategy. Supported values are
253            ``"standard_lc"`` (default) and ``"data"``.
254
255    Returns:
256        A centered matrix with the same shape as ``milk_matrix``.
257
258    Raises:
259        ValueError: If ``day_mean_method`` is not supported.
260    """
261    if day_mean_method == "standard_lc":
262        day_mean = standard_lc
263    elif day_mean_method == "data":
264        day_mean = np.nanmean(milk_matrix, axis=0)
265    else:
266        raise ValueError("day_mean_method must be 'standard_lc' or 'data'.")
267
268    return milk_matrix - day_mean