Improve constraint and update solver test to use new AddParentDataset

Author: bouweandela

packages/climate-ref-core/src/climate_ref_core/constraints.py

@@ -503,8 +503,8 @@ def apply(self, group: pd.DataFrame, data_catalog: pd.DataFrame) -> pd.DataFrame
             if set(all_parent_facets).issubset(data_catalog.keys()):
                 break
         else:
-            # No matching parent_facets
-            return group
+            msg = "Missing facets in data catalog to determine parent dataset"
+            raise ValueError(msg)
 
         # Remove datasets that do not have all parent facets set.
         valid_group = group[all_parent_facets].drop_duplicates().dropna(axis="columns")
@@ -519,15 +519,15 @@ def apply(self, group: pd.DataFrame, data_catalog: pd.DataFrame) -> pd.DataFrame
                 # Add the latest version of the dataset to the selection.
                 parent_dataset = data_catalog[parent_select]
                 parent_dataset = parent_dataset[parent_dataset["version"] == parent_dataset["version"].max()]
-                select[parent_dataset.index] = True
+                select.loc[parent_dataset.index] = True
             else:
-                # Drop the dataset if no parent dataset is found.
+                # Drop the child dataset if no parent dataset is found.
                 logger.debug(
                     f"Constraint {self} not satisfied because no parent dataset found for "
                     f"{', '.join(f'{k}={v}' for k, v in dataset.items())}"
                 )
-                dataset_facets = {k: (dataset[k],) for k in dataset}
-                dataset_select = valid_group[list(dataset_facets)].isin(dataset_facets).all(axis="columns")
-                select[dataset_select] = False
+                child_facets = {k: (dataset[k],) for k in dataset}
+                child_select = valid_group[list(child_facets)].isin(child_facets).all(axis="columns")
+                select.loc[child_select[child_select].index] = False
 
         return data_catalog[select]

packages/climate-ref/tests/unit/test_solver.py

@@ -20,7 +20,7 @@
     solve_executions,
     solve_required_executions,
 )
-from climate_ref_core.constraints import AddSupplementaryDataset, RequireFacets, SelectParentExperiment
+from climate_ref_core.constraints import AddParentDataset, AddSupplementaryDataset, RequireFacets
 from climate_ref_core.datasets import SourceDatasetType
 from climate_ref_core.diagnostics import DataRequirement, FacetFilter
 
@@ -188,35 +188,40 @@ def test_solver_build_from_db(self, solver):
             DataRequirement(
                 source_type=SourceDatasetType.CMIP6,
                 filters=(FacetFilter(facets={"variable_id": ("tas", "pr")}),),
-                constraints=(SelectParentExperiment(),),
+                constraints=(AddParentDataset(),),
                 group_by=("variable_id", "experiment_id"),
             ),
             pd.DataFrame(
                 {
-                    "variable_id": ["tas", "tas"],
                     "experiment_id": ["ssp119", "historical"],
+                    "grid_label": ["gn", "gn"],
                     "parent_experiment_id": ["historical", "none"],
+                    "parent_source_id": ["A", "A"],
+                    "parent_variant_label": ["r1i1p1f1", "none"],
+                    "source_id": ["A", "A"],
+                    "table_id": ["Amon", "Amon"],
+                    "variable_id": ["tas", "tas"],
+                    "variant_label": ["r1i1p1f1", "r1i1p1f1"],
+                    "version": ["v20210101", "v20220101"],
                 }
             ),
             {
                 (("variable_id", "tas"), ("experiment_id", "ssp119")): pd.DataFrame(
                     {
+                        "experiment_id": ["ssp119", "historical"],
+                        "grid_label": ["gn", "gn"],
+                        "parent_experiment_id": ["historical", "none"],
+                        "parent_source_id": ["A", "A"],
+                        "parent_variant_label": ["r1i1p1f1", "none"],
+                        "source_id": ["A", "A"],
+                        "table_id": ["Amon", "Amon"],
                         "variable_id": ["tas", "tas"],
-                        "experiment_id": ["historical", "ssp119"],
-                    },
-                    # The order of the rows is not guaranteed
-                    index=[1, 0],
-                ),
-                (("variable_id", "tas"), ("experiment_id", "historical")): pd.DataFrame(
-                    {
-                        "variable_id": ["tas", "tas"],
-                        "experiment_id": ["historical"],
+                        "variant_label": ["r1i1p1f1", "r1i1p1f1"],
+                        "version": ["v20210101", "v20220101"],
                     },
-                    # The order of the rows is not guaranteed
-                    index=[1, 0],
+                    index=[0, 1],
                 ),
             },
-            marks=[pytest.mark.xfail(reason="Parent experiment not implemented")],
             id="parent",
         ),
         pytest.param(