move gradient from IMAS to IMASDD

ProjectTorreyPines/SOLPS2ctrl.jl#50 (comment)

src/math.jl

@@ -135,4 +135,150 @@ function clip_1(f, x::Real, x1::Real, y1::T) where {T<:Real}
     else
         return f(x)::T
     end
+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 gradient. The returned gradient has the same shape as the input array.
+
+`method` of the gradient can be one of [:third_order, :second_order, :central, :backward, :forward]
+
+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