diff --git a/include/subexpr.h b/include/subexpr.h index 53fca59..cdd2930 100644 --- a/include/subexpr.h +++ b/include/subexpr.h @@ -113,7 +113,7 @@ typedef struct hstack_expr expr base; expr **args; int n_args; - CSR_matrix *CSR_work; /* for summing Hessians of children */ + matrix **child_hess; /* args' Hessians: the operands of the sum */ } hstack_expr; /* Elementwise multiplication */ diff --git a/src/atoms/affine/hstack.c b/src/atoms/affine/hstack.c index c1b5e51..be45c4f 100644 --- a/src/atoms/affine/hstack.c +++ b/src/atoms/affine/hstack.c @@ -16,9 +16,9 @@ * limitations under the License. */ #include "atoms/affine.h" -#include "int_utils.h" -#include "sparsediff_linalg/csr_kernels.h" +#include "sparsediff_linalg/CSR_matrix.h" #include "sparsediff_linalg/matrix.h" +#include "sparsediff_linalg/matrix_sum.h" #include "sparsediff_linalg/tracked_alloc.h" #include "subexpr.h" #include @@ -106,52 +106,30 @@ static void eval_jacobian_impl(expr *node) static void wsum_hess_init_impl(expr *node) { - /* initialize children's hessians */ + /* initialize children's hessians; this node's is their sum */ hstack_expr *hnode = (hstack_expr *) node; - int nnz = 0; + hnode->child_hess = (matrix **) sp_malloc(hnode->n_args * sizeof(matrix *)); for (int i = 0; i < hnode->n_args; i++) { wsum_hess_init(hnode->args[i]); - nnz += hnode->args[i]->wsum_hess->nnz; + hnode->child_hess[i] = hnode->args[i]->wsum_hess; } - - /* worst-case scenario the nnz of node->wsum_hess is the sum of children's - nnz, capped by the output cell count */ - int nnz_ub = MIN(nnz, sat_mul_int(node->n_vars, node->n_vars)); - CSR_matrix *H = new_CSR_matrix(node->n_vars, node->n_vars, nnz_ub); - hnode->CSR_work = new_CSR_matrix(node->n_vars, node->n_vars, nnz_ub); - - /* fill sparsity pattern */ - for (int i = 0; i < hnode->n_args; i++) - { - matrix *child_hess = hnode->args[i]->wsum_hess; - copy_CSR_matrix(H, hnode->CSR_work); - sum_csr_alloc(hnode->CSR_work, child_hess->to_csr(child_hess), H); - } - - /* trim both buffers to the final pattern size; CSR_work must be re-synced - from H first, since its row pointers still describe an older, smaller - pattern */ - CSR_trim(H); - copy_CSR_matrix(H, hnode->CSR_work); - CSR_trim(hnode->CSR_work); - node->wsum_hess = new_sparse_matrix(H); + node->wsum_hess = sum_n_matrices_alloc(hnode->child_hess, hnode->n_args); } static void wsum_hess_eval(expr *node, const double *w) { hstack_expr *hnode = (hstack_expr *) node; - CSR_matrix *H = node->wsum_hess->to_csr(node->wsum_hess); int row_offset = 0; - memset(H->x, 0, H->nnz * sizeof(double)); + /* add each child's hessian right after evaluating it: args may share a + subexpression, whose later evaluation would overwrite earlier values */ + sum_n_matrices_fill_zero(node->wsum_hess); for (int i = 0; i < hnode->n_args; i++) { expr *child = hnode->args[i]; eval_wsum_hess(child, w + row_offset); - copy_CSR_matrix(H, hnode->CSR_work); - sum_csr_fill_values(hnode->CSR_work, - child->wsum_hess->to_csr(child->wsum_hess), H); + sum_n_matrices_fill_add(child->wsum_hess, i, node->wsum_hess); row_offset += child->size; } } @@ -192,8 +170,8 @@ static void free_type_data(expr *node) hnode->args[i] = NULL; } - free_CSR_matrix(hnode->CSR_work); - hnode->CSR_work = NULL; + sp_free(hnode->child_hess); + hnode->child_hess = NULL; sp_free(hnode->args); hnode->args = NULL; } diff --git a/tests/all_tests.c b/tests/all_tests.c index 61b40d1..f7ec62c 100644 --- a/tests/all_tests.c +++ b/tests/all_tests.c @@ -319,6 +319,7 @@ int main(void) mu_run_test(test_wsum_hess_rel_entr_scalar_vector, tests_run); mu_run_test(test_wsum_hess_hstack, tests_run); mu_run_test(test_wsum_hess_hstack_matrix, tests_run); + mu_run_test(test_wsum_hess_hstack_shared_arg, tests_run); mu_run_test(test_wsum_hess_vstack_vectors, tests_run); mu_run_test(test_wsum_hess_vstack_matrix, tests_run); mu_run_test(test_wsum_hess_index_log, tests_run); diff --git a/tests/wsum_hess/affine/test_hstack.h b/tests/wsum_hess/affine/test_hstack.h index 2b6666b..0f2db28 100644 --- a/tests/wsum_hess/affine/test_hstack.h +++ b/tests/wsum_hess/affine/test_hstack.h @@ -204,3 +204,31 @@ const char *test_wsum_hess_hstack_matrix(void) free_expr(hstack_node); return 0; } + +const char *test_wsum_hess_hstack_shared_arg(void) +{ + /* hstack([log(x), log(x)]) with the same log(x) node in both slots: each + * slot gets its own weights, so H = diag(-(w[k] + w[k + 3]) / x_k^2). */ + double u_vals[3] = {1.0, 2.0, 3.0}; + double w[6] = {1.0, 2.0, 3.0, 4.0, 5.0, 6.0}; + + expr *x = new_variable(3, 1, 0, 3); + expr *log_x = new_log(x); + expr *args[2] = {log_x, log_x}; + expr *hstack_node = new_hstack(args, 2, 3); + + hstack_node->forward(hstack_node, u_vals); + jacobian_init(hstack_node); + wsum_hess_init(hstack_node); + eval_wsum_hess(hstack_node, w); + + int expected_p[4] = {0, 1, 2, 3}; + int expected_i[3] = {0, 1, 2}; + double expected_x[3] = {-5.0, -7.0 / 4.0, -9.0 / 9.0}; + mu_assert("sparsity fail", + cmp_sparsity(hstack_node->wsum_hess, expected_p, expected_i, 3, 3)); + mu_assert("vals fail", cmp_values(hstack_node->wsum_hess, expected_x, 3)); + + free_expr(hstack_node); + return 0; +}