Skip to content

Commit

Permalink
fix edge cases and mistake in solve routine
Browse files Browse the repository at this point in the history
  • Loading branch information
@RSchwan
RSchwan committed Feb 9, 2025
1 parent 7a51452 commit f7e4f06
Show file tree
Hide file tree
Showing 8 changed files with 115 additions and 50 deletions.
22 changes: 22 additions & 0 deletions include/piqp/sparse/blocksparse/block_kkt.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -102,6 +102,28 @@ struct BlockKKT

return *this;
}

void print()
{
for (std::size_t i = 0; i < D.size(); i++) {
if (D[i]) {
printf("D%zu:\n", i);
D[i]->print();
}
}
for (std::size_t i = 0; i < B.size(); i++) {
if (B[i]) {
printf("B%zu:\n", i);
B[i]->print();
}
}
for (std::size_t i = 0; i < E.size(); i++) {
if (E[i]) {
printf("E%zu:\n", i);
E[i]->print();
}
}
}
};

} // namespace sparse
Expand Down
23 changes: 23 additions & 0 deletions include/piqp/sparse/blocksparse/block_mat.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -28,6 +28,7 @@ struct BlockMat
// [ A_{3,3} A_{3,4} A_{3,N} ]
// [ ... ... ]
// [ A_{N-2,N-2} A_{N-2,N-1} A_{N-2,N}]
// [ A_{N-1,N-1} A_{N-1,N}]

// Row permutation from original matrix to block matrix.
// For example, if the original matrix b = A * x, then
Expand Down Expand Up @@ -118,6 +119,28 @@ struct BlockMat

return *this;
}

void print()
{
for (std::size_t i = 0; i < D.size(); i++) {
if (D[i]) {
printf("D%zu:\n", i);
D[i]->print();
}
}
for (std::size_t i = 0; i < B.size(); i++) {
if (B[i]) {
printf("B%zu:\n", i);
B[i]->print();
}
}
for (std::size_t i = 0; i < E.size(); i++) {
if (E[i]) {
printf("E%zu:\n", i);
E[i]->print();
}
}
}
};

} // namespace sparse
Expand Down
7 changes: 7 additions & 0 deletions include/piqp/sparse/blocksparse/block_vec.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -143,6 +143,13 @@ struct BlockVec
other(perm_inv(i++)) = 0;
}
}

void print()
{
for (BlasfeoVec& i : x) {
i.print();
}
}
};

} // namespace sparse
Expand Down
84 changes: 43 additions & 41 deletions include/piqp/sparse/blocksparse_stage_kkt.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -567,8 +567,8 @@ class BlocksparseStageKKT : public KKTSystem<T>
current_block_info.off_diag_block_size = 0;
}

if (at_end) {
// finalize last block
// finalize last block before arrow if it exists
if (at_end && current_block_info.diag_block_size > 0) {
block_info.push_back({current_block_info.diag_block_start, current_block_info.diag_block_size, current_block_info.off_diag_block_size});
assert(current_block_info.off_diag_block_size == 0);

Expand Down Expand Up @@ -684,14 +684,14 @@ class BlocksparseStageKKT : public KKTSystem<T>
if (init) {
A_block.perm.resize(sAT.cols());
A_block.perm_inv.resize(sAT.cols());
A_block.D.resize(N - 2);
A_block.D.resize(N - 1);
A_block.B.resize(N - 2);
A_block.E.resize(N - 2);
A_block.E.resize(N - 1);

// keep track on the current fill status of each block
A_block.block_row_sizes.resize(Eigen::Index(N - 2));
A_block.block_row_sizes.resize(Eigen::Index(N - 1));
A_block.block_row_sizes.setZero();
block_fill_counter.resize(Eigen::Index(N - 2));
block_fill_counter.resize(Eigen::Index(N - 1));
}

// Iterating over a csc transposed matrix corresponds
Expand All @@ -709,7 +709,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
I j = A_row_it.index();
std::size_t block_index = 0;
// find the corresponding block
while (block_info[block_index].start + block_info[block_index].diag_size <= j && block_index + 1 < N - 2) { block_index++; }
while (block_info[block_index].start + block_info[block_index].diag_size <= j && block_index + 1 < N - 1) { block_index++; }
A_block.block_row_sizes(Eigen::Index(block_index))++;
}
}
Expand All @@ -719,7 +719,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
if (init) {
block_row_acc.resize(A_block.block_row_sizes.rows() + 1);
block_row_acc[0] = 0;
for (Eigen::Index i = 0; i < Eigen::Index(N - 2); i++) {
for (Eigen::Index i = 0; i < Eigen::Index(N - 1); i++) {
block_row_acc[i + 1] = block_row_acc[i] + A_block.block_row_sizes[i];
}
}
Expand All @@ -739,7 +739,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
{
I j = A_row_it.index();
// find the corresponding block
while (block_info[block_index].start + block_info[block_index].diag_size <= j && block_index + 1 < N - 2) { block_index++; }
while (block_info[block_index].start + block_info[block_index].diag_size <= j && block_index + 1 < N - 1) { block_index++; }
block_i = block_fill_counter(Eigen::Index(block_index))++;
if (init) {
A_block.perm[i] = block_row_acc[Eigen::Index(block_index)] + block_i;
Expand Down Expand Up @@ -861,7 +861,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
sD.D[i]->setZero();
}

if (i < N - 2 && sA.D[i] && sB.D[i]) {
if (sA.D[i] && sB.D[i]) {
if (allocate) {
if (!sD.D[i]) {
int m = sA.D[i]->rows();
Expand Down Expand Up @@ -897,7 +897,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
sD.D[N-1]->setZero();
}

for (std::size_t i = 0; i < N - 2; i++)
for (std::size_t i = 0; i < N - 1; i++)
{
if (sA.E[i] && sB.E[i]) {
if (allocate) {
Expand Down Expand Up @@ -952,7 +952,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
sD.E[i]->setZero();
}

if (i < N - 2 && sA.E[i] && sB.D[i]) {
if (sA.E[i] && sB.D[i]) {
if (allocate) {
if (!sD.E[i]) {
int m = sA.E[i]->rows();
Expand Down Expand Up @@ -1232,14 +1232,14 @@ class BlocksparseStageKKT : public KKTSystem<T>
#ifdef PIQP_HAS_OPENMP
#pragma omp for nowait
#endif
for (std::size_t i = 0; i < N - 2; i++)
for (std::size_t i = 0; i < N - 1; i++)
{
if (sA.D[i]) {
// sD.D = sA.D * diag(sB)
blasfeo_dgemm_nd(1.0, *sA.D[i], sB.x[i], 0.0, *sD.D[i], *sD.D[i]);
}

if (sA.B[i]) {
if (i < N - 2 && sA.B[i]) {
// sD.B = sA.B * diag(sB)
blasfeo_dgemm_nd(1.0, *sA.B[i], sB.x[i], 0.0, *sD.B[i], *sD.B[i]);
}
Expand Down Expand Up @@ -1421,12 +1421,12 @@ class BlocksparseStageKKT : public KKTSystem<T>
// z = beta * y + alpha * A * x
// here it's assumed that the sparsity of the block matrix
// is transposed without the blocks individually transposed
// A = [A_{1,1} ]
// [A_{1,2} A_{2,2} ]
// [ A_{2,3} A_{3,3} ]
// [ A_{3,4} A_{4,4} A_{N-2,N-2}]
// [ ... A_{N-2,N-1}]
// [A_{1,N} A_{2,N} A_{3,N} ... A_{N-4,N} A_{N-3,N} A_{N-2,N} ]
// A = [A_{1,1} ]
// [A_{1,2} A_{2,2} ]
// [ A_{2,3} A_{3,3} ]
// [ A_{3,4} A_{4,4} ]
// [ ... A_{N-2,N-1} A_{N-1,N-1}]
// [A_{1,N} A_{2,N} A_{3,N} ... A_{N-4,N} A_{N-3,N} A_{N-1,N} ]
void block_t_gemv_n(double alpha, BlockMat<I>& sA, BlockVec& x, double beta, BlockVec& y, BlockVec& z)
{
// z = beta * y
Expand All @@ -1445,7 +1445,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
#endif
for (std::size_t i = 0; i < N - 1; i++)
{
if (i < N - 2 && sA.D[i]) {
if (sA.D[i]) {
int m = sA.D[i]->rows();
int n = sA.D[i]->cols();
assert(x.x[i].rows() == n && "size mismatch");
Expand All @@ -1470,7 +1470,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
#endif
if (arrow_width > 0)
{
for (std::size_t i = 0; i < N - 2; i++)
for (std::size_t i = 0; i < N - 1; i++)
{
if (sA.E[i]) {
int m = sA.E[i]->rows();
Expand All @@ -1494,12 +1494,12 @@ class BlocksparseStageKKT : public KKTSystem<T>
// z = beta * y + alpha * A^T * x
// here it's assumed that the sparsity of the block matrix
// is transposed without the blocks individually transposed
// A = [A_{1,1} ]
// [A_{1,2} A_{2,2} ]
// [ A_{2,3} A_{3,3} ]
// [ A_{3,4} A_{4,4} A_{N-2,N-2}]
// [ ... A_{N-2,N-1}]
// [A_{1,N} A_{2,N} A_{3,N} ... A_{N-4,N} A_{N-3,N} A_{N-2,N} ]
// A = [A_{1,1} ]
// [A_{1,2} A_{2,2} ]
// [ A_{2,3} A_{3,3} ]
// [ A_{3,4} A_{4,4} ]
// [ ... A_{N-2,N-1} A_{N-1,N-1}]
// [A_{1,N} A_{2,N} A_{3,N} ... A_{N-4,N} A_{N-3,N} A_{N-2,N} ]
void block_t_gemv_t(double alpha, BlockMat<I>& sA, BlockVec& x, double beta, BlockVec& y, BlockVec& z)
{
// z = beta * y
Expand All @@ -1516,7 +1516,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
#ifdef PIQP_HAS_OPENMP
#pragma omp for nowait
#endif
for (std::size_t i = 0; i < N - 2; i++)
for (std::size_t i = 0; i < N - 1; i++)
{
if (sA.D[i]) {
int m = sA.D[i]->rows();
Expand All @@ -1527,7 +1527,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
blasfeo_dgemv_t(m, n, alpha, sA.D[i]->ref(), 0, 0, x.x[i].ref(), 0, 1.0, z.x[i].ref(), 0, z.x[i].ref(), 0);
}

if (sA.B[i]) {
if (i < N - 2 && sA.B[i]) {
int m = sA.B[i]->rows();
int n = sA.B[i]->cols();
assert(x.x[i+1].rows() >= m && "size mismatch");
Expand Down Expand Up @@ -1555,12 +1555,12 @@ class BlocksparseStageKKT : public KKTSystem<T>
// z_t = beta_t * y_t + alpha_t * A^T * x_t
// here it's assumed that the sparsity of the block matrix
// is transposed without the blocks individually transposed
// A = [A_{1,1} ]
// [A_{1,2} A_{2,2} ]
// [ A_{2,3} A_{3,3} ]
// [ A_{3,4} A_{4,4} A_{N-2,N-2}]
// [ ... A_{N-2,N-1}]
// [A_{1,N} A_{2,N} A_{3,N} ... A_{N-4,N} A_{N-3,N} A_{N-2,N} ]
// A = [A_{1,1} ]
// [A_{1,2} A_{2,2} ]
// [ A_{2,3} A_{3,3} ]
// [ A_{3,4} A_{4,4} ]
// [ ... A_{N-2,N-1} A_{N-1,N-1}]
// [A_{1,N} A_{2,N} A_{3,N} ... A_{N-4,N} A_{N-3,N} A_{N-2,N} ]
void block_t_gemv_nt(double alpha_n, double alpha_t, BlockMat<I>& sA, BlockVec& x_n, BlockVec& x_t,
double beta_n, double beta_t, BlockVec& y_n, BlockVec& y_t, BlockVec& z_n, BlockVec& z_t)
{
Expand All @@ -1571,7 +1571,7 @@ class BlocksparseStageKKT : public KKTSystem<T>

std::size_t N = block_info.size();
I arrow_width = block_info.back().diag_size;
for (std::size_t i = 0; i < N - 2; i++)
for (std::size_t i = 0; i < N - 1; i++)
{
if (sA.D[i]) {
int m = sA.D[i]->rows();
Expand All @@ -1585,7 +1585,7 @@ class BlocksparseStageKKT : public KKTSystem<T>
blasfeo_dgemv_nt(m, n, alpha_n, alpha_t, sA.D[i]->ref(), 0, 0, x_n.x[i].ref(), 0, x_t.x[i].ref(), 0, 1.0, 1.0, z_n.x[i].ref(), 0, z_t.x[i].ref(), 0, z_n.x[i].ref(), 0, z_t.x[i].ref(), 0);
}

if (sA.B[i]) {
if (i < N - 2 && sA.B[i]) {
int m = sA.B[i]->rows();
int n = sA.B[i]->cols();
assert(x_n.x[i].rows() == n && "size mismatch");
Expand Down Expand Up @@ -1633,11 +1633,13 @@ class BlocksparseStageKKT : public KKTSystem<T>
assert(kkt_fac.D[i]->rows() >= m && "size mismatch");
assert(b_and_x.x[i-1].rows() == n && "size mismatch");
assert(b_and_x.x[i].rows() >= m && "size mismatch");
// y_i = L_i^{-1} * (b_i - C_{i-1} * y_{i-1})
// y_i = b_i - C_{i-1} * y_{i-1}
blasfeo_dgemv_n(m, n, -1.0, kkt_fac.B[i-1]->ref(), 0, 0, b_and_x.x[i-1].ref(), 0, 1.0, b_and_x.x[i].ref(), 0, b_and_x.x[i].ref(), 0);
m = kkt_fac.D[i]->rows();
blasfeo_dtrsv_lnn(m, kkt_fac.D[i]->ref(), 0, 0, b_and_x.x[i].ref(), 0, b_and_x.x[i].ref(), 0);
}
m = kkt_fac.D[i]->rows();
assert(b_and_x.x[i].rows() == m && "size mismatch");
// y_i = L_i^{-1} * y_i
blasfeo_dtrsv_lnn(m, kkt_fac.D[i]->ref(), 0, 0, b_and_x.x[i].ref(), 0, b_and_x.x[i].ref(), 0);
}

if (arrow_width > 0)
Expand Down
5 changes: 5 additions & 0 deletions include/piqp/utils/blasfeo_mat.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -112,6 +112,11 @@ class BlasfeoMat
}

blasfeo_dmat* ref() { return &mat; }

void print()
{
blasfeo_print_dmat(rows(), cols(), ref(), 0, 0);
}
};

} // namespace piqp
Expand Down
5 changes: 5 additions & 0 deletions include/piqp/utils/blasfeo_vec.hpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -111,6 +111,11 @@ class BlasfeoVec
}

blasfeo_dvec* ref() { return &vec; }

void print()
{
blasfeo_print_dvec(rows(), ref(), 0);
}
};

} // namespace piqp
Expand Down
Binary file added tests/data/small_sparse_dual_inf.mat
View file
Binary file not shown.
19 changes: 10 additions & 9 deletions tests/src/sparse/blocksparse_stage_kkt_test.cpp
View file
Original file line number Diff line number Diff line change
Expand Up @@ -58,14 +58,14 @@ void test_solve_multiply(Data<T, I>& data, KKT1& kkt1, KKT2& kkt2)
delta_x_2, delta_y_2, delta_z_2, delta_z_lb_2, delta_z_ub_2, delta_s_2, delta_s_lb_2, delta_s_ub_2);
PIQP_EIGEN_MALLOC_ALLOWED();

ASSERT_TRUE(delta_x_1.isApprox(delta_x_2, 1e-5));
ASSERT_TRUE(delta_y_1.isApprox(delta_y_2, 1e-5));
ASSERT_TRUE(delta_z_1.isApprox(delta_z_2, 1e-5));
ASSERT_TRUE(delta_z_lb_1.head(data.n_lb).isApprox(delta_z_lb_2.head(data.n_lb), 1e-5));
ASSERT_TRUE(delta_z_ub_1.head(data.n_ub).isApprox(delta_z_ub_2.head(data.n_ub), 1e-5));
ASSERT_TRUE(delta_s_1.isApprox(delta_s_2, 1e-5));
ASSERT_TRUE(delta_s_lb_1.head(data.n_lb).isApprox(delta_s_lb_2.head(data.n_lb), 1e-5));
ASSERT_TRUE(delta_s_ub_1.head(data.n_ub).isApprox(delta_s_ub_2.head(data.n_ub), 1e-5));
ASSERT_TRUE(delta_x_1.isApprox(delta_x_2, 1e-8));
ASSERT_TRUE(delta_y_1.isApprox(delta_y_2, 1e-8));
ASSERT_TRUE(delta_z_1.isApprox(delta_z_2, 1e-8));
ASSERT_TRUE(delta_z_lb_1.head(data.n_lb).isApprox(delta_z_lb_2.head(data.n_lb), 1e-8));
ASSERT_TRUE(delta_z_ub_1.head(data.n_ub).isApprox(delta_z_ub_2.head(data.n_ub), 1e-8));
ASSERT_TRUE(delta_s_1.isApprox(delta_s_2, 1e-8));
ASSERT_TRUE(delta_s_lb_1.head(data.n_lb).isApprox(delta_s_lb_2.head(data.n_lb), 1e-8));
ASSERT_TRUE(delta_s_ub_1.head(data.n_ub).isApprox(delta_s_ub_2.head(data.n_ub), 1e-8));

Vec<T> rhs_x_sol_1(data.n);
Vec<T> rhs_y_sol_1(data.p);
Expand Down Expand Up @@ -194,5 +194,6 @@ TEST_P(BlocksparseStageKKTTest, FactorizeSolveSQP)
}

INSTANTIATE_TEST_SUITE_P(FromFolder, BlocksparseStageKKTTest,
::testing::Values("small_dense", "chain_mass_sqp", "robot_arm_sqp",
::testing::Values("small_sparse_dual_inf", "small_dense",
"chain_mass_sqp", "robot_arm_sqp",
"robot_arm_sqp_constr_perm", "robot_arm_sqp_no_global"));

0 comments on commit f7e4f06

Please sign in to comment.