@@ -211,6 +211,9 @@ namespace BasisClasses
211211
212212 void Spherical::initialize ()
213213 {
214+ // Identifier
215+ //
216+ BasisID = " Spherical" ;
214217
215218 // Assign some defaults
216219 //
@@ -329,12 +332,11 @@ namespace BasisClasses
329332
330333 // Initialize covariance
331334 //
332- if (pcavar) init_covariance ( );
335+ if (pcavar) enableCoefCovariance ( true , sampT );
333336
334337 // Set spherical coordinates
335338 //
336339 coordinates = Coord::Spherical;
337- BasisID = " Spherical" ;
338340 }
339341
340342 void Spherical::init_covariance ()
@@ -1989,6 +1991,7 @@ namespace BasisClasses
19891991 " tkcum" ,
19901992 " tk_type" ,
19911993 " cachename" ,
1994+ " pcavar" ,
19921995 " nint" ,
19931996 " totalCovar" ,
19941997 " fullCovar"
@@ -2006,6 +2009,40 @@ namespace BasisClasses
20062009 initialize ();
20072010 }
20082011
2012+ void FlatDisk::init_covariance ()
2013+ {
2014+ if (pcavar) {
2015+
2016+ meanV.resize (sampT);
2017+ for (auto & v : meanV) {
2018+ v.resize (mmax+1 );
2019+ for (auto & vec : v) vec.resize (nmax);
2020+ }
2021+
2022+ covrV.resize (sampT);
2023+ for (auto & v : covrV) {
2024+ v.resize (mmax+1 );
2025+ for (auto & mat : v) mat.resize (nmax, nmax);
2026+ }
2027+
2028+ sampleCounts.resize (sampT);
2029+ sampleMasses.resize (sampT);
2030+
2031+ zero_covariance ();
2032+ }
2033+ }
2034+
2035+ void FlatDisk::zero_covariance ()
2036+ {
2037+ for (int T=0 ; T<sampT; T++) {
2038+ for (auto & v : meanV[T]) v.setZero ();
2039+ for (auto & v : covrV[T]) v.setZero ();
2040+ }
2041+
2042+ sampleCounts.setZero ();
2043+ sampleMasses.setZero ();
2044+ }
2045+
20092046 void FlatDisk::initialize ()
20102047 {
20112048 // Basis identifier
@@ -2057,6 +2094,10 @@ namespace BasisClasses
20572094 if (conf[" NO_M1" ]) NO_M1 = conf[" NO_M1" ].as <bool >();
20582095 if (conf[" EVEN_M" ]) EVEN_M = conf[" EVEN_M" ].as <bool >();
20592096 if (conf[" M0_ONLY" ]) M0_only = conf[" M0_ONLY" ].as <bool >();
2097+ if (conf[" pcavar" ]) pcavar = conf[" pcavar" ].as <bool >();
2098+ if (conf[" subsamp" ]) sampT = conf[" subsamp" ].as <int >();
2099+
2100+ sampT = std::max (1 , sampT); // Sanity
20602101 }
20612102 catch (YAML::Exception & error) {
20622103 if (myid==0 ) std::cout << " Error parsing parameter stanza for <"
@@ -2124,6 +2165,11 @@ namespace BasisClasses
21242165 // Set cylindrical coordindates
21252166 //
21262167 coordinates = Coord::Cylindrical;
2168+
2169+ // Initialize covariance
2170+ //
2171+ if (pcavar) enableCoefCovariance (true , sampT);
2172+
21272173 }
21282174
21292175 void FlatDisk::reset_coefs (void )
@@ -2132,6 +2178,7 @@ namespace BasisClasses
21322178 totalMass = 0.0 ;
21332179 used = 0 ;
21342180 G = 1.0 ;
2181+ if (pcavar) zero_covariance ();
21352182 }
21362183
21372184
@@ -2251,9 +2298,21 @@ namespace BasisClasses
22512298
22522299 ortho->get_pot (potd[tid], R, 0.0 );
22532300
2301+ // Sample index for pcavar
2302+ int T = 0 ;
2303+ if (pcavar) {
2304+ T = used % sampT;
2305+ sampleCounts (T) += 1 ;
2306+ sampleMasses (T) += mass;
2307+ }
2308+
22542309 // M loop
22552310 for (int m=0 , moffset=0 ; m<=mmax; m++) {
22562311
2312+ double ccos = cos (phi*m);
2313+ double ssin = sin (phi*m);
2314+
2315+ // Coefficient accumulation
22572316 if (m==0 ) {
22582317 for (int n=0 ; n<nmax; n++) {
22592318 expcoef (moffset, n) += potd[tid](m, n)* mass * norm0;
@@ -2262,14 +2321,21 @@ namespace BasisClasses
22622321 moffset++;
22632322 }
22642323 else {
2265- double ccos = cos (phi*m);
2266- double ssin = sin (phi*m);
22672324 for (int n=0 ; n<nmax; n++) {
22682325 expcoef (moffset , n) += ccos * potd[tid](m, n) * mass * norm1;
22692326 expcoef (moffset+1 , n) += ssin * potd[tid](m, n) * mass * norm1;
22702327 }
22712328 moffset+=2 ;
22722329 }
2330+
2331+ // Covariance and subsample accumulation
2332+ if (pcavar) {
2333+ Eigen::VectorXcd g = std::complex <double >(ccos, ssin) *
2334+ potd[tid].row (m).transpose () * norm0;
2335+
2336+ meanV[T][m].noalias () += g * mass;
2337+ covrV[T][m].noalias () += g * g.adjoint () * mass;
2338+ }
22732339 }
22742340 }
22752341
@@ -3893,7 +3959,8 @@ namespace BasisClasses
38933959
38943960 // Initialize covariance
38953961 //
3896- if (pcavar) init_covariance ();
3962+ if (pcavar) enableCoefCovariance (true , sampT);
3963+
38973964 }
38983965
38993966 Cube::Cube (const std::string& confstr) : BiorthBasis(confstr, " cube" )
@@ -3902,7 +3969,7 @@ namespace BasisClasses
39023969
39033970 // Initialize covariance
39043971 //
3905- if (pcavar) init_covariance ( );
3972+ if (pcavar) enableCoefCovariance ( true , sampT );
39063973 }
39073974
39083975 void Cube::initialize ()
@@ -5120,6 +5187,14 @@ namespace BasisClasses
51205187 file.createAttribute <double >(" hcyl" , HighFive::DataSpace::From (hcyl)).write (hcyl);
51215188 }
51225189
5190+ void FlatDisk::writeCovarH5Params (HighFive::File& file)
5191+ {
5192+ file.createAttribute <int >(" mmax" , HighFive::DataSpace::From (mmax)).write (mmax);
5193+ file.createAttribute <int >(" nmax" , HighFive::DataSpace::From (nmax)).write (nmax);
5194+ file.createAttribute <double >(" rcylmin" , HighFive::DataSpace::From (rcylmin)).write (rcylmin);
5195+ file.createAttribute <double >(" rcylmax" , HighFive::DataSpace::From (rcylmax)).write (rcylmax);
5196+ }
5197+
51235198 void Cube::writeCovarH5Params (HighFive::File& file)
51245199 {
51255200 file.createAttribute <int >(" nminx" , HighFive::DataSpace::From (nminx)).write (nminx);
0 commit comments