changeset 1022:bbbc31953367 refactor/sbpoperators/inflation

Introduce an inflate function in lazy tensors and use it in volume_operator and boundary_operator
author Jonatan Werpers <jonatan@werpers.com>
date Fri, 18 Mar 2022 16:57:00 +0100
parents 83046af6143a
children 52f07c77299d
files src/LazyTensors/lazy_tensor_operations.jl src/SbpOperators/boundaryops/boundary_operator.jl src/SbpOperators/volumeops/volume_operator.jl test/LazyTensors/lazy_tensor_operations_test.jl
diffstat 4 files changed, 41 insertions(+), 19 deletions(-) [+]
line wrap: on
line diff
--- a/src/LazyTensors/lazy_tensor_operations.jl	Wed Mar 16 18:39:00 2022 +0100
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Fri Mar 18 16:57:00 2022 +0100
@@ -413,6 +413,19 @@
 ⊗(a::TensorMapping, b::TensorMapping) = LazyOuterProduct(a,b)
 
 
+"""
+    inflate(tm, sz, dir)
+
+Inflate `tm` with identity tensors in all directions `d` for `d != dir`.
+
+# TODO: Describe when it is useful
+"""
+function inflate(tm::TensorMapping, sz, dir)
+    Is = IdentityMapping{eltype(tm)}.(sz)
+    parts = Base.setindex(Is, tm, dir)
+    return foldl(⊗, parts)
+end
+
 function check_domain_size(tm::TensorMapping, sz)
     if domain_size(tm) != sz
         throw(SizeMismatch(tm,sz))
--- a/src/SbpOperators/boundaryops/boundary_operator.jl	Wed Mar 16 18:39:00 2022 +0100
+++ b/src/SbpOperators/boundaryops/boundary_operator.jl	Fri Mar 18 16:57:00 2022 +0100
@@ -12,19 +12,10 @@
 function boundary_operator(grid::EquidistantGrid, closure_stencil, boundary::CartesianBoundary)
     #TODO:Check that dim(boundary) <= Dim?
 
-    # Create 1D boundary operator
-    r = region(boundary)
     d = dim(boundary)
-    op = BoundaryOperator(restrict(grid, d), closure_stencil, r)
+    op = BoundaryOperator(restrict(grid, d), closure_stencil, region(boundary))
 
-    # Create 1D IdentityMappings for each coordinate direction
-    one_d_grids = restrict.(Ref(grid), Tuple(1:dimension(grid)))
-    Is = IdentityMapping{eltype(grid)}.(size.(one_d_grids))
-
-    # Formulate the correct outer product sequence of the identity mappings and
-    # the boundary operator
-    parts = Base.setindex(Is, op, d)
-    return foldl(⊗, parts)
+    return LazyTensors.inflate(op, size(grid), d)
 end
 
 """
--- a/src/SbpOperators/volumeops/volume_operator.jl	Wed Mar 16 18:39:00 2022 +0100
+++ b/src/SbpOperators/volumeops/volume_operator.jl	Fri Mar 18 16:57:00 2022 +0100
@@ -12,16 +12,10 @@
 function volume_operator(grid::EquidistantGrid, inner_stencil, closure_stencils, parity, direction)
     #TODO: Check that direction <= Dim?
 
-    # Create 1D volume operator in along coordinate direction
     op = VolumeOperator(restrict(grid, direction), inner_stencil, closure_stencils, parity)
-    # Create 1D IdentityMappings for each coordinate direction
-    one_d_grids = restrict.(Ref(grid), Tuple(1:dimension(grid)))
-    Is = IdentityMapping{eltype(grid)}.(size.(one_d_grids))
-    # Formulate the correct outer product sequence of the identity mappings and
-    # the volume operator
-    parts = Base.setindex(Is, op, direction)
-    return foldl(⊗, parts)
+    return LazyTensors.inflate(op, size(grid), direction)
 end
+# TBD: Should the inflation happen here or should we remove this method and do it at the caller instead?
 
 """
     VolumeOperator{T,N,M,K} <: TensorOperator{T,1}
@@ -59,3 +53,4 @@
     r = getregion(i, closure_size(op), op.size[1])
     return LazyTensors.apply(op, v, Index(i, r))
 end
+# TODO: Move this to LazyTensors when we have the region communication down.
--- a/test/LazyTensors/lazy_tensor_operations_test.jl	Wed Mar 16 18:39:00 2022 +0100
+++ b/test/LazyTensors/lazy_tensor_operations_test.jl	Fri Mar 18 16:57:00 2022 +0100
@@ -527,3 +527,26 @@
     end
 
 end
+
+@testset "inflate" begin
+    struct ScalingOperator{T,D} <: TensorMapping{T,D,D}
+        λ::T
+        size::NTuple{D,Int}
+    end
+
+    LazyTensors.apply(m::ScalingOperator{T,D}, v, I::Vararg{Any,D}) where {T,D} = m.λ*v[I...]
+    LazyTensors.range_size(m::ScalingOperator) = m.size
+    LazyTensors.domain_size(m::ScalingOperator) = m.size
+
+
+    I = LazyTensors.inflate(IdentityMapping(),(3,4,5,6), 2)
+    @test I isa TensorMapping{Float64, 3,3}
+    @test range_size(I) == (3,5,6)
+    @test domain_size(I) == (3,5,6)
+
+    # TODO: More tests
+
+    # Check that the dir is inbounds
+
+    # tm = ScalingOperator(2., (4,))
+end