Mercurial > repos > public > sbplib_julia
annotate src/Grids/tensor_grid.jl @ 1257:198ccda331a6 refactor/grids
Remove range dim as a type paratmeter on Grid as it is already encoded in T if available
author | Jonatan Werpers <jonatan@werpers.com> |
---|---|
date | Thu, 23 Feb 2023 11:38:19 +0100 |
parents | 3fc78ad26d03 |
children | a4ddae8b5d49 |
rev | line source |
---|---|
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
|
1 struct TensorGrid{T,D,GT<:NTuple{N,Grid} where N} <: Grid{T,D} |
1222 | 2 grids::GT |
3 | |
4 function TensorGrid(gs...) | |
5 T = eltype(gs[1]) # All gs should have the same T | |
6 D = sum(ndims,gs) | |
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 return new{T,D,typeof(gs)}(gs) |
1222 | 9 end |
10 end | |
11 | |
1256
3fc78ad26d03
Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents:
1251
diff
changeset
|
12 # Indexing interface |
3fc78ad26d03
Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents:
1251
diff
changeset
|
13 # TODO |
3fc78ad26d03
Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents:
1251
diff
changeset
|
14 # Iteration interface |
3fc78ad26d03
Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents:
1251
diff
changeset
|
15 # TODO |
3fc78ad26d03
Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents:
1251
diff
changeset
|
16 |
3fc78ad26d03
Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents:
1251
diff
changeset
|
17 |
1222 | 18 function Base.size(g::TensorGrid) |
1251
6f75f2d2bf5c
Qualify package when using split_tuple and concatenate_tuple
Jonatan Werpers <jonatan@werpers.com>
parents:
1236
diff
changeset
|
19 return LazyTensors.concatenate_tuples(size.(g.grids)...) |
1222 | 20 end |
21 | |
22 function Base.getindex(g::TensorGrid, I...) | |
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 IndexStyle(::TensorGrid) = IndexCartesian() | |
32 | |
33 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
|
34 szs = LazyTensors.concatenate_tuples(size.(g.grids)...) |
1222 | 35 return CartesianIndices(szs) |
36 end | |
37 | |
38 | |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
39 struct TensorBoundary{N, BID<:BoundaryIdentifier} <: BoundaryIdentifier end |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
40 grid_id(::TensorBoundary{N, BID}) where {N, BID} = N |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
41 boundary_id(::TensorBoundary{N, BID}) where {N, BID} = BID() |
1222 | 42 |
43 | |
44 """ | |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
45 boundary_identifiers(::TensorGrid) |
1222 | 46 |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
47 Returns a tuple containing the boundary identifiers for the grid. |
1222 | 48 """ |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
49 function boundary_identifiers(g::TensorGrid) |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
50 n = length(g.grids) |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
51 per_grid = map(eachindex(g.grids)) do i |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
52 return map(bid -> TensorBoundary{i, bid}(), boundary_identifiers(g.grids[i])) |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
53 end |
1251
6f75f2d2bf5c
Qualify package when using split_tuple and concatenate_tuple
Jonatan Werpers <jonatan@werpers.com>
parents:
1236
diff
changeset
|
54 return LazyTensors.concatenate_tuples(per_grid...) |
1222 | 55 end |
56 | |
57 | |
58 """ | |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
59 boundary_grid(grid::TensorGrid, id::TensorBoundary) |
1222 | 60 |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
61 The grid for the boundary specified by `id`. |
1222 | 62 """ |
1236
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
63 function boundary_grid(g::TensorGrid, bid::TensorBoundary) |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
64 local_boundary_grid = boundary_grid(g.grids[grid_id(bid)], boundary_id(bid)) |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
65 new_grids = Base.setindex(g.grids, local_boundary_grid, grid_id(bid)) |
95e294576c2a
Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents:
1222
diff
changeset
|
66 return TensorGrid(new_grids...) |
1222 | 67 end |