invrs-opt 0.3.2__py3-none-any.whl → 0.10.3__py3-none-any.whl

invrs_opt/optimizers/wrapped_optax.py

@@ -0,0 +1,347 @@
+"""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 param_base,
+    filter_project,
+    gaussian_levelset,
+    pixel,
+)
+
+PyTree = Any
+WrappedOptaxState = Tuple[int, 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[param_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_fn(params: PyTree) -> WrappedOptaxState:
+        """Initializes the optimization state."""
+        latent_params = _init_latents(params)
+        _, latents = param_base.partition_density_metadata(latent_params)
+        return (
+            0,  # step
+            _params_from_latent_params(latent_params),  # params
+            latent_params,  # latent params
+            opt.init(latents),  # opt state
+        )
+
+    def params_fn(state: WrappedOptaxState) -> PyTree:
+        """Returns the parameters for the given `state`."""
+        _, params, _, _ = state
+        return params
+
+    def update_fn(
+        *,
+        grad: PyTree,
+        value: jnp.ndarray,
+        params: PyTree,
+        state: WrappedOptaxState,
+    ) -> WrappedOptaxState:
+        """Updates the state."""
+        del params
+
+        step, params, latent_params, opt_state = state
+        metadata, latents = param_base.partition_density_metadata(latent_params)
+
+        def _params_from_latents(latents: PyTree) -> PyTree:
+            latent_params = param_base.combine_density_metadata(metadata, latents)
+            return _params_from_latent_params(latent_params)
+
+        def _constraint_loss_latents(latents: PyTree) -> jnp.ndarray:
+            latent_params = param_base.combine_density_metadata(metadata, latents)
+            return _constraint_loss(latent_params)
+
+        _, vjp_fn = jax.vjp(_params_from_latents, latents)
+        (latents_grad,) = vjp_fn(grad)
+
+        if not (
+            tree_util.tree_structure(latents_grad)
+            == tree_util.tree_structure(latents)  # type: ignore[operator]
+        ):
+            raise ValueError(
+                f"Tree structure of `latents_grad` was different than expected, got \n"
+                f"{tree_util.tree_structure(latents_grad)} but expected \n"
+                f"{tree_util.tree_structure(latents)}."
+            )
+
+        constraint_loss_grad = jax.grad(_constraint_loss_latents)(latents)
+        latents_grad = tree_util.tree_map(
+            lambda a, b: a + b, latents_grad, constraint_loss_grad
+        )
+
+        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)
+        params = _params_from_latent_params(latent_params)
+        return (step + 1, params, latent_params, opt_state)
+
+    # -------------------------------------------------------------------------
+    # Functions related to the density parameterization.
+    # -------------------------------------------------------------------------
+
+    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_latent_params(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 _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(
+            _leaf_apply_updates,
+            updates,
+            params,
+            is_leaf=_is_parameterized_density,
+        )
+
+    # -------------------------------------------------------------------------
+    # Functions related to the constraints to be minimized.
+    # -------------------------------------------------------------------------
+
+    def _constraint_loss(latent_params: PyTree) -> jnp.ndarray:
+        def _constraint_loss_leaf(
+            leaf: param_base.ParameterizedDensity2DArray,
+        ) -> jnp.ndarray:
+            constraints = density_parameterization.constraints(leaf)
+            constraints = tree_util.tree_map(
+                lambda x: jnp.sum(jnp.maximum(x, 0.0) ** 2),
+                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))
+
+    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, param_base.ParameterizedDensity2DArray)
+
+
+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,208 @@
+"""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, partition_utils, types
+
+Array = jnp.ndarray | onp.ndarray[Any, Any]
+PyTree = Any
+
+
+@dataclasses.dataclass
+class ParameterizedDensity2DArray:
+    """Stores latents and metadata for a parameterized density array."""
+
+    latents: "LatentsBase"
+    metadata: Optional["MetadataBase"]
+
+
+class LatentsBase:
+    """Base class for latents of a parameterized density array."""
+
+    pass
+
+
+class MetadataBase:
+    """Base class for metadata of a parameterized density array."""
+
+    pass
+
+
+tree_util.register_dataclass(
+    ParameterizedDensity2DArray,
+    data_fields=["latents", "metadata"],
+    meta_fields=[],
+)
+json_utils.register_custom_type(ParameterizedDensity2DArray)
+
+
+def partition_density_metadata(tree: PyTree) -> Tuple[PyTree, PyTree]:
+    """Splits a pytree with parameterized densities into metadata from latents."""
+    metadata, latents = partition_utils.partition(
+        tree,
+        select_fn=lambda x: isinstance(x, MetadataBase),
+        is_leaf=_is_metadata_or_none,
+    )
+    return metadata, latents
+
+
+def combine_density_metadata(metadata: PyTree, latents: PyTree) -> PyTree:
+    """Combines pytrees containing metadata and latents."""
+    return partition_utils.combine(metadata, latents, is_leaf=_is_metadata_or_none)
+
+
+def _is_metadata_or_none(leaf: Any) -> bool:
+    """Return `True` if `leaf` is `None` or density metadata."""
+    return leaf is None or isinstance(leaf, MetadataBase)
+
+
+@dataclasses.dataclass
+class Density2DParameterization:
+    """Stores `(from_density, to_density, constraints, update)` function triple."""
+
+    from_density: "FromDensityFn"
+    to_density: "ToDensityFn"
+    constraints: "ConstraintsFn"
+    update: "UpdateFn"
+
+
+class FromDensityFn(Protocol):
+    """Generate the latent representation of a density array."""
+
+    def __call__(self, density: types.Density2DArray) -> ParameterizedDensity2DArray:
+        ...
+
+
+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:
+        ...
+
+
+class UpdateFn(Protocol):
+    """Performs the required update of a parameterized density for the given step."""
+
+    def __call__(
+        self,
+        params: PyTree,
+        updates: PyTree,
+        value: jnp.ndarray,
+        step: int,
+    ) -> PyTree:
+        ...
+
+
+@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)
+
+    @classmethod
+    def from_density(self, density: types.Density2DArray) -> "Density2DMetadata":
+        density_metadata_dict = dataclasses.asdict(density)
+        del density_metadata_dict["array"]
+        return Density2DMetadata(**density_metadata_dict)
+
+
+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,138 @@
+"""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 FilterProjectParams(base.ParameterizedDensity2DArray):
+    """Stores parameters for the filter-project parameterization."""
+
+    latents: "FilterProjectLatents"
+    metadata: "FilterProjectMetadata"
+
+
+@dataclasses.dataclass
+class FilterProjectLatents(base.LatentsBase):
+    """Stores latent parameters for the filter-project parameterization.
+
+    Attributes:s
+        latent_density: The latent variable from which the density is obtained.
+    """
+
+    latent_density: types.Density2DArray
+
+
+@dataclasses.dataclass
+class FilterProjectMetadata(base.MetadataBase):
+    """Stores metadata for the filter-project parameterization.
+
+    Attributes:
+        beta: Determines the sharpness of the thresholding operation.
+    """
+
+    beta: float
+
+
+tree_util.register_dataclass(
+    FilterProjectParams,
+    data_fields=["latents", "metadata"],
+    meta_fields=[],
+)
+tree_util.register_dataclass(
+    FilterProjectLatents,
+    data_fields=["latent_density"],
+    meta_fields=[],
+)
+tree_util.register_dataclass(
+    FilterProjectMetadata,
+    data_fields=[],
+    meta_fields=["beta"],
+)
+json_utils.register_custom_type(FilterProjectParams)
+json_utils.register_custom_type(FilterProjectLatents)
+json_utils.register_custom_type(FilterProjectMetadata)
+
+
+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) -> FilterProjectParams:
+        """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)
+        latent_density = density = dataclasses.replace(density, array=latent_array)
+        return FilterProjectParams(
+            latents=FilterProjectLatents(latent_density=latent_density),
+            metadata=FilterProjectMetadata(beta=beta),
+        )
+
+    def to_density_fn(params: FilterProjectParams) -> types.Density2DArray:
+        """Return a density from the latent parameters."""
+        latent_density = params.latents.latent_density
+        beta = params.metadata.beta
+
+        transformed = types.symmetrize_density(latent_density)
+        transformed = transforms.density_gaussian_filter_and_tanh(transformed, 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: FilterProjectParams) -> jnp.ndarray:
+        """Computes constraints associated with the params."""
+        del params
+        return jnp.asarray(0.0)
+
+    def update_fn(
+        params: FilterProjectParams,
+        updates: FilterProjectParams,
+        value: jnp.ndarray,
+        step: int,
+    ) -> FilterProjectParams:
+        """Perform updates to `params` required for the given `step`."""
+        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,
+        from_density=from_density_fn,
+        constraints=constraints_fn,
+        update=update_fn,
+    )