From a4b9dbcf11f7f97359521df4358d0508c44f0b34 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 13 Nov 2023 18:53:09 -0500 Subject: [PATCH 1/4] add flag to PlasmaVesselDistance to avoid recomputing surface coordinates when the surface is fixed during optimization --- desc/objectives/_geometry.py | 45 +++++++++++++++++++++++++++++------- 1 file changed, 37 insertions(+), 8 deletions(-) diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index abd113c958..2b097823ec 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -401,6 +401,11 @@ class PlasmaVesselDistance(_Objective): 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 @@ class PlasmaVesselDistance(_Objective): 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. + True 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 @@ def __init__( surface_grid=None, plasma_grid=None, use_softmin=False, + surface_fixed=True, alpha=1.0, name="plasma-vessel distance", ): @@ -476,6 +488,7 @@ def __init__( 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 @@ def build(self, use_jit=True, verbose=1): "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 @@ def compute(self, equil_params, surface_params, constants=None): 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 ) From 7c1488d1d4a326d4b59cf61a786c385e8de6e6a9 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Mon, 13 Nov 2023 18:56:55 -0500 Subject: [PATCH 2/4] update changelog --- CHANGELOG.md | 5 ++++- 1 file changed, 4 insertions(+), 1 deletion(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 289b677f85..9835bf9128 100644 --- a/CHANGELOG.md +++ b/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 True by default. v0.10.2 ------- From d97fa84474cde83e44f3a37d10f0d49b16e36308 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 14 Nov 2023 13:44:52 -0500 Subject: [PATCH 3/4] change surface_fixed flag to be false by default --- CHANGELOG.md | 2 +- desc/objectives/_geometry.py | 4 ++-- tests/test_objective_funs.py | 6 +++++- 3 files changed, 8 insertions(+), 4 deletions(-) diff --git a/CHANGELOG.md b/CHANGELOG.md index 9835bf9128..40a95ca21c 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -6,7 +6,7 @@ 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 True by default. +optimization. Set to False by default. v0.10.2 ------- diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 2b097823ec..121df1302c 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -450,7 +450,7 @@ class PlasmaVesselDistance(_Objective): 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. - True by default. + 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. @@ -478,7 +478,7 @@ def __init__( surface_grid=None, plasma_grid=None, use_softmin=False, - surface_fixed=True, + surface_fixed=False, alpha=1.0, name="plasma-vessel distance", ): diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index c3445d940e..af93dc58b2 100644 --- a/tests/test_objective_funs.py +++ b/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)) From c08b2cccd6436fdb44974aeecfc604a9747d1241 Mon Sep 17 00:00:00 2001 From: Dario Panici Date: Tue, 14 Nov 2023 21:17:01 -0500 Subject: [PATCH 4/4] update docstring --- desc/objectives/_geometry.py | 9 +++++---- 1 file changed, 5 insertions(+), 4 deletions(-) diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 121df1302c..2ee17dec3b 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -401,10 +401,11 @@ class PlasmaVesselDistance(_Objective): 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: By default, assumes the surface is not fixed and its coordinates are computed + at every iteration, for example if the winding surface you compare to is part of the + optimization and thus changing. + If the bounding surface is fixed, set surface_fixed=True to precompute the surface + coordinates and improve the efficiency of the calculation 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