.

include/sctl/sph_harm.hpp

@@ -36,7 +36,7 @@ template <class Real> class SphericalHarmonics{
 
     static void SHC2Grid(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, Long Nt_out, Long Np_out, Vector<Real>* X_out, Vector<Real>* X_theta_out=nullptr, Vector<Real>* X_phi_out=nullptr);
 
-    static void SHCEval(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, const Vector<Real>& theta_phi_in, Vector<Real>& X_out);
+    static void SHCEval(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, const Vector<Real>& cos_theta_phi_in, Vector<Real>& X_out);
 
     static void SHC2Pole(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, Vector<Real>& P_out);
 
@@ -48,130 +48,17 @@ template <class Real> class SphericalHarmonics{
 
     static void VecSHC2Grid(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, Long Nt_out, Long Np_out, Vector<Real>& X_out);
 
-    static void VecSHCEval(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, const Vector<Real>& theta_phi_in, Vector<Real>& X_out);
+    static void VecSHCEval(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, const Vector<Real>& cos_theta_phi_in, Vector<Real>& X_out);
 
-    static void StokesEvalSL(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, const Vector<Real>& theta_phi_in, Vector<Real>& X_out);
+    static void StokesEvalSL(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, const Vector<Real>& coord_in, Vector<Real>& X_out);
 
-    static void StokesEvalDL(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, const Vector<Real>& theta_phi_in, Vector<Real>& X_out);
+    static void StokesEvalDL(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, const Vector<Real>& coord_in, Vector<Real>& X_out);
 
 
-    static void test3() {
-      int k=0, n=1, m=0;
-      for (int k=0;k<3;k++) {
-        for (int n=0;n<3;n++) {
-          for (int m=0;m<=n;m++) {
-
-      int p=5;
-      int dof = 1;
-      int Nt = p+1;
-      int Np = 2*p+2;
-      int Ngrid = Nt * Np;
-      int Ncoeff = (p + 1) * (p + 2);
-      Vector<Real> sin_theta, cos_theta = LegendreNodes(Nt-1);
-      for (Long i=0;i<cos_theta.Dim();i++) sin_theta.PushBack(sqrt<Real>(1-cos_theta[i]*cos_theta[i]));
-
-      auto print_coeff = [&](Vector<Real> S) {
-        Long idx=0;
-        Long dof = S.Dim() / Ncoeff;
-        for (Long k=0;k<dof;k++) {
-          for (Long n=0;n<=p;n++) {
-            std::cout<<Vector<Real>(2*n+2, S.begin()+idx);
-            idx+=2*n+2;
-          }
-        }
-        std::cout<<'\n';
-      };
-
-      Vector<Real> coeff(dof*Ncoeff); coeff=0;
-      auto write_coeff = [&](Vector<Real>& coeff, Complex<Real> c, Long k, Long n, Long m) {
-        if (0 <= m && m <= n && n <= p) {
-          Long idx = k*Ncoeff + n*(n+1)+ 2*m;
-          coeff[idx+0] = c.real;
-          coeff[idx+1] = c.imag;
-        }
-      };
-      write_coeff(coeff, Complex<Real>(1,1.3),0,n,m);
-      //print_coeff(coeff);
-
-      Vector<Real> Y, Yt, Yp;
-      SHC2Grid(coeff, SHCArrange::ROW_MAJOR, p, Nt, Np, &Y, &Yt, &Yp);
-      //std::cout<<Matrix<Real>(Nt, Np, Y.begin())<<'\n';
-      //std::cout<<Matrix<Real>(Nt, Np, Yt.begin())<<'\n';
-      //std::cout<<Matrix<Real>(Nt, Np, Yp.begin())<<'\n';
-
-      Vector<Real> gradYt = Yt;
-      Vector<Real> gradYp = Yp;
-      for (Long i=0;i<Nt;i++) {
-        for (Long j=0;j<Np;j++) {
-          gradYp[i*Np+j] /= sin_theta[i];
-        }
-      }
-
-      Vector<Real> Vr, Vt, Vp;
-      Vector<Real> Wr, Wt, Wp;
-      Vector<Real> Xr, Xt, Xp;
-
-      Vr = Y*(-n-1);
-      Vt = gradYt;
-      Vp = gradYp;
-
-      Wr = Y*n;
-      Wt = gradYt;
-      Wp = gradYp;
-
-      Xr = Y*0;
-      Xt = gradYp;
-      Xp = gradYt * (-1);
-
-      Vector<Real> SS(COORD_DIM * Ngrid);
-      SS=0;
-      if (k == 0) {
-        Vector<Real>(Ngrid, SS.begin() + 0*Ngrid, false) = Vr;
-        Vector<Real>(Ngrid, SS.begin() + 1*Ngrid, false) = Vt;
-        Vector<Real>(Ngrid, SS.begin() + 2*Ngrid, false) = Vp;
-      }
-      if (k == 1) {
-        Vector<Real>(Ngrid, SS.begin() + 0*Ngrid, false) = Wr;
-        Vector<Real>(Ngrid, SS.begin() + 1*Ngrid, false) = Wt;
-        Vector<Real>(Ngrid, SS.begin() + 2*Ngrid, false) = Wp;
-      }
-      if (k == 2) {
-        Vector<Real>(Ngrid, SS.begin() + 0*Ngrid, false) = Xr;
-        Vector<Real>(Ngrid, SS.begin() + 1*Ngrid, false) = Xt;
-        Vector<Real>(Ngrid, SS.begin() + 2*Ngrid, false) = Xp;
-      }
-
-      Vector<Real> SSS;
-      {
-        Vector<Real> coeff(COORD_DIM*Ncoeff);
-        coeff=0;
-        write_coeff(coeff, Complex<Real>(1,1.3),k,n,m);
-        //print_coeff(coeff);
-        VecSHC2Grid(coeff, SHCArrange::ROW_MAJOR, p, Nt, Np, SSS);
-      }
-
-      //std::cout<<Matrix<Real>(COORD_DIM*Nt, Np, SS.begin())<<'\n';
-      //std::cout<<Matrix<Real>(COORD_DIM*Nt, Np, SSS.begin())<<'\n';
-      //std::cout<<Matrix<Real>(COORD_DIM*Nt, Np, SSS.begin()) - Matrix<Real>(COORD_DIM*Nt, Np, SS.begin())<<'\n';
-
-      auto err=SSS-SS;
-      Real max_err=0;
-      for (auto x:err) max_err=std::max(max_err, fabs(x));
-      std::cout<<max_err<<' ';
-
-          }
-          std::cout<<'\n';
-        }
-        std::cout<<'\n';
-      }
-
-      Clear();
-    }
-
-    static void test2() {
-      int p = 6;
+    static void test_stokes() {
+      int p = 4;
       int dof = 3;
-      int Nt = p+1, Np = 2*p+2;
+      int Nt = p+1, Np = 2*p+1;
 
       auto print_coeff = [&](Vector<Real> S) {
         Long idx=0;
@@ -186,82 +73,40 @@ template <class Real> class SphericalHarmonics{
 
       Vector<Real> f(dof * Nt * Np);
       { // Set f
-        Vector<Real> leg_nodes = LegendreNodes(Nt-1);
         for (Long i = 0; i < Nt; i++) {
           for (Long j = 0; j < Np; j++) {
-            Real cos_theta = leg_nodes[i];
-            Real sin_theta = sqrt<Real>(1-cos_theta*cos_theta);
-            Real phi = 2 * const_pi<Real>() * j / Np;
-            Real r = 1;
-
-            Real x = r * cos_theta;
-            Real y = r * sin_theta * cos<Real>(phi);
-            Real z = r * sin_theta * sin<Real>(phi);
-
-            // Unit vectors in polar coordinates
-            Real xr = cos_theta;
-            Real yr = sin_theta * cos<Real>(phi);
-            Real zr = sin_theta * sin<Real>(phi);
-
-            Real xt = - sin_theta;
-            Real yt = cos_theta * cos<Real>(phi);
-            Real zt = cos_theta * sin<Real>(phi);
-
-            Real xp = 0;
-            Real yp = - sin<Real>(phi);
-            Real zp = cos<Real>(phi);
-
-            ////////////////////////////////////
-
-            Real fx = 0;
-            Real fy = 0;
-            Real fz = 1;
-
-            f[(0 * Nt + i) * Np + j] = (fx * xr + fy * yr + fz * zr);
-            f[(1 * Nt + i) * Np + j] = (fx * xt + fy * yt + fz * zt);
-            f[(2 * Nt + i) * Np + j] = (fx * xp + fy * yp + fz * zp);
-
-            f[(0 * Nt + i) * Np + j] = 0; //sin(phi) * sin_theta;
-            f[(1 * Nt + i) * Np + j] = 1; //sin(phi) * cos_theta; // * sin_theta;
-            f[(2 * Nt + i) * Np + j] = 1; //cos(phi)            ; // * sin_theta;
+            f[(0 * Nt + i) * Np + j] = 3;
+            f[(1 * Nt + i) * Np + j] = 2;
+            f[(2 * Nt + i) * Np + j] = 1;
           }
         }
       }
 
       Vector<Real> f_coeff;
       Grid2VecSHC(f, Nt, Np, p, f_coeff, sctl::SHCArrange::ROW_MAJOR);
+      print_coeff(f_coeff);
 
-      if(0){
-        Vector<Real> f_, f_coeff_, f__;
-        SHC2Grid(f_coeff, sctl::SHCArrange::ROW_MAJOR, p, Nt, Np, &f_);
-        Grid2SHC(f_, Nt, Np, p, f_coeff_, sctl::SHCArrange::ROW_MAJOR);
-        SHC2Grid(f_coeff_, sctl::SHCArrange::ROW_MAJOR, p, Nt, Np, &f__);
-        std::cout<<Matrix<Real>(dof*Nt, Np, f.begin())-Matrix<Real>(dof*Nt, Np, f_.begin())<<'\n';
-        std::cout<<Matrix<Real>(dof*Nt, Np, f_.begin())-Matrix<Real>(dof*Nt, Np, f__.begin())<<'\n';
-      }
-
-      if(0)
       for (Long i = 0; i < 20; i++) { // Evaluate
-        Vector<Real> r_cos_theta_phi;
-        r_cos_theta_phi.PushBack(drand48() + (i<10?0:2)); // [1, inf)
-        r_cos_theta_phi.PushBack(drand48() * 2 - 1); // [-1, 1]
-        r_cos_theta_phi.PushBack(drand48() * 2 * const_pi<Real>()); // [0, 2*pi]
-
         Vector<Real> Df;
-        StokesEvalDL(f_coeff, sctl::SHCArrange::ROW_MAJOR, p, r_cos_theta_phi, Df);
-        //VecSHCEval(f_coeff, sctl::SHCArrange::ROW_MAJOR, p, r_cos_theta_phi, Df);
-        //std::cout<<r_cos_theta_phi;
-        std::cout<<Df[0]*Df[0] + Df[1]*Df[1] + Df[2]*Df[2]<<'\n';
+        Vector<Real> x(3);
+        x[0] = drand48();
+        x[1] = drand48();
+        x[2] = drand48();
+        Real R = sqrt<Real>(x[0]*x[0]+x[1]*x[1]+x[2]*x[2]);
+        x[0]/=R*(i+0.5)/10;
+        x[1]/=R*(i+0.5)/10;
+        x[2]/=R*(i+0.5)/10;
+
+        StokesEvalDL(f_coeff, sctl::SHCArrange::ROW_MAJOR, p, x, Df);
+        std::cout<<Df+1e-10;
       }
-
-      print_coeff(f_coeff);
       Clear();
     }
 
     static void test() {
-      int p = 4;
-      int dof = 3;
-      int Nt = p+1, Np = 2*p;
+      int p = 3;
+      int dof = 1;
+      int Nt = p+1, Np = 2*p+1;
 
       auto print_coeff = [&](Vector<Real> S) {
         Long idx=0;
@@ -291,29 +136,18 @@ template <class Real> class SphericalHarmonics{
       int Ncoeff = (p + 1) * (p + 1);
       Vector<Real> Xcoeff(dof * Ncoeff), Xgrid;
       for (int i=0;i<Xcoeff.Dim();i++) Xcoeff[i]=i+1;
-      Xcoeff=0;
-      Xcoeff[0]=1;
-      //Xcoeff[12]=1;
-      //for (int i=0*Ncoeff;i<1*Ncoeff;i++) Xcoeff[i]=i+1;
-
-      VecSHC2Grid(Xcoeff, sctl::SHCArrange::COL_MAJOR_NONZERO, p, Nt, Np, Xgrid);
-      std::cout<<"Y0_=[\n";
-      std::cout<<Matrix<Real>(Nt, Np, Xgrid.begin()+Nt*Np*0)<<"];\n";
-      std::cout<<"Y1_=[\n";
-      std::cout<<Matrix<Real>(Nt, Np, Xgrid.begin()+Nt*Np*1)<<"];\n";
-      std::cout<<"Y2_=[\n";
-      std::cout<<Matrix<Real>(Nt, Np, Xgrid.begin()+Nt*Np*2)<<"];\n";
-
-      if (0) {
+
+      SHC2Grid(Xcoeff, sctl::SHCArrange::COL_MAJOR_NONZERO, p, Nt, Np, &Xgrid);
+      std::cout<<Matrix<Real>(Nt*dof, Np, Xgrid.begin())<<'\n';
+
+      {
         Vector<Real> val;
-        VecSHCEval(Xcoeff, sctl::SHCArrange::COL_MAJOR_NONZERO, p, theta_phi, val);
-        //StokesEvalSL(Xcoeff, sctl::SHCArrange::COL_MAJOR_NONZERO, p, r_theta_phi, val);
+        SHCEval(Xcoeff, sctl::SHCArrange::COL_MAJOR_NONZERO, p, theta_phi, val);
         Matrix<Real>(dof, val.Dim()/dof, val.begin(), false) = Matrix<Real>(val.Dim()/dof, dof, val.begin()).Transpose();
-        std::cout<<Matrix<Real>(val.Dim()/Np, Np, val.begin())<<'\n';
-        std::cout<<Matrix<Real>(val.Dim()/Np, Np, val.begin()) - Matrix<Real>(Nt*dof, Np, Xgrid.begin())<<'\n';
+        std::cout<<Matrix<Real>(val.Dim()/Np, Np, val.begin()) - Matrix<Real>(Nt*dof, Np, Xgrid.begin())+1e-10<<'\n';
       }
 
-      Grid2VecSHC(Xgrid, Nt, Np, p, Xcoeff, sctl::SHCArrange::ROW_MAJOR);
+      Grid2SHC(Xgrid, Nt, Np, p, Xcoeff, sctl::SHCArrange::ROW_MAJOR);
       print_coeff(Xcoeff);
 
       //SphericalHarmonics<Real>::WriteVTK("test", nullptr, &Xcoeff, sctl::SHCArrange::ROW_MAJOR, p, 32);

include/sctl/sph_harm.txx

@@ -615,7 +615,7 @@ template <class Real> void SphericalHarmonics<Real>::VecSHCEval(const Vector<Rea
   }
 }
 
-template <class Real> void SphericalHarmonics<Real>::StokesEvalSL(const Vector<Real>& S, SHCArrange arrange, Long p0, const Vector<Real>& r_cos_theta_phi, Vector<Real>& X) {
+template <class Real> void SphericalHarmonics<Real>::StokesEvalSL(const Vector<Real>& S, SHCArrange arrange, Long p0, const Vector<Real>& coord, Vector<Real>& X) {
   Long M = (p0+1) * (p0+1);
 
   Long dof;
@@ -634,18 +634,19 @@ template <class Real> void SphericalHarmonics<Real>::StokesEvalSL(const Vector<R
   assert(B1.Dim(0) == dof);
 
   Long N;
-  Vector<Real> R;
   Matrix<Real> SHBasis;
+  Vector<Real> R, cos_theta_phi;
   { // Set N, R, SHBasis
-    N = r_cos_theta_phi.Dim() / COORD_DIM;
-    assert(r_cos_theta_phi.Dim() == N * COORD_DIM);
+    N = coord.Dim() / COORD_DIM;
+    assert(coord.Dim() == N * COORD_DIM);
 
     R.ReInit(N);
-    Vector<Real> cos_theta_phi(2 * N);
+    cos_theta_phi.ReInit(2 * N);
     for (Long i = 0; i < N; i++) { // Set R, cos_theta_phi
-      R[i] = r_cos_theta_phi[i * COORD_DIM + 0];
-      cos_theta_phi[i * 2 + 0] = r_cos_theta_phi[i * COORD_DIM + 1];
-      cos_theta_phi[i * 2 + 1] = r_cos_theta_phi[i * COORD_DIM + 2];
+      ConstIterator<Real> x = coord.begin() + i * COORD_DIM;
+      R[i] = sqrt<Real>(x[0]*x[0] + x[1]*x[1] + x[2]*x[2]);
+      cos_theta_phi[i * 2 + 0] = x[2] / R[i];
+      cos_theta_phi[i * 2 + 1] = atan2(x[1], x[0]); // TODO: works only for float and double
     }
     VecSHBasisEval(p0, cos_theta_phi, SHBasis);
     assert(SHBasis.Dim(1) == COORD_DIM * M);
@@ -729,17 +730,32 @@ template <class Real> void SphericalHarmonics<Real>::StokesEvalSL(const Vector<R
     if (X.Dim() != N * dof * COORD_DIM) X.ReInit(N * dof * COORD_DIM);
     for (Long k0 = 0; k0 < N; k0++) {
       for (Long k1 = 0; k1 < dof; k1++) {
+        StaticArray<Real,COORD_DIM> in;
         for (Long j = 0; j < COORD_DIM; j++) {
-          Real X_ = 0;
-          for (Long i = 0; i < COORD_DIM * M; i++) X_ += B1[k1][i] * StokesOp[k0 * COORD_DIM + j][i];
-          X[(k0 * dof + k1) * COORD_DIM + j] = X_;
+          in[j] = 0;
+          for (Long i = 0; i < COORD_DIM * M; i++) {
+            in[j] += B1[k1][i] * StokesOp[k0 * COORD_DIM + j][i];
+          }
         }
+
+        StaticArray<Real,9> M;
+        Real cos_theta = cos_theta_phi[k0 * 2 + 0];
+        Real sin_theta = sqrt<Real>(1 - cos_theta * cos_theta);
+        Real cos_phi = cos(cos_theta_phi[k0 * 2 + 1]);
+        Real sin_phi = sin(cos_theta_phi[k0 * 2 + 1]);
+        M[0] = sin_theta*cos_phi; M[1] = sin_theta*sin_phi; M[2] = cos_theta;
+        M[3] = cos_theta*cos_phi; M[4] = cos_theta*sin_phi; M[5] =-sin_theta;
+        M[6] =          -sin_phi; M[7] =           cos_phi; M[8] =         0;
+
+        X[(k0 * dof + k1) * COORD_DIM + 0] = M[0] * in[0] + M[3] * in[1] + M[6] * in[2];
+        X[(k0 * dof + k1) * COORD_DIM + 1] = M[1] * in[0] + M[4] * in[1] + M[7] * in[2];
+        X[(k0 * dof + k1) * COORD_DIM + 2] = M[2] * in[0] + M[5] * in[1] + M[8] * in[2];
       }
     }
   }
 }
 
-template <class Real> void SphericalHarmonics<Real>::StokesEvalDL(const Vector<Real>& S, SHCArrange arrange, Long p0, const Vector<Real>& r_cos_theta_phi, Vector<Real>& X) {
+template <class Real> void SphericalHarmonics<Real>::StokesEvalDL(const Vector<Real>& S, SHCArrange arrange, Long p0, const Vector<Real>& coord, Vector<Real>& X) {
   Long M = (p0+1) * (p0+1);
 
   Long dof;
@@ -758,18 +774,19 @@ template <class Real> void SphericalHarmonics<Real>::StokesEvalDL(const Vector<R
   assert(B1.Dim(0) == dof);
 
   Long N;
-  Vector<Real> R;
   Matrix<Real> SHBasis;
+  Vector<Real> R, cos_theta_phi;
   { // Set N, R, SHBasis
-    N = r_cos_theta_phi.Dim() / COORD_DIM;
-    assert(r_cos_theta_phi.Dim() == N * COORD_DIM);
+    N = coord.Dim() / COORD_DIM;
+    assert(coord.Dim() == N * COORD_DIM);
 
     R.ReInit(N);
-    Vector<Real> cos_theta_phi(2 * N);
+    cos_theta_phi.ReInit(2 * N);
     for (Long i = 0; i < N; i++) { // Set R, cos_theta_phi
-      R[i] = r_cos_theta_phi[i * COORD_DIM + 0];
-      cos_theta_phi[i * 2 + 0] = r_cos_theta_phi[i * COORD_DIM + 1];
-      cos_theta_phi[i * 2 + 1] = r_cos_theta_phi[i * COORD_DIM + 2];
+      ConstIterator<Real> x = coord.begin() + i * COORD_DIM;
+      R[i] = sqrt<Real>(x[0]*x[0] + x[1]*x[1] + x[2]*x[2]);
+      cos_theta_phi[i * 2 + 0] = x[2] / R[i];
+      cos_theta_phi[i * 2 + 1] = atan2(x[1], x[0]); // TODO: works only for float and double
     }
     VecSHBasisEval(p0, cos_theta_phi, SHBasis);
     assert(SHBasis.Dim(1) == COORD_DIM * M);
@@ -874,11 +891,26 @@ template <class Real> void SphericalHarmonics<Real>::StokesEvalDL(const Vector<R
     if (X.Dim() != N * dof * COORD_DIM) X.ReInit(N * dof * COORD_DIM);
     for (Long k0 = 0; k0 < N; k0++) {
       for (Long k1 = 0; k1 < dof; k1++) {
+        StaticArray<Real,COORD_DIM> in;
         for (Long j = 0; j < COORD_DIM; j++) {
-          Real X_ = 0;
-          for (Long i = 0; i < COORD_DIM * M; i++) X_ += B1[k1][i] * StokesOp[k0 * COORD_DIM + j][i];
-          X[(k0 * dof + k1) * COORD_DIM + j] = X_;
+          in[j] = 0;
+          for (Long i = 0; i < COORD_DIM * M; i++) {
+            in[j] += B1[k1][i] * StokesOp[k0 * COORD_DIM + j][i];
+          }
         }
+
+        StaticArray<Real,9> M;
+        Real cos_theta = cos_theta_phi[k0 * 2 + 0];
+        Real sin_theta = sqrt<Real>(1 - cos_theta * cos_theta);
+        Real cos_phi = cos(cos_theta_phi[k0 * 2 + 1]);
+        Real sin_phi = sin(cos_theta_phi[k0 * 2 + 1]);
+        M[0] = sin_theta*cos_phi; M[1] = sin_theta*sin_phi; M[2] = cos_theta;
+        M[3] = cos_theta*cos_phi; M[4] = cos_theta*sin_phi; M[5] =-sin_theta;
+        M[6] =          -sin_phi; M[7] =           cos_phi; M[8] =         0;
+
+        X[(k0 * dof + k1) * COORD_DIM + 0] = M[0] * in[0] + M[3] * in[1] + M[6] * in[2];
+        X[(k0 * dof + k1) * COORD_DIM + 1] = M[1] * in[0] + M[4] * in[1] + M[7] * in[2];
+        X[(k0 * dof + k1) * COORD_DIM + 2] = M[2] * in[0] + M[5] * in[1] + M[8] * in[2];
       }
     }
   }
@@ -1753,7 +1785,7 @@ template <class Real> const std::vector<Matrix<Real>>& SphericalHarmonics<Real>:
 
       Vector<Real> theta(Ngrid);
       for(Long i=0;i<theta.Dim();i++){ // Set theta
-        theta[i]=atan2(Mcoord1[1][i],Mcoord1[2][i]);
+        theta[i]=atan2(Mcoord1[1][i],Mcoord1[2][i]); // TODO: works only for float and double
       }
 
       Matrix<Real> Mcoef2grid(Ncoef, Ngrid);