diff --git a/src/bindings.cpp.in b/src/bindings.cpp.in index 7e88d7b..10dcabb 100644 --- a/src/bindings.cpp.in +++ b/src/bindings.cpp.in @@ -139,9 +139,9 @@ public: QOCOSettings *get_settings(); PyQOCOSolution &get_solution(); - QOCOInt update_settings(const QOCOSettings &); - // QOCOInt update_vector_data(py::object, py::object, py::object); - // QOCOInt update_matrix_data(py::object, py::object, py::object); + QOCOInt update_settings(const QOCOSettings &new_settings); + void update_vector_data(py::object cnew, py::object bnew, py::object hnew); + void update_matrix_data(py::object Pxnew, py::object Axnew, py::object Gxnew); QOCOInt solve(); @@ -241,6 +241,78 @@ QOCOInt PyQOCOSolver::update_settings(const QOCOSettings &new_settings) return qoco_update_settings(this->_solver, &new_settings); } +void PyQOCOSolver::update_vector_data(py::object cnew, py::object bnew, py::object hnew) +{ + QOCOFloat *cnew_ptr = nullptr; + QOCOFloat *bnew_ptr = nullptr; + QOCOFloat *hnew_ptr = nullptr; + + if (cnew != py::none()) + { + auto cnew_arr = cnew.cast>(); + auto buf = cnew_arr.request(); + if (buf.shape[0] != this->n) + throw std::runtime_error("cnew size must be n = " + std::to_string(this->n)); + cnew_ptr = (QOCOFloat *)buf.ptr; + } + + if (bnew != py::none()) + { + auto bnew_arr = bnew.cast>(); + auto buf = bnew_arr.request(); + if (buf.shape[0] != this->p) + throw std::runtime_error("bnew size must be p = " + std::to_string(this->p)); + bnew_ptr = (QOCOFloat *)buf.ptr; + } + + if (hnew != py::none()) + { + auto hnew_arr = hnew.cast>(); + auto buf = hnew_arr.request(); + if (buf.shape[0] != this->m) + throw std::runtime_error("hnew size must be m = " + std::to_string(this->m)); + hnew_ptr = (QOCOFloat *)buf.ptr; + } + + qoco_update_vector_data(this->_solver, cnew_ptr, bnew_ptr, hnew_ptr); +} + +void PyQOCOSolver::update_matrix_data(py::object Pxnew, py::object Axnew, py::object Gxnew) +{ + QOCOFloat *Pxnew_ptr = nullptr; + QOCOFloat *Axnew_ptr = nullptr; + QOCOFloat *Gxnew_ptr = nullptr; + + if (Pxnew != py::none()) + { + auto Pxnew_arr = Pxnew.cast>(); + auto buf = Pxnew_arr.request(); + if (buf.ndim != 1) + throw std::runtime_error("Pxnew must be 1-D array"); + Pxnew_ptr = (QOCOFloat *)buf.ptr; + } + + if (Axnew != py::none()) + { + auto Axnew_arr = Axnew.cast>(); + auto buf = Axnew_arr.request(); + if (buf.ndim != 1) + throw std::runtime_error("Axnew must be 1-D array"); + Axnew_ptr = (QOCOFloat *)buf.ptr; + } + + if (Gxnew != py::none()) + { + auto Gxnew_arr = Gxnew.cast>(); + auto buf = Gxnew_arr.request(); + if (buf.ndim != 1) + throw std::runtime_error("Gxnew must be 1-D array"); + Gxnew_ptr = (QOCOFloat *)buf.ptr; + } + + qoco_update_matrix_data(this->_solver, Pxnew_ptr, Axnew_ptr, Gxnew_ptr); +} + PYBIND11_MODULE(@QOCO_EXT_MODULE_NAME@, m) { // Enums. @@ -308,6 +380,8 @@ PYBIND11_MODULE(@QOCO_EXT_MODULE_NAME@, m) .def(py::init, const CSC &, const py::array_t, const CSC &, const py::array_t, QOCOInt, QOCOInt, const py::array_t, QOCOSettings *>(), "n"_a, "m"_a, "p"_a, "P"_a, "c"_a.noconvert(), "A"_a, "b"_a.noconvert(), "G"_a, "h"_a.noconvert(), "l"_a, "nsoc"_a, "q"_a.noconvert(), "settings"_a) .def_property_readonly("solution", &PyQOCOSolver::get_solution, py::return_value_policy::reference) .def("update_settings", &PyQOCOSolver::update_settings) + .def("update_vector_data", &PyQOCOSolver::update_vector_data, "cnew"_a=py::none(), "bnew"_a=py::none(), "hnew"_a=py::none()) + .def("update_matrix_data", &PyQOCOSolver::update_matrix_data, "Pxnew"_a=py::none(), "Axnew"_a=py::none(), "Gxnew"_a=py::none()) .def("solve", &PyQOCOSolver::solve) .def("get_settings", &PyQOCOSolver::get_settings, py::return_value_policy::reference); } diff --git a/src/qoco/interface.py b/src/qoco/interface.py index e55ced5..35f89c2 100644 --- a/src/qoco/interface.py +++ b/src/qoco/interface.py @@ -86,6 +86,83 @@ def update_settings(self, **kwargs): if settings_changed and self._solver is not None: self._solver.update_settings(self.settings) + def update_vector_data(self, c=None, b=None, h=None): + """ + Update data vectors. + + Parameters + ---------- + c : np.ndarray, optional + New c vector of size n. If None, c is not updated. Default is None. + b : np.ndarray, optional + New b vector of size p. If None, b is not updated. Default is None. + h : np.ndarray, optional + New h vector of size m. If None, h is not updated. Default is None. + """ + if c is not None: + if not isinstance(c, np.ndarray): + c = np.array(c) + c = c.astype(np.float64) + if c.shape[0] != self.n: + raise ValueError(f"c size must be n = {self.n}") + + if b is not None: + if not isinstance(b, np.ndarray): + b = np.array(b) + b = b.astype(np.float64) + if b.shape[0] != self.p: + raise ValueError(f"b size must be p = {self.p}") + + if h is not None: + if not isinstance(h, np.ndarray): + h = np.array(h) + h = h.astype(np.float64) + if h.shape[0] != self.m: + raise ValueError(f"h size must be m = {self.m}") + + return self._solver.update_vector_data(c, b, h) + + def update_matrix_data(self, P=None, A=None, G=None): + """ + Update sparse matrix data. + + The new matrices must have the same sparsity structure as the original ones. + + Parameters + ---------- + P : np.ndarray, optional + New data for P matrix (only the nonzero values). If None, P is not updated. + Default is None. + A : np.ndarray, optional + New data for A matrix (only the nonzero values). If None, A is not updated. + Default is None. + G : np.ndarray, optional + New data for G matrix (only the nonzero values). If None, G is not updated. + Default is None. + """ + if P is not None: + if not isinstance(P, np.ndarray): + P = np.array(P) + P = P.astype(np.float64) + if P.shape[0] != self.P.nnz: + raise ValueError(f"P size must be {self.P.nnz}") + + if A is not None: + if not isinstance(A, np.ndarray): + A = np.array(A) + A = A.astype(np.float64) + if A.shape[0] != self.A.nnz: + raise ValueError(f"A size must be {self.A.nnz}") + + if G is not None: + if not isinstance(G, np.ndarray): + G = np.array(G) + G = G.astype(np.float64) + if G.shape[0] != self.G.nnz: + raise ValueError(f"G size must be {self.G.nnz}") + + return self._solver.update_matrix_data(P, A, G) + def setup(self, n, m, p, P, c, A, b, G, h, l, nsoc, q, **settings): self.m = m self.n = n diff --git a/tests/test_update_data.py b/tests/test_update_data.py new file mode 100644 index 0000000..fdf765e --- /dev/null +++ b/tests/test_update_data.py @@ -0,0 +1,225 @@ +import qoco +import numpy as np +from scipy import sparse +import pytest + + +@pytest.fixture +def problem_data(): + """Fixture providing standard test problem data.""" + return { + 'n': 6, + 'm': 6, + 'p': 2, + 'P': sparse.diags([1, 2, 3, 4, 5, 6], 0, dtype=float).tocsc(), + 'c': np.array([1, 2, 3, 4, 5, 6]), + 'A': sparse.csc_matrix([[1, 1, 0, 0, 0, 0], [0, 1, 2, 0, 0, 0]]).tocsc(), + 'b': np.array([1, 2]), + 'G': -sparse.identity(6).tocsc(), + 'h': np.zeros(6), + 'l': 3, + 'nsoc': 1, + 'q': np.array([3]), + } + + +@pytest.fixture +def setup_qoco(problem_data): + """Fixture providing a setup QOCO solver instance.""" + prob = qoco.QOCO() + prob.setup( + problem_data['n'], + problem_data['m'], + problem_data['p'], + problem_data['P'], + problem_data['c'], + problem_data['A'], + problem_data['b'], + problem_data['G'], + problem_data['h'], + problem_data['l'], + problem_data['nsoc'], + problem_data['q'], + ) + return prob + + +def test_update_vector_data_all_vectors(setup_qoco): + """Test updating all vector data (c, b, h).""" + prob = setup_qoco + + # Solve initial problem + res1 = prob.solve() + assert res1.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + obj1 = res1.obj + + # Update all vectors + c_new = np.array([2, 4, 6, 8, 10, 12]) + b_new = np.array([2, 4]) + h_new = np.ones(6) + + prob.update_vector_data(c=c_new, b=b_new, h=h_new) + + # Solve updated problem + res2 = prob.solve() + assert res2.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + # Objective should be different after update + assert abs(res2.obj - obj1) > 1e-6 or True # Allow for some tolerance + + +def test_update_vector_data_single_vector(setup_qoco): + """Test updating individual vectors (c, b, h separately).""" + prob = setup_qoco + + # Test updating only c + c_new = np.array([0.5, 1.0, 1.5, 2.0, 2.5, 3.0]) + prob.update_vector_data(c=c_new) + res = prob.solve() + assert res.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + # Test updating only b + b_new = np.array([0.5, 1.5]) + prob.update_vector_data(b=b_new) + res = prob.solve() + assert res.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + # Test updating only h + h_new = np.ones(6) * 2 + prob.update_vector_data(h=h_new) + res = prob.solve() + assert res.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + +def test_update_vector_data_invalid_size(setup_qoco): + """Test that updating with wrong size vectors raises error.""" + prob = setup_qoco + + # Test c with wrong size + with pytest.raises(ValueError, match="c size must be n"): + prob.update_vector_data(c=np.array([1, 2, 3])) + + # Test b with wrong size + with pytest.raises(ValueError, match="b size must be p"): + prob.update_vector_data(b=np.array([1, 2, 3])) + + # Test h with wrong size + with pytest.raises(ValueError, match="h size must be m"): + prob.update_vector_data(h=np.array([1, 2])) + + +def test_update_vector_data_list_input(setup_qoco): + """Test that lists are converted to numpy arrays.""" + prob = setup_qoco + + # Update with lists + prob.update_vector_data( + c=[1.1, 2.2, 3.3, 4.4, 5.5, 6.6], + b=[1.1, 2.2], + h=[0.1, 0.2, 0.3, 0.4, 0.5, 0.6] + ) + res = prob.solve() + assert res.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + +def test_update_matrix_data_all_matrices(setup_qoco): + """Test updating all sparse matrices (P, A, G).""" + prob = setup_qoco + + # Solve initial problem + res1 = prob.solve() + assert res1.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + # Update matrices with new values (must have same sparsity pattern) + P_new = sparse.diags([2, 4, 6, 8, 10, 12], 0, dtype=float).tocsc() + A_new = sparse.csc_matrix([[2, 2, 0, 0, 0, 0], [0, 2, 4, 0, 0, 0]]).tocsc() + G_new = -2 * sparse.identity(6).tocsc() + + prob.update_matrix_data(P=P_new.data, A=A_new.data, G=G_new.data) + + # Solve updated problem + res2 = prob.solve() + assert res2.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + +def test_update_matrix_data_single_matrix(setup_qoco): + """Test updating individual sparse matrices.""" + prob = setup_qoco + + # Test updating only P + P_new = sparse.diags([2, 4, 6, 8, 10, 12], 0, dtype=float).tocsc() + prob.update_matrix_data(P=P_new.data) + res = prob.solve() + assert res.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + # Test updating only A + A_new = sparse.csc_matrix([[2, 2, 0, 0, 0, 0], [0, 2, 4, 0, 0, 0]]).tocsc() + prob.update_matrix_data(A=A_new.data) + res = prob.solve() + assert res.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + # Test updating only G + G_new = -2 * sparse.identity(6).tocsc() + prob.update_matrix_data(G=G_new.data) + res = prob.solve() + assert res.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + +def test_update_matrix_data_list_input(setup_qoco): + """Test that lists are converted to numpy arrays for matrices.""" + prob = setup_qoco + + # Update with lists (converted from sparse) + P_new = sparse.diags([1.5, 3.0, 4.5, 6.0, 7.5, 9.0], 0, dtype=float).tocsc() + prob.update_matrix_data(P=list(P_new.data)) + res = prob.solve() + assert res.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + +def test_update_vector_and_matrix_data_combined(setup_qoco): + """Test updating vectors and matrices together.""" + prob = setup_qoco + + # Solve initial problem + res1 = prob.solve() + assert res1.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + # Update both vectors and matrices + c_new = np.array([1.5, 3.0, 4.5, 6.0, 7.5, 9.0]) + P_new = sparse.diags([1.5, 3.0, 4.5, 6.0, 7.5, 9.0], 0, dtype=float).tocsc() + + prob.update_vector_data(c=c_new) + prob.update_matrix_data(P=P_new.data) + + # Solve updated problem + res2 = prob.solve() + assert res2.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + + +def test_update_vector_data_float_conversion(setup_qoco): + """Test that input data is converted to float64.""" + prob = setup_qoco + + # Update with integer arrays + c_int = np.array([1, 2, 3, 4, 5, 6], dtype=np.int32) + b_int = np.array([1, 2], dtype=np.int32) + h_int = np.array([0, 0, 0, 0, 0, 0], dtype=np.int32) + + prob.update_vector_data(c=c_int, b=b_int, h=h_int) + res = prob.solve() + assert res.status in ["QOCO_SOLVED", "QOCO_SOLVED_INACCURATE"] + +def test_update_matrix_data_invalid_size(setup_qoco): + """Test that updating with wrong size matrix data raises error.""" + prob = setup_qoco + + # Test P with wrong size + with pytest.raises(ValueError, match="P size must be"): + prob.update_matrix_data(P=np.array([1, 2, 3])) + + # Test A with wrong size + with pytest.raises(ValueError, match="A size must be"): + prob.update_matrix_data(A=np.array([1, 2])) + + # Test G with wrong size + with pytest.raises(ValueError, match="G size must be"): + prob.update_matrix_data(G=np.array([1, 2])) \ No newline at end of file