add test for scalar case to be consistent with vector mult

Author: Transurgeon

src/atoms/affine/parameter.c

@@ -64,7 +64,8 @@ expr *new_parameter(int d1, int d2, int param_id, int n_vars, const double *valu
     expr *node = &pnode->base;
     init_expr(node, d1, d2, n_vars, forward, jacobian_init_impl, eval_jacobian,
               is_affine, wsum_hess_init_impl, eval_wsum_hess, NULL);
-
+    
+    // TODO we should assert that the values array has the correct size.
     pnode->param_id = param_id;
 
     if (values == NULL)

src/atoms/affine/scalar_mult.c

@@ -30,9 +30,14 @@ static void forward(expr *node, const double *u)
     expr *child = node->left;
     scalar_mult_expr *snode = (scalar_mult_expr *) node;
 
-    /* call forward for param_source expr tree
-     ex: broadcast(param) or promote(const)*/
-    // snode->param_source->forward(snode->param_source, NULL);
+    /* call forward for param_source expr tree (this extra logic is needed
+       in case the parameter is a broadcast or promote node which needs to refresh
+       its values) */
+    if (snode->base.needs_parameter_refresh)
+    {
+        snode->param_source->forward(snode->param_source, NULL);
+        snode->base.needs_parameter_refresh = false;
+    }
 
     double a = snode->param_source->value[0];
 
@@ -125,5 +130,8 @@ expr *new_scalar_mult(expr *param_node, expr *child)
     mult_node->param_source = param_node;
     expr_retain(param_node);
 
+    /* special case for handling broadcasting of constants correctly */
+    mult_node->base.needs_parameter_refresh = true;
+
     return node;
 }

tests/all_tests.c

@@ -366,6 +366,8 @@ int main(void)
     mu_run_test(test_constant_promote_vector_mult, tests_run);
     mu_run_test(test_param_broadcast_vector_mult, tests_run);
     mu_run_test(test_param_promote_vector_mult, tests_run);
+    mu_run_test(test_const_sum_scalar_mult, tests_run);
+    mu_run_test(test_param_sum_scalar_mult, tests_run);
 #endif /* PROFILE_ONLY */
 
 #ifdef PROFILE_ONLY

tests/problem/test_param_broadcast.h

@@ -175,4 +175,86 @@ const char *test_param_promote_vector_mult(void)
     free_problem(prob);
     return 0;
 }
+
+const char *test_const_sum_scalar_mult(void)
+{
+    int n = 6;
+
+    /* minimize sum(x) subject to sum(c) * x, with c constant */
+    expr *x = new_variable(1, 1, 0, n);
+    expr *objective = new_sum(x, -1);
+    double c_vals[3] = {1.0, 2.0, 3.0};
+    expr *c = new_parameter(1, 3, PARAM_FIXED, n, c_vals);
+    expr *c_sum = new_sum(c, -1);
+    expr *constraint = new_scalar_mult(c_sum, x);
+    expr *constraints[1] = {constraint};
+    problem *prob = new_problem(objective, constraints, 1, false);
+
+    problem_init_derivatives(prob);
+
+    /* point for evaluating */
+    double x_vals[1] = {4.0};
+
+    problem_constraint_forward(prob, x_vals);
+    double constrs[1] = {6.0 * 4.0};
+    problem_jacobian(prob);
+    double jac_x[1] = {6.0};
+    mu_assert("vals fail", cmp_double_array(prob->constraint_values, constrs, 1));
+    mu_assert("vals fail", cmp_double_array(prob->jacobian->x, jac_x, 1));
+
+    mu_assert("check_jacobian failed",
+              check_jacobian_num(constraint, x_vals, NUMERICAL_DIFF_DEFAULT_H));
+    free_problem(prob);
+    return 0;
+}
+
+const char *test_param_sum_scalar_mult(void)
+{
+    int n = 6;
+
+    /* minimize sum(x) subject to sum(p) * x, with p parameter */
+    expr *x = new_variable(1, 1, 0, n);
+    expr *objective = new_sum(x, -1);
+    double c_vals[3] = {1.0, 2.0, 3.0};
+    expr *c = new_parameter(1, 3, 0, n, c_vals);
+    expr *c_sum = new_sum(c, -1);
+    expr *constraint = new_scalar_mult(c_sum, x);
+    expr *constraints[1] = {constraint};
+    problem *prob = new_problem(objective, constraints, 1, false);
+
+    expr *param_nodes[1] = {c};
+    problem_register_params(prob, param_nodes, 1);
+    problem_init_derivatives(prob);
+
+    /* point for evaluating */
+    double x_vals[1] = {4.0};
+
+    problem_constraint_forward(prob, x_vals);
+    double constrs[1] = {6.0 * 4.0};
+    problem_jacobian(prob);
+    double jac_x[1] = {6.0};
+    mu_assert("vals fail", cmp_double_array(prob->constraint_values, constrs, 1));
+    mu_assert("vals fail", cmp_double_array(prob->jacobian->x, jac_x, 1));
+
+    mu_assert("check_jacobian failed",
+              check_jacobian_num(constraint, x_vals, NUMERICAL_DIFF_DEFAULT_H));
+    
+    /* second iteration after updating parameter */
+    double theta[3] = {5.0, 4.0, 3.0};
+    problem_update_params(prob, theta);
+    problem_constraint_forward(prob, x_vals);
+    problem_jacobian(prob);
+    double updated_constrs[1] = {12.0 * 4.0};
+    double updated_jac_x[1] = {12.0};
+    mu_assert("vals fail",
+              cmp_double_array(prob->constraint_values, updated_constrs, 1));
+    mu_assert("vals fail", cmp_double_array(prob->jacobian->x, updated_jac_x, 1));
+
+    mu_assert("check_jacobian failed",
+              check_jacobian_num(constraint, x_vals, NUMERICAL_DIFF_DEFAULT_H));
+
+    free_problem(prob);
+    return 0;
+}
+
 #endif /* TEST_PARAM_BROADCAST_H */