invrs-opt 0.10.6__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.6"
+__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)
@@ -380,10 +397,9 @@ def parameterized_lbfgsb(
         _, vjp_fn = jax.vjp(_params_from_latents, latents)
         (latents_grad,) = vjp_fn(grad)
 
-        if not (
-            tree_util.tree_structure(latents_grad)
-            == tree_util.tree_structure(latents)  # type: ignore[operator]
-        ):
+        treedef = tree_util.tree_structure(latents_grad)
+        expected_treedef = tree_util.tree_structure(latents)
+        if not treedef == expected_treedef:  # type: ignore[operator]
             raise ValueError(
                 f"Tree structure of `latents_grad` was different than expected, got \n"
                 f"{tree_util.tree_structure(latents_grad)} but expected \n"
@@ -405,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.
@@ -502,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)
@@ -218,10 +236,9 @@ def parameterized_wrapped_optax(
         _, vjp_fn = jax.vjp(_params_from_latents, latents)
         (latents_grad,) = vjp_fn(grad)
 
-        if not (
-            tree_util.tree_structure(latents_grad)
-            == tree_util.tree_structure(latents)  # type: ignore[operator]
-        ):
+        treedef = tree_util.tree_structure(latents_grad)
+        expected_treedef = tree_util.tree_structure(latents)
+        if not treedef == expected_treedef:  # type: ignore[operator]
             raise ValueError(
                 f"Tree structure of `latents_grad` was different than expected, got \n"
                 f"{tree_util.tree_structure(latents_grad)} but expected \n"
@@ -233,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/parameterization/base.py

@@ -124,6 +124,22 @@ class Density2DMetadata:
         self.periodic = tuple(self.periodic)
         self.symmetries = tuple(self.symmetries)
 
+    def __eq__(self, other: Any) -> bool:
+        if not isinstance(other, Density2DMetadata):
+            return False
+        if not (
+            self.lower_bound == other.lower_bound
+            and self.upper_bound == other.upper_bound
+            and _arrays_equal_or_both_none(self.fixed_solid, other.fixed_solid)
+            and _arrays_equal_or_both_none(self.fixed_void, other.fixed_void)
+            and self.minimum_width == other.minimum_width
+            and self.minimum_spacing == other.minimum_spacing
+            and self.periodic == other.periodic
+            and self.symmetries == other.symmetries
+        ):
+            return False
+        return True
+
     @classmethod
     def from_density(self, density: types.Density2DArray) -> "Density2DMetadata":
         density_metadata_dict = dataclasses.asdict(density)
@@ -131,6 +147,21 @@ class Density2DMetadata:
         return Density2DMetadata(**density_metadata_dict)
 
 
+def _arrays_equal_or_both_none(a: Optional[Array], b: Optional[Array]) -> bool:
+    """Return `True` if `a` and `b` are equal arrays or both `None`."""
+    if (a is None, b is None) not in ((True, True), (False, False)):
+        return False
+    if a is None and b is None:
+        return True
+    assert isinstance(a, onp.ndarray)
+    assert isinstance(b, onp.ndarray)
+    if a.dtype != b.dtype:
+        return False
+    if a.shape != b.shape:
+        return False
+    return bool(onp.all(a == b))
+
+
 def _flatten_density_2d_metadata(
     metadata: Density2DMetadata,
 ) -> Tuple[

invrs_opt/parameterization/gaussian_levelset.py

@@ -388,7 +388,7 @@ def _phi_from_params(
     assert array.shape[-2] % s_factor == 0
     assert array.shape[-1] % s_factor == 0
     array = symmetry.symmetrize(array, tuple(example_density.symmetries))
-    return array
+    return jnp.asarray(array)
 
 
 # -----------------------------------------------------------------------------

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

@@ -1,6 +1,6 @@
-Metadata-Version: 2.2
+Metadata-Version: 2.4
 Name: invrs_opt
-Version: 0.10.6
+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>
@@ -531,6 +531,7 @@ Requires-Dist: bump-my-version; extra == "dev"
 Requires-Dist: darglint; extra == "dev"
 Requires-Dist: mypy; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
+Dynamic: license-file
 
 # invrs-opt - Optimization algorithms for inverse design
 ![Continuous integration](https://github.com/invrs-io/opt/actions/workflows/build-ci.yml/badge.svg)

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

@@ -1,20 +1,20 @@
-invrs_opt/__init__.py,sha256=4uq04SwbUcX0hVAhaX0-GJ4R4tRJAAXfIzzfU-VcalA,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=WP6ouVtLaXSwJBh7CSzWR7rnRdHZuSmOr57TKF4UxMg,36659
-invrs_opt/optimizers/wrapped_optax.py,sha256=781-8v_TlHsGaQF9Se9_iOEvtOLOr-BesTLudYarSlg,13685
+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=jSwrEO86lGkYQG5gWsHvcIMWpZnnbdiKpn--2qaU02g,5362
+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=PDvjdgBzklRTCUoBpo4ZMcmXeTkn0BpZEzQj7ojtYGE,24813
+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.6.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-invrs_opt-0.10.6.dist-info/METADATA,sha256=40XUQ7i3S4nYRGUqkEVOwKrJwGtpvdwmpD9ojFpAeAM,32816
-invrs_opt-0.10.6.dist-info/WHEEL,sha256=nn6H5-ilmfVryoAQl3ZQ2l8SH5imPWFpm1A5FgEuFV4,91
-invrs_opt-0.10.6.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
-invrs_opt-0.10.6.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,,

{invrs_opt-0.10.6.dist-info → invrs_opt-0.11.0.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (75.8.1)
+Generator: setuptools (80.9.0)
 Root-Is-Purelib: true
 Tag: py3-none-any