invrs-opt 0.9.4__py3-none-any.whl → 0.10.1__py3-none-any.whl

invrs_opt/__init__.py

@@ -3,7 +3,7 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.9.4"
+__version__ = "v0.10.1"
 __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

@@ -317,7 +317,10 @@ def parameterized_lbfgsb(
         latent_params = _init_latents(params)
         metadata, latents = param_base.partition_density_metadata(latent_params)
         latents, jax_lbfgsb_state = jax.pure_callback(
-            _init_state_pure, _example_state(latents, maxcor), latents
+            _init_state_pure,
+            _example_state(latents, maxcor),
+            latents,
+            vmap_method="sequential",
         )
         latent_params = param_base.combine_density_metadata(metadata, latents)
         return (
@@ -335,7 +338,7 @@ def parameterized_lbfgsb(
     def update_fn(
         *,
         grad: PyTree,
-        value: float,
+        value: jnp.ndarray,
         params: PyTree,
         state: LbfgsbState,
     ) -> LbfgsbState:
@@ -346,15 +349,18 @@ def parameterized_lbfgsb(
             flat_latent_grad: PyTree,
             value: jnp.ndarray,
             jax_lbfgsb_state: JaxLbfgsbDict,
-        ) -> Tuple[PyTree, NumpyLbfgsbDict]:
+        ) -> Tuple[NDArray, 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: NDArray
+            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 +401,22 @@ 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,
+            vmap_method="sequential",
         )
-        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 +445,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

@@ -229,10 +229,20 @@ def gaussian_levelset(
             pad_pixels=pad_pixels,
         )
 
-    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,
@@ -509,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.4.dist-info → invrs_opt-0.10.1.dist-info}/METADATA

@@ -1,10 +1,10 @@
 Metadata-Version: 2.1
 Name: invrs_opt
-Version: 0.9.4
+Version: 0.10.1
 Summary: Algorithms for inverse design
 Author-email: "Martin F. Schubert" <mfschubert@gmail.com>
 Maintainer-email: "Martin F. Schubert" <mfschubert@gmail.com>
-License: GNU LESSER GENERAL PUBLIC LICENSE
+License:                   GNU LESSER GENERAL PUBLIC LICENSE
                                Version 2.1, February 1999
         
          Copyright (C) 1991, 1999 Free Software Foundation, Inc.
@@ -513,7 +513,7 @@ Keywords: topology,optimization,jax,inverse design
 Requires-Python: >=3.7
 Description-Content-Type: text/markdown
 License-File: LICENSE
-Requires-Dist: jax<=0.4.35
+Requires-Dist: jax>=0.4.35
 Requires-Dist: jaxlib
 Requires-Dist: numpy
 Requires-Dist: requests
@@ -521,19 +521,19 @@ Requires-Dist: optax
 Requires-Dist: scipy
 Requires-Dist: totypes
 Requires-Dist: types-requests
-Provides-Extra: dev
-Requires-Dist: bump-my-version; extra == "dev"
-Requires-Dist: darglint; extra == "dev"
-Requires-Dist: mypy; extra == "dev"
-Requires-Dist: pre-commit; extra == "dev"
 Provides-Extra: tests
 Requires-Dist: parameterized; extra == "tests"
 Requires-Dist: pytest; extra == "tests"
 Requires-Dist: pytest-cov; extra == "tests"
 Requires-Dist: pytest-subtests; extra == "tests"
+Provides-Extra: dev
+Requires-Dist: bump-my-version; extra == "dev"
+Requires-Dist: darglint; extra == "dev"
+Requires-Dist: mypy; extra == "dev"
+Requires-Dist: pre-commit; extra == "dev"
 
 # invrs-opt - Optimization algorithms for inverse design
-`v0.9.4`
+`v0.10.1`
 
 ## Overview
 

invrs_opt-0.10.1.dist-info/RECORD

@@ -0,0 +1,20 @@
+invrs_opt/__init__.py,sha256=bd6NWb7l6z0v6kQif3vfX3xgxdlZrp9L9EFt_Ds20UU,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=kIluuCL1TjkNrdiwqsHEzqz7PEnPa0txpGa3rAw1GXU,36983
+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.1.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
+invrs_opt-0.10.1.dist-info/METADATA,sha256=gCJPUnv5b3erC7R4Fb-Hlhbo9uAjYK0fgTYL0qFg7Jg,32661
+invrs_opt-0.10.1.dist-info/WHEEL,sha256=PZUExdf71Ui_so67QXpySuHtCi3-J3wvF4ORK6k_S8U,91
+invrs_opt-0.10.1.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
+invrs_opt-0.10.1.dist-info/RECORD,,

{invrs_opt-0.9.4.dist-info → invrs_opt-0.10.1.dist-info}/WHEEL

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

invrs_opt-0.9.4.dist-info/RECORD

@@ -1,20 +0,0 @@
-invrs_opt/__init__.py,sha256=bVIs0NxPxNuRUOypBIE68qx-SA1lFCmqmo5cqvRpxmU,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=Ka4hW_OLxUaIPHQsyIOlryG7i1mC-LTIgiQQhCPwHPk,25626
-invrs_opt/parameterization/pixel.py,sha256=AwC4GBNNOysdICvYHv_D2tZdqJmYiRzOUZNq_-R9Z70,1617
-invrs_opt/parameterization/transforms.py,sha256=8GzaIsUuuXvMCLiqAEEfxmi9qE9KqHzbuTj_m0GjH3w,8216
-invrs_opt-0.9.4.dist-info/LICENSE,sha256=IMF9i4xIpgCADf0U-V1cuf9HBmqWQd3qtI3FSuyW4zE,26526
-invrs_opt-0.9.4.dist-info/METADATA,sha256=SOt6aECrbTQNI0QuznoPrWKzUdMeahaST0G4Q7sccxE,32641
-invrs_opt-0.9.4.dist-info/WHEEL,sha256=OVMc5UfuAQiSplgO0_WdW7vXVGAt9Hdd6qtN4HotdyA,91
-invrs_opt-0.9.4.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
-invrs_opt-0.9.4.dist-info/RECORD,,