sph_harm.txx 56 KB

12345678910111213141516171819202122232425262728293031323334353637383940414243444546474849505152535455565758596061626364656667686970717273747576777879808182838485868788899091929394959697989910010110210310410510610710810911011111211311411511611711811912012112212312412512612712812913013113213313413513613713813914014114214314414514614714814915015115215315415515615715815916016116216316416516616716816917017117217317417517617717817918018118218318418518618718818919019119219319419519619719819920020120220320420520620720820921021121221321421521621721821922022122222322422522622722822923023123223323423523623723823924024124224324424524624724824925025125225325425525625725825926026126226326426526626726826927027127227327427527627727827928028128228328428528628728828929029129229329429529629729829930030130230330430530630730830931031131231331431531631731831932032132232332432532632732832933033133233333433533633733833934034134234334434534634734834935035135235335435535635735835936036136236336436536636736836937037137237337437537637737837938038138238338438538638738838939039139239339439539639739839940040140240340440540640740840941041141241341441541641741841942042142242342442542642742842943043143243343443543643743843944044144244344444544644744844945045145245345445545645745845946046146246346446546646746846947047147247347447547647747847948048148248348448548648748848949049149249349449549649749849950050150250350450550650750850951051151251351451551651751851952052152252352452552652752852953053153253353453553653753853954054154254354454554654754854955055155255355455555655755855956056156256356456556656756856957057157257357457557657757857958058158258358458558658758858959059159259359459559659759859960060160260360460560660760860961061161261361461561661761861962062162262362462562662762862963063163263363463563663763863964064164264364464564664764864965065165265365465565665765865966066166266366466566666766866967067167267367467567667767867968068168268368468568668768868969069169269369469569669769869970070170270370470570670770870971071171271371471571671771871972072172272372472572672772872973073173273373473573673773873974074174274374474574674774874975075175275375475575675775875976076176276376476576676776876977077177277377477577677777877978078178278378478578678778878979079179279379479579679779879980080180280380480580680780880981081181281381481581681781881982082182282382482582682782882983083183283383483583683783883984084184284384484584684784884985085185285385485585685785885986086186286386486586686786886987087187287387487587687787887988088188288388488588688788888989089189289389489589689789889990090190290390490590690790890991091191291391491591691791891992092192292392492592692792892993093193293393493593693793893994094194294394494594694794894995095195295395495595695795895996096196296396496596696796896997097197297397497597697797897998098198298398498598698798898999099199299399499599699799899910001001100210031004100510061007100810091010101110121013101410151016101710181019102010211022102310241025102610271028102910301031103210331034103510361037103810391040104110421043104410451046104710481049105010511052105310541055105610571058105910601061106210631064106510661067106810691070107110721073107410751076107710781079108010811082108310841085108610871088108910901091109210931094109510961097109810991100110111021103110411051106110711081109111011111112111311141115111611171118111911201121112211231124112511261127112811291130113111321133113411351136113711381139114011411142114311441145114611471148114911501151115211531154115511561157115811591160116111621163116411651166116711681169117011711172117311741175117611771178117911801181118211831184118511861187118811891190119111921193119411951196119711981199120012011202120312041205120612071208120912101211121212131214121512161217121812191220122112221223122412251226122712281229123012311232123312341235123612371238123912401241124212431244124512461247124812491250125112521253125412551256125712581259126012611262126312641265126612671268126912701271127212731274127512761277127812791280128112821283128412851286128712881289129012911292129312941295129612971298129913001301130213031304130513061307130813091310131113121313131413151316131713181319132013211322132313241325132613271328132913301331133213331334133513361337133813391340134113421343134413451346134713481349135013511352135313541355135613571358135913601361136213631364136513661367136813691370137113721373137413751376137713781379138013811382138313841385138613871388138913901391139213931394139513961397139813991400140114021403140414051406140714081409141014111412141314141415141614171418141914201421142214231424142514261427142814291430143114321433143414351436143714381439144014411442144314441445144614471448144914501451145214531454145514561457145814591460146114621463146414651466146714681469147014711472147314741475147614771478147914801481148214831484148514861487148814891490149114921493149414951496149714981499150015011502150315041505150615071508150915101511151215131514151515161517151815191520152115221523152415251526152715281529153015311532153315341535153615371538153915401541154215431544154515461547154815491550155115521553155415551556155715581559156015611562156315641565156615671568156915701571157215731574157515761577157815791580158115821583158415851586158715881589159015911592159315941595159615971598159916001601160216031604160516061607160816091610161116121613161416151616161716181619
  1. #include SCTL_INCLUDE(legendre_rule.hpp)
  2. // TODO: Replace work vectors with dynamic-arrays
  3. namespace SCTL_NAMESPACE {
  4. template <class Real> void SphericalHarmonics<Real>::Grid2SHC(const Vector<Real>& X, Long Nt, Long Np, Long p1, Vector<Real>& S, SHCArrange arrange){
  5. const auto& Mf = OpFourierInv(Np);
  6. assert(Mf.Dim(0) == Np);
  7. const std::vector<Matrix<Real>>& Ml = SphericalHarmonics<Real>::MatLegendreInv(Nt-1,p1);
  8. assert((Long)Ml.size() == p1+1);
  9. Long N = X.Dim() / (Np*Nt);
  10. assert(X.Dim() == N*Np*Nt);
  11. Vector<Real> B0((2*p1+1) * N*Nt);
  12. #pragma omp parallel
  13. { // B0 <-- Transpose(FFT(X))
  14. Integer tid=omp_get_thread_num();
  15. Integer omp_p=omp_get_num_threads();
  16. Long a=(tid+0)*N*Nt/omp_p;
  17. Long b=(tid+1)*N*Nt/omp_p;
  18. Vector<Real> buff(Mf.Dim(1));
  19. Long fft_coeff_len = std::min(buff.Dim(), 2*p1+2);
  20. Matrix<Real> B0_(2*p1+1, N*Nt, B0.begin(), false);
  21. const Matrix<Real> MX(N * Nt, Np, (Iterator<Real>)X.begin(), false);
  22. for (Long i = a; i < b; i++) {
  23. { // buff <-- FFT(Xi)
  24. const Vector<Real> Xi(Np, (Iterator<Real>)X.begin() + Np * i, false);
  25. Mf.Execute(Xi, buff);
  26. }
  27. { // B0 <-- Transpose(buff)
  28. B0_[0][i] = buff[0]; // skipping buff[1] == 0
  29. for (Long j = 2; j < fft_coeff_len; j++) B0_[j-1][i] = buff[j];
  30. for (Long j = fft_coeff_len; j < 2*p1+2; j++) B0_[j-1][i] = 0;
  31. }
  32. }
  33. }
  34. Vector<Real> B1(N*(p1+1)*(p1+1));
  35. #pragma omp parallel
  36. { // Evaluate Legendre polynomial
  37. Integer tid=omp_get_thread_num();
  38. Integer omp_p=omp_get_num_threads();
  39. Long offset0=0;
  40. Long offset1=0;
  41. for (Long i = 0; i < p1+1; i++) {
  42. Long N_ = (i==0 ? N : 2*N);
  43. Matrix<Real> Min (N_, Nt , B0.begin()+offset0, false);
  44. Matrix<Real> Mout(N_, p1+1-i, B1.begin()+offset1, false);
  45. { // Mout = Min * Ml[i] // split between threads
  46. Long a=(tid+0)*N_/omp_p;
  47. Long b=(tid+1)*N_/omp_p;
  48. if (a < b) {
  49. Matrix<Real> Min_ (b-a, Min .Dim(1), Min [a], false);
  50. Matrix<Real> Mout_(b-a, Mout.Dim(1), Mout[a], false);
  51. Matrix<Real>::GEMM(Mout_,Min_,Ml[i]);
  52. }
  53. }
  54. offset0+=Min .Dim(0)*Min .Dim(1);
  55. offset1+=Mout.Dim(0)*Mout.Dim(1);
  56. }
  57. assert(offset0 == B0.Dim());
  58. assert(offset1 == B1.Dim());
  59. }
  60. B1 *= 1 / sqrt<Real>(4 * const_pi<Real>() * Np); // Scaling to match Zydrunas Fortran code.
  61. if (arrange == SHCArrange::ALL) { // S <-- Rearrange(B1)
  62. Long M = 2*(p1+1)*(p1+1);
  63. if(S.Dim() != N * M) S.ReInit(N * M);
  64. #pragma omp parallel
  65. { // S <-- Rearrange(B1)
  66. Integer tid=omp_get_thread_num();
  67. Integer omp_p=omp_get_num_threads();
  68. Long a=(tid+0)*N/omp_p;
  69. Long b=(tid+1)*N/omp_p;
  70. for (Long i = a; i < b; i++) {
  71. Long offset = 0;
  72. for (Long j = 0; j < p1+1; j++) {
  73. Long len = p1+1 - j;
  74. if (1) { // Set Real(S_n^m) for m=j and n=j..p
  75. ConstIterator<Real> B_ = B1.begin() + i*len + N*offset;
  76. Iterator<Real> S_ = S .begin() + i*M + j*(p1+1)*2 + j*2 + 0;
  77. for (Long k = 0; k < len; k++) S_[k * (p1+1)*2] = B_[k];
  78. offset += len;
  79. }
  80. if (j) { // Set Imag(S_n^m) for m=j and n=j..p
  81. ConstIterator<Real> B_ = B1.begin() + i*len + N*offset;
  82. Iterator<Real> S_ = S .begin() + i*M + j*(p1+1)*2 + j*2 + 1;
  83. for (Long k = 0; k < len; k++) S_[k * (p1+1)*2] = B_[k];
  84. offset += len;
  85. } else {
  86. Iterator<Real> S_ = S .begin() + i*M + j*(p1+1)*2 + j*2 + 1;
  87. for (Long k = 0; k < len; k++) S_[k * (p1+1)*2] = 0;
  88. }
  89. }
  90. }
  91. }
  92. }
  93. if (arrange == SHCArrange::ROW_MAJOR) { // S <-- Rearrange(B1)
  94. Long M = (p1+1)*(p1+2);
  95. if(S.Dim() != N * M) S.ReInit(N * M);
  96. #pragma omp parallel
  97. { // S <-- Rearrange(B1)
  98. Integer tid=omp_get_thread_num();
  99. Integer omp_p=omp_get_num_threads();
  100. Long a=(tid+0)*N/omp_p;
  101. Long b=(tid+1)*N/omp_p;
  102. for (Long i = a; i < b; i++) {
  103. Long offset = 0;
  104. for (Long j = 0; j < p1+1; j++) {
  105. Long len = p1+1 - j;
  106. if (1) { // Set Real(S_n^m) for m=j and n=j..p
  107. ConstIterator<Real> B_ = B1.begin() + i*len + N*offset;
  108. Iterator<Real> S_ = S .begin() + i*M + 0;
  109. for (Long k=0;k<len;k++) S_[(j+k)*(j+k+1) + 2*j] = B_[k];
  110. offset += len;
  111. }
  112. if (j) { // Set Imag(S_n^m) for m=j and n=j..p
  113. ConstIterator<Real> B_ = B1.begin() + i*len + N*offset;
  114. Iterator<Real> S_ = S .begin() + i*M + 1;
  115. for (Long k=0;k<len;k++) S_[(j+k)*(j+k+1) + 2*j] = B_[k];
  116. offset += len;
  117. } else {
  118. Iterator<Real> S_ = S .begin() + i*M + 1;
  119. for (Long k=0;k<len;k++) S_[(j+k)*(j+k+1) + 2*j] = 0;
  120. }
  121. }
  122. }
  123. }
  124. }
  125. if (arrange == SHCArrange::COL_MAJOR_NONZERO) { // S <-- Rearrange(B1)
  126. Long M = (p1+1)*(p1+1);
  127. if(S.Dim() != N * M) S.ReInit(N * M);
  128. #pragma omp parallel
  129. { // S <-- Rearrange(B1)
  130. Integer tid=omp_get_thread_num();
  131. Integer omp_p=omp_get_num_threads();
  132. Long a=(tid+0)*N/omp_p;
  133. Long b=(tid+1)*N/omp_p;
  134. for (Long i = a; i < b; i++) {
  135. Long offset = 0;
  136. for (Long j = 0; j < p1+1; j++) {
  137. Long len = p1+1 - j;
  138. if (1) { // Set Real(S_n^m) for m=j and n=j..p
  139. ConstIterator<Real> B_ = B1.begin() + i*len + N*offset;
  140. Iterator<Real> S_ = S .begin() + i*M + offset;
  141. for (Long k = 0; k < len; k++) S_[k] = B_[k];
  142. offset += len;
  143. }
  144. if (j) { // Set Imag(S_n^m) for m=j and n=j..p
  145. ConstIterator<Real> B_ = B1.begin() + i*len + N*offset;
  146. Iterator<Real> S_ = S .begin() + i*M + offset;
  147. for (Long k = 0; k < len; k++) S_[k] = B_[k];
  148. offset += len;
  149. }
  150. }
  151. }
  152. }
  153. }
  154. }
  155. template <class Real> void SphericalHarmonics<Real>::SHC2Grid(const Vector<Real>& S, SHCArrange arrange, Long p0, Long Nt, Long Np, Vector<Real>* X, Vector<Real>* X_phi, Vector<Real>* X_theta){
  156. const auto& Mf = OpFourier(Np);
  157. assert(Mf.Dim(1) == Np);
  158. const std::vector<Matrix<Real>>& Ml =SphericalHarmonics<Real>::MatLegendre (p0,Nt-1);
  159. const std::vector<Matrix<Real>>& Mdl=SphericalHarmonics<Real>::MatLegendreGrad(p0,Nt-1);
  160. assert((Long)Ml .size() == p0+1);
  161. assert((Long)Mdl.size() == p0+1);
  162. Long M, N;
  163. { // Set M, N
  164. if (arrange == SHCArrange::ALL) M = 2*(p0+1)*(p0+1);
  165. if (arrange == SHCArrange::ROW_MAJOR) M = (p0+1)*(p0+2);
  166. if (arrange == SHCArrange::COL_MAJOR_NONZERO) M = (p0+1)*(p0+1);
  167. N = S.Dim() / M;
  168. assert(S.Dim() == N * M);
  169. }
  170. Vector<Real> B0(N*(p0+1)*(p0+1));
  171. if (arrange == SHCArrange::ALL) { // B0 <-- Rearrange(S)
  172. #pragma omp parallel
  173. { // B0 <-- Rearrange(S)
  174. Integer tid=omp_get_thread_num();
  175. Integer omp_p=omp_get_num_threads();
  176. Long a=(tid+0)*N/omp_p;
  177. Long b=(tid+1)*N/omp_p;
  178. for (Long i = a; i < b; i++) {
  179. Long offset = 0;
  180. for (Long j = 0; j < p0+1; j++) {
  181. Long len = p0+1 - j;
  182. if (1) { // Get Real(S_n^m) for m=j and n=j..p
  183. Iterator<Real> B_ = B0.begin() + i*len + N*offset;
  184. ConstIterator<Real> S_ = S .begin() + i*M + j*(p0+1)*2 + j*2 + 0;
  185. for (Long k = 0; k < len; k++) B_[k] = S_[k * (p0+1)*2];
  186. offset += len;
  187. }
  188. if (j) { // Get Imag(S_n^m) for m=j and n=j..p
  189. Iterator<Real> B_ = B0.begin() + i*len + N*offset;
  190. ConstIterator<Real> S_ = S .begin() + i*M + j*(p0+1)*2 + j*2 + 1;
  191. for (Long k = 0; k < len; k++) B_[k] = S_[k * (p0+1)*2];
  192. offset += len;
  193. }
  194. }
  195. }
  196. }
  197. }
  198. if (arrange == SHCArrange::ROW_MAJOR) { // B0 <-- Rearrange(S)
  199. #pragma omp parallel
  200. { // B0 <-- Rearrange(S)
  201. Integer tid=omp_get_thread_num();
  202. Integer omp_p=omp_get_num_threads();
  203. Long a=(tid+0)*N/omp_p;
  204. Long b=(tid+1)*N/omp_p;
  205. for (Long i = a; i < b; i++) {
  206. Long offset = 0;
  207. for (Long j = 0; j < p0+1; j++) {
  208. Long len = p0+1 - j;
  209. if (1) { // Get Real(S_n^m) for m=j and n=j..p
  210. Iterator<Real> B_ = B0.begin() + i*len + N*offset;
  211. ConstIterator<Real> S_ = S .begin() + i*M + 0;
  212. for (Long k=0;k<len;k++) B_[k] = S_[(j+k)*(j+k+1) + 2*j];
  213. offset += len;
  214. }
  215. if (j) { // Get Imag(S_n^m) for m=j and n=j..p
  216. Iterator<Real> B_ = B0.begin() + i*len + N*offset;
  217. ConstIterator<Real> S_ = S .begin() + i*M + 1;
  218. for (Long k=0;k<len;k++) B_[k] = S_[(j+k)*(j+k+1) + 2*j];
  219. offset += len;
  220. }
  221. }
  222. }
  223. }
  224. }
  225. if (arrange == SHCArrange::COL_MAJOR_NONZERO) { // B0 <-- Rearrange(S)
  226. #pragma omp parallel
  227. { // B0 <-- Rearrange(S)
  228. Integer tid=omp_get_thread_num();
  229. Integer omp_p=omp_get_num_threads();
  230. Long a=(tid+0)*N/omp_p;
  231. Long b=(tid+1)*N/omp_p;
  232. for (Long i = a; i < b; i++) {
  233. Long offset = 0;
  234. for (Long j = 0; j < p0+1; j++) {
  235. Long len = p0+1 - j;
  236. if (1) { // Get Real(S_n^m) for m=j and n=j..p
  237. Iterator<Real> B_ = B0.begin() + i*len + N*offset;
  238. ConstIterator<Real> S_ = S .begin() + i*M + offset;
  239. for (Long k = 0; k < len; k++) B_[k] = S_[k];
  240. offset += len;
  241. }
  242. if (j) { // Get Imag(S_n^m) for m=j and n=j..p
  243. Iterator<Real> B_ = B0.begin() + i*len + N*offset;
  244. ConstIterator<Real> S_ = S .begin() + i*M + offset;
  245. for (Long k = 0; k < len; k++) B_[k] = S_[k];
  246. offset += len;
  247. }
  248. }
  249. }
  250. }
  251. }
  252. B0 *= sqrt<Real>(4 * const_pi<Real>() * Np); // Scaling to match Zydrunas Fortran code.
  253. if(X && X ->Dim()!=N*Np*Nt) X ->ReInit(N*Np*Nt);
  254. if(X_theta && X_theta->Dim()!=N*Np*Nt) X_theta->ReInit(N*Np*Nt);
  255. if(X_phi && X_phi ->Dim()!=N*Np*Nt) X_phi ->ReInit(N*Np*Nt);
  256. Vector<Real> B1(N*(2*p0+1)*Nt);
  257. if(X || X_phi){
  258. #pragma omp parallel
  259. { // Evaluate Legendre polynomial
  260. Integer tid=omp_get_thread_num();
  261. Integer omp_p=omp_get_num_threads();
  262. Long offset0=0;
  263. Long offset1=0;
  264. for(Long i=0;i<p0+1;i++){
  265. Long N_ = (i==0 ? N : 2*N);
  266. Matrix<Real> Min (N_, p0+1-i, B0.begin()+offset0, false);
  267. Matrix<Real> Mout(N_, Nt , B1.begin()+offset1, false);
  268. { // Mout = Min * Ml[i] // split between threads
  269. Long a=(tid+0)*N_/omp_p;
  270. Long b=(tid+1)*N_/omp_p;
  271. if(a<b){
  272. Matrix<Real> Min_ (b-a, Min .Dim(1), Min [a], false);
  273. Matrix<Real> Mout_(b-a, Mout.Dim(1), Mout[a], false);
  274. Matrix<Real>::GEMM(Mout_,Min_,Ml[i]);
  275. }
  276. }
  277. offset0+=Min .Dim(0)*Min .Dim(1);
  278. offset1+=Mout.Dim(0)*Mout.Dim(1);
  279. }
  280. }
  281. #pragma omp parallel
  282. { // Transpose and evaluate Fourier
  283. Integer tid=omp_get_thread_num();
  284. Integer omp_p=omp_get_num_threads();
  285. Long a=(tid+0)*N*Nt/omp_p;
  286. Long b=(tid+1)*N*Nt/omp_p;
  287. Vector<Real> buff(Mf.Dim(0)); buff = 0;
  288. Long fft_coeff_len = std::min(buff.Dim(), 2*p0+2);
  289. Matrix<Real> B1_(2*p0+1, N*Nt, B1.begin(), false);
  290. for (Long i = a; i < b; i++) {
  291. { // buff <-- Transpose(B1)
  292. buff[0] = B1_[0][i];
  293. buff[1] = 0;
  294. for (Long j = 2; j < fft_coeff_len; j++) buff[j] = B1_[j-1][i];
  295. for (Long j = fft_coeff_len; j < buff.Dim(); j++) buff[j] = 0;
  296. }
  297. { // X <-- FFT(buff)
  298. Vector<Real> Xi(Np, X->begin() + Np * i, false);
  299. Mf.Execute(buff, Xi);
  300. }
  301. if(X_phi){ // Evaluate Fourier gradient
  302. { // buff <-- Transpose(B1)
  303. buff[0] = 0;
  304. buff[1] = 0;
  305. for (Long j = 2; j < fft_coeff_len; j++) buff[j] = B1_[j-1][i];
  306. for (Long j = fft_coeff_len; j < buff.Dim(); j++) buff[j] = 0;
  307. for (Long j = 1; j < buff.Dim()/2; j++) {
  308. Real x = buff[2*j+0];
  309. Real y = buff[2*j+1];
  310. buff[2*j+0] = -j*y;
  311. buff[2*j+1] = j*x;
  312. }
  313. }
  314. { // X_phi <-- FFT(buff)
  315. Vector<Real> Xi(Np, X_phi->begin() + Np * i, false);
  316. Mf.Execute(buff, Xi);
  317. }
  318. }
  319. }
  320. }
  321. }
  322. if(X_theta){
  323. #pragma omp parallel
  324. { // Evaluate Legendre gradient
  325. Integer tid=omp_get_thread_num();
  326. Integer omp_p=omp_get_num_threads();
  327. Long offset0=0;
  328. Long offset1=0;
  329. for(Long i=0;i<p0+1;i++){
  330. Long N_ = (i==0 ? N : 2*N);
  331. Matrix<Real> Min (N_, p0+1-i, B0.begin()+offset0, false);
  332. Matrix<Real> Mout(N_, Nt , B1.begin()+offset1, false);
  333. { // Mout = Min * Mdl[i] // split between threads
  334. Long a=(tid+0)*N_/omp_p;
  335. Long b=(tid+1)*N_/omp_p;
  336. if(a<b){
  337. Matrix<Real> Min_ (b-a, Min .Dim(1), Min [a], false);
  338. Matrix<Real> Mout_(b-a, Mout.Dim(1), Mout[a], false);
  339. Matrix<Real>::GEMM(Mout_,Min_,Mdl[i]);
  340. }
  341. }
  342. offset0+=Min .Dim(0)*Min .Dim(1);
  343. offset1+=Mout.Dim(0)*Mout.Dim(1);
  344. }
  345. }
  346. #pragma omp parallel
  347. { // Transpose and evaluate Fourier
  348. Integer tid=omp_get_thread_num();
  349. Integer omp_p=omp_get_num_threads();
  350. Long a=(tid+0)*N*Nt/omp_p;
  351. Long b=(tid+1)*N*Nt/omp_p;
  352. Vector<Real> buff(Mf.Dim(0)); buff = 0;
  353. Long fft_coeff_len = std::min(buff.Dim(), 2*p0+2);
  354. Matrix<Real> B1_(2*p0+1, N*Nt, B1.begin(), false);
  355. for (Long i = a; i < b; i++) {
  356. { // buff <-- Transpose(B1)
  357. buff[0] = B1_[0][i];
  358. buff[1] = 0;
  359. for (Long j = 2; j < fft_coeff_len; j++) buff[j] = B1_[j-1][i];
  360. for (Long j = fft_coeff_len; j < buff.Dim(); j++) buff[j] = 0;
  361. }
  362. { // Xi <-- FFT(buff)
  363. Vector<Real> Xi(Np, X_theta->begin() + Np * i, false);
  364. Mf.Execute(buff, Xi);
  365. }
  366. }
  367. }
  368. }
  369. }
  370. template <class Real> void SphericalHarmonics<Real>::SHC2Pole(const Vector<Real>& S, SHCArrange arrange, Long p0, Vector<Real>& P){
  371. Vector<Real> QP[2];
  372. { // Set QP // TODO: store these weights
  373. Vector<Real> x(1), alp;
  374. const Real SQRT2PI = sqrt<Real>(4 * const_pi<Real>());
  375. for (Long i = 0; i < 2; i++) {
  376. x = (i ? -1 : 1);
  377. LegPoly(alp, x, p0);
  378. QP[i].ReInit(p0 + 1, alp.begin());
  379. QP[i] *= SQRT2PI;
  380. }
  381. }
  382. Long M, N;
  383. { // Set M, N
  384. if (arrange == SHCArrange::ALL) M = 2*(p0+1)*(p0+1);
  385. if (arrange == SHCArrange::ROW_MAJOR) M = (p0+1)*(p0+2);
  386. if (arrange == SHCArrange::COL_MAJOR_NONZERO) M = (p0+1)*(p0+1);
  387. N = S.Dim() / M;
  388. assert(S.Dim() == N * M);
  389. }
  390. if(P.Dim() != N * 2) P.ReInit(N * 2);
  391. if (arrange == SHCArrange::ALL) {
  392. #pragma omp parallel
  393. { // Compute pole
  394. Integer tid = omp_get_thread_num();
  395. Integer omp_p = omp_get_num_threads();
  396. Long a = (tid + 0) * N / omp_p;
  397. Long b = (tid + 1) * N / omp_p;
  398. for (Long i = a; i < b; i++) {
  399. Real P_[2] = {0, 0};
  400. for (Long j = 0; j < p0 + 1; j++) {
  401. P_[0] += S[i*M + j*(p0+1)*2] * QP[0][j];
  402. P_[1] += S[i*M + j*(p0+1)*2] * QP[1][j];
  403. }
  404. P[2*i+0] = P_[0];
  405. P[2*i+1] = P_[1];
  406. }
  407. }
  408. }
  409. if (arrange == SHCArrange::ROW_MAJOR) {
  410. #pragma omp parallel
  411. { // Compute pole
  412. Integer tid = omp_get_thread_num();
  413. Integer omp_p = omp_get_num_threads();
  414. Long a = (tid + 0) * N / omp_p;
  415. Long b = (tid + 1) * N / omp_p;
  416. for (Long i = a; i < b; i++) {
  417. Long idx = 0;
  418. Real P_[2] = {0, 0};
  419. for (Long j = 0; j < p0 + 1; j++) {
  420. P_[0] += S[i*M+idx] * QP[0][j];
  421. P_[1] += S[i*M+idx] * QP[1][j];
  422. idx += 2*(j+1);
  423. }
  424. P[2*i+0] = P_[0];
  425. P[2*i+1] = P_[1];
  426. }
  427. }
  428. }
  429. if (arrange == SHCArrange::COL_MAJOR_NONZERO) {
  430. #pragma omp parallel
  431. { // Compute pole
  432. Integer tid = omp_get_thread_num();
  433. Integer omp_p = omp_get_num_threads();
  434. Long a = (tid + 0) * N / omp_p;
  435. Long b = (tid + 1) * N / omp_p;
  436. for (Long i = a; i < b; i++) {
  437. Real P_[2] = {0, 0};
  438. for (Long j = 0; j < p0 + 1; j++) {
  439. P_[0] += S[i*M+j] * QP[0][j];
  440. P_[1] += S[i*M+j] * QP[1][j];
  441. }
  442. P[2*i+0] = P_[0];
  443. P[2*i+1] = P_[1];
  444. }
  445. }
  446. }
  447. }
  448. template <class Real> void SphericalHarmonics<Real>::WriteVTK(const char* fname, const Vector<Real>* S, const Vector<Real>* v_ptr, SHCArrange arrange, Long p0, Long p1, Real period, const Comm& comm){
  449. typedef double VTKReal;
  450. Vector<Real> SS;
  451. if (S == nullptr) {
  452. Integer p = 2;
  453. Integer Ncoeff = (p + 1) * (p + 1);
  454. Vector<Real> SSS(COORD_DIM * Ncoeff), SSS_grid;
  455. SSS.SetZero();
  456. SSS[1+0*p+0*Ncoeff] = sqrt<Real>(2.0)/sqrt<Real>(3.0);
  457. SSS[1+1*p+1*Ncoeff] = 1/sqrt<Real>(3.0);
  458. SSS[1+2*p+2*Ncoeff] = 1/sqrt<Real>(3.0);
  459. SphericalHarmonics<Real>::SHC2Grid(SSS, SHCArrange::COL_MAJOR_NONZERO, p, p+1, 2*p+2, &SSS_grid);
  460. SphericalHarmonics<Real>::Grid2SHC(SSS_grid, p+1, 2*p+2, p0, SS, arrange);
  461. S = &SS;
  462. }
  463. Vector<Real> X, Xp, V, Vp;
  464. { // Upsample X
  465. const Vector<Real>& X0=*S;
  466. SphericalHarmonics<Real>::SHC2Grid(X0, arrange, p0, p1+1, 2*p1, &X);
  467. SphericalHarmonics<Real>::SHC2Pole(X0, arrange, p0, Xp);
  468. }
  469. if(v_ptr){ // Upsample V
  470. const Vector<Real>& X0=*v_ptr;
  471. SphericalHarmonics<Real>::SHC2Grid(X0, arrange, p0, p1+1, 2*p1, &V);
  472. SphericalHarmonics<Real>::SHC2Pole(X0, arrange, p0, Vp);
  473. }
  474. std::vector<VTKReal> point_coord;
  475. std::vector<VTKReal> point_value;
  476. std::vector<int32_t> poly_connect;
  477. std::vector<int32_t> poly_offset;
  478. { // Set point_coord, point_value, poly_connect
  479. Long N_ves = X.Dim()/(2*p1*(p1+1)*COORD_DIM); // Number of vesicles
  480. assert(Xp.Dim() == N_ves*2*COORD_DIM);
  481. for(Long k=0;k<N_ves;k++){ // Set point_coord
  482. Real C[COORD_DIM]={0,0,0};
  483. if(period>0){
  484. for(Integer l=0;l<COORD_DIM;l++) C[l]=0;
  485. for(Long i=0;i<p1+1;i++){
  486. for(Long j=0;j<2*p1;j++){
  487. for(Integer l=0;l<COORD_DIM;l++){
  488. C[l]+=X[j+2*p1*(i+(p1+1)*(l+k*COORD_DIM))];
  489. }
  490. }
  491. }
  492. for(Integer l=0;l<COORD_DIM;l++) C[l]+=Xp[0+2*(l+k*COORD_DIM)];
  493. for(Integer l=0;l<COORD_DIM;l++) C[l]+=Xp[1+2*(l+k*COORD_DIM)];
  494. for(Integer l=0;l<COORD_DIM;l++) C[l]/=2*p1*(p1+1)+2;
  495. for(Integer l=0;l<COORD_DIM;l++) C[l]=(round(C[l]/period))*period;
  496. }
  497. for(Long i=0;i<p1+1;i++){
  498. for(Long j=0;j<2*p1;j++){
  499. for(Integer l=0;l<COORD_DIM;l++){
  500. point_coord.push_back(X[j+2*p1*(i+(p1+1)*(l+k*COORD_DIM))]-C[l]);
  501. }
  502. }
  503. }
  504. for(Integer l=0;l<COORD_DIM;l++) point_coord.push_back(Xp[0+2*(l+k*COORD_DIM)]-C[l]);
  505. for(Integer l=0;l<COORD_DIM;l++) point_coord.push_back(Xp[1+2*(l+k*COORD_DIM)]-C[l]);
  506. }
  507. if(v_ptr) {
  508. Long data__dof = V.Dim() / (2*p1*(p1+1));
  509. for(Long k=0;k<N_ves;k++){ // Set point_value
  510. for(Long i=0;i<p1+1;i++){
  511. for(Long j=0;j<2*p1;j++){
  512. for(Long l=0;l<data__dof;l++){
  513. point_value.push_back(V[j+2*p1*(i+(p1+1)*(l+k*data__dof))]);
  514. }
  515. }
  516. }
  517. for(Long l=0;l<data__dof;l++) point_value.push_back(Vp[0+2*(l+k*data__dof)]);
  518. for(Long l=0;l<data__dof;l++) point_value.push_back(Vp[1+2*(l+k*data__dof)]);
  519. }
  520. }
  521. for(Long k=0;k<N_ves;k++){
  522. for(Long j=0;j<2*p1;j++){
  523. Long i0= 0;
  524. Long i1=p1;
  525. Long j0=((j+0) );
  526. Long j1=((j+1)%(2*p1));
  527. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*(p1+1)+0);
  528. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*i0+j0);
  529. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*i0+j1);
  530. poly_offset.push_back(poly_connect.size());
  531. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*(p1+1)+1);
  532. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*i1+j0);
  533. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*i1+j1);
  534. poly_offset.push_back(poly_connect.size());
  535. }
  536. for(Long i=0;i<p1;i++){
  537. for(Long j=0;j<2*p1;j++){
  538. Long i0=((i+0) );
  539. Long i1=((i+1) );
  540. Long j0=((j+0) );
  541. Long j1=((j+1)%(2*p1));
  542. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*i0+j0);
  543. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*i1+j0);
  544. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*i1+j1);
  545. poly_connect.push_back((2*p1*(p1+1)+2)*k + 2*p1*i0+j1);
  546. poly_offset.push_back(poly_connect.size());
  547. }
  548. }
  549. }
  550. }
  551. Integer np = comm.Size();
  552. Integer myrank = comm.Rank();
  553. std::vector<VTKReal>& coord=point_coord;
  554. std::vector<VTKReal>& value=point_value;
  555. std::vector<int32_t>& connect=poly_connect;
  556. std::vector<int32_t>& offset=poly_offset;
  557. Long pt_cnt=coord.size()/COORD_DIM;
  558. Long poly_cnt=poly_offset.size();
  559. // Open file for writing.
  560. std::stringstream vtufname;
  561. vtufname<<fname<<"_"<<std::setfill('0')<<std::setw(6)<<myrank<<".vtp";
  562. std::ofstream vtufile;
  563. vtufile.open(vtufname.str().c_str());
  564. if(vtufile.fail()) return;
  565. bool isLittleEndian;
  566. { // Set isLittleEndian
  567. uint16_t number = 0x1;
  568. uint8_t *numPtr = (uint8_t*)&number;
  569. isLittleEndian=(numPtr[0] == 1);
  570. }
  571. // Proceed to write to file.
  572. Long data_size=0;
  573. vtufile<<"<?xml version=\"1.0\"?>\n";
  574. if(isLittleEndian) vtufile<<"<VTKFile type=\"PolyData\" version=\"0.1\" byte_order=\"LittleEndian\">\n";
  575. else vtufile<<"<VTKFile type=\"PolyData\" version=\"0.1\" byte_order=\"BigEndian\">\n";
  576. //===========================================================================
  577. vtufile<<" <PolyData>\n";
  578. vtufile<<" <Piece NumberOfPoints=\""<<pt_cnt<<"\" NumberOfVerts=\"0\" NumberOfLines=\"0\" NumberOfStrips=\"0\" NumberOfPolys=\""<<poly_cnt<<"\">\n";
  579. //---------------------------------------------------------------------------
  580. vtufile<<" <Points>\n";
  581. vtufile<<" <DataArray type=\"Float"<<sizeof(VTKReal)*8<<"\" NumberOfComponents=\""<<COORD_DIM<<"\" Name=\"Position\" format=\"appended\" offset=\""<<data_size<<"\" />\n";
  582. data_size+=sizeof(uint32_t)+coord.size()*sizeof(VTKReal);
  583. vtufile<<" </Points>\n";
  584. //---------------------------------------------------------------------------
  585. if(value.size()){ // value
  586. vtufile<<" <PointData>\n";
  587. vtufile<<" <DataArray type=\"Float"<<sizeof(VTKReal)*8<<"\" NumberOfComponents=\""<<value.size()/pt_cnt<<"\" Name=\""<<"value"<<"\" format=\"appended\" offset=\""<<data_size<<"\" />\n";
  588. data_size+=sizeof(uint32_t)+value.size()*sizeof(VTKReal);
  589. vtufile<<" </PointData>\n";
  590. }
  591. //---------------------------------------------------------------------------
  592. vtufile<<" <Polys>\n";
  593. vtufile<<" <DataArray type=\"Int32\" Name=\"connectivity\" format=\"appended\" offset=\""<<data_size<<"\" />\n";
  594. data_size+=sizeof(uint32_t)+connect.size()*sizeof(int32_t);
  595. vtufile<<" <DataArray type=\"Int32\" Name=\"offsets\" format=\"appended\" offset=\""<<data_size<<"\" />\n";
  596. data_size+=sizeof(uint32_t)+offset.size() *sizeof(int32_t);
  597. vtufile<<" </Polys>\n";
  598. //---------------------------------------------------------------------------
  599. vtufile<<" </Piece>\n";
  600. vtufile<<" </PolyData>\n";
  601. //===========================================================================
  602. vtufile<<" <AppendedData encoding=\"raw\">\n";
  603. vtufile<<" _";
  604. int32_t block_size;
  605. block_size=coord.size()*sizeof(VTKReal); vtufile.write((char*)&block_size, sizeof(int32_t)); vtufile.write((char*)&coord [0], coord.size()*sizeof(VTKReal));
  606. if(value.size()){ // value
  607. block_size=value.size()*sizeof(VTKReal); vtufile.write((char*)&block_size, sizeof(int32_t)); vtufile.write((char*)&value [0], value.size()*sizeof(VTKReal));
  608. }
  609. block_size=connect.size()*sizeof(int32_t); vtufile.write((char*)&block_size, sizeof(int32_t)); vtufile.write((char*)&connect[0], connect.size()*sizeof(int32_t));
  610. block_size=offset .size()*sizeof(int32_t); vtufile.write((char*)&block_size, sizeof(int32_t)); vtufile.write((char*)&offset [0], offset .size()*sizeof(int32_t));
  611. vtufile<<"\n";
  612. vtufile<<" </AppendedData>\n";
  613. //===========================================================================
  614. vtufile<<"</VTKFile>\n";
  615. vtufile.close();
  616. if(myrank) return;
  617. std::stringstream pvtufname;
  618. pvtufname<<fname<<".pvtp";
  619. std::ofstream pvtufile;
  620. pvtufile.open(pvtufname.str().c_str());
  621. if(pvtufile.fail()) return;
  622. pvtufile<<"<?xml version=\"1.0\"?>\n";
  623. pvtufile<<"<VTKFile type=\"PPolyData\">\n";
  624. pvtufile<<" <PPolyData GhostLevel=\"0\">\n";
  625. pvtufile<<" <PPoints>\n";
  626. pvtufile<<" <PDataArray type=\"Float"<<sizeof(VTKReal)*8<<"\" NumberOfComponents=\""<<COORD_DIM<<"\" Name=\"Position\"/>\n";
  627. pvtufile<<" </PPoints>\n";
  628. if(value.size()){ // value
  629. pvtufile<<" <PPointData>\n";
  630. pvtufile<<" <PDataArray type=\"Float"<<sizeof(VTKReal)*8<<"\" NumberOfComponents=\""<<value.size()/pt_cnt<<"\" Name=\""<<"value"<<"\"/>\n";
  631. pvtufile<<" </PPointData>\n";
  632. }
  633. {
  634. // Extract filename from path.
  635. std::stringstream vtupath;
  636. vtupath<<'/'<<fname;
  637. std::string pathname = vtupath.str();
  638. auto found = pathname.find_last_of("/\\");
  639. std::string fname_ = pathname.substr(found+1);
  640. for(Integer i=0;i<np;i++) pvtufile<<" <Piece Source=\""<<fname_<<"_"<<std::setfill('0')<<std::setw(6)<<i<<".vtp\"/>\n";
  641. }
  642. pvtufile<<" </PPolyData>\n";
  643. pvtufile<<"</VTKFile>\n";
  644. pvtufile.close();
  645. }
  646. template <class Real> void SphericalHarmonics<Real>::LegPolyDeriv(Vector<Real>& poly_val, const Vector<Real>& X, Long degree){
  647. Long N = X.Dim();
  648. Long Npoly = (degree + 1) * (degree + 2) / 2;
  649. if (poly_val.Dim() != N * Npoly) {
  650. poly_val.ReInit(N * Npoly);
  651. }
  652. Vector<Real> leg_poly(Npoly * N);
  653. LegPoly(leg_poly, X, degree);
  654. for(Long m=0;m<=degree;m++){
  655. for(Long n=0;n<=degree;n++) if(m<=n){
  656. const Real* Pn =&leg_poly[0];
  657. const Real* Pn_=&leg_poly[0];
  658. if((m+0)<=(n+0)) Pn =&leg_poly[N*(((degree*2-abs(m+0)+1)*abs(m+0))/2+(n+0))];
  659. if((m+1)<=(n+0)) Pn_=&leg_poly[N*(((degree*2-abs(m+1)+1)*abs(m+1))/2+(n+0))];
  660. Real* Hn =&poly_val[N*(((degree*2-abs(m+0)+1)*abs(m+0))/2+(n+0))];
  661. Real c1=(abs(m+0)<=(n+0)?1.0:0)*m;
  662. Real c2=(abs(m+1)<=(n+0)?1.0:0)*sqrt(n+m+1)*sqrt(n>m?n-m:1);
  663. for(Long i=0;i<N;i++){
  664. Hn[i]=-(c1*X[i]*Pn[i]+c2*sqrt(1-X[i]*X[i])*Pn_[i])/sqrt(1-X[i]*X[i]);
  665. }
  666. }
  667. }
  668. }
  669. template <class Real> void SphericalHarmonics<Real>::LegPoly(Vector<Real>& poly_val, const Vector<Real>& X, Long degree){
  670. Long N = X.Dim();
  671. Long Npoly = (degree + 1) * (degree + 2) / 2;
  672. if (poly_val.Dim() != Npoly * N) {
  673. poly_val.ReInit(Npoly * N);
  674. }
  675. Real fact=1.0/(Real)sqrt(4*M_PI);
  676. std::vector<Real> u(N);
  677. for(Long n=0;n<N;n++){
  678. u[n]=sqrt(1-X[n]*X[n]);
  679. if(X[n]*X[n]>1.0) u[n]=0;
  680. poly_val[n]=fact;
  681. }
  682. Long idx = 0;
  683. Long idx_nxt = 0;
  684. for(Long i=1;i<=degree;i++){
  685. idx_nxt += N*(degree-i+2);
  686. Real c=(i==1?sqrt(3.0/2.0):1);
  687. if(i>1)c*=sqrt((Real)(2*i+1)/(2*i));
  688. for(Long n=0;n<N;n++){
  689. poly_val[idx_nxt+n]=-poly_val[idx+n]*u[n]*c;
  690. }
  691. idx = idx_nxt;
  692. }
  693. idx = 0;
  694. for(Long m=0;m<degree;m++){
  695. for(Long n=0;n<N;n++){
  696. Real pmm=0;
  697. Real pmmp1=poly_val[idx+n];
  698. Real pll;
  699. for(Long ll=m+1;ll<=degree;ll++){
  700. Real a=sqrt(((Real)(2*ll-1)*(2*ll+1))/((ll-m)*(ll+m)));
  701. Real b=sqrt(((Real)(2*ll+1)*(ll+m-1)*(ll-m-1))/((ll-m)*(ll+m)*(2*ll-3)));
  702. pll=X[n]*a*pmmp1-b*pmm;
  703. pmm=pmmp1;
  704. pmmp1=pll;
  705. poly_val[idx+N*(ll-m)+n]=pll;
  706. }
  707. }
  708. idx+=N*(degree-m+1);
  709. }
  710. }
  711. template <class Real> const Vector<Real>& SphericalHarmonics<Real>::LegendreNodes(Long p){
  712. assert(p<SCTL_SHMAXDEG);
  713. assert(MatrixStore().Qx_.size() == SCTL_SHMAXDEG);
  714. Vector<Real>& Qx=MatrixStore().Qx_[p];
  715. if(!Qx.Dim()){
  716. Vector<double> qx1(p+1);
  717. Vector<double> qw1(p+1);
  718. cgqf(p+1, 1, 0.0, 0.0, -1.0, 1.0, &qx1[0], &qw1[0]);
  719. assert(typeid(Real) == typeid(double) || typeid(Real) == typeid(float)); // TODO: works only for float and double
  720. if (Qx.Dim() != p+1) Qx.ReInit(p+1);
  721. for (Long i = 0; i < p + 1; i++) Qx[i] = -qx1[i];
  722. }
  723. return Qx;
  724. }
  725. template <class Real> const Vector<Real>& SphericalHarmonics<Real>::LegendreWeights(Long p){
  726. assert(p<SCTL_SHMAXDEG);
  727. assert(MatrixStore().Qw_.size() == SCTL_SHMAXDEG);
  728. Vector<Real>& Qw=MatrixStore().Qw_[p];
  729. if(!Qw.Dim()){
  730. Vector<double> qx1(p+1);
  731. Vector<double> qw1(p+1);
  732. cgqf(p+1, 1, 0.0, 0.0, -1.0, 1.0, &qx1[0], &qw1[0]);
  733. assert(typeid(Real) == typeid(double) || typeid(Real) == typeid(float)); // TODO: works only for float and double
  734. if (Qw.Dim() != p+1) Qw.ReInit(p+1);
  735. for (Long i = 0; i < p + 1; i++) Qw[i] = qw1[i];
  736. }
  737. return Qw;
  738. }
  739. template <class Real> const Vector<Real>& SphericalHarmonics<Real>::SingularWeights(Long p1){
  740. assert(p1<SCTL_SHMAXDEG);
  741. assert(MatrixStore().Sw_.size() == SCTL_SHMAXDEG);
  742. Vector<Real>& Sw=MatrixStore().Sw_[p1];
  743. if(!Sw.Dim()){
  744. const Vector<Real>& qx1 = LegendreNodes(p1);
  745. const Vector<Real>& qw1 = LegendreWeights(p1);
  746. std::vector<Real> Yf(p1+1,0);
  747. { // Set Yf
  748. Vector<Real> x0(1); x0=1.0;
  749. Vector<Real> alp0((p1+1)*(p1+2)/2);
  750. LegPoly(alp0, x0, p1);
  751. Vector<Real> alp((p1+1) * (p1+1)*(p1+2)/2);
  752. LegPoly(alp, qx1, p1);
  753. for(Long j=0;j<p1+1;j++){
  754. for(Long i=0;i<p1+1;i++){
  755. Yf[i]+=4*M_PI/(2*j+1) * alp0[j] * alp[j*(p1+1)+i];
  756. }
  757. }
  758. }
  759. Sw.ReInit(p1+1);
  760. for(Long i=0;i<p1+1;i++){
  761. Sw[i]=(qw1[i]*M_PI/p1)*Yf[i]/cos(acos(qx1[i])/2);
  762. }
  763. }
  764. return Sw;
  765. }
  766. template <class Real> const Matrix<Real>& SphericalHarmonics<Real>::MatFourier(Long p0, Long p1){
  767. assert(p0<SCTL_SHMAXDEG && p1<SCTL_SHMAXDEG);
  768. assert(MatrixStore().Mf_ .size() == SCTL_SHMAXDEG*SCTL_SHMAXDEG);
  769. Matrix<Real>& Mf =MatrixStore().Mf_ [p0*SCTL_SHMAXDEG+p1];
  770. if(!Mf.Dim(0)){
  771. const Real SQRT2PI=sqrt(2*M_PI);
  772. { // Set Mf
  773. Matrix<Real> M(2*p0,2*p1);
  774. for(Long j=0;j<2*p1;j++){
  775. M[0][j]=SQRT2PI*1.0;
  776. for(Long k=1;k<p0;k++){
  777. M[2*k-1][j]=SQRT2PI*cos(j*k*M_PI/p1);
  778. M[2*k-0][j]=SQRT2PI*sin(j*k*M_PI/p1);
  779. }
  780. M[2*p0-1][j]=SQRT2PI*cos(j*p0*M_PI/p1);
  781. }
  782. Mf=M;
  783. }
  784. }
  785. return Mf;
  786. }
  787. template <class Real> const Matrix<Real>& SphericalHarmonics<Real>::MatFourierInv(Long p0, Long p1){
  788. assert(p0<SCTL_SHMAXDEG && p1<SCTL_SHMAXDEG);
  789. assert(MatrixStore().Mfinv_ .size() == SCTL_SHMAXDEG*SCTL_SHMAXDEG);
  790. Matrix<Real>& Mf =MatrixStore().Mfinv_ [p0*SCTL_SHMAXDEG+p1];
  791. if(!Mf.Dim(0)){
  792. const Real INVSQRT2PI=1.0/sqrt(2*M_PI)/p0;
  793. { // Set Mf
  794. Matrix<Real> M(2*p0,2*p1);
  795. M.SetZero();
  796. if(p1>p0) p1=p0;
  797. for(Long j=0;j<2*p0;j++){
  798. M[j][0]=INVSQRT2PI*0.5;
  799. for(Long k=1;k<p1;k++){
  800. M[j][2*k-1]=INVSQRT2PI*cos(j*k*M_PI/p0);
  801. M[j][2*k-0]=INVSQRT2PI*sin(j*k*M_PI/p0);
  802. }
  803. M[j][2*p1-1]=INVSQRT2PI*cos(j*p1*M_PI/p0);
  804. }
  805. if(p1==p0) for(Long j=0;j<2*p0;j++) M[j][2*p1-1]*=0.5;
  806. Mf=M;
  807. }
  808. }
  809. return Mf;
  810. }
  811. template <class Real> const FFT<Real>& SphericalHarmonics<Real>::OpFourier(Long Np){
  812. assert(Np<SCTL_SHMAXDEG);
  813. auto& Mf =MatrixStore().Mfftinv_ [Np];
  814. if(!Mf.Dim(0)){
  815. StaticArray<Long,1> fft_dim = {Np};
  816. Mf.Setup(FFT_Type::C2R, 1, Vector<Long>(1,fft_dim,false));
  817. }
  818. return Mf;
  819. }
  820. template <class Real> const FFT<Real>& SphericalHarmonics<Real>::OpFourierInv(Long Np){
  821. assert(Np<SCTL_SHMAXDEG);
  822. auto& Mf =MatrixStore().Mfft_ [Np];
  823. if(!Mf.Dim(0)){
  824. StaticArray<Long,1> fft_dim = {Np};
  825. Mf.Setup(FFT_Type::R2C, 1, Vector<Long>(1,fft_dim,false));
  826. }
  827. return Mf;
  828. }
  829. template <class Real> const Matrix<Real>& SphericalHarmonics<Real>::MatFourierGrad(Long p0, Long p1){
  830. assert(p0<SCTL_SHMAXDEG && p1<SCTL_SHMAXDEG);
  831. assert(MatrixStore().Mdf_.size() == SCTL_SHMAXDEG*SCTL_SHMAXDEG);
  832. Matrix<Real>& Mdf=MatrixStore().Mdf_[p0*SCTL_SHMAXDEG+p1];
  833. if(!Mdf.Dim(0)){
  834. const Real SQRT2PI=sqrt(2*M_PI);
  835. { // Set Mdf_
  836. Matrix<Real> M(2*p0,2*p1);
  837. for(Long j=0;j<2*p1;j++){
  838. M[0][j]=SQRT2PI*0.0;
  839. for(Long k=1;k<p0;k++){
  840. M[2*k-1][j]=-SQRT2PI*k*sin(j*k*M_PI/p1);
  841. M[2*k-0][j]= SQRT2PI*k*cos(j*k*M_PI/p1);
  842. }
  843. M[2*p0-1][j]=-SQRT2PI*p0*sin(j*p0*M_PI/p1);
  844. }
  845. Mdf=M;
  846. }
  847. }
  848. return Mdf;
  849. }
  850. template <class Real> const std::vector<Matrix<Real>>& SphericalHarmonics<Real>::MatLegendre(Long p0, Long p1){
  851. assert(p0<SCTL_SHMAXDEG && p1<SCTL_SHMAXDEG);
  852. assert(MatrixStore().Ml_ .size() == SCTL_SHMAXDEG*SCTL_SHMAXDEG);
  853. std::vector<Matrix<Real>>& Ml =MatrixStore().Ml_ [p0*SCTL_SHMAXDEG+p1];
  854. if(!Ml.size()){
  855. const Vector<Real>& qx1 = LegendreNodes(p1);
  856. Vector<Real> alp(qx1.Dim()*(p0+1)*(p0+2)/2);
  857. LegPoly(alp, qx1, p0);
  858. Ml.resize(p0+1);
  859. auto ptr = alp.begin();
  860. for(Long i=0;i<=p0;i++){
  861. Ml[i].ReInit(p0+1-i, qx1.Dim(), ptr);
  862. ptr+=Ml[i].Dim(0)*Ml[i].Dim(1);
  863. }
  864. }
  865. return Ml;
  866. }
  867. template <class Real> const std::vector<Matrix<Real>>& SphericalHarmonics<Real>::MatLegendreInv(Long p0, Long p1){
  868. assert(p0<SCTL_SHMAXDEG && p1<SCTL_SHMAXDEG);
  869. assert(MatrixStore().Mlinv_ .size() == SCTL_SHMAXDEG*SCTL_SHMAXDEG);
  870. std::vector<Matrix<Real>>& Ml =MatrixStore().Mlinv_ [p0*SCTL_SHMAXDEG+p1];
  871. if(!Ml.size()){
  872. const Vector<Real>& qx1 = LegendreNodes(p0);
  873. const Vector<Real>& qw1 = LegendreWeights(p0);
  874. Vector<Real> alp(qx1.Dim()*(p1+1)*(p1+2)/2);
  875. LegPoly(alp, qx1, p1);
  876. Ml.resize(p1+1);
  877. auto ptr = alp.begin();
  878. for(Long i=0;i<=p1;i++){
  879. Ml[i].ReInit(qx1.Dim(), p1+1-i);
  880. Matrix<Real> M(p1+1-i, qx1.Dim(), ptr, false);
  881. for(Long j=0;j<p1+1-i;j++){ // Transpose and weights
  882. for(Long k=0;k<qx1.Dim();k++){
  883. Ml[i][k][j]=M[j][k]*qw1[k]*2*M_PI;
  884. }
  885. }
  886. ptr+=Ml[i].Dim(0)*Ml[i].Dim(1);
  887. }
  888. }
  889. return Ml;
  890. }
  891. template <class Real> const std::vector<Matrix<Real>>& SphericalHarmonics<Real>::MatLegendreGrad(Long p0, Long p1){
  892. assert(p0<SCTL_SHMAXDEG && p1<SCTL_SHMAXDEG);
  893. assert(MatrixStore().Mdl_.size() == SCTL_SHMAXDEG*SCTL_SHMAXDEG);
  894. std::vector<Matrix<Real>>& Mdl=MatrixStore().Mdl_[p0*SCTL_SHMAXDEG+p1];
  895. if(!Mdl.size()){
  896. const Vector<Real>& qx1 = LegendreNodes(p1);
  897. Vector<Real> alp(qx1.Dim()*(p0+1)*(p0+2)/2);
  898. LegPolyDeriv(alp, qx1, p0);
  899. Mdl.resize(p0+1);
  900. auto ptr = alp.begin();
  901. for(Long i=0;i<=p0;i++){
  902. Mdl[i].ReInit(p0+1-i, qx1.Dim(), ptr);
  903. ptr+=Mdl[i].Dim(0)*Mdl[i].Dim(1);
  904. }
  905. }
  906. return Mdl;
  907. }
  908. template <class Real> const std::vector<Matrix<Real>>& SphericalHarmonics<Real>::MatRotate(Long p0){
  909. std::vector<std::vector<Long>> coeff_perm(p0+1);
  910. { // Set coeff_perm
  911. for(Long n=0;n<=p0;n++) coeff_perm[n].resize(std::min(2*n+1,2*p0));
  912. Long itr=0;
  913. for(Long i=0;i<2*p0;i++){
  914. Long m=(i+1)/2;
  915. for(Long n=m;n<=p0;n++){
  916. coeff_perm[n][i]=itr;
  917. itr++;
  918. }
  919. }
  920. }
  921. assert(p0<SCTL_SHMAXDEG);
  922. assert(MatrixStore().Mr_.size() == SCTL_SHMAXDEG);
  923. std::vector<Matrix<Real>>& Mr=MatrixStore().Mr_[p0];
  924. if(!Mr.size()){
  925. const Real SQRT2PI=sqrt(2*M_PI);
  926. Long Ncoef=p0*(p0+2);
  927. Long Ngrid=2*p0*(p0+1);
  928. Long Naleg=(p0+1)*(p0+2)/2;
  929. Matrix<Real> Mcoord0(3,Ngrid);
  930. const Vector<Real>& x=LegendreNodes(p0);
  931. for(Long i=0;i<p0+1;i++){ // Set Mcoord0
  932. for(Long j=0;j<2*p0;j++){
  933. Mcoord0[0][i*2*p0+j]=x[i];
  934. Mcoord0[1][i*2*p0+j]=sqrt(1-x[i]*x[i])*sin(M_PI*j/p0);
  935. Mcoord0[2][i*2*p0+j]=sqrt(1-x[i]*x[i])*cos(M_PI*j/p0);
  936. }
  937. }
  938. for(Long l=0;l<p0+1;l++){ // For each rotation angle
  939. Matrix<Real> Mcoord1;
  940. { // Rotate coordinates
  941. Matrix<Real> M(COORD_DIM, COORD_DIM);
  942. Real cos_=-x[l];
  943. Real sin_=-sqrt(1.0-x[l]*x[l]);
  944. M[0][0]= cos_; M[0][1]=0; M[0][2]=-sin_;
  945. M[1][0]= 0; M[1][1]=1; M[1][2]= 0;
  946. M[2][0]= sin_; M[2][1]=0; M[2][2]= cos_;
  947. Mcoord1=M*Mcoord0;
  948. }
  949. Matrix<Real> Mleg(Naleg, Ngrid);
  950. { // Set Mleg
  951. const Vector<Real> Vcoord1(Mcoord1.Dim(0)*Mcoord1.Dim(1), Mcoord1.begin(), false);
  952. Vector<Real> Vleg(Mleg.Dim(0)*Mleg.Dim(1), Mleg.begin(), false);
  953. LegPoly(Vleg, Vcoord1, p0);
  954. }
  955. Vector<Real> theta(Ngrid);
  956. for(Long i=0;i<theta.Dim();i++){ // Set theta
  957. theta[i]=atan2(Mcoord1[1][i],Mcoord1[2][i]);
  958. }
  959. Matrix<Real> Mcoef2grid(Ncoef, Ngrid);
  960. { // Build Mcoef2grid
  961. Long offset0=0;
  962. Long offset1=0;
  963. for(Long i=0;i<p0+1;i++){
  964. Long len=p0+1-i;
  965. { // P * cos
  966. for(Long j=0;j<len;j++){
  967. for(Long k=0;k<Ngrid;k++){
  968. Mcoef2grid[offset1+j][k]=SQRT2PI*Mleg[offset0+j][k]*cos(i*theta[k]);
  969. }
  970. }
  971. offset1+=len;
  972. }
  973. if(i!=0 && i!=p0){ // P * sin
  974. for(Long j=0;j<len;j++){
  975. for(Long k=0;k<Ngrid;k++){
  976. Mcoef2grid[offset1+j][k]=SQRT2PI*Mleg[offset0+j][k]*sin(i*theta[k]);
  977. }
  978. }
  979. offset1+=len;
  980. }
  981. offset0+=len;
  982. }
  983. assert(offset0==Naleg);
  984. assert(offset1==Ncoef);
  985. }
  986. Vector<Real> Vcoef2coef(Ncoef*Ncoef);
  987. Vector<Real> Vcoef2grid(Ncoef*Ngrid, Mcoef2grid[0], false);
  988. Grid2SHC(Vcoef2grid, p0+1, 2*p0, p0, Vcoef2coef, SHCArrange::COL_MAJOR_NONZERO);
  989. Matrix<Real> Mcoef2coef(Ncoef, Ncoef, Vcoef2coef.begin(), false);
  990. for(Long n=0;n<=p0;n++){ // Create matrices for fast rotation
  991. Matrix<Real> M(coeff_perm[n].size(),coeff_perm[n].size());
  992. for(Long i=0;i<(Long)coeff_perm[n].size();i++){
  993. for(Long j=0;j<(Long)coeff_perm[n].size();j++){
  994. M[i][j]=Mcoef2coef[coeff_perm[n][i]][coeff_perm[n][j]];
  995. }
  996. }
  997. Mr.push_back(M);
  998. }
  999. }
  1000. }
  1001. return Mr;
  1002. }
  1003. template <class Real> void SphericalHarmonics<Real>::SHC2GridTranspose(const Vector<Real>& X, Long p0, Long p1, Vector<Real>& S){
  1004. Matrix<Real> Mf =SphericalHarmonics<Real>::MatFourier(p1,p0).Transpose();
  1005. std::vector<Matrix<Real>> Ml =SphericalHarmonics<Real>::MatLegendre(p1,p0);
  1006. for(Long i=0;i<(Long)Ml.size();i++) Ml[i]=Ml[i].Transpose();
  1007. assert(p1==(Long)Ml.size()-1);
  1008. assert(p0==Mf.Dim(0)/2);
  1009. assert(p1==Mf.Dim(1)/2);
  1010. Long N=X.Dim()/(2*p0*(p0+1));
  1011. assert(N*2*p0*(p0+1)==X.Dim());
  1012. if(S.Dim()!=N*(p1*(p1+2))) S.ReInit(N*(p1*(p1+2)));
  1013. Vector<Real> B0, B1;
  1014. B0.ReInit(N* p1*(p1+2));
  1015. B1.ReInit(N*2*p1*(p0+1));
  1016. #pragma omp parallel
  1017. { // Evaluate Fourier and transpose
  1018. Integer tid=omp_get_thread_num();
  1019. Integer omp_p=omp_get_num_threads();
  1020. Long a=(tid+0)*N*(p0+1)/omp_p;
  1021. Long b=(tid+1)*N*(p0+1)/omp_p;
  1022. const Long block_size=16;
  1023. Matrix<Real> B2(block_size,2*p1);
  1024. for(Long i0=a;i0<b;i0+=block_size){
  1025. Long i1=std::min(b,i0+block_size);
  1026. const Matrix<Real> Min (i1-i0,2*p0, (Iterator<Real>)X.begin()+i0*2*p0, false);
  1027. Matrix<Real> Mout(i1-i0,2*p1, B2.begin(), false);
  1028. Matrix<Real>::GEMM(Mout, Min, Mf);
  1029. for(Long i=i0;i<i1;i++){
  1030. for(Long j=0;j<2*p1;j++){
  1031. B1[j*N*(p0+1)+i]=B2[i-i0][j];
  1032. }
  1033. }
  1034. }
  1035. }
  1036. #pragma omp parallel
  1037. { // Evaluate Legendre polynomial
  1038. Integer tid=omp_get_thread_num();
  1039. Integer omp_p=omp_get_num_threads();
  1040. Long offset0=0;
  1041. Long offset1=0;
  1042. for(Long i=0;i<p1+1;i++){
  1043. Long N0=2*N;
  1044. if(i==0 || i==p1) N0=N;
  1045. Matrix<Real> Min (N0, p0+1 , B1.begin()+offset0, false);
  1046. Matrix<Real> Mout(N0, p1+1-i, B0.begin()+offset1, false);
  1047. { // Mout = Min * Ml[i] // split between threads
  1048. Long a=(tid+0)*N0/omp_p;
  1049. Long b=(tid+1)*N0/omp_p;
  1050. if(a<b){
  1051. Matrix<Real> Min_ (b-a, Min .Dim(1), Min [a], false);
  1052. Matrix<Real> Mout_(b-a, Mout.Dim(1), Mout[a], false);
  1053. Matrix<Real>::GEMM(Mout_,Min_,Ml[i]);
  1054. }
  1055. }
  1056. offset0+=Min .Dim(0)*Min .Dim(1);
  1057. offset1+=Mout.Dim(0)*Mout.Dim(1);
  1058. }
  1059. }
  1060. #pragma omp parallel
  1061. { // S <-- Rearrange(B0)
  1062. Integer tid=omp_get_thread_num();
  1063. Integer omp_p=omp_get_num_threads();
  1064. Long a=(tid+0)*N/omp_p;
  1065. Long b=(tid+1)*N/omp_p;
  1066. for(Long i=a;i<b;i++){
  1067. Long offset=0;
  1068. for(Long j=0;j<2*p1;j++){
  1069. Long len=p1+1-(j+1)/2;
  1070. Real* B_=&B0[i*len+N*offset];
  1071. Real* S_=&S[i*p1*(p1+2)+offset];
  1072. for(Long k=0;k<len;k++) S_[k]=B_[k];
  1073. offset+=len;
  1074. }
  1075. }
  1076. }
  1077. }
  1078. template <class Real> void SphericalHarmonics<Real>::RotateAll(const Vector<Real>& S, Long p0, Long dof, Vector<Real>& S_){
  1079. const std::vector<Matrix<Real>>& Mr=MatRotate(p0);
  1080. std::vector<std::vector<Long>> coeff_perm(p0+1);
  1081. { // Set coeff_perm
  1082. for(Long n=0;n<=p0;n++) coeff_perm[n].resize(std::min(2*n+1,2*p0));
  1083. Long itr=0;
  1084. for(Long i=0;i<2*p0;i++){
  1085. Long m=(i+1)/2;
  1086. for(Long n=m;n<=p0;n++){
  1087. coeff_perm[n][i]=itr;
  1088. itr++;
  1089. }
  1090. }
  1091. }
  1092. Long Ncoef=p0*(p0+2);
  1093. Long N=S.Dim()/Ncoef/dof;
  1094. assert(N*Ncoef*dof==S.Dim());
  1095. if(S_.Dim()!=N*dof*Ncoef*p0*(p0+1)) S_.ReInit(N*dof*Ncoef*p0*(p0+1));
  1096. const Matrix<Real> S0(N*dof, Ncoef, (Iterator<Real>)S.begin(), false);
  1097. Matrix<Real> S1(N*dof*p0*(p0+1), Ncoef, S_.begin(), false);
  1098. #pragma omp parallel
  1099. { // Construct all p0*(p0+1) rotations
  1100. Integer tid=omp_get_thread_num();
  1101. Integer omp_p=omp_get_num_threads();
  1102. Matrix<Real> B0(dof*p0,Ncoef); // memory buffer
  1103. std::vector<Matrix<Real>> Bi(p0+1), Bo(p0+1); // memory buffers
  1104. for(Long i=0;i<=p0;i++){ // initialize Bi, Bo
  1105. Bi[i].ReInit(dof*p0,coeff_perm[i].size());
  1106. Bo[i].ReInit(dof*p0,coeff_perm[i].size());
  1107. }
  1108. Long a=(tid+0)*N/omp_p;
  1109. Long b=(tid+1)*N/omp_p;
  1110. for(Long i=a;i<b;i++){
  1111. for(Long d=0;d<dof;d++){
  1112. for(Long j=0;j<p0;j++){
  1113. Long offset=0;
  1114. for(Long k=0;k<p0+1;k++){
  1115. Real r[2]={cos(k*j*M_PI/p0),-sin(k*j*M_PI/p0)}; // exp(i*k*theta)
  1116. Long len=p0+1-k;
  1117. if(k!=0 && k!=p0){
  1118. for(Long l=0;l<len;l++){
  1119. Real x[2];
  1120. x[0]=S0[i*dof+d][offset+len*0+l];
  1121. x[1]=S0[i*dof+d][offset+len*1+l];
  1122. B0[j*dof+d][offset+len*0+l]=x[0]*r[0]-x[1]*r[1];
  1123. B0[j*dof+d][offset+len*1+l]=x[0]*r[1]+x[1]*r[0];
  1124. }
  1125. offset+=2*len;
  1126. }else{
  1127. for(Long l=0;l<len;l++){
  1128. B0[j*dof+d][offset+l]=S0[i*dof+d][offset+l];
  1129. }
  1130. offset+=len;
  1131. }
  1132. }
  1133. assert(offset==Ncoef);
  1134. }
  1135. }
  1136. { // Fast rotation
  1137. for(Long k=0;k<dof*p0;k++){ // forward permutation
  1138. for(Long l=0;l<=p0;l++){
  1139. for(Long j=0;j<(Long)coeff_perm[l].size();j++){
  1140. Bi[l][k][j]=B0[k][coeff_perm[l][j]];
  1141. }
  1142. }
  1143. }
  1144. for(Long t=0;t<=p0;t++){
  1145. for(Long l=0;l<=p0;l++){ // mat-vec
  1146. Matrix<Real>::GEMM(Bo[l],Bi[l],Mr[t*(p0+1)+l]);
  1147. }
  1148. Matrix<Real> Mout(dof*p0,Ncoef, S1[(i*(p0+1)+t)*dof*p0], false);
  1149. for(Long k=0;k<dof*p0;k++){ // reverse permutation
  1150. for(Long l=0;l<=p0;l++){
  1151. for(Long j=0;j<(Long)coeff_perm[l].size();j++){
  1152. Mout[k][coeff_perm[l][j]]=Bo[l][k][j];
  1153. }
  1154. }
  1155. }
  1156. }
  1157. }
  1158. }
  1159. }
  1160. }
  1161. template <class Real> void SphericalHarmonics<Real>::RotateTranspose(const Vector<Real>& S_, Long p0, Long dof, Vector<Real>& S){
  1162. std::vector<Matrix<Real>> Mr=MatRotate(p0);
  1163. for(Long i=0;i<(Long)Mr.size();i++) Mr[i]=Mr[i].Transpose();
  1164. std::vector<std::vector<Long>> coeff_perm(p0+1);
  1165. { // Set coeff_perm
  1166. for(Long n=0;n<=p0;n++) coeff_perm[n].resize(std::min(2*n+1,2*p0));
  1167. Long itr=0;
  1168. for(Long i=0;i<2*p0;i++){
  1169. Long m=(i+1)/2;
  1170. for(Long n=m;n<=p0;n++){
  1171. coeff_perm[n][i]=itr;
  1172. itr++;
  1173. }
  1174. }
  1175. }
  1176. Long Ncoef=p0*(p0+2);
  1177. Long N=S_.Dim()/Ncoef/dof/(p0*(p0+1));
  1178. assert(N*Ncoef*dof*(p0*(p0+1))==S_.Dim());
  1179. if(S.Dim()!=N*dof*Ncoef*p0*(p0+1)) S.ReInit(N*dof*Ncoef*p0*(p0+1));
  1180. Matrix<Real> S0(N*dof*p0*(p0+1), Ncoef, S.begin(), false);
  1181. const Matrix<Real> S1(N*dof*p0*(p0+1), Ncoef, (Iterator<Real>)S_.begin(), false);
  1182. #pragma omp parallel
  1183. { // Transpose all p0*(p0+1) rotations
  1184. Integer tid=omp_get_thread_num();
  1185. Integer omp_p=omp_get_num_threads();
  1186. Matrix<Real> B0(dof*p0,Ncoef); // memory buffer
  1187. std::vector<Matrix<Real>> Bi(p0+1), Bo(p0+1); // memory buffers
  1188. for(Long i=0;i<=p0;i++){ // initialize Bi, Bo
  1189. Bi[i].ReInit(dof*p0,coeff_perm[i].size());
  1190. Bo[i].ReInit(dof*p0,coeff_perm[i].size());
  1191. }
  1192. Long a=(tid+0)*N/omp_p;
  1193. Long b=(tid+1)*N/omp_p;
  1194. for(Long i=a;i<b;i++){
  1195. for(Long t=0;t<p0+1;t++){
  1196. Long idx0=(i*(p0+1)+t)*p0*dof;
  1197. { // Fast rotation
  1198. const Matrix<Real> Min(p0*dof, Ncoef, (Iterator<Real>)S1[idx0], false);
  1199. for(Long k=0;k<dof*p0;k++){ // forward permutation
  1200. for(Long l=0;l<=p0;l++){
  1201. for(Long j=0;j<(Long)coeff_perm[l].size();j++){
  1202. Bi[l][k][j]=Min[k][coeff_perm[l][j]];
  1203. }
  1204. }
  1205. }
  1206. for(Long l=0;l<=p0;l++){ // mat-vec
  1207. Matrix<Real>::GEMM(Bo[l],Bi[l],Mr[t*(p0+1)+l]);
  1208. }
  1209. for(Long k=0;k<dof*p0;k++){ // reverse permutation
  1210. for(Long l=0;l<=p0;l++){
  1211. for(Long j=0;j<(Long)coeff_perm[l].size();j++){
  1212. B0[k][coeff_perm[l][j]]=Bo[l][k][j];
  1213. }
  1214. }
  1215. }
  1216. }
  1217. for(Long j=0;j<p0;j++){
  1218. for(Long d=0;d<dof;d++){
  1219. Long idx1=idx0+j*dof+d;
  1220. Long offset=0;
  1221. for(Long k=0;k<p0+1;k++){
  1222. Real r[2]={cos(k*j*M_PI/p0),sin(k*j*M_PI/p0)}; // exp(i*k*theta)
  1223. Long len=p0+1-k;
  1224. if(k!=0 && k!=p0){
  1225. for(Long l=0;l<len;l++){
  1226. Real x[2];
  1227. x[0]=B0[j*dof+d][offset+len*0+l];
  1228. x[1]=B0[j*dof+d][offset+len*1+l];
  1229. S0[idx1][offset+len*0+l]=x[0]*r[0]-x[1]*r[1];
  1230. S0[idx1][offset+len*1+l]=x[0]*r[1]+x[1]*r[0];
  1231. }
  1232. offset+=2*len;
  1233. }else{
  1234. for(Long l=0;l<len;l++){
  1235. S0[idx1][offset+l]=B0[j*dof+d][offset+l];
  1236. }
  1237. offset+=len;
  1238. }
  1239. }
  1240. assert(offset==Ncoef);
  1241. }
  1242. }
  1243. }
  1244. }
  1245. }
  1246. }
  1247. template <class Real> void SphericalHarmonics<Real>::StokesSingularInteg(const Vector<Real>& S, Long p0, Long p1, Vector<Real>* SLMatrix, Vector<Real>* DLMatrix){
  1248. Long Ngrid=2*p0*(p0+1);
  1249. Long Ncoef= p0*(p0+2);
  1250. Long Nves=S.Dim()/(Ngrid*COORD_DIM);
  1251. if(SLMatrix) SLMatrix->ReInit(Nves*(Ncoef*COORD_DIM)*(Ncoef*COORD_DIM));
  1252. if(DLMatrix) DLMatrix->ReInit(Nves*(Ncoef*COORD_DIM)*(Ncoef*COORD_DIM));
  1253. Long BLOCK_SIZE=(Long)6e9/((3*2*p1*(p1+1))*(3*2*p0*(p0+1))*2*8); // Limit memory usage to 6GB
  1254. BLOCK_SIZE=std::min<Long>(BLOCK_SIZE,omp_get_max_threads());
  1255. BLOCK_SIZE=std::max<Long>(BLOCK_SIZE,1);
  1256. for(Long a=0;a<Nves;a+=BLOCK_SIZE){
  1257. Long b=std::min(a+BLOCK_SIZE, Nves);
  1258. Vector<Real> _SLMatrix, _DLMatrix;
  1259. if(SLMatrix) _SLMatrix.ReInit((b-a)*(Ncoef*COORD_DIM)*(Ncoef*COORD_DIM), SLMatrix->begin()+a*(Ncoef*COORD_DIM)*(Ncoef*COORD_DIM), false);
  1260. if(DLMatrix) _DLMatrix.ReInit((b-a)*(Ncoef*COORD_DIM)*(Ncoef*COORD_DIM), DLMatrix->begin()+a*(Ncoef*COORD_DIM)*(Ncoef*COORD_DIM), false);
  1261. const Vector<Real> _S ((b-a)*(Ngrid*COORD_DIM) , (Iterator<Real>)S.begin()+a*(Ngrid*COORD_DIM), false);
  1262. if(SLMatrix && DLMatrix) StokesSingularInteg_< true, true>(_S, p0, p1, _SLMatrix, _DLMatrix);
  1263. else if(SLMatrix) StokesSingularInteg_< true, false>(_S, p0, p1, _SLMatrix, _DLMatrix);
  1264. else if(DLMatrix) StokesSingularInteg_<false, true>(_S, p0, p1, _SLMatrix, _DLMatrix);
  1265. }
  1266. }
  1267. template <class Real> template <bool SLayer, bool DLayer> void SphericalHarmonics<Real>::StokesSingularInteg_(const Vector<Real>& X0, Long p0, Long p1, Vector<Real>& SL, Vector<Real>& DL){
  1268. Profile::Tic("Rotate");
  1269. Vector<Real> S0, S;
  1270. SphericalHarmonics<Real>::Grid2SHC(X0, p0+1, 2*p0, p0, S0, SHCArrange::COL_MAJOR_NONZERO);
  1271. SphericalHarmonics<Real>::RotateAll(S0, p0, COORD_DIM, S);
  1272. Profile::Toc();
  1273. Profile::Tic("Upsample");
  1274. Vector<Real> X, X_theta, X_phi, trg;
  1275. SphericalHarmonics<Real>::SHC2Grid(S, SHCArrange::COL_MAJOR_NONZERO, p0, p1+1, 2*p1, &X, &X_phi, &X_theta);
  1276. SphericalHarmonics<Real>::SHC2Pole(S, SHCArrange::COL_MAJOR_NONZERO, p0, trg);
  1277. Profile::Toc();
  1278. Profile::Tic("Stokes");
  1279. Vector<Real> SL0, DL0;
  1280. { // Stokes kernel
  1281. //Long M0=2*p0*(p0+1);
  1282. Long M1=2*p1*(p1+1);
  1283. Long N=trg.Dim()/(2*COORD_DIM);
  1284. assert(X.Dim()==M1*COORD_DIM*N);
  1285. if(SLayer && SL0.Dim()!=N*2*6*M1) SL0.ReInit(2*N*6*M1);
  1286. if(DLayer && DL0.Dim()!=N*2*6*M1) DL0.ReInit(2*N*6*M1);
  1287. const Vector<Real>& qw=SphericalHarmonics<Real>::SingularWeights(p1);
  1288. const Real scal_const_dl = 3.0/(4.0*M_PI);
  1289. const Real scal_const_sl = 1.0/(8.0*M_PI);
  1290. static Real eps=-1;
  1291. if(eps<0){
  1292. eps=1;
  1293. while(eps*(Real)0.5+(Real)1.0>1.0) eps*=0.5;
  1294. }
  1295. #pragma omp parallel
  1296. {
  1297. Integer tid=omp_get_thread_num();
  1298. Integer omp_p=omp_get_num_threads();
  1299. Long a=(tid+0)*N/omp_p;
  1300. Long b=(tid+1)*N/omp_p;
  1301. for(Long i=a;i<b;i++){
  1302. for(Long t=0;t<2;t++){
  1303. Real tx, ty, tz;
  1304. { // Read target coordinates
  1305. tx=trg[i*2*COORD_DIM+0*2+t];
  1306. ty=trg[i*2*COORD_DIM+1*2+t];
  1307. tz=trg[i*2*COORD_DIM+2*2+t];
  1308. }
  1309. for(Long j0=0;j0<p1+1;j0++){
  1310. for(Long j1=0;j1<2*p1;j1++){
  1311. Long s=2*p1*j0+j1;
  1312. Real dx, dy, dz;
  1313. { // Compute dx, dy, dz
  1314. dx=tx-X[(i*COORD_DIM+0)*M1+s];
  1315. dy=ty-X[(i*COORD_DIM+1)*M1+s];
  1316. dz=tz-X[(i*COORD_DIM+2)*M1+s];
  1317. }
  1318. Real nx, ny, nz;
  1319. { // Compute source normal
  1320. Real x_theta=X_phi[(i*COORD_DIM+0)*M1+s];
  1321. Real y_theta=X_phi[(i*COORD_DIM+1)*M1+s];
  1322. Real z_theta=X_phi[(i*COORD_DIM+2)*M1+s];
  1323. Real x_phi=X_theta[(i*COORD_DIM+0)*M1+s];
  1324. Real y_phi=X_theta[(i*COORD_DIM+1)*M1+s];
  1325. Real z_phi=X_theta[(i*COORD_DIM+2)*M1+s];
  1326. nx=(y_theta*z_phi-z_theta*y_phi);
  1327. ny=(z_theta*x_phi-x_theta*z_phi);
  1328. nz=(x_theta*y_phi-y_theta*x_phi);
  1329. }
  1330. Real area_elem=1.0;
  1331. if(SLayer){ // Compute area_elem
  1332. area_elem=sqrt(nx*nx+ny*ny+nz*nz);
  1333. }
  1334. Real rinv, rinv2;
  1335. { // Compute rinv, rinv2
  1336. Real r2=dx*dx+dy*dy+dz*dz;
  1337. rinv=1.0/sqrt(r2);
  1338. if(r2<=eps) rinv=0;
  1339. rinv2=rinv*rinv;
  1340. }
  1341. if(DLayer){
  1342. Real rinv5=rinv2*rinv2*rinv;
  1343. Real r_dot_n_rinv5=scal_const_dl*qw[j0*t+(p1-j0)*(1-t)] * (nx*dx+ny*dy+nz*dz)*rinv5;
  1344. DL0[((i*2+t)*6+0)*M1+s]=dx*dx*r_dot_n_rinv5;
  1345. DL0[((i*2+t)*6+1)*M1+s]=dx*dy*r_dot_n_rinv5;
  1346. DL0[((i*2+t)*6+2)*M1+s]=dx*dz*r_dot_n_rinv5;
  1347. DL0[((i*2+t)*6+3)*M1+s]=dy*dy*r_dot_n_rinv5;
  1348. DL0[((i*2+t)*6+4)*M1+s]=dy*dz*r_dot_n_rinv5;
  1349. DL0[((i*2+t)*6+5)*M1+s]=dz*dz*r_dot_n_rinv5;
  1350. }
  1351. if(SLayer){
  1352. Real area_rinv =scal_const_sl*qw[j0*t+(p1-j0)*(1-t)] * area_elem*rinv;
  1353. Real area_rinv2=area_rinv*rinv2;
  1354. SL0[((i*2+t)*6+0)*M1+s]=area_rinv+dx*dx*area_rinv2;
  1355. SL0[((i*2+t)*6+1)*M1+s]= dx*dy*area_rinv2;
  1356. SL0[((i*2+t)*6+2)*M1+s]= dx*dz*area_rinv2;
  1357. SL0[((i*2+t)*6+3)*M1+s]=area_rinv+dy*dy*area_rinv2;
  1358. SL0[((i*2+t)*6+4)*M1+s]= dy*dz*area_rinv2;
  1359. SL0[((i*2+t)*6+5)*M1+s]=area_rinv+dz*dz*area_rinv2;
  1360. }
  1361. }
  1362. }
  1363. }
  1364. }
  1365. }
  1366. Profile::Add_FLOP(20*(2*p1)*(p1+1)*2*N);
  1367. if(SLayer) Profile::Add_FLOP((19+6)*(2*p1)*(p1+1)*2*N);
  1368. if(DLayer) Profile::Add_FLOP( 22 *(2*p1)*(p1+1)*2*N);
  1369. }
  1370. Profile::Toc();
  1371. Profile::Tic("UpsampleTranspose");
  1372. Vector<Real> SL1, DL1;
  1373. SphericalHarmonics<Real>::SHC2GridTranspose(SL0, p1, p0, SL1);
  1374. SphericalHarmonics<Real>::SHC2GridTranspose(DL0, p1, p0, DL1);
  1375. Profile::Toc();
  1376. Profile::Tic("RotateTranspose");
  1377. Vector<Real> SL2, DL2;
  1378. SphericalHarmonics<Real>::RotateTranspose(SL1, p0, 2*6, SL2);
  1379. SphericalHarmonics<Real>::RotateTranspose(DL1, p0, 2*6, DL2);
  1380. Profile::Toc();
  1381. Profile::Tic("Rearrange");
  1382. Vector<Real> SL3, DL3;
  1383. { // Transpose
  1384. Long Ncoef=p0*(p0+2);
  1385. Long Ngrid=2*p0*(p0+1);
  1386. { // Transpose SL2
  1387. Long N=SL2.Dim()/(6*Ncoef*Ngrid);
  1388. SL3.ReInit(N*COORD_DIM*Ncoef*COORD_DIM*Ngrid);
  1389. #pragma omp parallel
  1390. {
  1391. Integer tid=omp_get_thread_num();
  1392. Integer omp_p=omp_get_num_threads();
  1393. Matrix<Real> B(COORD_DIM*Ncoef,Ngrid*COORD_DIM);
  1394. Long a=(tid+0)*N/omp_p;
  1395. Long b=(tid+1)*N/omp_p;
  1396. for(Long i=a;i<b;i++){
  1397. Matrix<Real> M0(Ngrid*6, Ncoef, SL2.begin()+i*Ngrid*6*Ncoef, false);
  1398. for(Long k=0;k<Ncoef;k++){ // Transpose
  1399. for(Long j=0;j<Ngrid;j++){ // TODO: needs blocking
  1400. B[k+Ncoef*0][j*COORD_DIM+0]=M0[j*6+0][k];
  1401. B[k+Ncoef*1][j*COORD_DIM+0]=M0[j*6+1][k];
  1402. B[k+Ncoef*2][j*COORD_DIM+0]=M0[j*6+2][k];
  1403. B[k+Ncoef*0][j*COORD_DIM+1]=M0[j*6+1][k];
  1404. B[k+Ncoef*1][j*COORD_DIM+1]=M0[j*6+3][k];
  1405. B[k+Ncoef*2][j*COORD_DIM+1]=M0[j*6+4][k];
  1406. B[k+Ncoef*0][j*COORD_DIM+2]=M0[j*6+2][k];
  1407. B[k+Ncoef*1][j*COORD_DIM+2]=M0[j*6+4][k];
  1408. B[k+Ncoef*2][j*COORD_DIM+2]=M0[j*6+5][k];
  1409. }
  1410. }
  1411. Matrix<Real> M1(Ncoef*COORD_DIM, COORD_DIM*Ngrid, SL3.begin()+i*COORD_DIM*Ncoef*COORD_DIM*Ngrid, false);
  1412. for(Long k=0;k<B.Dim(0);k++){ // Rearrange
  1413. for(Long j0=0;j0<COORD_DIM;j0++){
  1414. for(Long j1=0;j1<p0+1;j1++){
  1415. for(Long j2=0;j2<p0;j2++) M1[k][((j0*(p0+1)+ j1)*2+0)*p0+j2]=B[k][((j1*p0+j2)*2+0)*COORD_DIM+j0];
  1416. for(Long j2=0;j2<p0;j2++) M1[k][((j0*(p0+1)+p0-j1)*2+1)*p0+j2]=B[k][((j1*p0+j2)*2+1)*COORD_DIM+j0];
  1417. }
  1418. }
  1419. }
  1420. }
  1421. }
  1422. }
  1423. { // Transpose DL2
  1424. Long N=DL2.Dim()/(6*Ncoef*Ngrid);
  1425. DL3.ReInit(N*COORD_DIM*Ncoef*COORD_DIM*Ngrid);
  1426. #pragma omp parallel
  1427. {
  1428. Integer tid=omp_get_thread_num();
  1429. Integer omp_p=omp_get_num_threads();
  1430. Matrix<Real> B(COORD_DIM*Ncoef,Ngrid*COORD_DIM);
  1431. Long a=(tid+0)*N/omp_p;
  1432. Long b=(tid+1)*N/omp_p;
  1433. for(Long i=a;i<b;i++){
  1434. Matrix<Real> M0(Ngrid*6, Ncoef, DL2.begin()+i*Ngrid*6*Ncoef, false);
  1435. for(Long k=0;k<Ncoef;k++){ // Transpose
  1436. for(Long j=0;j<Ngrid;j++){ // TODO: needs blocking
  1437. B[k+Ncoef*0][j*COORD_DIM+0]=M0[j*6+0][k];
  1438. B[k+Ncoef*1][j*COORD_DIM+0]=M0[j*6+1][k];
  1439. B[k+Ncoef*2][j*COORD_DIM+0]=M0[j*6+2][k];
  1440. B[k+Ncoef*0][j*COORD_DIM+1]=M0[j*6+1][k];
  1441. B[k+Ncoef*1][j*COORD_DIM+1]=M0[j*6+3][k];
  1442. B[k+Ncoef*2][j*COORD_DIM+1]=M0[j*6+4][k];
  1443. B[k+Ncoef*0][j*COORD_DIM+2]=M0[j*6+2][k];
  1444. B[k+Ncoef*1][j*COORD_DIM+2]=M0[j*6+4][k];
  1445. B[k+Ncoef*2][j*COORD_DIM+2]=M0[j*6+5][k];
  1446. }
  1447. }
  1448. Matrix<Real> M1(Ncoef*COORD_DIM, COORD_DIM*Ngrid, DL3.begin()+i*COORD_DIM*Ncoef*COORD_DIM*Ngrid, false);
  1449. for(Long k=0;k<B.Dim(0);k++){ // Rearrange
  1450. for(Long j0=0;j0<COORD_DIM;j0++){
  1451. for(Long j1=0;j1<p0+1;j1++){
  1452. for(Long j2=0;j2<p0;j2++) M1[k][((j0*(p0+1)+ j1)*2+0)*p0+j2]=B[k][((j1*p0+j2)*2+0)*COORD_DIM+j0];
  1453. for(Long j2=0;j2<p0;j2++) M1[k][((j0*(p0+1)+p0-j1)*2+1)*p0+j2]=B[k][((j1*p0+j2)*2+1)*COORD_DIM+j0];
  1454. }
  1455. }
  1456. }
  1457. }
  1458. }
  1459. }
  1460. }
  1461. Profile::Toc();
  1462. Profile::Tic("Grid2SHC");
  1463. SphericalHarmonics<Real>::Grid2SHC(SL3, p0+1, 2*p0, p0, SL, SHCArrange::COL_MAJOR_NONZERO);
  1464. SphericalHarmonics<Real>::Grid2SHC(DL3, p0+1, 2*p0, p0, DL, SHCArrange::COL_MAJOR_NONZERO);
  1465. Profile::Toc();
  1466. }
  1467. } // end namespace