Mercurial > repos > public > sbplib_julia
annotate src/Grids/tensor_grid.jl @ 1435:29adc06e5b56 feature/grids/curvilinear
Merge bugfix/grids/tensor_grid_length
author | Jonatan Werpers <jonatan@werpers.com> |
---|---|
date | Thu, 07 Sep 2023 16:49:45 +0200 |
parents | af73340a8f0e 48e16efaac7a |
children | 433245fd33c0 |
rev | line source |
---|---|
1434 | 1 # TODO: Check this file and other grids for duplicate implementation of general methods implemented for Grid |
2 | |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
3 """ |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
4 TensorGrid{T,D} <: Grid{T,D} |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
5 |
1338
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
6 A grid constructed as the tensor product of other grids. |
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
7 |
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
8 Currently only supports grids with the `HasShape`-trait. |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
9 """ |
1257
198ccda331a6
Remove range dim as a type paratmeter on Grid as it is already encoded in T if available
Jonatan Werpers <jonatan@werpers.com>
parents:
1256
diff
changeset
|
10 struct TensorGrid{T,D,GT<:NTuple{N,Grid} where N} <: Grid{T,D} |
1222 | 11 grids::GT |
12 | |
13 function TensorGrid(gs...) | |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
14 T = mapreduce(eltype, combined_coordinate_vector_type, gs) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
15 D = sum(ndims, gs) |
1222 | 16 |
1257
198ccda331a6
Remove range dim as a type paratmeter on Grid as it is already encoded in T if available
Jonatan Werpers <jonatan@werpers.com>
parents:
1256
diff
changeset
|
17 return new{T,D,typeof(gs)}(gs) |
1222 | 18 end |
19 end | |
20 | |
1256
3fc78ad26d03
Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents:
1251
diff
changeset
|
21 # Indexing interface |
1386
4d628c83987e
Add method `getindex(::Grid,::CartesianIndex)` with a default implementation
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
22 function Base.getindex(g::TensorGrid, I::Vararg{Int}) |
1222 | 23 szs = ndims.(g.grids) |
24 | |
1251
6f75f2d2bf5c
Qualify package when using split_tuple and concatenate_tuple
Jonatan Werpers <jonatan@werpers.com>
parents:
1236
diff
changeset
|
25 Is = LazyTensors.split_tuple(I, szs) |
1222 | 26 ps = map((g,I)->SVector(g[I...]), g.grids, Is) |
27 | |
28 return vcat(ps...) | |
29 end | |
30 | |
31 function Base.eachindex(g::TensorGrid) | |
1251
6f75f2d2bf5c
Qualify package when using split_tuple and concatenate_tuple
Jonatan Werpers <jonatan@werpers.com>
parents:
1236
diff
changeset
|
32 szs = LazyTensors.concatenate_tuples(size.(g.grids)...) |
1222 | 33 return CartesianIndices(szs) |
34 end | |
35 | |
1399
adbdbf1c3480
Replace `firstindex` and `lastindex` with implementation of `axes`
Jonatan Werpers <jonatan@werpers.com>
parents:
1394
diff
changeset
|
36 function Base.axes(g::TensorGrid, d) |
1390
47931bef8471
Implement `firstindex` and `lastindex` for `TensorGrid`
Jonatan Werpers <jonatan@werpers.com>
parents:
1388
diff
changeset
|
37 i, ld = grid_and_local_dim_index(ndims.(g.grids), d) |
1399
adbdbf1c3480
Replace `firstindex` and `lastindex` with implementation of `axes`
Jonatan Werpers <jonatan@werpers.com>
parents:
1394
diff
changeset
|
38 return axes(g.grids[i], ld) |
1390
47931bef8471
Implement `firstindex` and `lastindex` for `TensorGrid`
Jonatan Werpers <jonatan@werpers.com>
parents:
1388
diff
changeset
|
39 end |
47931bef8471
Implement `firstindex` and `lastindex` for `TensorGrid`
Jonatan Werpers <jonatan@werpers.com>
parents:
1388
diff
changeset
|
40 |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
41 # Iteration interface |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
42 Base.iterate(g::TensorGrid) = iterate(Iterators.product(g.grids...)) |> _iterate_combine_coords |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
43 Base.iterate(g::TensorGrid, state) = iterate(Iterators.product(g.grids...), state) |> _iterate_combine_coords |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
44 _iterate_combine_coords(::Nothing) = nothing |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
45 _iterate_combine_coords((next,state)) = combine_coordinates(next...), state |
1222 | 46 |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
47 Base.IteratorSize(::Type{<:TensorGrid{<:Any, D}}) where D = Base.HasShape{D}() |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
48 Base.eltype(::Type{<:TensorGrid{T}}) where T = T |
1423
48e16efaac7a
Fix implementation of `length(::TensorGrid)`
Jonatan Werpers <jonatan@werpers.com>
parents:
1418
diff
changeset
|
49 Base.length(g::TensorGrid) = prod(length, g.grids) |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
50 Base.size(g::TensorGrid) = LazyTensors.concatenate_tuples(size.(g.grids)...) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
51 |
1222 | 52 |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
53 refine(g::TensorGrid, r::Int) = mapreduce(g->refine(g,r), TensorGrid, g.grids) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
54 coarsen(g::TensorGrid, r::Int) = mapreduce(g->coarsen(g,r), TensorGrid, g.grids) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
55 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
56 """ |
1330
5f05a708d730
grid.l: More documentation
Jonatan Werpers <jonatan@werpers.com>
parents:
1289
diff
changeset
|
57 TensorGridBoundary{N, BID} <: BoundaryIdentifier |
1338
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
58 |
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
59 A boundary identifier for a tensor grid. `N` Specifies which grid in the |
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
60 tensor product and `BID` which boundary on that grid. |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
61 """ |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
62 struct TensorGridBoundary{N, BID} <: BoundaryIdentifier end |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
63 grid_id(::TensorGridBoundary{N, BID}) where {N, BID} = N |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
64 boundary_id(::TensorGridBoundary{N, BID}) where {N, BID} = BID() |
1222 | 65 |
66 """ | |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
67 boundary_identifiers(g::TensorGrid) |
1222 | 68 |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
69 Returns a tuple containing the boundary identifiers of `g`. |
1222 | 70 """ |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
71 function boundary_identifiers(g::TensorGrid) |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
72 per_grid = map(eachindex(g.grids)) do i |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
73 return map(bid -> TensorGridBoundary{i, typeof(bid)}(), boundary_identifiers(g.grids[i])) |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
74 end |
1251
6f75f2d2bf5c
Qualify package when using split_tuple and concatenate_tuple
Jonatan Werpers <jonatan@werpers.com>
parents:
1236
diff
changeset
|
75 return LazyTensors.concatenate_tuples(per_grid...) |
1222 | 76 end |
77 | |
78 | |
79 """ | |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
80 boundary_grid(g::TensorGrid, id::TensorGridBoundary) |
1222 | 81 |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
82 The grid for the boundary of `g` specified by `id`. |
1222 | 83 """ |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
84 function boundary_grid(g::TensorGrid, id::TensorGridBoundary) |
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
85 local_boundary_grid = boundary_grid(g.grids[grid_id(id)], boundary_id(id)) |
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
86 new_grids = Base.setindex(g.grids, local_boundary_grid, grid_id(id)) |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
87 return TensorGrid(new_grids...) |
1222 | 88 end |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
89 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
90 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
91 function combined_coordinate_vector_type(coordinate_types...) |
1270
dcbac783e4c1
Factor out functions for getting the type and number of components in a type
Jonatan Werpers <jonatan@werpers.com>
parents:
1266
diff
changeset
|
92 combined_coord_length = mapreduce(_ncomponents, +, coordinate_types) |
1289
3b7ebd135918
Remove _component_type and replace with eltype
Jonatan Werpers <jonatan@werpers.com>
parents:
1270
diff
changeset
|
93 combined_coord_type = mapreduce(eltype, promote_type, coordinate_types) |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
94 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
95 if combined_coord_length == 1 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
96 return combined_coord_type |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
97 else |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
98 return SVector{combined_coord_length, combined_coord_type} |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
99 end |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
100 end |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
101 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
102 function combine_coordinates(coords...) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
103 return mapreduce(SVector, vcat, coords) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
104 end |
1388
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
105 |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
106 """ |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
107 grid_and_local_dim_index(nds, d) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
108 |
1415
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
109 Given a tuple of number of dimensions `nds`, and a global dimension index `d`, |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
110 calculate which grid index, and local dimension, `d` corresponds to. |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
111 |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
112 `nds` would come from broadcasting `ndims` on the grids tuple of a |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
113 `TensorGrid`. If you are interested in a dimension `d` of a tensor grid `g` |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
114 ```julia |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
115 gi, ldi = grid_and_local_dim_index(ndims.(g.grids), d) |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
116 ``` |
1418
e82240df974d
Correct spelling
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1415
diff
changeset
|
117 tells you which grid it belongs to (`gi`) and which index it is at within that |
1415
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
118 grid (`ldi`). |
1388
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
119 """ |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
120 function grid_and_local_dim_index(nds, d) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
121 I = findfirst(>=(d), cumsum(nds)) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
122 |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
123 if I == 1 |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
124 return (1, d) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
125 else |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
126 return (I, d-cumsum(nds)[I-1]) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
127 end |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
128 end |