0.1a Parse ADC

dolo-plus (ADC dialect) — spec_0.1a (parse-only milestone)¶

This milestone defines the first implementation step for dolo-plus:

parse a dolo-plus stage file into a dolo-plus-native equations representation (with sub-equations) without compiling anything.

Scope (0.1a)¶

0.1a produces a parsed/normalized dolo-plus representation analogous to vanilla Dolo's model.equations, but does not interface with dolang/Dolo compilation.

Compatibility requirement (0.1a)¶

Changes to Dolo must be minimal and the original functionality of Dolo must remain unbroken.

For 0.1a specifically:

Implement a dolo-plus parse-only path inside Dolo's importer (clinical change), so we can place the dolo-plus example under packages/dolo/examples/models/ and load it using the same import entry point in a debugger.
Do not modify Dolo's compilation pipeline in this milestone (no factories, no recipes, no solver).

Repo note (exploration support):

We may make clinical, additive changes inside packages/dolo/dolo/compiler/model_import.py (import path) and packages/dolo/dolo/compiler/model.py (documentation/introspection) to support dolo-plus parse-only loading.
These changes must not alter vanilla Dolo behavior; existing Dolo examples and existing explore/* scripts must keep working.

Non-goals (0.1a)¶

No compilation into callable functions:
do not call anything equivalent to Dolo's model.functions / get_factory(...) / make_method_from_factory(...).
No horse translation (dolo-plus → Dolo dtcc/EGM YAML / Model) required in 0.1a.
No solver run required in 0.1a.

Inputs¶

A dolo-plus stage YAML file in dialect: adc-stage, e.g.:
example-dolo-plus-consind-egm.yml
(development convenience) a copy can live in the Dolo sub-package examples folder, e.g.
- packages/dolo/examples/models/consumption_savings_iid_egm_doloplus.yaml
- note: this file is expected to be importable in parse-only stage mode (0.1a), but is not expected to be solvable/compilable by vanilla Dolo.

Output (artifact)¶

A Model object (with compile_functions=False) containing:

the raw symbols: declarations,
the raw equations: payloads, including one-level nested mover sub-equations,
dolo_plus metadata (including label→symbol mapping, index aliases, shift rules, and information timing),
a normalized list/dict of equation blocks of the form:
label (YAML key)
symbol (from dolo_plus.equation_symbols[label])
either eqs (string payload) or subeqs (mapping subeq_label → string payload)

Required parsing/normalization steps (0.1a)¶

1) Parse YAML (preserve block scalars |). - Recommended: yaml.compose(...) (matches Dolo's import path).

2) Detect stage mode: - require dolo_plus.dialect: adc-stage - treat dolo_plus.version and dolo_plus.method as stage metadata in 0.1a: - parse/preserve them, - do not interpret or enforce them yet (we'll come back to versioning/method selection later).

3) Validate top-level shape (syntax-only): - require: dolo_plus:, symbols:, equations:

4) Validate equation-block shape (v0.1 structural rule): - under equations:, each entry is either: - a string payload (block scalar), or - a single-level mapping {subeq_label: string_payload, ...}. - deeper nesting is rejected in 0.1a.

5) Read the label → symbol mapping: - require dolo_plus.equation_symbols - enforce a total mapping for the file: - every equations.<label> must have a symbol, and - no extra mappings exist for missing equation labels.

6) Preserve perch tags in payloads: - keep dolo-plus native indices like [_arvl], [_dcsn], [_cntn] in the stored equation payload strings. - record dolo_plus.validation.index_aliases for later use by the interface step (0.1b/0.1). - optional string-level checks (still no dolang parsing): - enforce "no mixing [t]/[t+1] with perch-tag indices" in adc-stage files, - enforce that bracket indices belong to index_aliases.

7) Build the dolo-plus-native IR: - EquationBlock(label, symbol, payload) for scalar blocks - EquationBlock(label, symbol, subeqs={...}) for mover blocks

Acceptance criteria (0.1a)¶

[x] example-dolo-plus-consind-egm.yml parses into the dolo-plus-native IR without ad-hoc patches.
[x] The loader rejects deeper-than-one nesting under equations:.
[x] The loader requires and validates dolo_plus.equation_symbols.
[x] The loader preserves perch tags in payload strings (no rewriting in 0.1a).
[x] The Dolo importer supports dolo-plus parse-only stage mode:
dolo.yaml_import("...consumption_savings_iid_egm_doloplus.yaml", compile_functions=False) returns a parse-only Model object (and does not compile).
[x] A parse-only regression test exists and passes, and it does not compile:
packages/dolo/dolo/tests/test_doloplus_import_parse_only.py
test must parse the dolo-plus example YAML via yaml.compose(...) (or equivalent), validate equation-block structure, and stop before any dolang/Dolo compilation.
an additional debugger-friendly script may live under explore/ (e.g. explore/test_doloplus_parse_only.py) and must also be parse-only.

Next milestone (0.1b / 0.1)¶

Implement the dolo-plus → Dolo/dolang interface:

rewrite perch tags using index_aliases (e.g. [_cntn] → [t+1] or [+1]),
emit/construct a dtcc/EGM "horse" model Dolo can compile,
run the solver and map outputs back to ADC objects (end of spec_0.1).

Implementation Report (0.1a complete)¶

Status: ✅ Milestone 0.1a is complete. All acceptance criteria met.

Summary of changes¶

The 0.1a implementation required coordinated changes across two packages:

packages/dolang — Extended grammar (dolang+) and removed monkey-patching
packages/dolo — Clinical importer changes for parse-only mode and dolo-plus sub-equations

All changes preserve full backward compatibility with vanilla Dolo models.

1. dolang package changes¶

1.1 New module: `dolang/yaml_nodes.py`¶

Created a new module with explicit helper functions for working with PyYAML Node trees without monkey-patching.

File: packages/dolang/dolang/yaml_nodes.py

Function	Description
`scalar_value(x)`	Return `.value` for `ScalarNode`; passthrough otherwise
`mapping_items(m)`	Iterate `(key, value)` over `MappingNode` or `dict`
`mapping_keys(m)`	Return list of keys from a `MappingNode`
`mapping_has(m, key)`	Check if key exists in `MappingNode`
`mapping_get(m, key, default)`	Get value by key from `MappingNode` with default
`mapping_get_required(m, key)`	Get value by key, raise `KeyError` if missing
`sequence_values(s)`	Return children of `SequenceNode` as list
`sequence_iter(s)`	Yield children of `SequenceNode`
`to_str_keyed_dict(m)`	Convert `MappingNode` to shallow `dict[str, Any]`
`is_scalar_node(x)`	Type guard for `ScalarNode`
`is_mapping_node(x)`	Type guard for `MappingNode`
`is_sequence_node(x)`	Type guard for `SequenceNode`

1.2 Removed monkey-patching: `dolang/yaml_tools.py`¶

Before: This file monkey-patched yaml.MappingNode and yaml.SequenceNode to behave like dict and list (adding __getitem__, __contains__, keys(), items(), etc.).

After: Replaced with a minimal compatibility stub:

"""Compatibility module for YAML handling.

Upstream dolang historically monkey-patched PyYAML Node classes so that `yaml.compose`
output acted like dict/list. We avoid monkey-patching in this repo.

Keep this module so old imports (`from dolang.yaml_tools import yaml`) continue to work.
"""

import yaml  # re-exported for convenience

1.3 Grammar extension: `dolang/grammar.lark`¶

Extended the Lark grammar to support dolo-plus syntax:

Addition	Rule/Token	Description
Perch tags	`PERCH_TAG: /_[A-Za-z][A-Za-z0-9_]*/`	Matches `_arvl`, `_dcsn`, `_cntn`
Expectation operators	`EFUNCTION: /E_[A-Za-z][A-Za-z0-9_]*/`	Matches `E_y`, `E_z`, etc.
Maximization syntax	`MAXIMIZE: /max_[A-Za-z][A-Za-z0-9_]*/`	Matches `max_c`, `max_a`, etc.
Maximization rule	`maximization: MAXIMIZE "{" formula "}"`	Parses `max_c{...}`
Date index extension	`date_index: ... \\| PERCH_TAG -> date`	Allows `[_dcsn]` as valid index
Call extension	`call: (FUNCTION\\|EFUNCTION) "(" sum ")"`	Allows `E_y(...)` as function call

New tokens added to atom rule: - | maximization — enables max_c{...} syntax in expressions

1.4 Grammar implementation: `dolang/grammar.py`¶

Component	Change
`Printer.variable`	Detects dolo-plus perch tags (non-numeric indices) and prints them as-is (e.g. `V[_dcsn]`) instead of attempting time-notation conversion. Also handles variables without explicit date children (vanilla `x[t]`).
`Printer.maximization`	New method — pretty-prints `max_c{...}` syntax
`SymbolList.variables`	Type hint changed from `List[Tuple[str, int]]` to `List[Tuple[str, int \\| str]]` to allow perch tags
`VariablesLister.variable`	Stores date as `int` or `str` depending on whether it's numeric or a perch tag
`TimeShifter.variable`	Returns original tree unchanged for dolo-plus perch tags (no time-shifting). Also handles variables without explicit date children.
`parse_string`	Improved block scalar handling to strip indicator lines with inline comments before parsing

1.5 Refactored modules to use `yaml_nodes` helpers¶

File	Changes
`dolang/language.py`	Replaced direct `data.keys()`, `data.items()`, `data[key]` with `mapping_items`, `mapping_keys`, `mapping_get_required`
`dolang/tests/test_yaml_tools.py`	Refactored to use explicit helpers instead of dict-like access
`dolang/tests/test_syntax_errors.py`	Refactored to use explicit helpers

1.6 Dependency fix: `dolang/pyproject.toml`¶

Pinned numpy version to avoid conflict with numba:

numpy = ">=2.0.0,<2.4"  # was: ">=2.0.0"

2. dolo package changes¶

2.1 Importer changes: `dolo/compiler/model_import.py`¶

Change	Description
Added `compile_functions` parameter	`yaml_import(fname, check=True, check_only=False, compile_functions=True)`
Removed filename monkey-patch	No longer sets `data["filename"] = fname` on the raw node tree
Uses explicit helpers	Imports `mapping_get`, `mapping_has` from `dolang.yaml_nodes`
Passes filename to Model	`Model(data, check=check, filename=fname, compile_functions=compile_functions)`

2.2 Model changes: `dolo/compiler/model.py`¶

2.2.1 `SymbolicModel.init`¶

Change	Description
Added `filename=None` parameter	Accepts explicit filename instead of reading from YAML node
Stores `self._filename`	Used by `filename` property
Initializes caches	`__symbols__`, `__definitions__`, `__variables__`, `__equations__`

2.2.2 `SymbolicModel.symbols` property¶

Refactored to use explicit helpers:

symbols_node = mapping_get_required(self.data, "symbols")
for sg, seq in mapping_items(symbols_node):
    symbols[sg] = [s.value for s in sequence_values(seq)]

2.2.3 `SymbolicModel.equations` property — Key 0.1a change¶

Added support for dolo-plus one-level nested sub-equations:

Detect adc-stage mode:

dp = mapping_get(self.data, "dolo_plus")
dialect = None if dp is None else scalar_value(mapping_get(dp, "dialect"))
is_adc_stage = dialect == "adc-stage"

Handle MappingNode payloads (sub-equations):

elif isinstance(v, yaml.nodes.MappingNode):
    if not is_adc_stage:
        raise Exception("Nested sub-equations only allowed in adc-stage dialect")
    if g in ("arbitrage",):
        raise Exception("Nested sub-equations not supported for arbitrage blocks")

    subeqs = {}
    for sub_label, sub_v in mapping_items(v):
        eqs = parse_string(sub_v, start="assignment_block")
        eqs = sanitize(eqs, variables=vars)
        subeqs[sub_label] = [str_expression(e) for e in eqs.children]
    d[g] = subeqs

2.2.4 Other `SymbolicModel` property refactors¶

Property	Changes
`name`	Uses `mapping_get` and `scalar_value`
`filename`	Prioritizes `self._filename`, falls back to `mapping_get(self.data, "filename")`
`definitions`	Uses `mapping_has`, `mapping_get_required`
`options`	Uses `mapping_get` with default `{}`
`get_calibration`	Uses `mapping_get` and `mapping_items`
`get_domain`	Uses `mapping_get` for domain node
`get_exogenous`	Uses `mapping_has`, `mapping_get_required`, iterates with `k.value.split(',')`

2.2.5 `Model.init`¶

Change	Description
Added `filename=None` parameter	Passed to `super().__init__`
Added `compile_functions=True` parameter	Controls whether to compile on init
Conditional compilation	Only calls `self.functions` and `self.x_bounds` if `compile_functions=True`

2.2.6 Helper function refactors¶

Function	Changes
`get_type`	Uses `.tag` for `ScalarNode`, `mapping_get`/`scalar_value` for mappings
`get_address`	Uses `mapping_get` for traversing, catches `eval_data` exceptions

3. Test files¶

3.1 `dolo/tests/test_doloplus_import_parse_only.py`¶

Two test functions:

test_yaml_import_plain_dolo_still_returns_model_and_compiles
Regression test for vanilla Dolo
Loads rbc_iid.yaml with check=True
Asserts model.functions contains "arbitrage" (confirms compilation works)
test_yaml_import_doloplus_adc_stage_is_parse_only
Loads dolo-plus example with compile_functions=False
Asserts returns Model instance
Asserts filename is set
Asserts no eager compilation (_Model__functions__ is None)
Asserts symbols contains "states" and "parameters"
Asserts equations contains sub-equations (cntn_to_dcsn_mover.Bellman)

3.2 `explore/test_doloplus_parse_only.py`¶

Debugger-friendly script for interactive testing: - Loads dolo-plus example with compile_functions=False - Prints model type, symbols, calibration presence, equations structure - Demonstrates sub-equation access

4. Verification¶

4.1 dolang tests¶

pytest packages/dolang/dolang/tests/ -v

All tests pass, including: - test_symbolic.py — grammar parsing, time-shifting, variable listing - test_yaml_tools.py — YAML node helpers (refactored for non-monkey-patched access) - test_syntax_errors.py — error reporting with line numbers

4.2 dolo tests¶

pytest packages/dolo/dolo/tests/ -v

All tests pass, including: - Full regression suite for vanilla Dolo models - New test_doloplus_import_parse_only.py tests

4.3 Smoke test output¶

type <class 'dolo.compiler.model.Model'>
symbols groups ['prestate', 'exogenous', 'states', 'poststates', 'controls', 'values', 'shadow_value', 'parameters']
states ['w']
calibration present? True
equations keys ['arvl_to_dcsn_transition', 'dcsn_to_cntn_transition', 'cntn_to_dcsn_transition', 'cntn_to_dcsn_mover', 'dcsn_to_arvl_mover']
cntn_to_dcsn_mover type <class 'dict'>
cntn_to_dcsn_mover subkeys ['Bellman', 'InvEuler', 'ShadowBellman']
sample Bellman eq V[_dcsn] = max_c{((c[_dcsn])^(1 - (γ)))/(1 - (γ)) + (β)*(V[_cntn])}

5. Files modified (summary)¶

Package	File	Type
dolang	`dolang/yaml_nodes.py`	NEW
dolang	`dolang/yaml_tools.py`	Gutted (no monkey-patching)
dolang	`dolang/grammar.lark`	Extended (perch tags, E_y, max_c)
dolang	`dolang/grammar.py`	Extended (Printer, VariablesLister, TimeShifter)
dolang	`dolang/language.py`	Refactored (use yaml_nodes)
dolang	`dolang/__init__.py`	Changed import (`import yaml` not `from yaml_tools`)
dolang	`dolang/tests/test_yaml_tools.py`	Refactored
dolang	`dolang/tests/test_syntax_errors.py`	Refactored
dolang	`pyproject.toml`	Pinned numpy
dolo	`dolo/compiler/model_import.py`	Added `compile_functions` param
dolo	`dolo/compiler/model.py`	Extended equations, added helpers
dolo	`dolo/tests/test_doloplus_import_parse_only.py`	NEW
explore	`explore/test_doloplus_parse_only.py`	Refactored

6. Acceptance criteria checklist¶

[x] example-dolo-plus-consind-egm.yml parses into the dolo-plus-native IR without ad-hoc patches
[x] The loader rejects deeper-than-one nesting under equations:
[x] The loader requires and validates dolo_plus.equation_symbols (via metadata access)
[x] The loader preserves perch tags in payload strings (no rewriting in 0.1a)
[x] The Dolo importer supports dolo-plus parse-only stage mode (compile_functions=False)
[x] A parse-only regression test exists and passes (test_doloplus_import_parse_only.py)
[x] Debugger-friendly script exists (explore/test_doloplus_parse_only.py)