Repository: CDCgov/cfa-viral-lineage-model

Issue:
Syncing lineages between model and eval data
I'm looking at the pair of `model.csv` and `eval.csv` for a 2022-01-01 forecast date, and the evaluation data has one more lineage than the model data, namely `22C`. We need to make sure that the set of lineages is consistent across modeling and evaluation.

Part of the awkwardness here is hindsight, in that we have said 22C is a lineage we want to model for all three of the forecast dates in `retrospective-forecasting/config` even though it emerges after the first calibration period. 

I think the generally correct thing to do is to make sure such lineages get labeled as "other" but there may be the possibility that we know we want to model a lineage but just don't happen to see it? (Something which has attracted attention elsewhere, but not been imported yet, for instance.)


Discussion:
swo: I generally really like where this is going. As you can see, I borrowed a lot of this thinking for https://github.com/CDCgov/cfa-immunization-uptake-projection/pull/3.

I think we're facing a design choice about what an "evaluator" is (ie how much it does), versus a "score".

My current leaning is toward this setup:

1. A score is a *function* that works at the level of one unit of analysis. It takes in one (set of) predictions (ie posterior samples), a matching set of real/true/observed data, and returns a single number.
    - For the typical forecasts, a unit of analysis is a date/geography combination.
    - For the "time to domination" predictions, the unit of analysis is the geography, over many dates.
2. An "evaluator" is either a function or an object, that does some level of things. There are a few options.

- Option A (minimum): The evaluator aggregates the unit-of-analysis-level scores into a model-level score. This could also be a function: that takes in a score function, a prediction dataset, and an evaluation (truth/target) dataset.
    - For typical forecasts, this might be a sum or a mean or something.
    - For time-to-domination, the aggregation is the trivial identity operation.
- Option B (medium), the evaluator takes one (or more?) fitted models and an evaluation (truth) dataset, tells the models to generate the needed predictions, then scores and aggregates. Again, this could be a function, but it's starting to feel more like an object, because it's doing message passing ("hey object, give me xyz predictions for xyz dates").
- Option C (maximum): At a maximum level, the evaluator takes the unfitted models, a calibration dataset, and an evaluation dataset. It fits the models, then does everything the medium-level evaluator would do.

The medium & maximum options are harder to conceptualize as solitary Python objects when we're talking about distributed computing. They're more like elements of the pipeline, that could be realized as Python objects if everything is done in one computer.

I do think there's also merit in having data object types (as per https://github.com/CDCgov/cfa-immunization-uptake-projection/pull/3/files). This was a score function can ask "is this thing the right type?" before doing anything with it. (If it wants to be really careful, it can force a `input_data.validate()` call, which someone noted somewhere is confusing because validation might require a `.collect()` call.) That might be mitigated by this "callback" architecture, where the evaluator is asking the fitted model for only exactly the stuff it needs.

One advantage here is that the evaluator might ask for `my_model.predict(my_dates, my_counts, metric="count")`. The model might error and say it doesn't know how to do counts. Or it might predict phi and then look at its superclass for a convert-phi-to-counts method.

(I guess, if the things the models will output are only ever phi or counts, then we don't need any of that. Or if the models will only ever output phi, and we will always convert to counts with multinomial sampling a la #50, then we don't need this.)
afmagee42: Also flagging for later that we may want to play with the choice of multinomial sampler here. In #50 I used numpyro's and it was somewhat ponderous for this many draws. There may not be much getting around that, but if we're going to do this many times on many datasets, it could be worth a few hour's testing between different library options and different batching schemes to save the down-stream computation.
swo: > Also flagging for later that we may want to play with the choice of multinomial sampler here. In #50 I used numpyro's and it was somewhat ponderous for this many draws. There may not be much getting around that, but if we're going to do this many times on many datasets, it could be worth a few hour's testing between different library options and different batching schemes to save the down-stream computation.

See #70 

Patch:
diff --git a/README.md b/README.md
index 5580414..95758a3 100644
--- a/README.md
+++ b/README.md
@@ -57,17 +57,6 @@ Rows are uniquely identified by `(fd_offset, division, lineage, sample_index)`.
 - [ ] (9) Study more on how to set priors on the logit scale to induce priors on the probability simplex
 - [ ] (10) Does our ability to identify "good" models change if we evaluate daily vs weekly predictions?
 
-| Sprint | Start Date | Target milestones | Notes                            |
-| ------ | ---------- | ----------------- | -------------------------------- |
-| L      | Jun 10     | 1                 |                                  |
-| M      | Jun 24     | 1, 2              |                                  |
-| N      | Jul 08     | 2, 3              |                                  |
-| O      | Jul 22     | 3                 |                                  |
-| P      | Aug 05     | 3, 4              | Thanasi at JSM for one week here |
-| Q      | Aug 19     | 4, 5, 8           |                                  |
-| R      | Sep 2      |                   |                                  |
-| S      | Sep 16     |                   |                                  |
-
 ### A ladder of (multinomial logistic regression) models for consideration
 
 It seems useful to start at the bottom of the ladder, both for debugging purposes, but also to get a sense of the added predictive power of each step.
diff --git a/linmod/data.py b/linmod/data.py
index 9b42408..65f977f 100755
--- a/linmod/data.py
+++ b/linmod/data.py
@@ -31,7 +31,7 @@ import polars as pl
 import yaml
 import zstandard
 
-from .utils import ValidPath, print_message
+from .utils import ValidPath, expand_grid, print_message
 
 DEFAULT_CONFIG = {
     "data": {
@@ -289,18 +289,39 @@ def main(cfg: Optional[dict]):
             .then(pl.col("lineage"))
             .otherwise(pl.lit("other"))
         )
+        .collect()
+    )
+
+    # Generate every combination of date-division-lineage, so that:
+    #  1. The evaluation dataset will be evaluation-ready, with 0 counts
+    #     where applicable
+    #  2. The modeling dataset will have every lineage of interest represented,
+    #     even if a lineage was only sampled in the evaluation period
+    observations_key = expand_grid(
+        date=full_df["date"].unique(),
+        division=full_df["division"].unique(),
+        lineage=full_df["lineage"].unique(),
     )
 
     eval_df = (
         full_df.group_by("lineage", "date", "division")
         .agg(count=pl.len())
+        .join(
+            observations_key,
+            on=("date", "division", "lineage"),
+            how="right",
+        )
         .with_columns(
-            fd_offset=(pl.col("date") - forecast_date).dt.total_days()
+            fd_offset=(pl.col("date") - forecast_date).dt.total_days(),
+            count=pl.col("count").fill_null(0),
         )
         .select("date", "fd_offset", "division", "lineage", "count")
-        .collect()
     )
 
+    assert (
+        eval_df.null_count().sum_horizontal().item() == 0
+    ), "Null values detected in evaluation dataset."
+
     eval_df.write_parquet(ValidPath(config["data"]["save_file"]["eval"]))
 
     print_message(" done.")
@@ -310,13 +331,24 @@ def main(cfg: Optional[dict]):
         full_df.filter(pl.col("date_submitted") <= forecast_date)
         .group_by("lineage", "date", "division")
         .agg(count=pl.len())
+        .join(
+            observations_key,
+            on=("date", "division", "lineage"),
+            how="right",
+        )
         .with_columns(
-            fd_offset=(pl.col("date") - forecast_date).dt.total_days()
+            fd_offset=(pl.col("date") - forecast_date).dt.total_days(),
+            count=pl.col("count").fill_null(0),
         )
         .select("date", "fd_offset", "division", "lineage", "count")
-        .collect()
+        # Remove division-days where no samples were collected, for brevity
+        .filter(pl.sum("count").over("date", "division") > 0)
     )
 
+    assert (
+        model_df.null_count().sum_horizontal().item() == 0
+    ), "Null values detected in modeling dataset."
+
     model_df.write_parquet(ValidPath(config["data"]["save_file"]["model"]))
 
     print_message(" done.")
diff --git a/linmod/eval.py b/linmod/eval.py
index 53fe51f..78cb55b 100644
--- a/linmod/eval.py
+++ b/linmod/eval.py
@@ -1,124 +1,264 @@
+import numpy as np
 import polars as pl
+from numpy.typing import ArrayLike
 
 from linmod.utils import pl_list_cycle, pl_norm
 
 
-def _merge_samples_and_data(samples, data):
-    r"""
-    Join the forecast samples and raw data dataframes, assuming they have the
-    standard format.
+def multinomial_count_sampler(
+    n: ArrayLike,
+    p: ArrayLike,
+    rng: np.random.Generator,
+) -> np.ndarray:
+    """
+    Samples from a `multinomial(n, p)` distribution.
 
-    Also compute the true proportions from the raw data.
+    Compatible shapes of `n` and `p` include:
+    - `n` is a scalar, `p` is a vector
+    - `n` is a vector, `p` is a matrix with rows corresponding to entries in `n`
     """
-    result = (
-        data.with_columns(
-            phi=(
-                pl.col("count") / pl.sum("count").over("fd_offset", "division")
-            ),
+
+    return rng.multinomial(n, p)
+
+
+class ProportionsEvaluator:
+    def __init__(self, samples: pl.LazyFrame, data: pl.LazyFrame):
+        """
+        `samples` should have the standard model output format.
+        `data` should have the standard model input format.
+        """
+
+        # Join the forecast samples and raw data dataframes.
+        # Also compute the true proportions from the raw data.
+        self.df = (
+            data.with_columns(
+                phi=(
+                    pl.col("count")
+                    / pl.sum("count").over("fd_offset", "division")
+                ),
+            )
+            .drop("count")
+            # NaN will appear on division-days where none of the lineages
+            # were sampled. We will filter these out.
+            .filter(pl.col("phi").is_not_nan())
+            .join(
+                samples,
+                on=("fd_offset", "division", "lineage"),
+                how="left",
+                suffix="_sampled",
+            )
         )
-        .drop("count")
-        .join(
-            samples,
-            on=("fd_offset", "division", "lineage"),
-            how="left",
-            suffix="_sampled",
+
+    def _mean_norm_per_division_day(self, p=1) -> pl.LazyFrame:
+        r"""
+        The expected norm of proportion forecast error for each division-day.
+
+        $E[ || f_{tg} - \phi_{tg} ||_p ]$
+
+        Returns a data frame with columns `(division, fd_offset, mean_norm)`.
+        """
+
+        return (
+            self.df.group_by("fd_offset", "division", "sample_index")
+            .agg(norm=pl_norm(pl.col("phi") - pl.col("phi_sampled"), p))
+            .group_by("fd_offset", "division")
+            .agg(mean_norm=pl.mean("norm"))
         )
-    )
 
-    return result
+    def mean_norm(self, p=1) -> float:
+        r"""
+        The expected norm of proportion forecast error, summed over all divisions
+        and days.
 
+        $\sum_{t, g} E[ || f_{tg} - \phi_{tg} ||_p ]$
+        """
 
-def proportions_mean_norm_per_division_day(samples, data, p=1):
-    r"""
-    The expected norm of proportion forecast error for each division-day.
+        return (
+            self._mean_norm_per_division_day(p=p)
+            .collect()
+            .get_column("mean_norm")
+            .sum()
+        )
 
-    $E[ || f_{tg} - \phi_{tg} ||_p ]$
+    def _energy_score_per_division_day(self, p=2) -> pl.LazyFrame:
+        r"""
+        Monte Carlo approximation to the energy score (multivariate generalization of
+        CRPS) of proportion forecasts for each division-day.
 
-    `samples` should have the standard model output format.
-    `data` should have the standard model input format.
+        $E[ || f_{tg} - \phi_{tg} ||_p ] - \frac{1}{2} E[ || f_{tg} - f_{tg}' ||_p ]$
 
-    Returns a DataFrame with columns `(division, fd_offset, mean_norm)`.
-    """
+        Returns a data frame with columns `(division, fd_offset, energy_score)`.
+        """
 
-    return (
-        _merge_samples_and_data(samples, data)
-        .group_by("fd_offset", "division", "sample_index")
-        .agg(norm=pl_norm(pl.col("phi") - pl.col("phi_sampled"), p))
-        .group_by("fd_offset", "division")
-        .agg(mean_norm=pl.mean("norm"))
-    )
+        return (
+            # First, we will gather the values of phi' we will use for (phi-phi')
+            self.df.group_by("date", "fd_offset", "division", "lineage")
+            .agg(pl.col("sample_index"), pl.col("phi"), pl.col("phi_sampled"))
+            .with_columns(replicate=pl_list_cycle(pl.col("phi_sampled"), 1))
+            .explode("sample_index", "phi", "phi_sampled", "replicate")
+            # Now we can compute the score
+            .group_by("fd_offset", "division", "sample_index")
+            .agg(
+                term1=pl_norm(pl.col("phi") - pl.col("phi_sampled"), p),
+                term2=pl_norm(
+                    (pl.col("phi_sampled") - pl.col("replicate")), p
+                ),
+            )
+            .group_by("fd_offset", "division")
+            .agg(
+                energy_score=pl.col("term1").mean()
+                - 0.5 * pl.col("term2").mean()
+            )
+        )
 
+    def energy_score(self, p=2) -> float:
+        r"""
+        The energy score of proportion forecasts, summed over all divisions and days.
 
-def proportions_mean_norm(
-    samples: pl.LazyFrame, data: pl.LazyFrame, p=1
-) -> float:
-    r"""
-    The expected norm of proportion forecast error, summed over all divisions and days.
+        $$
+        \sum_{t, g} E[ || f_{tg} - \phi_{tg} ||_p ]
+        - \frac{1}{2} E[ || f_{tg} - f_{tg}' ||_p ]
+        $$
+        """
 
-    $\sum_{t, g} E[ || f_{tg} - \phi_{tg} ||_p ]$
-    """
+        return (
+            self._energy_score_per_division_day(p=p)
+            .collect()
+            .get_column("energy_score")
+            .sum()
+        )
 
-    return (
-        proportions_mean_norm_per_division_day(samples, data, p=p)
-        .collect()
-        .get_column("mean_norm")
-        .sum()
-    )
 
+class CountsEvaluator:
+    _count_samplers = {
+        "multinomial": multinomial_count_sampler,
+    }
 
-def proportions_energy_score_per_division_day(samples, data, p=2):
-    r"""
-    Monte Carlo approximation to the energy score (multivariate generalization of CRPS)
-    of proportion forecasts for each division-day.
+    def __init__(
+        self,
+        samples: pl.LazyFrame,
+        data: pl.LazyFrame,
+        count_sampler: str = "multinomial",
+        seed: int = None,
+    ):
+        r"""
+        Evaluates count forecasts $\hat{Y}$ sampled from a specified observation model given model
+        proportion forecasts.
 
-    $E[ || f_{tg} - \phi_{tg} ||_p ] - \frac{1}{2} E[ || f_{tg} - \f_{tg}' ||_p ]$
+        `samples` should have the standard model output format.
+        `data` should have the standard model input format.
+        `count_sampler` should be one of the keys in `CountsEvaluator._count_samplers`.
+        `seed` is an optional random seed for the count sampler.
+        """
 
-    `samples` should have the standard model output format.
-    `data` should have the standard model input format.
+        assert count_sampler in type(self)._count_samplers, (
+            f"Count sampler '{count_sampler}' not found. "
+            f"Available samplers: {', '.join(type(self)._count_samplers)}"
+        )
+        count_sampler = type(self)._count_samplers[count_sampler]
 
-    Returns a DataFrame with columns `(division, fd_offset, energy_score)`.
-    """
+        rng = np.random.default_rng(seed)
 
-    # First, we will gather the values of phi' we will use for (phi-phi')
-    samples = (
-        samples.group_by("fd_offset", "division", "lineage")
-        .agg(pl.col("sample_index"), pl.col("phi"))
-        .with_columns(
-            sample_index=pl_list_cycle(pl.col("sample_index"), 1),
-            replicate=pl_list_cycle(pl.col("phi"), 1),
+        self.df = (
+            data.join(
+                samples.rename({"phi": "phi_sampled"}),
+                on=("fd_offset", "division", "lineage"),
+                how="left",
+            )
+            .group_by("date", "fd_offset", "division", "sample_index")
+            .agg(pl.col("lineage"), pl.col("phi_sampled"), pl.col("count"))
+            .with_columns(
+                count_sampled=pl.struct(
+                    pl.col("phi_sampled"), N=pl.col("count").list.sum()
+                ).map_elements(
+                    lambda struct: list(
+                        count_sampler(struct["N"], struct["phi_sampled"], rng)
+                    ),
+                    return_dtype=pl.List(pl.Int64),
+                )
+            )
+            .explode("lineage", "phi_sampled", "count", "count_sampled")
+            .drop("phi_sampled")
         )
-        .explode("sample_index", "phi", "replicate")
-    )
-
-    return (
-        _merge_samples_and_data(samples, data)
-        .group_by("fd_offset", "division", "sample_index")
-        .agg(
-            term1=pl_norm(pl.col("phi") - pl.col("phi_sampled"), p),
-            term2=pl_norm((pl.col("phi_sampled") - pl.col("replicate")), p),
+
+    def _mean_norm_per_division_day(self, p=1) -> pl.LazyFrame:
+        r"""
+        The expected norm of count forecast error for each division-day.
+
+        $E[ || \hat{Y}_{tg} - Y_{tg} ||_p ]$
+
+        Returns a data frame with columns `(division, fd_offset, mean_norm)`.
+        """
+
+        return (
+            self.df.group_by("fd_offset", "division", "sample_index")
+            .agg(norm=pl_norm(pl.col("count") - pl.col("count_sampled"), p))
+            .group_by("fd_offset", "division")
+            .agg(mean_norm=pl.mean("norm"))
         )
-        .group_by("fd_offset", "division")
-        .agg(
-            energy_score=pl.col("term1").mean() - 0.5 * pl.col("term2").mean()
+
+    def mean_norm(self, p=1) -> float:
+        r"""
+        The expected norm of count forecast error, summed over all divisions
+        and days.
+
+        $\sum_{t, g} E[ || \hat{Y}_{tg} - Y_{tg} ||_p ]$
+        """
+
+        return (
+            self._mean_norm_per_division_day(p=p)
+            .collect()
+            .get_column("mean_norm")
+            .sum()
         )
-    )
 
+    def _energy_score_per_division_day(self, p=2) -> pl.LazyFrame:
+        r"""
+        Monte Carlo approximation to the energy score (multivariate generalization of
+        CRPS) of count forecasts for each division-day.
 
-def proportions_energy_score(
-    samples: pl.LazyFrame, data: pl.LazyFrame, p=2
-) -> float:
-    r"""
-    The energy score of proportion forecasts, summed over all divisions and days.
+        $E[ || \hat{Y}_{tg} - Y_{tg} ||_p ] - \frac{1}{2} E[ || \hat{Y}_{tg} - \hat{Y}_{tg}' ||_p ]$
 
-    $$
-    \sum_{t, g} E[ || f_{tg} - \phi_{tg} ||_p ]
-    - \frac{1}{2} E[ || f_{tg} - f_{tg}' ||_p ]
-    $$
-    """
+        Returns a data frame with columns `(division, fd_offset, energy_score)`.
+        """
 
-    return (
-        proportions_energy_score_per_division_day(samples, data, p=p)
-        .collect()
-        .get_column("energy_score")
-        .sum()
-    )
+        return (
+            # First, we will gather the values of count' we will use for (count-count')
+            self.df.group_by("date", "fd_offset", "division", "lineage")
+            .agg(
+                pl.col("sample_index"),
+                pl.col("count"),
+                pl.col("count_sampled"),
+            )
+            .with_columns(replicate=pl_list_cycle(pl.col("count_sampled"), 1))
+            .explode("sample_index", "count", "count_sampled", "replicate")
+            # Now we can compute the score
+            .group_by("fd_offset", "division", "sample_index")
+            .agg(
+                term1=pl_norm(pl.col("count") - pl.col("count_sampled"), p),
+                term2=pl_norm(
+                    (pl.col("count_sampled") - pl.col("replicate")), p
+                ),
+            )
+            .group_by("fd_offset", "division")
+            .agg(
+                energy_score=pl.col("term1").mean()
+                - 0.5 * pl.col("term2").mean()
+            )
+        )
+
+    def energy_score(self, p=2) -> float:
+        r"""
+        The energy score of count forecasts, summed over all divisions and days.
+
+        $$
+        \sum_{t, g} E[ || \hat{Y}_{tg} - Y_{tg} ||_p ]
+        - \frac{1}{2} E[ || \hat{Y}_{tg} - \hat{Y}_{tg}' ||_p ]
+        $$
+        """
+        return (
+            self._energy_score_per_division_day(p=p)
+            .collect()
+            .get_column("energy_score")
+            .sum()
+        )
diff --git a/retrospective-forecasting/config/early-21L.yaml b/retrospective-forecasting/config/early-21L.yaml
index 4d3c95d..73dcea3 100644
--- a/retrospective-forecasting/config/early-21L.yaml
+++ b/retrospective-forecasting/config/early-21L.yaml
@@ -66,5 +66,6 @@ evaluation:
 
   # How should forecasts be evaluated?
   metrics:
-    - proportions_mean_norm
-    - proportions_energy_score
+    - ProportionsEvaluator
+    - CountsEvaluator:
+      - count_sampler: multinomial
diff --git a/retrospective-forecasting/config/late-21L.yaml b/retrospective-forecasting/config/late-21L.yaml
index 39d2274..44d4dec 100644
--- a/retrospective-forecasting/config/late-21L.yaml
+++ b/retrospective-forecasting/config/late-21L.yaml
@@ -66,5 +66,6 @@ evaluation:
 
   # How should forecasts be evaluated?
   metrics:
-    - proportions_mean_norm
-    - proportions_energy_score
+    - ProportionsEvaluator
+    - CountsEvaluator:
+      - count_sampler: multinomial
diff --git a/retrospective-forecasting/config/mid-21L.yaml b/retrospective-forecasting/config/mid-21L.yaml
index 2ab4fef..20bf414 100644
--- a/retrospective-forecasting/config/mid-21L.yaml
+++ b/retrospective-forecasting/config/mid-21L.yaml
@@ -66,5 +66,6 @@ evaluation:
 
   # How should forecasts be evaluated?
   metrics:
-    - proportions_mean_norm
-    - proportions_energy_score
+    - ProportionsEvaluator
+    - CountsEvaluator:
+      - count_sampler: multinomial
diff --git a/retrospective-forecasting/main.py b/retrospective-forecasting/main.py
index 61548d9..dd856b7 100755
--- a/retrospective-forecasting/main.py
+++ b/retrospective-forecasting/main.py
@@ -56,7 +56,7 @@ with open(plot_script_file, "w") as plot_script:
     plot_script.write("#/usr/bin/sh\n")
     for model_name in config["forecasting"]["models"]:
         print_message(f"Fitting {model_name} model...")
-        model_class = linmod.models.__dict__[model_name]
+        model_class = getattr(linmod.models, model_name)
         model = model_class(model_data)
 
         mcmc = MCMC(
@@ -69,42 +69,29 @@ with open(plot_script_file, "w") as plot_script:
         mcmc.run(jax.random.key(0))
 
         try:
+            convergence_config = config["forecasting"]["mcmc"]["convergence"]
+
             convergence = linmod.models.get_convergence(
                 mcmc,
                 ignore_nan_in=model.ignore_nan_in(),
                 drop_ignorable_nan=False,
             )
 
-            if (
-                config["forecasting"]["mcmc"]["convergence"]["report_mode"]
-                == "failing"
-            ):
+            if convergence_config["report_mode"] == "failing":
                 convergence = convergence.filter(
-                    (
-                        pl.col("n_eff")
-                        < config["forecasting"]["mcmc"]["convergence"][
-                            "ess_cutoff"
-                        ]
-                    )
-                    | (
-                        pl.col("r_hat")
-                        > config["forecasting"]["mcmc"]["convergence"][
-                            "psrf_cutoff"
-                        ]
-                    )
+                    (pl.col("n_eff") < convergence_config["ess_cutoff"])
+                    | (pl.col("r_hat") > convergence_config["psrf_cutoff"])
                 )
             convergence.write_parquet(
                 forecast_dir / f"convergence_{model_name}.parquet"
             )
 
-            if (
-                config["forecasting"]["mcmc"]["convergence"]["plot"]
-                and convergence.shape[0] > 0
-            ):
+            if convergence_config["plot"] and convergence.shape[0] > 0:
                 plot_dir = forecast_dir / ("convergence_" + model_name)
                 plots = linmod.models.plot_convergence(
                     mcmc, convergence["param"]
                 )
+
                 for plot, par in zip(plots, convergence["param"].to_list()):
                     plot.save(plot_dir / (par + ".png"), verbose=False)
 
@@ -169,7 +156,6 @@ with open(plot_script_file, "w") as plot_script:
         del model, mcmc, forecast
 
     # Load the full evaluation dataset
-
     eval_data = pl.read_parquet(config["data"]["save_file"]["eval"])
 
     viz_data = eval_data.filter(
@@ -181,44 +167,67 @@ with open(plot_script_file, "w") as plot_script:
     eval_dir = ValidPath(config["evaluation"]["save_dir"])
 
     scores = []
-    for metric_name in config["evaluation"]["metrics"]:
-        metric_function = linmod.eval.__dict__[metric_name]
 
-        for forecast_path in forecast_dir.glob("forecasts_*.parquet"):
-            model_name = forecast_path.stem.split("_")[1]
-            print_message(
-                f"Evaluating {model_name} model using {metric_name}...", end=""
+    for forecast_path in forecast_dir.glob("forecasts_*.parquet"):
+        model_name = forecast_path.stem.split("_")[1]
+        forecast = pl.scan_parquet(forecast_path)
+
+        for evaluator_config in config["evaluation"]["metrics"]:
+            if isinstance(evaluator_config, dict):
+                assert (
+                    len(evaluator_config) == 1
+                ), "Evaluator config is formatted incorrectly."
+
+                evaluator_config = list(evaluator_config.items())
+                evaluator_name = evaluator_config[0][0]
+                evaluator_args = {
+                    k: v for d in evaluator_config[0][1] for k, v in d.items()
+                }
+
+            else:
+                evaluator_name = evaluator_config
+                evaluator_args = {}
+
+            evaluator = getattr(linmod.eval, evaluator_name)(
+                samples=forecast,
+                data=eval_data.lazy(),
+                **evaluator_args,
             )
 
-            forecast = pl.read_parquet(forecast_path)
+            for metric_name, metric_function in vars(type(evaluator)).items():
+                if metric_name.startswith("_"):
+                    continue
 
-            forecast = forecast.lazy()
-            scores.append(
-                (
-                    metric_name,
-                    model_name,
-                    metric_function(forecast, eval_data.lazy()),
+                print_message(
+                    (
+                        f"Evaluating {model_name} model using "
+                        f"{evaluator_name}.{metric_name}..."
+                    ),
+                    end="",
                 )
-            )
 
-            if (
-                metric_name == config["evaluation"]["metrics"][0]
-            ):  # One plot per model
-                data_path = config["data"]["save_file"]["eval"]
-                png_path = (
-                    eval_dir / "visualizations" / f"eval_{model_name}.png"
-                )
-                plot_script.write(
+                scores.append(
                     (
-                        "python3 -m linmod.visualize "
-                        f"-f {forecast_path} "
-                        f"-d {data_path} "
-                        f"-p {png_path} "
-                        "-t eval\n"
+                        f"{evaluator_name}.{metric_name}",
+                        model_name,
+                        metric_function(evaluator),
                     )
                 )
 
-            print_message(" done.")
+                print_message(" done.")
+
+        data_path = config["data"]["save_file"]["eval"]
+        png_path = eval_dir / "visualizations" / f"eval_{model_name}.png"
+        plot_script.write(
+            (
+                "python3 -m linmod.visualize "
+                f"-f {forecast_path} "
+                f"-d {data_path} "
+                f"-p {png_path} "
+                "-t eval\n"
+            )
+        )
+
 
 print_message("Success!")
 
@@ -226,4 +235,4 @@ pl.DataFrame(
     scores,
     schema=["Metric", "Model", "Score"],
     orient="row",
-).write_parquet(eval_dir / "results.parquet")
+).write_csv(eval_dir / "results.csv")


Test patch:
diff --git a/linmod/tests/test_eval.py b/linmod/tests/test_eval.py
index 51767a5..33e1f5c 100644
--- a/linmod/tests/test_eval.py
+++ b/linmod/tests/test_eval.py
@@ -32,18 +32,27 @@ def _generate_fake_samples_and_data(
         lineage=range(num_lineages),
     )
 
-    data = data.with_columns(count=rng.integers(0, 100, data.shape[0]))
+    data = data.with_columns(
+        count=rng.integers(0, 100, data.shape[0]),
+        date=pl.col("fd_offset"),
+    )
 
     # Generate fake forecasts of population proportions
-    samples = eval._merge_samples_and_data(
-        expand_grid(
-            fd_offset=range(num_days),
-            division=range(num_divisions),
-            lineage=range(num_lineages),
-            sample_index=range(num_samples),
-        ),
-        data,
-    ).rename({"phi": "phi_mean"})
+    samples = expand_grid(
+        fd_offset=range(num_days),
+        division=range(num_divisions),
+        lineage=range(num_lineages),
+        sample_index=range(num_samples),
+    ).join(
+        data.with_columns(
+            phi_mean=(
+                pl.col("count") / pl.sum("count").over("fd_offset", "division")
+            ),
+        ).drop("count"),
+        on=("fd_offset", "division", "lineage"),
+        how="left",
+        suffix="_sampled",
+    )
 
     samples = samples.with_columns(
         phi=rng.normal(samples["phi_mean"], np.sqrt(sample_variance))
@@ -91,7 +100,7 @@ def test_proportions_mean_L1_norm(
     # $|X - \mu| \sim \text{half-normal}(\sigma)$, which has this mean).
 
     assert np.isclose(
-        eval.proportions_mean_norm(samples, data, p=1),
+        eval.ProportionsEvaluator(samples, data).mean_norm(p=1),
         np.sqrt(sample_variance * 2 / np.pi)
         * NUM_DAYS
         * NUM_DIVISIONS
@@ -157,9 +166,11 @@ def test_proportions_mean_L1_norm2():
         }
     )
 
-    result = eval.proportions_mean_norm_per_division_day(
-        samples, data, p=1
-    ).collect()
+    result = (
+        eval.ProportionsEvaluator(samples, data)
+        ._mean_norm_per_division_day(p=1)
+        .collect()
+    )
 
     assert_frame_equal(
         result,
@@ -232,7 +243,7 @@ def test_proportions_L1_energy_score(
     )
 
     assert np.isclose(
-        eval.proportions_energy_score(samples, data, p=1),
+        eval.ProportionsEvaluator(samples, data).energy_score(p=1),
         term1 - 0.5 * term2,
         atol=atol,
     )
@@ -335,9 +346,11 @@ def test_proportions_L1_energy_score2():
         }
     )
 
-    result = eval.proportions_energy_score_per_division_day(
-        samples, data, p=1
-    ).collect()
+    result = (
+        eval.ProportionsEvaluator(samples, data)
+        ._energy_score_per_division_day(p=1)
+        .collect()
+    )
 
     assert_frame_equal(
         result,