Merge pull request #2349 from EdgarGF93/fiberunits_numexpr

fiberunit numexpr

sandbox/test_grazingincidence.py

@@ -0,0 +1,47 @@
+from pyFAI.units import Unit, get_unit_fiber, to_unit
+from pyFAI.integrator.azimuthal import AzimuthalIntegrator
+from pyFAI.calibrant import get_calibrant
+from pyFAI.integrator.fiber import FiberIntegrator
+from pyFAI import detector_factory
+import time
+from pyFAI.gui.jupyter import subplots, display, plot2d
+import matplotlib.pyplot as plt
+import numpy
+from pyFAI.test.utilstest import UtilsTest
+import fabio
+from pyFAI import load
+
+if __name__ == "__main__":
+    fi_1 = FiberIntegrator(dist=0.1, poni1=0.1, poni2=0.1, detector=detector_factory("Eiger_4M"), wavelength=1e-10)
+    poni = UtilsTest.getimage("LaB6_5.poni")
+    fi_2 = load(filename=poni, type_="pyFAI.integrator.fiber.FiberIntegrator")
+
+    cal = get_calibrant("LaB6")
+    data_1 = cal.fake_calibration_image(ai=fi_1)
+    data_file = UtilsTest.getimage("Y6.edf")
+    data_2 = fabio.open(data_file).data
+
+    for fi, data in zip((fi_1, fi_2), (data_1, data_2)):
+        res2d_1 = fi.integrate2d_grazing_incidence(data=data, incident_angle=0.0, tilt_angle=0.0, sample_orientation=1)
+        res2d_2 = fi.integrate2d_grazing_incidence(data=data, incident_angle=0.2, tilt_angle=0.0, sample_orientation=1)
+        res2d_3 = fi.integrate2d_grazing_incidence(data=data, incident_angle=0.0, tilt_angle=-0.2, sample_orientation=1)
+        res2d_4 = fi.integrate2d_grazing_incidence(data=data, incident_angle=0.2, tilt_angle=-1.54, sample_orientation=1)
+
+        res2d_5 = fi.integrate2d_grazing_incidence(data=data, incident_angle=0.0, tilt_angle=0.0, sample_orientation=2)
+        res2d_6 = fi.integrate2d_grazing_incidence(data=data, incident_angle=0.0, tilt_angle=0.0, sample_orientation=3)
+        res2d_7 = fi.integrate2d_grazing_incidence(data=data, incident_angle=0.0, tilt_angle=0.0, sample_orientation=4)
+
+        fig, axes = subplots(ncols=4, nrows=2)
+        plot2d(res2d_1, ax=axes[0,0])
+        plot2d(res2d_2, ax=axes[0,1])
+        plot2d(res2d_3, ax=axes[0,2])
+        plot2d(res2d_4, ax=axes[0,3])
+        plot2d(res2d_5, ax=axes[1,0])
+        plot2d(res2d_6, ax=axes[1,1])
+        plot2d(res2d_7, ax=axes[1,2])
+        for ax in axes.ravel():
+            if len(ax.get_images()) == 0:
+                continue
+            ax.get_images()[0].set_cmap("viridis")
+            # ax.get_images()[0].set_clim(0,1500)
+        plt.show()

src/pyFAI/test/test_fiber_integrator.py

@@ -542,6 +542,52 @@ def test_sample_orientation_equivalence(self):
 
         self.assertLess((abs(res_so_4.intensity) - abs(res_so_3.intensity)).max(), threshold)
 
+    def test_numexpr_equations(self):
+        incident_angle = 0.2
+        tilt_angle = -1.0
+        sample_orientation = 3
+
+        qhorz = get_unit_fiber(name="qxgi_nm^-1", incident_angle=incident_angle, tilt_angle=tilt_angle, sample_orientation=sample_orientation)
+        qvert = get_unit_fiber(name="qygi_nm^-1", incident_angle=incident_angle, tilt_angle=tilt_angle, sample_orientation=sample_orientation)
+        qbeam = get_unit_fiber(name="qzgi_nm^-1", incident_angle=incident_angle, tilt_angle=tilt_angle, sample_orientation=sample_orientation)
+        qip = get_unit_fiber(name="qip_nm^-1", incident_angle=incident_angle, tilt_angle=tilt_angle, sample_orientation=sample_orientation)
+        qoop = get_unit_fiber(name="qoop_nm^-1", incident_angle=incident_angle, tilt_angle=tilt_angle, sample_orientation=sample_orientation)
+
+        self.fi.reset()
+        arr_qhorz_fast = self.fi.array_from_unit(unit=qhorz)
+        arr_qvert_fast = self.fi.array_from_unit(unit=qvert)
+        arr_qbeam_fast = self.fi.array_from_unit(unit=qbeam)
+        arr_qip_fast = self.fi.array_from_unit(unit=qip)
+        arr_qoop_fast = self.fi.array_from_unit(unit=qoop)
+        res2d_fast = self.fi.integrate2d_grazing_incidence(data=self.data, unit_ip=qip, unit_oop=qoop)
+
+        qhorz.formula = None
+        qhorz.equation = qhorz._equation
+        qvert.formula = None
+        qvert.equation = qvert._equation
+        qbeam.formula = None
+        qbeam.equation = qbeam._equation
+        qip.formula = None
+        qip.equation = qip._equation
+        qoop.formula = None
+        qoop.equation = qoop._equation
+
+        self.fi.reset()
+        arr_qhorz_slow = self.fi.array_from_unit(unit=qhorz)
+        arr_qvert_slow = self.fi.array_from_unit(unit=qvert)
+        arr_qbeam_slow = self.fi.array_from_unit(unit=qbeam)
+        arr_qip_slow = self.fi.array_from_unit(unit=qip)
+        arr_qoop_slow = self.fi.array_from_unit(unit=qoop)
+        res2d_slow = self.fi.integrate2d_grazing_incidence(data=self.data, unit_ip=qip, unit_oop=qoop)
+
+        self.assertAlmostEqual((arr_qhorz_fast - arr_qhorz_slow).max(), 0.0)
+        self.assertAlmostEqual((arr_qvert_fast - arr_qvert_slow).max(), 0.0)
+        self.assertAlmostEqual((arr_qbeam_fast - arr_qbeam_slow).max(), 0.0)
+        self.assertAlmostEqual((arr_qip_fast - arr_qip_slow).max(), 0.0)
+        self.assertAlmostEqual((arr_qoop_fast - arr_qoop_slow).max(), 0.0)
+        self.assertAlmostEqual((res2d_fast.intensity - res2d_slow.intensity).max(), 0.0)
+        self.assertAlmostEqual((res2d_fast.radial - res2d_slow.radial).max(), 0.0)
+        self.assertAlmostEqual((res2d_fast.azimuthal - res2d_slow.azimuthal).max(), 0.0)
 
 def suite():
     testsuite = unittest.TestSuite()

src/pyFAI/units.py

@@ -160,7 +160,7 @@ def __init__(self, name, scale=1, label=None, equation=None, formula=None,
             scale=scale,
             label=label,
             equation=equation,
-            formula=formula,
+            # formula=formula,
             center=center,
             corner=corner,
             delta=delta,
@@ -169,9 +169,49 @@ def __init__(self, name, scale=1, label=None, equation=None, formula=None,
             positive=positive,
             period=period,
         )
+        self.formula = formula
+        self.formula_so1 = formula
         self._incident_angle = incident_angle
         self._tilt_angle = tilt_angle
         self._sample_orientation = sample_orientation
+        self._update_ne_equation()
+
+    def _update_ne_equation(self):
+        if (numexpr is not None) and isinstance(self.formula, str):
+            signature = [(key, numpy.float64) for key in "xyzλπηχ" if key in self.formula]
+
+            formula_ = self.formula
+            if self._sample_orientation == 1:
+                ...
+            elif self._sample_orientation == 2:
+                formula_ = self.formula_so1
+                formula_ = formula_.replace('x', 'ψ').replace('y', 'ξ')
+                formula_ = formula_.replace('ψ', 'y').replace('ξ', 'x')
+            elif self._sample_orientation == 3:
+                formula_ = self.formula_so1.replace('x', '(-x)')
+            elif self._sample_orientation == 4:
+                formula_ = self.formula_so1
+                formula_ = formula_.replace('x', 'ψ').replace('y', 'ξ')
+                formula_ = formula_.replace('ψ', '(-y)').replace('ξ', '(-x)')
+            self.formula = formula_
+            ne_formula = numexpr.NumExpr(self.formula, signature)
+
+            def ne_equation(x, y, z=None, wavelength=None,
+                            incident_angle=self._incident_angle,
+                            tilt_angle=self._tilt_angle,
+                            sample_orientation=self._sample_orientation,
+                            ne_formula=ne_formula):
+                π = numpy.pi
+                λ = wavelength
+                η = self._incident_angle
+                χ = self._tilt_angle
+                ldict = locals()
+                args = tuple(ldict[i] for i in ne_formula.input_names)
+                return ne_formula(*args)
+
+            self.equation = ne_equation
+        else:
+            self.equation = self._equation
 
     def __repr__(self):
         return f"""
@@ -195,12 +235,15 @@ def sample_orientation(self):
 
     def set_incident_angle(self, value:float):
         self._incident_angle = value
+        self._update_ne_equation()
 
     def set_tilt_angle(self, value:float):
         self._tilt_angle = value
+        self._update_ne_equation()
 
     def set_sample_orientation(self, value: int):
         self._sample_orientation = value
+        self._update_ne_equation()
 
 
 RADIAL_UNITS = {}
@@ -618,6 +661,17 @@ def eq_q_total(x, y, z, wavelength, incident_angle=0.0, tilt_angle=0.0, sample_o
 formula_qx = "4.0e-9*π/λ*sin(arctan2(x, z)/2.0)"
 formula_qy = "4.0e-9*π/λ*sin(arctan2(y, z)/2.0)"
 
+formula_exit_angle = "arctan2(y,sqrt(z**2+x**2))"
+formula_exit_angle_horz = "arctan2(x,z)"
+formula_qbeam_lab = f"2.0e-9/λ*π*(cos({formula_exit_angle})*cos({formula_exit_angle_horz}) - 1)"
+formula_qhorz_lab = f"2.0e-9/λ*π*cos({formula_exit_angle})*sin({formula_exit_angle_horz})"
+formula_qvert_lab = f"2.0e-9/λ*π*sin({formula_exit_angle})"
+formula_qbeam_rot = f"cos(η)*({formula_qbeam_lab})+sin(η)*({formula_qvert_lab})"
+formula_qhorz_rot = f"cos(χ)*({formula_qhorz_lab})-sin(χ)*sin(η)*({formula_qbeam_lab})+sin(χ)*cos(η)*({formula_qvert_lab})"
+formula_qvert_rot = f"-sin(χ)*({formula_qhorz_lab})-cos(χ)*sin(η)*({formula_qbeam_lab})+cos(χ)*cos(η)*({formula_qvert_lab})"
+formula_qip = f"sqrt(({formula_qbeam_rot})**2+({formula_qhorz_rot})**2)*((({formula_qhorz_rot} > 0) * 2) - 1)"
+formula_qoop = formula_qvert_rot
+
 register_radial_unit("r_mm",
                      center="rArray",
                      delta="deltaR",
@@ -810,6 +864,7 @@ def eq_q_total(x, y, z, wavelength, incident_angle=0.0, tilt_angle=0.0, sample_o
 register_radial_fiber_unit("qxgi_nm^-1",
                      scale=1.0,
                      label=r"Scattering vector $q_x$ ($nm^{-1}$)",
+                     formula=formula_qhorz_rot,
                      equation=eq_qhorz_gi,
                      short_name="qxgi",
                      unit_symbol="nm^{-1}",
@@ -818,6 +873,7 @@ def eq_q_total(x, y, z, wavelength, incident_angle=0.0, tilt_angle=0.0, sample_o
 register_radial_fiber_unit("qygi_nm^-1",
                      scale=1.0,
                      label=r"Scattering vector $q_y$ ($nm^{-1}$)",
+                     formula=formula_qvert_rot,
                      equation=eq_qvert_gi,
                      short_name="qygi",
                      unit_symbol="nm^{-1}",
@@ -826,6 +882,7 @@ def eq_q_total(x, y, z, wavelength, incident_angle=0.0, tilt_angle=0.0, sample_o
 register_radial_fiber_unit("qzgi_nm^-1",
                      scale=1.0,
                      label=r"Scattering vector $q_z$ ($nm^{-1}$)",
+                     formula=formula_qbeam_rot,
                      equation=eq_qbeam_gi,
                      short_name="qzgi",
                      unit_symbol="nm^{-1}",
@@ -834,6 +891,7 @@ def eq_q_total(x, y, z, wavelength, incident_angle=0.0, tilt_angle=0.0, sample_o
 register_radial_fiber_unit("qip_nm^-1",
                      scale=1.0,
                      label=r"Scattering vector $q_{IP}$ ($nm^{-1}$)",
+                     formula=formula_qip,
                      equation=eq_qip,
                      short_name="qip",
                      unit_symbol="nm^{-1}",
@@ -842,6 +900,7 @@ def eq_q_total(x, y, z, wavelength, incident_angle=0.0, tilt_angle=0.0, sample_o
 register_radial_fiber_unit("qoop_nm^-1",
                      scale=1.0,
                      label=r"Scattering vector $q_{OOP}$ ($nm^{-1}$)",
+                     formula=formula_qoop,
                      equation=eq_qoop,
                      short_name="qoop",
                      unit_symbol="nm^{-1}",