use arael::model::{Model, Param, SelfBlock};
use arael::simple_lm::LmProblem;
use arael_sym::E;
#[arael::function]
fn square(x: E) -> E {
x * x
}
#[arael::model]
#[arael(root, jacobian)]
#[arael(constraint(hb, name = "sq_x_eq_9", {
// residual = square(m.x) - 9.0, weighted by isigma
[(square(m.x) - 9.0) * m.isigma]
}))]
struct M {
x: Param<f64>,
isigma: f64,
hb: SelfBlock<M>,
}
fn make() -> (M, Vec<f64>) {
let mut m = M { x: Param::new(2.0), isigma: 1.0, hb: SelfBlock::new() };
let mut p = Vec::new();
m.serialize64(&mut p);
(m, p)
}
#[test]
fn square_user_fn_forward_and_grad() {
let (mut m, params) = make();
let cost = m.calc_cost(¶ms);
assert!((cost - 25.0).abs() < 1e-9, "cost {} != 25", cost);
let n = params.len();
let mut g = vec![0.0f64; n];
let mut h = vec![0.0f64; n * n];
m.calc_grad_hessian_dense(¶ms, &mut g, &mut h);
let eps = 1e-5_f64;
let mut pp = params.clone();
pp[0] += eps; let cp = m.calc_cost(&pp);
pp[0] -= 2.0 * eps; let cm = m.calc_cost(&pp);
let ng = (cp - cm) / (2.0 * eps);
assert!((g[0] - ng).abs() < 1e-3, "grad[0]: analytic={} numerical={}", g[0], ng);
assert!((g[0] - (-40.0)).abs() < 1e-6, "grad[0] = {} (expected -40)", g[0]);
}
#[test]
fn square_user_fn_runtime_callable() {
let x = arael_sym::symbol("x");
let e = square(x);
let s = format!("{}", e);
assert!(s.contains("x") && (s.contains("*") || s.contains("^")),
"square(x) printed as {:?}", s);
}