better-optimize 0.1.7__tar.gz → 0.2.0__tar.gz

{better_optimize-0.1.7 → better_optimize-0.2.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: better_optimize
-Version: 0.1.7
+Version: 0.2.0
 Summary: A drop-in replacement for scipy optimize functions with quality of life improvements
 License: MIT License
 Project-URL: Homepage, https://github.com/jessegrabowski/better_optimize

{better_optimize-0.1.7 → better_optimize-0.2.0}/better_optimize/_version.py

@@ -8,11 +8,11 @@ import json
 
 version_json = '''
 {
- "date": "2025-11-05T16:14:59+0100",
+ "date": "2025-12-13T05:36:24+0100",
  "dirty": false,
  "error": null,
- "full-revisionid": "af8fb6def8c11bc4b50fa0a53aadbe8901c390bf",
- "version": "0.1.7"
+ "full-revisionid": "487a1c99b33650b7ba7264b19373ec2d54ca2e1f",
+ "version": "0.2.0"
 }
 '''  # END VERSION_JSON
 

{better_optimize-0.1.7 → better_optimize-0.2.0}/better_optimize/wrapper.py

@@ -122,10 +122,15 @@ class ObjectiveWrapper:
         if not self.progressbar:
             return
 
-        if isinstance(value, list | tuple):
-            value = sum([x**2 for x in value])
+        value = np.asarray(value)
+        if self.root:
+            value = np.linalg.norm(value)
+        elif value.size != 1:
+            raise ValueError(
+                f"Objective function must return a scalar value for minimization, got {value.shape}"
+            )
         else:
-            value = (value**2).sum()
+            value = value.item()
 
         value_dict = {"f_value": value}
         if grad is not None:

{better_optimize-0.1.7 → better_optimize-0.2.0}/better_optimize.egg-info/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: better_optimize
-Version: 0.1.7
+Version: 0.2.0
 Summary: A drop-in replacement for scipy optimize functions with quality of life improvements
 License: MIT License
 Project-URL: Homepage, https://github.com/jessegrabowski/better_optimize

{better_optimize-0.1.7 → better_optimize-0.2.0}/better_optimize.egg-info/SOURCES.txt

@@ -16,6 +16,7 @@ better_optimize.egg-info/dependency_links.txt
 better_optimize.egg-info/requires.txt
 better_optimize.egg-info/top_level.txt
 tests/test_basinhopping.py
+tests/test_jax_integration.py
 tests/test_minimize.py
 tests/test_root.py
 tests/test_utilities.py

better_optimize-0.2.0/tests/test_jax_integration.py

@@ -0,0 +1,68 @@
+import numpy as np
+import pytest
+
+from numpy.testing import assert_allclose
+
+from better_optimize.basinhopping import basinhopping
+from better_optimize.minimize import minimize
+from better_optimize.root import root
+
+jax = pytest.importorskip("jax")
+jnp = pytest.importorskip("jax.numpy")
+
+
+@pytest.fixture(scope="module")
+def jax_1d_quadratic():
+    def f(x):
+        return jnp.sum((x - 1.0) ** 2)
+
+    return jax.value_and_grad(f)
+
+
+def test_minimize_with_jax_array(jax_1d_quadratic):
+    x0 = np.array([2.0])
+
+    res = minimize(jax_1d_quadratic, x0, method="L-BFGS-B")
+
+    assert res.success
+    assert_allclose(res.x, np.array([1.0]))
+
+    f_val, _ = jax_1d_quadratic(res.x)
+    assert isinstance(f_val, jax.Array)
+    assert_allclose(np.asarray(f_val), 0.0, atol=1e-6, rtol=1e-6)
+
+
+def test_root_with_jax_array():
+    def f(x):
+        return jnp.array(x) - 2.0
+
+    x0 = np.array([0.0])
+
+    res = root(f, x0, method="hybr")
+
+    assert res.success
+    assert_allclose(res.x, np.array([2.0]), atol=1e-4, rtol=1e-4)
+    f_val = f(res.x)
+    assert isinstance(f_val, jax.Array)
+    assert_allclose(np.asarray(f_val), 0.0, atol=1e-6, rtol=1e-6)
+
+
+def test_basinhopping_with_jax_array(jax_1d_quadratic):
+    x0 = np.array([3.0])
+
+    res = basinhopping(
+        jax_1d_quadratic,
+        x0=x0,
+        minimizer_kwargs={"method": "L-BFGS-B", "tol": 1e-6},
+        niter=20,
+        progressbar=False,
+    )
+
+    assert res.success
+    assert_allclose(res.x, np.array([1.0]), atol=1e-3, rtol=1e-3)
+
+    val, grad = jax_1d_quadratic(res.x)
+    assert isinstance(val, jax.Array)
+    assert isinstance(grad, jax.Array)
+    assert_allclose(np.asarray(val), 0.0, atol=1e-5, rtol=1e-5)
+    assert_allclose(np.asarray(grad), np.array([0.0]), atol=1e-5, rtol=1e-5)