Mercurial > repos > public > sbplib_julia
diff src/Grids/tensor_grid.jl @ 1395:bdcdbd4ea9cd feature/boundary_conditions
Merge with default. Comment out broken tests for boundary_conditions at sat
author | Vidar Stiernström <vidar.stiernstrom@it.uu.se> |
---|---|
date | Wed, 26 Jul 2023 21:35:50 +0200 |
parents | 42ecd4b3e215 |
children | 4d628c83987e c0208286234e |
line wrap: on
line diff
--- /dev/null Thu Jan 01 00:00:00 1970 +0000 +++ b/src/Grids/tensor_grid.jl Wed Jul 26 21:35:50 2023 +0200 @@ -0,0 +1,99 @@ +""" + TensorGrid{T,D} <: Grid{T,D} + +A grid constructed as the tensor product of other grids. + +Currently only supports grids with the `HasShape`-trait. +""" +struct TensorGrid{T,D,GT<:NTuple{N,Grid} where N} <: Grid{T,D} + grids::GT + + function TensorGrid(gs...) + T = mapreduce(eltype, combined_coordinate_vector_type, gs) + D = sum(ndims, gs) + + return new{T,D,typeof(gs)}(gs) + end +end + +# Indexing interface +function Base.getindex(g::TensorGrid, I...) + szs = ndims.(g.grids) + + Is = LazyTensors.split_tuple(I, szs) + ps = map((g,I)->SVector(g[I...]), g.grids, Is) + + return vcat(ps...) +end + +Base.getindex(g::TensorGrid, I::CartesianIndex) = g[Tuple(I)...] + +function Base.eachindex(g::TensorGrid) + szs = LazyTensors.concatenate_tuples(size.(g.grids)...) + return CartesianIndices(szs) +end + +# Iteration interface +Base.iterate(g::TensorGrid) = iterate(Iterators.product(g.grids...)) |> _iterate_combine_coords +Base.iterate(g::TensorGrid, state) = iterate(Iterators.product(g.grids...), state) |> _iterate_combine_coords +_iterate_combine_coords(::Nothing) = nothing +_iterate_combine_coords((next,state)) = combine_coordinates(next...), state + +Base.IteratorSize(::Type{<:TensorGrid{<:Any, D}}) where D = Base.HasShape{D}() +Base.eltype(::Type{<:TensorGrid{T}}) where T = T +Base.length(g::TensorGrid) = sum(length, g.grids) +Base.size(g::TensorGrid) = LazyTensors.concatenate_tuples(size.(g.grids)...) + + +refine(g::TensorGrid, r::Int) = mapreduce(g->refine(g,r), TensorGrid, g.grids) +coarsen(g::TensorGrid, r::Int) = mapreduce(g->coarsen(g,r), TensorGrid, g.grids) + +""" + TensorGridBoundary{N, BID} <: BoundaryIdentifier + +A boundary identifier for a tensor grid. `N` Specifies which grid in the +tensor product and `BID` which boundary on that grid. +""" +struct TensorGridBoundary{N, BID} <: BoundaryIdentifier end +grid_id(::TensorGridBoundary{N, BID}) where {N, BID} = N +boundary_id(::TensorGridBoundary{N, BID}) where {N, BID} = BID() + +""" + boundary_identifiers(g::TensorGrid) + +Returns a tuple containing the boundary identifiers of `g`. +""" +function boundary_identifiers(g::TensorGrid) + per_grid = map(eachindex(g.grids)) do i + return map(bid -> TensorGridBoundary{i, typeof(bid)}(), boundary_identifiers(g.grids[i])) + end + return LazyTensors.concatenate_tuples(per_grid...) +end + + +""" + boundary_grid(g::TensorGrid, id::TensorGridBoundary) + +The grid for the boundary of `g` specified by `id`. +""" +function boundary_grid(g::TensorGrid, id::TensorGridBoundary) + local_boundary_grid = boundary_grid(g.grids[grid_id(id)], boundary_id(id)) + new_grids = Base.setindex(g.grids, local_boundary_grid, grid_id(id)) + return TensorGrid(new_grids...) +end + + +function combined_coordinate_vector_type(coordinate_types...) + combined_coord_length = mapreduce(_ncomponents, +, coordinate_types) + combined_coord_type = mapreduce(eltype, promote_type, coordinate_types) + + if combined_coord_length == 1 + return combined_coord_type + else + return SVector{combined_coord_length, combined_coord_type} + end +end + +function combine_coordinates(coords...) + return mapreduce(SVector, vcat, coords) +end