Resolve matrix output dimensionality issues

HEXRD · Jan 30, 2024 · 9caa1ee · 9caa1ee
1 parent 7f56e54
commit 9caa1ee
Showing 1 changed file with 3 additions and 3 deletions.
diff --git a/hexrd/transforms/cpp_sublibrary/src/inverse_distortion.cpp b/hexrd/transforms/cpp_sublibrary/src/inverse_distortion.cpp
@@ -13,18 +13,18 @@ constexpr double TWO_OVER_SQRT_THREE = 1.154700538;
 
 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 = inputs * xt::view(inverted_radii, xt::newaxis());
+    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 = solutions * inputs * xt::view(inverted_radii, xt::newaxis());
+    xt::pyarray<double> results = xt::view(solutions, xt::all(), xt::newaxis()) * inputs * xt::view(inverted_radii, xt::all(), xt::newaxis());
 
     return results;
 }