invrs-opt 0.5.1__tar.gz → 0.6.0__tar.gz

{invrs_opt-0.5.1/src/invrs_opt.egg-info → invrs_opt-0.6.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: invrs_opt
-Version: 0.5.1
+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: mypy; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
 
 # invrs-opt - Optimization algorithms for inverse design
-`v0.5.1`
+`v0.6.0`
 
 ## Overview
 

{invrs_opt-0.5.1 → invrs_opt-0.6.0}/README.md

@@ -1,5 +1,5 @@
 # invrs-opt - Optimization algorithms for inverse design
-`v0.5.1`
+`v0.6.0`
 
 ## Overview
 

{invrs_opt-0.5.1 → invrs_opt-0.6.0}/pyproject.toml

@@ -1,7 +1,7 @@
 [project]
 
 name = "invrs_opt"
-version = "v0.5.1"
+version = "v0.6.0"
 description = "Algorithms for inverse design"
 keywords = ["topology", "optimization", "jax", "inverse design"]
 readme = "README.md"
@@ -20,6 +20,7 @@ dependencies = [
     "jaxlib",
     "numpy",
     "requests",
+    "optax",
     "scipy",
     "totypes",
     "types-requests",

{invrs_opt-0.5.1 → invrs_opt-0.6.0}/src/invrs_opt/__init__.py

@@ -3,8 +3,12 @@
 Copyright (c) 2023 The INVRS-IO authors.
 """
 
-__version__ = "v0.5.1"
+__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-0.5.1 → invrs_opt-0.6.0}/src/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-0.5.1 → invrs_opt-0.6.0}/src/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-0.5.1 → invrs_opt-0.6.0}/src/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
@@ -258,7 +257,7 @@ def transformed_lbfgsb(
         (
             latent_params,
             jax_lbfgsb_state,
-        ) = jax.pure_callback(  # type: ignore[attr-defined]
+        ) = jax.pure_callback(
             _init_pure,
             _example_state(params, maxcor),
             initialize_latent_fn(params),
@@ -304,7 +303,7 @@ def transformed_lbfgsb(
         (
             flat_latent_params,
             jax_lbfgsb_state,
-        ) = jax.pure_callback(  # type: ignore[attr-defined]
+        ) = jax.pure_callback(
             _update_pure,
             (flat_latent_grad, jax_lbfgsb_state),
             flat_latent_grad,
@@ -542,19 +541,19 @@ class ScipyLbfgsbState:
         """Converts a dictionary of jax arrays to a `ScipyLbfgsbState`."""
         state_dict = copy.deepcopy(state_dict)
         return ScipyLbfgsbState(
-            x=onp.asarray(state_dict["x"], dtype=onp.float64),
+            x=onp.array(state_dict["x"], dtype=onp.float64),
             converged=onp.asarray(state_dict["converged"], dtype=bool),
             _maxcor=int(state_dict["_maxcor"]),
             _line_search_max_steps=int(state_dict["_line_search_max_steps"]),
             _ftol=onp.asarray(state_dict["_ftol"], dtype=onp.float64),
             _gtol=onp.asarray(state_dict["_gtol"], dtype=onp.float64),
-            _wa=onp.asarray(state_dict["_wa"], onp.float64),
-            _iwa=onp.asarray(state_dict["_iwa"], dtype=FORTRAN_INT),
+            _wa=onp.array(state_dict["_wa"], onp.float64),
+            _iwa=onp.array(state_dict["_iwa"], dtype=FORTRAN_INT),
             _task=_s60_str_from_array(state_dict["_task"]),
             _csave=_s60_str_from_array(state_dict["_csave"]),
-            _lsave=onp.asarray(state_dict["_lsave"], dtype=FORTRAN_INT),
-            _isave=onp.asarray(state_dict["_isave"], dtype=FORTRAN_INT),
-            _dsave=onp.asarray(state_dict["_dsave"], dtype=onp.float64),
+            _lsave=onp.array(state_dict["_lsave"], dtype=FORTRAN_INT),
+            _isave=onp.array(state_dict["_isave"], dtype=FORTRAN_INT),
+            _dsave=onp.array(state_dict["_dsave"], dtype=onp.float64),
             _lower_bound=onp.asarray(state_dict["_lower_bound"], dtype=onp.float64),
             _upper_bound=onp.asarray(state_dict["_upper_bound"], dtype=onp.float64),
             _bound_type=onp.asarray(state_dict["_bound_type"], dtype=int),

invrs_opt-0.6.0/src/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.1 → invrs_opt-0.6.0/src/invrs_opt.egg-info}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: invrs_opt
-Version: 0.5.1
+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: mypy; extra == "dev"
 Requires-Dist: pre-commit; extra == "dev"
 
 # invrs-opt - Optimization algorithms for inverse design
-`v0.5.1`
+`v0.6.0`
 
 ## Overview
 

{invrs_opt-0.5.1 → invrs_opt-0.6.0}/src/invrs_opt.egg-info/SOURCES.txt

@@ -4,6 +4,7 @@ pyproject.toml
 src/invrs_opt/__init__.py
 src/invrs_opt/base.py
 src/invrs_opt/py.typed
+src/invrs_opt/transform.py
 src/invrs_opt.egg-info/PKG-INFO
 src/invrs_opt.egg-info/SOURCES.txt
 src/invrs_opt.egg-info/dependency_links.txt
@@ -14,5 +15,7 @@ src/invrs_opt/experimental/client.py
 src/invrs_opt/experimental/labels.py
 src/invrs_opt/lbfgsb/__init__.py
 src/invrs_opt/lbfgsb/lbfgsb.py
-src/invrs_opt/lbfgsb/transform.py
-tests/test_algos.py
+src/invrs_opt/wrapped_optax/__init__.py
+src/invrs_opt/wrapped_optax/wrapped_optax.py
+tests/test_algos.py
+tests/test_transform.py

{invrs_opt-0.5.1 → invrs_opt-0.6.0}/src/invrs_opt.egg-info/requires.txt

@@ -2,6 +2,7 @@ jax
 jaxlib
 numpy
 requests
+optax
 scipy
 totypes
 types-requests

{invrs_opt-0.5.1 → invrs_opt-0.6.0}/tests/test_algos.py

@@ -9,6 +9,7 @@ import unittest
 import jax
 import jax.numpy as jnp
 import numpy as onp
+import optax
 import parameterized
 from totypes import json_utils, symmetry, types
 
@@ -21,6 +22,8 @@ jax.config.update("jax_enable_x64", True)
 OPTIMIZERS = [
     invrs_opt.lbfgsb(maxcor=20, line_search_max_steps=100),
     invrs_opt.density_lbfgsb(maxcor=20, line_search_max_steps=100, beta=2.0),
+    invrs_opt.wrapped_optax(optax.adam(1e-2)),
+    invrs_opt.density_wrapped_optax(optax.adam(1e-2), beta=2.0),
 ]
 
 # Various parameter combinations tested in this module.

invrs_opt-0.6.0/tests/test_transform.py

@@ -0,0 +1,362 @@
+"""Tests for `transforms`.
+
+Copyright (c) 2023 The INVRS-IO authors.
+"""
+
+import dataclasses
+import unittest
+
+import jax
+import jax.numpy as jnp
+import numpy as onp
+from parameterized import parameterized
+from totypes import types
+
+from invrs_opt import transform
+
+
+class GaussianFilterTest(unittest.TestCase):
+    @parameterized.expand([[1, 5], [3, 3], [5, 1]])
+    def test_transformed_matches_expected(self, minimum_width, minimum_spacing):
+        array = onp.zeros((9, 9))
+        array[4, 4] = 9
+        density = types.Density2DArray(
+            array=array,
+            lower_bound=0,
+            upper_bound=1,
+            minimum_width=minimum_width,
+            minimum_spacing=minimum_spacing,
+        )
+        beta = 1
+        transformed = transform.density_gaussian_filter_and_tanh(density, beta=beta)
+        expected = onp.asarray(
+            [
+                [0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12],
+                [0.12, 0.12, 0.13, 0.14, 0.14, 0.14, 0.13, 0.12, 0.12],
+                [0.12, 0.13, 0.15, 0.22, 0.27, 0.22, 0.15, 0.13, 0.12],
+                [0.12, 0.14, 0.22, 0.48, 0.64, 0.48, 0.22, 0.14, 0.12],
+                [0.12, 0.14, 0.27, 0.64, 0.82, 0.64, 0.27, 0.14, 0.12],
+                [0.12, 0.14, 0.22, 0.48, 0.64, 0.48, 0.22, 0.14, 0.12],
+                [0.12, 0.13, 0.15, 0.22, 0.27, 0.22, 0.15, 0.13, 0.12],
+                [0.12, 0.12, 0.13, 0.14, 0.14, 0.14, 0.13, 0.12, 0.12],
+                [0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12, 0.12],
+            ]
+        )
+        onp.testing.assert_allclose(transformed.array, expected, rtol=0.05)
+
+    @parameterized.expand([[1, 1], [3, 1], [5, 1], [10, 1], [10, 0.5], [10, 2]])
+    def test_ones_density_yields_tanh_beta(self, length_scale, upper_bound):
+        array = onp.ones((20, 20)) * upper_bound
+        density = types.Density2DArray(
+            array=array,
+            lower_bound=0,
+            upper_bound=upper_bound,
+            minimum_width=length_scale,
+            minimum_spacing=length_scale,
+        )
+        beta = 1
+        transformed = transform.density_gaussian_filter_and_tanh(density, beta=beta)
+        onp.testing.assert_allclose(
+            transformed.array,
+            (1 + onp.tanh(beta)) * 0.5 * upper_bound,
+            rtol=0.01,
+        )
+
+    def test_batch_matches_single(self):
+        beta = 4
+        density = types.Density2DArray(
+            array=onp.arange(600).reshape((6, 10, 10)),
+            minimum_width=5,
+            minimum_spacing=5,
+        )
+        transformed = transform.density_gaussian_filter_and_tanh(density, beta=beta)
+        for i in range(6):
+            transformed_single = transform.density_gaussian_filter_and_tanh(
+                density=dataclasses.replace(
+                    density,
+                    array=density.array[i, :, :],
+                ),
+                beta=beta,
+            )
+            onp.testing.assert_allclose(
+                transformed.array[i, :, :], transformed_single.array
+            )
+
+    def test_periodic(self):
+        beta = 100
+        array = onp.zeros((5, 5))
+        array[0, 0] = 9
+
+        # No periodicity.
+        density = types.Density2DArray(
+            array,
+            minimum_spacing=3,
+            minimum_width=3,
+            periodic=(False, False),
+            lower_bound=0,
+            upper_bound=1,
+        )
+        transformed = transform.density_gaussian_filter_and_tanh(density, beta=beta)
+        expected = onp.asarray(
+            [
+                [1, 1, 1, 0, 0],
+                [1, 1, 0, 0, 0],
+                [1, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+            ]
+        )
+        onp.testing.assert_allclose(transformed.array, expected, atol=0.01)
+
+        # Periodic along the first axis.
+        density = dataclasses.replace(density, periodic=(True, False))
+        transformed = transform.density_gaussian_filter_and_tanh(density, beta=beta)
+        expected = onp.asarray(
+            [
+                [1, 1, 0, 0, 0],
+                [1, 1, 0, 0, 0],
+                [1, 0, 0, 0, 0],
+                [1, 0, 0, 0, 0],
+                [1, 1, 0, 0, 0],
+            ]
+        )
+        onp.testing.assert_allclose(transformed.array, expected, atol=0.01)
+
+        # Periodic along the second axis.
+        density = dataclasses.replace(density, periodic=(False, True))
+        transformed = transform.density_gaussian_filter_and_tanh(density, beta=beta)
+        expected = onp.asarray(
+            [
+                [1, 1, 1, 1, 1],
+                [1, 1, 0, 0, 1],
+                [0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+            ]
+        )
+        onp.testing.assert_allclose(transformed.array, expected, atol=0.01)
+
+        # Periodic along both axes.
+        density = dataclasses.replace(density, periodic=(True, True))
+        transformed = transform.density_gaussian_filter_and_tanh(density, beta=beta)
+        expected = onp.asarray(
+            [
+                [1, 1, 0, 0, 1],
+                [1, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+                [0, 0, 0, 0, 0],
+                [1, 0, 0, 0, 0],
+            ]
+        )
+        onp.testing.assert_allclose(transformed.array, expected, atol=0.01)
+
+
+class RescaleTest(unittest.TestCase):
+    @parameterized.expand([(-1.0, 1.0), (2.0, 3.0), (-0.5, -0.1)])
+    def test_normalized_array_from_density(self, lower_bound, upper_bound):
+        density = types.Density2DArray(
+            array=jnp.linspace(lower_bound, upper_bound).reshape((10, 5)),
+            lower_bound=lower_bound,
+            upper_bound=upper_bound,
+            fixed_solid=None,
+            fixed_void=None,
+            minimum_width=1,
+            minimum_spacing=1,
+        )
+        # Compute `array`, which should now have values between `-1` and `1`.
+        array = transform.normalized_array_from_density(density)
+        expected = jnp.linspace(-1.0, 1.0).reshape((10, 5))
+        onp.testing.assert_allclose(array, expected, rtol=1e-5)
+
+    @parameterized.expand([(-1.0, 1.0), (2.0, 3.0), (-0.5, -0.1)])
+    def test_rescale_array_for_density(self, lower_bound, upper_bound):
+        dummy_density = types.Density2DArray(
+            array=jnp.ones((2, 2)),
+            lower_bound=lower_bound,
+            upper_bound=upper_bound,
+            fixed_solid=None,
+            fixed_void=None,
+            minimum_width=1,
+            minimum_spacing=1,
+        )
+        array = jnp.linspace(-1.0, 1.0).reshape((10, 5))
+        rescaled = transform.rescale_array_for_density(array, dummy_density)
+        expected = jnp.linspace(lower_bound, upper_bound).reshape((10, 5))
+        onp.testing.assert_allclose(rescaled, expected, rtol=1e-5)
+
+
+class FixedPixelTest(unittest.TestCase):
+    @parameterized.expand(
+        [
+            [
+                jnp.asarray([[1, 0, 0, 0, 0]], dtype=bool),
+                jnp.asarray([[0, 0, 1, 1, 0]], dtype=bool),
+                jnp.asarray([[3.0, 0.0, -0.5, -0.5, 0.0]]),
+            ],
+            [
+                None,
+                jnp.asarray([[0, 0, 1, 1, 0]], dtype=bool),
+                jnp.asarray([[0.0, 0.0, -0.5, -0.5, 0.0]]),
+            ],
+            [
+                None,
+                None,
+                jnp.asarray([[0.0, 0.0, 0.0, 0.0, 0.0]]),
+            ],
+        ]
+    )
+    def test_apply_fixed_pixels(self, fixed_solid, fixed_void, expected):
+        density = types.Density2DArray(
+            array=jnp.zeros((1, 5)),
+            fixed_solid=fixed_solid,
+            fixed_void=fixed_void,
+            lower_bound=-0.5,
+            upper_bound=3.0,
+        )
+        density = transform.apply_fixed_pixels(density)
+        onp.testing.assert_array_equal(density.array, expected)
+
+
+class Pad2DTest(unittest.TestCase):
+    def test_pad2d_edge(self):
+        array = jnp.asarray(
+            [
+                [0, 1, 2, 3, 4],
+                [5, 6, 7, 8, 9],
+                [10, 11, 12, 13, 14],
+            ]
+        )
+        expected = jnp.asarray(
+            [
+                [0, 0, 1, 2, 3, 4, 4],
+                [0, 0, 1, 2, 3, 4, 4],
+                [5, 5, 6, 7, 8, 9, 9],
+                [10, 10, 11, 12, 13, 14, 14],
+                [10, 10, 11, 12, 13, 14, 14],
+            ]
+        )
+        padded = transform.pad2d(array, ((1, 1), (1, 1)), "edge")
+        padded_both_specified = transform.pad2d(
+            array, ((1, 1), (1, 1)), ("edge", "edge")
+        )
+        onp.testing.assert_array_equal(padded, expected)
+        onp.testing.assert_array_equal(padded_both_specified, expected)
+
+    def test_pad2d_wrap(self):
+        array = jnp.asarray(
+            [
+                [0, 1, 2, 3, 4],
+                [5, 6, 7, 8, 9],
+                [10, 11, 12, 13, 14],
+            ]
+        )
+        expected = jnp.asarray(
+            [
+                [14, 10, 11, 12, 13, 14, 10],
+                [4, 0, 1, 2, 3, 4, 0],
+                [9, 5, 6, 7, 8, 9, 5],
+                [14, 10, 11, 12, 13, 14, 10],
+                [4, 0, 1, 2, 3, 4, 0],
+            ]
+        )
+        padded = transform.pad2d(array, ((1, 1), (1, 1)), "wrap")
+        padded_both_specified = transform.pad2d(
+            array, ((1, 1), (1, 1)), ("wrap", "wrap")
+        )
+        onp.testing.assert_array_equal(padded, expected)
+        onp.testing.assert_array_equal(padded_both_specified, expected)
+
+    def test_pad2d_wrap_edge(self):
+        array = jnp.asarray(
+            [
+                [0, 1, 2, 3, 4],
+                [5, 6, 7, 8, 9],
+                [10, 11, 12, 13, 14],
+            ]
+        )
+        expected = jnp.asarray(
+            [
+                [10, 10, 11, 12, 13, 14, 14],
+                [0, 0, 1, 2, 3, 4, 4],
+                [5, 5, 6, 7, 8, 9, 9],
+                [10, 10, 11, 12, 13, 14, 14],
+                [0, 0, 1, 2, 3, 4, 4],
+            ]
+        )
+        padded = transform.pad2d(array, ((1, 1), (1, 1)), ("wrap", "edge"))
+        onp.testing.assert_array_equal(padded, expected)
+
+    def test_pad2d_edge_wrap(self):
+        array = jnp.asarray(
+            [
+                [0, 1, 2, 3, 4],
+                [5, 6, 7, 8, 9],
+                [10, 11, 12, 13, 14],
+            ]
+        )
+        expected = jnp.asarray(
+            [
+                [4, 0, 1, 2, 3, 4, 0],
+                [4, 0, 1, 2, 3, 4, 0],
+                [9, 5, 6, 7, 8, 9, 5],
+                [14, 10, 11, 12, 13, 14, 10],
+                [14, 10, 11, 12, 13, 14, 10],
+            ]
+        )
+        padded = transform.pad2d(array, ((1, 1), (1, 1)), ("edge", "wrap"))
+        onp.testing.assert_array_equal(padded, expected)
+
+    def test_pad2d_batch_dims(self):
+        array = jnp.arange(300).reshape((2, 1, 5, 10, 3))
+        pad_width = ((3, 4), (1, 2))
+        padded = transform.pad2d(array, pad_width, "wrap")
+        self.assertSequenceEqual(padded.shape, (2, 1, 5, 17, 6))
+        for i in range(array.shape[0]):
+            for j in range(array.shape[1]):
+                for k in range(array.shape[2]):
+                    onp.testing.assert_array_equal(
+                        padded[i, j, k, :, :],
+                        transform.pad2d(array[i, j, k, :, :], pad_width, "wrap"),
+                    )
+
+
+class PadWidthForKernelShapeTest(unittest.TestCase):
+    @parameterized.expand([[(2, 3)], [(4, 4)], [(5, 6)]])
+    def test_pad_width(self, kernel_shape):
+        # Checks that when padded by the computed amount, a convolution
+        # valid padding returns an array with the original shape.
+        kernel = jnp.ones(kernel_shape)
+        array = jnp.arange(77).reshape((7, 11)).astype(float)
+        pad_width = transform.pad_width_for_kernel_shape(kernel_shape)
+        padded = jnp.pad(array, pad_width)
+        y = jax.lax.conv_general_dilated(
+            lhs=padded[jnp.newaxis, jnp.newaxis, :, :],  # HCHW
+            rhs=kernel[jnp.newaxis, jnp.newaxis, :, :],  # OIHW
+            padding="VALID",
+            dimension_numbers=("NCHW", "OIHW", "NCHW"),
+            window_strides=(1, 1),
+        )
+        self.assertEqual(y.shape, (1, 1) + array.shape)
+
+    @parameterized.expand([[(2, 3)], [(4, 4)], [(5, 6)]])
+    def test_pad_width_offsets(self, kernel_shape):
+        kernel = onp.zeros(kernel_shape)
+        kernel[kernel_shape[0] // 2, kernel_shape[1] // 2] = 1
+        array = jnp.arange(77).reshape((7, 11)).astype(float)
+        pad_width = transform.pad_width_for_kernel_shape(kernel_shape)
+        padded = jnp.pad(array, pad_width)
+        y = jax.lax.conv_general_dilated(
+            lhs=padded[jnp.newaxis, jnp.newaxis, :, :],  # HCHW
+            rhs=kernel[jnp.newaxis, jnp.newaxis, :, :],  # OIHW
+            padding="VALID",
+            dimension_numbers=("NCHW", "OIHW", "NCHW"),
+            window_strides=(1, 1),
+        )
+        onp.testing.assert_array_equal(y[0, 0, ...], array)
+
+
+class GaussianKernelTest(unittest.TestCase):
+    @parameterized.expand([[2], [3], [4], [5]])
+    def test_gaussian_peak_on_gridpoint(self, fwhm_size_multiple):
+        kernel = transform._gaussian_kernel(1.0, fwhm_size_multiple)
+        self.assertEqual(kernel[kernel.shape[0] // 2, kernel.shape[1] // 2], 1.0)