Make quadratic jacobian merge entries for the same variable

Author: metab0t

lib/nleval.cpp

@@ -71,27 +71,69 @@ void LinearQuadraticEvaluator::analyze_jacobian_structure(size_t &m_jacobian_nnz
 		auto &f = quadratic_constraints[i];
 		auto row = quadratic_constraint_indices[i];
 		auto N = f.size();
+
+		Hashmap<IndexT, size_t> variable_to_jacobian_nnz;
+
 		for (size_t j = 0; j < N; j++)
 		{
 			auto x1 = f.variable_1s[j];
 			auto x2 = f.variable_2s[j];
 			if (x1 == x2)
 			{
-				m_jacobian_rows.push_back(row);
-				m_jacobian_cols.push_back(x1);
-				jacobian_linear_terms.emplace_back(2.0 * f.coefficients[j], x1, m_jacobian_nnz);
-				m_jacobian_nnz += 1;
+				auto result = variable_to_jacobian_nnz.insert({x1, m_jacobian_nnz});
+				auto iter = result.first;
+				auto has_inserted = result.second;
+
+				if (has_inserted)
+				{
+					m_jacobian_rows.push_back(row);
+					m_jacobian_cols.push_back(x1);
+					jacobian_linear_terms.emplace_back(2.0 * f.coefficients[j], x1, m_jacobian_nnz);
+					m_jacobian_nnz += 1;
+				}
+				else
+				{
+					auto nnz = iter->second;
+					jacobian_linear_terms.emplace_back(2.0 * f.coefficients[j], x1, nnz);
+				}
 			}
 			else
 			{
-				m_jacobian_rows.push_back(row);
-				m_jacobian_cols.push_back(x1);
-				jacobian_linear_terms.emplace_back(f.coefficients[j], x2, m_jacobian_nnz);
-				m_jacobian_nnz += 1;
-				m_jacobian_rows.push_back(row);
-				m_jacobian_cols.push_back(x2);
-				jacobian_linear_terms.emplace_back(f.coefficients[j], x1, m_jacobian_nnz);
-				m_jacobian_nnz += 1;
+				{
+					auto result = variable_to_jacobian_nnz.insert({x1, m_jacobian_nnz});
+					auto iter = result.first;
+					auto has_inserted = result.second;
+
+					if (has_inserted)
+					{
+						m_jacobian_rows.push_back(row);
+						m_jacobian_cols.push_back(x1);
+						jacobian_linear_terms.emplace_back(f.coefficients[j], x2, m_jacobian_nnz);
+						m_jacobian_nnz += 1;
+					}
+					else
+					{
+						auto nnz = iter->second;
+						jacobian_linear_terms.emplace_back(f.coefficients[j], x2, nnz);
+					}
+				}
+				{
+					auto result = variable_to_jacobian_nnz.insert({x2, m_jacobian_nnz});
+					auto iter = result.first;
+					auto has_inserted = result.second;
+					if (has_inserted)
+					{
+						m_jacobian_rows.push_back(row);
+						m_jacobian_cols.push_back(x2);
+						jacobian_linear_terms.emplace_back(f.coefficients[j], x1, m_jacobian_nnz);
+						m_jacobian_nnz += 1;
+					}
+					else
+					{
+						auto nnz = iter->second;
+						jacobian_linear_terms.emplace_back(f.coefficients[j], x1, nnz);
+					}
+				}
 			}
 		}
 		if (f.affine_part)
@@ -100,11 +142,25 @@ void LinearQuadraticEvaluator::analyze_jacobian_structure(size_t &m_jacobian_nnz
 			auto N = af.size();
 			for (size_t j = 0; j < N; j++)
 			{
-				m_jacobian_rows.push_back(row);
-				m_jacobian_cols.push_back(af.variables[j]);
-				jacobian_constants.emplace_back(af.coefficients[j], m_jacobian_nnz + j);
+				auto x = af.variables[j];
+
+				auto result = variable_to_jacobian_nnz.insert({x, m_jacobian_nnz});
+				auto iter = result.first;
+				auto has_inserted = result.second;
+
+				if (has_inserted)
+				{
+					m_jacobian_rows.push_back(row);
+					m_jacobian_cols.push_back(x);
+					jacobian_constants.emplace_back(af.coefficients[j], m_jacobian_nnz);
+					m_jacobian_nnz += 1;
+				}
+				else
+				{
+					auto nnz = iter->second;
+					jacobian_constants.emplace_back(af.coefficients[j], nnz);
+				}
 			}
-			m_jacobian_nnz += N;
 		}
 	}
 }
@@ -170,7 +226,6 @@ void LinearQuadraticEvaluator::analyze_sparse_gradient_structure(
 			{
 				size_t grad_index =
 				    add_gradient_column(x1, gradient_nnz, gradient_cols, gradient_index_map);
-				;
 				gradient_linear_terms.emplace_back(c, x2, grad_index);
 				grad_index =
 				    add_gradient_column(x2, gradient_nnz, gradient_cols, gradient_index_map);

tests/test_nlp_bilinear.py

@@ -0,0 +1,20 @@
+import pyoptinterface as poi
+
+import pytest
+
+
+def test_bilinear(ipopt_model_ctor):
+    model = ipopt_model_ctor()
+
+    x = model.add_m_variables(3, lb=0.0)
+
+    obj = -(x[0] + x[1]) * (x[0] + x[2])
+
+    expr = (x[0] + x[1]) * (x[0] + x[1]) + (x[0] + x[2]) * (x[0] + x[2])
+    model.add_quadratic_constraint(expr, poi.Eq, 4.0)
+
+    model.set_objective(obj)
+    model.optimize()
+
+    objective_value = model.get_model_attribute(poi.ModelAttribute.ObjectiveValue)
+    assert objective_value == pytest.approx(-2.0, abs=1e-8)