r/rust 5d ago

🛠️ project Announcing Basin: A Numerical Optimization Library for Rust

Hi!

I've been working on Basin, a numerical optimization library for Rust. It's heavily inspired by argmin; same overall shape (Executor driver loop, Solver/Problem trait split, per-solver State, pluggable termination). Basin exists because I wanted to push on a few specific design directions that were awkward to retrofit.

What's Different

  • Framework-level termination, bound to state shape. max_iter, the *_tolerance family, max_time, eval budgets all live on the Executor and compose across solvers. Each criterion binds on the minimum state it needs, so asking for a GradientTolerance on a derivative-free solver is a compile error, not a runtime surprise.
  • First-class constraints. Constraints describe the problem, so they live problem-side (not on the executor, never on state). A constrained problem handed to an unconstrained solver is a compile error; there are opt-in adapters (projection/log-barrier/augmented Lagrangian) to wrap unconstrained solvers. Box bounds and linear (in)equalities today; nonlinear is on the roadmap.
  • Backend-generic linear algebra. Solvers are generic over Vec<f64> (no features), nalgebra, ndarray, and faer. A small universal vector tier keeps first-order/derivative-free solvers backend-generic; a richer linalg tier holds matrix ops that LA-heavy solvers bound on by the minimum subset they need.
  • WASM in the default build. No BLAS/LAPACK, no threads, no std::time::Instant in default paths. CI enforces wasm32-unknown-unknown. Parallelism and BLAS-backed paths are opt-in features.

Current Solvers Supported

  • First-order/quasi-Newton: gradient descent (with momentum + pluggable line searches), BFGS, L-BFGS, L-BFGS-B
  • Derivative-free: Nelder--Mead, Brent
  • Nonlinear least squares: Gauss--Newton, Levenberg--Marquardt, trust-region-reflective
  • Global/stochastic: random search, CMA-ES, steady-state GA, memetic combinations
  • Constrained: projected gradient, bounded Nelder--Mead/L-BFGS-B/CMA-ES, log-barrier, augmented Lagrangian

Example

use basin::{BasicState, CostFunction, Executor, Gradient, GradientDescent, GradientTolerance};

struct Rosenbrock;

impl CostFunction for Rosenbrock {
    type Param = Vec<f64>;
    type Output = f64;
    type Error = std::convert::Infallible;

    fn cost(&self, x: &Vec<f64>) -> Result<f64, std::convert::Infallible> {
        Ok((1.0 - x[0]).powi(2) + 100.0 * (x[1] - x[0].powi(2)).powi(2))
    }
}

impl Gradient for Rosenbrock {
    type Gradient = Vec<f64>;

    fn gradient(&self, x: &Vec<f64>) -> Result<Vec<f64>, std::convert::Infallible> {
        Ok(vec![
            -2.0 * (1.0 - x[0]) - 400.0 * x[0] * (x[1] - x[0].powi(2)),
            200.0 * (x[1] - x[0].powi(2)),
        ])
    }
}

let result = Executor::new(Rosenbrock, GradientDescent::new(1e-3), BasicState::new(vec![-1.2, 1.0]))
    .max_iter(50_000)
    .terminate_on(GradientTolerance(1e-6))
    .run()
    .unwrap();

Links

Feedback is very welcome, especially on the trait surface, naming, and any solver/backend combinations you'd want that aren't there yet.

113 Upvotes

28 comments sorted by

View all comments

6

u/Zogzer 5d ago

We have been using argmin for a while now. Biggest difference from your example above is that your CostFunction trait doesn't return a Result of any type. Our cost functions are not perfect and sometimes this happens and we want to exit quickly, is there a reason to this change or are there some type tricks that let us do type Output = Result<f64, ..>;? Thanks

3

u/johlars 4d ago

Ah, yes, no, I decided not to go that route because I think for the majority of use cases it suffices to condition on NaN/+inf returns. Adding `Result` also means you'll need `Ok()` + `?` everywhere. What is your actual use case? You're right that this is a gap, and Basin currently lacks an escape hatch, but I think this should probably end up on the termination criterion side instead.

6

u/Zogzer 4d ago

Our use case is finance/HFT. We need to exit early in cases where we know things are not going to work where we are testing millions of cases in parallel.

Rust doesn't have exceptions so there is no way to bail, but as a result user callbacks should really return result types so you give the option to the user to do that if they want. You will never know what the user will want to do.

An unfortunate choice with argmin was that it was limited to the anyhow error type that makes it expensive to return. If you could add a default type Error = (); associated type to the traits and then return Result<Self::Output, Self::Error> from cost instead, that would allow the flexibility to do either and be a significant improvement over argmin as it is. You are right that you end up with Ok and ? in a few places, but that is expected in rust.

3

u/johlars 4d ago

Ah, I see. Maye you're right, and the idea of a typed error is really nice, and would remove the performance penalty. I like it!