nxs-analysis-tools 0.0.35__py3-none-any.whl → 0.0.36__py3-none-any.whl

This diff represents the content of publicly available package versions that have been released to one of the supported registries. The information contained in this diff is provided for informational purposes only and reflects changes between package versions as they appear in their respective public registries.

Potentially problematic release.


This version of nxs-analysis-tools might be problematic. Click here for more details.

@@ -3,20 +3,89 @@ Tools for generating single crystal pair distribution functions.
3
3
  """
4
4
  import time
5
5
  import os
6
+ import gc
7
+ import math
6
8
  from scipy import ndimage
9
+ import scipy
7
10
  import matplotlib.pyplot as plt
8
11
  from matplotlib.transforms import Affine2D
9
12
  from nexusformat.nexus import nxsave, NXroot, NXentry, NXdata, NXfield
10
13
  import numpy as np
11
- from .datareduction import plot_slice, reciprocal_lattice_params, Padder
14
+ from astropy.convolution import Kernel, convolve_fft
15
+ import pyfftw
16
+ from .datareduction import plot_slice, reciprocal_lattice_params, Padder, \
17
+ array_to_nxdata
18
+
12
19
 
13
20
  class Symmetrizer2D:
14
21
  """
15
22
  A class for symmetrizing 2D datasets.
23
+
24
+ The `Symmetrizer2D` class provides functionality to apply symmetry
25
+ operations such as rotation and mirroring to 2D datasets.
26
+
27
+ Attributes
28
+ ----------
29
+ mirror_axis : int or None
30
+ The axis along which mirroring is performed. Default is None, meaning
31
+ no mirroring is applied.
32
+ symmetrized : NXdata or None
33
+ The symmetrized dataset after applying the symmetrization operations.
34
+ Default is None until symmetrization is performed.
35
+ wedges : NXdata or None
36
+ The wedges extracted from the dataset based on the angular limits.
37
+ Default is None until symmetrization is performed.
38
+ rotations : int or None
39
+ The number of rotations needed to reconstruct the full dataset from
40
+ a single wedge. Default is None until parameters are set.
41
+ transform : Affine2D or None
42
+ The transformation matrix used for skewing and scaling the dataset.
43
+ Default is None until parameters are set.
44
+ mirror : bool or None
45
+ Indicates whether mirroring is performed during symmetrization.
46
+ Default is None until parameters are set.
47
+ skew_angle : float or None
48
+ The skew angle (in degrees) between the principal axes of the plane
49
+ to be symmetrized. Default is None until parameters are set.
50
+ theta_max : float or None
51
+ The maximum angle (in degrees) for symmetrization. Default is None
52
+ until parameters are set.
53
+ theta_min : float or None
54
+ The minimum angle (in degrees) for symmetrization. Default is None
55
+ until parameters are set.
56
+ wedge : NXdata or None
57
+ The dataset wedge used in the symmetrization process. Default is
58
+ None until symmetrization is performed.
59
+ symmetrization_mask : NXdata or None
60
+ The mask used for selecting the region of the dataset to be symmetrized.
61
+ Default is None until symmetrization is performed.
62
+
63
+ Methods
64
+ -------
65
+ __init__(**kwargs):
66
+ Initializes the Symmetrizer2D object and optionally sets the parameters
67
+ using `set_parameters`.
68
+ set_parameters(theta_min, theta_max, lattice_angle=90, mirror=True, mirror_axis=0):
69
+ Sets the parameters for the symmetrization operation, including angle limits,
70
+ lattice angle, and mirroring options.
71
+ symmetrize_2d(data):
72
+ Symmetrizes a 2D dataset based on the set parameters.
73
+ test(data, **kwargs):
74
+ Performs a test visualization of the symmetrization process, displaying the
75
+ original data, mask, wedge, and symmetrized result.
16
76
  """
17
77
  symmetrization_mask: NXdata
18
78
 
19
79
  def __init__(self, **kwargs):
80
+ """
81
+ Initializes the Symmetrizer2D object.
82
+
83
+ Parameters
84
+ ----------
85
+ **kwargs : dict, optional
86
+ Keyword arguments that can be passed to the `set_parameters` method to
87
+ set the symmetrization parameters during initialization.
88
+ """
20
89
  self.mirror_axis = None
21
90
  self.symmetrized = None
22
91
  self.wedges = None
@@ -32,7 +101,8 @@ class Symmetrizer2D:
32
101
 
33
102
  def set_parameters(self, theta_min, theta_max, lattice_angle=90, mirror=True, mirror_axis=0):
34
103
  """
35
- Sets the parameters for the symmetrization operation.
104
+ Sets the parameters for the symmetrization operation, and calculates the
105
+ required transformations and rotations.
36
106
 
37
107
  Parameters
38
108
  ----------
@@ -41,7 +111,8 @@ class Symmetrizer2D:
41
111
  theta_max : float
42
112
  The maximum angle in degrees for symmetrization.
43
113
  lattice_angle : float, optional
44
- The angle in degrees between the two principal axes of the plane to be symmetrized (default: 90).
114
+ The angle in degrees between the two principal axes of the plane to be
115
+ symmetrized (default: 90).
45
116
  mirror : bool, optional
46
117
  If True, perform mirroring during symmetrization (default: True).
47
118
  mirror_axis : int, optional
@@ -79,7 +150,8 @@ class Symmetrizer2D:
79
150
 
80
151
  def symmetrize_2d(self, data):
81
152
  """
82
- Symmetrizes a 2D dataset based on the set parameters.
153
+ Symmetrizes a 2D dataset based on the set parameters, applying padding
154
+ to prevent rotation cutoff and handling overlapping pixels.
83
155
 
84
156
  Parameters
85
157
  ----------
@@ -98,9 +170,10 @@ class Symmetrizer2D:
98
170
  t = self.transform
99
171
  rotations = self.rotations
100
172
 
101
- # Pad the dataset so that rotations don't get cutoff if they extend past the extent of the dataset
173
+ # Pad the dataset so that rotations don't get cutoff if they extend
174
+ # past the extent of the dataset
102
175
  p = Padder(data)
103
- padding = tuple([len(data[axis]) for axis in data.axes])
176
+ padding = tuple(len(data[axis]) for axis in data.axes)
104
177
  data_padded = p.pad(padding)
105
178
 
106
179
  # Define axes that span the plane to be transformed
@@ -113,18 +186,24 @@ class Symmetrizer2D:
113
186
  # Calculate the angle for each data point
114
187
  theta = np.arctan2(q1.reshape((-1, 1)), q2.reshape((1, -1)))
115
188
  # Create a boolean array for the range of angles
116
- symmetrization_mask = np.logical_and(theta >= theta_min * np.pi / 180, theta <= theta_max * np.pi / 180)
117
- self.symmetrization_mask = NXdata(NXfield(p.unpad(symmetrization_mask), name='mask'),
118
- (data[data.axes[0]], data[data.axes[1]]))
119
-
120
- self.wedge = NXdata(NXfield(p.unpad(counts * symmetrization_mask), name=data.signal),
121
- (data[data.axes[0]], data[data.axes[1]]))
189
+ symmetrization_mask = np.logical_and(theta >= theta_min * np.pi / 180,
190
+ theta <= theta_max * np.pi / 180)
191
+ self.symmetrization_mask = NXdata(NXfield(p.unpad(symmetrization_mask),
192
+ name='mask'),
193
+ (data[data.axes[0]], data[data.axes[1]])
194
+ )
195
+
196
+ self.wedge = NXdata(NXfield(p.unpad(counts * symmetrization_mask),
197
+ name=data.signal),
198
+ (data[data.axes[0]], data[data.axes[1]])
199
+ )
122
200
 
123
201
  # Scale and skew counts
124
202
  skew_angle_adj = 90 - self.skew_angle
125
203
  counts_skew = ndimage.affine_transform(counts,
126
204
  t.inverted().get_matrix()[:2, :2],
127
- offset=[counts.shape[0] / 2 * np.sin(skew_angle_adj * np.pi / 180), 0],
205
+ offset=[counts.shape[0] / 2
206
+ * np.sin(skew_angle_adj * np.pi / 180), 0],
128
207
  order=0,
129
208
  )
130
209
  scale1 = np.cos(skew_angle_adj * np.pi / 180)
@@ -143,7 +222,7 @@ class Symmetrizer2D:
143
222
 
144
223
  # Reconstruct full dataset from wedge
145
224
  reconstructed = np.zeros(counts.shape)
146
- for n in range(0, rotations):
225
+ for _ in range(0, rotations):
147
226
  # The following are attempts to combine images with minimal overlapping pixels
148
227
  reconstructed += wedge
149
228
  # reconstructed = np.where(reconstructed == 0, reconstructed + wedge, reconstructed)
@@ -156,28 +235,33 @@ class Symmetrizer2D:
156
235
  if mirror:
157
236
  # The following are attempts to combine images with minimal overlapping pixels
158
237
  reconstructed = np.where(reconstructed == 0,
159
- reconstructed + np.flip(reconstructed, axis=mirror_axis), reconstructed)
238
+ reconstructed + np.flip(reconstructed, axis=mirror_axis),
239
+ reconstructed)
160
240
  # reconstructed += np.flip(reconstructed, axis=0)
161
241
 
162
242
  # self.rotated_and_mirrored = NXdata(NXfield(reconstructed, name=data.signal),
163
243
  # (q1, q2))
164
244
 
165
245
  reconstructed = ndimage.affine_transform(reconstructed,
166
- Affine2D().scale(scale2, 1).inverted().get_matrix()[:2, :2],
246
+ Affine2D().scale(
247
+ scale2, 1
248
+ ).inverted().get_matrix()[:2, :2],
167
249
  offset=[-(1 - scale2) * counts.shape[
168
250
  0] / 2 / scale2, 0],
169
251
  order=0,
170
252
  )
171
253
  reconstructed = ndimage.affine_transform(reconstructed,
172
- Affine2D().scale(scale1,
173
- 1).inverted().get_matrix()[:2, :2],
254
+ Affine2D().scale(
255
+ scale1, 1
256
+ ).inverted().get_matrix()[:2, :2],
174
257
  offset=[-(1 - scale1) * counts.shape[
175
258
  0] / 2 / scale1, 0],
176
259
  order=0,
177
260
  )
178
261
  reconstructed = ndimage.affine_transform(reconstructed,
179
262
  t.get_matrix()[:2, :2],
180
- offset=[(-counts.shape[0] / 2 * np.sin(skew_angle_adj * np.pi / 180)),
263
+ offset=[(-counts.shape[0] / 2
264
+ * np.sin(skew_angle_adj * np.pi / 180)),
181
265
  0],
182
266
  order=0,
183
267
  )
@@ -185,7 +269,8 @@ class Symmetrizer2D:
185
269
  reconstructed_unpadded = p.unpad(reconstructed)
186
270
 
187
271
  # Fix any overlapping pixels by truncating counts to max
188
- reconstructed_unpadded[reconstructed_unpadded > data[data.signal].nxdata.max()] = data[data.signal].nxdata.max()
272
+ reconstructed_unpadded[reconstructed_unpadded > data[data.signal].nxdata.max()] \
273
+ = data[data.signal].nxdata.max()
189
274
 
190
275
  symmetrized = NXdata(NXfield(reconstructed_unpadded, name=data.signal),
191
276
  (data[data.axes[0]],
@@ -195,7 +280,8 @@ class Symmetrizer2D:
195
280
 
196
281
  def test(self, data, **kwargs):
197
282
  """
198
- Performs a test visualization of the symmetrization process.
283
+ Performs a test visualization of the symmetrization process to help assess
284
+ the effect of the parameters.
199
285
 
200
286
  Parameters
201
287
  ----------
@@ -209,12 +295,13 @@ class Symmetrizer2D:
209
295
  fig : Figure
210
296
  The matplotlib Figure object that contains the test visualization plot.
211
297
  axesarr : ndarray
212
- The numpy array of Axes objects representing the subplots in the test visualization.
298
+ The numpy array of Axes objects representing the subplots in the test
299
+ visualization.
213
300
 
214
301
  Notes
215
302
  -----
216
- This method uses the `symmetrize_2d` method to perform the symmetrization on the input data and visualize
217
- the process.
303
+ This method uses the `symmetrize_2d` method to perform the symmetrization on
304
+ the input data and visualize the process.
218
305
 
219
306
  The test visualization plot includes the following subplots:
220
307
  - Subplot 1: The original dataset.
@@ -224,7 +311,7 @@ class Symmetrizer2D:
224
311
 
225
312
  Example usage:
226
313
  ```
227
- s = Scissors()
314
+ s = Symmetrizer2D()
228
315
  s.set_parameters(theta_min, theta_max, skew_angle, mirror)
229
316
  s.test(data)
230
317
  ```
@@ -244,18 +331,24 @@ class Symmetrizer2D:
244
331
 
245
332
  class Symmetrizer3D:
246
333
  """
247
- A class to symmetrize 3D datasets.
334
+ A class to symmetrize 3D datasets by performing sequential 2D symmetrization on
335
+ different planes.
336
+
337
+ This class applies 2D symmetrization on the three principal planes of a 3D dataset,
338
+ effectively enhancing the symmetry of the data across all axes.
248
339
  """
249
340
 
250
341
  def __init__(self, data=None):
251
342
  """
252
- Initialize the Symmetrizer3D object.
343
+ Initialize the Symmetrizer3D object with an optional 3D dataset.
344
+
345
+ If data is provided, the corresponding q-vectors and planes are automatically
346
+ set up for symmetrization.
253
347
 
254
348
  Parameters
255
349
  ----------
256
350
  data : NXdata, optional
257
351
  The input 3D dataset to be symmetrized.
258
-
259
352
  """
260
353
 
261
354
  self.a, self.b, self.c, self.al, self.be, self.ga = [None] * 6
@@ -282,13 +375,12 @@ class Symmetrizer3D:
282
375
 
283
376
  def set_data(self, data):
284
377
  """
285
- Sets the data to be symmetrized.
378
+ Sets the 3D dataset to be symmetrized and updates the corresponding q-vectors and planes.
286
379
 
287
380
  Parameters
288
381
  ----------
289
382
  data : NXdata
290
383
  The input 3D dataset to be symmetrized.
291
-
292
384
  """
293
385
  self.data = data
294
386
  self.q1 = data[data.axes[0]]
@@ -303,20 +395,33 @@ class Symmetrizer3D:
303
395
  print("Plane 3: " + self.plane3)
304
396
 
305
397
  def set_lattice_params(self, lattice_params):
398
+ """
399
+ Sets the lattice parameters and calculates the reciprocal lattice parameters.
400
+
401
+ Parameters
402
+ ----------
403
+ lattice_params : tuple of float
404
+ The lattice parameters (a, b, c, alpha, beta, gamma) in real space.
405
+ """
306
406
  self.a, self.b, self.c, self.al, self.be, self.ga = lattice_params
307
407
  self.lattice_params = lattice_params
308
408
  self.reciprocal_lattice_params = reciprocal_lattice_params(lattice_params)
309
- self.a_star, self.b_star, self.c_star, self.al_star, self.be_star, self.ga_star = self.reciprocal_lattice_params
409
+ self.a_star, self.b_star, self.c_star, \
410
+ self.al_star, self.be_star, self.ga_star = self.reciprocal_lattice_params
310
411
 
311
412
  def symmetrize(self):
312
413
  """
313
- Perform the symmetrization of the 3D dataset.
414
+ Perform the symmetrization of the 3D dataset by sequentially applying
415
+ 2D symmetrization on the three principal planes.
416
+
417
+ This method symmetrizes the dataset on the three principal planes
418
+ (q1-q2, q1-q3, q2-q3) and handles any negative values that might result
419
+ from the symmetrization process.
314
420
 
315
421
  Returns
316
422
  -------
317
423
  symmetrized : NXdata
318
424
  The symmetrized 3D dataset.
319
-
320
425
  """
321
426
  starttime = time.time()
322
427
  data = self.data
@@ -324,48 +429,49 @@ class Symmetrizer3D:
324
429
  out_array = np.zeros(data[data.signal].shape)
325
430
 
326
431
  print('Symmetrizing ' + self.plane1 + ' planes...')
327
- for k in range(0, len(q3)):
328
- print('Symmetrizing ' + q3.nxname + '=' + "{:.02f}".format(q3[k]) + "...", end='\r')
432
+ for k, value in enumerate(q3):
433
+ print(f'Symmetrizing {q3.nxname}={value:.02f}...', end='\r')
329
434
  data_symmetrized = self.plane1symmetrizer.symmetrize_2d(data[:, :, k])
330
435
  out_array[:, :, k] = data_symmetrized[data.signal].nxdata
331
436
  print('\nSymmetrized ' + self.plane1 + ' planes.')
332
437
 
333
438
  print('Symmetrizing ' + self.plane2 + ' planes...')
334
- for j in range(0, len(q2)):
335
- print('Symmetrizing ' + q2.nxname + '=' + "{:.02f}".format(q2[j]) + "...", end='\r')
439
+ for j, value in enumerate(q2):
440
+ print(f'Symmetrizing {q2.nxname}={value:.02f}...', end='\r')
336
441
  data_symmetrized = self.plane2symmetrizer.symmetrize_2d(
337
- NXdata(NXfield(out_array[:, j, :], name=data.signal),
338
- (q1, q3)))
442
+ NXdata(NXfield(out_array[:, j, :], name=data.signal), (q1, q3))
443
+ )
339
444
  out_array[:, j, :] = data_symmetrized[data.signal].nxdata
340
445
  print('\nSymmetrized ' + self.plane2 + ' planes.')
341
446
 
342
447
  print('Symmetrizing ' + self.plane3 + ' planes...')
343
- for i in range(0, len(q1)):
344
- print('Symmetrizing ' + q1.nxname + '=' + "{:.02f}".format(q1[i]) + "...", end='\r')
448
+ for i, value in enumerate(q1):
449
+ print(f'Symmetrizing {q1.nxname}={value:.02f}...', end='\r')
345
450
  data_symmetrized = self.plane3symmetrizer.symmetrize_2d(
346
- NXdata(NXfield(out_array[i, :, :], name=data.signal),
347
- (q2, q3)))
451
+ NXdata(NXfield(out_array[i, :, :], name=data.signal), (q2, q3))
452
+ )
348
453
  out_array[i, :, :] = data_symmetrized[data.signal].nxdata
349
454
  print('\nSymmetrized ' + self.plane3 + ' planes.')
350
455
 
351
456
  out_array[out_array < 0] = 0
352
457
 
353
458
  stoptime = time.time()
354
- print("\nSymmetriztaion finished in " + "{:.02f}".format((stoptime - starttime) / 60) + " minutes.")
459
+ print(f"\nSymmetrization finished in {((stoptime - starttime) / 60):.02f} minutes.")
355
460
 
356
- self.symmetrized = NXdata(NXfield(out_array, name=data.signal), tuple([data[axis] for axis in data.axes]))
461
+ self.symmetrized = NXdata(NXfield(out_array, name=data.signal),
462
+ tuple(data[axis] for axis in data.axes))
357
463
 
358
464
  return self.symmetrized
359
465
 
360
466
  def save(self, fout_name=None):
361
467
  """
362
- Save the symmetrized dataset to a file.
468
+ Save the symmetrized dataset to a NeXus file.
363
469
 
364
470
  Parameters
365
471
  ----------
366
472
  fout_name : str, optional
367
- The name of the output file. If not provided, the default name 'symmetrized.nxs' will be used.
368
-
473
+ The name of the output file. If not provided,
474
+ the default name 'symmetrized.nxs' will be used.
369
475
  """
370
476
  print("Saving file...")
371
477
 
@@ -382,29 +488,34 @@ def generate_gaussian(H, K, L, amp, stddev, lattice_params, coeffs=None):
382
488
  """
383
489
  Generate a 3D Gaussian distribution.
384
490
 
491
+ This function creates a 3D Gaussian distribution in reciprocal space based
492
+ on the specified parameters.
493
+
385
494
  Parameters
386
495
  ----------
387
496
  H, K, L : ndarray
388
- Arrays specifying the values of H, K, and L coordinates.
497
+ Arrays specifying the values of H, K, and L coordinates in reciprocal space.
389
498
  amp : float
390
499
  Amplitude of the Gaussian distribution.
391
500
  stddev : float
392
501
  Standard deviation of the Gaussian distribution.
393
502
  lattice_params : tuple
394
- Tuple of lattice parameters (a, b, c, alpha, beta, gamma).
503
+ Tuple of lattice parameters (a, b, c, alpha, beta, gamma) for the
504
+ reciprocal lattice.
395
505
  coeffs : list, optional
396
- Coefficients for the Gaussian expression, including cross-terms between axes. Default is [1, 0, 1, 0, 1, 0],
397
- corresponding to (1*H**2 + 0*H*K + 1*K**2 + 0*K*L + 1*L**2 + 0*L*H)
506
+ Coefficients for the Gaussian expression, including cross-terms between axes.
507
+ Default is [1, 0, 1, 0, 1, 0],
508
+ corresponding to (1*H**2 + 0*H*K + 1*K**2 + 0*K*L + 1*L**2 + 0*L*H).
398
509
 
399
510
  Returns
400
511
  -------
401
512
  gaussian : ndarray
402
- 3D Gaussian distribution.
513
+ 3D Gaussian distribution array.
403
514
  """
404
515
  if coeffs is None:
405
516
  coeffs = [1, 0, 1, 0, 1, 0]
406
517
  a, b, c, al, be, ga = lattice_params
407
- a_, b_, c_, al_, be_, ga_ = reciprocal_lattice_params((a, b, c, al, be, ga))
518
+ a_, b_, c_, _, _, _ = reciprocal_lattice_params((a, b, c, al, be, ga))
408
519
  H, K, L = np.meshgrid(H, K, L, indexing='ij')
409
520
  gaussian = amp * np.exp(-(coeffs[0] * H ** 2 +
410
521
  coeffs[1] * (b_ * a_ / (a_ ** 2)) * H * K +
@@ -420,10 +531,88 @@ def generate_gaussian(H, K, L, amp, stddev, lattice_params, coeffs=None):
420
531
 
421
532
 
422
533
  class Puncher:
534
+ """
535
+ A class for applying masks to 3D datasets, typically for data processing in reciprocal space.
536
+
537
+ This class provides methods for setting data, applying masks, and generating
538
+ masks based on various criteria. It can be used to "punch" or modify datasets
539
+ by setting specific regions to NaN according to the mask.
540
+
541
+ Attributes
542
+ ----------
543
+ punched : NXdata, optional
544
+ The dataset with regions modified (punched) based on the mask.
545
+ data : NXdata, optional
546
+ The input dataset to be processed.
547
+ HH, KK, LL : ndarray
548
+ Meshgrid arrays representing the H, K, and L coordinates in reciprocal space.
549
+ mask : ndarray, optional
550
+ The mask used for identifying and modifying specific regions in the dataset.
551
+ reciprocal_lattice_params : tuple, optional
552
+ The reciprocal lattice parameters derived from the lattice parameters.
553
+ lattice_params : tuple, optional
554
+ The lattice parameters (a, b, c, alpha, beta, gamma).
555
+ a, b, c, al, be, ga : float
556
+ Individual lattice parameters.
557
+ a_star, b_star, c_star, al_star, be_star, ga_star : float
558
+ Individual reciprocal lattice parameters.
559
+
560
+ Methods
561
+ -------
562
+ set_data(data)
563
+ Sets the dataset to be processed and initializes the coordinate arrays and mask.
564
+ set_lattice_params(lattice_params)
565
+ Sets the lattice parameters and computes the reciprocal lattice parameters.
566
+ add_mask(maskaddition)
567
+ Adds regions to the current mask using a logical OR operation.
568
+ subtract_mask(masksubtraction)
569
+ Removes regions from the current mask using a logical AND NOT operation.
570
+ generate_bragg_mask(punch_radius, coeffs=None, thresh=None)
571
+ Generates a mask for Bragg peaks based on a Gaussian distribution in
572
+ reciprocal space.
573
+ generate_intensity_mask(thresh, radius, verbose=True)
574
+ Generates a mask based on intensity thresholds, including a spherical
575
+ region around high-intensity points.
576
+ generate_mask_at_coord(coordinate, punch_radius, coeffs=None, thresh=None)
577
+ Generates a mask centered at a specific coordinate in reciprocal space
578
+ with a specified radius.
579
+ punch()
580
+ Applies the mask to the dataset, setting masked regions to NaN.
581
+ """
582
+
423
583
  def __init__(self):
584
+ """
585
+ Initialize the Puncher object.
586
+
587
+ This method sets up the initial state of the Puncher instance, including
588
+ attributes for storing the dataset, lattice parameters, and masks.
589
+ It prepares the object for further data processing and masking operations.
590
+
591
+ Attributes
592
+ ----------
593
+ punched : NXdata, optional
594
+ The dataset with modified (punched) regions, initialized as None.
595
+ data : NXdata, optional
596
+ The input dataset to be processed, initialized as None.
597
+ HH, KK, LL : ndarray, optional
598
+ Arrays representing the H, K, and L coordinates in reciprocal space,
599
+ initialized as None.
600
+ mask : ndarray, optional
601
+ The mask for identifying and modifying specific regions in the dataset,
602
+ initialized as None.
603
+ reciprocal_lattice_params : tuple, optional
604
+ The reciprocal lattice parameters, initialized as None.
605
+ lattice_params : tuple, optional
606
+ The lattice parameters (a, b, c, alpha, beta, gamma),
607
+ initialized as None.
608
+ a, b, c, al, be, ga : float
609
+ Individual lattice parameters, initialized as None.
610
+ a_star, b_star, c_star, al_star, be_star, ga_star : float
611
+ Individual reciprocal lattice parameters, initialized as None.
612
+ """
424
613
  self.punched = None
425
614
  self.data = None
426
- self.H, self.K, self.L = [None] * 3
615
+ self.HH, self.KK, self.LL = [None] * 3
427
616
  self.mask = None
428
617
  self.reciprocal_lattice_params = None
429
618
  self.lattice_params = None
@@ -431,55 +620,178 @@ class Puncher:
431
620
  self.a_star, self.b_star, self.c_star, self.al_star, self.be_star, self.ga_star = [None] * 6
432
621
 
433
622
  def set_data(self, data):
623
+ """
624
+ Set the 3D dataset and initialize the mask if not already set.
625
+
626
+ Parameters
627
+ ----------
628
+ data : NXdata
629
+ The dataset to be processed.
630
+
631
+ Notes
632
+ -----
633
+ This method also sets up the H, K, and L coordinate grids for the dataset.
634
+ """
434
635
  self.data = data
435
636
  if self.mask is None:
436
637
  self.mask = np.zeros(data[data.signal].nxdata.shape)
437
- self.H, self.K, self.L = np.meshgrid(data[data.axes[0]], data[data.axes[1]], data[data.axes[2]], indexing='ij')
638
+ self.HH, self.KK, self.LL = np.meshgrid(data[data.axes[0]],
639
+ data[data.axes[1]],
640
+ data[data.axes[2]],
641
+ indexing='ij')
438
642
 
439
643
  def set_lattice_params(self, lattice_params):
644
+ """
645
+ Set the lattice parameters and compute the reciprocal lattice parameters.
646
+
647
+ Parameters
648
+ ----------
649
+ lattice_params : tuple
650
+ Tuple of lattice parameters (a, b, c, alpha, beta, gamma).
651
+ """
440
652
  self.a, self.b, self.c, self.al, self.be, self.ga = lattice_params
441
653
  self.lattice_params = lattice_params
442
654
  self.reciprocal_lattice_params = reciprocal_lattice_params(lattice_params)
443
- self.a_star, self.b_star, self.c_star, self.al_star, self.be_star, self.ga_star = self.reciprocal_lattice_params
655
+ self.a_star, self.b_star, self.c_star, \
656
+ self.al_star, self.be_star, self.ga_star = self.reciprocal_lattice_params
444
657
 
445
658
  def add_mask(self, maskaddition):
659
+ """
660
+ Add regions to the current mask using a logical OR operation.
661
+
662
+ Parameters
663
+ ----------
664
+ maskaddition : ndarray
665
+ The mask to be added.
666
+ """
446
667
  self.mask = np.logical_or(self.mask, maskaddition)
447
668
 
448
669
  def subtract_mask(self, masksubtraction):
670
+ """
671
+ Remove regions from the current mask using a logical AND NOT operation.
672
+
673
+ Parameters
674
+ ----------
675
+ masksubtraction : ndarray
676
+ The mask to be subtracted.
677
+ """
449
678
  self.mask = np.logical_and(self.mask, np.logical_not(masksubtraction))
450
679
 
451
680
  def generate_bragg_mask(self, punch_radius, coeffs=None, thresh=None):
681
+ """
682
+ Generate a mask for Bragg peaks.
683
+
684
+ Parameters
685
+ ----------
686
+ punch_radius : float
687
+ Radius for the Bragg peak mask.
688
+ coeffs : list, optional
689
+ Coefficients for the expression of the sphere to be removed around
690
+ each Bragg position, corresponding to coefficients for H, HK, K, KL, L, and LH terms.
691
+ Default is [1, 0, 1, 0, 1, 0].
692
+ thresh : float, optional
693
+ Intensity threshold for applying the mask.
694
+
695
+ Returns
696
+ -------
697
+ mask : ndarray
698
+ Boolean mask identifying the Bragg peaks.
699
+ """
452
700
  if coeffs is None:
453
701
  coeffs = [1, 0, 1, 0, 1, 0]
454
702
  data = self.data
455
- H, K, L = self.H, self.K, self.L
456
- a_, b_, c_, al_, be_, ga_ = self.reciprocal_lattice_params
703
+ H, K, L = self.HH, self.KK, self.LL
704
+ a_, b_, c_, _, _, _ = self.reciprocal_lattice_params
457
705
 
458
706
  mask = (coeffs[0] * (H - np.rint(H)) ** 2 +
459
707
  coeffs[1] * (b_ * a_ / (a_ ** 2)) * (H - np.rint(H)) * (K - np.rint(K)) +
460
708
  coeffs[2] * (b_ / a_) ** 2 * (K - np.rint(K)) ** 2 +
461
709
  coeffs[3] * (b_ * c_ / (a_ ** 2)) * (K - np.rint(K)) * (L - np.rint(L)) +
462
710
  coeffs[4] * (c_ / a_) ** 2 * (L - np.rint(L)) ** 2 +
463
- coeffs[5] * (c_ * a_ / (a_ ** 2)) * (L - np.rint(L)) * (H - np.rint(H))) < punch_radius ** 2
711
+ coeffs[5] * (c_ * a_ / (a_ ** 2)) * (L - np.rint(L)) * (H - np.rint(H))) \
712
+ < punch_radius ** 2
464
713
 
465
714
  if thresh:
466
715
  mask = np.logical_and(mask, data[data.signal] > thresh)
467
716
 
468
717
  return mask
469
718
 
719
+ def generate_intensity_mask(self, thresh, radius, verbose=True):
720
+ """
721
+ Generate a mask based on intensity thresholds.
722
+
723
+ Parameters
724
+ ----------
725
+ thresh : float
726
+ Intensity threshold for creating the mask.
727
+ radius : int
728
+ Radius around high-intensity points to include in the mask.
729
+ verbose : bool, optional
730
+ Whether to print progress information.
731
+
732
+ Returns
733
+ -------
734
+ mask : ndarray
735
+ Boolean mask highlighting regions with high intensity.
736
+ """
737
+ data = self.data
738
+ counts = data[data.signal].nxdata
739
+ mask = np.zeros(counts.shape)
740
+
741
+ print(f"Shape of data is {counts.shape}") if verbose else None
742
+ for i in range(counts.shape[0]):
743
+ for j in range(counts.shape[1]):
744
+ for k in range(counts.shape[2]):
745
+ if counts[i, j, k] > thresh:
746
+ # Set the pixels within the sphere to NaN
747
+ for x in range(max(i - radius, 0),
748
+ min(i + radius + 1, counts.shape[0])):
749
+ for y in range(max(j - radius, 0),
750
+ min(j + radius + 1, counts.shape[1])):
751
+ for z in range(max(k - radius, 0),
752
+ min(k + radius + 1, counts.shape[2])):
753
+ mask[x, y, z] = 1
754
+ print(f"Found high intensity at ({i}, {j}, {k}).\t\t", end='\r') \
755
+ if verbose else None
756
+ print("\nDone.")
757
+ return mask
758
+
470
759
  def generate_mask_at_coord(self, coordinate, punch_radius, coeffs=None, thresh=None):
760
+ """
761
+ Generate a mask centered at a specific coordinate.
762
+
763
+ Parameters
764
+ ----------
765
+ coordinate : tuple of float
766
+ Center coordinate (H, K, L) for the mask.
767
+ punch_radius : float
768
+ Radius for the mask.
769
+ coeffs : list, optional
770
+ Coefficients for the expression of the sphere to be removed around
771
+ each Bragg position,
772
+ corresponding to coefficients for H, HK, K, KL, L, and LH terms.
773
+ Default is [1, 0, 1, 0, 1, 0].
774
+ thresh : float, optional
775
+ Intensity threshold for applying the mask.
776
+
777
+ Returns
778
+ -------
779
+ mask : ndarray
780
+ Boolean mask for the specified coordinate.
781
+ """
471
782
  if coeffs is None:
472
783
  coeffs = [1, 0, 1, 0, 1, 0]
473
784
  data = self.data
474
- H, K, L = self.H, self.K, self.L
475
- a_, b_, c_, al_, be_, ga_ = self.reciprocal_lattice_params
785
+ H, K, L = self.HH, self.KK, self.LL
786
+ a_, b_, c_, _, _, _ = self.reciprocal_lattice_params
476
787
  centerH, centerK, centerL = coordinate
477
788
  mask = (coeffs[0] * (H - centerH) ** 2 +
478
789
  coeffs[1] * (b_ * a_ / (a_ ** 2)) * (H - centerH) * (K - centerK) +
479
790
  coeffs[2] * (b_ / a_) ** 2 * (K - centerK) ** 2 +
480
791
  coeffs[3] * (b_ * c_ / (a_ ** 2)) * (K - centerK) * (L - centerL) +
481
792
  coeffs[4] * (c_ / a_) ** 2 * (L - centerL) ** 2 +
482
- coeffs[5] * (c_ * a_ / (a_ ** 2)) * (L - centerL) * (H - centerH)) < punch_radius ** 2
793
+ coeffs[5] * (c_ * a_ / (a_ ** 2)) * (L - centerL) * (H - centerH)) \
794
+ < punch_radius ** 2
483
795
 
484
796
  if thresh:
485
797
  mask = np.logical_and(mask, data[data.signal] > thresh)
@@ -487,99 +799,743 @@ class Puncher:
487
799
  return mask
488
800
 
489
801
  def punch(self):
490
- data= self.data
491
- self.punched = NXdata(NXfield(np.where(self.mask, np.nan, data[data.signal].nxdata), name=data.signal),
492
- (data[data.axes[0]],data[data.axes[1]],data[data.axes[2]]))
493
- return self.punched
802
+ """
803
+ Apply the mask to the dataset, setting masked regions to NaN.
804
+
805
+ This method creates a new dataset where the masked regions are set to
806
+ NaN, effectively "punching" those regions.
807
+ """
808
+ data = self.data
809
+ self.punched = NXdata(NXfield(
810
+ np.where(self.mask, np.nan, data[data.signal].nxdata),
811
+ name=data.signal),
812
+ (data[data.axes[0]], data[data.axes[1]], data[data.axes[2]])
813
+ )
814
+
815
+
816
+ def _round_up_to_odd_integer(value):
817
+ """
818
+ Round up a given number to the nearest odd integer.
819
+
820
+ This function takes a floating-point value and rounds it up to the smallest
821
+ odd integer that is greater than or equal to the given value.
822
+
823
+ Parameters
824
+ ----------
825
+ value : float
826
+ The input floating-point number to be rounded up.
827
+
828
+ Returns
829
+ -------
830
+ int
831
+ The nearest odd integer greater than or equal to the input value.
832
+
833
+ Examples
834
+ --------
835
+ >>> _round_up_to_odd_integer(4.2)
836
+ 5
837
+
838
+ >>> _round_up_to_odd_integer(5.0)
839
+ 5
840
+
841
+ >>> _round_up_to_odd_integer(6.7)
842
+ 7
843
+ """
844
+ i = int(math.ceil(value))
845
+ if i % 2 == 0:
846
+ return i + 1
847
+
848
+ return i
849
+
850
+
851
+ class Gaussian3DKernel(Kernel):
852
+ """
853
+ Initialize a 3D Gaussian kernel.
854
+
855
+ This constructor creates a 3D Gaussian kernel with the specified
856
+ standard deviation and size. The Gaussian kernel is generated based on
857
+ the provided coefficients and is then normalized.
858
+
859
+ Parameters
860
+ ----------
861
+ stddev : float
862
+ The standard deviation of the Gaussian distribution, which controls
863
+ the width of the kernel.
864
+
865
+ size : tuple of int
866
+ The dimensions of the kernel, given as (x_dim, y_dim, z_dim).
867
+
868
+ coeffs : list of float, optional
869
+ Coefficients for the Gaussian expression.
870
+ The default is [1, 0, 1, 0, 1, 0], corresponding to the Gaussian form:
871
+ (1 * X^2 + 0 * X * Y + 1 * Y^2 + 0 * Y * Z + 1 * Z^2 + 0 * Z * X).
872
+
873
+ **kwargs : keyword arguments
874
+ Additional keyword arguments passed to the parent class constructor.
875
+
876
+ Raises
877
+ ------
878
+ ValueError
879
+ If the dimensions in `size` are not positive integers.
880
+
881
+ Notes
882
+ -----
883
+ The kernel is generated over a grid that spans twice the size of
884
+ each dimension, and the resulting array is normalized.
885
+ """
886
+ _separable = True
887
+ _is_bool = False
888
+
889
+ def __init__(self, stddev, size, coeffs=None):
890
+ if not coeffs:
891
+ coeffs = [1, 0, 1, 0, 1, 0]
892
+ x_dim, y_dim, z_dim = size
893
+ x = np.linspace(-x_dim, x_dim, int(x_dim) + 1)
894
+ y = np.linspace(-y_dim, y_dim, int(y_dim) + 1)
895
+ z = np.linspacein(-z_dim, z_dim, int(z_dim) + 1)
896
+ X, Y, Z = np.meshgrid(x, y, z)
897
+ array = np.exp(-(coeffs[0] * X ** 2 +
898
+ coeffs[1] * X * Y +
899
+ coeffs[2] * Y ** 2 +
900
+ coeffs[3] * Y * Z +
901
+ coeffs[4] * Z ** 2 +
902
+ coeffs[5] * Z * X) / (2 * stddev ** 2)
903
+ )
904
+ self._default_size = _round_up_to_odd_integer(stddev)
905
+ super().__init__(array)
906
+ self.normalize()
907
+ self._truncation = np.abs(1. - self._array.sum())
908
+
909
+
910
+ class Interpolator():
911
+ """
912
+ A class to perform data interpolation using convolution with a specified
913
+ kernel.
914
+
915
+ Attributes
916
+ ----------
917
+ interp_time : float or None
918
+ Time taken for the last interpolation operation. Defaults to None.
919
+
920
+ window : ndarray or None
921
+ Window function to be applied to the interpolated data. Defaults to None.
922
+
923
+ interpolated : ndarray or None
924
+ The result of the interpolation operation. Defaults to None.
925
+
926
+ data : NXdata or None
927
+ The dataset to be interpolated. Defaults to None.
928
+
929
+ kernel : ndarray or None
930
+ The kernel used for convolution during interpolation. Defaults to None.
931
+
932
+ tapered : ndarray or None
933
+ The interpolated data after applying the window function. Defaults to None.
934
+ """
494
935
 
495
- class PuncherHK:
496
936
  def __init__(self):
497
- self.punched = None
937
+ """
938
+ Initialize an Interpolator object.
939
+
940
+ Sets up an instance of the Interpolator class with the
941
+ following attributes initialized to None:
942
+
943
+ - interp_time
944
+ - window
945
+ - interpolated
946
+ - data
947
+ - kernel
948
+ - tapered
949
+
950
+ """
951
+ self.interp_time = None
952
+ self.window = None
953
+ self.interpolated = None
498
954
  self.data = None
499
- self.q1, self.q2 = [None] * 2
500
- self.mask = None
955
+ self.kernel = None
956
+ self.tapered = None
957
+
958
+ def set_data(self, data):
959
+ """
960
+ Set the dataset to be interpolated.
961
+
962
+ Parameters
963
+ ----------
964
+ data : NXdata
965
+ The dataset containing the data to be interpolated.
966
+ """
967
+ self.data = data
968
+
969
+ def set_kernel(self, kernel):
970
+ """
971
+ Set the kernel to be used for interpolation.
972
+
973
+ Parameters
974
+ ----------
975
+ kernel : ndarray
976
+ The kernel to be used for convolution during interpolation.
977
+ """
978
+ self.kernel = kernel
979
+
980
+ def interpolate(self):
981
+ """
982
+ Perform interpolation on the dataset using the specified kernel.
983
+
984
+ The interpolation is done by convolving the data with the kernel
985
+ using the `convolve_fft` function. Updates the `interpolated`
986
+ attribute with the result.
987
+
988
+ Prints the time taken for the interpolation process.
989
+
990
+ Returns
991
+ -------
992
+ None
993
+ """
994
+ start = time.time()
995
+
996
+ if self.interp_time:
997
+ print(f"Last interpolation took {self.interp_time / 60:.2f} minutes.")
998
+
999
+
1000
+ print("Running interpolation...")
1001
+ result = np.real(
1002
+ convolve_fft(self.data[self.data.signal].nxdata,
1003
+ self.kernel, allow_huge=True, return_fft=False))
1004
+ print("Interpolation finished.")
1005
+
1006
+ end = time.time()
1007
+ interp_time = end - start
1008
+
1009
+ print(f'Interpolation took {interp_time / 60:.2f} minutes.')
1010
+
1011
+ result[result < 0] = 0
1012
+ self.interpolated = array_to_nxdata(result, self.data)
1013
+
1014
+ def set_tukey_window(self, tukey_alphas=(1.0, 1.0, 1.0)):
1015
+ """
1016
+ Set a Tukey window function for data tapering.
1017
+
1018
+ Parameters
1019
+ ----------
1020
+ tukey_alphas : tuple of floats, optional
1021
+ The alpha parameters for the Tukey window in each
1022
+ dimension (H, K, L). Default is (1.0, 1.0, 1.0).
1023
+
1024
+ Notes
1025
+ -----
1026
+ The window function is generated based on the size of the dataset in each dimension.
1027
+ """
1028
+ data = self.data
1029
+ tukey_H = np.tile(
1030
+ scipy.signal.tukey(len(data[data.axes[0]]), alpha=tukey_alphas[0])[:, None, None],
1031
+ (1, len(data[data.axes[1]]), len(data[data.axes[2]]))
1032
+ )
1033
+ tukey_K = np.tile(
1034
+ scipy.signal.tukey(len(data[data.axes[1]]), alpha=tukey_alphas[1])[None, :, None],
1035
+ (len(data[data.axes[0]]), 1, len(data[data.axes[2]]))
1036
+ )
1037
+ window = tukey_H * tukey_K
1038
+
1039
+ del tukey_H, tukey_K
1040
+ gc.collect()
1041
+
1042
+ tukey_L = np.tile(
1043
+ scipy.signal.tukey(len(data[data.axes[2]]), alpha=tukey_alphas[2])[None, None, :],
1044
+ (len(data[data.axes[0]]), len(data[data.axes[1]]), 1))
1045
+ window = window * tukey_L
1046
+
1047
+ self.window = window
1048
+
1049
+ def set_hexagonal_tukey_window(self, tukey_alphas=(1.0, 1.0, 1.0, 1.0)):
1050
+ """
1051
+ Set a hexagonal Tukey window function for data tapering.
1052
+
1053
+ Parameters
1054
+ ----------
1055
+ tukey_alphas : tuple of floats, optional
1056
+ The alpha parameters for the Tukey window in each dimension and
1057
+ for the hexagonal truncation (H, HK, K, L).
1058
+ Default is (1.0, 1.0, 1.0, 1.0).
1059
+
1060
+ Notes
1061
+ -----
1062
+ The hexagonal Tukey window is applied to the dataset in a manner that
1063
+ preserves hexagonal symmetry.
1064
+
1065
+ """
1066
+ data = self.data
1067
+ H_ = data[data.axes[0]]
1068
+ K_ = data[data.axes[1]]
1069
+ L_ = data[data.axes[2]]
1070
+
1071
+ tukey_H = np.tile(
1072
+ scipy.signal.tukey(len(data[data.axes[0]]), alpha=tukey_alphas[0])[:, None, None],
1073
+ (1, len(data[data.axes[1]]), len(data[data.axes[2]]))
1074
+ )
1075
+ tukey_K = np.tile(
1076
+ scipy.signal.tukey(len(data[data.axes[1]]), alpha=tukey_alphas[1])[None, :, None],
1077
+ (len(data[data.axes[0]]), 1, len(data[data.axes[2]]))
1078
+ )
1079
+ window = tukey_H * tukey_K
1080
+
1081
+ del tukey_H, tukey_K
1082
+ gc.collect()
1083
+
1084
+ truncation = int((len(H_) - int(len(H_) * np.sqrt(2) / 2)) / 2)
1085
+
1086
+ tukey_HK = scipy.ndimage.rotate(
1087
+ np.tile(
1088
+ np.concatenate(
1089
+ (np.zeros(truncation)[:, None, None],
1090
+ scipy.signal.tukey(len(H_) - 2 * truncation,
1091
+ alpha=tukey_alphas[2])[:, None, None],
1092
+ np.zeros(truncation)[:, None, None])),
1093
+ (1, len(K_), len(L_))
1094
+ ),
1095
+ angle=45, reshape=False, mode='nearest',
1096
+ )[0:len(H_), 0:len(K_), :]
1097
+ tukey_HK = np.nan_to_num(tukey_HK)
1098
+ window = window * tukey_HK
1099
+
1100
+ del tukey_HK
1101
+ gc.collect()
1102
+
1103
+ tukey_L = np.tile(
1104
+ scipy.signal.tukey(len(data[data.axes[2]]), alpha=tukey_alphas[3])[None, None, :],
1105
+ (len(data[data.axes[0]]), len(data[data.axes[1]]), 1)
1106
+ )
1107
+ window = window * tukey_L
1108
+
1109
+ del tukey_L
1110
+ gc.collect()
1111
+
1112
+ self.window = window
1113
+
1114
+ def set_window(self, window):
1115
+ """
1116
+ Set a custom window function for data tapering.
1117
+
1118
+ Parameters
1119
+ ----------
1120
+ window : ndarray
1121
+ A custom window function to be applied to the interpolated data.
1122
+ """
1123
+ self.window = window
1124
+
1125
+ def apply_window(self):
1126
+ """
1127
+ Apply the window function to the interpolated data.
1128
+
1129
+ The window function, if set, is applied to the `interpolated` data
1130
+ to produce the `tapered` result.
1131
+
1132
+ Returns
1133
+ -------
1134
+ None
1135
+ """
1136
+ self.tapered = self.interpolated * self.window
1137
+
1138
+
1139
+ def fourier_transform_nxdata(data):
1140
+ """
1141
+ Perform a 3D Fourier Transform on the given NXdata object.
1142
+
1143
+ This function applies an inverse Fourier Transform to the input data
1144
+ using the `pyfftw` library to optimize performance. The result is a
1145
+ transformed array with spatial frequency components calculated along
1146
+ each axis.
1147
+
1148
+ Parameters
1149
+ ----------
1150
+ data : NXdata
1151
+ An NXdata object containing the data to be transformed. It should
1152
+ include the `signal` field for the data and `axes` fields
1153
+ specifying the coordinate axes.
1154
+
1155
+ Returns
1156
+ -------
1157
+ NXdata
1158
+ A new NXdata object containing the Fourier Transformed data. The
1159
+ result includes:
1160
+
1161
+ - `dPDF`: The transformed data array.
1162
+ - `x`, `y`, `z`: Arrays representing the frequency components along each axis.
1163
+
1164
+ Notes
1165
+ -----
1166
+ - The FFT is performed in two stages: first along the last dimension of
1167
+ the input array and then along the first two dimensions.
1168
+ - The function uses `pyfftw` for efficient computation of the Fourier
1169
+ Transform.
1170
+ - The output frequency components are computed based on the step sizes
1171
+ of the original data axes.
1172
+
1173
+ """
1174
+ start = time.time()
1175
+ print("Starting FFT.")
1176
+
1177
+ padded = data[data.signal].nxdata
1178
+
1179
+ fft_array = np.zeros(padded.shape)
1180
+ print("FFT on axes 1,2")
1181
+ for k in range(0, padded.shape[2]):
1182
+ fft_array[:, :, k] = np.real(
1183
+ np.fft.fftshift(
1184
+ pyfftw.interfaces.numpy_fft.ifftn(np.fft.fftshift(padded[:, :, k]),
1185
+ planner_effort='FFTW_MEASURE'))
1186
+ )
1187
+ print(f'k={k} ', end='\r')
1188
+
1189
+ print("FFT on axis 3")
1190
+ for i in range(0, padded.shape[0]):
1191
+ for j in range(0, padded.shape[1]):
1192
+ f_slice = fft_array[i, j, :]
1193
+ print(f'i={i} ', end='\r')
1194
+ fft_array[i, j, :] = np.real(
1195
+ np.fft.fftshift(
1196
+ pyfftw.interfaces.numpy_fft.ifftn(np.fft.fftshift(f_slice),
1197
+ planner_effort='FFTW_MEASURE')
1198
+ )
1199
+ )
1200
+
1201
+ end = time.time()
1202
+ print("FFT complete.")
1203
+ print('FFT took ' + str(end - start) + ' seconds.')
1204
+
1205
+ H_step = data[data.axes[0]].nxdata[1] - data[data.axes[0]].nxdata[0]
1206
+ K_step = data[data.axes[1]].nxdata[1] - data[data.axes[1]].nxdata[0]
1207
+ L_step = data[data.axes[2]].nxdata[1] - data[data.axes[2]].nxdata[0]
1208
+
1209
+ fft = NXdata(NXfield(fft_array, name='dPDF'),
1210
+ (NXfield(np.linspace(-0.5 / H_step, 0.5 / H_step, padded.shape[0]), name='x'),
1211
+ NXfield(np.linspace(-0.5 / K_step, 0.5 / K_step, padded.shape[1]), name='y'),
1212
+ NXfield(np.linspace(-0.5 / L_step, 0.5 / L_step, padded.shape[1]), name='z'),
1213
+ )
1214
+ )
1215
+ return fft
1216
+
1217
+
1218
+ class DeltaPDF:
1219
+ """
1220
+ A class for processing and analyzing 3D diffraction data using various\
1221
+ operations, including masking, interpolation, padding, and Fourier
1222
+ transformation.
1223
+
1224
+ Attributes
1225
+ ----------
1226
+ fft : NXdata or None
1227
+ The Fourier transformed data.
1228
+ data : NXdata or None
1229
+ The input diffraction data.
1230
+ lattice_params : tuple or None
1231
+ Lattice parameters (a, b, c, al, be, ga).
1232
+ reciprocal_lattice_params : tuple or None
1233
+ Reciprocal lattice parameters (a*, b*, c*, al*, be*, ga*).
1234
+ puncher : Puncher
1235
+ An instance of the Puncher class for generating masks and punching
1236
+ the data.
1237
+ interpolator : Interpolator
1238
+ An instance of the Interpolator class for interpolating and applying
1239
+ windows to the data.
1240
+ padder : Padder
1241
+ An instance of the Padder class for padding the data.
1242
+ mask : ndarray or None
1243
+ The mask used for data processing.
1244
+ kernel : Kernel or None
1245
+ The kernel used for interpolation.
1246
+ window : ndarray or None
1247
+ The window applied to the interpolated data.
1248
+ padded : ndarray or None
1249
+ The padded data.
1250
+ tapered : ndarray or None
1251
+ The data after applying the window.
1252
+ interpolated : NXdata or None
1253
+ The interpolated data.
1254
+ punched : NXdata or None
1255
+ The punched data.
1256
+
1257
+ """
1258
+
1259
+ def __init__(self):
1260
+ """
1261
+ Initialize a DeltaPDF object with default attributes.
1262
+ """
501
1263
  self.reciprocal_lattice_params = None
1264
+ self.fft = None
1265
+ self.data = None
502
1266
  self.lattice_params = None
503
- self.a, self.b, self.c, self.al, self.be, self.ga = [None] * 6
504
- self.a_star, self.b_star, self.c_star, self.al_star, self.be_star, self.ga_star = [None] * 6
1267
+ self.puncher = Puncher()
1268
+ self.interpolator = Interpolator()
1269
+ self.padder = Padder()
1270
+ self.mask = None
1271
+ self.kernel = None
1272
+ self.window = None
1273
+ self.padded = None
1274
+ self.tapered = None
1275
+ self.interpolated = None
1276
+ self.punched = None
505
1277
 
506
1278
  def set_data(self, data):
1279
+ """
1280
+ Set the input diffraction data and update the Puncher and Interpolator
1281
+ with the data.
1282
+
1283
+ Parameters
1284
+ ----------
1285
+ data : NXdata
1286
+ The diffraction data to be processed.
1287
+ """
507
1288
  self.data = data
508
- if self.mask is None:
509
- self.mask = np.zeros(data[data.signal].nxdata.shape)
510
- self.q1, self.q2, = np.meshgrid(data[data.axes[0]], data[data.axes[1]], indexing='ij')
1289
+ self.puncher.set_data(data)
1290
+ self.interpolator.set_data(data)
511
1291
 
512
1292
  def set_lattice_params(self, lattice_params):
513
- self.a, self.b, self.c, self.al, self.be, self.ga = lattice_params
1293
+ """
1294
+ Sets the lattice parameters and calculates the reciprocal lattice
1295
+ parameters.
1296
+
1297
+ Parameters
1298
+ ----------
1299
+ lattice_params : tuple of float
1300
+ The lattice parameters (a, b, c, alpha, beta, gamma) in real space.
1301
+ """
514
1302
  self.lattice_params = lattice_params
515
- self.reciprocal_lattice_params = reciprocal_lattice_params(lattice_params)
516
- self.a_star, self.b_star, self.c_star, self.al_star, self.be_star, self.ga_star = self.reciprocal_lattice_params
1303
+ self.puncher.set_lattice_params(lattice_params)
1304
+ self.reciprocal_lattice_params = self.puncher.reciprocal_lattice_params
517
1305
 
518
1306
  def add_mask(self, maskaddition):
519
- self.mask = np.logical_or(self.mask, maskaddition)
1307
+ """
1308
+ Add regions to the current mask using a logical OR operation.
1309
+
1310
+ Parameters
1311
+ ----------
1312
+ maskaddition : ndarray
1313
+ The mask to be added.
1314
+ """
1315
+ self.puncher.add_mask(maskaddition)
1316
+ self.mask = self.puncher.mask
520
1317
 
521
1318
  def subtract_mask(self, masksubtraction):
522
- self.mask = np.logical_and(self.mask, np.logical_not(masksubtraction))
1319
+ """
1320
+ Remove regions from the current mask using a logical AND NOT operation.
1321
+
1322
+ Parameters
1323
+ ----------
1324
+ masksubtraction : ndarray
1325
+ The mask to be subtracted.
1326
+ """
1327
+ self.puncher.subtract_mask(masksubtraction)
1328
+ self.mask = self.puncher.mask
523
1329
 
524
1330
  def generate_bragg_mask(self, punch_radius, coeffs=None, thresh=None):
525
- if coeffs is None:
526
- coeffs = [1, 0, 1]
527
- data = self.data
528
- q1, q2 = self.q1, self.q2
529
- a_, b_, c_, al_, be_, ga_ = self.reciprocal_lattice_params
1331
+ """
1332
+ Generate a mask for Bragg peaks.
530
1333
 
531
- mask = (coeffs[0] * (q1 - np.rint(q1)) ** 2 +
532
- coeffs[1] * (b_ * a_ / (a_ ** 2)) * (q1 - np.rint(q1)) * (q2 - np.rint(q2)) +
533
- coeffs[2] * (b_ / a_) ** 2 * (q2 - np.rint(q2)) ** 2) < punch_radius ** 2
1334
+ Parameters
1335
+ ----------
1336
+ punch_radius : float
1337
+ Radius for the Bragg peak mask.
1338
+ coeffs : list, optional
1339
+ Coefficients for the expression of the sphere to be removed
1340
+ around each Bragg position, corresponding to coefficients
1341
+ for H, HK, K, KL, L, and LH terms. Default is [1, 0, 1, 0, 1, 0].
1342
+ thresh : float, optional
1343
+ Intensity threshold for applying the mask.
534
1344
 
535
- if thresh:
536
- mask = np.logical_and(mask, data[data.signal] > thresh)
1345
+ Returns
1346
+ -------
1347
+ mask : ndarray
1348
+ Boolean mask identifying the Bragg peaks.
1349
+ """
1350
+ return self.puncher.generate_bragg_mask(punch_radius, coeffs, thresh)
537
1351
 
538
- return mask
1352
+ def generate_intensity_mask(self, thresh, radius, verbose=True):
1353
+ """
1354
+ Generate a mask based on intensity thresholds.
1355
+
1356
+ Parameters
1357
+ ----------
1358
+ thresh : float
1359
+ Intensity threshold for creating the mask.
1360
+ radius : int
1361
+ Radius around high-intensity points to include in the mask.
1362
+ verbose : bool, optional
1363
+ Whether to print progress information.
1364
+
1365
+ Returns
1366
+ -------
1367
+ mask : ndarray
1368
+ Boolean mask highlighting regions with high intensity.
1369
+ """
1370
+ return self.puncher.generate_intensity_mask(thresh, radius, verbose)
539
1371
 
540
1372
  def generate_mask_at_coord(self, coordinate, punch_radius, coeffs=None, thresh=None):
541
- if coeffs is None:
542
- coeffs = [1, 0, 1]
543
- data = self.data
544
- q1, q2 = self.q1, self.q2
545
- a_, b_, c_, al_, be_, ga_ = self.reciprocal_lattice_params
546
- centerH, centerK, centerL = coordinate
547
- mask = (coeffs[0] * (H - centerH) ** 2 +
548
- coeffs[1] * (b_ * a_ / (a_ ** 2)) * (H - centerH) * (K - centerK) +
549
- coeffs[2] * (b_ / a_) ** 2 * (K - centerK) ** 2 +
550
- coeffs[3] * (b_ * c_ / (a_ ** 2)) * (K - centerK) * (L - centerL) +
551
- coeffs[4] * (c_ / a_) ** 2 * (L - centerL) ** 2 +
552
- coeffs[5] * (c_ * a_ / (a_ ** 2)) * (L - centerL) * (H - centerH)) < punch_radius ** 2
1373
+ """
1374
+ Generate a mask centered at a specific coordinate.
553
1375
 
554
- if thresh:
555
- mask = np.logical_and(mask, data[data.signal] > thresh)
1376
+ Parameters
1377
+ ----------
1378
+ coordinate : tuple of float
1379
+ Center coordinate (H, K, L) for the mask.
1380
+ punch_radius : float
1381
+ Radius for the mask.
1382
+ coeffs : list, optional
1383
+ Coefficients for the expression of the sphere to be removed around
1384
+ each Bragg position, corresponding to coefficients for
1385
+ H, HK, K, KL, L, and LH terms. Default is [1, 0, 1, 0, 1, 0].
1386
+ thresh : float, optional
1387
+ Intensity threshold for applying the mask.
556
1388
 
557
- return mask
1389
+ Returns
1390
+ -------
1391
+ mask : ndarray
1392
+ Boolean mask for the specified coordinate.
1393
+ """
1394
+ return self.puncher.generate_mask_at_coord(coordinate, punch_radius, coeffs, thresh)
558
1395
 
559
1396
  def punch(self):
560
- data= self.data
561
- self.punched = NXdata(NXfield(np.where(self.mask, np.nan, data[data.signal].nxdata), name=data.signal),
562
- (data[data.axes[0]],data[data.axes[1]]))
563
- return self.punched
1397
+ """
1398
+ Apply the mask to the dataset, setting masked regions to NaN.
564
1399
 
1400
+ This method creates a new dataset where the masked regions are set to
1401
+ NaN, effectively "punching" those regions.
1402
+ """
1403
+ self.puncher.punch()
1404
+ self.punched = self.puncher.punched
1405
+ self.interpolator.set_data(self.punched)
565
1406
 
566
- # class Reducer():
567
- # pass
568
- #
569
- #
1407
+ def set_kernel(self, kernel):
1408
+ """
1409
+ Set the kernel to be used for interpolation.
570
1410
 
571
- class Interpolator():
572
- def __init__(self):
573
- self.data = None
1411
+ Parameters
1412
+ ----------
1413
+ kernel : ndarray
1414
+ The kernel to be used for convolution during interpolation.
1415
+ """
1416
+ self.interpolator.set_kernel(kernel)
1417
+ self.kernel = kernel
574
1418
 
575
- def set_data(self, data):
576
- self.data = data
1419
+ def interpolate(self):
1420
+ """
1421
+ Perform interpolation on the dataset using the specified kernel.
577
1422
 
578
- def set_kernel(self, kernel):
579
- self.kernel = kernel
1423
+ The interpolation is done by convolving the data with the kernel using
1424
+ the `convolve_fft` function. Updates the `interpolated` attribute with
1425
+ the result.
580
1426
 
581
- def generate_gaussian_kernel(self, amp, stddev, coeffs=None):
582
- pass
1427
+ Prints the time taken for the interpolation process.
583
1428
 
584
- def interpolate(self):
585
- pass
1429
+ Returns
1430
+ -------
1431
+ None
1432
+ """
1433
+ self.interpolator.interpolate()
1434
+ self.interpolated = self.interpolator.interpolated
1435
+
1436
+ def set_tukey_window(self, tukey_alphas=(1.0, 1.0, 1.0)):
1437
+ """
1438
+ Set a Tukey window function for data tapering.
1439
+
1440
+ Parameters
1441
+ ----------
1442
+ tukey_alphas : tuple of floats, optional
1443
+ The alpha parameters for the Tukey window in each dimension
1444
+ (H, K, L). Default is (1.0, 1.0, 1.0).
1445
+
1446
+ Notes
1447
+ -----
1448
+ The window function is generated based on the size of the dataset
1449
+ in each dimension.
1450
+ """
1451
+ self.interpolator.set_tukey_window(tukey_alphas)
1452
+ self.window = self.interpolator.window
1453
+
1454
+ def set_hexagonal_tukey_window(self, tukey_alphas=(1.0, 1.0, 1.0, 1.0)):
1455
+ """
1456
+ Set a hexagonal Tukey window function for data tapering.
1457
+
1458
+ Parameters
1459
+ ----------
1460
+ tukey_alphas : tuple of floats, optional
1461
+ The alpha parameters for the Tukey window in each dimension and
1462
+ for the hexagonal truncation (H, HK, K, L). Default is (1.0, 1.0, 1.0, 1.0).
1463
+
1464
+ Notes
1465
+ -----
1466
+ The hexagonal Tukey window is applied to the dataset in a manner that
1467
+ preserves hexagonal symmetry.
1468
+ """
1469
+ self.interpolator.set_hexagonal_tukey_window(tukey_alphas)
1470
+
1471
+ def set_window(self, window):
1472
+ """
1473
+ Set a custom window function for data tapering.
1474
+
1475
+ Parameters
1476
+ ----------
1477
+ window : ndarray
1478
+ A custom window function to be applied to the interpolated data.
1479
+ """
1480
+ self.interpolator.set_window(window)
1481
+
1482
+ def apply_window(self):
1483
+ """
1484
+ Apply the window function to the interpolated data.
1485
+
1486
+ The window function, if set, is applied to the `interpolated` data to
1487
+ produce the `tapered` result.
1488
+
1489
+ Returns
1490
+ -------
1491
+ None
1492
+ """
1493
+ self.interpolator.apply_window()
1494
+ self.tapered = self.interpolator.tapered
1495
+ self.padder.set_data(self.tapered)
1496
+
1497
+ def pad(self, padding):
1498
+ """
1499
+ Symmetrically pads the data with zero values.
1500
+
1501
+ Parameters
1502
+ ----------
1503
+ padding : tuple
1504
+ The number of zero-value pixels to add along each edge of the array.
1505
+
1506
+ Returns
1507
+ -------
1508
+ NXdata
1509
+ The padded data with symmetric zero padding.
1510
+ """
1511
+ self.padded = self.padder.pad(padding)
1512
+
1513
+ def perform_fft(self):
1514
+ """
1515
+ Perform a 3D Fourier Transform on the padded data.
1516
+
1517
+ This function applies an inverse Fourier Transform to the padded data
1518
+ using the `pyfftw` library to optimize performance. The result is a
1519
+ transformed array with spatial frequency components calculated along
1520
+ each axis.
1521
+
1522
+ Parameters
1523
+ ----------
1524
+ None
1525
+
1526
+ Returns
1527
+ -------
1528
+ None
1529
+
1530
+ Notes
1531
+ -----
1532
+ - The FFT is performed in two stages: first along the last dimension of
1533
+ the input array and then along the first two dimensions.
1534
+ - The function uses `pyfftw` for efficient computation of the Fourier
1535
+ Transform.
1536
+ - The output frequency components are computed based on the step sizes
1537
+ of the original data axes.
1538
+
1539
+
1540
+ """
1541
+ self.fft = fourier_transform_nxdata(self.padded)