invrs-opt 0.1.4__py3-none-any.whl → 0.2.0__py3-none-any.whl

invrs_opt/__init__.py

@@ -3,8 +3,8 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.1.4"
+__version__ = "v0.2.0"
 __author__ = "Martin F. Schubert <mfschubert@gmail.com>"
 
-from invrs_opt.lbfgsb.lbfgsb import lbfgsb as lbfgsb
 from invrs_opt.lbfgsb.lbfgsb import density_lbfgsb as density_lbfgsb
+from invrs_opt.lbfgsb.lbfgsb import lbfgsb as lbfgsb

invrs_opt/lbfgsb/lbfgsb.py

@@ -10,20 +10,19 @@ from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
 import jax
 import jax.numpy as jnp
 import numpy as onp
-from jax import flatten_util
-from jax import tree_util
+from jax import flatten_util, tree_util
 from scipy.optimize._lbfgsb_py import (  # type: ignore[import-untyped]
     _lbfgsb as scipy_lbfgsb,
 )
+from totypes import types
 
-from invrs_opt.lbfgsb import transform
 from invrs_opt import base
-from totypes import types
+from invrs_opt.lbfgsb import transform
 
 NDArray = onp.ndarray[Any, Any]
 PyTree = Any
 ElementwiseBound = Union[NDArray, Sequence[Optional[float]]]
-LbfgsbState = Tuple[PyTree, Dict[str, NDArray]]
+LbfgsbState = Tuple[PyTree, Dict[str, jnp.ndarray]]
 
 
 # Task message prefixes for the underlying L-BFGS-B implementation.
@@ -188,18 +187,24 @@ def transformed_lbfgsb(
 
     def init_fn(params: PyTree) -> LbfgsbState:
         """Initializes the optimization state."""
-        lower_bound = types.extract_lower_bound(params)
-        upper_bound = types.extract_upper_bound(params)
-        scipy_lbfgsb_state = ScipyLbfgsbState.init(
-            x0=_to_numpy(params),
-            lower_bound=_bound_for_params(lower_bound, params),
-            upper_bound=_bound_for_params(upper_bound, params),
-            maxcor=maxcor,
-            line_search_max_steps=line_search_max_steps,
+
+        def _init_pure(params: PyTree) -> LbfgsbState:
+            lower_bound = types.extract_lower_bound(params)
+            upper_bound = types.extract_upper_bound(params)
+            scipy_lbfgsb_state = ScipyLbfgsbState.init(
+                x0=_to_numpy(params),
+                lower_bound=_bound_for_params(lower_bound, params),
+                upper_bound=_bound_for_params(upper_bound, params),
+                maxcor=maxcor,
+                line_search_max_steps=line_search_max_steps,
+            )
+            latent_params = _to_pytree(scipy_lbfgsb_state.x, params)
+            params = transform_fn(latent_params)
+            return params, scipy_lbfgsb_state.to_jax()
+
+        return jax.pure_callback(  # type: ignore[no-any-return, attr-defined]
+            _init_pure, _example_state(params, maxcor), params
         )
-        latent_params = _to_pytree(scipy_lbfgsb_state.x, params)
-        params = transform_fn(latent_params)
-        return (params, dataclasses.asdict(scipy_lbfgsb_state))
 
     def params_fn(state: LbfgsbState) -> PyTree:
         """Returns the parameters for the given `state`."""
@@ -214,23 +219,30 @@ def transformed_lbfgsb(
         state: LbfgsbState,
     ) -> LbfgsbState:
         """Updates the state."""
-        del params
-        params, lbfgsb_state_dict = state
-        # Avoid in-place updates.
-        lbfgsb_state_dict = copy.deepcopy(lbfgsb_state_dict)
-        scipy_lbfgsb_state = ScipyLbfgsbState(
-            **lbfgsb_state_dict  # type: ignore[arg-type]
-        )
 
-        latent_params = _to_pytree(scipy_lbfgsb_state.x, params)
-        _, vjp_fn = jax.vjp(transform_fn, latent_params)
-        (latent_grad,) = vjp_fn(grad)
+        def _update_pure(
+            grad: PyTree, value: float, params: PyTree, state: LbfgsbState
+        ) -> LbfgsbState:
+            del params
+
+            params, jax_lbfgsb_state = state
+            scipy_lbfgsb_state = ScipyLbfgsbState.from_jax(jax_lbfgsb_state)
 
-        assert onp.size(value) == 1
-        scipy_lbfgsb_state.update(grad=_to_numpy(latent_grad), value=onp.asarray(value))
-        latent_params = _to_pytree(scipy_lbfgsb_state.x, params)
-        params = transform_fn(latent_params)
-        return (params, dataclasses.asdict(scipy_lbfgsb_state))
+            latent_params = _to_pytree(scipy_lbfgsb_state.x, params)
+            _, vjp_fn = jax.vjp(transform_fn, latent_params)
+            (latent_grad,) = vjp_fn(grad)
+
+            assert onp.size(value) == 1
+            scipy_lbfgsb_state.update(
+                grad=_to_numpy(latent_grad), value=onp.asarray(value)
+            )
+            latent_params = _to_pytree(scipy_lbfgsb_state.x, params)
+            params = transform_fn(latent_params)
+            return params, scipy_lbfgsb_state.to_jax()
+
+        return jax.pure_callback(  # type: ignore[no-any-return, attr-defined]
+            _update_pure, state, grad, value, params, state
+        )
 
     return base.Optimizer(
         init=init_fn,
@@ -331,6 +343,29 @@ def _bound_for_params(bound: PyTree, params: PyTree) -> ElementwiseBound:
     return bound_flat
 
 
+def _example_state(params: PyTree, maxcor: int) -> PyTree:
+    """Return an example state for the given `params` and `maxcor`."""
+    params_flat, _ = flatten_util.ravel_pytree(params)  # type: ignore[no-untyped-call]
+    n = params_flat.size
+    float_params = tree_util.tree_map(lambda x: jnp.asarray(x, dtype=float), params)
+    example_jax_lbfgsb_state = dict(
+        x=jnp.zeros(n, dtype=float),
+        _maxcor=jnp.zeros((), dtype=int),
+        _line_search_max_steps=jnp.zeros((), dtype=int),
+        _wa=jnp.ones(_wa_size(n=n, maxcor=maxcor), dtype=float),
+        _iwa=jnp.ones(n * 3, dtype=jnp.int32),  # Fortran int
+        _task=jnp.zeros(59, dtype=int),
+        _csave=jnp.zeros(59, dtype=int),
+        _lsave=jnp.zeros(4, dtype=jnp.int32),  # Fortran int
+        _isave=jnp.zeros(44, dtype=jnp.int32),  # Fortran int
+        _dsave=jnp.zeros(29, dtype=float),
+        _lower_bound=jnp.zeros(n, dtype=float),
+        _upper_bound=jnp.zeros(n, dtype=float),
+        _bound_type=jnp.zeros(n, dtype=int),
+    )
+    return float_params, example_jax_lbfgsb_state
+
+
 # ------------------------------------------------------------------------------
 # Wrapper for scipy's L-BFGS-B implementation.
 # ------------------------------------------------------------------------------
@@ -395,6 +430,44 @@ class ScipyLbfgsbState:
         _validate_array_dtype(self._upper_bound, onp.float64)
         _validate_array_dtype(self._bound_type, int)
 
+    def to_jax(self) -> Dict[str, jnp.ndarray]:
+        """Generates a dictionary of jax arrays defining the state."""
+        return dict(
+            x=jnp.asarray(self.x),
+            _maxcor=jnp.asarray(self._maxcor),
+            _line_search_max_steps=jnp.asarray(self._line_search_max_steps),
+            _wa=jnp.asarray(self._wa),
+            _iwa=jnp.asarray(self._iwa),
+            _task=_array_from_s60_str(self._task),
+            _csave=_array_from_s60_str(self._csave),
+            _lsave=jnp.asarray(self._lsave),
+            _isave=jnp.asarray(self._isave),
+            _dsave=jnp.asarray(self._dsave),
+            _lower_bound=jnp.asarray(self._lower_bound),
+            _upper_bound=jnp.asarray(self._upper_bound),
+            _bound_type=jnp.asarray(self._bound_type),
+        )
+
+    @classmethod
+    def from_jax(cls, state_dict: Dict[str, jnp.ndarray]) -> "ScipyLbfgsbState":
+        """Converts a dictionary of jax arrays to a `ScipyLbfgsbState`."""
+        state_dict = copy.deepcopy(state_dict)
+        return ScipyLbfgsbState(
+            x=onp.asarray(state_dict["x"], dtype=onp.float64),
+            _maxcor=int(state_dict["_maxcor"]),
+            _line_search_max_steps=int(state_dict["_line_search_max_steps"]),
+            _wa=onp.asarray(state_dict["_wa"], onp.float64),
+            _iwa=onp.asarray(state_dict["_iwa"], dtype=FORTRAN_INT),
+            _task=_s60_str_from_array(state_dict["_task"]),
+            _csave=_s60_str_from_array(state_dict["_csave"]),
+            _lsave=onp.asarray(state_dict["_lsave"], dtype=FORTRAN_INT),
+            _isave=onp.asarray(state_dict["_isave"], dtype=FORTRAN_INT),
+            _dsave=onp.asarray(state_dict["_dsave"], dtype=onp.float64),
+            _lower_bound=onp.asarray(state_dict["_lower_bound"], dtype=onp.float64),
+            _upper_bound=onp.asarray(state_dict["_upper_bound"], dtype=onp.float64),
+            _bound_type=onp.asarray(state_dict["_bound_type"], dtype=int),
+        )
+
     @classmethod
     def init(
         cls,
@@ -440,12 +513,12 @@ class ScipyLbfgsbState:
 
         # See initialization of internal variables in the `lbfgsb._minimize_lbfgsb`
         # function.
-        wa_shape = 2 * maxcor * n + 5 * n + 11 * maxcor**2 + 8 * maxcor
+        wa_size = _wa_size(n=n, maxcor=maxcor)
         state = ScipyLbfgsbState(
             x=onp.array(x0, onp.float64),
             _maxcor=maxcor,
             _line_search_max_steps=line_search_max_steps,
-            _wa=onp.zeros(wa_shape, onp.float64),
+            _wa=onp.zeros(wa_size, onp.float64),
             _iwa=onp.zeros(3 * n, FORTRAN_INT),
             _task=task,
             _csave=onp.zeros(1, "S60"),
@@ -510,6 +583,11 @@ class ScipyLbfgsbState:
                 break
 
 
+def _wa_size(n: int, maxcor: int) -> int:
+    """Return the size of the `wa` attribute of lbfgsb state."""
+    return 2 * maxcor * n + 5 * n + 11 * maxcor**2 + 8 * maxcor
+
+
 def _validate_array_dtype(x: NDArray, dtype: Union[type, str]) -> None:
     """Validates that `x` is an array with the specified `dtype`."""
     if not isinstance(x, onp.ndarray):
@@ -534,3 +612,19 @@ def _configure_bounds(
         onp.asarray(upper_bound_array, onp.float64),
         onp.asarray(bound_type),
     )
+
+
+def _array_from_s60_str(s60_str: NDArray) -> jnp.ndarray:
+    """Return a jax array for a numpy s60 string."""
+    assert s60_str.shape == (1,)
+    chars = [int(o) for o in s60_str[0]]
+    chars.extend([32] * (59 - len(chars)))
+    return jnp.asarray(chars, dtype=int)
+
+
+def _s60_str_from_array(array: jnp.ndarray) -> NDArray:
+    """Return a numpy s60 string for a jax array."""
+    return onp.asarray(
+        [b"".join(int(i).to_bytes(length=1, byteorder="big") for i in array)],
+        dtype="S60",
+    )

invrs_opt/lbfgsb/transform.py

@@ -8,7 +8,6 @@ from typing import Tuple, Union
 import jax
 import jax.numpy as jnp
 from jax import tree_util
-
 from totypes import types
 
 PadMode = Union[str, Tuple[str, str]]

{invrs_opt-0.1.4.dist-info → invrs_opt-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: invrs_opt
-Version: 0.1.4
+Version: 0.2.0
 Summary: Algorithms for inverse design
 Author-email: "Martin F. Schubert" <mfschubert@gmail.com>
 Maintainer-email: "Martin F. Schubert" <mfschubert@gmail.com>
@@ -47,7 +47,7 @@ Requires-Dist: pytest-cov ; extra == 'tests'
 Requires-Dist: pytest-subtests ; extra == 'tests'
 
 # invrs-opt - Optimization algorithms for inverse design
-`v0.1.4`
+`v0.2.0`
 
 ## Overview
 

invrs_opt-0.2.0.dist-info/RECORD

@@ -0,0 +1,11 @@
+invrs_opt/__init__.py,sha256=ZRbgvTG7O9G3ZUHOofIAqDuyNXZoD924wUWForaFZM4,309
+invrs_opt/base.py,sha256=dSX9QkMPzI8ROxy2cFNmMwk_89eQbk0rw94CzvLPQoY,907
+invrs_opt/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+invrs_opt/lbfgsb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+invrs_opt/lbfgsb/lbfgsb.py,sha256=POK422WPlfVZw63y5i8wkrbh18YzojW2L5nM-eue7PI,23528
+invrs_opt/lbfgsb/transform.py,sha256=gYRBHUfhdkzSeBLfncdroWs3PP_o5T2X0GfDhTc82Rs,5926
+invrs_opt-0.2.0.dist-info/LICENSE,sha256=ty6jHPvpyjHy6dbhnu6aDSY05bbl2jQTjnq9u1sBCfM,1078
+invrs_opt-0.2.0.dist-info/METADATA,sha256=vmoIcnI1tL3zvvGwpJBcPula-rIZ3HVbs7bzyvfpZls,3272
+invrs_opt-0.2.0.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
+invrs_opt-0.2.0.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
+invrs_opt-0.2.0.dist-info/RECORD,,

invrs_opt-0.1.4.dist-info/RECORD

@@ -1,11 +0,0 @@
-invrs_opt/__init__.py,sha256=ByNDfFD5Jc1XezSskVZH4ERm3pBqhBkcxS5V8qtKlHU,309
-invrs_opt/base.py,sha256=dSX9QkMPzI8ROxy2cFNmMwk_89eQbk0rw94CzvLPQoY,907
-invrs_opt/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-invrs_opt/lbfgsb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-invrs_opt/lbfgsb/lbfgsb.py,sha256=tlVU6wxaSCr2FSYfEe3RUaTFT9EFyyk4i71l508vNsI,19393
-invrs_opt/lbfgsb/transform.py,sha256=TjFSeWGqlJv8uY4jtgaZ38Z5hplSX5WSQfQzN8rMV5U,5927
-invrs_opt-0.1.4.dist-info/LICENSE,sha256=ty6jHPvpyjHy6dbhnu6aDSY05bbl2jQTjnq9u1sBCfM,1078
-invrs_opt-0.1.4.dist-info/METADATA,sha256=iVhKz15sEHlsssEtbCmvkOs3ij5h4XGz5I9gNyUY5D4,3272
-invrs_opt-0.1.4.dist-info/WHEEL,sha256=oiQVh_5PnQM0E3gPdiz09WCNmwiHDMaGer_elqB3coM,92
-invrs_opt-0.1.4.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
-invrs_opt-0.1.4.dist-info/RECORD,,