diff src/LazyTensors/lazy_tensor_operations.jl @ 929:e9dd43cbd127 feature/variable_derivatives

Merge performance/get_region_type_inference
author Jonatan Werpers <jonatan@werpers.com>
date Mon, 21 Feb 2022 10:34:48 +0100
parents b41180efb6c2
children 8b0ff2fddc32 f1c2a4fa0ee1
line wrap: on
line diff
--- a/src/LazyTensors/lazy_tensor_operations.jl	Fri Feb 18 08:03:37 2022 +0100
+++ b/src/LazyTensors/lazy_tensor_operations.jl	Mon Feb 21 10:34:48 2022 +0100
@@ -430,3 +430,28 @@
     print(io, "SizeMismatch: ")
     print(io, "domain size $(domain_size(err.tm)) of TensorMapping not matching size $(err.sz)")
 end
+
+function apply_with_region(op, v, boundary_width::Integer, dim_size::Integer, i)
+    if 0 < i <= boundary_width
+        return LazyTensors.apply(op,v,Index(i,Lower))
+    elseif boundary_width < i <= dim_size-boundary_width
+        return LazyTensors.apply(op,v,Index(i,Interior))
+    elseif dim_size-boundary_width < i <= dim_size
+        return LazyTensors.apply(op,v,Index(i,Upper))
+    else
+        error("Bounds error") # TODO: Make this more standard
+    end
+end
+# TBD: Can these methods be merge by having a function as an arguement instead?
+# TODO: Add inference test that show where things break and how this rewrite fixes it.
+function apply_transpose_with_region(op, v, boundary_width::Integer, dim_size::Integer, i)
+    if 0 < i <= boundary_width
+        return LazyTensors.apply_transpose(op,v,Index(i,Lower))
+    elseif boundary_width < i <= dim_size-boundary_width
+        return LazyTensors.apply_transpose(op,v,Index(i,Interior))
+    elseif dim_size-boundary_width < i <= dim_size
+        return LazyTensors.apply_transpose(op,v,Index(i,Upper))
+    else
+        error("Bounds error") # TODO: Make this more standard
+    end
+end