Update derivative of shape coefficients in edge

src/TEQUILA.jl

@@ -28,6 +28,10 @@ const x0_2 = zeros(2)
 const jump_success = @SVector[JuMP.OPTIMAL, JuMP.LOCALLY_SOLVED]
 const int_order = 5
 
+# used in surfaces_FE for update_edge_derivatives!, as well as in fitted_surfaces
+const δ_frac_2 = 0.5
+const δ_frac_3 = 0.75
+
 mutable struct QuadInfo{VR1<:AbstractVector{<:Real}, VSV1<:Vector{<:AbstractSparseVector},
                         MR1<:AbstractMatrix{<:Real}, VR2<:AbstractVector{<:Real},
                         MR2<:AbstractMatrix{<:Real}, MR3<:AbstractMatrix{<:Real}}

src/fit_MXH.jl

@@ -323,8 +323,9 @@ function fit_MXH!(flat, shot, lvl, Ψaxis, Raxis, Zaxis, ρaxis, Fi, Fo, P)
 
 end
 
-function update_shot!(shot::Shot, surfaces, Ψaxis)
-    shot.R0fe, shot.Z0fe, shot.ϵfe, shot.κfe, shot.c0fe, shot.cfe, shot.sfe = surfaces_FE(shot.ρ, surfaces)
+function update_shot!(shot::Shot, surfaces, Ψaxis, flat_δ2=nothing, flat_δ3=nothing)
+    shot.surfaces .= surfaces
+    shot.R0fe, shot.Z0fe, shot.ϵfe, shot.κfe, shot.c0fe, shot.cfe, shot.sfe = surfaces_FE(shot.ρ, surfaces, flat_δ2, flat_δ3)
     shot.C .= 0.0
     shot.C[2:2:end, 1] .= Ψaxis .* (1.0 .- shot.ρ.^2)
     shot.C[1:2:end, 1] .= -2.0 .* Ψaxis .* shot.ρ
@@ -335,19 +336,21 @@ end
 
 function refit!(shot::Shot, Ψaxis::Real, Raxis::Real, Zaxis::Real; debug::Bool=false, fit_fallback::Bool=true,)
     # Using ρ = rho poloidal (sqrt((Ψ-Ψaxis)/sqrt(Ψbnd-Ψaxis)))
-    local surfaces
+    local surfaces, flat_δ2, flat_δ3
     warn_concentric = false
     try
-        surfaces = fitted_surfaces(shot, Ψaxis, Raxis, Zaxis)
+        surfaces, flat_δ2, flat_δ3 = fitted_surfaces(shot, Ψaxis, Raxis, Zaxis)
     catch err
        (isa(err, InterruptException) || !fit_fallback) && rethrow(err)
        warn_concentric = true
        if debug
            println("    Warning: Fit fell back to concentric surfaces due to ", typeof(err))
        end
        surfaces = concentric_surfaces(shot, Raxis, Zaxis)
+       flat_δ2 = nothing
+       flat_δ3 = nothing
     end
-    update_shot!(shot, surfaces, Ψaxis)
+    update_shot!(shot, surfaces, Ψaxis, flat_δ2, flat_δ3)
     return shot, warn_concentric
 end
 
@@ -358,7 +361,9 @@ function fitted_surfaces(shot, Ψaxis, Raxis, Zaxis)
     surfaces = deepcopy(shot.surfaces)
 
     ρaxis, _ = ρθ_RZ(shot, Raxis, Zaxis)
-    Threads.@threads for k in 2:(shot.N - 1)
+    #Threads.@threads
+    # multithreading is broken here somehow. Maybe the tid pattern?
+    for k in 2:(shot.N - 1)
         lvl = Ψaxis * (1.0 - shot.ρ[k] ^ 2)
         tid = Threads.threadid()
         Fi = Fis[tid]
@@ -368,6 +373,21 @@ function fitted_surfaces(shot, Ψaxis, Raxis, Zaxis)
         fit_MXH!(flat, shot, lvl, Ψaxis, Raxis, Zaxis, ρaxis, Fi, Fo, P)
     end
 
+    # Now fit two more surfaces in the last grid region so we can update derivatives
+    tid = Threads.threadid()
+    δρ = shot.ρ[end] - shot.ρ[end-1]
+
+    flat_δ2 = zeros(2L+5)
+    ρ_δ2 = shot.ρ[end-1] + δ_frac_2 * δρ
+    lvl = Ψaxis * (1.0 - ρ_δ2 ^ 2)
+    fit_MXH!(flat_δ2, shot, lvl, Ψaxis, Raxis, Zaxis, ρaxis, Fis[tid], Fos[tid], Ps[tid])
+
+    flat_δ3  = zeros(2L+5)
+    ρ_δ3  = shot.ρ[end-1] + δ_frac_3 * δρ
+    lvl = Ψaxis * (1.0 - ρ_δ3 ^ 2)
+    fit_MXH!(flat_δ3, shot, lvl, Ψaxis, Raxis, Zaxis, ρaxis, Fis[tid], Fos[tid], Ps[tid])
+
+
     # Extrapolate or set to zero on-axis
     ρ2 = shot.ρ[2]
     ρ3 = shot.ρ[3]
@@ -378,7 +398,7 @@ function fitted_surfaces(shot, Ψaxis, Raxis, Zaxis)
     surfaces[4, 1] = h .* (surfaces[4, 2] .* ρ3 .- surfaces[4, 3] .* ρ2)
     surfaces[5, 1] = h .* (surfaces[5, 2] .* ρ3 .- surfaces[5, 3] .* ρ2)
     @views surfaces[6:end, 1] .= 0.0
-    return surfaces
+    return surfaces, flat_δ2, flat_δ3
 end
 
 

src/surfaces.jl

@@ -82,7 +82,7 @@ function surfaces_FE(ρ:: AbstractVector{<:Real}, surfaces:: AbstractVector{<:MX
     return R0fe, Z0fe, ϵfe, κfe, c0fe, cfe, sfe
 end
 
-function surfaces_FE(ρ:: AbstractVector{<:Real}, surfaces:: AbstractMatrix{<:Real} )
+function surfaces_FE(ρ:: AbstractVector{<:Real}, surfaces:: AbstractMatrix{<:Real}, flat_δ2::Nothing=nothing, flat_δ3::Nothing=nothing)
 
     rtype = typeof(ρ[1])
 
@@ -104,6 +104,47 @@ function surfaces_FE(ρ:: AbstractVector{<:Real}, surfaces:: AbstractMatrix{<:Re
     return R0fe, Z0fe, ϵfe, κfe, c0fe, cfe, sfe
 end
 
+function surfaces_FE(ρ:: AbstractVector{<:Real}, surfaces:: AbstractMatrix{<:Real}, flat_δ2::AbstractVector{<:Real}, flat_δ3::AbstractVector{<:Real})
+
+    R0fe, Z0fe, ϵfe, κfe, c0fe, cfe, sfe = surfaces_FE(ρ, surfaces)
+
+    ρ1 = ρ[end-1]
+    ρ4 = ρ[end]
+    δρ = ρ4 - ρ1
+
+    ρ2 = ρ1 + δ_frac_2 * δρ
+    ρ3 = ρ1 + δ_frac_3 * δρ
+    _, nu_ou2, nu_eu2, nu_ol2, nu_el2 = compute_bases(ρ, ρ2)
+    _, nu_ou3, nu_eu3, nu_ol3, nu_el3 = compute_bases(ρ, ρ3)
+    nus = (nu_ou2, nu_eu2, nu_ol2, nu_el2, nu_ou3, nu_eu3, nu_ol3, nu_el3)
+
+    update_edge_derivatives!(R0fe, flat_δ2[1], flat_δ3[1], nus...)
+    update_edge_derivatives!(Z0fe, flat_δ2[2], flat_δ3[2], nus...)
+    update_edge_derivatives!(ϵfe,  flat_δ2[3], flat_δ3[3], nus...)
+    update_edge_derivatives!(κfe,  flat_δ2[4], flat_δ3[4], nus...)
+    update_edge_derivatives!(c0fe, flat_δ2[5], flat_δ3[5], nus...)
+    M = length(cfe)
+    for m in eachindex(cfe)
+        update_edge_derivatives!(cfe[m], flat_δ2[5+m],   flat_δ3[5+m],   nus...)
+        update_edge_derivatives!(sfe[m], flat_δ2[5+m+M], flat_δ3[5+m+M], nus...)
+    end
+
+    return R0fe, Z0fe, ϵfe, κfe, c0fe, cfe, sfe
+end
+
+function update_edge_derivatives!(Yfe, Y2, Y3, nu_ou2, nu_eu2, nu_ol2, nu_el2, nu_ou3, nu_eu3, nu_ol3, nu_el3)
+
+    # deterimant of |nu_ou2   nu_ol2|
+    #               |nu_ou3   nu_ol3|
+    D = nu_ou2 * nu_ol3 - nu_ou3 * nu_ol2
+    b2 = Y2 - Yfe.coeffs[end-2] * nu_eu2 - Yfe.coeffs[end] * nu_el2
+    b3 = Y3 - Yfe.coeffs[end-2] * nu_eu3 - Yfe.coeffs[end] * nu_el3
+    Yfe.coeffs[end-3] = (b2 * nu_ol3 - b3 * nu_ol2) / D
+    Yfe.coeffs[end-1] = (nu_ou2 * b3 - nu_ou3 * b2) / D
+end
+
+
+
 function θ_at_RZ(shot::Shot, ρ::Real, R::Real, Z::Real; tid::Int = Threads.threadid())
     R0x, Z0x, ϵx, κx, c0x, cx, sx = compute_MXH(shot, ρ; tid)
     ax = R0x * ϵx