Skip to content

Commit

Permalink
Gamma_c (Nemov and Velasco et. al) Γ_c (#1042)
Browse files Browse the repository at this point in the history 
Resolves #522 . Automatically differentiable. Will be replaced by #1290 soon, so users may want to wait for that.
  • Loading branch information
@unalmis
unalmis authored Dec 13, 2024
2 parents 9ef0b38 + ce98913 commit 4bfa728
Show file tree
Hide file tree
Showing 15 changed files with 1,037 additions and 133 deletions.
107 changes: 85 additions & 22 deletions desc/compute/_basis_vectors.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -34,7 +34,7 @@ def _b(params, transforms, profiles, data, **kwargs):


@register_compute_fun(
name="e^rho", # ∇ρ is the same in (ρ,θ,ζ) and (ρ,α,ζ) coordinates.
name="e^rho", # ∇ρ is the same in any coordinate system.
label="\\mathbf{e}^{\\rho}",
units="m^{-1}",
units_long="inverse meters",
Expand Down Expand Up @@ -927,7 +927,7 @@ def _e_sup_theta_zz(params, transforms, profiles, data, **kwargs):


@register_compute_fun(
name="e^zeta", # ∇ζ is the same in (ρ,θ,ζ) and (ρ,α,ζ) coordinates.
name="e^zeta", # ∇ζ is the same in any coordinate system.
label="\\mathbf{e}^{\\zeta}",
units="m^{-1}",
units_long="inverse meters",
Expand Down Expand Up @@ -2449,10 +2449,10 @@ def _e_sub_theta_over_sqrt_g(params, transforms, profiles, data, **kwargs):
)
def _e_sub_vartheta_rp(params, transforms, profiles, data, **kwargs):
# constant ρ and ϕ
e_vartheta = (
data["e_theta"].T * data["phi_z"] - data["e_zeta"].T * data["phi_t"]
) / (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"])
data["e_theta_PEST"] = e_vartheta.T
data["e_theta_PEST"] = (
(data["e_theta"].T * data["phi_z"] - data["e_zeta"].T * data["phi_t"])
/ (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"])
).T
return data


Expand All @@ -2473,11 +2473,13 @@ def _e_sub_vartheta_rp(params, transforms, profiles, data, **kwargs):
)
def _e_sub_phi_rv(params, transforms, profiles, data, **kwargs):
# constant ρ and ϑ
e_phi = (
data["e_zeta"].T * data["theta_PEST_t"]
- data["e_theta"].T * data["theta_PEST_z"]
) / (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"])
data["e_phi|r,v"] = e_phi.T
data["e_phi|r,v"] = (
(
data["e_zeta"].T * data["theta_PEST_t"]
- data["e_theta"].T * data["theta_PEST_z"]
)
/ (data["theta_PEST_t"] * data["phi_z"] - data["theta_PEST_z"] * data["phi_t"])
).T
return data


Expand All @@ -2499,10 +2501,10 @@ def _e_sub_phi_rv(params, transforms, profiles, data, **kwargs):
def _e_sub_rho_vp(params, transforms, profiles, data, **kwargs):
# constant ϑ and ϕ
data["e_rho|v,p"] = (
data["e_rho"].T
- data["e_vartheta"].T * data["theta_PEST_r"]
- data["e_phi|r,v"].T * data["phi_r"]
).T
data["e_rho"]
- data["e_vartheta"] * data["theta_PEST_r"][:, jnp.newaxis]
- data["e_phi|r,v"] * data["phi_r"][:, jnp.newaxis]
)
return data


Expand Down Expand Up @@ -3206,7 +3208,28 @@ def _e_sub_zeta_zz(params, transforms, profiles, data, **kwargs):
data["Z_zzz"],
]
).T
return data


@register_compute_fun(
name="grad(phi)",
label="\\nabla \\phi",
units="m^{-1}",
units_long="Inverse meters",
description="Gradient of cylindrical toroidal angle ϕ.",
dim=3,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["R", "0"],
parameterization=[
"desc.equilibrium.equilibrium.Equilibrium",
"desc.geometry.surface.FourierRZToroidalSurface",
],
)
def _grad_phi(params, transforms, profiles, data, **kwargs):
data["grad(phi)"] = jnp.column_stack([data["0"], 1 / data["R"], data["0"]])
return data


Expand Down Expand Up @@ -3456,7 +3479,7 @@ def _e_sub_theta_rp(params, transforms, profiles, data, **kwargs):
label="\\mathbf{e}_{\\rho} |_{\\alpha, \\zeta}",
units="m",
units_long="meters",
description="Tangent vector along radial field line label",
description="Covariant radial basis vector in (ρ, α, ζ) Clebsch coordinates.",
dim=3,
params=[],
transforms={},
Expand All @@ -3477,7 +3500,7 @@ def _e_rho_az(params, transforms, profiles, data, **kwargs):
label="\\mathbf{e}_{\\alpha}",
units="m",
units_long="meters",
description="Tangent vector along poloidal field line label",
description="Covariant poloidal basis vector in (ρ, α, ζ) Clebsch coordinates.",
dim=3,
params=[],
transforms={},
Expand All @@ -3496,8 +3519,8 @@ def _e_alpha(params, transforms, profiles, data, **kwargs):
label="\\partial_{\\theta} \\mathbf{e}_{\\alpha}",
units="m",
units_long="meters",
description="Tangent vector along poloidal field line label, derivative wrt"
" DESC poloidal angle",
description="Covariant poloidal basis vector in (ρ, α, ζ) Clebsch coordinates,"
" derivative wrt DESC poloidal angle",
dim=3,
params=[],
transforms={},
Expand All @@ -3518,8 +3541,8 @@ def _e_alpha_t(params, transforms, profiles, data, **kwargs):
label="\\partial_{\\zeta} \\mathbf{e}_{\\alpha}",
units="m",
units_long="meters",
description="Tangent vector along poloidal field line label, "
"derivative wrt DESC toroidal angle",
description="Covariant poloidal basis vector in (ρ, α, ζ) Clebsch coordinates, "
"derivative wrt DESC toroidal angle at fixed ρ,θ.",
dim=3,
params=[],
transforms={},
Expand Down Expand Up @@ -3603,7 +3626,7 @@ def _e_zeta_ra_a(params, transforms, profiles, data, **kwargs):
units="m",
units_long="meters",
description="Tangent vector along (collinear to) field line, "
"derivative wrt DESC toroidal angle",
"derivative wrt DESC toroidal angle at fixed ρ,θ.",
dim=3,
params=[],
transforms={},
Expand Down Expand Up @@ -3679,3 +3702,43 @@ def _d_ell_d_zeta_z(params, transforms, profiles, data, **kwargs):
dot(data["(e_zeta|r,a)_z|r,a"], data["e_zeta|r,a"]) / data["|e_zeta|r,a|"]
)
return data


@register_compute_fun(
name="e_alpha|r,p",
label="\\mathbf{e}_{\\alpha} |_{\\rho, \\phi}",
units="m",
units_long="meters",
description="Covariant poloidal basis vector in (ρ, α, ϕ) Clebsch coordinates.",
dim=3,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["e_theta", "alpha_t", "e_zeta", "alpha_z", "phi_t", "phi_z"],
)
def _e_alpha_rp(params, transforms, profiles, data, **kwargs):
data["e_alpha|r,p"] = (
(data["e_theta"].T * data["phi_z"] - data["e_zeta"].T * data["phi_t"])
/ (data["alpha_t"] * data["phi_z"] - data["alpha_z"] * data["phi_t"])
).T
return data


@register_compute_fun(
name="|e_alpha|r,p|",
label="|\\mathbf{e}_{\\alpha} |_{\\rho, \\phi}|",
units="m",
units_long="meters",
description="Norm of covariant poloidal basis vector in (ρ, α, ϕ) Clebsch "
"coordinates.",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["e_alpha|r,p"],
)
def _e_alpha_rp_norm(params, transforms, profiles, data, **kwargs):
data["|e_alpha|r,p|"] = jnp.linalg.norm(data["e_alpha|r,p"], axis=-1)
return data
19 changes: 19 additions & 0 deletions desc/compute/_core.py
View file
Original file line number Diff line number Diff line change
Expand Up @@ -3153,6 +3153,25 @@ def _theta_PEST_rt(params, transforms, profiles, data, **kwargs):
return data


@register_compute_fun(
name="theta_PEST_rz",
label="\\partial_{\\rho \\zeta} \\vartheta",
units="rad",
units_long="radians",
description="PEST straight field line poloidal angular coordinate, derivative wrt "
"radial and DESC toroidal coordinate",
dim=1,
params=[],
transforms={},
profiles=[],
coordinates="rtz",
data=["lambda_rz"],
)
def _theta_PEST_rz(params, transforms, profiles, data, **kwargs):
data["theta_PEST_rz"] = data["lambda_rz"]
return data


@register_compute_fun(
name="theta_PEST_rrt",
label="\\partial_{\\rho \\rho \\theta} \\vartheta",
Expand Down
Loading

0 comments on commit 4bfa728

Please sign in to comment.