invrs-opt 0.1.3__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.3"
+__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/base.py

@@ -6,8 +6,6 @@ Copyright (c) 2023 The INVRS-IO authors.
 import dataclasses
 from typing import Any, Protocol
 
-from totypes import json_utils
-
 PyTree = Any
 
 
@@ -46,23 +44,3 @@ class Optimizer:
     init: InitFn
     params: ParamsFn
     update: UpdateFn
-
-
-# Additional custom types and prefixes used for serializing optimizer state.
-CUSTOM_TYPES_AND_PREFIXES = ()
-
-
-def serialize(tree: PyTree) -> str:
-    """Serializes a pytree into a string."""
-    return json_utils.json_from_pytree(
-        tree,
-        extra_custom_types_and_prefixes=CUSTOM_TYPES_AND_PREFIXES,
-    )
-
-
-def deserialize(serialized: str) -> PyTree:
-    """Restores a pytree from a string."""
-    return json_utils.pytree_from_json(
-        serialized,
-        extra_custom_types_and_prefixes=CUSTOM_TYPES_AND_PREFIXES,
-    )

invrs_opt/lbfgsb/lbfgsb.py

@@ -10,19 +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 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.
@@ -187,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`."""
@@ -213,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
 
-        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))
+            params, jax_lbfgsb_state = state
+            scipy_lbfgsb_state = ScipyLbfgsbState.from_jax(jax_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,
@@ -245,31 +258,25 @@ def transformed_lbfgsb(
 
 def _to_numpy(params: PyTree) -> NDArray:
     """Flattens a `params` pytree into a single rank-1 numpy array."""
-    leaves = tree_util.tree_leaves(params)
-    x_numpy: NDArray = onp.concatenate([onp.asarray(leaf).flatten() for leaf in leaves])
-    return x_numpy.astype(onp.float64)
+    x, _ = flatten_util.ravel_pytree(params)  # type: ignore[no-untyped-call]
+    return onp.asarray(x, dtype=onp.float64)
 
 
 def _to_pytree(x_flat: NDArray, params: PyTree) -> PyTree:
     """Restores a pytree from a flat numpy array using the structure of `params`.
 
+    Note that the returned pytree includes jax array leaves.
+
     Args:
         x_flat: The rank-1 numpy array to be restored.
         params: A pytree of parameters whose structure is replicated in the restored
             pytree.
 
     Returns:
-        The restored pytree.
+        The restored pytree, with jax array leaves.
     """
-    leaves, treedef = tree_util.tree_flatten(params)
-    indices_or_sections = onp.cumsum([onp.size(leaf) for leaf in leaves])
-    x_split_flat = onp.split(x_flat, indices_or_sections)
-    x_split = [x.reshape(onp.shape(leaf)) for x, leaf in zip(x_split_flat, leaves)]
-    x_split_jax = [
-        jnp.asarray(x) if isinstance(leaf, jnp.ndarray) else x
-        for x, leaf in zip(x_split, leaves)
-    ]
-    return tree_util.tree_unflatten(treedef, x_split_jax)
+    _, unflatten_fn = flatten_util.ravel_pytree(params)  # type: ignore[no-untyped-call]
+    return unflatten_fn(jnp.asarray(x_flat, dtype=float))
 
 
 def _bound_for_params(bound: PyTree, params: PyTree) -> ElementwiseBound:
@@ -320,6 +327,8 @@ def _bound_for_params(bound: PyTree, params: PyTree) -> ElementwiseBound:
 
     bound_flat = []
     for b, p in zip(bound_leaves, params_leaves):
+        if p is None:
+            continue
         if b is None or onp.isscalar(b) or onp.shape(b) == ():
             bound_flat += [b] * onp.size(p)
         else:
@@ -334,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.
 # ------------------------------------------------------------------------------
@@ -398,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,
@@ -443,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"),
@@ -513,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):
@@ -537,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.3.dist-info → invrs_opt-0.2.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
-Name: invrs-opt
-Version: 0.1.3
+Name: invrs_opt
+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.3`
+`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.3.dist-info → invrs_opt-0.2.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: bdist_wheel (0.41.3)
+Generator: bdist_wheel (0.42.0)
 Root-Is-Purelib: true
 Tag: py3-none-any
 

invrs_opt-0.1.3.dist-info/RECORD

@@ -1,11 +0,0 @@
-invrs_opt/__init__.py,sha256=IpAs-pDwW_mo2FnbNkDpsR-XuxZW6h5TwvlHkc8kCuE,309
-invrs_opt/base.py,sha256=dm5nzlO4KXFfuIfyHcTn9V1VCU6hAy1w3IA2vfzaQD8,1481
-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=nZqotWv9oGbY56UKSI3zcetexpwMoaDsvqwJSgXjvwc,19597
-invrs_opt/lbfgsb/transform.py,sha256=TjFSeWGqlJv8uY4jtgaZ38Z5hplSX5WSQfQzN8rMV5U,5927
-invrs_opt-0.1.3.dist-info/LICENSE,sha256=ty6jHPvpyjHy6dbhnu6aDSY05bbl2jQTjnq9u1sBCfM,1078
-invrs_opt-0.1.3.dist-info/METADATA,sha256=oXMNM5J_RzBpmnNB3IBvFhscJg6nQbyLWNMz-OljYHQ,3272
-invrs_opt-0.1.3.dist-info/WHEEL,sha256=Xo9-1PvkuimrydujYJAjF7pCkriuXBpUPEjma1nZyJ0,92
-invrs_opt-0.1.3.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
-invrs_opt-0.1.3.dist-info/RECORD,,