From fe4436a4b7711983fc7a20ca0f25155b4b743a4a Mon Sep 17 00:00:00 2001 From: Martin Schubert Date: Mon, 23 Sep 2024 11:45:47 -0700 Subject: [PATCH 1/2] Add compute fields functionality --- src/invrs_gym/challenges/diffract/common.py | 77 ++++++++++++++++--- .../diffract/metagrating_challenge.py | 2 + .../challenges/diffract/splitter_challenge.py | 2 + tests/challenges/diffract/test_common.py | 31 ++++++++ 4 files changed, 103 insertions(+), 9 deletions(-) diff --git a/src/invrs_gym/challenges/diffract/common.py b/src/invrs_gym/challenges/diffract/common.py index 5a35972..1a95169 100644 --- a/src/invrs_gym/challenges/diffract/common.py +++ b/src/invrs_gym/challenges/diffract/common.py @@ -26,6 +26,10 @@ THICKNESS_SPACER = "thickness_spacer" DENSITY = "density" +EFIELD = "efield" +HFIELD = "hfield" +FIELD_COORDINATES = "field_coordinates" + @dataclasses.dataclass class GratingSpec: @@ -55,9 +59,11 @@ class GratingSpec: permittivity_spacer: complex permittivity_substrate: complex + thickness_ambient: float thickness_cap: float | jnp.ndarray | types.BoundedArray thickness_grating: float | jnp.ndarray | types.BoundedArray thickness_spacer: float | jnp.ndarray | types.BoundedArray + thickness_substrate: float period_x: float period_y: float @@ -193,6 +199,7 @@ def response( *, wavelength: Optional[Union[float, jnp.ndarray]] = None, expansion: Optional[basis.Expansion] = None, + compute_fields: bool = False, ) -> Tuple[GratingResponse, base.AuxDict]: """Computes the response of the grating. @@ -204,6 +211,8 @@ def response( by the `init` method. wavelength: Optional wavelength to override the default in `sim_params`. expansion: Optional expansion to override the default `expansion`. + compute_fields: If `True`, computes and xz cross section for electric + and magnetic fields, which makes the calculation more expensive. Returns: The `(response, aux)` tuple. @@ -212,7 +221,7 @@ def response( expansion = self.expansion if wavelength is None: wavelength = self.sim_params.wavelength - transmission_efficiency, reflection_efficiency = grating_efficiency( + (transmission_efficiency, reflection_efficiency), aux = grating_efficiency( density=params, spec=self.spec, wavelength=jnp.asarray(wavelength), @@ -220,6 +229,7 @@ def response( azimuthal_angle=jnp.asarray(self.sim_params.azimuthal_angle), expansion=expansion, formulation=self.sim_params.formulation, + compute_fields=compute_fields, ) response = GratingResponse( wavelength=jnp.asarray(wavelength), @@ -229,7 +239,7 @@ def response( reflection_efficiency=reflection_efficiency, expansion=expansion, ) - return response, {} + return response, aux class GratingWithOptimizableThicknessComponent(base.Component): @@ -302,6 +312,7 @@ def response( *, wavelength: Optional[Union[float, jnp.ndarray]] = None, expansion: Optional[basis.Expansion] = None, + compute_fields: bool = False, ) -> Tuple[GratingResponse, base.AuxDict]: """Computes the response of the grating component. @@ -313,6 +324,8 @@ def response( the `init` method. wavelength: Optional wavelength to override the default in `sim_params`. expansion: Optional expansion to override the default `expansion`. + compute_fields: If `True`, computes and xz cross section for electric + and magnetic fields, which makes the calculation more expensive. Returns: The `(response, aux)` tuple. @@ -327,7 +340,7 @@ def response( thickness_grating=jnp.asarray(params[THICKNESS_GRATING].array), thickness_spacer=jnp.asarray(params[THICKNESS_SPACER].array), ) - transmission_efficiency, reflection_efficiency = grating_efficiency( + (transmission_efficiency, reflection_efficiency), aux = grating_efficiency( density=params[DENSITY], # type: ignore[arg-type] spec=spec, wavelength=jnp.asarray(wavelength), @@ -335,6 +348,7 @@ def response( azimuthal_angle=jnp.asarray(self.sim_params.azimuthal_angle), expansion=expansion, formulation=self.sim_params.formulation, + compute_fields=compute_fields, ) response = GratingResponse( wavelength=jnp.asarray(wavelength), @@ -344,7 +358,7 @@ def response( reflection_efficiency=reflection_efficiency, expansion=expansion, ) - return response, {} + return response, aux def seed_density(grid_shape: Tuple[int, int], **kwargs: Any) -> types.Density2DArray: @@ -401,7 +415,8 @@ def grating_efficiency( azimuthal_angle: jnp.ndarray, expansion: basis.Expansion, formulation: fmm.Formulation, -) -> Tuple[jnp.ndarray, jnp.ndarray]: + compute_fields: bool, +) -> Tuple[Tuple[jnp.ndarray, jnp.ndarray], base.AuxDict]: """Compute the per-order transmission and reflection efficiency for a grating. The excitation for the calculation are separate a TE- and TM-polarized plane waves @@ -416,6 +431,7 @@ def grating_efficiency( azimuthal_angle: The azimuthal angle of the excitation. expansion: Defines the Fourier expansion for the calculation. formulation: Defines the FMM formulation to be used. + compute_fields: If `True`, computes fields in an xz cross section. Returns: The per-order transmission and reflection efficiency, having shape @@ -463,14 +479,25 @@ def grating_efficiency( # Layer thicknesses for the ambient and substrate are set to zero; these do not # affect the result of the calculation. layer_thicknesses = ( - jnp.zeros(()), + jnp.asarray(spec.thickness_ambient), jnp.asarray(spec.thickness_cap), jnp.asarray(spec.thickness_grating), jnp.asarray(spec.thickness_spacer), - jnp.zeros(()), + jnp.asarray(spec.thickness_substrate), ) - s_matrix = scattering.stack_s_matrix(layer_solve_results, layer_thicknesses) + if compute_fields: + # If fields wanted, compute the full set of interior scattering matrices. + s_matrices_interior = scattering.stack_s_matrices_interior( + layer_solve_results=layer_solve_results, + layer_thicknesses=layer_thicknesses, + ) + s_matrix = s_matrices_interior[-1][0] + else: + s_matrix = scattering.stack_s_matrix( + layer_solve_results=layer_solve_results, + layer_thicknesses=layer_thicknesses, + ) n = expansion.num_terms assert tuple(expansion.basis_coefficients[0, :]) == (0, 0) @@ -511,4 +538,36 @@ def grating_efficiency( transmission_efficiency = bwd_flux_ambient / total_incident_flux reflection_efficiency = fwd_flux_substrate / total_incident_flux - return transmission_efficiency, reflection_efficiency + # ------------------------------------------------------------------------- + # Compute fields in an xz cross section. + # ------------------------------------------------------------------------- + + aux = {} + if compute_fields: + amplitudes_interior = fields.stack_amplitudes_interior( + s_matrices_interior=s_matrices_interior, + forward_amplitude_0_start=jnp.zeros_like(bwd_amplitude_substrate_end), + backward_amplitude_N_end=bwd_amplitude_substrate_end, + ) + x = jnp.linspace(0, spec.period_x, density.shape[0]) + y = jnp.ones_like(x) * spec.period_y / 2 + layer_znum = tuple( + [int(jnp.round(t / spec.grid_spacing) + 1) for t in layer_thicknesses] + ) + (ex, ey, ez), (hx, hy, hz), (x, y, z) = fields.stack_fields_3d_on_coordinates( + amplitudes_interior=amplitudes_interior, + layer_solve_results=layer_solve_results, + layer_thicknesses=layer_thicknesses, + layer_znum=layer_znum, + x=x, + y=y, + ) + aux.update( + { + EFIELD: (ex, ey, ez), + HFIELD: (hx, hy, hz), + FIELD_COORDINATES: (x, y, z), + } + ) + + return (transmission_efficiency, reflection_efficiency), aux diff --git a/src/invrs_gym/challenges/diffract/metagrating_challenge.py b/src/invrs_gym/challenges/diffract/metagrating_challenge.py index 888face..3dffba5 100644 --- a/src/invrs_gym/challenges/diffract/metagrating_challenge.py +++ b/src/invrs_gym/challenges/diffract/metagrating_challenge.py @@ -144,9 +144,11 @@ def _value_for_order( permittivity_encapsulation=(1.0 + 0.00001j) ** 2, permittivity_spacer=(1.45 + 0.0j) ** 2, permittivity_substrate=(1.45 + 0.0j) ** 2, + thickness_ambient=0.0, thickness_cap=0.0, thickness_grating=0.325, thickness_spacer=0.0, + thickness_substrate=0.0, period_x=float(1.050 / jnp.sin(jnp.deg2rad(50.0))), period_y=0.525, grid_spacing=0.0117, # Yields a grid shape of `(118, 45)`. diff --git a/src/invrs_gym/challenges/diffract/splitter_challenge.py b/src/invrs_gym/challenges/diffract/splitter_challenge.py index 63bbecd..0213823 100644 --- a/src/invrs_gym/challenges/diffract/splitter_challenge.py +++ b/src/invrs_gym/challenges/diffract/splitter_challenge.py @@ -262,9 +262,11 @@ def extract_orders_for_splitting( permittivity_encapsulation=(1.0 + 0.00001j) ** 2, permittivity_spacer=(1.0 + 0.0j) ** 2, permittivity_substrate=(1.0 + 0.0j) ** 2, + thickness_ambient=0.0, thickness_cap=types.BoundedArray(array=0.0, lower_bound=0.0, upper_bound=0.1), thickness_grating=types.BoundedArray(array=0.692, lower_bound=0.5, upper_bound=1.5), thickness_spacer=types.BoundedArray(array=0.0, lower_bound=0.0, upper_bound=0.1), + thickness_substrate=0.0, period_x=7.2, period_y=7.2, grid_spacing=0.04, # Yields a grid shape of `(180, 180)`. diff --git a/tests/challenges/diffract/test_common.py b/tests/challenges/diffract/test_common.py index 12053ea..91be561 100644 --- a/tests/challenges/diffract/test_common.py +++ b/tests/challenges/diffract/test_common.py @@ -21,9 +21,11 @@ permittivity_encapsulation=(1.0 + 0.00001j) ** 2, permittivity_spacer=(1.45 + 0.0j) ** 2, permittivity_substrate=(1.45 + 0.0j) ** 2, + thickness_ambient=0.0, thickness_cap=0.0, thickness_grating=0.325, thickness_spacer=0.0, + thickness_substrate=0.0, period_x=float(1.050 / jnp.sin(jnp.deg2rad(50.0))), period_y=0.525, grid_spacing=0.0117, @@ -36,9 +38,11 @@ permittivity_encapsulation=(1.0 + 0.00001j) ** 2, permittivity_spacer=(1.45 + 0.0j) ** 2, permittivity_substrate=(1.45 + 0.0j) ** 2, + thickness_ambient=0.0, thickness_cap=types.BoundedArray(array=0.0, lower_bound=0.0, upper_bound=0.1), thickness_grating=types.BoundedArray(array=0.6, lower_bound=0.5, upper_bound=1.5), thickness_spacer=types.BoundedArray(array=0.0, lower_bound=0.0, upper_bound=0.1), + thickness_substrate=0.0, period_x=float(1.050 / jnp.sin(jnp.deg2rad(50.0))), period_y=0.525, grid_spacing=0.0117, @@ -111,6 +115,19 @@ def test_multiple_wavelengths(self): (2, mc.expansion.num_terms, 2), ) + def test_compute_fields(self): + mc = common.SimpleGratingComponent( + spec=SIMPLE_GRATING_SPEC, + sim_params=LIGHTWEIGHT_SIM_PARAMS, + density_initializer=lambda _, seed_density: seed_density, + ) + params = mc.init(jax.random.PRNGKey(0)) + response, aux = mc.response(params, compute_fields=True) + self.assertSequenceEqual( + set(aux.keys()), + {common.EFIELD, common.HFIELD, common.FIELD_COORDINATES}, + ) + class GratingWithOptimizableThicknessComponentTest(unittest.TestCase): def test_can_jit_response(self): @@ -141,3 +158,17 @@ def test_multiple_wavelengths(self): response.transmission_efficiency.shape, (2, mc.expansion.num_terms, 2), ) + + def test_compute_fields(self): + mc = common.GratingWithOptimizableThicknessComponent( + spec=GRATING_WITH_THICKNESS_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)) + response, aux = mc.response(params, compute_fields=True) + self.assertSequenceEqual( + set(aux.keys()), + {common.EFIELD, common.HFIELD, common.FIELD_COORDINATES}, + ) From c466e6197fa9d8a886d61e93015290f59dc5748f Mon Sep 17 00:00:00 2001 From: Martin Schubert Date: Mon, 23 Sep 2024 11:59:14 -0700 Subject: [PATCH 2/2] Version updated from v1.3.1 to v1.4.0 --- .bumpversion.toml | 2 +- README.md | 2 +- pyproject.toml | 2 +- src/invrs_gym/__init__.py | 2 +- 4 files changed, 4 insertions(+), 4 deletions(-) diff --git a/.bumpversion.toml b/.bumpversion.toml index 7cf3f06..619b319 100644 --- a/.bumpversion.toml +++ b/.bumpversion.toml @@ -1,5 +1,5 @@ [tool.bumpversion] -current_version = "v1.3.1" +current_version = "v1.4.0" commit = true commit_args = "--no-verify" tag = true diff --git a/README.md b/README.md index 1c1077f..70b673f 100644 --- a/README.md +++ b/README.md @@ -1,5 +1,5 @@ # invrs-gym -`v1.3.1` +`v1.4.0` ## Overview The `invrs_gym` package is an open-source gym containing a diverse set of photonic design challenges, which are relevant for a wide range of applications such as AR/VR, optical networking, LIDAR, and others. diff --git a/pyproject.toml b/pyproject.toml index 7a9cc90..15c584b 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -1,7 +1,7 @@ [project] name = "invrs_gym" -version = "v1.3.1" +version = "v1.4.0" description = "A collection of inverse design challenges" keywords = ["topology", "optimization", "jax", "inverse design"] readme = "README.md" diff --git a/src/invrs_gym/__init__.py b/src/invrs_gym/__init__.py index 009f0ca..f57d79a 100644 --- a/src/invrs_gym/__init__.py +++ b/src/invrs_gym/__init__.py @@ -3,7 +3,7 @@ Copyright (c) 2023 The INVRS-IO authors. """ -__version__ = "v1.3.1" +__version__ = "v1.4.0" __author__ = "Martin F. Schubert " from invrs_gym import challenges as challenges