Merge branch 'bounce' into ku/fourier_bounce

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
 -------
 

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)

desc/compute/_basis_vectors.py

@@ -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
 
 

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",
@@ -2711,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"]
@@ -2785,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"]
@@ -2841,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"]
@@ -2890,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",
@@ -3113,75 +3081,6 @@ def _theta_PEST_zz(params, transforms, profiles, data, **kwargs):
     return data
 
 
-@register_compute_fun(
-    name="theta_r",
-    label="\\partial_{\\rho} \\theta",
-    units="rad",
-    units_long="radians",
-    description="Poloidal angular coordinate (geometric, not magnetic), "
-    "derivative wrt radial coordinate",
-    dim=1,
-    params=[],
-    transforms={},
-    profiles=[],
-    coordinates="t",
-    data=["0"],
-    parameterization=[
-        "desc.equilibrium.equilibrium.Equilibrium",
-        "desc.geometry.core.Surface",
-    ],
-)
-def _theta_r(params, transforms, profiles, data, **kwargs):
-    data["theta_r"] = data["0"]
-    return data
-
-
-@register_compute_fun(
-    name="theta_t",
-    label="\\partial_{\\theta} \\theta",
-    units="rad",
-    units_long="radians",
-    description="Poloidal angular coordinate (geometric, not magnetic), "
-    "derivative wrt poloidal coordinate",
-    dim=1,
-    params=[],
-    transforms={},
-    profiles=[],
-    coordinates="t",
-    data=["0"],
-    parameterization=[
-        "desc.equilibrium.equilibrium.Equilibrium",
-        "desc.geometry.core.Surface",
-    ],
-)
-def _theta_t(params, transforms, profiles, data, **kwargs):
-    data["theta_t"] = jnp.ones_like(data["0"])
-    return data
-
-
-@register_compute_fun(
-    name="theta_z",
-    label="\\partial_{\\zeta} \\theta",
-    units="rad",
-    units_long="radians",
-    description="Poloidal angular coordinate (geometric, not magnetic), "
-    "derivative wrt toroidal coordinate",
-    dim=1,
-    params=[],
-    transforms={},
-    profiles=[],
-    coordinates="t",
-    data=["0"],
-    parameterization=[
-        "desc.equilibrium.equilibrium.Equilibrium",
-        "desc.geometry.core.Surface",
-    ],
-)
-def _theta_z(params, transforms, profiles, data, **kwargs):
-    data["theta_z"] = data["0"]
-    return data
-
-
 @register_compute_fun(
     name="zeta",
     label="\\zeta",
@@ -3202,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

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": []},

desc/compute/_field.py

@@ -2935,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",

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