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

invrs_opt/__init__.py

@@ -3,8 +3,19 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.5.2"
+__version__ = "v0.7.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 import parameterization as parameterization
+
+from invrs_opt.optimizers.lbfgsb import (
+    density_lbfgsb as density_lbfgsb,
+    lbfgsb as lbfgsb,
+    levelset_lbfgsb as levelset_lbfgsb,
+)
+
+from invrs_opt.optimizers.wrapped_optax import (
+    density_wrapped_optax as density_wrapped_optax,
+    levelset_wrapped_optax as levelset_wrapped_optax,
+    wrapped_optax as wrapped_optax,
+)

invrs_opt/experimental/client.py

@@ -4,15 +4,14 @@ 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
-
+from invrs_opt.optimizers import base
 
 PyTree = Any
 StateToken = str

invrs_opt/{base.py → optimizers/base.py}

@@ -4,8 +4,12 @@ Copyright (c) 2023 The INVRS-IO authors.
 """
 
 import dataclasses
+import inspect
 from typing import Any, Protocol
 
+import optax  # type: ignore[import-untyped]
+from totypes import json_utils
+
 PyTree = Any
 
 
@@ -44,3 +48,9 @@ class Optimizer:
     init: InitFn
     params: ParamsFn
     update: UpdateFn
+
+
+# Register all optax state types for serialization.
+for name, obj in inspect.getmembers(optax):
+    if name.endswith("State") and isinstance(obj, type):
+        json_utils.register_custom_type(obj)

invrs_opt/{lbfgsb → optimizers}/lbfgsb.py

@@ -5,19 +5,25 @@ Copyright (c) 2023 The INVRS-IO authors.
 
 import copy
 import dataclasses
-from typing import Any, Callable, Dict, Optional, Sequence, Tuple, Union
+from typing import Any, Dict, Optional, Sequence, Tuple, Union
 
 import jax
 import jax.numpy as jnp
 import numpy as onp
+import optax  # type: ignore[import-untyped]
 from jax import flatten_util, tree_util
 from scipy.optimize._lbfgsb_py import (  # type: ignore[import-untyped]
     _lbfgsb as scipy_lbfgsb,
 )
 from totypes import types
 
-from invrs_opt import base
-from invrs_opt.lbfgsb import transform
+from invrs_opt.optimizers import base
+from invrs_opt.parameterization import (
+    base as parameterization_base,
+    filter_project,
+    gaussian_levelset,
+    pixel,
+)
 
 NDArray = onp.ndarray[Any, Any]
 PyTree = Any
@@ -35,10 +41,10 @@ UPDATE_IPRINT = -1
 
 # Maximum value for the `maxcor` parameter in the L-BFGS-B scheme.
 MAXCOR_MAX_VALUE = 100
-MAXCOR_DEFAULT = 20
-LINE_SEARCH_MAX_STEPS_DEFAULT = 100
-FTOL_DEFAULT = 0.0
-GTOL_DEFAULT = 0.0
+DEFAULT_MAXCOR = 20
+DEFAULT_LINE_SEARCH_MAX_STEPS = 100
+DEFAULT_FTOL = 0.0
+DEFAULT_GTOL = 0.0
 
 # Maps bound scenarios to integers.
 BOUNDS_MAP: Dict[Tuple[bool, bool], int] = {
@@ -52,175 +58,225 @@ FORTRAN_INT = scipy_lbfgsb.types.intvar.dtype
 
 
 def lbfgsb(
-    maxcor: int = MAXCOR_DEFAULT,
-    line_search_max_steps: int = LINE_SEARCH_MAX_STEPS_DEFAULT,
-    ftol: float = FTOL_DEFAULT,
-    gtol: float = GTOL_DEFAULT,
+    *,
+    maxcor: int = DEFAULT_MAXCOR,
+    line_search_max_steps: int = DEFAULT_LINE_SEARCH_MAX_STEPS,
+    ftol: float = DEFAULT_FTOL,
+    gtol: float = DEFAULT_GTOL,
 ) -> base.Optimizer:
-    """Return an optimizer implementing the standard L-BFGS-B algorithm.
+    """Optimizer implementing the standard L-BFGS-B algorithm.
 
-    This optimizer wraps scipy's implementation of the algorithm, and provides
-    a jax-style API to the scheme. The optimizer works with custom types such
-    as the `BoundedArray` to constrain the optimization variable.
+    The standard L-BFGS-B algorithm uses the direct pixel parameterization for density
+    arrays, which simply enforces that values are between the declared upper and lower
+    bounds of the density.
 
-    Example usage is as follows:
-
-        def fn(x):
-            leaves_sum_sq = [jnp.sum(y)**2 for y in tree_util.tree_leaves(x)]
-            return jnp.sum(jnp.asarray(leaves_sum_sq))
-
-        x0 = {
-            "a": jnp.ones((3,)),
-            "b": BoundedArray(
-                value=-jnp.ones((2, 5)),
-                lower_bound=-5,
-                upper_bound=5,
-            ),
-        }
-        opt = lbfgsb(maxcor=20, line_search_max_steps=100)
-        state = opt.init(x0)
-        for _ in range(10):
-            x = opt.params(state)
-            value, grad = jax.value_and_grad(fn)(x)
-            state = opt.update(grad, value, state)
-
-    While the algorithm can work with pytrees of jax arrays, numpy arrays can
-    also be used. Thus, e.g. the optimizer can directly be used with autograd.
-
-    When the optimization has converged (according to `ftol` or `gtol` criteria), the
-    optimizer simply returns the parameters which obtained the converged result. The
+    When an optimization is determined to have converged (by `ftol` or `gtol` criteria)
+    the optimizer `params` function will simply return the optimal parameters. The
     convergence can be queried by `is_converged(state)`.
 
     Args:
-        maxcor: The maximum number of variable metric corrections used to define
-            the limited memory matrix, in the L-BFGS-B scheme.
+        maxcor: The maximum number of variable metric corrections used to define the
+            limited memory matrix, in the L-BFGS-B scheme.
         line_search_max_steps: The maximum number of steps in the line search.
-        ftol: Tolerance for stopping criteria based on function values. See scipy
-            documentation for details.
-        gtol: Tolerance for stopping criteria based on gradient.
+        ftol: Convergence criteria based on function values. See scipy documentation
+            for details.
+        gtol: Convergence criteria based on gradient.
 
     Returns:
-        The `base.Optimizer`.
+        The `Optimizer` implementing the L-BFGS-B optimizer.
     """
-    return transformed_lbfgsb(
+    return parameterized_lbfgsb(
+        density_parameterization=None,
+        penalty=0.0,
         maxcor=maxcor,
         line_search_max_steps=line_search_max_steps,
         ftol=ftol,
         gtol=gtol,
-        transform_fn=lambda x: x,
-        initialize_latent_fn=lambda x: x,
     )
 
 
 def density_lbfgsb(
+    *,
     beta: float,
-    maxcor: int = MAXCOR_DEFAULT,
-    line_search_max_steps: int = LINE_SEARCH_MAX_STEPS_DEFAULT,
-    ftol: float = FTOL_DEFAULT,
-    gtol: float = GTOL_DEFAULT,
+    maxcor: int = DEFAULT_MAXCOR,
+    line_search_max_steps: int = DEFAULT_LINE_SEARCH_MAX_STEPS,
+    ftol: float = DEFAULT_FTOL,
+    gtol: float = DEFAULT_GTOL,
 ) -> base.Optimizer:
-    """Return an L-BFGS-B optimizer with additional transforms for density arrays.
-
-    Parameters that are of type `DensityArray2D` are represented as latent parameters
-    that are transformed (in the case where lower and upper bounds are `(-1, 1)`) by,
+    """Optimizer using L-BFGS-B algorithm with filter-project density parameterization.
 
-        transformed = tanh(beta * conv(density.array, gaussian_kernel)) / tanh(beta)
+    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").
 
-    where the kernel has a full-width at half-maximum determined by the minimum width
-    and spacing parameters of the `DensityArray2D`. Where the bounds differ, the
-    density is scaled before the transform is applied, and then unscaled afterwards.
-
-    When the optimization has converged (according to `ftol` or `gtol` criteria), the
-    optimizer simply returns the parameters which obtained the converged result. The
+    When an optimization is determined to have converged (by `ftol` or `gtol` criteria)
+    the optimizer `params` function will simply return the optimal parameters. The
     convergence can be queried by `is_converged(state)`.
 
     Args:
-        beta: Determines the steepness of the thresholding.
-        maxcor: The maximum number of variable metric corrections used to define
-            the limited memory matrix, in the L-BFGS-B scheme.
+        beta: Determines the sharpness of the thresholding operation.
+        maxcor: The maximum number of variable metric corrections used to define the
+            limited memory matrix, in the L-BFGS-B scheme.
         line_search_max_steps: The maximum number of steps in the line search.
-        ftol: Tolerance for stopping criteria based on function values. See scipy
-            documentation for details.
-        gtol: Tolerance for stopping criteria based on gradient.
+        ftol: Convergence criteria based on function values. See scipy documentation
+            for details.
+        gtol: Convergence criteria based on gradient.
 
     Returns:
-        The `base.Optimizer`.
+        The `Optimizer` implementing the L-BFGS-B optimizer.
     """
+    return parameterized_lbfgsb(
+        density_parameterization=filter_project.filter_project(beta=beta),
+        penalty=0.0,
+        maxcor=maxcor,
+        line_search_max_steps=line_search_max_steps,
+        ftol=ftol,
+        gtol=gtol,
+    )
 
-    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 levelset_lbfgsb(
+    *,
+    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,
+    maxcor: int = DEFAULT_MAXCOR,
+    line_search_max_steps: int = DEFAULT_LINE_SEARCH_MAX_STEPS,
+    ftol: float = DEFAULT_FTOL,
+    gtol: float = DEFAULT_GTOL,
+) -> base.Optimizer:
+    """Optimizer using L-BFGS-B algorithm 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).
+
+    When an optimization is determined to have converged (by `ftol` or `gtol` criteria)
+    the optimizer `params` function will simply return the optimal parameters. The
+    convergence can be queried by `is_converged(state)`.
 
-    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_lbfgsb(
+    Args:
+        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.
+        maxcor: The maximum number of variable metric corrections used to define the
+            limited memory matrix, in the L-BFGS-B scheme.
+        line_search_max_steps: The maximum number of steps in the line search.
+        ftol: Convergence criteria based on function values. See scipy documentation
+            for details.
+        gtol: Convergence criteria based on gradient.
+
+    Returns:
+        The `Optimizer` implementing the L-BFGS-B optimizer.
+    """
+    return parameterized_lbfgsb(
+        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,
+        ),
+        penalty=penalty,
         maxcor=maxcor,
         line_search_max_steps=line_search_max_steps,
         ftol=ftol,
         gtol=gtol,
-        transform_fn=transform_fn,
-        initialize_latent_fn=initialize_latent_fn,
     )
 
 
-def transformed_lbfgsb(
-    maxcor: int,
-    line_search_max_steps: int,
-    ftol: float,
-    gtol: float,
-    transform_fn: Callable[[PyTree], PyTree],
-    initialize_latent_fn: Callable[[PyTree], PyTree],
-) -> base.Optimizer:
-    """Construct an latent parameter L-BFGS-B optimizer.
+# -----------------------------------------------------------------------------
+# Base parameterized L-BFGS-B optimizer.
+# -----------------------------------------------------------------------------
 
-    The optimized parameters are termed latent parameters, from which the
-    actual parameters returned by the optimizer are obtained using the
-    `transform_fn`. In the simple case where this is just `lambda x: x` (i.e.
-    the identity), this is equivalent to the standard L-BFGS-B algorithm.
 
-    When the optimization has converged (according to `ftol` or `gtol` criteria), the
-    optimizer simply returns the parameters which obtained the converged result. The
-    convergence can be queried by `is_converged(state)`.
+def parameterized_lbfgsb(
+    density_parameterization: Optional[parameterization_base.Density2DParameterization],
+    penalty: float,
+    maxcor: int = DEFAULT_MAXCOR,
+    line_search_max_steps: int = DEFAULT_LINE_SEARCH_MAX_STEPS,
+    ftol: float = DEFAULT_FTOL,
+    gtol: float = DEFAULT_GTOL,
+) -> base.Optimizer:
+    """Optimizer using L-BFGS-B optimizer with specified density parameterization.
+
+    This optimizer wraps scipy's implementation of the algorithm, and provides
+    a jax-style API to the scheme. The optimizer works with custom types such
+    as the `BoundedArray` to constrain the optimization variable.
 
     Args:
-        maxcor: The maximum number of variable metric corrections used to define
-            the limited memory matrix, in the L-BFGS-B scheme.
+        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.
+        maxcor: The maximum number of variable metric corrections used to define the
+            limited memory matrix, in the L-BFGS-B scheme.
         line_search_max_steps: The maximum number of steps in the line search.
-        ftol: Tolerance for stopping criteria based on function values. See scipy
-            documentation for details.
-        gtol: Tolerance for stopping criteria based on gradient.
-        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.
+        ftol: Convergence criteria based on function values. See scipy documentation
+            for details.
+        gtol: Convergence criteria based on gradient.
 
     Returns:
         The `base.Optimizer`.
@@ -237,33 +293,73 @@ def transformed_lbfgsb(
             f"{line_search_max_steps}"
         )
 
+    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) or density_parameterization is None:
+                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(latent_params: PyTree) -> PyTree:
+        def _leaf_params_from_latents(leaf: Any) -> Any:
+            if not _is_parameterized_density(leaf) or density_parameterization is None:
+                return leaf
+            return density_parameterization.to_density(leaf)
+
+        return tree_util.tree_map(
+            _leaf_params_from_latents,
+            latent_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) -> LbfgsbState:
         """Initializes the optimization state."""
 
-        def _init_pure(params: PyTree) -> Tuple[PyTree, JaxLbfgsbDict]:
-            lower_bound = types.extract_lower_bound(params)
-            upper_bound = types.extract_upper_bound(params)
+        def _init_state_pure(latent_params: PyTree) -> Tuple[PyTree, JaxLbfgsbDict]:
+            lower_bound = types.extract_lower_bound(latent_params)
+            upper_bound = types.extract_upper_bound(latent_params)
             scipy_lbfgsb_state = ScipyLbfgsbState.init(
-                x0=_to_numpy(params),
-                lower_bound=_bound_for_params(lower_bound, params),
-                upper_bound=_bound_for_params(upper_bound, params),
+                x0=_to_numpy(latent_params),
+                lower_bound=_bound_for_params(lower_bound, latent_params),
+                upper_bound=_bound_for_params(upper_bound, latent_params),
                 maxcor=maxcor,
                 line_search_max_steps=line_search_max_steps,
                 ftol=ftol,
                 gtol=gtol,
             )
-            latent_params = _to_pytree(scipy_lbfgsb_state.x, params)
+            latent_params = _to_pytree(scipy_lbfgsb_state.x, latent_params)
             return latent_params, scipy_lbfgsb_state.to_jax()
 
-        (
-            latent_params,
-            jax_lbfgsb_state,
-        ) = jax.pure_callback(
-            _init_pure,
-            _example_state(params, maxcor),
-            initialize_latent_fn(params),
+        latent_params = _init_latents(params)
+        latent_params, jax_lbfgsb_state = jax.pure_callback(
+            _init_state_pure, _example_state(latent_params, maxcor), latent_params
         )
-        return transform_fn(latent_params), latent_params, jax_lbfgsb_state
+        return _params_from_latents(latent_params), latent_params, jax_lbfgsb_state
 
     def params_fn(state: LbfgsbState) -> PyTree:
         """Returns the parameters for the given `state`."""
@@ -295,16 +391,35 @@ def transformed_lbfgsb(
             return flat_latent_params, scipy_lbfgsb_state.to_jax()
 
         _, latent_params, jax_lbfgsb_state = state
-        _, vjp_fn = jax.vjp(transform_fn, latent_params)
+        _, 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_value,
+            constraint_loss_grad,
+        ) = jax.value_and_grad(
+            _constraint_loss
+        )(latent_params)
+        value += constraint_loss_value
+        latent_grad = tree_util.tree_map(
+            lambda a, b: a + b, latent_grad, constraint_loss_grad
+        )
+
         flat_latent_grad, unflatten_fn = flatten_util.ravel_pytree(
             latent_grad
         )  # type: ignore[no-untyped-call]
 
-        (
-            flat_latent_params,
-            jax_lbfgsb_state,
-        ) = jax.pure_callback(
+        flat_latent_params, jax_lbfgsb_state = jax.pure_callback(
             _update_pure,
             (flat_latent_grad, jax_lbfgsb_state),
             flat_latent_grad,
@@ -312,7 +427,7 @@ def transformed_lbfgsb(
             jax_lbfgsb_state,
         )
         latent_params = unflatten_fn(flat_latent_params)
-        return transform_fn(latent_params), latent_params, jax_lbfgsb_state
+        return _params_from_latents(latent_params), latent_params, jax_lbfgsb_state
 
     return base.Optimizer(
         init=init_fn,
@@ -336,6 +451,31 @@ def _is_density(leaf: Any) -> Any:
     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)
+
+
 def _to_numpy(params: PyTree) -> NDArray:
     """Flattens a `params` pytree into a single rank-1 numpy array."""
     x, _ = flatten_util.ravel_pytree(params)  # type: ignore[no-untyped-call]