refactor: replace raw constraint and variable counts with a map and getter methods

yemine0x01 · yemine0x01 · commit 2843e5634bdf · 2026-02-17T13:14:11.000-05:00
diff --git a/include/pyoptinterface/knitro_model.hpp b/include/pyoptinterface/knitro_model.hpp
@@ -555,15 +555,13 @@ class KNITROModel : public OnesideLinearConstraintMixin<KNITROModel>,
 	KNINT _variable_index(const VariableIndex &variable) const;
 	KNINT _constraint_index(const ConstraintIndex &constraint) const;
 
-	size_t n_vars = 0;
-	size_t n_cons = 0;
-	size_t n_lincons = 0;
-	size_t n_quadcons = 0;
-	size_t n_soccons = 0;
-	size_t n_nlcons = 0;
+	size_t get_num_vars() const;
+	size_t get_num_cons(std::optional<ConstraintType> type = std::nullopt) const;
 
   private:
 	// Member variables
+	size_t m_n_vars = 0;
+	std::unordered_map<ConstraintType, size_t> m_n_cons_map;
 	std::shared_ptr<LM_context> m_lm = nullptr;
 	std::unique_ptr<KN_context, KNITROFreeProblemT> m_kc = nullptr;
 
@@ -605,7 +603,7 @@ class KNITROModel : public OnesideLinearConstraintMixin<KNITROModel>,
 	template <typename F>
 	ConstraintIndex _add_constraint_impl(ConstraintType type,
 	                                     const std::tuple<double, double> &interval,
-	                                     const char *name, size_t *np, const F &setter)
+	                                     const char *name, const F &setter)
 	{
 		KNINT indexCon;
 		int error = knitro::KN_add_con(m_kc.get(), &indexCon);
@@ -632,11 +630,16 @@ class KNITROModel : public OnesideLinearConstraintMixin<KNITROModel>,
 
 		m_con_sense_flags[indexCon] = CON_UPBND;
 
-		n_cons++;
-		if (np != nullptr)
+		auto it = m_n_cons_map.find(type);
+		if (it != m_n_cons_map.end())
 		{
-			(*np)++;
+			it->second++;
 		}
+		else
+		{
+			m_n_cons_map[type] = 1;
+		}
+
 		m_is_dirty = true;
 
 		return constraint;
diff --git a/lib/knitro_model.cpp b/lib/knitro_model.cpp
@@ -203,7 +203,7 @@ VariableIndex KNITROModel::add_variable(VariableDomain domain, double lb, double
 		_set_name(knitro::KN_set_var_name, indexVar, name);
 	}
 
-	n_vars++;
+	m_n_vars++;
 	_mark_dirty();
 
 	return variable;
@@ -304,7 +304,7 @@ void KNITROModel::delete_variable(const VariableIndex &variable)
 	_set_value<KNINT, int>(knitro::KN_set_var_type, indexVar, KN_VARTYPE_CONTINUOUS);
 	_set_value<KNINT, double>(knitro::KN_set_var_lobnd, indexVar, -get_infinity());
 	_set_value<KNINT, double>(knitro::KN_set_var_upbnd, indexVar, get_infinity());
-	n_vars--;
+	m_n_vars--;
 	_mark_dirty();
 }
 
@@ -330,7 +330,7 @@ ConstraintIndex KNITROModel::add_linear_constraint(const ScalarAffineFunction &f
 	auto setter = [this, &function](const ConstraintIndex &constraint) {
 		_set_linear_constraint(constraint, function);
 	};
-	return _add_constraint_impl(ConstraintType::Linear, interval, name, &n_lincons, setter);
+	return _add_constraint_impl(ConstraintType::Linear, interval, name, setter);
 }
 
 ConstraintIndex KNITROModel::add_quadratic_constraint(const ScalarQuadraticFunction &function,
@@ -349,7 +349,7 @@ ConstraintIndex KNITROModel::add_quadratic_constraint(const ScalarQuadraticFunct
 	auto setter = [this, &function](const ConstraintIndex &constraint) {
 		_set_quadratic_constraint(constraint, function);
 	};
-	return _add_constraint_impl(ConstraintType::Quadratic, interval, name, &n_quadcons, setter);
+	return _add_constraint_impl(ConstraintType::Quadratic, interval, name, setter);
 }
 
 ConstraintIndex KNITROModel::add_second_order_cone_constraint(
@@ -362,17 +362,15 @@ ConstraintIndex KNITROModel::add_second_order_cone_constraint(
 			_set_second_order_cone_constraint_rotated(constraint, variables);
 		};
 		std::pair<double, double> interval = {0.0, get_infinity()};
-		return _add_constraint_impl(ConstraintType::SecondOrderCone, interval, name, &n_soccons,
-		                            setter);
+		return _add_constraint_impl(ConstraintType::SecondOrderCone, interval, name, setter);
 	}
 	else
 	{
 		auto setter = [this, &variables](const ConstraintIndex &constraint) {
 			_set_second_order_cone_constraint(constraint, variables);
 		};
 		std::pair<double, double> interval = {0.0, get_infinity()};
-		return _add_constraint_impl(ConstraintType::SecondOrderCone, interval, name, &n_soccons,
-		                            setter);
+		return _add_constraint_impl(ConstraintType::SecondOrderCone, interval, name, setter);
 	}
 }
 
@@ -389,7 +387,7 @@ ConstraintIndex KNITROModel::add_single_nl_constraint(ExpressionGraph &graph,
 		m_pending_outputs[&graph].cons.push_back(constraint);
 		m_need_to_add_callbacks = true;
 	};
-	return _add_constraint_impl(ConstraintType::NL, interval, name, &n_nlcons, setter);
+	return _add_constraint_impl(ConstraintType::NL, interval, name, setter);
 }
 
 ConstraintIndex KNITROModel::add_single_nl_constraint_sense_rhs(ExpressionGraph &graph,
@@ -514,24 +512,7 @@ void KNITROModel::delete_constraint(const ConstraintIndex &constraint)
 	_set_value<KNINT, double>(knitro::KN_set_con_lobnd, indexCon, -get_infinity());
 	_set_value<KNINT, double>(knitro::KN_set_con_upbnd, indexCon, get_infinity());
 
-	n_cons--;
-	switch (constraint.type)
-	{
-	case ConstraintType::Linear:
-		n_lincons--;
-		break;
-	case ConstraintType::Quadratic:
-		n_quadcons--;
-		break;
-	case ConstraintType::SecondOrderCone:
-		n_soccons--;
-		break;
-	case ConstraintType::NL:
-		n_nlcons--;
-		break;
-	default:
-		break;
-	}
+	m_n_cons_map[constraint.type]--;
 
 	auto it = m_soc_aux_cons.find(indexCon);
 	if (it != m_soc_aux_cons.end())
@@ -1007,6 +988,29 @@ bool KNITROModel::empty() const
 	return m_kc == nullptr;
 }
 
+size_t KNITROModel::get_num_vars() const
+{
+	return m_n_vars;
+}
+
+size_t KNITROModel::get_num_cons(std::optional<ConstraintType> type) const
+{
+	if (!type.has_value())
+	{
+		size_t total = 0;
+		for (const auto &[_, count] : m_n_cons_map)
+		{
+			total += count;
+		}
+		return total;
+	}
+	else
+	{
+		auto it = m_n_cons_map.find(type.value());
+		return it != m_n_cons_map.end() ? it->second : 0;
+	}
+}
+
 int KNITROModel::get_solve_status() const
 {
 	_check_dirty();
@@ -1025,12 +1029,8 @@ void KNITROModel::_check_dirty() const
 void KNITROModel::_reset_state()
 {
 	m_kc.reset();
-	n_vars = 0;
-	n_cons = 0;
-	n_lincons = 0;
-	n_quadcons = 0;
-	n_soccons = 0;
-	n_nlcons = 0;
+	m_n_vars = 0;
+	m_n_cons_map.clear();
 	m_soc_aux_cons.clear();
 	m_con_sense_flags.clear();
 	m_obj_flag = 0;
diff --git a/lib/knitro_model_ext.cpp b/lib/knitro_model_ext.cpp
@@ -1,4 +1,5 @@
 #include <nanobind/nanobind.h>
+#include <nanobind/stl/optional.h>
 #include <nanobind/stl/string.h>
 #include <nanobind/stl/vector.h>
 #include <nanobind/stl/tuple.h>
@@ -31,6 +32,7 @@ NB_MODULE(knitro_model_ext, m)
 	    .def(nb::init<const KNITROEnv &>())
 	    .def("init", nb::overload_cast<>(&KNITROModel::init))
 	    .def("init", nb::overload_cast<const KNITROEnv &>(&KNITROModel::init))
+
 	    // clang-format off
 		BIND_F(close)
 		BIND_F(get_infinity)
@@ -43,12 +45,8 @@ NB_MODULE(knitro_model_ext, m)
 		BIND_F(get_release)
 	    // clang-format on
 
-	    .def_ro("n_vars", &KNITROModel::n_vars)
-	    .def_ro("n_cons", &KNITROModel::n_cons)
-	    .def_ro("n_lincons", &KNITROModel::n_lincons)
-	    .def_ro("n_quadcons", &KNITROModel::n_quadcons)
-	    .def_ro("n_soccons", &KNITROModel::n_soccons)
-	    .def_ro("n_nlcons", &KNITROModel::n_nlcons)
+	    .def("number_of_variables", &KNITROModel::get_num_vars)
+	    .def("number_of_constraints", &KNITROModel::get_num_cons, nb::arg("type") = nb::none())
 
 	    .def("add_variable", &KNITROModel::add_variable,
 	         nb::arg("domain") = VariableDomain::Continuous, nb::arg("lb") = -KN_INFINITY,
@@ -162,10 +160,9 @@ NB_MODULE(knitro_model_ext, m)
 	         nb::arg("expr"), nb::arg("sense") = ObjectiveSense::Minimize)
 	    .def("_add_single_nl_objective", &KNITROModel::add_single_nl_objective, nb::arg("graph"),
 	         nb::arg("result"))
-	    .def("set_objective_coefficient", &KNITROModel::set_objective_coefficient,
-	         nb::arg("variable"), nb::arg("coefficient"))
 
 	    // clang-format off
+		BIND_F(set_objective_coefficient)
 		BIND_F(get_obj_value)
 		BIND_F(set_obj_sense)
 		BIND_F(get_obj_sense)
diff --git a/src/pyoptinterface/_src/knitro.py b/src/pyoptinterface/_src/knitro.py
@@ -17,7 +17,6 @@
 from .core_ext import (
     ConstraintIndex,
     ConstraintSense,
-    ConstraintType,
     ExprBuilder,
     ScalarAffineFunction,
     ScalarQuadraticFunction,
@@ -334,25 +333,6 @@ def supports_model_attribute(
         else:
             return attribute in model_attribute_get_func_map
 
-    def number_of_variables(self) -> int:
-        return self.n_vars
-
-    def number_of_constraints(
-        self, constraint_type: Union[ConstraintType, None] = None
-    ) -> int:
-        if constraint_type is None:
-            return self.n_cons
-        elif constraint_type == ConstraintType.Linear:
-            return self.n_lincons
-        elif constraint_type == ConstraintType.Quadratic:
-            return self.n_quadcons
-        elif constraint_type == ConstraintType.SecondOrderCone:
-            return self.n_soccons
-        elif constraint_type == ConstraintType.KNITRO_NL:
-            return self.n_nlcons
-        else:
-            raise ValueError(f"Unknown constraint type: {constraint_type}")
-
     @overload
     def add_linear_constraint(
         self,
