Merge pull request #607 from HEXRD/xtensor-support

Add xtensor support, remove Eigen

conda.recipe/meta.yaml

@@ -30,7 +30,8 @@ requirements:
     - numba
     - pybind11
     - xsimd
-    - eigen
+    - xtensor
+    - xtensor-python
 
   run:
     - appdirs

hexrd/transforms/cpp_sublibrary/Makefile

@@ -1,6 +1,6 @@
 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/
+CXXFLAGS = -O3 -Wall -shared -fPIC -funroll-loops -Wno-maybe-uninitialized -Wno-attributes -fopenmp
+INCLUDES = -I/${CONDA_PREFIX}/include -I/${CONDA_PREFIX}/include/python3.11/ -I/${CONDA_PREFIX}/lib/python3.11/site-packages/numpy/core/include
 SRCS_INVERSE_DISTORTION = src/inverse_distortion.cpp
 TARGET_DIR = ../../extensions/
 TARGET_NAME_INVERSE_DISTORTION = inverse_distortion

hexrd/transforms/cpp_sublibrary/src/inverse_distortion.cpp

@@ -1,31 +1,39 @@
+#define FORCE_IMPORT_ARRAY
+#define XTENSOR_USE_XSIMD
+
 #include <pybind11/pybind11.h>
-#include <pybind11/eigen.h>
-#include <Eigen/Core>
+#include <xtensor/xmath.hpp>
+#include <xtensor/xarray.hpp>
+#include <xtensor-python/pyarray.hpp>
+#include <xtensor/xview.hpp>
 
 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;
+xt::pyarray<double> ge_41rt_inverse_distortion(const xt::pyarray<double>& inputs, const double rhoMax, const xt::pyarray<double>& params) {
+    auto radii = xt::sqrt(xt::sum(xt::square(inputs), {1}));
+
+    if (xt::amax(radii)() < 1e-10) {
+        return xt::zeros<double>({inputs.shape()[0], inputs.shape()[1]});
+    }
+
+    auto inverted_radii = xt::pow(radii, -1);
+    xt::pyarray<double> cosines = xt::view(inputs, xt::all(), 0) * inverted_radii;
+    auto cosine_double_angles = 2 * xt::square(cosines) - 1;
+    auto cosine_quadruple_angles = 2 * xt::square(cosine_double_angles) - 1;
+    auto sqrt_p_is = rhoMax / xt::sqrt(-params[0] * cosine_double_angles - params[1] * cosine_quadruple_angles - params[2]);
+    auto solutions = TWO_OVER_SQRT_THREE * sqrt_p_is * xt::cos(xt::acos(N_THREE_HALVES_SQRT_3 * radii / sqrt_p_is) / 3 + FOUR_THIRDS_PI);
+    xt::pyarray<double> results = xt::view(solutions, xt::all(), xt::newaxis()) * inputs * xt::view(inverted_radii, xt::all(), xt::newaxis());
+
+    return results;
 }
 
 PYBIND11_MODULE(inverse_distortion, m)
 {
-  m.doc() = "HEXRD inverse distribution plugin";
+    xt::import_numpy();
+    m.doc() = "HEXRD inverse distribution plugin";
 
-  m.def("ge_41rt_inverse_distortion", &ge_41rt_inverse_distortion, "Inverse distortion for ge_41rt");
+    m.def("ge_41rt_inverse_distortion", &ge_41rt_inverse_distortion, "Inverse distortion for ge_41rt");
 }

scripts/install/install_build_dependencies.py

@@ -36,10 +36,68 @@ def download_and_extract_tgz(url, md5sum, path):
     temp_file.unlink()
 
 
+def download_xtensor(path):
+    url = 'https://github.com/xtensor-stack/xtensor/archive/refs/tags/0.24.7.tar.gz'  # noqa
+    md5sum = 'e21a14d679db71e92a703bccd3c5866a'
+    out_dir_name = 'xtensor-0.24.7'
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        download_and_extract_tgz(url, md5sum, temp_dir)
+
+        target_path = Path(path) / 'xtensor'
+        if target_path.exists():
+            shutil.rmtree(target_path)
+
+        os.makedirs(path, exist_ok=True)
+        shutil.move(str(Path(temp_dir) / out_dir_name / 'include/xtensor'),
+                    str(Path(path) / 'xtensor/xtensor'))
+
+    return str(target_path)
+
+
+def download_xtensor_python(path):
+    url = 'https://github.com/xtensor-stack/xtensor-python/archive/refs/tags/0.26.1.tar.gz'  # noqa
+    md5sum = '5d05edf71ac948dc620968229320c791'
+    out_dir_name = 'xtensor-python-0.26.1'
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        download_and_extract_tgz(url, md5sum, temp_dir)
+
+        target_path = Path(path) / 'xtensor-python'
+        if target_path.exists():
+            shutil.rmtree(target_path)
+
+        os.makedirs(path, exist_ok=True)
+        shutil.move(
+            str(Path(temp_dir) / out_dir_name / 'include/xtensor-python'),
+            str(Path(path) / 'xtensor-python/xtensor-python'))
+
+    return str(target_path)
+
+
+def download_xtl(path):
+    url = 'https://github.com/xtensor-stack/xtl/archive/refs/tags/0.7.7.tar.gz'  # noqa
+    md5sum = '6df56ae8bc30471f6773b3f18642c8ab'
+    out_dir_name = 'xtl-0.7.7'
+
+    with tempfile.TemporaryDirectory() as temp_dir:
+        download_and_extract_tgz(url, md5sum, temp_dir)
+
+        target_path = Path(path) / 'xtl'
+        if target_path.exists():
+            shutil.rmtree(target_path)
+
+        os.makedirs(path, exist_ok=True)
+        shutil.move(str(Path(temp_dir) / out_dir_name / 'include/xtl'),
+                    str(Path(path) / 'xtl/xtl'))
+
+    return str(target_path)
+
+
 def download_xsimd(path):
-    url = 'https://github.com/xtensor-stack/xsimd/archive/refs/tags/7.6.0.tar.gz'  # noqa
-    md5sum = '6e52c2af8b3cb4688993a0e70825f4e8'
-    out_dir_name = 'xsimd-7.6.0'
+    url = 'https://github.com/xtensor-stack/xsimd/archive/refs/tags/12.1.1.tar.gz'  # noqa
+    md5sum = 'e8887de343bd6036bdfa1f4a4752dc64'
+    out_dir_name = 'xsimd-12.1.1'
 
     with tempfile.TemporaryDirectory() as temp_dir:
         download_and_extract_tgz(url, md5sum, temp_dir)
@@ -87,14 +145,16 @@ def download_pybind11(path):
         shutil.move(str(Path(temp_dir) / out_dir_name / 'include/pybind11'),
                     str(Path(path) / 'pybind11/pybind11'))
 
-
     return str(target_path)
 
 
 INSTALL_FUNCTIONS = {
     'eigen3': download_eigen3,
-    'xsimd': download_xsimd,
     'pybind11': download_pybind11,
+    'xsimd': download_xsimd,
+    'xtensor': download_xtensor,
+    'xtensor-python': download_xtensor_python,
+    'xtl': download_xtl,
 }
 
 

setup.py

@@ -115,15 +115,17 @@ def get_cpp_extensions():
     cpp_transform_pkgdir = Path('hexrd') / 'transforms/cpp_sublibrary'
     src_files = [str(cpp_transform_pkgdir / 'src/inverse_distortion.cpp')]
 
-    extra_compile_args = ['-O3', '-Wall', '-shared', '-std=c++11',
+    extra_compile_args = ['-O3', '-Wall', '-shared', '-std=c++14',
                           '-funroll-loops']
     if not sys.platform.startswith('win'):
         extra_compile_args.append('-fPIC')
 
     # Define include directories
     include_dirs = [
-        get_include_path('eigen3'),
         get_include_path('xsimd'),
+        get_include_path('xtensor'),
+        get_include_path('xtensor-python'),
+        get_include_path('xtl'),
         get_pybind11_include_path(),
         numpy.get_include(),
     ]