bugfix for remove_zero_valued_subgraphs!

numericalEFT · Oct 30, 2024 · 0d3ce5e · 0d3ce5e
1 parent 79f254b
commit 0d3ce5e
Show file tree

Hide file tree

Showing 2 changed files with 9 additions and 0 deletions.
diff --git a/src/computational_graph/transform.jl b/src/computational_graph/transform.jl
@@ -399,6 +399,13 @@ function remove_zero_valued_subgraphs!(g::AbstractGraph)
     mask_zeros = findall(x -> x != zero(x), subg_fac)
     if isempty(mask_zeros)
         mask_zeros = [1]  # retain eldest(g) if all subfactors are zero
+    else
+        if g.operator == Prod
+            idx = findfirst(x -> x==zero(x), subg_fac)
+            if !isnothing(idx)
+                append!(mask_zeros, idx)
+            end
+        end
     end
     set_subgraphs!(g, subg[mask_zeros])
     set_subgraph_factors!(g, subg_fac[mask_zeros])

diff --git a/src/computational_graph/tree_properties.jl b/src/computational_graph/tree_properties.jl
@@ -94,6 +94,8 @@ end
 function has_zero_subfactors(g::AbstractGraph)
     if isleaf(g)
         return false  # convention: subgraph_factors = [] ⟹ subfactorless = false
+    elseif g.operator == Prod && 0 in subgraph_factors(g)
+        return true
     else
         return iszero(subgraph_factors(g))
     end