Skip to content

Commit 1e06c49

Browse files
author
Martin D. Weinberg
committed
Add missing write calls from covarStore
1 parent 2913fe5 commit 1e06c49

2 files changed

Lines changed: 148 additions & 8 deletions

File tree

expui/BiorthBasis.H

Lines changed: 67 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -520,6 +520,7 @@ namespace BasisClasses
520520
}
521521
}
522522
}
523+
covarStore->writeCoefCovariance(compname, runtag, covarData, time);
523524
} else {
524525
throw std::runtime_error("Spherical::writeCoefCovariance: covariance storage not initialized");
525526
}
@@ -745,6 +746,18 @@ namespace BasisClasses
745746
void computeAccel(double x, double y, double z,
746747
Eigen::Ref<Eigen::Vector3d> vstore);
747748

749+
//@{
750+
//! Covariance structures. First index is T, second is the
751+
//! harmonic m
752+
std::vector<std::vector<Eigen::VectorXcd>> meanV;
753+
std::vector<std::vector<Eigen::MatrixXcd>> covrV;
754+
755+
void init_covariance();
756+
void zero_covariance();
757+
void writeCovarH5Params(HighFive::File& file);
758+
int sampT = 100;
759+
//@}
760+
748761
public:
749762

750763
//! Constructor from YAML node
@@ -798,6 +811,58 @@ namespace BasisClasses
798811
return ret;
799812
}
800813

814+
//! Write coefficient covariance data to an HDF5 file
815+
virtual void writeCoefCovariance(const std::string& compname,
816+
const std::string& runtag, double time=0.0)
817+
{
818+
if (covarStore) {
819+
// Prepare covariance data
820+
int Tsize = meanV.size();
821+
int Lsize = meanV[0].size();
822+
int Nsize = meanV[0][0].size();
823+
824+
SubsampleCovariance::CovarData covarData;
825+
826+
std::get<0>(covarData).resize(Tsize);
827+
std::get<1>(covarData).resize(Tsize);
828+
std::get<2>(covarData) = SubsampleCovariance::CoefType(Tsize, Lsize, Nsize);
829+
std::get<3>(covarData) = SubsampleCovariance::CovrType(Tsize, Lsize, Nsize, Nsize);
830+
831+
for (int T=0; T<Tsize; T++) {
832+
std::get<0>(covarData)(T) = sampleCounts(T);
833+
std::get<1>(covarData)(T) = sampleMasses(T);
834+
for (int lm=0; lm<Lsize; lm++) {
835+
for (int n1=0; n1<Nsize; n1++) {
836+
std::get<2>(covarData)(T, lm, n1) = meanV[T][lm](n1);
837+
for (int n2=0; n2<Nsize; n2++) {
838+
std::get<3>(covarData)(T, lm, n1, n2) = covrV[T][lm](n1, n2);
839+
}
840+
}
841+
}
842+
}
843+
covarStore->writeCoefCovariance(compname, runtag, covarData, time);
844+
} else {
845+
throw std::runtime_error("FlatDisk::writeCoefCovariance: covariance storage not initialized");
846+
}
847+
}
848+
849+
//! Request and initialize covariance computation
850+
virtual void enableCoefCovariance(bool pcavar_in, int sampT_in=100,
851+
bool ftype=false,
852+
bool covr_tot=true,
853+
bool covar_in=true)
854+
{
855+
pcavar = pcavar_in;
856+
sampT = sampT_in;
857+
scovr = covr_tot;
858+
covar = covar_in;
859+
if (pcavar) {
860+
init_covariance();
861+
covarStore = std::make_shared<SubsampleCovariance>
862+
([this](HighFive::File& file){this->writeCovarH5Params(file);}, BasisID, ftype, scovr, covar);
863+
}
864+
}
865+
801866
};
802867

803868
/**
@@ -1122,7 +1187,7 @@ namespace BasisClasses
11221187
std::tie(std::get<2>(elem), std::get<3>(elem)) = sl->getCoefCovariance();
11231188
covarStore->writeCoefCovariance(compname, runtag, elem, time);
11241189
} else {
1125-
throw std::runtime_error("Spherical::writeCoefCovariance: covariance storage not initialized");
1190+
throw std::runtime_error("Cylindrical::writeCoefCovariance: covariance storage not initialized");
11261191
}
11271192
}
11281193

@@ -1479,7 +1544,7 @@ namespace BasisClasses
14791544
}
14801545
covarStore->writeCoefCovariance(compname, runtag, elem, time);
14811546
} else {
1482-
throw std::runtime_error("Spherical::writeCoefCovariance: covariance storage not initialized");
1547+
throw std::runtime_error("Cube::writeCoefCovariance: covariance storage not initialized");
14831548
}
14841549
}
14851550

expui/BiorthBasis.cc

Lines changed: 81 additions & 6 deletions
Original file line numberDiff line numberDiff line change
@@ -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

Comments
 (0)