Allow more deeply nested mappings in params while retaining tight typing.

[hmgaudecker]

Summary

Small update in order to work sensibly with GETTSIM functions.

• Add MappingLeaf, a JAX-pytree-compatible wrapper that prevents flatten_regime_namespace from recursing into its contents, allowing nested parameter dicts to be passed through the params pipeline as opaque leaves. Same for SequenceLeaf.
• Add an as_leaf function, wich is how these will typically be generated, make it available via lcm.params.as_leaf
• Teach _make_immutable / _make_mutable to handle MappingLeaf (converting inner data while preserving the wrapper)
• [somewhat unrelated, just because this bug crept up in the model ] Add reproducer test for BUG: State variable only used in transitions causes ValueError in vmap_1d at last regime period #236 (state used only in transitions with terminal regime)

Test plan

• test_mapping_leaf.py — construction, repr, equality, JAX pytree round-trips, immutable/mutable conversions, flatten_regime_namespace leaf behavior

• test_model.py::test_state_only_in_transitions_with_terminal_regime — reproducer for BUG: State variable only used in transitions causes ValueError in vmap_1d at last regime period #236.

  Somewhat unrelated, just because the bug crept up in the model where we are already using this branch.

  This passes out of the box; must have been fixed along the way somewhere.

[timmens]

Very nice! I can see how this is useful/necessary for GETTSIM etc.

Only two minor remarks. The user would need to import NestedMappingParams and use that in her params definition, correct? In this case:

1. Should we potentially find a name that fits better with what this object actually does, i.e., tells PyLCM to treat this Mapping as a leaf (not sure about the wording yet, but NestedMappingParams confused me)?
2. Should this object live in either the top-level namespace or in lcm.params?

[hmgaudecker]

@timmens @MImmesberger -- do you like lcm.params.MappingLeaf ?

[timmens]

I think that is much better. I have one question (that might be stupid because I don't understand the application that well), but is there a reason why we don't simply want:

lcm.params.Leaf

In the sense that I have never read "MappingLeaf" before, but I have read "Leaf" frequently. And anyone who worked with JAX before will ---hopefully--- directly understand what is meant by "Leaf".

Or is solving the "more general" problem simply not necessary here, and we can restrict ourselves completely to the Mapping type?

Another thought that occurred to me, don't we actually want to export something like a as-function?

gettsim_params = {"a": [0, 1], "b": jnp.arange(4), "c": -1231}

lcm_params = {"d": {"e": 2, "f": 3}, "gettsim_params": lcm.params.as_leaf(gettsim_params)}

where lcm.params.as_leaf takes care of the processing and returns an lcm.params.Leaf object?

[hmgaudecker]

is there a reason why we don't simply want lcm.params.Leaf

Yes, we want to have a mapping that is a leaf, useful for typing to distinguish this from other leaves. Since it was cheap and symmetric, I added SequenceLeaf, that hopefully makes it clear!

as_leaf() is a great idea, implemented!

[hmgaudecker]

@MImmesberger, can you switch to .as_leaf? That should be the idiomatic way of using it, I am thinking about the classes more in terms of typing.

[hmgaudecker]

@timmens -- I thought you had approved already, but either I misremember or it got lost because of subsequent changes. Anything left?

[timmens]

No, looks good to me!