invrs-opt 0.5.2__py3-none-any.whl → 0.6.0__py3-none-any.whl

invrs_opt/__init__.py

@@ -3,8 +3,12 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.5.2"
+__version__ = "v0.6.0"
 __author__ = "Martin F. Schubert <mfschubert@gmail.com>"
 
 from invrs_opt.lbfgsb.lbfgsb import density_lbfgsb as density_lbfgsb
 from invrs_opt.lbfgsb.lbfgsb import lbfgsb as lbfgsb
+from invrs_opt.wrapped_optax.wrapped_optax import (
+    density_wrapped_optax as density_wrapped_optax,
+)
+from invrs_opt.wrapped_optax.wrapped_optax import wrapped_optax as wrapped_optax

invrs_opt/base.py

@@ -6,6 +6,9 @@ Copyright (c) 2023 The INVRS-IO authors.
 import dataclasses
 from typing import Any, Protocol
 
+import optax  # type: ignore[import-untyped]
+from totypes import json_utils
+
 PyTree = Any
 
 
@@ -44,3 +47,8 @@ class Optimizer:
     init: InitFn
     params: ParamsFn
     update: UpdateFn
+
+
+# TODO: consider programatically registering all optax states here.
+json_utils.register_custom_type(optax.EmptyState)
+json_utils.register_custom_type(optax.ScaleByAdamState)

invrs_opt/experimental/client.py

@@ -4,16 +4,15 @@ Copyright (c) 2023 The INVRS-IO authors.
 """
 
 import json
-import requests
 import time
 from typing import Any, Dict, Optional
 
+import requests
 from totypes import json_utils
 
 from invrs_opt import base
 from invrs_opt.experimental import labels
 
-
 PyTree = Any
 StateToken = str
 

invrs_opt/lbfgsb/lbfgsb.py

@@ -16,8 +16,7 @@ from scipy.optimize._lbfgsb_py import (  # type: ignore[import-untyped]
 )
 from totypes import types
 
-from invrs_opt import base
-from invrs_opt.lbfgsb import transform
+from invrs_opt import base, transform
 
 NDArray = onp.ndarray[Any, Any]
 PyTree = Any

invrs_opt/wrapped_optax/wrapped_optax.py

@@ -0,0 +1,150 @@
+import dataclasses
+from typing import Any, Callable, Tuple
+
+import jax
+import jax.numpy as jnp
+import optax  # type: ignore[import-untyped]
+from jax import tree_util
+from totypes import types
+
+from invrs_opt import base, transform
+
+PyTree = Any
+WrappedOptaxState = Tuple[PyTree, PyTree, PyTree]
+
+
+def wrapped_optax(opt: optax.GradientTransformation) -> base.Optimizer:
+    """Return a wrapped optax optimizer."""
+    return transformed_wrapped_optax(
+        opt=opt,
+        transform_fn=lambda x: x,
+        initialize_latent_fn=lambda x: x,
+    )
+
+
+def density_wrapped_optax(
+    opt: optax.GradientTransformation,
+    beta: float,
+) -> base.Optimizer:
+    """Return a wrapped optax optimizer with transforms for density arrays."""
+
+    def transform_fn(tree: PyTree) -> PyTree:
+        return tree_util.tree_map(
+            lambda x: transform_density(x) if _is_density(x) else x,
+            tree,
+            is_leaf=_is_density,
+        )
+
+    def initialize_latent_fn(tree: PyTree) -> PyTree:
+        return tree_util.tree_map(
+            lambda x: initialize_latent_density(x) if _is_density(x) else x,
+            tree,
+            is_leaf=_is_density,
+        )
+
+    def transform_density(density: types.Density2DArray) -> types.Density2DArray:
+        transformed = types.symmetrize_density(density)
+        transformed = transform.density_gaussian_filter_and_tanh(transformed, beta=beta)
+        # Scale to ensure that the full valid range of the density array is reachable.
+        mid_value = (density.lower_bound + density.upper_bound) / 2
+        transformed = tree_util.tree_map(
+            lambda array: mid_value + (array - mid_value) / jnp.tanh(beta), transformed
+        )
+        return transform.apply_fixed_pixels(transformed)
+
+    def initialize_latent_density(
+        density: types.Density2DArray,
+    ) -> types.Density2DArray:
+        array = transform.normalized_array_from_density(density)
+        array = jnp.clip(array, -1, 1)
+        array *= jnp.tanh(beta)
+        latent_array = jnp.arctanh(array) / beta
+        latent_array = transform.rescale_array_for_density(latent_array, density)
+        return dataclasses.replace(density, array=latent_array)
+
+    return transformed_wrapped_optax(
+        opt=opt,
+        transform_fn=transform_fn,
+        initialize_latent_fn=initialize_latent_fn,
+    )
+
+
+def transformed_wrapped_optax(
+    opt: optax.GradientTransformation,
+    transform_fn: Callable[[PyTree], PyTree],
+    initialize_latent_fn: Callable[[PyTree], PyTree],
+) -> base.Optimizer:
+    """Return a wrapped optax optimizer for transformed latent parameters.
+
+    Args:
+        opt: The optax `GradientTransformation` to be wrapped.
+        transform_fn: Function which transforms the internal latent parameters to
+            the parameters returned by the optimizer.
+        initialize_latent_fn: Function which computes the initial latent parameters
+            given the initial parameters.
+
+    Returns:
+        The `base.Optimizer`.
+    """
+
+    def init_fn(params: PyTree) -> WrappedOptaxState:
+        """Initializes the optimization state."""
+        latent_params = initialize_latent_fn(_clip(params))
+        params = transform_fn(latent_params)
+        return params, latent_params, opt.init(latent_params)
+
+    def params_fn(state: WrappedOptaxState) -> PyTree:
+        """Returns the parameters for the given `state`."""
+        params, _, _ = state
+        return params
+
+    def update_fn(
+        *,
+        grad: PyTree,
+        value: float,
+        params: PyTree,
+        state: WrappedOptaxState,
+    ) -> WrappedOptaxState:
+        """Updates the state."""
+        del value
+
+        _, latent_params, opt_state = state
+        _, vjp_fn = jax.vjp(transform_fn, latent_params)
+        (latent_grad,) = vjp_fn(grad)
+
+        updates, opt_state = opt.update(latent_grad, opt_state)
+        latent_params = optax.apply_updates(params=latent_params, updates=updates)
+        latent_params = _clip(latent_params)
+        params = transform_fn(latent_params)
+        return params, latent_params, opt_state
+
+    return base.Optimizer(
+        init=init_fn,
+        params=params_fn,
+        update=update_fn,
+    )
+
+
+def _is_density(leaf: Any) -> Any:
+    """Return `True` if `leaf` is a density array."""
+    return isinstance(leaf, types.Density2DArray)
+
+
+def _is_custom_type(leaf: Any) -> bool:
+    """Return `True` if `leaf` is a recognized custom type."""
+    return isinstance(leaf, (types.BoundedArray, types.Density2DArray))
+
+
+def _clip(pytree: PyTree) -> PyTree:
+    """Clips leaves on `pytree` to their bounds."""
+
+    def _clip_fn(leaf: Any) -> Any:
+        if not _is_custom_type(leaf):
+            return leaf
+        if leaf.lower_bound is None and leaf.upper_bound is None:
+            return leaf
+        return tree_util.tree_map(
+            lambda x: jnp.clip(x, leaf.lower_bound, leaf.upper_bound), leaf
+        )
+
+    return tree_util.tree_map(_clip_fn, pytree, is_leaf=_is_custom_type)

{invrs_opt-0.5.2.dist-info → invrs_opt-0.6.0.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: invrs_opt
-Version: 0.5.2
+Version: 0.6.0
 Summary: Algorithms for inverse design
 Author-email: "Martin F. Schubert" <mfschubert@gmail.com>
 Maintainer-email: "Martin F. Schubert" <mfschubert@gmail.com>
@@ -34,6 +34,7 @@ Requires-Dist: jax
 Requires-Dist: jaxlib
 Requires-Dist: numpy
 Requires-Dist: requests
+Requires-Dist: optax
 Requires-Dist: scipy
 Requires-Dist: totypes
 Requires-Dist: types-requests
@@ -49,7 +50,7 @@ Requires-Dist: pytest-cov ; extra == 'tests'
 Requires-Dist: pytest-subtests ; extra == 'tests'
 
 # invrs-opt - Optimization algorithms for inverse design
-`v0.5.2`
+`v0.6.0`
 
 ## Overview
 

invrs_opt-0.6.0.dist-info/RECORD

@@ -0,0 +1,16 @@
+invrs_opt/__init__.py,sha256=35pvMpeqvJgU4DizUO5hTzeE9j93prbB9RMtcaoFYwg,496
+invrs_opt/base.py,sha256=FdQyPTlWGo03YztI3K2_QBN6Q-v0PeXv6XCyXu_uh_4,1160
+invrs_opt/py.typed,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+invrs_opt/transform.py,sha256=a_Saj9Wq4lvqCJBrg5L2Z9DZ2NVs1xqrHLqha90a9Ws,5971
+invrs_opt/experimental/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+invrs_opt/experimental/client.py,sha256=MqC_TguT9IGrG7WW54vwz6QQMylKkbCjHxFPIG9vQMA,4841
+invrs_opt/experimental/labels.py,sha256=dQDAMPyFMV6mXnMy295z8Ap205DRdVzysXny_Be8FmY,562
+invrs_opt/lbfgsb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+invrs_opt/lbfgsb/lbfgsb.py,sha256=pfrqCaOMco-eHUQe2q03hbla9D2TYqmMB-07jK4-5Ik,27792
+invrs_opt/wrapped_optax/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
+invrs_opt/wrapped_optax/wrapped_optax.py,sha256=-ke0MNCb2EB0ntlj5IHIHrvybOVF4m24DM6JI4_Ktcc,4974
+invrs_opt-0.6.0.dist-info/LICENSE,sha256=ty6jHPvpyjHy6dbhnu6aDSY05bbl2jQTjnq9u1sBCfM,1078
+invrs_opt-0.6.0.dist-info/METADATA,sha256=V4hpzjEovC2CU0IgUBMAgKzPfgqGvLlutQ9X5-FOuIk,3347
+invrs_opt-0.6.0.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
+invrs_opt-0.6.0.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
+invrs_opt-0.6.0.dist-info/RECORD,,

invrs_opt-0.5.2.dist-info/RECORD

@@ -1,14 +0,0 @@
-invrs_opt/__init__.py,sha256=TPyc9kznbXNbM67XmXpnFIGTmGFQmYzPojKRueSZelU,309
-invrs_opt/base.py,sha256=dSX9QkMPzI8ROxy2cFNmMwk_89eQbk0rw94CzvLPQoY,907
-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=td5o_YqqbcSypDrWCVrHGSoF8UxEdOLtKU0z9Dth9lA,4842
-invrs_opt/experimental/labels.py,sha256=dQDAMPyFMV6mXnMy295z8Ap205DRdVzysXny_Be8FmY,562
-invrs_opt/lbfgsb/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
-invrs_opt/lbfgsb/lbfgsb.py,sha256=EKNTDTwEHcO_JeOIckMHVer5hxYcaNuqq9EZgnFnWJk,27820
-invrs_opt/lbfgsb/transform.py,sha256=a_Saj9Wq4lvqCJBrg5L2Z9DZ2NVs1xqrHLqha90a9Ws,5971
-invrs_opt-0.5.2.dist-info/LICENSE,sha256=ty6jHPvpyjHy6dbhnu6aDSY05bbl2jQTjnq9u1sBCfM,1078
-invrs_opt-0.5.2.dist-info/METADATA,sha256=SgIEqMR9ybcipS1NnHi-KSVTSZ3BzWKrv1ctw4jfqcE,3326
-invrs_opt-0.5.2.dist-info/WHEEL,sha256=GJ7t_kWBFywbagK5eo9IoUwLW6oyOeTKmQ-9iHFVNxQ,92
-invrs_opt-0.5.2.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
-invrs_opt-0.5.2.dist-info/RECORD,,