<commit_before>#include <petscksp.h>
#include "AFS.h"
#include "AFS_ctmc_petsc.h"
#include "im_clam.h"
#include "sparseExp.h"
#include "cs.h"
#include "AFS_ctmc.h"
#include "nlopt.h"
#include <gsl/gsl_linalg.h>

//Functions for jointly computing CTMC transitions and embedded chain

//this function fills two petsc matrices: 1) embedded Discrete MC trans Mat, 2) CTMC trans mat
// also fills the rates vector

void fillPetscTransMats(afsStateSpace *S, double *topol, int *moveType, int *nzCount,
	int *dim1, int *dim2, double theta1, double theta2, double mig1, double mig2, Mat *DTMC, Mat *CTMC, gsl_vector *rates){
	int i, j,k, newnz, abcount=0;
	int N = S->nstates,abFlag;
	double  c0r[N],c1r[N],m0r[N],m1r[N],totR[N],rowSums[N], tmp;;
	
	gsl_vector_set_zero(rates);
	//set rates vector
	for(i=0;i<N;i++){
		c0r[i] = S->states[i]->aCounts[0] * (S->states[i]->aCounts[0]-1)  / theta1 ;
		if(theta2 == 0)
			c1r[i] = 0.0;
		else
			c1r[i] = S->states[i]->aCounts[1] * (S->states[i]->aCounts[1]-1)  / theta2;
	
		m0r[i] = S->states[i]->aCounts[0] * mig1 ;
		m1r[i] = S->states[i]->aCounts[1] * mig2 ;
		totR[i] = c0r[i] + c1r[i] + m0r[i] + m1r[i];
	
		
		rowSums[i] = 0.0;
		if(S->states[i]->nalleles > 1  && c0r[i] >= 0 && c1r[i] >= 0 && totR[i] != 0.0)
			gsl_vector_set(rates,i,1.0/totR[i]);
		else
			gsl_vector_set(rates,i,0);
	}

	newnz=0;
	for(k=0;k<*nzCount;k++){
		i = dim1[k];
		j= dim2[k];
//		printf("i: %d, j:%d, movetype: %d\n",i,j,moveType[k]);
		if(S->states[i]->nalleles ==1 || S->states[j]->nalleles ==1 )abFlag=1;
	//	if(S->states[i]->nalleles ==1 )abFlag=1;
		else abFlag = 0;
		switch(moveType[k]){
			case 0: //coal pop0
			if(totR[i]==0)tmp=0;
			else tmp = c0r[i] / totR[i] * topol[k];

			MatSetValue(*CTMC,i,j,tmp*totR[i],INSERT_VALUES);	
			rowSums[i] += tmp* totR[i];
			if(abFlag==0)MatSetValue(*DTMC,i,j,tmp,INSERT_VALUES);				
			break;
			case 1: //coal pop1
			if(totR[i]==0)tmp=0;
			else tmp = c1r[i] / totR[i] * topol[k];
			
			MatSetValue(*CTMC,i,j,tmp*totR[i],INSERT_VALUES);	
			rowSums[i] += tmp* totR[i];
			if(abFlag==0)MatSetValue(*DTMC,i,j,tmp,INSERT_VALUES);			
			break;
			case 2: //mig pop0
			if(totR[i]==0)tmp=0;
			else tmp = m0r[i] / totR[i] * topol[k];
		
			MatSetValue(*CTMC,i,j,tmp*totR[i],INSERT_VALUES);	
			rowSums[i] += tmp* totR[i];
			if(abFlag==0)MatSetValue(*DTMC,i,j,tmp,INSERT_VALUES);			
			break;
			case 3: //mig pop1
			if(totR[i]==0)tmp=0;
			else tmp = m1r[i] / totR[i] * topol[k];			
		
			MatSetValue(*CTMC,i,j,tmp*totR[i],INSERT_VALUES);	
			rowSums[i] += tmp* totR[i];
			if(abFlag==0)MatSetValue(*DTMC,i,j,tmp,INSERT_VALUES);		
			break;
			case -1:
			abcount+=1;
			break;			
		}
	}
	newnz=*nzCount ;
	//now add diagonal elements; topol pre-set to -1
	for(k=0;k<N;k++){
		if((rowSums[k]) != 0.0){
			MatSetValue(*CTMC,k,k,0.0-rowSums[k],INSERT_VALUES);				
			newnz++;
		}
	}
	*nzCount = newnz;
	/* Assemble the matrix */
  	MatAssemblyBegin(*CTMC,MAT_FINAL_ASSEMBLY);
  	MatAssemblyEnd(*CTMC,MAT_FINAL_ASSEMBLY);
  	MatAssemblyBegin(*DTMC,MAT_FINAL_ASSEMBLY);
  	MatAssemblyEnd(*DTMC,MAT_FINAL_ASSEMBLY);
}

//fillPetscCsparseTransMats--this function fills one petsc matrix with CTMC transmat and returns pointer to CSparse matrix for DTMC trans mat
// also fills the rates vector
//note that the DTMC matrix needs to be cleared before entry to this function and compressed afterward.

struct cs_di_sparse * fillPetscCsparseTransMats(afsStateSpace *S, double *topol, int *moveType, int *nzCount,
	int *dim1, int *dim2, double theta1, double theta2, double mig1, double mig2, Mat *CTMC, gsl_vector *rates){
	int i, j,k, newnz, abcount=0;<commit_msg>full petsc<commit_after>#include <petscksp.h>
#include "AFS.h"
#include "AFS_ctmc_petsc.h"
#include "im_clam.h"
#include "sparseExp.h"
#include "cs.h"
#include "AFS_ctmc.h"
#include "nlopt.h"
#include <gsl/gsl_linalg.h>

//Functions for jointly computing CTMC transitions and embedded chain

//this function fills two petsc matrices: 1) embedded Discrete MC trans Mat, 2) CTMC trans mat
// also fills the rates vector

void fillPetscTransMats(afsStateSpace *S, double *topol, int *moveType, int *nzCount,
	int *dim1, int *dim2, double theta1, double theta2, double mig1, double mig2, Mat *DTMC, Mat *CTMC, gsl_vector *rates){
	int i, j,k, newnz, abcount=0;
	int N = S->nstates,abFlag;
	double  c0r[N],c1r[N],m0r[N],m1r[N],totR[N],rowSums[N], tmp;;
	
	gsl_vector_set_zero(rates);
	//set rates vector
	for(i=0;i<N;i++){
		c0r[i] = S->states[i]->aCounts[0] * (S->states[i]->aCounts[0]-1)  / theta1 ;
		if(theta2 == 0)
			c1r[i] = 0.0;
		else
			c1r[i] = S->states[i]->aCounts[1] * (S->states[i]->aCounts[1]-1)  / theta2;
	
		m0r[i] = S->states[i]->aCounts[0] * mig1 ;
		m1r[i] = S->states[i]->aCounts[1] * mig2 ;
		totR[i] = c0r[i] + c1r[i] + m0r[i] + m1r[i];
	
		
		rowSums[i] = 0.0;
		if(S->states[i]->nalleles > 1  && c0r[i] >= 0 && c1r[i] >= 0 && totR[i] != 0.0)
			gsl_vector_set(rates,i,1.0/totR[i]);
		else
			gsl_vector_set(rates,i,0);
	}

	newnz=0;
	for(k=0;k<*nzCount;k++){
		i = dim1[k];
		j= dim2[k];
//		printf("i: %d, j:%d, movetype: %d\n",i,j,moveType[k]);
		if(S->states[i]->nalleles ==1 || S->states[j]->nalleles ==1 )abFlag=1;
	//	if(S->states[i]->nalleles ==1 )abFlag=1;
		else abFlag = 0;
		switch(moveType[k]){
			case 0: //coal pop0
			if(totR[i]==0)tmp=0;
			else tmp = c0r[i] / totR[i] * topol[k];

			MatSetValue(*CTMC,i,j,tmp*totR[i],INSERT_VALUES);	
			rowSums[i] += tmp* totR[i];
			if(abFlag==0)MatSetValue(*DTMC,i,j,tmp,INSERT_VALUES);				
			break;
			case 1: //coal pop1
			if(totR[i]==0)tmp=0;
			else tmp = c1r[i] / totR[i] * topol[k];
			
			MatSetValue(*CTMC,i,j,tmp*totR[i],INSERT_VALUES);	
			rowSums[i] += tmp* totR[i];
			if(abFlag==0)MatSetValue(*DTMC,i,j,tmp,INSERT_VALUES);			
			break;
			case 2: //mig pop0
			if(totR[i]==0)tmp=0;
			else tmp = m0r[i] / totR[i] * topol[k];
		
			MatSetValue(*CTMC,i,j,tmp*totR[i],INSERT_VALUES);	
			rowSums[i] += tmp* totR[i];
			if(abFlag==0)MatSetValue(*DTMC,i,j,tmp,INSERT_VALUES);			
			break;
			case 3: //mig pop1
			if(totR[i]==0)tmp=0;
			else tmp = m1r[i] / totR[i] * topol[k];			
		
			MatSetValue(*CTMC,i,j,tmp*totR[i],INSERT_VALUES);	
			rowSums[i] += tmp* totR[i];
			if(abFlag==0)MatSetValue(*DTMC,i,j,tmp,INSERT_VALUES);		
			break;
			case -1:
			abcount+=1;
			break;			
		}
	}
	newnz=*nzCount ;
	//now add diagonal elements; topol pre-set to -1
	for(k=0;k<N;k++){
		if((rowSums[k]) != 0.0){
			MatSetValue(*CTMC,k,k,0.0-rowSums[k],INSERT_VALUES);				
			newnz++;
		}
		MatSetValue(*DTMC,k,k,0.0,INSERT_VALUES);
	}
	*nzCount = newnz;
	/* Assemble the matrix */
  	MatAssemblyBegin(*CTMC,MAT_FINAL_ASSEMBLY);
  	MatAssemblyEnd(*CTMC,MAT_FINAL_ASSEMBLY);
  	MatAssemblyBegin(*DTMC,MAT_FINAL_ASSEMBLY);
  	MatAssemblyEnd(*DTMC,MAT_FINAL_ASSEMBLY);
}

//fillPetscCsparseTransMats--this function fills one petsc matrix with CTMC transmat and returns pointer to CSparse matrix for DTMC trans mat
// also fills the rates vector
//note that the DTMC matrix needs to be cleared before entry to this function and compressed afterward.

struct cs_di_sparse * fillPetscCsparseTransMats(afsStateSpace *S, double *topol, int *moveType, int *nzCount,
	int *dim1, int *dim2, double theta1, double theta2, double mig1, double mig2, Mat *CTMC, gsl_vector *rates){
	int i, j,k, newnz, abcount=0;
{'fraction_non_alphanumeric': 0.18783394985614468, 'fraction_numerical': 0.027675023975887107, 'mean_word_length': 4.236728837876614, 'pattern_counts': {'":': 0, '<': 13, '<?xml version=': 0, '>': 55, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '6462496', 'n_tokens_mistral': 3672, 'n_tokens_neox': 3337, 'n_words': 629}