invrs-opt 0.8.0__py3-none-any.whl → 0.8.1__py3-none-any.whl

invrs_opt/__init__.py

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

invrs_opt/optimizers/wrapped_optax.py

@@ -178,64 +178,15 @@ def parameterized_wrapped_optax(
     if density_parameterization is None:
         density_parameterization = pixel.pixel()
 
-    def _init_latents(params: PyTree) -> PyTree:
-        def _leaf_init_latents(leaf: Any) -> Any:
-            leaf = _clip(leaf)
-            if not _is_density(leaf):
-                return leaf
-            return density_parameterization.from_density(leaf)
-
-        return tree_util.tree_map(_leaf_init_latents, params, is_leaf=_is_custom_type)
-
-    def _params_from_latents(params: PyTree) -> PyTree:
-        def _leaf_params_from_latents(leaf: Any) -> Any:
-            if not _is_parameterized_density(leaf):
-                return leaf
-            return density_parameterization.to_density(leaf)
-
-        return tree_util.tree_map(
-            _leaf_params_from_latents,
-            params,
-            is_leaf=_is_parameterized_density,
-        )
-
-    def _constraint_loss(latent_params: PyTree) -> jnp.ndarray:
-        def _constraint_loss_leaf(
-            params: parameterization_base.ParameterizedDensity2DArrayBase,
-        ) -> jnp.ndarray:
-            constraints = density_parameterization.constraints(params)
-            constraints = tree_util.tree_map(
-                lambda x: jnp.sum(jnp.maximum(x, 0.0) ** 2),
-                constraints,
-            )
-            return jnp.sum(jnp.asarray(constraints))
-
-        losses = [0.0] + [
-            _constraint_loss_leaf(p)
-            for p in tree_util.tree_leaves(
-                latent_params, is_leaf=_is_parameterized_density
-            )
-            if _is_parameterized_density(p)
-        ]
-        return penalty * jnp.sum(jnp.asarray(losses))
-
-    def _update_parameterized_densities(latent_params: PyTree, step: int) -> PyTree:
-        def _update_leaf(leaf: Any) -> Any:
-            if not _is_parameterized_density(leaf):
-                return leaf
-            return density_parameterization.update(leaf, step)
-
-        return tree_util.tree_map(
-            _update_leaf,
-            latent_params,
-            is_leaf=_is_parameterized_density,
-        )
-
     def init_fn(params: PyTree) -> WrappedOptaxState:
         """Initializes the optimization state."""
         latent_params = _init_latents(params)
-        params = _params_from_latents(latent_params)
-        return 0, params, latent_params, opt.init(latent_params)
+        return (
+            0,  # step
+            _params_from_latents(latent_params),  # params
+            latent_params,  # latent params
+            opt.init(tree_util.tree_leaves(latent_params)),  # opt state
+        )
 
     def params_fn(state: WrappedOptaxState) -> PyTree:
         """Returns the parameters for the given `state`."""
@@ -252,7 +203,8 @@ def parameterized_wrapped_optax(
         """Updates the state."""
         del value, params
 
-        step, _, latent_params, opt_state = state
+        step, params, latent_params, opt_state = state
+
         _, vjp_fn = jax.vjp(_params_from_latents, latent_params)
         (latent_grad,) = vjp_fn(grad)
 
@@ -271,14 +223,85 @@ def parameterized_wrapped_optax(
             lambda a, b: a + b, latent_grad, constraint_loss_grad
         )
 
-        updates, opt_state = opt.update(
-            updates=latent_grad, state=opt_state, params=latent_params
+        updates_leaves, opt_state = opt.update(
+            updates=tree_util.tree_leaves(latent_grad),
+            state=opt_state,
+            params=tree_util.tree_leaves(latent_params),
+        )
+        latent_params_leaves = optax.apply_updates(
+            params=tree_util.tree_leaves(latent_params),
+            updates=updates_leaves,
         )
-        latent_params = optax.apply_updates(params=latent_params, updates=updates)
+        latent_params = tree_util.tree_unflatten(
+            treedef=tree_util.tree_structure(latent_params),
+            leaves=latent_params_leaves,
+        )
+
         latent_params = _clip(latent_params)
-        latent_params = _update_parameterized_densities(latent_params, step)
+        latent_params = _update_parameterized_densities(latent_params, step + 1)
         params = _params_from_latents(latent_params)
-        return step + 1, params, latent_params, opt_state
+        return (step + 1, params, latent_params, opt_state)
+
+    # -------------------------------------------------------------------------
+    # Functions related to the density parameterization.
+    # -------------------------------------------------------------------------
+
+    def _init_latents(params: PyTree) -> PyTree:
+        def _leaf_init_latents(leaf: Any) -> Any:
+            leaf = _clip(leaf)
+            if not _is_density(leaf):
+                return leaf
+            return density_parameterization.from_density(leaf)
+
+        return tree_util.tree_map(_leaf_init_latents, params, is_leaf=_is_custom_type)
+
+    def _params_from_latents(params: PyTree) -> PyTree:
+        def _leaf_params_from_latents(leaf: Any) -> Any:
+            if not _is_parameterized_density(leaf):
+                return leaf
+            return density_parameterization.to_density(leaf)
+
+        return tree_util.tree_map(
+            _leaf_params_from_latents,
+            params,
+            is_leaf=_is_parameterized_density,
+        )
+
+    def _update_parameterized_densities(latent_params: PyTree, step: int) -> PyTree:
+        def _update_leaf(leaf: Any) -> Any:
+            if not _is_parameterized_density(leaf):
+                return leaf
+            return density_parameterization.update(leaf, step)
+
+        return tree_util.tree_map(
+            _update_leaf,
+            latent_params,
+            is_leaf=_is_parameterized_density,
+        )
+
+    # -------------------------------------------------------------------------
+    # Functions related to the constraints to be minimized.
+    # -------------------------------------------------------------------------
+
+    def _constraint_loss(latent_params: PyTree) -> jnp.ndarray:
+        def _constraint_loss_leaf(
+            params: parameterization_base.ParameterizedDensity2DArrayBase,
+        ) -> jnp.ndarray:
+            constraints = density_parameterization.constraints(params)
+            constraints = tree_util.tree_map(
+                lambda x: jnp.sum(jnp.maximum(x, 0.0) ** 2),
+                constraints,
+            )
+            return jnp.sum(jnp.asarray(constraints))
+
+        losses = [0.0] + [
+            _constraint_loss_leaf(p)
+            for p in tree_util.tree_leaves(
+                latent_params, is_leaf=_is_parameterized_density
+            )
+            if _is_parameterized_density(p)
+        ]
+        return penalty * jnp.sum(jnp.asarray(losses))
 
     return base.Optimizer(init=init_fn, params=params_fn, update=update_fn)
 

{invrs_opt-0.8.0.dist-info → invrs_opt-0.8.1.dist-info}/METADATA

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

{invrs_opt-0.8.0.dist-info → invrs_opt-0.8.1.dist-info}/RECORD

@@ -1,4 +1,4 @@
-invrs_opt/__init__.py,sha256=kTrg48iZu7i5OlH0Nqtfh_wBn3be9u2eZgJBRqUr6uQ,585
+invrs_opt/__init__.py,sha256=tjbabqQKp5y1I5eqDa7onsy7E58P9nTk_Ifn12oMfZM,585
 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=t4XxnditYbM9DWZeyBPj0Sa2acvkikT0ybhUdmH2r-Y,4852
@@ -6,15 +6,15 @@ invrs_opt/experimental/labels.py,sha256=dQDAMPyFMV6mXnMy295z8Ap205DRdVzysXny_Be8
 invrs_opt/optimizers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 invrs_opt/optimizers/base.py,sha256=-wNwH0J475y8FzB5aLAkc_1602LvYeF4Hddr9OiBkDY,1276
 invrs_opt/optimizers/lbfgsb.py,sha256=2NcyFllUM5CrjOZLEBocYAlCrkWt1fiRGJrxg05waog,35149
-invrs_opt/optimizers/wrapped_optax.py,sha256=zPC2j_KkfF2RqeorxB38ovuUsg1SNwdxwhjA7gvOMC4,12387
+invrs_opt/optimizers/wrapped_optax.py,sha256=c4iSsXHF2cJZaZhdxLpyRy26q5IxBTo4-6sY8IOLT38,13259
 invrs_opt/parameterization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 invrs_opt/parameterization/base.py,sha256=k2VXZorMnkmaIWWY_G2CvS0NtxYojfvKlPccDGZyGNA,3779
 invrs_opt/parameterization/filter_project.py,sha256=D2w8xrg34V8ysIbbr1RPvegM5WoLdz8QKUCGi80ieOI,3466
 invrs_opt/parameterization/gaussian_levelset.py,sha256=DllvpkBpVxuDFQFu941f8v_Xh2D0m9Eh-JxBH5afNHU,25221
 invrs_opt/parameterization/pixel.py,sha256=yWdGpf0x6Om_-7CYOcHTgNK9UF-XGAPM7GbOBNijnBw,1305
 invrs_opt/parameterization/transforms.py,sha256=8GzaIsUuuXvMCLiqAEEfxmi9qE9KqHzbuTj_m0GjH3w,8216
-invrs_opt-0.8.0.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-invrs_opt-0.8.0.dist-info/METADATA,sha256=WSzpZwByiBqOZfiwSGvaVaFa5TUAPj4b-mr2qZRGLq0,32633
-invrs_opt-0.8.0.dist-info/WHEEL,sha256=HiCZjzuy6Dw0hdX5R3LCFPDmFS4BWl8H-8W39XfmgX4,91
-invrs_opt-0.8.0.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
-invrs_opt-0.8.0.dist-info/RECORD,,
+invrs_opt-0.8.1.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
+invrs_opt-0.8.1.dist-info/METADATA,sha256=4qLsuDLp3pNn5C1TTPnF0qwWps7MeN_Uc3rKsBLbPrk,32633
+invrs_opt-0.8.1.dist-info/WHEEL,sha256=HiCZjzuy6Dw0hdX5R3LCFPDmFS4BWl8H-8W39XfmgX4,91
+invrs_opt-0.8.1.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
+invrs_opt-0.8.1.dist-info/RECORD,,