forked from grailbio/bio
-
Notifications
You must be signed in to change notification settings - Fork 0
/
basestrand.go
422 lines (388 loc) · 13.3 KB
/
basestrand.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
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
// Copyright 2020 Grail Inc.
//
// Licensed under the Apache License, Version 2.0 (the "License");
// you may not use this file except in compliance with the License.
// You may obtain a copy of the License at
//
// http://www.apache.org/licenses/LICENSE-2.0
//
// Unless required by applicable law or agreed to in writing, software
// distributed under the License is distributed on an "AS IS" BASIS,
// WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
// See the License for the specific language governing permissions and
// limitations under the License.
package snp
import (
"bytes"
"context"
"encoding/binary"
"fmt"
"io"
"strconv"
"strings"
"sync"
"github.com/Schaudge/grailbase/file"
"github.com/Schaudge/grailbase/log"
"github.com/Schaudge/grailbase/recordio"
"github.com/Schaudge/grailbase/recordio/recordiozstd"
"github.com/Schaudge/grailbase/tsv"
"github.com/Schaudge/grailbio/pileup"
)
const (
refNamesHeader = "RefNames"
trailerVersion = 1
)
// BaseStrandPile represents a single pileup entry with a count for every
// (base, strand) tuple.
//
// - Pos is zero-based; it is necessary to add 1 when converting to most text
// formats (but not BED).
// - In Counts[][], base is the major dimension, with pileup.BaseA=0, C=1, G=2,
// T=3. Strand is the minor dimension, with strandFwd=0 and strandRev=1.
// TODO(cchang): strandFwd=0, strandRev=1 is inconsistent with bio/pileup's
// internal representation (which has None=0). We have enough other code at
// this point with Fwd=0, Rev=1 that it's probably time to change
// bio/pileup's representation to match that.
type BaseStrandPile struct {
RefID uint32
Pos uint32
Counts [pileup.NBase][2]uint32
}
// Minimal implementation; can optimize later if we're using this on large
// pileups.
// WriteBaseStrandsRio writes the given BaseStrand-pileup entries to the given
// writer, using recordio.
func WriteBaseStrandsRio(piles []BaseStrandPile, refNames []string, out io.Writer) error {
// recordiozstd.Init() is called in singleton.go's init().
recordWriter := recordio.NewWriter(out, recordio.WriterOpts{
Marshal: marshalBaseStrand,
Transformers: []string{recordiozstd.Name},
})
// could error out if refNames is empty
recordWriter.AddHeader(refNamesHeader, strings.Join(refNames, "\000"))
recordWriter.AddHeader(recordio.KeyTrailer, true)
for i := range piles {
recordWriter.Append(&piles[i])
}
recordWriter.SetTrailer(baseStrandsRioTrailer(len(piles)))
return recordWriter.Finish()
}
func baseStrandsRioTrailer(numPiles int) []byte {
var buffer bytes.Buffer
if err := binary.Write(&buffer, binary.LittleEndian, int64(trailerVersion)); err != nil {
panic("couldn't write trailer version")
}
if err := binary.Write(&buffer, binary.LittleEndian, int64(numPiles)); err != nil {
panic("couldn't write numPiles to trailer")
}
return buffer.Bytes()
}
func parseBaseStrandsTrailer(trailer []byte) (int64, error) {
r := bytes.NewReader(trailer)
var version, numPiles int64
if err := binary.Read(r, binary.LittleEndian, &version); err != nil {
return 0, err
}
if version != trailerVersion {
return 0, fmt.Errorf("unrecognized trailer version: got %d, want %d", version, trailerVersion)
}
if err := binary.Read(r, binary.LittleEndian, &numPiles); err != nil {
return 0, err
}
return numPiles, nil
}
func marshalBaseStrand(scratch []byte, p interface{}) ([]byte, error) {
t := scratch
if len(t) < 40 {
t = make([]byte, 40)
}
// As of Go 1.12, the bounds-check-eliminator is not smart enough to
// determine that len(t) is always 40 if we move t = scratch[:40] in an
// else-branch above.
t = t[:40]
pile := p.(*BaseStrandPile)
binary.LittleEndian.PutUint32(t[:4], pile.RefID)
binary.LittleEndian.PutUint32(t[4:8], pile.Pos)
// I tried a few different ways to write a double-loop here, but the
// bounds-check-eliminator did not handle any of them well (on every
// inner-loop iteration, there were three separate t[] bounds-checks:
// offset < 0? offset > offset+4? offset+4 > len(t)?)
//
// The code below is handled decently.
//
// The trailing 0/1 indexes will be replaced with enums when bio/pileup as a
// whole has been switched over to strandFwd=0, strandRev=1, strandNone=2.
binary.LittleEndian.PutUint32(t[8:12], pile.Counts[pileup.BaseA][0])
binary.LittleEndian.PutUint32(t[12:16], pile.Counts[pileup.BaseA][1])
binary.LittleEndian.PutUint32(t[16:20], pile.Counts[pileup.BaseC][0])
binary.LittleEndian.PutUint32(t[20:24], pile.Counts[pileup.BaseC][1])
binary.LittleEndian.PutUint32(t[24:28], pile.Counts[pileup.BaseG][0])
binary.LittleEndian.PutUint32(t[28:32], pile.Counts[pileup.BaseG][1])
binary.LittleEndian.PutUint32(t[32:36], pile.Counts[pileup.BaseT][0])
binary.LittleEndian.PutUint32(t[36:40], pile.Counts[pileup.BaseT][1])
return t, nil
}
// ReadBaseStrandsRio reads BaseStrand piles from a recordio file written by
// WriteBaseStrandsRio.
func ReadBaseStrandsRio(rs io.ReadSeeker) (piles []BaseStrandPile, refNames []string, err error) {
var unmarshaller BaseStrandUnmarshaller
scanner := recordio.NewScanner(rs, recordio.ScannerOpts{
Unmarshal: unmarshaller.UnmarshalBaseStrand,
})
if len(scanner.Trailer()) != 0 {
var numPiles int64
if numPiles, err = parseBaseStrandsTrailer(scanner.Trailer()); err != nil {
return
}
unmarshaller.init(numPiles)
}
hdr := scanner.Header()
for _, kv := range hdr {
switch kv.Key {
case refNamesHeader:
packedRefNames := kv.Value.(string)
refNames = strings.Split(packedRefNames, "\000")
// Cannot return an error on unrecognized key since recordio can write its own.
}
}
for scanner.Scan() {
piles = append(piles, *scanner.Get().(*BaseStrandPile))
}
err = scanner.Err()
return
}
// BaseStrandUnmarshaller is used to allocate memory in large blocks during
// unmarshalling, to prevent contention with other goroutines.
type BaseStrandUnmarshaller struct {
piles []BaseStrandPile
offset int
}
func (b *BaseStrandUnmarshaller) init(size int64) {
if b.piles != nil {
panic("tried to initialize when already initialized")
}
b.piles = make([]BaseStrandPile, size)
}
func (b *BaseStrandUnmarshaller) UnmarshalBaseStrand(in []byte) (out interface{}, err error) {
in = in[:40] // help the bounds-checker
if b.offset == len(b.piles) {
b.piles = append(b.piles, BaseStrandPile{})
}
pile := &b.piles[b.offset]
b.offset++
pile.RefID = binary.LittleEndian.Uint32(in[:4])
pile.Pos = binary.LittleEndian.Uint32(in[4:8])
pile.Counts[pileup.BaseA][0] = binary.LittleEndian.Uint32(in[8:12])
pile.Counts[pileup.BaseA][1] = binary.LittleEndian.Uint32(in[12:16])
pile.Counts[pileup.BaseC][0] = binary.LittleEndian.Uint32(in[16:20])
pile.Counts[pileup.BaseC][1] = binary.LittleEndian.Uint32(in[20:24])
pile.Counts[pileup.BaseG][0] = binary.LittleEndian.Uint32(in[24:28])
pile.Counts[pileup.BaseG][1] = binary.LittleEndian.Uint32(in[28:32])
pile.Counts[pileup.BaseT][0] = binary.LittleEndian.Uint32(in[32:36])
pile.Counts[pileup.BaseT][1] = binary.LittleEndian.Uint32(in[36:40])
return pile, nil
}
// WriteBaseStrandToTSV writes a []BaseStrandPile as a TSV.
func WriteBaseStrandToTSV(piles []BaseStrandPile, refNames []string, w io.Writer) (err error) {
// Note that this is slightly different from the .basestrand.tsv format:
// there is no REF base column, or '#' marking the header line.
outTSV := tsv.NewWriter(w)
outTSV.WriteString("CHROM\tPOS\tA+\tA-\tC+\tC-\tG+\tG-\tT+\tT-")
if err = outTSV.EndLine(); err != nil {
return
}
for _, pile := range piles {
outTSV.WriteString(refNames[pile.RefID])
outTSV.WriteUint32(pile.Pos)
outTSV.WriteUint32(pile.Counts[pileup.BaseA][0])
outTSV.WriteUint32(pile.Counts[pileup.BaseA][1])
outTSV.WriteUint32(pile.Counts[pileup.BaseC][0])
outTSV.WriteUint32(pile.Counts[pileup.BaseC][1])
outTSV.WriteUint32(pile.Counts[pileup.BaseG][0])
outTSV.WriteUint32(pile.Counts[pileup.BaseG][1])
outTSV.WriteUint32(pile.Counts[pileup.BaseT][0])
outTSV.WriteUint32(pile.Counts[pileup.BaseT][1])
if err = outTSV.EndLine(); err != nil {
return
}
}
if err = outTSV.Flush(); err != nil {
return
}
return
}
// WriteBaseStrandsRioAsTSV converts the given recordio pileup to TSV.
func WriteBaseStrandsRioAsTSV(ctx context.Context, path string, w io.Writer) error {
file, err := file.Open(ctx, path)
if err != nil {
return err
}
piles, refNames, err2 := ReadBaseStrandsRio(file.Reader(ctx))
if err2 != nil {
return err2
}
return WriteBaseStrandToTSV(piles, refNames, w)
}
// BaseStrandTsvRow represents a single row of a basestrand.tsv file.
type BaseStrandTsvRow struct {
Chr string `tsv:"#CHROM"` // Chromosome
Pos int64 `tsv:"POS"` // Position in chromosome
Ref string `tsv:"REF"` // Reference base
FwdA int64 `tsv:"A+"` // A count on the forward strand
RevA int64 `tsv:"A-"` // A count on the reverse strand
FwdC int64 `tsv:"C+"` // C count on the forward strand
RevC int64 `tsv:"C-"` // C count on the reverse strand
FwdG int64 `tsv:"G+"` // G count on the forward strand
RevG int64 `tsv:"G-"` // G count on the reverse strand
FwdT int64 `tsv:"T+"` // T count on the forward strand
RevT int64 `tsv:"T-"` // T count on the reverse strand
}
// ReadBaseStrandTsv reads a basestrand.tsv file from the given io.Reader.
func ReadBaseStrandTsv(r io.Reader) ([]BaseStrandTsvRow, error) {
tsvReader := tsv.NewReader(r)
tsvReader.Comment = '#'
rows := make([]BaseStrandTsvRow, 0)
for {
var row BaseStrandTsvRow
if err := tsvReader.Read(&row); err != nil {
if err == io.EOF {
break
} else {
return nil, err
}
}
rows = append(rows, row)
}
return rows, nil
}
// WriteBaseStrandTsv writes a basestrand.tsv file to the given writer
func WriteBaseStrandTsv(rows []BaseStrandTsvRow, writer io.Writer) error {
tsvWriter := tsv.NewRowWriter(writer)
for _, row := range rows {
if err := tsvWriter.Write(&row); err != nil {
return err
}
}
return tsvWriter.Flush()
}
// singleStrandBaseStrandTsvRow represents a single row of a strand.<fwd/rev>.tsv file.
type singleStrandBaseStrandTsvRow struct {
Chr string `tsv:"CHROM"` // Chromosome
Pos int64 `tsv:"POS"` // Position in chromosome
Depth int64 `tsv:"DEPTH"` // Depth
Ref string `tsv:"REF"` // Reference base
A int64 `tsv:"A"` // A count
C int64 `tsv:"C"` // C count
G int64 `tsv:"G"` // G count
T int64 `tsv:"T"` // T count
N int64 `tsv:"N"` // N count
Ins int64 `tsv:"INS"` // Insertion
Del int64 `tsv:"DEL"` // Deletion
}
// ReadSingleStrandBaseStrandTsv reads strand specific strand.<fwd/rev>.snp.tsv
// files from the given io.Reader.
func ReadSingleStrandBaseStrandTsv(forward, reverse io.Reader) ([]BaseStrandTsvRow, error) {
fwdReader := tsv.NewReader(forward)
fwdReader.HasHeaderRow = true
fwdReader.UseHeaderNames = true
revReader := tsv.NewReader(reverse)
revReader.HasHeaderRow = true
revReader.UseHeaderNames = true
rows := make([]BaseStrandTsvRow, 0)
for {
var fwdRow, revRow singleStrandBaseStrandTsvRow
// Read in the forward strand.
if err := fwdReader.Read(&fwdRow); err != nil {
if err == io.EOF {
break
} else {
return nil, err
}
}
// Read in the reverse strand.
if err := revReader.Read(&revRow); err != nil {
if err == io.EOF {
break
} else {
return nil, err
}
}
// Intersect the rows.
// Check that the chromosome position, and reference are equal.
if fwdRow.Chr != revRow.Chr || fwdRow.Pos != revRow.Pos {
return nil, fmt.Errorf("expected equal chromosome position but got fwd: %s %d and rev: %s %d",
fwdRow.Chr, fwdRow.Pos, revRow.Chr, revRow.Pos)
}
if fwdRow.Ref != revRow.Ref {
return nil, fmt.Errorf("expected equal reference alleles but got fwd: %s and rev: %s", fwdRow.Ref, revRow.Ref)
}
row := BaseStrandTsvRow{
Chr: fwdRow.Chr,
Pos: fwdRow.Pos,
Ref: fwdRow.Ref,
FwdA: fwdRow.A,
RevA: revRow.A,
FwdC: fwdRow.C,
RevC: revRow.C,
FwdG: fwdRow.G,
RevG: revRow.G,
FwdT: fwdRow.T,
RevT: revRow.T,
}
rows = append(rows, row)
}
return rows, nil
}
// ReadBaseStrandTsvIntoChannel reads a basestrand.tsv file from the given tsv.Reader into the given channel.
func ReadBaseStrandTsvIntoChannel(reader *tsv.Reader, c chan []BaseStrandTsvRow, bufferLen int, fileName string, wg *sync.WaitGroup) {
wg.Add(1)
index := 0
buffer := make([]BaseStrandTsvRow, 0, bufferLen)
var row BaseStrandTsvRow
for {
if err := reader.Read(&row); err != nil {
if err == io.EOF {
c <- buffer
break
} else {
log.Panicf("could not parse BaseStrandTsvRow file error %v", err.Error())
}
}
if len(buffer) == cap(buffer) {
c <- buffer
buffer = make([]BaseStrandTsvRow, 0, bufferLen)
}
buffer = append(buffer, row)
index++
if index%1000000 == 0 {
log.Info.Printf("Processed %v Lines in %v", index, fileName)
}
}
close(c)
wg.Done()
log.Info.Printf("Finished reading %v", fileName)
}
// ChrId returns the index for the given chromosome string in bio/pileup
// format.
func ChrId(chr string) (int, error) {
// In our current set of reference files, the chromosome ID is prefixed with
// "chr".
chrId := chr[3:]
switch chrId {
case "X":
return 23, nil
case "Y":
return 24, nil
case "M":
return 25, nil
default:
i, err := strconv.Atoi(chrId)
if err != nil {
return 0, err
}
if i <= 0 || i >= 23 {
return 0, fmt.Errorf("%d is not a valid chromosome id", i)
}
return i, nil
}
}