diff --git a/CHANGELOG.md b/CHANGELOG.md index 705e6c7e10..aab80a4173 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,6 +1,10 @@ Changelog ========= +New Features + +- Add ``use_signed_distance`` flag to ``PlasmaVesselDistance`` which will use a signed distance as the target, which is positive when the plasma is inside of the vessel surface and negative if the plasma is outside of the vessel surface, to allow optimizer to distinguish if the equilbrium surface exits the vessel surface and guard against it by targeting a positive signed distance. + v0.12.1 ------- diff --git a/desc/coils.py b/desc/coils.py index aacb4f19a3..f184365918 100644 --- a/desc/coils.py +++ b/desc/coils.py @@ -1596,7 +1596,9 @@ def flatten_coils(coilset): with open(coilsFilename, "w") as f: f.writelines(lines) - def to_FourierPlanar(self, N=10, grid=None, basis="xyz", name=""): + def to_FourierPlanar( + self, N=10, grid=None, basis="xyz", name="", check_intersection=True + ): """Convert all coils to FourierPlanarCoil. Note that some types of coils may not be representable in this basis. @@ -1614,6 +1616,8 @@ def to_FourierPlanar(self, N=10, grid=None, basis="xyz", name=""): Coordinate system for center and normal vectors. Default = 'xyz'. name : str Name for this coilset. + check_intersection: bool + Whether or not to check the coils in the new coilset for intersections. Returns ------- @@ -1623,9 +1627,17 @@ def to_FourierPlanar(self, N=10, grid=None, basis="xyz", name=""): """ coils = [coil.to_FourierPlanar(N=N, grid=grid, basis=basis) for coil in self] - return self.__class__(*coils, NFP=self.NFP, sym=self.sym, name=name) + return self.__class__( + *coils, + NFP=self.NFP, + sym=self.sym, + name=name, + check_intersection=check_intersection, + ) - def to_FourierRZ(self, N=10, grid=None, NFP=None, sym=False, name=""): + def to_FourierRZ( + self, N=10, grid=None, NFP=None, sym=False, name="", check_intersection=True + ): """Convert all coils to FourierRZCoil representaion. Note that some types of coils may not be representable in this basis. @@ -1644,6 +1656,8 @@ def to_FourierRZ(self, N=10, grid=None, NFP=None, sym=False, name=""): Whether the curve is stellarator-symmetric or not. Default is False. name : str Name for this coilset. + check_intersection: bool + Whether or not to check the coils in the new coiilset for intersections. Returns ------- @@ -1652,9 +1666,15 @@ def to_FourierRZ(self, N=10, grid=None, NFP=None, sym=False, name=""): """ coils = [coil.to_FourierRZ(N=N, grid=grid, NFP=NFP, sym=sym) for coil in self] - return self.__class__(*coils, NFP=self.NFP, sym=self.sym, name=name) + return self.__class__( + *coils, + NFP=self.NFP, + sym=self.sym, + name=name, + check_intersection=check_intersection, + ) - def to_FourierXYZ(self, N=10, grid=None, s=None, name=""): + def to_FourierXYZ(self, N=10, grid=None, s=None, name="", check_intersection=True): """Convert all coils to FourierXYZCoil representation. Parameters @@ -1669,6 +1689,8 @@ def to_FourierXYZ(self, N=10, grid=None, s=None, name=""): normalized arclength. name : str Name for the new CoilSet. + check_intersection: bool + Whether or not to check the coils in the new coiilset for intersections. Returns ------- @@ -1678,9 +1700,17 @@ def to_FourierXYZ(self, N=10, grid=None, s=None, name=""): """ coils = [coil.to_FourierXYZ(N, grid, s) for coil in self] - return self.__class__(*coils, NFP=self.NFP, sym=self.sym, name=name) + return self.__class__( + *coils, + NFP=self.NFP, + sym=self.sym, + name=name, + check_intersection=check_intersection, + ) - def to_SplineXYZ(self, knots=None, grid=None, method="cubic", name=""): + def to_SplineXYZ( + self, knots=None, grid=None, method="cubic", name="", check_intersection=True + ): """Convert all coils to SplineXYZCoil representation. Parameters @@ -1704,6 +1734,8 @@ def to_SplineXYZ(self, knots=None, grid=None, method="cubic", name=""): - `'catmull-rom'`: C1 cubic centripetal "tension" splines name : str Name for the new CoilSet. + check_intersection: bool + Whether or not to check the coils in the new coiilset for intersections. Returns ------- @@ -1712,7 +1744,13 @@ def to_SplineXYZ(self, knots=None, grid=None, method="cubic", name=""): """ coils = [coil.to_SplineXYZ(knots, grid, method) for coil in self] - return self.__class__(*coils, NFP=self.NFP, sym=self.sym, name=name) + return self.__class__( + *coils, + NFP=self.NFP, + sym=self.sym, + name=name, + check_intersection=check_intersection, + ) def is_self_intersecting(self, grid=None, tol=None): """Check if any coils in the CoilSet intersect. @@ -2001,7 +2039,9 @@ def compute_magnetic_field( return B - def to_FourierPlanar(self, N=10, grid=None, basis="xyz", name=""): + def to_FourierPlanar( + self, N=10, grid=None, basis="xyz", name="", check_intersection=False + ): """Convert all coils to FourierPlanarCoil representation. Note that some types of coils may not be representable in this basis. @@ -2019,6 +2059,8 @@ def to_FourierPlanar(self, N=10, grid=None, basis="xyz", name=""): Coordinate system for center and normal vectors. Default = 'xyz'. name : str Name for the new MixedCoilSet. + check_intersection: bool + Whether or not to check the coils in the new coiilset for intersections. Returns ------- @@ -2028,9 +2070,11 @@ def to_FourierPlanar(self, N=10, grid=None, basis="xyz", name=""): """ coils = [coil.to_FourierPlanar(N=N, grid=grid, basis=basis) for coil in self] - return self.__class__(*coils, name=name) + return self.__class__(*coils, name=name, check_intersection=check_intersection) - def to_FourierRZ(self, N=10, grid=None, NFP=None, sym=False, name=""): + def to_FourierRZ( + self, N=10, grid=None, NFP=None, sym=False, name="", check_intersection=True + ): """Convert all coils to FourierRZCoil representation. Note that some types of coils may not be representable in this basis. @@ -2049,6 +2093,8 @@ def to_FourierRZ(self, N=10, grid=None, NFP=None, sym=False, name=""): Whether the curve is stellarator-symmetric or not. Default is False. name : str Name for the new MixedCoilSet. + check_intersection: bool + Whether or not to check the coils in the new coiilset for intersections. Returns ------- @@ -2057,9 +2103,9 @@ def to_FourierRZ(self, N=10, grid=None, NFP=None, sym=False, name=""): """ coils = [coil.to_FourierRZ(N=N, grid=grid, NFP=NFP, sym=sym) for coil in self] - return self.__class__(*coils, name=name) + return self.__class__(*coils, name=name, check_intersection=check_intersection) - def to_FourierXYZ(self, N=10, grid=None, s=None, name=""): + def to_FourierXYZ(self, N=10, grid=None, s=None, name="", check_intersection=True): """Convert all coils to FourierXYZCoil representation. Parameters @@ -2074,6 +2120,8 @@ def to_FourierXYZ(self, N=10, grid=None, s=None, name=""): normalized arclength. name : str Name for the new MixedCoilSet. + check_intersection: bool + Whether or not to check the coils in the new coiilset for intersections. Returns ------- @@ -2083,9 +2131,11 @@ def to_FourierXYZ(self, N=10, grid=None, s=None, name=""): """ coils = [coil.to_FourierXYZ(N, grid, s) for coil in self] - return self.__class__(*coils, name=name) + return self.__class__(*coils, name=name, check_intersection=check_intersection) - def to_SplineXYZ(self, knots=None, grid=None, method="cubic", name=""): + def to_SplineXYZ( + self, knots=None, grid=None, method="cubic", name="", check_intersection=True + ): """Convert all coils to SplineXYZCoil representation. Parameters @@ -2109,6 +2159,8 @@ def to_SplineXYZ(self, knots=None, grid=None, method="cubic", name=""): - `'catmull-rom'`: C1 cubic centripetal "tension" splines name : str Name for the new MixedCoilSet. + check_intersection: bool + Whether or not to check the coils in the new coiilset for intersections. Returns ------- @@ -2117,7 +2169,7 @@ def to_SplineXYZ(self, knots=None, grid=None, method="cubic", name=""): """ coils = [coil.to_SplineXYZ(knots, grid, method) for coil in self] - return self.__class__(*coils, name=name) + return self.__class__(*coils, name=name, check_intersection=check_intersection) def __add__(self, other): if isinstance(other, (CoilSet, MixedCoilSet)): diff --git a/desc/compute/__init__.py b/desc/compute/__init__.py index 416825cc55..87f03068bb 100644 --- a/desc/compute/__init__.py +++ b/desc/compute/__init__.py @@ -94,3 +94,30 @@ def _build_data_index(): _build_data_index() + + +def set_tier(name, p): + """Determine how deep in the dependency tree a given name is. + + tier of 0 means no dependencies on other data, + tier of 1 means it depends on only tier 0 stuff, + tier of 2 means it depends on tier 0 and tier 1, etc etc. + + Designed such that if you compute things in the order determined by tiers, + all dependencies will always be computed in the correct order. + """ + if "tier" in data_index[p][name]: + return + if len(data_index[p][name]["full_with_axis_dependencies"]["data"]) == 0: + data_index[p][name]["tier"] = 0 + else: + thistier = 0 + for name1 in data_index[p][name]["full_with_axis_dependencies"]["data"]: + set_tier(name1, p) + thistier = max(thistier, data_index[p][name1]["tier"]) + data_index[p][name]["tier"] = thistier + 1 + + +for par in data_index.keys(): + for name in data_index[par]: + set_tier(name, par) diff --git a/desc/compute/_basis_vectors.py b/desc/compute/_basis_vectors.py index 414c014182..8fca1346d2 100644 --- a/desc/compute/_basis_vectors.py +++ b/desc/compute/_basis_vectors.py @@ -12,7 +12,7 @@ from desc.backend import jnp from .data_index import register_compute_fun -from .utils import cross, safediv +from .utils import cross, dot, safediv @register_compute_fun( @@ -34,7 +34,7 @@ def _b(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e^rho", + name="e^rho", # ∇ρ is the same in (ρ,θ,ζ) and (ρ,α,ζ) coordinates. label="\\mathbf{e}^{\\rho}", units="m^{-1}", units_long="inverse meters", @@ -927,7 +927,7 @@ def _e_sup_theta_zz(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="e^zeta", + name="e^zeta", # ∇ζ is the same in (ρ,θ,ζ) and (ρ,α,ζ) coordinates. label="\\mathbf{e}^{\\zeta}", units="m^{-1}", units_long="inverse meters", @@ -1431,7 +1431,6 @@ def _e_sub_rho(params, transforms, profiles, data, **kwargs): # At the magnetic axis, this function returns the multivalued map whose # image is the set { 𝐞ᵨ | ρ=0 }. data["e_rho"] = jnp.array([data["R_r"], data["R"] * data["omega_r"], data["Z_r"]]).T - return data @@ -3396,7 +3395,7 @@ def _n_zeta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["e_theta", "e_phi", "phi_t"], + data=["e_theta", "e_phi|r,t", "phi_t"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.surface.FourierRZToroidalSurface", @@ -3404,5 +3403,235 @@ def _n_zeta(params, transforms, profiles, data, **kwargs): ], ) def _e_sub_theta_rp(params, transforms, profiles, data, **kwargs): - data["e_theta|r,p"] = data["e_theta"] - (data["e_phi"].T * data["phi_t"]).T + data["e_theta|r,p"] = data["e_theta"] - (data["e_phi|r,t"].T * data["phi_t"]).T + return data + + +@register_compute_fun( + name="e_rho|a,z", + label="\\mathbf{e}_{\\rho} |_{\\alpha, \\zeta}", + units="m", + units_long="meters", + description="Tangent vector along radial field line label", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_rho", "e_alpha", "alpha_r"], +) +def _e_rho_az(params, transforms, profiles, data, **kwargs): + # constant α and ζ + data["e_rho|a,z"] = ( + data["e_rho"] - data["e_alpha"] * data["alpha_r"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="e_alpha", + label="\\mathbf{e}_{\\alpha}", + units="m", + units_long="meters", + description="Tangent vector along poloidal field line label", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "alpha_t"], +) +def _e_alpha(params, transforms, profiles, data, **kwargs): + # constant ρ and ζ + data["e_alpha"] = data["e_theta"] / data["alpha_t"][:, jnp.newaxis] + return data + + +@register_compute_fun( + name="e_alpha_t", + label="\\partial_{\\theta} \\mathbf{e}_{\\alpha}", + units="m", + units_long="meters", + description="Tangent vector along poloidal field line label, derivative wrt" + " DESC poloidal angle", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "alpha_t", "e_theta_t", "alpha_tt"], +) +def _e_alpha_t(params, transforms, profiles, data, **kwargs): + data["e_alpha_t"] = ( + data["e_theta_t"] / data["alpha_t"][:, jnp.newaxis] + - data["e_theta"] * (data["alpha_tt"] / data["alpha_t"] ** 2)[:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="e_alpha_z", + label="\\partial_{\\zeta} \\mathbf{e}_{\\alpha}", + units="m", + units_long="meters", + description="Tangent vector along poloidal field line label, " + "derivative wrt DESC toroidal angle", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_theta", "alpha_t", "e_theta_z", "alpha_tz"], +) +def _e_alpha_z(params, transforms, profiles, data, **kwargs): + data["e_alpha_z"] = ( + data["e_theta_z"] / data["alpha_t"][:, jnp.newaxis] + - data["e_theta"] * (data["alpha_tz"] / data["alpha_t"] ** 2)[:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="e_zeta|r,a", # Same as B/(B⋅∇ζ). + label="\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha} " + "= \\frac{\\mathbf{B}}{\\mathbf{B} \\cdot \\nabla \\zeta}", + units="m", + units_long="meters", + description="Tangent vector along (collinear to) field line", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta", "e_alpha", "alpha_z"], +) +def _e_zeta_ra(params, transforms, profiles, data, **kwargs): + data["e_zeta|r,a"] = ( + data["e_zeta"] - data["e_alpha"] * data["alpha_z"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="(e_zeta|r,a)_t", + label="\\partial_{\\theta} (\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha})", + units="m", + units_long="meters", + description="Tangent vector along (collinear to) field line, " + "derivative wrt DESC poloidal angle", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta_t", "e_alpha", "alpha_z", "e_alpha_t", "alpha_tz"], +) +def _e_zeta_ra_t(params, transforms, profiles, data, **kwargs): + data["(e_zeta|r,a)_t"] = data["e_zeta_t"] - ( + data["e_alpha_t"] * data["alpha_z"][:, jnp.newaxis] + + data["e_alpha"] * data["alpha_tz"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="(e_zeta|r,a)_a", + label="\\partial_{\\alpha} (\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha})", + units="m", + units_long="meters", + description="Tangent vector along (collinear to) field line, derivative " + "wrt field line poloidal label", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["(e_zeta|r,a)_t", "alpha_t"], +) +def _e_zeta_ra_a(params, transforms, profiles, data, **kwargs): + data["(e_zeta|r,a)_a"] = data["(e_zeta|r,a)_t"] / data["alpha_t"][:, jnp.newaxis] + return data + + +@register_compute_fun( + name="(e_zeta|r,a)_z", + label="\\partial_{\\zeta} (\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha})", + units="m", + units_long="meters", + description="Tangent vector along (collinear to) field line, " + "derivative wrt DESC toroidal angle", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta_z", "e_alpha", "alpha_z", "e_alpha_z", "alpha_zz"], +) +def _e_zeta_ra_z(params, transforms, profiles, data, **kwargs): + data["(e_zeta|r,a)_z"] = data["e_zeta_z"] - ( + data["e_alpha_z"] * data["alpha_z"][:, jnp.newaxis] + + data["e_alpha"] * data["alpha_zz"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="(e_zeta|r,a)_z|r,a", + label="\\partial_{\\zeta} (\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha}) " + "|_{\\rho, \\alpha}", + units="m", + units_long="meters", + description="Curvature vector along field line", + dim=3, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["(e_zeta|r,a)_z", "(e_zeta|r,a)_a", "alpha_z"], +) +def _e_zeta_z_ra(params, transforms, profiles, data, **kwargs): + data["(e_zeta|r,a)_z|r,a"] = ( + data["(e_zeta|r,a)_z"] + - data["(e_zeta|r,a)_a"] * data["alpha_z"][:, jnp.newaxis] + ) + return data + + +@register_compute_fun( + name="|e_zeta|r,a|", # Often written as |dℓ/dζ| = |B/(B⋅∇ζ)|. + label="|\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha}|" + " = \\frac{|\\mathbf{B}|}{|\\mathbf{B} \\cdot \\nabla \\zeta|}", + units="m", + units_long="meters", + description="Differential length along field line", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["e_zeta|r,a"], +) +def _d_ell_d_zeta(params, transforms, profiles, data, **kwargs): + data["|e_zeta|r,a|"] = jnp.linalg.norm(data["e_zeta|r,a"], axis=-1) + return data + + +@register_compute_fun( + name="|e_zeta|r,a|_z|r,a", + label="\\partial_{\\zeta} |\\mathbf{e}_{\\zeta} |_{\\rho, \\alpha}| " + "|_{\\rho, \\alpha}", + units="m", + units_long="meters", + description="Differential length along field line, derivative along field line", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["|e_zeta|r,a|", "(e_zeta|r,a)_z|r,a", "e_zeta|r,a"], +) +def _d_ell_d_zeta_z(params, transforms, profiles, data, **kwargs): + data["|e_zeta|r,a|_z|r,a"] = ( + dot(data["(e_zeta|r,a)_z|r,a"], data["e_zeta|r,a"]) / data["|e_zeta|r,a|"] + ) return data diff --git a/desc/compute/_core.py b/desc/compute/_core.py index a157d91412..7bc6e8acb3 100644 --- a/desc/compute/_core.py +++ b/desc/compute/_core.py @@ -22,12 +22,37 @@ "desc.geometry.core.Surface", "desc.geometry.core.Curve", ], + aliases=["rho_t", "rho_z", "theta_r", "theta_z", "zeta_r", "zeta_t"], ) def _0(params, transforms, profiles, data, **kwargs): data["0"] = jnp.zeros(transforms["grid"].num_nodes) return data +@register_compute_fun( + name="1", + label="1", + units="~", + units_long="None", + description="Ones", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="rtz", + data=[], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + "desc.geometry.core.Curve", + ], + aliases=["rho_r", "theta_t", "zeta_z"], +) +def _1(params, transforms, profiles, data, **kwargs): + data["1"] = jnp.ones(transforms["grid"].num_nodes) + return data + + @register_compute_fun( name="R", label="R", @@ -1475,10 +1500,10 @@ def _Z_zzz(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["theta_PEST", "zeta", "iota"], + data=["theta_PEST", "phi", "iota"], ) def _alpha(params, transforms, profiles, data, **kwargs): - data["alpha"] = (data["theta_PEST"] - data["iota"] * data["zeta"]) % (2 * jnp.pi) + data["alpha"] = (data["theta_PEST"] - data["iota"] * data["phi"]) % (2 * jnp.pi) return data @@ -1518,7 +1543,43 @@ def _alpha_r(params, transforms, profiles, data, **kwargs): data=["theta_PEST_t", "phi_t", "iota"], ) def _alpha_t(params, transforms, profiles, data, **kwargs): - data["alpha_t"] = data["theta_PEST_t"] + data["iota"] * data["phi_t"] + data["alpha_t"] = data["theta_PEST_t"] - data["iota"] * data["phi_t"] + return data + + +@register_compute_fun( + name="alpha_tt", + label="\\partial_{\\theta \\theta} \\alpha", + units="~", + units_long="None", + description="Field line label, second derivative wrt poloidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["theta_PEST_tt", "phi_tt", "iota"], +) +def _alpha_tt(params, transforms, profiles, data, **kwargs): + data["alpha_tt"] = data["theta_PEST_tt"] - data["iota"] * data["phi_tt"] + return data + + +@register_compute_fun( + name="alpha_tz", + label="\\partial_{\\theta \\zeta} \\alpha", + units="~", + units_long="None", + description="Field line label, derivative wrt poloidal and toroidal coordinates", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["theta_PEST_tz", "phi_tz", "iota"], +) +def _alpha_tz(params, transforms, profiles, data, **kwargs): + data["alpha_tz"] = data["theta_PEST_tz"] - data["iota"] * data["phi_tz"] return data @@ -1540,6 +1601,24 @@ def _alpha_z(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="alpha_zz", + label="\\partial_{\\zeta \\zeta} \\alpha", + units="~", + units_long="None", + description="Field line label, second derivative wrt toroidal coordinate", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["theta_PEST_zz", "phi_zz", "iota"], +) +def _alpha_zz(params, transforms, profiles, data, **kwargs): + data["alpha_zz"] = data["theta_PEST_zz"] - data["iota"] * data["phi_zz"] + return data + + @register_compute_fun( name="lambda", label="\\lambda", @@ -2657,6 +2736,10 @@ def _phi(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["omega_r"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _phi_r(params, transforms, profiles, data, **kwargs): data["phi_r"] = data["omega_r"] @@ -2731,6 +2814,10 @@ def _phi_rz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["omega_t"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _phi_t(params, transforms, profiles, data, **kwargs): data["phi_t"] = data["omega_t"] @@ -2787,6 +2874,10 @@ def _phi_tz(params, transforms, profiles, data, **kwargs): profiles=[], coordinates="rtz", data=["omega_z"], + parameterization=[ + "desc.equilibrium.equilibrium.Equilibrium", + "desc.geometry.core.Surface", + ], ) def _phi_z(params, transforms, profiles, data, **kwargs): data["phi_z"] = 1 + data["omega_z"] @@ -2836,75 +2927,6 @@ def _rho(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="rho_r", - label="\\partial_{\\rho} \\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - + "of the toroidal flux, derivative wrt radial coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="r", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _rho_r(params, transforms, profiles, data, **kwargs): - data["rho_r"] = jnp.ones_like(data["0"]) - return data - - -@register_compute_fun( - name="rho_t", - label="\\partial_{\\theta} \\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - "of the toroidal flux, derivative wrt poloidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="r", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _rho_t(params, transforms, profiles, data, **kwargs): - data["rho_t"] = data["0"] - return data - - -@register_compute_fun( - name="rho_z", - label="\\partial_{\\zeta} \\rho", - units="~", - units_long="None", - description="Radial coordinate, proportional to the square root " - "of the toroidal flux, derivative wrt toroidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="r", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _rho_z(params, transforms, profiles, data, **kwargs): - data["rho_z"] = data["0"] - return data - - @register_compute_fun( name="theta", label="\\theta", @@ -2984,90 +3006,78 @@ def _theta_PEST_t(params, transforms, profiles, data, **kwargs): @register_compute_fun( - name="theta_PEST_z", - label="\\partial_{\\zeta} \\vartheta", + name="theta_PEST_tt", + label="\\partial_{\\theta \\theta} \\vartheta", units="rad", units_long="radians", - description="PEST straight field line poloidal angular coordinate, derivative wrt " - "toroidal coordinate", + description="PEST straight field line poloidal angular coordinate, second " + "derivative wrt poloidal coordinate", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["lambda_z"], + data=["lambda_tt"], ) -def _theta_PEST_z(params, transforms, profiles, data, **kwargs): - data["theta_PEST_z"] = data["lambda_z"] +def _theta_PEST_tt(params, transforms, profiles, data, **kwargs): + data["theta_PEST_tt"] = data["lambda_tt"] return data @register_compute_fun( - name="theta_r", - label="\\partial_{\\rho} \\theta", + name="theta_PEST_tz", + label="\\partial_{\\theta \\zeta} \\vartheta", units="rad", units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - "derivative wrt radial coordinate", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "poloidal and toroidal coordinates", dim=1, params=[], transforms={}, profiles=[], - coordinates="t", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], + coordinates="rtz", + data=["lambda_tz"], ) -def _theta_r(params, transforms, profiles, data, **kwargs): - data["theta_r"] = data["0"] +def _theta_PEST_tz(params, transforms, profiles, data, **kwargs): + data["theta_PEST_tz"] = data["lambda_tz"] return data @register_compute_fun( - name="theta_t", - label="\\partial_{\\theta} \\theta", + name="theta_PEST_z", + label="\\partial_{\\zeta} \\vartheta", units="rad", units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " - "derivative wrt poloidal coordinate", + description="PEST straight field line poloidal angular coordinate, derivative wrt " + "toroidal coordinate", dim=1, params=[], transforms={}, profiles=[], - coordinates="t", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], + coordinates="rtz", + data=["lambda_z"], ) -def _theta_t(params, transforms, profiles, data, **kwargs): - data["theta_t"] = jnp.ones_like(data["0"]) +def _theta_PEST_z(params, transforms, profiles, data, **kwargs): + data["theta_PEST_z"] = data["lambda_z"] return data @register_compute_fun( - name="theta_z", - label="\\partial_{\\zeta} \\theta", + name="theta_PEST_zz", + label="\\partial_{\\zeta \\zeta} \\vartheta", units="rad", units_long="radians", - description="Poloidal angular coordinate (geometric, not magnetic), " + description="PEST straight field line poloidal angular coordinate, second " "derivative wrt toroidal coordinate", dim=1, params=[], transforms={}, profiles=[], - coordinates="t", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], + coordinates="rtz", + data=["lambda_zz"], ) -def _theta_z(params, transforms, profiles, data, **kwargs): - data["theta_z"] = data["0"] +def _theta_PEST_zz(params, transforms, profiles, data, **kwargs): + data["theta_PEST_zz"] = data["lambda_zz"] return data @@ -3091,69 +3101,3 @@ def _theta_z(params, transforms, profiles, data, **kwargs): def _zeta(params, transforms, profiles, data, **kwargs): data["zeta"] = transforms["grid"].nodes[:, 2] return data - - -@register_compute_fun( - name="zeta_r", - label="\\partial_{\\rho} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate derivative, wrt radial coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="z", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _zeta_r(params, transforms, profiles, data, **kwargs): - data["zeta_r"] = data["0"] - return data - - -@register_compute_fun( - name="zeta_t", - label="\\partial_{\\theta} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate, derivative wrt poloidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="z", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _zeta_t(params, transforms, profiles, data, **kwargs): - data["zeta_t"] = data["0"] - return data - - -@register_compute_fun( - name="zeta_z", - label="\\partial_{\\zeta} \\zeta", - units="rad", - units_long="radians", - description="Toroidal angular coordinate, derivative wrt toroidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="z", - data=["0"], - parameterization=[ - "desc.equilibrium.equilibrium.Equilibrium", - "desc.geometry.core.Surface", - ], -) -def _zeta_z(params, transforms, profiles, data, **kwargs): - data["zeta_z"] = jnp.ones_like(data["0"]) - return data diff --git a/desc/compute/_curve.py b/desc/compute/_curve.py index 818da18092..2e96e7a767 100644 --- a/desc/compute/_curve.py +++ b/desc/compute/_curve.py @@ -31,7 +31,10 @@ def _s(params, transforms, profiles, data, **kwargs): label="ds", units="~", units_long="None", - description="Spacing of curve parameter", + description=( + "Quadrature weights for integration along the curve," + + " i.e. an alias for ``grid.spacing[:,2]``" + ), dim=1, params=[], transforms={"grid": []}, diff --git a/desc/compute/_field.py b/desc/compute/_field.py index c3c95d86a7..e53f033464 100644 --- a/desc/compute/_field.py +++ b/desc/compute/_field.py @@ -483,6 +483,49 @@ def _B_sup_zeta_z(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="B^zeta_a", + label="\\partial_{\\alpha} B^{\\zeta}", + units="T \\cdot m^{-1}", + units_long="Tesla / meter", + description=( + "Contravariant toroidal component of magnetic field, derivative wrt field" + " line poloidal label" + ), + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["B^zeta_t", "alpha_t"], +) +def _B_sup_zeta_a(params, transforms, profiles, data, **kwargs): + # constant ρ and ζ + data["B^zeta_a"] = data["B^zeta_t"] / data["alpha_t"] + return data + + +@register_compute_fun( + name="B^zeta_z|r,a", + label="\\partial_{\\zeta} B^{\\zeta} |_{\\rho, \\alpha}", + units="T \\cdot m^{-1}", + units_long="Tesla / meter", + description=( + "Contravariant toroidal component of magnetic field, derivative along field" + " line" + ), + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["B^zeta_z", "B^zeta_a", "alpha_z"], +) +def _B_sup_zeta_z_ra(params, transforms, profiles, data, **kwargs): + data["B^zeta_z|r,a"] = data["B^zeta_z"] - data["B^zeta_a"] * data["alpha_z"] + return data + + @register_compute_fun( name="B_z", label="\\partial_{\\zeta} \\mathbf{B}", @@ -2295,6 +2338,43 @@ def _B_mag_z(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="|B|_a", + label="\\partial_{\\alpha} (|\\mathbf{B}|) |_{\\rho, \\zeta}", + units="T", + units_long="Tesla", + description="Magnitude of magnetic field, derivative wrt field line angle", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["|B|_t", "alpha_t"], +) +def _B_mag_alpha(params, transforms, profiles, data, **kwargs): + # constant ρ and ζ + data["|B|_a"] = data["|B|_t"] / data["alpha_t"] + return data + + +@register_compute_fun( + name="|B|_z|r,a", + label="\\partial_{\\zeta} (|\\mathbf{B}|) |_{\\rho, \\alpha}", + units="T", + units_long="Tesla", + description="Magnitude of magnetic field, derivative along field line", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["|B|_z", "|B|_a", "alpha_z"], +) +def _B_mag_z_constant_rho_alpha(params, transforms, profiles, data, **kwargs): + data["|B|_z|r,a"] = data["|B|_z"] - data["|B|_a"] * data["alpha_z"] + return data + + @register_compute_fun( name="|B|_rr", label="\\partial_{\\rho\\rho} |\\mathbf{B}|", @@ -2855,7 +2935,7 @@ def _gradB2(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="|grad(|B|^2)|/2mu0", - label="|\\nabla |B|^{2}/(2\\mu_0)|", + label="|\\nabla |B|^{2}|/(2\\mu_0)", units="N \\cdot m^{-3}", units_long="Newton / cubic meter", description="Magnitude of magnetic pressure gradient", @@ -3322,7 +3402,7 @@ def _kappa(params, transforms, profiles, data, **kwargs): label="\\kappa_n", units="m^{-1}", units_long="Inverse meters", - description="Normal curvature vector of magnetic field lines", + description="Normal curvature of magnetic field lines", dim=1, params=[], transforms={}, @@ -3340,7 +3420,7 @@ def _kappa_n(params, transforms, profiles, data, **kwargs): label="\\kappa_g", units="m^{-1}", units_long="Inverse meters", - description="Geodesic curvature vector of magnetic field lines", + description="Geodesic curvature of magnetic field lines", dim=1, params=[], transforms={}, diff --git a/desc/compute/_geometry.py b/desc/compute/_geometry.py index d4e8e7ff0e..5e38a5c89b 100644 --- a/desc/compute/_geometry.py +++ b/desc/compute/_geometry.py @@ -625,32 +625,14 @@ def _curvature_k1_rho(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["g_tt", "g_tz", "g_zz", "L_sff_rho", "M_sff_rho", "N_sff_rho"], + data=["curvature_k1_rho"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.surface.FourierRZToroidalSurface", ], ) def _curvature_k2_rho(params, transforms, profiles, data, **kwargs): - # following notation from - # https://en.wikipedia.org/wiki/Parametric_surface - E = data["g_tt"] - F = data["g_tz"] - G = data["g_zz"] - L = data["L_sff_rho"] - M = data["M_sff_rho"] - N = data["N_sff_rho"] - a = E * G - F**2 - b = 2 * 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) - # In the axis limit, the matrix of the first fundamental form is singular. - # The diagonal of the shape operator becomes unbounded, - # so the eigenvalues do not exist. - data["curvature_k1_rho"] = jnp.maximum(r1, r2) - data["curvature_k2_rho"] = jnp.minimum(r1, r2) - return data + return data # noqa: unused dependency @register_compute_fun( @@ -804,25 +786,10 @@ def _curvature_k1_theta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["g_rr", "g_rz", "g_zz", "L_sff_theta", "M_sff_theta", "N_sff_theta"], + data=["curvature_k1_theta"], ) def _curvature_k2_theta(params, transforms, profiles, data, **kwargs): - # following notation from - # https://en.wikipedia.org/wiki/Parametric_surface - E = data["g_zz"] - F = data["g_rz"] - G = data["g_rr"] - L = data["L_sff_theta"] - M = data["M_sff_theta"] - N = data["N_sff_theta"] - a = E * G - F**2 - b = 2 * 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) - data["curvature_k1_theta"] = jnp.maximum(r1, r2) - data["curvature_k2_theta"] = jnp.minimum(r1, r2) - return data + return data # noqa: unused dependency @register_compute_fun( @@ -989,32 +956,14 @@ def _curvature_k1_zeta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["g_rr", "g_rt", "g_tt", "L_sff_zeta", "M_sff_zeta", "N_sff_zeta"], + data=["curvature_k1_zeta"], parameterization=[ "desc.equilibrium.equilibrium.Equilibrium", "desc.geometry.surface.ZernikeRZToroidalSection", ], ) def _curvature_k2_zeta(params, transforms, profiles, data, **kwargs): - # following notation from - # https://en.wikipedia.org/wiki/Parametric_surface - E = data["g_rr"] - F = data["g_rt"] - G = data["g_tt"] - L = data["L_sff_zeta"] - M = data["M_sff_zeta"] - N = data["N_sff_zeta"] - a = E * G - F**2 - b = 2 * 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) - # In the axis limit, the matrix of the first fundamental form is singular. - # The diagonal of the shape operator becomes unbounded, - # so the eigenvalues do not exist. - data["curvature_k1_zeta"] = jnp.maximum(r1, r2) - data["curvature_k2_zeta"] = jnp.minimum(r1, r2) - return data + return data # noqa: unused dependency @register_compute_fun( diff --git a/desc/compute/_metric.py b/desc/compute/_metric.py index 3a842e9378..96228ffc06 100644 --- a/desc/compute/_metric.py +++ b/desc/compute/_metric.py @@ -14,7 +14,7 @@ from desc.backend import jnp from .data_index import register_compute_fun -from .utils import cross, dot, safediv, safenorm +from .utils import cross, dot, safediv, safenorm, surface_averages @register_compute_fun( @@ -59,6 +59,27 @@ def _sqrtg_pest(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="sqrt(g)_Clebsch", + label="\\sqrt{g}_{\\text{Clebsch}}", + units="m^{3}", + units_long="cubic meters", + description="Jacobian determinant of Clebsch field line coordinate system (ρ,α,ζ)" + " where ζ is the DESC toroidal coordinate.", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["sqrt(g)", "alpha_t"], +) +def _sqrtg_clebsch(params, transforms, profiles, data, **kwargs): + # Same as dot(data["e_rho|a,z"], cross(data["e_alpha"], data["e_zeta|r,a"])), but + # more efficient as it avoids computing radial derivative of alpha and hence iota. + data["sqrt(g)_Clebsch"] = data["sqrt(g)"] / data["alpha_t"] + return data + + @register_compute_fun( name="|e_theta x e_zeta|", label="|\\mathbf{e}_{\\theta} \\times \\mathbf{e}_{\\zeta}|", @@ -229,6 +250,27 @@ def _e_rho_x_e_theta(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="|e_rho x e_alpha|", + label="|\\mathbf{e}_{\\rho} \\times \\mathbf{e}_{\\alpha}|", + units="m^{2}", + units_long="square meters", + description="2D Jacobian determinant for constant zeta surface in Clebsch " + "field line coordinates (ρ,α,ζ) where ζ is the DESC toroidal coordinate.", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["|e_rho x e_theta|", "alpha_t"], +) +def _e_rho_x_e_alpha(params, transforms, profiles, data, **kwargs): + # Same as safenorm(cross(data["e_rho|a,z"], data["e_alpha"]), axis=-1), but more + # efficient as it avoids computing radial derivative of iota and stream functions. + data["|e_rho x e_alpha|"] = data["|e_rho x e_theta|"] / jnp.abs(data["alpha_t"]) + return data + + @register_compute_fun( name="|e_rho x e_theta|_r", label="\\partial_{\\rho} |\\mathbf{e}_{\\rho} \\times \\mathbf{e}_{\\theta}|", @@ -1331,7 +1373,7 @@ def _g_sup_zz(params, transforms, profiles, data, **kwargs): label="g^{\\rho\\theta}", units="m^{-2}", units_long="inverse square meters", - description="Radial/Poloidal element of contravariant metric tensor", + description="Radial/Poloidal (ρ, θ) element of contravariant metric tensor", dim=1, params=[], transforms={}, @@ -1759,6 +1801,32 @@ def _gradrho(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="<|grad(rho)|>", # same as S(r) / V_r(r) + label="\\langle \\vert \\nabla \\rho \\vert \\rangle", + units="m^{-1}", + units_long="inverse meters", + description="Magnitude of contravariant radial basis vector, flux surface average", + dim=1, + params=[], + transforms={"grid": []}, + profiles=[], + coordinates="r", + data=["|grad(rho)|", "sqrt(g)"], + axis_limit_data=["sqrt(g)_r"], + resolution_requirement="tz", +) +def _gradrho_norm_fsa(params, transforms, profiles, data, **kwargs): + data["<|grad(rho)|>"] = surface_averages( + transforms["grid"], + data["|grad(rho)|"], + sqrt_g=transforms["grid"].replace_at_axis( + data["sqrt(g)"], lambda: data["sqrt(g)_r"], copy=True + ), + ) + return data + + @register_compute_fun( name="|grad(psi)|", label="|\\nabla\\psi|", @@ -1847,12 +1915,12 @@ def _gradzeta(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["|B|", "b", "grad(alpha)", "grad(|B|)"], + data=["|B|^2", "b", "grad(alpha)", "grad(|B|)"], ) def _gbdrift(params, transforms, profiles, data, **kwargs): data["gbdrift"] = ( 1 - / data["|B|"] ** 2 + / data["|B|^2"] * dot(data["b"], cross(data["grad(|B|)"], data["grad(alpha)"])) ) return data @@ -1874,11 +1942,11 @@ def _gbdrift(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="rtz", - data=["p_r", "psi_r", "|B|", "gbdrift"], + data=["p_r", "psi_r", "|B|^2", "gbdrift"], ) def _cvdrift(params, transforms, profiles, data, **kwargs): dp_dpsi = mu_0 * data["p_r"] / data["psi_r"] - data["cvdrift"] = 1 / data["|B|"] ** 2 * dp_dpsi + data["gbdrift"] + data["cvdrift"] = 1 / data["|B|^2"] * dp_dpsi + data["gbdrift"] return data @@ -1892,16 +1960,16 @@ def _cvdrift(params, transforms, profiles, data, **kwargs): units="1/(T-m^{2})", units_long="inverse Tesla meters^2", description="Radial component of the geometric part of the curvature drift" - + " used for local stability analyses, Gamma_c, epsilon_eff etc.", + + " used for local stability analyses, gyrokinetics, Gamma_c.", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", - data=["|B|", "b", "e^rho", "grad(|B|)"], + data=["|B|^2", "b", "e^rho", "grad(|B|)"], ) def _cvdrift0(params, transforms, profiles, data, **kwargs): data["cvdrift0"] = ( - 1 / data["|B|"] ** 2 * (dot(data["b"], cross(data["grad(|B|)"], data["e^rho"]))) + 1 / data["|B|^2"] * (dot(data["b"], cross(data["grad(|B|)"], data["e^rho"]))) ) return data diff --git a/desc/compute/_omnigenity.py b/desc/compute/_omnigenity.py index 469beae364..892800d48e 100644 --- a/desc/compute/_omnigenity.py +++ b/desc/compute/_omnigenity.py @@ -429,22 +429,21 @@ def _omni_angle(params, transforms, profiles, data, **kwargs): @register_compute_fun( name="theta_B", - label="(\\theta_{B},\\zeta_{B})", + label="\\theta_{B}", units="rad", units_long="radians", - description="Boozer angular coordinates", + description="Boozer poloidal angle", dim=1, params=[], transforms={}, profiles=[], coordinates="rtz", data=["alpha", "h"], - aliases=["zeta_B"], parameterization="desc.magnetic_fields._core.OmnigenousField", helicity="tuple: Type of quasisymmetry, (M,N). Default (1,0)", iota="float: Value of rotational transform on the Omnigenous surface. Default 1.0", ) -def _omni_map(params, transforms, profiles, data, **kwargs): +def _omni_map_theta_B(params, transforms, profiles, data, **kwargs): M, N = kwargs.get("helicity", (1, 0)) iota = kwargs.get("iota", 1) @@ -476,6 +475,24 @@ def _omni_map(params, transforms, profiles, data, **kwargs): return data +@register_compute_fun( + name="zeta_B", + label="\\zeta_{B}", + units="rad", + units_long="radians", + description="Boozer toroidal angle", + dim=1, + params=[], + transforms={}, + profiles=[], + coordinates="rtz", + data=["theta_B"], + parameterization="desc.magnetic_fields._core.OmnigenousField", +) +def _omni_map_zeta_B(params, transforms, profiles, data, **kwargs): + return data # noqa: unused dependency + + @register_compute_fun( name="|B|", label="|\\mathbf{B}|", diff --git a/desc/compute/_profiles.py b/desc/compute/_profiles.py index ccb701f072..1cfbc2f23f 100644 --- a/desc/compute/_profiles.py +++ b/desc/compute/_profiles.py @@ -82,10 +82,10 @@ def _psi_r(params, transforms, profiles, data, **kwargs): transforms={}, profiles=[], coordinates="r", - data=["rho"], + data=["1"], ) def _psi_rr(params, transforms, profiles, data, **kwargs): - data["psi_rr"] = params["Psi"] * jnp.ones_like(data["rho"]) / jnp.pi + data["psi_rr"] = data["1"] * params["Psi"] / jnp.pi return data diff --git a/desc/compute/_surface.py b/desc/compute/_surface.py index 5fb84fa060..d37f82337e 100644 --- a/desc/compute/_surface.py +++ b/desc/compute/_surface.py @@ -118,66 +118,6 @@ def _phi_Surface(params, transforms, profiles, data, **kwargs): return data -@register_compute_fun( - name="phi_r", - label="\\partial_{\\rho} \\phi", - units="rad", - units_long="radians", - description="Toroidal angle in lab frame, derivative wrt radial coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["e_rho"], - parameterization="desc.geometry.core.Surface", -) -def _phi_r_Surface(params, transforms, profiles, data, **kwargs): - coords_r = data["e_rho"] - data["phi_r"] = coords_r[:, 1] - return data - - -@register_compute_fun( - name="phi_t", - label="\\partial_{\\theta} \\phi", - units="rad", - units_long="radians", - description="Toroidal angle in lab frame, derivative wrt poloidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["e_theta"], - parameterization="desc.geometry.core.Surface", -) -def _phi_t_Surface(params, transforms, profiles, data, **kwargs): - coords_t = data["e_theta"] - data["phi_t"] = coords_t[:, 1] - return data - - -@register_compute_fun( - name="phi_z", - label="\\partial_{\\zeta} \\phi", - units="rad", - units_long="radians", - description="Toroidal angle in lab frame, derivative wrt toroidal coordinate", - dim=1, - params=[], - transforms={}, - profiles=[], - coordinates="rtz", - data=["e_zeta"], - parameterization="desc.geometry.core.Surface", -) -def _phi_z_Surface(params, transforms, profiles, data, **kwargs): - coords_z = data["e_zeta"] - data["phi_z"] = coords_z[:, 1] - return data - - @register_compute_fun( name="Z", label="Z", diff --git a/desc/compute/data_index.py b/desc/compute/data_index.py index bfd408487c..26341ec587 100644 --- a/desc/compute/data_index.py +++ b/desc/compute/data_index.py @@ -105,7 +105,11 @@ def register_compute_fun( # noqa: C901 or `'desc.equilibrium.Equilibrium'`. resolution_requirement : str Resolution requirements in coordinates. I.e. "r" expects radial resolution - in the grid, "rtz" expects a grid with radial, poloidal, and toroidal resolution. + in the grid. Likewise, "rtz" is shorthand for "rho, theta, zeta" and indicates + the computation expects a grid with radial, poloidal, and toroidal resolution. + If the computation simply performs pointwise operations, instead of a + reduction (such as integration) over a coordinate, then an empty string may + be used to indicate no requirements. source_grid_requirement : dict Attributes of the source grid that the compute function requires. Also assumes dependencies were computed on such a grid. @@ -128,6 +132,8 @@ def register_compute_fun( # noqa: C901 source_grid_requirement = {} if not isinstance(parameterization, (tuple, list)): parameterization = [parameterization] + if not isinstance(aliases, (tuple, list)): + aliases = [aliases] deps = { "params": params, @@ -172,10 +178,11 @@ def _decorator(func): if name in data_index[base_class]: if p == data_index[base_class][name]["parameterization"]: raise ValueError( - f"Already registered function with parameterization {p} and name {name}." + f"Already registered function with parameterization {p}" + f" and name {name}." ) - # if it was already registered from a parent class, we prefer - # the child class. + # if it was already registered from a parent class, we + # prefer the child class. inheritance_order = [base_class] + superclasses if inheritance_order.index(p) > inheritance_order.index( data_index[base_class][name]["parameterization"] diff --git a/desc/compute/utils.py b/desc/compute/utils.py index 4567637be3..92c41a000f 100644 --- a/desc/compute/utils.py +++ b/desc/compute/utils.py @@ -4,7 +4,6 @@ import inspect import numpy as np -from termcolor import colored from desc.backend import cond, execute_on_cpu, fori_loop, jnp, put from desc.grid import ConcentricGrid, Grid, LinearGrid @@ -43,7 +42,7 @@ def compute(parameterization, names, params, transforms, profiles, data=None, ** Type of object to compute for, eg Equilibrium, Curve, etc. names : str or array-like of str Name(s) of the quantity(s) to compute. - params : dict of ndarray + params : dict[str, jnp.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 @@ -51,7 +50,7 @@ def compute(parameterization, names, params, transforms, profiles, data=None, ** profiles : dict of Profile Profile objects for pressure, iota, current, etc. Defaults to attributes of self - data : dict of ndarray + data : dict[str, jnp.ndarray] Data computed so far, generally output from other compute functions. Any vector v = v¹ R̂ + v² ϕ̂ + v³ Ẑ should be given in components v = [v¹, v², v³] where R̂, ϕ̂, Ẑ are the normalized basis vectors @@ -212,7 +211,7 @@ def get_data_deps(keys, obj, has_axis=False, basis="rpz", data=None): Whether the grid to compute on has a node on the magnetic axis. basis : {"rpz", "xyz"} Basis of computed quantities. - data : dict of ndarray + data : dict[str, jnp.ndarray] Data computed so far, generally output from other compute functions Returns @@ -287,7 +286,7 @@ def _get_deps(parameterization, names, deps, data=None, has_axis=False, check_fu Name(s) of the quantity(s) to compute. deps : set[str] Dependencies gathered so far. - data : dict of ndarray or None + data : dict[str, jnp.ndarray] Data computed so far, generally output from other compute functions. has_axis : bool Whether the grid to compute on has a node on the magnetic axis. @@ -375,7 +374,7 @@ def get_derivs(keys, obj, has_axis=False, basis="rpz"): Returns ------- - derivs : dict of list of int + derivs : dict[list, str] Orders of derivatives needed to compute key. Keys for R, Z, L, etc @@ -465,7 +464,7 @@ def get_params(keys, obj, has_axis=False, basis="rpz"): Returns ------- - params : list of str or dict of ndarray + params : list[str] or dict[str, jnp.ndarray] Parameters needed to compute key. If eq is None, returns a list of the names of params needed otherwise, returns a dict of ndarray with keys for R_lmn, Z_lmn, etc. @@ -624,13 +623,13 @@ def has_dependencies(parameterization, qty, params, transforms, profiles, data): Type of thing we're checking dependencies for. eg desc.equilibrium.Equilibrium qty : str Name of something from the data index. - params : dict of ndarray + params : dict[str, jnp.ndarray] Dictionary of parameters we have. - transforms : dict of Transform + transforms : dict[str, Transform] Dictionary of transforms we have. - profiles : dict of Profile + profiles : dict[str, Profile] Dictionary of profiles we have. - data : dict of ndarray + data : dict[str, jnp.ndarray] Dictionary of what we've computed so far. Returns @@ -988,8 +987,10 @@ def line_integrals( line_label != "poloidal" and isinstance(grid, ConcentricGrid), msg="ConcentricGrid should only be used for poloidal line integrals.", ) - msg = colored("Correctness not guaranteed on grids with duplicate nodes.", "yellow") - warnif(isinstance(grid, LinearGrid) and grid.endpoint, msg=msg) + warnif( + isinstance(grid, LinearGrid) and grid.endpoint, + msg="Correctness not guaranteed on grids with duplicate nodes.", + ) # Generate a new quantity q_prime which is zero everywhere # except on the fixed surface, on which q_prime takes the value of q. # Then forward the computation to surface_integrals(). @@ -1075,11 +1076,8 @@ def surface_integrals_map(grid, surface_label="rho", expand_out=True, tol=1e-14) surface_label = grid.get_label(surface_label) warnif( surface_label == "poloidal" and isinstance(grid, ConcentricGrid), - msg=colored( - "Integrals over constant poloidal surfaces" - " are poorly defined for ConcentricGrid.", - "yellow", - ), + msg="Integrals over constant poloidal surfaces" + " are poorly defined for ConcentricGrid.", ) unique_size, inverse_idx, spacing, has_endpoint_dupe, has_idx = _get_grid_surface( grid, surface_label diff --git a/desc/equilibrium/coords.py b/desc/equilibrium/coords.py index 33d2bfd497..2fda119f04 100644 --- a/desc/equilibrium/coords.py +++ b/desc/equilibrium/coords.py @@ -1,17 +1,19 @@ """Functions for mapping between flux, sfl, and real space coordinates.""" import functools -import warnings import numpy as np -from termcolor import colored from desc.backend import fori_loop, jit, jnp, put, root, root_scalar, vmap from desc.compute import compute as compute_fun from desc.compute import data_index, get_data_deps, get_profiles, get_transforms from desc.grid import ConcentricGrid, Grid, LinearGrid, QuadratureGrid from desc.transform import Transform -from desc.utils import setdefault +from desc.utils import check_posint, errorif, setdefault, warnif + + +def _periodic(x, period): + return jnp.where(jnp.isfinite(period), x % period, x) def map_coordinates( # noqa: C901 @@ -21,43 +23,45 @@ def map_coordinates( # noqa: C901 outbasis=("rho", "theta", "zeta"), guess=None, params=None, - period=(np.inf, np.inf, np.inf), + period=None, tol=1e-6, maxiter=30, full_output=False, **kwargs, ): - """Given coordinates in inbasis, compute corresponding coordinates in outbasis. + """Transform coordinates given in ``inbasis`` to ``outbasis``. - First solves for the computational coordinates that correspond to inbasis, then - evaluates outbasis at those locations. + Solves for the computational coordinates that correspond to ``inbasis``, + then evaluates ``outbasis`` at those locations. - Speed can often be significantly improved by providing a reasonable initial guess. - The default is a nearest neighbor search on a grid. + Performance can often improve significantly given a reasonable initial guess. Parameters ---------- eq : Equilibrium - Equilibrium to use - coords : ndarray, shape(k,3) + Equilibrium to use. + coords : ndarray + Shape (k, 3). 2D array of input coordinates. Each row is a different point in space. inbasis, outbasis : tuple of str - Labels for input and output coordinates, eg ("R", "phi", "Z") or + Labels for input and output coordinates, e.g. ("R", "phi", "Z") or ("rho", "alpha", "zeta") or any combination thereof. Labels should be the - same as the compute function data key - guess : None or ndarray, shape(k,3) + same as the compute function data key. + guess : jnp.ndarray + Shape (k, 3). Initial guess for the computational coordinates ['rho', 'theta', 'zeta'] - corresponding to coords in inbasis. If None, heuristics are used based on - inbasis or a nearest neighbor search on a grid. + corresponding to ``coords`` in ``inbasis``. If not given, then heuristics + based on ``inbasis`` or a nearest neighbor search on a grid may be used. + In general, this must be given to be compatible with JIT. params : dict - Values of equilibrium parameters to use, eg eq.params_dict + Values of equilibrium parameters to use, e.g. ``eq.params_dict``. period : tuple of float - Assumed periodicity for each quantity in inbasis. - Use np.inf to denote no periodicity. + Assumed periodicity for each quantity in ``inbasis``. + Use ``np.inf`` to denote no periodicity. tol : float Stopping tolerance. - maxiter : int > 0 - Maximum number of Newton iterations + maxiter : int + Maximum number of Newton iterations. full_output : bool, optional If True, also return a tuple where the first element is the residual from the root finding and the second is the number of iterations. @@ -67,44 +71,67 @@ def map_coordinates( # noqa: C901 Returns ------- - coords : ndarray, shape(k,3) - Coordinates mapped from inbasis to outbasis. Values of NaN will be returned - for coordinates where root finding did not succeed, possibly because the - coordinate is not in the plasma volume. + out : jnp.ndarray + Shape (k, 3). + Coordinates mapped from ``inbasis`` to ``outbasis``. Values of NaN will be + returned for coordinates where root finding did not succeed, possibly + because the coordinate is not in the plasma volume. info : tuple 2 element tuple containing residuals and number of iterations - for each point. Only returned if ``full_output`` is True - - Notes - ----- - ``guess`` must be given for this function to be compatible with ``jit``. + for each point. Only returned if ``full_output`` is True. """ + check_posint(maxiter, allow_none=False) + errorif( + not np.isfinite(tol) or tol <= 0, + ValueError, + f"tol must be a positive float, got {tol}", + ) + params = setdefault(params, eq.params_dict) inbasis = tuple(inbasis) outbasis = tuple(outbasis) - assert ( - np.isfinite(maxiter) and maxiter > 0 - ), f"maxiter must be a positive integer, got {maxiter}" - assert np.isfinite(tol) and tol > 0, f"tol must be a positive float, got {tol}" - basis_derivs = tuple([f"{X}_{d}" for X in inbasis for d in ("r", "t", "z")]) + # TODO: make this work for permutations of in/out basis + if outbasis == ("rho", "theta", "zeta"): + # TODO: get iota if not supplied using below logic + if inbasis == ("rho", "alpha", "zeta") and "iota" in kwargs: + return _map_clebsch_coordinates( + coords, + kwargs.pop("iota"), + params["L_lmn"], + eq.L_basis, + guess[:, 1] if guess is not None else None, + tol, + maxiter, + full_output, + **kwargs, + ) + if inbasis == ("rho", "theta_PEST", "zeta"): + return _map_PEST_coordinates( + coords, + params["L_lmn"], + eq.L_basis, + guess[:, 1] if guess is not None else None, + tol, + maxiter, + full_output, + **kwargs, + ) + + basis_derivs = tuple(f"{X}_{d}" for X in inbasis for d in ("r", "t", "z")) for key in basis_derivs: - assert ( - key in data_index["desc.equilibrium.equilibrium.Equilibrium"] - ), f"don't have recipe to compute partial derivative {key}" + errorif( + key not in data_index["desc.equilibrium.equilibrium.Equilibrium"], + NotImplementedError, + f"don't have recipe to compute partial derivative {key}", + ) rhomin = kwargs.pop("rhomin", tol / 10) - kwargs.setdefault("tol", tol) - kwargs.setdefault("maxiter", maxiter) - period = np.asarray(period) - - def periodic(x): - return jnp.where(jnp.isfinite(period), x % period, x) - - coords = periodic(coords) + warnif(period is None, msg="Assuming no periodicity.") + period = np.asarray(setdefault(period, (np.inf, np.inf, np.inf))) + coords = _periodic(coords, period) - params = setdefault(params, eq.params_dict) - profiles = get_profiles(inbasis + basis_derivs, eq, None) + profiles = get_profiles(inbasis + basis_derivs, eq) p = "desc.equilibrium.equilibrium.Equilibrium" names = inbasis + basis_derivs + outbasis deps = list(set(get_data_deps(names, obj=p) + list(names))) @@ -132,7 +159,7 @@ def compute(y, basis): @jit def residual(y, coords): xk = compute(y, inbasis) - r = periodic(xk) - periodic(coords) + r = _periodic(xk, period) - _periodic(coords, period) return jnp.where((r > period / 2) & jnp.isfinite(period), -period + r, r) @jit @@ -155,13 +182,27 @@ def fixup(y, *args): if yk is None: yk = _initial_guess_heuristic(yk, coords, inbasis, eq, profiles) if yk is None: - yk = _initial_guess_nn_search(yk, coords, inbasis, eq, period, compute) + yk = _initial_guess_nn_search(coords, inbasis, eq, period, compute) yk = fixup(yk) vecroot = jit( - vmap(lambda x0, *p: root(residual, x0, jac=jac, args=p, fixup=fixup, **kwargs)) + vmap( + lambda x0, *p: root( + residual, + x0, + jac=jac, + args=p, + fixup=fixup, + tol=tol, + maxiter=maxiter, + **kwargs, + ) + ) ) + # See description here + # https://github.com/PlasmaControl/DESC/pull/504#discussion_r1194172532 + # except we make sure properly handle periodic coordinates. yk, (res, niter) = vecroot(yk, coords) out = compute(yk, outbasis) @@ -205,22 +246,19 @@ def _initial_guess_heuristic(yk, coords, inbasis, eq, profiles): iota = profiles["iota"](rho) theta = (alpha + iota * zeta) % (2 * jnp.pi) - yk = jnp.array([rho, theta, zeta]).T + yk = jnp.column_stack([rho, theta, zeta]) return yk -def _initial_guess_nn_search(yk, coords, inbasis, eq, period, compute): +def _initial_guess_nn_search(coords, inbasis, eq, period, compute): # nearest neighbor search on dense grid yg = ConcentricGrid(eq.L_grid, eq.M_grid, max(eq.N_grid, eq.M_grid)).nodes xg = compute(yg, inbasis) idx = jnp.zeros(len(coords)).astype(int) coords = jnp.asarray(coords) - def periodic(x): - return jnp.where(jnp.isfinite(period), x % period, x) - def _distance_body(i, idx): - d = periodic(coords[i]) - periodic(xg) + d = _periodic(coords[i], period) - _periodic(xg, period) d = jnp.where((d > period / 2) & jnp.isfinite(period), period - d, d) distance = jnp.linalg.norm(d, axis=-1) k = jnp.argmin(distance) @@ -231,24 +269,36 @@ def _distance_body(i, idx): return yg[idx] -def compute_theta_coords( - eq, flux_coords, L_lmn=None, tol=1e-6, maxiter=20, full_output=False, **kwargs +# TODO: decide later whether to assume given phi instead of zeta. +def _map_PEST_coordinates( + coords, + L_lmn, + L_basis, + guess, + tol=1e-6, + maxiter=30, + full_output=False, + **kwargs, ): - """Find theta_DESC for given straight field line theta_PEST. + """Find θ (theta_DESC) for given straight field line ϑ (theta_PEST). Parameters ---------- - eq : Equilibrium - Equilibrium to use - flux_coords : ndarray, shape(k,3) - 2d array of flux coordinates [rho,theta*,zeta]. Each row is a different - point in space. - L_lmn : ndarray - spectral coefficients for lambda. Defaults to eq.L_lmn + coords : ndarray + Shape (k, 3). + Straight field line PEST coordinates [ρ, ϑ, ϕ]. Assumes ζ = ϕ. + Each row is a different point in space. + L_lmn : jnp.ndarray + Spectral coefficients for lambda. + L_basis : Basis + Spectral basis for lambda. + guess : jnp.ndarray + Shape (k, ). + Optional initial guess for the computational coordinates. tol : float Stopping tolerance. - maxiter : int > 0 - maximum number of Newton iterations + maxiter : int + Maximum number of Newton iterations. full_output : bool, optional If True, also return a tuple where the first element is the residual from the root finding and the second is the number of iterations. @@ -258,37 +308,39 @@ def compute_theta_coords( Returns ------- - coords : ndarray, shape(k,3) - coordinates [rho,theta,zeta]. + out : ndarray + Shape (k, 3). + DESC computational coordinates [ρ, θ, ζ]. info : tuple - 2 element tuple containing residuals and number of iterations - for each point. Only returned if ``full_output`` is True - """ - kwargs.setdefault("maxiter", maxiter) - kwargs.setdefault("tol", tol) + 2 element tuple containing residuals and number of iterations for each point. + Only returned if ``full_output`` is True. - if L_lmn is None: - L_lmn = eq.L_lmn - rho, theta_star, zeta = flux_coords.T + """ + rho, theta_PEST, zeta = coords.T + theta_PEST = theta_PEST % (2 * np.pi) + # Assume λ=0 for initial guess. + guess = setdefault(guess, theta_PEST) + # Root finding for θₖ such that r(θₖ) = ϑₖ(ρ, θₖ, ζ) − ϑ = 0. def rootfun(theta_DESC, theta_PEST, rho, zeta): - nodes = jnp.atleast_2d( - jnp.array([rho.squeeze(), theta_DESC.squeeze(), zeta.squeeze()]) + nodes = jnp.array( + [rho.squeeze(), theta_DESC.squeeze(), zeta.squeeze()], ndmin=2 ) - A = eq.L_basis.evaluate(nodes) + A = L_basis.evaluate(nodes) lmbda = A @ L_lmn - theta_PESTk = theta_DESC + lmbda - r = (theta_PESTk % (2 * np.pi)) - (theta_PEST % (2 * np.pi)) + theta_PEST_k = (theta_DESC + lmbda) % (2 * np.pi) + r = theta_PEST_k - theta_PEST # r should be between -pi and pi r = jnp.where(r > np.pi, r - 2 * np.pi, r) r = jnp.where(r < -np.pi, r + 2 * np.pi, r) return r.squeeze() def jacfun(theta_DESC, theta_PEST, rho, zeta): - nodes = jnp.atleast_2d( - jnp.array([rho.squeeze(), theta_DESC.squeeze(), zeta.squeeze()]) + # Valid everywhere except θ such that θ+λ = k 2π where k ∈ ℤ. + nodes = jnp.array( + [rho.squeeze(), theta_DESC.squeeze(), zeta.squeeze()], ndmin=2 ) - A1 = eq.L_basis.evaluate(nodes, (0, 1, 0)) + A1 = L_basis.evaluate(nodes, (0, 1, 0)) lmbda_t = jnp.dot(A1, L_lmn) return 1 + lmbda_t.squeeze() @@ -298,15 +350,121 @@ def fixup(x, *args): vecroot = jit( vmap( lambda x0, *p: root_scalar( - rootfun, x0, jac=jacfun, args=p, fixup=fixup, **kwargs + rootfun, + x0, + jac=jacfun, + args=p, + fixup=fixup, + tol=tol, + maxiter=maxiter, + **kwargs, ) ) ) - theta_DESC, (res, niter) = vecroot(theta_star, theta_star, rho, zeta) + theta_DESC, (res, niter) = vecroot(guess, theta_PEST, rho, zeta) - nodes = jnp.array([rho, jnp.atleast_1d(theta_DESC.squeeze()), zeta]).T + out = jnp.column_stack([rho, jnp.atleast_1d(theta_DESC.squeeze()), zeta]) + + if full_output: + return out, (res, niter) + return out + + +# TODO: decide later whether to assume given phi instead of zeta. +def _map_clebsch_coordinates( + coords, + iota, + L_lmn, + L_basis, + guess=None, + tol=1e-6, + maxiter=30, + full_output=False, + **kwargs, +): + """Find θ for given Clebsch field line poloidal label α. + + Parameters + ---------- + coords : ndarray + Shape (k, 3). + Clebsch field line coordinates [ρ, α, ζ]. Assumes ζ = ϕ. + Each row is a different point in space. + iota : ndarray + Shape (k, ). + Rotational transform on each node. + L_lmn : jnp.ndarray + Spectral coefficients for lambda. + L_basis : Basis + Spectral basis for lambda. + guess : jnp.ndarray + Shape (k, ). + Optional initial guess for the computational coordinates. + tol : float + Stopping tolerance. + maxiter : int + Maximum number of Newton iterations. + full_output : bool, optional + If True, also return a tuple where the first element is the residual from + the root finding and the second is the number of iterations. + kwargs : dict, optional + Additional keyword arguments to pass to ``root_scalar`` such as ``maxiter_ls``, + ``alpha``. + + Returns + ------- + out : ndarray + Shape (k, 3). + DESC computational coordinates [ρ, θ, ζ]. + info : tuple + 2 element tuple containing residuals and number of iterations for each point. + Only returned if ``full_output`` is True. + + """ + rho, alpha, zeta = coords.T + if guess is None: + # Assume λ=0 for initial guess. + guess = (alpha + iota * zeta) % (2 * np.pi) + + # Root finding for θₖ such that r(θₖ) = αₖ(ρ, θₖ, ζ) − α = 0. + def rootfun(theta, alpha, rho, zeta, iota): + nodes = jnp.array([rho.squeeze(), theta.squeeze(), zeta.squeeze()], ndmin=2) + A = L_basis.evaluate(nodes) + lmbda = A @ L_lmn + alpha_k = theta + lmbda - iota * zeta + r = (alpha_k - alpha) % (2 * np.pi) + # r should be between -pi and pi + r = jnp.where(r > np.pi, r - 2 * np.pi, r) + r = jnp.where(r < -np.pi, r + 2 * np.pi, r) + return r.squeeze() + + def jacfun(theta, alpha, rho, zeta, iota): + # Valid everywhere except θ such that θ+λ = k 2π where k ∈ ℤ. + nodes = jnp.array([rho.squeeze(), theta.squeeze(), zeta.squeeze()], ndmin=2) + A1 = L_basis.evaluate(nodes, (0, 1, 0)) + lmbda_t = jnp.dot(A1, L_lmn) + return 1 + lmbda_t.squeeze() + + def fixup(x, *args): + return x % (2 * np.pi) + + vecroot = jit( + vmap( + lambda x0, *p: root_scalar( + rootfun, + x0, + jac=jacfun, + args=p, + fixup=fixup, + tol=tol, + maxiter=maxiter, + **kwargs, + ) + ) + ) + theta, (res, niter) = vecroot(guess, alpha, rho, zeta, iota) + out = jnp.column_stack([rho, jnp.atleast_1d(theta.squeeze()), zeta]) - out = nodes if full_output: return out, (res, niter) return out @@ -354,11 +512,7 @@ def is_nested(eq, grid=None, R_lmn=None, Z_lmn=None, L_lmn=None, msg=None): data = compute_fun( "desc.equilibrium.equilibrium.Equilibrium", "sqrt(g)_PEST", - params={ - "R_lmn": R_lmn, - "Z_lmn": Z_lmn, - "L_lmn": L_lmn, - }, + params={"R_lmn": R_lmn, "Z_lmn": Z_lmn, "L_lmn": L_lmn}, transforms=transforms, profiles={}, # no profiles needed ) @@ -366,25 +520,18 @@ def is_nested(eq, grid=None, R_lmn=None, Z_lmn=None, L_lmn=None, msg=None): nested = jnp.all( jnp.sign(data["sqrt(g)_PEST"][0]) == jnp.sign(data["sqrt(g)_PEST"]) ) - if not nested: - if msg == "auto": - warnings.warn( - colored( - "WARNING: Flux surfaces are no longer nested, exiting early. " - + "Automatic continuation method failed, consider specifying " - + "continuation steps manually", - "yellow", - ) - ) - elif msg == "manual": - warnings.warn( - colored( - "WARNING: Flux surfaces are no longer nested, exiting early." - + "Consider taking smaller perturbation/resolution steps " - + "or reducing trust radius", - "yellow", - ) - ) + warnif( + not nested and msg is not None, + RuntimeWarning, + "Flux surfaces are no longer nested, exiting early. " + + { + "auto": "Automatic continuation method failed, consider specifying " + "continuation steps manually.", + "manual": "Consider taking smaller perturbation/resolution steps " + "or reducing trust radius.", + None: "", + }[msg], + ) return nested @@ -505,7 +652,9 @@ def to_sfl( return eq_sfl -def get_rtz_grid(eq, radial, poloidal, toroidal, coordinates, period, jitable=True): +def get_rtz_grid( + eq, radial, poloidal, toroidal, coordinates, period, jitable=True, **kwargs +): """Return DESC grid in rtz (rho, theta, zeta) coordinates from given coordinates. Create a tensor-product grid from the given coordinates, and return the same grid @@ -548,20 +697,70 @@ def get_rtz_grid(eq, radial, poloidal, toroidal, coordinates, period, jitable=Tr "v": "theta_PEST", "a": "alpha", "z": "zeta", + "p": "phi", } - rtz_nodes = eq.map_coordinates( + rtz_nodes = map_coordinates( + eq, grid.nodes, inbasis=[inbasis[char] for char in coordinates], outbasis=("rho", "theta", "zeta"), period=period, + **kwargs, ) + idx = {} + if inbasis[coordinates[0]] == "rho": + # Should work as long as inbasis radial coordinate is + # single variable, monotonic increasing function of rho. + idx["_unique_rho_idx"] = grid.unique_rho_idx + idx["_inverse_rho_idx"] = grid.inverse_rho_idx desc_grid = Grid( nodes=rtz_nodes, coordinates="rtz", source_grid=grid, sort=False, jitable=jitable, - _unique_rho_idx=grid.unique_rho_idx, - _inverse_rho_idx=grid.inverse_rho_idx, + **idx, ) return desc_grid + + +# TODO: deprecated, remove eventually +def compute_theta_coords( + eq, flux_coords, L_lmn=None, tol=1e-6, maxiter=20, full_output=False, **kwargs +): + """Find θ (theta_DESC) for given straight field line ϑ (theta_PEST). + + Parameters + ---------- + eq : Equilibrium + Equilibrium to use. + flux_coords : ndarray + Shape (k, 3). + Straight field line PEST coordinates [ρ, ϑ, ϕ]. Assumes ζ = ϕ. + Each row is a different point in space. + L_lmn : ndarray + Spectral coefficients for lambda. Defaults to ``eq.L_lmn``. + tol : float + Stopping tolerance. + maxiter : int + Maximum number of Newton iterations. + full_output : bool, optional + If True, also return a tuple where the first element is the residual from + the root finding and the second is the number of iterations. + kwargs : dict, optional + Additional keyword arguments to pass to ``root_scalar`` such as + ``maxiter_ls``, ``alpha``. + + Returns + ------- + coords : ndarray + Shape (k, 3). + DESC computational coordinates [ρ, θ, ζ]. + info : tuple + 2 element tuple containing residuals and number of iterations for each + point. Only returned if ``full_output`` is True. + + """ + return eq.compute_theta_coords( + flux_coords, L_lmn, tol, maxiter, full_output, **kwargs + ) diff --git a/desc/equilibrium/equilibrium.py b/desc/equilibrium/equilibrium.py index 22d5947d1a..41d205a595 100644 --- a/desc/equilibrium/equilibrium.py +++ b/desc/equilibrium/equilibrium.py @@ -2,13 +2,11 @@ import copy import numbers -import warnings from collections.abc import MutableSequence import numpy as np from scipy import special from scipy.constants import mu_0 -from termcolor import colored from desc.backend import execute_on_cpu, jnp from desc.basis import FourierZernikeBasis, fourier, zernike_radial @@ -53,7 +51,7 @@ ) from ..compute.data_index import is_0d_vol_grid, is_1dr_rad_grid, is_1dz_tor_grid -from .coords import compute_theta_coords, is_nested, map_coordinates, to_sfl +from .coords import is_nested, map_coordinates, to_sfl from .initial_guess import set_initial_guess from .utils import parse_axis, parse_profile, parse_surface @@ -355,10 +353,10 @@ def __init__( p = getattr(self, profile) if hasattr(p, "change_resolution"): p.change_resolution(max(p.basis.L, self.L)) - if isinstance(p, PowerSeriesProfile) and p.sym != "even": - warnings.warn( - colored(f"{profile} profile is not an even power series.", "yellow") - ) + warnif( + isinstance(p, PowerSeriesProfile) and p.sym != "even", + msg=f"{profile} profile is not an even power series.", + ) # ensure number of field periods agree before setting guesses eq_NFP = self.NFP @@ -824,7 +822,7 @@ def compute( # noqa: C901 profiles : dict of Profile Profile objects for pressure, iota, current, etc. Defaults to attributes of self - data : dict of ndarray + data : dict[str, jnp.ndarray] Data computed so far, generally output from other compute functions. Any vector v = v¹ R̂ + v² ϕ̂ + v³ Ẑ should be given in components v = [v¹, v², v³] where R̂, ϕ̂, Ẑ are the normalized basis vectors @@ -977,10 +975,9 @@ def need_src(name): for dep in deps: req = data_index[p][dep]["resolution_requirement"] coords = data_index[p][dep]["coordinates"] - msg = lambda direction: colored( + msg = lambda direction: ( f"Dependency {dep} may require more {direction}" - " resolution to compute accurately.", - "yellow", + " resolution to compute accurately." ) warnif( # if need more radial resolution @@ -1142,46 +1139,52 @@ def need_src(name): ) return data - def map_coordinates( # noqa: C901 + def map_coordinates( self, coords, inbasis, outbasis=("rho", "theta", "zeta"), guess=None, params=None, - period=(np.inf, np.inf, np.inf), + period=None, tol=1e-6, maxiter=30, full_output=False, **kwargs, ): - """Given coordinates in inbasis, compute corresponding coordinates in outbasis. + """Transform coordinates given in ``inbasis`` to ``outbasis``. - First solves for the computational coordinates that correspond to inbasis, then - evaluates outbasis at those locations. + Solves for the computational coordinates that correspond to ``inbasis``, + then evaluates ``outbasis`` at those locations. + + Performance can often improve significantly given a reasonable initial guess. Parameters ---------- - coords : ndarray, shape(k,3) - 2D array of input coordinates. Each row is a different - point in space. + eq : Equilibrium + Equilibrium to use. + coords : ndarray + Shape (k, 3). + 2D array of input coordinates. Each row is a different point in space. inbasis, outbasis : tuple of str - Labels for input and output coordinates, eg ("R", "phi", "Z") or + Labels for input and output coordinates, e.g. ("R", "phi", "Z") or ("rho", "alpha", "zeta") or any combination thereof. Labels should be the - same as the compute function data key - guess : None or ndarray, shape(k,3) + same as the compute function data key. + guess : jnp.ndarray + Shape (k, 3). Initial guess for the computational coordinates ['rho', 'theta', 'zeta'] - corresponding to coords in inbasis. If None, heuristics are used based on - in basis and a nearest neighbor search on a coarse grid. + corresponding to ``coords`` in ``inbasis``. If not given, then heuristics + based on ``inbasis`` or a nearest neighbor search on a grid may be used. + In general, this must be given to be compatible with JIT. params : dict - Values of equilibrium parameters to use, eg eq.params_dict + Values of equilibrium parameters to use, e.g. ``eq.params_dict``. period : tuple of float - Assumed periodicity for each quantity in inbasis. - Use np.inf to denote no periodicity. + Assumed periodicity for each quantity in ``inbasis``. + Use ``np.inf`` to denote no periodicity. tol : float Stopping tolerance. - maxiter : int > 0 - Maximum number of Newton iterations + maxiter : int + Maximum number of Newton iterations. full_output : bool, optional If True, also return a tuple where the first element is the residual from the root finding and the second is the number of iterations. @@ -1191,15 +1194,14 @@ def map_coordinates( # noqa: C901 Returns ------- - coords : ndarray, shape(k,3) - Coordinates mapped from inbasis to outbasis. + out : jnp.ndarray + Shape (k, 3). + Coordinates mapped from ``inbasis`` to ``outbasis``. Values of NaN will be + returned for coordinates where root finding did not succeed, possibly + because the coordinate is not in the plasma volume. info : tuple 2 element tuple containing residuals and number of iterations - for each point. Only returned if ``full_output`` is True - - Notes - ----- - ``guess`` must be given for this function to be compatible with ``jit``. + for each point. Only returned if ``full_output`` is True. """ return map_coordinates( @@ -1219,21 +1221,20 @@ def map_coordinates( # noqa: C901 def compute_theta_coords( self, flux_coords, L_lmn=None, tol=1e-6, maxiter=20, full_output=False, **kwargs ): - """Find theta_DESC for given straight field line theta_PEST. + """Find θ (theta_DESC) for given straight field line ϑ (theta_PEST). Parameters ---------- - eq : Equilibrium - Equilibrium to use - flux_coords : ndarray, shape(k,3) - 2d array of flux coordinates [rho,theta*,zeta]. Each row is a different - point in space. + flux_coords : ndarray + Shape (k, 3). + Straight field line PEST coordinates [ρ, ϑ, ϕ]. Assumes ζ = ϕ. + Each row is a different point in space. L_lmn : ndarray - spectral coefficients for lambda. Defaults to eq.L_lmn + Spectral coefficients for lambda. Defaults to ``eq.L_lmn``. tol : float Stopping tolerance. - maxiter : int > 0 - maximum number of Newton iterations + maxiter : int + Maximum number of Newton iterations. full_output : bool, optional If True, also return a tuple where the first element is the residual from the root finding and the second is the number of iterations. @@ -1243,18 +1244,23 @@ def compute_theta_coords( Returns ------- - coords : ndarray, shape(k,3) - coordinates [rho,theta,zeta]. + coords : ndarray + Shape (k, 3). + DESC computational coordinates [ρ, θ, ζ]. info : tuple - 2 element tuple containing residuals and number of iterations - for each point. Only returned if ``full_output`` is True + 2 element tuple containing residuals and number of iterations for each + point. Only returned if ``full_output`` is True. + """ - return compute_theta_coords( + warnif(True, DeprecationWarning, msg="Use map_coordinates instead.") + return map_coordinates( self, flux_coords, - L_lmn=L_lmn, - maxiter=maxiter, + inbasis=("rho", "theta_PEST", "zeta"), + outbasis=("rho", "theta", "zeta"), + params=self.params_dict if L_lmn is None else {"L_lmn": L_lmn}, tol=tol, + maxiter=maxiter, full_output=full_output, **kwargs, ) @@ -2031,22 +2037,20 @@ def solve( if not isinstance(constraints, (list, tuple)): constraints = tuple([constraints]) - if self.N > self.N_grid or self.M > self.M_grid or self.L > self.L_grid: - warnings.warn( - colored( - "Equilibrium has one or more spectral resolutions " - + "greater than the corresponding collocation grid resolution! " - + "This is not recommended and may result in poor convergence. " - + "Set grid resolutions to be higher, (i.e. eq.N_grid=2*eq.N) " - + "to avoid this warning.", - "yellow", - ) - ) - if self.bdry_mode == "poincare": - raise NotImplementedError( - "Solving equilibrium with poincare XS as BC is not supported yet " - + "on master branch." - ) + warnif( + self.N > self.N_grid or self.M > self.M_grid or self.L > self.L_grid, + msg="Equilibrium has one or more spectral resolutions " + + "greater than the corresponding collocation grid resolution! " + + "This is not recommended and may result in poor convergence. " + + "Set grid resolutions to be higher, (i.e. eq.N_grid=2*eq.N) " + + "to avoid this warning.", + ) + errorif( + self.bdry_mode == "poincare", + NotImplementedError, + "Solving equilibrium with poincare XS as BC is not supported yet " + + "on master branch.", + ) things, result = optimizer.optimize( self, diff --git a/desc/grid.py b/desc/grid.py index aad2c86f33..06ce1329ae 100644 --- a/desc/grid.py +++ b/desc/grid.py @@ -715,9 +715,8 @@ def __init__( self._N = self.num_nodes errorif(len(kwargs), ValueError, f"Got unexpected kwargs {kwargs.keys()}") - @classmethod + @staticmethod def create_meshgrid( - cls, nodes, spacing=None, coordinates="rtz", @@ -790,7 +789,7 @@ def create_meshgrid( a.size, ) inverse_c_idx = jnp.tile(unique_c_idx, a.size * b.size) - return cls( + return Grid( nodes=nodes, spacing=spacing, weights=weights, diff --git a/desc/objectives/_geometry.py b/desc/objectives/_geometry.py index 7ae828d95e..f72e9f7605 100644 --- a/desc/objectives/_geometry.py +++ b/desc/objectives/_geometry.py @@ -2,16 +2,16 @@ import numpy as np -from desc.backend import jnp -from desc.compute import get_profiles, get_transforms, rpz2xyz +from desc.backend import jnp, vmap +from desc.compute import get_profiles, get_transforms, rpz2xyz, xyz2rpz from desc.compute.utils import _compute as compute_fun from desc.compute.utils import safenorm from desc.grid import LinearGrid, QuadratureGrid -from desc.utils import Timer, warnif +from desc.utils import Timer, errorif, parse_argname_change, warnif from .normalization import compute_scaling_factors from .objective_funs import _Objective -from .utils import softmin +from .utils import check_if_points_are_inside_perimeter, softmin class AspectRatio(_Objective): @@ -519,10 +519,10 @@ class PlasmaVesselDistance(_Objective): points on surface corresponding to the grid that the plasma-vessel distance is evaluated at, which can cause cusps or regions of very large curvature. - NOTE: When use_softmin=True, ensures that alpha*values passed in is + NOTE: When use_softmin=True, ensures that softmin_alpha*values passed in is at least >1, otherwise the softmin will return inaccurate approximations of the minimum. Will automatically multiply array values by 2 / min_val if the min - of alpha*array is <1. This is to avoid inaccuracies that arise when values <1 + of softmin_alpha*array is <1. This is to avoid inaccuracies when values <1 are present in the softmin, which can cause inaccurate mins or even incorrect signs of the softmin versus the actual min. @@ -565,20 +565,26 @@ class PlasmaVesselDistance(_Objective): Defaults to ``LinearGrid(M=eq.M_grid, N=eq.N_grid)``. use_softmin: bool, optional Use softmin or hard min. - surface_fixed: bool, optional - Whether the surface the distance from the plasma is computed to - is fixed or not. If True, the surface is fixed and its coordinates are - precomputed, which saves on computation time during optimization, and - self.things = [eq] only. - If False, the surface coordinates are computed at every iteration. + use_signed_distance: bool, optional + Whether to use absolute value of distance or a signed distance, with d + being positive if the plasma is inside of the bounding surface, and + negative if outside of the bounding surface. + NOTE: ``plasma_grid`` and ``surface_grid`` must have the same + toroidal angle values for signed distance to be used. + eq_fixed, surface_fixed: bool, optional + Whether the eq/surface is fixed or not. If True, the eq/surface is fixed + and its coordinates are precomputed, which saves on computation time during + optimization, and self.things = [surface]/[eq] only. + If False, the eq/surface coordinates are computed at every iteration. False by default, so that self.things = [eq, surface] - alpha: float, optional - Parameter used for softmin. The larger alpha, the closer the softmin - approximates the hardmin. softmin -> hardmin as alpha -> infinity. - if alpha*array < 1, the underlying softmin will automatically multiply - the array by 2/min_val to ensure that alpha*array>1. Making alpha larger - than this minimum value will make the softmin a more accurate approximation - of the true min. + Both cannot be True. + softmin_alpha: float, optional + Parameter used for softmin. The larger softmin_alpha, the closer the softmin + approximates the hardmin. softmin -> hardmin as softmin_alpha -> infinity. + if softmin_alpha*array < 1, the underlying softmin will automatically multiply + the array by 2/min_val to ensure that softmin_alpha*array>1. Making + softmin_alpha larger than this minimum value will make the softmin a + more accurate approximation of the true min. name : str, optional Name of the objective function. """ @@ -601,9 +607,12 @@ def __init__( surface_grid=None, plasma_grid=None, use_softmin=False, + eq_fixed=False, surface_fixed=False, - alpha=1.0, + softmin_alpha=1.0, name="plasma-vessel distance", + use_signed_distance=False, + **kwargs, ): if target is None and bounds is None: bounds = (1, np.inf) @@ -611,10 +620,29 @@ def __init__( self._surface_grid = surface_grid self._plasma_grid = plasma_grid self._use_softmin = use_softmin + self._use_signed_distance = use_signed_distance self._surface_fixed = surface_fixed - self._alpha = alpha + self._eq_fixed = eq_fixed + self._eq = eq + errorif( + eq_fixed and surface_fixed, ValueError, "Cannot fix both eq and surface" + ) + + self._softmin_alpha = parse_argname_change( + softmin_alpha, kwargs, "alpha", "softmin_alpha" + ) + errorif( + len(kwargs) != 0, + AssertionError, + f"PlasmaVesselDistance got unexpected keyword argument: {kwargs.keys()}", + ) + things = [] + if not eq_fixed: + things.append(eq) + if not surface_fixed: + things.append(surface) super().__init__( - things=[eq, self._surface] if not surface_fixed else [eq], + things=things, target=target, bounds=bounds, weight=weight, @@ -636,11 +664,15 @@ def build(self, use_jit=True, verbose=1): Level of output. """ - eq = self.things[0] - surface = self._surface if self._surface_fixed else self.things[1] - # if things[1] is different than self._surface, update self._surface - if surface != self._surface: - self._surface = surface + if self._eq_fixed: + eq = self._eq + surface = self.things[0] + elif self._surface_fixed: + eq = self.things[0] + surface = self._surface + else: + eq = self.things[0] + surface = self.things[1] if self._surface_grid is None: surface_grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP) else: @@ -660,6 +692,18 @@ def build(self, use_jit=True, verbose=1): "Plasma grid includes interior points, should be rho=1.", ) + # TODO: How to use with generalized toroidal angle? + errorif( + self._use_signed_distance + and not np.allclose( + plasma_grid.nodes[plasma_grid.unique_zeta_idx, 2], + surface_grid.nodes[surface_grid.unique_zeta_idx, 2], + ), + ValueError, + "Plasma grid and surface grid must contain points only at the " + "same zeta values in order to use signed distance", + ) + self._dim_f = surface_grid.num_nodes self._equil_data_keys = ["R", "phi", "Z"] self._surface_data_keys = ["x"] @@ -714,7 +758,15 @@ def build(self, use_jit=True, verbose=1): )["x"] surface_coords = rpz2xyz(surface_coords) self._constants["surface_coords"] = surface_coords - + elif self._eq_fixed: + data_eq = compute_fun( + self._eq, + self._equil_data_keys, + params=self._eq.params_dict, + transforms=equil_transforms, + profiles=equil_profiles, + ) + self._constants["data_equil"] = data_eq timer.stop("Precomputing transforms") if verbose > 1: timer.disp("Precomputing transforms") @@ -725,14 +777,15 @@ def build(self, use_jit=True, verbose=1): super().build(use_jit=use_jit, verbose=verbose) - def compute(self, equil_params, surface_params=None, constants=None): + def compute(self, params_1, params_2=None, constants=None): """Compute plasma-surface distance. Parameters ---------- - equil_params : dict - Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict - surface_params : dict + params_1 : dict + Dictionary of equilibrium degrees of freedom, eg Equilibrium.params_dict, + if eq_fixed is False, else the surface degrees of freedom + params_2 : dict Dictionary of surface degrees of freedom, eg Surface.params_dict Only needed if self._surface_fixed = False constants : dict @@ -747,14 +800,25 @@ def compute(self, equil_params, surface_params=None, constants=None): """ if constants is None: constants = self.constants - data = compute_fun( - "desc.equilibrium.equilibrium.Equilibrium", - self._equil_data_keys, - params=equil_params, - transforms=constants["equil_transforms"], - profiles=constants["equil_profiles"], - ) - plasma_coords = rpz2xyz(jnp.array([data["R"], data["phi"], data["Z"]]).T) + if self._eq_fixed: + surface_params = params_1 + elif self._surface_fixed: + equil_params = params_1 + else: + equil_params = params_1 + surface_params = params_2 + if not self._eq_fixed: + data = compute_fun( + "desc.equilibrium.equilibrium.Equilibrium", + self._equil_data_keys, + params=equil_params, + transforms=constants["equil_transforms"], + profiles=constants["equil_profiles"], + ) + else: + data = constants["data_equil"] + plasma_coords_rpz = jnp.array([data["R"], data["phi"], data["Z"]]).T + plasma_coords = rpz2xyz(plasma_coords_rpz) if self._surface_fixed: surface_coords = constants["surface_coords"] else: @@ -766,12 +830,54 @@ def compute(self, equil_params, surface_params=None, constants=None): profiles={}, )["x"] surface_coords = rpz2xyz(surface_coords) - d = safenorm(plasma_coords[:, None, :] - surface_coords[None, :, :], axis=-1) + + diff_vec = plasma_coords[:, None, :] - surface_coords[None, :, :] + d = safenorm(diff_vec, axis=-1) + + point_signs = jnp.ones(surface_coords.shape[0]) + if self._use_signed_distance: + surface_coords_rpz = xyz2rpz(surface_coords) + + plasma_coords_rpz = plasma_coords_rpz.reshape( + constants["equil_transforms"]["grid"].num_zeta, + constants["equil_transforms"]["grid"].num_theta, + 3, + ) + surface_coords_rpz = surface_coords_rpz.reshape( + constants["surface_transforms"]["grid"].num_zeta, + constants["surface_transforms"]["grid"].num_theta, + 3, + ) + + # loop over zeta planes + def fun(plasma_pts_at_zeta_plane, surface_pts_at_zeta_plane): + plasma_pts_at_zeta_plane = jnp.vstack( + (plasma_pts_at_zeta_plane, plasma_pts_at_zeta_plane[0, :]) + ) + + pt_sign = check_if_points_are_inside_perimeter( + plasma_pts_at_zeta_plane[:, 0], + plasma_pts_at_zeta_plane[:, 2], + surface_pts_at_zeta_plane[:, 0], + surface_pts_at_zeta_plane[:, 2], + ) + + return pt_sign + + point_signs = vmap(fun, in_axes=0)( + plasma_coords_rpz, surface_coords_rpz + ).flatten() + # at end here, point_signs is either +/- 1 with + # positive meaning the surface pt + # is outside the plasma and -1 if the surface pt is + # inside the plasma if self._use_softmin: # do softmin - return jnp.apply_along_axis(softmin, 0, d, self._alpha) + return ( + jnp.apply_along_axis(softmin, 0, d, self._softmin_alpha) * point_signs + ) else: # do hardmin - return d.min(axis=0) + return d.min(axis=0) * point_signs class MeanCurvature(_Objective): diff --git a/desc/objectives/normalization.py b/desc/objectives/normalization.py index 255eb346f4..0d996eed2f 100644 --- a/desc/objectives/normalization.py +++ b/desc/objectives/normalization.py @@ -38,24 +38,32 @@ def get_lowest_mode(basis, coeffs): scales["R0"] = R00 scales["a"] = np.sqrt(np.abs(R10 * Z10)) + scales["Psi"] = abs(thing.Psi) scales["A"] = np.pi * scales["a"] ** 2 scales["V"] = 2 * np.pi * scales["R0"] * scales["A"] - scales["B_T"] = abs(thing.Psi) / scales["A"] - iota_avg = np.mean(np.abs(thing.get_profile("iota")(np.linspace(0, 1, 20)))) - if np.isclose(iota_avg, 0): - scales["B_P"] = scales["B_T"] - else: - scales["B_P"] = scales["B_T"] * iota_avg - scales["B"] = np.sqrt(scales["B_T"] ** 2 + scales["B_P"] ** 2) - scales["I"] = scales["B_P"] * 2 * np.pi / mu_0 - scales["p"] = scales["B"] ** 2 / (2 * mu_0) - scales["W"] = scales["p"] * scales["V"] + scales["B"] = scales["Psi"] / scales["A"] * 1.25 + B_pressure = scales["B"] ** 2 / (2 * mu_0) + scales["I"] = scales["B"] * 2 * np.pi / mu_0 + scales["W"] = B_pressure * scales["V"] scales["J"] = scales["B"] / scales["a"] / mu_0 - scales["F"] = scales["p"] / scales["a"] + scales["F"] = B_pressure / scales["a"] scales["f"] = scales["F"] * scales["V"] - scales["Psi"] = abs(thing.Psi) - scales["n"] = 1e19 - scales["T"] = scales["p"] / (scales["n"] * elementary_charge) + + if thing.pressure is not None: + p0 = float(thing.pressure(0)[0]) + else: + scales["n"] = float( + ((thing.atomic_number(0) + 1) / 2 * thing.electron_density(0))[0] + ) + scales["T"] = np.mean( + [thing.electron_temperature(0), thing.ion_temperature(0)] + ) + p0 = elementary_charge * 2 * scales["n"] * scales["T"] + if p0 < 1: # vacuum + scales["p"] = B_pressure + else: + scales["p"] = p0 + scales["W_p"] = scales["p"] * scales["V"] / 2 elif isinstance(thing, FourierRZToroidalSurface): R00 = thing.R_lmn[thing.R_basis.get_idx(M=0, N=0)] diff --git a/desc/objectives/utils.py b/desc/objectives/utils.py index 924e8cfb43..5167573f51 100644 --- a/desc/objectives/utils.py +++ b/desc/objectives/utils.py @@ -9,7 +9,7 @@ from desc.utils import Index, errorif, flatten_list, svd_inv_null, unique_list, warnif -def factorize_linear_constraints(objective, constraint): # noqa: C901 +def factorize_linear_constraints(objective, constraint, x_scale="auto"): # noqa: C901 """Compute and factorize A to get pseudoinverse and nullspace. Given constraints of the form Ax=b, factorize A to find a particular solution xp @@ -22,6 +22,10 @@ def factorize_linear_constraints(objective, constraint): # noqa: C901 Objective function to optimize. constraint : ObjectiveFunction Objective function of linear constraints to enforce. + x_scale : array_like or ``'auto'``, optional + Characteristic scale of each variable. Setting ``x_scale`` is equivalent + to reformulating the problem in scaled variables ``xs = x / x_scale``. + If set to ``'auto'``, the scale is determined from the initial state vector. Returns ------- @@ -33,6 +37,8 @@ def factorize_linear_constraints(objective, constraint): # noqa: C901 Combined RHS vector. Z : ndarray Null space operator for full combined A such that A @ Z == 0. + D : ndarray + Scale of the full state vector x, as set by the parameter ``x_scale``. unfixed_idx : ndarray Indices of x that correspond to non-fixed values. project, recover : function @@ -130,32 +136,53 @@ def factorize_linear_constraints(objective, constraint): # noqa: C901 ) A = A[unfixed_rows][:, unfixed_idx] b = b[unfixed_rows] + unfixed_idx = indices_idx + fixed_idx = np.delete(np.arange(xp.size), unfixed_idx) + + # compute x_scale if not provided + if x_scale == "auto": + x_scale = objective.x(*objective.things) + errorif( + x_scale.shape != xp.shape, + ValueError, + "x_scale must be the same size as the full state vector. " + + f"Got size {x_scale.size} for state vector of size {xp.size}.", + ) + D = np.where(np.abs(x_scale) < 1e2, 1, np.abs(x_scale)) + + # null space & particular solution + A = A * D[None, unfixed_idx] if A.size: - Ainv_full, Z = svd_inv_null(A) + A_inv, Z = svd_inv_null(A) else: - Ainv_full = A.T + A_inv = A.T Z = np.eye(A.shape[1]) - Ainv_full = jnp.asarray(Ainv_full) - Z = jnp.asarray(Z) - b = jnp.asarray(b) - xp = put(xp, unfixed_idx, Ainv_full @ b) + xp = put(xp, unfixed_idx, A_inv @ b) + xp = put(xp, fixed_idx, ((1 / D) * xp)[fixed_idx]) + + # cast to jnp arrays xp = jnp.asarray(xp) + A = jnp.asarray(A) + b = jnp.asarray(b) + Z = jnp.asarray(Z) + D = jnp.asarray(D) @jit - def project(x): + def project(x_full): """Project a full state vector into the reduced optimization vector.""" - x_reduced = Z.T @ ((x - xp)[unfixed_idx]) + x_reduced = Z.T @ ((1 / D) * x_full - xp)[unfixed_idx] return jnp.atleast_1d(jnp.squeeze(x_reduced)) @jit def recover(x_reduced): """Recover the full state vector from the reduced optimization vector.""" dx = put(jnp.zeros(objective.dim_x), unfixed_idx, Z @ x_reduced) - return jnp.atleast_1d(jnp.squeeze(xp + dx)) + x_full = D * (xp + dx) + return jnp.atleast_1d(jnp.squeeze(x_full)) # check that all constraints are actually satisfiable - params = objective.unpack_state(xp, False) + params = objective.unpack_state(D * xp, False) for con in constraint.objectives: xpi = [params[i] for i, t in enumerate(objective.things) if t in con.things] y1 = con.compute_unscaled(*xpi) @@ -197,7 +224,7 @@ def recover(x_reduced): "or be due to floating point error.", ) - return xp, A, b, Z, unfixed_idx, project, recover + return xp, A, b, Z, D, unfixed_idx, project, recover def softmax(arr, alpha): @@ -295,3 +322,78 @@ def _parse_callable_target_bounds(target, bounds, x): if bounds is not None and callable(bounds[1]): bounds = (bounds[0], bounds[1](x)) return target, bounds + + +def check_if_points_are_inside_perimeter(R, Z, Rcheck, Zcheck): + """Function to check if the given points is inside the given polyognal perimeter. + + Rcheck, Zcheck are the points to check, and R, Z define the perimeter + in which to check. This function assumes that all points are in the same + plane. Function will return an array of signs (+/- 1), with positive sign meaning + the point is inside of the given perimeter, and a negative sign meaning the point + is outside of the given perimeter. + + NOTE: it does not matter if the input coordinates are cylindrical (R,Z) or + cartesian (X,Y), these are equivalent as long as they are in the same phi plane. + This function will work even if points are not in the same phi plane, but the + input coordinates must then be the equivalent of cartesian coordinates for whatever + plane the points lie in. + + Algorithm based off of "An Incremental Angle Point in Polygon Test", + K. Weiler, https://doi.org/10.1016/B978-0-12-336156-1.50012-4 + + Parameters + ---------- + R,Z : ndarray + 1-D arrays of coordinates of the points defining the polygonal + perimeter. The function will determine if the check point is inside + or outside of this perimeter. These should form a closed curve. + Rcheck, Zcheck : ndarray + coordinates of the points being checked if they are inside or outside of the + given perimeter. + + Returns + ------- + pt_sign : ndarray of {-1,1} + Integers corresponding to if the given point is inside or outside of the given + perimeter, with pt_sign[i]>0 meaning the point given by Rcheck[i], Zcheck[i] is + inside of the given perimeter, and a negative sign meaning the point is outside + of the given perimeter. + + """ + # R Z are the perimeter points + # Rcheck Zcheck are the points being checked for whether + # or not they are inside the check + + Rbool = R[:, None] > Rcheck + Zbool = Z[:, None] > Zcheck + # these are now size (Ncheck, Nperimeter) + quadrants = jnp.zeros_like(Rbool) + quadrants = jnp.where(jnp.logical_and(jnp.logical_not(Rbool), Zbool), 1, quadrants) + quadrants = jnp.where( + jnp.logical_and(jnp.logical_not(Rbool), jnp.logical_not(Zbool)), + 2, + quadrants, + ) + quadrants = jnp.where(jnp.logical_and(Rbool, jnp.logical_not(Zbool)), 3, quadrants) + deltas = quadrants[1:, :] - quadrants[0:-1, :] + deltas = jnp.where(deltas == 3, -1, deltas) + deltas = jnp.where(deltas == -3, 1, deltas) + # then flip sign if the R intercept is > Rcheck and the + # quadrant flipped over a diagonal + b = (Z[1:] / R[1:] - Z[0:-1] / R[0:-1]) / (Z[1:] - Z[0:-1]) + Rint = Rcheck[:, None] - b * (R[1:] - R[0:-1]) / (Z[1:] - Z[0:-1]) + deltas = jnp.where( + jnp.logical_and(jnp.abs(deltas) == 2, Rint.T > Rcheck), + -deltas, + deltas, + ) + pt_sign = jnp.sum(deltas, axis=0) + # positive distance if the check pt is inside the perimeter, else + # negative distance is assigned + # pt_sign = 0 : Means point is OUTSIDE of the perimeter, + # assign positive distance + # pt_sign = +/-4: Means point is INSIDE perimeter, so + # assign negative distance + pt_sign = jnp.where(jnp.isclose(pt_sign, 0), 1, -1) + return pt_sign diff --git a/desc/optimize/_constraint_wrappers.py b/desc/optimize/_constraint_wrappers.py index 93bf5385ae..8d4e480161 100644 --- a/desc/optimize/_constraint_wrappers.py +++ b/desc/optimize/_constraint_wrappers.py @@ -103,6 +103,7 @@ def build(self, use_jit=None, verbose=1): self._A, self._b, self._Z, + self._D, self._unfixed_idx, self._project, self._recover, @@ -113,10 +114,8 @@ def build(self, use_jit=None, verbose=1): self._dim_x = self._objective.dim_x self._dim_x_reduced = self._Z.shape[1] - # equivalent matrix for A[unfixed_idx]@Z == A@unfixed_idx_mat - self._unfixed_idx_mat = ( - jnp.eye(self._objective.dim_x)[:, self._unfixed_idx] @ self._Z - ) + # equivalent matrix for A[unfixed_idx] @ D @ Z == A @ unfixed_idx_mat + self._unfixed_idx_mat = jnp.diag(self._D)[:, self._unfixed_idx] @ self._Z self._built = True timer.stop("Linear constraint projection build") @@ -261,7 +260,7 @@ def grad(self, x_reduced, constants=None): """ x = self.recover(x_reduced) df = self._objective.grad(x, constants) - return df[self._unfixed_idx] @ self._Z + return df[self._unfixed_idx] @ (self._Z * self._D[self._unfixed_idx, None]) def hess(self, x_reduced, constants=None): """Compute Hessian of self.compute_scalar. @@ -281,13 +280,19 @@ def hess(self, x_reduced, constants=None): """ x = self.recover(x_reduced) df = self._objective.hess(x, constants) - return self._Z.T @ df[self._unfixed_idx, :][:, self._unfixed_idx] @ self._Z + return ( + (self._Z.T * (1 / self._D)[None, self._unfixed_idx]) + @ df[self._unfixed_idx, :][:, self._unfixed_idx] + @ (self._Z * self._D[self._unfixed_idx, None]) + ) def _jac(self, x_reduced, constants=None, op="scaled"): x = self.recover(x_reduced) if self._objective._deriv_mode == "blocked": fun = getattr(self._objective, "jac_" + op) - return fun(x, constants)[:, self._unfixed_idx] @ self._Z + return fun(x, constants)[:, self._unfixed_idx] @ ( + self._Z * self._D[self._unfixed_idx, None] + ) v = self._unfixed_idx_mat df = getattr(self._objective, "jvp_" + op)(v.T, x, constants) @@ -401,7 +406,7 @@ def jvp_unscaled(self, v, x_reduced, constants=None): def _vjp(self, v, x_reduced, constants=None, op="vjp_scaled"): x = self.recover(x_reduced) df = getattr(self._objective, op)(v, x, constants) - return df[self._unfixed_idx] @ self._Z + return df[self._unfixed_idx] @ (self._Z * self._D[self._unfixed_idx, None]) def vjp_scaled(self, v, x_reduced, constants=None): """Compute vector-Jacobian product of self.compute_scaled. @@ -533,8 +538,8 @@ def _set_eq_state_vector(self): self._args.remove(arg) linear_constraint = ObjectiveFunction(self._linear_constraints) linear_constraint.build() - _, A, _, self._Z, self._unfixed_idx, _, _ = factorize_linear_constraints( - self._constraint, linear_constraint + _, _, _, self._Z, self._D, self._unfixed_idx, _, _ = ( + factorize_linear_constraints(self._constraint, linear_constraint) ) # dx/dc - goes from the full state to optimization variables for eq @@ -618,14 +623,14 @@ def build(self, use_jit=None, verbose=1): # noqa: C901 ) self._dimx_per_thing = [t.dim_x for t in self.things] - # equivalent matrix for A[unfixed_idx]@Z == A@unfixed_idx_mat + # equivalent matrix for A[unfixed_idx] @ D @ Z == A @ unfixed_idx_mat self._unfixed_idx_mat = jnp.eye(self._objective.dim_x) self._unfixed_idx_mat = jnp.split( self._unfixed_idx_mat, np.cumsum([t.dim_x for t in self.things]), axis=-1 ) - self._unfixed_idx_mat[self._eq_idx] = ( - self._unfixed_idx_mat[self._eq_idx][:, self._unfixed_idx] @ self._Z - ) + self._unfixed_idx_mat[self._eq_idx] = self._unfixed_idx_mat[self._eq_idx][ + :, self._unfixed_idx + ] @ (self._Z * self._D[self._unfixed_idx, None]) self._unfixed_idx_mat = np.concatenate( [np.atleast_2d(foo) for foo in self._unfixed_idx_mat], axis=-1 ) @@ -1018,7 +1023,8 @@ def jvp_unscaled(self, v, x, constants=None): @functools.partial(jit, static_argnames=("self", "op")) def _jvp_f(self, xf, dc, constants, op): Fx = getattr(self._constraint, "jac_" + op)(xf, constants) - Fx_reduced = Fx[:, self._unfixed_idx] @ self._Z + # TODO: replace with self._unfixed_idx_mat? + Fx_reduced = Fx @ jnp.diag(self._D)[:, self._unfixed_idx] @ self._Z Fc = Fx @ (self._dxdc @ dc) Fxh = Fx_reduced cutoff = jnp.finfo(Fxh.dtype).eps * max(Fxh.shape) diff --git a/desc/optimize/tr_subproblems.py b/desc/optimize/tr_subproblems.py index c9149d2bb3..b107bca0a4 100644 --- a/desc/optimize/tr_subproblems.py +++ b/desc/optimize/tr_subproblems.py @@ -421,7 +421,7 @@ def trust_region_step_exact_qr( p_newton = solve_triangular_regularized(R, -Q.T @ f) else: Q, R = qr(J.T, mode="economic") - p_newton = Q @ solve_triangular_regularized(R.T, f, lower=True) + p_newton = Q @ solve_triangular_regularized(R.T, -f, lower=True) def truefun(*_): return p_newton, False, 0.0 diff --git a/desc/perturbations.py b/desc/perturbations.py index 3fa622b248..03529a1ed1 100644 --- a/desc/perturbations.py +++ b/desc/perturbations.py @@ -185,7 +185,7 @@ def perturb( # noqa: C901 - FIXME: break this up into simpler pieces if verbose > 0: print("Factorizing linear constraints") timer.start("linear constraint factorize") - xp, _, _, Z, unfixed_idx, project, recover = factorize_linear_constraints( + xp, _, _, Z, D, unfixed_idx, project, recover = factorize_linear_constraints( objective, constraint ) timer.stop("linear constraint factorize") @@ -291,7 +291,7 @@ def perturb( # noqa: C901 - FIXME: break this up into simpler pieces print("Computing df") timer.start("df computation") Jx = objective.jac_scaled_error(x) - Jx_reduced = Jx[:, unfixed_idx] @ Z @ scale + Jx_reduced = Jx @ jnp.diag(D)[:, unfixed_idx] @ Z @ scale RHS1 = objective.jvp_scaled(tangents, x) if include_f: f = objective.compute_scaled_error(x) @@ -388,9 +388,7 @@ def perturb( # noqa: C901 - FIXME: break this up into simpler pieces con.update_target(eq_new) constraint = ObjectiveFunction(constraints) constraint.build(verbose=verbose) - xp, _, _, Z, unfixed_idx, project, recover = factorize_linear_constraints( - objective, constraint - ) + _, _, _, _, _, _, _, recover = factorize_linear_constraints(objective, constraint) # update other attributes dx_reduced = dx1_reduced + dx2_reduced + dx3_reduced @@ -547,7 +545,7 @@ def optimal_perturb( # noqa: C901 - FIXME: break this up into simpler pieces constraint = ObjectiveFunction(constraints) constraint.build(verbose=verbose) - _, _, _, Z, unfixed_idx, project, recover = factorize_linear_constraints( + _, _, _, Z, D, unfixed_idx, project, recover = factorize_linear_constraints( objective_f, constraint ) @@ -564,7 +562,7 @@ def optimal_perturb( # noqa: C901 - FIXME: break this up into simpler pieces dx2_reduced = 0 # dx/dx_reduced - dxdx_reduced = jnp.eye(eq.dim_x)[:, unfixed_idx] @ Z + dxdx_reduced = jnp.diag(D)[:, unfixed_idx] @ Z # dx/dc dxdc = [] @@ -612,8 +610,8 @@ def optimal_perturb( # noqa: C901 - FIXME: break this up into simpler pieces timer.disp("dg computation") # projections onto optimization space - Fx_reduced = Fx[:, unfixed_idx] @ Z - Gx_reduced = Gx[:, unfixed_idx] @ Z + Fx_reduced = Fx @ jnp.diag(D)[:, unfixed_idx] @ Z + Gx_reduced = Gx @ jnp.diag(D)[:, unfixed_idx] @ Z Fc = Fx @ dxdc Gc = Gx @ dxdc @@ -752,9 +750,7 @@ def optimal_perturb( # noqa: C901 - FIXME: break this up into simpler pieces con.update_target(eq_new) constraint = ObjectiveFunction(constraints) constraint.build(verbose=verbose) - _, _, _, Z, unfixed_idx, project, recover = factorize_linear_constraints( - objective_f, constraint - ) + _, _, _, _, _, _, _, recover = factorize_linear_constraints(objective_f, constraint) # update other attributes dx_reduced = dx1_reduced + dx2_reduced diff --git a/desc/plotting.py b/desc/plotting.py index 551e4fe469..fdcb9a393c 100644 --- a/desc/plotting.py +++ b/desc/plotting.py @@ -16,7 +16,7 @@ from desc.backend import sign from desc.basis import fourier, zernike_radial_poly -from desc.coils import CoilSet +from desc.coils import CoilSet, _Coil from desc.compute import data_index, get_transforms from desc.compute.utils import _parse_parameterization, surface_averages_map from desc.equilibrium.coords import map_coordinates @@ -2394,6 +2394,11 @@ def plot_coils(coils, grid=None, fig=None, return_data=False, **kwargs): ValueError, f"plot_coils got unexpected keyword argument: {kwargs.keys()}", ) + errorif( + not isinstance(coils, _Coil), + ValueError, + "Expected `coils` to be of type `_Coil`, instead got type" f" {type(coils)}", + ) if not isinstance(lw, (list, tuple)): lw = [lw] @@ -2756,22 +2761,14 @@ def plot_boozer_surface( # default grids if grid_compute is None: - if eq_switch: - grid_kwargs = { - "rho": rho, - "M": 4 * thing.M, - "N": 4 * thing.N, - "NFP": thing.NFP, - "endpoint": False, - } - else: - grid_kwargs = { - "rho": rho, - "M": 50, - "N": 50, - "NFP": thing.NFP, - "endpoint": False, - } + # grid_compute only used for Equilibrium, not OmnigenousField + grid_kwargs = { + "rho": rho, + "M": 4 * getattr(thing, "M", 1), + "N": 4 * getattr(thing, "N", 1), + "NFP": thing.NFP, + "endpoint": False, + } grid_compute = _get_grid(**grid_kwargs) if grid_plot is None: grid_kwargs = { @@ -2779,7 +2776,7 @@ def plot_boozer_surface( "theta": 91, "zeta": 91, "NFP": thing.NFP, - "endpoint": True, + "endpoint": eq_switch, } grid_plot = _get_grid(**grid_kwargs) @@ -2816,7 +2813,7 @@ def plot_boozer_surface( warnings.simplefilter("ignore") data = thing.compute( ["theta_B", "zeta_B", "|B|"], - grid=grid_compute, + grid=grid_plot, helicity=thing.helicity, iota=iota, ) diff --git a/desc/utils.py b/desc/utils.py index 36b58f2281..a1df551229 100644 --- a/desc/utils.py +++ b/desc/utils.py @@ -532,7 +532,7 @@ def errorif(cond, err=ValueError, msg=""): just AssertionError. """ if cond: - raise err(msg) + raise err(colored(msg, "red")) class ResolutionWarning(UserWarning): @@ -544,7 +544,7 @@ class ResolutionWarning(UserWarning): def warnif(cond, err=UserWarning, msg=""): """Throw a warning if condition is met.""" if cond: - warnings.warn(msg, err) + warnings.warn(colored(msg, "yellow"), err) def check_nonnegint(x, name="", allow_none=True): diff --git a/desc/vmec.py b/desc/vmec.py index 7a68f52de9..a0212dbd31 100644 --- a/desc/vmec.py +++ b/desc/vmec.py @@ -192,7 +192,7 @@ def load( constraints = maybe_add_self_consistency(eq, constraints) objective = ObjectiveFunction(constraints) objective.build(verbose=0) - _, _, _, _, _, project, recover = factorize_linear_constraints( + _, _, _, _, _, _, project, recover = factorize_linear_constraints( objective, objective ) args = objective.unpack_state(recover(project(objective.x(eq))), False)[0] @@ -1687,7 +1687,7 @@ def root_fun(theta): axis=-1, ) theta_star_k = theta + lmbda # theta* = theta + lambda - err = theta_star - theta_star_k + err = theta_star - theta_star_k # FIXME: mod by 2pi return err out = optimize.root( @@ -1751,7 +1751,9 @@ def compute_coord_surfaces(cls, equil, vmec_data, Nr=10, Nt=8, Nz=None, **kwargs # angle in PEST-like flux coordinates # find theta angles corresponding to desired theta* angles - v_grid = Grid(equil.compute_theta_coords(t_grid.nodes)) + v_grid = Grid( + equil.map_coordinates(t_grid.nodes, inbasis=("rho", "theta_PEST", "zeta")) + ) r_coords_desc = equil.compute(["R", "Z"], grid=r_grid) v_coords_desc = equil.compute(["R", "Z"], grid=v_grid) diff --git a/docs/adding_compute_funs.rst b/docs/adding_compute_funs.rst index 66e3c2c047..67662cb571 100644 --- a/docs/adding_compute_funs.rst +++ b/docs/adding_compute_funs.rst @@ -27,6 +27,7 @@ The full code is below: coordinates="rtz", data=["sqrt(g)", "B_zeta_t", "B_theta_z"], axis_limit_data=["sqrt(g)_r", "B_zeta_rt", "B_theta_rz"], + resolution_requirement="", parameterization="desc.equilibrium.equilibrium.Equilibrium", ) def _J_sup_rho(params, transforms, profiles, data, **kwargs): @@ -96,6 +97,12 @@ metadata about the quantity. The necessary fields are detailed below: the quantities mentioned above to evaluate the magnetic axis limit. These dependencies are specified in ``axis_limit_data``. The dependencies specified in this list are marked to be computed only when there is a node at the magnetic axis. +* ``resolution_requirement``: Resolution requirements in coordinates. + I.e. "r" expects radial resolution in the grid. Likewise, "rtz" is shorthand for + "rho, theta, zeta" and indicates the computation expects a grid with radial, + poloidal, and toroidal resolution. If the computation simply performs + pointwise operations, instead of a reduction (such as integration) over a + coordinate, then an empty string may be used to indicate no requirements. * ``parameterization``: what sorts of DESC objects is this function for. Most functions will just be for ``Equilibrium``, but some methods may also be for ``desc.geometry.core.Curve``, or specific types eg ``desc.geometry.curve.FourierRZCurve``. If a quantity is computed differently @@ -148,8 +155,8 @@ if ``False`` is printed, then the limit of the quantity does not evaluate as fin The tests automatically detect this, so no further action is needed from developers in this case. -The second step is to run the ``test_compute_everything`` test located in the ``tests/test_compute_funs.py`` file. -This can be done with the command :console:`pytest -k test_compute_everything tests/test_compute_funs.py`. +The second step is to run the ``test_compute_everything`` test located in the ``tests/test_compute_everything.py`` file. +This can be done with the command :console:`pytest tests/test_compute_everything.py`. This test is a regression test to ensure that compute quantities in each new update of DESC do not differ significantly from previous versions of DESC. Since the new quantity did not exist in previous versions of DESC, one must run this test diff --git a/docs/write_variables.py b/docs/write_variables.py index eec46136f1..13b15b9061 100644 --- a/docs/write_variables.py +++ b/docs/write_variables.py @@ -78,6 +78,12 @@ def write_csv(parameterization): The keyword argument ``basis='xyz'`` can be used to convert the variables into Cartesian coordinates :math:`(X,Y,Z)`. ``basis`` must be one of ``{'rpz', 'xyz'}``. +Our convention to denote partial derivatives is an underscore followed by the first +letter of the coordinate that the partial derivative is taken with respect to. Unless +otherwise specified or implied by the variable name, these partial derivatives are +those of the DESC :math:`\rho, \theta, \zeta` coordinate system. For example, ``|B|_z`` +is :math:`(\partial \vert B \vert / \partial\zeta)|_{\rho, \theta}`. + """ block = """ diff --git a/tests/inputs/master_compute_data_rpz.pkl b/tests/inputs/master_compute_data_rpz.pkl index f230d7a816..8a9086d237 100644 Binary files a/tests/inputs/master_compute_data_rpz.pkl and b/tests/inputs/master_compute_data_rpz.pkl differ diff --git a/tests/inputs/master_compute_data_xyz.pkl b/tests/inputs/master_compute_data_xyz.pkl deleted file mode 100644 index 0de07ede29..0000000000 Binary files a/tests/inputs/master_compute_data_xyz.pkl and /dev/null differ diff --git a/tests/test_bootstrap.py b/tests/test_bootstrap.py index 34ffe08e22..811480bbaf 100644 --- a/tests/test_bootstrap.py +++ b/tests/test_bootstrap.py @@ -1589,7 +1589,7 @@ def test_bootstrap_optimization_comparison_qa(): objective=objective, constraints=constraints, optimizer="proximal-lsq-exact", - maxiter=4, + maxiter=5, gtol=1e-16, verbose=3, ) @@ -1622,5 +1622,5 @@ def test_bootstrap_optimization_comparison_qa(): grid.compress(data2[""]), grid.compress(data2[" Redl"]), rtol=1.8e-2 ) np.testing.assert_allclose( - grid.compress(data1[""]), grid.compress(data2[""]), rtol=1.8e-2 + grid.compress(data1[""]), grid.compress(data2[""]), rtol=1.9e-2 ) diff --git a/tests/test_coils.py b/tests/test_coils.py index ffd6558e8b..704ad5f761 100644 --- a/tests/test_coils.py +++ b/tests/test_coils.py @@ -542,9 +542,9 @@ def test_convert_type(self): # MixedCoilSet coils1 = MixedCoilSet.linspaced_linear(coil, displacement=[0, 0, 2.5], n=4) - coils2 = coils1.to_SplineXYZ(grid=grid) - coils3 = coils1.to_FourierXYZ(grid=grid) - coils4 = coils1.to_FourierPlanar(grid=grid) + coils2 = coils1.to_SplineXYZ(grid=grid, check_intersection=False) + coils3 = coils1.to_FourierXYZ(grid=grid, check_intersection=False) + coils4 = coils1.to_FourierPlanar(grid=grid, check_intersection=False) assert isinstance(coils2, MixedCoilSet) assert isinstance(coils3, MixedCoilSet) assert isinstance(coils4, MixedCoilSet) @@ -580,9 +580,9 @@ def test_convert_type(self): # CoilSet coil = coils3[0] # FourierXYZCoil coils5 = CoilSet.linspaced_linear(coil, displacement=[0, 0, 3.5], n=6) - coils6 = coils5.to_SplineXYZ(grid=grid) - coils7 = coils5.to_FourierRZ(grid=grid) - coils8 = coils5.to_FourierPlanar(grid=grid) + coils6 = coils5.to_SplineXYZ(grid=grid, check_intersection=False) + coils7 = coils5.to_FourierRZ(grid=grid, check_intersection=False) + coils8 = coils5.to_FourierPlanar(grid=grid, check_intersection=False) assert isinstance(coils6, CoilSet) assert isinstance(coils7, CoilSet) assert isinstance(coils8, CoilSet) diff --git a/tests/test_compute_everything.py b/tests/test_compute_everything.py new file mode 100644 index 0000000000..aff0345e8f --- /dev/null +++ b/tests/test_compute_everything.py @@ -0,0 +1,252 @@ +"""Test that the computations on this branch agree with those on master.""" + +import pickle +import warnings + +import numpy as np +import pytest + +from desc.coils import FourierPlanarCoil, FourierRZCoil, FourierXYZCoil, SplineXYZCoil +from desc.compute import data_index, xyz2rpz, xyz2rpz_vec +from desc.examples import get +from desc.geometry import ( + FourierPlanarCurve, + FourierRZCurve, + FourierRZToroidalSurface, + FourierXYZCurve, + ZernikeRZToroidalSection, +) +from desc.grid import LinearGrid +from desc.magnetic_fields import ( + CurrentPotentialField, + FourierCurrentPotentialField, + OmnigenousField, +) +from desc.utils import ResolutionWarning, errorif + + +def _compare_against_master( + p, data, master_data, error=False, update_master_data=False +): + + for name in data[p]: + if p in master_data and name in master_data[p]: + try: + np.testing.assert_allclose( + actual=data[p][name], + desired=master_data[p][name], + atol=1e-10, + rtol=1e-10, + err_msg=f"Parameterization: {p}. Name: {name}.", + ) + except AssertionError as e: + error = True + print(e) + else: # update master data with new compute quantity + update_master_data = True + + return error, update_master_data + + +def _xyz_to_rpz(data, name): + if name in ["x", "center"]: + res = xyz2rpz(data[name]) + else: + res = xyz2rpz_vec(data[name], phi=data["phi"]) + return res + + +def _compare_against_rpz(p, data, data_rpz, coordinate_conversion_func): + for name in data: + if data_index[p][name]["dim"] != 3: + continue + res = coordinate_conversion_func(data, name) - data_rpz[name] + errorif( + not np.all( + ( + np.isclose(res, 0, rtol=1e-8, atol=1e-8) + | np.isclose(np.abs(res[:, 1]), 2 * np.pi, rtol=1e-8, atol=1e-8)[ + :, np.newaxis + ] + ) + | (~np.isfinite(data_rpz[name])) + ), + AssertionError, + msg=f"Parameterization: {p}. Name: {name}. Residual {res}", + ) + + +@pytest.mark.unit +def test_compute_everything(): + """Test that the computations on this branch agree with those on master. + + Also make sure we can compute everything without errors. Computed quantities + are both in "rpz" and "xyz" basis. + """ + elliptic_cross_section_with_torsion = { + "R_lmn": [10, 1, 0.2], + "Z_lmn": [-2, -0.2], + "modes_R": [[0, 0], [1, 0], [0, 1]], + "modes_Z": [[-1, 0], [0, -1]], + } + things = { + # equilibria + "desc.equilibrium.equilibrium.Equilibrium": get("W7-X"), + # curves + "desc.geometry.curve.FourierXYZCurve": FourierXYZCurve( + X_n=[5, 10, 2], Y_n=[1, 2, 3], Z_n=[-4, -5, -6] + ), + "desc.geometry.curve.FourierRZCurve": FourierRZCurve( + R_n=[10, 1, 0.2], Z_n=[-2, -0.2], modes_R=[0, 1, 2], modes_Z=[-1, -2], NFP=2 + ), + "desc.geometry.curve.FourierPlanarCurve": FourierPlanarCurve( + center=[10, 1, 3], normal=[1, 2, 3], r_n=[1, 2, 3], modes=[0, 1, 2] + ), + "desc.geometry.curve.SplineXYZCurve": FourierXYZCurve( + X_n=[5, 10, 2], Y_n=[1, 2, 3], Z_n=[-4, -5, -6] + ).to_SplineXYZ(grid=LinearGrid(N=50)), + # surfaces + "desc.geometry.surface.FourierRZToroidalSurface": FourierRZToroidalSurface( + **elliptic_cross_section_with_torsion + ), + "desc.geometry.surface.ZernikeRZToroidalSection": ZernikeRZToroidalSection( + **elliptic_cross_section_with_torsion + ), + # magnetic fields + "desc.magnetic_fields._current_potential.CurrentPotentialField": CurrentPotentialField( # noqa:E501 + **elliptic_cross_section_with_torsion, + potential=lambda theta, zeta, G: G * zeta / 2 / np.pi, + potential_dtheta=lambda theta, zeta, G: np.zeros_like(theta), + potential_dzeta=lambda theta, zeta, G: G * np.ones_like(theta) / 2 / np.pi, + params={"G": 1e7}, + ), + "desc.magnetic_fields._current_potential.FourierCurrentPotentialField": ( + FourierCurrentPotentialField( + **elliptic_cross_section_with_torsion, I=0, G=1e7 + ) + ), + "desc.magnetic_fields._core.OmnigenousField": OmnigenousField( + L_B=0, + M_B=4, + L_x=0, + M_x=1, + N_x=1, + NFP=2, + helicity=(0, 2), + B_lm=np.array([0.8, 0.9, 1.1, 1.2]), + x_lmn=np.array([0, -np.pi / 8, 0, np.pi / 8, 0, np.pi / 4]), + ), + # coils + "desc.coils.FourierRZCoil": FourierRZCoil( + R_n=[10, 1, 0.2], Z_n=[-2, -0.2], modes_R=[0, 1, 2], modes_Z=[-1, -2], NFP=2 + ), + "desc.coils.FourierXYZCoil": FourierXYZCoil( + X_n=[5, 10, 2], Y_n=[1, 2, 3], Z_n=[-4, -5, -6] + ), + "desc.coils.FourierPlanarCoil": FourierPlanarCoil( + current=5, + center=[10, 1, 3], + normal=[1, 2, 3], + r_n=[1, 2, 3], + modes=[0, 1, 2], + ), + "desc.coils.SplineXYZCoil": SplineXYZCoil( + current=5, X=[5, 10, 2, 5], Y=[1, 2, 3, 1], Z=[-4, -5, -6, -4] + ), + } + assert things.keys() == data_index.keys(), ( + f"Missing the parameterization {data_index.keys() - things.keys()}" + f" to test against master." + ) + # use this low resolution grid for equilibria to reduce file size + eqgrid = LinearGrid( + L=9, + M=5, + N=5, + NFP=things["desc.equilibrium.equilibrium.Equilibrium"].NFP, + sym=things["desc.equilibrium.equilibrium.Equilibrium"].sym, + axis=True, + ) + curvegrid1 = LinearGrid(N=10) + curvegrid2 = LinearGrid(N=10, NFP=2) + fieldgrid = LinearGrid( + L=2, + M=4, + N=5, + NFP=things["desc.magnetic_fields._core.OmnigenousField"].NFP, + sym=False, + axis=True, + ) + grid = { + "desc.equilibrium.equilibrium.Equilibrium": {"grid": eqgrid}, + "desc.geometry.curve.FourierXYZCurve": {"grid": curvegrid1}, + "desc.geometry.curve.FourierRZCurve": {"grid": curvegrid2}, + "desc.geometry.curve.FourierPlanarCurve": {"grid": curvegrid1}, + "desc.geometry.curve.SplineXYZCurve": {"grid": curvegrid1}, + "desc.magnetic_fields._core.OmnigenousField": {"grid": fieldgrid}, + } + + with open("tests/inputs/master_compute_data_rpz.pkl", "rb") as file: + master_data_rpz = pickle.load(file) + this_branch_data_rpz = {} + update_master_data_rpz = False + error_rpz = False + + # some things can't compute "phi" and therefore can't convert to XYZ basis + no_xyz_things = ["desc.magnetic_fields._core.OmnigenousField"] + + with warnings.catch_warnings(): + # Max resolution of master_compute_data.pkl limited by GitHub file + # size cap at 100 mb, so can't hit suggested resolution for some things. + warnings.filterwarnings("ignore", category=ResolutionWarning) + for p in things: + names = { + name + for name in data_index[p] + # Skip these quantities as they should be covered in other tests. + if not data_index[p][name]["source_grid_requirement"] + } + this_branch_data_rpz[p] = things[p].compute( + list(names), **grid.get(p, {}), basis="rpz" + ) + # make sure we can compute everything + assert this_branch_data_rpz[p].keys() == names, ( + f"Parameterization: {p}. Can't compute " + + f"{names - this_branch_data_rpz[p].keys()}." + ) + # compare data against master branch + error_rpz, update_master_data_rpz = _compare_against_master( + p, + this_branch_data_rpz, + master_data_rpz, + error_rpz, + update_master_data_rpz, + ) + + # test compute in XYZ basis + if p in no_xyz_things: + continue + # remove quantities that are not implemented in the XYZ basis + # TODO: generalize this instead of hard-coding for "grad(B)" & dependencies + names_xyz = ( + names - {"grad(B)", "|grad(B)|", "L_grad(B)"} + if "grad(B)" in names + else names + ) + this_branch_data_xyz = things[p].compute( + list(names_xyz), **grid.get(p, {}), basis="xyz" + ) + assert this_branch_data_xyz.keys() == names_xyz, ( + f"Parameterization: {p}. Can't compute " + + f"{names_xyz - this_branch_data_xyz.keys()}." + ) + _compare_against_rpz( + p, this_branch_data_xyz, this_branch_data_rpz[p], _xyz_to_rpz + ) + + if not error_rpz and update_master_data_rpz: + # then update the master compute data + with open("tests/inputs/master_compute_data_rpz.pkl", "wb") as file: + # remember to git commit this file + pickle.dump(this_branch_data_rpz, file) + assert not error_rpz diff --git a/tests/test_compute_funs.py b/tests/test_compute_funs.py index 21c768e215..43c3d81449 100644 --- a/tests/test_compute_funs.py +++ b/tests/test_compute_funs.py @@ -1,32 +1,16 @@ """Tests for compute functions.""" -import copy -import pickle -import warnings - import numpy as np import pytest from scipy.signal import convolve2d -from desc.coils import FourierPlanarCoil, FourierRZCoil, FourierXYZCoil, SplineXYZCoil -from desc.compute import data_index, rpz2xyz_vec +from desc.compute import rpz2xyz_vec +from desc.compute.utils import dot from desc.equilibrium import Equilibrium from desc.examples import get -from desc.geometry import ( - FourierPlanarCurve, - FourierRZCurve, - FourierRZToroidalSurface, - FourierXYZCurve, - ZernikeRZToroidalSection, -) +from desc.geometry import FourierRZToroidalSurface from desc.grid import LinearGrid from desc.io import load -from desc.magnetic_fields import ( - CurrentPotentialField, - FourierCurrentPotentialField, - OmnigenousField, -) -from desc.utils import ResolutionWarning # convolve kernel is reverse of FD coeffs FD_COEF_1_2 = np.array([-1 / 2, 0, 1 / 2])[::-1] @@ -1052,8 +1036,8 @@ def test_magnetic_pressure_gradient(DummyStellarator): num_rho = 110 grid = LinearGrid(NFP=eq.NFP, rho=num_rho) drho = grid.nodes[1, 0] - data = eq.compute(["|B|", "grad(|B|^2)_rho"], grid=grid) - B2_r = np.convolve(data["|B|"] ** 2, FD_COEF_1_4, "same") / drho + data = eq.compute(["|B|^2", "grad(|B|^2)_rho"], grid=grid) + B2_r = np.convolve(data["|B|^2"], FD_COEF_1_4, "same") / drho np.testing.assert_allclose( data["grad(|B|^2)_rho"][3:-2], B2_r[3:-2], @@ -1065,8 +1049,8 @@ def test_magnetic_pressure_gradient(DummyStellarator): num_theta = 90 grid = LinearGrid(NFP=eq.NFP, theta=num_theta) dtheta = grid.nodes[1, 1] - data = eq.compute(["|B|", "grad(|B|^2)_theta"], grid=grid) - B2_t = np.convolve(data["|B|"] ** 2, FD_COEF_1_4, "same") / dtheta + data = eq.compute(["|B|^2", "grad(|B|^2)_theta"], grid=grid) + B2_t = np.convolve(data["|B|^2"], FD_COEF_1_4, "same") / dtheta np.testing.assert_allclose( data["grad(|B|^2)_theta"][2:-2], B2_t[2:-2], @@ -1078,8 +1062,8 @@ def test_magnetic_pressure_gradient(DummyStellarator): num_zeta = 90 grid = LinearGrid(NFP=eq.NFP, zeta=num_zeta) dzeta = grid.nodes[1, 2] - data = eq.compute(["|B|", "grad(|B|^2)_zeta"], grid=grid) - B2_z = np.convolve(data["|B|"] ** 2, FD_COEF_1_4, "same") / dzeta + data = eq.compute(["|B|^2", "grad(|B|^2)_zeta"], grid=grid) + B2_z = np.convolve(data["|B|^2"], FD_COEF_1_4, "same") / dzeta np.testing.assert_allclose( data["grad(|B|^2)_zeta"][2:-2], B2_z[2:-2], @@ -1150,209 +1134,6 @@ def test_boozer_transform(): ) -@pytest.mark.unit -def test_compute_everything(): - """Test that the computations on this branch agree with those on master. - - Also make sure we can compute everything without errors. Computed quantities - are both in "rpz" and "xyz" basis. - """ - elliptic_cross_section_with_torsion = { - "R_lmn": [10, 1, 0.2], - "Z_lmn": [-2, -0.2], - "modes_R": [[0, 0], [1, 0], [0, 1]], - "modes_Z": [[-1, 0], [0, -1]], - } - things = { - # equilibria - "desc.equilibrium.equilibrium.Equilibrium": get("W7-X"), - # curves - "desc.geometry.curve.FourierXYZCurve": FourierXYZCurve( - X_n=[5, 10, 2], Y_n=[1, 2, 3], Z_n=[-4, -5, -6] - ), - "desc.geometry.curve.FourierRZCurve": FourierRZCurve( - R_n=[10, 1, 0.2], Z_n=[-2, -0.2], modes_R=[0, 1, 2], modes_Z=[-1, -2], NFP=2 - ), - "desc.geometry.curve.FourierPlanarCurve": FourierPlanarCurve( - center=[10, 1, 3], normal=[1, 2, 3], r_n=[1, 2, 3], modes=[0, 1, 2] - ), - "desc.geometry.curve.SplineXYZCurve": FourierXYZCurve( - X_n=[5, 10, 2], Y_n=[1, 2, 3], Z_n=[-4, -5, -6] - ).to_SplineXYZ(grid=LinearGrid(N=50)), - # surfaces - "desc.geometry.surface.FourierRZToroidalSurface": FourierRZToroidalSurface( - **elliptic_cross_section_with_torsion - ), - "desc.geometry.surface.ZernikeRZToroidalSection": ZernikeRZToroidalSection( - **elliptic_cross_section_with_torsion - ), - # magnetic fields - "desc.magnetic_fields._current_potential.CurrentPotentialField": CurrentPotentialField( # noqa:E501 - **elliptic_cross_section_with_torsion, - potential=lambda theta, zeta, G: G * zeta / 2 / np.pi, - potential_dtheta=lambda theta, zeta, G: np.zeros_like(theta), - potential_dzeta=lambda theta, zeta, G: G * np.ones_like(theta) / 2 / np.pi, - params={"G": 1e7}, - ), - "desc.magnetic_fields._current_potential.FourierCurrentPotentialField": ( - FourierCurrentPotentialField( - **elliptic_cross_section_with_torsion, I=0, G=1e7 - ) - ), - "desc.magnetic_fields._core.OmnigenousField": OmnigenousField( - L_B=0, - M_B=4, - L_x=0, - M_x=1, - N_x=1, - NFP=2, - helicity=(0, 2), - B_lm=np.array([0.8, 0.9, 1.1, 1.2]), - x_lmn=np.array([0, -np.pi / 8, 0, np.pi / 8, 0, np.pi / 4]), - ), - # coils - "desc.coils.FourierRZCoil": FourierRZCoil( - R_n=[10, 1, 0.2], Z_n=[-2, -0.2], modes_R=[0, 1, 2], modes_Z=[-1, -2], NFP=2 - ), - "desc.coils.FourierXYZCoil": FourierXYZCoil( - X_n=[5, 10, 2], Y_n=[1, 2, 3], Z_n=[-4, -5, -6] - ), - "desc.coils.FourierPlanarCoil": FourierPlanarCoil( - current=5, - center=[10, 1, 3], - normal=[1, 2, 3], - r_n=[1, 2, 3], - modes=[0, 1, 2], - ), - "desc.coils.SplineXYZCoil": SplineXYZCoil( - current=5, X=[5, 10, 2, 5], Y=[1, 2, 3, 1], Z=[-4, -5, -6, -4] - ), - } - assert things.keys() == data_index.keys(), ( - f"Missing the parameterization {data_index.keys() - things.keys()}" - f" to test against master." - ) - # use this low resolution grid for equilibria to reduce file size - eqgrid = LinearGrid( - L=9, - M=5, - N=5, - NFP=things["desc.equilibrium.equilibrium.Equilibrium"].NFP, - sym=things["desc.equilibrium.equilibrium.Equilibrium"].sym, - axis=True, - ) - curvegrid1 = LinearGrid(N=10) - curvegrid2 = LinearGrid(N=10, NFP=2) - fieldgrid = LinearGrid( - L=2, - M=4, - N=5, - NFP=things["desc.magnetic_fields._core.OmnigenousField"].NFP, - sym=False, - axis=True, - ) - grid = { - "desc.equilibrium.equilibrium.Equilibrium": {"grid": eqgrid}, - "desc.geometry.curve.FourierXYZCurve": {"grid": curvegrid1}, - "desc.geometry.curve.FourierRZCurve": {"grid": curvegrid2}, - "desc.geometry.curve.FourierPlanarCurve": {"grid": curvegrid1}, - "desc.geometry.curve.SplineXYZCurve": {"grid": curvegrid1}, - "desc.magnetic_fields._core.OmnigenousField": {"grid": fieldgrid}, - } - - with open("tests/inputs/master_compute_data_rpz.pkl", "rb") as file: - master_data_rpz = pickle.load(file) - with open("tests/inputs/master_compute_data_xyz.pkl", "rb") as file: - master_data_xyz = pickle.load(file) - this_branch_data_rpz = {} - this_branch_data_xyz = {} - update_master_data_rpz = False - update_master_data_xyz = False - error_rpz = False - error_xyz = False - - # some things can't compute "phi" and therefore can't convert to XYZ basis - no_xyz_things = ["desc.magnetic_fields._core.OmnigenousField"] - - for p in things: - with warnings.catch_warnings(): - # Max resolution of master_compute_data.pkl limited by GitHub file - # size cap at 100 mb, so can't hit suggested resolution for some things. - warnings.filterwarnings("ignore", category=ResolutionWarning) - this_branch_data_rpz[p] = things[p].compute( - list(data_index[p].keys()), **grid.get(p, {}), basis="rpz" - ) - # make sure we can compute everything - assert this_branch_data_rpz[p].keys() == data_index[p].keys(), ( - f"Parameterization: {p}. Can't compute " - + f"{data_index[p].keys() - this_branch_data_rpz[p].keys()}." - ) - # compare data against master branch - for name in this_branch_data_rpz[p]: - if p in master_data_rpz and name in master_data_rpz[p]: - try: - np.testing.assert_allclose( - actual=this_branch_data_rpz[p][name], - desired=master_data_rpz[p][name], - atol=1e-10, - rtol=1e-10, - err_msg=f"Parameterization: {p}. Name: {name}.", - ) - except AssertionError as e: - error_rpz = True - print(e) - else: # update master data with new compute quantity - update_master_data_rpz = True - - # test compute in XYZ basis - if p not in no_xyz_things: - # remove quantities that are not implemented in the XYZ basis - # TODO: generalize this instead of hard-coding for "grad(B)" & dependencies - data_index_xyz = copy.deepcopy(data_index) - if "grad(B)" in list(data_index[p].keys()): - del data_index_xyz[p]["grad(B)"] - del data_index_xyz[p]["|grad(B)|"] - del data_index_xyz[p]["L_grad(B)"] - with warnings.catch_warnings(): - # Max resolution of master_compute_data.pkl limited by GitHub file - # size cap at 100 mb, so can't hit suggested resolution for some things. - warnings.filterwarnings("ignore", category=ResolutionWarning) - this_branch_data_xyz[p] = things[p].compute( - list(data_index_xyz[p].keys()), **grid.get(p, {}), basis="xyz" - ) - assert this_branch_data_xyz[p].keys() == data_index_xyz[p].keys(), ( - f"Parameterization: {p}. Can't compute " - + f"{data_index_xyz[p].keys() - this_branch_data_xyz[p].keys()}." - ) - # compare data against master branch - for name in this_branch_data_xyz[p]: - if p in master_data_xyz and name in master_data_xyz[p]: - try: - np.testing.assert_allclose( - actual=this_branch_data_xyz[p][name], - desired=master_data_xyz[p][name], - atol=1e-10, - rtol=1e-10, - err_msg=f"Parameterization: {p}. Name: {name}.", - ) - except AssertionError as e: - error_xyz = True - print(e) - else: # update master data with new compute quantity - update_master_data_xyz = True - - # update the master compute data, if necessary - # remember to git commit these files - if not error_rpz and update_master_data_rpz: - with open("tests/inputs/master_compute_data_rpz.pkl", "wb") as file: - pickle.dump(this_branch_data_rpz, file) - if not error_xyz and update_master_data_xyz: - with open("tests/inputs/master_compute_data_xyz.pkl", "wb") as file: - pickle.dump(this_branch_data_xyz, file) - assert not error_rpz - assert not error_xyz - - @pytest.mark.unit def test_compute_averages(): """Test that computing averages uses the correct grid.""" @@ -1738,3 +1519,35 @@ def test_surface_equilibrium_geometry(): rtol=3e-13, atol=1e-13, ) + + +@pytest.mark.unit +def test_parallel_grad(): + """Test geometric and physical methods of computing parallel gradients agree.""" + eq = get("W7-X") + with pytest.warns(UserWarning, match="Reducing radial"): + eq.change_resolution(2, 2, 2, 4, 4, 4) + data = eq.compute( + [ + "e_zeta|r,a", + "B", + "B^zeta", + "|B|_z|r,a", + "grad(|B|)", + "|e_zeta|r,a|_z|r,a", + "B^zeta_z|r,a", + "|B|", + ], + ) + np.testing.assert_allclose(data["e_zeta|r,a"], (data["B"].T / data["B^zeta"]).T) + np.testing.assert_allclose( + data["|B|_z|r,a"], dot(data["grad(|B|)"], data["e_zeta|r,a"]) + ) + np.testing.assert_allclose( + data["|e_zeta|r,a|_z|r,a"], + data["|B|_z|r,a"] / np.abs(data["B^zeta"]) + - data["|B|"] + * data["B^zeta_z|r,a"] + * np.sign(data["B^zeta"]) + / data["B^zeta"] ** 2, + ) diff --git a/tests/test_data_index.py b/tests/test_data_index.py index b582ff878c..ebb3df2479 100644 --- a/tests/test_data_index.py +++ b/tests/test_data_index.py @@ -11,121 +11,127 @@ from desc.utils import errorif -class TestDataIndex: - """Tests for things related to data_index.""" - - @staticmethod - def get_matches(fun, pattern): - """Return all matches of ``pattern`` in source code of function ``fun``.""" - src = inspect.getsource(fun) - # attempt to remove any decorator functions - # (currently works without this filter, but better to be defensive) +def _get_matches(fun, pattern, ignore_comments=True): + """Return all matches of ``pattern`` in source code of function ``fun``.""" + src = inspect.getsource(fun) + if ignore_comments: + # remove any decorator functions src = src.partition("def ")[2] - # attempt to remove comments + # remove comments src = "\n".join(line.partition("#")[0] for line in src.splitlines()) - matches = pattern.findall(src) - matches = {s.strip().strip('"') for s in matches} - return matches + matches = pattern.findall(src) + matches = {s.strip().strip('"') for s in matches} + return matches + + +def _get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"): + """Get parameterization of thing computed by function ``fun``.""" + pattern = re.compile(r'parameterization=(?:\[([^]]+)]|"([^"]+)")') + decorator = inspect.getsource(fun).partition("def ")[0] + matches = pattern.findall(decorator) + # if list was found, split strings in list, else string was found so get that + matches = [match[0].split(",") if match[0] else [match[1]] for match in matches] + # flatten the list + matches = {s.strip().strip('"') for sublist in matches for s in sublist} + matches.discard("") + return matches if matches else {default} - @staticmethod - def get_parameterization(fun, default="desc.equilibrium.equilibrium.Equilibrium"): - """Get parameterization of thing computed by function ``fun``.""" - pattern = re.compile(r'parameterization=(?:\[([^]]+)]|"([^"]+)")') - decorator = inspect.getsource(fun).partition("def ")[0] - matches = pattern.findall(decorator) - # if list was found, split strings in list, else string was found so get that - matches = [match[0].split(",") if match[0] else [match[1]] for match in matches] - # flatten the list - matches = {s.strip().strip('"') for sublist in matches for s in sublist} - matches.discard("") - return matches if matches else {default} - @pytest.mark.unit - def test_data_index_deps(self): - """Ensure developers do not add extra (or forget needed) dependencies. +@pytest.mark.unit +def test_data_index_deps(): + """Ensure developers do not add extra (or forget needed) dependencies. - The regular expressions used in this test will fail to detect the data - dependencies in the source code of compute functions if the query to - the key in the data dictionary is split across multiple lines. - To avoid failing this test unnecessarily in this case, try to refactor - code by wrapping the query to the key in the data dictionary inside a - parenthesis. + The regular expressions used in this test will fail to detect the data + dependencies in the source code of compute functions if the query to + the key in the data dictionary is split across multiple lines. + To avoid failing this test unnecessarily in this case, try to refactor + code by wrapping the query to the key in the data dictionary inside a + parenthesis. - Examples - -------- - .. code-block:: python + Examples + -------- + .. code-block:: python - # Don't do this. - x_square = data[ - "x" - ] ** 2 - # Either do this - x_square = ( - data["x"] - ) ** 2 - # or do this - x_square = data["x"] ** 2 + # Don't do this. + x_square = data[ + "x" + ] ** 2 + # Either do this + x_square = ( + data["x"] + ) ** 2 + # or do this + x_square = data["x"] ** 2 - """ - queried_deps = {p: {} for p in _class_inheritance} + """ + queried_deps = {p: {} for p in _class_inheritance} - pattern_names = re.compile(r"(? inheritance_order.index( + data_index[base_class][register_name][ + "parameterization" + ] + ): + continue + queried_deps[base_class][register_name] = deps + aliases = data_index[base_class][register_name]["aliases"] + for alias in aliases: + queried_deps[base_class][alias] = deps - for p in data_index: - for name, val in data_index[p].items(): - err_msg = f"Parameterization: {p}. Name: {name}." - deps = val["dependencies"] - data = set(deps["data"]) - axis_limit_data = set(deps["axis_limit_data"]) - profiles = set(deps["profiles"]) - params = set(deps["params"]) - # assert no duplicate dependencies - assert len(data) == len(deps["data"]), err_msg - assert len(axis_limit_data) == len(deps["axis_limit_data"]), err_msg - assert data.isdisjoint(axis_limit_data), err_msg - assert len(profiles) == len(deps["profiles"]), err_msg - assert len(params) == len(deps["params"]), err_msg - # assert correct dependencies are queried - # TODO: conversion from rpz to xyz is taken out of actual function - # registration because of this data["phi"] is not queried in - # the source code but actually needed for the computation. This - # is a temporary fix until we have a better way to automatically - # handle this. - assert queried_deps[p][name]["data"].issubset( - data | axis_limit_data - ), err_msg - errorif( - name not in queried_deps[p], - AssertionError, - "Did you reuse the function name (i.e. def_...) for" - f" '{name}' for some other quantity?", - ) - assert queried_deps[p][name]["profiles"] == profiles, err_msg - assert queried_deps[p][name]["params"] == params, err_msg + for p in data_index: + for name, val in data_index[p].items(): + err_msg = f"Parameterization: {p}. Name: {name}." + deps = val["dependencies"] + data = set(deps["data"]) + axis_limit_data = set(deps["axis_limit_data"]) + profiles = set(deps["profiles"]) + params = set(deps["params"]) + # assert no duplicate dependencies + assert len(data) == len(deps["data"]), err_msg + assert len(axis_limit_data) == len(deps["axis_limit_data"]), err_msg + assert data.isdisjoint(axis_limit_data), err_msg + assert len(profiles) == len(deps["profiles"]), err_msg + assert len(params) == len(deps["params"]), err_msg + errorif( + name not in queried_deps[p], + AssertionError, + "Did you reuse the function name (i.e. def_...) for" + f" '{name}' for some other quantity?", + ) + # assert correct dependencies are queried + if not queried_deps[p][name]["ignore"]: + assert queried_deps[p][name]["data"] == data | axis_limit_data, err_msg + assert queried_deps[p][name]["profiles"] == profiles, err_msg + assert queried_deps[p][name]["params"] == params, err_msg diff --git a/tests/test_equilibrium.py b/tests/test_equilibrium.py index 0b1af2572e..52958a6a9b 100644 --- a/tests/test_equilibrium.py +++ b/tests/test_equilibrium.py @@ -15,13 +15,14 @@ from desc.io import InputReader from desc.objectives import ForceBalance, ObjectiveFunction, get_equilibrium_objective from desc.profiles import PowerSeriesProfile +from desc.utils import errorif from .utils import area_difference, compute_coords @pytest.mark.unit -def test_compute_theta_coords(): - """Test root finding for theta(theta*,lambda(theta)).""" +def test_map_PEST_coordinates(): + """Test root finding for theta(theta_PEST,lambda(theta)).""" eq = get("DSHAPE_CURRENT") with pytest.warns(UserWarning, match="Reducing radial"): eq.change_resolution(3, 3, 0, 6, 6, 0) @@ -34,7 +35,7 @@ def test_compute_theta_coords(): flux_coords = nodes.copy() flux_coords[:, 1] += coords["lambda"] - geom_coords = eq.compute_theta_coords(flux_coords) + geom_coords = eq.map_coordinates(flux_coords, inbasis=("rho", "theta_PEST", "zeta")) geom_coords = np.array(geom_coords) # catch difference between 0 and 2*pi @@ -49,17 +50,33 @@ def test_map_coordinates(): """Test root finding for (rho,theta,zeta) for common use cases.""" # finding coordinates along a single field line eq = get("NCSX") - with pytest.warns(UserWarning, match="Reducing radial"): - eq.change_resolution(3, 3, 3, 6, 6, 6) n = 100 coords = np.array([np.ones(n), np.zeros(n), np.linspace(0, 10 * np.pi, n)]).T - out = eq.map_coordinates( + grid = LinearGrid(rho=1, M=eq.M_grid, N=eq.N_grid, sym=eq.sym, NFP=eq.NFP) + iota = grid.compress(eq.compute("iota", grid=grid)["iota"]) + iota = np.broadcast_to(iota, shape=(n,)) + + tol = 1e-5 + out_1 = eq.map_coordinates( + coords, inbasis=["rho", "alpha", "zeta"], iota=iota, tol=tol + ) + assert np.isfinite(out_1).all() + out_2 = eq.map_coordinates( coords, - ["rho", "alpha", "zeta"], - ["rho", "theta", "zeta"], - period=(np.inf, 2 * np.pi, 10 * np.pi), + inbasis=["rho", "alpha", "zeta"], + period=(np.inf, 2 * np.pi, np.inf), + tol=tol, + ) + assert np.isfinite(out_2).all() + diff = out_1 - out_2 + errorif( + not np.all( + np.isclose(diff, 0, atol=2 * tol) + | np.isclose(np.abs(diff), 2 * np.pi, atol=2 * tol) + ), + AssertionError, + f"diff: {diff}", ) - assert not np.any(np.isnan(out)) eq = get("DSHAPE") @@ -136,7 +153,9 @@ def foo(params, in_coords): @jax.jit def bar(L_lmn): - geom_coords = eq.compute_theta_coords(flux_coords, L_lmn) + geom_coords = eq.map_coordinates( + flux_coords, inbasis=("rho", "theta_PEST", "zeta") + ) return geom_coords J1 = jax.jit(jax.jacfwd(bar))(eq.params_dict["L_lmn"]) diff --git a/tests/test_examples.py b/tests/test_examples.py index a8bff8a538..6bae9dd16b 100644 --- a/tests/test_examples.py +++ b/tests/test_examples.py @@ -440,7 +440,6 @@ def test_NAE_QSC_solve(): orig_Rax_val = eq.axis.R_n orig_Zax_val = eq.axis.Z_n - eq_fit = eq.copy() eq_lambda_fixed_0th_order = eq.copy() eq_lambda_fixed_1st_order = eq.copy() @@ -472,7 +471,6 @@ def test_NAE_QSC_solve(): xtol=1e-6, constraints=constraints, ) - grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) grid_axis = LinearGrid(rho=0.0, theta=0.0, N=eq.N_grid, NFP=eq.NFP) # Make sure axis is same for eqq, string in zip( @@ -482,24 +480,15 @@ def test_NAE_QSC_solve(): np.testing.assert_array_almost_equal(orig_Rax_val, eqq.axis.R_n, err_msg=string) np.testing.assert_array_almost_equal(orig_Zax_val, eqq.axis.Z_n, err_msg=string) - # Make sure surfaces of solved equilibrium are similar near axis as QSC - rho_err, theta_err = area_difference_desc(eqq, eq_fit) + # Make sure iota of solved equilibrium is same on-axis as QSC - np.testing.assert_allclose(rho_err[:, 0:-6], 0, atol=1.5e-2, err_msg=string) - np.testing.assert_allclose(theta_err[:, 0:-6], 0, atol=1e-3, err_msg=string) - - # Make sure iota of solved equilibrium is same near axis as QSC - - iota = grid.compress(eqq.compute("iota", grid=grid)["iota"]) + iota = eqq.compute("iota", grid=LinearGrid(rho=0.0))["iota"] np.testing.assert_allclose(iota[0], qsc.iota, atol=1e-5, err_msg=string) - np.testing.assert_allclose(iota[1:10], qsc.iota, atol=1e-3, err_msg=string) - # check lambda to match near axis - # Evaluate lambda near the axis + # check lambda to match on-axis data_nae = eqq.compute(["lambda", "|B|"], grid=grid_axis) lam_nae = data_nae["lambda"] - # Reshape to form grids on theta and phi phi = np.squeeze(grid_axis.nodes[:, 2]) np.testing.assert_allclose( @@ -526,7 +515,6 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): orig_Rax_val = eq.axis.R_n orig_Zax_val = eq.axis.Z_n - eq_fit = eq.copy() eq_lambda_fixed_0th_order = eq.copy() eq_lambda_fixed_1st_order = eq.copy() @@ -558,7 +546,6 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): xtol=1e-6, constraints=constraints, ) - grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) grid_axis = LinearGrid(rho=0.0, theta=0.0, N=eq.N_grid, NFP=eq.NFP) # Make sure axis is same for eqq, string in zip( @@ -568,18 +555,11 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): np.testing.assert_array_almost_equal(orig_Rax_val, eqq.axis.R_n, err_msg=string) np.testing.assert_array_almost_equal(orig_Zax_val, eqq.axis.Z_n, err_msg=string) - # Make sure surfaces of solved equilibrium are similar near axis as QSC - rho_err, theta_err = area_difference_desc(eqq, eq_fit) - - np.testing.assert_allclose(rho_err[:, 0:-6], 0, atol=1.5e-2, err_msg=string) - np.testing.assert_allclose(theta_err[:, 0:-6], 0, atol=1e-3, err_msg=string) + # Make sure iota of solved equilibrium is same on axis as QSC - # Make sure iota of solved equilibrium is same near axis as QSC - - iota = grid.compress(eqq.compute("iota", grid=grid)["iota"]) + iota = eqq.compute("iota", grid=LinearGrid(rho=0.0))["iota"] np.testing.assert_allclose(iota[0], qsc.iota, atol=1e-5, err_msg=string) - np.testing.assert_allclose(iota[1:10], qsc.iota, rtol=5e-3, err_msg=string) ### check lambda to match on axis # Evaluate lambda on the axis @@ -603,18 +583,16 @@ def test_NAE_QSC_solve_near_axis_based_off_eq(): def test_NAE_QIC_solve(): """Test O(rho) NAE QIC constraints solve.""" # get Qic example - qic = Qic.from_paper("QI NFP2 r2", nphi=301, order="r1") + qic = Qic.from_paper("QI NFP2 r2", nphi=199, order="r1") qic.lasym = False # don't need to consider stellarator asym for order 1 constraints ntheta = 75 r = 0.01 - N = 11 + N = 9 eq = Equilibrium.from_near_axis(qic, r=r, L=7, M=7, N=N, ntheta=ntheta) orig_Rax_val = eq.axis.R_n orig_Zax_val = eq.axis.Z_n - eq_fit = eq.copy() - # this has all the constraints we need, cs = get_NAE_constraints(eq, qic, order=1) @@ -629,75 +607,24 @@ def test_NAE_QIC_solve(): np.testing.assert_array_almost_equal(orig_Rax_val, eq.axis.R_n) np.testing.assert_array_almost_equal(orig_Zax_val, eq.axis.Z_n) - # Make sure surfaces of solved equilibrium are similar near axis as QIC - rho_err, theta_err = area_difference_desc(eq, eq_fit) - - np.testing.assert_allclose(rho_err[:, 0:3], 0, atol=5e-2) - # theta error isn't really an indicator of near axis behavior - # since it's computed over the full radius, but just indicates that - # eq is similar to eq_fit - np.testing.assert_allclose(theta_err, 0, atol=5e-2) - # Make sure iota of solved equilibrium is same near axis as QIC - grid = LinearGrid(L=10, M=20, N=20, NFP=eq.NFP, sym=True, axis=False) - iota = grid.compress(eq.compute("iota", grid=grid)["iota"]) + iota = eq.compute("iota", grid=LinearGrid(rho=0.0))["iota"] np.testing.assert_allclose(iota[0], qic.iota, rtol=1e-5) - np.testing.assert_allclose(iota[1:10], qic.iota, rtol=1e-3) - - # check lambda to match near axis - grid_2d_05 = LinearGrid(rho=np.array(1e-6), M=50, N=50, NFP=eq.NFP, endpoint=True) - # Evaluate lambda near the axis - data_nae = eq.compute("lambda", grid=grid_2d_05) - lam_nae = data_nae["lambda"] + grid_axis = LinearGrid(rho=0.0, theta=0.0, zeta=qic.phi, NFP=eq.NFP) + phi = grid_axis.nodes[:, 2].squeeze() - # Reshape to form grids on theta and phi - zeta = ( - grid_2d_05.nodes[:, 2] - .reshape( - (grid_2d_05.num_theta, grid_2d_05.num_rho, grid_2d_05.num_zeta), order="F" - ) - .squeeze() - ) - - lam_nae = lam_nae.reshape( - (grid_2d_05.num_theta, grid_2d_05.num_rho, grid_2d_05.num_zeta), order="F" - ) - - phi = np.squeeze(zeta[0, :]) - lam_nae = np.squeeze(lam_nae[:, 0, :]) + # check lambda to match on-axis + lam_nae = eq.compute("lambda", grid=grid_axis)["lambda"] - lam_av_nae = np.mean(lam_nae, axis=0) np.testing.assert_allclose( - lam_av_nae, -qic.iota * qic.nu_spline(phi), atol=1e-4, rtol=1e-2 + lam_nae, -qic.iota * qic.nu_spline(phi), atol=1e-4, rtol=1e-2 ) # check |B| on axis - - grid = LinearGrid(M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, rho=np.array(1e-6)) - # Evaluate B modes near the axis - data_nae = eq.compute(["|B|_mn", "B modes"], grid=grid) - modes = data_nae["B modes"] - B_mn_nae = data_nae["|B|_mn"] - # Evaluate B on an angular grid - theta = np.linspace(0, 2 * np.pi, 150) - phi = np.linspace(0, 2 * np.pi, qic.nphi) - th, ph = np.meshgrid(theta, phi) - B_nae = np.zeros((qic.nphi, 150)) - - for i, (l, m, n) in enumerate(modes): - if m >= 0 and n >= 0: - B_nae += B_mn_nae[i] * np.cos(m * th) * np.cos(n * ph) - elif m >= 0 > n: - B_nae += -B_mn_nae[i] * np.cos(m * th) * np.sin(n * ph) - elif m < 0 <= n: - B_nae += -B_mn_nae[i] * np.sin(m * th) * np.cos(n * ph) - elif m < 0 and n < 0: - B_nae += B_mn_nae[i] * np.sin(m * th) * np.sin(n * ph) - # Eliminate the poloidal angle to focus on the toroidal behavior - B_av_nae = np.mean(B_nae, axis=1) - np.testing.assert_allclose(B_av_nae, np.ones(np.size(phi)) * qic.B0, atol=2e-2) + B_nae = eq.compute(["|B|"], grid=grid_axis)["|B|"] + np.testing.assert_allclose(B_nae, qic.B0, atol=1e-3) @pytest.mark.unit @@ -1058,7 +985,7 @@ def test_non_eq_optimization(): use_softmin=True, surface_grid=grid, plasma_grid=grid, - alpha=5000, + softmin_alpha=5000, ) objective = ObjectiveFunction((obj,)) optimizer = Optimizer("lsq-auglag") @@ -1630,3 +1557,102 @@ def circle_constraint(params): abs(surf_opt.Z_lmn[surf_opt.Z_basis.get_idx(M=-1, N=0)]) + offset, rtol=2e-2, ) + + +@pytest.mark.slow +@pytest.mark.regression +@pytest.mark.optimize +def test_signed_PlasmaVesselDistance(): + """Tests that signed distance works with surface optimization.""" + eq = get("HELIOTRON") + eq.change_resolution(M=2, N=2) + + surf = eq.surface.copy() + surf.change_resolution(M=1, N=1) + + target_dist = -0.25 + + grid = LinearGrid(M=10, N=4, NFP=eq.NFP) + obj = PlasmaVesselDistance( + surface=surf, + eq=eq, + target=target_dist, + surface_grid=grid, + plasma_grid=grid, + use_signed_distance=True, + eq_fixed=True, + ) + objective = ObjectiveFunction((obj,)) + + optimizer = Optimizer("lsq-exact") + (surf,), _ = optimizer.optimize( + (surf,), objective, verbose=3, maxiter=60, ftol=1e-8, xtol=1e-9 + ) + + np.testing.assert_allclose( + obj.compute(*obj.xs(surf)), target_dist, atol=1e-2, err_msg="Using hardmin" + ) + + # with softmin + surf = eq.surface.copy() + surf.change_resolution(M=1, N=1) + obj = PlasmaVesselDistance( + surface=surf, + eq=eq, + target=target_dist, + surface_grid=grid, + plasma_grid=grid, + use_signed_distance=True, + use_softmin=True, + softmin_alpha=100, + eq_fixed=True, + ) + objective = ObjectiveFunction((obj,)) + + optimizer = Optimizer("lsq-exact") + (surf,), _ = optimizer.optimize( + (surf,), + objective, + verbose=3, + maxiter=60, + ftol=1e-8, + xtol=1e-9, + ) + + np.testing.assert_allclose( + obj.compute(*obj.xs(surf)), target_dist, atol=1e-2, err_msg="Using softmin" + ) + + # with changing eq + eq = Equilibrium(M=1, N=1) + surf = eq.surface.copy() + surf.change_resolution(M=1, N=1) + grid = LinearGrid(M=20, N=8, NFP=eq.NFP) + + obj = PlasmaVesselDistance( + surface=surf, + eq=eq, + target=target_dist, + surface_grid=grid, + plasma_grid=grid, + use_signed_distance=True, + ) + objective = ObjectiveFunction(obj) + + optimizer = Optimizer("lsq-exact") + (eq, surf), _ = optimizer.optimize( + (eq, surf), + objective, + constraints=(FixParameters(surf),), + verbose=3, + maxiter=60, + ftol=1e-8, + xtol=1e-9, + ) + + np.testing.assert_allclose( + obj.compute(*obj.xs(eq, surf)), + target_dist, + atol=1e-2, + err_msg="allowing eq to change", + ) diff --git a/tests/test_linear_objectives.py b/tests/test_linear_objectives.py index 167765cd83..f67f934a55 100644 --- a/tests/test_linear_objectives.py +++ b/tests/test_linear_objectives.py @@ -451,7 +451,7 @@ def test_correct_indexing_passed_modes(): constraint = ObjectiveFunction(constraints, use_jit=False) constraint.build() - xp, A, b, Z, unfixed_idx, project, recover = factorize_linear_constraints( + xp, A, b, Z, D, unfixed_idx, project, recover = factorize_linear_constraints( objective, constraint ) @@ -461,8 +461,8 @@ def test_correct_indexing_passed_modes(): atol = 2e-15 np.testing.assert_allclose(x1, x2, atol=atol) np.testing.assert_allclose(A @ xp[unfixed_idx], b, atol=atol) - np.testing.assert_allclose(A @ x1[unfixed_idx], b, atol=atol) - np.testing.assert_allclose(A @ x2[unfixed_idx], b, atol=atol) + np.testing.assert_allclose(A @ (x1[unfixed_idx] / D[unfixed_idx]), b, atol=atol) + np.testing.assert_allclose(A @ (x2[unfixed_idx] / D[unfixed_idx]), b, atol=atol) np.testing.assert_allclose(A @ Z, 0, atol=atol) @@ -514,7 +514,7 @@ def test_correct_indexing_passed_modes_and_passed_target(): constraint = ObjectiveFunction(constraints, use_jit=False) constraint.build() - xp, A, b, Z, unfixed_idx, project, recover = factorize_linear_constraints( + xp, A, b, Z, D, unfixed_idx, project, recover = factorize_linear_constraints( objective, constraint ) @@ -524,8 +524,8 @@ def test_correct_indexing_passed_modes_and_passed_target(): atol = 2e-15 np.testing.assert_allclose(x1, x2, atol=atol) np.testing.assert_allclose(A @ xp[unfixed_idx], b, atol=atol) - np.testing.assert_allclose(A @ x1[unfixed_idx], b, atol=atol) - np.testing.assert_allclose(A @ x2[unfixed_idx], b, atol=atol) + np.testing.assert_allclose(A @ (x1[unfixed_idx] / D[unfixed_idx]), b, atol=atol) + np.testing.assert_allclose(A @ (x2[unfixed_idx] / D[unfixed_idx]), b, atol=atol) np.testing.assert_allclose(A @ Z, 0, atol=atol) @@ -574,7 +574,7 @@ def test_correct_indexing_passed_modes_axis(): constraint = ObjectiveFunction(constraints, use_jit=False) constraint.build() - xp, A, b, Z, unfixed_idx, project, recover = factorize_linear_constraints( + xp, A, b, Z, D, unfixed_idx, project, recover = factorize_linear_constraints( objective, constraint ) @@ -584,8 +584,8 @@ def test_correct_indexing_passed_modes_axis(): atol = 2e-15 np.testing.assert_allclose(x1, x2, atol=atol) np.testing.assert_allclose(A @ xp[unfixed_idx], b, atol=atol) - np.testing.assert_allclose(A @ x1[unfixed_idx], b, atol=atol) - np.testing.assert_allclose(A @ x2[unfixed_idx], b, atol=atol) + np.testing.assert_allclose(A @ (x1[unfixed_idx] / D[unfixed_idx]), b, atol=atol) + np.testing.assert_allclose(A @ (x2[unfixed_idx] / D[unfixed_idx]), b, atol=atol) np.testing.assert_allclose(A @ Z, 0, atol=atol) @@ -703,7 +703,7 @@ def test_correct_indexing_passed_modes_and_passed_target_axis(): constraint = ObjectiveFunction(constraints, use_jit=False) constraint.build() - xp, A, b, Z, unfixed_idx, project, recover = factorize_linear_constraints( + xp, A, b, Z, D, unfixed_idx, project, recover = factorize_linear_constraints( objective, constraint ) @@ -713,8 +713,8 @@ def test_correct_indexing_passed_modes_and_passed_target_axis(): atol = 2e-15 np.testing.assert_allclose(x1, x2, atol=atol) np.testing.assert_allclose(A @ xp[unfixed_idx], b, atol=atol) - np.testing.assert_allclose(A @ x1[unfixed_idx], b, atol=atol) - np.testing.assert_allclose(A @ x2[unfixed_idx], b, atol=atol) + np.testing.assert_allclose(A @ (x1[unfixed_idx] / D[unfixed_idx]), b, atol=atol) + np.testing.assert_allclose(A @ (x2[unfixed_idx] / D[unfixed_idx]), b, atol=atol) np.testing.assert_allclose(A @ Z, 0, atol=atol) diff --git a/tests/test_objective_funs.py b/tests/test_objective_funs.py index 9fcc12634b..d6f67a5ffc 100644 --- a/tests/test_objective_funs.py +++ b/tests/test_objective_funs.py @@ -640,6 +640,7 @@ def test_plasma_vessel_distance(self): ) obj.build() d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert d.size == obj.dim_f assert abs(d.min() - (a_s - a_p)) < 1e-14 assert abs(d.max() - (a_s - a_p)) < surf_grid.spacing[0, 1] * a_p @@ -679,6 +680,7 @@ def test_plasma_vessel_distance(self): ) obj.build() d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert d.size == obj.dim_f assert np.all(np.abs(d) < a_s - a_p) # for large enough alpha, should be same as actual min @@ -688,7 +690,7 @@ def test_plasma_vessel_distance(self): surface_grid=surf_grid, surface=surface, use_softmin=True, - alpha=100, + softmin_alpha=100, ) obj.build() d = obj.compute_unscaled(*obj.xs(eq, surface)) @@ -1137,6 +1139,89 @@ def test(eq, field, correct_value, rtol=1e-14, grid=None): test(eq, field, psi_from_field) + @pytest.mark.unit + def test_signed_plasma_vessel_distance(self): + """Test calculation of signed distance from plasma to vessel.""" + R0 = 10.0 + a_p = 1.0 + a_s = 2.0 + # default eq has R0=10, a=1 + eq = Equilibrium(M=3, N=2) + # surface with same R0, a=2, so true d=1 for all pts + surface = FourierRZToroidalSurface( + R_lmn=[R0, a_s], Z_lmn=[-a_s], modes_R=[[0, 0], [1, 0]], modes_Z=[[-1, 0]] + ) + grid = LinearGrid(M=5, N=6) + obj = PlasmaVesselDistance( + eq=eq, + surface_grid=grid, + plasma_grid=grid, + surface=surface, + use_signed_distance=True, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert obj.dim_f == d.size + np.testing.assert_allclose(d, a_s - a_p) + + # ensure that it works (dimension-wise) when compute_scaled is called + _ = obj.compute_scaled(*obj.xs(eq, surface)) + + # For plasma outside surface, should get signed distance + a_s = 0.5 * a_p + surface = FourierRZToroidalSurface( + R_lmn=[R0, a_s], + Z_lmn=[-a_s], + modes_R=[[0, 0], [1, 0]], + modes_Z=[[-1, 0]], + ) + grid = LinearGrid(M=5, N=6) + obj = PlasmaVesselDistance( + eq=eq, + surface_grid=grid, + plasma_grid=grid, + surface=surface, + use_signed_distance=True, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert obj.dim_f == d.size + np.testing.assert_allclose(d, a_s - a_p) + + # ensure it works with different sized grids (poloidal resolution different) + grid = LinearGrid(M=5, N=6) + obj = PlasmaVesselDistance( + eq=eq, + surface_grid=grid, + plasma_grid=LinearGrid(M=10, N=6), + surface=surface, + use_signed_distance=True, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert obj.dim_f == d.size + assert abs(d.max() - (-a_s)) < 1e-14 + assert abs(d.min() - (-a_s)) < grid.spacing[0, 1] * a_s + + # ensure it works with different sized grids (poloidal resolution different) + # and using softmin (with deprecated name alpha) + grid = LinearGrid(M=5, N=6) + with pytest.raises(FutureWarning): + obj = PlasmaVesselDistance( + eq=eq, + surface_grid=grid, + plasma_grid=LinearGrid(M=10, N=6), + surface=surface, + use_signed_distance=True, + use_softmin=True, + alpha=4000, + ) + obj.build() + d = obj.compute_unscaled(*obj.xs(eq, surface)) + assert obj.dim_f == d.size + assert abs(d.max() - (-a_s)) < 1e-14 + assert abs(d.min() - (-a_s)) < grid.spacing[0, 1] * a_s + @pytest.mark.regression def test_derivative_modes(): diff --git a/tests/test_optimizer.py b/tests/test_optimizer.py index 7c5f810d98..71cc515a58 100644 --- a/tests/test_optimizer.py +++ b/tests/test_optimizer.py @@ -5,6 +5,7 @@ import numpy as np import pytest from numpy.random import default_rng +from scipy.constants import mu_0 from scipy.optimize import ( BFGS, NonlinearConstraint, @@ -32,6 +33,7 @@ FixParameters, FixPressure, FixPsi, + FixSumCoilCurrent, ForceBalance, GenericObjective, MagneticWell, @@ -1350,3 +1352,30 @@ def test_quad_flux_with_surface_current_field(): (field_modular_opt,), result = opt.optimize( field, objective=obj, constraints=constraints, maxiter=1, copy=True ) + + +@pytest.mark.unit +def test_optimize_coil_currents(DummyCoilSet): + """Tests optimization takes step sizes proportional to variable scales.""" + eq = desc.examples.get("precise_QH") + coils = load(load_from=str(DummyCoilSet["output_path_sym"]), file_format="hdf5") + grid = LinearGrid(rho=1.0, M=eq.M_grid, N=eq.N_grid, NFP=eq.NFP, sym=eq.sym) + current = 2 * np.pi * eq.compute("G", grid=grid)["G"][0] / mu_0 + for coil in coils: + coil.current = current / coils.num_coils + + objective = ObjectiveFunction(QuadraticFlux(eq=eq, field=coils, vacuum=True)) + constraints = FixSumCoilCurrent(coils) + optimizer = Optimizer("lsq-exact") + [coils_opt], _ = optimizer.optimize( + things=coils, + objective=objective, + constraints=constraints, + verbose=2, + copy=True, + ) + # check that optimized coil currents changed by more than 15% from initial values + np.testing.assert_array_less( + np.asarray(coils.current) * 0.15, + np.abs(np.asarray(coils_opt.current) - np.asarray(coils.current)), + ) diff --git a/tests/test_plotting.py b/tests/test_plotting.py index 843e125319..1351b46c93 100644 --- a/tests/test_plotting.py +++ b/tests/test_plotting.py @@ -801,6 +801,8 @@ def test_plot_coils(): coil.rotate(angle=np.pi / N) coils = CoilSet.linspaced_angular(coil, I, [0, 0, 1], np.pi / NFP, N // NFP // 2) coils2 = MixedCoilSet.from_symmetry(coils, NFP, True) + with pytest.raises(ValueError, match="Expected `coils`"): + plot_coils("not a coil") fig, data = plot_coils(coils2, return_data=True) def flatten_coils(coilset): @@ -839,7 +841,7 @@ def test_plot_b_mag(): rhoa = rho * np.ones_like(zeta) c = np.vstack([rhoa, thetas, zeta]).T - coords = eq.compute_theta_coords(c) + coords = eq.map_coordinates(c, inbasis=("rho", "theta_PEST", "zeta")) grid = Grid(coords) # compute |B| normalized in the usual flux tube way diff --git a/tests/utils.py b/tests/utils.py index 85960f804d..8799bfdbf2 100644 --- a/tests/utils.py +++ b/tests/utils.py @@ -36,7 +36,9 @@ def compute_coords(equil, Nr=10, Nt=8, Nz=None): # angle in PEST-like flux coordinates # find theta angles corresponding to desired theta* angles - v_grid = Grid(equil.compute_theta_coords(t_grid.nodes)) + v_grid = Grid( + equil.map_coordinates(t_grid.nodes, inbasis=("rho", "theta_PEST", "zeta")) + ) r_coords = equil.compute(["R", "Z"], grid=r_grid) v_coords = equil.compute(["R", "Z"], grid=v_grid)