Allow quicksum for high dimensional numpy array directly without writing poi.quicksum(x.flat)

Author: metab0t

src/pyoptinterface/_src/aml.py

@@ -4,6 +4,8 @@
 from collections.abc import Collection
 from typing import Tuple, Union
 
+import numpy as np
+
 
 def make_variable_ndarray(
     model,
@@ -105,6 +107,8 @@ def f(*args):
 def quicksum_(expr: ExprBuilder, terms, f=None):
     if isinstance(terms, dict):
         iter = terms.values()
+    elif isinstance(terms, np.ndarray):
+        iter = terms.flat
     else:
         iter = terms
     if f:

tests/test_matrix_api.py

@@ -27,3 +27,15 @@ def test_matrix_api(model_interface):
     model.optimize()
     obj_value = model.get_model_attribute(poi.ModelAttribute.ObjectiveValue)
     assert obj_value == approx(-N * ub)
+
+
+def test_quicksum_ndarray(model_interface):
+    model = model_interface
+
+    N = 10
+    x = model.add_m_variables((N, 2 * N), lb=1.0, ub=3.0)
+    obj = poi.quicksum(x)
+    model.set_objective(obj)
+    model.optimize()
+    obj_value = model.get_model_attribute(poi.ModelAttribute.ObjectiveValue)
+    assert obj_value == approx(2 * N**2)