invrs-opt 0.7.2__py3-none-any.whl → 0.8.0__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.0"
 __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:
@@ -205,7 +205,7 @@ def parameterized_wrapped_optax(
         ) -> 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))
@@ -219,15 +219,27 @@ def parameterized_wrapped_optax(
         ]
         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 params, latent_params, opt.init(latent_params)
+        return 0, params, latent_params, opt.init(latent_params)
 
     def params_fn(state: WrappedOptaxState) -> PyTree:
         """Returns the parameters for the given `state`."""
-        params, _, _ = state
+        _, params, _, _ = state
         return params
 
     def update_fn(
@@ -240,7 +252,7 @@ def parameterized_wrapped_optax(
         """Updates the state."""
         del value, params
 
-        _, latent_params, opt_state = state
+        step, _, latent_params, opt_state = state
         _, vjp_fn = jax.vjp(_params_from_latents, latent_params)
         (latent_grad,) = vjp_fn(grad)
 
@@ -264,8 +276,9 @@ def parameterized_wrapped_optax(
         )
         latent_params = optax.apply_updates(params=latent_params, updates=updates)
         latent_params = _clip(latent_params)
+        latent_params = _update_parameterized_densities(latent_params, step)
         params = _params_from_latents(latent_params)
-        return params, latent_params, opt_state
+        return step + 1, params, latent_params, opt_state
 
     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.0.dist-info}/METADATA

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

invrs_opt-0.8.0.dist-info/RECORD

@@ -0,0 +1,20 @@
+invrs_opt/__init__.py,sha256=kTrg48iZu7i5OlH0Nqtfh_wBn3be9u2eZgJBRqUr6uQ,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=zPC2j_KkfF2RqeorxB38ovuUsg1SNwdxwhjA7gvOMC4,12387
+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.7.2.dist-info → invrs_opt-0.8.0.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,,