Mapping over simulation collection is working

Author: AstroPatty

src/opencosmo/collection/simulation/simulation.py

@@ -82,14 +82,30 @@ def make_schema(self) -> DataSchema:
 
         return schema
 
-    def __map(self, method, *args, **kwargs):
+    def __map(self, method, *args, construct=True, **kwargs):
         """
         This type of collection will only ever be constructed if all the underlying
         datasets have the same data type, so it is always safe to map operations
         across all of them.
         """
-        output = {k: getattr(v, method)(*args, **kwargs) for k, v in self.items()}
-        return SimulationCollection(output)
+        regular_kwargs = {}
+        mapped_kwargs = {}
+        known_datasets = set(self.keys())
+        for name, value in kwargs.items():
+            if isinstance(value, dict) and set(value.keys()) == known_datasets:
+                mapped_kwargs[name] = value
+            else:
+                regular_kwargs[name] = value
+
+        output = {}
+        for name, dataset in self.items():
+            dataset_mapped_kwargs = {key: kw[name] for key, kw in mapped_kwargs.items()}
+            output[name] = getattr(dataset, method)(
+                *args, **regular_kwargs, **dataset_mapped_kwargs
+            )
+        if construct:
+            return SimulationCollection(output)
+        return output
 
     def __map_attribute(self, attribute):
         return {k: getattr(v, attribute) for k, v in self.items()}
@@ -321,21 +337,16 @@ def evaluate(
         else:
             datasets = list(datasets)
 
-        results = {
-            ds_name: self[ds_name].evaluate(
-                func,
-                vectorize=vectorize,
-                insert=insert,
-                format=format,
-                **evaluate_kwargs,
-            )
-            for ds_name in datasets
-        }
-        if not insert:
-            return results
-        else:
-            output = {**self, **results}
-            return SimulationCollection(output)
+        results = self.__map(
+            "evaluate",
+            func,
+            vectorize=vectorize,
+            insert=insert,
+            format=format,
+            construct=insert,
+            **evaluate_kwargs,
+        )
+        return results
 
     def with_units(self, convention: str) -> Self:
         """

src/opencosmo/dataset/visit.py

@@ -72,7 +72,7 @@ def __make_output(
         first_values = {name: first_values}
     storage = {}
     for name, value in first_values.items():
-        shape = (n,)
+        shape: tuple[int, ...] = (n,)
         dtype = type(value)
         if isinstance(value, np.ndarray):
             shape = shape + value.shape

test/test_collection.py

@@ -539,6 +539,37 @@ def fof_px(fof_halo_mass, fof_halo_com_vx):
         )
 
 
+def test_simulation_collection_evaluate_map_kwarg(multi_path):
+    collection = oc.open(multi_path)
+
+    def fof_px(fof_halo_mass, fof_halo_com_vx, random_value, other_value):
+        return fof_halo_mass * fof_halo_com_vx * random_value / other_value
+
+    random_data = {
+        key: np.random.randint(0, 10, len(ds)) for key, ds in collection.items()
+    }
+    random_val = 3.0
+
+    output = collection.evaluate(
+        fof_px,
+        vectorize=True,
+        insert=False,
+        format="numpy",
+        random_value=random_data,
+        other_value=random_val,
+    )
+    for ds_name, ds in collection.items():
+        assert "fof_px" not in ds.columns
+        data = ds.select(["fof_halo_mass", "fof_halo_com_vx"]).get_data("numpy")
+        assert np.all(
+            output[ds_name]["fof_px"]
+            == data["fof_halo_mass"]
+            * data["fof_halo_com_vx"]
+            * random_data[ds_name]
+            / random_val
+        )
+
+
 def test_simulation_collection_add(multi_path):
     collection = oc.open(multi_path)
     ds_name = next(iter(collection.keys()))