diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index cbc54561c9..88b653a084 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -20,22 +20,30 @@ @register_compute_fun( name="B_theta_mn", label="B_{\\theta, m, n}", - units="T \\cdot m}", + units="T \\cdot m", units_long="Tesla * meters", description="Fourier coefficients for covariant poloidal component of " "magnetic field.", dim=1, params=[], - transforms={"B": [[0, 0, 0]]}, + transforms={"B": [[0, 0, 0]], "grid": []}, profiles=[], coordinates="rtz", data=["B_theta"], + resolution_requirement="tz", + grid_requirement={"is_meshgrid": True}, M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", - resolution_requirement="tz", ) def _B_theta_mn(params, transforms, profiles, data, **kwargs): - data["B_theta_mn"] = transforms["B"].fit(data["B_theta"]) + B_theta = transforms["grid"].meshgrid_reshape(data["B_theta"], "rtz") + + def fitfun(x): + return transforms["B"].fit(x.flatten(order="F")) + + B_theta_mn = vmap(fitfun)(B_theta) + # modes stored as shape(rho, mn) flattened + data["B_theta_mn"] = B_theta_mn.flatten() return data @@ -43,7 +51,7 @@ def _B_theta_mn(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="B_phi_mn", label="B_{\\phi, m, n}", - units="T \\cdot m}", + units="T \\cdot m", units_long="Tesla * meters", description="Fourier coefficients for covariant toroidal component of " "magnetic field in (ρ,θ,ϕ) coordinates.", @@ -53,13 +61,21 @@ def _B_theta_mn(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["B_phi|r,t"], - M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", - N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", resolution_requirement="tz", + grid_requirement={"is_meshgrid": True}, aliases="B_zeta_mn", # TODO: remove when phi != zeta + M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", + N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _B_phi_mn(params, transforms, profiles, data, **kwargs): - data["B_phi_mn"] = transforms["B"].fit(data["B_phi|r,t"]) + B_phi = transforms["grid"].meshgrid_reshape(data["B_phi|r,t"], "rtz") + + def fitfun(x): + return transforms["B"].fit(x.flatten(order="F")) + + B_zeta_mn = vmap(fitfun)(B_phi) + # modes stored as shape(rho, mn) flattened + data["B_phi_mn"] = B_zeta_mn.flatten() return data @@ -72,15 +88,16 @@ def _B_phi_mn(params, transforms, profiles, data, **kwargs): + "Boozer Coordinates'", dim=1, params=[], - transforms={"w": [[0, 0, 0]], "B": [[0, 0, 0]]}, + transforms={"w": [[0, 0, 0]], "B": [[0, 0, 0]], "grid": []}, profiles=[], coordinates="rtz", data=["B_theta_mn", "B_phi_mn"], + grid_requirement={"is_meshgrid": True}, M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _w_mn(params, transforms, profiles, data, **kwargs): - w_mn = jnp.zeros((transforms["w"].basis.num_modes,)) + 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] wm = transforms["w"].basis.modes[:, 1] @@ -89,15 +106,19 @@ def _w_mn(params, transforms, profiles, data, **kwargs): mask_t = (Bm[:, None] == -wm) & (Bn[:, None] == wn) & (wm != 0) mask_z = (Bm[:, None] == wm) & (Bn[:, None] == -wn) & (wm == 0) & (wn != 0) - num_t = (mask_t @ sign(wn)) * data["B_theta_mn"] + num_t = (mask_t @ sign(wn)) * data["B_theta_mn"].reshape( + (transforms["grid"].num_rho, -1) + ) den_t = mask_t @ jnp.abs(wm) - num_z = (mask_z @ sign(wm)) * data["B_phi_mn"] + num_z = (mask_z @ sign(wm)) * data["B_phi_mn"].reshape( + (transforms["grid"].num_rho, -1) + ) den_z = mask_z @ jnp.abs(NFP * wn) - w_mn = jnp.where(mask_t.any(axis=0), mask_t.T @ safediv(num_t, den_t), w_mn) - w_mn = jnp.where(mask_z.any(axis=0), mask_z.T @ safediv(num_z, den_z), w_mn) + w_mn = jnp.where(mask_t.any(axis=0), (mask_t.T @ safediv(num_t, den_t).T).T, w_mn) + w_mn = jnp.where(mask_z.any(axis=0), (mask_z.T @ safediv(num_z, den_z).T).T, w_mn) - data["w_Boozer_mn"] = w_mn + data["w_Boozer_mn"] = w_mn.flatten() return data @@ -110,16 +131,22 @@ def _w_mn(params, transforms, profiles, data, **kwargs): + "'Transformation from VMEC to Boozer Coordinates'", dim=1, params=[], - transforms={"w": [[0, 0, 0]]}, + transforms={"w": [[0, 0, 0]], "grid": []}, profiles=[], coordinates="rtz", data=["w_Boozer_mn"], resolution_requirement="tz", + grid_requirement={"is_meshgrid": True}, M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _w(params, transforms, profiles, data, **kwargs): - data["w_Boozer"] = transforms["w"].transform(data["w_Boozer_mn"]) + grid = transforms["grid"] + 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") + w = jnp.moveaxis(w, 0, 1) + data["w_Boozer"] = w.flatten(order="F") return data @@ -132,16 +159,24 @@ def _w(params, transforms, profiles, data, **kwargs): + "'Transformation from VMEC to Boozer Coordinates', poloidal derivative", dim=1, params=[], - transforms={"w": [[0, 1, 0]]}, + transforms={"w": [[0, 1, 0]], "grid": []}, profiles=[], coordinates="rtz", data=["w_Boozer_mn"], resolution_requirement="tz", + grid_requirement={"is_meshgrid": True}, M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _w_t(params, transforms, profiles, data, **kwargs): - data["w_Boozer_t"] = transforms["w"].transform(data["w_Boozer_mn"], dt=1) + grid = transforms["grid"] + 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) + w_t = vmap(fun)(w_mn) # shape(rho, theta*zeta) + w_t = w_t.reshape((grid.num_rho, grid.num_theta, grid.num_zeta), order="F") + w_t = jnp.moveaxis(w_t, 0, 1) + data["w_Boozer_t"] = w_t.flatten(order="F") return data @@ -154,16 +189,24 @@ def _w_t(params, transforms, profiles, data, **kwargs): + "'Transformation from VMEC to Boozer Coordinates', toroidal derivative", dim=1, params=[], - transforms={"w": [[0, 0, 1]]}, + transforms={"w": [[0, 0, 1]], "grid": []}, profiles=[], coordinates="rtz", data=["w_Boozer_mn"], resolution_requirement="tz", + grid_requirement={"is_meshgrid": True}, M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _w_z(params, transforms, profiles, data, **kwargs): - data["w_Boozer_z"] = transforms["w"].transform(data["w_Boozer_mn"], dz=1) + grid = transforms["grid"] + 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) + w_z = vmap(fun)(w_mn) # shape(rho, theta*zeta) + w_z = w_z.reshape((grid.num_rho, grid.num_theta, grid.num_zeta), order="F") + w_z = jnp.moveaxis(w_z, 0, 1) + data["w_Boozer_z"] = w_z.flatten(order="F") return data @@ -290,21 +333,38 @@ def _sqrtg_B(params, transforms, profiles, data, **kwargs): description="Boozer harmonics of magnetic field", dim=1, params=[], - transforms={"B": [[0, 0, 0]]}, + transforms={"B": [[0, 0, 0]], "grid": []}, profiles=[], coordinates="rtz", data=["sqrt(g)_B", "|B|", "rho", "theta_B", "zeta_B"], + resolution_requirement="tz", + grid_requirement={"is_meshgrid": True}, M_booz="int: Maximum poloidal mode number for Boozer harmonics. Default 2*eq.M", N_booz="int: Maximum toroidal mode number for Boozer harmonics. Default 2*eq.N", ) def _B_mn(params, transforms, profiles, data, **kwargs): - nodes = jnp.array([data["rho"], data["theta_B"], data["zeta_B"]]).T norm = 2 ** (3 - jnp.sum((transforms["B"].basis.modes == 0), axis=1)) - data["|B|_mn"] = ( - norm # 1 if m=n=0, 2 if m=0 or n=0, 4 if m!=0 and n!=0 - * (transforms["B"].basis.evaluate(nodes).T @ (data["sqrt(g)_B"] * data["|B|"])) - / transforms["B"].grid.num_nodes + grid = transforms["grid"] + + def fun(rho, theta_B, zeta_B, sqrtg_B, B): + # this fits Boozer modes on a single surface + nodes = jnp.array([rho, theta_B, zeta_B]).T + B_mn = ( + norm # 1 if m=n=0, 2 if m=0 or n=0, 4 if m!=0 and n!=0 + * (transforms["B"].basis.evaluate(nodes).T @ (sqrtg_B * B)) + / transforms["B"].grid.num_nodes + ) + return B_mn + + def reshape(x): + return grid.meshgrid_reshape(x, "rtz").reshape((grid.num_rho, -1)) + + rho, theta_B, zeta_B, sqrtg_B, B = map( + reshape, + (data["rho"], data["theta_B"], data["zeta_B"], data["sqrt(g)_B"], data["|B|"]), ) + B_mn = vmap(fun)(rho, theta_B, zeta_B, sqrtg_B, B) + data["|B|_mn"] = B_mn.flatten() return data diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index 26341ec587..f8f30fa36d 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -63,6 +63,7 @@ def register_compute_fun( # noqa: C901 aliases=None, parameterization="desc.equilibrium.equilibrium.Equilibrium", resolution_requirement="", + grid_requirement=None, source_grid_requirement=None, **kwargs, ): @@ -110,6 +111,11 @@ def register_compute_fun( # noqa: C901 If the computation simply performs pointwise operations, instead of a reduction (such as integration) over a coordinate, then an empty string may be used to indicate no requirements. + grid_requirement : dict + Attributes of the grid that the compute function requires. + Also assumes dependencies were computed on such a grid. + As an example, quantities that require tensor product grids over 2 or more + coordinates may specify ``grid_requirement={"is_meshgrid": True}``. source_grid_requirement : dict Attributes of the source grid that the compute function requires. Also assumes dependencies were computed on such a grid. @@ -130,6 +136,8 @@ def register_compute_fun( # noqa: C901 aliases = [] if source_grid_requirement is None: source_grid_requirement = {} + if grid_requirement is None: + grid_requirement = {} if not isinstance(parameterization, (tuple, list)): parameterization = [parameterization] if not isinstance(aliases, (tuple, list)): @@ -168,6 +176,7 @@ def _decorator(func): "dependencies": deps, "aliases": aliases, "resolution_requirement": resolution_requirement, + "grid_requirement": grid_requirement, "source_grid_requirement": source_grid_requirement, } for p in parameterization: diff --git a/desc/compute/utils.py b/desc/compute/utils.py index 0c6e2f7de3..8d05bdfb88 100644 --- a/desc/compute/utils.py +++ b/desc/compute/utils.py @@ -33,7 +33,9 @@ def _parse_parameterization(p): return module + "." + klass.__qualname__ -def compute(parameterization, names, params, transforms, profiles, data=None, **kwargs): +def compute( # noqa: C901 + parameterization, names, params, transforms, profiles, data=None, **kwargs +): """Compute the quantity given by name on grid. Parameters @@ -88,6 +90,15 @@ def compute(parameterization, names, params, transforms, profiles, data=None, ** if "grid" in transforms: def check_fun(name): + reqs = data_index[p][name]["grid_requirement"] + for req in reqs: + errorif( + not hasattr(transforms["grid"], req) + or reqs[req] != getattr(transforms["grid"], req), + AttributeError, + f"Expected grid with '{req}:{reqs[req]}' to compute {name}.", + ) + reqs = data_index[p][name]["source_grid_requirement"] errorif( reqs and not hasattr(transforms["grid"], "source_grid"), @@ -517,6 +528,7 @@ def get_transforms( """ from desc.basis import DoubleFourierSeries + from desc.grid import LinearGrid from desc.transform import Transform method = "jitable" if jitable or kwargs.get("method") == "jitable" else "auto" @@ -556,8 +568,15 @@ def get_transforms( ) transforms[c] = c_transform elif c == "B": # used for Boozer transform + # assume grid is a meshgrid but only care about a single surface + if grid.num_rho > 1: + theta = grid.nodes[grid.unique_theta_idx, 1] + zeta = grid.nodes[grid.unique_zeta_idx, 2] + grid_B = LinearGrid(theta=theta, zeta=zeta, NFP=grid.NFP, sym=grid.sym) + else: + grid_B = grid transforms["B"] = Transform( - grid, + grid_B, DoubleFourierSeries( M=kwargs.get("M_booz", 2 * obj.M), N=kwargs.get("N_booz", 2 * obj.N), @@ -570,8 +589,15 @@ def get_transforms( method=method, ) elif c == "w": # used for Boozer transform + # assume grid is a meshgrid but only care about a single surface + if grid.num_rho > 1: + theta = grid.nodes[grid.unique_theta_idx, 1] + zeta = grid.nodes[grid.unique_zeta_idx, 2] + grid_w = LinearGrid(theta=theta, zeta=zeta, NFP=grid.NFP, sym=grid.sym) + else: + grid_w = grid transforms["w"] = Transform( - grid, + grid_w, DoubleFourierSeries( M=kwargs.get("M_booz", 2 * obj.M), N=kwargs.get("N_booz", 2 * obj.N), diff --git a/desc/objectives/_omnigenity.py b/desc/objectives/_omnigenity.py index 05f08356c0..1eb7a8da6e 100644 --- a/desc/objectives/_omnigenity.py +++ b/desc/objectives/_omnigenity.py @@ -47,7 +47,7 @@ class QuasisymmetryBoozer(_Objective): reverse mode and forward over reverse mode respectively. grid : Grid, optional Collocation grid containing the nodes to evaluate at. - Must be a LinearGrid with a single flux surface and sym=False. + Must be a LinearGrid with sym=False. Defaults to ``LinearGrid(M=M_booz, N=N_booz)``. helicity : tuple, optional Type of quasi-symmetry (M, N). Default = quasi-axisymmetry (1, 0). @@ -122,12 +122,6 @@ def build(self, use_jit=True, verbose=1): grid = self._grid errorif(grid.sym, ValueError, "QuasisymmetryBoozer grid must be non-symmetric") - errorif( - grid.num_rho != 1, - ValueError, - "QuasisymmetryBoozer grid must be on a single surface. " - "To target multiple surfaces, use multiple objectives.", - ) warnif( grid.num_theta < 2 * eq.M, RuntimeWarning, @@ -195,7 +189,7 @@ def compute(self, params, constants=None): Returns ------- f : ndarray - Quasi-symmetry flux function error at each node (T^3). + Symmetry breaking harmonics of B (T). """ if constants is None: @@ -207,8 +201,11 @@ def compute(self, params, constants=None): transforms=constants["transforms"], profiles=constants["profiles"], ) - B_mn = constants["matrix"] @ data["|B|_mn"] - return B_mn[constants["idx"]] + B_mn = data["|B|_mn"].reshape((constants["transforms"]["grid"].num_rho, -1)) + B_mn = constants["matrix"] @ B_mn.T + # output order = (rho, mn).flatten(), ie all the surfaces concatenated + # one after the other + return B_mn[constants["idx"]].T.flatten() @property def helicity(self): diff --git a/desc/plotting.py b/desc/plotting.py index 520d1c9524..b3c8fdeb17 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -2614,7 +2614,7 @@ def plot_boozer_modes( # noqa: C901 elif np.isscalar(rho) and rho > 1: rho = np.linspace(1, 0, num=rho, endpoint=False) - B_mn = np.array([[]]) + rho = np.sort(rho) M_booz = kwargs.pop("M_booz", 2 * eq.M) N_booz = kwargs.pop("N_booz", 2 * eq.N) linestyle = kwargs.pop("ls", "-") @@ -2632,16 +2632,15 @@ def plot_boozer_modes( # noqa: C901 else: matrix, modes = ptolemy_linear_transform(basis.modes) - for i, r in enumerate(rho): - grid = LinearGrid(M=2 * eq.M_grid, N=2 * eq.N_grid, NFP=eq.NFP, rho=np.array(r)) - transforms = get_transforms( - "|B|_mn", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz - ) - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - data = eq.compute("|B|_mn", grid=grid, transforms=transforms) - b_mn = np.atleast_2d(matrix @ data["|B|_mn"]) - B_mn = np.vstack((B_mn, b_mn)) if B_mn.size else b_mn + grid = LinearGrid(M=2 * eq.M_grid, N=2 * eq.N_grid, NFP=eq.NFP, rho=rho) + transforms = get_transforms( + "|B|_mn", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz + ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + data = eq.compute("|B|_mn", grid=grid, transforms=transforms) + B_mn = data["|B|_mn"].reshape((len(rho), -1)) + B_mn = np.atleast_2d(matrix @ B_mn.T).T zidx = np.where((modes[:, 1:] == np.array([[0, 0]])).all(axis=1))[0] if norm: @@ -3010,6 +3009,7 @@ def plot_qs_error( # noqa: 16 fxn too complex rho = np.linspace(1, 0, num=20, endpoint=False) elif np.isscalar(rho) and rho > 1: rho = np.linspace(1, 0, num=rho, endpoint=False) + rho = np.sort(rho) fig, ax = _format_ax(ax, figsize=kwargs.pop("figsize", None)) @@ -3027,119 +3027,92 @@ def plot_qs_error( # noqa: 16 fxn too complex R0 = data["R0"] B0 = np.mean(data["|B|"] * data["sqrt(g)"]) / np.mean(data["sqrt(g)"]) - f_B = np.array([]) - f_C = np.array([]) - f_T = np.array([]) - plot_data = {} - for i, r in enumerate(rho): - grid = LinearGrid(M=2 * eq.M_grid, N=2 * eq.N_grid, NFP=eq.NFP, rho=np.array(r)) - if fB: - transforms = get_transforms( - "|B|_mn", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz - ) - if i == 0: # only need to do this once for the first rho surface - matrix, modes, idx = ptolemy_linear_transform( - transforms["B"].basis.modes, - helicity=helicity, - NFP=transforms["B"].basis.NFP, - ) - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - data = eq.compute( - ["|B|_mn", "B modes"], grid=grid, transforms=transforms - ) - B_mn = matrix @ data["|B|_mn"] - f_b = np.sqrt(np.sum(B_mn[idx] ** 2)) / np.sqrt(np.sum(B_mn**2)) - f_B = np.append(f_B, f_b) - if fC: - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - data = eq.compute("f_C", grid=grid, helicity=helicity) - f_c = ( - np.mean(np.abs(data["f_C"]) * data["sqrt(g)"]) - / np.mean(data["sqrt(g)"]) - / B0**3 - ) - f_C = np.append(f_C, f_c) - if fT: - with warnings.catch_warnings(): - warnings.simplefilter("ignore") - data = eq.compute("f_T", grid=grid) - f_t = ( - np.mean(np.abs(data["f_T"]) * data["sqrt(g)"]) - / np.mean(data["sqrt(g)"]) - * R0**2 - / B0**4 - ) - f_T = np.append(f_T, f_t) + plot_data = {"rho": rho} - plot_data["f_B"] = f_B - plot_data["f_C"] = f_C - plot_data["f_T"] = f_T - plot_data["rho"] = rho + grid = LinearGrid(M=2 * eq.M_grid, N=2 * eq.N_grid, NFP=eq.NFP, rho=rho) + names = [] + if fB: + names += ["|B|_mn"] + transforms = get_transforms( + "|B|_mn", obj=eq, grid=grid, M_booz=M_booz, N_booz=N_booz + ) + matrix, modes, idx = ptolemy_linear_transform( + transforms["B"].basis.modes, + helicity=helicity, + NFP=transforms["B"].basis.NFP, + ) + if fC or fT: + names += ["sqrt(g)"] + if fC: + names += ["f_C"] + if fT: + names += ["f_T"] - if log: - if fB: - ax.semilogy( - rho, - f_B, - ls=ls[0 % len(ls)], - c=colors[0 % len(colors)], - marker=markers[0 % len(markers)], - label=labels[0 % len(labels)], - lw=lw[0 % len(lw)], - ) - if fC: - ax.semilogy( - rho, - f_C, - ls=ls[1 % len(ls)], - c=colors[1 % len(colors)], - marker=markers[1 % len(markers)], - label=labels[1 % len(labels)], - lw=lw[1 % len(lw)], - ) - if fT: - ax.semilogy( - rho, - f_T, - ls=ls[2 % len(ls)], - c=colors[2 % len(colors)], - marker=markers[2 % len(markers)], - label=labels[2 % len(labels)], - lw=lw[2 % len(lw)], - ) - else: - if fB: - ax.plot( - rho, - f_B, - ls=ls[0 % len(ls)], - c=colors[0 % len(colors)], - marker=markers[0 % len(markers)], - label=labels[0 % len(labels)], - lw=lw[0 % len(lw)], - ) - if fC: - ax.plot( - rho, - f_C, - ls=ls[1 % len(ls)], - c=colors[1 % len(colors)], - marker=markers[1 % len(markers)], - label=labels[1 % len(labels)], - lw=lw[1 % len(lw)], - ) - if fT: - ax.plot( - rho, - f_T, - ls=ls[2 % len(ls)], - c=colors[2 % len(colors)], - marker=markers[2 % len(markers)], - label=labels[2 % len(labels)], - lw=lw[2 % len(lw)], - ) + with warnings.catch_warnings(): + warnings.simplefilter("ignore") + data = eq.compute( + names, grid=grid, M_booz=M_booz, N_booz=N_booz, helicity=helicity + ) + + if fB: + B_mn = data["|B|_mn"].reshape((len(rho), -1)) + B_mn = (matrix @ B_mn.T).T + f_B = np.sqrt(np.sum(B_mn[:, idx] ** 2, axis=-1)) / np.sqrt( + np.sum(B_mn**2, axis=-1) + ) + plot_data["f_B"] = f_B + if fC: + sqrtg = grid.meshgrid_reshape(data["sqrt(g)"], "rtz") + f_C = grid.meshgrid_reshape(data["f_C"], "rtz") + f_C = ( + np.mean(np.abs(f_C) * sqrtg, axis=(1, 2)) + / np.mean(sqrtg, axis=(1, 2)) + / B0**3 + ) + plot_data["f_C"] = f_C + if fT: + sqrtg = grid.meshgrid_reshape(data["sqrt(g)"], "rtz") + f_T = grid.meshgrid_reshape(data["f_T"], "rtz") + f_T = ( + np.mean(np.abs(f_T) * sqrtg, axis=(1, 2)) + / np.mean(sqrtg, axis=(1, 2)) + * R0**2 + / B0**4 + ) + plot_data["f_T"] = f_T + + plot_op = ax.semilogy if log else ax.plot + + if fB: + plot_op( + rho, + f_B, + ls=ls[0 % len(ls)], + c=colors[0 % len(colors)], + marker=markers[0 % len(markers)], + label=labels[0 % len(labels)], + lw=lw[0 % len(lw)], + ) + if fC: + plot_op( + rho, + f_C, + ls=ls[1 % len(ls)], + c=colors[1 % len(colors)], + marker=markers[1 % len(markers)], + label=labels[1 % len(labels)], + lw=lw[1 % len(lw)], + ) + if fT: + plot_op( + rho, + f_T, + ls=ls[2 % len(ls)], + c=colors[2 % len(colors)], + marker=markers[2 % len(markers)], + label=labels[2 % len(labels)], + lw=lw[2 % len(lw)], + ) ax.set_xlabel(_AXIS_LABELS_RTZ[0], fontsize=xlabel_fontsize) if ylabel: diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index eef5bbf2f6..d72778328e 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_compute_funs.py b/tests/test_compute_funs.py index 43c3d81449..d9f85f4339 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -1134,6 +1134,24 @@ def test_boozer_transform(): ) +@pytest.mark.unit +def test_boozer_transform_multiple_surfaces(): + """Test that computing over multiple surfaces is the same as over 1 at a time.""" + eq = get("HELIOTRON") + grid1 = LinearGrid(rho=0.6, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP) + grid2 = LinearGrid(rho=0.8, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP) + grid3 = LinearGrid(rho=np.array([0.6, 0.8]), M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP) + data1 = eq.compute("|B|_mn", grid=grid1, M_booz=eq.M, N_booz=eq.N) + data2 = eq.compute("|B|_mn", grid=grid2, M_booz=eq.M, N_booz=eq.N) + data3 = eq.compute("|B|_mn", grid=grid3, M_booz=eq.M, N_booz=eq.N) + np.testing.assert_allclose( + data1["|B|_mn"], data3["|B|_mn"].reshape((grid3.num_rho, -1))[0] + ) + np.testing.assert_allclose( + data2["|B|_mn"], data3["|B|_mn"].reshape((grid3.num_rho, -1))[1] + ) + + @pytest.mark.unit def test_compute_averages(): """Test that computing averages uses the correct grid.""" diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 97cb25f45b..1619e5c846 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -367,6 +367,51 @@ def test_qh_boozer(self): # should have the same values up until then np.testing.assert_allclose(f[idx_f][:120], B_mn[idx_B][:120]) + @pytest.mark.unit + def test_qh_boozer_multiple_surfaces(self): + """Test for computing Boozer error on multiple surfaces.""" + eq = get("WISTELL-A") # WISTELL-A is optimized for QH symmetry + helicity = (1, -eq.NFP) + M_booz = eq.M + N_booz = eq.N + grid1 = LinearGrid(rho=0.5, M=2 * eq.M, N=2 * eq.N, NFP=eq.NFP, sym=False) + grid2 = LinearGrid(rho=1.0, M=2 * eq.M, N=2 * eq.N, NFP=eq.NFP, sym=False) + grid3 = LinearGrid( + rho=np.array([0.5, 1.0]), M=2 * eq.M, N=2 * eq.N, NFP=eq.NFP, sym=False + ) + + obj1 = QuasisymmetryBoozer( + helicity=helicity, + M_booz=M_booz, + N_booz=N_booz, + grid=grid1, + normalize=False, + eq=eq, + ) + obj2 = QuasisymmetryBoozer( + helicity=helicity, + M_booz=M_booz, + N_booz=N_booz, + grid=grid2, + normalize=False, + eq=eq, + ) + obj3 = QuasisymmetryBoozer( + helicity=helicity, + M_booz=M_booz, + N_booz=N_booz, + grid=grid3, + normalize=False, + eq=eq, + ) + obj1.build() + obj2.build() + obj3.build() + f1 = obj1.compute_unscaled(*obj1.xs(eq)) + f2 = obj2.compute_unscaled(*obj2.xs(eq)) + f3 = obj3.compute_unscaled(*obj3.xs(eq)) + np.testing.assert_allclose(f3, np.concatenate([f1, f2]), atol=1e-14) + @pytest.mark.unit def test_qs_twoterm(self): """Test calculation of two term QS metric.""" @@ -441,11 +486,6 @@ def test_qs_boozer_grids(self): with pytest.raises(ValueError): QuasisymmetryBoozer(eq=eq, grid=grid).build() - # multiple flux surfaces - grid = LinearGrid(M=eq.M, N=eq.N, NFP=eq.NFP, rho=[0.25, 0.5, 0.75, 1]) - with pytest.raises(ValueError): - QuasisymmetryBoozer(eq=eq, grid=grid).build() - @pytest.mark.unit def test_mercier_stability(self): """Test calculation of mercier stability criteria."""