WIP

Author: marcorudolphflex

tests/test_components/autograd/numerical/test_autograd_polyslab_trianglemesh_numerical.py

@@ -16,6 +16,7 @@
     dimension_permutation,
     finite_difference,
     make_base_simulation,
+    run_parameter_simulations,
     squeeze_dimension,
 )
 from tidy3d import config
@@ -113,34 +114,34 @@ def make_polyslab_geometry(params, box_center, axis: int) -> td.PolySlab:
     return td.PolySlab(vertices=tuple(vertices), slab_bounds=slab_bounds, axis=axis)
 
 
-def run_parameter_simulations(
+def run_parameter_simulations2(
     parameter_sets: list[anp.ndarray],
     make_geometry,
     box_center,
-    tag: str,
     base_sim: td.Simulation,
     fom,
     tmp_path,
     *,
     local_gradient: bool,
 ):
-    simulation_dict = {}
+    simulations = []
 
-    for idx, param_values in enumerate(parameter_sets):
+    for param_values in parameter_sets:
         geometry = make_geometry(param_values, box_center)
         structure = td.Structure(
             geometry=geometry,
             medium=td.Medium(permittivity=PERMITTIVITY),
         )
         sim = base_sim.updated_copy(structures=[structure], validate=False)
-        simulation_dict[f"sim_{idx}"] = sim
+        simulations.append(sim)
 
-    sim_data_map = web.run(
-        simulation_dict,
+    sim_data = web.run(
+        simulations,
         local_gradient=local_gradient,
+        path=tmp_path,
         verbose=VERBOSE,
     )
-    sim_fom = [fom(sim_data_map[key]) for key in simulation_dict]
+    sim_fom = [fom(data) for data in sim_data]
     if len(sim_fom) == 1:
         sim_fom = sim_fom[0]
     return sim_fom

tidy3d/components/data/data_array.py

@@ -78,18 +78,28 @@ class DataArray(xr.DataArray):
     _data_attrs: dict[str, str] = {}
 
     def __init__(self, data, *args: Any, **kwargs: Any) -> None:
+        # convert numpy object arrays that contain autograd boxes; keep other types as-is
+        data = self._maybe_convert_object_boxes(data)
+
         # if data is a vanilla autograd box, convert to our box
         if isbox(data) and not is_tidy_box(data):
             data = TidyArrayBox.from_arraybox(data)
         # do the same for xr.Variable or xr.DataArray type
-        elif (
-            isinstance(data, (xr.Variable, xr.DataArray))
-            and isbox(data.data)
-            and not is_tidy_box(data.data)
-        ):
-            data.data = TidyArrayBox.from_arraybox(data.data)
+        elif isinstance(data, (xr.Variable, xr.DataArray)):
+            if isbox(data.data) and not is_tidy_box(data.data):
+                data.data = TidyArrayBox.from_arraybox(data.data)
         super().__init__(data, *args, **kwargs)
 
+    @staticmethod
+    def _maybe_convert_object_boxes(data):
+        """Convert object arrays of autograd boxes into ArrayBox instances."""
+
+        if isinstance(data, np.ndarray) and data.dtype == np.object_ and data.size:
+            # only convert if at least one element is an autograd tracer
+            if any(isbox(item) for item in data.flat):
+                return anp.array(data.tolist())
+        return data
+
     @classmethod
     def __get_validators__(cls):
         """Validators that get run when :class:`.DataArray` objects are added to pydantic models."""