🐛 fix branching parameter update semantics and finalize 2.0.1 release notes

Author: LaurenzBeck

CHANGELOG.md

@@ -8,6 +8,15 @@ This project uses [*towncrier*](https://towncrier.readthedocs.io/) to build the
 
 <!-- towncrier release notes start -->
 
+## [2.0.1](https://github.com/Infineon/StreamGen/tree/2.0.1) - 2026-03-31
+
+### 🐛 Fixed
+
+- fixed `SamplingTree.update` and `SamplingTree.set_update_step` so shared parameters in branching trees are updated exactly once per step.
+- added regression tests for branching update behavior to ensure schedule progression stays consistent.
+- stabilized plotting tests by forcing a non-interactive matplotlib backend during test runs.
+
+
 ## [2.0.0](https://github.com/Infineon/StreamGen/tree/2.0.0) - 2026-03-05
 
 ### ➖ Removed

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "streamgen"
-version = "2.0.0"
+version = "2.0.1"
 description = "🌌 a framework for generating streams of labeled data."
 authors = [
     {name = "Laurenz Farthofer", email = "laurenz@hey.com"}

streamgen/samplers/tree.py

@@ -325,10 +325,33 @@ def collect(self, num_samples: int, strategy: SamplingStrategy | SamplingStrateg
 
         return tuple(map(self.collate_func, zip(*samples, strict=True))) if self.collate_func else samples
 
+    def _get_unique_parameters(self) -> list[Parameter]:
+        """Collect every unique parameter object referenced by the tree nodes."""
+        unique_parameters = []
+        seen_parameter_ids = set()
+
+        for node in anytree.PreOrderIter(self.root):
+            if isinstance(node, BranchingNode) and node.probs is not None:
+                parameter_id = id(node.probs)
+                if parameter_id not in seen_parameter_ids:
+                    seen_parameter_ids.add(parameter_id)
+                    unique_parameters.append(node.probs)
+
+            if isinstance(node, TransformNode) and node.params is not None:
+                for parameter_name in node.params.parameter_names:
+                    parameter = node.params[parameter_name]
+                    parameter_id = id(parameter)
+                    if parameter_id in seen_parameter_ids:
+                        continue
+                    seen_parameter_ids.add(parameter_id)
+                    unique_parameters.append(parameter)
+
+        return unique_parameters
+
     def update(self) -> None:
         """🆙 updates every parameter."""
-        for node in self.nodes:
-            node.update()
+        for parameter in self._get_unique_parameters():
+            parameter.update()
 
     def set_update_step(self, idx: int) -> None:
         """🕐 updates every parameter to a certain update step using `param[idx]`.
@@ -339,8 +362,8 @@ def set_update_step(self, idx: int) -> None:
         Returns:
             None: this function mutates `self`
         """
-        for node in self.nodes:
-            node.set_update_step(idx)
+        for parameter in self._get_unique_parameters():
+            parameter[idx]
 
     def get_params(self) -> ParameterStore | None:
         """⚙️ collects parameters from every node.

tests/conftest.py

@@ -1 +1,6 @@
 """🗃️ fixtures available in all tests."""
+
+import matplotlib
+
+# Use a non-interactive backend in tests to avoid GUI/Tk dependencies.
+matplotlib.use("Agg", force=True)

tests/unit/streamgen_branch_update_regression_test.py

@@ -0,0 +1,115 @@
+"""🧪 regression tests for branching update behavior."""
+# ruff: noqa: S101, D103, ANN001, ANN201
+
+from collections import Counter
+
+from streamgen.nodes import TransformNode
+from streamgen.parameter import Parameter
+from streamgen.samplers.tree import SamplingTree
+from streamgen.transforms import noop
+
+
+# ---------------------------------------------------------------------------- #
+# *                             helper functions                               #
+# ---------------------------------------------------------------------------- #
+
+
+def _build_minimal_branch_tree() -> SamplingTree:
+    """Builds a small tree where one scoped parameter is attached multiple times.
+
+    Layout:
+        root -> BranchingNode(left|right) -> shared_scope
+
+    The node after the branching node is deep-copied into each branch by
+    streamgen's tree construction shorthand. All copies fetch params from the
+    same scope name ("shared_scope").
+    """
+    nodes = [
+        TransformNode(noop, name="root"),
+        {
+            "name": "selector",
+            "left": [TransformNode(noop, name="left_leaf")],
+            "right": [TransformNode(noop, name="right_leaf")],
+        },
+        TransformNode(noop, name="shared_scope"),
+    ]
+
+    params = {
+        "shared_scope": {
+            "sweep": {
+                "schedule": [10, 20, 30],
+                "strategy": "hold",
+            },
+        },
+    }
+
+    return SamplingTree(nodes=nodes, params=params, rng=0)
+
+
+def _collect_schedule_values(mode: str) -> list[int]:
+    """Collects two update steps from either tree.update or tree.params.update."""
+    tree = _build_minimal_branch_tree()
+    parameter = tree.get_params()["shared_scope.sweep"]
+    values = [int(parameter.value)]
+
+    for _ in range(2):
+        if mode == "tree":
+            tree.update()
+        elif mode == "params":
+            tree.params.update()
+        else:
+            raise ValueError(f"unknown mode: {mode}")
+        values.append(int(parameter.value))
+
+    return values
+
+
+# ---------------------------------------------------------------------------- #
+# *                                   tests                                    #
+# ---------------------------------------------------------------------------- #
+
+
+def test_tree_update_calls_each_parameter_once_in_branching_tree(monkeypatch) -> None:
+    """Tests that tree.update touches each parameter object once."""
+    tree = _build_minimal_branch_tree()
+    target = tree.get_params()["shared_scope.sweep"]
+
+    call_counts: Counter[int] = Counter()
+    original_update = Parameter.update
+
+    def wrapped_update(self, *args, **kwargs):
+        call_counts[id(self)] += 1
+        return original_update(self, *args, **kwargs)
+
+    monkeypatch.setattr(Parameter, "update", wrapped_update)
+    tree.update()
+
+    assert call_counts[id(target)] == 1
+
+
+def test_parameter_store_update_calls_each_parameter_once(monkeypatch) -> None:
+    """Tests that ParameterStore.update touches each parameter object once."""
+    tree = _build_minimal_branch_tree()
+    target = tree.get_params()["shared_scope.sweep"]
+
+    call_counts: Counter[int] = Counter()
+    original_update = Parameter.update
+
+    def wrapped_update(self, *args, **kwargs):
+        call_counts[id(self)] += 1
+        return original_update(self, *args, **kwargs)
+
+    monkeypatch.setattr(Parameter, "update", wrapped_update)
+    tree.params.update()
+
+    assert call_counts[id(target)] == 1
+
+
+def test_tree_update_matches_parameter_store_update_progression() -> None:
+    """Tests that tree.update progresses one step and matches params.update."""
+    tree_values = _collect_schedule_values("tree")
+    params_values = _collect_schedule_values("params")
+
+    assert tree_values == [10, 20, 30]
+    assert params_values == [10, 20, 30]
+    assert tree_values == params_values