diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index f828b868d9..108224f564 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -158,7 +158,7 @@ def _compute_A_of_z(grid, data, extrap=False, mean=False, expand_out=False): return jnp.mean(grid.compress(data["A(z)"], surface_label="zeta")) max_rho = jnp.max(data["rho"]) - if isinstance(grid, QuadratureGrid) or "n_rho" not in data: # TODO(#1761) + if isinstance(grid, QuadratureGrid): assert extrap A = surface_integrals( grid, @@ -179,17 +179,16 @@ def _compute_A_of_z(grid, data, extrap=False, mean=False, expand_out=False): # on constant ζ surface is 1. Then choose v = (w - w^ζ). # OR make sure that integration contour is along constant ϕ surface # using source_grid_requirement="rtp" and use e_rho|t,p and e_theta|r,p. - n = data["n_rho"] - n = n.at[:, 1].set(0) - n = n / jnp.linalg.norm(n, axis=-1)[:, jnp.newaxis] - A = jnp.abs( - line_integrals( - grid, - data["Z"] * n[:, 2] * safenorm(data["e_theta"], axis=-1), - line_label="theta", - fix_surface=("rho", max_rho), - expand_out=False, - ) + dl = jnp.sqrt(data["g_tt"]) + t = data["e_theta"] / dl[:, jnp.newaxis] + n = -data["e_theta_t"] + dot(data["e_theta_t"], t)[:, jnp.newaxis] * t + n = n[:, 2] / jnp.linalg.norm(n, axis=-1) + A = line_integrals( + grid, + dl * n * data["Z"], + line_label="theta", + fix_surface=("rho", max_rho), + expand_out=False, ) if extrap: # To approximate area at ρ ~ 1, we scale by ρ⁻², assuming the integrand @@ -216,12 +215,16 @@ def _compute_A_of_z(grid, data, extrap=False, mean=False, expand_out=False): transforms={"grid": []}, profiles=[], coordinates="z", - data=["Z", "n_rho", "e_theta", "rho", "|e_rho x e_theta|"], - parameterization=["desc.equilibrium.equilibrium.Equilibrium"], + data=["rho", "Z", "e_theta", "e_theta_t", "g_tt", "|e_rho x e_theta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], resolution_requirement="t", grid_requirement={"sym": False}, ) def _A_of_z(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym # noqa: unused dependency data["A(z)"] = _compute_A_of_z( transforms["grid"], data, extrap=True, expand_out=True @@ -241,7 +244,7 @@ def _A_of_z(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="", - data=["Z", "n_rho", "e_theta", "rho", "|e_rho x e_theta|"], + data=["rho", "Z", "e_theta", "e_theta_t", "g_tt", "|e_rho x e_theta|"], parameterization=["desc.equilibrium.equilibrium.Equilibrium"], resolution_requirement="tz", ) @@ -251,31 +254,6 @@ def _A(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="A(z)", - label="A(\\zeta)", - units="m^{2}", - units_long="square meters", - description="Area of enclosed cross-section (enclosed constant zeta surface), " - "extrapolated to last closed flux surface", - dim=1, - params=[], - transforms={"grid": []}, - profiles=[], - coordinates="z", - data=["rho", "|e_rho x e_theta|"], - parameterization=["desc.geometry.surface.ZernikeRZToroidalSection"], - resolution_requirement="rt", - grid_requirement={"sym": False}, -) -def _A_of_z_cross_section_surface(params, transforms, profiles, data, **kwargs): - # noqa: unused dependency - data["A(z)"] = _compute_A_of_z( - transforms["grid"], data, extrap=True, expand_out=True - ) - return data - - @register_compute_fun( name="A", label="A", @@ -288,13 +266,14 @@ def _A_of_z_cross_section_surface(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="", - data=["rho", "|e_rho x e_theta|"], + data=["rho", "Z", "e_theta", "e_theta_t", "g_tt", "|e_rho x e_theta|"], parameterization=["desc.geometry.surface.ZernikeRZToroidalSection"], - resolution_requirement="rt", + resolution_requirement="t", # Needs to be False since resolution requirement lacks toroidal resolution. grid_requirement={"sym": False}, ) def _A_cross_section_surface(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym # noqa: unused dependency data["A"] = _compute_A_of_z(transforms["grid"], data, extrap=True, mean=True) return data @@ -311,12 +290,13 @@ def _A_cross_section_surface(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="z", - data=["Z", "n_rho", "e_theta", "rho"], + data=["rho", "Z", "e_theta", "e_theta_t", "g_tt"], parameterization=["desc.geometry.surface.FourierRZToroidalSurface"], resolution_requirement="t", grid_requirement={"sym": False}, ) def _A_of_z_flux_surface(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym # noqa: unused dependency data["A(z)"] = _compute_A_of_z( transforms["grid"], data, extrap=False, expand_out=True @@ -335,7 +315,7 @@ def _A_of_z_flux_surface(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="", - data=["Z", "n_rho", "e_theta", "rho"], + data=["rho", "Z", "e_theta", "e_theta_t", "g_tt"], parameterization=["desc.geometry.surface.FourierRZToroidalSurface"], resolution_requirement="tz", ) @@ -562,6 +542,7 @@ def _R0_over_a(params, transforms, profiles, data, **kwargs): grid_requirement={"sym": False}, ) def _perimeter_of_z(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym max_rho = jnp.max(data["rho"]) data["perimeter(z)"] = ( line_integrals( @@ -595,6 +576,7 @@ def _perimeter_of_z(params, transforms, profiles, data, **kwargs): grid_requirement={"sym": False}, ) def _perimeter_of_z_flux_surface(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym max_rho = jnp.max(data["rho"]) data["perimeter(z)"] = line_integrals( transforms["grid"], diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index c7612e7767..00dc24b30a 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -38,6 +38,7 @@ N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _B_theta_mn(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym B_theta = transforms["grid"].meshgrid_reshape(data["B_theta"], "rtz") def fitfun(x): @@ -70,6 +71,7 @@ def fitfun(x): N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _B_phi_mn(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym B_phi = transforms["grid"].meshgrid_reshape(data["B_phi|r,t"], "rtz") def fitfun(x): @@ -99,6 +101,7 @@ def fitfun(x): N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _w_mn(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym w_mn = jnp.zeros((transforms["grid"].num_rho, transforms["w"].basis.num_modes)) Bm = transforms["B"].basis.modes[:, 1] Bn = transforms["B"].basis.modes[:, 2] @@ -144,6 +147,7 @@ def _w_mn(params, transforms, profiles, data, **kwargs): ) def _w(params, transforms, profiles, data, **kwargs): grid = transforms["grid"] + assert not grid.sym w_mn = data["w_Boozer_mn"].reshape((grid.num_rho, -1)) w = vmap(transforms["w"].transform)(w_mn) # shape(rho, theta*zeta) w = w.reshape((grid.num_rho, grid.num_theta, grid.num_zeta), order="F") @@ -172,6 +176,7 @@ def _w(params, transforms, profiles, data, **kwargs): ) def _w_t(params, transforms, profiles, data, **kwargs): grid = transforms["grid"] + assert not grid.sym w_mn = data["w_Boozer_mn"].reshape((grid.num_rho, -1)) # need to close over dt which can't be vmapped fun = lambda x: transforms["w"].transform(x, dt=1) @@ -202,6 +207,7 @@ def _w_t(params, transforms, profiles, data, **kwargs): ) def _w_z(params, transforms, profiles, data, **kwargs): grid = transforms["grid"] + assert not grid.sym w_mn = data["w_Boozer_mn"].reshape((grid.num_rho, -1)) # need to close over dz which can't be vmapped fun = lambda x: transforms["w"].transform(x, dz=1) @@ -256,6 +262,7 @@ def _nu(params, transforms, profiles, data, **kwargs): N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _nu_B_mn(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym norm = data["Boozer transform modes norm"] grid = transforms["grid"] @@ -431,6 +438,7 @@ def _sqrtg_B(params, transforms, profiles, data, **kwargs): N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _sqrtg_Boozer_mn(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym norm = data["Boozer transform modes norm"] grid = transforms["grid"] @@ -488,6 +496,7 @@ def reshape(x): aliases=["|B|_mn"], ) def _B_mn(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym norm = data["Boozer transform modes norm"] grid = transforms["grid"] @@ -544,6 +553,7 @@ def reshape(x): N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _R_mn(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym norm = data["Boozer transform modes norm"] grid = transforms["grid"] @@ -600,6 +610,7 @@ def reshape(x): N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _Z_mn(params, transforms, profiles, data, **kwargs): + assert not transforms["grid"].sym norm = data["Boozer transform modes norm"] grid = transforms["grid"] diff --git a/desc/objectives/_stability.py b/desc/objectives/_stability.py index b0b04d53a8..8f60b72a06 100644 --- a/desc/objectives/_stability.py +++ b/desc/objectives/_stability.py @@ -5,7 +5,7 @@ from desc.backend import jnp from desc.compute import get_profiles, get_transforms from desc.compute.utils import _compute as compute_fun -from desc.grid import LinearGrid +from desc.grid import LinearGrid, QuadratureGrid from desc.utils import ResolutionWarning, Timer, errorif, setdefault, warnif from .normalization import compute_scaling_factors @@ -402,6 +402,7 @@ class BallooningStability(_Objective): _static_attrs = _Objective._static_attrs + [ "_iota_keys", + "_a_keys", "_Neigvals", "_nturns", "_nzetaperturn", @@ -479,7 +480,8 @@ def build(self, use_jit=True, verbose=1): Level of output. """ - self._iota_keys = ["iota", "iota_r", "shear", "a"] + self._iota_keys = ["iota", "iota_r", "shear"] + self._a_keys = ["a"] eq = self.things[0] iota_grid = LinearGrid( @@ -492,7 +494,9 @@ def build(self, use_jit=True, verbose=1): ) assert not iota_grid.axis.size self._dim_f = iota_grid.num_rho - self._add_lcfs - transforms = get_transforms(self._iota_keys, eq, iota_grid) + iota_transforms = get_transforms(self._iota_keys, eq, iota_grid) + a_grid = QuadratureGrid(eq.L_grid, eq.M_grid, eq.N_grid, eq.NFP) + a_transforms = get_transforms(self._a_keys, eq, a_grid) profiles = get_profiles( self._iota_keys + ["ideal ballooning lambda"], eq, iota_grid ) @@ -508,7 +512,9 @@ def build(self, use_jit=True, verbose=1): +self._nturns * self._nzetaperturn, ), "zeta0": self._zeta0, - "iota_transforms": transforms, + "iota_transforms": iota_transforms, + "a_transforms": a_transforms, + "a_profiles": get_profiles(self._a_keys, eq, a_grid), "profiles": profiles, "quad_weights": 1.0, } @@ -541,6 +547,13 @@ def compute(self, params, constants=None): constants["iota_transforms"], constants["profiles"], ) + a_data = compute_fun( + eq, + self._a_keys, + params, + constants["a_transforms"], + constants["a_profiles"], + ) iota_grid = constants["iota_transforms"]["grid"] def get(key): @@ -564,7 +577,7 @@ def get(key): if (key != "iota" and key != "a") } data["iota"] = grid.expand(iota) - data["a"] = iota_data["a"] + data["a"] = a_data["a"] data = compute_fun( eq, ["ideal ballooning lambda"], diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index 6b9750fcbc..26b4350da5 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 95f297bc2b..16a95e42b8 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -2113,7 +2113,7 @@ def test_objective_against_compute_bounce(self): obj.build() assert obj._hyperparam["num_well"] == opts["num_well"] np.testing.assert_allclose( - obj.compute(eq.params_dict), grid.compress(data[names[0]]) + obj.compute(eq.params_dict), grid.compress(data[names[0]]), atol=1e-4 ) obj = GammaC(eq, grid=obj_grid, nufft_eps=1e-7, X=X, Y=Y, **opts) obj.build()