improvements to flattening (#103)

Author: Krastanov

CHANGELOG.md

@@ -1,5 +1,9 @@
 # News
 
+## v0.4.7 - 2025-02-25
+
+- Improvements to default flattening of expressions having nested sums and scalings.
+
 ## v0.4.6 - 2025-01-18
 
 - Migrate `express` functionality and representation types to QuantumInterface.

Project.toml

@@ -1,7 +1,7 @@
 name = "QuantumSymbolics"
 uuid = "efa7fd63-0460-4890-beb7-be1bbdfbaeae"
 authors = ["QuantumSymbolics.jl contributors"]
-version = "0.4.6"
+version = "0.4.7"
 
 [deps]
 Latexify = "23fbe1c1-3f47-55db-b15f-69d7ec21a316"

src/QSymbolicsBase/basic_ops_homogeneous.jl

@@ -85,6 +85,7 @@ julia> k₁ + k₂
     _arguments_precomputed
 end
 function SAdd{S}(d) where S
+    isempty(d) && return SZero{S}()
     terms = [c*obj for (obj,c) in d]
     length(d)==1 ? first(terms) : SAdd{S}(d,Set(terms),terms)
 end
@@ -99,7 +100,7 @@ function Base.:(+)(x::Symbolic{T}, xs::Vararg{Symbolic{T}, N}) where {T<:QObj, N
     xs = collect(xs)
     f = first(xs)
     nonzero_terms = filter!(x->!iszero(x),xs)
-    isempty(nonzero_terms) ? f : SAdd{T}(countmap_flatten(nonzero_terms, SScaled{T}))
+    isempty(nonzero_terms) ? f : SAdd{T}(countmap_flatten(nonzero_terms, SAdd{T}, SScaled{T}))
 end
 basis(x::SAdd) = basis(first(x.dict).first)
 

src/QSymbolicsBase/utils.jl

@@ -34,15 +34,30 @@ function countmap(samples) # A simpler version of StatsBase.countmap, because St
     counts
 end
 
-function countmap_flatten(samples, flattenhead)
+function countmap_flatten(samples, flattenadd, flattenmul)
     counts = Dict{Any,Any}()
     for s in samples
-        if isexpr(s) && s isa flattenhead # TODO Could you use the TermInterface `operation` here instead of `flattenhead`?
+        if s isa flattenadd
+            for (term,coef) in pairs(s.dict)
+                counts[term] = get(counts, term, 0)+coef
+            end
+        elseif s isa flattenmul
             coef, term = arguments(s)
-            counts[term] = get(counts, term, 0)+coef
+            if term isa flattenadd
+                for (_term,_coef) in pairs(term.dict)
+                    counts[_term] = get(counts, _term, 0)+coef*_coef
+                end
+            else
+                counts[term] = get(counts, term, 0)+coef
+            end
         else
             counts[s] = get(counts, s, 0)+1
         end
     end
+    for (term,coef) in pairs(counts)
+        if iszero(coef)===true # iszero might return symbolic expressions instead of true/false # TODO make into a proper function like isdefinitelyzero, see whether upstream Symbolics has it
+            delete!(counts, term)
+        end
+    end
     counts
-end
+end

test/test_sym_expressions.jl

@@ -2,13 +2,24 @@
     @test +(Z1) == Z1
     @test +(Z) == Z
     @test isequal(Z1 - Z2, Z1 + (-Z2))
-    @test_broken isequal(Z1 - 2*Z2 + 2*X1, -2*Z2 + Z1 + 2*X1)
-    @test_broken isequal(Z1 - 2*Z2 + 2*X1, Z1 + 2*(-Z2+X1))
+    @test isequal(Z1 - 2*Z2 + 2*X1, -2*Z2 + Z1 + 2*X1)
+    @test isequal(Z1 - 2*Z2 + 2*X1, Z1 + 2*(-Z2+X1))
 
     state1 = XBasisState(1, SpinBasis(1//2))
     state2 = XBasisState(1, SpinBasis(1//2))
     state3 = XBasisState(2, SpinBasis(1//2))
 
     @test isequal(state1, state2)
     @test !isequal(state1, state3)
+
+    @op A
+    @op B
+
+    @test isequal(A+B+A, 2A+B)
+    @test isequal(A+B-A, B)
+    @test A+B-A === B
+    @test 0*A === A-A-2A+2A
+    @test isequal(A+B-A+B, 2B)
+    @test isequal(2A-3B+2(A+B), 4A-B)
+    @test isequal(2A-3B+2(A+2B), 4A+B)
 end