xtensor_math

Complex.H

#ifndef Complex_H
#define Complex_H

#include <complex>
typedef std::complex<double> Complex;
typedef std::complex<double> complex_t;
typedef std::complex<float> complex64_t;
typedef std::complex<long double> complex256_t;

// decomplexify ---------------------------------------------------------------
template <typename T>
struct decomplexify_impl
{
  typedef T type;
};

template <typename ELT>
struct decomplexify_impl<std::complex<ELT> >
{
  typedef ELT type;
};

#include <type_traits>

template<typename T>
struct decomplexify
{
  typedef typename std::remove_reference<T>::type TT;
  typedef typename decomplexify_impl<TT>::type type;
};


// complexify -----------------------------------------------------------------
template <typename T>
struct complexify
{
  typedef std::complex<T> type;
};

template <typename ELT>
struct complexify<std::complex<ELT> >
{
  typedef std::complex<ELT> type;
};

#endif

test_xtensor_math.py

import numpy as np
from xtensor_math import mixer

M = mixer (0.0)
u = np.ones (100, dtype=complex)
v = M(u)

xtensor_math.cc

#define FORCE_IMPORT_ARRAY
#define _GNU_SOURCE
#define HAVE_CBLAS 1

#include <cmath>
#include <xtensor/xmath.hpp>
#include <xtensor/xtensor.hpp>
#include <xtensor/xarray.hpp>
#include "xtensor-python/pytensor.hpp"
#include "xtensor-python/pyarray.hpp"
#include "pybind11/pybind11.h"
#include <xtensor/xtensor_simd.hpp>
#include <xtensor/xexpression.hpp>
#include <xtensor/xcomplex.hpp>
#include "xtensor/xnorm.hpp"
#include "xtensor-python/pyvectorize.hpp"
#include "xtensor-blas/xlinalg.hpp"
#include "Complex.H"

namespace py = pybind11;

template <class E>
auto xsincos(E&& inp)
{
    return xt::make_lambda_xfunction([](auto in)
                              { std::array<decltype(in), 2> res; sincos(in, &res[0], &res[1]); return res;},
                             std::forward<E>(inp));
}

py::tuple pysincos (xt::pytensor<double,1> const& x) {
  xt::xtensor<std::array<double,2>, 1> res = xsincos (x);
  xt::pytensor<double,1> c = xt::pytensor<double,1>::from_shape(res.shape());
  xt::pytensor<double,1> s = xt::pytensor<double,1>::from_shape(res.shape());
  for (int i = 0; i < x.size(); ++i) {
    s[i] = res[i][0];
    c[i] = res[i][1];
  }
  return py::make_tuple(s, c);
}

py::tuple pysincos2 (xt::pytensor<double,1> const& x) {
  xt::pytensor<double,1> c = xt::pytensor<double,1>::from_shape(x.shape());
  xt::pytensor<double,1> s = xt::pytensor<double,1>::from_shape(x.shape());
  for (int i = 0; i < x.size(); ++i) {
    sincos (x[i], &s[i], &c[i]);
  }
  return py::make_tuple(s, c);
}

// Either of the below will work
template<typename T, typename U, typename V>
void xxsincos(xt::xcontainer<T> const& in, xt::xcontainer<U>& sin, xt::xcontainer<V>& cos)
// template<class E>
// void xxsincos(E const& in, E& sin, E& cos)
// Alternatively, can use xexpression, but then need
// auto& sin = sine.derived_cast()...
{
  constexpr std::size_t simd_size = xsimd::simd_type<double>::size;

  using simd_batch = xsimd::batch<double>;
  simd_batch b_sin, b_cos;
  int end = in.size()/simd_size * simd_size;
  std::size_t i = 0;
  for (; i < end; i += simd_size)
    {
      auto b_in = in.template load_simd<xsimd::unaligned_mode>(i);
      std::tie(b_sin, b_cos) = xsimd::sincos(b_in);
      sin.template store_simd<xsimd::unaligned_mode>(i, b_sin);
      cos.template store_simd<xsimd::unaligned_mode>(i, b_cos);
    }
  for (; i < in.size(); ++i) {
    ::sincos (in[i], &sin[i], &cos[i]);
  }
}

py::tuple pyxsincos(xt::pyarray<double>const& in)
{
  xt::pyarray<double> sin = xt::pyarray<double>::from_shape(in.shape());
  xt::pyarray<double> cos = xt::pyarray<double>::from_shape(in.shape());
  xxsincos (in, sin, cos);
  return py::make_tuple (std::move(sin), std::move(cos));
}

typedef std::complex<double> complex_t;

struct mixer {
  double delta;
  double phase;
  mixer(double _freq, double _phase = 0) : delta (_freq*2*M_PI), phase (_phase) {}

  auto generate2 (xt::pytensor<complex_t,1> const& in) {
    int N = in.size();
    double endphase = phase + delta*N;
    xt::xtensor<double,1> p = xt::linspace (phase, endphase, N, false);
    xt::xtensor<double,1> sin = xt::xtensor<double,1>::from_shape(in.shape());
    xt::xtensor<double,1> cos = xt::xtensor<double,1>::from_shape(in.shape());
    xxsincos (p, sin, cos);
    xt::pytensor<complex_t,1> ret = xt::pytensor<complex_t,1>::from_shape(in.shape());
    xt::real(ret) = cos * xt::real(in) - sin * xt::imag(in);
    xt::imag(ret) = cos * xt::imag(in) + sin * xt::real(in);
    phase = endphase;
    return py::make_tuple (ret, p);
  }

  auto operator()(xt::pytensor<complex_t,1> const& in) {
    int N = in.size();
    double endphase = phase + delta*N;
    xt::pytensor<double,1> p = xt::linspace (phase, endphase, N, false);
    xt::xtensor<double,1> sin = xt::xtensor<double,1>::from_shape(in.shape());
    xt::xtensor<double,1> cos = xt::xtensor<double,1>::from_shape(in.shape());
    xxsincos (p, sin, cos);
    xt::pytensor<complex_t,1> ret = xt::pytensor<complex_t,1>::from_shape(in.shape());
    xt::real(ret) = cos * xt::real(in) - sin * xt::imag(in);
    xt::imag(ret) = cos * xt::imag(in) + sin * xt::real(in);
    phase = endphase;
    return ret;
  }
};

// Either of the below will work
template<typename T>
void xxtanh(xt::xcontainer<T> const& in, xt::xcontainer<T>& out)
// template<class E>
// void xxsincos(E const& in, E& sin, E& cos)
// Alternatively, can use xexpression, but then need
// auto& sin = sine.derived_cast()...
{
  using simd_batch = xsimd::batch<double>;
  simd_batch b_tanh;
  int end = in.size()/4 * 4;
  std::size_t i = 0;
  for (; i < end; i += 4)
    {
      auto b_in = in.template load_simd<xsimd::unaligned_mode>(i);
      b_tanh = xsimd::tanh(b_in);
      out.template store_simd<xsimd::unaligned_mode>(i, b_tanh);
    }
  for (; i < in.size(); ++i) {
    out[i] = ::tanh (in[i]);
  }
}

xt::pyarray<double> pyxtanh(xt::pyarray<double>const& in)
{
  xt::pyarray<double> out = xt::pyarray<double>::from_shape(in.shape());
  xxtanh (in, out);
  return out;
}

inline double real (double x) { return x; }
inline double imag (double x) { return 0; }
inline double conj (double x) { return x; }
inline double norm (double x) { return x*x; }

template <class T>
struct order2 {
  typedef typename decomplexify<T>::type scalar_t;
  int N;
  T sumx;
  scalar_t sumx2;

  order2() : N(0), sumx(0), sumx2(0) {}

  order2(int _N, T _sumx, scalar_t _sumx2) :
	     N (_N), sumx (_sumx), sumx2 (_sumx2) {} // pickle
  
  void call1 (T x) {
    sumx += x;
    sumx2 += real(x)*real(x) + imag(x)*imag(x);
    N++;
  }
  
  T mean () const {
    return sumx / scalar_t (N);
  }

  scalar_t var () const {
    return sumx2 / scalar_t (N) - norm (sumx / scalar_t (N));
  }

  scalar_t rms () const {
    return std::sqrt (var());
  }
};

template <class T>
using ndarray = xt::pyarray<T, xt::layout_type::row_major>;

template <class T, int N>
using ndtensor = xt::pytensor<T, N, xt::layout_type::row_major>;

PYBIND11_MODULE (xtensor_math, m) {
  xt::import_numpy();

  m.def ("simd_size", []() { return xsimd::simd_type<double>::size; });

  m.def ("sincos", &pysincos);
  m.def ("sincos2", &pysincos2);
  m.def ("xsincos", &pyxsincos);

  py::class_<mixer>(m, "mixer")
    .def (py::init<double,double>(),
	  py::arg("freq"),
	  py::arg("phase")=0
	  )
    .def ("__call__", &mixer::operator())
    .def ("Generate2", &mixer::generate2) // imitate old behavior
    .def_readwrite ("delta", &mixer::delta)
    .def_readwrite ("phase", &mixer::phase)
    .def_property ("freq", [](mixer const& m) { return m.delta/(2*M_PI); },
		   [](mixer & m, double f) { m.delta = 2 * M_PI * f; }
		   )
    ;
  m.def("mag_sqr", [](xt::pyarray<double> const& a) {
      // print ("double");
      return py::vectorize([](double x) {return x * x;})(a);
    },
    py::arg("in").noconvert()
    );
  m.def("mag_sqr", [](xt::pyarray<float> const& a) {
      // print ("float");
      return py::vectorize([](float x) {return x * x;})(a);
    },
    py::arg("in").noconvert()
    );
  m.def("mag_sqr", [](xt::pyarray<std::complex<double>> const&  a) {
      // print ("complex");
      return py::vectorize([](std::complex<double> x) {return real(x) * real(x) + imag(x) * imag(x);})(a);
    },
    py::arg("in").noconvert()
    );
  m.def("mag_sqr", [](xt::pyarray<std::complex<float>> const& a) {
      return py::vectorize([](std::complex<float> x) {return real(x) * real(x) + imag(x) * imag(x);})(a);
    },
    py::arg("in").noconvert()
    );
  m.def("mag_sqr", [](std::complex<double> x) { return real(x) * real(x) + imag(x) * imag(x);}, py::arg("in"));
  m.def("mag_sqr", [](double x) { return x * x; }, py::arg("in"));
  //m.def("mag_sqr", [](float x) { return x * x; }, py::arg("in").noconvert());
  //m.def("mag_sqr", [](std::complex<float> x) { return real(x) * real(x) + imag(x) * imag(x);}, py::arg("in").noconvert());

  m.def ("norm_2", [](std::complex<double> const& x) { return xt::norm_l2(x); },
	 py::arg("in").noconvert()
	 );

  m.def ("norm_2", [](double const& x) { return xt::norm_l2(x); },
	 py::arg("in").noconvert()
	 );

  m.def ("norm_2",  [](xt::pyarray<std::complex<double>> const& x) { return std::real(xt::norm_l2 (x)()); },
        py::arg("in").noconvert()
	 );
  m.def("norm_2", [](xt::pyarray<double> const& x) { return xt::norm_l2 (x)(); },
        py::arg("in").noconvert()
        );
  m.def ("norm_2",  [](xt::pyarray<std::complex<float>> const& x) { return std::real(xt::norm_l2 (x)()); },
        py::arg("in").noconvert()
	 );
  m.def("norm_2", [](xt::pyarray<float> const& x) { return xt::norm_l2 (x)(); },
        py::arg("in").noconvert()
        );

  m.def ("norm_sq", [](std::complex<double> const& x) { return xt::norm_sq(x); },
	 py::arg("in").noconvert()
	 );

  m.def ("norm_sq", [](double const& x) { return xt::norm_sq(x); },
	 py::arg("in").noconvert()
	 );
  
  m.def ("norm_sq", [](xt::pyarray<double> const& x) { return xt::norm_sq (x)(); },
	 py::arg("in").noconvert()
	 );

  m.def ("norm_sq", [](xt::pyarray<double> const& x, std::vector<size_t>const& ax) { return xt::norm_sq (x, ax, xt::evaluation_strategy::immediate)(); },
	 py::arg("in").noconvert(),
	 py::arg("axes")
	 );

  m.def ("norm_sq", [](xt::pyarray<double> const& x, size_t ax) {
    auto shape = std::vector<size_t>{ax};
    return xt::norm_sq (x, shape, xt::evaluation_strategy::immediate)(); },
    py::arg("in").noconvert(),
    py::arg("axes")
    );

  m.def ("norm_sq", [](xt::pyarray<std::complex<double>> const& x) { return xt::norm_sq (x)(); },
	 py::arg("in").noconvert()
	 );

  m.def ("norm_sq", [](xt::pyarray<std::complex<double>> const& x, std::vector<size_t>const& ax) { return xt::norm_sq (x, ax, xt::evaluation_strategy::immediate)(); },
	 py::arg("in").noconvert(),
	 py::arg("axes")
	 );

  m.def ("norm_sq", [](xt::pyarray<std::complex<double>> const& x, size_t ax) {
    auto shape = std::vector<size_t>{ax};
    return xt::norm_sq (x, shape, xt::evaluation_strategy::immediate)(); },
    py::arg("in").noconvert(),
    py::arg("axes")
    );

  m.def("tanh", [](xt::pyarray<double> const& x) { return xt::pyarray<double>(xt::tanh (x)); },
        py::arg("in").noconvert()
        );
  m.def("expm1", [](xt::pyarray<double> const& x) { return xt::pyarray<double>(xt::expm1 (x)); },
        py::arg("in").noconvert()
        );
  m.def("log1p", [](xt::pyarray<double> const& x) { return xt::pyarray<double>(xt::log1p (x)); },
        py::arg("in").noconvert()
        );

  m.def("xtanh", [](xt::pyarray<double> const & x) { return pyxtanh(x); },
        py::arg("in").noconvert()
        );
	  
  m.def("arg", [](xt::pyarray<std::complex<double>> const& a) {
    return xt::pyarray<double>(xt::arg (a)); },
	py::arg("in").noconvert()
	);


  m.def("dot", [](xt::pytensor<std::complex<double>,1> const& a, xt::pytensor<std::complex<double>,1> const& b) { return xt::linalg::vdot (b, a); },
        py::arg("a").noconvert(), py::arg("b").noconvert()
	);

  m.def("dot", [](xt::pytensor<double,1> const& a, xt::pytensor<double,1> const& b) { return xt::linalg::vdot (b, a); },
        py::arg("a").noconvert(), py::arg("b").noconvert()
	);

  py::class_<order2<std::complex<double>>> (m, "stat2nd_complex")
    .def (py::init<>())
    .def (py::init<int, std::complex<double>, double>(),"for pickle")
    .def ("__iadd__", [](order2<std::complex<double>> & O, std::complex<double> x) {
      O.sumx += x; O.sumx2 += xt::norm_sq(x); O.N++; return O; })
			
    .def ("__iadd__", [](order2<std::complex<double>> & O, ndtensor<std::complex<double>,1> const& v) {
			O.sumx += xt::sum (v)();
			O.sumx2 += xt::norm_sq (v)();
			O.N += v.shape()[0];
			return O;
		      },
      py::arg("v").noconvert())
    .def (py::pickle (
		      [](const order2<std::complex<double>> &o) {
			py::dict d;
			d["sumx"] = o.sumx;
			d["N"] = o.N;
			d["sumx2"] = o.sumx2;
			return d;
		      },
		      [](py::dict d) {
			return new order2<std::complex<double>> (d["N"].cast<int>(), d["sumx"].cast<std::complex<double>>(), d["sumx2"].cast<double>());
		      }))
    .def_property_readonly ("mean", &order2<std::complex<double>>::mean)
    .def_property_readonly ("var", &order2<std::complex<double>>::var)
    .def_property_readonly ("rms", &order2<std::complex<double>>::rms)
    .def_readonly ("N", &order2<std::complex<double>>::N)
    .def_readonly ("sumx", &order2<std::complex<double>>::sumx)
    .def_readonly ("sumx2", &order2<std::complex<double>>::sumx2)
    ;

  py::class_<order2<double>> (m, "stat2nd_double")
    .def (py::init<>())
    .def (py::init<int, double, double>(),"for pickle")
    .def ("__iadd__", [](order2<double> & O, double x) {
      O.sumx += x; O.sumx2 += xt::norm_sq(x); O.N++; return O;})
    .def ("__iadd__", [](order2<double> & O, ndtensor<double,1> const& v) {
			O.sumx += xt::sum (v)();
			O.sumx2 += xt::norm_sq (v)();
			O.N += v.shape()[0];
			return O;
		      },
      py::arg("v").noconvert())
    .def (py::pickle (
		      [](const order2<double> &o) {
			py::dict d;
			d["sumx"] = o.sumx;
			d["N"] = o.N;
			d["sumx2"] = o.sumx2;
			return d;
		      },
		      [](py::dict d) {
			return new order2<double> (d["N"].cast<int>(), d["sumx"].cast<double>(), d["sumx2"].cast<double>());
		      }))
    .def_property_readonly ("mean", &order2<double>::mean)
    .def_property_readonly ("var", &order2<double>::var)
    .def_property_readonly ("rms", &order2<double>::rms)
    .def_readonly ("N", &order2<double>::N)
    .def_readonly ("sumx", &order2<double>::sumx)
    .def_readonly ("sumx2", &order2<double>::sumx2)
    ;

}