Update equations to match physics and boozer fit for general toroidal…

… angle (#1177)

Resolves #1176.

- [x] Update equations to use poloidal and toroidal angles rather than a
particular choice of stream function. No changes were made in places
where a particular choice of stream function is necessary to simplify
the computation.
- [x] Genearlizes boozer fit for arbitrary toroidal angle

desc/compute/_core.py

@@ -2764,6 +2764,25 @@ def _phi_rr(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="phi_rrz",
+    label="\\partial_{\\rho \\rho \\zeta} \\phi",
+    units="rad",
+    units_long="radians",
+    description="Toroidal angle in lab frame, second derivative wrt radial coordinate "
+    "and first wrt DESC toroidal coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["omega_rrz"],
+)
+def _phi_rrz(params, transforms, profiles, data, **kwargs):
+    data["phi_rrz"] = data["omega_rrz"]
+    return data
+
+
 @register_compute_fun(
     name="phi_rt",
     label="\\partial_{\\rho \\theta} \\phi",
@@ -2783,6 +2802,25 @@ def _phi_rt(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="phi_rtz",
+    label="\\partial_{\\rho \\theta \\zeta} \\phi",
+    units="rad",
+    units_long="radians",
+    description="Toroidal angle in lab frame, third derivative wrt radial, "
+    "poloidal, and toroidal coordinates",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["omega_rtz"],
+)
+def _phi_rtz(params, transforms, profiles, data, **kwargs):
+    data["phi_rtz"] = data["omega_rtz"]
+    return data
+
+
 @register_compute_fun(
     name="phi_rz",
     label="\\partial_{\\rho \\zeta} \\phi",
@@ -2802,6 +2840,25 @@ def _phi_rz(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="phi_rzz",
+    label="\\partial_{\\rho \\zeta \\zeta} \\phi",
+    units="rad",
+    units_long="radians",
+    description="Toroidal angle in lab frame, first derivative wrt radial and "
+    "second derivative wrt DESC toroidal coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["omega_rzz"],
+)
+def _phi_rzz(params, transforms, profiles, data, **kwargs):
+    data["phi_rzz"] = data["omega_rzz"]
+    return data
+
+
 @register_compute_fun(
     name="phi_t",
     label="\\partial_{\\theta} \\phi",
@@ -2843,12 +2900,31 @@ def _phi_tt(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="phi_ttz",
+    label="\\partial_{\\theta \\theta \\zeta} \\phi",
+    units="rad",
+    units_long="radians",
+    description="Toroidal angle in lab frame, second derivative wrt poloidal "
+    "coordinate and first derivative wrt toroidal coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["omega_ttz"],
+)
+def _phi_ttz(params, transforms, profiles, data, **kwargs):
+    data["phi_ttz"] = data["omega_ttz"]
+    return data
+
+
 @register_compute_fun(
     name="phi_tz",
     label="\\partial_{\\theta \\zeta} \\phi",
     units="rad",
     units_long="radians",
-    description="Toroidal angle in lab frame, second derivative wrt poloidal and "
+    description="Toroidal angle in lab frame, derivative wrt poloidal and "
     "toroidal coordinate",
     dim=1,
     params=[],
@@ -2862,6 +2938,25 @@ def _phi_tz(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="phi_tzz",
+    label="\\partial_{\\theta \\zeta \\zeta} \\phi",
+    units="rad",
+    units_long="radians",
+    description="Toroidal angle in lab frame, derivative wrt poloidal coordinate and "
+    "second derivative wrt toroidal coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["omega_tzz"],
+)
+def _phi_tzz(params, transforms, profiles, data, **kwargs):
+    data["phi_tzz"] = data["omega_tzz"]
+    return data
+
+
 @register_compute_fun(
     name="phi_z",
     label="\\partial_{\\zeta} \\phi",
@@ -2903,6 +2998,25 @@ def _phi_zz(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="phi_zzz",
+    label="\\partial_{\\zeta \\zeta \\zeta} \\phi",
+    units="rad",
+    units_long="radians",
+    description="Toroidal angle in lab frame, third derivative wrt toroidal "
+    "coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["omega_zzz"],
+)
+def _phi_zzz(params, transforms, profiles, data, **kwargs):
+    data["phi_zzz"] = data["omega_zzz"]
+    return data
+
+
 @register_compute_fun(
     name="rho",
     label="\\rho",
@@ -2986,6 +3100,83 @@ def _theta_PEST_r(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="theta_PEST_rt",
+    label="\\partial_{\\rho \\theta} \\vartheta",
+    units="rad",
+    units_long="radians",
+    description="PEST straight field line poloidal angular coordinate, derivative wrt "
+    "radial and DESC poloidal coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["lambda_rt"],
+)
+def _theta_PEST_rt(params, transforms, profiles, data, **kwargs):
+    data["theta_PEST_rt"] = data["lambda_rt"]
+    return data
+
+
+@register_compute_fun(
+    name="theta_PEST_rrt",
+    label="\\partial_{\\rho \\rho \\theta} \\vartheta",
+    units="rad",
+    units_long="radians",
+    description="PEST straight field line poloidal angular coordinate, second "
+    "derivative wrt radial coordinate and first derivative wrt DESC poloidal "
+    "coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["lambda_rrt"],
+)
+def _theta_PEST_rrt(params, transforms, profiles, data, **kwargs):
+    data["theta_PEST_rrt"] = data["lambda_rrt"]
+    return data
+
+
+@register_compute_fun(
+    name="theta_PEST_rtz",
+    label="\\partial_{\\rho \\theta \\zeta} \\vartheta",
+    units="rad",
+    units_long="radians",
+    description="PEST straight field line poloidal angular coordinate, derivative wrt "
+    "radial and DESC poloidal and toroidal coordinates",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["lambda_rtz"],
+)
+def _theta_PEST_rtz(params, transforms, profiles, data, **kwargs):
+    data["theta_PEST_rtz"] = data["lambda_rtz"]
+    return data
+
+
+@register_compute_fun(
+    name="theta_PEST_rtt",
+    label="\\partial_{\\rho \\theta \\theta} \\vartheta",
+    units="rad",
+    units_long="radians",
+    description="PEST straight field line poloidal angular coordinate, derivative wrt "
+    "radial coordinate once and DESC poloidal coordinate twice",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["lambda_rtt"],
+)
+def _theta_PEST_rtt(params, transforms, profiles, data, **kwargs):
+    data["theta_PEST_rtt"] = data["lambda_rtt"]
+    return data
+
+
 @register_compute_fun(
     name="theta_PEST_t",
     label="\\partial_{\\theta} \\vartheta",
@@ -3024,6 +3215,25 @@ def _theta_PEST_tt(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="theta_PEST_ttt",
+    label="\\partial_{\\theta \\theta \\theta} \\vartheta",
+    units="rad",
+    units_long="radians",
+    description="PEST straight field line poloidal angular coordinate, third "
+    "derivative wrt poloidal coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["lambda_ttt"],
+)
+def _theta_PEST_ttt(params, transforms, profiles, data, **kwargs):
+    data["theta_PEST_ttt"] = data["lambda_ttt"]
+    return data
+
+
 @register_compute_fun(
     name="theta_PEST_tz",
     label="\\partial_{\\theta \\zeta} \\vartheta",
@@ -3043,6 +3253,25 @@ def _theta_PEST_tz(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="theta_PEST_tzz",
+    label="\\partial_{\\theta \\zeta \\zeta} \\vartheta",
+    units="rad",
+    units_long="radians",
+    description="PEST straight field line poloidal angular coordinate, derivative wrt "
+    "poloidal coordinate once and toroidal coordinate twice",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["lambda_tzz"],
+)
+def _theta_PEST_tzz(params, transforms, profiles, data, **kwargs):
+    data["theta_PEST_tzz"] = data["lambda_tzz"]
+    return data
+
+
 @register_compute_fun(
     name="theta_PEST_z",
     label="\\partial_{\\zeta} \\vartheta",
@@ -3081,6 +3310,25 @@ def _theta_PEST_zz(params, transforms, profiles, data, **kwargs):
     return data
 
 
+@register_compute_fun(
+    name="theta_PEST_ttz",
+    label="\\partial_{\\theta \\theta \\zeta} \\vartheta",
+    units="rad",
+    units_long="radians",
+    description="PEST straight field line poloidal angular coordinate, second "
+    "derivative wrt poloidal coordinate and derivative wrt toroidal coordinate",
+    dim=1,
+    params=[],
+    transforms={},
+    profiles=[],
+    coordinates="rtz",
+    data=["lambda_ttz"],
+)
+def _theta_PEST_ttz(params, transforms, profiles, data, **kwargs):
+    data["theta_PEST_ttz"] = data["lambda_ttz"]
+    return data
+
+
 @register_compute_fun(
     name="zeta",
     label="\\zeta",