handle pow(cx_mat, mat)

2022-09-14 14:26:50 +10:00
parent 17a983dd59
commit c6f29dc15c
3 changed files with 101 additions and 3 deletions
@@ -27,7 +27,7 @@ arma_inline
 typename
 enable_if2
  <
-  ( is_arma_type<T1>::value && is_arma_type<T2>::value && is_same_type<typename T1::elem_type, typename T2::elem_type>::value ),
+  ( is_arma_type<T1>::value && is_arma_type<T2>::value && is_same_type<typename T1::elem_type, typename T2::elem_type>::yes ),
  const Glue<T1, T2, glue_powext>
  >::result
 pow
@@ -43,11 +43,31 @@ pow



+template<typename T1, typename T2>
+arma_warn_unused
+arma_inline
+typename
+enable_if2
+  <
+  ( is_arma_type<T1>::value && is_arma_type<T2>::value && is_cx<typename T1::elem_type>::yes && is_same_type<typename T1::pod_type, typename T2::elem_type>::yes ),
+  const mtGlue<typename T1::elem_type, T1, T2, glue_powext_cx>
+  >::result
+pow
+  (
+  const T1& X,
+  const T2& Y
+  )
+  {
+  arma_extra_debug_sigprint();
+  
+  return mtGlue<typename T1::elem_type, T1, T2, glue_powext_cx>(X, Y);
+  }
+
+
+
 // TODO:    mat = pow(mat.each_col(),    mat)  (no promotion)
 // TODO:    mat = pow(mat.each_row(),    mat)  (no promotion)

-// TODO: cx_mat = pow(cx_mat,            mat)  (promotion; ? implement via mtGlue to allow preservation of vector type info)
-
 // TODO: cx_mat = pow(cx_mat.each_col(), mat)  (promotion)
 // TODO: cx_mat = pow(cx_mat.each_row(), mat)  (promotion)

@@ -1,3 +1,4 @@
+
 // SPDX-License-Identifier: Apache-2.0
 // 
 // Copyright 2008-2016 Conrad Sanderson (http://conradsanderson.id.au)
@@ -34,4 +35,16 @@ class glue_powext



+class glue_powext_cx
+  : public traits_glue_or
+  {
+  public:
+  
+  template<typename T1, typename T2> inline static void apply(Mat<typename T1::elem_type>& out, const mtGlue<typename T1::elem_type,T1,T2,glue_powext_cx>& X);
+  
+  template<typename T> inline static void apply(Mat< std::complex<T> >& out, const Mat< std::complex<T> >& A, const Mat<T>& B);
+  };
+
+
+
 //! @}
@@ -83,4 +83,69 @@ glue_powext::apply(Mat<eT>& out, const Mat<eT>& A, const Mat<eT>& B)



+//
+
+
+
+template<typename T1, typename T2>
+inline
+void
+glue_powext_cx::apply(Mat<typename T1::elem_type>& out, const mtGlue<typename T1::elem_type, T1, T2, glue_powext_cx>& X)
+  {
+  arma_extra_debug_sigprint();
+  
+  typedef typename T1::elem_type eT;
+  typedef typename T1::pod_type   T;
+  
+  const quasi_unwrap<T1> UA(X.A);
+  const quasi_unwrap<T2> UB(X.B);
+  
+  const Mat<eT>& A = UA.M;
+  const Mat< T>& B = UB.M;
+  
+  arma_debug_assert_same_size(A, B, "element-wise pow()");
+  
+  // TODO: investigate use of openmp
+  
+  if(UA.is_alias(out) && (UA.has_subview))
+    {
+    Mat<eT> tmp;
+    
+    glue_powext_cx::apply(tmp, A, B);
+    
+    out.steal_mem(tmp);
+    }
+  else
+    {
+    glue_powext_cx::apply(out, A, B);
+    }
+  }
+
+
+
+template<typename T>
+inline
+void
+glue_powext_cx::apply(Mat< std::complex<T> >& out, const Mat< std::complex<T> >& A, const Mat<T>& B)
+  {
+  arma_extra_debug_sigprint();
+  
+  typedef typename std::complex<T> eT;
+  
+  out.set_size(A.n_rows, A.n_cols);
+  
+  const uword N = out.n_elem;
+  
+        eT* out_mem = out.memptr();
+  const eT*   A_mem =   A.memptr();
+  const  T*   B_mem =   B.memptr();
+  
+  for(uword i=0; i<N; ++i)
+    {
+    out_mem[i] = std::pow(A_mem[i], B_mem[i]);
+    }
+  }
+
+
+
 //! @}