Add operator simplification for unary negation and flatten nested AddedOperators

Author: albertomercurio

src/basic.jl

@@ -286,8 +286,18 @@ for T in SCALINGNUMBERTYPES[2:end]
     end
 end
 
-Base.:-(L::AbstractSciMLOperator) = ScaledOperator(-true, L)
 Base.:+(L::AbstractSciMLOperator) = L
+Base.:-(L::AbstractSciMLOperator) = ScaledOperator(-true, L)
+
+# Special cases for constant scalars. These simplify the structure when applicable
+function Base.:-(L::ScaledOperator)
+    isconstant(L.λ) && return ScaledOperator(-L.λ, L.L)
+    return ScaledOperator(L.λ, -L.L) # Try to propagate the rule
+end
+function Base.:-(L::MatrixOperator)
+    isconstant(L) && return MatrixOperator(-L.A)
+    return ScaledOperator(-true, L) # Going back to the generic case
+end
 
 function Base.convert(::Type{AbstractMatrix}, L::ScaledOperator)
     convert(Number, L.λ) * convert(AbstractMatrix, L.L)
@@ -428,9 +438,11 @@ struct AddedOperator{T,
 
     function AddedOperator(ops)
         @assert !isempty(ops)
-        _check_AddedOperator_sizes(ops)
-        T = mapreduce(eltype, promote_type, ops)
-        new{T, typeof(ops)}(ops)
+        # Flatten nested AddedOperators
+        ops_flat = _flatten_added_operators(ops)
+        _check_AddedOperator_sizes(ops_flat)
+        T = mapreduce(eltype, promote_type, ops_flat)
+        new{T, typeof(ops_flat)}(ops_flat)
     end
 end
 
@@ -440,6 +452,25 @@ end
 
 AddedOperator(L::AbstractSciMLOperator) = L
 
+# Helper function to flatten nested AddedOperators
+@generated function _flatten_added_operators(ops::Tuple)
+    exprs = ()
+    for i in 1:length(ops.parameters)
+        T = ops.parameters[i]
+        if T <: AddedOperator
+            # If this element is an AddedOperator, unpack its ops
+            exprs = (exprs..., :(ops[$i].ops...))
+        else
+            # Otherwise, keep the element as-is
+            exprs = (exprs..., :(ops[$i]))
+        end
+    end
+    
+    return quote
+        tuple($(exprs...))
+    end
+end
+
 @generated function _check_AddedOperator_sizes(ops::Tuple)
     ops_types = ops.parameters
     N = length(ops_types)

test/basic.jl

@@ -141,6 +141,35 @@ end
     end
 end
 
+@testset "Unary +/-" begin
+    A = MatrixOperator(rand(N, N))
+    v = rand(N, K)
+
+    # Test unary +
+    @test +A === A
+    
+    # Test unary - on constant MatrixOperator (simplified to MatrixOperator)
+    minusA = -A
+    @test minusA isa MatrixOperator
+    @test minusA * v ≈ -A.A * v
+    
+    # Test unary - on non-constant MatrixOperator (falls back to ScaledOperator)
+    func(A, u, p, t) = t
+    timeDepA = MatrixOperator(A.A; update_func = func)
+    minusTimeDepA = -timeDepA
+    @test minusTimeDepA isa ScaledOperator  # Can't simplify, uses generic fallback
+    
+    # Test unary - on non-constant ScaledOperator (propagates to inner operator)
+    timeDepScaled = ScalarOperator(0.0, func) * A
+    @test !isconstant(timeDepScaled)
+    minusTimeDepScaled = -timeDepScaled
+    @test minusTimeDepScaled.λ isa ScalarOperator && minusTimeDepScaled.L isa MatrixOperator # Propagates negation to inner MatrixOperator
+    
+    # Test double negation
+    @test -(-A) isa MatrixOperator
+    @test (-(-A)) * v ≈ A * v
+end
+
 @testset "ScaledOperator" begin
     A = rand(N, N)
     D = Diagonal(rand(N))
@@ -269,13 +298,24 @@ end
     w_orig = copy(v)
     @test mul!(v, op, u, α, β) ≈ α * (A + B) * u + β * w_orig
 
-    # ensure AddedOperator doesn't nest
+    # Test flattening of nested AddedOperators via direct constructor
     A = MatrixOperator(rand(N, N))
-    L = A + (A + A) + A
+    B = MatrixOperator(rand(N, N))
+    C = MatrixOperator(rand(N, N))
+    
+    # Create nested structure: (A + B) is an AddedOperator
+    AB = A + B
+    @test AB isa AddedOperator
+    
+    # When we create AddedOperator((AB, C)), it should flatten
+    L = AddedOperator((AB, C))
     @test L isa AddedOperator
-    for op in L.ops
-        @test !isa(op, AddedOperator)
-    end
+    @test length(L.ops) == 3  # Should have A, B, C (not AB and C)
+    @test all(op -> !isa(op, AddedOperator), L.ops)
+    
+    # Verify correctness
+    test_vec = rand(N, K)
+    @test L * test_vec ≈ (A + B + C) * test_vec
 
     ## Time-Dependent Coefficients
     for T in (Float32, Float64, ComplexF32, ComplexF64)