Add extra operators for fields and TensorFields (#7)

PTsolvers · Feb 12, 2024 · 189c061 · 189c061
1 parent d3a602a
commit 189c061
Show file tree

Hide file tree

Showing 4 changed files with 52 additions and 2 deletions.
diff --git a/src/Fields/Fields.jl b/src/Fields/Fields.jl
@@ -1,6 +1,6 @@
 module Fields
 
-export AbstractField, Field, VectorField
+export AbstractField, Field, VectorField, TensorField
 export ConstantField, ZeroField, OneField, ValueField
 export FunctionField
 export location, halo, interior, set!

diff --git a/src/Fields/abstract_field.jl b/src/Fields/abstract_field.jl
@@ -27,6 +27,14 @@ function Base.show(io::IO, field::AbstractField{T,N,L}) where {T,N,L}
 end
 
 # grid operators
+@propagate_inbounds @add_cartesian function GridOperators.left(f::AbstractField, dim, I::Vararg{Integer,N}) where {N}
+    GridOperators.left(f, flip(location(f, dim)), dim, I...)
+end
+
+@propagate_inbounds @add_cartesian function GridOperators.right(f::AbstractField, dim, I::Vararg{Integer,N}) where {N}
+    GridOperators.right(f, flip(location(f, dim)), dim, I...)
+end
+
 @propagate_inbounds @add_cartesian function GridOperators.δ(f::AbstractField, dim, I::Vararg{Integer,N}) where {N}
     GridOperators.δ(f, flip(location(f, dim)), dim, I...)
 end
@@ -37,10 +45,20 @@ end
 
 # 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
+        @propagate_inbounds @add_cartesian function GridOperators.$left(f::AbstractField, I::Vararg{Integer,N}) where {N}
+            GridOperators.left(f, flip(location(f, Val($dim))), Val($dim), I...)
+        end
+
+        @propagate_inbounds @add_cartesian function GridOperators.$right(f::AbstractField, I::Vararg{Integer,N}) where {N}
+            GridOperators.right(f, flip(location(f, Val($dim))), Val($dim), I...)
+        end
+
         @propagate_inbounds @add_cartesian function GridOperators.$δ(f::AbstractField, I::Vararg{Integer,N}) where {N}
             GridOperators.δ(f, flip(location(f, Val($dim))), Val($dim), I...)
         end

diff --git a/src/Fields/field.jl b/src/Fields/field.jl
@@ -101,3 +101,19 @@ function VectorField(backend::Backend, grid::StructuredGrid{N}, args...; kwargs.
     end
     return NamedTuple{names}(values)
 end
+
+# tensor fields
+TensorField(backend::Backend, grid::StructuredGrid{2}, args...; kwargs...) = (
+    xx = Field(backend, grid, Center(), args...; kwargs...),
+    yy = Field(backend, grid, Center(), args...; kwargs...),
+    xy = Field(backend, grid, Vertex(), args...; kwargs...)
+)
+
+TensorField(backend::Backend, grid::StructuredGrid{3}, args...; kwargs...) = (
+    xx = Field(backend, grid, Center(), args...; kwargs...),
+    yy = Field(backend, grid, Center(), args...; kwargs...),
+    zz = Field(backend, grid, Center(), args...; kwargs...),
+    xy = Field(backend, grid, (Vertex(), Vertex(), Center()), args...; kwargs...),
+    xz = Field(backend, grid, (Vertex(), Center(), Vertex()), args...; kwargs...),
+    yz = Field(backend, grid, (Center(), Vertex(), Vertex()), args...; kwargs...)
+)
diff --git a/src/GridOperators/GridOperators.jl b/src/GridOperators/GridOperators.jl
@@ -34,11 +34,27 @@ end
 
 # 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 $δ, $∂
+        export $δ, $∂, $left, $right
+
+        """
+            $($left)(f, loc, I)
+
+        "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...)
+
+        """
+            $($right)(f, loc, I)
+
+        "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...)
 
         """
             $($δ)(f, loc, I)