feat: Autograd support for FieldProjectionKSpaceMonitor

Author: yaugenst-flex

CHANGELOG.md

@@ -7,6 +7,9 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ## [Unreleased]
 
+### Added
+- Autograd support for local field projections using `FieldProjectionKSpaceMonitor`.
+
 ### Fixed
 - Regression in local field projection leading to incorrect results for `far_field_approx=True`.
 

tests/test_components/test_autograd.py

@@ -1003,7 +1003,7 @@ def objective(args):
 
 
 @pytest.mark.parametrize("far_field_approx", [True, False])
-@pytest.mark.parametrize("projection_type", ["angular", "cartesian"])
+@pytest.mark.parametrize("projection_type", ["angular", "cartesian", "kspace"])
 @pytest.mark.parametrize("sim_2d", [True, False])
 class TestFieldProjection:
     @staticmethod
@@ -1047,6 +1047,20 @@ def setup(far_field_approx, projection_type, sim_2d):
                 far_field_approx=far_field_approx,
                 name="far_field",
             )
+        elif projection_type == "kspace":
+            ux = np.linspace(-0.7, 0.7, 2)
+            uy = np.linspace(-0.7, 0.7, 3)
+            monitor_far = td.FieldProjectionKSpaceMonitor(
+                center=monitor.center,
+                size=monitor.size,
+                freqs=monitor.freqs,
+                ux=ux,
+                uy=uy,
+                proj_axis=1,
+                proj_distance=r_proj,
+                far_field_approx=far_field_approx,
+                name="far_field",
+            )
 
         sim = SIM_BASE.updated_copy(monitors=[monitor])
 

tidy3d/components/field_projection.py

@@ -759,7 +759,7 @@ def _project_fields_kspace(
 
         # compute projected fields for the dataset associated with each monitor
         field_names = ("Er", "Etheta", "Ephi", "Hr", "Htheta", "Hphi")
-        fields = [np.zeros((len(ux), len(uy), 1, len(freqs)), dtype=complex) for _ in field_names]
+        fields = np.zeros((len(field_names), len(ux), len(uy), 1, len(freqs)), dtype=complex)
 
         medium = monitor.medium if monitor.medium else self.medium
         k = AbstractFieldProjectionData.wavenumber(medium=medium, frequency=freqs)
@@ -793,16 +793,17 @@ def _project_fields_kspace(
                             currents=currents,
                             medium=medium,
                         )
-                        for field, _field in zip(fields, _fields):
-                            field = add_at(field, [i, j, 0, idx_f], _field * phase[idx_f])
-
+                        where = (slice(None), i, j, 0, idx_f)
+                        _fields = anp.reshape(_fields, fields[where].shape)
+                        fields = add_at(fields, where, _fields * phase[idx_f])
                 else:
                     _x, _y, _z = monitor.sph_2_car(monitor.proj_distance, theta, phi)
                     _fields = self._fields_for_surface_exact(
                         x=_x, y=_y, z=_z, surface=surface, currents=currents, medium=medium
                     )
-                    for field, _field in zip(fields, _fields):
-                        field = add_at(field, [i, j, 0], _field)
+                    where = (slice(None), i, j, 0)
+                    _fields = anp.reshape(_fields, fields[where].shape)
+                    fields = add_at(fields, where, _fields)
 
         coords = {
             "ux": np.array(monitor.ux),