Skip to content

Commit

Permalink
move gradient from IMAS to IMASDD
Browse files Browse the repository at this point in the history 
ProjectTorreyPines/SD4SOLPS.jl#50 (comment)
  • Loading branch information
@orso82
orso82 committed Mar 21, 2024
1 parent 98bffab commit 982ddec
Showing 1 changed file with 0 additions and 146 deletions.
146 changes: 0 additions & 146 deletions src/math.jl
View file
Original file line number Diff line number Diff line change
Expand Up @@ -26,152 +26,6 @@ function to_range(vector::AbstractVector{<:Real})
return range(vector[1], vector[end], N)
end

function gradient(arr::AbstractVector; method::Symbol=:second_order)
return gradient(1:length(arr), arr; method)
end

"""
gradient(coord::AbstractVector{C}, arr::AbstractVector{A}; method::Symbol=:second_order) where {C<:Real, A<:Real}
The finite difference `method` of the gradient can be one of [:third_order, :second_order, :central, :backward, :forward]
The returned gradient has the same shape as the input array https://numpy.org/doc/stable/reference/generated/numpy.gradient.html
For `:central` the gradient is computed using second order accurate central differences in the interior points and first order accurate one-sides (forward or backward) differences at the boundaries.
For `:second_order` the gradient is computed using second order accurate central differences in the interior points, and 2nd order differences at the boundaries.
For `:third_order` the gradient is computed from the cubic spline passing through the points
"""
function gradient(coord::AbstractVector{C}, arr::AbstractVector{A}; method::Symbol=:second_order) where {C<:Real,A<:Real}
grad = Array{promote_type(A, C)}(undef, length(arr))
return gradient!(grad, coord, arr; method)
end

"""
gradient!(grad::AbstractVector, coord::AbstractVector, arr::AbstractVector; method::Symbol=:second_order)
In place version of gradient(coord::AbstractVector, arr::AbstractVector; method::Symbol=:second_order)
"""
function gradient!(grad::Union{AbstractVector,SubArray{<:Real,1}}, coord::AbstractVector, arr::Union{AbstractVector,SubArray{<:Real,1}}; method::Symbol=:second_order)
np = length(arr)
@assert length(grad) == np "The length of your grad vector (length = $(length(grad))) is not equal to the length of your arr (length = $np)"
@assert length(coord) == np "The length of your coord (length = $(length(coord))) is not equal to the length of your arr (length = $np)"

if np < 3 && method == :second_order
method = :central
end

if method ∈ [:central, :backward, :forward]
# Forward difference at the beginning
grad[1] = (arr[2] - arr[1]) / (coord[2] - coord[1])
# backward difference at the end
grad[end] = (arr[end] - arr[end-1]) / (coord[end] - coord[end-1])
end

if method == :third_order
itp = DataInterpolations.CubicSpline(arr, coord)
for k in eachindex(arr)
grad[k] = DataInterpolations.derivative(itp, coord[k])
end

elseif method ∈ [:central, :second_order]
# Central difference in interior using numpy method
for p in 2:np-1
hs = coord[p] - coord[p-1]
fs = arr[p-1]
hd = coord[p+1] - coord[p]
fd = arr[p+1]
grad[p] = (hs^2 * fd + (hd^2 - hs^2) * arr[p] - hd^2 * fs) / (hs * hd * (hd + hs))
end
if method == :second_order
# Derived using Numerical Mathematics, section 10.10 and lecture notes by A. Yew
# and checked against formula from A. Yew for the case of equal spacing.
# Numerical Mathematics: A. Quarteroni, R. Sacco, F. Saleri, Springer (2007)
# https://sites.math.washington.edu/~morrow/464_17/sacco%20saleri%20numerical.pdf
# A. Yew: Lecture notes for APMA 0160 at Brown University
# https://www.dam.brown.edu/people/alcyew/handouts/numdiff.pdf
c = coord[2] - coord[1]
d = coord[3] - coord[2]
f = arr[1]
g = arr[2]
h = arr[3]
ccd = c / (c + d)
ccd2 = c^2 / (c + d)^2
grad[1] = (-f * (1 - ccd2) + g - h * ccd2) / (c * (1 - ccd))
c = coord[end-1] - coord[end]
d = coord[end-2] - coord[end-1]
f = arr[end]
g = arr[end-1]
h = arr[end-2]
ccd = c / (c + d)
ccd2 = c^2 / (c + d)^2
grad[end] = (-f * (1 - ccd2) + g - h * ccd2) / (c * (1 - ccd))
end

elseif method == :backward
for p in 2:np-1
grad[p] = (arr[p] - arr[p-1]) / (coord[p] - coord[p-1])
end

elseif method == :forward
for p in 2:np-1
grad[p] = (arr[p+1] - arr[p]) / (coord[p+1] - coord[p])
end

else
error("difference method $(method) doesn't exist in gradient function")
end

return grad
end

function gradient(mat::Matrix; method::Symbol=:second_order)
return gradient(1:size(mat)[1], 1:size(mat)[2], mat; method)
end

function gradient(mat::Matrix, dim::Int; method::Symbol=:second_order)
return gradient(1:size(mat)[1], 1:size(mat)[2], mat, dim; method)
end

"""
gradient(coord1::AbstractVector, coord2::AbstractVector, mat::Matrix, dim::Int; method::Symbol=:second_order)
Finite difference method of the gradient: [:second_order, :central, :backward, :forward]
Can be applied to either the first (dim=1) or second (dim=2) dimension
"""
function gradient(coord1::AbstractVector, coord2::AbstractVector, mat::Matrix, dim::Int; method::Symbol=:second_order)
nrows, ncols = size(mat)
d = Matrix{eltype(mat)}(undef, nrows, ncols)
if dim == 1
for i in 1:ncols
gradient!(@views(d[:, i]), coord1, @views(mat[:, i]); method)
end
return d
elseif dim == 2
for i in 1:nrows
gradient!(@views(d[i, :]), coord2, @views(mat[i, :]); method)
end
return d
else
throw(ArgumentError("dim should be either 1 or 2"))
end
end

"""
gradient(coord1::AbstractVector, coord2::AbstractVector, mat::Matrix)
Finite difference method of the gradient: [:second_order, :central, :backward, :forward]
Computes the gradient in both dimensions
"""
function gradient(coord1::AbstractVector, coord2::AbstractVector, mat::Matrix; method::Symbol=:second_order)
d1 = gradient(coord1, coord2, mat, 1; method)
d2 = gradient(coord1, coord2, mat, 2; method)
return d1, d2
end

"""
meshgrid(x::AbstractVector{T}, y::AbstractVector{T}) where {T}
Expand Down

0 comments on commit 982ddec

Please sign in to comment.