⚡ Optimized imports in xtl.diffraction & xtl.math

xtl.cli.diffraction & xtl.cli.math
- Optimized imports throughout to improve startup performance

Author: dtriand

src/xtl/cli/diffraction/__init__.py

@@ -2,12 +2,16 @@
 
 from .geometry import app as geometry_app
 from .plot import app as plot_app
+# from .sum import app as sum_app
 from .integrate import app as integrate_app
 from .correlate import app as correlate_app
+# from .array import app as array_app
 
 
 app = typer.Typer(name='xtl.diffraction', help='Utilities for diffraction data')
 app.add_typer(geometry_app)
 app.add_typer(plot_app)
+# app.add_typer(sum_app)
 app.add_typer(integrate_app)
 app.add_typer(correlate_app)
+# app.add_typer(array_app)

src/xtl/cli/diffraction/cli_utils.py

@@ -1,12 +1,12 @@
 from enum import Enum
 import re
-from typing import Optional
+from typing import Optional, TYPE_CHECKING
 
-from pyFAI.geometry import Geometry
-from pyFAI.detectors import Detector
+if TYPE_CHECKING:
+    from pyFAI.geometry import Geometry
 
-from xtl.diffraction.images.images import Image
-from xtl.units.crystallography.radial import RadialUnit, RadialUnitType
+    from xtl.diffraction.images.images import Image
+    from xtl.units.crystallography.radial import RadialUnit
 
 
 class ZScale(Enum):
@@ -26,7 +26,7 @@ class IntegrationRadialUnits(Enum):
     Q_NM = 'q_nm'
 
 
-def get_image_frames(images: list[str]) -> list[Image]:
+def get_image_frames(images: list[str]) -> list['Image']:
     opened_images = []
     for i, img in enumerate(images):
         parts = img.split(':')
@@ -42,6 +42,8 @@ def get_image_frames(images: list[str]) -> list[Image]:
         else:
             raise ValueError(f'Invalid image format for image [{i}]: {img!r}')
 
+        from xtl.diffraction.images.images import Image
+
         image = Image()
         try:
             image.open(file=file, frame=frame, is_eager=False)
@@ -52,11 +54,14 @@ def get_image_frames(images: list[str]) -> list[Image]:
     return opened_images
 
 
-def get_geometry_from_header(header: str) -> Geometry:
+def get_geometry_from_header(header: str) -> 'Geometry':
     """
     Return a pyFAI Geometry object from the header of an NPX file written by
     AzimuthalCrossCorrelatorQQ_1 or Integrator.
     """
+    from pyFAI.geometry import Geometry
+    from pyFAI.detectors import Detector
+
     lines = []
     for line in header.splitlines():
         if line.startswith('pyFAI.Geometry'):
@@ -89,7 +94,9 @@ def get_geometry_from_header(header: str) -> Geometry:
     return Geometry(**kwargs)
 
 
-def get_radial_units_from_header(header: str) -> Optional[RadialUnit]:
+def get_radial_units_from_header(header: str) -> Optional['RadialUnit']:
+    from xtl.units.crystallography.radial import RadialUnit, RadialUnitType
+
     for line in header.splitlines():
         if line.startswith('pyFAI.AzimuthalIntegrator.unit'):
             units = line.split(':')[-1].strip()

src/xtl/cli/diffraction/correlate/fft.py

@@ -1,17 +1,9 @@
-from functools import partial
 from pathlib import Path
 
-import matplotlib.pyplot as plt
-from matplotlib.colors import Normalize, LogNorm, SymLogNorm
-import numpy as np
 import typer
 
 from xtl.cli.cliio import Console, epilog
-from xtl.cli.utils import Timer
-from xtl.cli.diffraction.cli_utils import get_geometry_from_header, get_radial_units_from_header, ZScale
-from xtl.exceptions.utils import Catcher
-from xtl.files.npx import NpxFile
-from xtl.units.crystallography.radial import RadialUnitType, RadialValue
+from xtl.cli.diffraction.cli_utils import ZScale
 
 
 app = typer.Typer()
@@ -60,12 +52,16 @@ def cli_diffraction_correlate_fft(
         cli.print('Select only one parameter to calculate FFT (--2theta, --q)', style='red')
         raise typer.Abort()
 
+    from xtl.units.crystallography.radial import RadialUnitType, RadialValue
     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
+    from xtl.exceptions.utils import Catcher
+    from xtl.files.npx import NpxFile
+
     with Catcher(echo_func=cli.print, traceback_func=cli.print_traceback) as catcher:
         acc = NpxFile.load(ccf_file)
         for key in ['radial', 'delta', 'ccf']:
@@ -91,6 +87,8 @@ def cli_diffraction_correlate_fft(
             raise typer.Abort()
 
     # Get geometry from CCF file
+    from xtl.cli.diffraction.cli_utils import get_geometry_from_header, get_radial_units_from_header
+
     with Catcher(echo_func=cli.print, traceback_func=cli.print_traceback) as catcher:
         geometry = get_geometry_from_header(acc.header)
         if verbose > 2:
@@ -127,6 +125,9 @@ def cli_diffraction_correlate_fft(
         raise typer.Abort()
 
     # Get the radial index of the selection
+    import numpy as np
+    from xtl.cli.utils import Timer
+
     ccf_i = np.argmin(np.abs(acc.data['radial'] - selection.value))
 
     # Calculate the FFT of the CCF
@@ -154,6 +155,9 @@ def cli_diffraction_correlate_fft(
                 f'{d.name.latex}={d.value:.2f} {d.units.latex}')
 
     # Prepare plots
+    import matplotlib.pyplot as plt
+    from matplotlib.colors import Normalize, LogNorm, SymLogNorm
+
     fig = plt.figure('XCCA overview', figsize=(16 / 1.2, 9 / 1.2))
     fig.suptitle(f'{ccf_file.name}\n{subtitle}')
     gs = fig.add_gridspec(2, 3, wspace=0.2,)
@@ -196,6 +200,8 @@ def cli_diffraction_correlate_fft(
         if zmax is not None:
             vmax = zmax
 
+    from functools import partial
+
     if zscale == ZScale.LINEAR:
         norm = partial(Normalize, clip=False)
     elif zscale == ZScale.LOG:

src/xtl/cli/diffraction/correlate/qq.py

@@ -1,15 +1,9 @@
 from pathlib import Path
 
-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 get_image_frames, IntegrationErrorModel, IntegrationRadialUnits
-from xtl.diffraction.images.correlators import AzimuthalCrossCorrelatorQQ_1
-from xtl.exceptions.utils import Catcher
+from xtl.cli.diffraction.cli_utils import IntegrationErrorModel, IntegrationRadialUnits
 
 import warnings
 
@@ -56,7 +50,10 @@ def cli_diffraction_correlate_qq(
     cli = Console(verbose=verbose, debug=debug)
     input_images = images
 
+    from xtl.exceptions.utils import Catcher
+
     with Catcher(echo_func=cli.print, traceback_func=cli.print_traceback) as catcher:
+        from xtl.cli.diffraction.cli_utils import get_image_frames
         images = get_image_frames(input_images)
     if catcher.raised:
         cli.print(f'Error: Failed to read all images', style='red')
@@ -84,6 +81,11 @@ def cli_diffraction_correlate_qq(
         'error_model': error_model.value if error_model != IntegrationErrorModel.NONE else None,
     }
 
+    from xtl.cli.utils import Timer
+    from xtl.diffraction.images.correlators import AzimuthalCrossCorrelatorQQ_1
+    from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn, \
+        MofNCompleteColumn
+
     with (Catcher(echo_func=cli.print, traceback_func=cli.print_traceback),
           Progress(SpinnerColumn(), *Progress.get_default_columns(),
                    TimeElapsedColumn(), MofNCompleteColumn(),
@@ -115,7 +117,7 @@ def cli_diffraction_correlate_qq(
                 with warnings.catch_warnings():
                     warnings.filterwarnings('ignore')
                     accf.correlate(points_radial=points_radial, points_azimuthal=points_azimuthal,
-                                  units_radial=units_radial.value, method=0)
+                                   units_radial=units_radial.value, method=0)
 
             ccf_file = output_dir / f'{dataset_name}_ccf.npx'
             accf.save(ccf_file, overwrite=overwrite)
@@ -129,6 +131,9 @@ def cli_diffraction_correlate_qq(
                 progress.console.print(f'Saved 2D integration results to {ai2_file}')
 
             # Prepare plots
+            import matplotlib.pyplot as plt
+            import numpy as np
+
             fig = plt.figure('XCCA overview', figsize=(16 / 1.2, 9 / 1.2))
             gs0 = fig.add_gridspec(1, 2, wspace=0.2,
                                    width_ratios=[1.2, 2])  # outer grid (1x2)

src/xtl/cli/diffraction/geometry.py

@@ -1,16 +1,18 @@
 from pathlib import Path
+from typing import TYPE_CHECKING
 
-from pyFAI.detectors import ALL_DETECTORS, Detector
-from pyFAI.geometry import Geometry
 import typer
+if TYPE_CHECKING:
+    from pyFAI.detectors import ALL_DETECTORS, Detector
+    from pyFAI.geometry import Geometry
 
 from xtl.cli.cliio import Console
 
 
 app = typer.Typer()
 
 
-def get_detector_info(detector: Detector) -> dict:
+def get_detector_info(detector: 'Detector') -> dict:
     detector_info = {}
     detector_info['Detector name'] = detector.name
     if isinstance(detector.MANUFACTURER, list):
@@ -31,6 +33,8 @@ def get_detector_info(detector: Detector) -> dict:
 
 
 def get_detectors_list() -> dict:
+    from pyFAI.detectors import ALL_DETECTORS
+
     detectors = {}
     for alias, detector in ALL_DETECTORS.items():
         detector = detector()
@@ -50,7 +54,7 @@ def get_detectors_list() -> dict:
     return {k: v['info'] for k, v in detectors.items()}
 
 
-def get_geometry_info(geometry: Geometry) -> dict:
+def get_geometry_info(geometry: 'Geometry') -> dict:
     return {f'{k}': f'{v}' for k, v in geometry.get_config().items() if k not in ['poni_version', 'detector_config']}
 
 
@@ -63,15 +67,20 @@ def cli_diffraction_geometry():
     while detector is None:
         answer = typer.prompt('Enter detector name (? for list)').lower()
         if answer in ['detector', 'custom']:
+            from pyFAI.detectors import Detector
+
             cli.print('Initializing custom detector...')
             pixel1 = typer.prompt('Enter pixel size (horizontal) [pixel1 in \u03bcm]', default=50., type=float)
             pixel2 = typer.prompt('Enter pixel size (vertical) [pixel2 in \u03bcm]', default=50., type=float)
             shape1 = typer.prompt('Enter number of pixels (horizontal) [shape[1] in px]', default=1024, type=int)
             shape2 = typer.prompt('Enter number of pixels (vertical) [shape[0] in px', default=1024, type=int)
             detector = Detector(pixel1=pixel1/1e6, pixel2=pixel2/1e6, max_shape=(shape2, shape1))
-        elif answer in ALL_DETECTORS:
-            detector: Detector = ALL_DETECTORS[answer]()
         else:
+            from pyFAI.detectors import ALL_DETECTORS
+            if answer in ALL_DETECTORS:
+                detector: 'Detector' = ALL_DETECTORS[answer]()
+                continue
+
             detector_list = get_detectors_list()
             headers = list(list(detector_list.values())[0].keys())
             cli.print_table([d.values() for d in detector_list.values()], headers=headers,
@@ -89,6 +98,8 @@ def cli_diffraction_geometry():
             detector = None
 
     # Build Geometry object
+    from pyFAI.geometry import Geometry
+
     geometry = None
     while geometry is None:
         poni2 = typer.prompt('Enter point of normal incidence along horizontal (x) axis [poni2 in px]', default=0., type=float)

src/xtl/cli/diffraction/integrate/integrate_1d.py

@@ -1,11 +1,10 @@
 from pathlib import Path
 
-import matplotlib.pyplot as plt
-from rich.progress import Progress, SpinnerColumn, TimeElapsedColumn, MofNCompleteColumn
+from rich.progress import Progress
 import typer
 
 from xtl.cli.cliio import Console, epilog
-from xtl.cli.diffraction.cli_utils import get_image_frames, IntegrationErrorModel, IntegrationRadialUnits
+from xtl.cli.diffraction.cli_utils import IntegrationErrorModel, IntegrationRadialUnits
 
 
 app = typer.Typer()
@@ -48,6 +47,8 @@ def cli_diffraction_integrate_1d(
     input_images = images
 
     try:
+        from xtl.cli.diffraction.cli_utils import get_image_frames
+
         images = get_image_frames(input_images)
     except ValueError as e:
         cli.print_traceback(e)
@@ -78,6 +79,8 @@ def cli_diffraction_integrate_1d(
         'error_model': error_model.value if error_model != IntegrationErrorModel.NONE else None,
     }
 
+    import matplotlib.pyplot as plt
+
     output_dir = output_dir.expanduser().resolve()
     fig, ax = plt.subplots()
     xscale = 'log' if xlog else 'linear'

src/xtl/cli/diffraction/plot/plot_2d.py

@@ -1,13 +1,10 @@
 import math
 from pathlib import Path
 
-import matplotlib.pyplot as plt
-from mpl_toolkits.axes_grid1 import AxesGrid
-import numpy as np
 import typer
 
 from xtl.cli.cliio import Console, epilog
-from xtl.cli.diffraction.cli_utils import get_image_frames, ZScale
+from xtl.cli.diffraction.cli_utils import ZScale
 
 
 app = typer.Typer()
@@ -48,12 +45,18 @@ def cli_diffraction_plot_2d(
     input_images = images
 
     try:
+        from xtl.cli.diffraction.cli_utils import get_image_frames
         images = get_image_frames(input_images)
     except ValueError as e:
         cli.print_traceback(e)
         cli.print(f'Error: Failed to read all images', style='red')
         raise typer.Abort()
 
+    # Imports
+    import matplotlib.pyplot as plt
+    from mpl_toolkits.axes_grid1 import AxesGrid
+    import numpy as np
+
     # Initialize figure
     fig = plt.figure()
     nimages = len(images)

src/xtl/cli/math/spacing.py

@@ -2,7 +2,6 @@
 import typer
 
 from xtl.cli.cliio import Console
-from xtl.exceptions.utils import Catcher
 from xtl.units.crystallography.radial import RadialUnitType, RadialValue, RadialUnit
 
 app = typer.Typer()
@@ -98,6 +97,8 @@ def cli_math_spacing(
         cli.print('Wavelength required to calculate all units', style='red')
         raise typer.Abort()
 
+    from xtl.exceptions.utils import Catcher
+
     r = RadialValue(value=value, type=from_type)
     with Catcher(echo_func=cli.print, traceback_func=cli.print_traceback, silent=True) as catcher:
         if to_type: