Add ShiftedOperator for shifted linear operators (#394)

Introduces the ShiftedOperator type to represent operators of the form H + σI, supporting efficient multiplication and adjoint/transpose operations. Includes comprehensive tests for real symmetric, complex non-Hermitian, and mutable shift scenarios. The new operator is integrated into the main module.

Author: farhadrclass

src/LinearOperators.jl

@@ -22,6 +22,7 @@ include("DiagonalHessianApproximation.jl")
 include("linalg.jl")
 include("special-operators.jl")
 include("TimedOperators.jl")
+include("shifted_operators.jl")
 
 # Utilities
 include("utilities.jl")

src/shifted_operators.jl

@@ -0,0 +1,104 @@
+export ShiftedOperator
+
+"A data type to hold information relative to shifted operators."
+mutable struct ShiftedData{T, OpH}
+  H::OpH
+  σ::T
+  function ShiftedData{T, OpH}(H::OpH, σ::T) where {T, OpH}
+    size(H, 1) == size(H, 2) || throw(DimensionMismatch("Operator H must be square."))
+    new{T, OpH}(H, σ)
+  end
+end
+
+ShiftedData(H::OpH, σ::T) where {T, OpH} = ShiftedData{T, OpH}(H, σ)
+
+# Forward product: y = α(H + σI)x + βy
+function shifted_prod!(y, data::ShiftedData, x, α, β)
+  # y = α * H * x + β * y
+  mul!(y, data.H, x, α, β)
+
+  # y = y + (α * σ) * x
+  if !(iszero(data.σ) || iszero(α))
+    axpy!(α * data.σ, x, y)
+  end
+  return y
+end
+
+# Transpose product: y = α(Hᵀ + σI)x + βy
+function shifted_tprod!(y, data::ShiftedData, x, α, β)
+  # y = α * Hᵀ * x + β * y
+  mul!(y, transpose(data.H), x, α, β)
+
+  # y = y + (α * σ) * x
+  if !(iszero(data.σ) || iszero(α))
+    axpy!(α * data.σ, x, y)
+  end
+  return y
+end
+
+# Conjugate transpose (adjoint) product: y = α(Hᴴ + conj(σ)I)x + βy
+function shifted_ctprod!(y, data::ShiftedData, x, α, β)
+  # y = α * Hᴴ * x + β * y
+  mul!(y, adjoint(data.H), x, α, β)
+
+  # y = y + (α * conj(σ)) * x
+  if !(iszero(data.σ) || iszero(α))
+    axpy!(α * conj(data.σ), x, y)
+  end
+  return y
+end
+
+"""
+    ShiftedOperator(H, σ=0)
+
+Construct a linear operator representing `op = H + σI`.
+"""
+mutable struct ShiftedOperator{T, OpH, F, Ft, Fct} <: AbstractLinearOperator{T}
+  nrow::Int
+  ncol::Int
+  symmetric::Bool
+  hermitian::Bool
+  prod!::F      # Closure for op * x
+  tprod!::Ft    # Closure for transpose(op) * x
+  ctprod!::Fct  # Closure for adjoint(op) * x
+  data::ShiftedData{T, OpH}
+  nprod::Int    # Internal counter for products
+  ntprod::Int   # Internal counter for transpose products
+  nctprod::Int  # Internal counter for adjoint products
+end
+
+function ShiftedOperator(H::OpH, σ_in::Number = zero(eltype(H))) where {OpH}
+  T = eltype(H)
+  σ = convert(T, σ_in)  # Enforces that σ matches the element type of H
+
+  data = ShiftedData(H, σ)
+
+  prod! = (y, x, α, β) -> shifted_prod!(y, data, x, α, β)
+  tprod! = (y, x, α, β) -> shifted_tprod!(y, data, x, α, β)
+  ctprod! = (y, x, α, β) -> shifted_ctprod!(y, data, x, α, β)
+
+  n = size(H, 1)
+
+  is_sym = issymmetric(H)
+  is_herm = ishermitian(H) && isreal(σ)
+
+  return ShiftedOperator(n, n, is_sym, is_herm, prod!, tprod!, ctprod!, data, 0, 0, 0)
+end
+
+size(op::ShiftedOperator) = (op.nrow, op.ncol)
+issymmetric(op::ShiftedOperator) = op.symmetric
+ishermitian(op::ShiftedOperator) = op.hermitian && isreal(op.data.σ)
+
+has_args5(op::ShiftedOperator) = true
+use_prod5!(op::ShiftedOperator) = true
+
+isallocated5(op::ShiftedOperator) = true
+
+storage_type(op::ShiftedOperator{T}) where {T} = storage_type(op.data.H)
+
+function reset!(op::ShiftedOperator)
+  op.nprod = 0
+  op.ntprod = 0
+  op.nctprod = 0
+  return op
+end

test/runtests.jl

@@ -16,6 +16,7 @@ include("test_normest.jl")
 include("test_diag.jl")
 include("test_chainrules.jl")
 include("test_solve_shifted_system.jl")
+include("test_shifted_operator.jl")
 include("gpu/test_S_kwarg.jl")
 include("gpu/jlarrays.jl")
 if Sys.isapple() && occursin("arm64", Sys.MACHINE)

test/test_shifted_operator.jl

@@ -0,0 +1,127 @@
+@testset "ShiftedOperator Tests" begin
+  @testset "Real Symmetric (Float64)" begin
+    n = 5
+    H_dense = rand(n, n)
+    H_dense = H_dense + H_dense'
+
+    H_op = LinearOperator(H_dense; symmetric = true, hermitian = true)
+
+    σ = 2.0
+    op = ShiftedOperator(H_op, σ)
+
+    A_ref = H_dense + σ * I
+
+    x = rand(n)
+    y = zeros(n)
+
+    @test size(op) == (n, n)
+    @test issymmetric(op) == true
+    @test ishermitian(op) == true
+    @test op.data.σ == σ
+
+    mul!(y, op, x)
+    @test y ≈ A_ref * x
+
+    y_orig = rand(n)
+    y = copy(y_orig)
+    α, β = 0.5, -1.0
+    mul!(y, op, x, α, β)
+    @test y ≈ α * (A_ref * x) + β * y_orig
+
+    op_t = transpose(op)
+
+    y_t = zeros(n)
+    mul!(y_t, op_t, x)
+    @test y_t ≈ transpose(A_ref) * x
+  end
+
+  @testset "Complex Non-Hermitian" begin
+    n = 5
+    H_dense = rand(ComplexF64, n, n)
+    H_op = LinearOperator(H_dense)
+
+    σ = 1.0 + 2.0im
+    op = ShiftedOperator(H_op, σ)
+
+    @test !issymmetric(op)
+    @test !ishermitian(op)
+
+    A_ref = H_dense + σ * I
+    x = rand(ComplexF64, n)
+
+    op_c = adjoint(op)
+
+    y_c = zeros(ComplexF64, n)
+    mul!(y_c, op_c, x)
+    @test y_c ≈ adjoint(A_ref) * x
+  end
+
+  @testset "Mutation (Updating Sigma)" begin
+    n = 3
+    H_dense = rand(n, n)
+    op = ShiftedOperator(LinearOperator(H_dense), 1.0)
+
+    x = ones(n)
+    y1 = op * x
+
+    op.data.σ = 10.0
+
+    y2 = op * x
+
+    @test !(y1 ≈ y2)
+    @test y2 ≈ (H_dense + 10.0*I) * x
+  end
+
+  @testset "Mutation (Dynamic Hermitian Check)" begin
+    n = 3
+    H_dense = rand(ComplexF64, n, n)
+    H_dense = H_dense + H_dense' 
+    H_op = LinearOperator(H_dense; symmetric=false, hermitian=true)
+
+    σ = 2.0
+    op = ShiftedOperator(H_op, σ)
+    @test ishermitian(op)
+
+    op.data.σ = 2.0 + 1.0im
+    @test !ishermitian(op)
+    
+    op.data.σ = 3.0
+    @test ishermitian(op)
+  end
+
+  @testset "Strict Type Constraint" begin
+    H = LinearOperator(rand(Float32, 5, 5))
+    σ = 1.0
+    op = ShiftedOperator(H, σ)
+
+    @test eltype(op) == Float32
+    @test op.data.σ isa Float32
+
+    x = rand(Float32, 5)
+    y = op * x
+    @test eltype(y) == Float32
+  end
+
+  @testset "Coverage & Utilities" begin
+    n = 5
+    H = LinearOperator(rand(n, n))
+    x = rand(n)
+    y = zeros(n)
+    op = ShiftedOperator(H, 2.0)
+
+    mul!(y, transpose(op), x, 0.5, 1.0) 
+    
+    mul!(y, transpose(op), x, 0.0, 1.0)
+
+    op_zero = ShiftedOperator(H, 0.0)
+    mul!(y, transpose(op_zero), x)
+    
+    @test LinearOperators.isallocated5(op) == true
+  
+    @test LinearOperators.storage_type(op) == LinearOperators.storage_type(H)
+    
+    op.nprod = 10
+    reset!(op)
+    @test op.nprod == 0
+  end
+end