Abstract

Generic solver for fix-point (Bellman) equations. This module implements a generic solver for finding solutions to fix-point (Bellman) equations. The module may be used for finding solutions to dynamic programming problems. The implementation is based on John Rust's poly-algorithm, which combines the strategies of successive applications and the Newton-Kantorovich method.

The module is enriched with a version of the solver that uses Futhark's support for automatic differentiation (forward mode AD) to compute the Jacobian for the passed fix-point function.

For more details, see Rust's description in Section 4 of https://editorialexpress.com/jrust/sdp/ndp.pdf and a demonstration of its use in https://elsman.com/pdf/dpsolve-2025-09-27.pdf.

Synopsis

local module type dpsolve = {
type t
type mat [n]
type param = {max_fxpiter: i64, pi_max: i64, pi_tol: t, sa_max: i64, sa_min: i64, sa_tol: t, tol_ratio: t}
val default: param
val sa [m]: (f: ([m]t -> [m]t)) -> (v: [m]t) -> (p: param) -> (b: t) -> {conv: bool, iter: i64, res: [m]t, rtol: t, tol: t}
val nk [m]: (f: ([m]t -> ([m]t, mat [m]))) -> (v: [m]t) -> (p: param) -> {conv: bool, iter: i64, jac: mat [m], res: [m]t, tol: t}
val poly [m]: (f: ([m]t -> ([m]t, mat [m]))) -> (v: [m]t) -> (p: param) -> (b: t) -> {conv: bool, iter_nk: i64, iter_sa: i64, jac: mat [m], res: [m]t, rtrips: i64, tol: t}
val polyad [m]: (f: ([m]t -> [m]t)) -> (v: [m]t) -> (p: param) -> (b: t) -> {conv: bool, iter_nk: i64, iter_sa: i64, res: [m]t, rtrips: i64, tol: t}
}
module type ols_jac = {
type t
type mat [n]
val eye: (n: i64) -> mat [n]
val zero: (n: i64) -> mat [n]
val sub [n]: mat [n] -> mat [n] -> mat [n]
val ols [n]: mat [n] -> [n]t -> [n]t
val wrapj [n]: ([n]t -> [n]t) -> [n]t -> ([n]t, mat [n])
}
module mk_dpsolve: (T: real) -> (ols_jac: ols_jac with t = T.t) -> dpsolve with t = T.t with mat [n] = ols_jac.mat [n]
module mk_dpsolve_dense: (T: real) -> dpsolve with t = T.t with mat [n] = [n][n]T.t

Description

local module type dpsolve

Module type specifying generic solvers for Bellman equations. Notice that this module type is declared local, which means that it cannot directly be referenced by name from the outside. By not exposing the module type, the module may be extended with new members in minor versions.

type t

The scalar type.

type mat [n]

Type of square matrices of size n (e.g., square or dense), with t as scalar type.

type param = {max_fxpiter: i64, pi_max: i64, pi_tol: t, sa_max: i64, sa_min: i64, sa_tol: t, tol_ratio: t}

The type of solver parameters.

val default: param

The default parameter value

val sa [m]: (f: ([m]t -> [m]t)) -> (v: [m]t) -> (p: param) -> (b: t) -> {conv: bool, iter: i64, res: [m]t, rtol: t, tol: t}

Find a fix-point for the function f using successive approximation with the initial guess v, parameter p, and beta-value b. The tolerance, the minimal, and the maximal number of iterations can be adjusted by altering the parameter p. The argument b is the beta discount factor, which allows for stopping early (when the relative tolerance is close to b). The function returns a record containing an approximate fix-point, a boolean specifying whether the algorithm converged (according to the values in p), the number of iterations used, and finally, the tolerance and the relative tolerance of the last two fix-point approximations and the last two tolerances, respectively (maximum of each dimension).

val nk [m]: (f: ([m]t -> ([m]t, mat [m]))) -> (v: [m]t) -> (p: param) -> {conv: bool, iter: i64, jac: mat [m], res: [m]t, tol: t}

Find a fix-point for the function f using Newton-Kantorovich iterations with the initial guess v, and parameter p. The tolerance, the minimal, and the maximal number of iterations can be adjusted by altering the parameter p. The function f should return a pair of a new next approximation and the Jacobian matrix for the function f relative to the argument given. The function returns a record containing an approximate fix-point, a Jacobian matrix for the fix-point, a boolean specifying whether the algorithm converged (according to the values in p), the number of iterations used, and finally, the tolerance of the last two fix-point approximations (maximum of each dimension).

val poly [m]: (f: ([m]t -> ([m]t, mat [m]))) -> (v: [m]t) -> (p: param) -> (b: t) -> {conv: bool, iter_nk: i64, iter_sa: i64, jac: mat [m], res: [m]t, rtrips: i64, tol: t}

Find a fix-point for the function f using a combination of successive approximation iterations and Newton-Kantorovich iterations. The initial guess is v and the parameter p is passed to the calls to sa and nk. The function f should return a pair of a new next approximation and the Jacobian matrix for the function f relative to the argument given. The function returns a record containing an approximate fix-point, a Jacobian matrix for the fix-point, a boolean specifying whether the algorithm converged (according to the values in p), the number of iterations used for the total sa iterations, the total nk iterations, the number of round-trips, and the tolerance of the last two fix-point approximations (maximum of each dimension).

val polyad [m]: (f: ([m]t -> [m]t)) -> (v: [m]t) -> (p: param) -> (b: t) -> {conv: bool, iter_nk: i64, iter_sa: i64, res: [m]t, rtrips: i64, tol: t}

Find a fix-point for the function f using a combination of successive approximation iterations and Newton-Kantorovich iterations. The initial guess is v and the parameter p is passed to the calls to sa and nk. The function uses forward-mode automatic differentiation to compute the Jacobian matrix relative to the argument given to f. The function returns a record containing an approximate fix-point, a boolean specifying whether the algorithm converged (according to the values in p), the number of iterations used for the total sa iterations, the total nk iterations, the number of round-trips, and the tolerance of the last two fix-point approximations (maximum of each dimension).

module type ols_jac

Module type specifying a linear equations solver and functionality for lifting a function to return also the partial derivative (i.e., the Jacobian matrix) relative to the argument.

type t

The scalar type.

type mat [n]

Type of square matrices of size n (e.g., square or dense), with t as scalar type.

val eye: (n: i64) -> mat [n]

The identity matrix of size n x n.

val zero: (n: i64) -> mat [n]

The zero matrix of size n x n.

val sub [n]: mat [n] -> mat [n] -> mat [n]

Element-wise subtraction.

val ols [n]: mat [n] -> [n]t -> [n]t

Solve systems of linear equations Ax = b with respect to x.

val wrapj [n]: ([n]t -> [n]t) -> [n]t -> ([n]t, mat [n])

Augment function result with its partial derivative (Jacobian) relative to the function argument. We have wrap f x = (f x, Df x), where Df x is the Jacobian matrix for f at x.

module mk_dpsolve

Parameterised module for creating generic (e.g., non-linear) solvers. The module is parameterised over an instance of real, such as the module f64 and a compatible module for solving linear equations and computing partial derivatives for a function. One possible instance uses a dense linear equations solver on f64 (e.g., lu f64) and Futhark's AD support.

module mk_dpsolve_dense

Parameterised module for creating generic (e.g., non-linear) solvers using dense Jacobians. The module is parameterised over an instance of real, such as the module f64.