invrs-opt 0.7.2__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.7.2"
+__version__ = "v0.8.1"
 __author__ = "Martin F. Schubert <mfschubert@gmail.com>"
 
 from invrs_opt import parameterization as parameterization

invrs_opt/optimizers/lbfgsb.py

@@ -29,7 +29,7 @@ NDArray = onp.ndarray[Any, Any]
 PyTree = Any
 ElementwiseBound = Union[NDArray, Sequence[Optional[float]]]
 JaxLbfgsbDict = Dict[str, jnp.ndarray]
-LbfgsbState = Tuple[PyTree, PyTree, JaxLbfgsbDict]
+LbfgsbState = Tuple[int, PyTree, PyTree, JaxLbfgsbDict]
 
 
 # Task message prefixes for the underlying L-BFGS-B implementation.
@@ -323,7 +323,7 @@ def parameterized_lbfgsb(
         ) -> jnp.ndarray:
             constraints = density_parameterization.constraints(params)
             constraints = tree_util.tree_map(
-                lambda x: jnp.sum(jnp.maximum(x, 0.0)),
+                lambda x: jnp.sum(jnp.maximum(x, 0.0) ** 2),
                 constraints,
             )
             return jnp.sum(jnp.asarray(constraints))
@@ -337,6 +337,18 @@ def parameterized_lbfgsb(
         ]
         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) -> LbfgsbState:
         """Initializes the optimization state."""
 
@@ -359,11 +371,11 @@ def parameterized_lbfgsb(
         latent_params, jax_lbfgsb_state = jax.pure_callback(
             _init_state_pure, _example_state(latent_params, maxcor), latent_params
         )
-        return _params_from_latents(latent_params), latent_params, jax_lbfgsb_state
+        return 0, _params_from_latents(latent_params), latent_params, jax_lbfgsb_state
 
     def params_fn(state: LbfgsbState) -> PyTree:
         """Returns the parameters for the given `state`."""
-        params, _, _ = state
+        _, params, _, _ = state
         return params
 
     def update_fn(
@@ -390,7 +402,7 @@ def parameterized_lbfgsb(
             flat_latent_params = jnp.asarray(scipy_lbfgsb_state.x)
             return flat_latent_params, scipy_lbfgsb_state.to_jax()
 
-        _, latent_params, jax_lbfgsb_state = state
+        step, _, latent_params, jax_lbfgsb_state = state
         _, vjp_fn = jax.vjp(_params_from_latents, latent_params)
         (latent_grad,) = vjp_fn(grad)
 
@@ -427,7 +439,9 @@ def parameterized_lbfgsb(
             jax_lbfgsb_state,
         )
         latent_params = unflatten_fn(flat_latent_params)
-        return _params_from_latents(latent_params), latent_params, jax_lbfgsb_state
+        latent_params = _update_parameterized_densities(latent_params, step)
+        params = _params_from_latents(latent_params)
+        return step + 1, params, latent_params, jax_lbfgsb_state
 
     return base.Optimizer(
         init=init_fn,
@@ -438,7 +452,7 @@ def parameterized_lbfgsb(
 
 def is_converged(state: LbfgsbState) -> jnp.ndarray:
     """Returns `True` if the optimization has converged."""
-    return state[2]["converged"]
+    return state[3]["converged"]
 
 
 # ------------------------------------------------------------------------------

invrs_opt/optimizers/wrapped_optax.py

@@ -20,7 +20,7 @@ from invrs_opt.parameterization import (
 )
 
 PyTree = Any
-WrappedOptaxState = Tuple[PyTree, PyTree, PyTree]
+WrappedOptaxState = Tuple[int, PyTree, PyTree, PyTree]
 
 
 def wrapped_optax(opt: optax.GradientTransformation) -> base.Optimizer:
@@ -178,56 +178,19 @@ 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)),
-                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 init_fn(params: PyTree) -> WrappedOptaxState:
         """Initializes the optimization state."""
         latent_params = _init_latents(params)
-        params = _params_from_latents(latent_params)
-        return 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`."""
-        params, _, _ = state
+        _, params, _, _ = state
         return params
 
     def update_fn(
@@ -240,7 +203,8 @@ def parameterized_wrapped_optax(
         """Updates the state."""
         del value, params
 
-        _, 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)
 
@@ -259,13 +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 = tree_util.tree_unflatten(
+            treedef=tree_util.tree_structure(latent_params),
+            leaves=latent_params_leaves,
         )
-        latent_params = optax.apply_updates(params=latent_params, updates=updates)
+
         latent_params = _clip(latent_params)
+        latent_params = _update_parameterized_densities(latent_params, step + 1)
         params = _params_from_latents(latent_params)
-        return 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/parameterization/base.py

@@ -44,6 +44,13 @@ class ConstraintsFn(Protocol):
         ...
 
 
+class UpdateFn(Protocol):
+    """Performs the required update of a parameterized density for the given step."""
+
+    def __call__(self, params: PyTree, step: int) -> PyTree:
+        ...
+
+
 @dataclasses.dataclass
 class Density2DParameterization:
     """Stores `(from_density, to_density, constraints)` function triple."""
@@ -51,6 +58,7 @@ class Density2DParameterization:
     from_density: FromDensityFn
     to_density: ToDensityFn
     constraints: ConstraintsFn
+    update: UpdateFn
 
 
 @dataclasses.dataclass

invrs_opt/parameterization/filter_project.py

@@ -85,8 +85,14 @@ def filter_project(beta: float) -> base.Density2DParameterization:
         del params
         return jnp.asarray(0.0)
 
+    def update_fn(params: FilterAndProjectParams, step: int) -> FilterAndProjectParams:
+        """Perform updates to `params` required for the given `step`."""
+        del step
+        return params
+
     return base.Density2DParameterization(
         to_density=to_density_fn,
         from_density=from_density_fn,
         constraints=constraints_fn,
+        update=update_fn,
     )

invrs_opt/parameterization/gaussian_levelset.py

@@ -214,7 +214,7 @@ def gaussian_levelset(
             return params, state
 
         state = init_optimizer.init(params)
-        params, state = jax.lax.fori_loop(
+        params, _ = jax.lax.fori_loop(
             0, init_steps, body_fun=step_fn, init_val=(params, state)
         )
 
@@ -269,10 +269,16 @@ def gaussian_levelset(
         )
         return constraints / length_scale**2
 
+    def update_fn(params: GaussianLevelsetParams, step: int) -> GaussianLevelsetParams:
+        """Perform updates to `params` required for the given `step`."""
+        del step
+        return params
+
     return base.Density2DParameterization(
         to_density=to_density_fn,
         from_density=from_density_fn,
         constraints=constraints_fn,
+        update=update_fn,
     )
 
 

invrs_opt/parameterization/pixel.py

@@ -38,8 +38,13 @@ def pixel() -> base.Density2DParameterization:
         del params
         return jnp.asarray(0.0)
 
+    def update_fn(params: PixelParams, step: int) -> PixelParams:
+        del step
+        return params
+
     return base.Density2DParameterization(
         from_density=from_density_fn,
         to_density=to_density_fn,
         constraints=constraints_fn,
+        update=update_fn,
     )

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: invrs_opt
-Version: 0.7.2
+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.7.2`
+`v0.8.1`
 
 ## Overview
 

invrs_opt-0.8.1.dist-info/RECORD

@@ -0,0 +1,20 @@
+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
+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=-wNwH0J475y8FzB5aLAkc_1602LvYeF4Hddr9OiBkDY,1276
+invrs_opt/optimizers/lbfgsb.py,sha256=2NcyFllUM5CrjOZLEBocYAlCrkWt1fiRGJrxg05waog,35149
+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.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,,

{invrs_opt-0.7.2.dist-info → invrs_opt-0.8.1.dist-info}/WHEEL

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

invrs_opt-0.7.2.dist-info/RECORD

@@ -1,20 +0,0 @@
-invrs_opt/__init__.py,sha256=_hiwGthyPapJUx8-nP7JrJYEFX3mJYTsdUh_X0MyhEQ,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
-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=-wNwH0J475y8FzB5aLAkc_1602LvYeF4Hddr9OiBkDY,1276
-invrs_opt/optimizers/lbfgsb.py,sha256=d7i02NZZ3yYdJg7wkERDMPpdBD3GaRPhmQexGrZPz_Y,34597
-invrs_opt/optimizers/wrapped_optax.py,sha256=L836gwzWbwxPNWh8Y7PSgFHV4PZluSklnbh3BR5djlc,11859
-invrs_opt/parameterization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-invrs_opt/parameterization/base.py,sha256=QTnhOfMYbDchZOFzk9graryMd6rYHlyd7E2T1TtucB8,3570
-invrs_opt/parameterization/filter_project.py,sha256=XCPqQ2ECv7DDTLRtVGJePfnjKYB2XndI6DssSr-4MZw,3239
-invrs_opt/parameterization/gaussian_levelset.py,sha256=IGdQl3XMEHYNxNPSoKmUD7ZJu_fImIZTqqqmvilfmag,24998
-invrs_opt/parameterization/pixel.py,sha256=4qCYDUCcFPr8W94whX5YFllrXgyZbPBVIJBf8m_Dv4k,1173
-invrs_opt/parameterization/transforms.py,sha256=8GzaIsUuuXvMCLiqAEEfxmi9qE9KqHzbuTj_m0GjH3w,8216
-invrs_opt-0.7.2.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-invrs_opt-0.7.2.dist-info/METADATA,sha256=jx_MxaNkrJRUBRubRlDIytDRznP_1P9sWRHrqihaKi8,32633
-invrs_opt-0.7.2.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
-invrs_opt-0.7.2.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
-invrs_opt-0.7.2.dist-info/RECORD,,