annotate src/Grids/tensor_grid.jl @ 1256:3fc78ad26d03 refactor/grids

Add notes and todos about interface implementations for grids
author Jonatan Werpers <jonatan@werpers.com>
date Wed, 22 Feb 2023 22:38:54 +0100
parents 6f75f2d2bf5c
children 198ccda331a6
Ignore whitespace changes - Everywhere: Within whitespace: At end of lines:
rev   line source
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
1 struct TensorGrid{T,D,RD,GT<:NTuple{N,Grid} where N} <: Grid{T,D,RD}
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
2 grids::GT
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
3
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
4 function TensorGrid(gs...)
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
5 T = eltype(gs[1]) # All gs should have the same T
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
6 D = sum(ndims,gs)
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
7 RD = sum(nrangedims, gs)
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
8
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
9 return new{T,D,RD,typeof(gs)}(gs)
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
10 end
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
11 end
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
12
1256
3fc78ad26d03 Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents: 1251
diff changeset
13 # Indexing interface
3fc78ad26d03 Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents: 1251
diff changeset
14 # TODO
3fc78ad26d03 Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents: 1251
diff changeset
15 # Iteration interface
3fc78ad26d03 Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents: 1251
diff changeset
16 # TODO
3fc78ad26d03 Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents: 1251
diff changeset
17
3fc78ad26d03 Add notes and todos about interface implementations for grids
Jonatan Werpers <jonatan@werpers.com>
parents: 1251
diff changeset
18
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
19 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
20 return LazyTensors.concatenate_tuples(size.(g.grids)...)
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
21 end
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
22
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
23 function Base.getindex(g::TensorGrid, I...)
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
24 szs = ndims.(g.grids)
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
25
1251
6f75f2d2bf5c Qualify package when using split_tuple and concatenate_tuple
Jonatan Werpers <jonatan@werpers.com>
parents: 1236
diff changeset
26 Is = LazyTensors.split_tuple(I, szs)
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
27 ps = map((g,I)->SVector(g[I...]), g.grids, Is)
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
28
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
29 return vcat(ps...)
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
30 end
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
31
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
32 IndexStyle(::TensorGrid) = IndexCartesian()
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
33
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
34 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
35 szs = LazyTensors.concatenate_tuples(size.(g.grids)...)
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
36 return CartesianIndices(szs)
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
37 end
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
38
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
39
1236
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
40 struct TensorBoundary{N, BID<:BoundaryIdentifier} <: BoundaryIdentifier end
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
41 grid_id(::TensorBoundary{N, BID}) where {N, BID} = N
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
42 boundary_id(::TensorBoundary{N, BID}) where {N, BID} = BID()
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
43
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
44
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
45 """
1236
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
46 boundary_identifiers(::TensorGrid)
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
47
1236
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
48 Returns a tuple containing the boundary identifiers for the grid.
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
49 """
1236
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
50 function boundary_identifiers(g::TensorGrid)
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
51 n = length(g.grids)
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
52 per_grid = map(eachindex(g.grids)) do i
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
53 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
54 end
1251
6f75f2d2bf5c Qualify package when using split_tuple and concatenate_tuple
Jonatan Werpers <jonatan@werpers.com>
parents: 1236
diff changeset
55 return LazyTensors.concatenate_tuples(per_grid...)
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
56 end
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
57
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
58
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
59 """
1236
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
60 boundary_grid(grid::TensorGrid, id::TensorBoundary)
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
61
1236
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
62 The grid for the boundary specified by `id`.
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
63 """
1236
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
64 function boundary_grid(g::TensorGrid, bid::TensorBoundary)
95e294576c2a Implement boundary methods for TensorGrid
Jonatan Werpers <jonatan@werpers.com>
parents: 1222
diff changeset
65 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
66 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
67 return TensorGrid(new_grids...)
1222
5f677cd6f0b6 Start refactoring
Jonatan Werpers <jonatan@werpers.com>
parents:
diff changeset
68 end