feat: axis masking via cross-axis constraints (#31) by jc-macdonald · Pull Request #40 · ACCIDDA/op_system

jc-macdonald · 2026-03-25T12:09:12Z

Adds a constraints block to filter Cartesian product expansion, avoiding empty compartments.

Issues

Constraint parser and validation (_normalize_constraints) #36 Constraint parser and validation ✅
Constraint filtering for state, alias, and transition expansion #37 Constraint filtering for state, alias, and transition expansion
Constraint-aware sum_over/integrate_over and end-to-end tests #38 Constraint-aware sum_over/integrate_over and end-to-end tests
Document constraints block in README and guides #39 Document constraints block in README and guides

Refs #31

Add cross-axis constraint normalization infrastructure: - _validate_constraint_axes(): validates axes list per constraint rule - _resolve_constraint_mode(): determines allow/exclude mode - _validate_constraint_rule(): validates individual rule mappings - _normalize_constraints(): main entry point, parses and validates the constraints block from a spec dict Each constraint rule must reference at least two axes, specify either 'allow' or 'exclude' (not both), and only reference known axes/coords. Also adds PLC2701 to test per-file-ignores in pyproject.toml to allow importing private helpers for direct unit testing. Tests: 14 new tests covering happy paths and error cases. Refs: #31

TimothyWillard

Made a few comments where I think improvements would be nice. I also think some of this validation stuff can be simplified in the future by leveraging pydantic.

TimothyWillard · 2026-03-26T17:04:57Z

+def _validate_constraint_axes(
+    rule_axes_raw: object,
+    *,
+    idx: int,
+    axis_names: set[str],
+) -> tuple[list[str], set[str]]:
+    """Validate and return the axes list for a single constraint rule.
+
+    Returns:
+        Tuple of (ordered axis name list, axis name set).
+    """
+    if not isinstance(rule_axes_raw, (list, tuple)) or len(rule_axes_raw) < 2:
+        _raise_invalid_rhs_spec(
+            detail=(
+                f"constraints[{idx}].axes must be a list of at least two axis names"
+            )
+        )
+    rule_axes: list[str] = []
+    seen: set[str] = set()
+    for j, ax_name in enumerate(rule_axes_raw):
+        if not isinstance(ax_name, str) or not ax_name.strip():
+            _raise_invalid_rhs_spec(
+                detail=f"constraints[{idx}].axes[{j}] must be a non-empty string"
+            )
+        ax_s = ax_name.strip()
+        if ax_s not in axis_names:
+            _raise_invalid_rhs_spec(
+                detail=f"constraints[{idx}].axes references unknown axis {ax_s!r}"
+            )
+        if ax_s in seen:
+            _raise_invalid_rhs_spec(
+                detail=f"constraints[{idx}].axes contains duplicate axis {ax_s!r}"
+            )
+        seen.add(ax_s)
+        rule_axes.append(ax_s)
+    return rule_axes, seen


I think the annotation rule_axes_raw: object isn't quite right, I think it could be rule_axes_raw: Sequence[str]? See https://docs.python.org/3/library/collections.abc.html#collections.abc.Sequence.

I think it would be helpful if the docstring could also contain a description of the arguments, I know this is not publicly exported, but helpful for future development work.

Similar to (2), I think this is a prime candidate of a function that could have an examples section. In fact since this function is just doing some string manipulation and checking I think the doctest could replace unit tests for this particular function.

TimothyWillard · 2026-03-26T17:07:13Z

+        if not isinstance(ax_name, str) or not ax_name.strip():
+            _raise_invalid_rhs_spec(
+                detail=f"constraints[{idx}].axes[{j}] must be a non-empty string"
+            )
+        ax_s = ax_name.strip()


Suggested change

if not isinstance(ax_name, str) or not ax_name.strip():

_raise_invalid_rhs_spec(

detail=f"constraints[{idx}].axes[{j}] must be a non-empty string"

)

ax_s = ax_name.strip()

if not isinstance(ax_name, str) or not (ax_s := ax_name.strip()):

_raise_invalid_rhs_spec(

detail=f"constraints[{idx}].axes[{j}] must be a non-empty string"

)

To avoid doing this operation twice. Or just move the ax_s = ax_name.strip() line before this conditional.

TimothyWillard · 2026-03-26T17:08:29Z

+def _resolve_constraint_mode(
+    entry_map: Mapping[str, Any], *, idx: int
+) -> tuple[str, list[Any]]:
+    """Determine allow/exclude mode and return raw rules list.
+
+    Returns:
+        Tuple of (mode string, raw rule list).
+    """


I think my same comments about the _validate_constraint_axes docstring also apply here. See https://github.com/ACCIDDA/op_system/pull/40/changes#r2996408053.

TimothyWillard · 2026-03-26T17:13:37Z

+    Args:
+        rule: Raw rule object (expected to be a mapping).
+        label: Human-readable label for error messages (e.g.
+            ``constraints[0].allow[1]``).


Not a big deal, but this project uses mkdocs so documentation is rendered as markdown not rst so double backticks not required. Not applicable here since this is a private function and won't have it's documentation rendered, but just a heads up.

TimothyWillard · 2026-03-26T17:16:40Z

+        out.append({
+            "axes": tuple(rule_axes),
+            "mode": mode,
+            "rules": tuple(validated_rules),
+        })


Since the keys are fixed this actually seems like a prime place for using either a typed dict or named tuple (depending on if the entries need to be mutable or not). This would provide typing safety around accessing members of the object. See https://typing.python.org/en/latest/spec/typeddict.html or https://typing.python.org/en/latest/spec/namedtuples.html.

TimothyWillard · 2026-03-26T17:19:29Z

+    assert len(result) == 1
+    rule = result[0]
+    assert rule["axes"] == ("age", "vax")
+    assert rule["mode"] == "exclude"
+    assert rule["rules"][0] == {"age": ["u65"], "vax": ["dose1", "dose2"]}


I think this can be simplified down to 1 assert statement with the exact expected output object. Pytest introspection abilities can handle traversing dicts/lists to check for nested equality and point out differences.

- Thread 1: annotate rule_axes_raw as Sequence[str], add Args/Examples to _validate_constraint_axes docstring - Thread 2: use walrus operator to avoid double .strip() - Thread 3: add Args to _resolve_constraint_mode docstring - Thread 4: single backticks (Markdown) instead of double (rst) - Thread 5: replace plain dict with ConstraintRule NamedTuple - Thread 6: collapse test asserts into single expected-object comparison

jc-macdonald · 2026-04-07T06:18:05Z

All six review threads have been addressed in 894ed52:

Thread 1 — rule_axes_raw annotated as Sequence[str]; added Args and Examples sections to _validate_constraint_axes docstring.
Thread 2 — Walrus operator to avoid double .strip() call.
Thread 3 — Added Args section to _resolve_constraint_mode docstring.
Thread 4 — Single backticks (Markdown) replacing double backticks (rst) in all constraint docstrings.
Thread 5 — Introduced ConstraintRule NamedTuple replacing plain dict return type.
Thread 6 — Collapsed multi-assert tests into single expected-object comparisons.

Root CI passes clean (ruff, mypy, 62/62 tests, docs). The provider mypy failure (StateChangeEnum attr-defined) is pre-existing and tracked in #51.

jc-macdonald added this to the Axis Masking (Cross-Axis Constraints) milestone Mar 25, 2026

jc-macdonald self-assigned this Mar 25, 2026

jc-macdonald added the enhancement New feature or request label Mar 25, 2026

jc-macdonald requested review from TimothyWillard and pearsonca March 25, 2026 12:10

jc-macdonald mentioned this pull request Mar 25, 2026

Constraint parser and validation (_normalize_constraints) #36

Closed

TimothyWillard approved these changes Mar 26, 2026

View reviewed changes

jc-macdonald marked this pull request as ready for review April 7, 2026 06:16

jc-macdonald mentioned this pull request Apr 7, 2026

Fix provider package mypy error: StateChangeEnum not found in flepimop2.typing #51

Closed

jc-macdonald merged commit 5fb995f into main Apr 7, 2026
5 checks passed

jc-macdonald deleted the feature/31-axis-masking branch April 7, 2026 06:33

jc-macdonald mentioned this pull request Apr 7, 2026

InsightNet April 2026 — Tracking Issue ACCIDDA/flepimop2#200

Closed

61 tasks

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

feat: axis masking via cross-axis constraints (#31)#40

feat: axis masking via cross-axis constraints (#31)#40
jc-macdonald merged 2 commits intomainfrom
feature/31-axis-masking

jc-macdonald commented Mar 25, 2026

Uh oh!

TimothyWillard left a comment

Uh oh!

TimothyWillard Mar 26, 2026

Uh oh!

TimothyWillard Mar 26, 2026

Uh oh!

TimothyWillard Mar 26, 2026

Uh oh!

TimothyWillard Mar 26, 2026

Uh oh!

TimothyWillard Mar 26, 2026

Uh oh!

TimothyWillard Mar 26, 2026

Uh oh!

jc-macdonald commented Apr 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants

Conversation

jc-macdonald commented Mar 25, 2026

Issues

Uh oh!

TimothyWillard left a comment

Choose a reason for hiding this comment

Uh oh!

TimothyWillard Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

TimothyWillard Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

TimothyWillard Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

TimothyWillard Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

TimothyWillard Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

TimothyWillard Mar 26, 2026

Choose a reason for hiding this comment

Uh oh!

jc-macdonald commented Apr 7, 2026

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

2 participants