Skip to content

Commit 8d84867

Browse files
JonasBreulingJonasBreuling
authored
Add bindings for update_vector_data and update_matrix_data. (#4)
1 parent 6a6f53b commit 8d84867

4 files changed

Lines changed: 432 additions & 4 deletions

File tree

‎CMakeLists.txt‎

Lines changed: 1 addition & 1 deletion
Original file line numberDiff line numberDiff line change
@@ -23,7 +23,7 @@ include(FetchContent)
2323
FetchContent_Declare(
2424
qoco
2525
GIT_REPOSITORY https://github.com/qoco-org/qoco.git
26-
GIT_TAG 34d04f9654852567f3c70fd03ee4b22352f844f4
26+
GIT_TAG main
2727
)
2828

2929
list(POP_BACK CMAKE_MESSAGE_INDENT)

‎src/bindings.cpp.in‎

Lines changed: 77 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -139,9 +139,9 @@ public:
139139
QOCOSettings *get_settings();
140140
PyQOCOSolution &get_solution();
141141

142-
QOCOInt update_settings(const QOCOSettings &);
143-
// QOCOInt update_vector_data(py::object, py::object, py::object);
144-
// QOCOInt update_matrix_data(py::object, py::object, py::object);
142+
QOCOInt update_settings(const QOCOSettings &new_settings);
143+
void update_vector_data(py::object cnew, py::object bnew, py::object hnew);
144+
void update_matrix_data(py::object Pxnew, py::object Axnew, py::object Gxnew);
145145

146146
QOCOInt solve();
147147

@@ -241,6 +241,78 @@ QOCOInt PyQOCOSolver::update_settings(const QOCOSettings &new_settings)
241241
return qoco_update_settings(this->_solver, &new_settings);
242242
}
243243

244+
void PyQOCOSolver::update_vector_data(py::object cnew, py::object bnew, py::object hnew)
245+
{
246+
QOCOFloat *cnew_ptr = nullptr;
247+
QOCOFloat *bnew_ptr = nullptr;
248+
QOCOFloat *hnew_ptr = nullptr;
249+
250+
if (cnew != py::none())
251+
{
252+
auto cnew_arr = cnew.cast<py::array_t<QOCOFloat>>();
253+
auto buf = cnew_arr.request();
254+
if (buf.shape[0] != this->n)
255+
throw std::runtime_error("cnew size must be n = " + std::to_string(this->n));
256+
cnew_ptr = (QOCOFloat *)buf.ptr;
257+
}
258+
259+
if (bnew != py::none())
260+
{
261+
auto bnew_arr = bnew.cast<py::array_t<QOCOFloat>>();
262+
auto buf = bnew_arr.request();
263+
if (buf.shape[0] != this->p)
264+
throw std::runtime_error("bnew size must be p = " + std::to_string(this->p));
265+
bnew_ptr = (QOCOFloat *)buf.ptr;
266+
}
267+
268+
if (hnew != py::none())
269+
{
270+
auto hnew_arr = hnew.cast<py::array_t<QOCOFloat>>();
271+
auto buf = hnew_arr.request();
272+
if (buf.shape[0] != this->m)
273+
throw std::runtime_error("hnew size must be m = " + std::to_string(this->m));
274+
hnew_ptr = (QOCOFloat *)buf.ptr;
275+
}
276+
277+
qoco_update_vector_data(this->_solver, cnew_ptr, bnew_ptr, hnew_ptr);
278+
}
279+
280+
void PyQOCOSolver::update_matrix_data(py::object Pxnew, py::object Axnew, py::object Gxnew)
281+
{
282+
QOCOFloat *Pxnew_ptr = nullptr;
283+
QOCOFloat *Axnew_ptr = nullptr;
284+
QOCOFloat *Gxnew_ptr = nullptr;
285+
286+
if (Pxnew != py::none())
287+
{
288+
auto Pxnew_arr = Pxnew.cast<py::array_t<QOCOFloat>>();
289+
auto buf = Pxnew_arr.request();
290+
if (buf.ndim != 1)
291+
throw std::runtime_error("Pxnew must be 1-D array");
292+
Pxnew_ptr = (QOCOFloat *)buf.ptr;
293+
}
294+
295+
if (Axnew != py::none())
296+
{
297+
auto Axnew_arr = Axnew.cast<py::array_t<QOCOFloat>>();
298+
auto buf = Axnew_arr.request();
299+
if (buf.ndim != 1)
300+
throw std::runtime_error("Axnew must be 1-D array");
301+
Axnew_ptr = (QOCOFloat *)buf.ptr;
302+
}
303+
304+
if (Gxnew != py::none())
305+
{
306+
auto Gxnew_arr = Gxnew.cast<py::array_t<QOCOFloat>>();
307+
auto buf = Gxnew_arr.request();
308+
if (buf.ndim != 1)
309+
throw std::runtime_error("Gxnew must be 1-D array");
310+
Gxnew_ptr = (QOCOFloat *)buf.ptr;
311+
}
312+
313+
qoco_update_matrix_data(this->_solver, Pxnew_ptr, Axnew_ptr, Gxnew_ptr);
314+
}
315+
244316
PYBIND11_MODULE(@QOCO_EXT_MODULE_NAME@, m)
245317
{
246318
// Enums.
@@ -308,6 +380,8 @@ PYBIND11_MODULE(@QOCO_EXT_MODULE_NAME@, m)
308380
.def(py::init<QOCOInt, QOCOInt, QOCOInt, const CSC &, const py::array_t<QOCOFloat>, const CSC &, const py::array_t<QOCOFloat>, const CSC &, const py::array_t<QOCOFloat>, QOCOInt, QOCOInt, const py::array_t<QOCOInt>, 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)
309381
.def_property_readonly("solution", &PyQOCOSolver::get_solution, py::return_value_policy::reference)
310382
.def("update_settings", &PyQOCOSolver::update_settings)
383+
.def("update_vector_data", &PyQOCOSolver::update_vector_data, "cnew"_a=py::none(), "bnew"_a=py::none(), "hnew"_a=py::none())
384+
.def("update_matrix_data", &PyQOCOSolver::update_matrix_data, "Pxnew"_a=py::none(), "Axnew"_a=py::none(), "Gxnew"_a=py::none())
311385
.def("solve", &PyQOCOSolver::solve)
312386
.def("get_settings", &PyQOCOSolver::get_settings, py::return_value_policy::reference);
313387
}

‎src/qoco/interface.py‎

Lines changed: 77 additions & 0 deletions
Original file line numberDiff line numberDiff line change
@@ -86,6 +86,83 @@ def update_settings(self, **kwargs):
8686
if settings_changed and self._solver is not None:
8787
self._solver.update_settings(self.settings)
8888

89+
def update_vector_data(self, c=None, b=None, h=None):
90+
"""
91+
Update data vectors.
92+
93+
Parameters
94+
----------
95+
c : np.ndarray, optional
96+
New c vector of size n. If None, c is not updated. Default is None.
97+
b : np.ndarray, optional
98+
New b vector of size p. If None, b is not updated. Default is None.
99+
h : np.ndarray, optional
100+
New h vector of size m. If None, h is not updated. Default is None.
101+
"""
102+
if c is not None:
103+
if not isinstance(c, np.ndarray):
104+
c = np.array(c)
105+
c = c.astype(np.float64)
106+
if c.shape[0] != self.n:
107+
raise ValueError(f"c size must be n = {self.n}")
108+
109+
if b is not None:
110+
if not isinstance(b, np.ndarray):
111+
b = np.array(b)
112+
b = b.astype(np.float64)
113+
if b.shape[0] != self.p:
114+
raise ValueError(f"b size must be p = {self.p}")
115+
116+
if h is not None:
117+
if not isinstance(h, np.ndarray):
118+
h = np.array(h)
119+
h = h.astype(np.float64)
120+
if h.shape[0] != self.m:
121+
raise ValueError(f"h size must be m = {self.m}")
122+
123+
return self._solver.update_vector_data(c, b, h)
124+
125+
def update_matrix_data(self, P=None, A=None, G=None):
126+
"""
127+
Update sparse matrix data.
128+
129+
The new matrices must have the same sparsity structure as the original ones.
130+
131+
Parameters
132+
----------
133+
P : np.ndarray, optional
134+
New data for P matrix (only the nonzero values). If None, P is not updated.
135+
Default is None.
136+
A : np.ndarray, optional
137+
New data for A matrix (only the nonzero values). If None, A is not updated.
138+
Default is None.
139+
G : np.ndarray, optional
140+
New data for G matrix (only the nonzero values). If None, G is not updated.
141+
Default is None.
142+
"""
143+
if P is not None:
144+
if not isinstance(P, np.ndarray):
145+
P = np.array(P)
146+
P = P.astype(np.float64)
147+
if P.shape[0] != self.P.nnz:
148+
raise ValueError(f"P size must be {self.P.nnz}")
149+
150+
if A is not None:
151+
if not isinstance(A, np.ndarray):
152+
A = np.array(A)
153+
A = A.astype(np.float64)
154+
if A.shape[0] != self.A.nnz:
155+
raise ValueError(f"A size must be {self.A.nnz}")
156+
157+
if G is not None:
158+
if not isinstance(G, np.ndarray):
159+
G = np.array(G)
160+
G = G.astype(np.float64)
161+
if G.shape[0] != self.G.nnz:
162+
raise ValueError(f"G size must be {self.G.nnz}")
163+
164+
return self._solver.update_matrix_data(P, A, G)
165+
89166
def setup(self, n, m, p, P, c, A, b, G, h, l, nsoc, q, **settings):
90167
self.m = m
91168
self.n = n

0 commit comments

Comments
 (0)