Add functions for converting between cartesian and curvilinear coordi…

…nates.

src/analytic/analytic_equilibrium.jl

@@ -8,22 +8,30 @@ abstract type AnalyticPerturbation <: AnalyticField end
 
 get_functions(::AnalyticField) = (;)
 
-periodicity(x::AbstractVector, ::AnalyticField) = zero(x)
+x¹(::AbstractVector{T}, ::AnalyticField) where {T} = error("x¹() not implemented for ", ET)
+x²(::AbstractVector{T}, ::AnalyticField) where {T} = error("x²() not implemented for ", ET)
+x³(::AbstractVector{T}, ::AnalyticField) where {T} = error("x³() not implemented for ", ET)
 
-from_cartesian(::AbstractVector, ::AnalyticField) = error("from_cartesian() not implemented for ", ET)
-to_cartesian(::AbstractVector, ::AnalyticField) = error("from_cartesian() not implemented for ", ET)
+ξ¹(::AbstractVector{T}, ::AnalyticField) where {T} = error("ξ¹() not implemented for ", ET)
+ξ²(::AbstractVector{T}, ::AnalyticField) where {T} = error("ξ²() not implemented for ", ET)
+ξ³(::AbstractVector{T}, ::AnalyticField) where {T} = error("ξ³() not implemented for ", ET)
 
 J(::AbstractVector, ::AnalyticField) = error("J() not implemented for ", ET)
 
-g₁₁(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
+g₁₁(::AbstractVector{T}, ::AnalyticField) where {T} = one(T)
 g₁₂(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
 g₁₃(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
 g₂₁(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
-g₂₂(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
+g₂₂(::AbstractVector{T}, ::AnalyticField) where {T} = one(T)
 g₂₃(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
 g₃₁(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
 g₃₂(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
-g₃₃(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
+g₃₃(::AbstractVector{T}, ::AnalyticField) where {T} = one(T)
+
+periodicity(x::AbstractVector, ::AnalyticField) = zero(x)
+
+from_cartesian(x::AbstractVector, equ::AnalyticField) = [ξ¹(x,equ), ξ²(x,equ), ξ³(x,equ)]
+to_cartesian(ξ::AbstractVector, equ::AnalyticField) = [x¹(ξ,equ), x²(ξ,equ), x³(ξ,equ)]
 
 A₁(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
 A₂(::AbstractVector{T}, ::AnalyticField) where {T} = zero(T)
@@ -201,7 +209,6 @@ function generate_equilibrium_functions(equ::AnalyticEquilibrium, pert::Analytic
     Evec = [get_vector_component(E¹, ginv, i) for i in 1:3]
     symprint("Evec", Evec, output, 2)
 
-
     # collect all functions to generate code for
     functions = Dict{String,Any}()
     indices   = ["₁", "₂", "₃"]
@@ -211,6 +218,20 @@ function generate_equilibrium_functions(equ::AnalyticEquilibrium, pert::Analytic
         functions[string(f[1])] = f[2](x, equ)
     end
 
+    # cartesian coordinates
+    functions["x¹"] = x¹(x, equ)
+    functions["x²"] = x²(x, equ)
+    functions["x³"] = x³(x, equ)
+
+    # curvilinear coordinates
+    functions["ξ¹"] = ξ¹(x, equ)
+    functions["ξ²"] = ξ²(x, equ)
+    functions["ξ³"] = ξ³(x, equ)
+
+    # coordinate conversion functions
+    functions["from_cartesian"] = from_cartesian(x, equ)
+    functions["to_cartesian"] = to_cartesian(x, equ)
+
     functions["J"] = Jdet
     functions["B"] = Babs
     functions["φ"] = φ⁰
@@ -254,6 +275,18 @@ function generate_equilibrium_functions(equ::AnalyticEquilibrium, pert::Analytic
 end
 
 
+function replace_expr!(e, old, new)
+    for (i,a) in enumerate(e.args)
+        if a==old
+            e.args[i] = new
+        elseif a isa Expr
+            replace_expr!(a, old, new)
+        end
+    end
+    e
+end
+
+
 """
 Generate code for evaluating analytic equilibria.
 """
@@ -297,6 +330,7 @@ function generate_equilibrium_code(equ, pert=ZeroPerturbation(); output=0)
         output ≥ 1 ? println("Generating function ", key) : nothing
 
         f_body = convert(Expr, f_expr)
+        replace_expr!(f_body, :atan2, :atan)
         output ≥ 2 ? println("   ", f_body) : nothing
 
         f_code = quote

src/analytic/axisymmetric_tokamak_cylindrical.jl

@@ -1,3 +1,5 @@
+using SymEngine
+
 @doc raw"""
 Axisymmetric tokamak equilibrium in (R,Z,ϕ) coordinates with covariant
 components of the vector potential given by
@@ -52,6 +54,14 @@ A₁(x::AbstractVector, equ::AxisymmetricTokamakCylindrical) = + equ.B₀ * equ.
 A₂(x::AbstractVector, equ::AxisymmetricTokamakCylindrical) = - equ.B₀ * equ.R₀ * log(R(x,equ) / equ.R₀) / 2
 A₃(x::AbstractVector, equ::AxisymmetricTokamakCylindrical) = - equ.B₀ * r²(x,equ) / equ.q₀ / 2
 
+x¹(ξ::AbstractVector, equ::AxisymmetricTokamakCylindrical) = X(ξ,equ)
+x²(ξ::AbstractVector, equ::AxisymmetricTokamakCylindrical) = Y(ξ,equ)
+x³(ξ::AbstractVector, equ::AxisymmetricTokamakCylindrical) = Z(ξ,equ)
+
+ξ¹(x::AbstractVector, equ::AxisymmetricTokamakCylindrical) = sqrt(x[1]^2 + x[2]^2)
+ξ²(x::AbstractVector, equ::AxisymmetricTokamakCylindrical) = x[3]
+ξ³(x::AbstractVector, equ::AxisymmetricTokamakCylindrical) = atan(x[2], x[1])
+
 g₁₁(x::AbstractVector, equ::AxisymmetricTokamakCylindrical) = one(eltype(x))
 g₂₂(x::AbstractVector, equ::AxisymmetricTokamakCylindrical) = one(eltype(x))
 g₃₃(x::AbstractVector, equ::AxisymmetricTokamakCylindrical) = R(x, equ)^2

src/analytic/axisymmetric_tokamak_toroidal.jl

@@ -1,3 +1,5 @@
+using SymEngine
+
 @doc raw"""
 Axisymmetric tokamak equilibrium in (r,θ,ϕ) coordinates with covariant
 components of the vector potential given by
@@ -51,6 +53,14 @@ A₁(x::AbstractVector, equ::AxisymmetricTokamakToroidal) = zero(eltype(x))
 A₂(x::AbstractVector, equ::AxisymmetricTokamakToroidal) = + equ.B₀ * equ.R₀ / cos(θ(x,equ))^2 * ( r(x,equ) * cos(θ(x,equ)) - equ.R₀ * log(R(x,equ) / equ.R₀) )
 A₃(x::AbstractVector, equ::AxisymmetricTokamakToroidal) = - equ.B₀ * r(x,equ)^2 / equ.q₀ / 2
 
+x¹(ξ::AbstractVector, equ::AxisymmetricTokamakToroidal) = X(ξ,equ)
+x²(ξ::AbstractVector, equ::AxisymmetricTokamakToroidal) = Y(ξ,equ)
+x³(ξ::AbstractVector, equ::AxisymmetricTokamakToroidal) = Z(ξ,equ)
+
+ξ¹(x::AbstractVector, equ::AxisymmetricTokamakToroidal) = sqrt((sqrt(x[1]^2 + x[2]^2)-equ.R₀)^2 + x[3]^2)
+ξ²(x::AbstractVector, equ::AxisymmetricTokamakToroidal) = atan(x[3], sqrt(x[1]^2 + x[2]^2)-equ.R₀)
+ξ³(x::AbstractVector, equ::AxisymmetricTokamakToroidal) = atan(x[2], x[1])
+
 g₁₁(x::AbstractVector, equ::AxisymmetricTokamakToroidal) = one(eltype(x))
 g₂₂(x::AbstractVector, equ::AxisymmetricTokamakToroidal) = r(x, equ)^2
 g₃₃(x::AbstractVector, equ::AxisymmetricTokamakToroidal) = R(x, equ)^2
@@ -66,7 +76,6 @@ function periodicity(x::AbstractVector, equ::AxisymmetricTokamakToroidal)
     return p
 end
 
-
 macro axisymmetric_tokamak_equilibrium_toroidal(R₀, B₀, q₀)
     generate_equilibrium_code(AxisymmetricTokamakToroidal(R₀, B₀, q₀); output=false)
 end

src/analytic/cartesian_equilibrium.jl

@@ -11,6 +11,14 @@ Z(x::AbstractVector, ::CartesianField) = x[3]
 
 J(x::AbstractVector, ::CartesianField) = one(eltype(x))
 
+x¹(ξ::AbstractVector, ::CartesianField) = ξ[1]
+x²(ξ::AbstractVector, ::CartesianField) = ξ[2]
+x³(ξ::AbstractVector, ::CartesianField) = ξ[3]
+
+ξ¹(x::AbstractVector, ::CartesianField) = x[1]
+ξ²(x::AbstractVector, ::CartesianField) = x[2]
+ξ³(x::AbstractVector, ::CartesianField) = x[3]
+
 g₁₁(x::AbstractVector, ::CartesianField) = one(eltype(x))
 g₂₂(x::AbstractVector, ::CartesianField) = one(eltype(x))
 g₃₃(x::AbstractVector, ::CartesianField) = one(eltype(x))

src/analytic/solovev_common.jl

@@ -15,18 +15,23 @@ r(x::AbstractVector, equ::AbstractSolovevEquilibrium) = sqrt(r²(x, equ))
 
 J(x::AbstractVector, equ::AbstractSolovevEquilibrium) = R(x,equ)
 
-
 A₁(x::AbstractVector, equ::AbstractSolovevEquilibrium) = + equ.B₀ * equ.R₀ * x[2] / x[1] / 2
 A₂(x::AbstractVector, equ::AbstractSolovevEquilibrium) = - equ.B₀ * equ.R₀ * log(x[1]) / 2
 
+x¹(ξ::AbstractVector, equ::AbstractSolovevEquilibrium) = X(ξ,equ)
+x²(ξ::AbstractVector, equ::AbstractSolovevEquilibrium) = Y(ξ,equ)
+x³(ξ::AbstractVector, equ::AbstractSolovevEquilibrium) = Z(ξ,equ)
+
+ξ¹(x::AbstractVector, equ::AbstractSolovevEquilibrium) = sqrt(x[1]^2 + x[2]^2) / equ.R₀
+ξ²(x::AbstractVector, equ::AbstractSolovevEquilibrium) = x[3] / equ.R₀
+ξ³(x::AbstractVector, equ::AbstractSolovevEquilibrium) = atan(x[2], x[1])
+
 g₁₁(x::AbstractVector, equ::AbstractSolovevEquilibrium) = one(eltype(x))
 g₂₂(x::AbstractVector, equ::AbstractSolovevEquilibrium) = one(eltype(x))
 g₃₃(x::AbstractVector, equ::AbstractSolovevEquilibrium) = R(x, equ)^2
 
-
 get_functions(::AbstractSolovevEquilibrium) = (X=X, Y=Y, Z=Z, R=R, r=r, θ=θ, ϕ=ϕ, r²=r²)
 
-
 function periodicity(x::AbstractArray{T,1}, equ::AbstractSolovevEquilibrium) where {T <: Number}
     p = zero(x)
     p[3] = 2π

test/test_analytic.jl

@@ -64,6 +64,16 @@ function test_equilibrium(equ_mod, t, x)
     @test equ_mod.dA₃dx₁(x...) == equ_mod.dA₃dx₁(t,x)
     @test equ_mod.dA₃dx₂(x...) == equ_mod.dA₃dx₂(t,x)
     @test equ_mod.dA₃dx₃(x...) == equ_mod.dA₃dx₃(t,x)
+
+    @test equ_mod.x¹(x...) == equ_mod.x¹(t,x)
+    @test equ_mod.x²(x...) == equ_mod.x²(t,x)
+    @test equ_mod.x³(x...) == equ_mod.x³(t,x)
+
+    @test equ_mod.ξ¹(x...) == equ_mod.ξ¹(t,x)
+    @test equ_mod.ξ²(x...) == equ_mod.ξ²(t,x)
+    @test equ_mod.ξ³(x...) == equ_mod.ξ³(t,x)
+
+    @test equ_mod.from_cartesian(t, equ_mod.to_cartesian(t, x)) == x
 end