From dfa7f367edd9171a154acdaa92e2d78b4135efcc Mon Sep 17 00:00:00 2001
From: Alexander Condello <arcondello@gmail.com>
Date: Mon, 27 Mar 2023 15:58:21 -0700
Subject: [PATCH] Improve the performance of removing variables from
 expressions

---
 dimod/include/dimod/abc.h                     |  2 +-
 dimod/include/dimod/expression.h              | 27 ++++++++++++++-----
 dimod/include/dimod/quadratic_model.h         | 16 +++++++++++
 ...ariables-performance-f13a65b6a16fa484.yaml | 12 +++++++++
 4 files changed, 49 insertions(+), 8 deletions(-)
 create mode 100644 releasenotes/notes/expression-fix-variables-performance-f13a65b6a16fa484.yaml

diff --git a/dimod/include/dimod/abc.h b/dimod/include/dimod/abc.h
index 88f6f7560..193799260 100644
--- a/dimod/include/dimod/abc.h
+++ b/dimod/include/dimod/abc.h
@@ -707,7 +707,7 @@ void QuadraticModelBase<bias_type, index_type>::fix_variable(index_type v, T ass
     add_offset(assignment * linear(v));
 
     // finally remove v
-    remove_variable(v);
+    QuadraticModelBase<bias_type, index_type>::remove_variable(v);
 }
 
 template <class bias_type, class index_type>
diff --git a/dimod/include/dimod/expression.h b/dimod/include/dimod/expression.h
index a4b71ec3e..7f79126b5 100644
--- a/dimod/include/dimod/expression.h
+++ b/dimod/include/dimod/expression.h
@@ -419,7 +419,20 @@ void Expression<bias_type, index_type>::clear() {
 template <class bias_type, class index_type>
 template <class T>
 void Expression<bias_type, index_type>::fix_variable(index_type v, T assignment) {
-    throw std::logic_error("not implemented - fix_variable");
+    assert(v >= 0 && static_cast<size_type>(v) < parent_->num_variables());
+
+    auto vit = indices_.find(v);
+    if (vit == indices_.end()) return;  // nothing to remove
+
+    // remove the biases
+    base_type::fix_variable(vit->second, assignment);
+
+    // update the indices
+    auto it = variables_.erase(variables_.begin() + vit->second);
+    indices_.erase(vit);
+    for (; it != variables_.end(); ++it) {
+        indices_[*it] -= 1;
+    }
 }
 
 template <class bias_type, class index_type>
@@ -592,6 +605,8 @@ bool Expression<bias_type, index_type>::remove_interaction(index_type u, index_t
 
 template <class bias_type, class index_type>
 void Expression<bias_type, index_type>::remove_variable(index_type v) {
+    assert(v >= 0 && static_cast<size_type>(v) < parent_->num_variables());
+
     auto vit = indices_.find(v);
     if (vit == indices_.end()) return;  // nothing to remove
 
@@ -599,12 +614,10 @@ void Expression<bias_type, index_type>::remove_variable(index_type v) {
     base_type::remove_variable(vit->second);
 
     // update the indices
-    variables_.erase(variables_.begin() + vit->second);
-
-    // indices is no longer valid, so remake
-    indices_.clear();
-    for (size_type ui = 0; ui < variables_.size(); ++ui) {
-        indices_[variables_[ui]] = ui;
+    auto it = variables_.erase(variables_.begin() + vit->second);
+    indices_.erase(vit);
+    for (; it != variables_.end(); ++it) {
+        indices_[*it] -= 1;
     }
 }
 
diff --git a/dimod/include/dimod/quadratic_model.h b/dimod/include/dimod/quadratic_model.h
index df80b14a5..316bb2a1b 100644
--- a/dimod/include/dimod/quadratic_model.h
+++ b/dimod/include/dimod/quadratic_model.h
@@ -64,6 +64,15 @@ class QuadraticModel : public abc::QuadraticModelBase<Bias, Index> {
     /// Change the vartype of `v`, updating the biases appropriately.
     void change_vartype(Vartype vartype, index_type v);
 
+    /**
+     * Remove variable `v` from the model by fixing its value.
+     *
+     * Note that this causes a reindexing, where all variables above `v` have
+     * their index reduced by one.
+     */
+    template <class T>
+    void fix_variable(index_type v, T assignment);
+
     /// Return the lower bound on variable ``v``.
     bias_type lower_bound(index_type v) const;
 
@@ -229,6 +238,13 @@ void QuadraticModel<bias_type, index_type>::change_vartype(Vartype vartype, inde
     }
 }
 
+template <class bias_type, class index_type>
+template <class T>
+void QuadraticModel<bias_type, index_type>::fix_variable(index_type v, T assignment) {
+    base_type::fix_variable(v, assignment);
+    varinfo_.erase(varinfo_.begin() + v);
+}
+
 template <class bias_type, class index_type>
 bias_type QuadraticModel<bias_type, index_type>::lower_bound(index_type v) const {
     // even though v is unused, we need this to conform the the QuadraticModelBase API
diff --git a/releasenotes/notes/expression-fix-variables-performance-f13a65b6a16fa484.yaml b/releasenotes/notes/expression-fix-variables-performance-f13a65b6a16fa484.yaml
new file mode 100644
index 000000000..b63f9c6b7
--- /dev/null
+++ b/releasenotes/notes/expression-fix-variables-performance-f13a65b6a16fa484.yaml
@@ -0,0 +1,12 @@
+---
+features:
+  - |
+    Improve the performance of fixing and removing variables from constrained
+    quadratic model expressions.
+  - |
+    Implement the ``Expression::fix_variable()`` C++ method. Previously it would
+    throw ``std::logic_error("not implemented - fix_variable")``.
+upgrade:
+  - |
+    Add an overload to the C++ ``QuadraticModel::remove_variable()`` method.
+    This is binary compatible, but it makes ``&remove_variable`` ambiguous.