Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Jump to bottom

use b-splines to calculate moments; add spline keyword #3

Merged
merged 1 commit into from
Jun 12, 2023
Merged

use b-splines to calculate moments; add spline keyword #3

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
1 commit
Select commit
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
1 change: 1 addition & 0 deletions Project.toml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -4,6 +4,7 @@ authors = ["Orso Meneghini <orso82@gmail.com>"]
version = "0.1.0"

[deps]
Dierckx = "39dd38d3-220a-591b-8e3c-4c3a8c710a94"
LinearAlgebra = "37e2e46d-f89d-539d-b4ee-838fcccc9c8e"
Optim = "429524aa-4258-5aef-a3af-852621145aeb"
RecipesBase = "3cdcf5f2-1ef4-517c-9805-6587b60abb01"
124 changes: 82 additions & 42 deletions src/MXH.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -97,17 +97,34 @@ function copy_MXH!(mxh1::MXH, mxh2::MXH)
end

"""
MXH_moment(f, w, d)
MXH_moment_trapz(f, w, d)

This does Int[f.w]/Int[w.w]
This does Int[f.w]/Int[w.w] using Trapezoidal method
If w is a pure Fourier mode, this gives the Fourier coefficient
"""
function MXH_moment(f::AbstractVector{<:Real}, w::AbstractVector{<:Real}, d::AbstractVector{<:Real})
# Could probably be replaced by some Julia trapz
function MXH_moment_trapz(f::AbstractVector{<:Real}, w::AbstractVector{<:Real}, dx::AbstractVector{<:Real})
# Could probably be replaced by Trapz.jl
N = length(f)
@assert length(w) == length(d) == N
@inbounds s0 = sum((f[i] * w[i] + f[i+1] * w[i+1]) * d[i] for i in 1:(N-1))
@inbounds s1 = sum((w[i]^2 + w[i+1]^2) * d[i] for i in 1:(N-1))
@assert length(w) == length(dx) == N
@inbounds s0 = sum((f[i] * w[i] + f[i+1] * w[i+1]) * dx[i] for i in 1:(N-1))
@inbounds s1 = sum((w[i]^2 + w[i+1]^2) * dx[i] for i in 1:(N-1))
res = s0 / s1
return res
end

"""
MXH_moment_spline(f, w, x)

This does Int[f.w]/Int[w.w] using B-splines
If w is a pure Fourier mode, this gives the Fourier coefficient
"""
function MXH_moment_spline(f::AbstractVector{<:Real}, w::AbstractVector{<:Real}, x::AbstractVector{<:Real})
@assert length(w) == length(x) == length(f)

spl0 = Spline1D(x, f .* w)
spl1 = Spline1D(x, w.^2)
s0 = integrate(spl0,x[begin],x[end])
s1 = integrate(spl1,x[begin],x[end])
res = s0 / s1
return res
end
Expand Down Expand Up @@ -184,25 +201,25 @@ function find_extrema(R, Z)
end

"""
MXH(pr::Vector{T}, pz::Vector{T}, MXH_modes::Integer) where {T<:Real}
MXH(pr::Vector{T}, pz::Vector{T}, MXH_modes::Integer; spline=false) where {T<:Real}

Compute Fourier coefficients for Miller-extended-harmonic representation:

R(r,θ) = R(r) + a(r)*cos(θᵣ(r,θ)) where θᵣ(r,θ) = θ + C₀(r) + sum[Cᵢ(r)*cos(i*θ) + Sᵢ(r)*sin(i*θ)]
Z(r,θ) = Z(r) - κ(r)*a(r)*sin(θ)

Where pr,pz are the flux surface coordinates and MXH_modes is the number of modes
Where pr,pz are the flux surface coordinates and MXH_modes is the number of modes. Spline keyword indicates to use spline integration for modes
"""
function MXH(pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real}, MXH_modes::Integer=5; θ=nothing, Δθᵣ=nothing, dθ=nothing, Fm=nothing, optimize_fit=false)
function MXH(pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real}, MXH_modes::Integer=5; θ=nothing, Δθᵣ=nothing, dθ=nothing, Fm=nothing, optimize_fit=false, spline=false)
Copy link
Member

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Should we make spline=true the default? @lstagner do you have a sense of how the spline approach affects performance?

Copy link
Contributor Author

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

its a somewhat significant slowdown.

N = 2 spline = false:  0.000039 seconds (25 allocations: 12.375 KiB)
N = 2 spline = true:   0.000336 seconds (115 allocations: 347.688 KiB)
N = 4 spline = false:  0.000073 seconds (25 allocations: 12.406 KiB)
N = 4 spline = true:   0.000438 seconds (187 allocations: 615.969 KiB)
N = 6 spline = false:  0.000051 seconds (25 allocations: 12.438 KiB)
N = 6 spline = true:   0.000759 seconds (259 allocations: 884.250 KiB)
N = 8 spline = false:  0.000067 seconds (25 allocations: 12.469 KiB)
N = 8 spline = true:   0.000908 seconds (331 allocations: 1.126 MiB)
N = 10 spline = false:  0.000065 seconds (25 allocations: 12.500 KiB)
N = 10 spline = true:   0.000979 seconds (403 allocations: 1.388 MiB)

sin_coeffs = zeros(MXH_modes)
cos_coeffs = zeros(MXH_modes)
mxh = MXH(0.0, 0.0, 0.0, 0.0, 0.0, cos_coeffs, sin_coeffs)
return MXH!(mxh, pr, pz; θ, Δθᵣ, dθ, Fm, optimize_fit)
return MXH!(mxh, pr, pz; θ, Δθᵣ, dθ, Fm, optimize_fit, spline)
end

function MXH!(mxh::MXH, pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real};
θ=nothing, Δθᵣ=nothing, dθ=nothing, Fm=nothing,
optimize_fit=false)
optimize_fit=false, spline=false)
rmin = 0.0
rmax = 0.0
zmin = 0.0
Expand All @@ -223,7 +240,7 @@ function MXH!(mxh::MXH, pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real};
Z0 = 0.5 * (zmax + zmin)
a = 0.5 * (rmax - rmin)
b = 0.5 * (zmax - zmin)
return MXH!(mxh, pr, pz, R0, Z0, a, b, θ, Δθᵣ, dθ, Fm, optimize_fit)
return MXH!(mxh, pr, pz, R0, Z0, a, b, θ, Δθᵣ, dθ, Fm, optimize_fit, spline)
end

function clockwise!(pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real})
Expand Down Expand Up @@ -327,36 +344,65 @@ function MXH_angles!(θ::AbstractVector{<:Real}, Δθᵣ::AbstractVector{<:Real}
end
end

function MXH_coeffs!(sin_coeffs::AbstractVector{<:Real}, cos_coeffs::AbstractVector{<:Real},
function MXH_coeffs_trapz!(c0::Ref{<:Real}, sin_coeffs::AbstractVector{<:Real}, cos_coeffs::AbstractVector{<:Real},
θ::AbstractVector{<:Real}, Δθᵣ::AbstractVector{<:Real}, dθ::AbstractVector{<:Real};
Fm::Union{AbstractVector{<:Real},Nothing}=nothing)
@assert length(sin_coeffs) == length(cos_coeffs)
Fm === nothing && (Fm = similar(θ))

@inbounds @views for j in eachindex(dθ)[1:end-1]
dθ[j] = θ[j+1] - θ[j]
dθ[j] < 0 && (dθ[j] += 2π)
end
dθ[end] = dθ[1]

Fm .= 1.0
c0[] = MXH_moment_trapz(Δθᵣ, Fm, dθ)

@inbounds for m in eachindex(sin_coeffs)
Fm .= sin.(m .* θ)
sin_coeffs[m] = MXH_moment(Δθᵣ, Fm, dθ)
sin_coeffs[m] = MXH_moment_trapz(Δθᵣ, Fm, dθ)

Fm .= cos.(m .* θ)
cos_coeffs[m] = MXH_moment(Δθᵣ, Fm, dθ)
cos_coeffs[m] = MXH_moment_trapz(Δθᵣ, Fm, dθ)
end
end

function MXH_coeffs_spline!(c0::Ref{<:Real}, sin_coeffs::AbstractVector{<:Real}, cos_coeffs::AbstractVector{<:Real},
θ::AbstractVector{<:Real}, Δθᵣ::AbstractVector{<:Real}, dθ::AbstractVector{<:Real};
Fm::Union{AbstractVector{<:Real},Nothing}=nothing)
@assert length(sin_coeffs) == length(cos_coeffs)
Fm === nothing && (Fm = similar(θ))

Fm .= 1.0
c0[] = MXH_moment_spline(Δθᵣ, Fm, θ)

@inbounds for m in eachindex(sin_coeffs)
Fm .= sin.(m .* θ)
sin_coeffs[m] = MXH_moment_spline(Δθᵣ, Fm, θ)

Fm .= cos.(m .* θ)
cos_coeffs[m] = MXH_moment_spline(Δθᵣ, Fm, θ)
end
end

function MXH(pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real}, R0::Real, Z0::Real, a::Real, b::Real, MXH_modes::Integer;
θ=nothing, Δθᵣ=nothing, dθ=nothing, Fm=nothing)
θ=nothing, Δθᵣ=nothing, dθ=nothing, Fm=nothing, optimize_fit=false, spline=false)

sin_coeffs = zeros(MXH_modes)
cos_coeffs = zeros(MXH_modes)
mxh = MXH(0.0, 0.0, 0.0, 0.0, 0.0, cos_coeffs, sin_coeffs)
return MXH!(mxh, pr, pz, R0, Z0, a, b, θ, Δθᵣ, dθ, Fm)
return MXH!(mxh, pr, pz, R0, Z0, a, b, θ, Δθᵣ, dθ, Fm, optimize_fit, spline)
end

function MXH!(mxh::MXH, pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real}, R0::Real, Z0::Real, a::Real, b::Real,
θ::Nothing=nothing, Δθᵣ::Nothing=nothing, dθ::Nothing=nothing, Fm::Nothing=nothing, optimize_fit=false)
MXH!(mxh, pr, pz, R0, Z0, a, b, similar(pr), similar(pr), similar(pr), similar(pr), optimize_fit)
θ::Nothing=nothing, Δθᵣ::Nothing=nothing, dθ::Nothing=nothing, Fm::Nothing=nothing, optimize_fit=false, spline=false)
MXH!(mxh, pr, pz, R0, Z0, a, b, similar(pr), similar(pr), similar(pr), similar(pr), optimize_fit, spline)
end

function MXH!(mxh::MXH, pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real}, R0::Real, Z0::Real, a::Real, b::Real,
θ::AbstractVector{<:Real}, Δθᵣ::AbstractVector{<:Real}, dθ::AbstractVector{<:Real}, Fm::AbstractVector{<:Real}, optimize_fit=false)
θ::AbstractVector{<:Real}, Δθᵣ::AbstractVector{<:Real}, dθ::AbstractVector{<:Real}, Fm::AbstractVector{<:Real},
optimize_fit=false, spline=false)

@assert length(pr) == length(pz)

Expand All @@ -370,16 +416,13 @@ function MXH!(mxh::MXH, pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real},
# Calculate angles with proper branches
MXH_angles!(θ, Δθᵣ, pr, pz, R0, Z0, a, b)

@inbounds @views for j in eachindex(dθ)[1:end-1]
dθ[j] = θ[j+1] - θ[j]
dθ[j] < 0 && (dθ[j] += 2π)
c0 = Ref(0.0)
if spline
@views MXH_coeffs_spline!(c0, mxh.s, mxh.c, θ, Δθᵣ, dθ; Fm)
else
@views MXH_coeffs_trapz!(c0, mxh.s, mxh.c, θ, Δθᵣ, dθ; Fm)
end
dθ[end] = dθ[1]

Fm .= 1.0 # cos(0 * θ)
mxh.c0 = MXH_moment(Δθᵣ, Fm, dθ)

MXH_coeffs!(mxh.s, mxh.c, θ, Δθᵣ, dθ; Fm)
mxh.c0 = c0[]

if optimize_fit
flat = flat_coeffs(mxh)
Expand Down Expand Up @@ -500,7 +543,7 @@ end

function fit_flattened!(flat::AbstractVector{<:Real}, pr::AbstractVector{<:Real}, pz::AbstractVector{<:Real},
θ::AbstractVector{<:Real}, Δθᵣ::AbstractVector{<:Real}, dθ::AbstractVector{<:Real}, Fm::AbstractVector{<:Real};
rmin::Real=minimum(pr), rmax::Real=maximum(pr), zmin::Real=minimum(pz), zmax::Real=maximum(pz))
rmin::Real=minimum(pr), rmax::Real=maximum(pr), zmin::Real=minimum(pz), zmax::Real=maximum(pz), spline::Bool=false)

@assert length(pr) == length(pz)

Expand All @@ -519,17 +562,14 @@ function fit_flattened!(flat::AbstractVector{<:Real}, pr::AbstractVector{<:Real}
# Calculate angles with proper branches
MXH_angles!(θ, Δθᵣ, pr, pz, R0, Z0, a, b)

@inbounds @views for j in eachindex(dθ)[1:end-1]
dθ[j] = θ[j+1] - θ[j]
dθ[j] < 0 && (dθ[j] += 2π)
end
dθ[end] = dθ[1]

Fm .= 1.0 # cos(0 * θ)
flat[5] = MXH_moment(Δθᵣ, Fm, dθ)

L = (length(flat) - 5) ÷ 2
@views MXH_coeffs!(flat[(6+L):(5+2L)], flat[6:(5+L)], θ, Δθᵣ, dθ; Fm)
c0 = Ref(0.0)
if spline
@views MXH_coeffs_spline!(c0,flat[(6+L):(5+2L)], flat[6:(5+L)], θ, Δθᵣ, dθ; Fm)
else
@views MXH_coeffs_trapz!(c0,flat[(6+L):(5+2L)], flat[6:(5+L)], θ, Δθᵣ, dθ; Fm)
end
flat[5] = c0[]
return flat
end

Expand Down Expand Up @@ -656,4 +696,4 @@ function (mxh::MXH)(θ::Real)
r = R_MXH(θ, mxh, a)
z = Z_MXH(θ, mxh, a)
return r, z
end
end
3 changes: 2 additions & 1 deletion src/MillerExtendedHarmonic.jl
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -3,6 +3,7 @@ module MillerExtendedHarmonic
using RecipesBase
import LinearAlgebra: dot
using Optim
using Dierckx

const halfpi = 0.5 * π

Expand All @@ -17,4 +18,4 @@ export Jacobian, ∇ρ, ∇ρ2, ∇θ, ∇θ2, gρρ, gρθ, gθθ, gρρ_gρθ,
include("relative.jl")
export in_surface, nearest_angle

end # module
end # module