+ and - for DiagTrav (#142)

* + and - for DiagTrav * overload /, \ and - for DiagTrav
JuliaLinearAlgebra · Jan 28, 2025 · 9ed8340 · 9ed8340 · dlfivefifty · Jan 28, 2025
1 parent f6b5f1c
commit 9ed8340
Show file tree

Hide file tree

Showing 3 changed files with 28 additions and 1 deletion.
diff --git a/Project.toml b/Project.toml
@@ -1,7 +1,7 @@
 name = "LazyBandedMatrices"
 uuid = "d7e5e226-e90b-4449-9968-0f923699bf6f"
 authors = ["Sheehan Olver <[email protected]>"]
-version = "0.11.2"
+version = "0.11.3"
 
 [deps]
 ArrayLayouts = "4c555306-a7a7-4459-81d9-ec55ddd5c99a"

diff --git a/src/blockkron.jl b/src/blockkron.jl
@@ -212,6 +212,19 @@ end
 invdiagtrav(a) = InvDiagTrav(a)
 invdiagtrav(a::DiagTrav) = a.array
 
+-(A::DiagTrav) = DiagTrav(-A.array)
+
+for op in (:+, :-)
+    @eval $op(A::DiagTrav, B::DiagTrav) = DiagTrav($op(A.array, B.array))
+end
+
+for op in (:*, :/)
+    @eval $op(A::DiagTrav, k::Number) = DiagTrav($op(A.array, k))
+end
+
+for op in (:*, :\)
+    @eval $op(k::Number, A::DiagTrav) = DiagTrav($op(k, A.array))
+end
 
 """
     KronTrav(A,B,C...)

diff --git a/test/test_blockkron.jl b/test/test_blockkron.jl
@@ -34,6 +34,20 @@ LinearAlgebra.factorize(A::MyLazyArray) = factorize(A.data)
         A = [1 2; 3 4; 5 6]
         @test DiagTrav(A) == [1, 3, 2, 5, 4]
 
+        @test DiagTrav(A) + DiagTrav(A) isa DiagTrav
+        @test DiagTrav(A) - DiagTrav(A) isa DiagTrav
+        @test 2DiagTrav(A) isa DiagTrav
+        @test 2\DiagTrav(A) isa DiagTrav
+        @test DiagTrav(A)*2 isa DiagTrav
+        @test DiagTrav(A)/2 isa DiagTrav
+        @test -DiagTrav(A) isa DiagTrav
+
+
+        @test DiagTrav(A) + DiagTrav(A) == 2DiagTrav(A) == DiagTrav(A)*2 == 2Vector(DiagTrav(A))
+        @test DiagTrav(A) - DiagTrav(A) == 0DiagTrav(A) == 0Vector(DiagTrav(A))
+        @test 2\DiagTrav(A) == DiagTrav(A)/2 == Vector(DiagTrav(A))/2
+        @test -DiagTrav(A) == -Vector(DiagTrav(A))
+
         @test resize!(DiagTrav(A), Block(2)) == [1, 3,2]
 
         A = DiagTrav(randn(3,3,3))