PoC: Python bindings #340

jjbayer · 2023-03-19T13:32:16Z

This is a PoC for adding Python bindings to argmin. See PR comments for open questions.

Currently, only the Newton solver is implemented.

TODO:

codecov-commenter · 2023-03-19T13:44:54Z

Codecov Report

Attention: 168 lines in your changes are missing coverage. Please review.

Comparison is base (e9bebb2) 90.12% compared to head (ee48f38) 89.36%.

Files	Patch %	Lines
crates/argmin-py/src/executor.rs	0.00%	36 Missing ⚠️
crates/argmin-py/src/problem.rs	0.00%	27 Missing ⚠️
crates/argmin-py/src/solver.rs	0.00%	17 Missing ⚠️
crates/argmin-py/src/lib.rs	0.00%	7 Missing ⚠️
crates/argmin/src/core/test_utils.rs	0.00%	3 Missing ⚠️
crates/argmin/src/solver/brent/brentopt.rs	0.00%	3 Missing ⚠️
crates/argmin/src/solver/brent/brentroot.rs	0.00%	3 Missing ⚠️
crates/argmin/src/solver/conjugategradient/cg.rs	0.00%	3 Missing ⚠️
...rgmin/src/solver/conjugategradient/nonlinear_cg.rs	0.00%	3 Missing ⚠️
...n/src/solver/gaussnewton/gaussnewton_linesearch.rs	0.00%	3 Missing ⚠️
... and 22 more

Additional details and impacted files

@@            Coverage Diff             @@
##             main     #340      +/-   ##
==========================================
- Coverage   90.12%   89.36%   -0.76%     
==========================================
  Files         162      166       +4     
  Lines       19549    19720     +171     
==========================================
+ Hits        17618    17623       +5     
- Misses       1931     2097     +166

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

jjbayer · 2023-03-19T13:54:40Z

argmin/src/core/executor.rs

@@ -26,7 +26,7 @@ pub struct Executor<O, S, I> {
    /// Storage for observers
    observers: Observers<I>,
    /// Checkpoint
-    checkpoint: Option<Box<dyn Checkpoint<S, I>>>,
+    checkpoint: Option<Box<dyn Checkpoint<S, I> + Send>>,


I added these bounds because PyClass must be Send.

Sounds also good to me. Would it make sense to add the Send bound to the Checkpoint trait?

jjbayer · 2023-03-19T14:04:56Z

argmin-py/Cargo.toml

+argmin_testfunctions = "0.1.1"
+argmin = {path="../argmin", default-features=false, features=[]}
+argmin-math = {path="../argmin-math", features=["ndarray_latest-serde"]}
+ndarray-linalg = { version = "0.16", features = ["netlib"] }
+ndarray = { version = "0.15", features = ["serde-1"] }


This section requires cleanup, I am not sure what's the best configuration of features.

This looks mostly fine. I think in the long run at least the serde1 feature of argmin can be enabled because that would allow checkpointing. But I guess checkpointing will need more work anyways.
At some point we will have to decide which BLAS backend to use. This is probably mostly a platform issue (only Intel-MKL works on Linux, Windows and Mac) and a licensing issue since the compiled code will be packaged into a python module.

jjbayer · 2023-03-19T14:07:37Z

argmin-py/src/problem.rs

+        let args = PyTuple::new(py, [param.to_pyarray(py)]);
+        let pyresult = callable.call(py, args, Default::default())?;
+        let pyarray = pyresult.extract::<&numpy::PyArray<Scalar, OutputDimension>>(py)?;
+        // TODO: try to get ownership instead of cloning
+        Ok(pyarray.to_owned_array())


I am unsure what's the overhead of calling to_pyarray and extract for every evaluation of the gradient, hessian etc. Probably needs benchmarks.

to_owned_array makes a copy of the data, this should not be necessary.

I just had a new idea: You're currently using the ndarray backend which requires transitioning between numpy arrays and ndarray arrays. Instead we could add a new math backend based on PyArray, which would mean that numpy would do all the heavy lifting. I'm not sure whether numpy or ndarray is faster and I haven't really thought this through either.

I somehow thought that it would be possible to use the underlying memory in both Rust and Python without copying but that doesn't seem to be the case.

Regarding point 2: I agree. I assumed that there is also into_owned_array but that does not seem to be the case.

I somehow thought that it would be possible to use the underlying memory in both Rust and Python without copying but that doesn't seem to be the case.

I'll give it a try!

jjbayer · 2023-03-19T14:22:50Z

Requires #341.

argmin-py/src/executor.rs

stefan-k · 2023-03-19T20:29:43Z

argmin-py/src/executor.rs

+            let param = kwargs
+                .get_item("param")
+                .map(|x| x.extract::<&PyArray1>())
+                .map_or(Ok(None), |r| r.map(Some))?;
+            let max_iters = kwargs
+                .get_item("max_iters")
+                .map(|x| x.extract())
+                .map_or(Ok(None), |r| r.map(Some))?;


I really like how you transformed the somewhat weird way one needs to set the initial state in argmin (using a closure) into a very pythonic **kwargs thing. This may however turn into quite a chore given that there are multiple kinds of state (IterState and PopulationState) with lots of methods. However, I would label this problem low priority for now.

Makes sense. I might rewrite this when I add more solvers to the python extension.

stefan-k · 2023-03-19T20:39:33Z

argmin-py/src/solver.rs

+}
+
+pub enum DynamicSolver {
+    // NOTE: I tried using a Box<dyn Solver<> here, but Solver is not object safe.


Too bad!! That's what I was hoping for, but I tend to forget about object safety. I don't think it'll be possible to make it object safe :(

stefan-k · 2023-03-19T20:42:45Z

argmin-py/src/solver.rs

+}
+
+impl core::Solver<Problem, IterState> for DynamicSolver {
+    // TODO: make this a trait method so we can return a dynamic


Sounds good to me! We could have both for backwards compatibility, right? The default impl of the name method would then just return self.NAME.

It seems I was able to solve two problems at once: When I remove the associated constant, Solver becomes object-safe, so we can create trait objects for it. Let me know if you objects.

Nice! I thought the generics would also be a problem for object safety, but it's great if this isn't the case. Sounds good to me!

stefan-k · 2023-03-19T20:43:46Z

argmin-py/src/types.rs

+
+//! Base types for the Python extension.
+
+pub type Scalar = f64; // TODO: allow complex numbers


Complex numbers would be great but I wouldn't give this a high priority.

argmin-py/src/executor.rs

jjbayer · 2023-03-26T12:51:31Z

argmin-py/src/executor.rs

+            let param = kwargs
+                .get_item("param")
+                .map(|x| x.extract::<&PyArray1>())
+                .map_or(Ok(None), |r| r.map(Some))?;
+            let max_iters = kwargs
+                .get_item("max_iters")
+                .map(|x| x.extract())
+                .map_or(Ok(None), |r| r.map(Some))?;


Makes sense. I might rewrite this when I add more solvers to the python extension.

jjbayer · 2023-03-26T13:08:47Z

argmin-py/src/problem.rs

+        let args = PyTuple::new(py, [param.to_pyarray(py)]);
+        let pyresult = callable.call(py, args, Default::default())?;
+        let pyarray = pyresult.extract::<&numpy::PyArray<Scalar, OutputDimension>>(py)?;
+        // TODO: try to get ownership instead of cloning
+        Ok(pyarray.to_owned_array())


I somehow thought that it would be possible to use the underlying memory in both Rust and Python without copying but that doesn't seem to be the case.

I'll give it a try!

jjbayer · 2023-03-26T13:09:50Z

argmin-py/src/solver.rs

+}
+
+impl core::Solver<Problem, IterState> for DynamicSolver {
+    // TODO: make this a trait method so we can return a dynamic


It seems I was able to solve two problems at once: When I remove the associated constant, Solver becomes object-safe, so we can create trait objects for it. Let me know if you objects.

jjbayer · 2023-03-26T13:11:34Z

media/book/src/implementing_solver.md

@@ -1,7 +1,7 @@
 # Implementing a solver

 In this section we are going to implement the Landweber solver, which essentially is a special form of gradient descent.
-In iteration \\( k \\), the new parameter vector \\( x_{k+1} \\) is calculated from the previous parameter vector \\( x_k \\) and the gradient at \\( x_k \\) according to the following update rule:
+In iteration \\( k \\), the new parameter vector \\( x\_{k+1} \\) is calculated from the previous parameter vector \\( x_k \\) and the gradient at \\( x_k \\) according to the following update rule:


My autoformatter made most changes to this file.

I wasn't sure if the change in line 4 would render correctly but surprisingly it does.

Python extension: First draft

4c94c51

self-review

bc2e840

jjbayer commented Mar 19, 2023

View reviewed changes

jjbayer requested a review from stefan-k March 19, 2023 14:22

stefan-k reviewed Mar 19, 2023

View reviewed changes

argmin-py/src/executor.rs Outdated Show resolved Hide resolved

stefan-k reviewed Mar 19, 2023

View reviewed changes

Joris Bayer added 4 commits March 26, 2023 13:13

Merge remote-tracking branch 'argmin/main' into argmin-py

4c0c79c

Make Checkpoint and Observe Send

007a13c

Dynamic solver name

f24ea5a

Object-safe Solver

b669e48

jjbayer commented Mar 26, 2023

View reviewed changes

Joris Bayer added 9 commits March 26, 2023 15:14

fmt

f096301

license

839ef5f

Merge branch 'main' into argmin-py

72677e6

repair dependencies

1be6014

Upgrade dependencies

d430569

Format

4cbf7a0

Checkin empty README

61f157e

Exclude argmin-py from wasm build

52030e0

More formatting

ee48f38

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

PoC: Python bindings #340

PoC: Python bindings #340

jjbayer commented Mar 19, 2023 •

edited

codecov-commenter commented Mar 19, 2023 •

edited

jjbayer Mar 19, 2023

stefan-k Mar 19, 2023

jjbayer Mar 19, 2023

stefan-k Mar 19, 2023

jjbayer Mar 19, 2023

stefan-k Mar 19, 2023

jjbayer Mar 26, 2023

jjbayer commented Mar 19, 2023

stefan-k Mar 19, 2023

jjbayer Mar 26, 2023

stefan-k Mar 19, 2023

stefan-k Mar 19, 2023

jjbayer Mar 26, 2023

stefan-k Mar 26, 2023

stefan-k Mar 19, 2023

jjbayer Mar 26, 2023

jjbayer Mar 26, 2023

jjbayer Mar 26, 2023

jjbayer Mar 26, 2023

stefan-k Mar 26, 2023


		//! Base types for the Python extension.

		pub type Scalar = f64; // TODO: allow complex numbers

PoC: Python bindings #340

Are you sure you want to change the base?

PoC: Python bindings #340

Conversation

jjbayer commented Mar 19, 2023 • edited

codecov-commenter commented Mar 19, 2023 • edited

Codecov Report

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jjbayer commented Mar 19, 2023

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

Choose a reason for hiding this comment

jjbayer commented Mar 19, 2023 •

edited

codecov-commenter commented Mar 19, 2023 •

edited