Mercurial > repos > public > sbplib_julia
annotate src/Grids/tensor_grid.jl @ 1572:157c43966b06 feature/grids/manifolds
Add some tests and implement parameterspace for ConcreteChart
author | Jonatan Werpers <jonatan@werpers.com> |
---|---|
date | Thu, 25 Apr 2024 14:13:41 +0200 |
parents | d9d9ab18cdfc |
children | b02917bcd7d5 |
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}() |
1423
48e16efaac7a
Fix implementation of `length(::TensorGrid)`
Jonatan Werpers <jonatan@werpers.com>
parents:
1418
diff
changeset
|
48 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
|
49 Base.size(g::TensorGrid) = LazyTensors.concatenate_tuples(size.(g.grids)...) |
1437
e3a80ef08d09
Implement size(g,d) for EquidistantGrid and TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1418
diff
changeset
|
50 Base.size(g::TensorGrid, d) = size(g)[d] |
1266
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 """ | |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
79 boundary_grid(g::TensorGrid, id::TensorGridBoundary) |
1222 | 80 |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
81 The grid for the boundary of `g` specified by `id`. |
1222 | 82 """ |
1347
08f06bfacd5c
Fix typos and formatting of documentation
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1338
diff
changeset
|
83 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
|
84 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
|
85 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
|
86 return TensorGrid(new_grids...) |
1222 | 87 end |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
88 |
1465
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
89 function boundary_indices(g::TensorGrid, id::TensorGridBoundary) |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
90 per_grid_ind = map(g.grids) do g |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
91 ntuple(i->:, ndims(g)) |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
92 end |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
93 |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
94 local_b_ind = boundary_indices(g.grids[grid_id(id)], boundary_id(id)) |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
95 b_ind = Base.setindex(per_grid_ind, local_b_ind, grid_id(id)) |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
96 |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
97 return LazyTensors.concatenate_tuples(b_ind...) |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
98 end |
88e738d807cb
Implement boundary_indices for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1441
diff
changeset
|
99 |
1266
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
100 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
|
101 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
|
102 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
|
103 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
104 if combined_coord_length == 1 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
105 return combined_coord_type |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
106 else |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
107 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
|
108 end |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
109 end |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
110 |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
111 function combine_coordinates(coords...) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
112 return mapreduce(SVector, vcat, coords) |
a4ddae8b5d49
Add tests for TensorGrid and make them pass
Jonatan Werpers <jonatan@werpers.com>
parents:
1257
diff
changeset
|
113 end |
1388
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
114 |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
115 """ |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
116 grid_and_local_dim_index(nds, d) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
117 |
1415
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
118 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
|
119 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
|
120 |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
121 `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
|
122 `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
|
123 ```julia |
18e21601da2d
Attempt to clearify docs for `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1399
diff
changeset
|
124 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
|
125 ``` |
1418
e82240df974d
Correct spelling
Vidar Stiernström <vidar.stiernstrom@it.uu.se>
parents:
1415
diff
changeset
|
126 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
|
127 grid (`ldi`). |
1388
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
128 """ |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
129 function grid_and_local_dim_index(nds, d) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
130 I = findfirst(>=(d), cumsum(nds)) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
131 |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
132 if I == 1 |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
133 return (1, d) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
134 else |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
135 return (I, d-cumsum(nds)[I-1]) |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
136 end |
c0208286234e
Add `grid_and_local_dim_index`
Jonatan Werpers <jonatan@werpers.com>
parents:
1349
diff
changeset
|
137 end |