-
Notifications
You must be signed in to change notification settings - Fork 5
/
wrightFisher.go
131 lines (118 loc) · 3.33 KB
/
wrightFisher.go
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
package popgen
import (
"fmt"
"io"
"strconv"
"github.com/vertgenlab/gonomics/exception"
"github.com/vertgenlab/gonomics/fasta"
"github.com/vertgenlab/gonomics/fileio"
"github.com/vertgenlab/gonomics/vcf"
)
type WrightFisherSettings struct {
PopSize int
MutRate float64
NumGen int
GenomeSize int
RFitness float64
GcContent float64
InitFreq string
FitnessString string
SetSeed int64
Verbose bool
Fasta bool
Vcf bool
AncestralAllele string
}
type WrightFisherPopData struct {
Fasta []fasta.Fasta
Vcf vcf.Vcf
Meta []string
Freq [][][]float64
Ancestral []string //Ancestral states of each allele on each site (aka replicates)
Settings WrightFisherSettings
}
/*
WriteTSV writes a .tsv file based on WrightFisherPopData, including:
Comments starts with # that include metadata about the parameters of the simulation
Header line: Gen Site Freq.A Freq.C Freq.G Freq.T Ancestral
Frequencies table.
*/
func WriteTSV(outFile string, wf WrightFisherPopData) {
file := fileio.EasyCreate(outFile)
header := []string{
"Gen",
"Site",
"Freq.A",
"Freq.C",
"Freq.G",
"Freq.T",
"Ancestral",
}
writeMeta(file, wf.Meta) // Write metadata
writeEachLine(file, header) // Write the header line
writeToFileHandle(file, floatArrayToString(wf)) // Write the main frequencies table
exception.PanicOnErr(file.Close()) // Close the file
}
/*
floatArrayToString converts the all frequency array from wf to a 2D slice of string
2D array, zero-based, [generation and site][base].
*/
func floatArrayToString(wf WrightFisherPopData) [][]string {
// nrow = product of number of generation and number of site
nrow := (wf.Settings.NumGen + 1) * wf.Settings.GenomeSize
ncol := 7
answer := make([][]string, nrow)
var t, s, r, b int // iterator for generation, site, row, and base
for t = 0; t <= wf.Settings.NumGen; t++ {
for s = 0; s < wf.Settings.GenomeSize; s++ {
r = (t * wf.Settings.GenomeSize) + s // Keeps track of index based on generation and site
answer[r] = make([]string, ncol)
answer[r][0] = strconv.Itoa(t)
answer[r][1] = strconv.Itoa(s)
for b = 0; b < 4; b++ {
answer[r][2+b] = strconv.FormatFloat(wf.Freq[t][s][b], 'f', 5, 64)
}
answer[r][6] = wf.Ancestral[s]
}
}
return answer
}
/*
writeToFileHandle feeds one slice of string at a time into writeEachLine().
*/
func writeToFileHandle(file io.Writer, records [][]string) {
for _, rec := range records {
writeEachLine(file, rec)
}
}
/*
writeEachLine writes a tab-separated line based on elements in []string.
*/
func writeEachLine(file io.Writer, rec []string) {
var err error
for i := 0; i < len(rec); i++ {
if i == len(rec)-1 { // Ending of each line
_, err = fmt.Fprintf(file, "%s\n", rec[i])
exception.PanicOnErr(err)
} else {
_, err = fmt.Fprintf(file, "%s\t", rec[i])
exception.PanicOnErr(err)
}
}
}
/*
writeMeta writes metadata based on the []string of metadata
Separate each entry with ":" instead of a tab.
*/
func writeMeta(file io.Writer, rec []string) {
var err error
for i := 0; i < len(rec); i++ {
if i == len(rec)-1 {
_, err = fmt.Fprintf(file, "%s\n", rec[i])
exception.PanicOnErr(err)
} else {
_, err = fmt.Fprintf(file, "%s:", rec[i])
exception.PanicOnErr(err)
}
}
}