qudit_inst/numerical/
function.rs

1//! This module defines traits for cost functions, gradients, Hessians, residual functions, and Jacobians.
2use faer::{Col, ColMut, Mat, MatMut, Row, RowMut};
3use qudit_core::RealScalar;
4
5pub trait Function {
6    fn num_params(&self) -> usize;
7}
8
9/// A trait for cost functions.
10pub trait CostFunction<R: RealScalar>: Function {
11    /// Calculates the cost for the given parameters.
12    fn cost(&mut self, params: &[R]) -> R;
13}
14
15/// A trait for gradients.
16pub trait Gradient<R: RealScalar>: CostFunction<R> {
17    fn allocate_gradient(&self) -> Row<R> {
18        Row::zeros(self.num_params())
19    }
20
21    /// Calculates the gradient for the given parameters.
22    fn gradient_into(&mut self, params: &[R], grad_out: RowMut<R>);
23
24    /// Calculates the cost and gradient for the given parameters.
25    fn cost_and_gradient_into(&mut self, params: &[R], grad_out: RowMut<R>) -> R {
26        let cost = self.cost(params);
27        self.gradient_into(params, grad_out);
28        cost
29    }
30
31    // TODO: non-into methods that allocate buffers and return owned values
32    // TODO: batch_cost_and_gradient.. and others
33}
34
35/// A trait for Hessians.
36pub trait Hessian<R: RealScalar>: Gradient<R> {
37    fn allocate_hessian(&self) -> Mat<R> {
38        Mat::zeros(self.num_params(), self.num_params())
39    }
40
41    /// Calculates the Hessian for the given parameters.
42    fn hessian_into(&mut self, params: &[R], hess_out: MatMut<R>);
43
44    /// Calculates the cost, gradient, and Hessian for the given parameters.
45    fn cost_gradient_and_hessian_into(
46        &mut self,
47        params: &[R],
48        grad_out: RowMut<R>,
49        hess_out: MatMut<R>,
50    ) -> R {
51        let cost = self.cost(params);
52        self.gradient_into(params, grad_out);
53        self.hessian_into(params, hess_out);
54        cost
55    }
56}
57
58/// A trait for residual functions.
59pub trait ResidualFunction<R: RealScalar>: Function {
60    fn num_residuals(&self) -> usize;
61
62    fn allocate_residual(&self) -> Col<R> {
63        Col::zeros(self.num_residuals())
64    }
65
66    /// Calculates the residuals for the given parameters.
67    fn residuals_into(&mut self, params: &[R], residuals_out: ColMut<R>);
68}
69
70/// A trait for Jacobians.
71pub trait Jacobian<R: RealScalar>: ResidualFunction<R> {
72    fn allocate_jacobian(&self) -> Mat<R> {
73        Mat::zeros(self.num_residuals(), self.num_params())
74    }
75
76    /// Calculates the Jacobian for the given parameters.
77    fn jacobian_into(&mut self, params: &[R], jacobian_out: MatMut<R>);
78
79    /// Calculates the residuals and Jacobian for the given parameters.
80    fn residuals_and_jacobian_into(
81        &mut self,
82        params: &[R],
83        residuals_out: ColMut<R>,
84        jacobian_out: MatMut<R>,
85    ) {
86        self.residuals_into(params, residuals_out);
87        self.jacobian_into(params, jacobian_out);
88    }
89}