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

invrs_opt/{transform.py → parameterization/transforms.py}

@@ -62,6 +62,20 @@ def density_gaussian_filter_and_tanh(
     return transformed_density
 
 
+def _gaussian_kernel(fwhm: float, fwhm_size_multiple: float) -> jnp.ndarray:
+    """Returns a Gaussian kernel with the specified full-width at half-maximum."""
+    with jax.ensure_compile_time_eval():
+        kernel_size = max(1, int(jnp.ceil(fwhm * fwhm_size_multiple)))
+    # Ensure the kernel size is odd, so that there is always a central pixel which will
+    # contain the peak value of the Gaussian.
+    kernel_size += (kernel_size + 1) % 2
+    d = jnp.arange(0.5, kernel_size) - kernel_size / 2
+    x = d[:, jnp.newaxis]
+    y = d[jnp.newaxis, :]
+    sigma = fwhm / (2 * jnp.sqrt(2 * jnp.log(2)))
+    return jnp.exp(-(x**2 + y**2) / (2 * sigma**2))
+
+
 def normalized_array_from_density(density: types.Density2DArray) -> jnp.ndarray:
     """Returns an array with values scaled to the range `(-1, 1)`."""
     value_mid = (density.upper_bound + density.lower_bound) / 2
@@ -154,15 +168,66 @@ def _pad_mode_for_density(density: types.Density2DArray) -> Union[str, Tuple[str
     )
 
 
-def _gaussian_kernel(fwhm: float, fwhm_size_multiple: float) -> jnp.ndarray:
-    """Returns a Gaussian kernel with the specified full-width at half-maximum."""
+def resample(
+    x: jnp.ndarray,
+    shape: Tuple[int, ...],
+    method: jax.image.ResizeMethod = jax.image.ResizeMethod.LINEAR,
+) -> jnp.ndarray:
+    """Resamples `x` to have the specified `shape`.
+
+    The algorithm first upsamples `x` so that the pixels in the output image are
+    comprised of an integer number of pixels in the upsampled `x`, and then
+    performs box downsampling.
+
+    Args:
+        x: The array to be resampled.
+        shape: The shape of the output array.
+        method: The method used to resize `x` prior to box downsampling.
+
+    Returns:
+        The resampled array.
+    """
+    if x.ndim != len(shape):
+        raise ValueError(
+            f"`shape` must have length matching number of dimensions in `x`, "
+            f"but got {shape} when `x` had shape {x.shape}."
+        )
+
     with jax.ensure_compile_time_eval():
-        kernel_size = max(1, int(jnp.ceil(fwhm * fwhm_size_multiple)))
-    # Ensure the kernel size is odd, so that there is always a central pixel which will
-    # contain the peak value of the Gaussian.
-    kernel_size += (kernel_size + 1) % 2
-    d = jnp.arange(0.5, kernel_size) - kernel_size / 2
-    x = d[:, jnp.newaxis]
-    y = d[jnp.newaxis, :]
-    sigma = fwhm / (2 * jnp.sqrt(2 * jnp.log(2)))
-    return jnp.exp(-(x**2 + y**2) / (2 * sigma**2))
+        factor = [int(jnp.ceil(dx / d)) for dx, d in zip(x.shape, shape)]
+        upsampled_shape = tuple([d * f for d, f in zip(shape, factor)])
+
+    x_upsampled = jax.image.resize(
+        image=x,
+        shape=upsampled_shape,
+        method=method,
+    )
+
+    return box_downsample(x_upsampled, shape)
+
+
+def box_downsample(x: jnp.ndarray, shape: Tuple[int, ...]) -> jnp.ndarray:
+    """Downsamples `x` to a coarser resolution array using box downsampling.
+
+    Box downsampling forms nonoverlapping windows and simply averages the
+    pixels within each window. For example, downsampling `(0, 1, 2, 3, 4, 5)`
+    with a factor of `2` yields `(0.5, 2.5, 4.5)`.
+
+    Args:
+        x: The array to be downsampled.
+        shape: The shape of the output array; each axis dimension must evenly
+            divide the corresponding axis dimension in `x`.
+
+    Returns:
+        The output array with shape `shape`.
+    """
+    if x.ndim != len(shape) or any([(d % s) != 0 for d, s in zip(x.shape, shape)]):
+        raise ValueError(
+            f"Each axis of `shape` must evenly divide the corresponding axis "
+            f"dimension in `x`, but got shape {shape} when `x` has shape "
+            f"{x.shape}."
+        )
+    shape = sum([(s, d // s) for d, s in zip(x.shape, shape)], ())
+    axes = list(range(1, 2 * x.ndim, 2))
+    x = x.reshape(shape)
+    return jnp.mean(x, axis=axes)

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

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: invrs_opt
-Version: 0.6.0
+Version: 0.7.0
 Summary: Algorithms for inverse design
 Author-email: "Martin F. Schubert" <mfschubert@gmail.com>
 Maintainer-email: "Martin F. Schubert" <mfschubert@gmail.com>
@@ -39,18 +39,18 @@ 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'
+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'
+Requires-Dist: parameterized; extra == "tests"
+Requires-Dist: pytest; extra == "tests"
+Requires-Dist: pytest-cov; extra == "tests"
+Requires-Dist: pytest-subtests; extra == "tests"
 
 # invrs-opt - Optimization algorithms for inverse design
-`v0.6.0`
+`v0.7.0`
 
 ## Overview
 

invrs_opt-0.7.0.dist-info/RECORD

@@ -0,0 +1,20 @@
+invrs_opt/__init__.py,sha256=cRGTL0pWiEmPY26dG0bfKaS_RV5I5N3FdFkKLM8Z-Wk,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=ol9skiS4_itx6pmqTx49fqqRZGzHHSxjKPKguy6602s,1199
+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=OYXPNC3GIGqCxea-u7rTXCOAEERLoict9kqnjYvGGto,24838
+invrs_opt/parameterization/pixel.py,sha256=4qCYDUCcFPr8W94whX5YFllrXgyZbPBVIJBf8m_Dv4k,1173
+invrs_opt/parameterization/transforms.py,sha256=8GzaIsUuuXvMCLiqAEEfxmi9qE9KqHzbuTj_m0GjH3w,8216
+invrs_opt-0.7.0.dist-info/LICENSE,sha256=ty6jHPvpyjHy6dbhnu6aDSY05bbl2jQTjnq9u1sBCfM,1078
+invrs_opt-0.7.0.dist-info/METADATA,sha256=EBEC3SgXD2hwdZ5s-z2dydbr5oSPSsgB34iMSdhaLeE,3339
+invrs_opt-0.7.0.dist-info/WHEEL,sha256=R0nc6qTxuoLk7ShA2_Y-UWkN8ZdfDBG2B6Eqpz2WXbs,91
+invrs_opt-0.7.0.dist-info/top_level.txt,sha256=hOziS2uJ_NgwaW9yhtOfeuYnm1X7vobPBcp_6eVWKfM,10
+invrs_opt-0.7.0.dist-info/RECORD,,

{invrs_opt-0.6.0.dist-info → invrs_opt-0.7.0.dist-info}/WHEEL

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

invrs_opt/wrapped_optax/wrapped_optax.py

@@ -1,150 +0,0 @@
-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.6.0.dist-info/RECORD

@@ -1,16 +0,0 @@
-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,,