Repository: CDCgov/cfa-immunization-uptake-projection
Pull request #141: Refactor Hill Model
Author: login: eschrom
Created: 2025-03-14T23:07:04Z | Merged: 2025-03-20T16:13:51Z

The Hill model is refactored to handle arbitrary grouping factors (including but not limited to season), separate observation from latent true uptake, etc. See docs/model_details.md for the full description of this model.

This is currently a work in progress. In particular, I am running into `Segmentation fault (core dumped)` when I run the model on flu data across 14 seasons and 54 states/jurisdictions. I did not expect this to be a large-enough data set or a complex-enough model to crash my CPU, so more investigation must follow.


Diff:
diff --git a/docs/model_details.md b/docs/model_details.md
index dd984665..94fbad34 100644
--- a/docs/model_details.md
+++ b/docs/model_details.md
@@ -89,7 +89,7 @@ At a high level, the Hill model is structured as follows:
 \end{align*}
 ```
 
-More details on the Hill model are given below.
+For the Hill model, $t$ is rescaled by dividing by 365, so that $t$ represents the proportion of a season elapsed. This is not a problem because $t$ need not be an integer here, and it helps keep all parameters on a similar scale. More details on the Hill model are given below.
 
 ## Observation Layer
 
@@ -122,9 +122,10 @@ and similarly for $H$.
 
 ```math
 \begin{align*}
-&A \sim \text{Beta}(\text{mean = }0.4, \text{ shape = }0.2) \\
-&H \sim \text{Gamma}(\text{shape = }100, \text{ scale = }1) \\
-&\sigma_{A_{G_i}} \sim \text{Exponential}(\text{mean = }0.1) ~\forall~i~\text{ in } 1, ..., I \\
-&\sigma_{H_{G_i}} \sim \text{Exponential}(\text{mean = }0.1) ~\forall~i~\text{ in } 1, ..., I \\
+&A \sim \text{Beta}(\text{shape1 = }15.0, \text{ shape2 = }20.0) \\
+&H \sim \text{Beta}(\text{shape1 = }25.0, \text{ shape2 = }50.0) \\
+&\sigma_{A_{G_i}} \sim \text{Exponential}(\text{rate = }40.0) ~\forall~i~\text{ in } 1, ..., I \\
+&\sigma_{H_{G_i}} \sim \text{Exponential}(\text{rate = }40.0) ~\forall~i~\text{ in } 1, ..., I \\
+&n \sim \text{Gamma}(\text{shape = }20.0, \text{ rate = }5.0) \\
 \end{align*}
 ```
diff --git a/iup/models.py b/iup/models.py
index 09d296ee..694f2537 100644
--- a/iup/models.py
+++ b/iup/models.py
@@ -7,7 +7,7 @@
 import numpyro.distributions as dist
 import polars as pl
 from jax import random
-from numpyro.infer import MCMC, NUTS, Predictive
+from numpyro.infer import MCMC, NUTS, Predictive, init_to_sample
 from typing_extensions import Self
 
 import iup.utils
@@ -650,30 +650,38 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
@@ -681,27 +689,46 @@ def hill(
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
             )
-            H_s = numpyro.sample(
-                "H_s",
-                dist.TruncatedNormal(H, sig_H, low=H_low, high=H_high),
-                sample_shape=(len(np.unique(season)),),
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
             )
-            mu = A_s[season] * (elapsed**n) / (H_s[season] ** n + elapsed**n)
+            # Draw deviations from the overall average A and H for each level
+            # of each grouping factor and scale them by the characteristic spread
+            # for each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
+            ) * np.repeat(A_sigs, np.array(num_group_levels))
+            H_devs = H_devs = numpyro.sample(
+                "H_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
+            ) * np.repeat(H_sigs, np.array(num_group_levels))
+            # Across all data points, look up the A and H deviations due to the
+            # grouping factors. Sum across grouping factors and include the overall
+            # average A and H to get the final total A and H for each datum
+            A_tot = np.sum(A_devs[groups], axis=1) + A
+            H_tot = np.sum(H_devs[groups], axis=1) + H
+            # Calculate the postulated latent true uptake given the time elapsed at
+            # each datum, accounting for the final total A and H values
+            mu = A_tot * (elapsed**n) / (H_tot**n + elapsed**n)
         else:
+            # Without grouping factors, use the same A and H across all data
             mu = A * (elapsed**n) / (H**n + elapsed**n)
-        sig = numpyro.sample("sig", dist.Exponential(sig_mn))
-        numpyro.sample("obs", dist.Normal(mu, sig), obs=cum_uptake)
+        # Consider the observations to be a sample with empirically known std dev,
+        # centered on the postulated latent true uptake.
+        if std_dev is None:
+            std_dev = 0.0000005
+        numpyro.sample(
+            "obs", dist.TruncatedNormal(mu, std_dev, low=0, high=1), obs=cum_uptake
+        )
 
     @staticmethod
     def augment_data(
@@ -711,6 +738,30 @@ def augment_data(
         groups: List[str] | None,
         rollouts: List[dt.date] | None,
     ) -> CumulativeUptakeData:
+        """
+        Format preprocessed data for fitting a Hill model.
+
+        Parameters:
+        data: CumulativeUptakeData
+            training data for fitting a Hill model
+         season_start_month: int
+            first month of the overwinter disease season
+        season_start_day: int
+            first day of the first month of the overwinter disease season
+        groups: List[str, ] | None - UNUSED HERE
+            column name(s) for grouping factors
+        rollouts: List[dt.date] | None - UNUSED HERE
+            rollout dates for each season in the training data
+
+        Returns:
+            Cumulative uptake data ready for fitting a Hill model.
+
+        Details
+        The following steps are required to prepare preprocessed data
+        for fitting a linear incident uptake model:
+        - Add an extra columns for time elapsed since rollout, in days
+        - Rescale this time elapsed to a proportion of the year
+        """
         data = CumulativeUptakeData(
             data.with_columns(
                 elapsed=iup.utils.date_to_elapsed(
@@ -718,6 +769,7 @@ def augment_data(
                     season_start_month,
                     season_start_day,
                 )
+                / 365
             )
         )
 
@@ -766,19 +818,28 @@ def fit(
         """
         self.group_combos = extract_group_combos(data, groups)
 
-        if groups is None:
-            groups = []
-
-        if "season" in groups:
-            season = data["season"].to_numpy()
-            unique_seasons = np.unique(season)
-            season_to_index = {s: i for i, s in enumerate(unique_seasons)}
-            season = np.array([season_to_index[s] for s in season])
-            self.season = season
+        # Tranform the levels of the grouping factors into numeric codes
+        if groups is not None:
+            self.num_group_factors = len(groups)
+            self.value_to_index = iup.utils.map_value_to_index(data.select(groups))
+            group_codes = iup.utils.value_to_index(
+                data.select(groups), self.value_to_index
+            )
+            self.num_group_levels = iup.utils.count_unique_values(group_codes)
         else:
-            season = None
+            group_codes = None
+            self.num_group_factors = 0
+            self.num_group_levels = [0]
+
+        # Prepare the data to be fed to the model. Must be numpy arrays.
+        # Cannot have zero as a standard deviation.
+        elapsed = data["elapsed"].to_numpy()
+        cum_uptake = data["estimate"].to_numpy()
+        std_dev = np.where(
+            data["sdev"].to_numpy() == 0, 0.0001, data["sdev"].to_numpy()
+        )
 
-        self.kernel = NUTS(self.model)
+        self.kernel = NUTS(self.model, init_strategy=init_to_sample)
         self.mcmc = MCMC(
             self.kernel,
             num_warmup=mcmc["num_warmup"],
@@ -788,16 +849,20 @@ def fit(
 
         self.mcmc.run(
             self.rng_key,
-            cum_uptake=data["estimate"].to_numpy(),
-            elapsed=data["elapsed"].to_numpy(),
-            season=season,
-            n_low=params["n_low"],
-            n_high=params["n_high"],
-            A_low=params["A_low"],
-            A_high=params["A_high"],
-            H_low=params["H_low"],
-            H_high=params["H_high"],
-            sig_mn=params["sig_mn"],
+            elapsed=elapsed,
+            cum_uptake=cum_uptake,
+            std_dev=std_dev,
+            groups=group_codes,
+            num_group_factors=self.num_group_factors,
+            num_group_levels=self.num_group_levels,
+            A_shape1=params["A_shape1"],
+            A_shape2=params["A_shape2"],
+            A_sig=params["A_sig"],
+            H_shape1=params["H_shape1"],
+            H_shape2=params["H_shape2"],
+            H_sig=params["H_sig"],
+            n_shape=params["n_shape"],
+            n_rate=params["n_rate"],
         )
 
         return self
@@ -828,6 +893,7 @@ def augment_scaffold(
             elapsed=iup.utils.date_to_elapsed(
                 pl.col("time_end"), season_start_month, season_start_day
             )
+            / 365
         ).drop("estimate")
 
         return scaffold
@@ -887,12 +953,22 @@ def predict(
         )
 
         predictive = Predictive(self.model, self.mcmc.get_samples())
+
         if groups is not None:
+            # Make a numpy array of numeric codes for grouping factor levels
+            # that matches the same codes used when fitting the model
+            group_codes = iup.utils.value_to_index(
+                scaffold.select(groups), self.value_to_index
+            )
+
+            # Make a prediction-machine from the fit model
             predictions = np.array(
                 predictive(
                     self.rng_key,
                     elapsed=scaffold["elapsed"].to_numpy(),
-                    season=np.repeat(self.season.max(), scaffold.shape[0]),
+                    groups=group_codes,
+                    num_group_factors=self.num_group_factors,
+                    num_group_levels=self.num_group_levels,
                 )["obs"]
             ).transpose()
         else:
@@ -996,6 +1072,7 @@ def build_scaffold(
             test_data.filter((pl.col("time_end").is_between(start_date, end_date)))
             .select(["time_end", "season"])
             .with_columns(estimate=pl.lit(0.0))
+            .unique()
         )
     # If there are no test data, use exactly the dates that were provided
     else:
@@ -1017,6 +1094,14 @@ def build_scaffold(
         )
 
     if group_combos is not None:
-        scaffold = scaffold.join(group_combos, how="cross")
+        # Even if season is a grouping factor, predictions should not
+        # be made for every season
+        if "season" in group_combos.columns:
+            group_combos = group_combos.drop("season").unique()
+
+        # Only include grouping factors in the scaffold if season
+        # wasn't the only grouping factor
+        if group_combos.shape[1] > 0:
+            scaffold = scaffold.join(group_combos, how="cross")
 
     return scaffold
diff --git a/iup/utils.py b/iup/utils.py
index d3594b04..ca7634a3 100644
--- a/iup/utils.py
+++ b/iup/utils.py
@@ -1,3 +1,6 @@
+from typing import List
+
+import numpy as np
 import polars as pl
 
 
@@ -180,3 +183,78 @@ def date_to_elapsed(
 
         # return the number of days from season start to each date
         return (date_col - season_start).dt.total_days()
+
+
+def map_value_to_index(groups: pl.DataFrame) -> dict:
+    """
+    Choose a numeric index for each level of each grouping factor in a data frame.
+
+    Parameters
+    groups: pl.DataFrame
+        levels of grouping factors (cols) for multiple data points (rows)
+
+    Returns
+    dict
+        Dictionary of dictionaries: for each grouping factor, a dictionary mapping levels to numeric codes
+    """
+    mapping = {}
+    for i in range(groups.shape[1]):
+        col_name = groups.columns[i]
+        unique_values = groups.select(col_name).unique().to_series().to_list()
+        mapping[col_name] = {v: j for j, v in enumerate(unique_values)}
+
+    return mapping
+
+
+def value_to_index(groups: pl.DataFrame, mapping: dict, unique=True) -> np.ndarray:
+    """
+    Replace each level of each grouping factor in a data frame, using a pre-determined mapping.
+
+    Parameters
+    groups: pl.DataFrame
+        levels of grouping factors (cols) for multiple data points (rows)
+    mapping: dict
+        mapping of each level of each grouping factor to a numeric code
+    unique: bool
+        whether numeric codes should be unique across grouping factors
+
+    Returns
+    np.ndarray
+        array of group levels but with numeric codes instead of level names
+
+    Details
+    If unique is False, numeric codes will be reused across grouping factors.
+    If unique is True (default), numeric codes will be used only once across grouping factors.
+    The keys of mapping must match the column names of groups.
+    """
+    assert set(mapping.keys()) == set(groups.columns), (
+        "Keys of mapping do not match grouping factor names."
+    )
+
+    for i in range(groups.shape[1]):
+        col_name = groups.columns[i]
+        groups = groups.with_columns(
+            pl.col(col_name).replace(mapping[col_name]).cast(pl.UInt8).alias(col_name)
+        )
+
+    array = groups.to_numpy()
+
+    if unique:
+        unique_counts = count_unique_values(array)
+        array = array + np.cumsum(np.array([0] + unique_counts[:-1]))
+
+    return array
+
+
+def count_unique_values(array: np.ndarray) -> List[int,]:
+    """
+    Count unique values in each column of an array
+
+    Parameters
+    array: np.ndarray
+
+    Returns
+    List[int,]
+        Number of unique values in each column of the array
+    """
+    return [len(np.unique(array[:, i])) for i in range(array.shape[1])]
diff --git a/scripts/config_template.yaml b/scripts/config_template.yaml
index 8feb3cba..a4680976 100644
--- a/scripts/config_template.yaml
+++ b/scripts/config_template.yaml
@@ -15,7 +15,7 @@ data:
   # keep only these data columns
   keep: [estimate, time_end, lci, uci]
   # use these columns as grouping factors. Almost always use at least "season", but use None if there are truly none.
-  groups:
+  groups: [season]
 
 # Timeframe for the longest desired forecast
 forecast_timeframe:
@@ -44,19 +44,18 @@ models:
   - name: HillModel
     seed: 0
     params:
-      n_low: 1.0
-      n_high: 5.0
-      A_low: 0.0
-      A_high: 1.0
-      A_sig: 0.1
-      H_low: 10.0
-      H_high: 180.0
-      H_sig: 10.0
-      sig_mn: 0.1
+      A_shape1: 60.0
+      A_shape2: 80.0
+      A_sig: 40.0
+      H_shape1: 125.0
+      H_shape2: 300.0
+      H_sig: 40.0
+      n_shape: 20.0
+      n_rate: 5.0
     mcmc:
       num_warmup: 1000
-      num_samples: 100
-      num_chains: 4
+      num_samples: 1000
+      num_chains: 1
 
 # score metrics
 score_funs: [mspe, mean_bias, eos_abe]
diff --git a/scripts/preprocess.py b/scripts/preprocess.py
index 4ff91586..2c15e1b1 100644
--- a/scripts/preprocess.py
+++ b/scripts/preprocess.py
@@ -34,6 +34,11 @@ def preprocess(
         )
         .with_columns(sdev=pl.max_horizontal("uci", "lci") / st.norm.ppf(0.975, 0, 1))
         .drop(["lci", "uci"])
+        .with_columns(
+            sdev=pl.when(pl.col("sdev") < 0.0000005)
+            .then(0.0000005)
+            .otherwise(pl.col("sdev"))
+        )
     )
 
     if groups is not None:
diff --git a/tests/test_generic_model.py b/tests/test_generic_model.py
index a6740b51..301cd15c 100644
--- a/tests/test_generic_model.py
+++ b/tests/test_generic_model.py
@@ -62,11 +62,13 @@ def test_build_scaffold_handles_no_test_data():
     interval = "7d"
     output = iup.models.build_scaffold(start_date, end_date, interval, None, None, 9, 1)
 
-    assert [d.strftime("%Y-%m-%d") for d in output["time_end"].to_list()] == [
-        "2020-01-03",
-        "2020-01-10",
-        "2020-01-17",
-    ]
+    assert set([d.strftime("%Y-%m-%d") for d in output["time_end"].to_list()]) == set(
+        [
+            "2020-01-03",
+            "2020-01-10",
+            "2020-01-17",
+        ]
+    )
 
 
 def test_build_scaffold_handles_test_data(frame):
@@ -81,10 +83,12 @@ def test_build_scaffold_handles_test_data(frame):
         start_date, end_date, interval, frame, None, 9, 1
     )
 
-    assert [d.strftime("%Y-%m-%d") for d in output["time_end"].to_list()] == [
-        "2020-01-07",
-        "2020-01-14",
-    ]
+    assert set([d.strftime("%Y-%m-%d") for d in output["time_end"].to_list()]) == set(
+        [
+            "2020-01-07",
+            "2020-01-14",
+        ]
+    )
 
 
 def test_build_scaffold_handles_groups():
diff --git a/tests/test_hill_model.py b/tests/test_hill_model.py
index a95e75d9..21884443 100644
--- a/tests/test_hill_model.py
+++ b/tests/test_hill_model.py
@@ -25,6 +25,7 @@ def frame():
             ],
             "estimate": [0.0, 0.0, 0.1, 0.01, 0.3, 0.03, 0.4, 0.04],
             "season": "2019/2020",
+            "sdev": [0.01] * 8,
         }
     ).with_columns(time_end=pl.col("time_end").str.strptime(pl.Date, "%Y-%m-%d"))
 
@@ -40,15 +41,14 @@ def params():
     """
 
     params = {
-        "n_low": 1.0,
-        "n_high": 5.0,
-        "A_low": 0.0,
-        "A_high": 1.0,
-        "A_sig": 1.0,
-        "H_low": 10.0,
-        "H_high": 180.0,
-        "H_sig": 1.0,
-        "sig_mn": 1.0,
+        "A_shape1": 15.0,
+        "A_shape2": 20.0,
+        "A_sig": 40.0,
+        "H_shape1": 25.0,
+        "H_shape2": 50.0,
+        "H_sig": 40.0,
+        "n_shape": 20.0,
+        "n_rate": 5.0,
     }
 
     return params
@@ -70,7 +70,7 @@ def test_augment_data(frame):
     """
     output = iup.models.HillModel.augment_data(frame, 9, 1, None, None)
 
-    assert output["elapsed"].to_list() == [
+    assert [round(i * 365, 1) for i in output["elapsed"].to_list()] == [
         121.0,
         121.0,
         128.0,
@@ -115,7 +115,7 @@ def test_augment_scaffold(frame):
     output = iup.models.HillModel.augment_scaffold(frame, 9, 1)
 
     assert output.shape[1] == frame.shape[1]
-    assert output["elapsed"].to_list() == [
+    assert [round(i * 365, 1) for i in output["elapsed"].to_list()] == [
         121.0,
         121.0,
         128.0,


Reviews:
user: login: afmagee42
body: I've tried to point out a few places that could be responsible for either slow gradient computations or model misspecifcation. But I can't be sure how of the probability of success.

I'm also starting to wonder if a matrix-based one-hot-encoding of the group factors is worth pursuing. It would require a decently-sized matrix, but it would potentially let us get both A and H done via the same matrix multiplications rather than separately `jnp.repeat`ing them.
state: APPROVED
author_association: NONE
submitted_at: 2025-03-19T19:28:07Z
user: login: eschrom
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T12:43:37Z
user: login: eschrom
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T12:45:49Z
user: login: eschrom
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T12:49:05Z
user: login: eschrom
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T12:51:06Z
user: login: eschrom
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T13:39:59Z
user: login: eschrom
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T13:49:47Z
user: login: eschrom
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T13:53:14Z
user: login: eschrom
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T13:56:19Z
user: login: afmagee42
state: COMMENTED
author_association: NONE
submitted_at: 2025-03-20T13:58:52Z
user: login: eschrom
state: COMMENTED
author_association: COLLABORATOR
submitted_at: 2025-03-20T14:38:19Z

Review comments:
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
path: iup/models.py
user: login: afmagee42
body: I think that the `hill` method is over-commented.

As an example: when there's a `numpyro.sample()` statement, it is clear by context that we are drawing something, so we don't need to annotate that. And when the variables are labeled in a clear fashion (which these are), we don't need to say as much about what they're doing.
created_at: 2025-03-19T16:49:54Z
author_association: NONE
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
 
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
+            )
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
+            )
+            # Draw a sample of the centered deviation from the overall average
+            # A and H for each level of each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            H_s = numpyro.sample(
-                "H_s",
-                dist.TruncatedNormal(H, sig_H, low=H_low, high=H_high),
-                sample_shape=(len(np.unique(season)),),
+            H_devs = numpyro.sample(
+                "H_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            mu = A_s[season] * (elapsed**n) / (H_s[season] ** n + elapsed**n)
+            # Scale the centered deviation for each level of each grouping factor
+            # by the characteristic spread for each grouping factor
+            A_devs = A_devs * jnp.repeat(
+                A_sigs,
+                jnp.array(num_group_levels),
path: iup/models.py
user: login: afmagee42
body: I don't think lines 718 or 723 need to be `jnp.array(num_group_levels)` and could be `np.array(num_group_levels)`. My gut tells me that `jnp.concat` and `jnp.repeat` are likely pain points for speed, so putting fewer things inside that JAX might think it needs to track might eke out a slightly higher samples/second.
created_at: 2025-03-19T17:27:37Z
author_association: NONE
diff_hunk: @@ -122,9 +122,10 @@ and similarly for $H$.
 
 ```math
 \begin{align*}
-&A \sim \text{Beta}(\text{mean = }0.4, \text{ shape = }0.2) \\
-&H \sim \text{Gamma}(\text{shape = }100, \text{ scale = }1) \\
-&\sigma_{A_{G_i}} \sim \text{Exponential}(\text{mean = }0.1) ~\forall~i~\text{ in } 1, ..., I \\
-&\sigma_{H_{G_i}} \sim \text{Exponential}(\text{mean = }0.1) ~\forall~i~\text{ in } 1, ..., I \\
+&A \sim \text{Beta}(\text{shape1 = }15.0, \text{ shape2 = }20.0) \\
+&H \sim \text{Beta}(\text{shape1 = }25.0, \text{ shape2 = }50.0) \\
+&\sigma_{A_{G_i}} \sim \text{Exponential}(\text{mean = }40.0) ~\forall~i~\text{ in } 1, ..., I \\
+&\sigma_{H_{G_i}} \sim \text{Exponential}(\text{mean = }40.0) ~\forall~i~\text{ in } 1, ..., I \\
path: docs/model_details.md
user: login: afmagee42
body: This says Exponential(mean = 40).

The config says `A_sig: 40.0`

The model says `dist.Exponential(A_sig)`.

Since NumPyro uses the rate (https://num.pyro.ai/en/stable/distributions.html#exponential), this seems a likely place for a cause of model unhappiness per https://github.com/CDCgov/cfa-immunization-uptake-projection/pull/141#issuecomment-2737459520?
created_at: 2025-03-19T17:35:37Z
author_association: NONE
diff_hunk: @@ -122,9 +122,10 @@ and similarly for $H$.
 
 ```math
 \begin{align*}
-&A \sim \text{Beta}(\text{mean = }0.4, \text{ shape = }0.2) \\
-&H \sim \text{Gamma}(\text{shape = }100, \text{ scale = }1) \\
-&\sigma_{A_{G_i}} \sim \text{Exponential}(\text{mean = }0.1) ~\forall~i~\text{ in } 1, ..., I \\
-&\sigma_{H_{G_i}} \sim \text{Exponential}(\text{mean = }0.1) ~\forall~i~\text{ in } 1, ..., I \\
+&A \sim \text{Beta}(\text{shape1 = }15.0, \text{ shape2 = }20.0) \\
+&H \sim \text{Beta}(\text{shape1 = }25.0, \text{ shape2 = }50.0) \\
+&\sigma_{A_{G_i}} \sim \text{Exponential}(\text{mean = }40.0) ~\forall~i~\text{ in } 1, ..., I \\
+&\sigma_{H_{G_i}} \sim \text{Exponential}(\text{mean = }40.0) ~\forall~i~\text{ in } 1, ..., I \\
path: docs/model_details.md
user: login: afmagee42
body: On the other hand, changing the config to 1/40 makes the model big mad, so, maybe it's just a typo here?
created_at: 2025-03-19T17:38:04Z
author_association: NONE
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
 
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
+            )
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
+            )
+            # Draw a sample of the centered deviation from the overall average
+            # A and H for each level of each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
path: iup/models.py
user: login: afmagee42
body: This would cause us so much less pain if JAX had support for ragged arrays
created_at: 2025-03-19T18:03:40Z
author_association: NONE
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
 
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
+            )
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
+            )
+            # Draw a sample of the centered deviation from the overall average
+            # A and H for each level of each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
path: iup/models.py
user: login: afmagee42
body: (Not saying you're doing anything wrong, just that if this were square or JAX had ragged support, we could declare the shape of this as a vector of vectors and multiply the `_sigs` right in)
created_at: 2025-03-19T18:04:57Z
author_association: NONE
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
 
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
+            )
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
+            )
+            # Draw a sample of the centered deviation from the overall average
+            # A and H for each level of each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            H_s = numpyro.sample(
-                "H_s",
-                dist.TruncatedNormal(H, sig_H, low=H_low, high=H_high),
-                sample_shape=(len(np.unique(season)),),
+            H_devs = numpyro.sample(
+                "H_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            mu = A_s[season] * (elapsed**n) / (H_s[season] ** n + elapsed**n)
+            # Scale the centered deviation for each level of each grouping factor
+            # by the characteristic spread for each grouping factor
+            A_devs = A_devs * jnp.repeat(
path: iup/models.py
user: login: afmagee42
body: I'm not sure how JAX handles reassignment of variable names, but I generally try to avoid it. Here, we could just slot in the `numpyro.sample()` statement and avoid the first declaration entirely (the non-expanded `_devs` will still be tracked in the output by dint of being declared inside a sample `.sample()` statement.
created_at: 2025-03-19T19:02:46Z
author_association: NONE
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
 
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
+            )
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
+            )
+            # Draw a sample of the centered deviation from the overall average
+            # A and H for each level of each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            H_s = numpyro.sample(
-                "H_s",
-                dist.TruncatedNormal(H, sig_H, low=H_low, high=H_high),
-                sample_shape=(len(np.unique(season)),),
+            H_devs = numpyro.sample(
+                "H_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            mu = A_s[season] * (elapsed**n) / (H_s[season] ** n + elapsed**n)
+            # Scale the centered deviation for each level of each grouping factor
+            # by the characteristic spread for each grouping factor
+            A_devs = A_devs * jnp.repeat(
+                A_sigs,
+                jnp.array(num_group_levels),
+                total_repeat_length=sum(num_group_levels),
+            )
+            H_devs = H_devs * jnp.repeat(
+                H_sigs,
+                jnp.array(num_group_levels),
+                total_repeat_length=sum(num_group_levels),
+            )
+            # Across all data points, look up the A and H deviations due to the
+            # grouping factors. Sum across grouping factors and include the overall
+            # average A and H to get the final total A and H for each datum
+            A_tot = jnp.sum(A_devs[groups], axis=1) + A
+            H_tot = jnp.sum(H_devs[groups], axis=1) + H
path: iup/models.py
user: login: afmagee42
body: My numpy indexing is weak, but are we sure this is doing what we want it to do?
created_at: 2025-03-19T19:06:02Z
author_association: NONE
diff_hunk: @@ -766,19 +830,54 @@ def fit(
         """
         self.group_combos = extract_group_combos(data, groups)
 
-        if groups is None:
-            groups = []
-
-        if "season" in groups:
-            season = data["season"].to_numpy()
-            unique_seasons = np.unique(season)
-            season_to_index = {s: i for i, s in enumerate(unique_seasons)}
-            season = np.array([season_to_index[s] for s in season])
-            self.season = season
+        # Tranform the levels of the grouping factors into numeric codes
+        # Make a dictionary of dictionaries that map levels to numeric codes
+        # Save this as a model attribute, along with the number of grouping factors
+        # and the number of levels for each grouping factor. These will be
+        # useful to have saved as model attributes during prediction.
+        if groups is not None:
+            group_codes = data.select(groups)
+            self.num_group_factors = group_codes.shape[1]
+            self.value_to_index = {}
+            for i in range(self.num_group_factors):
+                col_name = group_codes.columns[i]
+                unique_values = (
+                    group_codes.select(col_name).unique().to_series().to_list()
path: iup/models.py
user: login: afmagee42
body: Is the `.to_series()` necessary when it's being `.to_list()`ed anyways?
created_at: 2025-03-19T19:16:32Z
author_association: NONE
diff_hunk: @@ -766,19 +830,54 @@ def fit(
         """
         self.group_combos = extract_group_combos(data, groups)
 
-        if groups is None:
-            groups = []
-
-        if "season" in groups:
-            season = data["season"].to_numpy()
-            unique_seasons = np.unique(season)
-            season_to_index = {s: i for i, s in enumerate(unique_seasons)}
-            season = np.array([season_to_index[s] for s in season])
-            self.season = season
+        # Tranform the levels of the grouping factors into numeric codes
+        # Make a dictionary of dictionaries that map levels to numeric codes
+        # Save this as a model attribute, along with the number of grouping factors
+        # and the number of levels for each grouping factor. These will be
+        # useful to have saved as model attributes during prediction.
path: iup/models.py
user: login: afmagee42
body: We definitely want to write "give me the groups as integers" just once, and then call it elsewhere.

It's a general notion, and you'll want this in the refactored linear incident uptake model as well as these two places, which crosses the magical number of three at which point refactoring to pull it out is advised.
created_at: 2025-03-19T19:25:30Z
author_association: NONE
diff_hunk: @@ -766,19 +830,54 @@ def fit(
         """
         self.group_combos = extract_group_combos(data, groups)
 
-        if groups is None:
-            groups = []
-
-        if "season" in groups:
-            season = data["season"].to_numpy()
-            unique_seasons = np.unique(season)
-            season_to_index = {s: i for i, s in enumerate(unique_seasons)}
-            season = np.array([season_to_index[s] for s in season])
-            self.season = season
+        # Tranform the levels of the grouping factors into numeric codes
+        # Make a dictionary of dictionaries that map levels to numeric codes
+        # Save this as a model attribute, along with the number of grouping factors
+        # and the number of levels for each grouping factor. These will be
+        # useful to have saved as model attributes during prediction.
+        if groups is not None:
+            group_codes = data.select(groups)
+            self.num_group_factors = group_codes.shape[1]
+            self.value_to_index = {}
+            for i in range(self.num_group_factors):
+                col_name = group_codes.columns[i]
+                unique_values = (
+                    group_codes.select(col_name).unique().to_series().to_list()
+                )
+                self.value_to_index[col_name] = {
+                    v: j for j, v in enumerate(unique_values)
+                }
+                group_codes = group_codes.with_columns(
+                    pl.col(col_name)
+                    .replace(self.value_to_index[col_name])
+                    .cast(pl.Int8)
path: iup/models.py
user: login: afmagee42
body: Might as well cast these to unsigned ints, since they're going to have to be positive to be used as indexes?
created_at: 2025-03-19T19:27:28Z
author_association: NONE
diff_hunk: @@ -122,9 +122,10 @@ and similarly for $H$.
 
 ```math
 \begin{align*}
-&A \sim \text{Beta}(\text{mean = }0.4, \text{ shape = }0.2) \\
-&H \sim \text{Gamma}(\text{shape = }100, \text{ scale = }1) \\
-&\sigma_{A_{G_i}} \sim \text{Exponential}(\text{mean = }0.1) ~\forall~i~\text{ in } 1, ..., I \\
-&\sigma_{H_{G_i}} \sim \text{Exponential}(\text{mean = }0.1) ~\forall~i~\text{ in } 1, ..., I \\
+&A \sim \text{Beta}(\text{shape1 = }15.0, \text{ shape2 = }20.0) \\
+&H \sim \text{Beta}(\text{shape1 = }25.0, \text{ shape2 = }50.0) \\
+&\sigma_{A_{G_i}} \sim \text{Exponential}(\text{mean = }40.0) ~\forall~i~\text{ in } 1, ..., I \\
+&\sigma_{H_{G_i}} \sim \text{Exponential}(\text{mean = }40.0) ~\forall~i~\text{ in } 1, ..., I \\
path: docs/model_details.md
user: login: eschrom
body: Just a typo in doc, now fixed. The rate should be 40 and the mean 0.025.
created_at: 2025-03-20T12:43:37Z
author_association: COLLABORATOR
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
path: iup/models.py
user: login: eschrom
body: I'll keep the comments until everyone who wants to look through this PR has had a chance, then I'll remove them.
created_at: 2025-03-20T12:45:49Z
author_association: COLLABORATOR
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
 
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
+            )
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
+            )
+            # Draw a sample of the centered deviation from the overall average
+            # A and H for each level of each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
path: iup/models.py
user: login: eschrom
body: At first I assumed JAX did support ragged arrays and got burned!
created_at: 2025-03-20T12:49:05Z
author_association: COLLABORATOR
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
 
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
+            )
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
+            )
+            # Draw a sample of the centered deviation from the overall average
+            # A and H for each level of each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            H_s = numpyro.sample(
-                "H_s",
-                dist.TruncatedNormal(H, sig_H, low=H_low, high=H_high),
-                sample_shape=(len(np.unique(season)),),
+            H_devs = numpyro.sample(
+                "H_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            mu = A_s[season] * (elapsed**n) / (H_s[season] ** n + elapsed**n)
+            # Scale the centered deviation for each level of each grouping factor
+            # by the characteristic spread for each grouping factor
+            A_devs = A_devs * jnp.repeat(
path: iup/models.py
user: login: eschrom
body: I didn't realize the non-expanded `_devs` would still be tracked, but seeing as they are named by `.sample()`, this makes total sense. Good suggestion!
created_at: 2025-03-20T12:51:06Z
author_association: COLLABORATOR
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
 
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
+            )
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
+            )
+            # Draw a sample of the centered deviation from the overall average
+            # A and H for each level of each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            H_s = numpyro.sample(
-                "H_s",
-                dist.TruncatedNormal(H, sig_H, low=H_low, high=H_high),
-                sample_shape=(len(np.unique(season)),),
+            H_devs = numpyro.sample(
+                "H_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            mu = A_s[season] * (elapsed**n) / (H_s[season] ** n + elapsed**n)
+            # Scale the centered deviation for each level of each grouping factor
+            # by the characteristic spread for each grouping factor
+            A_devs = A_devs * jnp.repeat(
+                A_sigs,
+                jnp.array(num_group_levels),
path: iup/models.py
user: login: eschrom
body: I tested this out and numpy does work in place of jax.numpy - happy to make the switch.
created_at: 2025-03-20T13:39:59Z
author_association: COLLABORATOR
diff_hunk: @@ -650,58 +651,96 @@ def __init__(self, seed: int):
     @staticmethod
     def hill(
         elapsed,
-        season=None,
         cum_uptake=None,
-        n_low=1.0,
-        n_high=5.0,
-        A_low=0.0,
-        A_high=1.0,
-        A_sig=1.0,
-        H_low=10.0,
-        H_high=180.0,
-        H_sig=1.0,
-        sig_mn=1.0,
+        std_dev=None,
+        groups=None,
+        num_group_factors=0,
+        num_group_levels=[0],
+        A_shape1=15.0,
+        A_shape2=20.0,
+        A_sig=40.0,
+        H_shape1=25.0,
+        H_shape2=50.0,
+        H_sig=40.0,
+        n_shape=20.0,
+        n_rate=5.0,
     ):
         """
         Fit a Hill model on training data.
 
         Parameters
         elapsed: np.array
-            days elapsed since the season start for each data point
-        season: np.array | None
-            numeric code for the season each data point belongs to
+            column of days elapsed since the season start for each data point
         cum_uptake: np.array | None
-            cumulative uptake measured at each data point
+            column of cumulative uptake measured at each data point
+        std_dev: np.array | None
+            column of standard deviations in cumulative uptake estimate
+        groups: np.array | None
+            numeric codes for groups: row = data point, col = grouping factor
+        num_group_factors: Int
+            number of grouping factors
+        num_group_levels: List[Int,]
+            number of unique levels of each grouping factor
         other parameters: float
-            means and standard deviations to specify the prior distributions
+            parameters to specify the prior distributions
 
         Returns
         Nothing
 
         Details
         Provides the model structure and priors for a Hill model.
         """
-        n = numpyro.sample("n", dist.Uniform(n_low, n_high))
-        A = numpyro.sample("A", dist.Uniform(A_low, A_high))
-        sig_A = numpyro.sample("sig_A", dist.Exponential(A_sig))
-        H = numpyro.sample("H", dist.Uniform(H_low, H_high))
-        sig_H = numpyro.sample("sig_H", dist.Exponential(H_sig))
-        if season is not None:
-            A_s = numpyro.sample(
-                "A_s",
-                dist.TruncatedNormal(A, sig_A, low=A_low, high=A_high),
-                sample_shape=(len(np.unique(season)),),
+        # Sample the overall average value for each Hill function parameter
+        A = numpyro.sample("A", dist.Beta(A_shape1, A_shape2))
+        H = numpyro.sample("H", dist.Beta(H_shape1, H_shape2))
+        n = numpyro.sample("n", dist.Gamma(n_shape, n_rate))
+        # If grouping factors are given, find the specific A and H for each datum
+        if groups is not None:
+            # Draw a sample of the spread among levels for each grouping factor
+            A_sigs = numpyro.sample(
+                "A_sigs", dist.Exponential(A_sig), sample_shape=(num_group_factors,)
+            )
+            H_sigs = numpyro.sample(
+                "H_sigs", dist.Exponential(H_sig), sample_shape=(num_group_factors,)
+            )
+            # Draw a sample of the centered deviation from the overall average
+            # A and H for each level of each grouping factor
+            A_devs = numpyro.sample(
+                "A_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            H_s = numpyro.sample(
-                "H_s",
-                dist.TruncatedNormal(H, sig_H, low=H_low, high=H_high),
-                sample_shape=(len(np.unique(season)),),
+            H_devs = numpyro.sample(
+                "H_devs", dist.Normal(0, 1), sample_shape=(sum(num_group_levels),)
             )
-            mu = A_s[season] * (elapsed**n) / (H_s[season] ** n + elapsed**n)
+            # Scale the centered deviation for each level of each grouping factor
+            # by the characteristic spread for each grouping factor
+            A_devs = A_devs * jnp.repeat(
+                A_sigs,
+                jnp.array(num_group_levels),
+                total_repeat_length=sum(num_group_levels),
+            )
+            H_devs = H_devs * jnp.repeat(
+                H_sigs,
+                jnp.array(num_group_levels),
+                total_repeat_length=sum(num_group_levels),
+            )
+            # Across all data points, look up the A and H deviations due to the
+            # grouping factors. Sum across grouping factors and include the overall
+            # average A and H to get the final total A and H for each datum
+            A_tot = jnp.sum(A_devs[groups], axis=1) + A
+            H_tot = jnp.sum(H_devs[groups], axis=1) + H
path: iup/models.py
user: login: eschrom
body: Yes sir. It took a little experimenting to convince myself, but in the end I was impressed by how convenient this turns out to be. Here is a little demo:
```
import numpy as np
A_devs = np.array([-0.2, -0.1, 0.1, 0.2])
A_devs
#     array([-0.2, -0.1,  0.1,  0.2])
groups = np.array([[0,2],[0,3],[1,2],[1,3]])
groups
#     array([[0, 2],
#            [0, 3],
#            [1, 2],
#            [1, 3]])
A_devs[groups]
#     array([[-0.2,  0.1],
#            [-0.2,  0.2],
#            [-0.1,  0.1],
#            [-0.1,  0.2]])
np.sum(A_devs[groups], axis=1) + 4
#     array([3.9, 4. , 4. , 4.1])
```
created_at: 2025-03-20T13:49:47Z
author_association: COLLABORATOR
diff_hunk: @@ -766,19 +830,54 @@ def fit(
         """
         self.group_combos = extract_group_combos(data, groups)
 
-        if groups is None:
-            groups = []
-
-        if "season" in groups:
-            season = data["season"].to_numpy()
-            unique_seasons = np.unique(season)
-            season_to_index = {s: i for i, s in enumerate(unique_seasons)}
-            season = np.array([season_to_index[s] for s in season])
-            self.season = season
+        # Tranform the levels of the grouping factors into numeric codes
+        # Make a dictionary of dictionaries that map levels to numeric codes
+        # Save this as a model attribute, along with the number of grouping factors
+        # and the number of levels for each grouping factor. These will be
+        # useful to have saved as model attributes during prediction.
+        if groups is not None:
+            group_codes = data.select(groups)
+            self.num_group_factors = group_codes.shape[1]
+            self.value_to_index = {}
+            for i in range(self.num_group_factors):
+                col_name = group_codes.columns[i]
+                unique_values = (
+                    group_codes.select(col_name).unique().to_series().to_list()
path: iup/models.py
user: login: eschrom
body: Weirdly yes - you cannot directly `.to_list()` a polars data frame column.
created_at: 2025-03-20T13:53:14Z
author_association: COLLABORATOR
diff_hunk: @@ -766,19 +830,54 @@ def fit(
         """
         self.group_combos = extract_group_combos(data, groups)
 
-        if groups is None:
-            groups = []
-
-        if "season" in groups:
-            season = data["season"].to_numpy()
-            unique_seasons = np.unique(season)
-            season_to_index = {s: i for i, s in enumerate(unique_seasons)}
-            season = np.array([season_to_index[s] for s in season])
-            self.season = season
+        # Tranform the levels of the grouping factors into numeric codes
+        # Make a dictionary of dictionaries that map levels to numeric codes
+        # Save this as a model attribute, along with the number of grouping factors
+        # and the number of levels for each grouping factor. These will be
+        # useful to have saved as model attributes during prediction.
+        if groups is not None:
+            group_codes = data.select(groups)
+            self.num_group_factors = group_codes.shape[1]
+            self.value_to_index = {}
+            for i in range(self.num_group_factors):
+                col_name = group_codes.columns[i]
+                unique_values = (
+                    group_codes.select(col_name).unique().to_series().to_list()
+                )
+                self.value_to_index[col_name] = {
+                    v: j for j, v in enumerate(unique_values)
+                }
+                group_codes = group_codes.with_columns(
+                    pl.col(col_name)
+                    .replace(self.value_to_index[col_name])
+                    .cast(pl.Int8)
path: iup/models.py
user: login: eschrom
body: Fair enough - done.
created_at: 2025-03-20T13:56:19Z
author_association: COLLABORATOR
diff_hunk: @@ -766,19 +830,54 @@ def fit(
         """
         self.group_combos = extract_group_combos(data, groups)
 
-        if groups is None:
-            groups = []
-
-        if "season" in groups:
-            season = data["season"].to_numpy()
-            unique_seasons = np.unique(season)
-            season_to_index = {s: i for i, s in enumerate(unique_seasons)}
-            season = np.array([season_to_index[s] for s in season])
-            self.season = season
+        # Tranform the levels of the grouping factors into numeric codes
+        # Make a dictionary of dictionaries that map levels to numeric codes
+        # Save this as a model attribute, along with the number of grouping factors
+        # and the number of levels for each grouping factor. These will be
+        # useful to have saved as model attributes during prediction.
+        if groups is not None:
+            group_codes = data.select(groups)
+            self.num_group_factors = group_codes.shape[1]
+            self.value_to_index = {}
+            for i in range(self.num_group_factors):
+                col_name = group_codes.columns[i]
+                unique_values = (
+                    group_codes.select(col_name).unique().to_series().to_list()
path: iup/models.py
user: login: afmagee42
body: Oh, right, you're `.select`ing it, and select returns a dataframe. You could pull it straight to series instead:
```
group_codes[col_name].unique().to_list()
```
created_at: 2025-03-20T13:58:51Z
author_association: NONE
diff_hunk: @@ -766,19 +830,54 @@ def fit(
         """
         self.group_combos = extract_group_combos(data, groups)
 
-        if groups is None:
-            groups = []
-
-        if "season" in groups:
-            season = data["season"].to_numpy()
-            unique_seasons = np.unique(season)
-            season_to_index = {s: i for i, s in enumerate(unique_seasons)}
-            season = np.array([season_to_index[s] for s in season])
-            self.season = season
+        # Tranform the levels of the grouping factors into numeric codes
+        # Make a dictionary of dictionaries that map levels to numeric codes
+        # Save this as a model attribute, along with the number of grouping factors
+        # and the number of levels for each grouping factor. These will be
+        # useful to have saved as model attributes during prediction.
path: iup/models.py
user: login: eschrom
body: Good call - I just made these into utils and it looks much nicer.
created_at: 2025-03-20T14:38:19Z
author_association: COLLABORATOR

PR comments:
user: login: eschrom
created_at: 2025-03-18T21:37:40Z
body: The refactored Hill model is ready for review! See the [model details doc](https://github.com/CDCgov/cfa-immunization-uptake-projection/blob/main/docs/model_details.md) for a description of the intended model structure (and note that there are updates to that doc in this PR).

I have tested the model using season and state as grouping factors: it runs successfully and produces group-specific forecasts in the right ballpark!

That said, I would still like to direct your attention to several caveats and questions:
- A few changes in this PR (e.g. adjustments to scaffold building) might cause the linear incident uptake model to break. Let's not worry about that, since the linear incident uptake model will be replaced by a refined autoregressive model soon anyway.
- The refactored Hill model accommodates arbitrary grouping factors (each with arbitrary levels) without using plates! Is my approach sensible and general? If so, should we document this elsewhere for use in other projects?
- My approach requires a little bookkeeping of numeric codes for group levels, to insure that group levels are interpretted identically between the tasks of fitting and prediction. This bookkeeping currently appears in `HillModel.fit()` and `HillModel.predict()`. Should this bookkeeping be outsourced to separate, modular functions, perhaps in `utils`?
- The accommodation of arbitrary grouping factors in this PR will cause downstream evaluation and plotting to break. @Fuhan-Yang, would you like to address that in a separate future PR, or would you like me to try to address that here?
- Despite rescaling some parameters to keep them all within the same order of magnitude and choosing fairly targeted priors, MCMC is still _very_ slow (~1-2 it/s, for both warmup and sampling). Have I missed an obvious inefficiency or misspecification? Or do I just have unrealistic expectations of speed when 2 different parameters take on different values for each of 14 seasons and 54 states?
- The standard deviation for the truncated normal likelihood cannot be 0. This is a problem in two contexts: 1) a handful of the data points report an observed standard deviation of zero, and 2) there is no observed standard deviation when performing prediction. I didn't know what to do in either case ... so I just shamefully declared a standard deviation of 0.0001 in both cases. What are the right ways to handle these cases?
author_association: COLLABORATOR
user: login: eschrom
created_at: 2025-03-19T13:45:37Z
body: Here is a diagram to explain the bookkeeping on the levels of each grouping factor:
![image](https://github.com/user-attachments/assets/ba8f81b6-80b3-4486-a582-ecd0f8c7e5a5)



author_association: COLLABORATOR
user: login: afmagee42
created_at: 2025-03-19T17:19:18Z
body: > * Despite rescaling some parameters to keep them all within the same order of magnitude and choosing fairly targeted priors, MCMC is still _very_ slow (~1-2 it/s, for both warmup and sampling). Have I missed an obvious inefficiency or misspecification? Or do I just have unrealistic expectations of speed when 2 different parameters take on different values for each of 14 seasons and 54 states?

There are many potential causes of slow MCMC. I'm thinking this might be a misspecification problem, because running this model, I'm seeing this:

warmup:   5%|▍        | 95/2000 [00:15<05:09,  1.16it/s, 1023 steps of size 4.36e-04. acc. prob=0.74]

This tells me that the MCMC is very unhappy: its acceptance probability is nowhere near where it wants to be and it's taking a ton of steps (presumably to try to get there).

(Gelman's) Folk Theorem interpretation would be that the MCMC is trying to tell us the model (encompassing everything from the choice of fixed hyperparameters through prior distribution families to the likelihood) is not a good fit.

Playing this out a bit farther, I'm also seeing the reported steps taken swing wildly from iteration to iteration, which strikes me as odd.
author_association: NONE
user: login: afmagee42
created_at: 2025-03-19T17:47:41Z
body: > * The standard deviation for the truncated normal likelihood cannot be 0. This is a problem in two contexts: 1) a handful of the data points report an observed standard deviation of zero, and 2) there is no observed standard deviation when performing prediction. I didn't know what to do in either case ... so I just shamefully declared a standard deviation of 0.0001 in both cases. What are the right ways to handle these cases?

I make no claim on "right," but "more defensible" maybe.
- If the SDs don't show any obvious trends throughout a season, you could:
  - Replacing small values with
    - Some small quantile of the non-zero SDs
    - Assume that 0 is a rounding error and that the variance is basically the largest possible value that would get rounded to 0. So, if all SDs are reported to the 4th decimal place, take 0 to mean 0.00005.
  - Predict using the (square root  of the) mean (variance (variances are additive, not SDs)) would seem better. Or the median (which you can take on the variance scale, quantiles are preserved by monotonic transforms).
- Trend or not, you can also make the model do the work: model the SDs (well, model the _variances_ and then transform them) somehow. That will give us something for predictions too.
author_association: NONE
user: login: eschrom
created_at: 2025-03-19T21:46:30Z
body: > (Gelman's) Folk Theorem interpretation would be that the MCMC is trying to tell us the model (encompassing everything from the choice of fixed hyperparameters through prior distribution families to the likelihood) is not a good fit.

I ran the model without any grouping factors, to fit the three Hill function parameters sans hierarchy. Here are the 2022/2023 data vs. the model trained on all seasons from 2009/2010 through 2022/2023:

![image](https://github.com/user-attachments/assets/3d892d72-ad27-4036-a969-fca2d26a8e31)

One reason for the dodgy fit is because the curve represents the average of all seasons, whereas the points represent just one season. But another reason is that after ~February, the Hill function asymptotes whereas real uptake continues to creep up. 

This could be a clue that the Hill function is just not a good fit. I doubt that's the whole story behind the super-slow MCMC with grouping factors included (for the overall model, MCMC hit a crisp 653it/s), but I thought it worth documenting here.
author_association: COLLABORATOR
user: login: afmagee42
created_at: 2025-03-20T13:12:04Z
body: > This could be a clue that the Hill function is just not a good fit. I doubt that's the whole story behind the super-slow MCMC with grouping factors included (for the overall model, MCMC hit a crisp 653it/s), but I thought it worth documenting here.

Agreed, there are degrees of badness and plenty bad models fit easily enough. This does suggest some experiments, though:
- Season-only grouping (pooled data across all states for all seasons or only one state)
- State-only grouping (for one season or all seasons pooled)
- Half the states, half the seasons (or some other intermediate form)
author_association: NONE
user: login: afmagee42
created_at: 2025-03-20T13:48:24Z
body: I'm also, perhaps belatedly since it was staring me in the face, wondering if there's a possibility this is related to the problem that required such a drastic model rewrite in [linmod](https://github.com/CDCgov/cfa-viral-lineage-model/pull/66). 

To try to condense that essay down at lease slightly:
- In linmod there was a model where `parameter[i] = overall_location + deviation[i]`. Here it's `parameter[i,j] = overall_location + deviation1[i] + deviation2[j]`.
- With only one "observation" of `parameter[i]` (or `[i,j]`) you can't do this in purely likelihood-based inference, it's not identifiable (is the overall location 10, or are all the deviations 10?). 
- Such a model is perfectly valid in Bayesian contexts because the prior ensures there is identifiability (the Normal(0, s^2 ) prior on the deviations means that we won't set all deviations to 10 when we could simply set the overall location to 10). 
- _But_ the parameter exhibit strong correlation in the posterior and this makes the geometry poorly-suited to handling via a diagonal mass matrix. _And_ the dense mass matrix was too much to learn. It worked about 3 runs in 4, but the failure mode in the 4th was even worse than the correlations.

If this _is_ the same problem, there are two solutions:
- Give in and reduce the number of free parameters. Either:
  - Surrender to asymmetry and choose a "reference" group combo which does not get a deviation.
  - Remove the overall `A` and `H`.
- Try a block-diagonal mass matrix. Specify that `A` and the `A_devs` are a block, as are `H` and the `H_devs`.

The block-diagonal approach _might_ be possible without the loop required in `linmod`. But it might still be a bit too much to learn. Back of the envelope math says it still requires learning almost half the off-diagonal elements of the mass matrix. (Taking `n = sum(num_group_levels)` and `g = len(num_group_levels)`, the block-diagonal matrix has about $n^2$ non-zero off-diagonal elements to learn, while the full-dense has about $0.5 (2n + 2g + 3)^2$.) If it is too much, you would probably need to re-write the model so that you can specify, separately, correlations within grouping factors. But since the sizes of the grouping variables are so wildly uneven, this isn't going to be all that big of a reduction...
author_association: NONE
user: login: eschrom
created_at: 2025-03-20T15:17:11Z
body: > * If the SDs don't show any obvious trends throughout a season, you could:
>   
>   * Replacing small values with
>     
>     * Some small quantile of the non-zero SDs
>     * Assume that 0 is a rounding error and that the variance is basically the largest possible value that would get rounded to 0. So, if all SDs are reported to the 4th decimal place, take 0 to mean 0.00005.
>   * Predict using the (square root  of the) mean (variance (variances are additive, not SDs)) would seem better. Or the median (which you can take on the variance scale, quantiles are preserved by monotonic transforms).
> * Trend or not, you can also make the model do the work: model the SDs (well, model the _variances_ and then transform them) somehow. That will give us something for predictions too.

SDs (well, more precisely, the mysterious 95% confidence intervals) in the data are rounded to the nearest 0.000001 (millionth), so I've replaced zeros with 0.0000005, preferring not to add complication to the model until the current slowness is solved.
author_association: COLLABORATOR
user: login: eschrom
created_at: 2025-03-20T16:04:02Z
body: Good news! And an embarrassing admission...

My replacement of the SDs with a small number was only working if no observed SDs were given at all. For the >80 instances of SD = 0 in the observed data, SD was not being replaced by a small number. Rather than complaining, numpyro was just grinding through an impossible search for parameter values that could hit all these >80 "error-less" observations perfectly.

This was the cause of the slow speed. Now that I have corrected my approach to replace these >80 SDs of 0 with a small number, numpyro rips through 1000 warmups + 1000 samples in ~5s, even with 14 seasons and 54 states.
author_association: COLLABORATOR
user: login: eschrom
created_at: 2025-03-20T16:13:35Z
body: Thanks @afmagee42 for your review and insightful discussion! All the points I originally drew our attention to are now solved, except that plotting/eval must be updated to account for grouping factors. @Fuhan-Yang and I can tackle that in a future PR. Thanks again!
author_association: COLLABORATOR