phaserEM 0.1__py3-none-any.whl

phaser/utils/_cuda_kernels.py

@@ -0,0 +1,213 @@
+
+import functools
+import typing as t
+
+import cupy  # pyright: ignore[reportMissingImports]
+import numpy
+
+# grid
+#  block
+#    thread
+# blocks contain threads, a grid contains blocks
+# max 1024 threads/block, maximum per axis (1024, 1024, 64)
+# max grid size per axis ~2^32
+# x is the fast axis (warp axis)
+
+# TODO:
+#   optimize block size, optimize warps?
+
+def set_cutouts(obj: cupy.ndarray, cutouts: cupy.ndarray, start_idxs: cupy.ndarray):
+    assert obj.ndim >= 2
+
+    kernel = _get_cutout_kernel(obj.dtype, 'set')
+    if kernel is NotImplemented:
+        raise NotImplementedError()
+
+    # output array must be the correct dtype
+    assert obj.dtype == cutouts.dtype
+    # output array must be contiguous
+    assert obj.flags.c_contiguous
+
+    # cutouts must be the correct shape
+    assert cutouts.shape[:-2] == (*start_idxs.shape[:-1], *obj.shape[:-2])
+    assert obj.shape[:-2] == cutouts.shape[-obj.ndim:-2]
+    assert start_idxs.shape[-1] == 2
+
+    n_cutouts = numpy.prod((1, *start_idxs.shape[:-1]))
+    n_objects = numpy.prod((1, *obj.shape[:-2]))
+
+    block = (128, 8, 1)
+    shape = (cutouts.shape[-1], cutouts.shape[-2], n_cutouts * n_objects)
+    grid = tuple((s + b - 1) // b for (s, b) in zip(shape, block))
+
+    start_idxs = cupy.ascontiguousarray(start_idxs.astype(cupy.uint64))
+
+    default = numpy.bool_(0).astype(obj.dtype)
+    args = (obj, cupy.ascontiguousarray(cutouts), start_idxs, n_cutouts, n_objects, *obj.shape[-2:], *cutouts.shape[-2:], default)
+    kernel(grid, block, args)
+    return cutouts
+
+
+def get_cutouts(obj: cupy.ndarray, start_idxs: cupy.ndarray, cutout_shape: t.Tuple[int, int], default: cupy.generic) -> cupy.ndarray:
+    assert obj.ndim >= 2
+
+    kernel = _get_cutout_kernel(obj.dtype, 'get')
+    if kernel is NotImplemented:
+        raise NotImplementedError()
+
+    # check start_idxs shape
+    assert start_idxs.shape[-1] == 2
+    cutouts = cupy.empty((*start_idxs.shape[:-1], *obj.shape[:-2], *cutout_shape), dtype=obj.dtype)
+    # output array must be contiguous
+    assert cutouts.flags.c_contiguous
+
+    n_cutouts = numpy.prod((1, *start_idxs.shape[:-1]))
+    n_objects = numpy.prod((1, *obj.shape[:-2]))
+
+    block = (128, 8, 1)
+    shape = (cutouts.shape[-1], cutouts.shape[-2], n_cutouts * n_objects)
+    grid = tuple((s + b - 1) // b for (s, b) in zip(shape, block))
+
+    start_idxs = cupy.ascontiguousarray(start_idxs.astype(cupy.uint64))
+
+    args = (cupy.ascontiguousarray(obj), cutouts, start_idxs, n_cutouts, n_objects, *obj.shape[-2:], *cutouts.shape[-2:], default)
+    kernel(grid, block, args)
+    return cutouts
+
+
+def add_cutouts(obj: cupy.ndarray, cutouts: cupy.ndarray, start_idxs: cupy.ndarray):
+    assert obj.ndim >= 2
+
+    kernel = _get_cutout_kernel(obj.dtype, 'add')
+    if kernel is NotImplemented:
+        raise NotImplementedError()
+
+    # output array must be the correct dtype
+    assert obj.dtype == cutouts.dtype
+    # output array must be contiguous
+    assert obj.flags.c_contiguous
+
+    # cutouts must be the correct shape
+    assert cutouts.shape[:-2] == (*start_idxs.shape[:-1], *obj.shape[:-2])
+    assert obj.shape[:-2] == cutouts.shape[-obj.ndim:-2]
+    assert start_idxs.shape[-1] == 2
+
+    n_cutouts = numpy.prod((1, *start_idxs.shape[:-1]))
+    n_objects = numpy.prod((1, *obj.shape[:-2]))
+
+    block = (128, 8, 1)
+    shape = (cutouts.shape[-1], cutouts.shape[-2], n_cutouts * n_objects)
+    grid = tuple((s + b - 1) // b for (s, b) in zip(shape, block))
+
+    start_idxs = cupy.ascontiguousarray(start_idxs.astype(cupy.uint64))
+
+    default = numpy.bool_(0).astype(obj.dtype)
+    args = (obj, cupy.ascontiguousarray(cutouts), start_idxs, n_cutouts, n_objects, *obj.shape[-2:], *cutouts.shape[-2:], default)
+    kernel(grid, block, args)
+    return cutouts
+
+
+_CUTOUT_OPERATIONS = {
+    'get': "cutouts[cutout_idx] = obj[obj_idx];",
+    'add': "atomicAdd(&obj[obj_idx], cutouts[cutout_idx]);",
+    # TODO does this need to be atomic?
+    'set': "obj[obj_idx] = cutouts[cutout_idx];",
+}
+
+_CUTOUT_CONST = {
+    'get': (True, False),
+    'set': (False, True),
+    'add': (False, True),
+}
+
+# cupy/_core/_scalar.pyx
+_DTYPE_TO_KERNEL_TYPE: t.Dict[t.Type[numpy.generic], str] = {
+    numpy.bool_: "bool",
+    numpy.int8: "signed char",
+    numpy.uint8: "unsigned char",
+    numpy.int16: "signed short",
+    numpy.uint16: "unsigned short",
+    numpy.int32: "signed int",
+    numpy.uint32: "unsigned int",
+    numpy.int64: "signed long long",
+    numpy.uint64: "unsigned long long",
+    numpy.float16: "float16",
+    numpy.float32: "float",
+    numpy.float64: "double",
+    numpy.complex64: "complex<float>",
+    numpy.complex128: "complex<double>",
+}
+
+@functools.cache
+def _get_cutout_kernel(dtype: numpy.dtype, operation: str) -> cupy.RawKernel:
+    try:
+        ty = _DTYPE_TO_KERNEL_TYPE[dtype.type]
+        op = _CUTOUT_OPERATIONS[operation]
+        const_s = tuple("const " if c else "" for c in _CUTOUT_CONST[operation])
+    except KeyError:
+        return NotImplemented
+
+    if operation == 'add':
+        if ty not in {
+                "float16", "float", "double", "complex<float>", "complex<double>",
+                "signed int", "unsigned int", "unsigned long long",
+            }:
+            return NotImplemented
+        # double atomicAdd requires compute capability >=6.x
+        device = cupy.cuda.Device()
+        if ty in ('complex<double>', 'double') and int(device.compute_capability[0]) < 6:
+            return NotImplemented
+
+    kernel_name = f"{operation}_cutouts_{dtype.type.__name__}"
+
+    kernel = cupy.RawKernel(rf"""
+#include <cupy/complex.cuh>
+
+#define T {ty}
+
+// support atomicAdd for complex datatypes
+// we can just add real and complex parts separately
+template<typename U>
+__device__ inline void atomicAdd(complex<U> *address, complex<U> val) {{
+    atomicAdd(&reinterpret_cast<U*>(address)[0], val.real());
+    atomicAdd(&reinterpret_cast<U*>(address)[1], val.imag());
+}}
+
+extern "C" __global__
+void {kernel_name}({const_s[0]}T *obj, {const_s[1]}T *cutouts, const long long *start_idxs, long long n_cutouts, long long n_objects,
+                         long long obj_shape_y, long long obj_shape_x, long long cutout_shape_y, long long cutout_shape_x, T default_value) {{
+    size_t          j = blockIdx.x * blockDim.x + threadIdx.x;
+    size_t          i = blockIdx.y * blockDim.y + threadIdx.y;
+    size_t cutout_num = blockIdx.z * blockDim.z + threadIdx.z;
+
+    size_t     object_num = cutout_num % n_objects;
+    size_t pos_cutout_num = cutout_num / n_objects;
+
+    
+
+    // check if our thread is outside the requested cutout
+    if (pos_cutout_num >= n_cutouts || object_num >= n_objects || i >= cutout_shape_y || j >= cutout_shape_x) {{
+        return;
+    }}
+
+    // compute cutout index
+    size_t cutout_idx = cutout_shape_y * cutout_shape_x * cutout_num + cutout_shape_x * i + j;
+
+    // offset to get object location
+    long long obj_i = i + start_idxs[2*pos_cutout_num];
+    long long obj_j = j + start_idxs[2*pos_cutout_num + 1];
+
+    // check that it's in bounds, get default value if not
+    if (obj_i < 0 || obj_i >= obj_shape_y || obj_j < 0 || obj_j >= obj_shape_x) {{
+        {'cutouts[cutout_idx] = default_value;' if op == 'get' else ''}
+        return;
+    }}
+
+    // compute object index
+    size_t obj_idx = obj_shape_y * obj_shape_x * object_num + obj_shape_x * obj_i + obj_j;
+    // and perform our operation
+    {op}
+}}
+""", kernel_name)
+    kernel.compile()
+    return kernel

phaser/utils/_jax_kernels.py

@@ -0,0 +1,98 @@
+
+from functools import partial
+import typing as t
+
+from numpy.typing import ArrayLike
+import jax  # pyright: ignore[reportMissingImports]
+import jax.numpy as jnp    # pyright: ignore[reportMissingImports]
+
+
+def to_nd(arr: jax.Array, n: int) -> jax.Array:
+    if arr.ndim > n:
+        arr = arr.reshape(-1, *arr.shape[arr.ndim - n + 1:])
+    elif arr.ndim < n:
+        arr = jax.lax.expand_dims(arr, [0] * (n - arr.ndim))
+
+    return arr
+
+to_2d = partial(to_nd, n=2)
+to_3d = partial(to_nd, n=3)
+
+
+def _flatten_ndim(arr: jax.Array, n: int) -> jax.Array:
+    return arr.reshape(-1, *arr.shape[n+1:])
+
+
+@jax.jit
+def set_cutouts(obj: jax.Array, cutouts: jax.Array, start_idxs: jax.Array) -> jax.Array:
+    def inner(obj: jax.Array, v: t.Tuple[jax.Array, jax.Array]) -> t.Tuple[jax.Array, None]:
+        (cutouts, start_idxs) = v
+        obj = jax.vmap(lambda slice, cutouts: jax.lax.dynamic_update_slice(slice, cutouts, start_idxs))(obj, to_3d(cutouts))
+        return (obj, None)
+
+    start_idxs = jnp.atleast_2d(start_idxs)
+    (cutouts, start_idxs) = map(lambda a: _flatten_ndim(a, start_idxs.ndim - 2), (cutouts, start_idxs))
+    return jax.lax.scan(inner, to_3d(obj), (cutouts, start_idxs))[0].reshape(obj.shape)
+
+
+@jax.jit
+def add_cutouts(obj: jax.Array, cutouts: jax.Array, start_idxs: jax.Array) -> jax.Array:
+    zero_obj = jax.numpy.zeros_like(obj, shape=obj.shape[-2:])
+
+    def inner(obj: jax.Array, v: t.Tuple[jax.Array, jax.Array]) -> t.Tuple[jax.Array, None]:
+        (cutouts, start_idxs) = v
+        obj += jax.vmap(lambda cutouts: jax.lax.dynamic_update_slice(zero_obj, cutouts, start_idxs))(to_3d(cutouts))
+        return (obj, None)
+
+    start_idxs = jnp.atleast_2d(start_idxs)
+    (cutouts, start_idxs) = map(lambda a: _flatten_ndim(a, start_idxs.ndim - 2), (cutouts, start_idxs))
+    return jax.lax.scan(inner, to_3d(obj), (cutouts, start_idxs))[0].reshape(obj.shape)
+
+
+@partial(jax.jit, static_argnums=2)
+def get_cutouts(obj: jax.Array, start_idxs: jax.Array, cutout_shape: t.Tuple[int, int]) -> jax.Array:
+    return jax.vmap(jax.vmap(lambda start_idx, obj: jax.lax.dynamic_slice(obj, start_idx, cutout_shape), (None, 0)), (0, None))(
+        to_2d(start_idxs), to_3d(obj)
+    ).reshape((*start_idxs.shape[:-1], *obj.shape[:-2], *cutout_shape))
+
+
+@partial(jax.jit, static_argnums=0)
+def outer(ufunc: t.Any, x: jax.Array, y: jax.Array) -> jax.Array:
+    return jax.vmap(jax.vmap(ufunc, (None, 0)), (0, None))(x, y)
+
+
+@partial(jax.jit, static_argnames=('output_shape', 'order', 'mode', 'cval'))
+def affine_transform(
+    input: jax.Array,
+    matrix: ArrayLike,
+    offset: t.Optional[ArrayLike] = None,
+    output_shape: t.Optional[t.Tuple[int, ...]] = None,
+    order: int = 1,
+    mode: str = 'constant',
+    cval: t.Any = 0.0,
+) -> jax.Array:
+    import jax.scipy.ndimage
+    jax_mode = {'grid-constant': 'constant', 'grid-wrap': 'wrap'}.get(mode, mode)
+
+    if output_shape is None:
+        output_shape = input.shape
+    n_axes = len(output_shape)  # num axes to transform over
+
+    indices = jnp.indices(output_shape, dtype=float)
+
+    matrix = jnp.array(matrix)
+    if matrix.shape == (n_axes + 1, n_axes + 1):
+        # homogenous transform matrix
+        coords = jnp.tensordot(
+            matrix, jnp.stack((*indices, jnp.ones_like(indices[0])), axis=0), axes=1
+        )[:-1]
+    elif matrix.shape == (n_axes,):
+        coords = (indices.T * matrix + jnp.array(offset)).T
+    else:
+        raise ValueError(f"Expected matrix of shape ({n_axes + 1}, {n_axes + 1}) or ({n_axes},), instead got shape {matrix.shape}")
+
+    coords += jnp.finfo(coords.dtype).eps
+
+    return jax.vmap(
+        lambda a: jax.scipy.ndimage.map_coordinates(a, tuple(coords), order=order, mode=jax_mode, cval=cval),
+    )(to_nd(input, n_axes + 1)).reshape((*input.shape[:-n_axes], *output_shape))

phaser/utils/analysis.py

@@ -0,0 +1,263 @@
+import typing as t
+
+import numpy
+from numpy.typing import ArrayLike, NDArray
+
+from phaser.utils.num import fft2, ifft2, abs2, Sampling, get_array_module, at, to_numpy, NumT
+from phaser.utils.object import ObjectSampling
+from phaser.utils.image import remove_linear_ramp, translation_matrix, rotation_matrix
+from phaser.state import ObjectState
+from phaser.types import cast_length
+
+
+def split_image(img: NDArray[NumT]) -> t.Tuple[NDArray[NumT], NDArray[NumT], NDArray[NumT], NDArray[NumT]]:
+    """
+    Split image into 4 subimages (upper left, upper right, lower left, lower right).
+    """
+    shape = img.shape[-2:]
+    # crop image divisible by two
+    shape = tuple(s - s%2 for s in shape)
+    img = img[..., *(slice(0, s) for s in shape)]
+
+    return cast_length((
+        img[..., slice(row_start, shape[0], 2), slice(col_start, shape[1], 2)]
+        for row_start in (0, 1)
+        for col_start in (0, 1)
+    ), 4)
+
+
+def fourier_correlate(img1: NDArray[NumT], img2: NDArray[NumT]) -> NDArray[numpy.complex128]:
+    assert img1.shape == img2.shape
+    from skimage.filters import window
+
+    xp = get_array_module(img1, img2)
+
+    win = xp.array(window('hann', img1.shape))
+    img1_fft = numpy.fft.fftshift(fft2(img1.astype(numpy.float64) * win), axes=(-2, -1))
+    img2_fft = xp.conj(numpy.fft.fftshift(fft2(img2.astype(numpy.float64) * win), axes=(-2, -1)))
+    fft1_mag = abs2(img1_fft)
+    fft2_mag = abs2(img2_fft)
+
+    with xp.errstate(invalid='ignore'):
+        corr = img1_fft * img2_fft / xp.sqrt(fft1_mag * fft2_mag)
+    return xp.nan_to_num(corr, nan=1.)
+
+
+def _calc_r_binning(
+    shape: t.Tuple[int, ...], r_spacing: float = 1.0,
+    xp: t.Any = None
+) -> NDArray[numpy.int64]:
+    if xp is None or t.TYPE_CHECKING:
+        xp = numpy
+
+    c_y, c_x = tuple(int(s//2) for s in shape[-2:])
+    y, x = xp.indices(shape[-2:])
+    r = xp.sqrt((y - c_y)**2 + (x - c_x)**2)
+    return xp.floor(r / r_spacing).astype(numpy.int64).ravel()
+
+
+def contrast_transfer(
+    img1: NDArray[NumT], img2: NDArray[NumT], *,
+    inscribed: bool = True, r_spacing: float = 1.0
+) -> t.Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
+    assert img1.shape == img2.shape
+    from skimage.filters import window
+
+    xp = get_array_module(img1, img2)
+
+    win = xp.array(window(('tukey', 0.3), img1.shape), dtype=numpy.float64)
+    img1_fft = xp.fft.fftshift(fft2(img1.astype(numpy.float64) * win), axes=(-2, -1))
+    img2_fft = xp.conj(xp.fft.fftshift(fft2(img2.astype(numpy.float64) * win), axes=(-2, -1)))
+    fft1_mag = abs2(img1_fft)
+    fft2_mag = abs2(img2_fft)
+
+    r_i = _calc_r_binning(img1.shape, r_spacing, xp=xp)
+    mag1_count = xp.bincount(r_i, fft1_mag.ravel())
+    mag2_count = xp.bincount(r_i, fft2_mag.ravel())
+
+    with numpy.errstate(invalid='ignore', divide='ignore'):
+        vals = to_numpy(xp.sqrt(mag1_count / mag2_count))
+    vals = numpy.nan_to_num(vals, nan=1.)
+
+    rs = numpy.linspace(0., len(vals) * r_spacing, len(vals), endpoint=False, dtype=numpy.float64)
+    freq = rs / numpy.sqrt(numpy.prod(img1.shape[-2:]))
+
+    if inscribed:
+        # crop to inner radii
+        n_r = int(numpy.floor(min(img1.shape[-2:]) / (2.*r_spacing)))
+        return (vals[:n_r], rs[:n_r], freq[:n_r])
+    return (vals, rs, freq)
+
+
+def fourier_ring_correlate(
+    img1: NDArray[NumT], img2: NDArray[NumT], *,
+    inscribed: bool = True, r_spacing: float = 1.0
+) -> t.Tuple[NDArray[numpy.float64], NDArray[numpy.float64], NDArray[numpy.float64]]:
+    assert img1.shape == img2.shape
+    from skimage.filters import window
+
+    xp = get_array_module(img1, img2)
+
+    win = xp.array(window(('tukey', 0.3), img1.shape), dtype=numpy.float64)
+    img1_fft = xp.fft.fftshift(fft2(img1.astype(numpy.float64) * win), axes=(-2, -1))
+    img2_fft = xp.conj(xp.fft.fftshift(fft2(img2.astype(numpy.float64) * win), axes=(-2, -1)))
+    fft1_mag = abs2(img1_fft)
+    fft2_mag = abs2(img2_fft)
+
+    r_i = _calc_r_binning(img1.shape, r_spacing, xp=xp)
+    # sum real and imaginary separately
+    real_count = xp.bincount(r_i, (img1_fft * xp.conj(img2_fft)).real.ravel())
+    mag1_count = xp.bincount(r_i, fft1_mag.ravel())
+    mag2_count = xp.bincount(r_i, fft2_mag.ravel())
+
+    with numpy.errstate(invalid='ignore', divide='ignore'):
+        vals = to_numpy(real_count / xp.sqrt(mag1_count * mag2_count))
+    vals = numpy.nan_to_num(vals, posinf=0., nan=1.)
+
+    #vals = numpy.bincount(r_i.ravel(), corr.ravel()) / numpy.bincount(r_i.ravel())
+    rs = numpy.linspace(0., len(vals) * r_spacing, len(vals), endpoint=False, dtype=numpy.float64)
+    freq = rs / numpy.sqrt(numpy.prod(img1.shape[-2:]))
+
+    if inscribed:
+        # crop to inner radii
+        n_r = int(numpy.floor(min(img1.shape[-2:]) / (2.*r_spacing)))
+        return (vals[:n_r], rs[:n_r], freq[:n_r])
+    return (vals, rs, freq)
+
+
+def frc_intersect_threshold(frc: NDArray[numpy.floating], freq: NDArray[numpy.floating], threshold: float) -> t.Tuple[float, float]:
+    diff = frc - threshold
+    lastdiff = numpy.roll(diff, shift=1)
+    # find negative zero crossing
+    try:
+        i = numpy.nonzero((diff[1:] < 0) & (lastdiff[1:] > 0))[0][0]
+    except IndexError:
+        raise ValueError("No crossing found when evaluating FRC resolution")
+
+    m = (diff[i+1] - diff[i]) / (freq[i+1] - freq[i])
+    m_f = (frc[i+1] - frc[i]) / (freq[i+1] - freq[i])
+    x_d = diff[i]/m
+    return (float(freq[i] - x_d), float(frc[i] - x_d * m_f))
+
+
+def _cross_correlate(x: NDArray[numpy.floating], y: NDArray[numpy.floating], max_shift: float) -> t.Tuple[float, float]:
+    xp = get_array_module(x, y)
+
+    samp = Sampling(tuple(x.shape), sampling=(1.0, 1.0))  # type: ignore
+    yy, xx = samp.real_grid(xp=xp)
+
+    cross_corr = ifft2(fft2(x) * xp.conj(fft2(y))).real
+    # limit maximum shift
+    cross_corr = at(cross_corr, yy**2 + xx**2 > max_shift**2).set(numpy.nan)  # type: ignore
+
+    max_i = xp.nanargmax(cross_corr)
+    y = yy.ravel()[max_i]
+    x = xx.ravel()[max_i]
+    return (float(y), float(x))
+
+
+def align_object_to_ground_truth(
+    object: ObjectState,
+    ground_truth: NDArray[numpy.floating],
+    ground_truth_sampling: ArrayLike,
+    rotation_angle: float = 0.0,
+    refinement_niter: int = 0,
+    order: int = 1,
+) -> t.Tuple[NDArray[numpy.floating], NDArray[numpy.floating]]:
+    """
+    ground_truth: Ground truth phase (in radians, or radians/angstrom for multislice data)
+
+    Returns a tuple (object, ground_truth)
+    """
+    xp = get_array_module(object.data, ground_truth)
+    ground_truth = xp.asarray(ground_truth)
+    object_phase = xp.angle(xp.asarray(object.data))
+
+    # normalize multislice objects to radians/angstrom
+    if len(object.thicknesses):
+        object_phase /= object.thicknesses[:, None, None]
+
+    # remove linear ramp
+    object_roi = object.sampling.get_region_mask(xp=xp)
+    object_phase = remove_linear_ramp(object_phase, object_roi)
+    object_phase -= xp.nanquantile(object_phase[..., object_roi], 0.01, axis=-1)[:, None, None]
+    # and get average
+    object_mean = xp.mean(object_phase, axis=0)
+
+    # initially center ground truth on object roi
+    obj_center = object.sampling.get_region_center()
+    ground_truth_corner = obj_center - numpy.array(ground_truth.shape) * ground_truth_sampling / 2.
+    ground_truth_samp = ObjectSampling(
+        t.cast(t.Tuple[int, int], ground_truth.shape),
+        sampling=ground_truth_sampling, corner=ground_truth_corner,
+        region_min=object.sampling.region_min, region_max=object.sampling.region_max,
+    )
+    crop = ground_truth_samp.get_region_crop(pad=-0.05 * object.sampling.get_region_extent())
+
+    # perform initial resampling
+    upsamp_obj = object.sampling.resample(object_mean, ground_truth_samp, cval=0., rotation=rotation_angle, order=order)
+
+    # cross correlate
+    max_shift = float(numpy.min(
+        (ground_truth_samp.extent - object.sampling.get_region_extent()) / (3.0 * ground_truth_samp.sampling)
+    ))
+    if max_shift < 0:
+        raise ValueError("Error: Ground truth extent smaller than object extent")
+    #print(f"max_shift: {max_shift:.3f} px ({max_shift * ground_truth_samp.sampling[0]:.3f} angstrom)")
+
+    shift = numpy.array(_cross_correlate(upsamp_obj[tuple(crop)], ground_truth[tuple(crop)], max_shift))
+
+    # shift ground truth to match experiment
+    ground_truth_samp = ObjectSampling(
+        tuple(ground_truth_samp.shape),
+        sampling=ground_truth_sampling, corner=ground_truth_corner + shift * ground_truth_sampling,
+        region_min=object.sampling.region_min, region_max=object.sampling.region_max, 
+    )
+    crop = ground_truth_samp.get_region_crop()
+
+    if refinement_niter == 0:
+        # perform final upsampling
+        upsamp_obj = object.sampling.resample(object_phase, ground_truth_samp, cval=0., rotation=rotation_angle, order=order)
+        return upsamp_obj[(slice(None), *crop)], ground_truth[tuple(crop)]
+
+    import scipy.optimize
+
+    def _make_affine(mat: NDArray[numpy.floating]) -> NDArray[numpy.floating]:
+        (a, b, c, d, e, f) = mat
+        affine = translation_matrix(obj_center) @ numpy.array([
+            [a, b, e],
+            [c, d, f],
+            [0., 0., 1],
+        ]) @ rotation_matrix(rotation_angle) @ translation_matrix(-obj_center)
+        return affine
+
+    def align_and_correlate(mat: NDArray[numpy.floating]) -> NDArray[numpy.floating]:
+        affine = _make_affine(mat)
+        upsamp_obj = object.sampling.resample(object_mean, ground_truth_samp, cval=0., affine=affine, order=order)
+        return to_numpy(upsamp_obj[tuple(crop)] - ground_truth[tuple(crop)]).ravel()
+
+    max_shift_refine = max_shift * ground_truth_samp.sampling[0] / 5.
+
+    # (a, b, c, d, e, f)
+    init_mat = numpy.array([1., 0., 0., 1., 0., 0.])
+    min_bound = numpy.array([0.9, -0.1, -0.1, 0.9, -max_shift_refine, -max_shift_refine])
+    max_bound = numpy.array([1.1, 0.1, 0.1, 1.1, max_shift_refine, max_shift_refine])
+
+    result = scipy.optimize.least_squares(align_and_correlate, init_mat, bounds=(min_bound, max_bound),
+                                          method='dogbox', max_nfev=refinement_niter, xtol=1e-4)
+    print(f"""\
+Refinement result: {result.message}
+    nfev: {result.nfev}
+    matrix: {result.x}""")
+
+    affine = _make_affine(result.x)
+    (a, b, c, d, e, f) = result.x
+    affine = translation_matrix(obj_center) @ numpy.array([
+        [a, b, e],
+        [c, d, f],
+        [0., 0., 1],
+    ]) @ rotation_matrix(rotation_angle) @ translation_matrix(-obj_center)
+
+    # perform final upsampling
+    upsamp_obj = object.sampling.resample(object_phase, ground_truth_samp, cval=0., affine=affine, order=order)
+    return upsamp_obj[(slice(None), *crop)], ground_truth[tuple(crop)]