♻️ Refactored xtl.diffraction correlate fft CLI

dtriand · dtriand · commit a0904a018163 · 2025-03-09T23:38:41.000+01:00
xtl.cli.diffraction.correlate.fft:cli_diffraction_correlate_fft()
- Refactored to make use of the new radial unit converters and
  dataclasses
- New `--fill` option to replace NaN values in the CCF to allow
  calculation of the FFT (will be reworked in the future)

xtl.cli.diffraction.correlate.qq:cli_diffraction_correlate_qq()
- Fixed the y-axis label on the average CCF plot from CCF_χ to CCF_Δ

xtl.cli.diffraction:cli_utils
- Removed the `units_repr()` function that is no longer necessary
- Updated `get_radial_units_from_header()` to return a `RadialUnit`
  instance
diff --git a/src/xtl/cli/diffraction/cli_utils.py b/src/xtl/cli/diffraction/cli_utils.py
@@ -6,6 +6,7 @@
 from pyFAI.detectors import Detector
 
 from xtl.diffraction.images.images import Image
+from xtl.units.crystallography.radial import RadialUnit, RadialUnitType
 
 
 class ZScale(Enum):
@@ -88,20 +89,12 @@ def get_geometry_from_header(header: str) -> Geometry:
     return Geometry(**kwargs)
 
 
-def get_radial_units_from_header(header: str) -> Optional[IntegrationRadialUnits]:
+def get_radial_units_from_header(header: str) -> Optional[RadialUnit]:
     for line in header.splitlines():
         if line.startswith('pyFAI.AzimuthalIntegrator.unit'):
             units = line.split(':')[-1].strip()
             if units in ['2th_deg', '2theta', '2th']:
-                return IntegrationRadialUnits.TWOTHETA_DEG
+                return RadialUnit.from_type(RadialUnitType.TWOTHETA_DEG)
             elif units in ['q_nm', 'q', 'q_nm^-1']:
-                return IntegrationRadialUnits.Q_NM
-    return None
-
-
-def units_repr(units: IntegrationRadialUnits) -> Optional[tuple[str, str]]:
-    if units == IntegrationRadialUnits.TWOTHETA_DEG:
-        return '2\u03b8', '\u00b0'
-    elif units == IntegrationRadialUnits.Q_NM:
-        return 'q', 'nm\u207B\u00B9'
+                return RadialUnit.from_type(RadialUnitType.Q_NM)
     return None
diff --git a/src/xtl/cli/diffraction/correlate/fft.py b/src/xtl/cli/diffraction/correlate/fft.py
@@ -1,20 +1,15 @@
 from pathlib import Path
-from typing import Optional
 
 import matplotlib.pyplot as plt
 import numpy as np
-from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn, MofNCompleteColumn
 import typer
 
 from xtl.cli.cliio import Console, epilog
 from xtl.cli.utils import Timer
-from xtl.cli.diffraction.cli_utils import (IntegrationRadialUnits, get_geometry_from_header,
-                                           get_radial_units_from_header, units_repr)
-from xtl.diffraction.images.correlators import AzimuthalCrossCorrelatorQQ_1
+from xtl.cli.diffraction.cli_utils import get_geometry_from_header, get_radial_units_from_header
 from xtl.exceptions.utils import Catcher
 from xtl.files.npx import NpxFile
-
-import warnings
+from xtl.units.crystallography.radial import RadialUnitType, RadialValue
 
 
 app = typer.Typer()
@@ -25,13 +20,18 @@ def cli_diffraction_correlate_fft(
         ccf_file: Path = typer.Argument(..., metavar='ccf.npx', help='CCF file to analyze'),
         ai2_file: Path = typer.Argument(None, metavar='ai2.npx', help='Azimuthal integration file'),
         # Selection parameters
-        selection_2theta: float = typer.Option(None, '-tt', '--2theta', help='2\u03b8 angle to inspect (in \u00b0)',
+        selection_2theta: float = typer.Option(None, '-t', '--2theta', help='2\u03b8 angle to inspect (in \u00b0)',
                                                rich_help_panel='Selection parameters'),
         selection_q: float = typer.Option(None, '-q', '--q', help='Q value to inspect (in nm\u207B\u00B9)',
                                           rich_help_panel='Selection parameters'),
         # FFT parameters
         no_coeffs: int = typer.Option(24, '-n', '--no-coeffs', help='Number of Fourier components to display',
                                       rich_help_panel='FFT parameters'),
+        fill_value: float = typer.Option(None, '--fill', help='Fill NaN values in CCF with a given value',
+                                         rich_help_panel='FFT parameters'),
+        # Plotting parameters
+        # polar_plots: bool = typer.Option(False, '--polar', help='Plot CCF and azimuthal integration in polar coordinates',
+        #                                  rich_help_panel='Plotting parameters'),
         # Debugging
         verbose: int = typer.Option(0, '-v', '--verbose', count=True, help='Print additional information',
                                     rich_help_panel='Debugging'),
@@ -47,8 +47,11 @@ def cli_diffraction_correlate_fft(
     elif selection_2theta is not None and selection_q is not None:
         cli.print('Select only one parameter to calculate FFT (--2theta, --q)', style='red')
         raise typer.Abort()
-    selection_units = IntegrationRadialUnits.Q_NM if selection_q is not None else IntegrationRadialUnits.TWOTHETA_DEG
-    selection = selection_q if selection_q is not None else selection_2theta
+
+    if selection_2theta is not None:
+        selection = RadialValue(value=selection_2theta, type=RadialUnitType.TWOTHETA_DEG)
+    else:
+        selection = RadialValue(value=selection_q, type=RadialUnitType.Q_NM)
 
     # Load CCF data
     with Catcher(echo_func=cli.print, traceback_func=cli.print_traceback) as catcher:
@@ -86,89 +89,104 @@ def cli_diffraction_correlate_fft(
         raise typer.Abort()
 
     # Get the radial units of the CCF
-    radial_units = get_radial_units_from_header(acc.header)
-    if radial_units is None:
+    r = get_radial_units_from_header(acc.header)
+    if r is None:
         cli.print('Error: Failed to get radial units from CCF file header', style='red')
         raise typer.Abort()
-    u = units_repr(radial_units)
-    cli.print(f'Radial units in CCF file: {u[0]} ({u[1]})')
+    cli.print(f'Radial units in CCF file: {r.name.latex} ({r.unit.latex})')
 
     # Convert selection units to the units of CCF if necessary
-    if radial_units != selection_units:
-        previous_selection = selection
-        if radial_units == IntegrationRadialUnits.Q_NM:  # Convert 2theta to q
-            selection = 4 * np.pi / (geometry.wavelength / 1e-10) * np.sin(np.radians(selection / 2))
-            selection *= 10  # Convert 1/A to nm^-1
-        elif radial_units == IntegrationRadialUnits.TWOTHETA_DEG:  # Convert q to 2theta
-            selection /= 10  # Convert nm^-1 to 1/A
-            selection = 2 * np.degrees(np.arcsin(selection * (geometry.wavelength / 1e-10) / (4 * np.pi)))
-        else:
-            cli.print(f'Error: Invalid radial units {radial_units!r}', style='red')
-            raise typer.Abort()
-        u0 = units_repr(selection_units)
-        u1 = units_repr(radial_units)
-        cli.print(f'Converted selection from {u0[0]}={previous_selection:.4f} ({u0[1]}) '
-                  f'to {u1[0]}={selection:.4f} ({u1[1]}) using \u03bb={geometry.wavelength / 1e-10:.4f} (\u212b)')
+    wavelength = geometry.wavelength / 1e-10
+    if r != selection.units:
+        with Catcher(echo_func=cli.print, traceback_func=cli.print_traceback) as catcher:
+            n0, v0, u0 = selection.name.latex, selection.value, selection.units.latex
+            selection = selection.to(units=r, wavelength=wavelength)
+            n1, v1, u1 = selection.name.latex, selection.value, selection.units.latex
+
+            u0 = u0 if u0 == '\u00b0' else f' {u0}'
+            u1 = u1 if u1 == '\u00b0' else f' {u1}'
+            cli.print(f'Converted selection from {n0}={v0:.4f}{u0} '
+                      f'to {n1}={v1:.4f}{u1} using \u03bb={wavelength:.4f} \u212b')
 
     # Check if selection is within the radial range of the CCF
     radial_min, radial_max = acc.data['radial'].min(), acc.data['radial'].max()
-    if selection < radial_min or selection > radial_max:
+    if selection.value < radial_min or selection.value > radial_max:
         cli.print(f'Error: Selection is outside the radial range of the CCF: {[radial_min, radial_max]}', style='red')
         raise typer.Abort()
 
     # Get the radial index of the selection
-    ccf_i = np.argmin(np.abs(acc.data['radial'] - selection))
+    ccf_i = np.argmin(np.abs(acc.data['radial'] - selection.value))
 
     # Calculate the FFT of the CCF
     if verbose:
-        cli.print(f'Calculating FFT(CCF)... ')
-    with Timer(silent=verbose<=1, echo_func=cli.print):
+        cli.print(f'Calculating FFT... ')
+    with Timer(silent=verbose<1, echo_func=cli.print):
         ccf = acc.data['ccf']
         ccf_delta = ccf[:, ccf_i]
+        if np.any(np.isnan(ccf_delta)):
+            if fill_value is not None:
+                cli.print(f'Replacing missing values in CCF with {fill_value}', style='yellow')
+                ccf_delta = np.nan_to_num(ccf_delta, nan=fill_value)
+            else:
+                cli.print('Warning: NaN values found in CCF data. FFT will be empty. '
+                          'Use --fill to replace the NaNs.', style='yellow')
         fc = np.fft.fft(ccf_delta) / len(ccf_delta)
         fc = np.abs(fc)
 
+    # Calculate selection in 2theta, q and d
+    tth = selection.to(RadialUnitType.TWOTHETA_DEG, wavelength=wavelength)
+    q = selection.to(RadialUnitType.Q_NM, wavelength=wavelength)
+    d = selection.to(RadialUnitType.D_A, wavelength=wavelength)
+    subtitle = (f'{tth.name.latex}={tth.value:.4f}{tth.units.latex} | '
+                f'{q.name.latex}={q.value:.4f} {q.units.latex} | '
+                f'{d.name.latex}={d.value:.2f} {d.units.latex}')
+
     # Prepare plots
     fig = plt.figure('XCCA overview', figsize=(16 / 1.2, 9 / 1.2))
-    fig.suptitle(f'{ccf_file.name} {u[0]}={selection:.4f} {u[1]}')
+    fig.suptitle(f'{ccf_file.name}\n{subtitle}')
     gs = fig.add_gridspec(2, 3, wspace=0.2,)
     ax0 = fig.add_subplot(gs[0, 0])  # CCF 2D
     ax1 = fig.add_subplot(gs[1, 0])  # Azimuthal integration 2D
     ax2 = fig.add_subplot(gs[0, 1])  # CCF 1D
+    # ax2 = fig.add_subplot(gs[0, 1], projection='polar')  # CCF 1D
     ax3 = fig.add_subplot(gs[1, 1])  # Azimuthal integration 1D
+    # ax3 = fig.add_subplot(gs[1, 1], projection='polar')  # Azimuthal integration 1D
     ax4 = fig.add_subplot(gs[:, 2])  # FFT
 
     ccf, radial, delta = acc.data['ccf'], acc.data['radial'], acc.data['delta']
     ax0.imshow(ccf, origin='lower', aspect='auto', interpolation='nearest', cmap='Spectral',
                #norm=norm(vmin=zmin, vmax=zmax),
                extent=(radial.min(), radial.max(), delta.min(), delta.max()))
-    ax0.vlines(selection, delta.min(), delta.max(), 'r', '--')
+    ax0.vlines(selection.value, delta.min(), delta.max(), 'r', '--')
     ax0.set_title('2D Cross-correlation function')
-    ax0.set_ylabel(f'Angular offset / \u0394 ({u[1]})')
-    ax0.set_xlabel(f'Radial angle / {u[0]} ({u[1]})')
+    ax0.set_ylabel(f'\u0394 (\u00b0)')
+    ax0.set_xlabel(f'{selection.name.latex} ({selection.units.latex})')
 
-    ax2.plot(delta, ccf[:, ccf_i])
+    ax2.plot(delta, ccf[:, ccf_i], color='xkcd:light brown')
+    # ax2.plot(delta * np.pi / 180., ccf[:, ccf_i], color='xkcd:light brown')
     ax2.set_title('1D Cross-correlation function')
     ax2.set_ylabel('Cross-correlation function')
-    ax2.set_xlabel(f'Angular offset / \u0394 ({u[1]})')
+    ax2.set_xlabel('\u03c7 (\u00b0)')
 
     if has_ai2:
         intensities, radial, azimuthal = ai2.data['intensities'], ai2.data['radial'], ai2.data['azimuthal']
         ax1.imshow(intensities, origin='lower', aspect='auto', interpolation='nearest', cmap='magma',
                    #norm=norm(vmin=zmin, vmax=zmax),
                    extent=(radial.min(), radial.max(), azimuthal.min(), azimuthal.max()))
-        ax1.vlines(selection, azimuthal.min(), azimuthal.max(), 'r', '--')
+        ax1.vlines(selection.value, azimuthal.min(), azimuthal.max(), 'r', '--')
         ax1.set_title('2D Azimuthal integration')
-        ax1.set_ylabel(f'Azimuthal angle / \u03c7 {u[1]}')
-        ax1.set_xlabel(f'Radial angle / {u[0]} ({u[1]})')
+        ax1.set_ylabel('\u03c7 (\u00b0)')
+        ax1.set_xlabel(f'{selection.name.latex} ({selection.units.latex})')
 
-        ax3.plot(azimuthal, intensities[:, ccf_i])
+        ax3.plot(azimuthal, intensities[:, ccf_i], color='xkcd:crimson')
+        # ax3.plot(azimuthal * np.pi / 180., intensities[:, ccf_i], color='xkcd:crimson')
+        # ax3.set_ylim(intensities[:, ccf_i].min(), intensities[:, ccf_i].max())
         ax3.set_title('1D Azimuthal integration')
         ax3.set_ylabel('Intensity')
-        ax3.set_xlabel(f'Azimuthal angle / \u03c7 ({u[1]})')
+        ax3.set_xlabel('\u03c7 (\u00b0)')
 
     ax4.bar(range(1, no_coeffs + 1), fc[1:no_coeffs+1])
-    ax4.set_title('FFT(CCF)')
+    ax4.set_title('$\u2131\{\mathrm{CCF}(\u0394)\}$')
     ax4.set_ylabel('Distribution')
     ax4.set_xlabel('Fourier coefficient')
 
diff --git a/src/xtl/cli/diffraction/correlate/qq.py b/src/xtl/cli/diffraction/correlate/qq.py
@@ -23,6 +23,7 @@ def cli_diffraction_correlate_qq(
         geometry: Path = typer.Option(..., '-g', '--geometry', help='Geometry .PONI file',
                                       exists=True),
         # mask: Path = typer.Option(None, '-m', '--mask', help='Mask file'),
+        # blemishes: Path = typer.Option(None, '-b', '--blemishes', help='Blemishes file'),
         # Integration parameters
         points_radial: int = typer.Option(300, '-pR', '--points-radial', help='Number of points along the radial axis',
                                           min=50, rich_help_panel='Integration parameters'),
@@ -181,7 +182,7 @@ def cli_diffraction_correlate_qq(
                 ccf_mean = np.nanmean(accf.ccf, axis=0)
             ax4.plot(img.ai1.results.radial, ccf_mean, color='xkcd:plum')
             ax4.set_xlabel(accf.units_radial_repr)
-            ax4.set_ylabel('\u27e8CCF\u27e9$_{\u03c7}$')
+            ax4.set_ylabel('\u27e8CCF\u27e9$_{\u0394}$')
             ax4.set_title('Average CCF')
 
             fig.suptitle(f'{img.file.name} frame #{img.frame}', y=0.95)
