jpayne@68: package prok;
jpayne@68: 
jpayne@68: import java.util.Arrays;
jpayne@68: 
jpayne@68: import dna.AminoAcid;
jpayne@68: import shared.KillSwitch;
jpayne@68: import shared.Parse;
jpayne@68: import shared.Tools;
jpayne@68: import structures.ByteBuilder;
jpayne@68: 
jpayne@68: /**
jpayne@68:  * Stores frame-relative kmer counts for a type of genomic feature, such as a coding start site.
jpayne@68:  * @author Brian Bushnell
jpayne@68:  * @date Sep 24, 2018
jpayne@68:  */
jpayne@68: public class FrameStats {
jpayne@68: 	
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 	/*----------------        Initialization        ----------------*/
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 	
jpayne@68: 	public FrameStats(String name_, int k_, int frames_, int leftOffset_){
jpayne@68: 		name=name_;
jpayne@68: 		k=k_;
jpayne@68: 		mask=~((-1)<<(2*k));
jpayne@68: 		frames=frames_;
jpayne@68: 		kMax=1<<(2*k);
jpayne@68: 		invFrames=1.0f/frames;
jpayne@68: 		leftOffset=leftOffset_;
jpayne@68: 		
jpayne@68: 		probs=new float[frames][kMax];
jpayne@68: 		countsTrue=new long[frames][kMax];
jpayne@68: 		countsFalse=new long[frames][kMax];
jpayne@68: 		counts=new long[][][] {countsFalse, countsTrue};
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 	/*----------------           Methods            ----------------*/
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 	
jpayne@68: 	public void add(int kmer, int frame, int valid){
jpayne@68: 		counts[valid][frame][kmer]++;
jpayne@68: 		validSums[valid]++;
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public boolean compatibleWith(FrameStats fs) {
jpayne@68: 		return fs.name.equals(name) && fs.leftOffset==leftOffset && fs.k==k && fs.frames==frames;
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public void clear() {
jpayne@68: 		Tools.fill(counts, 0);
jpayne@68: 		Arrays.fill(validSums, 0);
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public void setFrom(FrameStats fs) {
jpayne@68: 		assert(compatibleWith(fs)) : name+", "+frames+", "+fs.name+", "+fs.frames;
jpayne@68: 		clear();
jpayne@68: 		add(fs);
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public void add(FrameStats fs){
jpayne@68: 		assert(fs.name.equals(name));
jpayne@68: 		assert(fs.leftOffset==leftOffset);
jpayne@68: 		assert(fs.k==k);
jpayne@68: 		assert(fs.frames==frames) : name+", "+frames+", "+fs.name+", "+fs.frames;
jpayne@68: //		for(int x=0; x<counts.length; x++) {
jpayne@68: //			for(int y=0; y<counts[x].length; y++) {
jpayne@68: //				for(int z=0; z<counts[x][y].length; z++) {
jpayne@68: //					assert(fs.counts[x][y][z]>=0) : counts[x][y][z]+", "+fs.counts[x][y][z]+", "+fs.name;
jpayne@68: //					assert(counts[x][y][z]>=0) : counts[x][y][z]+", "+fs.counts[x][y][z];
jpayne@68: //				}
jpayne@68: //			}
jpayne@68: //		}
jpayne@68: 
jpayne@68: 		Tools.add(counts, fs.counts);
jpayne@68: 		Tools.add(validSums, fs.validSums);
jpayne@68: //		for(int x=0; x<counts.length; x++) {
jpayne@68: //			for(int y=0; y<counts[x].length; y++) {
jpayne@68: //				for(int z=0; z<counts[x][y].length; z++) {
jpayne@68: //					assert(fs.counts[x][y][z]>=0) : counts[x][y][z]+", "+fs.counts[x][y][z];
jpayne@68: //					assert(counts[x][y][z]>=0) : counts[x][y][z]+", "+fs.counts[x][y][z];
jpayne@68: //				}
jpayne@68: //			}
jpayne@68: //		}
jpayne@68: //		calculate();
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public void multiplyBy(double mult) {
jpayne@68: 		Tools.multiplyBy(counts, mult);
jpayne@68: 		Tools.multiplyBy(validSums, mult);
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	void calculate(){
jpayne@68: 		average=(float)((validSums[1]+1.0)/(validSums[0]+validSums[1]+1.0));
jpayne@68: 		invAvg=1.0f/average;
jpayne@68: 		
jpayne@68: 		for(int a=0; a<frames; a++){
jpayne@68: 			for(int b=0; b<kMax; b++){
jpayne@68: 				long t=countsTrue[a][b];
jpayne@68: 				long f=countsFalse[a][b];
jpayne@68: 				probs[a][b]=(float)(t/(t+f+1.0))*invAvg;
jpayne@68: 			}
jpayne@68: 		}
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public float scorePoint(int point, byte[] bases){
jpayne@68: 		final int mask=~((-1)<<(2*k));
jpayne@68: 		
jpayne@68: 		int kmer=0;
jpayne@68: 		int len=0;
jpayne@68: 		float score=0;
jpayne@68: 
jpayne@68: //		outstream.println("k="+k);
jpayne@68: //		outstream.println("mask="+mask);
jpayne@68: 		
jpayne@68: 		int start=point-leftOffset;
jpayne@68: 		for(int i=start, frame=0-k+1; i<bases.length && frame<frames; i++, frame++){
jpayne@68: 			byte b=(i>=0 ? bases[i] : (byte)'A');
jpayne@68: 			int x=AminoAcid.baseToNumber[b];
jpayne@68: 			kmer=((kmer<<2)|x)&mask;
jpayne@68: 			
jpayne@68: //			outstream.println("b="+(char)b+", kmer="+kmer+", len="+(len+1)+", frame="+frame);
jpayne@68: 			
jpayne@68: 			if(x>=0){
jpayne@68: 				len++;
jpayne@68: 				if(len>=k){
jpayne@68: 					float prob=probs[frame][kmer];
jpayne@68: 					float dif=prob-0.99f;
jpayne@68: 					score+=dif;
jpayne@68: 					
jpayne@68: //					if(name.equals("startStats")){
jpayne@68: //						System.err.println("frame="+frame+" kmer="+AminoAcid.kmerToString(kmer, k)+
jpayne@68: //								String.format(Locale.ROOT, " prob=%.4f\tdif=%.4f\tscore=%.4f", prob, dif, score)+
jpayne@68: //								"\tvalid="+counts[1][frame][kmer]+"\tinvalid="+counts[0][frame][kmer]);
jpayne@68: //					}
jpayne@68: 				}
jpayne@68: 			}else{len=0;}
jpayne@68: 		}
jpayne@68: 		
jpayne@68: 		return score*invFrames;
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	void processCDSFrames(byte[] bases, byte[] validFrames){
jpayne@68: 		if(!ProkObject.callCDS){return;}
jpayne@68: 		int kmer=0;
jpayne@68: 		int len=0;
jpayne@68: 		
jpayne@68: 		for(int i=0; i<bases.length; i++){
jpayne@68: 			byte b=bases[i];
jpayne@68: 			int x=AminoAcid.baseToNumber[b];
jpayne@68: 			kmer=((kmer<<2)|x)&mask;
jpayne@68: 			if(x>=0){
jpayne@68: 				len++;
jpayne@68: 				if(len>=k){
jpayne@68: 					int vf=validFrames[i];
jpayne@68: 					for(int frame=0; frame<frames; frame++){
jpayne@68: 						int valid=vf&1;
jpayne@68: 						add(kmer, frame, valid);
jpayne@68: 						//For CDS start (0-based) of 189, i=192, j=189, vf=1, frame=0 - all as expected.
jpayne@68: //						assert(valid==0) : "vf="+vf+", frame="+frame+", len="+len+", kmer="+AminoAcid.kmerToString(kmer, k)+", i="+i+", j="+j;
jpayne@68: 						vf=(vf>>1);
jpayne@68: 					}
jpayne@68: 				}
jpayne@68: 			}else{len=0;}
jpayne@68: 		}
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	void processPoint(byte[] bases, int point, int valid){
jpayne@68: 		
jpayne@68: 		//Degenerate cases where the point is at the end, possibly from a truncated gene.
jpayne@68: 		if(point<3){return;}
jpayne@68: 		if(point>=bases.length-3){return;}
jpayne@68: 		
jpayne@68: 		int kmer=0;
jpayne@68: 		int len=0;
jpayne@68: 
jpayne@68: //		outstream.println("k="+k);
jpayne@68: //		outstream.println("mask="+mask);
jpayne@68: 		
jpayne@68: 		int start=point-leftOffset;
jpayne@68: 		
jpayne@68: 		int i=start, frame=0-k+1;
jpayne@68: 		while(i<0){i++; frame++;}
jpayne@68: 		for(; i<bases.length && frame<frames; i++, frame++){
jpayne@68: 			byte b=bases[i];
jpayne@68: 			int x=AminoAcid.baseToNumber[b];
jpayne@68: 			kmer=((kmer<<2)|x)&mask;
jpayne@68: 			
jpayne@68: //			outstream.println("b="+(char)b+", kmer="+kmer+", len="+(len+1)+", frame="+frame);
jpayne@68: 			
jpayne@68: 			if(x>=0){
jpayne@68: 				len++;
jpayne@68: 				if(len>=k){
jpayne@68: 					add(kmer, frame, valid);
jpayne@68: 				}
jpayne@68: 			}else{len=0;}
jpayne@68: 		}
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 	/*----------------         Text Methods         ----------------*/
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 
jpayne@68: 	public void parseData(byte[] line) {
jpayne@68: 		int a=0, b=0;
jpayne@68: 		final int valid, frame;
jpayne@68: 		
jpayne@68: 		while(b<line.length && line[b]!='\t'){b++;}
jpayne@68: 		assert(b>a) : "Missing field 0: "+new String(line);
jpayne@68: 		valid=Parse.parseInt(line, a, b);
jpayne@68: 		b++;
jpayne@68: 		a=b;
jpayne@68: 		
jpayne@68: 		while(b<line.length && line[b]!='\t'){b++;}
jpayne@68: 		assert(b>a) : "Missing field 1: "+new String(line);
jpayne@68: 		frame=Parse.parseInt(line, a, b);
jpayne@68: 		b++;
jpayne@68: 		a=b;
jpayne@68: 		
jpayne@68: 		assert(valid==0 || valid==1);
jpayne@68: 		assert(frame>=0 && frame<frames);
jpayne@68: 		try {
jpayne@68: 			final long[] row=counts[valid][frame];
jpayne@68: 			long sum=0;
jpayne@68: 			for(int kmer=0; kmer<row.length; kmer++){
jpayne@68: 				while(b<line.length && line[b]!='\t'){b++;}
jpayne@68: 				assert(b>a) : "Missing field 1: "+new String(line);
jpayne@68: 				long count=Parse.parseLong(line, a, b);
jpayne@68: 				b++;
jpayne@68: 				a=b;
jpayne@68: 				row[kmer]=count;
jpayne@68: 				sum+=count;
jpayne@68: 			}
jpayne@68: 			validSums[valid]+=sum;
jpayne@68: 		} catch (Exception e) {
jpayne@68: 			System.err.println(new String(line)+"\n"+name);
jpayne@68: 			assert(false) : e;
jpayne@68: 		}
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	@Override
jpayne@68: 	public String toString(){
jpayne@68: 		return appendTo(new ByteBuilder()).toString();
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public ByteBuilder appendTo(ByteBuilder bb){
jpayne@68: 		bb.append("#name\t").append(name).nl();
jpayne@68: 		bb.append("#k\t").append(k).nl();
jpayne@68: 		bb.append("#frames\t").append(frames).nl();
jpayne@68: 		bb.append("#offset\t").append(leftOffset).nl();
jpayne@68: 		bb.append("#valid\tframe");
jpayne@68: 		for(int i=0; i<kMax; i++){bb.tab().append(AminoAcid.kmerToString(i, k));}
jpayne@68: 		bb.nl();
jpayne@68: 		for(int a=0; a<2; a++){//valid
jpayne@68: 			for(int b=0; b<frames; b++){
jpayne@68: 				bb.append(a);
jpayne@68: 				bb.tab().append(b);
jpayne@68: 				for(int c=0; c<kMax; c++){
jpayne@68: 					bb.tab().append(counts[a][b][c]);
jpayne@68: 				}
jpayne@68: 				bb.nl();
jpayne@68: 			}
jpayne@68: 		}
jpayne@68: 		return bb;
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	public ByteBuilder append0(ByteBuilder bb){
jpayne@68: 		bb.append("#name\t").append(name).nl();
jpayne@68: 		bb.append("#k\t").append(k).nl();
jpayne@68: 		bb.append("#frames\t").append(frames).nl();
jpayne@68: 		bb.append("#offset\t").append(leftOffset).nl();
jpayne@68: 		bb.append("#valid\tframe");
jpayne@68: 		for(int i=0; i<kMax; i++){bb.tab().append(AminoAcid.kmerToString(i, k));}
jpayne@68: 		bb.nl();
jpayne@68: 		for(int a=0; a<2; a++){//valid
jpayne@68: 			for(int b=0; b<frames; b++){
jpayne@68: 				bb.append(a);
jpayne@68: 				bb.tab().append(b);
jpayne@68: 				for(int c=0; c<kMax; c++){
jpayne@68: 					bb.tab().append(0);
jpayne@68: 				}
jpayne@68: 				bb.nl();
jpayne@68: 			}
jpayne@68: 		}
jpayne@68: 		return bb;
jpayne@68: 	}
jpayne@68: 	
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 	/*----------------            Fields            ----------------*/
jpayne@68: 	/*--------------------------------------------------------------*/
jpayne@68: 
jpayne@68: 	public final String name;
jpayne@68: 	public final int k;
jpayne@68: 	public final int mask;
jpayne@68: 	public final int frames;
jpayne@68: 	public final int kMax;
jpayne@68: 	public final float invFrames;
jpayne@68: 	public final int leftOffset;
jpayne@68: 	public int rightOffset() {return frames-leftOffset-1;}
jpayne@68: 	
jpayne@68: 	public final float[][] probs;
jpayne@68: 	public final long[][] countsTrue;
jpayne@68: 	public final long[][] countsFalse;
jpayne@68: 	public final long[][][] counts;
jpayne@68: 
jpayne@68: 	public final long[] validSums=KillSwitch.allocLong1D(2);
jpayne@68: 	private float average=-1;
jpayne@68: 	private float invAvg=-1;
jpayne@68: 	
jpayne@68: }