refc: Make adjoint helpers private

Author: yaugenst-flex

CHANGELOG.md

@@ -13,6 +13,7 @@ and this project adheres to [Semantic Versioning](https://semver.org/spec/v2.0.0
 
 ### Changed
 - Relaxed bounds checking of path integrals during `WavePort` validation.
+- Internal adjoint helper methods are now prefixed with an underscore to separate them from the public API.
 
 ## [2.8.4] - 2025-05-15
 

tests/test_components/test_autograd.py

@@ -175,7 +175,7 @@ def emulated_run_fwd(simulation, task_name, **run_kwargs) -> td.SimulationData:
                 sim_original = simulation
                 sim_fields_keys = run_kwargs["sim_fields_keys"]
                 # add gradient monitors and make combined simulation
-                sim_combined = sim_original.with_adjoint_monitors(sim_fields_keys)
+                sim_combined = sim_original._with_adjoint_monitors(sim_fields_keys)
                 sim_data_combined = run_emulated(sim_combined, task_name=task_name)
 
                 # store both original and fwd data aux_data
@@ -1099,12 +1099,12 @@ def objective(*params):
 
         sim_full_static = sim_full_traced.to_static()
 
-        sim_fields = sim_full_traced.strip_traced_fields()
+        sim_fields = sim_full_traced._strip_traced_fields()
 
         # note: there is one traced structure in SIM_FULL already with 6 fields + 1 = 7
         assert len(sim_fields) == 10
 
-        sim_traced = sim_full_static.insert_traced_fields(sim_fields)
+        sim_traced = sim_full_static._insert_traced_fields(sim_fields)
 
         assert sim_traced == sim_full_traced
 
@@ -1135,7 +1135,7 @@ def test_sim_fields_io(structure_key, tmp_path):
     s = make_structures(params0)[structure_key]
     s = s.updated_copy(geometry=s.geometry.updated_copy(center=(2, 2, 2), size=(0, 0, 0)))
     sim_full_traced = SIM_FULL.updated_copy(structures=list(SIM_FULL.structures) + [s])
-    sim_fields = sim_full_traced.strip_traced_fields()
+    sim_fields = sim_full_traced._strip_traced_fields()
 
     field_map = FieldMap.from_autograd_field_map(sim_fields)
     field_map_file = join(tmp_path, "test_sim_fields.hdf5.gz")
@@ -1434,7 +1434,7 @@ def J(eps):
 
     monkeypatch.setattr(
         td.PoleResidue,
-        "derivative_eps_complex_volume",
+        "_derivative_eps_complex_volume",
         lambda self, E_der_map, bounds, freqs: dJ_deps,
     )
 
@@ -1467,7 +1467,7 @@ def J(eps):
         bounds_intersect=((-1, -1, -1), (1, 1, 1)),
     )
 
-    grads_computed = pr.compute_derivatives(derivative_info=info)
+    grads_computed = pr._compute_derivatives(derivative_info=info)
 
     def f(eps_inf, poles):
         eps = td.PoleResidue._eps_model(eps_inf, poles, freq)
@@ -1549,7 +1549,7 @@ def J(eps):
         bounds_intersect=((-1, -1, -1), (1, 1, 1)),
     )
 
-    grads_computed = pr.compute_derivatives(derivative_info=info)
+    grads_computed = pr._compute_derivatives(derivative_info=info)
 
     poles_complex = [
         (np.array(a.values, dtype=complex), np.array(c.values, dtype=complex)) for a, c in poles
@@ -2141,7 +2141,7 @@ def objective(params):
             structure_traced = make_structures(params)["medium"]
             sim = SIM_BASE.updated_copy(structures=[structure_traced], monitors=monitors)
             data = run(sim, task_name="adjoint_freq_test")
-            assert data.simulation.freqs_adjoint == [FREQ0]
+            assert data.simulation._freqs_adjoint == [FREQ0]
             return anp.sum(data["field"].flux.values)
 
         return objective

tidy3d/components/base.py

@@ -948,7 +948,7 @@ def make_json_compatible(json_string: str) -> str:
         json_string = make_json_compatible(json_string)
         return json_string
 
-    def strip_traced_fields(
+    def _strip_traced_fields(
         self, starting_path: tuple[str] = (), include_untraced_data_arrays: bool = False
     ) -> AutogradFieldMap:
         """Extract a dictionary mapping paths in the model to the data traced by ``autograd``.
@@ -1004,7 +1004,7 @@ def handle_value(x: Any, path: tuple[str, ...]) -> None:
         # convert the resulting field_mapping to an autograd-traced dictionary
         return dict_ag(field_mapping)
 
-    def insert_traced_fields(self, field_mapping: AutogradFieldMap) -> Tidy3dBaseModel:
+    def _insert_traced_fields(self, field_mapping: AutogradFieldMap) -> Tidy3dBaseModel:
         """Recursively insert a map of paths to autograd-traced fields into a copy of this obj."""
 
         self_dict = self.dict()
@@ -1037,7 +1037,7 @@ def to_static(self) -> Tidy3dBaseModel:
         """Version of object with all autograd-traced fields removed."""
 
         # get dictionary of all traced fields
-        field_mapping = self.strip_traced_fields()
+        field_mapping = self._strip_traced_fields()
 
         # shortcut to just return self if no tracers found, for performance
         if not field_mapping:
@@ -1047,7 +1047,7 @@ def to_static(self) -> Tidy3dBaseModel:
         field_mapping_static = {key: get_static(val) for key, val in field_mapping.items()}
 
         # insert the static values into a copy of self
-        return self.insert_traced_fields(field_mapping_static)
+        return self._insert_traced_fields(field_mapping_static)
 
     @classmethod
     def add_type_field(cls) -> None:

tidy3d/components/data/monitor_data.py

@@ -166,7 +166,7 @@ def _updated(self, update: Dict) -> MonitorData:
         data_dict.update(update)
         return type(self).parse_obj(data_dict)
 
-    def make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[Source]:
+    def _make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[Source]:
         """Generate adjoint sources for this ``MonitorData`` instance."""
 
         # TODO: if there's data in the MonitorData, but no adjoint source, then
@@ -1235,7 +1235,7 @@ def to_zbf(
                 e_flat = e.values.flatten(order="C")
                 # Interweave real and imaginary parts
                 e_values = np.ravel(np.column_stack((e_flat.real, e_flat.imag)))
-                fout.write(struct.pack(f"<{2 * n_x*n_y}d", *e_values))
+                fout.write(struct.pack(f"<{2 * n_x * n_y}d", *e_values))
 
         return e_x, e_y
 
@@ -1333,7 +1333,7 @@ def to_source(
             field_dataset=dataset, source_time=source_time, center=center, size=size, **kwargs
         )
 
-    def make_adjoint_sources(
+    def _make_adjoint_sources(
         self, dataset_names: list[str], fwidth: float
     ) -> List[CustomCurrentSource]:
         """Converts a :class:`.FieldData` to a list of adjoint current or point sources."""
@@ -2068,14 +2068,14 @@ def _check_fields_stored(self, components: list[EMField]):
                 "include the mode field profiles in the corresponding 'ModeData'."
             )
 
-    def make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[ModeSource]:
+    def _make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[ModeSource]:
         """Get all adjoint sources for the ``ModeMonitorData``."""
 
         adjoint_sources = []
 
         for name in dataset_names:
             if name == "amps":
-                adjoint_sources += self.make_adjoint_sources_amps(fwidth=fwidth)
+                adjoint_sources += self._make_adjoint_sources_amps(fwidth=fwidth)
             elif not np.all(self.n_complex.values == 0.0):
                 log.warning(
                     f"Can't create adjoint source for 'ModeData.{type(self)}.{name}'. "
@@ -2086,7 +2086,7 @@ def make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[
 
         return adjoint_sources
 
-    def make_adjoint_sources_amps(self, fwidth: float) -> list[ModeSource]:
+    def _make_adjoint_sources_amps(self, fwidth: float) -> list[ModeSource]:
         """Generate adjoint sources for ``ModeMonitorData.amps``."""
 
         coords = self.amps.coords
@@ -2102,12 +2102,12 @@ def make_adjoint_sources_amps(self, fwidth: float) -> list[ModeSource]:
                     if self.get_amplitude(amp_single) == 0.0:
                         continue
 
-                    adjoint_source = self.adjoint_source_amp(amp=amp_single, fwidth=fwidth)
+                    adjoint_source = self._adjoint_source_amp(amp=amp_single, fwidth=fwidth)
                     adjoint_sources.append(adjoint_source)
 
         return adjoint_sources
 
-    def adjoint_source_amp(self, amp: DataArray, fwidth: float) -> ModeSource:
+    def _adjoint_source_amp(self, amp: DataArray, fwidth: float) -> ModeSource:
         """Generate an adjoint ``ModeSource`` for a single amplitude."""
 
         monitor = self.monitor
@@ -2266,7 +2266,7 @@ class FluxData(MonitorData):
         ..., title="Flux", description="Flux values in the frequency-domain."
     )
 
-    def make_adjoint_sources(
+    def _make_adjoint_sources(
         self, dataset_names: list[str], fwidth: float
     ) -> List[Union[CustomCurrentSource, PointDipole]]:
         """Converts a :class:`.FieldData` to a list of adjoint current or point sources."""
@@ -2606,7 +2606,7 @@ def radar_cross_section(self) -> DataArray:
 
         return self.make_data_array(data=rcs_data)
 
-    def make_adjoint_sources(
+    def _make_adjoint_sources(
         self, dataset_names: list[str], fwidth: float
     ) -> List[Union[CustomCurrentSource, PointDipole]]:
         """Error if server-side field projection is used for autograd"""
@@ -2764,7 +2764,7 @@ def _check_integration_suitability(self):
                 "There are not enough sampling points along `theta` or `phi` for accurate integration. "
                 f"Currently, {len(self.theta)} samples for `theta` and {len(self.phi)} samples for `phi`. "
                 f"Consider using, at the very least, {MIN_ANGULAR_SAMPLES_SPHERE} samples for `theta` and "
-                f"{2*MIN_ANGULAR_SAMPLES_SPHERE} samples for `phi`."
+                f"{2 * MIN_ANGULAR_SAMPLES_SPHERE} samples for `phi`."
             )
         self._check_coords_sorted(self.theta, "theta")
         self._check_coords_sorted(self.phi, "phi")
@@ -3359,13 +3359,13 @@ def _make_dataset(self, fields: Tuple[np.ndarray, ...], keys: Tuple[str, ...]) -
 
     """ Autograd code """
 
-    def make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[PlaneWave]:
+    def _make_adjoint_sources(self, dataset_names: list[str], fwidth: float) -> list[PlaneWave]:
         """Get all adjoint sources for the ``DiffractionMonitor.amps``."""
 
         # NOTE: everything just goes through `.amps`, any post-processing is encoded in E-fields
-        return self.make_adjoint_sources_amps(fwidth=fwidth)
+        return self._make_adjoint_sources_amps(fwidth=fwidth)
 
-    def make_adjoint_sources_amps(self, fwidth: float) -> list[PlaneWave]:
+    def _make_adjoint_sources_amps(self, fwidth: float) -> list[PlaneWave]:
         """Make adjoint sources for outputs that depend on DiffractionData.`amps`."""
 
         amps = self.amps

tidy3d/components/data/sim_data.py

@@ -994,7 +994,7 @@ def source_spectrum_fn(freqs):
 
         return self.copy(update=dict(simulation=simulation, data=data_normalized))
 
-    def split_adjoint_data(self: SimulationData, num_mnts_original: int) -> tuple[list, list]:
+    def _split_adjoint_data(self: SimulationData, num_mnts_original: int) -> tuple[list, list]:
         """Split data list into original, adjoint field, and adjoint permittivity."""
 
         data_all = list(self.data)
@@ -1008,11 +1008,11 @@ def split_adjoint_data(self: SimulationData, num_mnts_original: int) -> tuple[li
 
         return data_original, data_adjoint
 
-    def split_original_fwd(self, num_mnts_original: int) -> Tuple[SimulationData, SimulationData]:
+    def _split_original_fwd(self, num_mnts_original: int) -> Tuple[SimulationData, SimulationData]:
         """Split this simulation data into original and fwd data from number of original mnts."""
 
         # split the data and monitors into the original ones & adjoint gradient ones (for 'fwd')
-        data_original, data_fwd = self.split_adjoint_data(num_mnts_original=num_mnts_original)
+        data_original, data_fwd = self._split_adjoint_data(num_mnts_original=num_mnts_original)
         monitors_orig, monitors_fwd = split_list(self.simulation.monitors, index=num_mnts_original)
 
         # reconstruct the simulation data for the user, using original sim, and data for original mnts
@@ -1033,7 +1033,7 @@ def split_original_fwd(self, num_mnts_original: int) -> Tuple[SimulationData, Si
 
         return sim_data_original, sim_data_fwd
 
-    def make_adjoint_sims(
+    def _make_adjoint_sims(
         self,
         data_vjp_paths: set[tuple],
         adjoint_monitors: list[Monitor],
@@ -1046,15 +1046,15 @@ def make_adjoint_sims(
         sim_original = self.simulation
 
         # generate the adjoint sources {mnt_name : list[Source]}
-        sources_adj_dict = self.make_adjoint_sources(data_vjp_paths=data_vjp_paths)
+        sources_adj_dict = self._make_adjoint_sources(data_vjp_paths=data_vjp_paths)
         if not sources_adj_dict:
             return []
 
         adj_srcs = []
         for src_list in sources_adj_dict.values():
             adj_srcs += list(src_list)
 
-        adjoint_source_infos = self.process_adjoint_sources(adj_srcs=adj_srcs)
+        adjoint_source_infos = self._process_adjoint_sources(adj_srcs=adj_srcs)
 
         if not adjoint_source_infos:
             return []
@@ -1095,7 +1095,7 @@ def make_adjoint_sims(
 
         return adj_sims
 
-    def make_adjoint_sources(self, data_vjp_paths: set[tuple]) -> dict[str, SourceType]:
+    def _make_adjoint_sources(self, data_vjp_paths: set[tuple]) -> dict[str, SourceType]:
         """Generate all of the non-zero sources for the adjoint simulation given the VJP data."""
 
         # map of index into 'self.data' to the list of datasets we need adjoint sources for
@@ -1107,8 +1107,8 @@ def make_adjoint_sources(self, data_vjp_paths: set[tuple]) -> dict[str, SourceTy
         sources_adj_all = defaultdict(list)
         for data_index, dataset_names in adj_src_map.items():
             mnt_data = self.data[data_index]
-            sources_adj = mnt_data.make_adjoint_sources(
-                dataset_names=dataset_names, fwidth=self.fwidth_adj
+            sources_adj = mnt_data._make_adjoint_sources(
+                dataset_names=dataset_names, fwidth=self._fwidth_adj
             )
             sources_adj_all[mnt_data.monitor.name] = sources_adj
             log.info(
@@ -1118,12 +1118,12 @@ def make_adjoint_sources(self, data_vjp_paths: set[tuple]) -> dict[str, SourceTy
         return sources_adj_all
 
     @property
-    def fwidth_adj(self) -> float:
+    def _fwidth_adj(self) -> float:
         # fwidth of forward pass, try as default for adjoint
         normalize_index_fwd = self.simulation.normalize_index or 0
         return self.simulation.sources[normalize_index_fwd].source_time.fwidth
 
-    def process_adjoint_sources(self, adj_srcs: list[SourceType]) -> list[AdjointSourceInfo]:
+    def _process_adjoint_sources(self, adj_srcs: list[SourceType]) -> list[AdjointSourceInfo]:
         """Compute list of final sources along with a post run normalization for adj fields."""
         # dictionary mapping hash of sources with same freq dependence to list of time-dependencies
         hashes_to_sources = defaultdict(None)
@@ -1157,7 +1157,7 @@ def process_adjoint_sources(self, adj_srcs: list[SourceType]) -> list[AdjointSou
             for src_hash, src_times in hashes_to_src_times.items():
                 base_src = hashes_to_sources[src_hash]
                 group = [base_src.updated_copy(source_time=src_time) for src_time in src_times]
-                processed_srcs, post_norm = self.process_adjoint_sources_broadband(group)
+                processed_srcs, post_norm = self._process_adjoint_sources_broadband(group)
                 adjoint_infos.append(
                     AdjointSourceInfo(
                         sources=processed_srcs, post_norm=post_norm, normalize_sim=True
@@ -1167,7 +1167,7 @@ def process_adjoint_sources(self, adj_srcs: list[SourceType]) -> list[AdjointSou
         log.info(f"Created {len(adjoint_infos)} adjoint source groups.")
         return adjoint_infos
 
-    def process_adjoint_sources_broadband(
+    def _process_adjoint_sources_broadband(
         self, adj_srcs: list[SourceType]
     ) -> tuple[list[SourceType], xr.DataArray]:
         """Process adjoint sources for the case of several sources at the same freq."""
@@ -1211,10 +1211,10 @@ def _make_post_norm_amps(adj_srcs: list[SourceType]) -> xr.DataArray:
         amps_complex = np.array(amps_complex)
         return xr.DataArray(amps_complex, coords=coords)
 
-    def get_adjoint_data(self, structure_index: int, data_type: str) -> MonitorDataType:
+    def _get_adjoint_data(self, structure_index: int, data_type: str) -> MonitorDataType:
         """Grab the field or permittivity data for a given structure index."""
 
-        monitor_name = Structure.get_monitor_name(index=structure_index, data_type=data_type)
+        monitor_name = Structure._get_monitor_name(index=structure_index, data_type=data_type)
         return self[monitor_name]
 
     def to_mat_file(self, fname: str, **kwargs):