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.

110 Upvotes

28 comments sorted by

22

u/AlivePrune3013 5d ago

nice work on this! the compile-time constraint checking is pretty clever approach - definitely beats finding out at runtime that you messed up solver choice

been wanting something that works better with different la backends so this timing is perfect. how's the performance compared to argmin in your benchmarks? and any plans for adding more global optimization methods beyond cma-es?

the wasm support being default is smart move too, makes it way easier to use in web projects without having to worry about feature flags

9

u/johlars 5d ago

It's seems to be slightly faster than argmin (https://basin.bz/benchmarks/competitors/).

About global optimizers, yes, definitely. Next up is some kind of differential algorithm, I think, but please open a feature request if you have anything in mind that you'd like.

9

u/v_0ver 5d ago

This is a great initiative. I’m currently developing a type-driven SSM framework myself, and I’m excited to see the infrastructure for scientific computing expanding in Rust.

But why edition = "2021" instead of 2024?

16

u/johlars 4d ago

The original reason was that I needed this to work on the MSRV supported by all the machines on CRAN (R package repository), which was 1.81 at the time I started working on Basin. But it's been lifted to 1.91 (edition 2024-compatible) so I can and will bump it :)

5

u/Careful-Nothing-2432 4d ago

Wonder if Claude did that, I’ve noticed it tends to default to setting edition 2021 whenever I have it spin up a new rust workspace for me

6

u/Zogzer 4d 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.

7

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!

1

u/w1th0utnam3 4d ago

Would it be an option to implement what you are describing as `FallibleCostFunction`, require this by the solver and then provide an implementation with `Error = ()` for `T: CostFunction`? Then users who don't need errors can stick with plain `f64` return types etc.? Though I'm not sure for which code the author is concerned about additional `Ok` and `?`.

2

u/Zogzer 4d ago

Ofc it is possible, but you could also go the other way and have InfallibleCostFunction where CostFunction is impl for T: InfallibleCostFunction.

In practice however, I don't think anyone using rust is scared of an additional Ok(). If anything, it's the expected and idiomatic pattern.

2

u/johlars 4d ago

I considered this too, as well as `TryCostFunction`, but in the end I didn't really like the growth in traits and two-pronged path. In the end, the additional `Ok()`, `?`, `unwrap()` and return type verbosity seems like an okay burden to bear given how common it anyway is in Rust. A bonus is that it now looks more like argmin, which might help anyone who wants to migrate.

1

u/johlars 4d ago

Okay, this is a breaking change of course, but I went ahead and did it in https://github.com/jolars/basin/commit/6b46d570e651bc3ce74a0d989e9c16a1bb4f5f1c

Please do take a look and see if this meets your expectations. I haven't released yet.

2

u/Zogzer 3d ago

Yeah, this looks cool and we would give it a go over argmin for sure. The use of std::convert::Infallible is very nice.

1

u/johlars 3d ago

Great! Just remember that Basin is alpha-state, undergoing rapid changes, and nowhere near as battle-tested as argmin is.

4

u/plump_unhappiness 5d ago

The compile-time constraint checking is a solid design move, beats debugging that stuff in production. Curious how the performance stacks up against argmin on the benchmarks you've got linked.

2

u/raoul_lu 5d ago

Hi, this looks great and you seem like you have put a lot of thought into this !
As a mathematician, I'm always happy to hear any projects that have at least something to do with maths too ^^

Anyways, I was wondering why you choose only the MIT license instead of the double license of argmin (and as is the default for most rust projects)?

2

u/johlars 4d ago

To be honest, I didn't at all think about this, but I'd be happy to dual-license it. I'm the sole contributor at this point, so it will most likely not be hard to do, I think.

3

u/johlars 3d ago

It's now dual-licensed! 🚀

1

u/raoul_lu 2d ago

Nice :)

1

u/SamPost 4d ago

Very interesting! How much was Claude able to do from just the posted .md files and rules and such, and how much fine tuning did you have to do?

And I see the Claude skill to ingest numerical research papers, but I didn't see the list of papers. Did I miss that?

1

u/johlars 4d ago

Hm, well, so far it has depended a bit on how well-known the algorithm is. Some algorithms (BFGS etc.) are so well-known that there wasn't really any need to ingest a paper. Other algorithms have been licensed and packaged so that I could do a straight port (L-BFGS-B) and maintain iteration-wise identical behavior, in which case it's also been easy to do via Claude.

But for the more obscure algorithms (CMA-ES with injected solvers etc), it was been tremendously helpful to have either an arXiv latex source or a PDF paper I could parse with Marker and feed to Claude. The papers are ignored on git, so they won't show up there.

Even so, I think it's been essential to have another project (https://eunoia.bz/) that I've been building in tandem with Basin, and switching from argmin and another optimization crate to Basin revealed a lot issues and was very helpful in knowing what I needed from Basin.

1

u/mtj23 3d ago

Great work! I am definitely going to look into using this in some of my projects.

I have a question that I've never had a chance to ask of someone who understands numerical solvers. In my work I do a lot of analysis of 3D shapes, and for the past 15 years I've used various Levenberg Marquardt implementations to perform best-fit alignments between shapes in 2D and 3D. I've been happy with its performance and results.

There's a particular type of measurement in metrology (floating line and surface profiles in ISO 1101) where you're trying to find the minimum zone around an expected surface that encloses all points in a set, and you can move the points around to do so. It ends up being the equivalent of trying to minimize the absolute value of the worst residual.

Do you have a recommendation for an algorithm or set of algorithms I should try for this type of problem? The parameters are either 3 (2d) or 6 (3d) values composed to form an isometry, the samples are a large number of points that get transformed, and the residuals and jacobian are calculated from closest distance searches to some kind of fixed reference geometry.

1

u/johlars 3d ago

Hm, off the top of my head: have you tried a smooth approximation of the loss? That should enable you to keep using the LM algorithm.

1

u/mtj23 3d ago

Help me figure out if I understand you. Is the idea that I would apply some sort of smoothing function to the individual residuals, reweighting them? And that would effectively transform the sum-of-squares into a minimax problem?

1

u/johlars 3d ago

No, sorry, I thought you had a mini-max problem already and wanted a solver for that, in which case you could try using a smooth approximation of that mini-max problem and solve it with LM.

-1

u/[deleted] 4d ago

[deleted]

3

u/johlars 4d ago

Well, I guess that depends slightly on the use case, no?