invrs-opt 0.10.7__py3-none-any.whl → 0.11.0__py3-none-any.whl

invrs_opt/__init__.py

@@ -3,7 +3,7 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.10.7"
+__version__ = "v0.11.0"
 __author__ = "Martin F. Schubert <mfschubert@gmail.com>"
 
 from invrs_opt import parameterization as parameterization

invrs_opt/optimizers/lbfgsb.py

@@ -16,7 +16,7 @@ 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 totypes import json_utils, types
 
 from invrs_opt.optimizers import base
 from invrs_opt.parameterization import (
@@ -31,7 +31,26 @@ PyTree = Any
 ElementwiseBound = Union[NDArray, Sequence[Optional[float]]]
 NumpyLbfgsbDict = Dict[str, NDArray]
 JaxLbfgsbDict = Dict[str, jnp.ndarray]
-LbfgsbState = Tuple[int, PyTree, PyTree, JaxLbfgsbDict]
+
+
+@dataclasses.dataclass
+class LbfgsbState:
+    """Stores the state of the L-BFGS-B optimizer."""
+
+    step: int
+    params: PyTree
+    latent_params: PyTree
+    opt_state: JaxLbfgsbDict
+
+
+tree_util.register_dataclass(
+    LbfgsbState,
+    data_fields=["step", "params", "latent_params", "opt_state"],
+    meta_fields=[],
+)
+
+
+json_utils.register_custom_type(LbfgsbState)
 
 
 # Task message prefixes for the underlying L-BFGS-B implementation.
@@ -327,17 +346,16 @@ def parameterized_lbfgsb(
             latents,
         )
         latent_params = param_base.combine_density_metadata(metadata, latents)
-        return (
-            0,  # step
-            _params_from_latent_params(latent_params),  # params
-            latent_params,  # latent params
-            jax_lbfgsb_state,  # opt state
+        return LbfgsbState(
+            step=0,
+            params=_params_from_latent_params(latent_params),
+            latent_params=latent_params,
+            opt_state=jax_lbfgsb_state,
         )
 
     def params_fn(state: LbfgsbState) -> PyTree:
         """Returns the parameters for the given `state`."""
-        _, params, _, _ = state
-        return params
+        return state.params
 
     def update_fn(
         *,
@@ -366,8 +384,7 @@ def parameterized_lbfgsb(
             flat_latent_updates = updated_flat_latent_params - flat_latent_params
             return flat_latent_updates, scipy_lbfgsb_state.to_dict()
 
-        step, _, latent_params, jax_lbfgsb_state = state
-        metadata, latents = param_base.partition_density_metadata(latent_params)
+        metadata, latents = param_base.partition_density_metadata(state.latent_params)
 
         def _params_from_latents(latents: PyTree) -> PyTree:
             latent_params = param_base.combine_density_metadata(metadata, latents)
@@ -404,23 +421,28 @@ def parameterized_lbfgsb(
             latents_grad
         )  # type: ignore[no-untyped-call]
 
-        flat_latent_updates, jax_lbfgsb_state = callback_sequential(
+        flat_latent_updates, opt_state = callback_sequential(
             _update_pure,
-            (flat_latents_grad, jax_lbfgsb_state),
+            (flat_latents_grad, state.opt_state),
             flat_latents_grad,
             value,
-            jax_lbfgsb_state,
+            state.opt_state,
         )
         latent_updates = unflatten_fn(flat_latent_updates)
         latent_params = _apply_updates(
-            params=latent_params,
+            params=state.latent_params,
             updates=param_base.combine_density_metadata(metadata, latent_updates),
             value=value,
-            step=step,
+            step=state.step,
         )
         latent_params = _clip(latent_params)
         params = _params_from_latent_params(latent_params)
-        return step + 1, params, latent_params, jax_lbfgsb_state
+        return LbfgsbState(
+            step=state.step + 1,
+            params=params,
+            latent_params=latent_params,
+            opt_state=opt_state,
+        )
 
     # -------------------------------------------------------------------------
     # Functions related to the density parameterization.
@@ -501,7 +523,7 @@ def parameterized_lbfgsb(
 
 def is_converged(state: LbfgsbState) -> jnp.ndarray:
     """Returns `True` if the optimization has converged."""
-    return state[3]["converged"]
+    return state.opt_state["converged"]
 
 
 # ------------------------------------------------------------------------------

invrs_opt/optimizers/wrapped_optax.py

@@ -3,13 +3,14 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-from typing import Any, Optional, Tuple
+import dataclasses
+from typing import Any, Optional
 
 import jax
 import jax.numpy as jnp
 import optax  # type: ignore[import-untyped]
 from jax import tree_util
-from totypes import types
+from totypes import json_utils, types
 
 from invrs_opt.optimizers import base
 from invrs_opt.parameterization import (
@@ -20,7 +21,26 @@ from invrs_opt.parameterization import (
 )
 
 PyTree = Any
-WrappedOptaxState = Tuple[int, PyTree, PyTree, PyTree]
+
+
+@dataclasses.dataclass
+class WrappedOptaxState:
+    """Stores the state of a wrapped optax optimizer."""
+
+    step: int
+    params: PyTree
+    latent_params: PyTree
+    opt_state: PyTree
+
+
+tree_util.register_dataclass(
+    WrappedOptaxState,
+    data_fields=["step", "params", "latent_params", "opt_state"],
+    meta_fields=[],
+)
+
+
+json_utils.register_custom_type(WrappedOptaxState)
 
 
 def wrapped_optax(opt: optax.GradientTransformation) -> base.Optimizer:
@@ -182,17 +202,16 @@ def parameterized_wrapped_optax(
         """Initializes the optimization state."""
         latent_params = _init_latents(params)
         _, latents = param_base.partition_density_metadata(latent_params)
-        return (
-            0,  # step
-            _params_from_latent_params(latent_params),  # params
-            latent_params,  # latent params
-            opt.init(latents),  # opt state
+        return WrappedOptaxState(
+            step=0,
+            params=_params_from_latent_params(latent_params),
+            latent_params=latent_params,
+            opt_state=opt.init(latents),
         )
 
     def params_fn(state: WrappedOptaxState) -> PyTree:
         """Returns the parameters for the given `state`."""
-        _, params, _, _ = state
-        return params
+        return state.params
 
     def update_fn(
         *,
@@ -204,8 +223,7 @@ def parameterized_wrapped_optax(
         """Updates the state."""
         del params
 
-        step, params, latent_params, opt_state = state
-        metadata, latents = param_base.partition_density_metadata(latent_params)
+        metadata, latents = param_base.partition_density_metadata(state.latent_params)
 
         def _params_from_latents(latents: PyTree) -> PyTree:
             latent_params = param_base.combine_density_metadata(metadata, latents)
@@ -232,16 +250,23 @@ def parameterized_wrapped_optax(
             lambda a, b: a + b, latents_grad, constraint_loss_grad
         )
 
-        latent_updates, opt_state = opt.update(latents_grad, opt_state, params=latents)
+        latent_updates, opt_state = opt.update(
+            latents_grad, state.opt_state, params=latents
+        )
         latent_params = _apply_updates(
-            params=latent_params,
+            params=state.latent_params,
             updates=param_base.combine_density_metadata(metadata, latent_updates),
             value=value,
-            step=step,
+            step=state.step,
         )
         latent_params = _clip(latent_params)
         params = _params_from_latent_params(latent_params)
-        return (step + 1, params, latent_params, opt_state)
+        return WrappedOptaxState(
+            step=state.step + 1,
+            params=params,
+            latent_params=latent_params,
+            opt_state=opt_state,
+        )
 
     # -------------------------------------------------------------------------
     # Functions related to the density parameterization.

{invrs_opt-0.10.7.dist-info → invrs_opt-0.11.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: invrs_opt
-Version: 0.10.7
+Version: 0.11.0
 Summary: Algorithms for inverse design
 Author-email: "Martin F. Schubert" <mfschubert@gmail.com>
 Maintainer-email: "Martin F. Schubert" <mfschubert@gmail.com>

{invrs_opt-0.10.7.dist-info → invrs_opt-0.11.0.dist-info}/RECORD

@@ -1,20 +1,20 @@
-invrs_opt/__init__.py,sha256=4ZhV7JC3OdPDzBPTgWsWk8NNPrlg_Hthf7M4BTmUQ4g,586
+invrs_opt/__init__.py,sha256=CQM3bUeeV9fgJgAee7NFy1B2HJyw2uhNSzoKynmVyOQ,586
 invrs_opt/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 invrs_opt/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 invrs_opt/experimental/client.py,sha256=tbtH13FrA65XmTZfTO71CxJ78jeAEj3Zf85R-MTwbiU,4909
 invrs_opt/experimental/labels.py,sha256=dQDAMPyFMV6mXnMy295z8Ap205DRdVzysXny_Be8FmY,562
 invrs_opt/optimizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 invrs_opt/optimizers/base.py,sha256=uFfkN2LwWzAtwh6ktwWNy2iHNOY-sW3JzI46iSFkgok,1306
-invrs_opt/optimizers/lbfgsb.py,sha256=uyeVoTq1ZtCLILw5Rv0ALNRiEXMxhEqdVElukmeuoZ4,36693
-invrs_opt/optimizers/wrapped_optax.py,sha256=v-pezJDDdtUckh-jL33Wj-H6a_a3xqrxkE9tOBWcCNI,13719
+invrs_opt/optimizers/lbfgsb.py,sha256=QB8lD02sMr-2V0d_k4UB8Y7SOlNektp0vzwrJId0u44,37059
+invrs_opt/optimizers/wrapped_optax.py,sha256=hftrGCpg4kVv2NarZAVHLg6Gkk87zcc4_yyfvSvQHTo,14161
 invrs_opt/parameterization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 invrs_opt/parameterization/base.py,sha256=GFubMydcww6MXlGNkxHGMH_sOCyj9M5R-iwZYTqCo2I,6539
 invrs_opt/parameterization/filter_project.py,sha256=XL3HTEBLrF-q_75TjhOWLNdfUOSEEjKcoM7Qj844QpQ,4590
 invrs_opt/parameterization/gaussian_levelset.py,sha256=bmVU1We92zPpNIJI8sCq3OCeHZw6emZl86unJYwnWbc,24826
 invrs_opt/parameterization/pixel.py,sha256=YWkyBhfYtzI8cQ-M90PAZqRAbabwVaUh0UiYIGegQHI,1955
 invrs_opt/parameterization/transforms.py,sha256=mqDKuAg4wpSL9kh0oYKxtSoH0mHOQeKG1RND2fJSYaU,9441
-invrs_opt-0.10.7.dist-info/licenses/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-invrs_opt-0.10.7.dist-info/METADATA,sha256=2uHpZfOyDdaMlLtnUoYuj6rGcG-T3KYchXU-oORyRQY,32838
-invrs_opt-0.10.7.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
-invrs_opt-0.10.7.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
-invrs_opt-0.10.7.dist-info/RECORD,,
+invrs_opt-0.11.0.dist-info/licenses/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
+invrs_opt-0.11.0.dist-info/METADATA,sha256=M6IzHHGIctk49mBxnswyKcvE3BikRoqZz_19UDOqB9c,32838
+invrs_opt-0.11.0.dist-info/WHEEL,sha256=_zCd3N1l69ArxyTb8rzEoP9TpbYXkqRFSNOD5OuxnTs,91
+invrs_opt-0.11.0.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
+invrs_opt-0.11.0.dist-info/RECORD,,