Align observation process API with user's mental model

[cdc-mitzimorris]

Problem

The current observation process API requires users to manage timeline offsets between their data and the model's internal representation. This creates confusion and potential for off-by-one errors.

User's mental model

• "I have 90 days of hospital data, days 0-89"
• "I provide 90 observations"
• "I get back 90 days of inferred infections"

Current implementation

The model internally requires an initialization period (n_init) to satisfy the renewal equation and delay convolutions:

• Infections array: 133 days total (43 init + 90 user days)
• Predicted counts: First 42 days are NaN (insufficient convolution history), giving 91 valid days
• User's observations: 90 days, mapping to model days 43-132

Where: n_init = max(len(generation_interval), len(delay_pmf_1), len(delay_pmf_2), ...)

The issue

NaN filtering alone gives 91 valid days (42-132), but user has 90 observations (43-132). The observation process needs to know where user's data starts — not just filter NaN.

Proposed solution

Add an n_init parameter to observation sample() methods that specifies where the user's observation period begins.

Counts.sample() revised signature

def sample(
    self,
    infections: ArrayLike,
    obs: ArrayLike | None = None,
    n_init: int | None = None,
) -> ObservationSample:
    predicted_counts = self._predicted_obs(infections)
    self._deterministic("predicted", predicted_counts)

    if n_init is None:
        n_init = self.min_valid_index

    predicted_obs = predicted_counts[n_init:]

    if obs is not None and len(obs) != len(predicted_obs):
        raise ValueError(
            f"obs length {len(obs)} != predicted length {len(predicted_obs)}. "
            f"Expected {len(predicted_obs)} observations for days {n_init}+."
        )

    observed = self.noise.sample(
        name=self._sample_site_name("obs"),
        predicted=predicted_obs,
        obs=obs,
    )

    return ObservationSample(observed=observed, predicted=predicted_counts)

Usage standalone (no ModelBuilder)

• User has 90 days of data
• Counts uses default n_init = min_valid_index = 42
• call: counts.sample(infections=infections_132, obs=obs_90)
• return: predicted_counts[42:] has length 90, matches obs

Usage with ModelBuilder

• Model computes n_init = 43 (max of all lookbacks)
• Model passes n_init to observation
• call: counts.sample(infections=infections_133, obs=obs_90, n_init=43)
• return: predicted_counts[43:] has length 90, matches obs

Summary of changes

• Counts.sample(): Add n_init parameter, slice predicted_counts[n_init:]
• CountsBySubpop.sample(): Add n_init parameter, offset times internally
• Measurements.sample(): Add n_init parameter, offset times internally
• MultiSignalModel.sample(): Pass n_init to each observation's sample()
• User API: User provides observations in natural coordinates (day 0 = first data point)

User experience after changes

# User's data
hospital_obs = jnp.array([...])  # 90 days, indexed 0-89

# With ModelBuilder (recommended)
model.fit(
    n_days=90,
    hospital={"obs": hospital_obs},  # No offset needed
)

# Standalone Counts (advanced)
counts.sample(
    infections=full_infections,
    obs=hospital_obs,
    n_init=43,
)

Related

• PR Add generic observation processes which combine the convolution with the noise model. #644: Generic observation processes

[dylanhmorris]

tagging @damonbayer for awareness / discussion

[dylanhmorris]

Thanks for this, @cdc-mitzimorris! Splitting this into a comment on the specific extra day issue (which I think we can avoid separately) and a comment on the more general proposal.

Extra day question

Good point about the potential for an extra day (or more generally time unit). I think we can and should avoid this problem with how we compute the minimum n_init. You wrote:
n_init = max(len(generation_interval), len(delay_pmf_1), len(delay_pmf_2), ...)

But it is safe to use:
``n_init = max(len(generation_interval), len(delay_pmf_1) - 1, len(delay_pmf_2) - 1, ...)`

This should prevent the extra day, no?

Details

What we want to calculate is
n_init = max({max_possible_generation_length}, {max_possible_observation_delay_1}, {max_possible_observation_delay_2}, ...)

The -1s for the observation delays stem from the fact that we allow 0-day delays from infection to observation but not from infection to transmission. Lots of discussion about what the least surprising and most ergonomic convention in renewal models is. See e.g. here. Something to think about/discuss in #490.

But for now PyRenew uses the 1-indexed generation interval PMF/0-indexed delay PMF convention.

[dylanhmorris]

General question: how to handle alignment between infections and observations in observation classes

That said, resolving the extra day issue doesn't fully answer the design question you've outlined. I think that question boils down to "how do we handle alignment between infections and observations in observation classes like Counts?"

I agree that the goal should be "make it as easy and error-proof as possible for a user to line up inferred infections, predicted observations, and actual observations".

As in your API example, we can create utilities at the model / model_builder level to handle alignment for the user:

model.fit(
   n_days=90,
   hospital={"obs": hospital_obs},  # No offset needed
)

The question is how to handle it at the observation level. I think there are two good options

Option 1: ObservationProcess.sample() takes an offset argument

This is the proposal above. I'd suggest naming it just offset rather than n_init for generality. This approach is similar to what's currently done in Pyrenew-HEW using the offsets returned by compute_delay_ascertained_incidence.

Option 2: ObservationProcess.sample() returns a prediction vectors and accepts observation vectors that are necessarily of the same length as the infections vector (but padded).

This was a proposal we discussed face to face earlier this week. Would require careful padding/unpadding on behalf of the user. I think it has the following pros and cons relative to option 1:

Pros

• Each ObservationProcess has a clear internal clock. That is, it uses the same indexing scheme for all timeseries vectors it gets passed.
• Removes need for switches around sparse versus dense observation vectors within ObservationProcesses (e.g. a times argument). Sparsity handled upstream.

Cons

• Somewhat trickier to implement safely.
• Extraneous and potentially inefficient padding and unpadding operations.
• Need a padding value. nan is probably the best choice, but that means we can't use integer arrays even where that would make sense (since nan is a floating-point-only concept in JAX)

Pro or con depending on your perspective

Requires/enables a common way of representing the following in prediction and observation vectors

• Not-predicted-and-not-observed days. That is, the first offset or n_init days of the prediction vector and the observation vector).
• Predicted-but-not-observed days (e.g. a date for which we have missing or excluded outlier value in our hospital admissions timeseries)
• observed-but-not-predicted quantities (e.g. a date prior to the first predicted admissions value for which we have a reported admissions count).

As @damonbayer and I mentioned the other day, being able to represent predicted-but-not-observed quantities would be super nice to have. Being able to represent the other two things does create some bug opportunities, and the benefits are less obvious.

Thoughts? I tend to favor trying Option 2 because I think the benefits will outweigh the costs, but am very persuadable.

[cdc-mitzimorris]

Option 2: ObservationProcess.sample() returns a prediction vectors and accepts observation vectors that are necessarily of the same length as the infections vector (but padded).

will be implemented as part of PR #650

[dylanhmorris]

Closed by #650