Mercurial > repos > public > sbplib_julia
annotate src/Grids/tensor_grid.jl @ 1415:18e21601da2d bugfix/grids/complete_interface_impl
Attempt to clearify docs for `grid_and_local_dim_index`
author | Jonatan Werpers <jonatan@werpers.com> |
---|---|
date | Wed, 23 Aug 2023 15:35:01 +0200 |
parents | adbdbf1c3480 |
children | e82240df974d |
rev | line source |
---|---|
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
1 """ |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
2 TensorGrid{T,D} <: Grid{T,D} |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
3 |
1338
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
4 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
|
5 |
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
6 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
|
7 """ |
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
|
8 struct TensorGrid{T,D,GT<:NTuple{N,Grid} where N} <: Grid{T,D} |
1222 | 9 grids::GT |
10 | |
11 function TensorGrid(gs...) | |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
12 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
|
13 D = sum(ndims, gs) |
1222 | 14 |
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
|
15 return new{T,D,typeof(gs)}(gs) |
1222 | 16 end |
17 end | |
18 | |
1256
3fc78ad26d03
Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents:
1251
diff
changeset
|
19 # Indexing interface |
1386
4d628c83987e
Add method `getindex(::Grid,::CartesianIndex)` with a default implementation
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
20 function Base.getindex(g::TensorGrid, I::Vararg{Int}) |
1222 | 21 szs = ndims.(g.grids) |
22 | |
1251
6f75f2d2bf5c
Qualify package when using split_tuple and concatenate_tuple
Jonatan Werpers <jonatan@werpers.com>
parents:
1236
diff
changeset
|
23 Is = LazyTensors.split_tuple(I, szs) |
1222 | 24 ps = map((g,I)->SVector(g[I...]), g.grids, Is) |
25 | |
26 return vcat(ps...) | |
27 end | |
28 | |
29 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
|
30 szs = LazyTensors.concatenate_tuples(size.(g.grids)...) |
1222 | 31 return CartesianIndices(szs) |
32 end | |
33 | |
1399
adbdbf1c3480
Replace `firstindex` and `lastindex` with implementation of `axes`
Jonatan Werpers <jonatan@werpers.com>
parents:
1394
diff
changeset
|
34 function Base.axes(g::TensorGrid, d) |
1390
47931bef8471
Implement `firstindex` and `lastindex` for `TensorGrid`
Jonatan Werpers <jonatan@werpers.com>
parents:
1388
diff
changeset
|
35 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
|
36 return axes(g.grids[i], ld) |
1390
47931bef8471
Implement `firstindex` and `lastindex` for `TensorGrid`
Jonatan Werpers <jonatan@werpers.com>
parents:
1388
diff
changeset
|
37 end |
47931bef8471
Implement `firstindex` and `lastindex` for `TensorGrid`
Jonatan Werpers <jonatan@werpers.com>
parents:
1388
diff
changeset
|
38 |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
39 # Iteration interface |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
40 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
|
41 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
|
42 _iterate_combine_coords(::Nothing) = nothing |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
43 _iterate_combine_coords((next,state)) = combine_coordinates(next...), state |
1222 | 44 |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
45 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
|
46 Base.eltype(::Type{<:TensorGrid{T}}) where T = T |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
47 Base.length(g::TensorGrid) = sum(length, g.grids) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
48 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
|
49 |
1222 | 50 |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
51 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
|
52 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
|
53 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
54 """ |
1330
5f05a708d730
grid.l: More documentation
Jonatan Werpers <jonatan@werpers.com>
parents:
1289
diff
changeset
|
55 TensorGridBoundary{N, BID} <: BoundaryIdentifier |
1338
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
56 |
5604676d8426
Docs in tensor_grid.jl
Jonatan Werpers <jonatan@werpers.com>
parents:
1337
diff
changeset
|
57 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
|
58 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
|
59 """ |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
60 struct TensorGridBoundary{N, BID} <: BoundaryIdentifier end |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
61 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
|
62 boundary_id(::TensorGridBoundary{N, BID}) where {N, BID} = BID() |
1222 | 63 |
64 """ | |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
65 boundary_identifiers(g::TensorGrid) |
1222 | 66 |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
67 Returns a tuple containing the boundary identifiers of `g`. |
1222 | 68 """ |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
69 function boundary_identifiers(g::TensorGrid) |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
70 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
|
71 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
|
72 end |
1251
6f75f2d2bf5c
Qualify package when using split_tuple and concatenate_tuple
Jonatan Werpers <jonatan@werpers.com>
parents:
1236
diff
changeset
|
73 return LazyTensors.concatenate_tuples(per_grid...) |
1222 | 74 end |
75 | |
76 | |
77 """ | |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
78 boundary_grid(g::TensorGrid, id::TensorGridBoundary) |
1222 | 79 |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
80 The grid for the boundary of `g` specified by `id`. |
1222 | 81 """ |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
82 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
|
83 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
|
84 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
|
85 return TensorGrid(new_grids...) |
1222 | 86 end |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
87 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
88 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
89 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
|
90 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
|
91 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
|
92 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
93 if combined_coord_length == 1 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
94 return combined_coord_type |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
95 else |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
96 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
|
97 end |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
98 end |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
99 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
100 function combine_coordinates(coords...) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
101 return mapreduce(SVector, vcat, coords) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
102 end |
1388
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
103 |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
104 """ |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
105 grid_and_local_dim_index(nds, d) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
106 |
1415
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
107 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
|
108 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
|
109 |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
110 `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
|
111 `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
|
112 ```julia |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
113 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
|
114 ``` |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
115 tells you which grid it belongs to (`gi`) and wich index it is at within that |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
116 grid (`ldi`). |
1388
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
117 """ |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
118 function grid_and_local_dim_index(nds, d) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
119 I = findfirst(>=(d), cumsum(nds)) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
120 |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
121 if I == 1 |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
122 return (1, d) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
123 else |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
124 return (I, d-cumsum(nds)[I-1]) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
125 end |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
126 # TBD: Is there a cleaner way to compute this? |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
127 end |