Fix bug in the order in contraction_cost, fix header bugs, and missing includes (#46)

s-mandra · web-flow · commit 4e278e1c7164 · 2026-03-08T09:37:32.000-07:00
This PR fixes a bug related to the order of arguments in the C++
`contraction_cost` function. This bug was subtle because it wasn't
triggering any error and existing unit tests failed to catch it; indeed,
for the implemented `SimpleCost`, the order of arguments in C++ does not
actually matter. Specifically, `SimpleCost` only depends on the total
number of indices, counted once. In principle, the total number of
indices should be computed as `inds_in1 | inds_in2`. However, it was
previously computed as `inds_in1 | inds_out`. Since `inds_out =
(inds_in1 ^ inds_in2) | hyper_inds`, it is easy to show that `inds_in1 |
inds_in2 == inds_in1 | inds_out` (or any other permutation).

Other changes:

- Rename `inds_A`, `inds_B`, and `inds_C` to `inds_out`, `inds_in1`, and
`inds_in2` in both the C++ and Python cost models for better clarity and
consistency.
- Fix copy-paste errors in `isinf` and `signbit` pybind11 bindings for
fixed_float, and correct the return type of `operator/=` to `FixedFloat
&amp;`.
- Prevent potential division by zero and domain errors in the
Metropolis-Hastings probability calculation (`mh.hpp`) and in the
relative error validation blocks (`optimizer.hpp`).
- Add missing `&lt;stdexcept&gt;`, `&lt;sstream&gt;`, and `&lt;variant&gt;` standard
library includes across various headers (`assert.hpp`, `globals.hpp`,
`node.hpp`, `tree.hpp`).
diff --git a/include/tnco/assert.hpp b/include/tnco/assert.hpp
@@ -15,6 +15,7 @@
 #pragma once
 
 #include <sstream>
+#include <stdexcept>
 
 namespace tnco {
 
diff --git a/include/tnco/fixed_float.hpp b/include/tnco/fixed_float.hpp
@@ -104,7 +104,7 @@ struct FixedFloat {
     return *this;
   }
 
-  auto operator/=(const FixedFloat &x) -> FixedFloat {
+  auto operator/=(const FixedFloat &x) -> FixedFloat & {
     mpfr_div(_x, _x, x._x, rounding);
     return *this;
   }
@@ -575,9 +575,10 @@ void init(py::module &m, const std::string &name) {
            [](const self_type &self, const long double &other) -> auto {
              return self >= other;
            })
-      .def("isinf", [](const self_type &self) -> auto { return isnan(self); })
+      .def("isinf", [](const self_type &self) -> auto { return isinf(self); })
       .def("isnan", [](const self_type &self) -> auto { return isnan(self); })
-      .def("signbit", [](const self_type &self) -> auto { return isnan(self); })
+      .def("signbit",
+           [](const self_type &self) -> auto { return signbit(self); })
       .def(py::pickle(
           [](const self_type &self) -> std::string { return dump(self); },
           [](const std::string &state) -> auto {
diff --git a/include/tnco/globals.hpp b/include/tnco/globals.hpp
@@ -15,6 +15,7 @@
 #pragma once
 
 #include <random>
+#include <sstream>
 
 #include "bitset.hpp"
 #include "ctree.hpp"
diff --git a/include/tnco/node.hpp b/include/tnco/node.hpp
@@ -19,6 +19,7 @@
 
 #include <array>
 #include <cstdint>
+#include <variant>
 
 namespace tnco::node {
 
diff --git a/include/tnco/optimize/finite_width/cost_model/base.hpp b/include/tnco/optimize/finite_width/cost_model/base.hpp
@@ -73,9 +73,9 @@ struct CostModel {
     return std::numeric_limits<width_type>::signaling_NaN();
   }
 
-  [[nodiscard]] virtual auto contraction_cost(const bitset_type &inds_A,
-                                              const bitset_type &inds_B,
-                                              const bitset_type &inds_C,
+  [[nodiscard]] virtual auto contraction_cost(const bitset_type &inds_in1,
+                                              const bitset_type &inds_in2,
+                                              const bitset_type &inds_out,
                                               const dims_type &dims,
                                               const bitset_type &slices) const
       -> cost_type {
@@ -123,8 +123,8 @@ void init(py::module &m, const std::string &name) {
       .def_readonly("max_width", &self_type::max_width)
       .def("width", &self_type::width, "inds"_a, "dims"_a)
       .def("delta_width", &self_type::delta_width, "inds"_a, "dims"_a, "pos"_a)
-      .def("contraction_cost", &self_type::contraction_cost, "inds_A"_a,
-           "inds_B"_a, "inds_C"_a, "dims"_a, "slices"_a)
+      .def("contraction_cost", &self_type::contraction_cost, "inds_in1"_a,
+           "inds_in2"_a, "inds_out"_a, "dims"_a, "slices"_a)
       .def_property_readonly("width_type",
                              [](const self_type &self) -> auto {
                                return type_to_str<width_type>();
diff --git a/include/tnco/optimize/finite_width/cost_model/simple.hpp b/include/tnco/optimize/finite_width/cost_model/simple.hpp
@@ -121,23 +121,23 @@ struct CostModel : cost_model::base::CostModel<T...> {
         dims);
   }
 
-  [[nodiscard]] auto contraction_cost(const bitset_type &inds_A,
-                                      const bitset_type &inds_B,
-                                      const bitset_type &inds_C,
+  [[nodiscard]] auto contraction_cost(const bitset_type &inds_in1,
+                                      const bitset_type &inds_in2,
+                                      const bitset_type &inds_out,
                                       const dims_type &dims,
                                       const bitset_type &slices) const
 #ifdef NDEBUG
       noexcept
 #endif
       -> cost_type override {
-    ASSERT(std::size(inds_A) == std::size(inds_B) &&
-               std::size(inds_A) == std::size(inds_C) &&
-               std::size(inds_A) == std::size(slices),
+    ASSERT(std::size(inds_in1) == std::size(inds_in2) &&
+               std::size(inds_in1) == std::size(inds_out) &&
+               std::size(inds_in1) == std::size(slices),
            "Indices  and slices must have the same size.");
-    ASSERT(inds_C.is_subset_of(inds_A | inds_B),
-           "'inds_C' must be a subset of 'inds_A | inds_B'.");
+    ASSERT(inds_out.is_subset_of(inds_in1 | inds_in2),
+           "'inds_out' must be a subset of 'inds_in1 | inds_in2'.");
     return std::visit(
-        [inds = inds_A | inds_B | slices](auto &&dims) -> auto {
+        [inds = inds_in1 | inds_in2 | slices](auto &&dims) -> auto {
           return infinite_memory::cost_model::simple::get_cost<cost_type>(inds,
                                                                           dims);
         },
diff --git a/include/tnco/optimize/finite_width/cost_model/simple_sparse_inds.hpp b/include/tnco/optimize/finite_width/cost_model/simple_sparse_inds.hpp
@@ -133,23 +133,23 @@ struct CostModel : simple::CostModel<T...> {
         dims);
   }
 
-  [[nodiscard]] auto contraction_cost(const bitset_type &inds_A,
-                                      const bitset_type &inds_B,
-                                      const bitset_type &inds_C,
+  [[nodiscard]] auto contraction_cost(const bitset_type &inds_in1,
+                                      const bitset_type &inds_in2,
+                                      const bitset_type &inds_out,
                                       const dims_type &dims,
                                       const bitset_type &slices) const
 #ifdef NDEBUG
       noexcept
 #endif
       -> cost_type override {
-    ASSERT(std::size(inds_A) == std::size(inds_B) &&
-               std::size(inds_A) == std::size(inds_C) &&
-               std::size(inds_A) == std::size(slices),
+    ASSERT(std::size(inds_in1) == std::size(inds_in2) &&
+               std::size(inds_in1) == std::size(inds_out) &&
+               std::size(inds_in1) == std::size(slices),
            "Indices  and slices must have the same size.");
-    ASSERT(inds_C.is_subset_of(inds_A | inds_B),
-           "'inds_C' must be a subset of 'inds_A | inds_B'.");
+    ASSERT(inds_out.is_subset_of(inds_in1 | inds_in2),
+           "'inds_out' must be a subset of 'inds_in1 | inds_in2'.");
     return std::visit(
-        [inds = inds_A | inds_B | slices, &sparse_inds = this->sparse_inds,
+        [inds = inds_in1 | inds_in2 | slices, &sparse_inds = this->sparse_inds,
          &n_projs = this->n_projs](auto &&dims) -> auto {
           return infinite_memory::cost_model::simple_sparse_inds::get_cost<
               cost_type>(inds, sparse_inds, n_projs, dims);
diff --git a/include/tnco/optimize/finite_width/greedy/optimizer.hpp b/include/tnco/optimize/finite_width/greedy/optimizer.hpp
@@ -188,9 +188,9 @@ struct Optimizer : optimize::optimizer::Optimizer {
       if (new_sliced_width_B <= p_cmodel->max_width) {
         // Get new costs for B and A
         const auto new_ccost_A =
-            cmodel.contraction_cost(inds_A, new_inds_B, inds_E, dims, slices);
+            cmodel.contraction_cost(new_inds_B, inds_E, inds_A, dims, slices);
         const auto new_ccost_B =
-            cmodel.contraction_cost(new_inds_B, inds_D, inds_C, dims, slices);
+            cmodel.contraction_cost(inds_D, inds_C, new_inds_B, dims, slices);
 
         // Get the delta cost
         const auto delta_cost =
@@ -343,11 +343,16 @@ struct Optimizer : optimize::optimizer::Optimizer {
     }
 
     // Check final cost
-    if (const auto rel_error = abs((log(total_cost) - log(get_total_cost())) /
-                                   log(get_total_cost()));
-        rel_error > 1e-2) {
-      throw std::logic_error("Total cost is not properly cached (rel. error: " +
-                             tnco::to_string(rel_error * 100) + "%)");
+    if (const auto actual_total_cost = get_total_cost();
+        actual_total_cost > 0) {
+      if (const auto rel_error =
+              abs((log(total_cost) - log(actual_total_cost)) /
+                  log(actual_total_cost));
+          rel_error > 1e-2) {
+        throw std::logic_error(
+            "Total cost is not properly cached (rel. error: " +
+            tnco::to_string(rel_error * 100) + "%)");
+      }
     }
 #endif
 
diff --git a/include/tnco/optimize/infinite_memory/cost_model/base.hpp b/include/tnco/optimize/infinite_memory/cost_model/base.hpp
@@ -45,9 +45,9 @@ struct CostModel {
   auto operator=(CostModel &&) -> CostModel & = delete;
   virtual ~CostModel() = default;
 
-  [[nodiscard]] virtual auto contraction_cost(const bitset_type &inds_A,
-                                              const bitset_type &inds_B,
-                                              const bitset_type &inds_C,
+  [[nodiscard]] virtual auto contraction_cost(const bitset_type &inds_in1,
+                                              const bitset_type &inds_in2,
+                                              const bitset_type &inds_out,
                                               const dims_type &dims) const
       -> cost_type {
     /*
@@ -87,8 +87,8 @@ void init(py::module &m, const std::string &name) {
            [](const self_type &self, const self_type &other) -> auto {
              return true;
            })
-      .def("contraction_cost", &self_type::contraction_cost, "inds_A"_a,
-           "inds_B"_a, "inds_C"_a, "dims"_a)
+      .def("contraction_cost", &self_type::contraction_cost, "inds_in1"_a,
+           "inds_in2"_a, "inds_out"_a, "dims"_a)
       .def_property_readonly("cost_type", [](const self_type &self) -> auto {
         return type_to_str<cost_type>();
       });
diff --git a/include/tnco/optimize/infinite_memory/cost_model/simple.hpp b/include/tnco/optimize/infinite_memory/cost_model/simple.hpp
@@ -62,21 +62,21 @@ struct CostModel : cost_model::base::CostModel<T...> {
   using bitset_type = typename base_type::bitset_type;
   using dims_type = typename base_type::dims_type;
 
-  [[nodiscard]] auto contraction_cost(const bitset_type &inds_A,
-                                      const bitset_type &inds_B,
-                                      const bitset_type &inds_C,
+  [[nodiscard]] auto contraction_cost(const bitset_type &inds_in1,
+                                      const bitset_type &inds_in2,
+                                      const bitset_type &inds_out,
                                       const dims_type &dims) const
 #ifdef NDEBUG
       noexcept
 #endif
       -> cost_type override {
-    ASSERT(std::size(inds_A) == std::size(inds_B) &&
-               std::size(inds_A) == std::size(inds_C),
+    ASSERT(std::size(inds_in1) == std::size(inds_in2) &&
+               std::size(inds_in1) == std::size(inds_out),
            "Indices must have the same size.");
-    ASSERT(inds_C.is_subset_of(inds_A | inds_B),
-           "'inds_C' must be a subset of 'inds_A | inds_B'.");
+    ASSERT(inds_out.is_subset_of(inds_in1 | inds_in2),
+           "'inds_out' must be a subset of 'inds_in1 | inds_in2'.");
     return std::visit(
-        [inds = inds_A | inds_B](auto &&dims) -> auto {
+        [inds = inds_in1 | inds_in2](auto &&dims) -> auto {
           return get_cost<cost_type>(inds, dims);
         },
         dims);
diff --git a/include/tnco/optimize/infinite_memory/cost_model/simple_sparse_inds.hpp b/include/tnco/optimize/infinite_memory/cost_model/simple_sparse_inds.hpp
@@ -66,21 +66,21 @@ struct CostModel : simple::CostModel<T...> {
     }
   }
 
-  [[nodiscard]] auto contraction_cost(const bitset_type &inds_A,
-                                      const bitset_type &inds_B,
-                                      const bitset_type &inds_C,
+  [[nodiscard]] auto contraction_cost(const bitset_type &inds_in1,
+                                      const bitset_type &inds_in2,
+                                      const bitset_type &inds_out,
                                       const dims_type &dims) const
 #ifdef NDEBUG
       noexcept
 #endif
       -> cost_type override {
-    ASSERT(std::size(inds_A) == std::size(inds_B) &&
-               std::size(inds_A) == std::size(inds_C),
+    ASSERT(std::size(inds_in1) == std::size(inds_in2) &&
+               std::size(inds_in1) == std::size(inds_out),
            "Indices must have the same size.");
-    ASSERT(inds_C.is_subset_of(inds_A | inds_B),
-           "'inds_C' must be a subset of 'inds_A | inds_B'.");
+    ASSERT(inds_out.is_subset_of(inds_in1 | inds_in2),
+           "'inds_out' must be a subset of 'inds_in1 | inds_in2'.");
     return std::visit(
-        [inds = inds_A | inds_B, &sparse_inds = this->sparse_inds,
+        [inds = inds_in1 | inds_in2, &sparse_inds = this->sparse_inds,
          &n_projs = this->n_projs](auto &&dims) -> auto {
           return get_cost<cost_type>(inds, sparse_inds, n_projs, dims);
         },
diff --git a/include/tnco/optimize/infinite_memory/optimizer.hpp b/include/tnco/optimize/infinite_memory/optimizer.hpp
@@ -150,9 +150,9 @@ struct Optimizer : optimize::optimizer::Optimizer {
       auto &ccost_A = cost_cache.contraction_cost[pos_A];
       auto &ccost_B = cost_cache.contraction_cost[pos_B];
       const auto new_ccost_A =
-          cmodel.contraction_cost(inds_A, new_inds_B, inds_E, dims);
+          cmodel.contraction_cost(new_inds_B, inds_E, inds_A, dims);
       const auto new_ccost_B =
-          cmodel.contraction_cost(new_inds_B, inds_D, inds_C, dims);
+          cmodel.contraction_cost(inds_D, inds_C, new_inds_B, dims);
 
       // Get the delta cost
       const auto delta_cost = (new_ccost_B - ccost_B) + (new_ccost_A - ccost_A);
@@ -206,11 +206,16 @@ struct Optimizer : optimize::optimizer::Optimizer {
     }
 
     // Check final cost
-    if (const auto rel_error = abs((log(total_cost) - log(get_total_cost())) /
-                                   log(get_total_cost()));
-        rel_error > 1e-2) {
-      throw std::logic_error("Total cost is not properly cached (rel. error: " +
-                             tnco::to_string(rel_error * 100) + "%)");
+    if (const auto actual_total_cost = get_total_cost();
+        actual_total_cost > 0) {
+      if (const auto rel_error =
+              abs((log(total_cost) - log(actual_total_cost)) /
+                  log(actual_total_cost));
+          rel_error > 1e-2) {
+        throw std::logic_error(
+            "Total cost is not properly cached (rel. error: " +
+            tnco::to_string(rel_error * 100) + "%)");
+      }
     }
 #endif
   }
diff --git a/include/tnco/optimize/infinite_memory/utils.hpp b/include/tnco/optimize/infinite_memory/utils.hpp
@@ -44,10 +44,10 @@ struct CostCache {
             cache.contraction_cost[pos] = 0;
             cache.partial_cost[pos] = 0;
           } else {
-            const auto &inds_A = ctree.inds[pos];
-            const auto &inds_B = ctree.inds[node.children[0]];
-            const auto &inds_C = ctree.inds[node.children[1]];
-            const auto cost_A = ccost(inds_A, inds_B, inds_C, ctree.dims);
+            const auto &inds_out = ctree.inds[pos];
+            const auto &inds_in1 = ctree.inds[node.children[0]];
+            const auto &inds_in2 = ctree.inds[node.children[1]];
+            const auto cost_A = ccost(inds_in1, inds_in2, inds_out, ctree.dims);
             const auto &pcost_B = cache.partial_cost[node.children[0]];
             const auto &pcost_C = cache.partial_cost[node.children[1]];
             cache.contraction_cost[pos] = cost_A;
@@ -83,10 +83,10 @@ struct HyperCache {
       if (const auto &node = ctree.nodes[pos]; node.is_leaf()) {
         hyper_inds[pos] = bitset_type(n_inds);
       } else {
-        const auto &inds_A = ctree.inds[pos];
-        const auto &inds_B = ctree.inds[node.children[0]];
-        const auto &inds_C = ctree.inds[node.children[1]];
-        hyper_inds[pos] = inds_A & inds_B & inds_C;
+        const auto &inds_out = ctree.inds[pos];
+        const auto &inds_in1 = ctree.inds[node.children[0]];
+        const auto &inds_in2 = ctree.inds[node.children[1]];
+        hyper_inds[pos] = inds_out & inds_in1 & inds_in2;
       }
     }
   }
@@ -106,10 +106,10 @@ template <typename CostType, typename CTree, typename CCost>
   tnco::utils::traverse(
       ctree, [&ctree, &ccost, &total_cost = total_cost](auto &&pos) -> auto {
         if (const auto &node = ctree.nodes[pos]; !node.is_leaf()) {
-          const auto &inds_A = ctree.inds[pos];
-          const auto &inds_B = ctree.inds[node.children[0]];
-          const auto &inds_C = ctree.inds[node.children[1]];
-          total_cost += ccost(inds_A, inds_B, inds_C, ctree.dims);
+          const auto &inds_out = ctree.inds[pos];
+          const auto &inds_in1 = ctree.inds[node.children[0]];
+          const auto &inds_in2 = ctree.inds[node.children[1]];
+          total_cost += ccost(inds_in1, inds_in2, inds_out, ctree.dims);
         }
       });
   return total_cost;
diff --git a/include/tnco/optimize/prob/mh.hpp b/include/tnco/optimize/prob/mh.hpp
@@ -49,8 +49,13 @@ struct Probability final : prob::base::Probability<T...> {
 #endif
       -> cost_type override {
     ASSERT(delta_cost <= 0 || (delta_cost / old_cost) >= -1, "Domain error.");
-    return delta_cost <= 0 ? cost_type{1}
-                           : pow(1 + (delta_cost / old_cost), -beta);
+    if (delta_cost <= 0) {
+      return cost_type{1};
+    }
+    if (old_cost == 0) {
+      return cost_type{0};
+    }
+    return pow(1 + (delta_cost / old_cost), -beta);
   }
 
   [[nodiscard]] auto clone() const noexcept -> clone_ptr_type override {
diff --git a/include/tnco/tree.hpp b/include/tnco/tree.hpp
@@ -20,6 +20,7 @@
 #include <numeric>
 #include <stack>
 #include <tnco/assert.hpp>
+#include <variant>
 #include <vector>
 
 namespace tnco::tree {
diff --git a/tnco/optimize/finite_width/cost_model.py b/tnco/optimize/finite_width/cost_model.py
diff --git a/tnco/optimize/infinite_memory/cost_model.py b/tnco/optimize/infinite_memory/cost_model.py
