get_measurements_for_experiment

Author: dweindl

petab/v2/problem.py

@@ -1134,6 +1134,20 @@ def get_changes_for_period(
             )
         )
 
+    def get_measurements_for_experiment(
+        self, experiment: core.Experiment
+    ) -> list[core.Measurement]:
+        """Get the measurements for a given experiment.
+
+        :param experiment: The experiment to get the measurements for.
+        :return: A list of measurements for the given experiment.
+        """
+        return [
+            measurement
+            for measurement in self.measurements
+            if measurement.experiment_id == experiment.id
+        ]
+
 
 class ModelFile(BaseModel):
     """A file in the PEtab problem configuration."""

tests/v2/test_problem.py

@@ -223,3 +223,56 @@ def test_get_changes_for_period():
     )
     assert problem.get_changes_for_period(p1) == [ch1, ch2]
     assert problem.get_changes_for_period(p2) == [ch3]
+
+
+def test_get_measurements_for_experiment():
+    """Test getting measurements for an experiment."""
+    problem = Problem()
+    problem += Condition(
+        id="condition1",
+        changes=[Change(target_id="target1", target_value=1.0)],
+    )
+    problem += Condition(
+        id="condition2",
+        changes=[Change(target_id="target2", target_value=2.0)],
+    )
+
+    e1 = Experiment(
+        id="exp1",
+        periods=[
+            ExperimentPeriod(id="p1", time=0, condition_ids=["condition1"]),
+        ],
+    )
+    e2 = Experiment(
+        id="exp2",
+        periods=[
+            ExperimentPeriod(id="p2", time=1, condition_ids=["condition2"]),
+        ],
+    )
+    problem += e1
+    problem += e2
+
+    m1 = Measurement(
+        observable_id="observable1",
+        experiment_id="exp1",
+        time=0,
+        measurement=10.0,
+    )
+    m2 = Measurement(
+        observable_id="observable2",
+        experiment_id="exp1",
+        time=1,
+        measurement=20.0,
+    )
+    m3 = Measurement(
+        observable_id="observable3",
+        experiment_id="exp2",
+        time=1,
+        measurement=30.0,
+    )
+    problem += m1
+    problem += m2
+    problem += m3
+
+    assert problem.get_measurements_for_experiment(e1) == [m1, m2]
+    assert problem.get_measurements_for_experiment(e2) == [m3]