-
Notifications
You must be signed in to change notification settings - Fork 4
Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
- Loading branch information
Showing
13 changed files
with
275 additions
and
123 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,7 +1,7 @@ | ||
| name = "Chmy" | ||
| uuid = "33a72cf0-4690-46d7-b987-06506c2248b9" | ||
| authors = ["Ivan Utkin <[email protected]>, Ludovic Raess <[email protected]>, and contributors"] | ||
| version = "0.1.4" | ||
| version = "0.1.5" | ||
|
|
||
| [deps] | ||
| Adapt = "79e6a3ab-5dfb-504d-930d-738a2a938a0e" | ||
|
|
||
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -1,76 +1,35 @@ | ||
| module GridOperators | ||
|
|
||
| using Chmy | ||
| using Chmy.Grids | ||
| import Chmy.@add_cartesian | ||
|
|
||
| export left, right, δ, ∂, lerp | ||
|
|
||
| Base.@assume_effects :foldable p(::Dim{D}, I::Vararg{Integer,N}) where {D,N} = ntuple(i -> i == D ? I[i] + oneunit(I[i]) : I[i], Val(N)) | ||
| Base.@assume_effects :foldable m(::Dim{D}, I::Vararg{Integer,N}) where {D,N} = ntuple(i -> i == D ? I[i] - oneunit(I[i]) : I[i], Val(N)) | ||
|
|
||
| @add_cartesian left(f, ::Vertex, dim, I::Vararg{Integer,N}) where {N} = f[m(dim, I...)...] | ||
| @add_cartesian left(f, ::Center, dim, I::Vararg{Integer,N}) where {N} = f[I...] | ||
|
|
||
| @add_cartesian right(f, ::Vertex, dim, I::Vararg{Integer,N}) where {N} = f[I...] | ||
| @add_cartesian right(f, ::Center, dim, I::Vararg{Integer,N}) where {N} = f[p(dim, I...)...] | ||
|
|
||
| # finite difference | ||
| @add_cartesian δ(f, loc, dim, I::Vararg{Integer,N}) where {N} = right(f, loc, dim, I...) - left(f, loc, dim, I...) | ||
| export left, right, δ, ∂ | ||
|
|
||
| # partial derivative | ||
| @add_cartesian ∂(f, loc, grid, dim, I::Vararg{Integer,N}) where {N} = δ(f, loc, dim, I...) * iΔ(grid, loc, dim, I...) | ||
| export InterpolationRule, Linear, HarmonicLinear | ||
| export itp, lerp, hlerp | ||
|
|
||
| # interpolation | ||
| @add_cartesian lerp(f, ::Center, grid, dim, I::Vararg{Integer,N}) where {N} = eltype(f)(0.5) * (f[I...] + f[p(dim, I...)...]) | ||
| @add_cartesian function lerp(f, ::Vertex, grid, dim, I::Vararg{Integer,N}) where {N} | ||
| t = eltype(f)(0.5) * Δ(grid, Center(), dim, m(dim, I...)...) * iΔ(grid, Vertex(), dim, I...) | ||
| a = f[I...] | ||
| b = f[m(dim, I...)...] | ||
| return muladd(t, a - b, b) | ||
| end | ||
|
|
||
| # more efficient for uniform grids | ||
| @add_cartesian lerp(f, ::Vertex, grid::UniformGrid, dim, I::Vararg{Integer,N}) where {N} = eltype(f)(0.5) * (f[m(dim, I...)...] + f[I...]) | ||
|
|
||
| # operators on Cartesian grids | ||
| for (dim, coord) in enumerate((:x, :y, :z)) | ||
| left = Symbol(:left, coord) | ||
| right = Symbol(:right, coord) | ||
| δ = Symbol(:δ, coord) | ||
| ∂ = Symbol(:∂, coord) | ||
|
|
||
| @eval begin | ||
| export $δ, $∂, $left, $right | ||
| using Chmy | ||
| using Chmy.Grids | ||
| using Chmy.Fields | ||
|
|
||
| """ | ||
| $($left)(f, loc, I) | ||
| import Chmy.@add_cartesian | ||
|
|
||
| "left side" of a field (`[1:end-1]`) in $($(string(coord))) direction. | ||
| """ | ||
| @add_cartesian $left(f, loc, I::Vararg{Integer,N}) where {N} = left(f, loc, Val($dim), I...) | ||
| import Base: @propagate_inbounds, front | ||
|
|
||
| """ | ||
| $($right)(f, loc, I) | ||
| p(::Dim{D}, I::Vararg{Integer,N}) where {D,N} = ntuple(i -> i == D ? I[i] + oneunit(I[i]) : I[i], Val(N)) | ||
| m(::Dim{D}, I::Vararg{Integer,N}) where {D,N} = ntuple(i -> i == D ? I[i] - oneunit(I[i]) : I[i], Val(N)) | ||
|
|
||
| "right side" of a field (`[2:end]`) in $($(string(coord))) direction. | ||
| """ | ||
| @add_cartesian $right(f, loc, I::Vararg{Integer,N}) where {N} = right(f, loc, Val($dim), I...) | ||
| @add_cartesian il(loc::Vertex, from::Center, dim, I::Vararg{Integer,N}) where {N} = I | ||
| @add_cartesian il(loc::Center, from::Vertex, dim, I::Vararg{Integer,N}) where {N} = m(dim, I...) | ||
|
|
||
| """ | ||
| $($δ)(f, loc, I) | ||
| @add_cartesian ir(loc::Vertex, from::Center, dim, I::Vararg{Integer,N}) where {N} = p(dim, I...) | ||
| @add_cartesian ir(loc::Center, from::Vertex, dim, I::Vararg{Integer,N}) where {N} = I | ||
|
|
||
| Finite difference in $($(string(coord))) direction. | ||
| """ | ||
| @add_cartesian $δ(f, loc, I::Vararg{Integer,N}) where {N} = δ(f, loc, Val($dim), I...) | ||
| @add_cartesian il(loc::L, from::L, dim, I::Vararg{Integer,N}) where {N,L<:Location} = m(dim, I...) | ||
| @add_cartesian ir(loc::L, from::L, dim, I::Vararg{Integer,N}) where {N,L<:Location} = p(dim, I...) | ||
|
|
||
| """ | ||
| $($∂)(f, loc, grid, I) | ||
| @add_cartesian left(f, loc, from, dim, I::Vararg{Integer,N}) where {N} = f[il(loc, from, dim, I...)...] | ||
| @add_cartesian right(f, loc, from, dim, I::Vararg{Integer,N}) where {N} = f[ir(loc, from, dim, I...)...] | ||
|
|
||
| Directional partial derivative in $($(string(coord))) direction. | ||
| """ | ||
| @add_cartesian $∂(f, loc, grid, I::Vararg{Integer,N}) where {N} = ∂(f, loc, grid, Val($dim), I...) | ||
| end | ||
| end | ||
| include("partial_derivatives.jl") | ||
| include("interpolation.jl") | ||
| include("field_operators.jl") | ||
|
|
||
| end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,65 @@ | ||
| # staggered grid operators | ||
| @add_cartesian left(f::AbstractField, dim, I::Vararg{Integer,N}) where {N} = left(f, location(f, dim), flip(location(f, dim)), dim, I...) | ||
|
|
||
| @add_cartesian right(f::AbstractField, dim, I::Vararg{Integer,N}) where {N} = right(f, location(f, dim), flip(location(f, dim)), dim, I...) | ||
|
|
||
| @add_cartesian δ(f::AbstractField, dim, I::Vararg{Integer,N}) where {N} = δ(f, location(f, dim), flip(location(f, dim)), dim, I...) | ||
|
|
||
| @add_cartesian ∂(f::AbstractField, grid, dim, I::Vararg{Integer,N}) where {N} = ∂(f, location(f, dim), flip(location(f, dim)), grid, dim, I...) | ||
|
|
||
| # staggered operators on Cartesian grids | ||
| for (dim, coord) in enumerate((:x, :y, :z)) | ||
| _l = Symbol(:left, coord) | ||
| _r = Symbol(:right, coord) | ||
| _δ = Symbol(:δ, coord) | ||
| _∂ = Symbol(:∂, coord) | ||
|
|
||
| @eval begin | ||
| export $_δ, $_∂, $_l, $_r | ||
|
|
||
| """ | ||
| $($_l)(f, I) | ||
| "left side" of a field (`[1:end-1]`) in $($(string(coord))) direction. | ||
| """ | ||
| @add_cartesian function $_l(f::AbstractField, I::Vararg{Integer,N}) where {N} | ||
| left(f, Dim($dim), I...) | ||
| end | ||
|
|
||
| """ | ||
| $($_r)(f, I) | ||
| "right side" of a field (`[2:end]`) in $($(string(coord))) direction. | ||
| """ | ||
| @add_cartesian function $_r(f::AbstractField, I::Vararg{Integer,N}) where {N} | ||
| right(f, Dim($dim), I...) | ||
| end | ||
|
|
||
| """ | ||
| $($_δ)(f, I) | ||
| Finite difference in $($(string(coord))) direction. | ||
| """ | ||
| @add_cartesian function $_δ(f::AbstractField, I::Vararg{Integer,N}) where {N} | ||
| δ(f, Dim($dim), I...) | ||
| end | ||
|
|
||
| """ | ||
| $($_∂)(f, grid, I) | ||
| Directional partial derivative in $($(string(coord))) direction. | ||
| """ | ||
| @add_cartesian function $_∂(f::AbstractField, grid, I::Vararg{Integer,N}) where {N} | ||
| ∂(f, grid, Dim($dim), I...) | ||
| end | ||
| end | ||
| end | ||
|
|
||
| # covariant derivatives | ||
| @propagate_inbounds @generated function divg(V::NamedTuple{names,<:NTuple{N,AbstractField}}, grid::StructuredGrid{N}, I::Vararg{Integer,N}) where {names,N} | ||
| return :(Base.Cartesian.@ncall $N (+) D -> ∂(V[D], grid, Dim(D), I...)) | ||
| end | ||
|
|
||
| @propagate_inbounds function divg(V::NamedTuple{names,<:NTuple{N,AbstractField}}, grid::StructuredGrid{N}, I::CartesianIndex{N}) where {names,N} | ||
| return divg(V, grid, Tuple(I)...) | ||
| end |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,97 @@ | ||
| # 1D interpolation rules | ||
|
|
||
| abstract type InterpolationRule end | ||
|
|
||
| struct Linear <: InterpolationRule end | ||
| struct HarmonicLinear <: InterpolationRule end | ||
|
|
||
| itp_rule(::Linear, t, a, b) = muladd(t, b - a, a) | ||
| itp_rule(::HarmonicLinear, t, a, b) = inv(muladd(t, inv(b) - inv(a), inv(a))) | ||
|
|
||
| # recursive rule application | ||
|
|
||
| itp_impl(r, t, v0, v1) = itp_impl(Dim(length(t)), r, t, v0, v1) | ||
|
|
||
| itp_impl(::Dim{D}, r, t, v0, v1) where {D} = itp_rule(r, last(t), | ||
| itp_impl(Dim(D - 1), r, front(t), v0...), | ||
| itp_impl(Dim(D - 1), r, front(t), v1...)) | ||
|
|
||
| itp_impl(::Dim{1}, r, t, v0, v1) = itp_rule(r, last(t), v0, v1) | ||
|
|
||
| # interpolation weights | ||
|
|
||
| itp_weight(ax::AbstractAxis, ::Vertex, ::Center, i) = convert(eltype(ax), 0.5) | ||
| itp_weight(ax::AbstractAxis, ::Center, ::Vertex, i) = convert(eltype(ax), 0.5) * Δ(ax, Center(), i - oneunit(i)) * iΔ(ax, Vertex(), i) | ||
|
|
||
| # special rule for efficient interpolation on uniform axes | ||
| itp_weight(ax::UniformAxis, ::Center, ::Vertex, i) = convert(eltype(ax), 0.5) | ||
|
|
||
| # dimensions to interpolate | ||
|
|
||
| @generated function itp_dims(from::NTuple{N,Location}, to::NTuple{N,Location}) where {N} | ||
| dims = Tuple(Dim(D) for D in 1:N) | ||
| locs = zip(dims, from.instance, to.instance) | ||
|
|
||
| pred = (_, l1, l2) -> l1 !== l2 | ||
| filtered_locs = Iterators.filter(splat(pred), locs) | ||
|
|
||
| dims_r, from_r, to_r = (zip(filtered_locs...)...,) | ||
|
|
||
| return :($dims_r, $from_r, $to_r) | ||
| end | ||
|
|
||
| # interpolation knots | ||
|
|
||
| @propagate_inbounds itp_knots(f, from, to, dims, I) = itp_knots_impl(f, from, to, dims, I) | ||
|
|
||
| @propagate_inbounds itp_knots_impl(f, from, to, dims, I) = (itp_knots_impl(f, front(from), front(to), front(dims), il(last(from), last(to), last(dims), I...)), | ||
| itp_knots_impl(f, front(from), front(to), front(dims), ir(last(from), last(to), last(dims), I...))) | ||
|
|
||
| @propagate_inbounds itp_knots_impl(f, from::Tuple{}, to::Tuple{}, dims::Tuple{}, I) = f[I...] | ||
|
|
||
| # interpolation interface | ||
|
|
||
| """ | ||
| itp(f, to, r, grid, I...) | ||
| Interpolates the field `f` from its current location to the specified location(s) `to` using the given interpolation rule `r`. | ||
| The indices specify the position within the grid at location(s) `to`. | ||
| """ | ||
| @add_cartesian function itp(f::AbstractField, to::NTuple{N,Location}, r::InterpolationRule, grid::StructuredGrid{N}, I::Vararg{Integer,N}) where {N} | ||
| from = location(f) | ||
| if typeof(from) === typeof(to) | ||
| return f[I...] | ||
| end | ||
|
|
||
| dims_r, from_r, to_r = itp_dims(from, to) | ||
|
|
||
| t = ntuple(Val(length(dims_r))) do D | ||
| ax = axis(grid, dims_r[D]) | ||
| itp_weight(ax, from_r[D], to_r[D], I[D]) | ||
| end | ||
|
|
||
| v = itp_knots(f, from_r, to_r, dims_r, I) | ||
|
|
||
| return itp_impl(r, t, v...) | ||
| end | ||
|
|
||
| # version for repeated locations | ||
| @add_cartesian itp(f, to::Location, r, grid, I::Vararg{Integer,N}) where {N} = itp(f, expand_loc(Val(N), to), r, grid, I...) | ||
|
|
||
| # shortcuts for common use cases | ||
|
|
||
| """ | ||
| lerp(f, to, grid, I...) | ||
| Linearly interpolate values of a field `f` to location `to`. | ||
| """ | ||
| @add_cartesian lerp(f, to, grid, I::Vararg{Integer,N}) where {N} = itp(f, to, Linear(), grid, I...) | ||
|
|
||
| """ | ||
| hlerp(f, to, grid, I...) | ||
| Interpolate a field `f` to location to using harmonic linear interpolation rule. | ||
| `rule(t, v0, v1) = 1/(1/v0 + t * (1/v1 - 1/v0))` | ||
| """ | ||
| @add_cartesian hlerp(f, to, grid, I::Vararg{Integer,N}) where {N} = itp(f, to, HarmonicLinear(), grid, I...) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,5 @@ | ||
| # finite difference | ||
| @add_cartesian δ(f, loc, from, dim, I::Vararg{Integer,N}) where {N} = right(f, loc, from, dim, I...) - left(f, loc, from, dim, I...) | ||
|
|
||
| # partial derivative | ||
| @add_cartesian ∂(f, loc, from, grid, dim, I::Vararg{Integer,N}) where {N} = δ(f, loc, from, dim, I...) * iΔ(grid, loc, dim, I...) |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Oops, something went wrong.