invrs-opt 0.9.3__py3-none-any.whl → 0.10.0__py3-none-any.whl

invrs_opt/__init__.py

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

invrs_opt/experimental/client.py

@@ -132,7 +132,11 @@ def optimizer_client(
         response = json.loads(get_response.text)
         return json_utils.pytree_from_json(response[labels.PARAMS])
 
-    return base.Optimizer(init=init_fn, update=update_fn, params=params_fn)
+    return base.Optimizer(
+        init=init_fn,
+        update=update_fn,  # type: ignore[arg-type]
+        params=params_fn,
+    )
 
 
 # -----------------------------------------------------------------------------

invrs_opt/optimizers/base.py

@@ -7,6 +7,7 @@ import dataclasses
 import inspect
 from typing import Any, Protocol
 
+import jax.numpy as jnp
 import optax  # type: ignore[import-untyped]
 from totypes import json_utils
 
@@ -34,7 +35,7 @@ class UpdateFn(Protocol):
         self,
         *,
         grad: PyTree,
-        value: float,
+        value: jnp.ndarray,
         params: PyTree,
         state: PyTree,
     ) -> PyTree:

invrs_opt/optimizers/lbfgsb.py

@@ -335,7 +335,7 @@ def parameterized_lbfgsb(
     def update_fn(
         *,
         grad: PyTree,
-        value: float,
+        value: jnp.ndarray,
         params: PyTree,
         state: LbfgsbState,
     ) -> LbfgsbState:
@@ -349,12 +349,14 @@ def parameterized_lbfgsb(
         ) -> Tuple[PyTree, NumpyLbfgsbDict]:
             assert onp.size(value) == 1
             scipy_lbfgsb_state = ScipyLbfgsbState.from_jax(jax_lbfgsb_state)
+            flat_latent_params = scipy_lbfgsb_state.x.copy()
             scipy_lbfgsb_state.update(
                 grad=onp.array(flat_latent_grad, dtype=onp.float64),
                 value=onp.array(value, dtype=onp.float64),
             )
-            flat_latent_params = jnp.asarray(scipy_lbfgsb_state.x)
-            return flat_latent_params, scipy_lbfgsb_state.to_dict()
+            updated_flat_latent_params = scipy_lbfgsb_state.x
+            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)
@@ -395,16 +397,21 @@ def parameterized_lbfgsb(
             latents_grad
         )  # type: ignore[no-untyped-call]
 
-        flat_latents, jax_lbfgsb_state = jax.pure_callback(
+        flat_latent_updates, jax_lbfgsb_state = jax.pure_callback(
             _update_pure,
             (flat_latents_grad, jax_lbfgsb_state),
             flat_latents_grad,
             value,
             jax_lbfgsb_state,
         )
-        latents = unflatten_fn(flat_latents)
-        latent_params = param_base.combine_density_metadata(metadata, latents)
-        latent_params = _update_parameterized_densities(latent_params, step)
+        latent_updates = unflatten_fn(flat_latent_updates)
+        latent_params = _apply_updates(
+            params=latent_params,
+            updates=param_base.combine_density_metadata(metadata, latent_updates),
+            value=value,
+            step=step,
+        )
+        latent_params = _clip(latent_params)
         params = _params_from_latent_params(latent_params)
         return step + 1, params, latent_params, jax_lbfgsb_state
 
@@ -433,15 +440,24 @@ def parameterized_lbfgsb(
             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)
+    def _apply_updates(
+        params: PyTree,
+        updates: PyTree,
+        value: jnp.ndarray,
+        step: int,
+    ) -> PyTree:
+        def _leaf_apply_updates(update: Any, leaf: Any) -> Any:
+            if _is_parameterized_density(leaf):
+                return density_parameterization.update(
+                    params=leaf, updates=update, value=value, step=step
+                )
+            else:
+                return optax.apply_updates(params=leaf, updates=update)
 
         return tree_util.tree_map(
-            _update_leaf,
-            latent_params,
+            _leaf_apply_updates,
+            updates,
+            params,
             is_leaf=_is_parameterized_density,
         )
 

invrs_opt/optimizers/wrapped_optax.py

@@ -197,12 +197,12 @@ def parameterized_wrapped_optax(
     def update_fn(
         *,
         grad: PyTree,
-        value: float,
+        value: jnp.ndarray,
         params: PyTree,
         state: WrappedOptaxState,
     ) -> WrappedOptaxState:
         """Updates the state."""
-        del value, params
+        del params
 
         step, params, latent_params, opt_state = state
         metadata, latents = param_base.partition_density_metadata(latent_params)
@@ -233,12 +233,14 @@ def parameterized_wrapped_optax(
             lambda a, b: a + b, latents_grad, constraint_loss_grad
         )
 
-        updates, opt_state = opt.update(latents_grad, state=opt_state, params=latents)
-        latents = optax.apply_updates(params=latents, updates=updates)
-
-        latent_params = param_base.combine_density_metadata(metadata, latents)
+        latent_updates, opt_state = opt.update(latents_grad, opt_state, params=latents)
+        latent_params = _apply_updates(
+            params=latent_params,
+            updates=param_base.combine_density_metadata(metadata, latent_updates),
+            value=value,
+            step=step,
+        )
         latent_params = _clip(latent_params)
-        latent_params = _update_parameterized_densities(latent_params, step + 1)
         params = _params_from_latent_params(latent_params)
         return (step + 1, params, latent_params, opt_state)
 
@@ -267,15 +269,24 @@ def parameterized_wrapped_optax(
             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)
+    def _apply_updates(
+        params: PyTree,
+        updates: PyTree,
+        value: jnp.ndarray,
+        step: int,
+    ) -> PyTree:
+        def _leaf_apply_updates(update: Any, leaf: Any) -> Any:
+            if _is_parameterized_density(leaf):
+                return density_parameterization.update(
+                    params=leaf, updates=update, value=value, step=step
+                )
+            else:
+                return optax.apply_updates(params=leaf, updates=update)
 
         return tree_util.tree_map(
-            _update_leaf,
-            latent_params,
+            _leaf_apply_updates,
+            updates,
+            params,
             is_leaf=_is_parameterized_density,
         )
 

invrs_opt/parameterization/base.py

@@ -97,7 +97,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:
+    def __call__(
+        self,
+        params: PyTree,
+        updates: PyTree,
+        value: jnp.ndarray,
+        step: int,
+    ) -> PyTree:
         ...
 
 

invrs_opt/parameterization/filter_project.py

@@ -115,10 +115,20 @@ def filter_project(beta: float) -> base.Density2DParameterization:
         del params
         return jnp.asarray(0.0)
 
-    def update_fn(params: FilterProjectParams, step: int) -> FilterProjectParams:
+    def update_fn(
+        params: FilterProjectParams,
+        updates: FilterProjectParams,
+        value: jnp.ndarray,
+        step: int,
+    ) -> FilterProjectParams:
         """Perform updates to `params` required for the given `step`."""
-        del step
-        return params
+        del step, value
+        return FilterProjectParams(
+            latents=tree_util.tree_map(
+                lambda a, b: a + b, params.latents, updates.latents
+            ),
+            metadata=params.metadata,
+        )
 
     return base.Density2DParameterization(
         to_density=to_density_fn,

invrs_opt/parameterization/gaussian_levelset.py

@@ -218,38 +218,31 @@ def gaussian_levelset(
         pad_pixels: int = 2,
     ) -> jnp.ndarray:
         """Computes constraints associated with the params."""
-        length_scale_constraint, curvature_constraint = _levelset_constraints(
-            params,
-            beta=length_scale_constraint_beta,
+        return analytical_constraints(
+            params=params,
             length_scale_constraint_factor=length_scale_constraint_factor,
-            pad_pixels=pad_pixels,
-        )
-        fixed_pixel_constraint = _fixed_pixel_constraint(
-            params,
+            length_scale_constraint_beta=length_scale_constraint_beta,
+            length_scale_constraint_weight=length_scale_constraint_weight,
+            curvature_constraint_weight=curvature_constraint_weight,
+            fixed_pixel_constraint_weight=fixed_pixel_constraint_weight,
             mask_gradient=mask_gradient,
             pad_pixels=pad_pixels,
         )
 
-        constraints = jnp.stack(
-            [
-                length_scale_constraint * length_scale_constraint_weight,
-                curvature_constraint * curvature_constraint_weight,
-                fixed_pixel_constraint * fixed_pixel_constraint_weight,
-            ],
-            axis=-1,
-        )
-
-        # Normalize constraints to make them (somewhat) resolution-independent.
-        example_density = _example_density(params)
-        length_scale = 0.5 * (
-            example_density.minimum_spacing + example_density.minimum_width
-        )
-        return constraints / length_scale**2
-
-    def update_fn(params: GaussianLevelsetParams, step: int) -> GaussianLevelsetParams:
+    def update_fn(
+        params: GaussianLevelsetParams,
+        updates: GaussianLevelsetParams,
+        value: jnp.ndarray,
+        step: int,
+    ) -> GaussianLevelsetParams:
         """Perform updates to `params` required for the given `step`."""
-        del step
-        return params
+        del step, value
+        return GaussianLevelsetParams(
+            latents=tree_util.tree_map(
+                lambda a, b: a + b, params.latents, updates.latents
+            ),
+            metadata=params.metadata,
+        )
 
     return base.Density2DParameterization(
         to_density=to_density_fn,
@@ -403,6 +396,46 @@ def _phi_from_params(
 # -----------------------------------------------------------------------------
 
 
+def analytical_constraints(
+    params: GaussianLevelsetParams,
+    length_scale_constraint_factor: float,
+    length_scale_constraint_beta: float,
+    length_scale_constraint_weight: float,
+    curvature_constraint_weight: float,
+    fixed_pixel_constraint_weight: float,
+    mask_gradient: bool,
+    pad_pixels: int,
+) -> jnp.ndarray:
+    """Computes analytical levelset constraints associated with the params."""
+    length_scale_constraint, curvature_constraint = _levelset_constraints(
+        params,
+        beta=length_scale_constraint_beta,
+        length_scale_constraint_factor=length_scale_constraint_factor,
+        pad_pixels=pad_pixels,
+    )
+    fixed_pixel_constraint = _fixed_pixel_constraint(
+        params,
+        mask_gradient=mask_gradient,
+        pad_pixels=pad_pixels,
+    )
+
+    constraints = jnp.stack(
+        [
+            length_scale_constraint * length_scale_constraint_weight,
+            curvature_constraint * curvature_constraint_weight,
+            fixed_pixel_constraint * fixed_pixel_constraint_weight,
+        ],
+        axis=-1,
+    )
+
+    # Normalize constraints to make them (somewhat) resolution-independent.
+    example_density = _example_density(params)
+    length_scale = 0.5 * (
+        example_density.minimum_spacing + example_density.minimum_width
+    )
+    return constraints / length_scale**2
+
+
 def _fixed_pixel_constraint(
     params: GaussianLevelsetParams,
     mask_gradient: bool,
@@ -486,11 +519,13 @@ def _levelset_constraints(
     )
 
     d = minimum_length_scale * length_scale_constraint_factor
-    length_scale_constraint = (
-        jnp.abs(phi_vv) / (jnp.pi / d * jnp.abs(phi) + beta * phi_v) - jnp.pi / d
-    )
+    denom = jnp.pi / d * jnp.abs(phi) + beta * phi_v
+    denom_safe = jnp.where(jnp.isclose(phi_vv, 0.0), 1.0, denom)
+    length_scale_constraint = jnp.abs(phi_vv) / denom_safe - jnp.pi / d
+
+    curvature_denom_safe = jnp.where(jnp.isclose(phi_v, 0.0), 1.0, phi)
     curvature_constraint = (
-        jnp.abs(inverse_radius * jnp.arctan(phi_v / phi)) - jnp.pi / d
+        jnp.abs(inverse_radius * jnp.arctan(phi_v / curvature_denom_safe)) - jnp.pi / d
     )
 
     # Downsample so that constraints shape matches the density shape.

invrs_opt/parameterization/pixel.py

@@ -52,9 +52,20 @@ def pixel() -> base.Density2DParameterization:
         del params
         return jnp.asarray(0.0)
 
-    def update_fn(params: PixelParams, step: int) -> PixelParams:
-        del step
-        return params
+    def update_fn(
+        params: PixelParams,
+        updates: PixelParams,
+        value: jnp.ndarray,
+        step: int,
+    ) -> PixelParams:
+        """Perform updates to `params` required for the given `step`."""
+        del step, value
+        return PixelParams(
+            latents=tree_util.tree_map(
+                lambda a, b: a + b, params.latents, updates.latents
+            ),
+            metadata=params.metadata,
+        )
 
     return base.Density2DParameterization(
         from_density=from_density_fn,

{invrs_opt-0.9.3.dist-info → invrs_opt-0.10.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: invrs_opt
-Version: 0.9.3
+Version: 0.10.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.9.3`
+`v0.10.0`
 
 ## Overview
 

invrs_opt-0.10.0.dist-info/RECORD

@@ -0,0 +1,20 @@
+invrs_opt/__init__.py,sha256=8rVngIALR7klCVPflhoqcaf745snBRiTOkwQuq6xYvE,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=C85Ejvq9kJO6BorrLwxeaHLzMCQwFfJ7227JDHACCd0,36840
+invrs_opt/optimizers/wrapped_optax.py,sha256=781-8v_TlHsGaQF9Se9_iOEvtOLOr-BesTLudYarSlg,13685
+invrs_opt/parameterization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+invrs_opt/parameterization/base.py,sha256=jSwrEO86lGkYQG5gWsHvcIMWpZnnbdiKpn--2qaU02g,5362
+invrs_opt/parameterization/filter_project.py,sha256=XL3HTEBLrF-q_75TjhOWLNdfUOSEEjKcoM7Qj844QpQ,4590
+invrs_opt/parameterization/gaussian_levelset.py,sha256=-6foekLTFoZDtMKuoMEvdxMJt0_zTxrKNJo0Vn-Rv80,26073
+invrs_opt/parameterization/pixel.py,sha256=YWkyBhfYtzI8cQ-M90PAZqRAbabwVaUh0UiYIGegQHI,1955
+invrs_opt/parameterization/transforms.py,sha256=8GzaIsUuuXvMCLiqAEEfxmi9qE9KqHzbuTj_m0GjH3w,8216
+invrs_opt-0.10.0.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
+invrs_opt-0.10.0.dist-info/METADATA,sha256=LN2Csi1bX4q2iagAKnwNl-jXZxKQmkxtX-fYq9QX7hs,32643
+invrs_opt-0.10.0.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
+invrs_opt-0.10.0.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
+invrs_opt-0.10.0.dist-info/RECORD,,

{invrs_opt-0.9.3.dist-info → invrs_opt-0.10.0.dist-info}/WHEEL

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

invrs_opt-0.9.3.dist-info/RECORD

@@ -1,20 +0,0 @@
-invrs_opt/__init__.py,sha256=tCzzgk3fRw4_HT8YLg6RlXIrvv19H11TGJngEGL5Iyk,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=9DrmyCj4Ny04NFVRUTz3pJypbw_j5Gw4wpKfe0WKEv4,36336
-invrs_opt/optimizers/wrapped_optax.py,sha256=VXdCteT2kumqhP81l3p6QiEqwBffoUuJ3UjrAyX5ToA,13468
-invrs_opt/parameterization/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-invrs_opt/parameterization/base.py,sha256=BObzbz6efT2nBjib0_5BSdkCmFi2f0mcZ9VJYpDzO6Q,5278
-invrs_opt/parameterization/filter_project.py,sha256=7Jb8JVENmBTdx3-XmI-VRm4aMjxg_Wtin8tMKKKxWvQ,4309
-invrs_opt/parameterization/gaussian_levelset.py,sha256=Uagx7k69SWmass0YirD5JN8O4QDbwwKTBBjRfkIXvv8,24793
-invrs_opt/parameterization/pixel.py,sha256=AwC4GBNNOysdICvYHv_D2tZdqJmYiRzOUZNq_-R9Z70,1617
-invrs_opt/parameterization/transforms.py,sha256=8GzaIsUuuXvMCLiqAEEfxmi9qE9KqHzbuTj_m0GjH3w,8216
-invrs_opt-0.9.3.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-invrs_opt-0.9.3.dist-info/METADATA,sha256=eMtr2aqWLeT2y6dOzEyETbNWbbeTzh2iCcfGu4G2cuQ,32641
-invrs_opt-0.9.3.dist-info/WHEEL,sha256=GV9aMThwP_4oNCtvEC2ec3qUYutgWeAzklro_0m4WJQ,91
-invrs_opt-0.9.3.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
-invrs_opt-0.9.3.dist-info/RECORD,,