From c66de47f9589d08fdb14d82a58cd0fc06bcd3ed6 Mon Sep 17 00:00:00 2001 From: Abe Levitan Date: Tue, 1 Feb 2022 17:04:46 -0500 Subject: [PATCH] Misc updates --- CDTools/datasets/ptycho_2d_dataset.py | 2 +- CDTools/models/__init__.py | 1 + CDTools/models/fancy_ptycho.py | 34 +- CDTools/models/multimode_rpi.py | 359 +++++++++++++++++++++ CDTools/models/rpi.py | 7 +- CDTools/tools/analysis/analysis.py | 164 +++++++++- CDTools/tools/interactions/interactions.py | 7 +- CDTools/tools/plotting/plotting.py | 17 +- 8 files changed, 560 insertions(+), 31 deletions(-) create mode 100644 CDTools/models/multimode_rpi.py diff --git a/CDTools/datasets/ptycho_2d_dataset.py b/CDTools/datasets/ptycho_2d_dataset.py index f2c1c76..8959163 100644 --- a/CDTools/datasets/ptycho_2d_dataset.py +++ b/CDTools/datasets/ptycho_2d_dataset.py @@ -233,5 +233,5 @@ class Ptycho2DDataset(CDataset): else: cbar_title='Diffraction Intensity' - plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title) + return plotting.plot_nanomap_with_images(self.translations.detach().cpu(), get_images, values=nanomap_values, nanomap_units=units, image_title='Diffraction Pattern', image_colorbar_title=cbar_title) diff --git a/CDTools/models/__init__.py b/CDTools/models/__init__.py index 49dcca2..90f9e37 100644 --- a/CDTools/models/__init__.py +++ b/CDTools/models/__init__.py @@ -31,6 +31,7 @@ from CDTools.models.polarized_fancy_ptycho import PolarizedFancyPtycho from CDTools.models.bragg_2d_ptycho import Bragg2DPtycho from CDTools.models.multislice_2d_ptycho import Multislice2DPtycho from CDTools.models.rpi import RPI +from CDTools.models.multimode_rpi import MultimodeRPI # Still needs to be updated for the new complex numbers #from CDTools.models.s_matrix_ptycho import SMatrixPtycho diff --git a/CDTools/models/fancy_ptycho.py b/CDTools/models/fancy_ptycho.py index 6bbf2bf..b768012 100644 --- a/CDTools/models/fancy_ptycho.py +++ b/CDTools/models/fancy_ptycho.py @@ -278,8 +278,15 @@ class FancyPtycho(CDIModel): basis_prs = self.probe * self.probe_support[..., :, :] # Now we construct the probes for each shot from the basis probes - Ws = self.weights[index] - if len(self.weights[0].shape) == 0: + if self.weights is not None: + Ws = self.weights[index] + else: + try: + Ws = t.ones(len(index)) # I'm positive this introduced a bug + except: + Ws = 1 + + if self.weights is None or len(self.weights[0].shape) == 0: # If a purely stable coherent illumination is defined prs = Ws[..., None, None, None] * basis_prs else: @@ -291,11 +298,6 @@ class FancyPtycho(CDIModel): prs = t.sum(Ws[..., None, None] * basis_prs, axis=-3) if self.simulate_probe_translation: - #det_pix_trans = t.tensordot( - # translations, - # t.as_tensor(self.detector_geometry['basis'], - # dtype=t.float32), - # dims=1) det_pix_trans = tools.interactions.translations_to_pixel( self.detector_geometry['basis'], translations, @@ -363,7 +365,7 @@ class FancyPtycho(CDIModel): self.surface_normal = self.surface_normal.to(*args, **kwargs) - def sim_to_dataset(self, args_list): + def sim_to_dataset(self, args_list, calculation_width=None): # In the future, potentially add more control # over what metadata is saved (names, etc.) @@ -389,9 +391,23 @@ class FancyPtycho(CDIModel): wavelength = self.wavelength indices, translations = args_list + data = [] + len(indices) + if calculation_width is None: + calculation_width = len(indices) + index_chunks = [indices[i:i + calculation_width] + for i in range(0, len(indices), + calculation_width)] + translation_chunks = [translations[i:i + calculation_width] + for i in range(0, len(indices), + calculation_width)] + + # Then we simulate the results - data = self.forward(indices, translations).detach() + data = [self.forward(idx, trans).detach() + for idx, trans in zip(index_chunks, translation_chunks)] + data = t.cat(data, dim=0) # And finally, we make the dataset return Ptycho2DDataset( translations, data, diff --git a/CDTools/models/multimode_rpi.py b/CDTools/models/multimode_rpi.py new file mode 100644 index 0000000..2db9e21 --- /dev/null +++ b/CDTools/models/multimode_rpi.py @@ -0,0 +1,359 @@ +import torch as t +from CDTools.models import CDIModel +from CDTools import tools +from CDTools.tools import plotting as p +from CDTools.tools.interactions import RPI_interaction +from CDTools.tools import initializers +from scipy.ndimage.morphology import binary_dilation +import numpy as np +from copy import copy + +__all__ = ['MultimodeRPI'] + + + +__all__ = ['RPI'] + +class MultimodeRPI(CDIModel): + + @property + def obj(self): + return t.complex(self.obj_real, self.obj_imag) + + @property + def weights(self): + ws = t.complex(self.weights_real, self.weights_imag) + return ws / 10# / self.obj_real.size().numel() + + def __init__(self, wavelength, detector_geometry, probe_basis, + probe, obj_guess, detector_slice=None, + background=None, mask=None, saturation=None, + obj_support=None, oversampling=1, weight_matrix=False): + + super(MultimodeRPI, self).__init__() + + self.wavelength = t.tensor(wavelength) + self.detector_geometry = copy(detector_geometry) + + det_geo = self.detector_geometry + if hasattr(det_geo, 'distance'): + det_geo['distance'] = t.tensor(det_geo['distance']) + if hasattr(det_geo, 'basis'): + det_geo['basis'] = t.tensor(det_geo['basis']) + if hasattr(det_geo, 'corner'): + det_geo['corner'] = t.tensor(det_geo['corner']) + + self.probe_basis = t.tensor(probe_basis) + + scale_factor = t.tensor([probe.shape[-1]/obj_guess.shape[-1], + probe.shape[-2]/obj_guess.shape[-2]]) + self.obj_basis = self.probe_basis * scale_factor + self.detector_slice = detector_slice + + # Maybe something to include in a bit + # self.surface_normal = t.tensor(surface_normal) + + self.saturation = saturation + + if mask is None: + self.mask = mask + else: + self.mask = t.tensor(mask, dtype=t.bool) + + + self.probe = t.tensor(probe, dtype=t.complex64) + + obj_guess = t.tensor(obj_guess, dtype=t.complex64) + + self.obj_real = t.nn.Parameter(obj_guess.real) + self.obj_imag = t.nn.Parameter(obj_guess.imag) + + self.weights_real = t.nn.Parameter(t.eye(probe.shape[0])* 10)# * self.obj_real.size().numel()) + self.weights_imag = t.nn.Parameter(t.zeros(probe.shape[0])) + + if not weight_matrix: + self.weights_real.requires_grad=False + self.weights_imag.requires_grad=False + + # Wait for LBFGS to be updated for complex-valued parameters + # self.obj = t.nn.Parameter(obj_guess.to(t.float32)) + + if background is None: + if detector_slice is not None: + background = 1e-6 * t.ones( + self.probe[0][self.detector_slice].shape, + dtype=t.float32) + else: + background = 1e-6 * t.ones(self.probe[0].shape, + dtype=t.float32) + + self.background = t.tensor(background, dtype=t.float32) + + if obj_support is not None: + self.obj_support = obj_support + self.obj.data = self.obj * obj_support[None, ...] + else: + self.obj_support = t.ones_like(self.obj[0, ...]) + + self.oversampling = oversampling + + + @classmethod + def from_dataset(cls, dataset, probe, obj_size=None, background=None, mask=None, padding=0, n_modes=1, saturation=None, scattering_mode=None, oversampling=1, auto_center=False, initialization='random', opt_for_fft=False, weight_matrix=False, probe_threshold=0): + raise NotImplementedError() + + wavelength = dataset.wavelength + det_basis = dataset.detector_geometry['basis'] + det_shape = dataset[0][1].shape + distance = dataset.detector_geometry['distance'] + + # always do this on the cpu + get_as_args = dataset.get_as_args + dataset.get_as(device='cpu') + # We only need the patterns here, not the inputs associated with them. + _, patterns = dataset[:] + dataset.get_as(*get_as_args[0],**get_as_args[1]) + + # Set to none to avoid issues with things outside the detector + if auto_center: + center = tools.image_processing.centroid(t.sum(patterns,dim=0)) + else: + center = None + + # Then, generate the probe geometry from the dataset + ewg = tools.initializers.exit_wave_geometry + probe_basis, probe_shape, det_slice = ewg(det_basis, + det_shape, + wavelength, + distance, + center=center, + padding=padding, + opt_for_fft=opt_for_fft, + oversampling=oversampling) + + if not isinstance(probe,t.Tensor): + probe = t.as_tensor(probe) + + # Potentially need all of this orientation stuff later + + #if hasattr(dataset, 'sample_info') and \ + # dataset.sample_info is not None and \ + # 'orientation' in dataset.sample_info: + # surface_normal = dataset.sample_info['orientation'][2] + #else: + # surface_normal = np.array([0.,0.,1.]) + + # If this information is supplied when the function is called, + # then we override the information in the .cxi file + #if scattering_mode in {'t', 'transmission'}: + # surface_normal = np.array([0.,0.,1.]) + #elif scattering_mode in {'r', 'reflection'}: + # outgoing_dir = np.cross(det_basis[:,0], det_basis[:,1]) + # outgoing_dir /= np.linalg.norm(outgoing_dir) + # surface_normal = outgoing_dir + np.array([0.,0.,1.]) + # surface_normal /= np.linalg.norm(surface_normal) + + + if background is None and hasattr(dataset, 'background') \ + and dataset.background is not None: + background = t.sqrt(dataset.background) + elif background is not None: + background = t.sqrt(t.Tensor(background).to(dtype=t.float32)) + + det_geo = dataset.detector_geometry + + # If no mask is given, but one exists in the dataset, load it. + if mask is None and hasattr(dataset, 'mask') \ + and dataset.mask is not None: + mask = dataset.mask.to(t.bool) + + # Now we initialize the object + if obj_size is None: + # This is a standard size for a well-matched probe and detector + obj_size = (np.array(probe_shape) // 2).astype(int) + + if initialization.lower().strip() == 'random': + # I think something to do with the fact that the object is defined + # on a coarser grid needs to be accounted for here that is not + # accounted for yet + scale = t.sum(patterns[0]) / t.sum(t.abs(probe)**2) + obj_guess = scale * t.exp(2j * np.pi * t.rand([n_modes,]+obj_size)) + elif initialization.lower().strip() == 'spectral': + if background is not None: + obj_guess = initializers.RPI_spectral_init( + patterns[0], probe, obj_size, mask=mask, + background=background**2, n_modes=n_modes) + else: + obj_guess = initializers.RPI_spectral_init( + patterns[0], probe, obj_size, mask=mask, + n_modes=n_modes) + + else: + raise KeyError('Initialization "' + str(initialization) + \ + '" invalid - use "spectral" or "random"') + + probe_intensity = t.sqrt(t.sum(t.abs(probe)**2,axis=0)) + probe_fft = tools.propagators.far_field(probe_intensity) + pad0l = (probe.shape[-2] - obj_size[-2])//2 + pad0r = probe.shape[-2] - obj_size[-2] - pad0l + pad1l = (probe.shape[-1] - obj_size[-1])//2 + pad1r = probe.shape[-1] - obj_size[-1] - pad1l + probe_lr_fft = probe_fft[pad0l:-pad0r,pad1l:-pad1r] + probe_lr = t.abs(tools.propagators.inverse_far_field(probe_lr_fft)) + + obj_support = probe_lr > t.max(probe_lr) * probe_threshold + obj_support = t.as_tensor(binary_dilation(obj_support)) + + return cls(wavelength, det_geo, probe_basis, + probe, obj_guess, detector_slice=det_slice, + background=background, mask=mask, saturation=saturation, + obj_support=obj_support, oversampling=oversampling, + weight_matrix=weight_matrix) + + + def random_init(self, pattern): + scale = t.sum(pattern) / t.sum(t.abs(self.probe)**2) + self.obj.data = scale * t.exp( + 2j * np.pi * t.rand(self.obj.shape)).to( + dtype=self.obj.dtype, device=self.obj.device) + + def spectral_init(self, pattern): + if self.background is not None: + self.obj.data = initializers.RPI_spectral_init( + pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask, + background=self.background**2, n_modes=self.obj.shape[0]).to( + dtype=self.obj.dtype, device=self.obj.device) + else: + self.obj.data = initializers.RPI_spectral_init( + pattern, self.probe, self.obj.shape[-3:-1], mask=self.mask, + n_modes=self.obj.shape[0]).to( + dtype=self.obj.dtype, device=self.obj.device) + + # Needs work + def interaction(self, index, *args): + # including *args allows this to work with all sorts of datasets + # that might include other information in with the index in their + # "input" parameters (such as translations for a ptychography dataset). + # This makes it seamless to use such a dataset even though those + # extra arguments will not be used. + + + all_exit_waves = [] + + # Mix the probes with the weight matrix + prs = t.sum(self.weights[..., None, None] * self.probe, axis=-3) + + for i in range(self.probe.shape[0]): + pr = prs[i] + # Here we have a 3D probe (one single mode) + # and a 4D object (multiple modes mixing incoherently) + exit_waves = RPI_interaction(pr, + self.obj_support * self.obj[i]) + all_exit_waves.append(exit_waves.unsqueeze(0)) + + # This creates a bunch of modes generated from all possible combos + # of the probe and object modes all strung out along the first index + + output = t.cat(all_exit_waves) + + # If we have multiple indexes input, we unsqueeze and repeat the stack + # of wavefields enough times to simulate each requested index. This + # seems silly, but it enables (for example) one to do a reconstruction + # from a set of diffraction patterns that are all known to be from the + # same object. + try: + # will fail if index has no length, for example when index + # is just an int. In this case, we just do nothing instead + output = output.unsqueeze(0).repeat(1,len(index),1,1,1) + except TypeError: + pass + return output + + + def forward_propagator(self, wavefields): + return tools.propagators.far_field(wavefields) + + + def backward_propagator(self, wavefields): + return tools.propagators.inverse_far_field(wavefields) + + + def measurement(self, wavefields): + # Here I'm taking advantage of an undocumented feature in the + # incoherent_sum measurement function where it will work with + # a 4D wavefield array as well as a 5D array. + return tools.measurements.quadratic_background(wavefields, + self.background, + detector_slice=self.detector_slice, + measurement=tools.measurements.incoherent_sum, + saturation=self.saturation, + oversampling=self.oversampling) + + def loss(self, sim_data, real_data, mask=None): + return tools.losses.amplitude_mse(real_data, sim_data, mask=mask) + #return tools.losses.poisson_nll(real_data, sim_data, mask=mask) + + def regularizer(self, factors): + return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \ + + factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2) + + def to(self, *args, **kwargs): + super(MultimodeRPI, self).to(*args, **kwargs) + self.wavelength = self.wavelength.to(*args,**kwargs) + # move the detector geometry too + det_geo = self.detector_geometry + if hasattr(det_geo, 'distance'): + det_geo['distance'] = det_geo['distance'].to(*args,**kwargs) + if hasattr(det_geo, 'basis'): + det_geo['basis'] = det_geo['basis'].to(*args,**kwargs) + if hasattr(det_geo, 'corner'): + det_geo['corner'] = det_geo['corner'].to(*args,**kwargs) + + if self.mask is not None: + self.mask = self.mask.to(*args, **kwargs) + + self.probe = self.probe.to(*args,**kwargs) + self.probe_basis = self.probe_basis.to(*args,**kwargs) + self.obj_basis = self.obj_basis.to(*args,**kwargs) + self.obj_support = self.obj_support.to(*args,**kwargs) + self.background = self.background.to(*args, **kwargs) + + # Maybe include in a bit + #self.surface_normal = self.surface_normal.to(*args, **kwargs) + + def sim_to_dataset(self, args_list): + raise NotImplementedError('No sim to dataset yet, sorry!') + + plot_list = [ + ('Root Sum Squared Amplitude of all Probes', + lambda self, fig: p.plot_amplitude( + np.sqrt(np.sum((t.abs(t.sum(self.weights[..., None, None].detach() * self.probe, axis=-3))**2).cpu().numpy(),axis=0)), + fig=fig, basis=self.probe_basis)), + ('Object Amplitudes', + lambda self, fig: p.plot_amplitude(self.obj, fig=fig, + basis=self.obj_basis)), + ('Object Phases', + lambda self, fig: p.plot_phase(self.obj, fig=fig, + basis=self.obj_basis)) + ] + + + def save_results(self, dataset=None, full_obj=False): + # dataset is set as a kwarg here because it isn't needed, but the + # common pattern is to pass a dataset. This makes it okay if one + # continues to use that standard pattern + probe_basis = self.probe_basis.detach().cpu().numpy() + obj_basis = self.obj_basis.detach().cpu().numpy() + probe = self.probe.detach().cpu().numpy() + # Provide the option to save out the subdominant objects or + # just the dominant one + if full_obj: + obj = self.obj.detach().cpu().numpy() + else: + obj = self.obj[0].detach().cpu().numpy() + background = self.background.detach().cpu().numpy()**2 + + return {'probe_basis': probe_basis, 'obj_basis': obj_basis, + 'probe': probe,'obj': obj, + 'background': background} + diff --git a/CDTools/models/rpi.py b/CDTools/models/rpi.py index 6bd2e66..a2e0e51 100644 --- a/CDTools/models/rpi.py +++ b/CDTools/models/rpi.py @@ -329,8 +329,11 @@ class RPI(CDIModel): #return tools.losses.poisson_nll(real_data, sim_data, mask=mask) def regularizer(self, factors): - return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \ - + factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2) + if self.obj.shape[0] == 1: + return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) + else: + return factors[0] * t.sum(t.abs(self.obj[0,:,:])**2) \ + + factors[1] * t.sum(t.abs(self.obj[1:,:,:])**2) def to(self, *args, **kwargs): super(RPI, self).to(*args, **kwargs) diff --git a/CDTools/tools/analysis/analysis.py b/CDTools/tools/analysis/analysis.py index dd3e8fb..feefb11 100644 --- a/CDTools/tools/analysis/analysis.py +++ b/CDTools/tools/analysis/analysis.py @@ -457,7 +457,7 @@ def calc_deconvolved_cross_correlation(im1, im2, im_slice=None): return cor -def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.): +def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1., limit='side'): """Calculates a Fourier ring correlation between two images This function requires an input of a basis to allow for FRC calculations @@ -480,6 +480,8 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.): Number of bins to break the FRC up into snr : float The signal to noise ratio (for the combined information in both images) to return a threshold curve for. + limit : str + Default is 'side'. What is the highest frequency to calculate the FRC to? If 'side', it chooses the side of the Fourier transform, if 'corner' it goes fully to the corner. Returns ------- @@ -508,14 +510,14 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.): (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] if nbins is None: - nbins = np.max(im1[im_slice].shape) // 4 + nbins = np.max(im1[im_slice].shape) // 8 + f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) + f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)) + cor_fft = f1 * t.conj(f2) - cor_fft = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) * \ - t.fft.fftshift(t.conj(t.fft.fft2(im2[im_slice])),dim=(-1,-2)) - - F1 = t.abs(t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)))**2 - F2 = t.abs(t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)))**2 + F1 = t.abs(f1)**2 + F2 = t.abs(f2)**2 di = np.linalg.norm(basis[:,0]) @@ -527,21 +529,42 @@ def calc_frc(im1, im2, basis, im_slice=None, nbins=None, snr=1.): Js,Is = np.meshgrid(j_freqs,i_freqs) Rs = np.sqrt(Is**2+Js**2) + if limit.lower().strip() == 'side': + max_i = np.max(i_freqs) + max_j = np.max(j_freqs) + frc_range = [0, max(max_i,max_j)] + elif limit.lower().strip() == 'corner': + frc_range = [0, np.max(Rs)] + else: + raise ValueError('Invalid FRC limit: choose "side" or "corner"') + + numerator, bins = np.histogram(Rs, bins=nbins, range=frc_range, + weights=cor_fft.numpy()) + denominator_F1, bins = np.histogram(Rs, bins=nbins, range=frc_range, + weights=F1.detach().cpu().numpy()) + denominator_F2, bins = np.histogram(Rs, bins=nbins, range=frc_range, + weights=F2.detach().cpu().numpy()) + n_pix, bins = np.histogram(Rs, bins=nbins, range=frc_range) - numerator, bins = np.histogram(Rs,bins=nbins,weights=cor_fft.numpy()) - denominator_F1, bins = np.histogram(Rs,bins=nbins,weights=F1.detach().cpu().numpy()) - denominator_F2, bins = np.histogram(Rs,bins=nbins,weights=F2.detach().cpu().numpy()) - n_pix, bins = np.histogram(Rs,bins=nbins) - - frc = np.abs(numerator / np.sqrt(denominator_F1*denominator_F2)) + n_pix = n_pix / 4 # This is for an apodized image, apodized with a hann window + + frc = np.abs(numerator) / np.sqrt(denominator_F1*denominator_F2) # This moves from combined-image SNR to single-image SNR snr /= 2 - threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \ + # NOTE: I should update this to produce lots of different threshold curves + # sigma, 2sigma, 3sigma, traditional FRC 1-bit, my better one, n_pix, etc. + + threshold = (snr + (2 * np.sqrt(snr) + 1) / np.sqrt(n_pix)) / \ (1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix)) + my_threshold = np.sqrt(snr**2 + (2*snr**2 + 2*snr + 1)/n_pix) / \ + np.sqrt(snr**2 + 2 * snr + 1 + 2*snr**2 / n_pix) + + twosigma_threshold = 2/ np.sqrt(n_pix) + if not im_np: bins = t.tensor(bins) frc = t.tensor(frc) @@ -741,7 +764,14 @@ def calc_fidelity(fields_1, fields_2, dims=2): mult = fields_1.unsqueeze(-dims-2) * fields_2.unsqueeze(-dims-1).conj() sumdims = tuple(d - dims for d in range(dims)) mat = t.sum(mult,dim=sumdims) + + # Because I think this is the nuclear norm squared, I would like to swap + # Out the definition for this, but I need to test it before swapping. + # It also probably makes sense to implement sqrt_fidelity separately + # because that's more important + #return t.linalg.matrix_norm(mat, ord='nuc')**2 + # I think this is just the nuclear norm. svdvals = t.linalg.svdvals(mat) return t.sum(svdvals, dim=-1)**2 @@ -822,3 +852,109 @@ def calc_generalized_rms_error(fields_1, fields_2, normalize=False, dims=2): return t.sqrt(result) + +def calc_generalized_frc(fields_1, fields_2, basis, im_slice=None, nbins=None, snr=1.): + """Calculates a Fourier ring correlation between two images + + This function requires an input of a basis to allow for FRC calculations + to be related to physical units. + + Like other analysis functions, this can take input in numpy or pytorch, + and will return output in the respective format. + + Parameters + ---------- + im1 : array + The first image, a complex or real valued array + im2 : array + The first image, a complex or real valued array + basis : array + The basis for the images, defined as is standard for datasets + im_slice : slice + Default is from 3/8 to 5/8 across the image, a slice to use in the processing. + nbins : int + Number of bins to break the FRC up into + snr : float + The signal to noise ratio (for the combined information in both images) to return a threshold curve for. + + Returns + ------- + freqs : array + The frequencies associated with each FRC value + FRC : array + The FRC values + threshold : array + The threshold curve for comparison + + """ + + im_np = False + if isinstance(fields_1, np.ndarray): + fields_1 = t.as_tensor(fields_) + im_np = True + if isinstance(fields_2, np.ndarray): + fields_2 = t.as_tensor(fields_2) + im_np = True + + if isinstance(basis, np.ndarray): + basis = t.tensor(basis) + + if im_slice is None: + im_slice = np.s_[(im1.shape[0]//8)*3:(im1.shape[0]//8)*5, + (im1.shape[1]//8)*3:(im1.shape[1]//8)*5] + + if nbins is None: + nbins = np.max(fields_1[...,im_slice].shape[-2:]) // 4 + + + f1 = t.fft.fftshift(t.fft.fft2(im1[im_slice]),dim=(-1,-2)) + f2 = t.fft.fftshift(t.fft.fft2(im2[im_slice]),dim=(-1,-2)) + cor_fft = f1 * t.conj(f2) + + + F1 = t.abs(f1)**2 + F2 = t.abs(f2)**2 + + + di = np.linalg.norm(basis[:,0]) + dj = np.linalg.norm(basis[:,1]) + + i_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[0],d=di)) + j_freqs = fftpack.fftshift(fftpack.fftfreq(cor_fft.shape[1],d=dj)) + + Js,Is = np.meshgrid(j_freqs,i_freqs) + Rs = np.sqrt(Is**2+Js**2) + + # This line is used to get a set of bins that matches the logic + # used by np.histogram, so that this function will match the choices + # of bin edges that comes from the non-generalized version. This also + # gets us the count on the number of pixels per bin so we can calculate + # the threshold curve + n_pix, bins = np.histogram(Rs,bins=nbins) + + frc = [] + for i in range(len(bins)-1): + mask = t.logical_and(Rs=bins[i]) + masked_f1 = f1 * mask[...,:,:] + masked_f2 = f2 * mask[...,:,:] + numerator = t.sqrt(calc_fidelity(masked_f1, masked_f2)) + denominator_f1 = t.sqrt(calc_fidelity(masked_f1, masked_f1)) + denominator_f2 = t.sqrt(calc_fidelity(masked_f2, masked_f2)) + frc.append(numerator / t.sqrt((denominator_f1 * denominator_f2))) + + frc = np.array(frc) + + # This moves from combined-image SNR to single-image SNR + snr /= 2 + + threshold = (snr + (2 * snr + 1) / np.sqrt(n_pix)) / \ + (1 + snr + (2 * np.sqrt(snr)) / np.sqrt(n_pix)) + + if not im_np: + bins = t.tensor(bins) + frc = t.tensor(frc) + threshold = t.tensor(threshold) + + return bins[:-1], frc, threshold + + diff --git a/CDTools/tools/interactions/interactions.py b/CDTools/tools/interactions/interactions.py index 246698a..f89eac8 100644 --- a/CDTools/tools/interactions/interactions.py +++ b/CDTools/tools/interactions/interactions.py @@ -614,9 +614,12 @@ def RPI_interaction(probe, obj): # The far-field propagator is just a 2D FFT but with an fftshift fftobj = propagators.far_field(obj) # We calculate the padding that we need to do the upsampling - pad0l = (probe.shape[-2] - obj.shape[-2])//2 + # This is carefully set up to keep the zero-frequency pixel in the correct + # location as the overall shape changes. Don't mess with this without + # having thought about this carefully. + pad0l = probe.shape[-2]//2 - obj.shape[-2]//2 pad0r = probe.shape[-2] - obj.shape[-2] - pad0l - pad1l = (probe.shape[-1] - obj.shape[-1])//2 + pad1l = probe.shape[-1]//2 - obj.shape[-1]//2 pad1r = probe.shape[-1] - obj.shape[-1] - pad1l if obj.dim() == 2: diff --git a/CDTools/tools/plotting/plotting.py b/CDTools/tools/plotting/plotting.py index 1f77d6f..5f4d315 100644 --- a/CDTools/tools/plotting/plotting.py +++ b/CDTools/tools/plotting/plotting.py @@ -416,7 +416,7 @@ def plot_colorized(im, fig=None, basis=None, units='$\\mu$m', **kwargs): units=units, **kwargs) -def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwargs): +def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, invert_xaxis=True, **kwargs): """Plots a set of probe translations in a nicely formatted way Parameters @@ -429,6 +429,8 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa Default is um, units to report in (assuming input in m) lines : bool Whether to plot lines indicating the path taken + invert_xaxis : bool + Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective \\**kwargs All other args are passed to fig.add_subplot(111, \\**kwargs) @@ -453,6 +455,9 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa translations = translations * factor plt.plot(translations[:,0], translations[:,1],'k.') + if invert_xaxis: + plt.gca().invert_xaxis() + if lines: plt.plot(translations[:,0], translations[:,1],'b-', linewidth=0.5) plt.xlabel('X (' + units + ')') @@ -461,7 +466,7 @@ def plot_translations(translations, fig=None, units='$\\mu$m', lines=True, **kwa return fig -def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe'): +def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='probe', invert_xaxis=True): """Plots a set of nanomap data in a flexible way Parameters @@ -476,6 +481,8 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr Default is um, units to report in (assuming input in m) convention : str Default is 'probe', alternative is 'obj'. Whether the translations refer to the probe or object. + invert_xaxis : bool + Default is True. This flips the x axis to match the convention from .cxi files of viewing the image from the beam's perspective Returns ------- @@ -509,7 +516,9 @@ def plot_nanomap(translations, values, fig=None, units='$\\mu$m', convention='pr s /= 4 # A rough value to make the size work out plt.scatter(factor * trans[:,0],factor * trans[:,1],s=s,c=values) - + if invert_xaxis: + plt.gca().invert_xaxis() + plt.gca().set_facecolor('k') plt.xlabel('Translation x (' + units + ')') plt.ylabel('Translation y (' + units + ')') @@ -767,3 +776,5 @@ def plot_nanomap_with_images(translations, get_image_func, values=None, mask=Non # (like the nanomap dot sizes) that otherwise would change on the # first update update(0) + + return fig