feature[design]: allow design plugin to interface with pytorch via Re…

…sult

CHANGELOG.md

@@ -21,6 +21,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Mobility models for `SemiconductorMedium`: `ConstantMobilityModel`, `CaugheyThomasMobility`
 - Bandgap narrowing models for `SemiconductorMedium`: `SlotboomBandGapNarrowing`
 - Generation-recombination models for `SemiconductorMedium`: `ShockleyReedHallRecombination`, `RadiativeRecombination`, `AugerRecombination`
+- Accessors and length functions implemented for `Result` class in design plugin.
 
 ### Changed
 - `ModeMonitor` and `ModeSolverMonitor` now use the default `td.ModeSpec()` with `num_modes=1` when `mode_spec` is not provided.
@@ -34,6 +35,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 - Validation error in inverse design plugin when simulation has no sources by adding a source existence check before validating pixel size.
 - System-dependent floating-point precision issue in EMEGrid validation.
 - Fixed magnitude of gradient computation in `CustomMedium` by accounting properly for full volume element when permittivity data is defined over less dimensions than the medium.
+- Fixed key ordering in design plugin when returning `Result` from an optimization method run.
 
 ## [2.8.0rc1] - 2024-12-17
 

tests/test_plugins/test_design.py

@@ -35,6 +35,28 @@
     "MethodParticleSwarm": ["MethodParticleSwarm_2_2", "MethodParticleSwarm_1_4"],
 }
 
+ACCESSOR_DATA = [
+    dict(
+        dims=("a", "b", "c"),
+        coords=((1, 0, 1), (2, 3, 4), (5, 1, 2)),
+        values=(1, 2, 3),
+        multi_val=False,
+    ),
+    dict(
+        dims=("a", "b", "c"),
+        coords=((5, 2, 0), (1, 2, 3), (3, 4, 1)),
+        values=((3, 4), (4, 6), (6, 9)),
+        multi_val=True,
+    ),
+    dict(
+        dims=("x", "y", "z", "t"),
+        coords=((1, 2, 3, 4), (5, 6, 7, 8)),
+        values=(5, 3),
+        multi_val=False,
+    ),
+    dict(dims=tuple("a"), coords=(tuple("b"), tuple("c")), values=("d", "e"), multi_val=False),
+]
+
 
 def emulated_batch_run(simulations, path_dir: str = None, **kwargs):
     data_dict = {task_name: run_emulated(sim) for task_name, sim in simulations.simulations.items()}
@@ -782,3 +804,67 @@ def count_float_non_td_post(res):
 
     # Check that the result and the aux are the same value
     assert all(df["output"] == df["aux_key_0"])
+
+
+@pytest.mark.parametrize("result_data", ACCESSOR_DATA)
+def test_result_accessor_and_len(result_data):
+    """Test the accessor for Result lines up with regular indexing into the coordinates and values
+    and that the length matches that of the coordinates and values. Finally, since we have defined
+    accessors and length, ensure we can cast the Result class to a list and tuple"""
+    dims = result_data["dims"]
+    coords = result_data["coords"]
+    values = result_data["values"]
+
+    result = tdd.Result(dims=dims, values=values, coords=coords)
+
+    for idx in range(0, len(values)):
+        coord, value = result[idx]
+
+        assert all(coord == coords[idx])
+
+        if result_data["multi_val"]:
+            assert all(value == values[idx])
+        else:
+            assert value == values[idx]
+
+    assert len(result) == len(coords)
+    assert len(result) == len(values)
+
+    result_list = list(result)
+    result_tuple = tuple(result)
+
+
+def test_dim_coord_alignment():
+    """Test that the dims coming out of the optimization method align with the ordering of the coordinates.
+    The ordering of these dimensions from the optimization can differ from the initial ordering of dimensions
+    input into the DesignSpace."""
+    bay_opt = tdd.MethodBayOpt(initial_iter=5, n_iter=2, seed=1)
+
+    parameters_sweep = [
+        [
+            tdd.ParameterFloat(name="radius", span=(1, 5)),
+            tdd.ParameterFloat(name="alpha", span=(-5, -1)),
+        ],
+        [
+            tdd.ParameterFloat(name="alpha", span=(-5, -1)),
+            tdd.ParameterFloat(name="radius", span=(1, 5)),
+        ],
+    ]
+
+    for parameters in parameters_sweep:
+        design_space = tdd.DesignSpace(
+            parameters=parameters, method=bay_opt, name="bay opt", task_name="bayesian"
+        )
+
+        def opt_fn(radius, alpha):
+            return radius**2 - alpha * radius + alpha**2 - 2 * radius + 2 * alpha
+
+        result = design_space.run(opt_fn)
+
+        # get the location of these parameters according to the result
+        radius_parameter_location = result.dims.index("radius")
+        alpha_parameter_location = result.dims.index("alpha")
+
+        # make sure each of the coordinates (as given by their dimension) make sense with their allowed spans
+        assert all(coord[radius_parameter_location] > 0 for coord in result.coords)
+        assert all(coord[alpha_parameter_location] < 0 for coord in result.coords)

tidy3d/plugins/design/design.py

@@ -124,7 +124,7 @@ def _package_run_results(
         fn_args_coords_T = list(map(list, zip(*fn_args_coords)))
 
         return Result(
-            dims=self.dims,
+            dims=tuple(fn_args[0].keys()),
             values=fn_values,
             coords=fn_args_coords_T,
             fn_source=fn_source,

tidy3d/plugins/design/result.py

@@ -427,3 +427,16 @@ def add(self, fn_args: Dict[str, float], value: Any) -> Result:
             return new_result
 
         return self.updated_copy(values=new_values, coords=new_coords)
+
+    def __len__(self):
+        """Implement len function to return the number of items in the result."""
+
+        return len(self.coords)
+
+    def __getitem__(self, data_index):
+        """Implement the accessor function to index into the coordinates and values of the result."""
+
+        features = self.coords[data_index]
+        labels = self.values[data_index]
+
+        return np.array(features), np.array(labels)