add flag to PlasmaVesselDistance to avoid recomputing surface coordinates when the surface is fixed during optimization

CHANGELOG.md

@@ -3,7 +3,10 @@ Changelog
 
 - Adds ``desc.compat.rescale`` for rescaling equilibria to a specified size and field
 strength.
-
+ - Adds new keyword ``surface_fixed`` to ``PlasmaVesselDistance`` objective which says
+whether or not the surface comparing the distance from the plasma to is fixed or not.
+If True, then the surface coordinates can be precomputed, saving on computation during
+optimization. Set to False by default.
 
 v0.10.2
 -------

desc/objectives/_geometry.py

@@ -401,6 +401,11 @@
     will only be an upper bound on the minimum separation between the plasma and the
     surrounding surface.
 
+    NOTE: By default, assumes the surface is fixed and its coordinates are precomputed,
+    set surface_fixed=False to compute the surface coordinates at every iteration, for
+    example if the winding surface you compare to is part of the optimization and thus
+    changing.
+
     NOTE: for best results, use this objective in combination with either MeanCurvature
     or PrincipalCurvature, to penalize the tendency for the optimizer to only move the
     points on surface corresponding to the grid that the plasma-vessel distance
@@ -440,6 +445,12 @@
         Collocation grid containing the nodes to evaluate plasma geometry at.
     use_softmin: bool, optional
         Use softmin or hard min.
+    surface_fixed: bool, optional
+        Whether the surface the distance from the plasma is computed to
+        is fixed or not. If True, the surface is fixed and its coordinates are
+        precomputed, which saves on computation time during optimization.
+        If False, the surface coordinates are computed at every iteration.
+        False by default.
     alpha: float, optional
         Parameter used for softmin. The larger alpha, the closer the softmin
         approximates the hardmin. softmin -> hardmin as alpha -> infinity.
@@ -467,6 +478,7 @@
         surface_grid=None,
         plasma_grid=None,
         use_softmin=False,
+        surface_fixed=False,
         alpha=1.0,
         name="plasma-vessel distance",
     ):
@@ -476,6 +488,7 @@
         self._surface_grid = surface_grid
         self._plasma_grid = plasma_grid
         self._use_softmin = use_softmin
+        self._surface_fixed = surface_fixed
         self._alpha = alpha
         super().__init__(
             things=[eq, self._surface],
@@ -559,6 +572,19 @@
             "quad_weights": w,
         }
 
+        if self._surface_fixed:
+            # precompute the surface coordinates
+            # as the surface is fixed during the optimization
+            surface_coords = compute_fun(
+                self._surface,
+                self._surface_data_keys,
+                params=self._surface.params_dict,
+                transforms=surface_transforms,
+                profiles={},
+                basis="xyz",
+            )["x"]
+            self._constants["surface_coords"] = surface_coords
+
         timer.stop("Precomputing transforms")
         if verbose > 1:
             timer.disp("Precomputing transforms")
@@ -598,14 +624,17 @@
             profiles=constants["equil_profiles"],
         )
         plasma_coords = rpz2xyz(jnp.array([data["R"], data["phi"], data["Z"]]).T)
-        surface_coords = compute_fun(
-            self._surface,
-            self._surface_data_keys,
-            params=surface_params,
-            transforms=constants["surface_transforms"],
-            profiles={},
-            basis="xyz",
-        )["x"]
+        if self._surface_fixed:
+            surface_coords = constants["surface_coords"]
+        else:
+            surface_coords = compute_fun(
+                self._surface,
+                self._surface_data_keys,
+                params=surface_params,
+                transforms=constants["surface_transforms"],
+                profiles={},
+                basis="xyz",
+            )["x"]
         d = jnp.linalg.norm(
             plasma_coords[:, None, :] - surface_coords[None, :, :], axis=-1
         )

tests/test_objective_funs.py

@@ -642,7 +642,11 @@ def test_plasma_vessel_distance():
     surf_grid = LinearGrid(M=5, N=6)
     plas_grid = LinearGrid(M=10, N=6)
     obj = PlasmaVesselDistance(
-        eq=eq, plasma_grid=plas_grid, surface_grid=surf_grid, surface=surface
+        eq=eq,
+        plasma_grid=plas_grid,
+        surface_grid=surf_grid,
+        surface=surface,
+        surface_fixed=True,
     )
     obj.build()
     d = obj.compute_unscaled(*obj.xs(eq, surface))