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

invrs_opt/optimizers/wrapped_optax.py

@@ -0,0 +1,300 @@
+"""Defines a wrapper for optax optimizers.
+
+Copyright (c) 2023 The INVRS-IO authors.
+"""
+
+from typing import Any, Optional, 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.optimizers import base
+from invrs_opt.parameterization import (
+    base as parameterization_base,
+    filter_project,
+    gaussian_levelset,
+    pixel,
+)
+
+PyTree = Any
+WrappedOptaxState = Tuple[PyTree, PyTree, PyTree]
+
+
+def wrapped_optax(opt: optax.GradientTransformation) -> base.Optimizer:
+    """Return a wrapped optax optimizer."""
+    return parameterized_wrapped_optax(
+        opt=opt, penalty=0.0, density_parameterization=None
+    )
+
+
+def density_wrapped_optax(
+    opt: optax.GradientTransformation,
+    *,
+    beta: float,
+) -> base.Optimizer:
+    """Wrapped optax optimizer with filter-project density parameterization.
+
+    In the filter-project density parameterization, the optimization variable
+    associated with a density array is a latent density array; the density is obtained
+    by convolving (i.e. "filtering") the latent density with a Gaussian kernel having
+    full-width at half-maximum equal to the length scale (the mean of declared minimum
+    width and minimum spacing). Then, a tanh nonlinearity is used as a smooth threshold
+    operation ("projection").
+
+    Args:
+        opt: The optax optimizer to be wrapped.
+        beta: Determines the sharpness of the thresholding operation.
+
+    Returns:
+        The wrapped optax optimizer.
+    """
+    return parameterized_wrapped_optax(
+        opt=opt,
+        penalty=0.0,
+        density_parameterization=filter_project.filter_project(beta=beta),
+    )
+
+
+def levelset_wrapped_optax(
+    opt: optax.GradientTransformation,
+    *,
+    penalty: float,
+    length_scale_spacing_factor: float = (
+        gaussian_levelset.DEFAULT_LENGTH_SCALE_SPACING_FACTOR
+    ),
+    length_scale_fwhm_factor: float = (
+        gaussian_levelset.DEFAULT_LENGTH_SCALE_FWHM_FACTOR
+    ),
+    length_scale_constraint_factor: float = (
+        gaussian_levelset.DEFAULT_LENGTH_SCALE_CONSTRAINT_FACTOR
+    ),
+    smoothing_factor: int = gaussian_levelset.DEFAULT_SMOOTHING_FACTOR,
+    length_scale_constraint_beta: float = (
+        gaussian_levelset.DEFAULT_LENGTH_SCALE_CONSTRAINT_BETA
+    ),
+    length_scale_constraint_weight: float = (
+        gaussian_levelset.DEFAULT_LENGTH_SCALE_CONSTRAINT_WEIGHT
+    ),
+    curvature_constraint_weight: float = (
+        gaussian_levelset.DEFAULT_CURVATURE_CONSTRAINT_WEIGHT
+    ),
+    fixed_pixel_constraint_weight: float = (
+        gaussian_levelset.DEFAULT_FIXED_PIXEL_CONSTRAINT_WEIGHT
+    ),
+    init_optimizer: optax.GradientTransformation = (
+        gaussian_levelset.DEFAULT_INIT_OPTIMIZER
+    ),
+    init_steps: int = gaussian_levelset.DEFAULT_INIT_STEPS,
+) -> base.Optimizer:
+    """Wrapped optax optimizer with levelset density parameterization.
+
+    In the levelset parameterization, the optimization variable associated with a
+    density array is an array giving the amplitudes of Gaussian radial basis functions
+    that represent a levelset function over the domain of the density. In the levelset
+    parameterization, gradients are nonzero only at the edges of features, and in
+    general the topology of a solution does not change during the course of
+    optimization.
+
+    The spacing and full-width at half-maximum of the Gaussian basis functions gives
+    some amount of control over length scales. In addition, constraints associated with
+    length scale, radius of curvature, and deviation from fixed pixels are
+    automatically computed and penalized with a weight given by `penalty`. In general,
+    this helps ensure that features in an optimized density array violate the specified
+    constraints to a lesser degree. The constraints are based on "Analytical level set
+    fabrication constraints for inverse design," by D. Vercruysse et al. (2019).
+
+    Args:
+        opt: The optax optimizer to be wrapped.
+        penalty: The weight of the fabrication penalty, which combines length scale,
+            curvature, and fixed pixel constraints.
+        length_scale_spacing_factor: The number of levelset control points per unit of
+            minimum length scale (mean of density minimum width and minimum spacing).
+        length_scale_fwhm_factor: The ratio of Gaussian full-width at half-maximum to
+            the minimum length scale.
+        length_scale_constraint_factor: Multiplies the target length scale in the
+            levelset constraints. A value greater than 1 is pessimistic and drives the
+            solution to have a larger length scale (relative to smaller values).
+        smoothing_factor: For values greater than 1, the density is initially computed
+            at higher resolution and then downsampled, yielding smoother geometries.
+        length_scale_constraint_beta: Controls relaxation of the length scale
+            constraint near the zero level.
+        length_scale_constraint_weight: The weight of the length scale constraint in
+            the overall fabrication constraint peenalty.
+        curvature_constraint_weight: The weight of the curvature constraint.
+        fixed_pixel_constraint_weight: The weight of the fixed pixel constraint.
+        init_optimizer: The optimizer used in the initialization of the levelset
+            parameterization. At initialization, the latent parameters are optimized so
+            that the initial parameters match the binarized initial density.
+        init_steps: The number of optimization steps used in the initialization.
+
+    Returns:
+        The wrapped optax optimizer.
+    """
+    return parameterized_wrapped_optax(
+        opt=opt,
+        penalty=penalty,
+        density_parameterization=gaussian_levelset.gaussian_levelset(
+            length_scale_spacing_factor=length_scale_spacing_factor,
+            length_scale_fwhm_factor=length_scale_fwhm_factor,
+            length_scale_constraint_factor=length_scale_constraint_factor,
+            smoothing_factor=smoothing_factor,
+            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,
+            init_optimizer=init_optimizer,
+            init_steps=init_steps,
+        ),
+    )
+
+
+# -----------------------------------------------------------------------------
+# Base parameterized wrapped optax optimizer.
+# -----------------------------------------------------------------------------
+
+
+def parameterized_wrapped_optax(
+    opt: optax.GradientTransformation,
+    density_parameterization: Optional[parameterization_base.Density2DParameterization],
+    penalty: float,
+) -> base.Optimizer:
+    """Wrapped optax optimizer with specified density parameterization.
+
+    Args:
+        opt: The optax `GradientTransformation` to be wrapped.
+        density_parameterization: The parameterization to be used, or `None`. When no
+            parameterization is given, the direct pixel parameterization is used for
+            density arrays.
+        penalty: The weight of the scalar penalty formed from the constraints of the
+            parameterization.
+
+    Returns:
+        The `base.Optimizer`.
+    """
+
+    if density_parameterization is None:
+        density_parameterization = pixel.pixel()
+
+    def _init_latents(params: PyTree) -> PyTree:
+        def _leaf_init_latents(leaf: Any) -> Any:
+            leaf = _clip(leaf)
+            if not _is_density(leaf):
+                return leaf
+            return density_parameterization.from_density(leaf)
+
+        return tree_util.tree_map(_leaf_init_latents, params, is_leaf=_is_custom_type)
+
+    def _params_from_latents(params: PyTree) -> PyTree:
+        def _leaf_params_from_latents(leaf: Any) -> Any:
+            if not _is_parameterized_density(leaf):
+                return leaf
+            return density_parameterization.to_density(leaf)
+
+        return tree_util.tree_map(
+            _leaf_params_from_latents,
+            params,
+            is_leaf=_is_parameterized_density,
+        )
+
+    def _constraint_loss(latent_params: PyTree) -> jnp.ndarray:
+        def _constraint_loss_leaf(
+            params: parameterization_base.ParameterizedDensity2DArrayBase,
+        ) -> jnp.ndarray:
+            constraints = density_parameterization.constraints(params)
+            constraints = tree_util.tree_map(
+                lambda x: jnp.sum(jnp.maximum(x, 0.0)),
+                constraints,
+            )
+            return jnp.sum(jnp.asarray(constraints))
+
+        losses = [0.0] + [
+            _constraint_loss_leaf(p)
+            for p in tree_util.tree_leaves(
+                latent_params, is_leaf=_is_parameterized_density
+            )
+            if _is_parameterized_density(p)
+        ]
+        return penalty * jnp.sum(jnp.asarray(losses))
+
+    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)
+
+    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, params
+
+        _, latent_params, opt_state = state
+        _, vjp_fn = jax.vjp(_params_from_latents, latent_params)
+        (latent_grad,) = vjp_fn(grad)
+
+        if not (
+            tree_util.tree_structure(latent_grad)
+            == tree_util.tree_structure(latent_params)  # type: ignore[operator]
+        ):
+            raise ValueError(
+                f"Tree structure of `latent_grad` was different than expected, got \n"
+                f"{tree_util.tree_structure(latent_grad)} but expected \n"
+                f"{tree_util.tree_structure(latent_params)}."
+            )
+
+        constraint_loss_grad = jax.grad(_constraint_loss)(latent_params)
+        latent_grad = tree_util.tree_map(
+            lambda a, b: a + b, latent_grad, constraint_loss_grad
+        )
+
+        updates, opt_state = opt.update(
+            updates=latent_grad, state=opt_state, params=latent_params
+        )
+        latent_params = optax.apply_updates(params=latent_params, updates=updates)
+        latent_params = _clip(latent_params)
+        params = _params_from_latents(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_parameterized_density(leaf: Any) -> Any:
+    """Return `True` if `leaf` is a parameterized density array."""
+    return isinstance(leaf, parameterization_base.ParameterizedDensity2DArrayBase)
+
+
+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/parameterization/base.py

@@ -0,0 +1,148 @@
+"""Base types for density parameterizations.
+
+Copyright (c) 2023 The INVRS-IO authors.
+"""
+
+import dataclasses
+from typing import Any, Optional, Protocol, Sequence, Tuple
+
+import jax.numpy as jnp
+import numpy as onp
+from jax import tree_util
+from totypes import json_utils, types
+
+Array = jnp.ndarray | onp.ndarray[Any, Any]
+PyTree = Any
+
+
+class ParameterizedDensity2DArrayBase:
+    """Base class for parameterized density arrays."""
+
+    pass
+
+
+class FromDensityFn(Protocol):
+    """Generate the latent representation of a density array."""
+
+    def __call__(
+        self, density: types.Density2DArray
+    ) -> ParameterizedDensity2DArrayBase:
+        ...
+
+
+class ToDensityFn(Protocol):
+    """Generate a density from its latent representation."""
+
+    def __call__(self, params: PyTree) -> types.Density2DArray:
+        ...
+
+
+class ConstraintsFn(Protocol):
+    """Compute constraints for a latent representation of a density array."""
+
+    def __call__(self, params: PyTree) -> jnp.ndarray:
+        ...
+
+
+@dataclasses.dataclass
+class Density2DParameterization:
+    """Stores `(from_density, to_density, constraints)` function triple."""
+
+    from_density: FromDensityFn
+    to_density: ToDensityFn
+    constraints: ConstraintsFn
+
+
+@dataclasses.dataclass
+class Density2DMetadata:
+    """Stores the metadata of a `Density2DArray`."""
+
+    lower_bound: float
+    upper_bound: float
+    fixed_solid: Optional[Array]
+    fixed_void: Optional[Array]
+    minimum_width: int
+    minimum_spacing: int
+    periodic: Sequence[bool]
+    symmetries: Sequence[str]
+
+    def __post_init__(self) -> None:
+        self.periodic = tuple(self.periodic)
+        self.symmetries = tuple(self.symmetries)
+
+
+def _flatten_density_2d_metadata(
+    metadata: Density2DMetadata,
+) -> Tuple[
+    Tuple[()],
+    Tuple[
+        float,
+        float,
+        types.HashableWrapper,
+        types.HashableWrapper,
+        int,
+        int,
+        Sequence[bool],
+        Sequence[str],
+    ],
+]:
+    """Flattens a `Density2DMetadata` into children and auxilliary data."""
+    return (
+        (),
+        (
+            metadata.lower_bound,
+            metadata.upper_bound,
+            types.HashableWrapper(metadata.fixed_solid),
+            types.HashableWrapper(metadata.fixed_void),
+            metadata.minimum_width,
+            metadata.minimum_spacing,
+            metadata.periodic,
+            metadata.symmetries,
+        ),
+    )
+
+
+def _unflatten_density_2d_metadata(
+    aux: Tuple[
+        float,
+        float,
+        types.HashableWrapper,
+        types.HashableWrapper,
+        int,
+        int,
+        Sequence[bool],
+        Sequence[str],
+    ],
+    children: Tuple[()],
+) -> Density2DMetadata:
+    """Unflattens a flattened `Density2DMetadata`."""
+    del children
+    (
+        lower_bound,
+        upper_bound,
+        wrapped_fixed_solid,
+        wrapped_fixed_void,
+        minimum_width,
+        minimum_spacing,
+        periodic,
+        symmetries,
+    ) = aux
+    return Density2DMetadata(
+        lower_bound=lower_bound,
+        upper_bound=upper_bound,
+        fixed_solid=wrapped_fixed_solid.array,  # type: ignore[arg-type]
+        fixed_void=wrapped_fixed_void.array,  # type: ignore[arg-type]
+        minimum_width=minimum_width,
+        minimum_spacing=minimum_spacing,
+        periodic=tuple(periodic),
+        symmetries=tuple(symmetries),
+    )
+
+
+tree_util.register_pytree_node(
+    Density2DMetadata,
+    flatten_func=_flatten_density_2d_metadata,
+    unflatten_func=_unflatten_density_2d_metadata,
+)
+
+json_utils.register_custom_type(Density2DMetadata)

invrs_opt/parameterization/filter_project.py

@@ -0,0 +1,92 @@
+"""Defines filter-and-project density parameterization.
+
+Copyright (c) 2023 The INVRS-IO authors.
+"""
+
+import dataclasses
+
+import jax.numpy as jnp
+from jax import tree_util
+from totypes import json_utils, types
+
+from invrs_opt.parameterization import base, transforms
+
+
+@dataclasses.dataclass
+class FilterAndProjectParams(base.ParameterizedDensity2DArrayBase):
+    """Stores the latent parameters of the pixel parameterization.
+
+    Attributes:
+        latent_density: The latent variable from which the density is obtained.
+        beta: Determines the sharpness of the thresholding operation.
+    """
+
+    latent_density: types.Density2DArray
+    beta: float
+
+
+tree_util.register_dataclass(
+    FilterAndProjectParams,
+    data_fields=["latent_density"],
+    meta_fields=["beta"],
+)
+
+json_utils.register_custom_type(FilterAndProjectParams)
+
+
+def filter_project(beta: float) -> base.Density2DParameterization:
+    """Defines a filter-project parameterization for density arrays.
+
+    The `DensityArray2D` is represented as latent density array that is transformed by,
+
+        transformed = tanh(beta * conv(density.array, gaussian_kernel)) / tanh(beta)
+
+    where the kernel has a full-width at half-maximum determined by the minimum width
+    and spacing parameters of the `DensityArray2D`.
+
+    When the density lower and upper bounds are -1 and +1, this basic expression is
+    Where the bounds differ, the density is scaled before the transform is applied, and
+    then unscaled afterwards.
+
+    Args:
+        beta: Determines the sharpness of the thresholding operation.
+
+    Returns:
+        The `Density2DParameterization`.
+    """
+
+    def from_density_fn(density: types.Density2DArray) -> FilterAndProjectParams:
+        """Return latent parameters for the given `density`."""
+        array = transforms.normalized_array_from_density(density)
+        array = jnp.clip(array, -1, 1)
+        array *= jnp.tanh(beta)
+        latent_array = jnp.arctanh(array) / beta
+        latent_array = transforms.rescale_array_for_density(latent_array, density)
+        return FilterAndProjectParams(
+            latent_density=dataclasses.replace(density, array=latent_array),
+            beta=beta,
+        )
+
+    def to_density_fn(params: FilterAndProjectParams) -> types.Density2DArray:
+        """Return a density from the latent parameters."""
+        transformed = types.symmetrize_density(params.latent_density)
+        transformed = transforms.density_gaussian_filter_and_tanh(
+            transformed, beta=params.beta
+        )
+        # Scale to ensure that the full valid range of the density array is reachable.
+        mid_value = (transformed.lower_bound + transformed.upper_bound) / 2
+        transformed = tree_util.tree_map(
+            lambda array: mid_value + (array - mid_value) / jnp.tanh(beta), transformed
+        )
+        return transforms.apply_fixed_pixels(transformed)
+
+    def constraints_fn(params: FilterAndProjectParams) -> jnp.ndarray:
+        """Computes constraints associated with the params."""
+        del params
+        return jnp.asarray(0.0)
+
+    return base.Density2DParameterization(
+        to_density=to_density_fn,
+        from_density=from_density_fn,
+        constraints=constraints_fn,
+    )