diff --git a/CMakeLists.txt b/CMakeLists.txt index f0e887a..2011a02 100644 --- a/CMakeLists.txt +++ b/CMakeLists.txt @@ -23,7 +23,7 @@ include(FetchContent) FetchContent_Declare( qoco GIT_REPOSITORY https://github.com/qoco-org/qoco.git - GIT_TAG 34d04f9654852567f3c70fd03ee4b22352f844f4 + GIT_TAG main ) list(POP_BACK CMAKE_MESSAGE_INDENT) 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..37b0a38 --- /dev/null +++ b/tests/test_update_data.py @@ -0,0 +1,277 @@ +import qoco +import numpy as np +from scipy.sparse import csc_matrix +import pytest + +abstol = 1e-10 +reltol = 1e-10 + + +@pytest.fixture +def problem_data(): + """Fixture providing standard test problem data.""" + # problem dimensions + n = 6 + m = 6 + p = 2 + l = 3 + nsoc = 1 + + # initial problem data + P = csc_matrix(( + [1, 2, 3, 4, 5, 6], + [0, 1, 2, 3, 4, 5], + [0, 1, 2, 3, 4, 5, 6] + ), + shape=(n, n), + dtype=float, + ) + + A = csc_matrix(( + [1, 1, 1, 2], + [0, 0, 1, 1], + [0, 1, 3, 4, 4, 4, 4] + ), + shape=(p, n), + dtype=float, + ) + + G = csc_matrix(( + [-1, -1, -1, -1, -1, -1], + [0, 1, 2, 3, 4, 5], + [0, 1, 2, 3, 4, 5, 6] + ), + shape=(m, n), + dtype=float, + ) + + c = np.array([1, 2, 3, 4, 5, 6], dtype=float) + b = np.array([1, 2], dtype=float) + h = np.array([0, 0, 0, 0, 0, 0], dtype=float) + q = np.array([3], dtype=float) + + return { + "n": n, + "m": m, + "p": p, + "P": P, + "c": c, + "A": A, + "b": b, + "G": G, + "h": h, + "l": l, + "nsoc": nsoc, + "q": q, + } + + +@pytest.fixture +def problem(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"], + abstol=abstol, + reltol=reltol, + ) + return prob + + +def test_update_vector_data(problem_data, problem): + # setup and solve initial problem + res1 = problem.solve() + assert res1.status == "QOCO_SOLVED" + + # new vector data + c_new = np.array([0, 0, 0, 0, 0, 0], dtype=float) + b_new = np.array([4, 5], dtype=float) + h_new = np.array([1, 1, 1, 1, 1, 1], dtype=float) + + # update solver and solve again + problem.update_vector_data(c=c_new, b=b_new, h=h_new) + res2 = problem.solve() + assert res2.status == "QOCO_SOLVED" + + # reference solution + prob2 = qoco.QOCO() + prob2.setup( + problem_data["n"], + problem_data["m"], + problem_data["p"], + problem_data["P"], + c_new, + problem_data["A"], + b_new, + problem_data["G"], + h_new, + problem_data["l"], + problem_data["nsoc"], + problem_data["q"], + abstol=abstol, + reltol=reltol, + ) + res2_ref = prob2.solve() + assert res2_ref.status == "QOCO_SOLVED" + + assert np.allclose(res2.x, res2_ref.x) + assert np.allclose(res2.s, res2_ref.s) + assert np.allclose(res2.y, res2_ref.y) + assert np.allclose(res2.z, res2_ref.z) + assert np.allclose(res2.obj, res2_ref.obj) + assert np.allclose(res2.dres, res2_ref.dres) + assert np.allclose(res2.pres, res2_ref.pres) + assert np.allclose(res2.gap, res2_ref.gap) + assert np.allclose(res2.iters, res2_ref.iters) + + +def test_update_cost_matrix(problem_data, problem): + # setup and solve initial problem + res1 = problem.solve() + assert res1.status == "QOCO_SOLVED" + + # new matrix data + P_data_new = np.array([6, 5, 4, 3, 2, 1], dtype=float) + P_new = problem_data["P"].copy() + P_new.data = P_data_new.copy() + + # update solver and solve again + problem.update_matrix_data(P=P_data_new) + res2 = problem.solve() + assert res2.status == "QOCO_SOLVED" + + # reference solution + prob2 = qoco.QOCO() + prob2.setup( + problem_data["n"], + problem_data["m"], + problem_data["p"], + P_new, + problem_data["c"], + problem_data["A"], + problem_data["b"], + problem_data["G"], + problem_data["h"], + problem_data["l"], + problem_data["nsoc"], + problem_data["q"], + abstol=abstol, + reltol=reltol, + ) + res2_ref = prob2.solve() + assert res2_ref.status == "QOCO_SOLVED" + + assert np.allclose(res2.x, res2_ref.x) + assert np.allclose(res2.s, res2_ref.s) + assert np.allclose(res2.y, res2_ref.y) + assert np.allclose(res2.z, res2_ref.z) + assert np.allclose(res2.obj, res2_ref.obj) + assert np.allclose(res2.dres, res2_ref.dres) + assert np.allclose(res2.pres, res2_ref.pres) + assert np.allclose(res2.gap, res2_ref.gap) + assert np.allclose(res2.iters, res2_ref.iters) + + +def test_update_constraint_matrix(problem_data, problem): + # setup and solve initial problem + res1 = problem.solve() + assert res1.status == "QOCO_SOLVED" + + # new matrix data + A_data_new = np.array([1, 2, 3, 4], dtype=float) + A_new = problem_data["A"].copy() + A_new.data = A_data_new.copy() + + # update solver and solve again + problem.update_matrix_data(A=A_data_new) + res2 = problem.solve() + assert res2.status == "QOCO_SOLVED" + + # reference solution + prob2 = qoco.QOCO() + prob2.setup( + problem_data["n"], + problem_data["m"], + problem_data["p"], + problem_data["P"], + problem_data["c"], + A_new, + problem_data["b"], + problem_data["G"], + problem_data["h"], + problem_data["l"], + problem_data["nsoc"], + problem_data["q"], + abstol=abstol, + reltol=reltol, + ) + res2_ref = prob2.solve() + assert res2_ref.status == "QOCO_SOLVED" + + assert np.allclose(res2.x, res2_ref.x) + assert np.allclose(res2.s, res2_ref.s) + assert np.allclose(res2.y, res2_ref.y) + assert np.allclose(res2.z, res2_ref.z) + assert np.allclose(res2.obj, res2_ref.obj) + assert np.allclose(res2.dres, res2_ref.dres) + assert np.allclose(res2.pres, res2_ref.pres) + assert np.allclose(res2.gap, res2_ref.gap) + assert np.allclose(res2.iters, res2_ref.iters) + + +def test_update_soc_matrix(problem_data, problem): + # setup and solve initial problem + res1 = problem.solve() + assert res1.status == "QOCO_SOLVED" + + # new matrix data + G_data_new = np.array([-2, -2, -2, -2, -2, -2], dtype=float) + G_new = problem_data["G"].copy() + G_new.data = G_data_new.copy() + + # update solver and solve again + problem.update_matrix_data(G=G_data_new) + res2 = problem.solve() + assert res2.status == "QOCO_SOLVED" + + # reference solution + prob2 = qoco.QOCO() + prob2.setup( + problem_data["n"], + problem_data["m"], + problem_data["p"], + problem_data["P"], + problem_data["c"], + problem_data["A"], + problem_data["b"], + G_new, + problem_data["h"], + problem_data["l"], + problem_data["nsoc"], + problem_data["q"], + abstol=abstol, + reltol=reltol, + ) + res2_ref = prob2.solve() + assert res2_ref.status == "QOCO_SOLVED" + + assert np.allclose(res2.x, res2_ref.x) + assert np.allclose(res2.s, res2_ref.s) + assert np.allclose(res2.y, res2_ref.y) + assert np.allclose(res2.z, res2_ref.z) + assert np.allclose(res2.obj, res2_ref.obj) + assert np.allclose(res2.dres, res2_ref.dres) + assert np.allclose(res2.pres, res2_ref.pres) + assert np.allclose(res2.gap, res2_ref.gap) + assert np.allclose(res2.iters, res2_ref.iters)