pyTEMlib 0.2024.2.2__py2.py3-none-any.whl → 0.2024.6.0__py2.py3-none-any.whl

pyTEMlib/eels_dialog.py

@@ -1,6 +1,4 @@
 """
-QT dialog window for EELS compositional analysis
-
 Author: Gerd Duscher
 """
 
@@ -93,7 +91,7 @@ class CurveVisualizer(object):
             legline.set_alpha(0.2)
         self.fig.canvas.draw()
 
-def get_sidebar():
+def get_core_loss_sidebar():
     side_bar = ipywidgets.GridspecLayout(14, 3,width='auto', grid_gap="0px")
 
     
@@ -194,6 +192,7 @@ def get_sidebar():
     return side_bar
 
 
+
 class CompositionWidget(object):
     def __init__(self, datasets=None, key=None):
         
@@ -203,7 +202,7 @@ class CompositionWidget(object):
 
         
         self.model = []
-        self.sidebar = get_sidebar()
+        self.sidebar = get_core_loss_sidebar()
         
         self.set_dataset(key)
         
@@ -312,11 +311,12 @@ class CompositionWidget(object):
         reference_list =[('None', -1)]
         
         for index, key in enumerate(self.datasets.keys()):
-            if 'Reference' not in key:
-                if 'SPECTR' in self.datasets[key].data_type.name:
-                    spectrum_list.append((f'{key}: {self.datasets[key].title}', index)) 
-                    self.spectrum_keys_list.append(key)
-            reference_list.append((f'{key}: {self.datasets[key].title}', index))  
+            if '_rel' not in key:
+                if 'Reference' not in key :
+                    if 'SPECTR' in self.datasets[key].data_type.name:
+                        spectrum_list.append((f'{key}: {self.datasets[key].title}', index)) 
+                        self.spectrum_keys_list.append(key)
+                reference_list.append((f'{key}: {self.datasets[key].title}', index))  
         
         if set_key in self.spectrum_keys_list:
             self.key = set_key
@@ -324,10 +324,11 @@ class CompositionWidget(object):
             self.key = self.spectrum_keys_list[-1]
         self.dataset = self.datasets[self.key]
         
-        spec_dim = self.dataset.get_dimensions_by_type(sidpy.DimensionType.SPECTRAL)
-        self.spec_dim = self.dataset._axes[spec_dim[0]]
+        self.spec_dim = self.dataset.get_spectral_dims(return_axis=True)[0]
 
         self.energy_scale = self.spec_dim.values
+        self.dd = (self.energy_scale[0], self.energy_scale[1])
+
         self.dataset.metadata['experiment']['offset'] = self.energy_scale[0]
         self.dataset.metadata['experiment']['dispersion'] = self.energy_scale[1] - self.energy_scale[0]
         if 'edges' not in self.dataset.metadata or self.dataset.metadata['edges'] == {}:
@@ -634,7 +635,7 @@ class CompositionWidget(object):
                 raise ValueError('need a experiment parameter in metadata dictionary')
 
             eff_beta = eels.effective_collection_angle(self.energy_scale, alpha, beta, beam_kv)
-
+            eff_beta = beta
             self.low_loss = None
             if self.sidebar[12, 1].value:
                 for key in self.datasets.keys():

pyTEMlib/eels_dialog_utilities.py

@@ -1033,13 +1033,12 @@ class EdgesAtCursor(object):
             self.label.remove()
             self.line.remove()
         if event.button == 1:
-            self.label = self.ax.text(x, y_max, eels.find_major_edges(event.xdata, self.maximal_chemical_shift),
+            self.label = self.ax.text(x, y_max, eels.find_all_edges(event.xdata, self.maximal_chemical_shift, major_edges_only=True),
                                       verticalalignment='top')
             self.line, = self.ax.plot([x, x], [y_min, y_max], color='black')
         if event.button == 3:
             self.line, = self.ax.plot([x, x], [y_min, y_max], color='black')
-            self.label = self.ax.text(x, y_max, eels.find_all_edges(event.xdata, self.maximal_chemical_shift),
-                                      verticalalignment='top')
+            self.label = self.ax.text(x, y_max, eels.find_all_edges(event.xdata, self.maximal_chemical_shift), verticalalignment='top')
         self.ax.set_ylim(y_min, y_max)
 
     def mouse_move(self, event):

pyTEMlib/eels_tools.py

@@ -373,15 +373,30 @@ def shift_energy(dataset: sidpy.Dataset, shifts: np.ndarray) -> sidpy.Dataset:
 
 
 def align_zero_loss(dataset: sidpy.Dataset) -> sidpy.Dataset:
+    """
+    Shifts the energy axis of the input dataset to be aligned with the zero-loss peak.
+
+    Parameters:
+    -----------
+    dataset : sidpy.Dataset
+        The input dataset containing the energy axis to be aligned.
 
+    Returns:
+    --------
+    sidpy.Dataset
+        The dataset with the energy axis shifted to align the zero-loss peak.
+
+    """
     shifts = get_zero_loss_energy(dataset)
-    print(shifts, dataset)
+    # print(shifts, dataset)
     new_si = shift_energy(dataset, shifts)    
     new_si.metadata.update({'zero_loss': {'shifted': shifts}})
     return new_si
 
 
-def get_resolution_functions(dset: sidpy.Dataset, startFitEnergy: float=-1, endFitEnergy: float=+1,
+
+
+def get_resolution_functions(dataset: sidpy.Dataset, startFitEnergy: float=-1, endFitEnergy: float=+1,
                              n_workers: int=1, n_threads: int=8):
     """
     Analyze and fit low-loss EELS data within a specified energy range to determine zero-loss peaks.
@@ -392,33 +407,35 @@ def get_resolution_functions(dset: sidpy.Dataset, startFitEnergy: float=-1, endF
     from the dataset. The function handles both 2D and 3D datasets.
 
     Parameters:
-        dset: sidpy.Dataset
-            The dataset containing TEM spectral data.
-        startFitEnergy: float
-            The start energy of the fitting window.
-        endFitEnergy: float
-            The end energy of the fitting window.
-        n_workers: int, optional
-            The number of workers for parallel processing (default is 1).
-        n_threads: int, optional
-            The number of threads for parallel processing (default is 8).
+    -----------
+        dataset (sidpy.Dataset): The dataset containing TEM spectral data.
+        startFitEnergy (float): The start energy of the fitting window.
+        endFitEnergy (float): The end energy of the fitting window.
+        n_workers (int, optional): The number of workers for parallel processing (default is 1).
+        n_threads (int, optional): The number of threads for parallel processing (default is 8).
 
     Returns:
+    --------
         tuple: A tuple containing:
             - z_loss_dset (sidpy.Dataset): The dataset with added zero-loss peak information.
             - z_loss_params (numpy.ndarray): Array of parameters used for the zero-loss peak fitting.
 
     Raises:
+    -------
         ValueError: If the input dataset does not have the expected dimensions or format.
 
     Notes:
+    ------
         - The function expects `dset` to have specific dimensionalities and will raise an error if they are not met.
         - Parallel processing is employed to enhance performance, particularly for large datasets.
     """
-    energy = dset.get_spectral_dims(return_axis=True)[0].values
+    energy = dataset.get_spectral_dims(return_axis=True)[0].values
     start_fit_pixel = np.searchsorted(energy, startFitEnergy)
     end_fit_pixel = np.searchsorted(energy, endFitEnergy)
     guess_width = (endFitEnergy - startFitEnergy)/2
+    if end_fit_pixel - start_fit_pixel < 5:
+        start_fit_pixel -= 2
+        end_fit_pixel += 2
     
     def get_good_guess(zl_func, energy, spectrum):
         popt, pcov = curve_fit(zl_func, energy, spectrum,
@@ -428,21 +445,21 @@ def get_resolution_functions(dset: sidpy.Dataset, startFitEnergy: float=-1, endF
 
     fit_energy = energy[start_fit_pixel:end_fit_pixel]
     # get a good guess for the fit parameters
-    if len(dset.shape) == 3:
-        fit_dset = dset[:, :, start_fit_pixel:end_fit_pixel]
+    if len(dataset.shape) == 3:
+        fit_dset = dataset[:, :, start_fit_pixel:end_fit_pixel]
         guess_amplitude = np.sqrt(fit_dset.max())
         guess_params = get_good_guess(zl_func, fit_energy, fit_dset.sum(axis=(0, 1))/fit_dset.shape[0]/fit_dset.shape[1])
-    elif len(dset.shape) == 2:
-        fit_dset = dset[:, start_fit_pixel:end_fit_pixel]
+    elif len(dataset.shape) == 2:
+        fit_dset = dataset[:, start_fit_pixel:end_fit_pixel]
         fit_energy = energy[start_fit_pixel:end_fit_pixel]
         guess_amplitude = np.sqrt(fit_dset.max())
         guess_params = get_good_guess(zl_func, fit_energy, fit_dset.sum(axis=0)/fit_dset.shape[0])
-    elif len(dset.shape) == 1:
-        fit_dset = dset[start_fit_pixel:end_fit_pixel]
+    elif len(dataset.shape) == 1:
+        fit_dset = dataset[start_fit_pixel:end_fit_pixel]
         fit_energy = energy[start_fit_pixel:end_fit_pixel]
         guess_amplitude = np.sqrt(fit_dset.max())
         guess_params = get_good_guess(zl_func, fit_energy, fit_dset)
-        z_loss_dset = dset.copy()
+        z_loss_dset = dataset.copy()
         z_loss_dset *= 0.0
         z_loss_dset += zl_func(energy, *guess_params)
         if 'zero_loss' not in z_loss_dset.metadata:
@@ -450,11 +467,11 @@ def get_resolution_functions(dset: sidpy.Dataset, startFitEnergy: float=-1, endF
         z_loss_dset.metadata['zero_loss'].update({'startFitEnergy': startFitEnergy,
                                                   'endFitEnergy': endFitEnergy,
                                                   'fit_parameter': guess_params,
-                                                  'original_low_loss': dset.title})
+                                                  'original_low_loss': dataset.title})
         return z_loss_dset
     else:
         print('Error: need a spectrum or spectral image sidpy dataset')
-        print('Not dset.shape = ', dset.shape)
+        print('Not dset.shape = ', dataset.shape)
         return None
 
     # define guess function for SidFitter
@@ -466,7 +483,7 @@ def get_resolution_functions(dset: sidpy.Dataset, startFitEnergy: float=-1, endF
                                  return_cov=False, return_fit=False, return_std=False, km_guess=False, num_fit_parms=6)
     
     [z_loss_params] = zero_loss_fitter.do_fit()
-    z_loss_dset = dset.copy()
+    z_loss_dset = dataset.copy()
     z_loss_dset *= 0.0
 
     energy_grid = np.broadcast_to(energy.reshape((1, 1, -1)), (z_loss_dset.shape[0],
@@ -480,7 +497,7 @@ def get_resolution_functions(dset: sidpy.Dataset, startFitEnergy: float=-1, endF
     z_loss_dset.metadata['zero_loss'].update({'startFitEnergy': startFitEnergy,
                                               'endFitEnergy': endFitEnergy,
                                               'fit_parameter': z_loss_params,
-                                              'original_low_loss': dset.title})
+                                              'original_low_loss': dataset.title})
 
 
     return z_loss_dset
@@ -503,7 +520,7 @@ def drude_lorentz(eps_inf, leng, ep, eb, gamma, e, amplitude):
     return eps
 
 
-def fit_plasmon(dataset, startFitEnergy, endFitEnergy, plot_result=False, number_workers=4, number_threads=8):
+def fit_plasmon(dataset: Union[sidpy.Dataset, np.ndarray], startFitEnergy: float, endFitEnergy: float, plot_result: bool = False, number_workers: int = 4, number_threads: int = 8) -> Union[sidpy.Dataset, np.ndarray]:
     """
     Fit plasmon peak positions and widths in a TEM dataset using a Drude model.
 
@@ -540,7 +557,7 @@ def fit_plasmon(dataset, startFitEnergy, endFitEnergy, plot_result=False, number
         - If `plot_result` is True, the function plots Ep, Ew, and A as separate subplots.
     """
     # define Drude function for plasmon fitting
-    def energy_loss_function(E, Ep, Ew, A):
+    def energy_loss_function(E: np.ndarray, Ep: float, Ew: float, A: float) -> np.ndarray:
         E = E/E.max()
         eps = 1 - Ep**2/(E**2+Ew**2) + 1j * Ew * Ep**2/E/(E**2+Ew**2)
         elf = (-1/eps).imag
@@ -592,7 +609,6 @@ def drude_simulation(dset, e, ep, ew, tnm, eb):
     Gives probabilities of dielectric function eps relative to zero-loss integral (i0 = 1) per eV
     Details in R.F.Egerton: EELS in the Electron Microscope, 3rd edition, Springer 2011
 
-    # function drude(ep,ew,eb,epc,e0,beta,nn,tnm)
     # Given the plasmon energy (ep), plasmon fwhm (ew) and binding energy(eb),
     # this program generates:
     # EPS1, EPS2 from modified Eq. (3.40), ELF=Im(-1/EPS) from Eq. (3.42),
@@ -603,41 +619,10 @@ def drude_simulation(dset, e, ep, ew, tnm, eb):
     # Details in R.F.Egerton: EELS in the Electron Microscope, 3rd edition, Springer 2011
     # Version 10.11.26
 
-
-    b.7 drude Simulation of a Low-Loss Spectrum
-    The program DRUDE calculates a single-scattering plasmon-loss spectrum for
-    a specimen of a given thickness tnm (in nm), recorded with electrons of a
-    specified incident energy e0 by a spectrometer that accepts scattering up to a
-    specified collection semi-angle beta. It is based on the extended drude model
-    (Section 3.3.2), with a volume energy-loss function elf in accord with Eq. (3.64) and
-    a surface-scattering energy-loss function srelf as in Eq. (4.31). Retardation effects
-    and coupling between the two surface modes are not included. The surface term can
-    be made negligible by entering a large specimen thickness (tnm > 1000).
-    Surface intensity srfint and volume intensity volint are calculated from
-    Eqs. (4.31) and (4.26), respectively. The total spectral intensity ssd is written to
-    the file DRUDE.SSD, which can be used as input for KRAKRO. These intensities are
-    all divided by i0, to give relative probabilities (per eV). The real and imaginary parts
-    of the dielectric function are written to DRUDE.EPS and can be used for comparison
-    with the results of Kramers–Kronig analysis (KRAKRO.DAT).
-    Written output includes the surface-loss probability Ps, obtained by integrating
-    srfint (a value that relates to two surfaces but includes the negative begrenzungs
-    term), for comparison with the analytical integration represented by Eq. (3.77). The
-    volume-loss probability p_v is obtained by integrating volint and is used to calculate
-    the volume plasmon mean free path (lam = tnm/p_v). The latter is listed and
-    compared with the MFP obtained from Eq. (3.44), which represents analytical integration
-    assuming a zero-width plasmon peak. The total probability (Pt = p_v+Ps) is
-    calculated and used to evaluate the thickness (lam.Pt) that would be given by the formula
-    t/λ = ln(It/i0), ignoring the surface-loss probability. Note that p_v will exceed
-    1 for thicker specimens (t/λ > 1), since it represents the probability of plasmon
-    scattering relative to that of no inelastic scattering.
-    The command-line usage is drude(ep,ew,eb,epc,beta,e0,tnm,nn), where ep is the
-    plasmon energy, ew the plasmon width, eb the binding energy of the electrons (0 for
-    a metal), and nn is the number of channels in the output spectrum. An example of
-    the output is shown in Fig. b.1a,b.
-
     """
-    
-    epc = dset.energy_scale[1] - dset.energy_scale[0]  # input('ev per channel : ');
+    energy_scale = dset.get_spectral_dims(return_axis=True)[0].values
+
+    epc = energy_scale[1] - energy_scale[0]  # input('ev per channel : ');
     
     b = dset.metadata['collection_angle'] / 1000.  # rad
     epc = dset.energy_scale[1] - dset.energy_scale[0]  # input('ev per channel : ');
@@ -1465,31 +1450,31 @@ def fit_edges2(spectrum, energy_scale, edges):
 
 
     def model(xx, pp):
-        yy = pp[0] + x**pp[1] + pp[2] + pp[3] * xx + pp[4] * xx * xx
+        yy = pp[0] *  xx**pp[1] +  pp[2] + pp[3]* xx + pp[4] * xx * xx
         for i in range(number_of_edges):
             pp[i+5] = np.abs(pp[i+5])
             yy = yy + pp[i+5] * xsec[i, :]
         return yy
 
     def residuals(pp, xx, yy):
-        err = np.abs((yy - model(xx, pp)) * mask)  # / np.sqrt(np.abs(y))
+        err = np.abs((yy - model(xx, pp)) * mask)  / np.sqrt(np.abs(y))
         return err
 
     scale = y[100]
-    pin = np.array([A,r, 10., 1., 0.00] + [scale/5] * number_of_edges)
+    pin = np.array([A,-r, 10., 1., 0.00] + [scale/5] * number_of_edges)
     [p, _] = leastsq(residuals, pin, args=(x, y))
 
     for key in edges:
         if key.isdigit():
             edges[key]['areal_density'] = p[int(key)+5]
-
+    print(p)
     edges['model'] = {}
-    edges['model']['background'] = (background + p[6] + p[7] * x + p[8] * x * x)
-    edges['model']['background-poly_0'] = p[6]
-    edges['model']['background-poly_1'] = p[7]
-    edges['model']['background-poly_2'] = p[8]
-    edges['model']['background-A'] = A
-    edges['model']['background-r'] = r
+    edges['model']['background'] = ( p[0] * np.power(x, -p[1])+ p[2]+ x**p[3] +  p[4] * x * x)
+    edges['model']['background-poly_0'] = p[2]
+    edges['model']['background-poly_1'] = p[3]
+    edges['model']['background-poly_2'] = p[4]
+    edges['model']['background-A'] = p[0]
+    edges['model']['background-r'] = p[1]
     edges['model']['spectrum'] = model(x, p)
     edges['model']['blurred'] = blurred
     edges['model']['mask'] = mask

pyTEMlib/file_tools.py

@@ -731,7 +731,13 @@ def open_file(filename=None,  h5_group=None, write_hdf_file=False, sum_frames=Fa
             read_essential_metadata(dset)
             dset.metadata['filename'] = filename
             dataset_dict = {'Channel_000': dset}
-            
+
+        # Temporary Fix for dual eels spectra in dm files
+        # Todo: Fic in ScifyReaders
+        for dset in dataset_dict.values():   
+            if 'single_exposure_time' in dset.metadata['experiment']:
+                dset.metadata['experiment']['exposure_time'] = dset.metadata['experiment']['number_of_frames'] * \
+                                                               dset.metadata['experiment']['single_exposure_time']
         if write_hdf_file:
             h5_master_group = save_dataset(dataset_dict, filename=filename)
 

pyTEMlib/image_tools.py

@@ -50,6 +50,11 @@ from sklearn.cluster import DBSCAN
 
 from collections import Counter
 
+# center diff function
+from skimage.filters import threshold_otsu, sobel
+from scipy.optimize import leastsq
+from sklearn.cluster import DBSCAN
+
 
 _SimpleITK_present = True
 try:
@@ -275,6 +280,67 @@ def diffractogram_spots(dset, spot_threshold, return_center=True, eps=0.1):
     return spots, center
 
 
+def center_diffractogram(dset, return_plot = True, histogram_factor = None, smoothing = 1, min_samples = 100):
+    try:
+        diff = np.array(dset).T.astype(np.float16)
+        diff[diff < 0] = 0
+        
+        if histogram_factor is not None:
+            hist, bins = np.histogram(np.ravel(diff), bins=256, range=(0, 1), density=True)
+            threshold = threshold_otsu(diff, hist = hist * histogram_factor)
+        else:
+            threshold = threshold_otsu(diff)
+        binary = (diff > threshold).astype(float)
+        smoothed_image = ndimage.gaussian_filter(binary, sigma=smoothing) # Smooth before edge detection
+        smooth_threshold = threshold_otsu(smoothed_image)
+        smooth_binary = (smoothed_image > smooth_threshold).astype(float)
+        # Find the edges using the Sobel operator
+        edges = sobel(smooth_binary)
+        edge_points = np.argwhere(edges)
+
+        # Use DBSCAN to cluster the edge points
+        db = DBSCAN(eps=10, min_samples=min_samples).fit(edge_points)
+        labels = db.labels_
+        if len(set(labels)) == 1:
+            raise ValueError("DBSCAN clustering resulted in only one group, check the parameters.")
+
+        # Get the largest group of edge points
+        unique, counts = np.unique(labels, return_counts=True)
+        counts = dict(zip(unique, counts))
+        largest_group = max(counts, key=counts.get)
+        edge_points = edge_points[labels == largest_group]
+
+        # Fit a circle to the diffraction ring
+        def calc_distance(c, x, y):
+            Ri = np.sqrt((x - c[0])**2 + (y - c[1])**2)
+            return Ri - Ri.mean()
+        x_m = np.mean(edge_points[:, 1])
+        y_m = np.mean(edge_points[:, 0])
+        center_guess = x_m, y_m
+        center, ier = leastsq(calc_distance, center_guess, args=(edge_points[:, 1], edge_points[:, 0]))
+        mean_radius = np.mean(calc_distance(center, edge_points[:, 1], edge_points[:, 0])) + np.sqrt((edge_points[:, 1] - center[0])**2 + (edge_points[:, 0] - center[1])**2).mean()
+    
+    finally:
+        if return_plot:
+            fig, ax = plt.subplots(1, 4, figsize=(10, 4))
+            ax[0].set_title('Diffractogram')
+            ax[0].imshow(dset.T, cmap='viridis')
+            ax[1].set_title('Otsu Binary Image')
+            ax[1].imshow(binary, cmap='gray')
+            ax[2].set_title('Smoothed Binary Image')
+            ax[2].imshow(smooth_binary, cmap='gray')
+            ax[3].set_title('Edge Detection and Fitting')
+            ax[3].imshow(edges, cmap='gray')
+            ax[3].scatter(center[0], center[1], c='r', s=10)
+            circle = plt.Circle(center, mean_radius, color='red', fill=False)
+            ax[3].add_artist(circle)
+            for axis in ax:
+                axis.axis('off')
+            fig.tight_layout()
+          
+    return center
+
+
 def adaptive_fourier_filter(dset, spots, low_pass=3, reflection_radius=0.3):
     """
     Use spots in diffractogram for a Fourier Filter
@@ -375,16 +441,10 @@ def complete_registration(main_dataset, storage_channel=None):
 
     rigid_registered_dataset = rigid_registration(main_dataset)
 
-    if storage_channel is not None:
-        registration_channel = ft.log_results(storage_channel, rigid_registered_dataset)
-
+    
     print('Non-Rigid_Registration')
 
     non_rigid_registered = demon_registration(rigid_registered_dataset)
-    if storage_channel is not None:
-        registration_channel = ft.log_results(storage_channel, non_rigid_registered)
-
-    non_rigid_registered.h5_dataset = registration_channel
     return non_rigid_registered, rigid_registered_dataset
 
 
@@ -473,7 +533,7 @@ def demon_registration(dataset, verbose=False):
 ###############################
 # Rigid Registration New 05/09/2020
 
-def rigid_registration(dataset):
+def rigid_registration(dataset, sub_pixel=True):
     """
     Rigid registration of image stack with pixel accuracy
 
@@ -529,9 +589,13 @@ def rigid_registration(dataset):
         selection[frame_dim[0]] = slice(i, i+1)
         moving = dataset[tuple(selection)].squeeze().compute()
         fft_moving = np.fft.fft2(moving)
-        image_product = fft_fixed * fft_moving.conj()
-        cc_image = np.fft.fftshift(np.fft.ifft2(image_product))
-        shift = np.array(ndimage.maximum_position(cc_image.real))-cc_image.shape[0]/2
+        if sub_pixel:
+            shift = skimage.registration.phase_cross_correlation(fft_fixed, fft_moving, upsample_factor=1000,
+                                                                 space='fourier')[0]
+        else:    
+            image_product = fft_fixed * fft_moving.conj()
+            cc_image = np.fft.fftshift(np.fft.ifft2(image_product))
+            shift = np.array(ndimage.maximum_position(cc_image.real))-cc_image.shape[0]/2
         fft_fixed = fft_moving
         relative_drift.append(shift)
     rig_reg, drift = rig_reg_drift(dataset, relative_drift)
@@ -542,9 +606,18 @@ def rigid_registration(dataset):
     rigid_registered.source = dataset.title
     rigid_registered.metadata = {'analysis': 'rigid sub-pixel registration', 'drift': drift,
                                  'input_crop': input_crop, 'input_shape': dataset.shape[1:]}
-    rigid_registered.set_dimension(0, dataset._axes[frame_dim[0]])
-    rigid_registered.set_dimension(1, dataset._axes[spatial_dim[0]][input_crop[0]:input_crop[1]])
-    rigid_registered.set_dimension(2, dataset._axes[spatial_dim[1]][input_crop[2]:input_crop[3]])
+    rigid_registered.set_dimension(0, sidpy.Dimension(np.arange(rigid_registered.shape[0]), 
+                                          name='frame', units='frame', quantity='time',
+                                          dimension_type='temporal'))
+    
+    array_x = dataset._axes[spatial_dim[0]][input_crop[0]:input_crop[1]].values
+    rigid_registered.set_dimension(1, sidpy.Dimension(array_x,
+                                          'x', units='nm', quantity='Length',
+                                          dimension_type='spatial'))
+    array_y = dataset._axes[spatial_dim[1]][input_crop[2]:input_crop[3]].values
+    rigid_registered.set_dimension(2, sidpy.Dimension(array_y,
+                                          'y', units='nm', quantity='Length',
+                                          dimension_type='spatial'))
     return rigid_registered.rechunk({0: 'auto', 1: -1, 2: -1})
 
 
@@ -589,6 +662,7 @@ def rig_reg_drift(dset, rel_drift):
     rig_reg = np.zeros([dset.shape[frame_dim[0]], dset.shape[spatial_dim[0]], dset.shape[spatial_dim[1]]])
 
     # absolute drift
+    print(rel_drift)
     drift = np.array(rel_drift).copy()
 
     drift[0] = [0, 0]
@@ -1070,9 +1144,8 @@ def cartesian2polar(x, y, grid, r, t, order=3):
     return ndimage.map_coordinates(grid, np.array([new_ix, new_iy]), order=order).reshape(new_x.shape)
 
 
-def warp(diff):
-    """Takes a centered diffraction pattern (as a sidpy dataset)and warps it to a polar grid"""
-    """Centered diff can be produced with it.diffractogram_spots(return_center = True)"""
+def warp(diff, center):
+    """Takes a diffraction pattern (as a sidpy dataset)and warps it to a polar grid"""
 
     # Define original polar grid
     nx = np.shape(diff)[0]
@@ -1080,20 +1153,19 @@ def warp(diff):
 
     # Define center pixel
     pix2nm = np.gradient(diff.u.values)[0]
-    center_pixel = [abs(min(diff.u.values)), abs(min(diff.v.values))]//pix2nm
 
-    x = np.linspace(1, nx, nx, endpoint=True)-center_pixel[0]
-    y = np.linspace(1, ny, ny, endpoint=True)-center_pixel[1]
+    x = np.linspace(1, nx, nx, endpoint=True)-center[0]
+    y = np.linspace(1, ny, ny, endpoint=True)-center[1]
     z = diff
 
     # Define new polar grid
-    nr = int(min([center_pixel[0], center_pixel[1], diff.shape[0]-center_pixel[0], diff.shape[1]-center_pixel[1]])-1)
-    nt = 360*3
+    nr = int(min([center[0], center[1], diff.shape[0]-center[0], diff.shape[1]-center[1]])-1)
+    nt = 360 * 3
 
     r = np.linspace(1, nr, nr)
     t = np.linspace(0., np.pi, nt, endpoint=False)
 
-    return cartesian2polar(x, y, z, r, t, order=3)
+    return cartesian2polar(x, y, z, r, t, order=3).T
 
 
 def calculate_ctf(wavelength, cs, defocus, k):