Merge pull request #584 from HEXRD/inverse-distortion

Inverse distortion

conda.recipe/meta.yaml

@@ -14,19 +14,23 @@ build:
 requirements:
   build:
     # This is so that we can build cross-platform for osx-arm64
-    - {{ compiler('c') }}
+    - {{ compiler('cxx') }}
     - python {{ python }}                    # [build_platform != target_platform]
     - cross-python_{{ target_platform }}     # [build_platform != target_platform]
     - numpy >=1.20                           # [build_platform != target_platform]
     # Numba is only here to make sure we use a version of numpy that is compatible
     - numba                                  # [build_platform != target_platform]
+    - pybind11                               # [build_platform != target_platform]
   host:
     - python {{ python }}
     - numpy >=1.20
     - setuptools
     - setuptools_scm
     # Numba is only here to make sure we use a version of numpy that is compatible
     - numba
+    - pybind11
+    - xsimd
+    - eigen
 
   run:
     - appdirs

hexrd/distortion/ge_41rt.py

@@ -1,190 +1,113 @@
 """GE41RT Detector Distortion"""
-import numpy as np
+from typing import List
 
-from hexrd import constants as cnst
-from hexrd.constants import USE_NUMBA
-if USE_NUMBA:
-    import numba
+import numpy as np
+import numba
 
 from .distortionabc import DistortionABC
 from .registry import _RegisterDistortionClass
 from .utils import newton
 
+from hexrd import constants as cnst
+from hexrd.extensions import inverse_distortion
+
 RHO_MAX = 204.8  # max radius in mm for ge detector
 
-if USE_NUMBA:
-    @numba.njit(nogil=True, cache=True)
-    def _ge_41rt_inverse_distortion(out, in_, rhoMax, params):
-        maxiter = 100
-        prec = cnst.epsf
-
-        p0, p1, p2, p3, p4, p5 = params[0:6]
-        rxi = 1.0/rhoMax
-        for el in range(len(in_)):
-            xi, yi = in_[el, 0:2]
-            ri = np.sqrt(xi*xi + yi*yi)
-            if ri < cnst.sqrt_epsf:
-                ri_inv = 0.0
-            else:
-                ri_inv = 1.0/ri
-            sinni = yi*ri_inv
-            cosni = xi*ri_inv
-            ro = ri
-            cos2ni = cosni*cosni - sinni*sinni
-            sin2ni = 2*sinni*cosni
-            cos4ni = cos2ni*cos2ni - sin2ni*sin2ni
-            # newton solver iteration
-            for i in range(maxiter):
-                ratio = ri*rxi
-                fx = (p0*ratio**p3*cos2ni +
-                      p1*ratio**p4*cos4ni +
-                      p2*ratio**p5 + 1)*ri - ro  # f(x)
-                fxp = (p0*ratio**p3*cos2ni*(p3+1) +
-                       p1*ratio**p4*cos4ni*(p4+1) +
-                       p2*ratio**p5*(p5+1) + 1)  # f'(x)
-
-                delta = fx/fxp
-                ri = ri - delta
-                # convergence check for newton
-                if np.abs(delta) <= prec*np.abs(ri):
-                    break
-
-            xi = ri*cosni
-            yi = ri*sinni
-            out[el, 0] = xi
-            out[el, 1] = yi
-
-        return out
-
-    @numba.njit(nogil=True, cache=True)
-    def _ge_41rt_distortion(out, in_, rhoMax, params):
-        p0, p1, p2, p3, p4, p5 = params[0:6]
-        rxi = 1.0/rhoMax
-
-        for el in range(len(in_)):
-            xi, yi = in_[el, 0:2]
-            ri = np.sqrt(xi*xi + yi*yi)
-            if ri < cnst.sqrt_epsf:
-                ri_inv = 0.0
-            else:
-                ri_inv = 1.0/ri
-            sinni = yi*ri_inv
-            cosni = xi*ri_inv
-            cos2ni = cosni*cosni - sinni*sinni
-            sin2ni = 2*sinni*cosni
-            cos4ni = cos2ni*cos2ni - sin2ni*sin2ni
-            ratio = ri*rxi
-
-            ri = (p0*ratio**p3*cos2ni
-                  + p1*ratio**p4*cos4ni
-                  + p2*ratio**p5
-                  + 1)*ri
-            xi = ri*cosni
-            yi = ri*sinni
-            out[el, 0] = xi
-            out[el, 1] = yi
-
-        return out
-else:
-    # non-numba versions for the direct and inverse distortion
-    def _ge_41rt_inverse_distortion(out, in_, rhoMax, params):
-        maxiter = 100
-        prec = cnst.epsf
-
-        p0, p1, p2, p3, p4, p5 = params[0:6]
-        rxi = 1.0/rhoMax
-
-        xi, yi = in_[:, 0], in_[:, 1]
-        ri = np.sqrt(xi*xi + yi*yi)
-        # !!! adding fix TypeError when processings list of coords
-        zfix = []
-        if np.any(ri) < cnst.sqrt_epsf:
-            zfix = ri < cnst.sqrt_epsf
-            ri[zfix] = 1.0
-        ri_inv = 1.0/ri
-        ri_inv[zfix] = 0.
-
-        sinni = yi*ri_inv
-        cosni = xi*ri_inv
-        ro = ri
-        cos2ni = cosni*cosni - sinni*sinni
-        sin2ni = 2*sinni*cosni
-        cos4ni = cos2ni*cos2ni - sin2ni*sin2ni
-
-        # newton solver iteration
-        #
-        # FIXME: looks like we hae a problem here,
-        #        should iterate over single coord pairs?
-        for i in range(maxiter):
-            ratio = ri*rxi
-            fx = (p0*ratio**p3*cos2ni +
-                  p1*ratio**p4*cos4ni +
-                  p2*ratio**p5 + 1)*ri - ro  # f(x)
-            fxp = (p0*ratio**p3*cos2ni*(p3+1) +
-                   p1*ratio**p4*cos4ni*(p4+1) +
-                   p2*ratio**p5*(p5+1) + 1)  # f'(x)
-
-            delta = fx/fxp
-            ri = ri - delta
-
-            # convergence check for newton
-            if np.max(np.abs(delta/ri)) <= prec:
-                break
-
-        out[:, 0] = ri*cosni
-        out[:, 1] = ri*sinni
-
-        return out
-
-    def _ge_41rt_distortion(out, in_, rhoMax, params):
-        p0, p1, p2, p3, p4, p5 = params[0:6]
-        rxi = 1.0/rhoMax
-
-        xi, yi = in_[:, 0], in_[:, 1]
-
-        # !!! included fix on ValueError for array--like in_
-        ri = np.sqrt(xi*xi + yi*yi)
-        ri[ri < cnst.sqrt_epsf] = np.inf
-        ri_inv = 1.0/ri
-
-        sinni = yi*ri_inv
-        cosni = xi*ri_inv
-        cos2ni = cosni*cosni - sinni*sinni
-        sin2ni = 2*sinni*cosni
-        cos4ni = cos2ni*cos2ni - sin2ni*sin2ni
-        ratio = ri*rxi
-
-        ri = (p0*ratio**p3*cos2ni
-              + p1*ratio**p4*cos4ni
-              + p2*ratio**p5
-              + 1)*ri
-        out[:, 0] = ri*cosni
-        out[:, 1] = ri*sinni
-
-        return out
+
+# NOTE: Deprecated in favor of inverse_distortion.ge_41rt_inverse_distortion
+@numba.njit(nogil=True, cache=True, fastmath=True)
+def _ge_41rt_inverse_distortion(
+    inputs: np.ndarray[np.float64, np.float64],
+    rhoMax: float,
+    params: List[float],
+):
+    radii = np.hypot(inputs[:, 0], inputs[:, 1])
+    inverted_radii = np.reciprocal(radii)
+    cosines = inputs[:, 0] * inverted_radii
+    cosine_double_angles = 2 * np.square(cosines) - 1
+    cosine_quadruple_angles = 2 * np.square(cosine_double_angles) - 1
+    sqrt_p_is = rhoMax / np.sqrt(
+        -(
+            params[0] * cosine_double_angles
+            + params[1] * cosine_quadruple_angles
+            + params[2]
+        )
+    )
+    solutions = (
+        (2 / np.sqrt(3))
+        * sqrt_p_is
+        * np.cos(
+            np.arccos(((-3 * np.sqrt(3) / 2) * radii / sqrt_p_is)) / 3
+            + 4 * np.pi / 3
+        )
+    )
+
+    return solutions[:, None] * inputs * inverted_radii[:, None]
+
+
+@numba.njit(nogil=True, cache=True)
+def _ge_41rt_distortion(out, in_, rhoMax, params):
+    p0, p1, p2, p3, p4, p5 = params[0:6]
+    rxi = 1.0 / rhoMax
+
+    for el in range(len(in_)):
+        xi, yi = in_[el, 0:2]
+        ri = np.sqrt(xi * xi + yi * yi)
+        if ri < cnst.sqrt_epsf:
+            ri_inv = 0.0
+        else:
+            ri_inv = 1.0 / ri
+        sinni = yi * ri_inv
+        cosni = xi * ri_inv
+        cos2ni = cosni * cosni - sinni * sinni
+        sin2ni = 2 * sinni * cosni
+        cos4ni = cos2ni * cos2ni - sin2ni * sin2ni
+        ratio = ri * rxi
+
+        ri = (
+            p0 * ratio**p3 * cos2ni
+            + p1 * ratio**p4 * cos4ni
+            + p2 * ratio**p5
+            + 1
+        ) * ri
+        xi = ri * cosni
+        yi = ri * sinni
+        out[el, 0] = xi
+        out[el, 1] = yi
+
+    return out
 
 
 def _rho_scl_func_inv(ri, ni, ro, rx, p):
-    retval = (p[0]*(ri/rx)**p[3] * np.cos(2.0 * ni) +
-              p[1]*(ri/rx)**p[4] * np.cos(4.0 * ni) +
-              p[2]*(ri/rx)**p[5] + 1)*ri - ro
+    retval = (
+        p[0] * (ri / rx) ** p[3] * np.cos(2.0 * ni)
+        + p[1] * (ri / rx) ** p[4] * np.cos(4.0 * ni)
+        + p[2] * (ri / rx) ** p[5]
+        + 1
+    ) * ri - ro
     return retval
 
 
 def _rho_scl_dfunc_inv(ri, ni, ro, rx, p):
-    retval = p[0]*(ri/rx)**p[3] * np.cos(2.0 * ni) * (p[3] + 1) + \
-        p[1]*(ri/rx)**p[4] * np.cos(4.0 * ni) * (p[4] + 1) + \
-        p[2]*(ri/rx)**p[5] * (p[5] + 1) + 1
+    retval = (
+        p[0] * (ri / rx) ** p[3] * np.cos(2.0 * ni) * (p[3] + 1)
+        + p[1] * (ri / rx) ** p[4] * np.cos(4.0 * ni) * (p[4] + 1)
+        + p[2] * (ri / rx) ** p[5] * (p[5] + 1)
+        + 1
+    )
     return retval
 
 
 def inverse_distortion_numpy(rho0, eta0, rhoMax, params):
-    return newton(rho0, _rho_scl_func_inv, _rho_scl_dfunc_inv,
-                  (eta0, rho0, rhoMax, params))
+    return newton(
+        rho0,
+        _rho_scl_func_inv,
+        _rho_scl_dfunc_inv,
+        (eta0, rho0, rhoMax, params),
+    )
 
 
 class GE_41RT(DistortionABC, metaclass=_RegisterDistortionClass):
-
     maptype = "GE_41RT"
 
     def __init__(self, params, **kwargs):
@@ -201,10 +124,9 @@ def params(self, x):
 
     @property
     def is_trivial(self):
-        return \
-            self.params[0] == 0 and \
-            self.params[1] == 0 and \
-            self.params[2] == 0
+        return (
+            self.params[0] == 0 and self.params[1] == 0 and self.params[2] == 0
+        )
 
     def apply(self, xy_in):
         if self.is_trivial:
@@ -222,8 +144,7 @@ def apply_inverse(self, xy_in):
             return xy_in
         else:
             xy_in = np.asarray(xy_in, dtype=float)
-            xy_out = np.empty_like(xy_in)
-            _ge_41rt_inverse_distortion(
-                xy_out, xy_in, float(RHO_MAX), np.asarray(self.params)
+            xy_out = inverse_distortion.ge_41rt_inverse_distortion(
+                xy_in, float(RHO_MAX), np.asarray(self.params[:3])
             )
             return xy_out

hexrd/transforms/cpp_sublibrary/Makefile

@@ -0,0 +1,18 @@
+CXX = c++
+CXXFLAGS = -O3 -Wall -shared -fPIC -funroll-loops -Wno-maybe-uninitialized -fopenmp
+INCLUDES = -I/${CONDA_PREFIX}/include/eigen3/ -I/${CONDA_PREFIX}/include/xsimd -I/${CONDA_PREFIX}/include -I/${CONDA_PREFIX}/include/python3.11/
+SRCS_INVERSE_DISTORTION = src/inverse_distortion.cpp
+TARGET_DIR = ../../extensions/
+TARGET_NAME_INVERSE_DISTORTION = inverse_distortion
+EXTENSION_SUFFIX = $(shell python3-config --extension-suffix)
+
+all: inverse_distortion
+
+inverse_distortion: $(TARGET_DIR)$(TARGET_NAME_INVERSE_DISTORTION)$(EXTENSION_SUFFIX)
+
+$(TARGET_DIR)$(TARGET_NAME_INVERSE_DISTORTION)$(EXTENSION_SUFFIX): $(SRCS_INVERSE_DISTORTION)
+	$(CXX) $(CXXFLAGS) $(PYBIND_INCLUDES) $< -o $@ $(INCLUDES)
+
+.PHONY: clean inverse_distortion
+clean:
+	rm -f $(TARGET_DIR)$(TARGET_NAME_INVERSE_DISTORTION)$(EXTENSION_SUFFIX)

hexrd/transforms/cpp_sublibrary/src/inverse_distortion.cpp

@@ -0,0 +1,31 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/eigen.h>
+#include <Eigen/Core>
+
+namespace py = pybind11;
+constexpr double FOUR_THIRDS_PI = 4.1887902;
+constexpr double N_THREE_HALVES_SQRT_3 = -2.59807621;
+constexpr double TWO_OVER_SQRT_THREE = 1.154700538;
+
+Eigen::ArrayXXd ge_41rt_inverse_distortion(const Eigen::ArrayXXd& inputs, const double rhoMax, const Eigen::ArrayXd& params) {
+  Eigen::ArrayXd radii = inputs.matrix().rowwise().norm();
+  if (radii.maxCoeff() < 1e-10) {
+    return Eigen::ArrayXXd::Zero(inputs.rows(), inputs.cols());
+  }
+  Eigen::ArrayXd inverted_radii = radii.cwiseInverse();
+  Eigen::ArrayXd cosines = inputs.col(0) * inverted_radii;
+  Eigen::ArrayXd cosine_double_angles = 2*cosines.square() - 1;
+  Eigen::ArrayXd cosine_quadruple_angles = 2*cosine_double_angles.square() - 1;
+  Eigen::ArrayXd sqrt_p_is = rhoMax / (-params[0]*cosine_double_angles - params[1]*cosine_quadruple_angles - params[2]).sqrt();
+  Eigen::ArrayXd solutions = TWO_OVER_SQRT_THREE*sqrt_p_is*(acos(N_THREE_HALVES_SQRT_3*radii/sqrt_p_is)/3 + FOUR_THIRDS_PI).cos();
+  Eigen::ArrayXXd results = solutions.rowwise().replicate(inputs.cols()).array() * inputs * inverted_radii.rowwise().replicate(inputs.cols()).array();
+
+  return results;
+}
+
+PYBIND11_MODULE(inverse_distortion, m)
+{
+  m.doc() = "HEXRD inverse distribution plugin";
+
+  m.def("ge_41rt_inverse_distortion", &ge_41rt_inverse_distortion, "Inverse distortion for ge_41rt");
+}