Skip to content

Commit

Permalink
Merge branch 'master' into dp/from_input_kwargs
Browse files Browse the repository at this point in the history
  • Loading branch information
rahulgaur104 authored Aug 27, 2024
2 parents 71eeefd + fd73bdb commit 15180dd
Show file tree
Hide file tree
Showing 10 changed files with 44 additions and 54 deletions.
38 changes: 22 additions & 16 deletions desc/equilibrium/coords.py
Original file line number Diff line number Diff line change
Expand Up @@ -91,13 +91,28 @@ def map_coordinates( # noqa: C901
inbasis = tuple(inbasis)
outbasis = tuple(outbasis)

basis_derivs = tuple(f"{X}_{d}" for X in inbasis for d in ("r", "t", "z"))
for key in basis_derivs:
errorif(
key not in data_index["desc.equilibrium.equilibrium.Equilibrium"],
NotImplementedError,
f"don't have recipe to compute partial derivative {key}",
)

profiles = get_profiles(inbasis + basis_derivs, eq)

# TODO: make this work for permutations of in/out basis
if outbasis == ("rho", "theta", "zeta"):
# TODO: get iota if not supplied using below logic
if inbasis == ("rho", "alpha", "zeta") and "iota" in kwargs:
if inbasis == ("rho", "alpha", "zeta"):
if "iota" in kwargs:
iota = kwargs.pop("iota")
else:
if profiles["iota"] is None:
profiles["iota"] = eq.get_profile(["iota", "iota_r"], params=params)
iota = profiles["iota"].compute(Grid(coords, sort=False, jitable=True))
return _map_clebsch_coordinates(
coords,
kwargs.pop("iota"),
iota,
params["L_lmn"],
eq.L_basis,
guess[:, 1] if guess is not None else None,
Expand All @@ -118,29 +133,20 @@ def map_coordinates( # noqa: C901
**kwargs,
)

basis_derivs = tuple(f"{X}_{d}" for X in inbasis for d in ("r", "t", "z"))
for key in basis_derivs:
errorif(
key not in data_index["desc.equilibrium.equilibrium.Equilibrium"],
NotImplementedError,
f"don't have recipe to compute partial derivative {key}",
)
# do surface average to get iota once
if "iota" in profiles and profiles["iota"] is None:
profiles["iota"] = eq.get_profile(["iota", "iota_r"], params=params)
params["i_l"] = profiles["iota"].params

rhomin = kwargs.pop("rhomin", tol / 10)
warnif(period is None, msg="Assuming no periodicity.")
period = np.asarray(setdefault(period, (np.inf, np.inf, np.inf)))
coords = _periodic(coords, period)

profiles = get_profiles(inbasis + basis_derivs, eq)
p = "desc.equilibrium.equilibrium.Equilibrium"
names = inbasis + basis_derivs + outbasis
deps = list(set(get_data_deps(names, obj=p) + list(names)))

# do surface average to get iota once
if "iota" in profiles and profiles["iota"] is None:
profiles["iota"] = eq.get_profile(["iota", "iota_r"], params=params)
params["i_l"] = profiles["iota"].params

@functools.partial(jit, static_argnums=1)
def compute(y, basis):
grid = Grid(y, sort=False, jitable=True)
Expand Down
5 changes: 4 additions & 1 deletion desc/grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -1295,7 +1295,10 @@ def _create_nodes(self, L=1, M=1, N=1, NFP=1):
self._N = check_nonnegint(N, "N", False)
self._NFP = check_posint(NFP, "NFP", False)
self._period = (np.inf, 2 * np.pi, 2 * np.pi / self._NFP)
L = L + 1
# floor divide (L+2) by 2 bc only need (L+1)/2 points to
# integrate L-th order jacobi polynomial exactly, so this
# ensures we have enough pts for both odd and even L
L = (L + 2) // 2
M = 2 * M + 1
N = 2 * N + 1

Expand Down
Binary file modified tests/baseline/test_plot_b_mag.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified tests/baseline/test_plot_grid_quad.png
Loading
Sorry, something went wrong. Reload?
Sorry, we cannot display this file.
Sorry, this file is invalid so it cannot be displayed.
Binary file modified tests/inputs/master_compute_data_rpz.pkl
Binary file not shown.
2 changes: 1 addition & 1 deletion tests/test_bootstrap.py
Original file line number Diff line number Diff line change
Expand Up @@ -1196,7 +1196,7 @@ def test_BootstrapRedlConsistency_normalization(self):
)

# Results are not perfectly identical because ln(Lambda) is not quite invariant.
np.testing.assert_allclose(results, expected, rtol=2e-3)
np.testing.assert_allclose(results, expected, rtol=3e-3)

@pytest.mark.regression
def test_bootstrap_consistency_iota(self, TmpDir):
Expand Down
30 changes: 6 additions & 24 deletions tests/test_equilibrium.py
Original file line number Diff line number Diff line change
Expand Up @@ -15,7 +15,6 @@
from desc.io import InputReader
from desc.objectives import ForceBalance, ObjectiveFunction, get_equilibrium_objective
from desc.profiles import PowerSeriesProfile
from desc.utils import errorif

from .utils import area_difference, compute_coords

Expand Down Expand Up @@ -50,33 +49,16 @@ def test_map_coordinates():
"""Test root finding for (rho,theta,zeta) for common use cases."""
# finding coordinates along a single field line
eq = get("NCSX")
with pytest.warns(UserWarning, match="Reducing radial"):
eq.change_resolution(3, 3, 3, 6, 6, 6)
n = 100
coords = np.array([np.ones(n), np.zeros(n), np.linspace(0, 10 * np.pi, n)]).T
grid = LinearGrid(rho=1, M=eq.M_grid, N=eq.N_grid, sym=eq.sym, NFP=eq.NFP)
iota = grid.compress(eq.compute("iota", grid=grid)["iota"])
iota = np.broadcast_to(iota, shape=(n,))

tol = 1e-5
out_1 = eq.map_coordinates(
coords, inbasis=["rho", "alpha", "zeta"], iota=iota, tol=tol
)
assert np.isfinite(out_1).all()
out_2 = eq.map_coordinates(
out = eq.map_coordinates(
coords,
inbasis=["rho", "alpha", "zeta"],
period=(np.inf, 2 * np.pi, np.inf),
tol=tol,
)
assert np.isfinite(out_2).all()
diff = out_1 - out_2
errorif(
not np.all(
np.isclose(diff, 0, atol=2 * tol)
| np.isclose(np.abs(diff), 2 * np.pi, atol=2 * tol)
),
AssertionError,
f"diff: {diff}",
)
assert np.isfinite(out).all()

eq = get("DSHAPE")

Expand All @@ -89,9 +71,9 @@ def test_map_coordinates():

grid = Grid(np.vstack([rho, theta, zeta]).T, sort=False)
in_data = eq.compute(inbasis, grid=grid)
in_coords = np.stack([in_data[k] for k in inbasis], axis=-1)
in_coords = np.column_stack([in_data[k] for k in inbasis])
out_data = eq.compute(outbasis, grid=grid)
out_coords = np.stack([out_data[k] for k in outbasis], axis=-1)
out_coords = np.column_stack([out_data[k] for k in outbasis])

out = eq.map_coordinates(
in_coords,
Expand Down
10 changes: 6 additions & 4 deletions tests/test_grid.py
Original file line number Diff line number Diff line change
Expand Up @@ -482,14 +482,16 @@ def test_quadrature_grid(self):
N = 0
NFP = 1
grid_quad = QuadratureGrid(L, M, N, NFP)
roots, weights = special.js_roots(3, 2, 2)
roots, weights = special.js_roots(2, 2, 2)
quadrature_nodes = np.stack(
[
np.array([roots[0]] * 5 + [roots[1]] * 5 + [roots[2]] * 5),
np.array(
[0, 2 * np.pi / 5, 4 * np.pi / 5, 6 * np.pi / 5, 8 * np.pi / 5] * 3
[roots[0]] * grid_quad.num_theta + [roots[1]] * grid_quad.num_theta
),
np.zeros(15),
np.array(
[0, 2 * np.pi / 5, 4 * np.pi / 5, 6 * np.pi / 5, 8 * np.pi / 5] * 2
),
np.zeros(10),
]
).T
np.testing.assert_allclose(grid_quad.spacing.prod(axis=1), grid_quad.weights)
Expand Down
4 changes: 2 additions & 2 deletions tests/test_objective_funs.py
Original file line number Diff line number Diff line change
Expand Up @@ -208,8 +208,8 @@ def test(eq):
obj.build()
f = obj.compute_unscaled(*obj.xs(eq))
f_scaled = obj.compute_scaled_error(*obj.xs(eq))
np.testing.assert_allclose(f, 1.3 / 0.7, rtol=5e-3)
np.testing.assert_allclose(f_scaled, 2 * (1.3 / 0.7), rtol=5e-3)
np.testing.assert_allclose(f, 1.3 / 0.7, rtol=8e-3)
np.testing.assert_allclose(f_scaled, 2 * (1.3 / 0.7), rtol=8e-3)

test(get("HELIOTRON"))
test(get("HELIOTRON").surface)
Expand Down
9 changes: 3 additions & 6 deletions tests/test_plotting.py
Original file line number Diff line number Diff line change
Expand Up @@ -5,7 +5,6 @@
import matplotlib.pyplot as plt
import numpy as np
import pytest
from scipy.interpolate import interp1d

from desc.basis import (
DoubleFourierSeries,
Expand Down Expand Up @@ -823,17 +822,15 @@ def flatten_coils(coilset):
def test_plot_b_mag():
"""Test plot of |B| on longer field lines for gyrokinetic simulations."""
psi = 0.5
rho = np.sqrt(psi)
npol = 2
nzgrid = 128
alpha = 0
# compute and fit iota profile
# compute iota
eq = get("W7-X")
data = eq.compute("iota")
fi = interp1d(data["rho"], data["iota"])
iota = eq.compute("iota", grid=LinearGrid(rho=rho, NFP=eq.NFP))["iota"][0]

# get flux tube coordinate system
rho = np.sqrt(psi)
iota = fi(rho)
zeta = np.linspace(
-np.pi * npol / np.abs(iota), np.pi * npol / np.abs(iota), 2 * nzgrid + 1
)
Expand Down

0 comments on commit 15180dd

Please sign in to comment.