From 0be43041e31ca123decbfb0c5ede85ca086d19a4 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Tue, 27 Jun 2023 18:54:17 -0400 Subject: [PATCH 01/28] Break curve compute functions into standalone module --- desc/compute/_curve.py | 699 +++++++++++++++++++++++++++++++++++++++++ 1 file changed, 699 insertions(+) create mode 100644 desc/compute/_curve.py diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py new file mode 100644 index 0000000000..2271a95bb7 --- /dev/null +++ b/desc/compute/_curve.py @@ -0,0 +1,699 @@ +from desc.backend import jnp +from desc.compute.utils import _has_params, _has_transforms, has_dependencies +from desc.geometry.utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec + +curve_data_index = {} + + +def register_curve_compute_fun( + name, + label, + units, + units_long, + description, + dim, + params, + transforms, + data, + parameterization=None, + **kwargs, +): + """Decorator to wrap a function and add it to the list of things we can compute. + + Parameters + ---------- + name : str + Name of the quantity. This will be used as the key used to compute the + quantity in `compute` and its name in the data dictionary. + label : str + Title of the quantity in LaTeX format. + units : str + Units of the quantity in LaTeX format. + units_long : str + Full units without abbreviations. + description : str + Description of the quantity. + dim : int + Dimension of the quantity: 0-D (global qty), 1-D (local scalar qty), + or 3-D (local vector qty). + params : list of str + Parameters of equilibrium needed to compute quantity, eg "R_lmn", "Z_lmn" + transforms : dict + Dictionary of keys and derivative orders [rho, theta, zeta] for R, Z, etc. + data : list of str + Names of other items in the data index needed to compute qty. + parameterization: str or None + Name of curve types the method is valid for. eg 'FourierXYZCurve'. + None means it is valid for any curve type (calculation does not depend on + parameterization). + + Notes + ----- + Should only list *direct* dependencies. The full dependencies will be built + recursively at runtime using each quantity's direct dependencies. + """ + deps = { + "params": params, + "transforms": transforms, + "data": data, + "kwargs": list(kwargs.values()), + } + + def _decorator(func): + d = { + "label": label, + "units": units, + "units_long": units_long, + "description": description, + "fun": func, + "dim": dim, + "dependencies": deps, + } + if parameterization not in curve_data_index: + curve_data_index[parameterization] = {} + curve_data_index[parameterization][name] = d + return func + + return _decorator + + +def compute(names, params, transforms, 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. + params : dict of ndarray + Parameters from the equilibrium, such as R_lmn, Z_lmn, i_l, p_l, etc + 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. + + """ + if isinstance(names, str): + names = [names] + for name in names: + if name not in curve_data_index: + raise ValueError("Unrecognized value '{}'.".format(name)) + allowed_kwargs = {} # might have some in the future + bad_kwargs = kwargs.keys() - allowed_kwargs + if len(bad_kwargs) > 0: + raise ValueError(f"Unrecognized argument(s): {bad_kwargs}") + + for name in names: + assert _has_params(name, params), f"Don't have params to compute {name}" + assert _has_transforms( + name, transforms + ), f"Don't have transforms to compute {name}" + + if data is None: + data = {} + + data = _compute( + names, + params=params, + transforms=transforms, + data=data, + **kwargs, + ) + return data + + +def _compute(names, params, transforms, data=None, **kwargs): + """Same as above but without checking inputs for faster recursion.""" + for name in names: + if name in data: + # don't compute something that's already been computed + continue + if not has_dependencies(name, params, transforms, data): + # then compute the missing dependencies + data = _compute( + curve_data_index[name]["dependencies"]["data"], + params=params, + transforms=transforms, + data=data, + **kwargs, + ) + # now compute the quantity + data = curve_data_index[name]["fun"](params, transforms, data, **kwargs) + return data + + +@register_curve_compute_fun( + name="s", + label="s", + units="~", + units_long="None", + description="Curve parameter", + dim=3, + params=[], + transforms={"grid": []}, + data=[], +) +def _s(params, data, transforms, basis="rpz"): + data["s"] = transforms["grid"].nodes[:, 2] + return data + + +def _rotation_matrix_from_normal(normal): + nx, ny, nz = normal + nxny = jnp.sqrt(nx**2 + ny**2) + R = jnp.array( + [ + [ny / nxny, -nx / nxny, 0], + [nx * nx / nxny, ny * nz / nxny, -nxny], + [nx, ny, nz], + ] + ).T + R = jnp.where(nxny == 0, jnp.eye(3), R) + return R + + +@register_curve_compute_fun( + name="r", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve", + dim=3, + params=["r_n", "center", "normal"], + transforms={ + "r": [[0, 0, 0]], + "rotmat": [], + "shift": [], + }, + data=["s"], + parameterization="FourierPlanarCurve", +) +def _r_FourierPlanarCurve(params, data, transforms, basis="rpz"): + # create planar curve at z==0 + r = transforms["r"].transform(params["r_n"], dz=0) + Z = jnp.zeros_like(r) + X = r * jnp.cos(data["s"]) + Y = r * jnp.sin(data["s"]) + coords = jnp.array([X, Y, Z]).T + # rotate into place + R = _rotation_matrix_from_normal(params["normal"]) + coords = jnp.matmul(coords, R.T) + params["center"] + coords = jnp.matmul(coords, transforms["rotmat"].T) + transforms["shift"] + if basis.lower() == "rpz": + xyzcoords = jnp.array([X, Y, Z]).T + xyzcoords = jnp.matmul(xyzcoords, R.T) + params["center"] + xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] + x, y, z = xyzcoords.T + coords = xyz2rpz_vec(coords, x=x, y=y) + data["r"] = coords + return data + + +@register_curve_compute_fun( + name="r_s", + label="\\partial_{s} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, first derivative", + dim=3, + params=["r_n", "center", "normal"], + transforms={ + "r": [[0, 0, 0], [0, 0, 1]], + "rotmat": [], + "shift": [], + }, + data=["s"], + parameterization="FourierPlanarCurve", +) +def _r_s_FourierPlanarCurve(params, data, transforms, basis="rpz"): + r = transforms["r"].transform(params["r_n"], dz=0) + dr = transforms["r"].transform(params["r_n"], dz=1) + dX = dr * jnp.cos(data["s"]) - r * jnp.sin(data["s"]) + dY = dr * jnp.sin(data["s"]) + r * jnp.cos(data["s"]) + dZ = jnp.zeros_like(dX) + coords = jnp.array([dX, dY, dZ]).T + A = _rotation_matrix_from_normal(params["normal"]) + coords = jnp.matmul(coords, A.T) + coords = jnp.matmul(coords, transforms["rotmat"].T) + if basis.lower() == "rpz": + X = r * jnp.cos(data["s"]) + Y = r * jnp.sin(data["s"]) + Z = jnp.zeros_like(X) + xyzcoords = jnp.array([X, Y, Z]).T + xyzcoords = jnp.matmul(xyzcoords, A.T) + params["center"] + xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] + x, y, z = xyzcoords.T + coords = xyz2rpz_vec(coords, x=x, y=y) + data["r_s"] = coords + return data + + +@register_curve_compute_fun( + name="r_ss", + label="\\partial_{ss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, second derivative", + dim=3, + params=["r_n", "center", "normal"], + transforms={ + "r": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], + "rotmat": [], + "shift": [], + }, + data=["s"], + parameterization="FourierPlanarCurve", +) +def _r_ss_FourierPlanarCurve(params, data, transforms, basis="rpz"): + r = transforms["r"].transform(params["r_n"], dz=0) + dr = transforms["r"].transform(params["r_n"], dz=1) + d2r = transforms["r"].transform(params["r_n"], dz=2) + d2X = ( + d2r * jnp.cos(data["s"]) - 2 * dr * jnp.sin(data["s"]) - r * jnp.cos(data["s"]) + ) + d2Y = ( + d2r * jnp.sin(data["s"]) + 2 * dr * jnp.cos(data["s"]) - r * jnp.sin(data["s"]) + ) + d2Z = jnp.zeros_like(d2X) + coords = jnp.array([d2X, d2Y, d2Z]).T + A = _rotation_matrix_from_normal(params["normal"]) + coords = jnp.matmul(coords, A.T) + coords = jnp.matmul(coords, transforms["rotmat"].T) + if basis.lower() == "rpz": + X = r * jnp.cos(data["s"]) + Y = r * jnp.sin(data["s"]) + Z = jnp.zeros_like(X) + xyzcoords = jnp.array([X, Y, Z]).T + xyzcoords = jnp.matmul(xyzcoords, A.T) + params["center"] + xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] + x, y, z = xyzcoords.T + coords = xyz2rpz_vec(coords, x=x, y=y) + data["r_ss"] = coords + return data + + +@register_curve_compute_fun( + name="r_sss", + label="\\partial_{sss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, third derivative", + dim=3, + params=["r_n", "center", "normal"], + transforms={ + "r": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]], + "rotmat": [], + "shift": [], + }, + data=["s"], + parameterization="FourierPlanarCurve", +) +def _r_sss_FourierPlanarCurve(params, data, transforms, basis="rpz"): + r = transforms["r"].transform(params["r_n"], dz=0) + dr = transforms["r"].transform(params["r_n"], dz=1) + d2r = transforms["r"].transform(params["r_n"], dz=2) + d3r = transforms["r"].transform(params["r_n"], dz=3) + d3X = ( + d3r * jnp.cos(data["s"]) + - 3 * d2r * jnp.sin(data["s"]) + - 3 * dr * jnp.cos(data["s"]) + + r * jnp.sin(data["s"]) + ) + d3Y = ( + d3r * jnp.sin(data["s"]) + + 3 * d2r * jnp.cos(data["s"]) + - 3 * dr * jnp.sin(data["s"]) + - r * jnp.cos(data["s"]) + ) + d3Z = jnp.zeros_like(d3X) + coords = jnp.array([d3X, d3Y, d3Z]).T + A = _rotation_matrix_from_normal(params["normal"]) + coords = jnp.matmul(coords, A.T) + coords = jnp.matmul(coords, transforms["rotmat"].T) + if basis.lower() == "rpz": + X = r * jnp.cos(data["s"]) + Y = r * jnp.sin(data["s"]) + Z = jnp.zeros_like(X) + xyzcoords = jnp.array([X, Y, Z]).T + xyzcoords = jnp.matmul(xyzcoords, A.T) + params["center"] + xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] + x, y, z = xyzcoords.T + coords = xyz2rpz_vec(coords, x=x, y=y) + data["r_sss"] = coords + return data + + +@register_curve_compute_fun( + name="r", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve", + dim=3, + params=["R_n", "Z_n"], + transforms={ + "R": [[0, 0, 0]], + "Z": [[0, 0, 0]], + "grid": [], + "rotmat": [], + "shift": [], + }, + data=[], + parameterization="FourierRZCurve", +) +def _r_FourierRZCurve(params, data, transforms, basis="rpz"): + R = transforms["R"].transform(params["R_n"], dz=0) + Z = transforms["Z"].transform(params["Z_n"], dz=0) + phi = transforms["grid"].nodes[:, 2] + coords = jnp.stack([R, phi, Z], axis=1) + # convert to xyz for displacement and rotation + coords = rpz2xyz(coords) + coords = coords @ transforms["rotmat"].T + transforms["shift"][jnp.newaxis, :] + if basis.lower() == "rpz": + coords = xyz2rpz(coords) + data["r"] = coords + return data + + +@register_curve_compute_fun( + name="r_s", + label="\\partial_{s} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, first derivative", + dim=3, + params=["R_n", "Z_n"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1]], + "Z": [[0, 0, 1]], + "grid": [], + "rotmat": [], + }, + data=[], + parameterization="FourierRZCurve", +) +def _r_s_FourierRZCurve(params, data, transforms, basis="rpz"): + + R0 = transforms["R"].transform(params["R_n"], dz=0) + dR = transforms["R"].transform(params["R_n"], dz=1) + dZ = transforms["Z"].transform(params["Z_n"], dz=1) + dphi = R0 + coords = jnp.stack([dR, dphi, dZ], axis=1) + # convert to xyz for displacement and rotation + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + coords = coords @ transforms["rotmat"].T + if basis.lower() == "rpz": + coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_s"] = coords + return data + + +@register_curve_compute_fun( + name="r_ss", + label="\\partial_{ss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, second derivative", + dim=3, + params=["R_n", "Z_n"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], + "Z": [[0, 0, 2]], + "grid": [], + "rotmat": [], + }, + data=[], + parameterization="FourierRZCurve", +) +def _r_ss_FourierRZCurve(params, data, transforms, basis="rpz"): + R0 = transforms["R"].transform(params["R_n"], dz=0) + dR = transforms["R"].transform(params["R_n"], dz=1) + d2R = transforms["R"].transform(params["R_n"], dz=2) + d2Z = transforms["Z"].transform(params["Z_n"], dz=2) + R = d2R - R0 + Z = d2Z + # 2nd derivative wrt phi = 0 + phi = 2 * dR + coords = jnp.stack([R, phi, Z], axis=1) + # convert to xyz for displacement and rotation + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + coords = coords @ transforms["rotmat"].T + if basis.lower() == "rpz": + coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_ss"] = coords + return data + + +@register_curve_compute_fun( + name="r_sss", + label="\\partial_{sss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, third derivative", + dim=3, + params=["R_n", "Z_n"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]], + "Z": [[0, 0, 3]], + "grid": [], + "rotmat": [], + }, + data=[], + parameterization="FourierRZCurve", +) +def _r_sss_FourierRZCurve(params, data, transforms, basis="rpz"): + R0 = transforms["R"].transform(params["R_n"], dz=0) + dR = transforms["R"].transform(params["R_n"], dz=1) + d2R = transforms["R"].transform(params["R_n"], dz=2) + d3R = transforms["R"].transform(params["R_n"], dz=3) + d3Z = transforms["Z"].transform(params["Z_n"], dz=3) + R = d3R - 3 * dR + Z = d3Z + phi = 3 * d2R - R0 + coords = jnp.stack([R, phi, Z], axis=1) + # convert to xyz for displacement and rotation + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + coords = coords @ transforms["rotmat"].T + if basis.lower() == "rpz": + coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_sss"] = coords + return data + + +@register_curve_compute_fun( + name="r", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve", + dim=3, + params=["X_n", "Y_n", "Z_n"], + transforms={"r": [[0, 0, 0]], "rotmat": [], "shift": []}, + data=[], + parameterization="FourierXYZCurve", +) +def _r_FourierXYZCurve(params, data, transforms, basis="rpz"): + X = transforms["r"].transform(params["X_n"], dz=0) + Y = transforms["r"].transform(params["Y_n"], dz=0) + Z = transforms["r"].transform(params["Z_n"], dz=0) + coords = jnp.stack([X, Y, Z], axis=1) + coords = coords @ transforms["rotmat"].T + transforms["shift"][jnp.newaxis, :] + if basis.lower() == "rpz": + coords = xyz2rpz(coords) + data["r"] = coords + return data + + +@register_curve_compute_fun( + name="r_s", + label="\\partial_{s} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, first derivative", + dim=3, + params=["X_n", "Y_n", "Z_n"], + transforms={"r": [[0, 0, 1]], "rotmat": [], "shift": []}, + data=[], + parameterization="FourierXYZCurve", +) +def _r_s_FourierXYZCurve(params, data, transforms, basis="rpz"): + X = transforms["r"].transform(params["X_n"], dz=1) + Y = transforms["r"].transform(params["Y_n"], dz=1) + Z = transforms["r"].transform(params["Z_n"], dz=1) + coords = jnp.stack([X, Y, Z], axis=1) + coords = coords @ transforms["rotmat"].T + if basis.lower() == "rpz": + coords = xyz2rpz_vec( + coords, + x=coords[:, 1] + transforms["shift"][0], + y=coords[:, 1] + transforms["shift"][1], + ) + data["r_s"] = coords + return data + + +@register_curve_compute_fun( + name="r_ss", + label="\\partial_{ss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, second derivative", + dim=3, + params=["X_n", "Y_n", "Z_n"], + transforms={"r": [[0, 0, 2]], "rotmat": [], "shift": []}, + data=[], + parameterization="FourierXYZCurve", +) +def _r_ss_FourierXYZCurve(params, data, transforms, basis="rpz"): + X = transforms["r"].transform(params["X_n"], dz=2) + Y = transforms["r"].transform(params["Y_n"], dz=2) + Z = transforms["r"].transform(params["Z_n"], dz=2) + coords = jnp.stack([X, Y, Z], axis=1) + coords = coords @ transforms["rotmat"].T + if basis.lower() == "rpz": + coords = xyz2rpz_vec( + coords, + x=coords[:, 1] + transforms["shift"][0], + y=coords[:, 1] + transforms["shift"][1], + ) + data["r_ss"] = coords + return data + + +@register_curve_compute_fun( + name="r_sss", + label="\\partial_{sss} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along curve, third derivative", + dim=3, + params=["X_n", "Y_n", "Z_n"], + transforms={"r": [[0, 0, 3]], "rotmat": [], "shift": []}, + data=[], + parameterization="FourierXYZCurve", +) +def _r_sss_FourierXYZCurve(params, data, transforms, basis="rpz"): + X = transforms["r"].transform(params["X_n"], dz=3) + Y = transforms["r"].transform(params["Y_n"], dz=3) + Z = transforms["r"].transform(params["Z_n"], dz=3) + coords = jnp.stack([X, Y, Z], axis=1) + coords = coords @ transforms["rotmat"].T + if basis.lower() == "rpz": + coords = xyz2rpz_vec( + coords, + x=coords[:, 1] + transforms["shift"][0], + y=coords[:, 1] + transforms["shift"][1], + ) + data["r_sss"] = coords + return data + + +@register_curve_compute_fun( + name="frenet_tangent", + label="T", + units="~", + units_long="None", + description="Tangent unit vector to curve in Frenet-Serret frame", + dim=3, + params=[], + transforms={}, + data=["r_s"], +) +def _frenet_tangent(params, data, transforms, basis="rpz"): + data["frenet_tangent"] = data["r_s"] / jnp.linalg.norm(data["r_s"])[:, None] + return data + + +@register_curve_compute_fun( + name="frenet_normal", + label="N", + units="~", + units_long="None", + description="Normal unit vector to curve in Frenet-Serret frame", + dim=3, + params=[], + transforms={}, + data=["r_ss"], +) +def _frenet_normal(params, data, transforms, basis="rpz"): + data["frenet_normal"] = data["r_ss"] / jnp.linalg.norm(data["r_ss"])[:, None] + return data + + +@register_curve_compute_fun( + name="frenet_binormal", + label="B", + units="~", + units_long="None", + description="Binormal unit vector to curve in Frenet-Serret frame", + dim=3, + params=[], + transforms={"rotmat": []}, + data=["T", "N"], +) +def _frenet_binormal(params, data, transforms, basis="rpz"): + data["frenet_binormal"] = jnp.cross(data["T"], data["N"], axis=1) * jnp.linalg.det( + transforms["rotmat"] + ) + return data + + +@register_curve_compute_fun( + name="curvature", + label="\\kappa", + units="m^{-1}", + units_long="Inverse meters", + description="Scalar curvature of the curve", + dim=1, + params=[], + transforms={}, + data=["r_s", "r_ss"], +) +def _curvature(params, data, transforms, basis="rpz"): + dxn = jnp.linalg.norm(data["r_s"], axis=1)[:, jnp.newaxis] + data["curvature"] = jnp.linalg.norm( + jnp.cross(data["r_s"], data["r_ss"], axis=1) / dxn**3, axis=1 + ) + return data + + +@register_curve_compute_fun( + name="torsion", + label="\\tau", + units="m^{-1}", + units_long="Inverse meters", + description="Scalar torsion of the curve", + dim=1, + params=[], + transforms={}, + data=["r_s", "r_ss", "r_sss"], +) +def _torsion(params, data, transforms, basis="rpz"): + dxd2x = jnp.cross(data["r_s"], data["r_ss"], axis=1) + data["torsion"] = ( + jnp.sum(dxd2x * data["r_sss"], axis=1) + / jnp.linalg.norm(dxd2x, axis=1)[:, jnp.newaxis] ** 2 + ) + return data + + +@register_curve_compute_fun( + name="length", + label="L", + units="m", + units_long="meters", + description="Length of the curve", + dim=0, + params=[], + transforms={}, + data=["s", "r_s"], +) +def _length(params, data, transforms, basis="rpz"): + T = jnp.linalg.norm(data["r_s"], axis=1) + data["length"] = jnp.trapz(T, data["s"]) + return data From 314742ffccc041652a1410741b36a618003e60ab Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Wed, 5 Jul 2023 15:17:13 -0400 Subject: [PATCH 02/28] Update formula for cross section area to work for surfaces --- desc/compute/_geometry.py | 4 ++-- 1 file changed, 2 insertions(+), 2 deletions(-) diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index 698a850756..817c8c28de 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -97,13 +97,13 @@ def _V_rr_of_r(params, transforms, profiles, data, **kwargs): transforms={"grid": []}, profiles=[], coordinates="", - data=["sqrt(g)", "R"], + data=["|e_rho x e_theta|"], ) def _A(params, transforms, profiles, data, **kwargs): data["A"] = jnp.mean( surface_integrals( transforms["grid"], - jnp.abs(data["sqrt(g)"] / data["R"]), + data["|e_rho x e_theta|"], # equivalent to sqrt(g)/R for phi=zeta surface_label="zeta", expand_out=False, ) From c3711cd4e2ef2b5584dc9ebb8bf525c5e2dc8874 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 6 Jul 2023 00:11:29 -0400 Subject: [PATCH 03/28] Update curve compute funs to regular data_index --- desc/compute/_curve.py | 322 ++++++++++++++--------------------------- 1 file changed, 112 insertions(+), 210 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 2271a95bb7..11abb85cac 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -1,164 +1,24 @@ from desc.backend import jnp -from desc.compute.utils import _has_params, _has_transforms, has_dependencies from desc.geometry.utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec -curve_data_index = {} - - -def register_curve_compute_fun( - name, - label, - units, - units_long, - description, - dim, - params, - transforms, - data, - parameterization=None, - **kwargs, -): - """Decorator to wrap a function and add it to the list of things we can compute. - - Parameters - ---------- - name : str - Name of the quantity. This will be used as the key used to compute the - quantity in `compute` and its name in the data dictionary. - label : str - Title of the quantity in LaTeX format. - units : str - Units of the quantity in LaTeX format. - units_long : str - Full units without abbreviations. - description : str - Description of the quantity. - dim : int - Dimension of the quantity: 0-D (global qty), 1-D (local scalar qty), - or 3-D (local vector qty). - params : list of str - Parameters of equilibrium needed to compute quantity, eg "R_lmn", "Z_lmn" - transforms : dict - Dictionary of keys and derivative orders [rho, theta, zeta] for R, Z, etc. - data : list of str - Names of other items in the data index needed to compute qty. - parameterization: str or None - Name of curve types the method is valid for. eg 'FourierXYZCurve'. - None means it is valid for any curve type (calculation does not depend on - parameterization). - - Notes - ----- - Should only list *direct* dependencies. The full dependencies will be built - recursively at runtime using each quantity's direct dependencies. - """ - deps = { - "params": params, - "transforms": transforms, - "data": data, - "kwargs": list(kwargs.values()), - } - - def _decorator(func): - d = { - "label": label, - "units": units, - "units_long": units_long, - "description": description, - "fun": func, - "dim": dim, - "dependencies": deps, - } - if parameterization not in curve_data_index: - curve_data_index[parameterization] = {} - curve_data_index[parameterization][name] = d - return func - - return _decorator - - -def compute(names, params, transforms, 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. - params : dict of ndarray - Parameters from the equilibrium, such as R_lmn, Z_lmn, i_l, p_l, etc - 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. - - """ - if isinstance(names, str): - names = [names] - for name in names: - if name not in curve_data_index: - raise ValueError("Unrecognized value '{}'.".format(name)) - allowed_kwargs = {} # might have some in the future - bad_kwargs = kwargs.keys() - allowed_kwargs - if len(bad_kwargs) > 0: - raise ValueError(f"Unrecognized argument(s): {bad_kwargs}") - - for name in names: - assert _has_params(name, params), f"Don't have params to compute {name}" - assert _has_transforms( - name, transforms - ), f"Don't have transforms to compute {name}" - - if data is None: - data = {} - - data = _compute( - names, - params=params, - transforms=transforms, - data=data, - **kwargs, - ) - return data - - -def _compute(names, params, transforms, data=None, **kwargs): - """Same as above but without checking inputs for faster recursion.""" - for name in names: - if name in data: - # don't compute something that's already been computed - continue - if not has_dependencies(name, params, transforms, data): - # then compute the missing dependencies - data = _compute( - curve_data_index[name]["dependencies"]["data"], - params=params, - transforms=transforms, - data=data, - **kwargs, - ) - # now compute the quantity - data = curve_data_index[name]["fun"](params, transforms, data, **kwargs) - return data +from .data_index import register_compute_fun -@register_curve_compute_fun( +@register_compute_fun( name="s", label="s", units="~", units_long="None", - description="Curve parameter", + description="Curve parameter, on [0, 2pi)", dim=3, params=[], transforms={"grid": []}, + profiles=[], + coordinates="s", data=[], + parameterization="desc.geometry.Curve", ) -def _s(params, data, transforms, basis="rpz"): +def _s(params, transforms, profiles, data, **kwargs): data["s"] = transforms["grid"].nodes[:, 2] return data @@ -177,7 +37,7 @@ def _rotation_matrix_from_normal(normal): return R -@register_curve_compute_fun( +@register_compute_fun( name="r", label="\\mathbf{r}", units="m", @@ -190,10 +50,12 @@ def _rotation_matrix_from_normal(normal): "rotmat": [], "shift": [], }, + profiles=[], + coordinates="s", data=["s"], - parameterization="FourierPlanarCurve", + parameterization="desc.geometry.FourierPlanarCurve", ) -def _r_FourierPlanarCurve(params, data, transforms, basis="rpz"): +def _r_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): # create planar curve at z==0 r = transforms["r"].transform(params["r_n"], dz=0) Z = jnp.zeros_like(r) @@ -204,7 +66,7 @@ def _r_FourierPlanarCurve(params, data, transforms, basis="rpz"): R = _rotation_matrix_from_normal(params["normal"]) coords = jnp.matmul(coords, R.T) + params["center"] coords = jnp.matmul(coords, transforms["rotmat"].T) + transforms["shift"] - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": xyzcoords = jnp.array([X, Y, Z]).T xyzcoords = jnp.matmul(xyzcoords, R.T) + params["center"] xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] @@ -214,7 +76,7 @@ def _r_FourierPlanarCurve(params, data, transforms, basis="rpz"): return data -@register_curve_compute_fun( +@register_compute_fun( name="r_s", label="\\partial_{s} \\mathbf{r}", units="m", @@ -227,10 +89,12 @@ def _r_FourierPlanarCurve(params, data, transforms, basis="rpz"): "rotmat": [], "shift": [], }, + profiles=[], + coordinates="s", data=["s"], - parameterization="FourierPlanarCurve", + parameterization="desc.geometry.FourierPlanarCurve", ) -def _r_s_FourierPlanarCurve(params, data, transforms, basis="rpz"): +def _r_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): r = transforms["r"].transform(params["r_n"], dz=0) dr = transforms["r"].transform(params["r_n"], dz=1) dX = dr * jnp.cos(data["s"]) - r * jnp.sin(data["s"]) @@ -240,7 +104,7 @@ def _r_s_FourierPlanarCurve(params, data, transforms, basis="rpz"): A = _rotation_matrix_from_normal(params["normal"]) coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, transforms["rotmat"].T) - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": X = r * jnp.cos(data["s"]) Y = r * jnp.sin(data["s"]) Z = jnp.zeros_like(X) @@ -253,7 +117,7 @@ def _r_s_FourierPlanarCurve(params, data, transforms, basis="rpz"): return data -@register_curve_compute_fun( +@register_compute_fun( name="r_ss", label="\\partial_{ss} \\mathbf{r}", units="m", @@ -266,10 +130,12 @@ def _r_s_FourierPlanarCurve(params, data, transforms, basis="rpz"): "rotmat": [], "shift": [], }, + profiles=[], + coordinates="s", data=["s"], - parameterization="FourierPlanarCurve", + parameterization="desc.geometry.FourierPlanarCurve", ) -def _r_ss_FourierPlanarCurve(params, data, transforms, basis="rpz"): +def _r_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): r = transforms["r"].transform(params["r_n"], dz=0) dr = transforms["r"].transform(params["r_n"], dz=1) d2r = transforms["r"].transform(params["r_n"], dz=2) @@ -284,7 +150,7 @@ def _r_ss_FourierPlanarCurve(params, data, transforms, basis="rpz"): A = _rotation_matrix_from_normal(params["normal"]) coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, transforms["rotmat"].T) - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": X = r * jnp.cos(data["s"]) Y = r * jnp.sin(data["s"]) Z = jnp.zeros_like(X) @@ -297,7 +163,7 @@ def _r_ss_FourierPlanarCurve(params, data, transforms, basis="rpz"): return data -@register_curve_compute_fun( +@register_compute_fun( name="r_sss", label="\\partial_{sss} \\mathbf{r}", units="m", @@ -310,10 +176,12 @@ def _r_ss_FourierPlanarCurve(params, data, transforms, basis="rpz"): "rotmat": [], "shift": [], }, + profiles=[], + coordinates="s", data=["s"], - parameterization="FourierPlanarCurve", + parameterization="desc.geometry.FourierPlanarCurve", ) -def _r_sss_FourierPlanarCurve(params, data, transforms, basis="rpz"): +def _r_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): r = transforms["r"].transform(params["r_n"], dz=0) dr = transforms["r"].transform(params["r_n"], dz=1) d2r = transforms["r"].transform(params["r_n"], dz=2) @@ -335,7 +203,7 @@ def _r_sss_FourierPlanarCurve(params, data, transforms, basis="rpz"): A = _rotation_matrix_from_normal(params["normal"]) coords = jnp.matmul(coords, A.T) coords = jnp.matmul(coords, transforms["rotmat"].T) - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": X = r * jnp.cos(data["s"]) Y = r * jnp.sin(data["s"]) Z = jnp.zeros_like(X) @@ -348,7 +216,7 @@ def _r_sss_FourierPlanarCurve(params, data, transforms, basis="rpz"): return data -@register_curve_compute_fun( +@register_compute_fun( name="r", label="\\mathbf{r}", units="m", @@ -363,10 +231,12 @@ def _r_sss_FourierPlanarCurve(params, data, transforms, basis="rpz"): "rotmat": [], "shift": [], }, + profiles=[], + coordinates="s", data=[], - parameterization="FourierRZCurve", + parameterization="desc.geometry.FourierRZCurve", ) -def _r_FourierRZCurve(params, data, transforms, basis="rpz"): +def _r_FourierRZCurve(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_n"], dz=0) Z = transforms["Z"].transform(params["Z_n"], dz=0) phi = transforms["grid"].nodes[:, 2] @@ -374,13 +244,13 @@ def _r_FourierRZCurve(params, data, transforms, basis="rpz"): # convert to xyz for displacement and rotation coords = rpz2xyz(coords) coords = coords @ transforms["rotmat"].T + transforms["shift"][jnp.newaxis, :] - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz(coords) data["r"] = coords return data -@register_curve_compute_fun( +@register_compute_fun( name="r_s", label="\\partial_{s} \\mathbf{r}", units="m", @@ -394,10 +264,12 @@ def _r_FourierRZCurve(params, data, transforms, basis="rpz"): "grid": [], "rotmat": [], }, + profiles=[], + coordinates="s", data=[], - parameterization="FourierRZCurve", + parameterization="desc.geometry.FourierRZCurve", ) -def _r_s_FourierRZCurve(params, data, transforms, basis="rpz"): +def _r_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_n"], dz=0) dR = transforms["R"].transform(params["R_n"], dz=1) @@ -407,13 +279,13 @@ def _r_s_FourierRZCurve(params, data, transforms, basis="rpz"): # convert to xyz for displacement and rotation coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) coords = coords @ transforms["rotmat"].T - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) data["r_s"] = coords return data -@register_curve_compute_fun( +@register_compute_fun( name="r_ss", label="\\partial_{ss} \\mathbf{r}", units="m", @@ -427,10 +299,12 @@ def _r_s_FourierRZCurve(params, data, transforms, basis="rpz"): "grid": [], "rotmat": [], }, + profiles=[], + coordinates="s", data=[], - parameterization="FourierRZCurve", + parameterization="desc.geometry.FourierRZCurve", ) -def _r_ss_FourierRZCurve(params, data, transforms, basis="rpz"): +def _r_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_n"], dz=0) dR = transforms["R"].transform(params["R_n"], dz=1) d2R = transforms["R"].transform(params["R_n"], dz=2) @@ -443,13 +317,13 @@ def _r_ss_FourierRZCurve(params, data, transforms, basis="rpz"): # convert to xyz for displacement and rotation coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) coords = coords @ transforms["rotmat"].T - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) data["r_ss"] = coords return data -@register_curve_compute_fun( +@register_compute_fun( name="r_sss", label="\\partial_{sss} \\mathbf{r}", units="m", @@ -463,10 +337,12 @@ def _r_ss_FourierRZCurve(params, data, transforms, basis="rpz"): "grid": [], "rotmat": [], }, + profiles=[], + coordinates="s", data=[], - parameterization="FourierRZCurve", + parameterization="desc.geometry.FourierRZCurve", ) -def _r_sss_FourierRZCurve(params, data, transforms, basis="rpz"): +def _r_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_n"], dz=0) dR = transforms["R"].transform(params["R_n"], dz=1) d2R = transforms["R"].transform(params["R_n"], dz=2) @@ -479,13 +355,13 @@ def _r_sss_FourierRZCurve(params, data, transforms, basis="rpz"): # convert to xyz for displacement and rotation coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) coords = coords @ transforms["rotmat"].T - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) data["r_sss"] = coords return data -@register_curve_compute_fun( +@register_compute_fun( name="r", label="\\mathbf{r}", units="m", @@ -494,22 +370,24 @@ def _r_sss_FourierRZCurve(params, data, transforms, basis="rpz"): dim=3, params=["X_n", "Y_n", "Z_n"], transforms={"r": [[0, 0, 0]], "rotmat": [], "shift": []}, + profiles=[], + coordinates="s", data=[], - parameterization="FourierXYZCurve", + parameterization="desc.geometry.FourierXYZCurve", ) -def _r_FourierXYZCurve(params, data, transforms, basis="rpz"): +def _r_FourierXYZCurve(params, transforms, profiles, data, **kwargs): X = transforms["r"].transform(params["X_n"], dz=0) Y = transforms["r"].transform(params["Y_n"], dz=0) Z = transforms["r"].transform(params["Z_n"], dz=0) coords = jnp.stack([X, Y, Z], axis=1) coords = coords @ transforms["rotmat"].T + transforms["shift"][jnp.newaxis, :] - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz(coords) data["r"] = coords return data -@register_curve_compute_fun( +@register_compute_fun( name="r_s", label="\\partial_{s} \\mathbf{r}", units="m", @@ -518,16 +396,18 @@ def _r_FourierXYZCurve(params, data, transforms, basis="rpz"): dim=3, params=["X_n", "Y_n", "Z_n"], transforms={"r": [[0, 0, 1]], "rotmat": [], "shift": []}, + profiles=[], + coordinates="s", data=[], - parameterization="FourierXYZCurve", + parameterization="desc.geometry.FourierXYZCurve", ) -def _r_s_FourierXYZCurve(params, data, transforms, basis="rpz"): +def _r_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): X = transforms["r"].transform(params["X_n"], dz=1) Y = transforms["r"].transform(params["Y_n"], dz=1) Z = transforms["r"].transform(params["Z_n"], dz=1) coords = jnp.stack([X, Y, Z], axis=1) coords = coords @ transforms["rotmat"].T - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec( coords, x=coords[:, 1] + transforms["shift"][0], @@ -537,7 +417,7 @@ def _r_s_FourierXYZCurve(params, data, transforms, basis="rpz"): return data -@register_curve_compute_fun( +@register_compute_fun( name="r_ss", label="\\partial_{ss} \\mathbf{r}", units="m", @@ -546,16 +426,18 @@ def _r_s_FourierXYZCurve(params, data, transforms, basis="rpz"): dim=3, params=["X_n", "Y_n", "Z_n"], transforms={"r": [[0, 0, 2]], "rotmat": [], "shift": []}, + profiles=[], + coordinates="s", data=[], - parameterization="FourierXYZCurve", + parameterization="desc.geometry.FourierXYZCurve", ) -def _r_ss_FourierXYZCurve(params, data, transforms, basis="rpz"): +def _r_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): X = transforms["r"].transform(params["X_n"], dz=2) Y = transforms["r"].transform(params["Y_n"], dz=2) Z = transforms["r"].transform(params["Z_n"], dz=2) coords = jnp.stack([X, Y, Z], axis=1) coords = coords @ transforms["rotmat"].T - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec( coords, x=coords[:, 1] + transforms["shift"][0], @@ -565,7 +447,7 @@ def _r_ss_FourierXYZCurve(params, data, transforms, basis="rpz"): return data -@register_curve_compute_fun( +@register_compute_fun( name="r_sss", label="\\partial_{sss} \\mathbf{r}", units="m", @@ -574,16 +456,18 @@ def _r_ss_FourierXYZCurve(params, data, transforms, basis="rpz"): dim=3, params=["X_n", "Y_n", "Z_n"], transforms={"r": [[0, 0, 3]], "rotmat": [], "shift": []}, + profiles=[], + coordinates="s", data=[], - parameterization="FourierXYZCurve", + parameterization="desc.geometry.FourierXYZCurve", ) -def _r_sss_FourierXYZCurve(params, data, transforms, basis="rpz"): +def _r_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): X = transforms["r"].transform(params["X_n"], dz=3) Y = transforms["r"].transform(params["Y_n"], dz=3) Z = transforms["r"].transform(params["Z_n"], dz=3) coords = jnp.stack([X, Y, Z], axis=1) coords = coords @ transforms["rotmat"].T - if basis.lower() == "rpz": + if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec( coords, x=coords[:, 1] + transforms["shift"][0], @@ -593,57 +477,66 @@ def _r_sss_FourierXYZCurve(params, data, transforms, basis="rpz"): return data -@register_curve_compute_fun( +@register_compute_fun( name="frenet_tangent", - label="T", + label="\\mathbf{T}_{\\mathrm{Frenet-Serret}}", units="~", units_long="None", description="Tangent unit vector to curve in Frenet-Serret frame", dim=3, params=[], transforms={}, + profiles=[], + coordinates="s", data=["r_s"], + parameterization="desc.geometry.Curve", ) -def _frenet_tangent(params, data, transforms, basis="rpz"): +def _frenet_tangent(params, transforms, profiles, data, **kwargs): data["frenet_tangent"] = data["r_s"] / jnp.linalg.norm(data["r_s"])[:, None] return data -@register_curve_compute_fun( +@register_compute_fun( name="frenet_normal", - label="N", + label="\\mathbf{N}_{\\mathrm{Frenet-Serret}}", units="~", units_long="None", description="Normal unit vector to curve in Frenet-Serret frame", dim=3, params=[], transforms={}, + profiles=[], + coordinates="s", data=["r_ss"], + parameterization="desc.geometry.Curve", ) -def _frenet_normal(params, data, transforms, basis="rpz"): +def _frenet_normal(params, transforms, profiles, data, **kwargs): data["frenet_normal"] = data["r_ss"] / jnp.linalg.norm(data["r_ss"])[:, None] return data -@register_curve_compute_fun( +@register_compute_fun( name="frenet_binormal", - label="B", + label="\\mathrm{B}_{\\mathrm{Frenet-Serret}}", units="~", units_long="None", description="Binormal unit vector to curve in Frenet-Serret frame", dim=3, params=[], transforms={"rotmat": []}, + profiles=[], + coordinates="s", data=["T", "N"], + parameterization="desc.geometry.Curve", ) -def _frenet_binormal(params, data, transforms, basis="rpz"): +def _frenet_binormal(params, transforms, profiles, data, **kwargs): data["frenet_binormal"] = jnp.cross(data["T"], data["N"], axis=1) * jnp.linalg.det( transforms["rotmat"] ) return data -@register_curve_compute_fun( +@register_compute_fun( name="curvature", label="\\kappa", units="m^{-1}", @@ -652,9 +545,12 @@ def _frenet_binormal(params, data, transforms, basis="rpz"): dim=1, params=[], transforms={}, + profiles=[], + coordinates="s", data=["r_s", "r_ss"], + parameterization="desc.geometry.Curve", ) -def _curvature(params, data, transforms, basis="rpz"): +def _curvature(params, transforms, profiles, data, **kwargs): dxn = jnp.linalg.norm(data["r_s"], axis=1)[:, jnp.newaxis] data["curvature"] = jnp.linalg.norm( jnp.cross(data["r_s"], data["r_ss"], axis=1) / dxn**3, axis=1 @@ -662,7 +558,7 @@ def _curvature(params, data, transforms, basis="rpz"): return data -@register_curve_compute_fun( +@register_compute_fun( name="torsion", label="\\tau", units="m^{-1}", @@ -671,9 +567,12 @@ def _curvature(params, data, transforms, basis="rpz"): dim=1, params=[], transforms={}, + profiles=[], + coordinates="s", data=["r_s", "r_ss", "r_sss"], + parameterization="desc.geometry.Curve", ) -def _torsion(params, data, transforms, basis="rpz"): +def _torsion(params, transforms, profiles, data, **kwargs): dxd2x = jnp.cross(data["r_s"], data["r_ss"], axis=1) data["torsion"] = ( jnp.sum(dxd2x * data["r_sss"], axis=1) @@ -682,7 +581,7 @@ def _torsion(params, data, transforms, basis="rpz"): return data -@register_curve_compute_fun( +@register_compute_fun( name="length", label="L", units="m", @@ -691,9 +590,12 @@ def _torsion(params, data, transforms, basis="rpz"): dim=0, params=[], transforms={}, + profiles=[], + coordinates="s", data=["s", "r_s"], + parameterization="desc.geometry.Curve", ) -def _length(params, data, transforms, basis="rpz"): +def _length(params, transforms, profiles, data, **kwargs): T = jnp.linalg.norm(data["r_s"], axis=1) data["length"] = jnp.trapz(T, data["s"]) return data From 4ddf1da6ae40e4262745a2dd8a1620cd32198ffc Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 6 Jul 2023 00:11:45 -0400 Subject: [PATCH 04/28] Add basic surface compute funs --- desc/compute/_surface.py | 531 +++++++++++++++++++++++++++++++++++++++ 1 file changed, 531 insertions(+) create mode 100644 desc/compute/_surface.py diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py new file mode 100644 index 0000000000..a620a63c5d --- /dev/null +++ b/desc/compute/_surface.py @@ -0,0 +1,531 @@ +from desc.backend import jnp +from desc.geometry.utils import rpz2xyz, rpz2xyz_vec + +from .data_index import register_compute_fun + + +@register_compute_fun( + name="r", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 0, 0]], + "Z": [[0, 0, 0]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"]) + Z = transforms["Z"].transform(params["Z_lmn"]) + phi = transforms["grid"].nodes[:, 2] + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["r"] = coords + return data + + +@register_compute_fun( + name="r_r", + label="\\partial_{\\rho} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, radial derivative", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["r_r"] = coords + return data + + +@register_compute_fun( + name="r_t", + label="\\partial_{\\theta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, poloidal derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 1, 0]], + "Z": [[0, 1, 0]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dt=1) + Z = transforms["Z"].transform(params["Z_lmn"], dt=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_t"] = coords + return data + + +@register_compute_fun( + name="r_z", + label="\\partial_{\\zeta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, toroidal derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1]], + "Z": [[0, 0, 1]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R0 = transforms["R"].transform(params["R_lmn"], dz=0) + dR = transforms["R"].transform(params["R_lmn"], dz=1) + dZ = transforms["Z"].transform(params["Z_lmn"], dz=1) + dphi = R0 + coords = jnp.stack([dR, dphi, dZ], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_z"] = coords + return data + + +@register_compute_fun( + name="r_rr", + label="\\partial_{\\rho \\rho} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, second radial derivative", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["r_rr"] = coords + return data + + +@register_compute_fun( + name="r_tt", + label="\\partial_{\theta \\theta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, second poloidal derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 2, 0]], + "Z": [[0, 2, 0]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_tt_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dt=2) + Z = transforms["Z"].transform(params["Z_lmn"], dt=2) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_tt"] = coords + return data + + +@register_compute_fun( + name="r_zz", + label="\\partial_{\\zeta \\zeta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, toroidal derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], + "Z": [[0, 0, 2]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_zz_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R0 = transforms["R"].transform(params["R_lmn"], dz=0) + dR = transforms["R"].transform(params["R_lmn"], dz=1) + d2R = transforms["R"].transform(params["R_lmn"], dz=2) + d2Z = transforms["Z"].transform(params["Z_lmn"], dz=2) + dphi = 2 * dR + coords = jnp.stack([d2R - R0, dphi, d2Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_zz"] = coords + return data + + +@register_compute_fun( + name="r_rt", + label="\\partial_{\\rho \\theta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, radial/poloidal derivative", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_rt_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["r_rt"] = coords + return data + + +@register_compute_fun( + name="r_rz", + label="\\partial_{\\rho \\zeta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, radial/toroidal derivative", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_rz_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["r_rz"] = coords + return data + + +@register_compute_fun( + name="r_tz", + label="\\partial_{\\theta \\zeta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, poloidal/toroidal derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 0, 0], [0, 0, 1]], + "Z": [[0, 0, 1]], + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _r_tz_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + dR = transforms["R"].transform(params["R_lmn"], dt=1, dz=1) + dZ = transforms["Z"].transform(params["Z_lmn"], dt=1, dz=1) + dphi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([dR, dphi, dZ], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_tz"] = coords + return data + + +@register_compute_fun( + name="r", + label="\\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 0, 0]], + "Z": [[0, 0, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"]) + Z = transforms["Z"].transform(params["Z_lmn"]) + phi = transforms["grid"].nodes[:, 2] + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["r"] = coords + return data + + +@register_compute_fun( + name="r_r", + label="\\partial_{\\rho} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, radial derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[1, 0, 0]], + "Z": [[1, 0, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=1) + Z = transforms["Z"].transform(params["Z_lmn"], dr=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["r_r"] = coords + return data + + +@register_compute_fun( + name="r_t", + label="\\partial_{\\theta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, poloidal derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 1, 0]], + "Z": [[0, 1, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dt=1) + Z = transforms["Z"].transform(params["Z_lmn"], dt=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_t"] = coords + return data + + +@register_compute_fun( + name="r_z", + label="\\partial_{\\zeta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, toroidal derivative", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["r_z"] = coords + return data + + +@register_compute_fun( + name="r_rr", + label="\\partial_{\\rho \\rho} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, second radial derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[2, 0, 0]], + "Z": [[2, 0, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=2) + Z = transforms["Z"].transform(params["Z_lmn"], dr=2) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["r_rr"] = coords + return data + + +@register_compute_fun( + name="r_tt", + label="\\partial_{\theta \\theta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, second poloidal derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 2, 0]], + "Z": [[0, 2, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_tt_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dt=2) + Z = transforms["Z"].transform(params["Z_lmn"], dt=2) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["r_tt"] = coords + return data + + +@register_compute_fun( + name="r_zz", + label="\\partial_{\\zeta \\zeta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, toroidal derivative", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_zz_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["r_zz"] = coords + return data + + +@register_compute_fun( + name="r_rt", + label="\\partial_{\\rho \\theta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, radial/poloidal derivative", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[1, 1, 0]], + "Z": [[1, 1, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_rt_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=1, dt=1) + Z = transforms["Z"].transform(params["Z_lmn"], dr=1, dt=1) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz(coords) + data["r_rt"] = coords + return data + + +@register_compute_fun( + name="r_rz", + label="\\partial_{\\rho \\zeta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, radial/toroidal derivative", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_rz_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["r_rz"] = coords + return data + + +@register_compute_fun( + name="r_tz", + label="\\partial_{\\theta \\zeta} \\mathbf{r}", + units="m", + units_long="meters", + description="Position vector along surface, poloidal/toroidal derivative", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _r_tz_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["r_tz"] = coords + return data From abdc04a173c035bf69c067c8aa97b17ae8c2ad88 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 6 Jul 2023 00:21:55 -0400 Subject: [PATCH 05/28] Fix typos --- desc/compute/_curve.py | 4 ++-- desc/compute/_surface.py | 4 ++-- 2 files changed, 4 insertions(+), 4 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 11abb85cac..ce38de8424 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -517,7 +517,7 @@ def _frenet_normal(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="frenet_binormal", - label="\\mathrm{B}_{\\mathrm{Frenet-Serret}}", + label="\\mathbf{B}_{\\mathrm{Frenet-Serret}}", units="~", units_long="None", description="Binormal unit vector to curve in Frenet-Serret frame", @@ -526,7 +526,7 @@ def _frenet_normal(params, transforms, profiles, data, **kwargs): transforms={"rotmat": []}, profiles=[], coordinates="s", - data=["T", "N"], + data=["frenet_tangent", "frenet_normal"], parameterization="desc.geometry.Curve", ) def _frenet_binormal(params, transforms, profiles, data, **kwargs): diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index a620a63c5d..dd28e499fb 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -138,7 +138,7 @@ def _r_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs) @register_compute_fun( name="r_tt", - label="\\partial_{\theta \\theta} \\mathbf{r}", + label="\\partial_{\\theta \\theta} \\mathbf{r}", units="m", units_long="meters", description="Position vector along surface, second poloidal derivative", @@ -409,7 +409,7 @@ def _r_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs) @register_compute_fun( name="r_tt", - label="\\partial_{\theta \\theta} \\mathbf{r}", + label="\\partial_{\\theta \\theta} \\mathbf{r}", units="m", units_long="meters", description="Position vector along surface, second poloidal derivative", From b6bb2976b5cc3929cdfcf7e35d9bf2e05cc239ac Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 6 Jul 2023 00:22:13 -0400 Subject: [PATCH 06/28] Add curve/surface stuff to compute init --- desc/compute/__init__.py | 2 ++ desc/compute/data_index.py | 5 +++++ 2 files changed, 7 insertions(+) diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index 5cfba65308..feed595ecf 100644 --- a/desc/compute/__init__.py +++ b/desc/compute/__init__.py @@ -30,6 +30,7 @@ _basis_vectors, _bootstrap, _core, + _curve, _equil, _field, _geometry, @@ -37,6 +38,7 @@ _profiles, _qs, _stability, + _surface, ) from .data_index import data_index from .utils import ( diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index d5b873bb0e..e7ae8c16aa 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -111,30 +111,35 @@ def _decorator(func): "desc.geometry.FourierRZCurve", "FourierRZCurve", "desc.geometry.core.Curve", + "desc.geometry.Curve", "Curve", ], "desc.geometry.curve.FourierXYZCurve": [ "desc.geometry.FourierXYZCurve", "FourierXYZCurve", "desc.geometry.core.Curve", + "desc.geometry.Curve", "Curve", ], "desc.geometry.curve.FourierPlanarCurve": [ "desc.geometry.FourierPlanarCurve", "FourierPlanarCurve", "desc.geometry.core.Curve", + "desc.geometry.Curve", "Curve", ], "desc.geometry.surface.FourierRZToroidalSurface": [ "desc.geometry.FourierRZToroidalSurface", "FourierRZToroidalSurface", "desc.geometry.core.Surface", + "desc.geometry.Surface", "Surface", ], "desc.geometry.surface.ZernikeRZToroidalSection": [ "desc.geometry.ZernikeRZToroidalSection", "ZernikeRZToroidalSection", "desc.geometry.core.Surface", + "desc.geometry.Surface", "Surface", ], } From 42d6822eeec3c7ffdcfe28ee0201d0fd0abce305 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 21:03:48 -0400 Subject: [PATCH 07/28] Prevent overwriting of dependencies of shared functions --- desc/compute/data_index.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index e7ae8c16aa..1b9891a78e 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -84,7 +84,7 @@ def _decorator(func): flag = False for base_name, aliases in _aliases.items(): if p in aliases: - data_index[base_name][name] = d + data_index[base_name][name] = d.copy() flag = True if not flag: raise ValueError( From 0f33c183b35dcfded5232c22d04d4d2fbe25dece Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 21:04:26 -0400 Subject: [PATCH 08/28] Add rotmat and shift to transforms --- desc/compute/utils.py | 10 +++++++++- 1 file changed, 9 insertions(+), 1 deletion(-) diff --git a/desc/compute/utils.py b/desc/compute/utils.py index c43bc45973..b80e167a89 100644 --- a/desc/compute/utils.py +++ b/desc/compute/utils.py @@ -313,7 +313,10 @@ def get_params(keys, obj, has_axis=False, **kwargs): params = [] for key in deps: params += data_index[p][key]["dependencies"]["params"] - params = _sort_args(list(set(params))) + if p == "desc.equilibrium.equilibrium.Equilibrium": + # probably need some way to distinguish between params from different instances + # of the same class? + params = _sort_args(list(set(params))) if isinstance(obj, str) or inspect.isclass(obj): return params params = {name: np.atleast_1d(getattr(obj, name)).copy() for name in params} @@ -376,6 +379,11 @@ def get_transforms(keys, obj, grid, **kwargs): build=True, build_pinv=True, ) + elif c == "rotmat": + transforms["rotmat"] = obj.rotmat + elif c == "shift": + transforms["shift"] = obj.shift + return transforms From 15f95e8f54e4f6dedda5947e7b920969b82ba5c4 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 21:06:44 -0400 Subject: [PATCH 09/28] Rename bases for curve classes --- desc/compute/_curve.py | 56 ++++++++++++------ desc/geometry/curve.py | 125 ++++++++++++++++------------------------- 2 files changed, 88 insertions(+), 93 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index ce38de8424..fc4f9d63c6 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -369,16 +369,22 @@ def _r_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve", dim=3, params=["X_n", "Y_n", "Z_n"], - transforms={"r": [[0, 0, 0]], "rotmat": [], "shift": []}, + transforms={ + "X": [[0, 0, 0]], + "Y": [[0, 0, 0]], + "Z": [[0, 0, 0]], + "rotmat": [], + "shift": [], + }, profiles=[], coordinates="s", data=[], parameterization="desc.geometry.FourierXYZCurve", ) def _r_FourierXYZCurve(params, transforms, profiles, data, **kwargs): - X = transforms["r"].transform(params["X_n"], dz=0) - Y = transforms["r"].transform(params["Y_n"], dz=0) - Z = transforms["r"].transform(params["Z_n"], dz=0) + X = transforms["X"].transform(params["X_n"], dz=0) + Y = transforms["Y"].transform(params["Y_n"], dz=0) + Z = transforms["Z"].transform(params["Z_n"], dz=0) coords = jnp.stack([X, Y, Z], axis=1) coords = coords @ transforms["rotmat"].T + transforms["shift"][jnp.newaxis, :] if kwargs.get("basis", "rpz").lower() == "rpz": @@ -395,16 +401,22 @@ def _r_FourierXYZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, first derivative", dim=3, params=["X_n", "Y_n", "Z_n"], - transforms={"r": [[0, 0, 1]], "rotmat": [], "shift": []}, + transforms={ + "X": [[0, 0, 1]], + "Y": [[0, 0, 1]], + "Z": [[0, 0, 1]], + "rotmat": [], + "shift": [], + }, profiles=[], coordinates="s", data=[], parameterization="desc.geometry.FourierXYZCurve", ) def _r_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): - X = transforms["r"].transform(params["X_n"], dz=1) - Y = transforms["r"].transform(params["Y_n"], dz=1) - Z = transforms["r"].transform(params["Z_n"], dz=1) + X = transforms["X"].transform(params["X_n"], dz=1) + Y = transforms["Y"].transform(params["Y_n"], dz=1) + Z = transforms["Z"].transform(params["Z_n"], dz=1) coords = jnp.stack([X, Y, Z], axis=1) coords = coords @ transforms["rotmat"].T if kwargs.get("basis", "rpz").lower() == "rpz": @@ -425,16 +437,22 @@ def _r_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, second derivative", dim=3, params=["X_n", "Y_n", "Z_n"], - transforms={"r": [[0, 0, 2]], "rotmat": [], "shift": []}, + transforms={ + "X": [[0, 0, 2]], + "Y": [[0, 0, 2]], + "Z": [[0, 0, 2]], + "rotmat": [], + "shift": [], + }, profiles=[], coordinates="s", data=[], parameterization="desc.geometry.FourierXYZCurve", ) def _r_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): - X = transforms["r"].transform(params["X_n"], dz=2) - Y = transforms["r"].transform(params["Y_n"], dz=2) - Z = transforms["r"].transform(params["Z_n"], dz=2) + X = transforms["X"].transform(params["X_n"], dz=2) + Y = transforms["Y"].transform(params["Y_n"], dz=2) + Z = transforms["Z"].transform(params["Z_n"], dz=2) coords = jnp.stack([X, Y, Z], axis=1) coords = coords @ transforms["rotmat"].T if kwargs.get("basis", "rpz").lower() == "rpz": @@ -455,16 +473,22 @@ def _r_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): description="Position vector along curve, third derivative", dim=3, params=["X_n", "Y_n", "Z_n"], - transforms={"r": [[0, 0, 3]], "rotmat": [], "shift": []}, + transforms={ + "X": [[0, 0, 3]], + "Y": [[0, 0, 3]], + "Z": [[0, 0, 3]], + "rotmat": [], + "shift": [], + }, profiles=[], coordinates="s", data=[], parameterization="desc.geometry.FourierXYZCurve", ) def _r_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): - X = transforms["r"].transform(params["X_n"], dz=3) - Y = transforms["r"].transform(params["Y_n"], dz=3) - Z = transforms["r"].transform(params["Z_n"], dz=3) + X = transforms["X"].transform(params["X_n"], dz=3) + Y = transforms["Y"].transform(params["Y_n"], dz=3) + Z = transforms["Z"].transform(params["Z_n"], dz=3) coords = jnp.stack([X, Y, Z], axis=1) coords = coords @ transforms["rotmat"].T if kwargs.get("basis", "rpz").lower() == "rpz": diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index bc00e548b1..48017ee5ae 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -475,14 +475,19 @@ class FourierXYZCurve(Curve): Fourier coefficients for X, Y, Z modes : array-like mode numbers associated with X_n etc. - grid : Grid - default grid for computation name : str name for this curve """ - _io_attrs_ = Curve._io_attrs_ + ["_X_n", "_Y_n", "_Z_n", "_basis", "_transform"] + _io_attrs_ = Curve._io_attrs_ + [ + "_X_n", + "_Y_n", + "_Z_n", + "_X_basis", + "_Y_basis", + "_Z_basis", + ] def __init__( self, @@ -503,54 +508,44 @@ def __init__( assert issubclass(modes.dtype.type, np.integer) N = np.max(abs(modes)) - self._basis = FourierSeries(N, NFP=1, sym=False) - self._X_n = copy_coeffs(X_n, modes, self.basis.modes[:, 2]) - self._Y_n = copy_coeffs(Y_n, modes, self.basis.modes[:, 2]) - self._Z_n = copy_coeffs(Z_n, modes, self.basis.modes[:, 2]) - - if grid is None: - grid = LinearGrid(N=2 * N, endpoint=True) - self._grid = grid - self._transform = self._get_transforms(grid) + self._X_basis = FourierSeries(N, NFP=1, sym=False) + self._Y_basis = FourierSeries(N, NFP=1, sym=False) + self._Z_basis = FourierSeries(N, NFP=1, sym=False) + self._X_n = copy_coeffs(X_n, modes, self.X_basis.modes[:, 2]) + self._Y_n = copy_coeffs(Y_n, modes, self.Y_basis.modes[:, 2]) + self._Z_n = copy_coeffs(Z_n, modes, self.Z_basis.modes[:, 2]) @property - def basis(self): - """Spectral basis for Fourier series.""" - return self._basis + def X_basis(self): + """Spectral basis for X Fourier series.""" + return self._X_basis @property - def grid(self): - """Default grid for computation.""" - return self._grid + def Y_basis(self): + """Spectral basis for Y Fourier series.""" + return self._Y_basis - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif jnp.isscalar(new): - self._grid = LinearGrid(N=new, endpoint=True) - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._transform.grid = self.grid + @property + def Z_basis(self): + """Spectral basis for Z Fourier series.""" + return self._Z_basis @property def N(self): """Maximum mode number.""" - return self.basis.N + return max(self.X_basis.N, self.Y_basis.N, self.Z_basis.N) def change_resolution(self, N=None): """Change the maximum angular resolution.""" if (N is not None) and (N != self.N): modes_old = self.basis.modes - self.basis.change_resolution(N=N) + self.X_basis.change_resolution(N=N) + self.Y_basis.change_resolution(N=N) + self.Z_basis.change_resolution(N=N) self._transform = self._get_transforms(self.grid) - self.X_n = copy_coeffs(self.X_n, modes_old, self.basis.modes) - self.Y_n = copy_coeffs(self.Y_n, modes_old, self.basis.modes) - self.Z_n = copy_coeffs(self.Z_n, modes_old, self.basis.modes) + self.X_n = copy_coeffs(self.X_n, modes_old, self.X_basis.modes) + self.Y_n = copy_coeffs(self.Y_n, modes_old, self.Y_basis.modes) + self.Z_n = copy_coeffs(self.Z_n, modes_old, self.Z_basis.modes) def get_coeffs(self, n): """Get Fourier coefficients for given mode number(s).""" @@ -593,12 +588,12 @@ def X_n(self): @X_n.setter def X_n(self, new): - if len(new) == self._basis.num_modes: + if len(new) == self.X_basis.num_modes: self._X_n = jnp.asarray(new) else: raise ValueError( f"X_n should have the same size as the basis, got {len(new)} for " - + f"basis with {self._basis.num_modes} modes." + + f"basis with {self.X_basis.num_modes} modes." ) @property @@ -608,12 +603,12 @@ def Y_n(self): @Y_n.setter def Y_n(self, new): - if len(new) == self._basis.num_modes: + if len(new) == self.Y_basis.num_modes: self._Y_n = jnp.asarray(new) else: raise ValueError( f"Y_n should have the same size as the basis, got {len(new)} for " - + f"basis with {self._basis.num_modes} modes." + + f"basis with {self.Y_basis.num_modes} modes." ) @property @@ -623,12 +618,12 @@ def Z_n(self): @Z_n.setter def Z_n(self, new): - if len(new) == self._basis.num_modes: + if len(new) == self.Z_basis.num_modes: self._Z_n = jnp.asarray(new) else: raise ValueError( f"Z_n should have the same size as the basis, got {len(new)} for " - + f"basis with {self._basis.num_modes} modes." + + f"basis with {self.Z_basis.num_modes} modes." ) def _get_transforms(self, grid=None): @@ -839,8 +834,7 @@ class FourierPlanarCurve(Curve): "_r_n", "_center", "_normal", - "_basis", - "_transform", + "_r_basis", ] # Reference frame is centered at the origin with normal in the +Z direction. @@ -863,52 +857,29 @@ def __init__( assert issubclass(modes.dtype.type, np.integer) N = np.max(abs(modes)) - self._basis = FourierSeries(N, NFP=1, sym=False) - self._r_n = copy_coeffs(r_n, modes, self.basis.modes[:, 2]) + self._r_basis = FourierSeries(N, NFP=1, sym=False) + self._r_n = copy_coeffs(r_n, modes, self.r_basis.modes[:, 2]) self.normal = normal self.center = center - if grid is None: - grid = LinearGrid(N=2 * self.N, endpoint=True) - self._grid = grid - self._transform = self._get_transforms(grid) @property - def basis(self): + def r_basis(self): """Spectral basis for Fourier series.""" - return self._basis - - @property - def grid(self): - """Default grid for computation.""" - return self._grid - - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif jnp.isscalar(new): - self._grid = LinearGrid(N=new, endpoint=True) - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._transform.grid = self.grid + return self._r_basis @property def N(self): """Maximum mode number.""" - return self.basis.N + return self.r_basis.N def change_resolution(self, N=None): """Change the maximum angular resolution.""" if (N is not None) and (N != self.N): - modes_old = self.basis.modes - self.basis.change_resolution(N=N) + modes_old = self.r_basis.modes + self.r_basis.change_resolution(N=N) self._transform = self._get_transforms(self.grid) - self.r_n = copy_coeffs(self.r_n, modes_old, self.basis.modes) + self.r_n = copy_coeffs(self.r_n, modes_old, self.r_basis.modes) @property def center(self): @@ -945,12 +916,12 @@ def r_n(self): @r_n.setter def r_n(self, new): - if len(np.asarray(new)) == self._basis.num_modes: + if len(np.asarray(new)) == self.r_basis.num_modes: self._r_n = jnp.asarray(new) else: raise ValueError( f"r_n should have the same size as the basis, got {len(new)} for " - + f"basis with {self._basis.num_modes} modes." + + f"basis with {self.r_basis.num_modes} modes." ) def get_coeffs(self, n): From ca39f265a59394f8d1292e34f7f93dd4e34bf71f Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 21:08:27 -0400 Subject: [PATCH 10/28] Remove grid, transform, compute methods from curve classes --- desc/geometry/core.py | 8 +- desc/geometry/curve.py | 665 +---------------------------------------- 2 files changed, 7 insertions(+), 666 deletions(-) diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 0d9ab0e346..3ef8338ff8 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -13,7 +13,7 @@ class Curve(IOAble, ABC): """Abstract base class for 1D curves in 3D space.""" - _io_attrs_ = ["_name", "_grid", "shift", "rotmat"] + _io_attrs_ = ["_name", "shift", "rotmat"] def __init__(self, name=""): self.shift = jnp.array([0, 0, 0]) @@ -76,14 +76,14 @@ def __repr__(self): type(self).__name__ + " at " + str(hex(id(self))) - + " (name={}, grid={})".format(self.name, self.grid) + + " (name={})".format(self.name) ) class Surface(IOAble, ABC): """Abstract base class for 2d surfaces in 3d space.""" - _io_attrs_ = ["_name", "_grid", "_sym", "_L", "_M", "_N"] + _io_attrs_ = ["_name", "_sym", "_L", "_M", "_N"] @property def name(self): @@ -207,5 +207,5 @@ def __repr__(self): type(self).__name__ + " at " + str(hex(id(self))) - + " (name={}, grid={})".format(self.name, self.grid) + + " (name={})".format(self.name) ) diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index 48017ee5ae..88a7c0f4b3 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -6,12 +6,11 @@ from desc.backend import jnp, put from desc.basis import FourierSeries -from desc.grid import Grid, LinearGrid +from desc.grid import LinearGrid from desc.transform import Transform from desc.utils import copy_coeffs from .core import Curve -from .utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec __all__ = [ "FourierRZCurve", @@ -35,8 +34,6 @@ class FourierRZCurve(Curve): Number of field periods. sym : bool Whether to enforce stellarator symmetry. - grid : Grid - Default grid for computation. name : str Name for this curve. @@ -47,8 +44,6 @@ class FourierRZCurve(Curve): "_Z_n", "_R_basis", "_Z_basis", - "_R_transform", - "_Z_transform", "_sym", "_NFP", ] @@ -61,7 +56,6 @@ def __init__( modes_Z=None, NFP=1, sym="auto", - grid=None, name="", ): super().__init__(name) @@ -92,17 +86,12 @@ def __init__( NZ = np.max(abs(modes_Z)) N = max(NR, NZ) self._NFP = NFP - self._R_basis = FourierSeries(NR, NFP, sym="cos" if sym else False) - self._Z_basis = FourierSeries(NZ, NFP, sym="sin" if sym else False) + self._R_basis = FourierSeries(N, NFP, sym="cos" if sym else False) + self._Z_basis = FourierSeries(N, NFP, sym="sin" if sym else False) self._R_n = copy_coeffs(R_n, modes_R, self.R_basis.modes[:, 2]) self._Z_n = copy_coeffs(Z_n, modes_Z, self.Z_basis.modes[:, 2]) - if grid is None: - grid = LinearGrid(N=2 * N, NFP=self.NFP, endpoint=True) - self._grid = grid - self._R_transform, self._Z_transform = self._get_transforms(grid) - @property def sym(self): """Whether this curve has stellarator symmetry.""" @@ -130,26 +119,6 @@ def NFP(self, new): ), f"NFP should be a positive integer, got {type(new)}" self.change_resolution(NFP=new) - @property - def grid(self): - """Default grid for computation.""" - return self._grid - - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif jnp.isscalar(new): - self._grid = LinearGrid(N=new, endpoint=True) - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._R_transform.grid = self.grid - self._Z_transform.grid = self.grid - @property def N(self): """Maximum mode number.""" @@ -238,207 +207,6 @@ def Z_n(self, new): + f"basis with {self.Z_basis.num_modes} modes" ) - def _get_transforms(self, grid=None): - if grid is None: - return self._R_transform, self._Z_transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = np.linspace(0, 2 * np.pi, grid) - grid = np.atleast_1d(grid) - if grid.ndim == 1: - grid = np.pad(grid[:, np.newaxis], ((0, 0), (2, 0))) - grid = Grid(grid, sort=False) - R_transform = Transform( - grid, - self.R_basis, - derivs=np.array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), - ) - Z_transform = Transform( - grid, - self.Z_basis, - derivs=np.array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), - ) - return R_transform, Z_transform - - def compute_coordinates(self, R_n=None, Z_n=None, grid=None, dt=0, basis="rpz"): - """Compute values using specified coefficients. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. Defaults to self.R_n, self.Z_n - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - dt: int - derivative order to compute - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - R, phi, Z or x, y, z coordinates of the curve at specified grid locations - in phi. - - """ - assert basis.lower() in ["rpz", "xyz"] - if R_n is None: - R_n = self.R_n - if Z_n is None: - Z_n = self.Z_n - R_transform, Z_transform = self._get_transforms(grid) - if dt == 0: - R = R_transform.transform(R_n, dz=0) - Z = Z_transform.transform(Z_n, dz=0) - phi = R_transform.grid.nodes[:, 2] - coords = jnp.stack([R, phi, Z], axis=1) - elif dt == 1: - R0 = R_transform.transform(R_n, dz=0) - dR = R_transform.transform(R_n, dz=dt) - dZ = Z_transform.transform(Z_n, dz=dt) - dphi = R0 - coords = jnp.stack([dR, dphi, dZ], axis=1) - elif dt == 2: - R0 = R_transform.transform(R_n, dz=0) - dR = R_transform.transform(R_n, dz=1) - d2R = R_transform.transform(R_n, dz=2) - d2Z = Z_transform.transform(Z_n, dz=2) - R = d2R - R0 - Z = d2Z - # 2nd derivative wrt phi = 0 - phi = 2 * dR - coords = jnp.stack([R, phi, Z], axis=1) - elif dt == 3: - R0 = R_transform.transform(R_n, dz=0) - dR = R_transform.transform(R_n, dz=1) - d2R = R_transform.transform(R_n, dz=2) - d3R = R_transform.transform(R_n, dz=3) - d3Z = Z_transform.transform(Z_n, dz=3) - R = d3R - 3 * dR - Z = d3Z - phi = 3 * d2R - R0 - coords = jnp.stack([R, phi, Z], axis=1) - else: - raise NotImplementedError( - "Derivatives higher than 3 have not been implemented in " - + "cylindrical coordinates." - ) - # convert to xyz for displacement and rotation - if dt > 0: - coords = rpz2xyz_vec(coords, phi=R_transform.grid.nodes[:, 2]) - else: - coords = rpz2xyz(coords) - coords = coords @ self.rotmat.T + (self.shift[jnp.newaxis, :] * (dt == 0)) - if basis.lower() == "rpz": - if dt > 0: - coords = xyz2rpz_vec(coords, phi=R_transform.grid.nodes[:, 2]) - else: - coords = xyz2rpz(coords) - return coords - - def compute_frenet_frame(self, R_n=None, Z_n=None, grid=None, basis="rpz"): - """Compute Frenet frame vectors using specified coefficients. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. Defaults to self.R_n, self.Z_n - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - basis : {"rpz", "xyz"} - basis vectors to use for Frenet vector representation - - Returns - ------- - T, N, B : ndarrays, shape(k,3) - tangent, normal, and binormal vectors of the curve at specified grid - locations in phi - - """ - T = self.compute_coordinates(R_n, Z_n, grid, dt=1, basis=basis) - N = self.compute_coordinates(R_n, Z_n, grid, dt=2, basis=basis) - - T = T / jnp.linalg.norm(T, axis=1)[:, jnp.newaxis] - N = N / jnp.linalg.norm(N, axis=1)[:, jnp.newaxis] - B = jnp.cross(T, N, axis=1) * jnp.linalg.det(self.rotmat) - - return T, N, B - - def compute_curvature(self, R_n=None, Z_n=None, grid=None): - """Compute curvature using specified coefficients. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. Defaults to self.R_n, self.Z_n - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - kappa : ndarray, shape(k,) - curvature of the curve at specified grid locations in phi - - """ - dx = self.compute_coordinates(R_n, Z_n, grid, dt=1) - d2x = self.compute_coordinates(R_n, Z_n, grid, dt=2) - dxn = jnp.linalg.norm(dx, axis=1)[:, jnp.newaxis] - kappa = jnp.linalg.norm(jnp.cross(dx, d2x, axis=1) / dxn**3, axis=1) - return kappa - - def compute_torsion(self, R_n=None, Z_n=None, grid=None): - """Compute torsion using specified coefficients. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. Defaults to self.R_n, self.Z_n - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - tau : ndarray, shape(k,) - torsion of the curve at specified grid locations in phi - - """ - dx = self.compute_coordinates(R_n, Z_n, grid, dt=1) - d2x = self.compute_coordinates(R_n, Z_n, grid, dt=2) - d3x = self.compute_coordinates(R_n, Z_n, grid, dt=3) - dxd2x = jnp.cross(dx, d2x, axis=1) - tau = ( - jnp.sum(dxd2x * d3x, axis=1) - / jnp.linalg.norm(dxd2x, axis=1)[:, jnp.newaxis] ** 2 - ) - return tau - - def compute_length(self, R_n=None, Z_n=None, grid=None): - """Compute the length of the curve using specified nodes for quadrature. - - Parameters - ---------- - R_n, Z_n: array-like - Fourier coefficients for R, Z. If not given, defaults to values given - by R_n, Z_n attributes - grid : Grid or array-like - Toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - length : float - length of the curve approximated by quadrature - - """ - R_transform, Z_transform = self._get_transforms(grid) - T = self.compute_coordinates(R_n, Z_n, grid, dt=1) - T = jnp.linalg.norm(T, axis=1) - phi = R_transform.grid.nodes[:, 2] - return jnp.trapz(T, phi) - def to_FourierXYZCurve(self, N=None): """Convert to FourierXYZCurve representation. @@ -495,7 +263,6 @@ def __init__( Y_n=[0, 0, 0], Z_n=[-2, 0, 0], modes=None, - grid=None, name="", ): super().__init__(name) @@ -626,182 +393,6 @@ def Z_n(self, new): + f"basis with {self.Z_basis.num_modes} modes." ) - def _get_transforms(self, grid=None): - if grid is None: - return self._transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = np.linspace(0, 2 * np.pi, grid) - grid = np.atleast_1d(grid) - if grid.ndim == 1: - grid = np.pad(grid[:, np.newaxis], ((0, 0), (2, 0))) - grid = Grid(grid, sort=False) - transform = Transform( - grid, - self.basis, - derivs=np.array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), - ) - return transform - - def compute_coordinates( - self, X_n=None, Y_n=None, Z_n=None, grid=None, dt=0, basis="xyz" - ): - """Compute values using specified coefficients. - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - dt: int - derivative order to compute - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - X, Y, Z or R, phi, Z coordinates of the curve at specified grid locations - in phi. - - """ - assert basis.lower() in ["rpz", "xyz"] - if X_n is None: - X_n = self.X_n - if Y_n is None: - Y_n = self.Y_n - if Z_n is None: - Z_n = self.Z_n - - transform = self._get_transforms(grid) - X = transform.transform(X_n, dz=dt) - Y = transform.transform(Y_n, dz=dt) - Z = transform.transform(Z_n, dz=dt) - - coords = jnp.stack([X, Y, Z], axis=1) - coords = coords @ self.rotmat.T + (self.shift[jnp.newaxis, :] * (dt == 0)) - if basis.lower() == "rpz": - if dt > 0: - coords = xyz2rpz_vec( - coords, - x=coords[:, 1] + self.shift[0], - y=coords[:, 1] + self.shift[1], - ) - else: - coords = xyz2rpz(coords) - return coords - - def compute_frenet_frame( - self, X_n=None, Y_n=None, Z_n=None, grid=None, basis="xyz" - ): - """Compute Frenet frame vectors using specified coefficients. - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - basis : {"rpz", "xyz"} - basis vectors to use for Frenet vector representation - - Returns - ------- - T, N, B : ndarrays, shape(k,3) - tangent, normal, and binormal vectors of the curve at specified grid - locations - - """ - T = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=1, basis=basis) - N = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=2, basis=basis) - - T = T / jnp.linalg.norm(T, axis=1)[:, jnp.newaxis] - N = N / jnp.linalg.norm(N, axis=1)[:, jnp.newaxis] - B = jnp.cross(T, N, axis=1) * jnp.linalg.det(self.rotmat) - - return T, N, B - - def compute_curvature(self, X_n=None, Y_n=None, Z_n=None, grid=None): - """Compute curvature using specified coefficients. - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - kappa : ndarray, shape(k,) - curvature of the curve at specified grid locations in phi - - """ - dx = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=1) - d2x = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=2) - dxn = jnp.linalg.norm(dx, axis=1)[:, jnp.newaxis] - kappa = jnp.linalg.norm(jnp.cross(dx, d2x, axis=1) / dxn**3, axis=1) - return kappa - - def compute_torsion(self, X_n=None, Y_n=None, Z_n=None, grid=None): - """Compute torsion using specified coefficientsnp.empty((0, 3)). - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - tau : ndarray, shape(k,) - torsion of the curve at specified grid locations in phi - - """ - dx = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=1) - d2x = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=2) - d3x = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=3) - dxd2x = jnp.cross(dx, d2x, axis=1) - tau = ( - jnp.sum(dxd2x * d3x, axis=1) - / jnp.linalg.norm(dxd2x, axis=1)[:, jnp.newaxis] ** 2 - ) - return tau - - def compute_length(self, X_n=None, Y_n=None, Z_n=None, grid=None): - """Compute the length of the curve using specified nodes for quadrature. - - Parameters - ---------- - X_n, Y_n, Z_n: array-like - Fourier coefficients for X, Y, Z. If not given, defaults to values given - by X_n, Y_n, Z_n attributes - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - length : float - length of the curve approximated by quadrature - - """ - transform = self._get_transforms(grid) - T = self.compute_coordinates(X_n, Y_n, Z_n, grid, dt=1) - T = jnp.linalg.norm(T, axis=1) - theta = transform.grid.nodes[:, 2] - return jnp.trapz(T, theta) - # TODO: to_rz method for converting to FourierRZCurve representation # (might be impossible to parameterize with toroidal angle phi) @@ -823,8 +414,6 @@ class FourierPlanarCurve(Curve): Fourier coefficients for radius from center as function of polar angle modes : array-like mode numbers associated with r_n - grid : Grid - default grid for computation name : str name for this curve @@ -845,7 +434,6 @@ def __init__( normal=[0, 1, 0], r_n=2, modes=None, - grid=None, name="", ): super().__init__(name) @@ -942,250 +530,3 @@ def set_coeffs(self, n, r=None): idx = self.basis.get_idx(0, 0, nn) if rr is not None: self.r_n = put(self.r_n, idx, rr) - - def _normal_rotmat(self, normal=None): - """Rotation matrix to rotate z axis into plane normal.""" - nx, ny, nz = normal - nxny = jnp.sqrt(nx**2 + ny**2) - - R = jnp.array( - [ - [ny / nxny, -nx / nxny, 0], - [nx * nx / nxny, ny * nz / nxny, -nxny], - [nx, ny, nz], - ] - ).T - return jnp.where(nxny == 0, jnp.eye(3), R) - - def _get_transforms(self, grid=None): - if grid is None: - return self._transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = np.linspace(0, 2 * np.pi, grid) - grid = np.atleast_1d(grid) - if grid.ndim == 1: - grid = np.pad(grid[:, np.newaxis], ((0, 0), (2, 0))) - grid = Grid(grid, sort=False) - transform = Transform( - grid, - self.basis, - derivs=np.array([[0, 0, 0], [0, 0, 1], [0, 0, 2], [0, 0, 3]]), - ) - return transform - - def compute_coordinates( - self, center=None, normal=None, r_n=None, grid=None, dt=0, basis="xyz" - ): - """Compute values using specified coefficients. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - dt: int - derivative order to compute - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - X, Y, Z or R, phi, Z coordinates of the curve at specified grid locations - in theta. - - """ - assert basis.lower() in ["rpz", "xyz"] - if center is None: - center = self.center - if normal is None: - normal = self.normal - if r_n is None: - r_n = self.r_n - transform = self._get_transforms(grid) - r = transform.transform(r_n, dz=0) - t = transform.grid.nodes[:, -1] - Z = np.zeros_like(r) - - if dt == 0: - X = r * jnp.cos(t) - Y = r * jnp.sin(t) - coords = jnp.array([X, Y, Z]).T - elif dt == 1: - dr = transform.transform(r_n, dz=1) - dX = dr * jnp.cos(t) - r * jnp.sin(t) - dY = dr * jnp.sin(t) + r * jnp.cos(t) - coords = jnp.array([dX, dY, Z]).T - elif dt == 2: - dr = transform.transform(r_n, dz=1) - d2r = transform.transform(r_n, dz=2) - d2X = d2r * jnp.cos(t) - 2 * dr * jnp.sin(t) - r * jnp.cos(t) - d2Y = d2r * jnp.sin(t) + 2 * dr * jnp.cos(t) - r * jnp.sin(t) - coords = jnp.array([d2X, d2Y, Z]).T - elif dt == 3: - dr = transform.transform(r_n, dz=1) - d2r = transform.transform(r_n, dz=2) - d3r = transform.transform(r_n, dz=3) - d3X = ( - d3r * jnp.cos(t) - - 3 * d2r * jnp.sin(t) - - 3 * dr * jnp.cos(t) - + r * jnp.sin(t) - ) - d3Y = ( - d3r * jnp.sin(t) - + 3 * d2r * jnp.cos(t) - - 3 * dr * jnp.sin(t) - - r * jnp.cos(t) - ) - coords = jnp.array([d3X, d3Y, Z]).T - else: - raise NotImplementedError( - "Derivatives higher than 3 have not been implemented for planar curves." - ) - R = self._normal_rotmat(normal) - coords = jnp.matmul(coords, R.T) + (center * (dt == 0)) - coords = jnp.matmul(coords, self.rotmat.T) + (self.shift * (dt == 0)) - - if basis.lower() == "rpz": - X = r * jnp.cos(t) - Y = r * jnp.sin(t) - xyzcoords = jnp.array([X, Y, Z]).T - xyzcoords = jnp.matmul(xyzcoords, R.T) + center - xyzcoords = jnp.matmul(xyzcoords, self.rotmat.T) + self.shift - x, y, z = xyzcoords.T - coords = xyz2rpz_vec(coords, x=x, y=y) - - return coords - - def compute_frenet_frame( - self, center=None, normal=None, r_n=None, grid=None, basis="xyz" - ): - """Compute Frenet frame vectors using specified coefficients. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - basis : {"rpz", "xyz"} - basis vectors to use for Frenet vector representation - - Returns - ------- - T, N, B : ndarrays, shape(k,3) - tangent, normal, and binormal vectors of the curve at specified grid - locations in theta. - - """ - T = self.compute_coordinates(center, normal, r_n, grid, dt=1, basis=basis) - N = self.compute_coordinates(center, normal, r_n, grid, dt=2, basis=basis) - - T = T / jnp.linalg.norm(T, axis=1)[:, jnp.newaxis] - N = N / jnp.linalg.norm(N, axis=1)[:, jnp.newaxis] - B = jnp.cross(T, N, axis=1) * jnp.linalg.det(self.rotmat) - - return T, N, B - - def compute_curvature(self, center=None, normal=None, r_n=None, grid=None): - """Compute curvature using specified coefficients. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - kappa : ndarray, shape(k,) - curvature of the curve at specified grid locations in theta - - """ - dx = self.compute_coordinates(center, normal, r_n, grid, dt=1) - d2x = self.compute_coordinates(center, normal, r_n, grid, dt=2) - dxn = jnp.linalg.norm(dx, axis=1)[:, jnp.newaxis] - kappa = jnp.linalg.norm(jnp.cross(dx, d2x, axis=1) / dxn**3, axis=1) - return kappa - - def compute_torsion(self, center=None, normal=None, r_n=None, grid=None): - """Compute torsion using specified coefficients. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - tau : ndarray, shape(k,) - torsion of the curve at specified grid locations in phi - - """ - # torsion is zero for planar curves - transform = self._get_transforms(grid) - torsion = jnp.zeros_like(transform.grid.nodes[:, -1]) - return torsion - - def compute_length(self, center=None, normal=None, r_n=None, grid=None): - """Compute the length of the curve using specified nodes for quadrature. - - Parameters - ---------- - center : array-like, shape(3,) - x,y,z coordinates of center of curve. If not given, defaults to self.center - normal : array-like, shape(3,) - x,y,z components of normal vector to planar surface. If not given, defaults - to self.normal - r_n : array-like - Fourier coefficients for radius from center as function of polar angle. - If not given defaults to self.r_n - grid : Grid or array-like - dependent coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - length : float - length of the curve approximated by quadrature - - """ - transform = self._get_transforms(grid) - T = self.compute_coordinates(center, normal, r_n, grid, dt=1) - T = jnp.linalg.norm(T, axis=1) - theta = transform.grid.nodes[:, 2] - return jnp.trapz(T, theta) From b9bf5921f5cbaeea4e5bae0fa9e501a49f88ebf4 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 21:11:11 -0400 Subject: [PATCH 11/28] Fix norm axis in curve compute --- desc/compute/_curve.py | 28 +++++++++++++++------------- 1 file changed, 15 insertions(+), 13 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index fc4f9d63c6..a2fd809dee 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -2,6 +2,7 @@ from desc.geometry.utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec from .data_index import register_compute_fun +from .utils import cross, dot @register_compute_fun( @@ -516,7 +517,9 @@ def _r_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): parameterization="desc.geometry.Curve", ) def _frenet_tangent(params, transforms, profiles, data, **kwargs): - data["frenet_tangent"] = data["r_s"] / jnp.linalg.norm(data["r_s"])[:, None] + data["frenet_tangent"] = ( + data["r_s"] / jnp.linalg.norm(data["r_s"], axis=-1)[:, None] + ) return data @@ -535,7 +538,9 @@ def _frenet_tangent(params, transforms, profiles, data, **kwargs): parameterization="desc.geometry.Curve", ) def _frenet_normal(params, transforms, profiles, data, **kwargs): - data["frenet_normal"] = data["r_ss"] / jnp.linalg.norm(data["r_ss"])[:, None] + data["frenet_normal"] = ( + data["r_ss"] / jnp.linalg.norm(data["r_ss"], axis=-1)[:, None] + ) return data @@ -554,9 +559,9 @@ def _frenet_normal(params, transforms, profiles, data, **kwargs): parameterization="desc.geometry.Curve", ) def _frenet_binormal(params, transforms, profiles, data, **kwargs): - data["frenet_binormal"] = jnp.cross(data["T"], data["N"], axis=1) * jnp.linalg.det( - transforms["rotmat"] - ) + data["frenet_binormal"] = cross( + data["frenet_tangent"], data["frenet_normal"] + ) * jnp.linalg.det(transforms["rotmat"]) return data @@ -575,10 +580,8 @@ def _frenet_binormal(params, transforms, profiles, data, **kwargs): parameterization="desc.geometry.Curve", ) def _curvature(params, transforms, profiles, data, **kwargs): - dxn = jnp.linalg.norm(data["r_s"], axis=1)[:, jnp.newaxis] - data["curvature"] = jnp.linalg.norm( - jnp.cross(data["r_s"], data["r_ss"], axis=1) / dxn**3, axis=1 - ) + dxn = jnp.linalg.norm(data["r_s"], axis=-1)[:, jnp.newaxis] + data["curvature"] = jnp.linalg.norm(cross(data["r_s"], data["r_ss"]) / dxn**3) return data @@ -597,10 +600,9 @@ def _curvature(params, transforms, profiles, data, **kwargs): parameterization="desc.geometry.Curve", ) def _torsion(params, transforms, profiles, data, **kwargs): - dxd2x = jnp.cross(data["r_s"], data["r_ss"], axis=1) + dxd2x = cross(data["r_s"], data["r_ss"]) data["torsion"] = ( - jnp.sum(dxd2x * data["r_sss"], axis=1) - / jnp.linalg.norm(dxd2x, axis=1)[:, jnp.newaxis] ** 2 + dot(dxd2x, data["r_sss"]) / jnp.linalg.norm(dxd2x, axis=-1)[:, jnp.newaxis] ** 2 ) return data @@ -620,6 +622,6 @@ def _torsion(params, transforms, profiles, data, **kwargs): parameterization="desc.geometry.Curve", ) def _length(params, transforms, profiles, data, **kwargs): - T = jnp.linalg.norm(data["r_s"], axis=1) + T = jnp.linalg.norm(data["r_s"], axis=-1) data["length"] = jnp.trapz(T, data["s"]) return data From 4eb0fb6f683276cd65247e34bf3f2fa9b761c44a Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 21:11:38 -0400 Subject: [PATCH 12/28] Add compute function to curve base class, add tests --- desc/geometry/core.py | 70 +++++++++++++++++++++++++++++--------- tests/test_compute_funs.py | 16 +++++++++ 2 files changed, 70 insertions(+), 16 deletions(-) diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 3ef8338ff8..408cdcfcd3 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -5,6 +5,9 @@ import numpy as np from desc.backend import jnp +from desc.compute.utils import compute as compute_fun +from desc.compute.utils import get_params, get_transforms +from desc.grid import LinearGrid from desc.io import IOAble from .utils import reflection_matrix, rotation_matrix @@ -34,25 +37,60 @@ def name(self, new): def grid(self): """Grid: Nodes for computation.""" - @abstractmethod - def compute_coordinates(self, params=None, grid=None, dt=0): - """Compute real space coordinates on predefined grid.""" - - @abstractmethod - def compute_frenet_frame(self, params=None, grid=None): - """Compute Frenet frame on predefined grid.""" + def compute( + self, + names, + grid=None, + params=None, + transforms=None, + data=None, + **kwargs, + ): + """Compute the quantity given by name on grid. - @abstractmethod - def compute_curvature(self, params=None, grid=None): - """Compute curvature on predefined grid.""" + Parameters + ---------- + names : str or array-like of str + Name(s) of the quantity(s) to compute. + grid : Grid, optional + Grid of coordinates to evaluate at. Defaults to the Linear grid. + 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 - @abstractmethod - def compute_torsion(self, params=None, grid=None): - """Compute torsion on predefined grid.""" + Returns + ------- + data : dict of ndarray + Computed quantity and intermediate variables. - @abstractmethod - def compute_length(self, params=None, grid=None): - """Compute the length of the curve using specified nodes for quadrature.""" + """ + if isinstance(names, str): + names = [names] + if grid is None: + NFP = self.NFP if hasattr(self, "NFP") else 1 + grid = LinearGrid(N=2 * self.N + 5, NFP=NFP, endpoint=True) + + if params is None: + params = get_params(names, obj=self) + if transforms is None: + transforms = get_transforms(names, obj=self, grid=grid, **kwargs) + if data is None: + data = {} + profiles = {} + + data = compute_fun( + self, + names, + params=params, + transforms=transforms, + profiles=profiles, + data=data, + **kwargs, + ) + return data def translate(self, displacement=[0, 0, 0]): """Translate the curve by a rigid displacement in x, y, z.""" diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 6a7c29f98c..7684b2ee02 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -9,6 +9,7 @@ from desc.compute import data_index from desc.compute.utils import compress from desc.equilibrium import EquilibriaFamily, Equilibrium +from desc.geometry import FourierPlanarCurve, FourierRZCurve, FourierXYZCurve from desc.grid import LinearGrid, QuadratureGrid # convolve kernel is reverse of FD coeffs @@ -1120,6 +1121,21 @@ def test_compute_everything(): assert key in data +@pytest.mark.unit +def test_curve_compute_everything(): + """Make sure we can compute every curve thing without errors.""" + curves = { + "desc.geometry.curve.FourierXYZCurve": FourierXYZCurve(), + "desc.geometry.curve.FourierRZCurve": FourierRZCurve(), + "desc.geometry.curve.FourierPlanarCurve": FourierPlanarCurve(), + } + + for p, thing in curves.items(): + for key in data_index[p].keys(): + data = thing.compute(key) + assert key in data + + @pytest.mark.unit def test_compute_averages(): """Test that computing averages uses the correct grid.""" From d6be2e0b986f68e9b30e12a37f33922b534075b4 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 21:16:42 -0400 Subject: [PATCH 13/28] Rename curve position vector from r to x --- desc/compute/_curve.py | 97 +++++++++++++++++++++--------------------- 1 file changed, 48 insertions(+), 49 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index a2fd809dee..4ab7abcd5d 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -39,7 +39,7 @@ def _rotation_matrix_from_normal(normal): @register_compute_fun( - name="r", + name="x", label="\\mathbf{r}", units="m", units_long="meters", @@ -56,7 +56,7 @@ def _rotation_matrix_from_normal(normal): data=["s"], parameterization="desc.geometry.FourierPlanarCurve", ) -def _r_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): +def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): # create planar curve at z==0 r = transforms["r"].transform(params["r_n"], dz=0) Z = jnp.zeros_like(r) @@ -73,12 +73,12 @@ def _r_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] x, y, z = xyzcoords.T coords = xyz2rpz_vec(coords, x=x, y=y) - data["r"] = coords + data["x"] = coords return data @register_compute_fun( - name="r_s", + name="x_s", label="\\partial_{s} \\mathbf{r}", units="m", units_long="meters", @@ -95,7 +95,7 @@ def _r_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): data=["s"], parameterization="desc.geometry.FourierPlanarCurve", ) -def _r_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): +def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): r = transforms["r"].transform(params["r_n"], dz=0) dr = transforms["r"].transform(params["r_n"], dz=1) dX = dr * jnp.cos(data["s"]) - r * jnp.sin(data["s"]) @@ -114,12 +114,12 @@ def _r_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] x, y, z = xyzcoords.T coords = xyz2rpz_vec(coords, x=x, y=y) - data["r_s"] = coords + data["x_s"] = coords return data @register_compute_fun( - name="r_ss", + name="x_ss", label="\\partial_{ss} \\mathbf{r}", units="m", units_long="meters", @@ -136,7 +136,7 @@ def _r_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): data=["s"], parameterization="desc.geometry.FourierPlanarCurve", ) -def _r_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): +def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): r = transforms["r"].transform(params["r_n"], dz=0) dr = transforms["r"].transform(params["r_n"], dz=1) d2r = transforms["r"].transform(params["r_n"], dz=2) @@ -160,12 +160,12 @@ def _r_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] x, y, z = xyzcoords.T coords = xyz2rpz_vec(coords, x=x, y=y) - data["r_ss"] = coords + data["x_ss"] = coords return data @register_compute_fun( - name="r_sss", + name="x_sss", label="\\partial_{sss} \\mathbf{r}", units="m", units_long="meters", @@ -182,7 +182,7 @@ def _r_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): data=["s"], parameterization="desc.geometry.FourierPlanarCurve", ) -def _r_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): +def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): r = transforms["r"].transform(params["r_n"], dz=0) dr = transforms["r"].transform(params["r_n"], dz=1) d2r = transforms["r"].transform(params["r_n"], dz=2) @@ -213,12 +213,12 @@ def _r_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] x, y, z = xyzcoords.T coords = xyz2rpz_vec(coords, x=x, y=y) - data["r_sss"] = coords + data["x_sss"] = coords return data @register_compute_fun( - name="r", + name="x", label="\\mathbf{r}", units="m", units_long="meters", @@ -237,7 +237,7 @@ def _r_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierRZCurve", ) -def _r_FourierRZCurve(params, transforms, profiles, data, **kwargs): +def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_n"], dz=0) Z = transforms["Z"].transform(params["Z_n"], dz=0) phi = transforms["grid"].nodes[:, 2] @@ -247,12 +247,12 @@ def _r_FourierRZCurve(params, transforms, profiles, data, **kwargs): coords = coords @ transforms["rotmat"].T + transforms["shift"][jnp.newaxis, :] if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz(coords) - data["r"] = coords + data["x"] = coords return data @register_compute_fun( - name="r_s", + name="x_s", label="\\partial_{s} \\mathbf{r}", units="m", units_long="meters", @@ -270,8 +270,7 @@ def _r_FourierRZCurve(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierRZCurve", ) -def _r_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): - +def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_n"], dz=0) dR = transforms["R"].transform(params["R_n"], dz=1) dZ = transforms["Z"].transform(params["Z_n"], dz=1) @@ -282,12 +281,12 @@ def _r_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): coords = coords @ transforms["rotmat"].T if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_s"] = coords + data["x_s"] = coords return data @register_compute_fun( - name="r_ss", + name="x_ss", label="\\partial_{ss} \\mathbf{r}", units="m", units_long="meters", @@ -305,7 +304,7 @@ def _r_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierRZCurve", ) -def _r_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): +def _x_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_n"], dz=0) dR = transforms["R"].transform(params["R_n"], dz=1) d2R = transforms["R"].transform(params["R_n"], dz=2) @@ -320,12 +319,12 @@ def _r_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): coords = coords @ transforms["rotmat"].T if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_ss"] = coords + data["x_ss"] = coords return data @register_compute_fun( - name="r_sss", + name="x_sss", label="\\partial_{sss} \\mathbf{r}", units="m", units_long="meters", @@ -343,7 +342,7 @@ def _r_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierRZCurve", ) -def _r_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): +def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_n"], dz=0) dR = transforms["R"].transform(params["R_n"], dz=1) d2R = transforms["R"].transform(params["R_n"], dz=2) @@ -358,12 +357,12 @@ def _r_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): coords = coords @ transforms["rotmat"].T if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_sss"] = coords + data["x_sss"] = coords return data @register_compute_fun( - name="r", + name="x", label="\\mathbf{r}", units="m", units_long="meters", @@ -382,7 +381,7 @@ def _r_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierXYZCurve", ) -def _r_FourierXYZCurve(params, transforms, profiles, data, **kwargs): +def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs): X = transforms["X"].transform(params["X_n"], dz=0) Y = transforms["Y"].transform(params["Y_n"], dz=0) Z = transforms["Z"].transform(params["Z_n"], dz=0) @@ -390,12 +389,12 @@ def _r_FourierXYZCurve(params, transforms, profiles, data, **kwargs): coords = coords @ transforms["rotmat"].T + transforms["shift"][jnp.newaxis, :] if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz(coords) - data["r"] = coords + data["x"] = coords return data @register_compute_fun( - name="r_s", + name="x_s", label="\\partial_{s} \\mathbf{r}", units="m", units_long="meters", @@ -414,7 +413,7 @@ def _r_FourierXYZCurve(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierXYZCurve", ) -def _r_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): +def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): X = transforms["X"].transform(params["X_n"], dz=1) Y = transforms["Y"].transform(params["Y_n"], dz=1) Z = transforms["Z"].transform(params["Z_n"], dz=1) @@ -426,12 +425,12 @@ def _r_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): x=coords[:, 1] + transforms["shift"][0], y=coords[:, 1] + transforms["shift"][1], ) - data["r_s"] = coords + data["x_s"] = coords return data @register_compute_fun( - name="r_ss", + name="x_ss", label="\\partial_{ss} \\mathbf{r}", units="m", units_long="meters", @@ -450,7 +449,7 @@ def _r_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierXYZCurve", ) -def _r_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): +def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): X = transforms["X"].transform(params["X_n"], dz=2) Y = transforms["Y"].transform(params["Y_n"], dz=2) Z = transforms["Z"].transform(params["Z_n"], dz=2) @@ -462,12 +461,12 @@ def _r_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): x=coords[:, 1] + transforms["shift"][0], y=coords[:, 1] + transforms["shift"][1], ) - data["r_ss"] = coords + data["x_ss"] = coords return data @register_compute_fun( - name="r_sss", + name="x_sss", label="\\partial_{sss} \\mathbf{r}", units="m", units_long="meters", @@ -486,7 +485,7 @@ def _r_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierXYZCurve", ) -def _r_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): +def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): X = transforms["X"].transform(params["X_n"], dz=3) Y = transforms["Y"].transform(params["Y_n"], dz=3) Z = transforms["Z"].transform(params["Z_n"], dz=3) @@ -498,7 +497,7 @@ def _r_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): x=coords[:, 1] + transforms["shift"][0], y=coords[:, 1] + transforms["shift"][1], ) - data["r_sss"] = coords + data["x_sss"] = coords return data @@ -513,12 +512,12 @@ def _r_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="s", - data=["r_s"], + data=["x_s"], parameterization="desc.geometry.Curve", ) def _frenet_tangent(params, transforms, profiles, data, **kwargs): data["frenet_tangent"] = ( - data["r_s"] / jnp.linalg.norm(data["r_s"], axis=-1)[:, None] + data["x_s"] / jnp.linalg.norm(data["x_s"], axis=-1)[:, None] ) return data @@ -534,12 +533,12 @@ def _frenet_tangent(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="s", - data=["r_ss"], + data=["x_ss"], parameterization="desc.geometry.Curve", ) def _frenet_normal(params, transforms, profiles, data, **kwargs): data["frenet_normal"] = ( - data["r_ss"] / jnp.linalg.norm(data["r_ss"], axis=-1)[:, None] + data["x_ss"] / jnp.linalg.norm(data["x_ss"], axis=-1)[:, None] ) return data @@ -576,12 +575,12 @@ def _frenet_binormal(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="s", - data=["r_s", "r_ss"], + data=["x_s", "x_ss"], parameterization="desc.geometry.Curve", ) def _curvature(params, transforms, profiles, data, **kwargs): - dxn = jnp.linalg.norm(data["r_s"], axis=-1)[:, jnp.newaxis] - data["curvature"] = jnp.linalg.norm(cross(data["r_s"], data["r_ss"]) / dxn**3) + dxn = jnp.linalg.norm(data["x_s"], axis=-1)[:, jnp.newaxis] + data["curvature"] = jnp.linalg.norm(cross(data["x_s"], data["x_ss"]) / dxn**3) return data @@ -596,13 +595,13 @@ def _curvature(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="s", - data=["r_s", "r_ss", "r_sss"], + data=["x_s", "x_ss", "x_sss"], parameterization="desc.geometry.Curve", ) def _torsion(params, transforms, profiles, data, **kwargs): - dxd2x = cross(data["r_s"], data["r_ss"]) + dxd2x = cross(data["x_s"], data["x_ss"]) data["torsion"] = ( - dot(dxd2x, data["r_sss"]) / jnp.linalg.norm(dxd2x, axis=-1)[:, jnp.newaxis] ** 2 + dot(dxd2x, data["x_sss"]) / jnp.linalg.norm(dxd2x, axis=-1)[:, jnp.newaxis] ** 2 ) return data @@ -618,10 +617,10 @@ def _torsion(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="s", - data=["s", "r_s"], + data=["s", "x_s"], parameterization="desc.geometry.Curve", ) def _length(params, transforms, profiles, data, **kwargs): - T = jnp.linalg.norm(data["r_s"], axis=-1) + T = jnp.linalg.norm(data["x_s"], axis=-1) data["length"] = jnp.trapz(T, data["s"]) return data From 67fc53f500c7a2b6f277b2b2b9b19c77791c0ef0 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 22:21:57 -0400 Subject: [PATCH 14/28] Add basis to curve compute, fix basis conversion --- desc/compute/_curve.py | 74 +++++++++++++++++++++++------------------- desc/compute/utils.py | 2 +- 2 files changed, 42 insertions(+), 34 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 4ab7abcd5d..c36f09e424 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -55,6 +55,7 @@ def _rotation_matrix_from_normal(normal): coordinates="s", data=["s"], parameterization="desc.geometry.FourierPlanarCurve", + basis="basis", ) def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): # create planar curve at z==0 @@ -68,11 +69,7 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coords = jnp.matmul(coords, R.T) + params["center"] coords = jnp.matmul(coords, transforms["rotmat"].T) + transforms["shift"] if kwargs.get("basis", "rpz").lower() == "rpz": - xyzcoords = jnp.array([X, Y, Z]).T - xyzcoords = jnp.matmul(xyzcoords, R.T) + params["center"] - xyzcoords = jnp.matmul(xyzcoords, transforms["rotmat"].T) + transforms["shift"] - x, y, z = xyzcoords.T - coords = xyz2rpz_vec(coords, x=x, y=y) + coords = xyz2rpz(coords) data["x"] = coords return data @@ -94,6 +91,7 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=["s"], parameterization="desc.geometry.FourierPlanarCurve", + basis="basis", ) def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): r = transforms["r"].transform(params["r_n"], dz=0) @@ -135,6 +133,7 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=["s"], parameterization="desc.geometry.FourierPlanarCurve", + basis="basis", ) def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): r = transforms["r"].transform(params["r_n"], dz=0) @@ -181,6 +180,7 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=["s"], parameterization="desc.geometry.FourierPlanarCurve", + basis="basis", ) def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): r = transforms["r"].transform(params["r_n"], dz=0) @@ -236,6 +236,7 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=[], parameterization="desc.geometry.FourierRZCurve", + basis="basis", ) def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_n"], dz=0) @@ -269,6 +270,7 @@ def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=[], parameterization="desc.geometry.FourierRZCurve", + basis="basis", ) def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_n"], dz=0) @@ -303,6 +305,7 @@ def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=[], parameterization="desc.geometry.FourierRZCurve", + basis="basis", ) def _x_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_n"], dz=0) @@ -341,6 +344,7 @@ def _x_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=[], parameterization="desc.geometry.FourierRZCurve", + basis="basis", ) def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_n"], dz=0) @@ -380,6 +384,7 @@ def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=[], parameterization="desc.geometry.FourierXYZCurve", + basis="basis", ) def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs): X = transforms["X"].transform(params["X_n"], dz=0) @@ -402,8 +407,8 @@ def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs): dim=3, params=["X_n", "Y_n", "Z_n"], transforms={ - "X": [[0, 0, 1]], - "Y": [[0, 0, 1]], + "X": [[0, 0, 0], [0, 0, 1]], + "Y": [[0, 0, 0], [0, 0, 1]], "Z": [[0, 0, 1]], "rotmat": [], "shift": [], @@ -412,18 +417,19 @@ def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=[], parameterization="desc.geometry.FourierXYZCurve", + basis="basis", ) def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): - X = transforms["X"].transform(params["X_n"], dz=1) - Y = transforms["Y"].transform(params["Y_n"], dz=1) - Z = transforms["Z"].transform(params["Z_n"], dz=1) - coords = jnp.stack([X, Y, Z], axis=1) + dX = transforms["X"].transform(params["X_n"], dz=1) + dY = transforms["Y"].transform(params["Y_n"], dz=1) + dZ = transforms["Z"].transform(params["Z_n"], dz=1) + coords = jnp.stack([dX, dY, dZ], axis=1) coords = coords @ transforms["rotmat"].T if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec( coords, - x=coords[:, 1] + transforms["shift"][0], - y=coords[:, 1] + transforms["shift"][1], + x=transforms["X"].transform(params["X_n"]) + transforms["shift"][0], + y=transforms["Y"].transform(params["Y_n"]) + transforms["shift"][1], ) data["x_s"] = coords return data @@ -438,8 +444,8 @@ def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): dim=3, params=["X_n", "Y_n", "Z_n"], transforms={ - "X": [[0, 0, 2]], - "Y": [[0, 0, 2]], + "X": [[0, 0, 0], [0, 0, 2]], + "Y": [[0, 0, 0], [0, 0, 2]], "Z": [[0, 0, 2]], "rotmat": [], "shift": [], @@ -448,18 +454,19 @@ def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=[], parameterization="desc.geometry.FourierXYZCurve", + basis="basis", ) def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): - X = transforms["X"].transform(params["X_n"], dz=2) - Y = transforms["Y"].transform(params["Y_n"], dz=2) - Z = transforms["Z"].transform(params["Z_n"], dz=2) - coords = jnp.stack([X, Y, Z], axis=1) + d2X = transforms["X"].transform(params["X_n"], dz=2) + d2Y = transforms["Y"].transform(params["Y_n"], dz=2) + d2Z = transforms["Z"].transform(params["Z_n"], dz=2) + coords = jnp.stack([d2X, d2Y, d2Z], axis=1) coords = coords @ transforms["rotmat"].T if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec( coords, - x=coords[:, 1] + transforms["shift"][0], - y=coords[:, 1] + transforms["shift"][1], + x=transforms["X"].transform(params["X_n"]) + transforms["shift"][0], + y=transforms["Y"].transform(params["Y_n"]) + transforms["shift"][1], ) data["x_ss"] = coords return data @@ -474,8 +481,8 @@ def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): dim=3, params=["X_n", "Y_n", "Z_n"], transforms={ - "X": [[0, 0, 3]], - "Y": [[0, 0, 3]], + "X": [[0, 0, 0], [0, 0, 3]], + "Y": [[0, 0, 0], [0, 0, 3]], "Z": [[0, 0, 3]], "rotmat": [], "shift": [], @@ -484,18 +491,19 @@ def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): coordinates="s", data=[], parameterization="desc.geometry.FourierXYZCurve", + basis="basis", ) def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): - X = transforms["X"].transform(params["X_n"], dz=3) - Y = transforms["Y"].transform(params["Y_n"], dz=3) - Z = transforms["Z"].transform(params["Z_n"], dz=3) - coords = jnp.stack([X, Y, Z], axis=1) + d3X = transforms["X"].transform(params["X_n"], dz=3) + d3Y = transforms["Y"].transform(params["Y_n"], dz=3) + d3Z = transforms["Z"].transform(params["Z_n"], dz=3) + coords = jnp.stack([d3X, d3Y, d3Z], axis=1) coords = coords @ transforms["rotmat"].T if kwargs.get("basis", "rpz").lower() == "rpz": coords = xyz2rpz_vec( coords, - x=coords[:, 1] + transforms["shift"][0], - y=coords[:, 1] + transforms["shift"][1], + x=transforms["X"].transform(params["X_n"]) + transforms["shift"][0], + y=transforms["Y"].transform(params["Y_n"]) + transforms["shift"][1], ) data["x_sss"] = coords return data @@ -580,7 +588,9 @@ def _frenet_binormal(params, transforms, profiles, data, **kwargs): ) def _curvature(params, transforms, profiles, data, **kwargs): dxn = jnp.linalg.norm(data["x_s"], axis=-1)[:, jnp.newaxis] - data["curvature"] = jnp.linalg.norm(cross(data["x_s"], data["x_ss"]) / dxn**3) + data["curvature"] = jnp.linalg.norm( + cross(data["x_s"], data["x_ss"]) / dxn**3, axis=-1 + ) return data @@ -600,9 +610,7 @@ def _curvature(params, transforms, profiles, data, **kwargs): ) def _torsion(params, transforms, profiles, data, **kwargs): dxd2x = cross(data["x_s"], data["x_ss"]) - data["torsion"] = ( - dot(dxd2x, data["x_sss"]) / jnp.linalg.norm(dxd2x, axis=-1)[:, jnp.newaxis] ** 2 - ) + data["torsion"] = dot(dxd2x, data["x_sss"]) / jnp.linalg.norm(dxd2x, axis=-1) ** 2 return data diff --git a/desc/compute/utils.py b/desc/compute/utils.py index b80e167a89..3ce5a0342c 100644 --- a/desc/compute/utils.py +++ b/desc/compute/utils.py @@ -86,7 +86,7 @@ def compute(parameterization, names, params, transforms, profiles, data=None, ** for name in names: if name not in data_index[p]: raise ValueError("Unrecognized value '{}'.".format(name)) - allowed_kwargs = {"helicity", "M_booz", "N_booz", "gamma"} + allowed_kwargs = {"helicity", "M_booz", "N_booz", "gamma", "basis"} bad_kwargs = kwargs.keys() - allowed_kwargs if len(bad_kwargs) > 0: raise ValueError(f"Unrecognized argument(s): {bad_kwargs}") From e44e70d33458b2313eb41a6fba47a028de773dc3 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 22:24:09 -0400 Subject: [PATCH 15/28] Finish updating curve classes --- desc/geometry/core.py | 12 ++++++------ desc/geometry/curve.py | 43 ++++++++++++++++++++++-------------------- 2 files changed, 29 insertions(+), 26 deletions(-) diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 408cdcfcd3..a4339ca516 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -1,5 +1,6 @@ """Base classes for curves and surfaces.""" +import numbers from abc import ABC, abstractmethod import numpy as np @@ -32,11 +33,6 @@ def name(self): def name(self, new): self._name = new - @property - @abstractmethod - def grid(self): - """Grid: Nodes for computation.""" - def compute( self, names, @@ -52,8 +48,9 @@ def compute( ---------- names : str or array-like of str Name(s) of the quantity(s) to compute. - grid : Grid, optional + grid : Grid or int, optional Grid of coordinates to evaluate at. Defaults to the 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 @@ -72,6 +69,9 @@ def compute( if grid is None: NFP = self.NFP if hasattr(self, "NFP") else 1 grid = LinearGrid(N=2 * self.N + 5, NFP=NFP, endpoint=True) + if isinstance(grid, numbers.Integral): + NFP = self.NFP if hasattr(self, "NFP") else 1 + grid = LinearGrid(N=grid, NFP=NFP, endpoint=True) if params is None: params = get_params(names, obj=self) diff --git a/desc/geometry/curve.py b/desc/geometry/curve.py index 88a7c0f4b3..2ab0171695 100644 --- a/desc/geometry/curve.py +++ b/desc/geometry/curve.py @@ -143,11 +143,6 @@ def change_resolution(self, N=None, NFP=None, sym=None): self.Z_basis.change_resolution( N=N, NFP=self.NFP, sym="sin" if self.sym else self.sym ) - if hasattr(self.grid, "change_resolution"): - self.grid.change_resolution( - self.grid.L, self.grid.M, self.grid.N, self.NFP - ) - self._R_transform, self._Z_transform = self._get_transforms(self.grid) self.R_n = copy_coeffs(self.R_n, R_modes_old, self.R_basis.modes) self.Z_n = copy_coeffs(self.Z_n, Z_modes_old, self.Z_basis.modes) @@ -226,7 +221,7 @@ def to_FourierXYZCurve(self, N=None): N = max(self.R_basis.N, self.Z_basis.N) grid = LinearGrid(N=4 * N, NFP=1, sym=False) basis = FourierSeries(N=N, NFP=1, sym=False) - xyz = self.compute_coordinates(grid=grid, basis="xyz") + xyz = self.compute("x", grid=grid, basis="xyz")["x"] transform = Transform(grid, basis, build_pinv=True) X_n = transform.fit(xyz[:, 0]) Y_n = transform.fit(xyz[:, 1]) @@ -305,14 +300,15 @@ def N(self): def change_resolution(self, N=None): """Change the maximum angular resolution.""" if (N is not None) and (N != self.N): - modes_old = self.basis.modes + Xmodes_old = self.X_basis.modes + Ymodes_old = self.Y_basis.modes + Zmodes_old = self.Z_basis.modes self.X_basis.change_resolution(N=N) self.Y_basis.change_resolution(N=N) self.Z_basis.change_resolution(N=N) - self._transform = self._get_transforms(self.grid) - self.X_n = copy_coeffs(self.X_n, modes_old, self.X_basis.modes) - self.Y_n = copy_coeffs(self.Y_n, modes_old, self.Y_basis.modes) - self.Z_n = copy_coeffs(self.Z_n, modes_old, self.Z_basis.modes) + self.X_n = copy_coeffs(self.X_n, Xmodes_old, self.X_basis.modes) + self.Y_n = copy_coeffs(self.Y_n, Ymodes_old, self.Y_basis.modes) + self.Z_n = copy_coeffs(self.Z_n, Zmodes_old, self.Z_basis.modes) def get_coeffs(self, n): """Get Fourier coefficients for given mode number(s).""" @@ -321,11 +317,13 @@ def get_coeffs(self, n): Y = np.zeros_like(n).astype(float) Z = np.zeros_like(n).astype(float) - idx = np.where(n[:, np.newaxis] == self.basis.modes[:, 2]) + Xidx = np.where(n[:, np.newaxis] == self.X_basis.modes[:, 2]) + Yidx = np.where(n[:, np.newaxis] == self.Y_basis.modes[:, 2]) + Zidx = np.where(n[:, np.newaxis] == self.Z_basis.modes[:, 2]) - X[idx[0]] = self.X_n[idx[1]] - Y[idx[0]] = self.Y_n[idx[1]] - Z[idx[0]] = self.Z_n[idx[1]] + X[Xidx[0]] = self.X_n[Xidx[1]] + Y[Yidx[0]] = self.Y_n[Yidx[1]] + Z[Zidx[0]] = self.Z_n[Zidx[1]] return X, Y, Z def set_coeffs(self, n, X=None, Y=None, Z=None): @@ -339,12 +337,18 @@ def set_coeffs(self, n, X=None, Y=None, Z=None): X = np.broadcast_to(X, n.shape) Y = np.broadcast_to(Y, n.shape) Z = np.broadcast_to(Z, n.shape) - for nn, XX, YY, ZZ in zip(n, X, Y, Z): - idx = self.basis.get_idx(0, 0, nn) + for nn, XX in zip(n, X): + idx = self.X_basis.get_idx(0, 0, nn) if XX is not None: self.X_n = put(self.X_n, idx, XX) + + for nn, YY in zip(n, Y): + idx = self.Y_basis.get_idx(0, 0, nn) if YY is not None: self.Y_n = put(self.Y_n, idx, YY) + + for nn, ZZ in zip(n, Z): + idx = self.Z_basis.get_idx(0, 0, nn) if ZZ is not None: self.Z_n = put(self.Z_n, idx, ZZ) @@ -466,7 +470,6 @@ def change_resolution(self, N=None): if (N is not None) and (N != self.N): modes_old = self.r_basis.modes self.r_basis.change_resolution(N=N) - self._transform = self._get_transforms(self.grid) self.r_n = copy_coeffs(self.r_n, modes_old, self.r_basis.modes) @property @@ -517,7 +520,7 @@ def get_coeffs(self, n): n = np.atleast_1d(n).astype(int) r = np.zeros_like(n).astype(float) - idx = np.where(n[:, np.newaxis] == self.basis.modes[:, 2]) + idx = np.where(n[:, np.newaxis] == self.r_basis.modes[:, 2]) r[idx[0]] = self.r_n[idx[1]] return r @@ -527,6 +530,6 @@ def set_coeffs(self, n, r=None): n, r = np.atleast_1d(n), np.atleast_1d(r) r = np.broadcast_to(r, n.shape) for nn, rr in zip(n, r): - idx = self.basis.get_idx(0, 0, nn) + idx = self.r_basis.get_idx(0, 0, nn) if rr is not None: self.r_n = put(self.r_n, idx, rr) From d27ebdeb98b1f05421d5ef0f2b830974aca6ac79 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Thu, 13 Jul 2023 22:24:21 -0400 Subject: [PATCH 16/28] Update curve tests with new API --- tests/test_curves.py | 153 ++++++++++++++++++++++--------------------- 1 file changed, 80 insertions(+), 73 deletions(-) diff --git a/tests/test_curves.py b/tests/test_curves.py index bfc8ebb5cd..1e4e0058ec 100644 --- a/tests/test_curves.py +++ b/tests/test_curves.py @@ -14,46 +14,54 @@ class TestRZCurve: def test_length(self): """Test length of circular curve.""" c = FourierRZCurve() - np.testing.assert_allclose(c.compute_length(grid=20), 10 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 10 * 2 * np.pi + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_length(grid=20), 10 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 10 * 2 * np.pi + ) @pytest.mark.unit def test_curvature(self): """Test curvature of circular curve.""" c = FourierRZCurve() - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 10) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 10) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 10) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 10) @pytest.mark.unit def test_torsion(self): """Test torsion of circular curve.""" c = FourierRZCurve() - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose(c.compute("torsion", grid=20)["torsion"], 0) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose(c.compute("torsion", grid=20)["torsion"], 0) @pytest.mark.unit def test_frenet(self): """Test frenet-seret frame of circular curve.""" c = FourierRZCurve() - c.grid = 0 - T, N, B = c.compute_frenet_frame(basis="rpz") + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="rpz", grid=0 + ) + T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] np.testing.assert_allclose(T, np.array([[0, 1, 0]]), atol=1e-12) np.testing.assert_allclose(N, np.array([[-1, 0, 0]]), atol=1e-12) np.testing.assert_allclose(B, np.array([[0, 0, 1]]), atol=1e-12) c.rotate(angle=np.pi) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - c.grid = np.array([[0, 0, 0]]) - T, N, B = c.compute_frenet_frame(basis="xyz") + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="xyz", grid=0 + ) + T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] np.testing.assert_allclose(T, np.array([[0, 1, 0]]), atol=1e-12) np.testing.assert_allclose(N, np.array([[1, 0, 0]]), atol=1e-12) np.testing.assert_allclose(B, np.array([[0, 0, 1]]), atol=1e-12) @@ -62,14 +70,14 @@ def test_frenet(self): def test_coords(self): """Test lab frame coordinates of circular curve.""" c = FourierRZCurve() - x, y, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz").T + x, y, z = c.compute("x", grid=0, basis="xyz")["x"].T np.testing.assert_allclose(x, 10) np.testing.assert_allclose(y, 0) np.testing.assert_allclose(z, 0) c.rotate(angle=np.pi / 2) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - r, p, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="rpz").T + r, p, z = c.compute("x", grid=0, basis="rpz")["x"].T np.testing.assert_allclose(r, np.sqrt(1**2 + 9**2)) np.testing.assert_allclose(p, np.arctan2(-9, 1)) np.testing.assert_allclose(z, 1) @@ -78,9 +86,6 @@ def test_coords(self): def test_misc(self): """Test getting/setting misc attributes of FourierRZCurve.""" c = FourierRZCurve() - grid = LinearGrid(M=2, N=2) - c.grid = grid - assert grid.eq(c.grid) R, Z = c.get_coeffs(0) np.testing.assert_allclose(R, 10) @@ -121,18 +126,12 @@ def test_misc(self): assert c.NFP == 3 assert c.R_basis.NFP == 3 assert c.Z_basis.NFP == 3 - assert c.grid.NFP == 3 @pytest.mark.unit def test_asserts(self): """Test error checking when creating FourierRZCurve.""" with pytest.raises(ValueError): - c = FourierRZCurve(R_n=[]) - c = FourierRZCurve() - with pytest.raises(NotImplementedError): - c.compute_coordinates(dt=4) - with pytest.raises(TypeError): - c.grid = [1, 2, 3] + _ = FourierRZCurve(R_n=[]) @pytest.mark.unit def test_to_FourierXYZCurve(self): @@ -140,18 +139,23 @@ def test_to_FourierXYZCurve(self): rz = FourierRZCurve(R_n=[0, 10, 1], Z_n=[-1, 0, 0]) xyz = rz.to_FourierXYZCurve(N=2) + grid = LinearGrid(N=20, endpoint=True) + np.testing.assert_allclose( - rz.compute_curvature(), xyz.compute_curvature(grid=rz.grid) + rz.compute("curvature", grid=grid)["curvature"], + xyz.compute("curvature", grid=grid)["curvature"], ) np.testing.assert_allclose( - rz.compute_torsion(), xyz.compute_torsion(grid=rz.grid) + rz.compute("torsion", grid=grid)["torsion"], + xyz.compute("torsion", grid=grid)["torsion"], ) np.testing.assert_allclose( - rz.compute_length(), xyz.compute_length(grid=rz.grid) + rz.compute("length", grid=grid)["length"], + xyz.compute("length", grid=grid)["length"], ) np.testing.assert_allclose( - rz.compute_coordinates(basis="rpz"), - xyz.compute_coordinates(basis="rpz", grid=rz.grid), + rz.compute("x", grid=grid, basis="rpz")["x"], + xyz.compute("x", basis="rpz", grid=grid)["x"], atol=1e-12, ) @@ -163,46 +167,58 @@ class TestXYZCurve: def test_length(self): """Test length of circular curve.""" c = FourierXYZCurve() - np.testing.assert_allclose(c.compute_length(grid=20), 2 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 2 * 2 * np.pi + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_length(grid=20), 2 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 2 * 2 * np.pi + ) @pytest.mark.unit def test_curvature(self): """Test curvature of circular curve.""" c = FourierXYZCurve() - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 2) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 2) @pytest.mark.unit def test_torsion(self): """Test torsion of circular curve.""" c = FourierXYZCurve(modes=[-1, 0, 1]) - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose( + c.compute("torsion", grid=20)["torsion"], 0, atol=1e-12 + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose( + c.compute("torsion", grid=20)["torsion"], 0, atol=1e-12 + ) @pytest.mark.unit def test_frenet(self): """Test frenet-seret frame of circular curve.""" c = FourierXYZCurve() - c.grid = 0 - T, N, B = c.compute_frenet_frame(basis="rpz") + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="rpz", grid=0 + ) + T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] np.testing.assert_allclose(T, np.array([[0, 0, -1]]), atol=1e-12) np.testing.assert_allclose(N, np.array([[-1, 0, 0]]), atol=1e-12) np.testing.assert_allclose(B, np.array([[0, 1, 0]]), atol=1e-12) c.rotate(angle=np.pi) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - c.grid = np.array([0, 0, 0]) - T, N, B = c.compute_frenet_frame(basis="xyz") + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="xyz", grid=0 + ) + T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] np.testing.assert_allclose(T, np.array([[0, 0, -1]]), atol=1e-12) np.testing.assert_allclose(N, np.array([[1, 0, 0]]), atol=1e-12) np.testing.assert_allclose(B, np.array([[0, 1, 0]]), atol=1e-12) @@ -211,14 +227,14 @@ def test_frenet(self): def test_coords(self): """Test lab frame coordinates of circular curve.""" c = FourierXYZCurve() - x, y, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz").T + x, y, z = c.compute("x", grid=0, basis="xyz")["x"].T np.testing.assert_allclose(x, 12) np.testing.assert_allclose(y, 0) np.testing.assert_allclose(z, 0) c.rotate(angle=np.pi / 2) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - r, p, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="rpz").T + r, p, z = c.compute("x", grid=0, basis="rpz")["x"].T np.testing.assert_allclose(r, np.sqrt(1**2 + 11**2)) np.testing.assert_allclose(p, np.arctan2(-11, 1)) np.testing.assert_allclose(z, 1) @@ -227,9 +243,6 @@ def test_coords(self): def test_misc(self): """Test getting/setting misc attributes of FourierXYZCurve.""" c = FourierXYZCurve() - grid = LinearGrid(M=2, N=2) - c.grid = grid - assert grid.eq(c.grid) X, Y, Z = c.get_coeffs(0) np.testing.assert_allclose(X, 10) @@ -252,15 +265,6 @@ def test_misc(self): with pytest.raises(ValueError): c.Z_n = s.Z_n - @pytest.mark.unit - def test_asserts(self): - """Test error checking when creating FourierXYZCurve.""" - c = FourierXYZCurve() - with pytest.raises(KeyError): - c.compute_coordinates(dt=4) - with pytest.raises(TypeError): - c.grid = [1, 2, 3] - class TestPlanarCurve: """Tests for FourierPlanarCurve class.""" @@ -269,46 +273,58 @@ class TestPlanarCurve: def test_length(self): """Test length of circular curve.""" c = FourierPlanarCurve(modes=[0]) - np.testing.assert_allclose(c.compute_length(grid=20), 2 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 2 * 2 * np.pi + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_length(grid=20), 2 * 2 * np.pi) + np.testing.assert_allclose( + c.compute("length", grid=20)["length"], 2 * 2 * np.pi + ) @pytest.mark.unit def test_curvature(self): """Test curvature of circular curve.""" c = FourierPlanarCurve() - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 2) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_curvature(grid=20), 1 / 2) + np.testing.assert_allclose(c.compute("curvature", grid=20)["curvature"], 1 / 2) @pytest.mark.unit def test_torsion(self): """Test torsion of circular curve.""" c = FourierPlanarCurve() - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose( + c.compute("torsion", grid=20)["torsion"], 0, atol=1e-12 + ) c.translate([1, 1, 1]) c.rotate(angle=np.pi) c.flip([0, 1, 0]) - np.testing.assert_allclose(c.compute_torsion(grid=20), 0) + np.testing.assert_allclose( + c.compute("torsion", grid=20)["torsion"], 0, atol=1e-12 + ) @pytest.mark.unit def test_frenet(self): """Test frenet-seret frame of circular curve.""" c = FourierPlanarCurve() - c.grid = 0 - T, N, B = c.compute_frenet_frame(basis="xyz") + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="xyz", grid=0 + ) + T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] np.testing.assert_allclose(T, np.array([[0, 0, -1]]), atol=1e-12) np.testing.assert_allclose(N, np.array([[-1, 0, 0]]), atol=1e-12) np.testing.assert_allclose(B, np.array([[0, 1, 0]]), atol=1e-12) c.rotate(angle=np.pi) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - c.grid = np.array([0, 0, 0]) - T, N, B = c.compute_frenet_frame(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz") + data = c.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], basis="xyz", grid=0 + ) + T, N, B = data["frenet_tangent"], data["frenet_normal"], data["frenet_binormal"] np.testing.assert_allclose(T, np.array([[0, 0, -1]]), atol=1e-12) np.testing.assert_allclose(N, np.array([[1, 0, 0]]), atol=1e-12) np.testing.assert_allclose(B, np.array([[0, 1, 0]]), atol=1e-12) @@ -317,20 +333,18 @@ def test_frenet(self): def test_coords(self): """Test lab frame coordinates of circular curve.""" c = FourierPlanarCurve() - r, p, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="rpz").T + r, p, z = c.compute("x", grid=0, basis="rpz")["x"].T np.testing.assert_allclose(r, 12) np.testing.assert_allclose(p, 0) np.testing.assert_allclose(z, 0) - dr, dp, dz = c.compute_coordinates( - grid=np.array([[0.0, 0.0, 0.0]]), dt=3, basis="rpz" - ).T + dr, dp, dz = c.compute("x_sss", grid=0, basis="rpz")["x_sss"].T np.testing.assert_allclose(dr, 0) np.testing.assert_allclose(dp, 0) np.testing.assert_allclose(dz, 2) c.rotate(angle=np.pi / 2) c.flip([0, 1, 0]) c.translate([1, 1, 1]) - x, y, z = c.compute_coordinates(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz").T + x, y, z = c.compute("x", grid=0, basis="xyz")["x"].T np.testing.assert_allclose(x, 1) np.testing.assert_allclose(y, -11) np.testing.assert_allclose(z, 1) @@ -339,9 +353,6 @@ def test_coords(self): def test_misc(self): """Test getting/setting misc attributes of FourierPlanarCurve.""" c = FourierPlanarCurve() - grid = LinearGrid(M=2, N=2) - c.grid = grid - assert grid.eq(c.grid) r = c.get_coeffs(0) np.testing.assert_allclose(r, 2) @@ -369,10 +380,6 @@ def test_misc(self): def test_asserts(self): """Test error checking when creating FourierPlanarCurve.""" c = FourierPlanarCurve() - with pytest.raises(NotImplementedError): - c.compute_coordinates(dt=4) - with pytest.raises(TypeError): - c.grid = [1, 2, 3] with pytest.raises(ValueError): c.center = [4] with pytest.raises(ValueError): From ed1562920bd846c5320960a51327a5c3312097a1 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 15 Jul 2023 14:50:40 -0400 Subject: [PATCH 17/28] Add parameterization info to misc geometry compute funs --- desc/compute/_geometry.py | 95 +++++++++++++++++++++++++++++++++++++++ desc/compute/_metric.py | 44 ++++++++++++++++++ 2 files changed, 139 insertions(+) diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index 7116c0b4f7..6a31cb6f52 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -22,6 +22,34 @@ def _V(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="V", + label="V", + units="m^{3}", + units_long="cubic meters", + description="Volume", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["e_theta", "e_zeta", "x"], + parameterization="desc.geometry.surface.FourierRZToroidalSurface", +) +def _V_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + # divergence theorem: integral(dV div [0, 0, Z]) = integral(dS dot [0, 0, Z]) + data["V"] = jnp.max( # take max in case there are multiple surfaces for some reason + jnp.abs( + surface_integrals( + transforms["grid"], + cross(data["e_theta"], data["e_zeta"])[:, 2] * data["x"][:, 2], + expand_out=False, + ) + ) + ) + return data + + @register_compute_fun( name="V(r)", label="V(\\rho)", @@ -98,6 +126,10 @@ def _V_rr_of_r(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="", data=["|e_rho x e_theta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _A(params, transforms, profiles, data, **kwargs): data["A"] = jnp.mean( @@ -111,6 +143,9 @@ def _A(params, transforms, profiles, data, **kwargs): return data +# TODO: compute cross section area for toroidal surface using stokes? + + @register_compute_fun( name="S(r)", label="S(\\rho)", @@ -123,6 +158,10 @@ def _A(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="r", data=["|e_theta x e_zeta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _S_of_r(params, transforms, profiles, data, **kwargs): data["S(r)"] = surface_integrals(transforms["grid"], data["|e_theta x e_zeta|"]) @@ -245,6 +284,10 @@ def _a_major_over_a_minor(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_rho", "e_theta_t"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _L_sff_rho(params, transforms, profiles, data, **kwargs): data["L_sff_rho"] = dot(data["e_theta_t"], data["n_rho"]) @@ -263,6 +306,10 @@ def _L_sff_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_rho", "e_theta_z"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _M_sff_rho(params, transforms, profiles, data, **kwargs): data["M_sff_rho"] = dot(data["e_theta_z"], data["n_rho"]) @@ -281,6 +328,10 @@ def _M_sff_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_rho", "e_zeta_z"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _N_sff_rho(params, transforms, profiles, data, **kwargs): data["N_sff_rho"] = dot(data["e_zeta_z"], data["n_rho"]) @@ -299,6 +350,10 @@ def _N_sff_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["g_tt", "g_tz", "g_zz", "L_sff_rho", "M_sff_rho", "N_sff_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _curvature_k1_rho(params, transforms, profiles, data, **kwargs): # following notation from @@ -331,6 +386,10 @@ def _curvature_k1_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["g_tt", "g_tz", "g_zz", "L_sff_rho", "M_sff_rho", "N_sff_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _curvature_k2_rho(params, transforms, profiles, data, **kwargs): # following notation from @@ -363,6 +422,10 @@ def _curvature_k2_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["curvature_k1_rho", "curvature_k2_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _curvature_K_rho(params, transforms, profiles, data, **kwargs): data["curvature_K_rho"] = data["curvature_k1_rho"] * data["curvature_k2_rho"] @@ -381,6 +444,10 @@ def _curvature_K_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["curvature_k1_rho", "curvature_k2_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], ) def _curvature_H_rho(params, transforms, profiles, data, **kwargs): data["curvature_H_rho"] = (data["curvature_k1_rho"] + data["curvature_k2_rho"]) / 2 @@ -555,6 +622,10 @@ def _curvature_H_theta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_zeta", "e_rho_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _L_sff_zeta(params, transforms, profiles, data, **kwargs): data["L_sff_zeta"] = dot(data["e_rho_r"], data["n_zeta"]) @@ -573,6 +644,10 @@ def _L_sff_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_zeta", "e_rho_t"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _M_sff_zeta(params, transforms, profiles, data, **kwargs): data["M_sff_zeta"] = dot(data["e_rho_t"], data["n_zeta"]) @@ -591,6 +666,10 @@ def _M_sff_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["n_zeta", "e_theta_t"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _N_sff_zeta(params, transforms, profiles, data, **kwargs): data["N_sff_zeta"] = dot(data["e_theta_t"], data["n_zeta"]) @@ -609,6 +688,10 @@ def _N_sff_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["g_rr", "g_rt", "g_tt", "L_sff_zeta", "M_sff_zeta", "N_sff_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _curvature_k1_zeta(params, transforms, profiles, data, **kwargs): # following notation from @@ -641,6 +724,10 @@ def _curvature_k1_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["g_rr", "g_rt", "g_tt", "L_sff_zeta", "M_sff_zeta", "N_sff_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _curvature_k2_zeta(params, transforms, profiles, data, **kwargs): # following notation from @@ -673,6 +760,10 @@ def _curvature_k2_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["curvature_k1_zeta", "curvature_k2_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _curvature_K_zeta(params, transforms, profiles, data, **kwargs): data["curvature_K_zeta"] = data["curvature_k1_zeta"] * data["curvature_k2_zeta"] @@ -691,6 +782,10 @@ def _curvature_K_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["curvature_k1_zeta", "curvature_k2_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.ZernikeRZToroidalSection", + ], ) def _curvature_H_zeta(params, transforms, profiles, data, **kwargs): data["curvature_H_zeta"] = ( diff --git a/desc/compute/_metric.py b/desc/compute/_metric.py index f881c37ce9..a42cb17553 100644 --- a/desc/compute/_metric.py +++ b/desc/compute/_metric.py @@ -54,6 +54,10 @@ def _sqrtg_pest(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _e_theta_x_e_zeta(params, transforms, profiles, data, **kwargs): data["|e_theta x e_zeta|"] = jnp.linalg.norm( @@ -74,6 +78,10 @@ def _e_theta_x_e_zeta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _e_zeta_x_e_rho(params, transforms, profiles, data, **kwargs): data["|e_zeta x e_rho|"] = jnp.linalg.norm( @@ -94,6 +102,10 @@ def _e_zeta_x_e_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_theta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _e_rho_x_e_theta(params, transforms, profiles, data, **kwargs): data["|e_rho x e_theta|"] = jnp.linalg.norm( @@ -417,6 +429,10 @@ def _sqrtg_rz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_rr(params, transforms, profiles, data, **kwargs): data["g_rr"] = dot(data["e_rho"], data["e_rho"]) @@ -435,6 +451,10 @@ def _g_sub_rr(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_tt(params, transforms, profiles, data, **kwargs): data["g_tt"] = dot(data["e_theta"], data["e_theta"]) @@ -453,6 +473,10 @@ def _g_sub_tt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_zz(params, transforms, profiles, data, **kwargs): data["g_zz"] = dot(data["e_zeta"], data["e_zeta"]) @@ -471,6 +495,10 @@ def _g_sub_zz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_theta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_rt(params, transforms, profiles, data, **kwargs): data["g_rt"] = dot(data["e_rho"], data["e_theta"]) @@ -489,6 +517,10 @@ def _g_sub_rt(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_rz(params, transforms, profiles, data, **kwargs): data["g_rz"] = dot(data["e_rho"], data["e_zeta"]) @@ -507,6 +539,10 @@ def _g_sub_rz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_tz(params, transforms, profiles, data, **kwargs): data["g_tz"] = dot(data["e_theta"], data["e_zeta"]) @@ -526,6 +562,10 @@ def _g_sub_tz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_theta_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_tt_r(params, transforms, profiles, data, **kwargs): data["g_tt_r"] = 2 * dot(data["e_theta"], data["e_theta_r"]) @@ -545,6 +585,10 @@ def _g_sub_tt_r(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta", "e_theta_r", "e_zeta_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _g_sub_tz_r(params, transforms, profiles, data, **kwargs): data["g_tz_r"] = dot(data["e_theta_r"], data["e_zeta"]) + dot( From 5496dc0a9b7348a9e958b22ac68da83a410642ba Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 15 Jul 2023 14:51:27 -0400 Subject: [PATCH 18/28] Update surface compute data to match existing name conventions --- desc/compute/_surface.py | 380 +++++++++++++++++++++++++++------------ 1 file changed, 263 insertions(+), 117 deletions(-) diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index dd28e499fb..949eab733c 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -5,7 +5,7 @@ @register_compute_fun( - name="r", + name="x", label="\\mathbf{r}", units="m", units_long="meters", @@ -22,23 +22,23 @@ data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): +def _x_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"]) Z = transforms["Z"].transform(params["Z_lmn"]) phi = transforms["grid"].nodes[:, 2] coords = jnp.stack([R, phi, Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz(coords) - data["r"] = coords + data["x"] = coords return data @register_compute_fun( - name="r_r", - label="\\partial_{\\rho} \\mathbf{r}", + name="e_rho", + label="\\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Position vector along surface, radial derivative", + description="Covariant radial basis vector", dim=3, params=[], transforms={ @@ -49,18 +49,18 @@ def _r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): +def _e_rho_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["r_r"] = coords + data["e_rho"] = coords return data @register_compute_fun( - name="r_t", - label="\\partial_{\\theta} \\mathbf{r}", + name="e_theta", + label="\\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Position vector along surface, poloidal derivative", + description="Covariant poloidal basis vector", dim=3, params=["R_lmn", "Z_lmn"], transforms={ @@ -73,23 +73,23 @@ def _r_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): +def _e_theta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dt=1) Z = transforms["Z"].transform(params["Z_lmn"], dt=1) phi = jnp.zeros(transforms["grid"].num_nodes) coords = jnp.stack([R, phi, Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_t"] = coords + data["e_theta"] = coords return data @register_compute_fun( - name="r_z", - label="\\partial_{\\zeta} \\mathbf{r}", + name="e_zeta", + label="\\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Position vector along surface, toroidal derivative", + description="Covariant toroidal basis vector", dim=3, params=["R_lmn", "Z_lmn"], transforms={ @@ -102,7 +102,7 @@ def _r_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): +def _e_zeta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_lmn"], dz=0) dR = transforms["R"].transform(params["R_lmn"], dz=1) dZ = transforms["Z"].transform(params["Z_lmn"], dz=1) @@ -110,16 +110,16 @@ def _r_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.stack([dR, dphi, dZ], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_z"] = coords + data["e_zeta"] = coords return data @register_compute_fun( - name="r_rr", - label="\\partial_{\\rho \\rho} \\mathbf{r}", + name="e_rho_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Position vector along surface, second radial derivative", + description="Covariant radial basis vector, derivative wrt radial coordinate", dim=3, params=[], transforms={ @@ -130,18 +130,84 @@ def _r_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): +def _e_rho_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["r_rr"] = coords + data["e_rho_r"] = coords return data @register_compute_fun( - name="r_tt", - label="\\partial_{\\theta \\theta} \\mathbf{r}", + name="e_rho_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Position vector along surface, second poloidal derivative", + description="Covariant radial basis vector, derivative wrt poloidal angle", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _e_rho_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_rho_t"] = coords + return data + + +@register_compute_fun( + name="e_rho_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\rho}", + units="m", + units_long="meters", + description="Covariant radial basis vector, derivative wrt toroidal angle", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _e_rho_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_rho_z"] = coords + return data + + +@register_compute_fun( + name="e_theta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="tz", + data=[], + parameterization="desc.geometry.FourierRZToroidalSurface", +) +def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_theta_r"] = coords + return data + + +@register_compute_fun( + name="e_theta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt poloidal angle", dim=3, params=["R_lmn", "Z_lmn"], transforms={ @@ -154,28 +220,28 @@ def _r_rr_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_tt_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): +def _e_theta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dt=2) Z = transforms["Z"].transform(params["Z_lmn"], dt=2) phi = jnp.zeros(transforms["grid"].num_nodes) coords = jnp.stack([R, phi, Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_tt"] = coords + data["e_theta_t"] = coords return data @register_compute_fun( - name="r_zz", - label="\\partial_{\\zeta \\zeta} \\mathbf{r}", + name="e_theta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Position vector along surface, toroidal derivative", + description="Covariant poloidal basis vector, derivative wrt toroidal angle", dim=3, params=["R_lmn", "Z_lmn"], transforms={ - "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], - "Z": [[0, 0, 2]], + "R": [[0, 1, 1]], + "Z": [[0, 1, 1]], "grid": [], }, profiles=[], @@ -183,25 +249,23 @@ def _r_tt_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_zz_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): - R0 = transforms["R"].transform(params["R_lmn"], dz=0) - dR = transforms["R"].transform(params["R_lmn"], dz=1) - d2R = transforms["R"].transform(params["R_lmn"], dz=2) - d2Z = transforms["Z"].transform(params["Z_lmn"], dz=2) - dphi = 2 * dR - coords = jnp.stack([d2R - R0, dphi, d2Z], axis=1) +def _e_theta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + dR = transforms["R"].transform(params["R_lmn"], dt=1, dz=1) + dZ = transforms["Z"].transform(params["Z_lmn"], dt=1, dz=1) + dphi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([dR, dphi, dZ], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_zz"] = coords + data["e_theta_z"] = coords return data @register_compute_fun( - name="r_rt", - label="\\partial_{\\rho \\theta} \\mathbf{r}", + name="e_zeta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Position vector along surface, radial/poloidal derivative", + description="Covariant toroidal basis vector, derivative wrt radial coordinate", dim=3, params=[], transforms={ @@ -212,21 +276,23 @@ def _r_zz_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_rt_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): +def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["r_rt"] = coords + data["e_zeta_r"] = coords return data @register_compute_fun( - name="r_rz", - label="\\partial_{\\rho \\zeta} \\mathbf{r}", + name="e_zeta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Position vector along surface, radial/toroidal derivative", + description="Covariant toroidal basis vector, derivative wrt poloidal angle", dim=3, - params=[], + params=["R_lmn", "Z_lmn"], transforms={ + "R": [[0, 1, 1]], + "Z": [[0, 1, 1]], "grid": [], }, profiles=[], @@ -234,23 +300,28 @@ def _r_rt_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_rz_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): - coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["r_rz"] = coords +def _e_zeta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + dR = transforms["R"].transform(params["R_lmn"], dt=1, dz=1) + dZ = transforms["Z"].transform(params["Z_lmn"], dt=1, dz=1) + dphi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([dR, dphi, dZ], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_zeta_t"] = coords return data @register_compute_fun( - name="r_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{r}", + name="e_zeta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Position vector along surface, poloidal/toroidal derivative", + description="Covariant toroidal basis vector, derivative wrt toroidal angle", dim=3, params=["R_lmn", "Z_lmn"], transforms={ - "R": [[0, 0, 0], [0, 0, 1]], - "Z": [[0, 0, 1]], + "R": [[0, 0, 0], [0, 0, 1], [0, 0, 2]], + "Z": [[0, 0, 2]], "grid": [], }, profiles=[], @@ -258,19 +329,21 @@ def _r_rz_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.FourierRZToroidalSurface", ) -def _r_tz_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): - dR = transforms["R"].transform(params["R_lmn"], dt=1, dz=1) - dZ = transforms["Z"].transform(params["Z_lmn"], dt=1, dz=1) - dphi = jnp.zeros(transforms["grid"].num_nodes) - coords = jnp.stack([dR, dphi, dZ], axis=1) +def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): + R0 = transforms["R"].transform(params["R_lmn"], dz=0) + dR = transforms["R"].transform(params["R_lmn"], dz=1) + d2R = transforms["R"].transform(params["R_lmn"], dz=2) + d2Z = transforms["Z"].transform(params["Z_lmn"], dz=2) + dphi = 2 * dR + coords = jnp.stack([d2R - R0, dphi, d2Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_tz"] = coords + data["e_zeta_z"] = coords return data @register_compute_fun( - name="r", + name="x", label="\\mathbf{r}", units="m", units_long="meters", @@ -287,23 +360,23 @@ def _r_tz_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): +def _x_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"]) Z = transforms["Z"].transform(params["Z_lmn"]) phi = transforms["grid"].nodes[:, 2] coords = jnp.stack([R, phi, Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz(coords) - data["r"] = coords + data["x"] = coords return data @register_compute_fun( - name="r_r", - label="\\partial_{\\rho} \\mathbf{r}", + name="e_rho", + label="\\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Position vector along surface, radial derivative", + description="Covariant radial basis vector", dim=3, params=["R_lmn", "Z_lmn"], transforms={ @@ -316,23 +389,23 @@ def _r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): +def _e_rho_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dr=1) Z = transforms["Z"].transform(params["Z_lmn"], dr=1) phi = jnp.zeros(transforms["grid"].num_nodes) coords = jnp.stack([R, phi, Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz(coords) - data["r_r"] = coords + data["e_rho"] = coords return data @register_compute_fun( - name="r_t", - label="\\partial_{\\theta} \\mathbf{r}", + name="e_theta", + label="\\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Position vector along surface, poloidal derivative", + description="Covariant poloidal basis vector", dim=3, params=["R_lmn", "Z_lmn"], transforms={ @@ -345,23 +418,23 @@ def _r_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): +def _e_theta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dt=1) Z = transforms["Z"].transform(params["Z_lmn"], dt=1) phi = jnp.zeros(transforms["grid"].num_nodes) coords = jnp.stack([R, phi, Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_t"] = coords + data["e_theta"] = coords return data @register_compute_fun( - name="r_z", - label="\\partial_{\\zeta} \\mathbf{r}", + name="e_zeta", + label="\\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Position vector along surface, toroidal derivative", + description="Covariant toroidal basis vector", dim=3, params=[], transforms={ @@ -372,18 +445,18 @@ def _r_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): +def _e_zeta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["r_z"] = coords + data["e_zeta"] = coords return data @register_compute_fun( - name="r_rr", - label="\\partial_{\\rho \\rho} \\mathbf{r}", + name="e_rho_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Position vector along surface, second radial derivative", + description="Covariant radial basis vector, derivative wrt radial coordinate", dim=3, params=["R_lmn", "Z_lmn"], transforms={ @@ -396,28 +469,28 @@ def _r_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): +def _e_rho_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dr=2) Z = transforms["Z"].transform(params["Z_lmn"], dr=2) phi = jnp.zeros(transforms["grid"].num_nodes) coords = jnp.stack([R, phi, Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz(coords) - data["r_rr"] = coords + data["e_rho_r"] = coords return data @register_compute_fun( - name="r_tt", - label="\\partial_{\\theta \\theta} \\mathbf{r}", + name="e_rho_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Position vector along surface, second poloidal derivative", + description="Covariant radial basis vector, derivative wrt poloidal angle", dim=3, params=["R_lmn", "Z_lmn"], transforms={ - "R": [[0, 2, 0]], - "Z": [[0, 2, 0]], + "R": [[1, 1, 0]], + "Z": [[1, 1, 0]], "grid": [], }, profiles=[], @@ -425,23 +498,23 @@ def _r_rr_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_tt_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): - R = transforms["R"].transform(params["R_lmn"], dt=2) - Z = transforms["Z"].transform(params["Z_lmn"], dt=2) +def _e_rho_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dr=1, dt=1) + Z = transforms["Z"].transform(params["Z_lmn"], dr=1, dt=1) phi = jnp.zeros(transforms["grid"].num_nodes) coords = jnp.stack([R, phi, Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": - coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) - data["r_tt"] = coords + coords = rpz2xyz(coords) + data["e_rho_t"] = coords return data @register_compute_fun( - name="r_zz", - label="\\partial_{\\zeta \\zeta} \\mathbf{r}", + name="e_rho_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\rho}", units="m", units_long="meters", - description="Position vector along surface, toroidal derivative", + description="Covariant radial basis vector, derivative wrt toroidal angle", dim=3, params=[], transforms={ @@ -452,18 +525,18 @@ def _r_tt_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_zz_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): +def _e_rho_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["r_zz"] = coords + data["e_rho_z"] = coords return data @register_compute_fun( - name="r_rt", - label="\\partial_{\\rho \\theta} \\mathbf{r}", + name="e_theta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\theta}", units="m", units_long="meters", - description="Position vector along surface, radial/poloidal derivative", + description="Covariant poloidal basis vector, derivative wrt radial coordinate", dim=3, params=["R_lmn", "Z_lmn"], transforms={ @@ -476,23 +549,96 @@ def _r_zz_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_rt_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): +def _e_theta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dr=1, dt=1) Z = transforms["Z"].transform(params["Z_lmn"], dr=1, dt=1) phi = jnp.zeros(transforms["grid"].num_nodes) coords = jnp.stack([R, phi, Z], axis=1) if kwargs.get("basis", "rpz").lower() == "xyz": coords = rpz2xyz(coords) - data["r_rt"] = coords + data["e_theta_r"] = coords + return data + + +@register_compute_fun( + name="e_theta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt poloidal angle", + dim=3, + params=["R_lmn", "Z_lmn"], + transforms={ + "R": [[0, 2, 0]], + "Z": [[0, 2, 0]], + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _e_theta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + R = transforms["R"].transform(params["R_lmn"], dt=2) + Z = transforms["Z"].transform(params["Z_lmn"], dt=2) + phi = jnp.zeros(transforms["grid"].num_nodes) + coords = jnp.stack([R, phi, Z], axis=1) + if kwargs.get("basis", "rpz").lower() == "xyz": + coords = rpz2xyz_vec(coords, phi=transforms["grid"].nodes[:, 2]) + data["e_theta_t"] = coords + return data + + +@register_compute_fun( + name="e_theta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\theta}", + units="m", + units_long="meters", + description="Covariant poloidal basis vector, derivative wrt toroidal angle", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _e_theta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_theta_z"] = coords + return data + + +@register_compute_fun( + name="e_zeta_r", + label="\\partial_{\\rho} \\mathbf{e}_{\\zeta}", + units="m", + units_long="meters", + description="Covariant toroidal basis vector, derivative wrt radial coordinate", + dim=3, + params=[], + transforms={ + "grid": [], + }, + profiles=[], + coordinates="rt", + data=[], + parameterization="desc.geometry.ZernikeRZToroidalSection", +) +def _e_zeta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): + coords = jnp.zeros((transforms["grid"].num_nodes, 3)) + data["e_zeta_r"] = coords return data @register_compute_fun( - name="r_rz", - label="\\partial_{\\rho \\zeta} \\mathbf{r}", + name="e_zeta_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Position vector along surface, radial/toroidal derivative", + description="Covariant toroidal basis vector, derivative wrt poloidal angle", dim=3, params=[], transforms={ @@ -503,18 +649,18 @@ def _r_rt_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_rz_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): +def _e_zeta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["r_rz"] = coords + data["e_zeta_t"] = coords return data @register_compute_fun( - name="r_tz", - label="\\partial_{\\theta \\zeta} \\mathbf{r}", + name="e_zeta_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\zeta}", units="m", units_long="meters", - description="Position vector along surface, poloidal/toroidal derivative", + description="Covariant toroidal basis vector, derivative wrt toroidal angle", dim=3, params=[], transforms={ @@ -525,7 +671,7 @@ def _r_rz_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs) data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", ) -def _r_tz_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): +def _e_zeta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) - data["r_tz"] = coords + data["e_zeta_z"] = coords return data From fd6a0bf67596aee666f786e5967fbccfc8cc5900 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 15 Jul 2023 14:57:06 -0400 Subject: [PATCH 19/28] Add compute method to surface base class, add test --- desc/geometry/core.py | 62 ++++++++++++++++++++++++++++++++++++-- tests/test_compute_funs.py | 22 +++++++++++++- 2 files changed, 81 insertions(+), 3 deletions(-) diff --git a/desc/geometry/core.py b/desc/geometry/core.py index a4339ca516..cea436909f 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -8,7 +8,7 @@ from desc.backend import jnp from desc.compute.utils import compute as compute_fun from desc.compute.utils import get_params, get_transforms -from desc.grid import LinearGrid +from desc.grid import LinearGrid, QuadratureGrid from desc.io import IOAble from .utils import reflection_matrix, rotation_matrix @@ -49,7 +49,7 @@ def compute( 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 the Linear grid. + 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. @@ -194,6 +194,64 @@ def grid(self): def change_resolution(self, *args, **kwargs): """Change the maximum resolution.""" + 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, optional + Grid of coordinates to evaluate at. Defaults to a Linear grid for constant + rho surfaces or a Quadrature grid for constant zeta surfaces. + 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. + + """ + if isinstance(names, str): + names = [names] + if grid is None: + NFP = self.NFP if hasattr(self, "NFP") else 1 + if self.L == 0: + grid = LinearGrid(M=2 * self.M + 5, N=2 * self.N + 5, NFP=NFP) + else: + grid = QuadratureGrid(L=2 * self.L + 5, M=2 * self.M + 5, N=0, NFP=NFP) + if params is None: + params = get_params(names, obj=self) + if transforms is None: + transforms = get_transforms(names, obj=self, grid=grid, **kwargs) + if data is None: + data = {} + profiles = {} + + data = compute_fun( + self, + names, + params=params, + transforms=transforms, + profiles=profiles, + data=data, + **kwargs, + ) + return data + @abstractmethod def compute_coordinates(self, params=None, grid=None, dt=0, dz=0): """Compute coordinate values at specified nodes.""" diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 7684b2ee02..af8e4028d4 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -9,7 +9,13 @@ from desc.compute import data_index from desc.compute.utils import compress from desc.equilibrium import EquilibriaFamily, Equilibrium -from desc.geometry import FourierPlanarCurve, FourierRZCurve, FourierXYZCurve +from desc.geometry import ( + FourierPlanarCurve, + FourierRZCurve, + FourierRZToroidalSurface, + FourierXYZCurve, + ZernikeRZToroidalSection, +) from desc.grid import LinearGrid, QuadratureGrid # convolve kernel is reverse of FD coeffs @@ -1136,6 +1142,20 @@ def test_curve_compute_everything(): assert key in data +@pytest.mark.unit +def test_surface_compute_everything(): + """Make sure we can compute every surface thing without errors.""" + surfaces = { + "desc.geometry.surface.FourierRZToroidalSurface": FourierRZToroidalSurface(), + "desc.geometry.surface.ZernikeRZToroidalSection": ZernikeRZToroidalSection(), + } + + for p, thing in surfaces.items(): + for key in data_index[p].keys(): + data = thing.compute(key) + assert key in data + + @pytest.mark.unit def test_compute_averages(): """Test that computing averages uses the correct grid.""" From 2e7a558ccbef40648f6cbf6c939bfa0a2ffa6305 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 15 Jul 2023 15:14:52 -0400 Subject: [PATCH 20/28] Use surface label rather than resolution to determine grid type --- desc/geometry/core.py | 9 ++++++--- desc/geometry/surface.py | 8 ++++---- 2 files changed, 10 insertions(+), 7 deletions(-) diff --git a/desc/geometry/core.py b/desc/geometry/core.py index cea436909f..3684554b76 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -229,10 +229,13 @@ def compute( names = [names] if grid is None: NFP = self.NFP if hasattr(self, "NFP") else 1 - if self.L == 0: - grid = LinearGrid(M=2 * self.M + 5, N=2 * self.N + 5, NFP=NFP) - else: + if hasattr(self, "rho"): # constant rho surface + grid = LinearGrid( + rho=np.array(self.rho), M=2 * self.M + 5, N=2 * self.N + 5, NFP=NFP + ) + elif hasattr(self, "zeta"): # constant zeta surface grid = QuadratureGrid(L=2 * self.L + 5, M=2 * self.M + 5, N=0, NFP=NFP) + grid._nodes[:, 2] = self.zeta if params is None: params = get_params(names, obj=self) if transforms is None: diff --git a/desc/geometry/surface.py b/desc/geometry/surface.py index b834589f62..ade10ef3cb 100644 --- a/desc/geometry/surface.py +++ b/desc/geometry/surface.py @@ -34,7 +34,7 @@ class FourierRZToroidalSurface(Surface): sym : bool whether to enforce stellarator symmetry. Default is "auto" which enforces if modes are symmetric. If True, non-symmetric modes will be truncated. - rho : float (0,1) + rho : float [0,1] flux surface label for the toroidal surface grid : Grid default grid for computation @@ -593,7 +593,7 @@ class ZernikeRZToroidalSection(Surface): whether to enforce stellarator symmetry. Default is "auto" which enforces if modes are symmetric. If True, non-symmetric modes will be truncated. spectral_indexing : {``'ansi'``, ``'fringe'``} - Indexing method, default value = ``'fringe'`` + Indexing method, default value = ``'ansi'`` For L=0, all methods are equivalent and give a "chevron" shaped basis (only the outer edge of the zernike pyramid of width M). @@ -609,7 +609,7 @@ class ZernikeRZToroidalSection(Surface): decreasing size, ending in a diamond shape for L=2*M where the traditional fringe/U of Arizona indexing is recovered. For L > 2*M, adds chevrons to the bottom, making a hexagonal diamond - zeta : float (0,2pi) + zeta : float [0,2pi) toroidal angle for the section. grid : Grid default grid for computation @@ -639,7 +639,7 @@ def __init__( Z_lmn=None, modes_R=None, modes_Z=None, - spectral_indexing="fringe", + spectral_indexing="ansi", sym="auto", zeta=0.0, grid=None, From bed65cb16d464c6217f9a7611bc170ffc07829ca Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 15 Jul 2023 17:28:05 -0400 Subject: [PATCH 21/28] Add new method for computing outermost surface area --- desc/compute/_geometry.py | 32 ++++++++++++++++++++++++++++---- 1 file changed, 28 insertions(+), 4 deletions(-) diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index 6a31cb6f52..50c22ad0c5 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -146,6 +146,34 @@ def _A(params, transforms, profiles, data, **kwargs): # TODO: compute cross section area for toroidal surface using stokes? +@register_compute_fun( + name="S", + label="S", + units="m^{2}", + units_long="square meters", + description="Surface area of outermost flux surface", + dim=0, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="", + data=["|e_theta x e_zeta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.surface.FourierRZToroidalSurface", + ], +) +def _S(params, transforms, profiles, data, **kwargs): + data["S"] = jnp.max( + surface_integrals( + transforms["grid"], + data["|e_theta x e_zeta|"], + expand_out=False, + ) + ) + return data + + @register_compute_fun( name="S(r)", label="S(\\rho)", @@ -158,10 +186,6 @@ def _A(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="r", data=["|e_theta x e_zeta|"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.surface.FourierRZToroidalSurface", - ], ) def _S_of_r(params, transforms, profiles, data, **kwargs): data["S(r)"] = surface_integrals(transforms["grid"], data["|e_theta x e_zeta|"]) From 2647cca2e71f137a1724fd0aabc606e7fb009f65 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 15 Jul 2023 17:30:33 -0400 Subject: [PATCH 22/28] Remove grid, transform, old compute methods from surface classes --- desc/geometry/core.py | 50 ----- desc/geometry/surface.py | 408 ----------------------------------- desc/objectives/_geometry.py | 6 +- desc/plotting.py | 2 +- tests/test_configuration.py | 8 +- tests/test_surfaces.py | 68 +++--- 6 files changed, 39 insertions(+), 503 deletions(-) diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 3684554b76..12791fc6d4 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -185,11 +185,6 @@ def _flip_orientation(self): one[self.Z_basis.modes[:, 1] < 0] *= -1 self.Z_lmn *= one - @property - @abstractmethod - def grid(self): - """Grid: Nodes for computation.""" - @abstractmethod def change_resolution(self, *args, **kwargs): """Change the maximum resolution.""" @@ -255,51 +250,6 @@ def compute( ) return data - @abstractmethod - def compute_coordinates(self, params=None, grid=None, dt=0, dz=0): - """Compute coordinate values at specified nodes.""" - - @abstractmethod - def compute_normal(self, params=None, grid=None): - """Compute normal vectors to the surface on predefined grid.""" - - @abstractmethod - def compute_surface_area(self, params=None, grids=None): - """Compute surface area via quadrature.""" - - def compute_curvature(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute gaussian and mean curvature. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - fourier coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - K, H, k1, k2 : ndarray, shape(k,) - Gaussian, mean and 2 principle curvatures at points specified in grid. - - """ - # following notation from - # https://en.wikipedia.org/wiki/Parametric_surface#Curvature - E, F, G = self._compute_first_fundamental_form(R_lmn, Z_lmn, grid) - L, M, N = self._compute_second_fundamental_form(R_lmn, Z_lmn, grid) - # coeffs of quadratic eqn for determinant - a = E * G - F**2 - b = F * M - L * G - E * N - c = L * N - M**2 - r1 = (-b + jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) - r2 = (-b - jnp.sqrt(b**2 - 4 * a * c)) / (2 * a) - k1 = jnp.maximum(r1, r2) - k2 = jnp.minimum(r1, r2) - K = k1 * k2 - H = (k1 + k2) / 2 - return K, H, k1, k2 - def __repr__(self): """Get the string form of the object.""" return ( diff --git a/desc/geometry/surface.py b/desc/geometry/surface.py index ade10ef3cb..c7a31bdefc 100644 --- a/desc/geometry/surface.py +++ b/desc/geometry/surface.py @@ -7,13 +7,10 @@ from desc.backend import jnp, put, sign from desc.basis import DoubleFourierSeries, ZernikePolynomial -from desc.grid import Grid, LinearGrid from desc.io import InputReader -from desc.transform import Transform from desc.utils import copy_coeffs from .core import Surface -from .utils import rpz2xyz, rpz2xyz_vec __all__ = ["FourierRZToroidalSurface", "ZernikeRZToroidalSection"] @@ -36,8 +33,6 @@ class FourierRZToroidalSurface(Surface): modes are symmetric. If True, non-symmetric modes will be truncated. rho : float [0,1] flux surface label for the toroidal surface - grid : Grid - default grid for computation name : str name for this surface check_orientation : bool @@ -125,16 +120,6 @@ def __init__( self._flip_orientation() assert self._compute_orientation() == 1 - if grid is None: - grid = LinearGrid( - M=2 * self.M, - N=2 * self.N, - NFP=self.NFP, - rho=np.asarray(self.rho), - endpoint=True, - ) - self._grid = grid - self._R_transform, self._Z_transform = self._get_transforms(grid) self.name = name @property @@ -159,24 +144,6 @@ def Z_basis(self): """DoubleFourierSeries: Spectral basis for Z.""" return self._Z_basis - @property - def grid(self): - """Grid: Nodes for computation.""" - return self._grid - - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._R_transform.grid = self.grid - self._Z_transform.grid = self.grid - def change_resolution(self, *args, **kwargs): """Change the maximum poloidal and toroidal resolution.""" assert ( @@ -219,11 +186,6 @@ def change_resolution(self, *args, **kwargs): self.Z_basis.change_resolution( M=M, N=N, NFP=self.NFP, sym="sin" if self.sym else self.sym ) - if hasattr(self.grid, "change_resolution"): - self.grid.change_resolution( - self.grid.L, self.grid.M, self.grid.N, self.NFP - ) - self._R_transform, self._Z_transform = self._get_transforms(self.grid) self.R_lmn = copy_coeffs(self.R_lmn, R_modes_old, self.R_basis.modes) self.Z_lmn = copy_coeffs(self.Z_lmn, Z_modes_old, self.Z_basis.modes) self._M = M @@ -297,195 +259,6 @@ def set_coeffs(self, m, n=0, R=None, Z=None): idxZ = self.Z_basis.get_idx(0, mm, nn) self.Z_lmn = put(self.Z_lmn, idxZ, ZZ) - def _get_transforms(self, grid=None): - if grid is None: - return self._R_transform, self._Z_transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = LinearGrid( - M=grid, N=grid, rho=np.asarray(self.rho), NFP=self.NFP - ) - elif len(grid) == 2: - grid = LinearGrid( - M=grid[0], N=grid[1], rho=np.asarray(self.rho), NFP=self.NFP - ) - elif grid.shape[1] == 2: - grid = np.pad(grid, ((0, 0), (1, 0)), constant_values=self.rho) - grid = Grid(grid, sort=False) - else: - grid = Grid(grid, sort=False) - R_transform = Transform( - grid, - self.R_basis, - derivs=np.array( - [[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 0, 1], [0, 0, 2], [0, 1, 1]] - ), - ) - Z_transform = Transform( - grid, - self.Z_basis, - derivs=np.array( - [[0, 0, 0], [0, 1, 0], [0, 2, 0], [0, 0, 1], [0, 0, 2], [0, 1, 1]] - ), - ) - return R_transform, Z_transform - - def _compute_first_fundamental_form(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute coefficients for the first fundamental form.""" - rt = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - rz = self.compute_coordinates(R_lmn, Z_lmn, grid, dz=1) - E = jnp.sum(rt * rt, axis=-1) - F = jnp.sum(rt * rz, axis=-1) - G = jnp.sum(rz * rz, axis=-1) - return E, F, G - - def _compute_second_fundamental_form(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute coefficients for the second fundamental form.""" - rtt = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=2) - rtz = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1, dz=1) - rzz = self.compute_coordinates(R_lmn, Z_lmn, grid, dz=2) - n = self.compute_normal(R_lmn, Z_lmn, grid) - L = jnp.sum(rtt * n, axis=-1) - M = jnp.sum(rtz * n, axis=-1) - N = jnp.sum(rzz * n, axis=-1) - return L, M, N - - def compute_coordinates( - self, R_lmn=None, Z_lmn=None, grid=None, dt=0, dz=0, basis="rpz" - ): - """Compute values using specified coefficients. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - fourier coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - dt, dz: int - derivative order to compute in theta, zeta - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - R, phi, Z or x, y, z coordinates of the surface at points specified in grid - """ - assert basis.lower() in ["rpz", "xyz"] - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - if dz == 0: - R = R_transform.transform(R_lmn, dt=dt, dz=0) - Z = Z_transform.transform(Z_lmn, dt=dt, dz=0) - phi = R_transform.grid.nodes[:, 2] * (dt == 0) - coords = jnp.stack([R, phi, Z], axis=1) - elif dz == 1: - R0 = R_transform.transform(R_lmn, dt=dt, dz=0) - dR = R_transform.transform(R_lmn, dt=dt, dz=1) - dZ = Z_transform.transform(Z_lmn, dt=dt, dz=1) - dphi = R0 * (dt == 0) - coords = jnp.stack([dR, dphi, dZ], axis=1) - elif dz == 2: - R0 = R_transform.transform(R_lmn, dt=dt, dz=0) - dR = R_transform.transform(R_lmn, dt=dt, dz=1) - d2R = R_transform.transform(R_lmn, dt=dt, dz=2) - d2Z = Z_transform.transform(Z_lmn, dt=dt, dz=2) - R = d2R - R0 - Z = d2Z - # 2nd derivative wrt phi = 0 - phi = 2 * dR * (dt == 0) - coords = jnp.stack([R, phi, Z], axis=1) - elif dz == 3: - R0 = R_transform.transform(R_lmn, dt=dt, dz=0) - dR = R_transform.transform(R_lmn, dt=dt, dz=1) - d2R = R_transform.transform(R_lmn, dt=dt, dz=2) - d3R = R_transform.transform(R_lmn, dt=dt, dz=3) - d3Z = Z_transform.transform(Z_lmn, dt=dt, dz=3) - R = d3R - 3 * dR - Z = d3Z - phi = (3 * d2R - R0) * (dt == 0) - coords = jnp.stack([R, phi, Z], axis=1) - else: - raise NotImplementedError( - "Derivatives higher than 3 have not been implemented in " - + "cylindrical coordinates." - ) - if basis.lower() == "xyz": - if (dt > 0) or (dz > 0): - coords = rpz2xyz_vec(coords, phi=R_transform.grid.nodes[:, 2]) - else: - coords = rpz2xyz(coords) - return coords - - def compute_normal(self, R_lmn=None, Z_lmn=None, grid=None, basis="rpz"): - """Compute normal vector to surface on default grid. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - fourier coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - basis : {"rpz", "xyz"} - basis vectors to use for normal vector representation - - Returns - ------- - N : ndarray, shape(k,3) - normal vector to surface in specified coordinates - """ - assert basis.lower() in ["rpz", "xyz"] - - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - r_t = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - r_z = self.compute_coordinates(R_lmn, Z_lmn, grid, dz=1) - - N = jnp.cross(r_t, r_z, axis=1) - N = N / jnp.linalg.norm(N, axis=1)[:, jnp.newaxis] - if basis.lower() == "xyz": - phi = R_transform.grid.nodes[:, 2] - N = rpz2xyz_vec(N, phi=phi) - return N - - def compute_surface_area(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute surface area via quadrature. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - fourier coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,2pi) - - Returns - ------- - area : float - surface area - """ - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - r_t = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - r_z = self.compute_coordinates(R_lmn, Z_lmn, grid, dz=1) - - N = jnp.cross(r_t, r_z, axis=1) - return jnp.sum(R_transform.grid.weights * jnp.linalg.norm(N, axis=1)) - @classmethod def from_input_file(cls, path): """Create a surface from Fourier coefficients in a DESC or VMEC input file. @@ -611,8 +384,6 @@ class ZernikeRZToroidalSection(Surface): For L > 2*M, adds chevrons to the bottom, making a hexagonal diamond zeta : float [0,2pi) toroidal angle for the section. - grid : Grid - default grid for computation name : str name for this surface check_orientation : bool @@ -708,12 +479,6 @@ def __init__( self._flip_orientation() assert self._compute_orientation() == 1 - if grid is None: - grid = LinearGrid( - L=self.L, M=2 * self.M, zeta=np.asarray(self.zeta), endpoint=True - ) - self._grid = grid - self._R_transform, self._Z_transform = self._get_transforms(grid) self.name = name @property @@ -731,24 +496,6 @@ def Z_basis(self): """ZernikePolynomial: Spectral basis for Z.""" return self._Z_basis - @property - def grid(self): - """Grid: Nodes for computation.""" - return self._grid - - @grid.setter - def grid(self, new): - if isinstance(new, Grid): - self._grid = new - elif isinstance(new, (np.ndarray, jnp.ndarray)): - self._grid = Grid(new, sort=False) - else: - raise TypeError( - f"grid should be a Grid or subclass, or ndarray, got {type(new)}" - ) - self._R_transform.grid = self.grid - self._Z_transform.grid = self.grid - def change_resolution(self, *args, **kwargs): """Change the maximum radial and poloidal resolution.""" assert ( @@ -785,9 +532,6 @@ def change_resolution(self, *args, **kwargs): self.Z_basis.change_resolution( L=L, M=M, sym="sin" if self.sym else self.sym ) - if hasattr(self.grid, "change_resolution"): - self.grid.change_resolution(self.grid.L, self.grid.M, self.grid.N) - self._R_transform, self._Z_transform = self._get_transforms(self.grid) self.R_lmn = copy_coeffs(self.R_lmn, R_modes_old, self.R_basis.modes) self.Z_lmn = copy_coeffs(self.Z_lmn, Z_modes_old, self.Z_basis.modes) self._L = L @@ -860,155 +604,3 @@ def set_coeffs(self, l, m=0, R=None, Z=None): if ZZ is not None: idxZ = self.Z_basis.get_idx(ll, mm, 0) self.Z_lmn = put(self.Z_lmn, idxZ, ZZ) - - def _get_transforms(self, grid=None): - if grid is None: - return self._R_transform, self._Z_transform - if not isinstance(grid, Grid): - if np.isscalar(grid): - grid = LinearGrid(L=grid, M=grid, zeta=np.asarray(self.zeta)) - elif len(grid) == 2: - grid = LinearGrid(L=grid[0], M=grid[1], zeta=np.asarray(self.zeta)) - elif grid.shape[1] == 2: - grid = np.pad(grid, ((0, 0), (0, 1)), constant_values=self.zeta) - grid = Grid(grid, sort=False) - else: - grid = Grid(grid, sort=False) - R_transform = Transform( - grid, - self.R_basis, - derivs=np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [0, 1, 0], [0, 2, 0], [1, 1, 0]] - ), - ) - Z_transform = Transform( - grid, - self.Z_basis, - derivs=np.array( - [[0, 0, 0], [1, 0, 0], [2, 0, 0], [0, 1, 0], [0, 2, 0], [1, 1, 0]] - ), - ) - return R_transform, Z_transform - - def _compute_first_fundamental_form(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute coefficients for the first fundamental form.""" - rr = self.compute_coordinates(R_lmn, Z_lmn, grid, dr=1) - rt = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - E = jnp.sum(rr * rr, axis=-1) - F = jnp.sum(rr * rt, axis=-1) - G = jnp.sum(rt * rt, axis=-1) - return E, F, G - - def _compute_second_fundamental_form(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute coefficients for the second fundamental form.""" - rrr = self.compute_coordinates(R_lmn, Z_lmn, grid, dr=2) - rrt = self.compute_coordinates(R_lmn, Z_lmn, grid, dr=1, dt=1) - rtt = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=2) - n = self.compute_normal(R_lmn, Z_lmn, grid) - L = jnp.sum(rrr * n, axis=-1) - M = jnp.sum(rrt * n, axis=-1) - N = jnp.sum(rtt * n, axis=-1) - return L, M, N - - def compute_coordinates( - self, R_lmn=None, Z_lmn=None, grid=None, dr=0, dt=0, basis="rpz" - ): - """Compute values using specified coefficients. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - zernike coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,1)x(0,2pi) - dr, dt: int - derivative order to compute in rho, theta - basis : {"rpz", "xyz"} - coordinate system for returned points - - Returns - ------- - values : ndarray, shape(k,3) - R, phi, Z or x, y, z coordinates of the surface at points specified in grid - """ - assert basis.lower() in ["rpz", "xyz"] - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - R = R_transform.transform(R_lmn, dr=dr, dt=dt) - Z = Z_transform.transform(Z_lmn, dr=dr, dt=dt) - phi = R_transform.grid.nodes[:, 2] * (dr == 0) * (dt == 0) - coords = jnp.stack([R, phi, Z], axis=1) - if basis.lower() == "xyz": - if (dt > 0) or (dr > 0): - coords = rpz2xyz_vec(coords, phi=R_transform.grid.nodes[:, 2]) - else: - coords = rpz2xyz(coords) - return coords - - def compute_normal(self, R_lmn=None, Z_lmn=None, grid=None, basis="rpz"): - """Compute normal vector to surface on default grid. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - zernike coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,1)x(0,2pi) - basis : {"rpz", "xyz"} - basis vectors to use for normal vector representation - - Returns - ------- - N : ndarray, shape(k,3) - normal vector to surface in specified coordinates - """ - assert basis.lower() in ["rpz", "xyz"] - - R_transform, Z_transform = self._get_transforms(grid) - - phi = R_transform.grid.nodes[:, -1] - - # normal vector is a constant 1*phihat - N = jnp.array([jnp.zeros_like(phi), jnp.ones_like(phi), jnp.zeros_like(phi)]).T - - if basis.lower() == "xyz": - N = rpz2xyz_vec(N, phi=phi) - - return N - - def compute_surface_area(self, R_lmn=None, Z_lmn=None, grid=None): - """Compute surface area via quadrature. - - Parameters - ---------- - R_lmn, Z_lmn: array-like - zernike coefficients for R, Z. Defaults to self.R_lmn, self.Z_lmn - grid : Grid or array-like - toroidal coordinates to compute at. Defaults to self.grid - If an integer, assumes that many linearly spaced points in (0,1)x(0,2pi) - - Returns - ------- - area : float - surface area - - """ - if R_lmn is None: - R_lmn = self.R_lmn - if Z_lmn is None: - Z_lmn = self.Z_lmn - R_transform, Z_transform = self._get_transforms(grid) - - r_r = self.compute_coordinates(R_lmn, Z_lmn, grid, dr=1) - r_t = self.compute_coordinates(R_lmn, Z_lmn, grid, dt=1) - - N = jnp.cross(r_r, r_t, axis=1) - return jnp.sum(R_transform.grid.weights * jnp.linalg.norm(N, axis=1)) / ( - 2 * np.pi - ) diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 0356e78ebb..686eec47fb 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -534,9 +534,9 @@ def build(self, eq=None, use_jit=True, verbose=1): print("Precomputing transforms") timer.start("Precomputing transforms") - self._surface_coords = self._surface.compute_coordinates( - grid=surface_grid, basis="xyz" - ) + self._surface_coords = self._surface.compute( + "x", grid=surface_grid, basis="xyz" + )["x"] self._profiles = get_profiles( self._data_keys, obj=eq, diff --git a/desc/plotting.py b/desc/plotting.py index f9c1cf63d5..66d0474ce7 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -2057,7 +2057,7 @@ def flatten_coils(coilset): plot_data["Y"] = [] plot_data["Z"] = [] for i, coil in enumerate(coils_list): - x, y, z = coil.compute_coordinates(grid=grid, basis="xyz").T + x, y, z = coil.compute("x", grid=grid, basis="xyz")["x"].T plot_data["X"].append(x) plot_data["Y"].append(y) plot_data["Z"].append(z) diff --git a/tests/test_configuration.py b/tests/test_configuration.py index f3ae780645..de24a6f235 100644 --- a/tests/test_configuration.py +++ b/tests/test_configuration.py @@ -419,9 +419,7 @@ def test_get_rho_surface(self): rho = 0.5 surf = eq.get_surface_at(rho=rho) assert surf.rho == rho - np.testing.assert_allclose( - surf.compute_surface_area(), 4 * np.pi**2 * R0 * rho - ) + np.testing.assert_allclose(surf.compute("S")["S"], 4 * np.pi**2 * R0 * rho) @pytest.mark.unit def test_get_zeta_surface(self): @@ -430,7 +428,7 @@ def test_get_zeta_surface(self): surf = eq.get_surface_at(zeta=np.pi) assert surf.zeta == np.pi rho = 1 - np.testing.assert_allclose(surf.compute_surface_area(), np.pi * rho**2) + np.testing.assert_allclose(surf.compute("A")["A"], np.pi * rho**2) @pytest.mark.unit def test_get_theta_surface(self): @@ -458,7 +456,7 @@ def test_magnetic_axis(HELIOTRON_vac): grid = LinearGrid(N=3 * eq.N_grid, NFP=eq.NFP, rho=np.array(0.0)) data = eq.compute(["R", "Z"], grid=grid) - coords = axis.compute_coordinates(grid=grid) + coords = axis.compute("x", grid=grid)["x"] np.testing.assert_allclose(coords[:, 0], data["R"]) np.testing.assert_allclose(coords[:, 2], data["Z"]) diff --git a/tests/test_surfaces.py b/tests/test_surfaces.py index 5fe5a44285..010e4d72d3 100644 --- a/tests/test_surfaces.py +++ b/tests/test_surfaces.py @@ -17,33 +17,24 @@ def test_area(self): """Test calculation of surface area.""" s = FourierRZToroidalSurface() grid = LinearGrid(M=24, N=24) - s.grid = grid - area = 4 * np.pi**2 * 10 - np.testing.assert_allclose(s.compute_surface_area(), area) - np.testing.assert_allclose(s.compute_surface_area(grid=10), area) - np.testing.assert_allclose(s.compute_surface_area(grid=(10, 15)), area) + np.testing.assert_allclose(s.compute("S", grid=grid)["S"], area) @pytest.mark.unit def test_normal(self): """Test calculation of surface normal vector.""" s = FourierRZToroidalSurface() grid = LinearGrid(theta=np.pi / 2, zeta=np.pi) - s.grid = grid - N = s.compute_normal() + N = s.compute("n_rho", grid=grid)["n_rho"] np.testing.assert_allclose(N[0], [0, 0, -1], atol=1e-14) grid = LinearGrid(theta=0.0, zeta=0.0) - s.grid = grid - N = s.compute_normal(basis="xyz") + N = s.compute("n_rho", grid=grid)["n_rho"] np.testing.assert_allclose(N[0], [1, 0, 0], atol=1e-12) @pytest.mark.unit def test_misc(self): """Test getting/setting attributes of surface.""" c = FourierRZToroidalSurface() - grid = LinearGrid(L=0, M=2, N=2) - c.grid = grid - assert grid.eq(c.grid) R, Z = c.get_coeffs(0, 0) np.testing.assert_allclose(R, 10) @@ -82,7 +73,6 @@ def test_misc(self): assert c.NFP == 3 assert c.R_basis.NFP == 3 assert c.Z_basis.NFP == 3 - assert c.grid.NFP == 3 @pytest.mark.unit def test_from_input_file(self): @@ -133,12 +123,19 @@ def test_curvature(self): """Tests for gaussian, mean, principle curvatures.""" s = FourierRZToroidalSurface() grid = LinearGrid(theta=np.pi / 2, zeta=np.pi) - s.grid = grid - K, H, k1, k2 = s.compute_curvature() - np.testing.assert_allclose(K, 0) - np.testing.assert_allclose(H, -1 / 2) - np.testing.assert_allclose(k1, 0) - np.testing.assert_allclose(k2, -1) + data = s.compute( + [ + "curvature_K_rho", + "curvature_H_rho", + "curvature_k1_rho", + "curvature_k2_rho", + ], + grid=grid, + ) + np.testing.assert_allclose(data["curvature_K_rho"], 0) + np.testing.assert_allclose(data["curvature_H_rho"], -1 / 2) + np.testing.assert_allclose(data["curvature_k1_rho"], 0) + np.testing.assert_allclose(data["curvature_k2_rho"], -1) class TestZernikeRZToroidalSection: @@ -149,29 +146,21 @@ def test_area(self): """Test calculation of surface area.""" s = ZernikeRZToroidalSection() grid = LinearGrid(L=10, M=10) - s.grid = grid - area = np.pi * 1**2 - np.testing.assert_allclose(s.compute_surface_area(), area) - np.testing.assert_allclose(s.compute_surface_area(grid=15), area) - np.testing.assert_allclose(s.compute_surface_area(grid=(5, 5)), area) + np.testing.assert_allclose(s.compute("A", grid=grid)["A"], area) @pytest.mark.unit def test_normal(self): """Test calculation of surface normal vector.""" s = ZernikeRZToroidalSection() - grid = LinearGrid(L=8, M=4, N=0) - s.grid = grid - N = s.compute_normal(basis="xyz") + grid = LinearGrid(L=8, M=4, N=0, axis=False) + N = s.compute("n_zeta", grid=grid)["n_zeta"] np.testing.assert_allclose(N, np.broadcast_to([0, 1, 0], N.shape), atol=1e-12) @pytest.mark.unit def test_misc(self): """Test getting/setting surface attributes.""" c = ZernikeRZToroidalSection() - grid = LinearGrid(L=2, M=2, N=0) - c.grid = grid - assert grid.eq(c.grid) R, Z = c.get_coeffs(0, 0) np.testing.assert_allclose(R, 10) @@ -212,12 +201,19 @@ def test_curvature(self): """ s = ZernikeRZToroidalSection() grid = LinearGrid(theta=np.pi / 2, rho=0.5) - s.grid = grid - K, H, k1, k2 = s.compute_curvature() - np.testing.assert_allclose(K, 0) - np.testing.assert_allclose(H, 0) - np.testing.assert_allclose(k1, 0) - np.testing.assert_allclose(k2, 0) + data = s.compute( + [ + "curvature_K_zeta", + "curvature_H_zeta", + "curvature_k1_zeta", + "curvature_k2_zeta", + ], + grid=grid, + ) + np.testing.assert_allclose(data["curvature_K_zeta"], 0) + np.testing.assert_allclose(data["curvature_H_zeta"], 0) + np.testing.assert_allclose(data["curvature_k1_zeta"], 0) + np.testing.assert_allclose(data["curvature_k2_zeta"], 0) @pytest.mark.unit From 8f9b598201c0dd1d9719ec2d0efa41f782fff2ee Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 15 Jul 2023 17:33:28 -0400 Subject: [PATCH 23/28] Add basis kwarg to surface compute stuff --- desc/compute/_surface.py | 26 ++++++++++++++++++++++++++ 1 file changed, 26 insertions(+) diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index 949eab733c..c8bd99da2a 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -21,6 +21,7 @@ coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _x_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"]) @@ -48,6 +49,7 @@ def _x_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_rho_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -72,6 +74,7 @@ def _e_rho_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_theta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dt=1) @@ -101,6 +104,7 @@ def _e_theta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_zeta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_lmn"], dz=0) @@ -129,6 +133,7 @@ def _e_zeta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwarg coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_rho_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -151,6 +156,7 @@ def _e_rho_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_rho_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -173,6 +179,7 @@ def _e_rho_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_rho_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -195,6 +202,7 @@ def _e_rho_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -219,6 +227,7 @@ def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kw coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_theta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dt=2) @@ -248,6 +257,7 @@ def _e_theta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kw coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_theta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): dR = transforms["R"].transform(params["R_lmn"], dt=1, dz=1) @@ -275,6 +285,7 @@ def _e_theta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kw coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -299,6 +310,7 @@ def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_zeta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): dR = transforms["R"].transform(params["R_lmn"], dt=1, dz=1) @@ -328,6 +340,7 @@ def _e_zeta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa coordinates="tz", data=[], parameterization="desc.geometry.FourierRZToroidalSurface", + basis="basis", ) def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): R0 = transforms["R"].transform(params["R_lmn"], dz=0) @@ -359,6 +372,7 @@ def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _x_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"]) @@ -388,6 +402,7 @@ def _x_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_rho_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dr=1) @@ -417,6 +432,7 @@ def _e_rho_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_theta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dt=1) @@ -444,6 +460,7 @@ def _e_theta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_zeta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -468,6 +485,7 @@ def _e_zeta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwarg coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_rho_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dr=2) @@ -497,6 +515,7 @@ def _e_rho_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_rho_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dr=1, dt=1) @@ -524,6 +543,7 @@ def _e_rho_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_rho_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -548,6 +568,7 @@ def _e_rho_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_theta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dr=1, dt=1) @@ -577,6 +598,7 @@ def _e_theta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kw coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_theta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): R = transforms["R"].transform(params["R_lmn"], dt=2) @@ -604,6 +626,7 @@ def _e_theta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kw coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_theta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -626,6 +649,7 @@ def _e_theta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kw coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_zeta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -648,6 +672,7 @@ def _e_zeta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwa coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_zeta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) @@ -670,6 +695,7 @@ def _e_zeta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwa coordinates="rt", data=[], parameterization="desc.geometry.ZernikeRZToroidalSection", + basis="basis", ) def _e_zeta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): coords = jnp.zeros((transforms["grid"].num_nodes, 3)) From ec7b65747feeb4f6248d47553bb55f2a489ecc20 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 15 Jul 2023 18:14:00 -0400 Subject: [PATCH 24/28] Update coils to new compute API --- desc/coils.py | 125 +++++++++++++++++++++++-------------- desc/compute/data_index.py | 21 +++++++ tests/test_coils.py | 67 +++++++++++--------- 3 files changed, 135 insertions(+), 78 deletions(-) diff --git a/desc/coils.py b/desc/coils.py index 1b43aae624..099449c240 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -46,7 +46,7 @@ def current(self, new): assert jnp.isscalar(new) or new.size == 1 self._current = new - def compute_magnetic_field(self, coords, params={}, basis="rpz"): + def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): """Compute magnetic field at a set of points. The coil is discretized into a series of straight line segments, using @@ -64,6 +64,9 @@ def compute_magnetic_field(self, coords, params={}, basis="rpz"): parameters to pass to curve basis : {"rpz", "xyz"} basis for input coordinates and returned magnetic field + grid : Grid, int or None + Grid used to discretize coil. If an integer, uses that many equally spaced + points. Returns ------- @@ -76,8 +79,11 @@ def compute_magnetic_field(self, coords, params={}, basis="rpz"): coords = jnp.atleast_2d(coords) if basis == "rpz": coords = rpz2xyz(coords) - current = params.pop("current", self.current) - coil_coords = self.compute_coordinates(**params, basis="xyz") + if params is None: + current = self.current + else: + current = params.pop("current", self.current) + coil_coords = self.compute("x", params=params, grid=grid, basis="xyz")["x"] B = biot_savart(coords, coil_coords, current) if basis == "rpz": B = xyz2rpz_vec(B, x=coords[:, 0], y=coords[:, 1]) @@ -110,8 +116,6 @@ class FourierRZCoil(Coil, FourierRZCurve): number of field periods sym : bool whether to enforce stellarator symmetry - grid : Grid - default grid for computation name : str name for this coil """ @@ -127,10 +131,9 @@ def __init__( modes_Z=None, NFP=1, sym="auto", - grid=None, name="", ): - super().__init__(current, R_n, Z_n, modes_R, modes_Z, NFP, sym, grid, name) + super().__init__(current, R_n, Z_n, modes_R, modes_Z, NFP, sym, name) class FourierXYZCoil(Coil, FourierXYZCurve): @@ -144,8 +147,6 @@ class FourierXYZCoil(Coil, FourierXYZCurve): fourier coefficients for X, Y, Z modes : array-like mode numbers associated with X_n etc. - grid : Grid - default grid or computation name : str name for this coil @@ -160,10 +161,9 @@ def __init__( Y_n=[0, 0, 0], Z_n=[-2, 0, 0], modes=None, - grid=None, name="", ): - super().__init__(current, X_n, Y_n, Z_n, modes, grid, name) + super().__init__(current, X_n, Y_n, Z_n, modes, name) class FourierPlanarCoil(Coil, FourierPlanarCurve): @@ -185,8 +185,6 @@ class FourierPlanarCoil(Coil, FourierPlanarCurve): fourier coefficients for radius from center as function of polar angle modes : array-like mode numbers associated with r_n - grid : Grid - default grid for computation name : str name for this coil @@ -201,10 +199,9 @@ def __init__( normal=[0, 1, 0], r_n=2, modes=None, - grid=None, name="", ): - super().__init__(current, center, normal, r_n, modes, grid, name) + super().__init__(current, center, normal, r_n, modes, name) class CoilSet(Coil, MutableSequence): @@ -251,35 +248,65 @@ def current(self, new): for coil, cur in zip(self.coils, new): coil.current = cur - @property - def grid(self): - """Grid: nodes for computation.""" - return self.coils[0].grid - - @grid.setter - def grid(self, new): - for coil in self.coils: - coil.grid = new - - def compute_coordinates(self, *args, **kwargs): - """Compute real space coordinates using underlying curve method.""" - return [coil.compute_coordinates(*args, **kwargs) for coil in self.coils] - - def compute_frenet_frame(self, *args, **kwargs): - """Compute Frenet frame using underlying curve method.""" - return [coil.compute_frenet_frame(*args, **kwargs) for coil in self.coils] + def _make_arraylike(self, x): + if isinstance(x, dict): + x = [x] * len(self) + try: + len(x) + except TypeError: + x = [x] * len(self) + assert len(x) == len(self) + return x + + def compute( + self, + names, + grid=None, + params=None, + transforms=None, + data=None, + **kwargs, + ): + """Compute the quantity given by name on grid, for each coil in the coilset. - def compute_curvature(self, *args, **kwargs): - """Compute curvature using underlying curve method.""" - return [coil.compute_curvature(*args, **kwargs) for coil in self.coils] + Parameters + ---------- + names : str or array-like of str + Name(s) of the quantity(s) to compute. + grid : Grid or int or array-like, optional + Grid of coordinates to evaluate at. Defaults to a Linear grid. + If an integer, uses that many equally spaced points. + If array-like, should be 1 value per coil. + params : dict of ndarray or array-like + Parameters from the equilibrium. Defaults to attributes of self. + If array-like, should be 1 value per coil. + transforms : dict of Transform or array-like + Transforms for R, Z, lambda, etc. Default is to build from grid. + If array-like, should be 1 value per coil. + data : dict of ndarray or array-like + Data computed so far, generally output from other compute functions + If array-like, should be 1 value per coil. - def compute_torsion(self, *args, **kwargs): - """Compute torsion using underlying curve method.""" - return [coil.compute_torsion(*args, **kwargs) for coil in self.coils] + Returns + ------- + data : list of dict of ndarray + Computed quantity and intermediate variables, for each coil in the set. + List entries map to coils in coilset, each dict contains data for an + individual coil. - def compute_length(self, *args, **kwargs): - """Compute the length of the curve using underlying curve method.""" - return [coil.compute_length(*args, **kwargs) for coil in self.coils] + """ + grid = self._make_arraylike(grid) + params = self._make_arraylike(params) + transforms = self._make_arraylike(transforms) + data = self._make_arraylike(data) + return [ + coil.compute( + names, grid=grd, params=par, transforms=tran, data=dat, **kwargs + ) + for (coil, grd, par, tran, dat) in zip( + self.coils, grid, params, transforms, data + ) + ] def translate(self, *args, **kwargs): """Translate the coils along an axis.""" @@ -293,7 +320,7 @@ def flip(self, *args, **kwargs): """Flip the coils across a plane.""" [coil.flip(*args, **kwargs) for coil in self.coils] - def compute_magnetic_field(self, coords, params={}, basis="rpz"): + def compute_magnetic_field(self, coords, params=None, basis="rpz", grid=None): """Compute magnetic field at a set of points. Parameters @@ -305,18 +332,22 @@ def compute_magnetic_field(self, coords, params={}, basis="rpz"): or one for each member basis : {"rpz", "xyz"} basis for input coordinates and returned magnetic field + grid : Grid, int or None or array-like, optional + Grid used to discretize coil, either the same for all coils or one for each + member of the coilset. If an integer, uses that many equally spaced + points. Returns ------- field : ndarray, shape(n,3) magnetic field at specified points, in either rpz or xyz coordinates """ - if isinstance(params, dict): - params = [params] * len(self) - assert len(params) == len(self) + params = self._make_arraylike(params) + grid = self._make_arraylike(grid) + B = 0 - for coil, par in zip(self.coils, params): - B += coil.compute_magnetic_field(coords, par, basis) + for coil, par, grd in zip(self.coils, params, grid): + B += coil.compute_magnetic_field(coords, par, basis, grd) return B diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index 3eb9effcdb..024c4ab44c 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -142,6 +142,27 @@ def _decorator(func): "desc.geometry.Surface", "Surface", ], + "desc.coils.FourierRZCoil": [ + "desc.geometry.FourierRZCurve", + "FourierRZCurve", + "desc.geometry.core.Curve", + "desc.geometry.Curve", + "Curve", + ], + "desc.coils.FourierXYZCoil": [ + "desc.geometry.FourierXYZCurve", + "FourierXYZCurve", + "desc.geometry.core.Curve", + "desc.geometry.Curve", + "Curve", + ], + "desc.coils.FourierPlanarCoil": [ + "desc.geometry.FourierPlanarCurve", + "FourierPlanarCurve", + "desc.geometry.core.Curve", + "desc.geometry.Curve", + "Curve", + ], } data_index = {p: {} for p in _aliases.keys()} diff --git a/tests/test_coils.py b/tests/test_coils.py index f82e918f59..58391f9ad4 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -20,10 +20,9 @@ def test_biot_savart(self): By_true = 1e-7 * 2 * np.pi * R**2 * I / (y**2 + R**2) ** (3 / 2) B_true = np.array([0, By_true, 0]) coil = FourierXYZCoil(I) - coil.grid = LinearGrid(zeta=100, endpoint=True) - assert coil.grid.num_nodes == 100 + grid = LinearGrid(zeta=100, endpoint=True) B_approx = coil.compute_magnetic_field( - Grid([[10, y, 0], [10, -y, 0]]), basis="xyz" + Grid([[10, y, 0], [10, -y, 0]]), basis="xyz", grid=grid )[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @@ -55,9 +54,7 @@ def test_linspaced_linear(self): ) coils.current = I np.testing.assert_allclose(coils.current, I) - coils.grid = 32 - assert coils.grid.N == 32 - B_approx = coils.compute_magnetic_field([0, 0, z[-1]], basis="xyz")[0] + B_approx = coils.compute_magnetic_field([0, 0, z[-1]], basis="xyz", grid=32)[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @pytest.mark.unit @@ -71,9 +68,7 @@ def test_linspaced_angular(self): coil = FourierPlanarCoil() coil.current = I coils = CoilSet.linspaced_angular(coil, n=N) - coils.grid = 32 - assert all([coil.grid.N == 32 for coil in coils]) - B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz")[0] + B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz", grid=32)[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @pytest.mark.unit @@ -87,9 +82,7 @@ def test_from_symmetry(self): coil = FourierPlanarCoil() coils = CoilSet.linspaced_angular(coil, angle=np.pi / 2, n=N // 4) coils = CoilSet.from_symmetry(coils, NFP=4) - coils.grid = 32 - assert all([coil.grid.N == 32 for coil in coils]) - B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz")[0] + B_approx = coils.compute_magnetic_field([10, 0, 0], basis="rpz", grid=32)[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) # with stellarator symmetry @@ -99,10 +92,8 @@ def test_from_symmetry(self): coils = CoilSet.linspaced_angular( coil, I, [0, 0, 1], np.pi / NFP, N // NFP // 2 ) - coils.grid = 32 - assert coils.grid.N == 32 coils2 = CoilSet.from_symmetry(coils, NFP, True) - B_approx = coils2.compute_magnetic_field([10, 0, 0], basis="rpz")[0] + B_approx = coils2.compute_magnetic_field([10, 0, 0], basis="rpz", grid=32)[0] np.testing.assert_allclose(B_true, B_approx, rtol=1e-3, atol=1e-10) @pytest.mark.unit @@ -110,9 +101,20 @@ def test_properties(self): """Test getting/setting of CoilSet attributes.""" coil = FourierPlanarCoil() coils = CoilSet.linspaced_linear(coil, n=4) - coils.grid = np.array([[0.0, 0.0, 0.0]]) + data = coils.compute( + [ + "x", + "curvature", + "torsion", + "frenet_tangent", + "frenet_normal", + "frenet_binormal", + ], + grid=0, + basis="xyz", + ) np.testing.assert_allclose( - coils.compute_coordinates(), + [dat["x"] for dat in data], np.array( [ [12, 0, 0], @@ -122,12 +124,11 @@ def test_properties(self): ] ).reshape((4, 1, 3)), ) - np.testing.assert_allclose(coils.compute_curvature(), 1 / 2) - np.testing.assert_allclose(coils.compute_torsion(), 0) - TNB = coils.compute_frenet_frame(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz") - T = [foo[0] for foo in TNB] - N = [foo[1] for foo in TNB] - B = [foo[2] for foo in TNB] + np.testing.assert_allclose([dat["curvature"] for dat in data], 1 / 2) + np.testing.assert_allclose([dat["torsion"] for dat in data], 0) + T = [dat["frenet_tangent"] for dat in data] + N = [dat["frenet_normal"] for dat in data] + B = [dat["frenet_binormal"] for dat in data] np.testing.assert_allclose( T, np.array( @@ -164,16 +165,20 @@ def test_properties(self): ).reshape((4, 1, 3)), atol=1e-12, ) - coils.grid = 32 - np.testing.assert_allclose(coils.compute_length(), 2 * 2 * np.pi) + data = coils.compute("length", grid=32) + np.testing.assert_allclose([dat["length"] for dat in data], 2 * 2 * np.pi) coils.translate([1, 1, 1]) - np.testing.assert_allclose(coils.compute_length(), 2 * 2 * np.pi) + data = coils.compute("length", grid=32) + np.testing.assert_allclose([dat["length"] for dat in data], 2 * 2 * np.pi) coils.flip([1, 0, 0]) - coils.grid = np.array([[0.0, 0.0, 0.0]]) - TNB = coils.compute_frenet_frame(grid=np.array([[0.0, 0.0, 0.0]]), basis="xyz") - T = [foo[0] for foo in TNB] - N = [foo[1] for foo in TNB] - B = [foo[2] for foo in TNB] + data = coils.compute( + ["frenet_tangent", "frenet_normal", "frenet_binormal"], + grid=0, + basis="xyz", + ) + T = [dat["frenet_tangent"] for dat in data] + N = [dat["frenet_normal"] for dat in data] + B = [dat["frenet_binormal"] for dat in data] np.testing.assert_allclose( T, np.array( From 47f44dc50d812157164163b00e5dbfb605ca9fec Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Fri, 21 Jul 2023 14:27:50 -0400 Subject: [PATCH 25/28] Remove old io attributes --- desc/geometry/surface.py | 4 ---- 1 file changed, 4 deletions(-) diff --git a/desc/geometry/surface.py b/desc/geometry/surface.py index 521e3d5c0a..fe3992ac81 100644 --- a/desc/geometry/surface.py +++ b/desc/geometry/surface.py @@ -47,8 +47,6 @@ class FourierRZToroidalSurface(Surface): "_Z_lmn", "_R_basis", "_Z_basis", - "_R_transform", - "_Z_transform", "rho", "_NFP", ] @@ -398,8 +396,6 @@ class ZernikeRZToroidalSection(Surface): "_Z_lmn", "_R_basis", "_Z_basis", - "_R_transform", - "_Z_transform", "zeta", "_spectral_indexing", ] From b9f3954b83861fa1a6906929e7ce0bfa7e77bc39 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Fri, 28 Jul 2023 15:31:56 -0400 Subject: [PATCH 26/28] Update parameterization names --- desc/compute/_basis_vectors.py | 12 ++++++++ desc/compute/_curve.py | 38 ++++++++++++------------- desc/compute/_surface.py | 52 +++++++++++++++++----------------- 3 files changed, 57 insertions(+), 45 deletions(-) diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 22b74c55c5..a8a7e1eb32 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -666,6 +666,10 @@ def _b(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_theta", "e_zeta", "|e_theta x e_zeta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _n_rho(params, transforms, profiles, data, **kwargs): # equal to e^rho / |e^rho| but works correctly for surfaces as well that don't have @@ -688,6 +692,10 @@ def _n_rho(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_zeta", "|e_zeta x e_rho|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _n_theta(params, transforms, profiles, data, **kwargs): data["n_theta"] = ( @@ -708,6 +716,10 @@ def _n_theta(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["e_rho", "e_theta", "|e_rho x e_theta|"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _n_zeta(params, transforms, profiles, data, **kwargs): data["n_zeta"] = ( diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index c36f09e424..12e8557023 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -17,7 +17,7 @@ profiles=[], coordinates="s", data=[], - parameterization="desc.geometry.Curve", + parameterization="desc.geometry.core.Curve", ) def _s(params, transforms, profiles, data, **kwargs): data["s"] = transforms["grid"].nodes[:, 2] @@ -54,7 +54,7 @@ def _rotation_matrix_from_normal(normal): profiles=[], coordinates="s", data=["s"], - parameterization="desc.geometry.FourierPlanarCurve", + parameterization="desc.geometry.curve.FourierPlanarCurve", basis="basis", ) def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): @@ -90,7 +90,7 @@ def _x_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=["s"], - parameterization="desc.geometry.FourierPlanarCurve", + parameterization="desc.geometry.curve.FourierPlanarCurve", basis="basis", ) def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): @@ -132,7 +132,7 @@ def _x_s_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=["s"], - parameterization="desc.geometry.FourierPlanarCurve", + parameterization="desc.geometry.curve.FourierPlanarCurve", basis="basis", ) def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): @@ -179,7 +179,7 @@ def _x_ss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=["s"], - parameterization="desc.geometry.FourierPlanarCurve", + parameterization="desc.geometry.curve.FourierPlanarCurve", basis="basis", ) def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): @@ -235,7 +235,7 @@ def _x_sss_FourierPlanarCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=[], - parameterization="desc.geometry.FourierRZCurve", + parameterization="desc.geometry.curve.FourierRZCurve", basis="basis", ) def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs): @@ -269,7 +269,7 @@ def _x_FourierRZCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=[], - parameterization="desc.geometry.FourierRZCurve", + parameterization="desc.geometry.curve.FourierRZCurve", basis="basis", ) def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): @@ -304,7 +304,7 @@ def _x_s_FourierRZCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=[], - parameterization="desc.geometry.FourierRZCurve", + parameterization="desc.geometry.curve.FourierRZCurve", basis="basis", ) def _x_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): @@ -343,7 +343,7 @@ def _x_ss_FourierRZCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=[], - parameterization="desc.geometry.FourierRZCurve", + parameterization="desc.geometry.curve.FourierRZCurve", basis="basis", ) def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): @@ -383,7 +383,7 @@ def _x_sss_FourierRZCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=[], - parameterization="desc.geometry.FourierXYZCurve", + parameterization="desc.geometry.curve.FourierXYZCurve", basis="basis", ) def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs): @@ -416,7 +416,7 @@ def _x_FourierXYZCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=[], - parameterization="desc.geometry.FourierXYZCurve", + parameterization="desc.geometry.curve.FourierXYZCurve", basis="basis", ) def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): @@ -453,7 +453,7 @@ def _x_s_FourierXYZCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=[], - parameterization="desc.geometry.FourierXYZCurve", + parameterization="desc.geometry.curve.FourierXYZCurve", basis="basis", ) def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): @@ -490,7 +490,7 @@ def _x_ss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=[], - parameterization="desc.geometry.FourierXYZCurve", + parameterization="desc.geometry.curve.FourierXYZCurve", basis="basis", ) def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): @@ -521,7 +521,7 @@ def _x_sss_FourierXYZCurve(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=["x_s"], - parameterization="desc.geometry.Curve", + parameterization="desc.geometry.core.Curve", ) def _frenet_tangent(params, transforms, profiles, data, **kwargs): data["frenet_tangent"] = ( @@ -542,7 +542,7 @@ def _frenet_tangent(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=["x_ss"], - parameterization="desc.geometry.Curve", + parameterization="desc.geometry.core.Curve", ) def _frenet_normal(params, transforms, profiles, data, **kwargs): data["frenet_normal"] = ( @@ -563,7 +563,7 @@ def _frenet_normal(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=["frenet_tangent", "frenet_normal"], - parameterization="desc.geometry.Curve", + parameterization="desc.geometry.core.Curve", ) def _frenet_binormal(params, transforms, profiles, data, **kwargs): data["frenet_binormal"] = cross( @@ -584,7 +584,7 @@ def _frenet_binormal(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=["x_s", "x_ss"], - parameterization="desc.geometry.Curve", + parameterization="desc.geometry.core.Curve", ) def _curvature(params, transforms, profiles, data, **kwargs): dxn = jnp.linalg.norm(data["x_s"], axis=-1)[:, jnp.newaxis] @@ -606,7 +606,7 @@ def _curvature(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=["x_s", "x_ss", "x_sss"], - parameterization="desc.geometry.Curve", + parameterization="desc.geometry.core.Curve", ) def _torsion(params, transforms, profiles, data, **kwargs): dxd2x = cross(data["x_s"], data["x_ss"]) @@ -626,7 +626,7 @@ def _torsion(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="s", data=["s", "x_s"], - parameterization="desc.geometry.Curve", + parameterization="desc.geometry.core.Curve", ) def _length(params, transforms, profiles, data, **kwargs): T = jnp.linalg.norm(data["x_s"], axis=-1) diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index c8bd99da2a..18e88457b0 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -20,7 +20,7 @@ profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _x_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -48,7 +48,7 @@ def _x_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_rho_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -73,7 +73,7 @@ def _e_rho_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_theta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -103,7 +103,7 @@ def _e_theta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_zeta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -132,7 +132,7 @@ def _e_zeta_FourierRZToroidalSurface(params, transforms, profiles, data, **kwarg profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_rho_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -155,7 +155,7 @@ def _e_rho_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_rho_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -178,7 +178,7 @@ def _e_rho_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_rho_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -201,7 +201,7 @@ def _e_rho_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwar profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -226,7 +226,7 @@ def _e_theta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kw profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_theta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -256,7 +256,7 @@ def _e_theta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kw profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_theta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -284,7 +284,7 @@ def _e_theta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kw profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -309,7 +309,7 @@ def _e_zeta_r_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_zeta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -339,7 +339,7 @@ def _e_zeta_t_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa profiles=[], coordinates="tz", data=[], - parameterization="desc.geometry.FourierRZToroidalSurface", + parameterization="desc.geometry.surface.FourierRZToroidalSurface", basis="basis", ) def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwargs): @@ -371,7 +371,7 @@ def _e_zeta_z_FourierRZToroidalSurface(params, transforms, profiles, data, **kwa profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _x_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -401,7 +401,7 @@ def _x_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_rho_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -431,7 +431,7 @@ def _e_rho_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_theta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -459,7 +459,7 @@ def _e_theta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_zeta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -484,7 +484,7 @@ def _e_zeta_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwarg profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_rho_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -514,7 +514,7 @@ def _e_rho_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_rho_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -542,7 +542,7 @@ def _e_rho_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_rho_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -567,7 +567,7 @@ def _e_rho_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwar profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_theta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -597,7 +597,7 @@ def _e_theta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kw profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_theta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -625,7 +625,7 @@ def _e_theta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kw profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_theta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -648,7 +648,7 @@ def _e_theta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kw profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_zeta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -671,7 +671,7 @@ def _e_zeta_r_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwa profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_zeta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): @@ -694,7 +694,7 @@ def _e_zeta_t_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwa profiles=[], coordinates="rt", data=[], - parameterization="desc.geometry.ZernikeRZToroidalSection", + parameterization="desc.geometry.surface.ZernikeRZToroidalSection", basis="basis", ) def _e_zeta_z_ZernikeRZToroidalSection(params, transforms, profiles, data, **kwargs): From 88c3b84564bdad27e111fdbbdd42fad98ca900e7 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 5 Aug 2023 20:56:23 -0400 Subject: [PATCH 27/28] Move geometry utils to compute module to avoid circular imports --- desc/coils.py | 2 +- desc/compute/__init__.py | 1 + desc/compute/_curve.py | 2 +- desc/compute/_surface.py | 2 +- desc/{geometry/utils.py => compute/geom_utils.py} | 0 desc/geometry/core.py | 5 ++--- desc/magnetic_fields.py | 2 +- desc/objectives/_geometry.py | 3 +-- tests/test_geometry.py | 2 +- 9 files changed, 9 insertions(+), 10 deletions(-) rename desc/{geometry/utils.py => compute/geom_utils.py} (100%) diff --git a/desc/coils.py b/desc/coils.py index 099449c240..77b84a45bb 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -6,8 +6,8 @@ import numpy as np from desc.backend import jnp +from desc.compute import rpz2xyz, xyz2rpz_vec from desc.geometry import FourierPlanarCurve, FourierRZCurve, FourierXYZCurve -from desc.geometry.utils import rpz2xyz, xyz2rpz_vec from desc.grid import Grid from desc.magnetic_fields import MagneticField, biot_savart diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index feed595ecf..f2e214dcfa 100644 --- a/desc/compute/__init__.py +++ b/desc/compute/__init__.py @@ -41,6 +41,7 @@ _surface, ) from .data_index import data_index +from .geom_utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec from .utils import ( arg_order, compute, diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 12e8557023..d70b409f81 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -1,7 +1,7 @@ from desc.backend import jnp -from desc.geometry.utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec from .data_index import register_compute_fun +from .geom_utils import rpz2xyz, rpz2xyz_vec, xyz2rpz, xyz2rpz_vec from .utils import cross, dot diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index 18e88457b0..c6b3981e72 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -1,7 +1,7 @@ from desc.backend import jnp -from desc.geometry.utils import rpz2xyz, rpz2xyz_vec from .data_index import register_compute_fun +from .geom_utils import rpz2xyz, rpz2xyz_vec @register_compute_fun( diff --git a/desc/geometry/utils.py b/desc/compute/geom_utils.py similarity index 100% rename from desc/geometry/utils.py rename to desc/compute/geom_utils.py diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 12791fc6d4..1550011b40 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -6,13 +6,12 @@ import numpy as np from desc.backend import jnp -from desc.compute.utils import compute as compute_fun +from desc.compute import compute as compute_fun +from desc.compute.geom_utils import reflection_matrix, rotation_matrix from desc.compute.utils import get_params, get_transforms from desc.grid import LinearGrid, QuadratureGrid from desc.io import IOAble -from .utils import reflection_matrix, rotation_matrix - class Curve(IOAble, ABC): """Abstract base class for 1D curves in 3D space.""" diff --git a/desc/magnetic_fields.py b/desc/magnetic_fields.py index 442c10e278..f2315e96a5 100644 --- a/desc/magnetic_fields.py +++ b/desc/magnetic_fields.py @@ -6,8 +6,8 @@ from netCDF4 import Dataset from desc.backend import jit, jnp, odeint +from desc.compute import rpz2xyz_vec, xyz2rpz from desc.derivatives import Derivative -from desc.geometry.utils import rpz2xyz_vec, xyz2rpz from desc.grid import Grid from desc.interpolate import _approx_df, interp2d, interp3d from desc.io import IOAble diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 29ced9281a..058134e612 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -6,8 +6,7 @@ from desc.backend import jnp from desc.compute import compute as compute_fun -from desc.compute import get_params, get_profiles, get_transforms -from desc.geometry.utils import rpz2xyz +from desc.compute import get_params, get_profiles, get_transforms, rpz2xyz from desc.grid import LinearGrid, QuadratureGrid from desc.utils import Timer diff --git a/tests/test_geometry.py b/tests/test_geometry.py index 18622ad3c4..79c3f88040 100644 --- a/tests/test_geometry.py +++ b/tests/test_geometry.py @@ -3,7 +3,7 @@ import numpy as np import pytest -from desc.geometry.utils import ( +from desc.compute.geom_utils import ( rotation_matrix, rpz2xyz, rpz2xyz_vec, From a925e00cd099918163bb5d934f49a132eaeed452 Mon Sep 17 00:00:00 2001 From: Rory Conlin Date: Sat, 5 Aug 2023 21:00:02 -0400 Subject: [PATCH 28/28] Make geometry compute methods use default grid for global quantities --- desc/compute/_curve.py | 2 +- desc/compute/_geometry.py | 2 +- desc/geometry/core.py | 91 +++++++++++++++++++++++++++++++++++++-- 3 files changed, 89 insertions(+), 6 deletions(-) diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index d70b409f81..bb26246519 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -624,7 +624,7 @@ def _torsion(params, transforms, profiles, data, **kwargs): params=[], transforms={}, profiles=[], - coordinates="s", + coordinates="", data=["s", "x_s"], parameterization="desc.geometry.core.Curve", ) diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index a1b82ef769..fbc9561a7f 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -32,7 +32,7 @@ def _V(params, transforms, profiles, data, **kwargs): params=[], transforms={"grid": []}, profiles=[], - coordinates="r", + coordinates="", data=["e_theta", "e_zeta", "x"], parameterization="desc.geometry.surface.FourierRZToroidalSurface", ) diff --git a/desc/geometry/core.py b/desc/geometry/core.py index 1550011b40..82972ca41a 100644 --- a/desc/geometry/core.py +++ b/desc/geometry/core.py @@ -7,8 +7,14 @@ from desc.backend import jnp from desc.compute import compute as compute_fun +from desc.compute import data_index from desc.compute.geom_utils import reflection_matrix, rotation_matrix -from desc.compute.utils import get_params, get_transforms +from desc.compute.utils import ( + _parse_parameterization, + get_data_deps, + get_params, + get_transforms, +) from desc.grid import LinearGrid, QuadratureGrid from desc.io import IOAble @@ -80,6 +86,33 @@ def compute( data = {} profiles = {} + p = _parse_parameterization(self) + deps = list(set(get_data_deps(names, obj=p) + names)) + dep0d = [ + dep + for dep in deps + if (data_index[p][dep]["coordinates"] == "") and (dep not in data) + ] + calc0d = bool(len(dep0d)) + # see if the grid we're already using will work for desired qtys + if calc0d and (grid.N >= 2 * self.N + 5) and isinstance(grid, LinearGrid): + calc0d = False + + if calc0d: + grid0d = LinearGrid(N=2 * self.N + 5, NFP=NFP, endpoint=True) + data0d = compute_fun( + self, + dep0d, + params=params, + transforms=get_transforms(dep0d, obj=self, grid=grid0d, **kwargs), + profiles={}, + data=None, + **kwargs, + ) + # these should all be 0d quantities so don't need to compress/expand + data0d = {key: val for key, val in data0d.items() if key in dep0d} + data.update(data0d) + data = compute_fun( self, names, @@ -222,13 +255,15 @@ def compute( if isinstance(names, str): names = [names] if grid is None: - NFP = self.NFP if hasattr(self, "NFP") else 1 if hasattr(self, "rho"): # constant rho surface grid = LinearGrid( - rho=np.array(self.rho), M=2 * self.M + 5, N=2 * self.N + 5, NFP=NFP + rho=np.array(self.rho), + M=2 * self.M + 5, + N=2 * self.N + 5, + NFP=self.NFP, ) elif hasattr(self, "zeta"): # constant zeta surface - grid = QuadratureGrid(L=2 * self.L + 5, M=2 * self.M + 5, N=0, NFP=NFP) + grid = QuadratureGrid(L=2 * self.L + 5, M=2 * self.M + 5, N=0, NFP=1) grid._nodes[:, 2] = self.zeta if params is None: params = get_params(names, obj=self) @@ -238,6 +273,54 @@ def compute( data = {} profiles = {} + p = _parse_parameterization(self) + deps = list(set(get_data_deps(names, obj=p) + names)) + dep0d = [ + dep + for dep in deps + if (data_index[p][dep]["coordinates"] == "") and (dep not in data) + ] + calc0d = bool(len(dep0d)) + # see if the grid we're already using will work for desired qtys + if calc0d and hasattr(self, "rho"): # constant rho surface + if ( + (grid.N >= 2 * self.N + 5) + and (grid.M > 2 * self.M + 5) + and isinstance(grid, LinearGrid) + ): + calc0d = False + else: + grid0d = LinearGrid( + rho=np.array(self.rho), + M=2 * self.M + 5, + N=2 * self.N + 5, + NFP=self.NFP, + ) + elif calc0d and hasattr(self, "zeta"): # constant zeta surface + if ( + (grid.L >= self.L + 1) + and (grid.M > 2 * self.M + 5) + and isinstance(grid, QuadratureGrid) + ): + calc0d = False + else: + grid0d = QuadratureGrid(L=2 * self.L + 5, M=2 * self.M + 5, N=0, NFP=1) + grid0d._nodes[:, 2] = self.zeta + + if calc0d: + data0d = compute_fun( + self, + dep0d, + params=params, + transforms=get_transforms(dep0d, obj=self, grid=grid0d, **kwargs), + profiles={}, + data=None, + **kwargs, + ) + # these should all be 0d quantities so don't need to compress/expand + data0d = {key: val for key, val in data0d.items() if key in dep0d} + data.update(data0d) + data = compute_fun( self, names,