,

include/sctl/sph_harm.hpp

@@ -31,6 +31,7 @@ template <class Real> class SphericalHarmonics{
 
   public:
 
+    // Scalar Spherical Harmonics
     static void Grid2SHC(const Vector<Real>& X_in, Long Nt_in, Long Np_in, Long p_out, Vector<Real>& S_out, SHCArrange arrange_out);
 
     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);
@@ -42,16 +43,19 @@ template <class Real> class SphericalHarmonics{
     static void WriteVTK(const char* fname, const Vector<Real>* S, const Vector<Real>* f_val, SHCArrange arrange, Long p_in, Long p_out, Real period=0, const Comm& comm = Comm::World());
 
 
+    // Vector Spherical Harmonics
     static void Grid2VecSHC(const Vector<Real>& X_in, Long Nt_in, Long Np_in, Long p_out, Vector<Real>& S_out, SHCArrange arrange_out);
 
     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 StokesEvalSL(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, const Vector<Real>& theta_phi_in, Vector<Real>& X_out);
+
 
     static void test() {
       int p = 4;
-      int dof = 3;
+      int dof = 6;
       int Nt = p+1, Np = 2*p+1;
 
       auto print_coeff = [&](Vector<Real> S) {
@@ -65,11 +69,14 @@ template <class Real> class SphericalHarmonics{
         std::cout<<'\n';
       };
 
-      Vector<Real> theta_phi;
-      { // Set theta_phi
+      Vector<Real> r_theta_phi, theta_phi;
+      { // Set r_theta_phi, theta_phi
         Vector<Real> leg_nodes = LegendreNodes(Nt-1);
         for (Long i=0;i<Nt;i++) {
           for (Long j=0;j<Np;j++) {
+            r_theta_phi.PushBack(1);
+            r_theta_phi.PushBack(leg_nodes[i]);
+            r_theta_phi.PushBack(j * 2 * const_pi<Real>() / Np);
             theta_phi.PushBack(leg_nodes[i]);
             theta_phi.PushBack(j * 2 * const_pi<Real>() / Np);
           }
@@ -85,6 +92,7 @@ template <class Real> class SphericalHarmonics{
       {
         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);
         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()) - Matrix<Real>(Nt*dof, Np, Xgrid.begin())<<'\n';
       }

include/sctl/sph_harm.txx

@@ -562,6 +562,130 @@ 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) {
+  Long M = (p0+1) * (p0+1);
+
+  Long dof;
+  Matrix<Real> B1;
+  { // Set B1, dof
+    Vector<Real> B0;
+    SHCArrange1(S, arrange, p0, B0);
+    dof = B0.Dim() / M / COORD_DIM;
+    assert(B0.Dim() == dof * COORD_DIM * M);
+
+    B1.ReInit(dof, COORD_DIM * M);
+    Vector<Real> B1_(B1.Dim(0) * B1.Dim(1), B1.begin(), false);
+    SHCArrange0(B0, p0, B1_, SHCArrange::COL_MAJOR_NONZERO);
+  }
+  assert(B1.Dim(1) == COORD_DIM * M);
+  assert(B1.Dim(0) == dof);
+
+  Long N;
+  Vector<Real> R;
+  Matrix<Real> SHBasis;
+  { // Set N, R, SHBasis
+    N = r_cos_theta_phi.Dim() / COORD_DIM;
+    assert(r_cos_theta_phi.Dim() == N * COORD_DIM);
+
+    R.ReInit(N);
+    Vector<Real> cos_theta_phi(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];
+    }
+    VecSHBasisEval(p0, cos_theta_phi, SHBasis);
+    assert(SHBasis.Dim(1) == COORD_DIM * M);
+    assert(SHBasis.Dim(0) == N * COORD_DIM);
+  }
+
+  Matrix<Real> StokesOp(SHBasis.Dim(0), SHBasis.Dim(1));
+  for (Long i = 0; i < N; i++) { // Set StokesOp
+    for (Long m = 0; m <= p0; m++) {
+      for (Long n = m; n <= p0; n++) {
+        auto read_coeff = [&](Long n, Long m, Long k0, Long k1) {
+          Complex<Real> c;
+          if (0 <= m && m <= n && n <= p0 && 0 <= k0 && k0 < COORD_DIM && 0 <= k1 && k1 < COORD_DIM) {
+            Long idx = (2 * p0 - m + 2) * m - (m ? p0+1 : 0) + n;
+            c.real = SHBasis[i * COORD_DIM + k1][k0 * M + idx];
+            if (m) {
+              idx += (p0+1-m);
+              c.imag = SHBasis[i * COORD_DIM + k1][k0 * M + idx];
+            }
+          }
+          return c;
+        };
+        auto write_coeff = [&](Complex<Real> c, Long n, Long m, Long k0, Long k1) {
+          if (0 <= m && m <= n && n <= p0 && 0 <= k0 && k0 < COORD_DIM && 0 <= k1 && k1 < COORD_DIM) {
+            Long idx = (2 * p0 - m + 2) * m - (m ? p0+1 : 0) + n;
+            StokesOp[i * COORD_DIM + k1][k0 * M + idx] = c.real;
+            if (m) {
+              idx += (p0+1-m);
+              StokesOp[i * COORD_DIM + k1][k0 * M + idx] = c.imag;
+            }
+          }
+        };
+
+        auto Vr = read_coeff(n, m, 0, 0);
+        auto Vt = read_coeff(n, m, 0, 1);
+        auto Vp = read_coeff(n, m, 0, 2);
+
+        auto Wr = read_coeff(n, m, 1, 0);
+        auto Wt = read_coeff(n, m, 1, 1);
+        auto Wp = read_coeff(n, m, 1, 2);
+
+        auto Xr = read_coeff(n, m, 2, 0);
+        auto Xt = read_coeff(n, m, 2, 1);
+        auto Xp = read_coeff(n, m, 2, 2);
+
+
+        Real a,b;
+        a = n / (Real)((2*n+1) * (2*n+3)) * pow<Real>(R[i], -n-2);
+        Complex<Real> SVr = a * Vr;
+        Complex<Real> SVt = a * Vt;
+        Complex<Real> SVp = a * Vp;
+
+        a = (n+1) / (Real)((2*n+1) * (2*n-1)) * pow<Real>(R[i], -n);
+        b = n / (Real)(4*n+2) * (pow<Real>(R[i], -n-2) - pow<Real>(R[i], -n));
+        Complex<Real> SWr = a * Wr + b * Vr;
+        Complex<Real> SWt = a * Wt + b * Vt;
+        Complex<Real> SWp = a * Wp + b * Vp;
+
+        a = 1 / (Real)(2*n+1) * pow<Real>(R[i], -n-1);
+        Complex<Real> SXr = a * Xr;
+        Complex<Real> SXt = a * Xt;
+        Complex<Real> SXp = a * Xp;
+
+
+        write_coeff(SVr, n, m, 0, 0);
+        write_coeff(SVt, n, m, 0, 1);
+        write_coeff(SVp, n, m, 0, 2);
+
+        write_coeff(SWr, n, m, 1, 0);
+        write_coeff(SWt, n, m, 1, 1);
+        write_coeff(SWp, n, m, 1, 2);
+
+        write_coeff(SXr, n, m, 2, 0);
+        write_coeff(SXt, n, m, 2, 1);
+        write_coeff(SXp, n, m, 2, 2);
+      }
+    }
+  }
+
+  { // Set X
+    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++) {
+        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_;
+        }
+      }
+    }
+  }
+}
+