diff --git a/src/math.jl b/src/math.jl index ea1e37e9..5f92e1d6 100644 --- a/src/math.jl +++ b/src/math.jl @@ -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}