Skip to content

Commit

Permalink
Merge branch 'fieldline_compute' into bounce
Browse files Browse the repository at this point in the history
  • Loading branch information
@unalmis
unalmis committed Jul 10, 2024
2 parents d1981b2 + 73c6eaa commit 718ceb3
Show file tree
Hide file tree
Showing 10 changed files with 261 additions and 147 deletions.
18 changes: 15 additions & 3 deletions desc/coils.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -226,6 +226,12 @@ def compute_magnetic_field(
current = self.current
else:
current = params.pop("current", self.current)
if source_grid is None and hasattr(self, "NFP"):
# NFP=1 to ensure we have points along whole grid
# multiply by NFP in case the coil has NFP>1
# to ensure whole coil gets counted for the
# biot savart integration
source_grid = LinearGrid(N=2 * self.N * self.NFP + 5, NFP=1, endpoint=False)

if not params or not transforms:
data = self.compute(
Expand Down Expand Up @@ -1153,9 +1159,15 @@ def from_symmetry(cls, coils, NFP=1, sym=False):
"""
if not isinstance(coils, CoilSet):
coils = CoilSet(coils)

[_check_type(coil, coils[0]) for coil in coils]
try:
coils = CoilSet(coils)
except (TypeError, ValueError):
# likely there are multiple coil types,
# so make a MixedCoilSet
coils = MixedCoilSet(coils)
if not isinstance(coils, MixedCoilSet):
# only need to check this for a CoilSet, not MixedCoilSet
[_check_type(coil, coils[0]) for coil in coils]

# check toroidal extent of coils to be repeated
maxphi = 2 * np.pi / NFP / (sym + 1)
Expand Down
109 changes: 57 additions & 52 deletions desc/compute/_curve.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -646,38 +646,39 @@ def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs):
units_long="meters",
description="Position vector along curve",
dim=3,
params=["X", "Y", "Z", "knots", "rotmat", "shift"],
transforms={"method": []},
params=["X", "Y", "Z", "rotmat", "shift"],
transforms={"knots": []},
profiles=[],
coordinates="s",
data=["s"],
parameterization="desc.geometry.curve.SplineXYZCurve",
basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.",
)
def _x_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
xq = data["s"]

Xq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["X"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Yq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Y"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Zq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Z"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Expand All @@ -699,38 +700,39 @@ def _x_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
units_long="meters",
description="Position vector along curve, first derivative",
dim=3,
params=["X", "Y", "Z", "knots", "rotmat", "shift"],
transforms={"method": []},
params=["X", "Y", "Z", "rotmat", "shift"],
transforms={"knots": []},
profiles=[],
coordinates="s",
data=["s"],
parameterization="desc.geometry.curve.SplineXYZCurve",
basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.",
)
def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
xq = data["s"]

dXq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["X"],
method=transforms["method"],
method=kwargs["method"],
derivative=1,
period=2 * jnp.pi,
)
dYq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Y"],
method=transforms["method"],
method=kwargs["method"],
derivative=1,
period=2 * jnp.pi,
)
dZq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Z"],
method=transforms["method"],
method=kwargs["method"],
derivative=1,
period=2 * jnp.pi,
)
Expand All @@ -742,25 +744,25 @@ def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
# calculate the xy coordinates to rotate to rpz
Xq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["X"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Yq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Y"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Zq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Z"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Expand All @@ -783,38 +785,39 @@ def _x_s_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
units_long="meters",
description="Position vector along curve, second derivative",
dim=3,
params=["X", "Y", "Z", "knots", "rotmat", "shift"],
transforms={"method": []},
params=["X", "Y", "Z", "rotmat", "shift"],
transforms={"knots": []},
profiles=[],
coordinates="s",
data=["s"],
parameterization="desc.geometry.curve.SplineXYZCurve",
basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.",
)
def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
xq = data["s"]

d2Xq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["X"],
method=transforms["method"],
method=kwargs["method"],
derivative=2,
period=2 * jnp.pi,
)
d2Yq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Y"],
method=transforms["method"],
method=kwargs["method"],
derivative=2,
period=2 * jnp.pi,
)
d2Zq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Z"],
method=transforms["method"],
method=kwargs["method"],
derivative=2,
period=2 * jnp.pi,
)
Expand All @@ -826,25 +829,25 @@ def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
# calculate the xy coordinates to rotate to rpz
Xq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["X"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Yq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Y"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Zq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Z"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Expand All @@ -866,38 +869,39 @@ def _x_ss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
units_long="meters",
description="Position vector along curve, third derivative",
dim=3,
params=["X", "Y", "Z", "knots", "rotmat", "shift"],
transforms={"method": []},
params=["X", "Y", "Z", "rotmat", "shift"],
transforms={"knots": []},
profiles=[],
coordinates="s",
data=["s"],
parameterization="desc.geometry.curve.SplineXYZCurve",
basis="{'rpz', 'xyz'}: Basis for returned vectors, Default 'rpz'",
method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.",
)
def _x_sss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
xq = data["s"]

d3Xq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["X"],
method=transforms["method"],
method=kwargs["method"],
derivative=3,
period=2 * jnp.pi,
)
d3Yq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Y"],
method=transforms["method"],
method=kwargs["method"],
derivative=3,
period=2 * jnp.pi,
)
d3Zq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Z"],
method=transforms["method"],
method=kwargs["method"],
derivative=3,
period=2 * jnp.pi,
)
Expand All @@ -909,25 +913,25 @@ def _x_sss_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
# calculate the xy coordinates to rotate to rpz
Xq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["X"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Yq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Y"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Zq = interp1d(
xq,
params["knots"],
transforms["knots"],
params["Z"],
method=transforms["method"],
method=kwargs["method"],
derivative=0,
period=2 * jnp.pi,
)
Expand Down Expand Up @@ -1080,14 +1084,15 @@ def _length(params, transforms, profiles, data, **kwargs):
description="Length of the curve",
dim=0,
params=[],
transforms={"method": []},
transforms={},
profiles=[],
coordinates="",
data=["ds", "x", "x_s"],
parameterization="desc.geometry.curve.SplineXYZCurve",
method="Interpolation type, Default 'cubic'. See SplineXYZCurve docs for options.",
)
def _length_SplineXYZCurve(params, transforms, profiles, data, **kwargs):
if transforms["method"] == "nearest": # cannot use derivative method as deriv=0
if kwargs["method"] == "nearest": # cannot use derivative method as deriv=0
coords = data["x"]
if kwargs.get("basis", "rpz").lower() == "rpz":
coords = rpz2xyz(coords)
Expand Down
42 changes: 41 additions & 1 deletion desc/geometry/curve.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -822,7 +822,7 @@ def __init__(
knots = knots[:-1] if closed_flag else knots

self._knots = knots
self.method = method
self._method = method

@optimizable_parameter
@property
Expand Down Expand Up @@ -925,6 +925,46 @@ def method(self, new):
+ f"instead got unknown method {new} "
)

def compute(
self,
names,
grid=None,
params=None,
transforms=None,
data=None,
**kwargs,
):
"""Compute the quantity given by name on grid.
Parameters
----------
names : str or array-like of str
Name(s) of the quantity(s) to compute.
grid : Grid or int, optional
Grid of coordinates to evaluate at. Defaults to a Linear grid.
If an integer, uses that many equally spaced points.
params : dict of ndarray
Parameters from the equilibrium. Defaults to attributes of self.
transforms : dict of Transform
Transforms for R, Z, lambda, etc. Default is to build from grid
data : dict of ndarray
Data computed so far, generally output from other compute functions
Returns
-------
data : dict of ndarray
Computed quantity and intermediate variables.
"""
return super().compute(
names=names,
grid=grid,
params=params,
transforms=transforms,
data=data,
method=self._method,
**kwargs,
)

@classmethod
def from_values(cls, coords, knots=None, method="cubic", name="", basis="xyz"):
"""Create SplineXYZCurve from coordinate values.
Expand Down
Loading

0 comments on commit 718ceb3

Please sign in to comment.