From 1eb19eb0e31d484883434895bef65afc71954b0e Mon Sep 17 00:00:00 2001
From: Dariusz Murakowski <murakdar@mit.edu>
Date: Fri, 17 May 2013 10:27:10 -0400
Subject: [PATCH] Special two-site + two-allele (two-locus + binary)
 Hamiltonian with all parameters passed as command line arguments.

---
 hamiltonian.cpp | 113 ++++++++++++++++++++++++++++++++++++
 hamiltonian.h   |   9 +++
 population.cpp  |  42 ++++++++++++++
 population.h    |   1 +
 wf.cpp          | 173 +++++++++++++++++++++++++++++++++++++++-----------------
 wf.h            |  10 ++++
 6 files changed, 295 insertions(+), 53 deletions(-)

diff --git a/hamiltonian.cpp b/hamiltonian.cpp
index 842504a..593d9b0 100644
--- a/hamiltonian.cpp
+++ b/hamiltonian.cpp
@@ -292,6 +292,119 @@ double EpitopeHamiltonian::get_energy(const std::set<unsigned int> &mutated_site
 }
 
 
+/***********************************************
+ * my simple 2-site 2-allele system
+ ***********************************************/
 
+// Return the energy for a given set of mutated sites
+
+TwoSiteHamiltonian::TwoSiteHamiltonian(double h1, double h2, double J12) {
+
+    this->bh=1.0;
+    this->bJ=1.0;
+
+    this->size = 2;
+
+    J.resize(size);
+
+    for(int i=0; i<size; ++i) {
+        J[i].resize(size);
+    }
+
+    J[0][0] = h1;
+    J[1][1] = h2;
+    J[0][1] = J12;
+    J[1][0] = J[0][1];
+
+}
+
+TwoSiteHamiltonian::TwoSiteHamiltonian(std::string &FILENAME) {
+	
+    this->bh=1.0;
+    this->bJ=1.0;
+    
+	FILE* input;
+	input = fopen(FILENAME.c_str(),"r");
+	
+	typedef std::vector<int> Key;
+	std::map<Key,double> JIndex;
+    char o;
+	int site=0;
+    
+	fscanf(input,"%d",&size);
+	
+	while (fscanf(input,"%d",&site)==1) {
+        
+        Key inputKey(2,site);        
+		while (fscanf(input,"%c",&o)==1) {
+			
+			if (o==';') break;
+			
+			else {
+				
+				int neighborSite=0;
+				double neighborJ=0;
+				fscanf(input,"%d %c %lf",&neighborSite,&o,&neighborJ);
+				
+				if (neighborSite==site) {
+                    
+					inputKey[1]=site;
+                    JIndex[inputKey]=neighborJ;
+					
+				}
+				else {
+                
+                        inputKey[1]=neighborSite;
+                        JIndex[inputKey]=neighborJ;
+                        
+                }
+                //std::cout << site << " " << neighborSite << " " << neighborJ << "\n";
+                
+			}
+            
+		}
+        
+	}
+		
+	fclose(input);
+	Key inputKey(2,0);
+	J.resize(size);
+    
+	for(int i=0;i<size;i++) J[i].resize(size,0.0);
+    Graph.resize(size);
+    
+	for (int i=0;i<size;i++) {
+		
+        J[i].resize(size,0.0);
+        
+		for(int j=i;j<size;j++){
+        
+			inputKey[0]=i;
+			inputKey[1]=j;
+            
+			if (JIndex[inputKey]!=0.0) {
+            
+				J[i][j]+=JIndex[inputKey];
+				J[j][i]=J[i][j];
+                
+				if (J[i][j]!=0.0) {
+					
+                    Graph[i].push_back(j);	//Populate the adjacency list
+					if(j!=i) Graph[j].push_back(i);
+                    
+				}
+                
+			}
+            
+		}
+        
+	}
+
+    //for(int i=0; i<size; i++)
+        //for(int j=0; j<size; j++)
+            //std::cout << i << " " << j << " " << J[i][j] << "\n";
+            //printf("%d %d %lf\n",i,j,J[i][j]);
+    
+}
 
 
diff --git a/hamiltonian.h b/hamiltonian.h
index 90e2978..197f8ba 100644
--- a/hamiltonian.h
+++ b/hamiltonian.h
@@ -58,4 +58,13 @@ public:
 };
 
 
+class TwoSiteHamiltonian : public Hamiltonian {
+
+public:
+    
+    TwoSiteHamiltonian(std::string &FILENAME);
+    TwoSiteHamiltonian(double h1, double h2, double J12);
+    
+};
+
 #endif
diff --git a/population.cpp b/population.cpp
index 3c3cdc2..fed7284 100644
--- a/population.cpp
+++ b/population.cpp
@@ -41,6 +41,15 @@ Population::Population(const Hamiltonian &H, unsigned int N, double mu, const st
     
 	p     = 1.0-pow((1.0-mu),H.size);
     Eavg /= (double) N;
+
+
+    // if not given any particular thing
+    // then initialize to N wild type
+    if (initPop.size() == 0) {
+        Virus V(H, mu);
+        pop[V] = N;
+        Eavg = V.energy;
+    }
 	
 }
 
@@ -278,4 +287,37 @@ unsigned int Population::escape_variant(Hamiltonian &H, std::set<unsigned int> &
 
 
 
+// Output population to file
+// two-site two-allele
+
+void Population::write_two_site_population(FILE *output, const Hamiltonian &H, unsigned int generation) {
+
+	fprintf(output,"%d",generation);
+
+    //static const unsigned int v00_arr[] = {};
+    static const unsigned int v10_arr[] = {0};
+    static const unsigned int v01_arr[] = {1};
+    static const unsigned int v11_arr[] = {0,1};
+
+    //std::set<unsigned int> v00_set (v00_arr, v00_arr+sizeof(v00_arr)/sizeof(v00_arr[0]));
+    std::set<unsigned int> v10_set (v10_arr, v10_arr+sizeof(v10_arr)/sizeof(v10_arr[0]));
+    std::set<unsigned int> v01_set (v01_arr, v01_arr+sizeof(v01_arr)/sizeof(v01_arr[0]));
+    std::set<unsigned int> v11_set (v11_arr, v11_arr+sizeof(v11_arr)/sizeof(v11_arr[0]));
+
+    Virus v00_vir (H, mu);//, v00_set);
+    Virus v10_vir (H, mu, v10_set);
+    Virus v01_vir (H, mu, v01_set);
+    Virus v11_vir (H, mu, v11_set);
+
+    fprintf(output,"\t%d",pop[v00_vir]);
+    fprintf(output,"\t%d",pop[v10_vir]);
+    fprintf(output,"\t%d",pop[v01_vir]);
+    fprintf(output,"\t%d",pop[v11_vir]);
+    fprintf(output,"\n");
+
+	fflush(output);
+
+}
+
+
 
diff --git a/population.h b/population.h
index 7e7edd7..9374b9a 100644
--- a/population.h
+++ b/population.h
@@ -28,6 +28,7 @@ public:
     void next_generation(const Hamiltonian &H, gsl_rng* r, bool useRelative, bool useVerbose);
     void write_population(std::string filepath);
     void write_population(FILE *output, unsigned int generation);
+    void write_two_site_population(FILE *output, const Hamiltonian &H, unsigned int generation);
     unsigned int compute_population_size();
     void print_population_size();
     void mutate_population(const Hamiltonian &H, gsl_rng* r);
diff --git a/wf.cpp b/wf.cpp
index c89e797..01e4362 100644
--- a/wf.cpp
+++ b/wf.cpp
@@ -38,8 +38,6 @@ static unsigned sim_random_seed() {
 
 void run(RunParameters &r) {
     
-    printf("\n");
-    
     // Initialize RNG and set initial state, if importing from file
     
     static gsl_rng *rnd = gsl_rng_alloc(gsl_rng_taus2);	//pointer to random number generator state
@@ -50,74 +48,107 @@ void run(RunParameters &r) {
     
     r.setFiles();
     FILE *popout=fopen(r.trajectoryOutfile.c_str(),"w");        // <outfile>.dat
-    FILE *supout=fopen(r.supplementaryOutfile.c_str(),"w");     // <outfile>.sum
-    
-    if (r.importState) importState(r);
-    
-    
-    // Run (w/ targeted epitope)
+    FILE *supout=r.useTwoSite ? NULL : fopen(r.supplementaryOutfile.c_str(),"w");     // <outfile>.sum
     
-    if (r.useEpitope) {
+    if (r.importState) importState(r);      // <infile>.st
     
-        for (unsigned int n=0;n<r.num_runs;n++) {
-        
-            EpitopeHamiltonian H(r.couplingsInfile);
-            H.set_epitope(r.eWT, r.eMut, r.penalty);
-            H.set_temp(r.bh, r.bJ);
-            Population P(H, r.n, r.mu, r.initPop, r.initFrac);
+
+    if(r.useTwoSite) {
+
+            fprintf(popout,"gen\tV00\tV10\tV01\tV11\n");
+
+            for (unsigned int n=0;n<r.num_runs;n++) {
+
+                //fprintf(popout,"gen\tV00\tV10\tV01\tV11\n");
             
-            unsigned int i;
-            for (i=0; i<r.g; ++i) {
+                TwoSiteHamiltonian H(r.h1,r.h2,r.J12);
+                H.set_temp(r.bh, r.bJ);
+                Population P(H, r.n, r.mu, r.initPop, r.initFrac);
+                
+                P.write_two_site_population(popout,H,0);
+
+                unsigned int i;
+                for (i=0; i<r.g; ++i) {
+                
+                    P.next_generation(H, rnd, r.useRelative, r.useVerbose);
+                    P.write_two_site_population(popout,H,i+1);
+
+                    //if (r.runUntilEscape && P.escaped(H)) break;
+
+                }
+                
+            }
+
+    }
+
+
+    else {
         
-                P.next_generation(H, rnd, r.useRelative, r.useVerbose);
-                P.write_population(popout,i);
-            
-                if (r.runUntilEscape && P.escaped(H)) break;
+        // Run (w/ targeted epitope)
+        
+        if (r.useEpitope) {
+        
+            for (unsigned int n=0;n<r.num_runs;n++) {
             
-            }
+                EpitopeHamiltonian H(r.couplingsInfile);
+                H.set_epitope(r.eWT, r.eMut, r.penalty);
+                H.set_temp(r.bh, r.bJ);
+                Population P(H, r.n, r.mu, r.initPop, r.initFrac);
+                
+                unsigned int i;
+                for (i=0; i<r.g; ++i) {
             
-            std::set<unsigned int> esc_var;
-            unsigned int esc_num=P.escape_variant(H, esc_var);
+                    P.next_generation(H, rnd, r.useRelative, r.useVerbose);
+                    P.write_population(popout,i);
+                
+                    if (r.runUntilEscape && P.escaped(H)) break;
+                
+                }
+                
+                std::set<unsigned int> esc_var;
+                unsigned int esc_num=P.escape_variant(H, esc_var);
+                
+                fprintf(supout,"%d\t%d",i,esc_num);
+                for (std::set<unsigned int>::iterator iter=esc_var.begin(); iter!=esc_var.end(); ++iter) fprintf(supout,"\t%d",*iter);
+                fprintf(supout,"\n");
+                
+                fflush(supout);
             
-            fprintf(supout,"%d\t%d",i,esc_num);
-            for (std::set<unsigned int>::iterator iter=esc_var.begin(); iter!=esc_var.end(); ++iter) fprintf(supout,"\t%d",*iter);
-            fprintf(supout,"\n");
+            }
             
-            fflush(supout);
-        
         }
         
-    }
-    
-    
-    // Run (w/out targeted epitope)
-    
-    else {
-
-        //for (unsigned int n=0;n<r.num_runs;n++) {
         
-            Hamiltonian H(r.couplingsInfile);
-            H.set_temp(r.bh, r.bJ);
-            Population P(H, r.n, r.mu);
-            
-            unsigned int i;
-            for (i=0; i<r.g; ++i) {
+        // Run (w/out targeted epitope)
+        
+        else {
+
+            //for (unsigned int n=0;n<r.num_runs;n++) {
             
-                P.next_generation(H, rnd, r.useRelative, r.useVerbose);
-                P.write_population(popout,i);
+                Hamiltonian H(r.couplingsInfile);
+                H.set_temp(r.bh, r.bJ);
+                Population P(H, r.n, r.mu);
+                
+                unsigned int i;
+                for (i=0; i<r.g; ++i) {
+                
+                    P.next_generation(H, rnd, r.useRelative, r.useVerbose);
+                    P.write_population(popout,i);
 
-                if (r.runUntilEscape && P.escaped(H)) break;
+                    if (r.runUntilEscape && P.escaped(H)) break;
 
-            }
+                }
+                
+            //}
             
-        //}
-        
+        }
+
     }
     
     gsl_rng_free(rnd);  //Free up memory from random number generator
 
     fclose(popout);     // close file handles
-    fclose(supout);
+    if(!r.useTwoSite) fclose(supout);
 
 }
 
@@ -183,6 +214,21 @@ void importState(RunParameters &r) {
 
 }
 
+
+// insert into initial state of population
+
+void add_to_two_site_pop(RunParameters &r, bool s1, bool s2, double frac) {
+
+    r.initFrac.push_back(frac);
+
+    std::set<unsigned int> tempSet;
+    if (s1) tempSet.insert(0);
+    if (s2) tempSet.insert(1);
+    r.initPop.push_back(tempSet);
+
+}
+
+
 void usage()
 {
     std::cout <<
@@ -199,8 +245,16 @@ void usage()
 " -e  (int) (int)     targeted epitope spans sites [e#1, ..., e#2], assuming all WT\n"
 " -r                  flag to use relative energy condition for survival probability\n"
 " -esc                flag to run until escape observed (or max number of generations reached)\n"
-" -v                  flag for verbose output\n";
-    std::cout.flush();
+" -v                  flag for verbose output\n"
+"\n"
+" -2site              flag for two-site two-allele system\n"
+" -h1  (double)       value of field at site 1\n"
+" -h2  (double)       value of field at site 2\n"
+" -J12 (double)       value of coupling between sites 1 and 2\n"
+"                     (note the sign convention on these is P ~ exp(-H),\n"
+"                      H = sum_i h_i s_i + sum_{i<j} J_{ij} s_i s_j\n"
+"                      although internal representation is opposite)\n"
+;   std::cout.flush();
 }
 
 /*
@@ -218,7 +272,7 @@ void usage()
  *  -r                  flag to use relative energy condition for survival probability
  *  -esc                flag to run until escape observed (or max number of generations reached)
  *  -v                  flag for verbose output
- *
+ *  -2site              flag for two-site two-allele system
  */
 
 int main(int argc, char *argv[]) {
@@ -248,6 +302,19 @@ int main(int argc, char *argv[]) {
         else if (strcmp(argv[i],"-r")==0)   { r.useRelative=true;    }
         else if (strcmp(argv[i],"-esc")==0) { r.runUntilEscape=true; }
         else if (strcmp(argv[i],"-v")==0)   { r.useVerbose=true;     }
+        
+        else if (strcmp(argv[i],"-numruns")==0)  { if (++i==argc) break; else r.num_runs=(unsigned int)strtodouble(argv[i]); }
+
+        else if (strcmp(argv[i],"-2site")==0)   { r.useTwoSite=true;     }
+        else if (strcmp(argv[i],"-h1")==0)  { if (++i==argc) break; else r.h1=-strtodouble(argv[i]);  }
+        else if (strcmp(argv[i],"-h2")==0)  { if (++i==argc) break; else r.h2=-strtodouble(argv[i]);  }
+        else if (strcmp(argv[i],"-J12")==0) { if (++i==argc) break; else r.J12=-strtodouble(argv[i]); }
+
+        else if (strcmp(argv[i],"-p00")==0) { if (++i==argc) break; else add_to_two_site_pop(r,false,false,strtodouble(argv[i])); }
+        else if (strcmp(argv[i],"-p10")==0) { if (++i==argc) break; else add_to_two_site_pop(r,true ,false,strtodouble(argv[i])); }
+        else if (strcmp(argv[i],"-p01")==0) { if (++i==argc) break; else add_to_two_site_pop(r,false,true ,strtodouble(argv[i])); }
+        else if (strcmp(argv[i],"-p11")==0) { if (++i==argc) break; else add_to_two_site_pop(r,true ,true ,strtodouble(argv[i])); }
+
 
         else if (strcmp(argv[i],"-h")==0)   { usage(); return 0;    }
         
diff --git a/wf.h b/wf.h
index 7250345..ddb376e 100644
--- a/wf.h
+++ b/wf.h
@@ -93,6 +93,11 @@ public:
     bool useEpitope;            // Include selective pressure on an epitope 
     bool useRelative;           // If true, use relative energy for survival probability
     bool useVerbose;            // If true, print extra information while program is running
+
+    bool useTwoSite;            // If true, use two-site two-allele model
+    double h1;
+    double h2;
+    double J12;
     
     std::vector<std::set<unsigned int> > initPop;   // Initial population sequences
     std::vector<double>                  initFrac;  // Initial population fraction
@@ -127,6 +132,11 @@ public:
         useEpitope=false;
         useRelative=false;
         useVerbose=false;
+
+        useTwoSite = false;
+        h1  = 0.0;
+        h2  = 0.0;
+        J12 = 0.0;
         
     }
     void setFiles() {
-- 
2.7.4