evolution

python/misc.cc

@@ -30,12 +30,12 @@
 #include <vector>
 
 #include "mr_util/infra/mav.h"
+#include "mr_util/infra/transpose.h"
 #include "mr_util/math/fft.h"
 #include "mr_util/math/constants.h"
 #include "mr_util/math/gl_integrator.h"
 #include "mr_util/bindings/pybind_utils.h"
 #include "python/alm.h"
-#include "python/transpose.h"
 
 namespace mr {
 
@@ -180,6 +180,31 @@ py::array py_ascontiguousarray(const py::array &in)
   MR_fail("unsupported datatype");
   }
 
+template<typename T> py::array tphelp2(const py::array &in, py::array &out)
+  {
+  auto in2 = to_fmav<T>(in, false);
+  auto out2 = to_fmav<T>(out, true);
+  transpose(in2, out2, [](const T &in, T &out){out=in;});
+  return out;
+  }
+
+py::array py_transpose(const py::array &in, py::array &out)
+  {
+  if (isPyarr<float>(in))
+    return tphelp2<float>(in, out);
+  if (isPyarr<double>(in))
+    return tphelp2<double>(in, out);
+  if (isPyarr<complex<float>>(in))
+    return tphelp2<complex<float>>(in, out);
+  if (isPyarr<complex<double>>(in))
+    return tphelp2<complex<double>>(in, out);
+  if (isPyarr<int>(in))
+    return tphelp2<int>(in, out);
+  if (isPyarr<long>(in))
+    return tphelp2<long>(in, out);
+  MR_fail("unsupported datatype");
+  }
+
 
 const char *misc_DS = R"""(
 Various unsorted utilities
@@ -201,6 +226,7 @@ void add_misc(py::module &msup)
     "has_np"_a, "has_sp"_a, "out"_a=py::none());
 
   m.def("ascontiguousarray",&py_ascontiguousarray, "in"_a);
+  m.def("transpose",&py_transpose, "in"_a, "out"_a);
   }
 
 }

python/transpose.h → src/mr_util/infra/transpose.h

@@ -40,7 +40,7 @@ using namespace std;
 using shape_t=fmav_info::shape_t;
 using stride_t=fmav_info::stride_t;
 
-template<typename T1, typename T2> void rearrange(T1 &v, const T2 &idx)
+template<typename T1, typename T2> inline void rearrange(T1 &v, const T2 &idx)
   {
   T1 tmp(v);
   for (size_t i=0; i<idx.size(); ++i)
@@ -88,33 +88,35 @@ template<typename T, typename Func> void sthelper1(const T *in, T *out,
   for (size_t i=0; i<s0; ++i, in+=sti0, out+=sto0)
     func(*in, *out);
   }
+
 bool critical(ptrdiff_t s)
   {
   s = (s>=0) ? s : -s;
   return (s>4096) && ((s&(s-1))==0);
   }
-template<typename T, typename Func> void sthelper2(const T *in, T *out,
-  size_t s0, size_t s1, ptrdiff_t sti0, ptrdiff_t sti1, ptrdiff_t sto0,
-  ptrdiff_t sto1, Func func)
+
+template<typename T, typename Func> void sthelper2(const T * MRUTIL_RESTRICT in,
+  T * MRUTIL_RESTRICT out, size_t s0, size_t s1, ptrdiff_t sti0, ptrdiff_t sti1,
+  ptrdiff_t sto0, ptrdiff_t sto1, Func func)
   {
   if ((sti0<=sti1) && (sto0<=sto1)) // no need to block
     {
-    for (size_t i=0; i<s1; ++i, in+=sti1, out+=sto1)
+    for (size_t i1=0; i1<s1; ++i1, in+=sti1, out+=sto1)
       {
       auto pi0=in;
       auto po0=out;
-      for (size_t j=0; j<s0; ++j, pi0+=sti0, po0+=sto0)
+      for (size_t i0=0; i0<s0; ++i0, pi0+=sti0, po0+=sto0)
         func(*pi0, *po0);
       }
     return;
     }
   if ((sti0>=sti1) && (sto0>=sto1)) // no need to block
     {
-    for (size_t i=0; i<s0; ++i, in+=sti0, out+=sto0)
+    for (size_t i0=0; i0<s0; ++i0, in+=sti0, out+=sto0)
       {
       auto pi1=in;
       auto po1=out;
-      for (size_t j=0; j<s1; ++j, pi1+=sti1, po1+=sto1)
+      for (size_t i1=0; i1<s1; ++i1, pi1+=sti1, po1+=sto1)
         func(*pi1, *po1);
       }
     return;