Fix test and mypy

invrs-io · Nov 17, 2023 · 7fb8020 · 7fb8020
1 parent 36319e8
commit 7fb8020
Show file tree

Hide file tree

Showing 2 changed files with 10 additions and 10 deletions.
diff --git a/src/invrs_gym/challenges/sorter/common.py b/src/invrs_gym/challenges/sorter/common.py
@@ -229,9 +229,9 @@ def response(
 
         spec = dataclasses.replace(
             self.spec,
-            thickness_cap=jnp.asarray(params[THICKNESS_CAP].array),
-            thickness_metasurface=jnp.asarray(params[THICKNESS_METASURFACE].array),
-            thickness_spacer=jnp.asarray(params[THICKNESS_SPACER].array),
+            thickness_cap=params[THICKNESS_CAP],  # type: ignore[arg-type]
+            thickness_metasurface=params[THICKNESS_METASURFACE],  # type: ignore[arg-type]
+            thickness_spacer=params[THICKNESS_SPACER],  # type: ignore[arg-type]
         )
         return simulate_sorter(
             density=params[DENSITY_METASURFACE],  # type: ignore[arg-type]
@@ -375,9 +375,9 @@ def simulate_sorter(
 
     layer_thicknesses = [
         jnp.zeros(()),  # Ambient
-        jnp.asarray(spec.thickness_cap),
-        jnp.asarray(spec.thickness_metasurface),
-        jnp.asarray(spec.thickness_spacer),
+        spec.thickness_cap.array,
+        spec.thickness_metasurface.array,
+        spec.thickness_spacer.array,
         jnp.asarray(spec.offset_monitor_substrate),  # Substrate
     ]
 

diff --git a/tests/challenges/sorter/test_common.py b/tests/challenges/sorter/test_common.py
@@ -75,10 +75,10 @@ def test_density_has_expected_properties(self, minimum_width, minimum_spacing):
         self.assertEqual(
             set(params.keys()),
             {
-                "metasurface_density",
-                "metasurface_thickness",
-                "cap_thickness",
-                "spacer_thickness",
+                "density_metasurface",
+                "thickness_metasurface",
+                "thickness_cap",
+                "thickness_spacer",
             },
         )