Merge pull request invrs-io#99 from invrs-io/scan

Make use of new method for computing scattering matrices

Author: mfschubert

docs/notebooks/metalens_challenge.ipynb

@@ -151,7 +151,7 @@
     "\n",
     "ex, ey, ez = aux[\"efield\"]\n",
     "x, _, z = aux[\"field_coordinates\"]\n",
-    "xplot, zplot = onp.meshgrid(x[:, 0], z, indexing=\"ij\")\n",
+    "xplot, zplot = onp.meshgrid(x[0, :, 0], z, indexing=\"ij\")\n",
     "\n",
     "abs_field = onp.sqrt(onp.abs(ex) ** 2 + onp.abs(ey) ** 2 + onp.abs(ez) ** 2)\n",
     "\n",

src/invrs_gym/challenges/metalens/challenge.py

@@ -8,6 +8,8 @@
 from typing import Tuple
 
 from fmmax import fmm  # type: ignore[import-untyped]
+import jax
+from jax import nn
 from jax import numpy as jnp
 from totypes import symmetry, types
 
@@ -64,9 +66,8 @@ def loss(self, response: metalens_component.MetalensResponse) -> jnp.ndarray:
             enhancement = response.enhancement_ex
         else:
             enhancement = response.enhancement_ey
-        # The `sqrt` is used to help avoid situtations where one wavelength has
-        # strong enhancement while the others have less.
-        return -jnp.mean(jnp.sqrt(enhancement))
+
+        return soft_amax(-enhancement, scale=10.0)
 
     def distance_to_target(
         self, response: metalens_component.MetalensResponse
@@ -107,6 +108,24 @@ def metrics(
         return metrics
 
 
+def soft_amax(x: jnp.ndarray, scale: float) -> jnp.ndarray:
+    """A soft version of `amax`.
+
+    The softness is set by `scale`. For small values, the output is close to that of
+    `amax`, while for larger values it is closer to that of `mean`. This function can
+    be used to scalarize a vector objective in a manner related to the concept of
+    minimax optimization.
+
+    Args:
+        x: The array to be scalarized.
+        scale: The scale of smoothness.
+
+    Returns:
+        The scalarized array.
+    """
+    return jnp.sum(jax.lax.stop_gradient(nn.softmax(x / scale)) * x)
+
+
 METALENS_SPEC = metalens_component.MetalensSpec(
     permittivity_ambient=(1.0 + 0.0001j) ** 2,
     permittivity_metalens=(2.4 + 0.0001j) ** 2,

src/invrs_gym/challenges/metalens/component.py

@@ -372,30 +372,47 @@ def eigensolve_fn(permittivity: jnp.ndarray) -> fmm.LayerSolveResult:
         solve_result_ambient = eigensolve_fn(
             permittivity=jnp.full(spec.grid_shape, spec.permittivity_ambient)
         )
-        solve_results_metalens = [
-            eigensolve_fn(permittivity=p) for p in metalens_permittivities
-        ]
         solve_result_substrate = eigensolve_fn(
             permittivity=jnp.full(spec.grid_shape, spec.permittivity_substrate)
         )
-        layer_solve_results = (
-            [solve_result_ambient] + solve_results_metalens + [solve_result_substrate]
-        )
 
     if compute_fields:
         # If the field calculation is desired, compute the interior scattering
         # matrices. For each layer in the stack, the interior scattering matrices
         # consist of a pair of matrices, one for the substack below the layer, and
         # one for the substack above the layer.
+        solve_results_metalens = [
+            eigensolve_fn(permittivity=p) for p in metalens_permittivities
+        ]
+        layer_solve_results = (
+            [solve_result_ambient] + solve_results_metalens + [solve_result_substrate]
+        )
         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,
+        solve_results_metalens_batch = eigensolve_fn(
+            permittivity=jnp.asarray(metalens_permittivities)[:, jnp.newaxis, :, :]
+        )
+        # Merge with the ambient and substrate solve results to get the solve results
+        # for the full stack, needed by `stack_s_matrix_scan`.
+        stack_layer_solve_results = tree_util.tree_map(
+            lambda a, b, c: jnp.concatenate(
+                [
+                    a[jnp.newaxis, ...],
+                    jnp.broadcast_to(b, (num_layers,) + b.shape[1:]),
+                    c[jnp.newaxis, ...],
+                ]
+            ),
+            solve_result_ambient,
+            solve_results_metalens_batch,
+            solve_result_substrate,
+        )
+        s_matrix = scattering.stack_s_matrix_scan(
+            layer_solve_results=stack_layer_solve_results,
+            layer_thicknesses=jnp.asarray(layer_thicknesses),
         )
 
     # Compute the source, consisting of a smoothed step function.