Spec 0.1b: dolo-plus to Dolo Model Object Translator¶

Version: 0.1b (revised) Date: 2026-01-01 Status: Implemented Depends on: spec_0.1a (dolo-plus parsing), ../syntax-semantic-rules/doloplus-foundations.md

Implementation Summary¶

Location: packages/dolo/dolo/compiler/doloplus_translator/

File	Purpose
`core.py`	Main entry point `doloplus_to_dolo()`, translation pipeline
`transforms.py`	Pure functions for equation string transformations
`mappings.py`	Default mapping tables (symbol groups, perch offsets, etc.)
`__init__.py`	Public exports

Key commits: - 1b3e917: First translator implementation - 2710910: Refactor doloplus translator and add methodization support

Usage:

from dolo.compiler.model_import import yaml_import
from dolo.compiler.doloplus_translator import doloplus_to_dolo

# Load dolo-plus model (without compiling)
model_plus = yaml_import("model.yaml", compile_functions=False)

# Translate to vanilla Dolo
model_dolo = doloplus_to_dolo(
    model_plus,
    mapping_tables="path/to/mapping_tables.yaml"
)

1. Purpose¶

This spec defines a Python module that converts a parsed dolo-plus model object to a vanilla Dolo model object suitable for compilation by the unchanged Dolo solver (EGM/dtcc recipes).

Key constraint: The translation operates at the model representation level (on the parsed Model object), not at the YAML level. This follows the approach documented in the bellman-dev extraction (28 Dec 2025).

2. Scope¶

2.1 In Scope (v0.1b)¶

Symbol group translation (dolo-plus → Dolo)
Equation block translation with:
Perch tag → time subscript conversion (group-aware)
Mover sub-equation flattening
Expectation operator E_{y}(...) or E_{y,z}(...) handling (list-first)
Maximization operator max_{c}(...) handling (strip for EGM)
Slot normalization (dolo-plus -1/0/+1 → Dolo conventions)
Generation of Dolo expectations symbol group and expectation block

2.2 Out of Scope (v0.1b)¶

Multi-stage models (only single-stage stationary models)
Non-EGM solution methods
YAML-level transformation (translator works on parsed objects)
Branching stages
Non-stationary (finite-horizon) problems

3. Canonical Operator Syntax¶

3.1 Expectation Operator¶

The canonical expectation syntax is list-first, even for singletons:

E_{y}(objective)       # Single shock variable
E_{y,z}(objective)     # Multiple shock variables
E_{y,z,w}(objective)   # Three or more

Unicode form: 𝔼_{y,z}(objective) is also accepted.

Legacy form (accepted as input, normalized to canonical on output):

E_y(objective)         # Legacy single-variable form → E_{y}(objective)

Semantics for 0.1b: - E_{y,z}(...) must match the full exogenous shock vector relevant for the expectation block - If the variable list doesn't match the model's exogenous process, translation fails (fail fast) - This prevents silently-wrong translations due to filtration mismatches

3.2 Maximization Operator¶

The canonical maximization syntax mirrors expectation:

max_{c}(objective)     # Single control variable
max_{c,a}(objective)   # Multiple control variables

Legacy form (accepted as input, normalized to canonical on output):

max_c{objective}       # Legacy brace form → max_{c}(objective)

3.3 Grammar Integration¶

Both operators use a lexer-safe prefix tokenization approach to avoid ambiguity with NAME tokens.

4. Dolang Grammar Updates¶

4.1 The Lexer Trap Problem¶

A naive grammar rule like "E" "_" "{" ... creates a lexer ambiguity: - NAME matches E_ (underscore allowed in names) - Input E_{y}( tokenizes as NAME("E_") + {... — never parses as expectation

4.2 Solution: High-Priority Prefix Tokens¶

Define tokens that match the whole operator prefix:

// --- Expectation openers (HIGH priority to prevent NAME from grabbing E_ prefix) ---
_EOPEN.2: "E[" | "𝔼["
_EUSCORE.2: "E_{" | "𝔼_{"

// --- Maximization opener (HIGH priority) ---
_MAXUSCORE.2: "max_{"

// Legacy patterns (lower priority, for backwards compat)
EFUNCTION: /E_[A-Za-z][A-Za-z0-9_]*/
MAXIMIZE: /max_[A-Za-z][A-Za-z0-9_]*/

4.3 Grammar Rules¶

// Variable list helper (used by both E and max)
exp_var_list: NAME ("," NAME)*   -> exp_var_list
max_var_list: NAME ("," NAME)*   -> max_var_list

// Expectation rule with three forms:
// 1. Bracket form: E[...] or 𝔼[...]
// 2. Subscript form: E_{y,z}(...) or 𝔼_{y,z}(...)
// 3. Legacy function form: E_y(...) - normalized to expectation
expectation: _EOPEN formula "]"                              -> expectation
    | _EUSCORE exp_var_list "}" "(" formula ")"              -> expectation
    | EFUNCTION "(" formula ")"                              -> expectation

// Maximization rule with two forms:
// 1. Subscript form: max_{c}(...) or max_{c,a}(...)
// 2. Legacy brace form: max_c{...} - normalized to maximization
maximization: _MAXUSCORE max_var_list "}" "(" formula ")"    -> maximization
    | MAXIMIZE "{" formula "}"                               -> maximization

4.4 AST Shape Convention¶

Both operators produce AST nodes with consistent structure:

Expectation node: - Child 0: exp_var_list (list of variable names, possibly empty for E[...]) - Child 1: formula (the integrand)

Maximization node: - Child 0: max_var_list (list of control variables) - Child 1: formula (the objective)

This makes translation transforms straightforward — always iterate over a list.

5. Input/Output Contract¶

5.1 Input¶

A parsed dolo-plus Model object obtained via:

from dolo.compiler.model_import import yaml_import
model_plus = yaml_import(path, check=False, compile_functions=False)

The model must have: - dolo_plus.dialect == "adc-stage" - dolo_plus.version == "0.1" - Equation blocks using perch tags ([_arvl], [_dcsn], [_cntn])

5.2 Output¶

A vanilla Dolo Model object that: - Has only standard Dolo symbol groups (exogenous, states, controls, expectations, poststates, parameters) - Has equation blocks using time subscripts ([t-1], [t], [t+1]) - Can be compiled via model.functions / model.__original_gufunctions__ - Is solvable by vanilla Dolo EGM solver

5.3 Model Object Construction Note¶

When building the output Model object, do not put lists of equation strings directly into data['equations']. The vanilla Dolo YAML structure expects:

equations: { block_name: <ScalarNode multiline string> }

So build_dolo_model(...) should: 1. Construct a dict where equations[block] is a single multiline string (not a list) 2. Let the normal Dolo Model parser turn that into the internal list-of-equations form

This keeps us inside "minimal change to Dolo" and ensures correct sanitization.

6. Translation Rules¶

6.1 Symbol Group Translation¶

dolo-plus Group	Dolo Group	Notes
`prestate`	(absorbed into `poststates[-1]`)	Arrival state becomes previous-period poststate
`exogenous`	`exogenous`	Direct mapping
`states`	`states`	Direct mapping
`poststates`	`poststates`	Direct mapping
`controls`	`controls`	Direct mapping
`values`	(dropped)	Value objects are implicit in Dolo EGM
`shadow_value`	(dropped)	Shadow values become `expectations` (see 6.4)
`parameters`	`parameters`	Direct mapping

New group creation: - Create expectations: [mr] where mr is the expected marginal value object

6.2 Group-Aware Perch Tag → Time Subscript Conversion¶

The translator applies group-aware substitutions. The mapping depends on which symbol group a token belongs to:

6.2.1 States / Controls / Exogenous (decision-observed)¶

dolo-plus Perch Tag	Dolo Time Subscript
`[_arvl]`	`[t-1]`
`[_dcsn]`	`[t]`
`[_cntn]`	`[t+1]` (when referring to next-stage decision objects)

6.2.2 Poststates (end-of-period objects)¶

Dolo uses [t] for current-period poststate (not [t+1]):

dolo-plus	Dolo
`a[_cntn]`	`a[t]`
`a[_dcsn]`	(typically not used for poststates)

6.2.3 Prestate → Poststate Rename¶

dolo-plus prestate becomes Dolo's previous-period poststate:

dolo-plus	Dolo
`b[_arvl]`	`a[t-1]`

6.3 Implementation: Group-Aware Mapper¶

def perch_to_time(
    token: str,
    perch_tag: str,
    symbol_groups: dict[str, list[str]],
    prestate_rename: dict[str, str],
    index_aliases: dict[str, int]
) -> str:
    """
    Convert a dolo-plus token with perch tag to Dolo token with time subscript.

    Args:
        token: Variable name (e.g., 'a', 'c', 'b')
        perch_tag: Perch tag (e.g., '_cntn', '_dcsn', '_arvl')
        symbol_groups: dolo-plus symbol groups (to determine token's group)
        prestate_rename: Mapping for prestate symbols (e.g., {'b': 'a'})
        index_aliases: Perch tag → slot mapping

    Returns:
        Dolo-style indexed token (e.g., 'a[t]', 'c[t+1]')
    """
    slot = index_aliases[perch_tag]

    # Handle prestate rename
    if token in prestate_rename:
        token = prestate_rename[token]
        # Prestate at _arvl → poststate at [t-1]
        return f"{token}[t-1]"

    # Handle poststates specially: _cntn → [t] not [t+1]
    if token in symbol_groups.get('poststates', []):
        if perch_tag == '_cntn':
            return f"{token}[t]"

    # Default mapping
    if slot == 0:
        return f"{token}[t]"
    elif slot > 0:
        return f"{token}[t+{slot}]"
    else:
        return f"{token}[t{slot}]"

6.4 Equation Block Translation¶

Based on the mapping table from doloplus-foundations.md Section 5:

dolo-plus Block	dolo-plus Sub-Eq	Dolo Block	Transformation
`arvl_to_dcsn_transition` (g_ad)	—	`half_transition`	Perch→time; `b[_arvl]` → `a[t-1]`
`dcsn_to_cntn_transition` (g_de)	—	`auxiliary_direct_egm`	Perch→time; `a[_cntn]` → `a[t]`
`cntn_to_dcsn_transition` (g_ed)	—	`reverse_state`	Perch→time only
`cntn_to_dcsn_mover` (T_ed)	`InvEuler`	`direct_response_egm`	See Section 6.5
`cntn_to_dcsn_mover` (T_ed)	`ShadowBellman`	`expectation` (RHS)	Becomes integrand
`dcsn_to_arvl_mover` (T_da)	`ShadowBellman`	`expectation` (combined)	See Section 6.6
`cntn_to_dcsn_mover` (T_ed)	`Bellman`	(dropped)	Implicit in EGM
`dcsn_to_arvl_mover` (T_da)	`Bellman`	(dropped)	Implicit in EGM

6.5 InvEuler → direct_response_egm¶

dolo-plus input:

c[_cntn] = (β*dV[_cntn])^(-1/γ)

Transformation: 1. Replace dV[_cntn] with mr[t] (Dolo's expectation variable) 2. Note: The β inside may need adjustment depending on where discount is applied 3. Replace perch tags with time subscripts (group-aware)

Dolo output:

c[t] = (mr[t])^(-1/γ)

6.6 Expectation Block Synthesis¶

The Dolo expectation block requires combining information from two dolo-plus sub-equations. This is the core of the "Bellman operator factorization" approach.

6.6.1 Source Sub-Equations¶

T_ed[ShadowBellman]: Defines the within-stage integrand at the decision perch
```
dV[_dcsn] = (c[_dcsn])^(-γ)
```
T_da[ShadowBellman]: Defines the expectation + return factor
```
dV[_arvl] = r * E_{y}(dV[_dcsn])
```

6.6.2 Synthesis Algorithm¶

The synthesis produces the Dolo expectation block equation mr[t] = ...:

def synthesize_expectation_block(equations: dict, eq_symbols: dict) -> str:
    """
    Combine T_ed.ShadowBellman + T_da.ShadowBellman → Dolo expectation block.

    Key insight: T_ed[ShadowBellman] is a within-stage integrand definition.
    In EGM, the horse `expectation` block evaluates next-stage decision
    integrand and integrates it back.
    """
    # 1. Extract integrand from T_ed.ShadowBellman
    #    dV[_dcsn] = ϕ(...) → extract ϕ
    integrand = extract_rhs(equations['cntn_to_dcsn_mover']['ShadowBellman'])

    # 2. Apply inter-stage + within-stage shift:
    #    Decision-perch objects become next-period decision objects
    #    c[_dcsn] → c[t+1] in the Dolo output equation
    integrand_shifted = shift_to_next_period(integrand)

    # 3. Extract multiplicative factors from T_da.ShadowBellman
    #    dV[_arvl] = r * E_{y}(dV[_dcsn]) → extract 'r'
    factors = extract_factors(equations['dcsn_to_arvl_mover']['ShadowBellman'])

    # 4. Apply discount factor β
    #    Dolo's mr is the *discounted* expected marginal value
    discounted = f"β * ({integrand_shifted})"

    # 5. Multiply by return factors
    result = f"mr[t] = {discounted}"
    for factor in factors:
        result = f"{result} * {factor}"

    return result

6.6.3 Example Synthesis¶

Input: - T_ed.ShadowBellman: dV[_dcsn] = (c[_dcsn])^(-γ) - T_da.ShadowBellman: dV[_arvl] = r * E_{y}(dV[_dcsn])

Steps: 1. Extract integrand: (c[_dcsn])^(-γ) 2. Shift to next period: (c[t+1])^(-γ) 3. Extract factors: r 4. Apply discount: β * (c[t+1])^(-γ) 5. Multiply factors: β * (c[t+1])^(-γ) * r

Output:

mr[t] = β * (c[t+1])^(-γ) * r

6.7 Operator Handling¶

6.7.1 Expectation Operator `E_{y}(...)` / `E_{y,z}(...)`¶

The expectation operator is removed from the generated Dolo equation string; the expectation is performed algorithmically by Dolo's expectation block mechanism.

Important: The operator remains in the dolo-plus source and AST — we only strip it when generating the Dolo output.

# dolo-plus: dV[_arvl] = r * E_{y}(dV[_dcsn])
#
# The E_{y}(...) wrapper is absorbed into the block semantics.
# The integrand dV[_dcsn] is what goes into the RHS of the
# synthesized expectation equation.

For 0.1b, if the variable list in E_{y,z}(...) doesn't match the model's exogenous process, translation fails (fail fast to prevent silently-wrong translations).

6.7.2 Maximization Operator `max_{c}(...)`¶

The maximization operator is stripped entirely for EGM models; the Bellman sub-equation is dropped since EGM handles optimization implicitly.

# dolo-plus: V[_dcsn] = max_{c}((c[_dcsn]^(1-γ))/(1-γ) + β*V[_cntn])
# This equation is DROPPED in EGM translation

7. Transform Functions¶

7.1 Core Transforms (transforms.py)¶

# dolo/compiler/doloplus_translator/transforms.py

def perch_to_time(
    eq_str: str,
    symbol_groups: dict[str, list[str]],
    prestate_rename: dict[str, str],
    index_aliases: dict[str, int]
) -> str:
    """
    Convert perch tags [_arvl], [_dcsn], [_cntn] to time subscripts.
    Uses group-aware logic to handle poststates correctly.
    """
    ...

def strip_expectation_operator(eq_str: str) -> tuple[str, list[str] | None]:
    """
    Remove E_{var,var2,...}(...) operator, return (inner_expr, var_list).

    E.g., 'r * E_{y}(dV[_dcsn])' → ('r * dV[_dcsn]', ['y'])
    E.g., 'E_{y,z}(V + W)' → ('V + W', ['y', 'z'])

    Returns (eq_str, None) if no expectation operator found.

    For simple, non-nested forms, regex may be used.
    For nested cases like E_{y}(a + E_{z}(b)), use AST-based parsing.
    """
    ...

def strip_max_operator(eq_str: str) -> tuple[str, list[str] | None]:
    """
    Remove max_{c,a,...}(...) operator, return (inner_expr, control_vars).

    E.g., 'max_{c}(u(c) + V)' → ('u(c) + V', ['c'])
    E.g., 'max_{c,k}(r(c,k) + V)' → ('r(c,k) + V', ['c', 'k'])
    """
    ...

def rename_symbols(eq_str: str, mapping: dict[str, str]) -> str:
    """Rename symbols in equation string. E.g., {'b': 'a', 'dV': 'mr'}."""
    ...

def shift_timing(eq_str: str, shift: int) -> str:
    """Shift all time indices by `shift`. E.g., x[t] → x[t+1] with shift=1."""
    ...

def shift_to_next_period(eq_str: str) -> str:
    """
    Apply inter-stage shift for expectation synthesis.
    Decision-perch objects become next-period decision objects.
    c[_dcsn] → c[t+1], c[t] → c[t+1]
    """
    return shift_timing(eq_str, shift=1)

7.2 AST-Based Transforms (preferred)¶

from dolang.symbolic import parse_string
from dolang.grammar import Printer

class PerchToTimeTransformer(lark.Transformer):
    """
    Transform perch tags to time indices with group awareness.
    """
    def __init__(self, symbol_groups, prestate_rename, index_aliases):
        self.symbol_groups = symbol_groups
        self.prestate_rename = prestate_rename
        self.index_aliases = index_aliases

    def variable(self, children):
        name, date = children
        # Apply group-aware transformation
        ...

class ExpectationStripper(lark.Transformer):
    """
    Remove expectation operators, returning the inner expression
    and the list of shock variables.
    """
    def __init__(self):
        self.shock_vars = []

    def expectation(self, children):
        if len(children) == 2:
            # E_{var_list}(formula) form
            var_list, formula = children
            self.shock_vars = self._extract_var_list(var_list)
            return formula
        else:
            # E[formula] form (no explicit var list)
            return children[0]

8. Algorithm Outline¶

def translate_doloplus_to_dolo(model_plus: Model) -> Model:
    """
    Translate a parsed dolo-plus model to a vanilla Dolo model.

    Args:
        model_plus: Parsed dolo-plus Model (with compile_functions=False)

    Returns:
        model_dolo: Vanilla Dolo Model ready for compilation
    """

    # 1. Validate input is dolo-plus stage model
    validate_doloplus_dialect(model_plus)

    # 2. Extract metadata
    dp_meta = model_plus.data['dolo_plus']
    index_aliases = dp_meta['validation']['index_aliases']
    eq_symbols = dp_meta['equation_symbols']
    slot_map = dp_meta['slot_map']

    # 3. Build symbol group lookup for group-aware transforms
    symbol_groups = model_plus.symbols
    prestate_rename = build_prestate_rename_map(slot_map)

    # 4. Translate symbols
    symbols_dolo = translate_symbols(model_plus.symbols)

    # 5. Translate equations
    equations_dolo = {}

    for label, payload in model_plus.equations.items():
        canonical = eq_symbols[label]

        if is_transition(canonical):
            # g_ad, g_de, g_ed → half_transition, auxiliary_direct_egm, reverse_state
            block_name, eq_str = translate_transition(
                canonical, payload, symbol_groups, prestate_rename, index_aliases
            )
            equations_dolo[block_name] = eq_str

        elif is_mover(canonical):
            # T_ed, T_da → extract sub-equations
            for subeq_name, subeq_payload in payload.items():
                translated = translate_mover_subeq(
                    canonical, subeq_name, subeq_payload,
                    symbol_groups, prestate_rename, index_aliases
                )
                if translated:
                    block_name, eq_str = translated
                    equations_dolo[block_name] = eq_str

    # 6. Synthesize expectation block
    equations_dolo['expectation'] = synthesize_expectation_block(
        model_plus.equations, eq_symbols, symbol_groups, index_aliases
    )

    # 7. Build output Model
    # IMPORTANT: equations should be multiline strings, not lists
    model_dolo = build_dolo_model(
        symbols=symbols_dolo,
        equations=format_equations_as_strings(equations_dolo),
        calibration=model_plus.calibration,
        domain=model_plus.domain,
        exogenous=model_plus.exogenous,
        options=model_plus.options,
    )

    return model_dolo

9. File Structure & Module Architecture¶

9.1 Directory Layout¶

packages/dolo/dolo/
├── compiler/
│   ├── model.py
│   ├── model_import.py
│   ├── doloplus_translator/          # NEW: dolo-plus → Dolo translation subpackage
│   │   ├── __init__.py               # Public API
│   │   ├── transforms.py             # Reusable transforms (perch→time, etc.)
│   │   ├── translate_egm.py          # EGM-specific translator
│   │   ├── translate_vfi.py          # VFI-specific translator (future)
│   │   └── base.py                   # Base translator class & utilities
│   ├── recipes.py
│   └── recipes.yaml
└── tests/
    └── test_doloplus_translator.py

9.2 Module Responsibilities¶

`transforms.py` — Reusable Transforms¶

These are model-agnostic, reusable building blocks with list-first operator handling:

# dolo/compiler/doloplus_translator/transforms.py

def perch_to_time(
    eq_str: str,
    symbol_groups: dict[str, list[str]],
    prestate_rename: dict[str, str],
    index_aliases: dict[str, int]
) -> str:
    """Convert perch tags using group-aware logic."""
    ...

def strip_expectation_operator(eq_str: str) -> tuple[str, list[str] | None]:
    """
    Remove E_{var,var2,...}(...) operator.
    Returns (inner_expr, var_list) where var_list is always a list.
    """
    ...

def strip_max_operator(eq_str: str) -> tuple[str, list[str] | None]:
    """
    Remove max_{c,a,...}(...) operator.
    Returns (inner_expr, control_list) where control_list is always a list.
    """
    ...

`base.py` — Base Translator & Utilities¶

# dolo/compiler/doloplus_translator/base.py

from abc import ABC, abstractmethod
from dolo.compiler.model import Model

class DoloplusTranslator(ABC):
    """Base class for dolo-plus → Dolo translators."""

    def __init__(self, model_plus: Model):
        self.model_plus = model_plus
        self.validate()
        self._build_lookups()

    def _build_lookups(self):
        """Build symbol group and slot map lookups."""
        self.symbol_groups = self.model_plus.symbols
        dp_meta = self.model_plus.data['dolo_plus']
        self.index_aliases = dp_meta['validation']['index_aliases']
        self.prestate_rename = self._build_prestate_rename_map(dp_meta['slot_map'])

    @abstractmethod
    def translate(self) -> Model:
        """Translate to vanilla Dolo model. Subclasses implement this."""
        ...

`translate_egm.py` — EGM-Specific Translator¶

# dolo/compiler/doloplus_translator/translate_egm.py

from .base import DoloplusTranslator
from . import transforms

class EGMTranslator(DoloplusTranslator):
    """Translator for EGM-solvable dolo-plus models."""

    def translate(self) -> Model:
        """
        Translate dolo-plus stage model to vanilla Dolo EGM model.

        EGM-specific logic:
        - Drops Bellman sub-equations (implicit in EGM)
        - Synthesizes expectation block from ShadowBellman sub-eqs
        - Maps InvEuler → direct_response_egm
        - Uses group-aware perch→time conversion
        """
        ...

    def synthesize_expectation(self) -> str:
        """
        Combine T_ed.ShadowBellman + T_da.ShadowBellman → expectation.

        Implements the Bellman-Backus operator factorization:
        1. Extract integrand from T_ed.ShadowBellman (within-stage)
        2. Shift integrand to next period (c[_dcsn] → c[t+1])
        3. Apply discount factor β
        4. Multiply by return factors from T_da.ShadowBellman
        """
        ...

10. Test Cases¶

10.1 Operator Parsing Tests (dolang-level)¶

These should be in dolang's test suite — the translator depends on correct parsing.

def test_expectation_single_var():
    """Parse E_{y}(V[_dcsn] + 1) correctly."""
    tree = parse_string("E_{y}(V[_dcsn] + 1)")
    assert tree.data == "expectation"
    var_list = extract_var_list(tree.children[0])
    assert var_list == ["y"]

def test_expectation_multi_var():
    """Parse E_{y,z}(V + W) correctly."""
    tree = parse_string("E_{y,z}(V + W)")
    var_list = extract_var_list(tree.children[0])
    assert var_list == ["y", "z"]

def test_expectation_nested():
    """Parse nested expectations: E_{y}(a + E_{z}(b))."""
    tree = parse_string("E_{y}(a + E_{z}(b))")
    # Should have two expectation nodes in the tree
    ...

def test_max_single_var():
    """Parse max_{c}(u(c) + V) correctly."""
    tree = parse_string("max_{c}(u(c) + V)")
    assert tree.data == "maximization"
    var_list = extract_var_list(tree.children[0])
    assert var_list == ["c"]

def test_max_multi_var():
    """Parse max_{c,k}(r + V) correctly."""
    tree = parse_string("max_{c,k}(r + V)")
    var_list = extract_var_list(tree.children[0])
    assert var_list == ["c", "k"]

10.2 Transform Tests¶

def test_strip_expectation_returns_list():
    """strip_expectation_operator always returns a list."""
    eq, vars = strip_expectation_operator("E_{y}(x + 1)")
    assert vars == ["y"]  # List, not string

    eq, vars = strip_expectation_operator("E_{y,z}(x + 1)")
    assert vars == ["y", "z"]

def test_group_aware_poststate_mapping():
    """Poststates at _cntn map to [t], not [t+1]."""
    symbol_groups = {'poststates': ['a'], 'controls': ['c']}

    result = perch_to_time(
        "a[_cntn] = w - c[_dcsn]",
        symbol_groups=symbol_groups,
        prestate_rename={},
        index_aliases={'_cntn': 1, '_dcsn': 0}
    )

    assert "a[t]" in result      # NOT a[t+1]
    assert "c[t]" in result

10.3 Round-Trip Test¶

def test_roundtrip():
    """
    Load dolo-plus model, translate, compile, verify functions exist.
    """
    model_plus = yaml_import(DOLOPLUS_PATH, compile_functions=False)
    model_dolo = translate_doloplus_to_dolo(model_plus)

    # Should compile without error
    funcs = model_dolo.__original_gufunctions__

    # Required blocks exist
    assert 'half_transition' in model_dolo.equations
    assert 'reverse_state' in model_dolo.equations
    assert 'expectation' in model_dolo.equations
    assert 'direct_response_egm' in model_dolo.equations

10.4 Equation Content Tests¶

def test_half_transition_uses_time_subscripts():
    """g_ad translation uses time subscripts, not perch tags."""
    model_dolo = translate_doloplus_to_dolo(model_plus)
    eq = model_dolo.equations['half_transition'][0]

    assert '[t]' in eq or '[t-1]' in eq
    assert '[_arvl]' not in eq
    assert '[_dcsn]' not in eq

def test_expectation_synthesis():
    """Expectation block combines integrand + factor + discount."""
    model_dolo = translate_doloplus_to_dolo(model_plus)
    eq = model_dolo.equations['expectation'][0]

    # Should have mr[t] = β*(...)*r form
    assert 'mr[t]' in eq
    assert 'β' in eq or 'beta' in eq.lower()
    # Integrand should be shifted to t+1
    assert '[t+1]' in eq

def test_poststate_at_cntn_uses_t_not_t_plus_1():
    """Poststate a[_cntn] becomes a[t], not a[t+1]."""
    model_dolo = translate_doloplus_to_dolo(model_plus)

    # In auxiliary_direct_egm: a[_cntn] = w[_dcsn] - c[_dcsn]
    eq = model_dolo.equations['auxiliary_direct_egm'][0]
    assert 'a[t]' in eq
    assert 'a[t+1]' not in eq

11. Dependencies¶

dolang (for parsing/transforming equation strings with proper operator handling)
dolo.compiler.model (for Model class)
dolo.compiler.model_import (for yaml_import)
lark (for AST transformations, via dolang)

12. Open Questions (for future specs)¶

Non-EGM methods: How to translate for VFI or time-iteration solvers?
Multi-stage models: How to handle inter-stage linking when stages differ?
Branching stages: How to represent alternative shock realizations?
Shadow value semantics: What exactly does dV[slot] differentiate w.r.t.?
Multiple shocks same timing: How to handle E_{w1,w2}(...) when two shocks arrive together?

13. Structural Insights from bellman-dev¶

13.1 Model-Level Translation Pattern¶

The bellman-dev codebase confirms that translation should happen at the model representation level, not YAML level:

dolo-plus YAML → yaml_import() → Model Object → Translator → Dolo Model Object → Compile

13.2 EquationMap Pattern¶

bellman-dev uses an EquationMap class to map abstract equation types to concrete Dolo block names:

# From equation_map_dolo_egm.py
standard_mappings = [
    (EquationType.transition_arvl_to_dcsn, "half_transition"),
    (EquationType.transition_cntn_to_dcsn, "reverse_state"),
    (EquationType.expected, "expectation"),
    (EquationType.control_from_state_cntn, "direct_response_egm"),
]

13.3 Dolo Solver Function Handles¶

The EGM solver (egm.py) expects specific compiled functions:

funs = model.__original_gufunctions__
gt = funs["half_transition"]      # Arrival → decision transition
aτ = funs["reverse_state"]        # Continuation → decision (inverse)
τ  = funs["direct_response_egm"]  # EGM policy function
h  = funs["expectation"]          # Expectation integrand

13.4 Slot Normalization Insight¶

Concept	dolo-plus Slot	Dolo Slot
Current-period poststate	`+1` (via `_cntn`)	`0` (uses `[t]`)
Previous-period poststate	`-1` (via prestate at `_arvl`)	`-1`

14. References¶

doloplus-foundations.md Section 5 (mapping table)
packages/dolang/dolang/grammar.lark (operator grammar)
packages/dolo/examples/models/consumption_savings_iid_egm_doloplus.yaml
packages/dolo/examples/models/consumption_savings_iid_egm.yaml (target format)

15. Implementation Notes (Retrospective)¶

Added January 2026 after implementation complete.

15.1 Architecture Differences from Spec¶

The actual implementation differs from the proposed class-based architecture:

Spec Proposed	Actual Implementation
`DoloplusTranslator` base class	Functional API with `doloplus_to_dolo()`
`EGMTranslator`, `VFITranslator` subclasses	Single translator with external mapping tables
Hardcoded mappings in translator classes	YAML mapping tables (user-configurable)

Rationale: The mapping-table approach is more flexible and allows users to customize translation without modifying Python code.

15.2 Translation Pipeline¶

The implemented pipeline (in core.py):

1. validate_model()     → Check dolo_plus block, dialect, version
2. build_context()      → Extract symbol groups, index aliases, eq mappings
3. translate_symbols()  → Filter/transform symbol groups per mapping tables
4. translate_equations() → Transform equations using configurable pipelines
5. assemble_model()     → Build vanilla Dolo Model object

15.3 Key Functions¶

Function	Location	Purpose
`doloplus_to_dolo()`	`core.py`	Main entry point
`load_mapping_tables()`	`core.py`	Load YAML mapping config
`perch_to_time()`	`transforms.py`	Convert perch tags to time indices
`shift_perch_tags()`	`transforms.py`	Shift perch offsets for expectations
`rename_symbols()`	`transforms.py`	Apply symbol renames (e.g., `dV` → `mr`)
`extract_rhs()`	`transforms.py`	Get RHS of equation
`synthesize_expectation()`	`core.py`	Build expectation block from movers

15.4 Mapping Tables Structure¶

The mapping tables YAML must contain:

PERCH_TO_TIME_OFFSET:     # Required: perch tag → time offset
TRANSITION_TO_BLOCK:      # Required: transition name → Dolo block
SYMBOL_GROUPS_KEEP:       # Required: which groups to keep
PIPELINES:                # Required: transformation pipelines
  transitions: [...]
  movers: [...]
SYMBOL_RENAMES:           # Optional: symbol name substitutions
EXPECTATION_CONFIG:       # Optional: expectation synthesis config
ARBITRAGE_CONFIG:         # Optional: arbitrage synthesis config

15.5 Test Coverage¶

Tests located in: - packages/dolo/dolo/tests/test_doloplus_import_parse_only.py - Integration tested via example models in explore/ folder

15.6 Example Models¶

Working dolo-plus → Dolo translation examples: - packages/dolo/examples/models/doloplus/consumption_portfolio_b2_doloplus/ - packages/dolo/examples/models/doloplus/consumption_portfolio_b3_doloplus/