Pandas UI for initial conditions and Markov transition probabilities.

[hmgaudecker]

Fix #242.
Fix #259.

Summary

Add pandas interoperability utilities and simplify the simulation API so users work
with DataFrames and a single initial_conditions dict instead of raw JAX arrays.

Public API changes

New module (lcm.pandas_utils, re-exported from lcm):

• initial_conditions_from_dataframe(df, model=) — convert a DataFrame with a
  "regime" column to the initial_conditions dict expected by simulate() and
  solve_and_simulate(). Discrete state labels and the "regime" column are
  mapped to integer codes.
• transition_probs_from_series(series, model=, regime_name=, ...) —
  build transition probability arrays from a pandas Series with a named MultiIndex.
  Works for both state and regime transitions (inferred from next_<x>; automatically
  reorders index levels to match the function signature.
  Uses "age" (not "period") as the time-dimension level name.
• validate_transition_probs(probs_array, ...) — validate shape, bounds, and row-sums
  of transition probability arrays.

Simulation API simplification:

• model.simulate() and model.solve_and_simulate() now take a single
  initial_conditions: Mapping[str, Array] instead of separate initial_states +
  initial_regimes. The dict includes all state arrays plus "regime_id" (integer
  codes from model.regime_names_to_ids).
• Internal simulate() and validate_initial_conditions() updated accordingly.

@categorical now requires ordered=True/False — breaking change to the decorator
syntax (@categorical → @categorical(ordered=False)). The flag flows through to
pd.CategoricalDtype in simulation output, so ordered categoricals (e.g. education
levels) render correctly in pandas. Gives us a 1:1 mapping with Pandas categorical
dtypes.

Internal changes

• AST-based probs_array validation (error_handling.py): At model init, inspects
  transition function source to verify probs_array[...] indexing matches the parameter
  declaration order. Warns when the pattern is non-trivial (computed indices, multiple
  subscripts).
• Ordered category merging (regime_processing.py): _merge_ordered_categories()
  validates consistent ordering across regimes and merges category sets via topological
  sort.
• Terminal-period validation (Raise a meaningful error if a non-terminal regime is active in the terminal period. #277): Non-terminal regimes active in the terminal
  period now raise at model init.
• Simulation cleanup: Removed unused next_age parameter; moved regime transition
  probability validation from simulation to model init.
• Shared test models: tests/test_models/basic_discrete.py (discrete + continuous
  states, no stochastic transitions) and tests/test_models/regime_markov.py
  (MarkovTransition on regime transitions).

Documentation

• New user guide: Working with DataFrames and Series (pandas_interop.md)
• New explanation notebook: Stochastic Transitions (stochastic_transitions.ipynb)
• Updated Solving and Simulating: initial_conditions is now the primary API
• Updated tiny example and all example docs to use initial_conditions

Infrastructure

• Type-checking environment split: ty now runs in its own pixi environment
• AI instructions consolidated into .ai-instructions submodule

[hmgaudecker]

Code review

Found 1 issue:

1. _collect_all_state_names includes shock states despite docstring claiming otherwise. The function does state_names.update(regime.states.keys()) without filtering out _ShockGrid instances. For any model with shock states (Normal, Uniform, Tauchen, Rouwenhorst), initial_states_from_dataframe will incorrectly require those shock columns in the DataFrame — but shock states are drawn automatically and should never be user-supplied. The error message at line 78 even says "All non-shock states must be provided", confirming the intent. The fix is to filter with isinstance(grid, _ShockGrid) as done elsewhere in the codebase (e.g. regime.py:346).

pylcm/src/lcm/pandas_utils.py

Lines 402 to 414 in 46b6a99

	def _collect_all_state_names(
	regimes: Mapping[str, Regime],
	initial_regimes: list[str],
	) -> set[str]:
	"""Collect all non-shock state names from regimes present in initial_regimes."""
	state_names: set[str] = set()
	for regime_name in set(initial_regimes):
	regime = regimes[regime_name]
	state_names.update(regime.states.keys())
	# Always include age
	state_names.add("age")
	return state_names

🤖 Generated with Claude Code

_{- If this code review was useful, please react with 👍. Otherwise, react with 👎.}

[timmens]

Very very nice! No showstoppers, only a few comments on the interface.

[timmens]

We could also allow model.solve_and_simulate to accept a pandas.DataFrame as the initial conditions. And if it is a dataframe we call initial_conditions_from_dataframe?

The user would get a very similar error, and it would probably feel even more natural?

Alternatively, this could be made a model method.

[hmgaudecker]

Shoot, a comment of mine here did not make it through yesterday, it seems. Sorry about that.

We went down that rout of allowing all kinds of different inputs in ttsim and it is a pain. While I agree that would be nice, I really want to keep the interface so that it accepts exactly one thing here, code maintenance is hard enough. We can expect users to be well-versed enough to run the pre-processing steps themselves.

[timmens]

Consider making this a method on Model and inferring state_name from the Series' MultiIndex — the "next_" prefix on the outcome level already identifies the state (or "next_regime" for regime transitions). That would reduce the signature to:

model.transition_probs_from_series(series, regime_name="working")

if I am not mistaken.

[hmgaudecker]

See above for the model-method part. I am investigating the interface change right now!