From f4758449e6345f8ea67d506f71506c6c68c5ad48 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Sun, 19 May 2024 20:04:34 -0500 Subject: [PATCH 1/2] NEW: Use fixed confidence for position regularization --- src/tike/ptycho/position.py | 29 +++++----- src/tike/ptycho/solvers/lstsq.py | 99 ++++++++++++++++++++------------ 2 files changed, 77 insertions(+), 51 deletions(-) diff --git a/src/tike/ptycho/position.py b/src/tike/ptycho/position.py index 9855f4bd..92e00665 100644 --- a/src/tike/ptycho/position.py +++ b/src/tike/ptycho/position.py @@ -233,9 +233,12 @@ def astuple(self) -> tuple: self.t1, ) - def __call__(self, x: np.ndarray, gpu=False) -> np.ndarray: + def __call__(self, x: np.ndarray, gpu=False, shift=True) -> np.ndarray: xp = cp.get_array_module(x) - return (x @ self.asarray(xp)) + xp.array((self.t0, self.t1)) + result = x @ self.asarray(xp) + if shift: + result += xp.array((self.t0, self.t1)) + return result def estimate_global_transformation( @@ -665,19 +668,13 @@ def _affine_position_helper( scan, position_options: PositionOptions, max_error, - relax=0.1, + relax=0.9, ): predicted_positions = position_options.transform( - position_options.initial_scan) - err = predicted_positions - position_options.initial_scan - # constrain more the probes in flat regions - W = relax * (1 - (position_options.confidence / - (1 + position_options.confidence))) - # penalize positions that are further than max_error from origin; avoid travel larger than max error - W = cp.minimum(10 * relax, - W + cp.maximum(0, err - max_error)**2 / max_error**2) - # allow free movement in depenence on realibility and max allowed error - new_scan = scan * (1 - W) + W * predicted_positions + position_options.initial_scan, + shift=False, + ) + new_scan = scan * (1 - relax) + relax * predicted_positions return new_scan @@ -762,13 +759,15 @@ def gaussian_gradient( sigma=sigma, order=1, axis=-2, - mode='nearest', + mode="nearest", + truncate=6.0, ), cupyx.scipy.ndimage.gaussian_filter1d( -x, sigma=sigma, order=1, axis=-1, - mode='nearest', + mode="nearest", + truncate=6.0, ), ) diff --git a/src/tike/ptycho/solvers/lstsq.py b/src/tike/ptycho/solvers/lstsq.py index c8d6c1da..42a9d5ce 100644 --- a/src/tike/ptycho/solvers/lstsq.py +++ b/src/tike/ptycho/solvers/lstsq.py @@ -634,45 +634,72 @@ def keep_some_args_constant( if position_options: m = 0 + # # start section to compute position certainty metric # TODO: Try adjusting gradient sigma property - grad_x, grad_y = tike.ptycho.position.gaussian_gradient( - bpatches[blo:bhi]) - - # start section to compute position certainty metric - crop = probe.shape[-1] // 4 - total_illumination = op.diffraction.patch.fwd( - images=object_preconditioner, - positions=scan[lo:hi], - patch_width=probe.shape[-1], - )[:, crop:-crop, crop:-crop].real - - power = cp.abs(probe[0, 0, 0, crop:-crop, crop:-crop])**2 + # grad_x, grad_y = tike.ptycho.position.gaussian_gradient( + # cp.angle(bpatches[blo:bhi]), + # sigma=probe.shape[-1] / 6, + # ) + # # FIXME: Try filtering preconditioner first, then just pick the middle one + + # total_illumination = cupyx.scipy.ndimage.gaussian_filter( + # object_preconditioner, + # sigma=(probe.shape[-2] / 6, probe.shape[-1] / 6), + # order=0, + # mode='constant', + # cval=0.0, + # truncate=6.0, + # ) + # # TODO: Do the picking out of the array? + # total_illumination /= np.max(total_illumination) + + # total_illumination = op.diffraction.patch.fwd( + # total_illumination, + # positions=scan[lo:hi] + probe.shape[-1] / 2, + # patch_width=1, + # )[:, 0, 0] + + # grad_x = np.abs(grad_x) + # grad_y = np.abs(grad_y) + # grad_x /= np.max(grad_x) + # grad_y /= np.max(grad_y) + + # weighted_average = tike.ptycho.probe.gaussian_kernel2D( + # probe.shape[-1], + # probe.shape[-1] // 6, + # xp=cp, + # ) + # dX = cp.sum( + # grad_x[:, 0, 0, :, :] * weighted_average, + # axis=(-2, -1), + # keepdims=False, + # ) + # dY = cp.sum( + # grad_y[:, 0, 0, :, :] * weighted_average, + # axis=(-2, -1), + # keepdims=False, + # ) + + # total_variation = cp.stack( + # [total_illumination.real, total_illumination.real], + # axis=1, + # ) + # normalized_variation = ( + # total_variation + # / cp.max( + # total_variation, + # axis=0, + # ) + # + 1e-6 + # ) + # position_options.confidence[lo:hi] = scan[lo:hi] + # # end section to compute position certainty metric - dX = cp.mean( - cp.abs(grad_x[:, 0, 0, crop:-crop, crop:-crop]).real * - total_illumination * power, - axis=(-2, -1), - keepdims=False, - ) - dY = cp.mean( - cp.abs(grad_y[:, 0, 0, crop:-crop, crop:-crop]).real * - total_illumination * power, - axis=(-2, -1), - keepdims=False, + grad_x, grad_y = tike.ptycho.position.gaussian_gradient( + bpatches[blo:bhi], + sigma=0.333, ) - - total_variation = cp.sqrt(cp.stack( - [dX, dY], - axis=1, - )) - mean_variation = (cp.mean( - total_variation**4, - axis=0, - ) + 1e-6) - position_options.confidence[ - lo:hi] = total_variation**4 / mean_variation - # end section to compute position certainty metric - + crop = probe.shape[-1] // 4 position_update_numerator[lo:hi, ..., 0] = cp.sum( cp.real( cp.conj(grad_x[..., crop:-crop, crop:-crop] * From fd2bc637debb4c3f2c03deafaf31734b96e880e1 Mon Sep 17 00:00:00 2001 From: Daniel Ching Date: Sun, 19 May 2024 20:15:08 -0500 Subject: [PATCH 2/2] REF: Remove dead code --- src/tike/ptycho/solvers/lstsq.py | 62 -------------------------------- 1 file changed, 62 deletions(-) diff --git a/src/tike/ptycho/solvers/lstsq.py b/src/tike/ptycho/solvers/lstsq.py index 42a9d5ce..0fc90009 100644 --- a/src/tike/ptycho/solvers/lstsq.py +++ b/src/tike/ptycho/solvers/lstsq.py @@ -633,68 +633,6 @@ def keep_some_args_constant( if position_options: m = 0 - - # # start section to compute position certainty metric - # TODO: Try adjusting gradient sigma property - # grad_x, grad_y = tike.ptycho.position.gaussian_gradient( - # cp.angle(bpatches[blo:bhi]), - # sigma=probe.shape[-1] / 6, - # ) - # # FIXME: Try filtering preconditioner first, then just pick the middle one - - # total_illumination = cupyx.scipy.ndimage.gaussian_filter( - # object_preconditioner, - # sigma=(probe.shape[-2] / 6, probe.shape[-1] / 6), - # order=0, - # mode='constant', - # cval=0.0, - # truncate=6.0, - # ) - # # TODO: Do the picking out of the array? - # total_illumination /= np.max(total_illumination) - - # total_illumination = op.diffraction.patch.fwd( - # total_illumination, - # positions=scan[lo:hi] + probe.shape[-1] / 2, - # patch_width=1, - # )[:, 0, 0] - - # grad_x = np.abs(grad_x) - # grad_y = np.abs(grad_y) - # grad_x /= np.max(grad_x) - # grad_y /= np.max(grad_y) - - # weighted_average = tike.ptycho.probe.gaussian_kernel2D( - # probe.shape[-1], - # probe.shape[-1] // 6, - # xp=cp, - # ) - # dX = cp.sum( - # grad_x[:, 0, 0, :, :] * weighted_average, - # axis=(-2, -1), - # keepdims=False, - # ) - # dY = cp.sum( - # grad_y[:, 0, 0, :, :] * weighted_average, - # axis=(-2, -1), - # keepdims=False, - # ) - - # total_variation = cp.stack( - # [total_illumination.real, total_illumination.real], - # axis=1, - # ) - # normalized_variation = ( - # total_variation - # / cp.max( - # total_variation, - # axis=0, - # ) - # + 1e-6 - # ) - # position_options.confidence[lo:hi] = scan[lo:hi] - # # end section to compute position certainty metric - grad_x, grad_y = tike.ptycho.position.gaussian_gradient( bpatches[blo:bhi], sigma=0.333,