greg comments

Author: yaugenst-flex

tests/test_components/test_field_projection.py

@@ -705,7 +705,7 @@ def primitive(currents_in):
             idx_u,
             idx_v,
             is_2d=False,
-            idx_int_1d=None,
+            idx_integration_1d=None,
         )
 
     vjp_primitive, ans_primitive = make_vjp(primitive)(currents)
@@ -723,15 +723,15 @@ def primitive(currents_in):
     np.testing.assert_allclose(grad_primitive, grad_reference, rtol=1e-10, atol=1e-10)
 
 
-@pytest.mark.parametrize("idx_int_1d", [0, 1, 2])
-def test_far_field_integral_vjp_2d(idx_int_1d):
+@pytest.mark.parametrize("idx_integration_1d", [0, 1, 2])
+def test_far_field_integral_vjp_2d(idx_integration_1d):
     rng = np.random.default_rng(1)
     n_theta, n_phi = 5, 6
 
     n_x, n_y, n_z = 1, 1, 1
-    if idx_int_1d == 0:
+    if idx_integration_1d == 0:
         n_x = 4
-    elif idx_int_1d == 1:
+    elif idx_integration_1d == 1:
         n_y = 4
     else:
         n_z = 4
@@ -752,7 +752,7 @@ def test_far_field_integral_vjp_2d(idx_int_1d):
 
     def reference(currents_in):
         chunk = anp.einsum("xtp,ytp,zt,xyz->xyztp", phase_0, phase_1, phase_2, currents_in)
-        return FieldProjector.trapezoid(chunk, pts[idx_int_1d], idx_int_1d)
+        return FieldProjector.trapezoid(chunk, pts[idx_integration_1d], idx_integration_1d)
 
     def primitive(currents_in):
         return _far_field_integral(
@@ -764,7 +764,7 @@ def primitive(currents_in):
             0,
             1,
             is_2d=True,
-            idx_int_1d=idx_int_1d,
+            idx_integration_1d=idx_integration_1d,
         )
 
     vjp_primitive, ans_primitive = make_vjp(primitive)(currents)

tidy3d/components/field_projection.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 from collections.abc import Iterable
+from itertools import product
 from typing import Union
 
 import autograd.numpy as anp
@@ -50,7 +51,18 @@
 
 
 def _trapz_weights_1d(points: np.ndarray) -> np.ndarray:
-    """Trapezoidal integration weights for `trapz(y, x=points)`."""
+    """Trapezoidal integration weights for `trapz(y, x=points)`.
+
+    Parameters
+    ----------
+    points : np.ndarray
+        1D array of integration points.
+
+    Returns
+    -------
+    np.ndarray
+        Trapezoidal integration weights with shape ``(len(points),)``.
+    """
     points = np.asarray(points)
     num_points = points.size
     if num_points <= 1:
@@ -75,82 +87,92 @@ def _far_field_integral(
     idx_v: int,
     *,
     is_2d: bool,
-    idx_int_1d: Union[int, None],
+    idx_integration_1d: Union[int, None],
     weights: list[np.ndarray] | None = None,
 ) -> np.ndarray:
+    """Evaluate the separable far-field surface/line integral.
+
+    This helper computes the near-to-far integral using precomputed separable phase factors
+    and trapezoidal integration weights, with an implementation tailored for autograd.
+
+    Parameters
+    ----------
+    currents : np.ndarray
+        Complex surface current values on the monitor grid with shape ``(nx, ny, nz)``.
+    phase_0 : np.ndarray
+        Phase factor along the x-axis with shape ``(nx, n_theta, n_phi)``.
+    phase_1 : np.ndarray
+        Phase factor along the y-axis with shape ``(ny, n_theta, n_phi)``.
+    phase_2 : np.ndarray
+        Phase factor along the z-axis with shape ``(nz, n_theta)``.
+    pts : list[np.ndarray]
+        List of 1D coordinate arrays ``[x, y, z]`` matching the spatial axes of ``currents``.
+    idx_u : int
+        First surface axis index (0, 1, or 2) for 3D integration.
+    idx_v : int
+        Second surface axis index (0, 1, or 2) for 3D integration.
+    is_2d : bool
+        If ``True``, treat the source as a 1D line and integrate along ``idx_integration_1d``.
+    idx_integration_1d : int | None
+        Spatial axis index (0, 1, or 2) used for the 2D line integral.
+    weights : list[np.ndarray] | None
+        Optional trapezoidal weights for each axis. If ``None``, computed from ``pts``.
+
+    Returns
+    -------
+    np.ndarray
+        Integrated values as an array with trailing axes ``(n_theta, n_phi)``.
+    """
     if weights is None:
         weights = [_trapz_weights_1d(pt) for pt in pts]
 
     optimize = not (phase_0.shape[1] == 1 and phase_0.shape[2] == 1)
 
     if is_2d:
-        if idx_int_1d is None:
-            raise ValueError("Expected 'idx_int_1d' for 2D far-field projection.")
-
-        if idx_int_1d == 0:
-            return anp.einsum(
-                "xtp,ytp,zt,xyz,x->yztp",
-                phase_0,
-                phase_1,
-                phase_2,
-                currents,
-                weights[0],
-                optimize=optimize,
-            )
-        if idx_int_1d == 1:
-            return anp.einsum(
-                "xtp,ytp,zt,xyz,y->xztp",
-                phase_0,
-                phase_1,
-                phase_2,
-                currents,
-                weights[1],
-                optimize=optimize,
-            )
-        if idx_int_1d == 2:
-            return anp.einsum(
-                "xtp,ytp,zt,xyz,z->xytp",
-                phase_0,
-                phase_1,
-                phase_2,
-                currents,
-                weights[2],
-                optimize=optimize,
-            )
-        raise ValueError(f"Invalid 2D integration axis: '{idx_int_1d}'.")
+        if idx_integration_1d is None:
+            raise ValueError("Expected 'idx_integration_1d' for 2D far-field projection.")
+
+        if idx_integration_1d == 0:
+            equation = "xtp,ytp,zt,xyz,x->yztp"
+            weight = weights[0]
+        elif idx_integration_1d == 1:
+            equation = "xtp,ytp,zt,xyz,y->xztp"
+            weight = weights[1]
+        elif idx_integration_1d == 2:
+            equation = "xtp,ytp,zt,xyz,z->xytp"
+            weight = weights[2]
+        else:
+            raise ValueError(f"Invalid 2D integration axis: '{idx_integration_1d}'.")
 
-    integrated_axes = {idx_u, idx_v}
-    remaining_axis = ({0, 1, 2} - integrated_axes).pop()
-    if remaining_axis == 0:
-        return anp.einsum(
-            "xtp,ytp,zt,xyz,y,z->xtp",
-            phase_0,
-            phase_1,
-            phase_2,
-            currents,
-            weights[1],
-            weights[2],
-            optimize=optimize,
-        )
-    if remaining_axis == 1:
         return anp.einsum(
-            "xtp,ytp,zt,xyz,x,z->ytp",
+            equation,
             phase_0,
             phase_1,
             phase_2,
             currents,
-            weights[0],
-            weights[2],
+            weight,
             optimize=optimize,
         )
+
+    integrated_axes = {idx_u, idx_v}
+    remaining_axis = ({0, 1, 2} - integrated_axes).pop()
+    if remaining_axis == 0:
+        equation = "xtp,ytp,zt,xyz,y,z->xtp"
+        weights_uv = (weights[1], weights[2])
+    elif remaining_axis == 1:
+        equation = "xtp,ytp,zt,xyz,x,z->ytp"
+        weights_uv = (weights[0], weights[2])
+    else:
+        equation = "xtp,ytp,zt,xyz,x,y->ztp"
+        weights_uv = (weights[0], weights[1])
+
     return anp.einsum(
-        "xtp,ytp,zt,xyz,x,y->ztp",
+        equation,
         phase_0,
         phase_1,
         phase_2,
         currents,
-        weights[0],
-        weights[1],
+        *weights_uv,
         optimize=optimize,
     )
 
@@ -536,7 +558,7 @@ def _far_fields_for_surface(
         _, source_names = surface.monitor.pop_axis(("x", "y", "z"), axis=surface.axis)
 
         # integration dimension for 2d far field projection
-        idx_int_1d = None
+        idx_integration_1d = None
         zero_dim = [dim for dim, size in enumerate(self.sim_data.simulation.size) if size == 0]
         if self.is_2d_simulation:
             # Ensure zero_dim has a single element since {zero_dim} expects a value
@@ -545,7 +567,7 @@ def _far_fields_for_surface(
 
             zero_dim = zero_dim[0]
             integration_axis = {0, 1, 2} - {zero_dim, surface.axis}
-            idx_int_1d = integration_axis.pop()
+            idx_integration_1d = integration_axis.pop()
 
         idx_u, idx_v = idx_uv
         cmp_1, cmp_2 = source_names
@@ -588,7 +610,7 @@ def _far_fields_for_surface(
                 idx_u,
                 idx_v,
                 is_2d=self.is_2d_simulation,
-                idx_int_1d=idx_int_1d,
+                idx_integration_1d=idx_integration_1d,
                 weights=weights,
             )
 
@@ -803,15 +825,9 @@ def _project_fields_cartesian(
 
         total_points = len(x) * len(y) * len(z)
 
-        def iter_coords():
-            for _x in x:
-                for _y in y:
-                    for _z in z:
-                        yield _x, _y, _z
-
         point_fields = []
         for _x, _y, _z in track(
-            iter_coords(),
+            product(x, y, z),
             description="Computing projected fields",
             total=total_points,
             console=get_logging_console(),
@@ -905,14 +921,9 @@ def _project_fields_kspace(
 
         total_points = len(ux) * len(uy)
 
-        def iter_coords():
-            for _ux in ux:
-                for _uy in uy:
-                    yield _ux, _uy
-
         point_fields = []
         for _ux, _uy in track(
-            iter_coords(),
+            product(ux, uy),
             description="Computing projected fields",
             total=total_points,
             console=get_logging_console(),
@@ -924,30 +935,34 @@ def iter_coords():
                 for idx_f, frequency in enumerate(freqs):
                     fields_sum = anp.zeros((len(field_names),), dtype=complex)
                     for surface, currents in surface_currents:
-                        _fields = self._far_fields_for_surface(
+                        fields_surface = self._far_fields_for_surface(
                             frequency=frequency,
                             theta=theta,
                             phi=phi,
                             surface=surface,
                             currents=currents,
                             medium=medium,
                         )
-                        _fields = anp.reshape(_fields, fields_sum.shape)
-                        fields_sum = fields_sum + _fields * phase[idx_f]
+                        fields_surface = anp.reshape(fields_surface, fields_sum.shape)
+                        fields_sum = fields_sum + fields_surface * phase[idx_f]
                     fields_by_freq.append(fields_sum)
 
                 point_fields.append(anp.stack(fields_by_freq, axis=1))
-                continue
-
-            _x, _y, _z = monitor.sph_2_car(monitor.proj_distance, theta, phi)
-            fields_sum = anp.zeros((len(field_names), len(freqs)), dtype=complex)
-            for surface, currents in surface_currents:
-                _fields = self._fields_for_surface_exact(
-                    x=_x, y=_y, z=_z, surface=surface, currents=currents, medium=medium
-                )
-                _fields = anp.reshape(_fields, fields_sum.shape)
-                fields_sum = fields_sum + _fields
-            point_fields.append(fields_sum)
+            else:
+                x_obs, y_obs, z_obs = monitor.sph_2_car(monitor.proj_distance, theta, phi)
+                fields_sum = anp.zeros((len(field_names), len(freqs)), dtype=complex)
+                for surface, currents in surface_currents:
+                    fields_surface = self._fields_for_surface_exact(
+                        x=x_obs,
+                        y=y_obs,
+                        z=z_obs,
+                        surface=surface,
+                        currents=currents,
+                        medium=medium,
+                    )
+                    fields_surface = anp.reshape(fields_surface, fields_sum.shape)
+                    fields_sum = fields_sum + fields_surface
+                point_fields.append(fields_sum)
 
         stacked_fields = anp.stack(point_fields, axis=0)
         stacked_fields = anp.reshape(