new idea 0.5.2

Zilong-Li · Feb 15, 2024 · df32aca · df32aca
1 parent 9395b6f
commit df32aca
Show file tree

Hide file tree

Showing 3 changed files with 20 additions and 5 deletions.
diff --git a/src/admixture.cpp b/src/admixture.cpp
@@ -11,20 +11,23 @@ using namespace std;
 // the complexity of this version should be O(CC + 2CK)
 double Admixture::runOptimalWithBigAss(int ind, const std::unique_ptr<BigAss> & genome)
 {
-    MyArr2D kapa, Ekg;
+    MyArr2D kapa, Ekg, AE;
     MyArr1D iQ = MyArr1D::Zero(K);
     MyArr1D Hz(C);
     double norm = 0, llike = 0, tmp = 0;
     int c1, k1, s, c2, c12;
-    for(int ic = 0, g = 0, ss = 0; ic < genome->nchunks; ic++)
+    for(int ic = 0, g = 0, g2 = 0, ss = 0; ic < genome->nchunks; ic++)
     {
         const int S = genome->pos[ic].size();
         const int nGrids = grids[ic];
         Eigen::Map<const MyArr2D> gli(genome->gls[ic].data() + ind * S * 3, S, 3);
         Eigen::Map<const MyArr2D> P(genome->P[ic].data(), S, C);
         Eigen::Map<const MyArr2D> PI(genome->PI[ic].data(), C, nGrids);
         Eigen::Map<const MyArr2D> R(genome->R[ic].data(), 3, nGrids);
-        Eigen::Map<const MyArr2D> AE(genome->AE[ic].data(), C * C, nGrids);
+        // Eigen::Map<const MyArr2D> AE(genome->AE[ic].data(), C * C, nGrids);
+        AE.setZero(C * C, nGrids);
+        for(int i = 0; i < nGrids; i++, g2++)
+            AE.col(i) = (CF.col(g2).matrix() * CF.col(g2).transpose().matrix()).reshaped().array();
         const auto cl = get_cluster_likelihoods(gli, P, R, PI, AE, collapse.segment(ss, S));
         ss += S;
         kapa.setZero(C * K, nGrids); // C x K x M layout
@@ -144,6 +147,7 @@ void Admixture::initIteration()
 {
     Ekc.setZero(C * K, G);
     NormF.setZero(K, G);
+    getCF();
 }
 
 void Admixture::updateIteration()
@@ -152,6 +156,15 @@ void Admixture::updateIteration()
     protectPars();
 }
 
+void Admixture::getCF()
+{
+    CF.setZero(C, G);
+    for(int g = 0; g < G; g++)
+        for(int c = 0; c < C; c++)
+            for(int k = 0; k < K; k++) CF(c, g) += Q.row(k).sum() * F(k * C + c, g);
+    CF.rowwise() /= CF.colwise().sum(); // normalize
+}
+
 void Admixture::protectPars()
 {
     if(!nonewQ)

diff --git a/src/admixture.hpp b/src/admixture.hpp
@@ -43,13 +43,15 @@ class Admixture
     MyArr2D Q; // K x N
     MyArr2D Ekc; // (C * K) x M, expected number of alleles per c per k
     MyArr2D NormF; // K x M
+    MyArr2D CF; // C x M, cluster frequency
     Bool1D collapse;
     Int1D grids;
 
     void initIteration();
     void updateIteration();
     void protectPars();
     void constrainF();
+    void getCF();
     void setFlags(double, double, double, bool, bool, bool);
     void setStartPoint(const std::unique_ptr<BigAss> & genome, std::string qfile);
     double runNativeWithBigAss(int ind, const std::unique_ptr<BigAss> & genome);

diff --git a/src/main.cpp b/src/main.cpp
@@ -18,7 +18,7 @@ int main(int argc, char * argv[])
 {
     // ========= helper message and parameters parsing ===========================
 
-    const std::string VERSION{"0.5.1"};
+    const std::string VERSION{"0.5.2"};
 
     // below for catching ctrl+c, and dumping files
     struct sigaction sa;
@@ -198,7 +198,7 @@ int main(int argc, char * argv[])
         .scan<'i', int>();
     cmd_admix.add_argument("-i", "--iterations")
         .help("number of maximun EM iterations")
-        .default_value(2000)
+        .default_value(1000)
         .scan<'i', int>();
     cmd_admix.add_argument("-n", "--threads")
         .help("number of threads")