Add SHCEval

include/sctl/sph_harm.hpp

@@ -35,6 +35,8 @@ 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 SHC2Pole(const Vector<Real>& S_in, SHCArrange arrange_in, Long p_in, Vector<Real>& P_out);
 
     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());
@@ -47,25 +49,30 @@ template <class Real> class SphericalHarmonics{
     static void test() {
       int p = 3;
       int dof = 2;
+      int Nt = p+1, Np = 2*p+1;
+
+      auto print_coeff = [&](Vector<Real> S) {
+        Long idx=0;
+        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';
+      };
 
       int Ncoeff = (p + 1) * (p + 1);
       Vector<Real> Xcoeff(dof * Ncoeff);
       for (int i=0;i<Xcoeff.Dim();i++) Xcoeff[i]=i;
 
       Vector<Real> Xgrid;
-      int Nt = p+1, Np = 2*p+1;
       SHC2Grid(Xcoeff, sctl::SHCArrange::COL_MAJOR_NONZERO, p, Nt, Np, &Xgrid);
       Grid2SHC(Xgrid, Nt, Np, p, Xcoeff, sctl::SHCArrange::ROW_MAJOR);
-
-      int indx=0;
-      for (int i=0;i<dof;i++) {
-        for (int n=0;n<=p;n++){
-          std::cout<<Vector<Real>(2*n+2,Xcoeff.begin()+indx);
-          indx+=2*n+2;
-        }
-      }
+      print_coeff(Xcoeff);
 
       SphericalHarmonics<Real>::WriteVTK("test", nullptr, &Xcoeff, sctl::SHCArrange::ROW_MAJOR, p, 32);
+      Clear();
     }
 
     static void Clear() { MatrixStore().Resize(0); }
@@ -112,6 +119,9 @@ template <class Real> class SphericalHarmonics{
     static const std::vector<Matrix<Real>>& MatLegendreInv(Long p0, Long p1);
     static const std::vector<Matrix<Real>>& MatLegendreGrad(Long p0, Long p1);
 
+    // Evaluate all Spherical Harmonic basis functions up to order p at (theta, phi) coordinates.
+    static void SHBasisEval(Long p, const Vector<Real>& cos_theta_phi, Matrix<Real>& M);
+
     static const std::vector<Matrix<Real>>& MatRotate(Long p0);
 
     template <bool SLayer, bool DLayer> static void StokesSingularInteg_(const Vector<Real>& X0, Long p0, Long p1, Vector<Real>& SL, Vector<Real>& DL);

include/sctl/sph_harm.txx

@@ -426,6 +426,41 @@ template <class Real> void SphericalHarmonics<Real>::SHC2Grid_(const Vector<Real
   }
 }
 
+template <class Real> void SphericalHarmonics<Real>::SHCEval(const Vector<Real>& S, SHCArrange arrange, Long p0, const Vector<Real>& 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;
+    assert(B0.Dim() == dof * M);
+
+    B1.ReInit(dof, M);
+    Vector<Real> B1_(B1.Dim(0) * B1.Dim(1), B1.begin(), false);
+    SHCArrange0(B0, p0, B1_, SHCArrange::COL_MAJOR_NONZERO);
+  }
+  assert(B1.Dim(0) == dof);
+  assert(B1.Dim(1) == M);
+
+  Matrix<Real> SHBasis;
+  SHBasisEval(p0, cos_theta_phi, SHBasis);
+  assert(SHBasis.Dim(1) == M);
+  Long N = SHBasis.Dim(0);
+
+  if (X.Dim() != N*dof) X.ReInit(N * dof);
+  { // Set X
+    for (Long k0 = 0; k0 < N; k0++) {
+      for (Long k1 = 0; k1 < dof; k1++) {
+        Real X_ = 0;
+        for (Long i = 0; i < M; i++) X_ += B1[k1][i] * SHBasis[k0][i];
+        X[k0 * dof + k1] = X_;
+      }
+    }
+  }
+}
+
 template <class Real> void SphericalHarmonics<Real>::SHC2Pole(const Vector<Real>& S, SHCArrange arrange, Long p0, Vector<Real>& P){
   Vector<Real> QP[2];
   { // Set QP // TODO: store these weights
@@ -1172,6 +1207,51 @@ template <class Real> const std::vector<Matrix<Real>>& SphericalHarmonics<Real>:
   return Mdl;
 }
 
+template <class Real> void SphericalHarmonics<Real>::SHBasisEval(Long p0, const Vector<Real>& cos_theta_phi, Matrix<Real>& SHBasis) {
+  Long M = (p0+1) * (p0+1);
+  Long N = cos_theta_phi.Dim() / 2;
+  assert(cos_theta_phi.Dim() == N * 2);
+
+  Vector<Complex<Real>> exp_phi(N);
+  Matrix<Real> LegP((p0+1)*(p0+2)/2, N);
+  { // Set exp_phi, LegP
+    Vector<Real> cos_theta(N);
+    for (Long i = 0; i < N; i++) { // Set cos_theta, exp_phi
+      cos_theta[i] = cos_theta_phi[i*2+0];
+      exp_phi[i].real = cos(cos_theta_phi[i*2+1]);
+      exp_phi[i].imag = sin(cos_theta_phi[i*2+1]);
+    }
+
+    Vector<Real> alp(LegP.Dim(0) * LegP.Dim(1), LegP.begin(), false);
+    LegPoly(alp, cos_theta, p0);
+  }
+
+  { // Set SHBasis
+    SHBasis.ReInit(N, M);
+    Real s = 4 * sqrt<Real>(const_pi<Real>());
+    for (Long k0 = 0; k0 < N; k0++) {
+      Complex<Real> exp_phi_ = 1;
+      Complex<Real> exp_phi1 = exp_phi[k0];
+      for (Long m = 0; m <= p0; m++) {
+        for (Long n = m; n <= p0; n++) {
+          Long poly_idx = (2 * p0 - m + 1) * m / 2 + n;
+          Long basis_idx = (2 * p0 - m + 2) * m - (m ? p0+1 : 0) + n;
+          SHBasis[k0][basis_idx] = LegP[poly_idx][k0] * exp_phi_.real * s;
+          if (m) { // imaginary part
+            basis_idx += (p0+1-m);
+            SHBasis[k0][basis_idx] = -LegP[poly_idx][k0] * exp_phi_.imag * s;
+          } else {
+            SHBasis[k0][basis_idx] = SHBasis[k0][basis_idx] * 0.5;
+          }
+        }
+        exp_phi_ = exp_phi_ * exp_phi1;
+      }
+    }
+  }
+  assert(SHBasis.Dim(0) == N);
+  assert(SHBasis.Dim(1) == M);
+}
+
 template <class Real> const std::vector<Matrix<Real>>& SphericalHarmonics<Real>::MatRotate(Long p0){
   std::vector<std::vector<Long>> coeff_perm(p0+1);
   { // Set coeff_perm