From 6359b190f98f7683e2d14991a9bfbfd639759922 Mon Sep 17 00:00:00 2001 From: Heene Date: Sat, 11 Oct 2025 18:25:58 +0800 Subject: [PATCH] feat(scalar_fields_XZ): Accelerate rotation calculation using Numba - Introduce Numba's @njit decorator to accelerate the rotation function - Extract the original rotation calculation logic into an independent _rotate_numba function - Call the optimized Numba-accelerated function in class methods - Maintain original functionality while improving computational performance - Add parallel=True and nogil=True parameters to enable parallel computation - Preserve the mathematical correctness of the original rotation logic --- diffractio/scalar_fields_XZ.py | 14 +++++++++----- 1 file changed, 9 insertions(+), 5 deletions(-) diff --git a/diffractio/scalar_fields_XZ.py b/diffractio/scalar_fields_XZ.py index 48f48e1c..44c9a8f6 100755 --- a/diffractio/scalar_fields_XZ.py +++ b/diffractio/scalar_fields_XZ.py @@ -94,11 +94,19 @@ kernelRS, kernelRSinverse) from .scalar_masks_X import Scalar_mask_X from .scalar_sources_X import Scalar_source_X +from numba import njit copyreg.pickle(types.MethodType, _pickle_method, _unpickle_method) percentage_intensity_config = CONF_DRAWING['percentage_intensity'] +@njit(parallel=True, nogil=True) +def _rotate_numba(X, Z, angle, x0, z0): + """Numba-accelerated rotation computation""" + Xrot = x0 + (X - x0) * np.cos(angle) + (Z - z0) * np.sin(angle) + Zrot = z0 - (X - x0) * np.sin(angle) + (Z - z0) * np.cos(angle) + return Xrot, Zrot + class Scalar_field_XZ(): """Class for working with XZ scalar fields. @@ -300,11 +308,7 @@ def __rotate__(self, angle: float, position=None): # Definicion de la rotation x0, z0 = position - Xrot = x0 + (self.X - x0) * np.cos(angle) + (self.Z - - z0) * np.sin(angle) - Zrot = z0 - (self.X - x0) * np.sin(angle) + (self.Z - - z0) * np.cos(angle) - return Xrot, Zrot + return _rotate_numba(self.X, self.Z, angle, x0, z0) def size(self, verbose: bool = False):