Implement operations for UniformScaling

Project.toml

@@ -1,7 +1,7 @@
 name = "Tensorial"
 uuid = "98f94333-fa9f-48a9-ad80-1c66397b2b38"
 authors = ["Keita Nakamura <[email protected]>"]
-version = "0.3.1"
+version = "0.3.2"
 
 [deps]
 ForwardDiff = "f6369f11-7733-5829-9624-2563aa707210"

src/Tensorial.jl

@@ -2,7 +2,7 @@ module Tensorial
 
 using LinearAlgebra, Statistics
 # re-exports from LinearAlgebra and Statistics
-export ⋅, ×, dot, tr, det, norm, mean
+export ⋅, ×, dot, tr, det, norm, mean, I
 
 using StaticArrays
 using Base: @pure, @_inline_meta, @_propagate_inbounds_meta

src/ops.jl

@@ -19,11 +19,38 @@ end
 @inline Base.:+(x::AbstractTensor) = x
 @inline Base.:-(x::AbstractTensor) = _map(-, x)
 
-# error for standard multiplications
-function Base.:*(::Tensor, ::Tensor)
-    error("use `⋅` (`\\cdot`) for single contraction and `⊡` (`\\boxdot`) for double contraction instead of `*`")
+const SquareTensor{dim, T} = Union{AbstractTensor{Tuple{dim, dim}, T, 2}, AbstractTensor{Tuple{@Symmetry{dim, dim}}, T, 2}}
+
+@generated function _add_uniform(x::SquareTensor{dim}, λ::Real) where {dim}
+    S = promote_size(Size(x), Size(Symmetry(dim, dim)))
+    tocartesian = CartesianIndices(S)
+    exps = map(indices(S)) do i
+        i, j = Tuple(tocartesian[i])
+        ex = getindex_expr(:x, x, i, j)
+        return i == j ? :($ex + λ) : ex
+    end
+    TT = tensortype(S)
+    return quote
+        @_inline_meta
+        @inbounds $TT($(exps...))
+    end
 end
 
+@inline Base.:+(x::AbstractTensor, y::UniformScaling) = _add_uniform( x,  y.λ)
+@inline Base.:-(x::AbstractTensor, y::UniformScaling) = _add_uniform( x, -y.λ)
+@inline Base.:+(x::UniformScaling, y::AbstractTensor) = _add_uniform( y,  x.λ)
+@inline Base.:-(x::UniformScaling, y::AbstractTensor) = _add_uniform(-y,  x.λ)
+@inline dot(x::Union{AbstractVec, AbstractMat, SquareTensor}, y::UniformScaling) = x * y.λ
+@inline dot(x::UniformScaling, y::Union{AbstractVec, AbstractMat, SquareTensor}) = x.λ * y
+@inline double_contraction(x::SquareTensor, y::UniformScaling) = tr(x) * y.λ
+@inline double_contraction(x::UniformScaling, y::SquareTensor) = x.λ * tr(y)
+
+# error for standard multiplications
+error_multiply() = error("use `⋅` (`\\cdot`) for single contraction and `⊡` (`\\boxdot`) for double contraction instead of `*`")
+Base.:*(::AbstractTensor, ::AbstractTensor) = error_multiply()
+Base.:*(::AbstractTensor, ::UniformScaling) = error_multiply()
+Base.:*(::UniformScaling, ::AbstractTensor) = error_multiply()
+
 function contraction_exprs(S1::Size, S2::Size, ::Val{N}) where {N}
     S = contraction(S1, S2, Val(N))
     s1 = map(i -> EinsumIndex(:(Tuple(x)), i), independent_indices(S1))
@@ -135,8 +162,6 @@ julia> a = x ⋅ y
     end
 end
 
-const SquareTensor{dim, T} = Union{AbstractTensor{Tuple{dim, dim}, T, 2}, AbstractTensor{Tuple{@Symmetry{dim, dim}}, T, 2}}
-
 # tr/mean
 @inline function tr(x::SquareTensor{dim}) where {dim}
     sum(i -> @inbounds(x[i,i]), 1:dim)

test/ops.jl

@@ -202,3 +202,32 @@ end
         @test_throws Exception rotmat(Vec(1,0) => Vec(1,1)) # length of two vectors must be the same
     end
 end
+
+@testset "UniformScaling" begin
+    for T in (Float32, Float64)
+        for SquareTensorType in (Tensor{Tuple{3, 3}, T}, Tensor{Tuple{@Symmetry{3, 3}}, T})
+            x = rand(SquareTensorType)
+            @test (@inferred x + I)::SquareTensorType == x + one(x)
+            @test (@inferred x - I)::SquareTensorType == x - one(x)
+            @test (@inferred I + x)::SquareTensorType == one(x) + x
+            @test (@inferred I - x)::SquareTensorType == one(x) - x
+            y = rand(Mat{3, 4, T})
+            v = rand(Vec{3, T})
+            @test (@inferred x ⋅ I)::SquareTensorType == x ⋅ one(x)
+            @test (@inferred I ⋅ x)::SquareTensorType == one(x) ⋅ x
+            @test (@inferred y ⋅ I)::Mat{3, 4, T} == y ⋅ one(Mat{4, 4})
+            @test (@inferred I ⋅ y)::Mat{3, 4, T} == one(Mat{3, 3}) ⋅ y
+            @test (@inferred v ⋅ I)::Vec{3, T} == v ⋅ one(x)
+            @test (@inferred I ⋅ v)::Vec{3, T} == one(x) ⋅ v
+            @test (@inferred I ⊡ x)::T == one(x) ⊡ x
+            @test (@inferred x ⊡ I)::T == x ⊡ one(x)
+            # wrong input
+            @test_throws Exception x * I
+            @test_throws Exception y * I
+            @test_throws Exception v * I
+            @test_throws Exception I * x
+            @test_throws Exception I * y
+            @test_throws Exception I * v
+        end
+    end
+end