update grid docs (#24)

JuliaPlasma · Mar 22, 2024 · 8bad081 · 8bad081
1 parent 0256e0c
commit 8bad081
Show file tree

Hide file tree

Showing 18 changed files with 94 additions and 131 deletions.
diff --git a/benchmark/benchmarks.jl b/benchmark/benchmarks.jl
@@ -9,10 +9,10 @@ upper_bounds = ntuple(x -> 1.0, dimension)
 num_cells = ntuple(x -> 16, dimension)
 periodic = ntuple(x -> true, dimension)
 g = UniformCartesianGrid(lower_bounds, upper_bounds, num_cells, periodic)
-rho = Field(g, ParticleInCell.node, 1)
-phi = Field(g, ParticleInCell.node, 1)
-Eedge = Field(g, ParticleInCell.edge, 2)
-Enode = Field(g, ParticleInCell.node, 2)
+rho = Field(g, NodeOffset(), 1)
+phi = Field(g, NodeOffset(), 1)
+Eedge = Field(g, EdgeOffset(), 2)
+Enode = Field(g, NodeOffset(), 2)
 
 # Create a species
 n_particles = 100

diff --git a/logo/logo.jl b/logo/logo.jl
@@ -23,10 +23,10 @@ grid = UniformCartesianGrid((0.0,), (sim_length,), (num_cells,), (periodic,));
 
 field_dimension = 1
 lower_guard_cells = 10
-rho = Field(grid, ParticleInCell.node, field_dimension, lower_guard_cells)
-phi = Field(grid, ParticleInCell.node, field_dimension, lower_guard_cells)
-Eedge = Field(grid, ParticleInCell.edge, field_dimension, lower_guard_cells)
-Enode = Field(grid, ParticleInCell.node, field_dimension, lower_guard_cells);
+rho = Field(grid, NodeOffset(), field_dimension, lower_guard_cells)
+phi = Field(grid, NodeOffset(), field_dimension, lower_guard_cells)
+Eedge = Field(grid, NodeOffset(), field_dimension, lower_guard_cells)
+Enode = Field(grid, NodeOffset(), field_dimension, lower_guard_cells);
 
 num_particles_per_stream = div(num_cells * particles_per_cell, 2)
 particles_per_macro = nom_density * sim_length / num_particles_per_stream

diff --git a/src/ParticleInCell.jl b/src/ParticleInCell.jl
@@ -16,7 +16,7 @@ using LinearSolve
 import Base: eachindex
 
 include("grids/grid.jl")
-export AbstractGrid, UniformCartesianGrid, node, edge, face
+export AbstractGrid, UniformCartesianGrid
 
 include("species/species.jl")
 export AbstractSpecies,
@@ -35,7 +35,7 @@ export AbstractSpecies,
     electrons
 
 include("field.jl")
-export Field, num_elements
+export Field, NodeOffset, EdgeOffset, FaceOffset, CenterOffset, num_elements
 
 abstract type AbstractSimulationStep end
 function step!(::T) where {T<:AbstractSimulationStep}

diff --git a/src/communicate/communicate_field_test.jl b/src/communicate/communicate_field_test.jl
@@ -1,7 +1,7 @@
 @testitem "CommunicateGuardCells" tags = [:unit] begin
     @testset "1D" begin
         grid = UniformCartesianGrid((0.0,), (1.0,), (10,), (false,))
-        field = Field(grid, ParticleInCell.node, 1)
+        field = Field(grid, NodeOffset(), 1)
 
         step = CommunicateGuardCells(field)
 

diff --git a/src/field.jl b/src/field.jl
@@ -1,35 +1,34 @@
-abstract type AbstractField{T,N,D,G} end
+abstract type AbstractField{T,N,O,D,G} end
 
-# TODO: Consider making these singleton types to help with dispatch
-@enum GridOffset begin
-    node
-    edge
-    face
-end
+abstract type GridOffset end
+struct NodeOffset <: GridOffset end
+struct EdgeOffset <: GridOffset end
+struct FaceOffset <: GridOffset end
+struct CenterOffset <: GridOffset end
 
 """
     Field(grid, offset, num_elements,[ lower_guard_cells = 0,
         [upper_guard_cells = lower_guard_cells + 1]])
 
 Stores the values of a field.
 """
-struct Field{T,N,D,D1,G} <: AbstractField{T,N,D,G}
+struct Field{T,N,O,D,D1,G} <: AbstractField{T,N,O,D,G}
     values::Array{T,D1}
     grid::G
-    offset::GridOffset
+    offset::O
 
     lower_guard_cells::Int
     upper_guard_cells::Int
     index_offset::NTuple{D,Int}
 
     function Field(
         grid::G,
-        offset::GridOffset,
+        offset::O,
         num_elements::Int,
         lower_guard_cells = 0,
         upper_guard_cells = lower_guard_cells + 1,
         T = Float64,
-    ) where {D,G<:AbstractGrid{D}}
+    ) where {D,O<:GridOffset,G<:AbstractGrid{D}}
 
         index_offset = ntuple(_ -> lower_guard_cells + 1, D)
         values = zeros(
@@ -38,7 +37,7 @@ struct Field{T,N,D,D1,G} <: AbstractField{T,N,D,G}
             num_elements,
         )
 
-        new{T,num_elements,D,D + 1,G}(
+        new{T,num_elements,O,D,D + 1,G}(
             values,
             grid,
             offset,
@@ -51,10 +50,10 @@ end
 
 num_elements(f::Field{T,N}) where {T,N} = N
 
-function Base.show(io::IO, f::Field{T,N,D,D1,G}) where {T,N,D,D1,G}
+function Base.show(io::IO, f::Field{T,N,O,D,D1,G}) where {T,N,O,D,D1,G}
     print(
         io,
-        "Field(grid=$(f.grid), offset=$(f.offset),
+        "Field(grid=$(f.grid), offset=$(O),
 num_elements=$(N), T=$T",
     )
 end
@@ -77,35 +76,37 @@ end
 @inline Base.firstindex(f::Field) = firstindex(f.values)
 
 @inline function cell_index_to_cell_coords(f::Field, I)
-    # align_offset = f.offset == edge ? 0.5 : 0.0
-    # return Tuple(I) .- f.index_offset .+ align_offset
     return Tuple(I) .- f.index_offset
 end
 
-@inline function cell_index_to_cell_coords(f::Field{T,N}, I, dim) where {T,N}
-    node_cell_coords = cell_index_to_cell_coords(f, I)
+@inline function cell_index_to_cell_coords(f::Field{T,N,NodeOffset}, I, dim) where {T,N}
+    return cell_index_to_cell_coords(f, I)
+end
 
-    if f.offset == node
-        return node_cell_coords
-    elseif f.offset == edge
-        return node_cell_coords .+ 0.5 .* unit_vec(dim, Val(N))
-    elseif f.offset == face
-        return node_cell_coords .+ 0.5 .* orth_vec(dim, Val(N))
-    end
+@inline function cell_index_to_cell_coords(f::Field{T,N,EdgeOffset}, I, dim) where {T,N}
+    return cell_index_to_cell_coords(f, I) .+ 0.5 .* unit_vec(dim, Val(N))
+end
+
+@inline function cell_index_to_cell_coords(f::Field{T,N,FaceOffset}, I, dim) where {T,N}
+    return cell_index_to_cell_coords(f, I) .+ 0.5 .* orth_vec(dim, Val(N))
+end
+
+@inline function cell_index_to_cell_coords(f::Field{T,N,CenterOffset}, I, dim) where {T,N}
+    return cell_index_to_cell_coords(f, I) .+ 0.5
 end
 
 @inline function cell_coords_to_cell_index(f::Field, idxs)
     return CartesianIndex(Tuple(floor.(Int, idxs) .+ f.index_offset))
 end
 
+extra_width_func(f::Field{T,N,NodeOffset}) where {T,N} = 0
+extra_width_func(f) = 1
+
 @inline function phys_coords_to_cell_index_ittr(f::Field, xs, width::Int)
     cell_coords = phys_coords_to_cell_coords(f.grid, xs)
     cell_index = Tuple(cell_coords_to_cell_index(f, cell_coords))
 
-    extra_width = 0
-    if f.offset == edge || f.offset == face
-        extra_width = 1
-    end
+    extra_width = extra_width_func(f)
 
     # Need to add one to lower bound to account for the fact that the particle
     # is 'in' the cell at cell_index.

diff --git a/src/field_test.jl b/src/field_test.jl
@@ -5,10 +5,10 @@
         (10, 10, 10),
         (true, true, true),
     )
-    scalar_node_field = Field(grid, ParticleInCell.node, 1, 3)
-    vector_node_field = Field(grid, ParticleInCell.node, 3)
-    vector_edge_field = Field(grid, ParticleInCell.edge, 3)
-    vector_face_field = Field(grid, ParticleInCell.face, 3)
+    scalar_node_field = Field(grid, NodeOffset(), 1, 3)
+    vector_node_field = Field(grid, NodeOffset(), 3)
+    vector_edge_field = Field(grid, EdgeOffset(), 3)
+    vector_face_field = Field(grid, FaceOffset(), 3)
 
     @test num_elements(scalar_node_field) == 1
     @test num_elements(vector_node_field) == 3

diff --git a/src/field_utils.jl b/src/field_utils.jl
@@ -1,17 +1,22 @@
-struct FiniteDifferenceToEdges{F,NT} <: AbstractSimulationStep
-    nodal_field::F
-    edge_field::F
+struct FiniteDifferenceToEdges{F1,F2,NT} <: AbstractSimulationStep
+    nodal_field::F1
+    edge_field::F2
 
     edge_lengths::NT
 
-    function FiniteDifferenceToEdges(nodal_field::F, edge_field::F) where {F}
-        @assert nodal_field.offset == node
-        @assert edge_field.offset == edge
+    function FiniteDifferenceToEdges(
+        nodal_field::Field{T1,N1,NodeOffset},
+        edge_field::Field{T2,N2,EdgeOffset},
+    ) where {T1,T2,N1,N2}
         @assert edge_field.upper_guard_cells >= 1
 
         edge_lengths = cell_lengths(nodal_field.grid)
 
-        new{F,typeof(edge_lengths)}(nodal_field, edge_field, edge_lengths)
+        new{typeof(nodal_field),typeof(edge_field),typeof(edge_lengths)}(
+            nodal_field,
+            edge_field,
+            edge_lengths,
+        )
     end
 end
 
@@ -31,17 +36,18 @@ function step!(step::FiniteDifferenceToEdges{F}) where {T,D,F<:AbstractField{T,D
     end
 end
 
-struct AverageEdgesToNodes{F} <: AbstractSimulationStep
-    edge_field::F
-    nodal_field::F
+struct AverageEdgesToNodes{F1,F2} <: AbstractSimulationStep
+    edge_field::F1
+    nodal_field::F2
 
-    function AverageEdgesToNodes(edge_field::F, nodal_field::F) where {F}
-        @assert edge_field.offset == edge
-        @assert nodal_field.offset == node
+    function AverageEdgesToNodes(
+        edge_field::Field{T1,N1,EdgeOffset},
+        nodal_field::Field{T2,N2,NodeOffset},
+    ) where {T1,T2,N1,N2}
         @assert edge_field.lower_guard_cells >= 1
         @assert edge_field.upper_guard_cells >= 1
 
-        new{F}(edge_field, nodal_field)
+        new{typeof(edge_field),typeof(nodal_field)}(edge_field, nodal_field)
     end
 end
 

diff --git a/src/field_utils_test.jl b/src/field_utils_test.jl
@@ -1,8 +1,8 @@
 @testitem "FiniteDifferenceToEdges" tags = [:unit] begin
     @testset "1D" begin
         grid = UniformCartesianGrid((0.0,), (1.0,), (10,), (false,))
-        nodal_field = Field(grid, ParticleInCell.node, 1)
-        edge_field = Field(grid, ParticleInCell.edge, 1)
+        nodal_field = Field(grid, NodeOffset(), 1)
+        edge_field = Field(grid, EdgeOffset(), 1)
 
         step = FiniteDifferenceToEdges(nodal_field, edge_field)
         @test step.edge_lengths == (0.1,)
@@ -21,8 +21,8 @@ end
 @testitem "AverageEdgesToNodes" tags = [:unit] begin
     @testset "1D" begin
         grid = UniformCartesianGrid((0.0,), (1.0,), (10,), (false,))
-        nodal_field = Field(grid, ParticleInCell.node, 1, 1)
-        edge_field = Field(grid, ParticleInCell.edge, 1, 1)
+        nodal_field = Field(grid, NodeOffset(), 1, 1)
+        edge_field = Field(grid, EdgeOffset(), 1, 1)
 
         step = AverageEdgesToNodes(edge_field, nodal_field)
 

diff --git a/src/grids/grid.jl b/src/grids/grid.jl
@@ -9,7 +9,7 @@ numbering systems that can refer to a location in the simulation domain:
  2. The 'cell coordinates' of a point is the non-dimensional location of the point
     in units of cell lengths. This value can range from 0 to num_cells - 1,
     or outside this range if guard cells are included. The value will
-    typically have the type `NTuple{N, Int}`.
+    typically have the type `NTuple{N, Int}` or `NTuple{N, T}` with `T <: Real`.
  3. The 'cell index' of a point is the `CartesianIndex` that can be used to
     index into field arrays at that point. This value must strictly be
     confined to `axes(field.values)`, which, for any given dimension, will
@@ -51,12 +51,10 @@ orth_vec(i, ::Val{3}) = ntuple(j -> j == i ? 0 : 1, 3)
 orth_vec(::Any, ::Union{Val{1},Val{2}}) = error("orth_vec only makes sense in 3D")
 
 """
-    cell_coords_to_phys_coords(grid, idxs, [offset, component])
+    cell_coords_to_phys_coords(grid, idxs)
 
 Converts the cell coordinates `idxs` to a physical coordinate using the geometry
-specified in `grid`. Optionally, an offset and component can be specified to get
-the physical coordinates of a specific `edge` or `face` of the cell. The component
-argument is one-indexed.
+specified in `grid`.
 
 For more information on the different types of coordinate systems, see
 `AbstractGrid`.

diff --git a/src/grids/grid_test.jl b/src/grids/grid_test.jl
@@ -8,39 +8,10 @@
         )
 
         @test ParticleInCell.cell_coords_to_phys_coords(grid, (0, 0, 0)) == (0, 0, 0)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (0, 0, 0), node) == (0, 0, 0)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (0, 0, 0), edge, 1) ==
-              (0.05, 0, 0)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (0, 0, 0), face, 1) ==
-              (0, 0.05, 0.05)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (0, 0, 0), edge, 3) ==
-              (0, 0, 0.05)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (0, 0, 0), face, 3) ==
-              (0.05, 0.05, 0)
 
         @test ParticleInCell.cell_coords_to_phys_coords(grid, (5, 5, 5)) == (0.5, 0.5, 0.5)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (5, 5, 5), node) ==
-              (0.5, 0.5, 0.5)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (5, 5, 5), edge, 1) ==
-              (0.55, 0.5, 0.5)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (5, 5, 5), face, 1) ==
-              (0.5, 0.55, 0.55)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (5, 5, 5), edge, 3) ==
-              (0.5, 0.5, 0.55)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (5, 5, 5), face, 3) ==
-              (0.55, 0.55, 0.5)
 
         @test ParticleInCell.cell_coords_to_phys_coords(grid, (10, 10, 10)) == (1, 1, 1)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (10, 10, 10), node) ==
-              (1, 1, 1)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (10, 10, 10), edge, 1) ==
-              (1.05, 1, 1)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (10, 10, 10), face, 1) ==
-              (1, 1.05, 1.05)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (10, 10, 10), edge, 3) ==
-              (1, 1, 1.05)
-        @test ParticleInCell.cell_coords_to_phys_coords(grid, (10, 10, 10), face, 3) ==
-              (1.05, 1.05, 1)
     end
 
     @testset "phys_coords_to_cell_coords" begin

diff --git a/src/grids/uniform_cartesian_grid.jl b/src/grids/uniform_cartesian_grid.jl
@@ -19,23 +19,8 @@ end
     return (grid.upper_bounds .- grid.lower_bounds) ./ grid.num_cells
 end
 
-# TODO: inline this function once the offsets become singletons, and this
-# function becomes performant
-function cell_coords_to_phys_coords(
-    grid::UniformCartesianGrid{D},
-    idxs,
-    offset = node,
-    component::Int = 1,
-) where {D}
-    node_coords = grid.lower_bounds .+ idxs .* cell_lengths(grid)
-
-    if offset == node
-        return node_coords
-    elseif offset == edge
-        return node_coords .+ cell_lengths(grid) ./ 2 .* unit_vec(component, Val(D))
-    elseif offset == face
-        return node_coords .+ cell_lengths(grid) ./ 2 .* orth_vec(component, Val(D))
-    end
+@inline function cell_coords_to_phys_coords(grid::UniformCartesianGrid, idxs)
+    return grid.lower_bounds .+ idxs .* cell_lengths(grid)
 end
 
 @inline function phys_coords_to_cell_coords(grid::UniformCartesianGrid, xs)

diff --git a/src/interpolation.jl b/src/interpolation.jl
@@ -147,13 +147,15 @@ struct BSplineChargeInterpolation{S,F,IF} <: AbstractSimulationStep
 
     interp_func::IF
 
-    function BSplineChargeInterpolation(species::S, rho::F, bspline_degree::Int) where {S,F}
-        @assert rho.offset == node
-
+    function BSplineChargeInterpolation(
+        species::S,
+        rho::Field{T,N,NodeOffset},
+        bspline_degree::Int,
+    ) where {T,N,S}
         interp_width = div(bspline_degree, 2) + 1
         interp_func = select_bs_interp(bspline_degree)
 
-        new{S,F,typeof(interp_func)}(
+        new{S,typeof(rho),typeof(interp_func)}(
             species,
             rho,
             bspline_degree,

diff --git a/src/interpolation_test.jl b/src/interpolation_test.jl
@@ -11,8 +11,8 @@
     end
 
     g = UniformCartesianGrid((0.0,), (1.0,), (10,), (true,))
-    rho = Field(g, ParticleInCell.node, 1)
-    phi = Field(g, ParticleInCell.node, 1)
+    rho = Field(g, NodeOffset(), 1)
+    phi = Field(g, NodeOffset(), 1)
     s = VariableWeightSpecies(fill(0.5, 1, 1), fill(0.0, 1, 1), [1.0], 1.0, 1.0)
 
     bs_interp = BSplineChargeInterpolation(s, rho, 1)