lactationcurve.characteristics

293def calculate_characteristic(
294    dim,
295    milkrecordings,
296    model="wood",
297    characteristic="cumulative_milk_yield",
298    fitting="frequentist",
299    key=None,
300    parity=3,
301    breed="H",
302    continent="USA",
303    persistency_method="derived",
304    lactation_length=305,
305):
306    """Evaluate a lactation curve characteristic from observed test-day data.
307
308    This function fits the requested model (frequentist or Bayesian via MilkBot),
309    retrieves model parameters, and evaluates the requested characteristic using the
310    symbolic expression (if available), falling back to numeric methods when needed.
311
312    Args:
313        dim (Int): Days in milk (DIM).
314        milkrecordings (Float): Milk recordings (kg or lbs) for each DIM.
315        model (str): Model name. Supported for this function:
316            'milkbot', 'wood', 'wilmink', 'ali_schaeffer', 'fischer'.
317        characteristic (str): One of 'time_to_peak', 'peak_yield',
318            'cumulative_milk_yield', 'persistency'.
319        fitting (str): 'frequentist' (default) or 'bayesian'.
320        key (str | None): API key for MilkBot Bayesian fitting.
321        parity (Int): Parity of the cow; values above 3 are considered as 3 (Bayesian).
322        breed (str): 'H' (Holstein) or 'J' (Jersey) (Bayesian).
323        continent (str): 'USA', 'EU', or 'CHEN' (Bayesian).
324        persistency_method (str): 'derived' (average slope after peak; default) or 'literature'
325            (only for Wood and MilkBot).
326        lactation_length (Int | str): Horizon for persistency calculation:
327            305 (default), 'max' (use max DIM in data), or an integer.
328
329    Returns:
330        float: The requested characteristic value.
331
332    Raises:
333        Exception: If inputs are invalid, the model/characteristic is unsupported,
334            an API key is missing for Bayesian fitting, or the characteristic cannot be computed.
335    """
336    # check and prepare input
337    inputs = validate_and_prepare_inputs(
338        dim,
339        milkrecordings,
340        model=model,
341        fitting=fitting,
342        breed=breed,
343        parity=parity,
344        continent=continent,
345        persistency_method=persistency_method,
346        lactation_length=lactation_length,
347    )
348
349    dim = inputs.dim
350    milkrecordings = inputs.milkrecordings
351    model = inputs.model
352    fitting = inputs.fitting
353    breed = inputs.breed
354    parity = inputs.parity
355    continent = inputs.continent
356    persistency_method = inputs.persistency_method
357    lactation_length = inputs.lactation_length
358
359    # After validation, these fields are guaranteed non-None for the paths below
360    assert model is not None
361    assert fitting is not None
362
363    if model not in ["milkbot", "wood", "wilmink", "ali_schaeffer", "fischer"]:
364        msg = (
365            "this function currently only works for"
366            " the milkbot, wood, wilmink,"
367            " ali_schaeffer and fischer models"
368        )
369        raise Exception(msg)
370
371    characteristic_options: list[str] = [
372        "time_to_peak",
373        "peak_yield",
374        "cumulative_milk_yield",
375        "persistency",
376    ]
377
378    if characteristic in characteristic_options:
379        if fitting == "frequentist":
380            # Get fitted parameters from your fitting function
381            fitted_params = get_lc_parameters(dim, milkrecordings, model)
382            assert fitted_params is not None
383
384            if characteristic != "persistency":
385                # Try symbolic formula first
386                assert isinstance(lactation_length, int)
387                expr, params, fn = lactation_curve_characteristic_function(
388                    model, characteristic, lactation_length
389                )
390                with np.errstate(
391                    divide="ignore", invalid="ignore"
392                ):  # get rid of warnings for invalid operations
393                    value = fn(*fitted_params)
394
395                # If symbolic formula fails or is invalid (use numeric approach)
396                if (
397                    value is None
398                    or not np.isfinite(value)
399                    or (characteristic == "time_to_peak" and value <= 0)
400                ):
401                    assert parity is not None
402                    assert breed is not None
403                    assert continent is not None
404                    if characteristic == "time_to_peak":
405                        value = numeric_time_to_peak(
406                            dim,
407                            milkrecordings,
408                            model,
409                            fitting=fitting,
410                            key=key,
411                            parity=parity,
412                            breed=breed,
413                            continent=continent,
414                        )
415                    elif characteristic == "peak_yield":
416                        value = numeric_peak_yield(
417                            dim,
418                            milkrecordings,
419                            model,
420                            fitting=fitting,
421                            key=key,
422                            parity=parity,
423                            breed=breed,
424                            continent=continent,
425                        )
426                    elif characteristic == "cumulative_milk_yield":
427                        value = numeric_cumulative_yield(
428                            dim,
429                            milkrecordings,
430                            model,
431                            fitting=fitting,
432                            lactation_length=lactation_length,
433                            key=key,
434                            parity=parity,
435                            breed=breed,
436                            continent=continent,
437                        )
438
439            else:
440                if persistency_method == "derived":
441                    # find lactation length from data
442                    if lactation_length == "max":
443                        lactation_length = max(dim)
444                    elif isinstance(lactation_length, int):
445                        lactation_length = lactation_length
446                    else:
447                        lactation_length = 305
448                    value = persistency_fitted_curve(
449                        dim,
450                        milkrecordings,
451                        model,
452                        fitting="frequentist",
453                        lactation_length=lactation_length,
454                    )
455
456                else:
457                    if model == "wood":
458                        value = persistency_wood(fitted_params[1], fitted_params[2])
459
460                    elif model == "milkbot":
461                        assert len(fitted_params) >= 4
462                        value = persistency_milkbot(fitted_params[3])
463
464                    else:
465                        msg = (
466                            "Currently only the Wood model"
467                            " and MilkBot model have a"
468                            " separate model function from"
469                            " the literature integrated for"
470                            " persistency. if"
471                            ' persistency="derived" is'
472                            " selected, persistency can be"
473                            " calculated for every model as"
474                            " the average slope of the"
475                            " lactation after the peak."
476                        )
477                        raise Exception(msg)
478            try:
479                return float(value)
480            except ValueError:
481                raise Exception(
482                    "Could not compute characteristic, possibly due to invalid fitted parameters"
483                )
484
485        else:
486            if model == "milkbot":
487                if key is None:
488                    raise Exception("Key needed to use Bayesian fitting engine MilkBot")
489                else:
490                    assert parity is not None
491                    assert breed is not None
492                    assert continent is not None
493                    fitted_params_bayes = bayesian_fit_milkbot_single_lactation(
494                        dim, milkrecordings, key, parity, breed, continent
495                    )
496                    fitted_params_bayes = (
497                        fitted_params_bayes["scale"],
498                        fitted_params_bayes["ramp"],
499                        fitted_params_bayes["offset"],
500                        fitted_params_bayes["decay"],
501                    )
502
503                    if characteristic != "persistency":
504                        # Get the symbolic expression and model parameters
505                        expr, param_symbols, fn = lactation_curve_characteristic_function(
506                            model, characteristic
507                        )
508                        value = fn(*fitted_params_bayes)
509
510                    else:
511                        if persistency_method == "derived":
512                            # find lactation length from data
513                            if lactation_length == "max":
514                                lactation_length = max(dim)
515                            elif isinstance(lactation_length, int):
516                                lactation_length = lactation_length
517                            else:
518                                lactation_length = 305
519                            value = persistency_fitted_curve(
520                                dim,
521                                milkrecordings,
522                                model,
523                                fitting="bayesian",
524                                key=key,
525                                parity=parity,
526                                breed=breed,
527                                continent=continent,
528                                lactation_length=lactation_length,
529                            )
530
531                        else:
532                            value = persistency_milkbot(fitted_params_bayes[3])
533
534                    return float(value)
535            else:
536                raise Exception("Bayesian fitting is currently only implemented for MilkBot models")
537
538    else:
539        raise Exception("Unknown characteristic")

 93def lactation_curve_characteristic_function(
 94    model="wood", characteristic=None, lactation_length=305
 95):
 96    """Build (or fetch from cache) a symbolic expression and fast numeric function for an LCC.
 97
 98    This function derives the requested **lactation curve characteristic** for a given
 99    model using SymPy (derivative / root finding / integration). It returns the symbolic
100    expression, the tuple of parameter symbols (argument order), and a lambdified
101    numeric function that can be evaluated with numerical parameters.
102
103    The symbolic derivation and integration are done only once per (model, characteristic)
104    and then **cached** for reuse.
105
106    Args:
107        model (str): Model name. Options:
108            'milkbot', 'wood', 'wilmink', 'ali_schaeffer', 'fischer',
109            'brody', 'sikka', 'nelder', 'dhanoa', 'emmans', 'hayashi',
110            'rook', 'dijkstra', 'prasad'.
111        characteristic (str | None): Desired characteristic. Options:
112            'time_to_peak', 'peak_yield', 'cumulative_milk_yield', 'persistency'.
113            If `None` or unrecognized, a dict of all available characteristics is returned
114            (with `persistency` possibly `None` if derivation is not feasible).
115        lactation_length (int): Length of lactation in days used in persistency
116            computation (default 305).
117
118    Returns:
119        tuple:
120            expr: SymPy expression (or dict of expressions if `characteristic` is None).
121            params: Tuple of SymPy symbols for model parameters (argument order).
122            func: Lambdified numeric function `f(*params)` (or dict of functions).
123
124    Raises:
125        Exception: If the model is unknown, or if no positive real solution for
126            peak timing/yield exists where required.
127    """
128    # check the cache
129    storage = (model, characteristic)
130    if storage in _LCC_CACHE:
131        return (
132            _LCC_CACHE[storage]["expr"],
133            _LCC_CACHE[storage]["params"],
134            _LCC_CACHE[storage]["func"],
135        )
136
137    # make sure model is all lowercase
138    model = model.lower()
139
140    # define functions
141    if model == "brody":
142        # === BRODY 1 ===
143        a, b, k1, k2, t = symbols("a b k1 k2 t", real=True, positive=True)
144        function = a * exp(-k1 * t) - b * exp(-k2 * t)
145
146    elif model == "sikka":
147        # === SIKKA ===
148        a, b, c, t = symbols("a b c t", real=True, positive=True)
149        function = a * exp(b * t - c * t**2)
150
151    elif model == "fischer":
152        # === FISCHER ===
153        a, b, c, t = symbols("a b c t", real=True, positive=True)
154        function = a - b * t - a * exp(-c * t)
155
156    elif model == "nelder":
157        # === NELDER ===
158        a, b, c, t = symbols("a b c t", real=True, positive=True)
159        function = t / (a + b * t + c * t**2)
160
161    elif model == "wood":
162        # === WOOD ===
163        a, b, c, t = symbols("a b c t", real=True, positive=True)
164        function = a * t**b * exp(-c * t)
165
166    elif model == "dhanoa":
167        # === DHANOA ===
168        a, b, c, t = symbols("a b c t", real=True, positive=True)
169        function = a * t ** (b * c) * exp(-c * t)
170
171    elif model == "emmans":
172        # === EMMANS ===
173        a, b, c, d, t = symbols("a b c d t")
174        function = a * exp(-exp(d - b * t)) * exp(-c * t)
175
176    elif model == "ali_schaeffer":
177        # ====ALI=====
178        a, b, c, d, k, t = symbols("a b c d k t", real=True, positive=True)
179        function = (
180            a + b * (t / 305) + c * (t / 305) ** 2 + d * log(t / 305) + k * log((305 / t) ** 2)
181        )
182
183    elif model == "wilmink":
184        # === WILMINK ===
185        a, b, c, k, t = symbols("a b c k t", real=True, positive=True)
186        function = a + b * t + c * exp(-k * t)
187
188    elif model == "hayashi":
189        # === HAYASHI ===
190        a, b, c, d, t = symbols("a b c d t", real=True, positive=True)
191        function = b * (exp(-t / c) - exp(-t / (a * c)))
192
193    elif model == "rook":
194        # === ROOK ===
195        a, b, c, d, t = symbols("a b c d t", real=True, positive=True)
196        function = a * (1 / (1 + b / (c + t))) * exp(-d * t)
197
198    elif model == "dijkstra":
199        # === DIJKSTRA ===
200        a, b, c, d, t = symbols("a b c d t", real=True, positive=True)
201        function = a * exp((b * (1 - exp(-c * t)) / c) - d * t)
202
203    elif model == "prasad":
204        # === PRASAD ===
205        a, b, c, d, t = symbols("a b c d t", real=True, positive=True)
206        function = a + b * t + c * t**2 + d / t
207
208    elif model == "milkbot":
209        # === MILKBOT ===
210        a, b, c, d, t = symbols("a b c d t", real=True, positive=True)
211        function = a * (1 - exp((c - t) / b) / 2) * exp(-d * t)
212
213    else:
214        raise Exception("Unknown model")
215
216    # find derivative
217    fdiff = diff(function, t)
218    # solve derivative for when it is zero to find the function for time of peak
219    tpeak = solve(fdiff, t)
220
221    # define the end of lactation
222    T = lactation_length  # days in milk
223
224    # Persistency = average slope after peak, does not work for all models so therefore try except
225    persistency = None
226    try:
227        if tpeak:
228            tmp = (function.subs(t, T) - function.subs(t, tpeak[0])) / (T - tpeak[0])
229            tmp = tmp.cancel()  # light simplification
230            if is_valid_sympy_expr(tmp):
231                persistency = tmp
232    except Exception:
233        persistency = None
234
235    if characteristic != "cumulative_milk_yield":
236        if tpeak:  # Check if the list is not empty
237            peak_expr = simplify(function.subs(t, tpeak[0]))
238        else:
239            raise Exception("No positive real solution for time to peak and peak yield found")
240
241    # find function for cumulative milk yield of the lactation.
242    # A lactation length of 305 is the default, but this value can
243    # also be provided as a parameter in the characteristic function.
244    cum_my_expr = integrate(function, (t, 0, lactation_length))
245
246    # Sorted parameter list (exclude t)
247    params = tuple(
248        sorted([s for s in function.free_symbols if s.name != "t"], key=lambda x: x.name)
249    )
250
251    # ----------------------------------------------------
252    # Select requested characteristic
253    # ----------------------------------------------------
254    if characteristic == "time_to_peak":
255        expr = tpeak[0]
256    elif characteristic == "peak_yield":
257        expr = peak_expr
258    elif characteristic == "cumulative_milk_yield":
259        expr = cum_my_expr
260    elif characteristic == "persistency":
261        if persistency is None:
262            raise Exception("Persistency could not be computed symbolically")
263        expr = persistency
264    else:
265        # Return all four if None or 'all'
266        expr = {
267            "time_to_peak": tpeak[0],
268            "peak_yield": peak_expr,
269            "persistency": persistency,  # possibly None
270            "cumulative_milk_yield": cum_my_expr,
271        }
272
273    # ----------------------------------------------------
274    # Build fast numeric function with lambdify
275    # ----------------------------------------------------
276    if isinstance(expr, dict):
277        func = {
278            name: lambdify(params, ex, modules=["numpy", "scipy"])
279            for name, ex in expr.items()
280            if ex is not None
281        }
282    else:
283        func = lambdify(params, expr, modules=["numpy", "scipy"])
284
285    # ----------------------------------------------------
286    # Store in cache
287    # ----------------------------------------------------
288    _LCC_CACHE[storage] = {"expr": expr, "params": params, "func": func}
289
290    return expr, params, func

588def numeric_cumulative_yield(
589    dim, milkrecordings, model, fitting="frequentist", lactation_length=305, **kwargs
590) -> float:
591    """Compute cumulative milk yield numerically over a given horizon.
592
593    Adds up the fitted milk yield for the first `lactation_length` days of the
594    predicted yield curve.
595
596    Args:
597        dim: DIM values.
598        milkrecordings: Milk recordings (kg).
599        model: Model name.
600        fitting: 'frequentist' or 'bayesian'.
601        lactation_length: Number of days to integrate (default 305).
602        **kwargs: Additional arguments forwarded to `fit_lactation_curve`.
603
604    Returns:
605        float: Cumulative milk yield over the specified horizon.
606    """
607    y = fit_lactation_curve(dim, milkrecordings, model, fitting=fitting, **kwargs)
608    return float(np.trapezoid(y[:lactation_length], dx=1))

611def numeric_peak_yield(
612    dim,
613    milkrecordings,
614    model,
615    fitting="frequentist",
616    key=None,
617    parity=3,
618    breed="H",
619    continent="USA",
620) -> float:
621    """Compute peak yield numerically from the fitted curve.
622
623    Args:
624        dim: DIM values.
625        milkrecordings: Milk recordings (kg).
626        model: Model name.
627        fitting: 'frequentist' or 'bayesian'.
628        key: API key for MilkBot (Bayesian).
629        parity: Parity for Bayesian fitting.
630        breed: Breed for Bayesian fitting ('H' or 'J').
631        continent: Prior source for Bayesian ('USA', 'EU', 'CHEN').
632
633    Returns:
634        float: Maximum predicted milk yield.
635    """
636    # Fit the curve to get predicted milk yields
637    yields = fit_lactation_curve(
638        dim,
639        milkrecordings,
640        model,
641        fitting=fitting,
642        key=key,
643        parity=parity,
644        breed=breed,
645        continent=continent,
646    )
647    # Find the peak yield
648    peak_yield = np.max(yields)
649    return peak_yield

543def numeric_time_to_peak(
544    dim,
545    milkrecordings,
546    model,
547    fitting="frequentist",
548    key=None,
549    parity=3,
550    breed="H",
551    continent="USA",
552) -> int:
553    """Compute time to peak using a numeric approach.
554
555    Fits the curve (frequentist or Bayesian), evaluates the predicted yields,
556    and returns the DIM corresponding to the maximum predicted yield.
557
558    Args:
559        dim: DIM values.
560        milkrecordings: Milk recordings (kg).
561        model: Model name.
562        fitting: 'frequentist' or 'bayesian'.
563        key: API key for MilkBot (Bayesian).
564        parity: Parity for Bayesian fitting.
565        breed: Breed for Bayesian fitting ('H' or 'J').
566        continent: Prior source for Bayesian ('USA', 'EU', 'CHEN').
567
568    Returns:
569        int: DIM at which the curve attains its maximum (1-indexed).
570    """
571    # Fit the curve to get predicted milk yields
572    yields = fit_lactation_curve(
573        dim,
574        milkrecordings,
575        model,
576        fitting=fitting,
577        key=key,
578        parity=parity,
579        breed=breed,
580        continent=continent,
581    )
582    # Find the index of the peak yield
583    peak_idx = np.argmax(yields)
584    # Return the corresponding DIM
585    return int(peak_idx + 1)  # +1 because DIM starts at 1, not 0

677def persistency_fitted_curve(
678    dim,
679    milkrecordings,
680    model,
681    fitting="frequentist",
682    key=None,
683    parity=3,
684    breed="H",
685    continent="USA",
686    lactation_length=305,
687) -> float:
688    """Persistency as the average slope after peak until end of lactation (numeric).
689
690    This is the default approach because symbolic derivation is not feasible for
691    all models. It computes:
692        `(yield_at_end - yield_at_peak) / (lactation_length - time_to_peak)`
693
694    Args:
695        dim: DIM values.
696        milkrecordings: Milk recordings (kg).
697        model: Model name. Options include 'milkbot' (Bayesian or frequentist),
698            'wood', 'wilmink', 'ali_schaeffer', 'fischer'.
699        fitting: 'frequentist' or 'bayesian'.
700        key: API key (only for Bayesian fitting).
701        parity: Parity of the cow; values above 3 treated as 3 (Bayesian).
702        breed: 'H' or 'J' (Bayesian).
703        continent: 'USA', 'EU', or 'CHEN' (Bayesian).
704        lactation_length (int | str): 305 (default), 'max' (use max DIM), or integer.
705
706    Returns:
707        float: Average slope after the peak until end of lactation.
708    """
709    # calculate time to peak
710    t_peak = numeric_time_to_peak(
711        dim,
712        milkrecordings,
713        model,
714        fitting=fitting,
715        key=key,
716        parity=parity,
717        breed=breed,
718        continent=continent,
719    )
720
721    # determine lactation length
722    if lactation_length == "max":
723        lactation_length = max(dim)
724    elif isinstance(lactation_length, int):
725        lactation_length = lactation_length
726    else:
727        lactation_length = 305
728
729    # calculate milk yield at peak
730    yields = fit_lactation_curve(
731        dim,
732        milkrecordings,
733        model,
734        fitting=fitting,
735        key=key,
736        parity=parity,
737        breed=breed,
738        continent=continent,
739    )
740    peak_yield = yields[int(t_peak) - 1]  # -1 to prevent index error
741    # calculate milk yield at end of lactation
742    end_yield = yields[int(lactation_length) - 1]  # -1 to prevent index error
743
744    # calculate persistency
745    persistency = (end_yield - peak_yield) / (lactation_length - t_peak)
746    return persistency

665def persistency_milkbot(d) -> float:
666    """Persistency from the MilkBot model (Ehrlich, 2013): `Persistency = 0.693 / d`.
667
668    Args:
669        d (float): Parameter `d` of the MilkBot model.
670
671    Returns:
672        float: Persistency value from the MilkBot formula.
673    """
674    return 0.693 / d

652def persistency_wood(b, c) -> float:
653    """Persistency from Wood et al. (1984): `Persistency = -(b + 1) * ln(c)`.
654
655    Args:
656        b (float): Parameter `b` of the Wood model.
657        c (float): Parameter `c` of the Wood model.
658
659    Returns:
660        float: Persistency value from the Wood formula.
661    """
662    return float(-(b + 1) * ln(c))

 38def test_interval_method(
 39    df, days_in_milk_col=None, milking_yield_col=None, test_id_col=None, default_test_id=1
 40):
 41    """Compute 305-day total milk yield using the ICAR Test Interval Method.
 42
 43    The method applies:
 44    - Linear projection from calving to the first test day,
 45    - Trapezoidal integration between consecutive test days,
 46    - Linear projection from the last test day to DIM=305.
 47
 48    Args:
 49        df (pd.DataFrame): Input DataFrame with at least DaysInMilk, MilkingYield,
 50            and (optionally) TestId columns (names can be provided via arguments
 51            or matched via known aliases, case-insensitive).
 52        days_in_milk_col (str | None): Optional column name override for DaysInMilk.
 53        milking_yield_col (str | None): Optional column name override for MilkingYield.
 54        test_id_col (str | None): Optional column name override for TestId.
 55        default_test_id (Any): If TestId is missing, a new `TestId` column is created
 56            with this value.
 57
 58    Returns:
 59        pd.DataFrame: Two-column DataFrame with
 60            - "TestId": identifier per lactation,
 61            - "Total305Yield": computed total milk yield over 305 days.
 62
 63    Raises:
 64        ValueError: If required columns (DaysInMilk or MilkingYield) cannot be found.
 65
 66    Notes:
 67        - Records with DIM > 305 are dropped before computation.
 68        - At least two data points per TestId are required for trapezoidal integration;
 69          otherwise the lactation is skipped.
 70    """
 71    result = []
 72
 73    # create a bit more flexibility in naming the columns and
 74    # when only one lactation is put in without a testid
 75
 76    # Define accepted variations for each logical column
 77    # Accepted aliases (case-insensitive)
 78    aliases = {
 79        "DaysInMilk": ["daysinmilk", "dim", "testday"],
 80        "MilkingYield": ["milkingyield", "testdaymilkyield", "milkyield", "yield"],
 81        "TestId": ["animalid", "testid", "id"],
 82    }
 83
 84    # Create a mapping from lowercase to actual column names
 85    col_lookup = {col.lower(): col for col in df.columns}
 86
 87    def get_col_name(override, possible_names):
 88        """Return a matching actual column name from `df`, or `None` if not found.
 89
 90        Args:
 91            override (str | None): Explicit column name provided by the user.
 92            possible_names (list[str]): List of acceptable aliases (lowercase).
 93
 94        Returns:
 95            str | None: The actual column name present in `df`, or `None` if no match.
 96        """
 97        if override:
 98            return col_lookup.get(override.lower())
 99        for name in possible_names:
100            if name in col_lookup:
101                return col_lookup[name]
102        return None
103
104    # Resolve columns
105    dim_col = get_col_name(days_in_milk_col, aliases["DaysInMilk"])
106    if not dim_col:
107        raise ValueError("No DaysInMilk column found in DataFrame.")
108
109    my_col = get_col_name(milking_yield_col, aliases["MilkingYield"])
110    if not my_col:
111        raise ValueError("No MilkingYield column found in DataFrame.")
112
113    id_col = get_col_name(test_id_col, aliases["TestId"])
114    if not id_col:
115        id_col = "TestId"
116        df[id_col] = default_test_id
117
118    # Filter out records where Day > 305
119    df = df[df[dim_col] <= 305]
120
121    # Iterate over each lactation
122    for lactation in df[id_col].unique():
123        lactation_df = df[df[id_col] == lactation].copy()
124
125        # Sort by DaysInMilk ascending
126        lactation_df.sort_values(by=dim_col, ascending=True, inplace=True)
127
128        if len(lactation_df) < 2:
129            print(f"Skipping TestId {lactation}: not enough data points for interpolation.")
130            continue
131
132        # Start and end points
133        start = lactation_df.iloc[0]
134        end = lactation_df.iloc[-1]
135
136        # Start contribution
137        MY0 = start[dim_col] * start[my_col]
138
139        # End contribution
140        MYend = (306 - end[dim_col]) * end[my_col]
141
142        # Intermediate trapezoidal contributions
143        lactation_df["width"] = lactation_df[dim_col].diff().shift(-1)
144        lactation_df["avg_yield"] = (lactation_df[my_col] + lactation_df[my_col].shift(-1)) / 2
145        lactation_df["trapezoid_area"] = lactation_df["width"] * lactation_df["avg_yield"]
146
147        total_intermediate = lactation_df["trapezoid_area"].sum()
148
149        total_yield = MY0 + total_intermediate + MYend
150        result.append((lactation, total_yield))
151
152    return pd.DataFrame(result, columns=["TestId", "Total305Yield"])