Stochastic transitions

[RomeshA]

Creating new PR from #474 (since develop was deleted)

[RomeshA]

@claude review

[claude]

Given the inline findings, this PR needs human review before merging. Beyond those, I also checked the flagged candidate issue in JunctionCompartment.balance() where outflow_fractions /= total_outflow divides a plain Python list in place — this does not raise a TypeError because total_outflow is a numpy.float64 (from summing values pulled out of a numpy array), so numpy's reflected division kicks in and converts the list to an ndarray. Confirmed this behavior directly in Python, so it is not a real bug.

Extended reasoning...

Verified the ruled-out candidate (JunctionCompartment.balance list in-place division) by reproducing it directly: x = [np.float64(1.0), np.float64(2.0)]; x /= sum(x) succeeds because Python's /= falls back to numpy.float64.__rtruediv__, which treats the list as array-like and returns an ndarray — whereas the same operation with a plain Python float or int (e.g. an empty outlinks list, where sum([]) is a plain int) would raise. Since link.parameter.vals[ti] always returns numpy floats in the normal case, the existing code path is safe. This does not affect the verdict on the reported bugs, which stand as real issues requiring a fix.

[claude]

🔴 In ResidualJunctionCompartment.balance() (atomica/model.py), the scalar branch used when a residual junction has no duration group (the common case) never checks self.multinomial and always computes deterministic fractional outflows (net_inflow * frac). Only the vector/duration-group branch below it was updated to sample via self.rng_sampler.multinomial. In a stochastic run this junction silently receives non-integer link values that get truncated downstream, causing quiet population drift rather than a loud failure (contrast with TimedCompartment.resolve_outflows, which asserts self.multinomial == False for the case it doesn't support).

Extended reasoning...

The bug: ResidualJunctionCompartment.balance() has two full branches, chosen based on self.duration_group:

1. The scalar branch, entered when not self.duration_group (this is described in the method's own comment as "the common design pattern for cohort models where there are many junctions and no durations"). This branch computes outflow_fractions = [link.parameter.vals[ti] ...], normalizes them, then does outflows = [net_inflow * frac for frac in outflow_fractions] and assigns the (possibly fractional) result straight to link.vals[ti]. There is no reference to self.multinomial anywhere in this branch.
2. The vector branch, entered when self.duration_group is set, which this PR updated to add self.rng_sampler.multinomial(...) sampling, plus special-casing for a residual link, plus a sanity check that net_inflow is integer-valued when self.multinomial is on.

Population.build() passes multinomial=self.multinomial into every ResidualJunctionCompartment it constructs, so the compartment itself is fully aware it might be running in stochastic mode -- but the scalar branch simply never consults that flag.

Why this matters / how it manifests: A residual junction without a duration group is the ordinary case for junctions in cohort models (per the method's own docstring/comment). So when a user runs a model with Model.multinomial=True (the feature this PR adds), essentially every "plain" junction with a residual/default outflow link will go through the unguarded scalar branch and emit fractional flow values, even though the entire point of multinomial mode is that all flows should be sampled, integer counts of people.

Why nothing catches it: The scalar branch still conserves total flow (residual = net_inflow - sum(outflows)), so it doesn't crash and doesn't produce NaNs or negative values -- it just produces the deterministic (fractional) split instead of a stochastic (integer) one. Downstream, ordinary Compartment.resolve_outflows() in multinomial mode calls self.rng_sampler.multinomial(self[ti], probs), and numpy's Generator.multinomial silently accepts and truncates a non-integer n rather than raising -- verified directly:

rng.multinomial(3.7, [0.5, 0.5])  ->  array([2, 1])   # truncates 3.7 to 3, no warning/error

So the model keeps running to completion, but population conservation quietly degrades over time and the junction is never actually sampled stochastically, defeating the purpose of the new feature for a very common topology. This is exactly the kind of latent numerical error that a full test suite exercising every demo framework under multinomial=True (as test_functional_all_framework_databook does) will not catch, because there's no exception and no obviously-wrong output value to assert against.

Contrast with existing precedent in this same PR: TimedCompartment.resolve_outflows() handles an analogous "this code path doesn't support stochastic mode yet" situation by failing loudly: assert self.multinomial == False, "resolve_outflows needs to be updated to work with stochastic transitions". The scalar branch of ResidualJunctionCompartment.balance() has no equivalent guard, so instead of failing loudly on an unsupported configuration it fails silently on a supported and common one.

Step-by-step proof:

1. Build a model with multinomial=True, containing a ResidualJunctionCompartment with no duration_group, one parameterized outflow link with link.parameter.vals[ti] = 0.3 and one residual (no-parameter) outflow link.
2. Suppose net_inflow = 10 (an integer count of people) at timestep ti.
3. Execution enters if not self.duration_group: (the scalar branch) because no duration group is set.
4. outflow_fractions = [0.3, 0.0], total_outflow = 0.3 (not greater than 1, so no rescaling).
5. outflows = [10*0.3, 10*0.0] = [3.0, 0.0]; residual = 10 - 3.0 = 7.0.
6. The parameterized link gets link.vals[ti] = 3.0 and the residual link gets link.vals[ti] = 7.0 -- both integers here only by coincidence of the chosen input. As soon as the parameter fraction doesn't divide evenly into net_inflow (e.g. 0.33 instead of 0.3), outflows[0] = 3.3, a fractional value, is assigned directly with no sampling and no integer check, unlike the vector branch which would have raised on a non-integer net_inflow or drawn from rng_sampler.multinomial.
7. That fractional link.vals[ti] becomes the outflow into a downstream compartment; on the next timestep that compartment's Compartment.resolve_outflows() (multinomial mode) calls self.rng_sampler.multinomial(self[ti], probs) with a fractional self[ti], which numpy silently truncates -- no exception, just a quietly wrong (slightly reduced) population count propagating forward.

Suggested fix: Either implement multinomial sampling in the scalar branch analogous to the vector branch (drawing outflow counts from self.rng_sampler.multinomial(net_inflow, outflow_fractions) and assigning the residual to the residual link the same way the vector branch does), or, at minimum, add an explicit guard (e.g. assert not self.multinomial) so unsupported configurations fail loudly instead of silently drifting.

[claude]

🔴 In Program.sample() (atomica/programs.py:1237-1241), only capacity_constraint is sampled with the seeded rng_sampler; spend_data, unit_cost, saturation, and coverage are called with .sample(constant) and no rng_sampler, so they silently fall back to the unseeded global np.random inside TimeSeries.sample(). This defeats the reproducibility goal of this PR for any program with uncertainty on those fields — fix by passing rng_sampler to all four calls, matching the capacity_constraint call.

Extended reasoning...

The bug: Program.sample() in atomica/programs.py is supposed to perturb all five of a program's uncertain TimeSeries fields using the caller-supplied rng_sampler so that seeded runs are reproducible. However, only one of the five calls actually forwards rng_sampler:

self.spend_data = self.spend_data.sample(constant)
self.unit_cost = self.unit_cost.sample(constant)
self.capacity_constraint = self.capacity_constraint.sample(constant, rng_sampler)
self.saturation = self.saturation.sample(constant)
self.coverage = self.coverage.sample(constant)

Why it manifests: Looking at TimeSeries.sample() in atomica/utils.py, the signature is sample(self, constant=True, rng_sampler=None), and when rng_sampler is None it falls back with rng_sampler = np.random (the global, unseeded module). So for spend_data, unit_cost, saturation, and coverage, the perturbation draws from whatever global numpy random state happens to exist at call time, rather than the seeded generator that was threaded all the way down from Project.run_sampled_sims(rand_seed=...) → _run_sampled_sim → ProgramSet.sample(rng_sampler=...) → Program.sample(constant, rng_sampler=...).

Why nothing else catches this: ProgramSet.sample() correctly forwards rng_sampler into prog.sample(constant, rng_sampler=rng_sampler), and Covout.sample() also correctly uses rng_sampler (falling back to np.random only when none is supplied at its own call site, which is expected/documented default behavior at the top-level entry point). The omission is localized specifically to four of the five TimeSeries.sample() calls inside Program.sample() — there's no other code path that re-injects the seed for these fields.

Impact: For any program whose spend_data, unit_cost, saturation, or coverage has an entered sigma (uncertainty), the sampled program values used during a seeded run_sampled_sims will not be reproducible run-to-run, even though the user supplied rand_seed. This silently undermines the core reproducibility guarantee that this PR is introducing — it will not raise an error or crash, so this is likely to go unnoticed until someone tries to reproduce a specific seeded stochastic result and finds it differs.

Step-by-step proof:

1. User calls proj.run_sampled_sims(parset, progset, ..., rand_seed=42, n_samples=2).
2. Internally, for each sample, a seeded rng_attempt = np.random.default_rng(seed=rand_seeds[attempts]) is created in _run_sampled_sim (atomica/project.py), and passed as sampled_progset = progset.sample(rng_sampler=rng_attempt).
3. Inside ProgramSet.sample(), this correctly reaches each Program via prog.sample(constant, rng_sampler=rng_attempt).
4. Inside Program.sample(), self.capacity_constraint.sample(constant, rng_sampler) correctly uses rng_attempt, so its perturbation is deterministic given the seed.
5. But self.spend_data.sample(constant) is called with no rng_sampler argument, so inside TimeSeries.sample(), rng_sampler is None and it sets rng_sampler = np.random — the global module-level state. If a program has spend_data.sigma set, its perturbed value depends on the global numpy random state, which varies depending on what other code executed before this call (e.g., other samples/attempts, other library calls) — not solely on rand_seed=42.
6. Running the exact same run_sampled_sims(rand_seed=42, ...) call tw, in a fresh process each time will reproduce the capacity_constraint perturbations identically, but the spend_data/unit_cost/saturation/coverage perturbations will differ between the two runs whenever any global numpy randomness has been consumed differently beforehand (e.g. due to other sampled attempts, other programs, or any other numpy-using code executed earlier in the process) — breaking the reproducibility contract for exactly the fields most likely to carry uncertainty (spend and unit cost).

Fix: Pass rng_sampler through to all four remaining calls, e.g.:

self.spend_data = self.spend_data.sample(constant, rng_sampler)
self.unit_cost = self.unit_cost.sample(constant, rng_sampler)
self.capacity_constraint = self.capacity_constraint.sample(constant, rng_sampler)
self.saturation = self.saturation.sample(constant, rng_sampler)
self.coverage = self.coverage.sample(constant, rng_sampler)

[claude]

🔴 In TimeDependentConnections._write_pop_matrix (atomica/excel.py:866), the boolean gating cell is now computed as "Y" if (value.vals or value.assumption) else "N", which uses truthiness rather than presence-of-data. A transfer/interaction entry that has an explicit constant value of exactly 0.0 (empty vals, assumption == 0.0) will be written as N instead of Y, causing that entry to be silently dropped when the databook is re-read (e.g. after recalculation on open, or on any write→read round trip). The fix is to check presence of data (e.g. value.has_data) rather than truthiness, mirroring TimeSeries.has_data (assumption is not None).

Extended reasoning...

The bug: atomica/excel.py:866 was changed from value = "Y" if value else "N" to value = "Y" if (value.vals or value.assumption) else "N". The old code was actually a no-op bug of its own — TimeSeries defines __slots__ with no __bool__/__len__, so any TimeSeries instance is always truthy, meaning every cell was always written "Y". The new code fixes that non-writing issue but introduces a new one: it now uses truthiness of value.vals and value.assumption, rather than checking whether data is actually present. TimeSeries.has_data (the canonical way this check is done elsewhere in the codebase) explicitly uses self.assumption is not None, precisely because 0.0 is a legitimate value that must not be treated as "absent."

Where it manifests: Consider a transfer or interaction TimeSeries that has only a constant, explicitly-entered value of 0.0 (e.g., an interaction weight of 0 between two populations, or a transfer rate of exactly zero) and no time-varying values. In that case value.vals == [] (falsy) and value.assumption == 0.0 (falsy), so value.vals or value.assumption evaluates to False, and the cell is written as "N" instead of "Y".

Why this causes data loss: The Y/N matrix cell acts as a gate for the corresponding time-dependent data row. In write() (~line 703-708), gate_content() wraps the from/to population cells and OR-interaction cell in =IF(gating_cell="Y", ref, "..."), so when the gate is "N", those cells become "..." once Excel recalculates the formulas. On read, from_tables() (~line 560) explicitly skips any row where row[0].value == "...". So a zero-valued transfer/interaction entry that round-trips through a save/open/re-save (or is read back after Excel recalculates) will vanish from self.ts entirely — not just render as zero, but disappear as an entry.

Why nothing else catches this: There is no validation elsewhere that cross-checks the boolean matrix against the underlying TimeSeries presence of data — the boolean cell is the sole source of truth for whether a row is included on read, so a truthiness bug here directly translates into silent data loss with no error or warning.

Step-by-step proof:

1. Suppose a TDC entry for population pair (A, B) is entered in the databook as a constant value of 0.0 (no time-specific values) — i.e., ts.vals == [], ts.assumption == 0.0.
2. On _write_pop_matrix(), value.vals or value.assumption → [] or 0.0 → 0.0, which is falsy in Python, so content[from_idx, to_idx] = "N".
3. The boolean cell for (A, B) is written as "N" in the spreadsheet.
4. In write(), the row for (A, B) has its from/to-population cells (and the underlying value cell) wrapped in =IF(A2="Y", ..., "...") referencing that boolean cell.
5. When Excel opens/recalculates the sheet (or if the workbook is read back without a live Excel recalculation but xlsxwriter's cached formula values reflect the "N" gate), the from-population cell resolves to "...".
6. On from_tables(), the loop hits if row[0].value == "...": continue (skipping the row), so the (A, B) entry never gets included when reconstructing self.ts.
7. Net effect: the explicit 0.0 transfer/interaction value the user entered is silently gone after a write→read round trip, whereas a time-varying 0.0 (vals=[0.0]) would survive because a non-empty list is truthy.

Suggested fix: Replace the condition with a presence check, e.g. value.has_data (or equivalently value.vals or value.assumption is not None), matching the semantics already used by TimeSeries.has_data elsewhere in the codebase.

[claude]

🟡 In demo(), the fallback error for an unsupported which references the undefined name demos instead of demo_projects (renamed by this PR). This causes a NameError to mask the intended informative exception whenever an invalid project name is passed to demo().

Extended reasoning...

The bug: In atomica/demos.py, this PR renames the local options list to a module-level demo_projects list (also newly exported via __all__). The if which not in demo_projects check on line 51 correctly uses the new name, but the exception message on line 52 still refers to the old name via a variable called demos, which has never existed anywhere in the module:

if which not in demo_projects:
    raise Exception("Supported demonstration projects are:\n%s" % ("\n".join(demos)))

How it manifests: Python evaluates function arguments before the function call itself, so "\n".join(demos) is evaluated first. Since demos is not defined in the local scope, module scope, or as a builtin/import, this raises NameError: name 'demos' is not defined — and that NameError propagates instead of the intended Exception("Supported demonstration projects are:...").

Trigger path: Any call to at.demo(which='invalid_name') (or any which not present in demo_projects) hits this line. Nothing upstream validates which before this point — the default value was also changed from None to 'sir' in this PR, but that only affects the default argument; any explicit invalid string still reaches the broken exception line.

Step-by-step proof:

1. Call at.demo(which='nonexistent').
2. dtdict = sc.odict.fromkeys(demo_projects, 1.0) executes fine.
3. if which not in demo_projects: evaluates to True since 'nonexistent' isn't in the list.
4. Python begins evaluating the raise Exception(...) statement's arguments, starting with "\n".join(demos).
5. Name resolution for demos fails (not a local, not a global in demos.py, not a builtin) → NameError: name 'demos' is not defined is raised instead of the intended Exception.

Why nothing prevents this: There's no test coverage for the invalid-which error path (existing tests such as tests/test_tox_library.py only exercise valid demo names via at.demo_projects), so the mismatched rename slipped through.

Impact: Limited to the error-reporting path for invalid input — a real exception is still raised (integration/valid-usage code paths are unaffected), but the message is a confusing NameError rather than the helpful list of supported project names, which would be misleading for anyone debugging a typo'd demo name.

Fix: Change demos to demo_projects on that line:

raise Exception("Supported demonstration projects are:\n%s" % ("\n".join(demo_projects)))

[Copilot]

Pull request overview

This PR revives and extends the “stochastic transitions” work (originally PR #474) by adding a stochastic/multinomial transition mode to the model, plus RNG plumbing to make sampling reproducible, along with new/updated tests and a few ancillary fixes.

Changes:

• Add multinomial (“discrete people”) transition support in the model and pass seeded RNGs through Project.run_sim()/run_sampled_sims() into run_model()/Model.
• Add acceptance-criteria support for rejecting sampled runs that fall outside specified bounds.
• Add/adjust tox tests and minor updates in plotting, Excel output, and docs.

Reviewed changes

Copilot reviewed 15 out of 16 changed files in this pull request and generated 20 comments.

Show a summary per file

File	Description
tests/test_tox_markovchain.py	Adds functional + statistical comparison tests for stochastic transitions
tests/test_tox_library.py	Updates model parametrization to use at.demo_projects
tests/conftest.py	Attempts to force ordering of one expensive test
docs/examples/Uncertainty.ipynb	Seed/doc/metadata updates in the uncertainty example notebook
atomica/utils.py	Adds RNG-aware sampling and extends parallel_progress input handling; adds stochastic_rounding
atomica/system.py	Reserves constant keyword in framework code names
atomica/results.py	Refactors dataframe Excel writing and a couple of plotting/unit-message tweaks
atomica/project.py	Adds multinomial setting and RNG/acceptance-criteria plumbing into simulation entrypoints
atomica/programs.py	Threads RNG through program sampling (partially)
atomica/plotting.py	Extends color handling and threads colors into plot_series()
atomica/parameters.py	Threads RNG through parameter sampling
atomica/model.py	Implements multinomial flows in compartments/junctions and acceptance-criteria evaluation
atomica/migration.py	Trailing newline change only
atomica/excel.py	Changes TimeDependentConnections.__repr__ and boolean-matrix writing logic
atomica/demos.py	Introduces demo_projects list and updates demo() options logic

Comments suppressed due to low confidence (3)

atomica/project.py:537

• run_sampled_sims(..., acceptance_criteria=[]) uses a mutable default list which may be unintentionally shared between calls.

    def run_sampled_sims(self, parset, progset=None, progset_instructions=None, result_names=None, n_samples: int = 1, rand_seed: int = None, parallel=False, max_attempts=None, num_workers=None, acceptance_criteria=[]) -> list:

atomica/project.py:782

• _run_sampled_sim(..., acceptance_criteria=[]) uses a mutable default list which may be unintentionally shared between calls.

def _run_sampled_sim(proj, parset, progset, progset_instructions: list, result_names: list, max_attempts: int = None, rng_sampler=None, acceptance_criteria=[]):

atomica/model.py:2976

• run_model(..., acceptance_criteria=[]) uses a mutable default list which may be unintentionally shared between calls.

def run_model(settings, framework, parset: ParameterSet, progset: ProgramSet = None, program_instructions: ProgramInstructions = None, name: str = None, rng_sampler=None, acceptance_criteria=[]):

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.