Thickness upper bound

src/invrs_gym/challenges/diffract/common.py

@@ -43,7 +43,7 @@ class GratingSpec:
     permittivity_encapsulation: complex
     permittivity_substrate: complex
 
-    thickness_grating: float | jnp.ndarray
+    thickness_grating: float | jnp.ndarray | types.BoundedArray
 
     period_x: float
     period_y: float

src/invrs_gym/challenges/diffract/splitter_challenge.py

@@ -70,11 +70,6 @@ def __init__(
         self.thickness_initializer = thickness_initializer
         self.density_initializer = density_initializer
 
-        self.seed_thickness = types.BoundedArray(
-            array=self.spec.thickness_grating,
-            lower_bound=0.0,
-            upper_bound=None,
-        )
         self.seed_density = common.seed_density(
             self.sim_params.grid_shape, **seed_density_kwargs
         )
@@ -92,7 +87,9 @@ def init(self, key: jax.Array) -> Params:
         """Return the initial parameters for the diffractive splitter component."""
         key_thickness, key_density = jax.random.split(key)
         params = {
-            THICKNESS: self.thickness_initializer(key_thickness, self.seed_thickness),
+            THICKNESS: self.thickness_initializer(
+                key_thickness, self.spec.thickness_grating  # type: ignore[arg-type]
+            ),
             DENSITY: self.density_initializer(key_density, self.seed_density),
         }
         # Ensure that there are no weak types in the initial parameters.
@@ -303,7 +300,7 @@ def extract_orders_for_splitting(
     permittivity_grating=(1.46 + 0.00001j) ** 2,
     permittivity_encapsulation=(1.0 + 0.00001j) ** 2,
     permittivity_substrate=(1.0 + 0.0j) ** 2,
-    thickness_grating=0.692,
+    thickness_grating=types.BoundedArray(array=0.692, lower_bound=0.5, upper_bound=1.5),
     period_x=7.2,
     period_y=7.2,
 )

tests/challenges/diffract/test_metagrating_challenge.py

@@ -23,17 +23,6 @@
 
 
 class MetagratingComponentTest(unittest.TestCase):
-    def test_density_has_expected_properties(self):
-        mc = metagrating_challenge.MetagratingComponent(
-            spec=metagrating_challenge.METAGRATING_SPEC,
-            sim_params=LIGHTWEIGHT_SIM_PARAMS,
-            density_initializer=lambda _, seed_density: seed_density,
-        )
-        params = mc.init(jax.random.PRNGKey(0))
-        self.assertEqual(params.lower_bound, 0.0)
-        self.assertEqual(params.upper_bound, 1.0)
-        self.assertSequenceEqual(params.periodic, (True, True))
-
     def test_can_jit_response(self):
         mc = metagrating_challenge.MetagratingComponent(
             spec=metagrating_challenge.METAGRATING_SPEC,
@@ -91,10 +80,14 @@ def step_fn(params, state):
     @parameterized.expand([[1, 1], [2, 3]])
     def test_density_has_expected_attrs(self, min_width, min_spacing):
         mc = metagrating_challenge.metagrating(
-            minimum_width=min_width,
+            spec=metagrating_challenge.METAGRATING_SPEC,
+            sim_params=LIGHTWEIGHT_SIM_PARAMS,
+            density_initializer=lambda _, seed_density: seed_density,
             minimum_spacing=min_spacing,
+            minimum_width=min_width,
         )
         params = mc.component.init(jax.random.PRNGKey(0))
+
         self.assertEqual(params.lower_bound, 0.0)
         self.assertEqual(params.upper_bound, 1.0)
         self.assertSequenceEqual(params.periodic, (True, True))

tests/challenges/diffract/test_splitter_challenge.py

@@ -22,22 +22,6 @@
 
 
 class SplitterComponentTest(unittest.TestCase):
-    def test_density_has_expected_properties(self):
-        mc = splitter_challenge.DiffractiveSplitterComponent(
-            spec=splitter_challenge.DIFFRACTIVE_SPLITTER_SPEC,
-            sim_params=LIGHTWEIGHT_SIM_PARAMS,
-            thickness_initializer=lambda _, thickness: thickness,
-            density_initializer=lambda _, seed_density: seed_density,
-        )
-        params = mc.init(jax.random.PRNGKey(0))
-        self.assertEqual(params["density"].lower_bound, 0.0)
-        self.assertEqual(params["density"].upper_bound, 1.0)
-        self.assertSequenceEqual(params["density"].periodic, (True, True))
-        self.assertEqual(
-            params["thickness"].array,
-            splitter_challenge.DIFFRACTIVE_SPLITTER_SPEC.thickness_grating,
-        )
-
     def test_can_jit_response(self):
         mc = splitter_challenge.DiffractiveSplitterComponent(
             spec=splitter_challenge.DIFFRACTIVE_SPLITTER_SPEC,
@@ -97,8 +81,12 @@ def step_fn(params, state):
     @parameterized.expand([[1, 1], [2, 3]])
     def test_density_has_expected_attrs(self, min_width, min_spacing):
         mc = splitter_challenge.diffractive_splitter(
-            minimum_width=min_width,
+            spec=splitter_challenge.DIFFRACTIVE_SPLITTER_SPEC,
+            sim_params=LIGHTWEIGHT_SIM_PARAMS,
+            thickness_initializer=lambda _, thickness: thickness,
+            density_initializer=lambda _, seed_density: seed_density,
             minimum_spacing=min_spacing,
+            minimum_width=min_width,
         )
         params = mc.component.init(jax.random.PRNGKey(0))
 
@@ -113,5 +101,9 @@ def test_density_has_expected_attrs(self, min_width, min_spacing):
         self.assertIsNone(params["density"].fixed_solid)
         self.assertIsNone(params["density"].fixed_void)
 
-        self.assertEqual(params["thickness"].lower_bound, 0.0)
-        self.assertIsNone(params["thickness"].upper_bound)
+        self.assertEqual(
+            params["thickness"].array,
+            splitter_challenge.DIFFRACTIVE_SPLITTER_SPEC.thickness_grating.array,
+        )
+        self.assertEqual(params["thickness"].lower_bound, 0.5)
+        self.assertEqual(params["thickness"].upper_bound, 1.5)