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,152 +1,45 @@
 """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
 
-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
+from hexrd import constants as cnst
+from hexrd.extensions import inverse_distortion
 
-        xi, yi = in_[:, 0], in_[:, 1]
+RHO_MAX = 204.8  # max radius in mm for ge detector
 
-        # !!! included fix on ValueError for array--like in_
+# 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)
-        ri[ri < cnst.sqrt_epsf] = np.inf
-        ri_inv = 1.0/ri
-
+        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
@@ -158,11 +51,12 @@ def _ge_41rt_distortion(out, in_, rhoMax, params):
               + p1*ratio**p4*cos4ni
               + p2*ratio**p5
               + 1)*ri
-        out[:, 0] = ri*cosni
-        out[:, 1] = ri*sinni
-
-        return out
+        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) +
@@ -222,8 +116,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
+            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");
+}

hexrd/utils/json.py

@@ -0,0 +1,40 @@
+import base64
+import io
+import json
+
+import numpy as np
+
+ndarray_key = '!__hexrd_ndarray__'
+
+class NumpyEncoder(json.JSONEncoder):
+    def default(self, obj):
+        if isinstance(obj, np.ndarray):
+            # Write it as an npy file
+            with io.BytesIO() as bytes_io:
+                np.save(bytes_io, obj, allow_pickle=False)
+                data = bytes_io.getvalue()
+
+            return {
+                # Need to base64 encode it so it is json-valid
+                ndarray_key: base64.b64encode(data).decode('ascii')
+            }
+
+        # Let the base class default method raise the TypeError
+        return super().default(obj)
+
+
+class NumpyDecoder(json.JSONDecoder):
+    def __init__(self, *args, **kwargs):
+        kwargs = {
+            'object_hook': self.object_hook,
+            **kwargs,
+        }
+        super().__init__(*args, **kwargs)
+
+    def object_hook(self, obj):
+        if ndarray_key in obj:
+            data = base64.b64decode(obj[ndarray_key])
+            with io.BytesIO(data) as bytes_io:
+                return np.load(bytes_io)
+
+        return obj

pyproject.toml

@@ -1,2 +1,2 @@
 [build-system]
-requires = ["setuptools", "wheel", "numpy<1.27", "setuptools_scm[toml]"]
+requires = ["setuptools", "wheel", "numpy<1.27", "setuptools_scm[toml]", "pybind11>=2.11.0"]