Export hgvs one-hot numpy: -1 for missing / low quality tiles.
[lightning.git] / export.go
1 // Copyright (C) The Lightning Authors. All rights reserved.
2 //
3 // SPDX-License-Identifier: AGPL-3.0
4
5 package lightning
6
7 import (
8         "bufio"
9         "bytes"
10         "context"
11         "errors"
12         "flag"
13         "fmt"
14         "io"
15         "net/http"
16         _ "net/http/pprof"
17         "os"
18         "path/filepath"
19         "runtime"
20         "sort"
21         "strconv"
22         "strings"
23         "sync"
24         "time"
25
26         "git.arvados.org/arvados.git/sdk/go/arvados"
27         "github.com/arvados/lightning/hgvs"
28         "github.com/klauspost/pgzip"
29         "github.com/kshedden/gonpy"
30         "github.com/sirupsen/logrus"
31         log "github.com/sirupsen/logrus"
32 )
33
34 type tvVariant struct {
35         hgvs.Variant
36         librefs map[tileLibRef]bool
37 }
38
39 type outputFormat interface {
40         Filename() string
41         PadLeft() bool
42         Head(out io.Writer, cgs []CompactGenome) error
43         Print(out io.Writer, seqname string, varslice []tvVariant) error
44         Finish(outdir string, out io.Writer, seqname string) error
45         MaxGoroutines() int
46 }
47
48 var outputFormats = map[string]func() outputFormat{
49         "hgvs-numpy": func() outputFormat {
50                 return &formatHGVSNumpy{alleles: map[string][][]int8{}}
51         },
52         "hgvs-onehot": func() outputFormat { return formatHGVSOneHot{} },
53         "hgvs":        func() outputFormat { return formatHGVS{} },
54         "pvcf":        func() outputFormat { return formatPVCF{} },
55         "vcf":         func() outputFormat { return formatVCF{} },
56 }
57
58 type exporter struct {
59         outputFormat   outputFormat
60         outputPerChrom bool
61         compress       bool
62         maxTileSize    int
63         filter         filter
64 }
65
66 func (cmd *exporter) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
67         var err error
68         defer func() {
69                 if err != nil {
70                         fmt.Fprintf(stderr, "%s\n", err)
71                 }
72         }()
73         flags := flag.NewFlagSet("", flag.ContinueOnError)
74         flags.SetOutput(stderr)
75         pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
76         pprofdir := flags.String("pprof-dir", "", "write Go profile data to `directory` periodically")
77         runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
78         projectUUID := flags.String("project", "", "project `UUID` for output data")
79         priority := flags.Int("priority", 500, "container request priority")
80         refname := flags.String("ref", "", "reference genome `name`")
81         inputDir := flags.String("input-dir", ".", "input `directory`")
82         outputDir := flags.String("output-dir", ".", "output `directory`")
83         outputFormatStr := flags.String("output-format", "hgvs", "output `format`: hgvs, pvcf, or vcf")
84         outputBed := flags.String("output-bed", "", "also output bed `file`")
85         flags.BoolVar(&cmd.outputPerChrom, "output-per-chromosome", true, "output one file per chromosome")
86         flags.BoolVar(&cmd.compress, "z", false, "write gzip-compressed output files")
87         labelsFilename := flags.String("output-labels", "", "also output genome labels csv `file`")
88         flags.IntVar(&cmd.maxTileSize, "max-tile-size", 50000, "don't try to make annotations for tiles bigger than given `size`")
89         cmd.filter.Flags(flags)
90         err = flags.Parse(args)
91         if err == flag.ErrHelp {
92                 err = nil
93                 return 0
94         } else if err != nil {
95                 return 2
96         }
97         if flag.NArg() > 0 {
98                 err = fmt.Errorf("extra unparsed command line arguments: %q", flag.Args())
99                 return 2
100         }
101
102         if f, ok := outputFormats[*outputFormatStr]; !ok {
103                 err = fmt.Errorf("invalid output format %q", *outputFormatStr)
104                 return 2
105         } else {
106                 cmd.outputFormat = f()
107         }
108
109         if *pprof != "" {
110                 go func() {
111                         log.Println(http.ListenAndServe(*pprof, nil))
112                 }()
113         }
114         if *pprofdir != "" {
115                 go writeProfilesPeriodically(*pprofdir)
116         }
117
118         if !*runlocal {
119                 if *outputDir != "." {
120                         err = errors.New("cannot specify output directory in container mode: not implemented")
121                         return 1
122                 }
123                 runner := arvadosContainerRunner{
124                         Name:        "lightning export",
125                         Client:      arvados.NewClientFromEnv(),
126                         ProjectUUID: *projectUUID,
127                         RAM:         750000000000,
128                         VCPUs:       96,
129                         Priority:    *priority,
130                         APIAccess:   true,
131                 }
132                 err = runner.TranslatePaths(inputDir)
133                 if err != nil {
134                         return 1
135                 }
136                 if *outputBed != "" {
137                         if strings.Contains(*outputBed, "/") {
138                                 err = fmt.Errorf("cannot use -output-bed filename %q containing '/' char", *outputBed)
139                                 return 1
140                         }
141                         *outputBed = "/mnt/output/" + *outputBed
142                 }
143                 runner.Args = []string{"export", "-local=true",
144                         "-pprof", ":6000",
145                         "-pprof-dir", "/mnt/output",
146                         "-ref", *refname,
147                         "-output-format", *outputFormatStr,
148                         "-output-bed", *outputBed,
149                         "-output-labels", "/mnt/output/labels.csv",
150                         "-output-per-chromosome=" + fmt.Sprintf("%v", cmd.outputPerChrom),
151                         "-max-tile-size", fmt.Sprintf("%d", cmd.maxTileSize),
152                         "-input-dir", *inputDir,
153                         "-output-dir", "/mnt/output",
154                         "-z=" + fmt.Sprintf("%v", cmd.compress),
155                 }
156                 runner.Args = append(runner.Args, cmd.filter.Args()...)
157                 var output string
158                 output, err = runner.Run()
159                 if err != nil {
160                         return 1
161                 }
162                 fmt.Fprintln(stdout, output)
163                 return 0
164         }
165
166         var cgs []CompactGenome
167         tilelib := &tileLibrary{
168                 retainNoCalls:       true,
169                 retainTileSequences: true,
170                 compactGenomes:      map[string][]tileVariantID{},
171         }
172         err = tilelib.LoadDir(context.Background(), *inputDir, nil)
173         if err != nil {
174                 return 1
175         }
176
177         refseq, ok := tilelib.refseqs[*refname]
178         if !ok {
179                 err = fmt.Errorf("reference name %q not found in input; have %v", *refname, func() (names []string) {
180                         for name := range tilelib.refseqs {
181                                 names = append(names, name)
182                         }
183                         return
184                 }())
185                 return 1
186         }
187
188         log.Infof("filtering: %+v", cmd.filter)
189         cmd.filter.Apply(tilelib)
190
191         names := cgnames(tilelib)
192         for _, name := range names {
193                 cgs = append(cgs, CompactGenome{Name: name, Variants: tilelib.compactGenomes[name]})
194         }
195         if *labelsFilename != "" {
196                 log.Infof("writing labels to %s", *labelsFilename)
197                 var f *os.File
198                 f, err = os.OpenFile(*labelsFilename, os.O_CREATE|os.O_WRONLY, 0777)
199                 if err != nil {
200                         return 1
201                 }
202                 defer f.Close()
203                 for i, name := range names {
204                         _, err = fmt.Fprintf(f, "%d,%q,%q\n", i, trimFilenameForLabel(name), cmd.outputFormat.Filename())
205                         if err != nil {
206                                 err = fmt.Errorf("write %s: %w", *labelsFilename, err)
207                                 return 1
208                         }
209                 }
210                 err = f.Close()
211                 if err != nil {
212                         err = fmt.Errorf("close %s: %w", *labelsFilename, err)
213                         return 1
214                 }
215         }
216
217         var bedout io.Writer
218         var bedfile *os.File
219         var bedbufw *bufio.Writer
220         if *outputBed != "" {
221                 bedfile, err = os.OpenFile(*outputBed, os.O_CREATE|os.O_WRONLY, 0666)
222                 if err != nil {
223                         return 1
224                 }
225                 defer bedfile.Close()
226                 bedbufw = bufio.NewWriterSize(bedfile, 16*1024*1024)
227                 bedout = bedbufw
228         }
229
230         err = cmd.export(*outputDir, bedout, tilelib, refseq, cgs)
231         if err != nil {
232                 return 1
233         }
234         if bedout != nil {
235                 err = bedbufw.Flush()
236                 if err != nil {
237                         return 1
238                 }
239                 err = bedfile.Close()
240                 if err != nil {
241                         return 1
242                 }
243         }
244         return 0
245 }
246
247 func (cmd *exporter) export(outdir string, bedout io.Writer, tilelib *tileLibrary, refseq map[string][]tileLibRef, cgs []CompactGenome) error {
248         var seqnames []string
249         var missing []tileLibRef
250         for seqname, librefs := range refseq {
251                 seqnames = append(seqnames, seqname)
252                 for _, libref := range librefs {
253                         if libref.Variant != 0 && tilelib.TileVariantSequence(libref) == nil {
254                                 missing = append(missing, libref)
255                         }
256                 }
257         }
258         sort.Strings(seqnames)
259
260         if len(missing) > 0 {
261                 if limit := 100; len(missing) > limit {
262                         log.Warnf("first %d missing tiles: %v", limit, missing[:limit])
263                 } else {
264                         log.Warnf("missing tiles: %v", missing)
265                 }
266                 return fmt.Errorf("%d needed tiles are missing from library", len(missing))
267         }
268
269         outw := make([]io.WriteCloser, len(seqnames))
270         bedw := make([]io.WriteCloser, len(seqnames))
271
272         var merges sync.WaitGroup
273         merge := func(dst io.Writer, src []io.WriteCloser, label string) {
274                 var mtx sync.Mutex
275                 for i, seqname := range seqnames {
276                         pr, pw := io.Pipe()
277                         src[i] = pw
278                         merges.Add(1)
279                         seqname := seqname
280                         go func() {
281                                 defer merges.Done()
282                                 log.Infof("writing %s %s", seqname, label)
283                                 scanner := bufio.NewScanner(pr)
284                                 for scanner.Scan() {
285                                         mtx.Lock()
286                                         dst.Write(scanner.Bytes())
287                                         dst.Write([]byte{'\n'})
288                                         mtx.Unlock()
289                                 }
290                                 log.Infof("writing %s %s done", seqname, label)
291                         }()
292                 }
293         }
294         if cmd.outputPerChrom {
295                 for i, seqname := range seqnames {
296                         fnm := filepath.Join(outdir, strings.Replace(cmd.outputFormat.Filename(), ".", "."+seqname+".", 1))
297                         if cmd.compress {
298                                 fnm += ".gz"
299                         }
300                         f, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
301                         if err != nil {
302                                 return err
303                         }
304                         defer f.Close()
305                         log.Infof("writing %q", f.Name())
306                         outw[i] = f
307                         if cmd.compress {
308                                 z := pgzip.NewWriter(f)
309                                 defer z.Close()
310                                 outw[i] = z
311                         }
312                         err = cmd.outputFormat.Head(outw[i], cgs)
313                         if err != nil {
314                                 return err
315                         }
316                 }
317         } else {
318                 fnm := filepath.Join(outdir, cmd.outputFormat.Filename())
319                 if cmd.compress {
320                         fnm += ".gz"
321                 }
322                 f, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY|os.O_TRUNC, 0666)
323                 if err != nil {
324                         return err
325                 }
326                 defer f.Close()
327                 log.Infof("writing %q", fnm)
328                 var out io.Writer = f
329                 if cmd.compress {
330                         z := pgzip.NewWriter(out)
331                         defer z.Close()
332                         out = z
333                 }
334                 cmd.outputFormat.Head(out, cgs)
335                 merge(out, outw, "output")
336         }
337         if bedout != nil {
338                 merge(bedout, bedw, "bed")
339         }
340
341         throttle := throttle{Max: runtime.NumCPU()}
342         if max := cmd.outputFormat.MaxGoroutines(); max > 0 {
343                 throttle.Max = max
344         }
345         log.Infof("assembling %d sequences in %d goroutines", len(seqnames), throttle.Max)
346         for seqidx, seqname := range seqnames {
347                 seqidx, seqname := seqidx, seqname
348                 outw := outw[seqidx]
349                 bedw := bedw[seqidx]
350                 throttle.Acquire()
351                 go func() {
352                         defer throttle.Release()
353                         if bedw != nil {
354                                 defer bedw.Close()
355                         }
356                         outwb := bufio.NewWriterSize(outw, 8*1024*1024)
357                         eachVariant(bedw, tilelib.taglib.keylen, seqname, refseq[seqname], tilelib, cgs, cmd.outputFormat.PadLeft(), cmd.maxTileSize, func(varslice []tvVariant) {
358                                 err := cmd.outputFormat.Print(outwb, seqname, varslice)
359                                 throttle.Report(err)
360                         })
361                         err := cmd.outputFormat.Finish(outdir, outwb, seqname)
362                         throttle.Report(err)
363                         err = outwb.Flush()
364                         throttle.Report(err)
365                         err = outw.Close()
366                         throttle.Report(err)
367                 }()
368         }
369
370         merges.Wait()
371         throttle.Wait()
372         return throttle.Err()
373 }
374
375 // Align genome tiles to reference tiles, call callback func on each
376 // variant, and (if bedw is not nil) write tile coverage to bedw.
377 func eachVariant(bedw io.Writer, taglen int, seqname string, reftiles []tileLibRef, tilelib *tileLibrary, cgs []CompactGenome, padLeft bool, maxTileSize int, callback func(varslice []tvVariant)) {
378         t0 := time.Now()
379         progressbar := time.NewTicker(time.Minute)
380         defer progressbar.Stop()
381         var outmtx sync.Mutex
382         defer outmtx.Lock()
383         refpos := 0
384         variantAt := map[int][]tvVariant{} // variantAt[chromOffset][genomeIndex*2+phase]
385         for refstep, libref := range reftiles {
386                 select {
387                 case <-progressbar.C:
388                         var eta interface{}
389                         if refstep > 0 {
390                                 fin := t0.Add(time.Duration(float64(time.Now().Sub(t0)) * float64(len(reftiles)) / float64(refstep)))
391                                 eta = fmt.Sprintf("%v (%v)", fin.Format(time.RFC3339), fin.Sub(time.Now()))
392                         } else {
393                                 eta = "N/A"
394                         }
395                         log.Printf("exportSeq: %s: refstep %d of %d, %.0f/s, ETA %v", seqname, refstep, len(reftiles), float64(refstep)/time.Now().Sub(t0).Seconds(), eta)
396                 default:
397                 }
398                 diffs := map[tileLibRef][]hgvs.Variant{}
399                 refseq := tilelib.TileVariantSequence(libref)
400                 tagcoverage := 0 // number of times the start tag was found in genomes -- max is len(cgs)*2
401                 for cgidx, cg := range cgs {
402                         for phase := 0; phase < 2; phase++ {
403                                 var variant tileVariantID
404                                 if i := int(libref.Tag)*2 + phase; len(cg.Variants) > i {
405                                         variant = cg.Variants[i]
406                                 }
407                                 if variant > 0 {
408                                         tagcoverage++
409                                 }
410                                 if variant == libref.Variant || variant == 0 {
411                                         continue
412                                 }
413                                 glibref := tileLibRef{Tag: libref.Tag, Variant: variant}
414                                 vars, ok := diffs[glibref]
415                                 if !ok {
416                                         genomeseq := tilelib.TileVariantSequence(glibref)
417                                         if len(genomeseq) == 0 {
418                                                 // Hash is known but sequence
419                                                 // is not, e.g., retainNoCalls
420                                                 // was false during import
421                                                 continue
422                                         }
423                                         if len(genomeseq) > maxTileSize {
424                                                 continue
425                                         }
426                                         refSequence := refseq
427                                         // If needed, extend the
428                                         // reference sequence up to
429                                         // the tag at the end of the
430                                         // genomeseq sequence.
431                                         refstepend := refstep + 1
432                                         for refstepend < len(reftiles) && len(refSequence) >= taglen && !bytes.EqualFold(refSequence[len(refSequence)-taglen:], genomeseq[len(genomeseq)-taglen:]) && len(refSequence) <= maxTileSize {
433                                                 if &refSequence[0] == &refseq[0] {
434                                                         refSequence = append([]byte(nil), refSequence...)
435                                                 }
436                                                 refSequence = append(refSequence, tilelib.TileVariantSequence(reftiles[refstepend])...)
437                                                 refstepend++
438                                         }
439                                         // (TODO: handle no-calls)
440                                         if len(refSequence) <= maxTileSize {
441                                                 refstr := strings.ToUpper(string(refSequence))
442                                                 genomestr := strings.ToUpper(string(genomeseq))
443                                                 vars, _ = hgvs.Diff(refstr, genomestr, time.Second)
444                                         }
445                                         diffs[glibref] = vars
446                                 }
447                                 for _, v := range vars {
448                                         if padLeft {
449                                                 v = v.PadLeft()
450                                         }
451                                         v.Position += refpos
452                                         varslice := variantAt[v.Position]
453                                         if varslice == nil {
454                                                 varslice = make([]tvVariant, len(cgs)*2)
455                                                 variantAt[v.Position] = varslice
456                                         }
457                                         varslice[cgidx*2+phase].Variant = v
458                                         if varslice[cgidx*2+phase].librefs == nil {
459                                                 varslice[cgidx*2+phase].librefs = map[tileLibRef]bool{glibref: true}
460                                         } else {
461                                                 varslice[cgidx*2+phase].librefs[glibref] = true
462                                         }
463                                 }
464                         }
465                 }
466                 refpos += len(refseq) - taglen
467
468                 // Flush entries from variantAt that are behind
469                 // refpos. Flush all entries if this is the last
470                 // reftile of the path/chromosome.
471                 flushpos := make([]int, 0, len(variantAt))
472                 lastrefstep := refstep == len(reftiles)-1
473                 for pos := range variantAt {
474                         if lastrefstep || pos <= refpos {
475                                 flushpos = append(flushpos, pos)
476                         }
477                 }
478                 sort.Slice(flushpos, func(i, j int) bool { return flushpos[i] < flushpos[j] })
479                 flushvariants := make([][]tvVariant, len(flushpos))
480                 for i, pos := range flushpos {
481                         varslice := variantAt[pos]
482                         delete(variantAt, pos)
483                         // Check for uninitialized (zero-value)
484                         // elements in varslice
485                         for i := range varslice {
486                                 if varslice[i].Position != 0 {
487                                         // Not a zero-value element
488                                         continue
489                                 }
490                                 // Set the position so
491                                 // varslice[*].Position are all equal
492                                 varslice[i].Position = pos
493                                 // This could be either =ref or a
494                                 // missing/low-quality tile. Figure
495                                 // out which.
496                                 v := cgs[i/2].Variants[int(libref.Tag)*2+i%2]
497                                 if v < 1 || len(tilelib.TileVariantSequence(tileLibRef{Tag: libref.Tag, Variant: v})) == 0 {
498                                         // Missing/low-quality tile.
499                                         varslice[i].New = "-" // fasta "gap of indeterminate length"
500                                 }
501                         }
502                         flushvariants[i] = varslice
503                 }
504                 outmtx.Lock()
505                 go func() {
506                         defer outmtx.Unlock()
507                         for _, varslice := range flushvariants {
508                                 callback(varslice)
509                         }
510                 }()
511                 if bedw != nil && len(refseq) > 0 {
512                         tilestart := refpos - len(refseq) + taglen
513                         tileend := refpos
514                         if !lastrefstep {
515                                 tileend += taglen
516                         }
517                         thickstart := tilestart + taglen
518                         if refstep == 0 {
519                                 thickstart = 0
520                         }
521                         thickend := refpos
522
523                         // coverage score, 0 to 1000
524                         score := 1000
525                         if len(cgs) > 0 {
526                                 score = 1000 * tagcoverage / len(cgs) / 2
527                         }
528
529                         fmt.Fprintf(bedw, "%s %d %d %d %d . %d %d\n",
530                                 seqname, tilestart, tileend,
531                                 libref.Tag,
532                                 score,
533                                 thickstart, thickend)
534                 }
535         }
536 }
537
538 func bucketVarsliceByRef(varslice []tvVariant) map[string]map[string]int {
539         byref := map[string]map[string]int{}
540         for _, v := range varslice {
541                 if v.Ref == "" && v.New == "" {
542                         // =ref
543                         continue
544                 }
545                 if v.New == "-" {
546                         // no-call
547                         continue
548                 }
549                 alts := byref[v.Ref]
550                 if alts == nil {
551                         alts = map[string]int{}
552                         byref[v.Ref] = alts
553                 }
554                 alts[v.New]++
555         }
556         return byref
557 }
558
559 type formatVCF struct{}
560
561 func (formatVCF) MaxGoroutines() int                     { return 0 }
562 func (formatVCF) Filename() string                       { return "out.vcf" }
563 func (formatVCF) PadLeft() bool                          { return true }
564 func (formatVCF) Finish(string, io.Writer, string) error { return nil }
565 func (formatVCF) Head(out io.Writer, cgs []CompactGenome) error {
566         _, err := fmt.Fprint(out, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n")
567         return err
568 }
569 func (formatVCF) Print(out io.Writer, seqname string, varslice []tvVariant) error {
570         for ref, alts := range bucketVarsliceByRef(varslice) {
571                 altslice := make([]string, 0, len(alts))
572                 for alt := range alts {
573                         altslice = append(altslice, alt)
574                 }
575                 sort.Strings(altslice)
576
577                 info := "AC="
578                 for i, a := range altslice {
579                         if i > 0 {
580                                 info += ","
581                         }
582                         info += strconv.Itoa(alts[a])
583                 }
584                 _, err := fmt.Fprintf(out, "%s\t%d\t.\t%s\t%s\t.\t.\t%s\n", seqname, varslice[0].Position, ref, strings.Join(altslice, ","), info)
585                 if err != nil {
586                         return err
587                 }
588         }
589         return nil
590 }
591
592 type formatPVCF struct{}
593
594 func (formatPVCF) MaxGoroutines() int                     { return 0 }
595 func (formatPVCF) Filename() string                       { return "out.vcf" }
596 func (formatPVCF) PadLeft() bool                          { return true }
597 func (formatPVCF) Finish(string, io.Writer, string) error { return nil }
598 func (formatPVCF) Head(out io.Writer, cgs []CompactGenome) error {
599         fmt.Fprintln(out, `##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">`)
600         fmt.Fprintf(out, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT")
601         for _, cg := range cgs {
602                 fmt.Fprintf(out, "\t%s", cg.Name)
603         }
604         _, err := fmt.Fprintf(out, "\n")
605         return err
606 }
607
608 func (formatPVCF) Print(out io.Writer, seqname string, varslice []tvVariant) error {
609         for ref, alts := range bucketVarsliceByRef(varslice) {
610                 altslice := make([]string, 0, len(alts))
611                 for alt := range alts {
612                         altslice = append(altslice, alt)
613                 }
614                 sort.Strings(altslice)
615                 for i, a := range altslice {
616                         alts[a] = i + 1
617                 }
618                 _, err := fmt.Fprintf(out, "%s\t%d\t.\t%s\t%s\t.\t.\t.\tGT", seqname, varslice[0].Position, ref, strings.Join(altslice, ","))
619                 if err != nil {
620                         return err
621                 }
622                 for i := 0; i < len(varslice); i += 2 {
623                         v1, v2 := varslice[i], varslice[i+1]
624                         a1, a2 := alts[v1.New], alts[v2.New]
625                         if v1.Ref != ref {
626                                 // variant on allele 0 belongs on a
627                                 // different output line -- same
628                                 // chr,pos but different "ref" length
629                                 a1 = 0
630                         }
631                         if v2.Ref != ref {
632                                 a2 = 0
633                         }
634                         _, err := fmt.Fprintf(out, "\t%d/%d", a1, a2)
635                         if err != nil {
636                                 return err
637                         }
638                 }
639                 _, err = out.Write([]byte{'\n'})
640                 if err != nil {
641                         return err
642                 }
643         }
644         return nil
645 }
646
647 type formatHGVS struct{}
648
649 func (formatHGVS) MaxGoroutines() int                            { return 0 }
650 func (formatHGVS) Filename() string                              { return "out.tsv" }
651 func (formatHGVS) PadLeft() bool                                 { return false }
652 func (formatHGVS) Head(out io.Writer, cgs []CompactGenome) error { return nil }
653 func (formatHGVS) Finish(string, io.Writer, string) error        { return nil }
654 func (formatHGVS) Print(out io.Writer, seqname string, varslice []tvVariant) error {
655         for i := 0; i < len(varslice)/2; i++ {
656                 if i > 0 {
657                         out.Write([]byte{'\t'})
658                 }
659                 var1, var2 := varslice[i*2], varslice[i*2+1]
660                 if var1.New == "-" || var2.New == "-" {
661                         _, err := out.Write([]byte{'N'})
662                         if err != nil {
663                                 return err
664                         }
665                         continue
666                 }
667                 if var1.Variant == var2.Variant {
668                         if var1.Ref == var1.New {
669                                 _, err := out.Write([]byte{'.'})
670                                 if err != nil {
671                                         return err
672                                 }
673                         } else {
674                                 _, err := fmt.Fprintf(out, "%s:g.%s", seqname, var1.String())
675                                 if err != nil {
676                                         return err
677                                 }
678                         }
679                 } else {
680                         _, err := fmt.Fprintf(out, "%s:g.[%s];[%s]", seqname, var1.String(), var2.String())
681                         if err != nil {
682                                 return err
683                         }
684                 }
685         }
686         _, err := out.Write([]byte{'\n'})
687         return err
688 }
689
690 type formatHGVSOneHot struct{}
691
692 func (formatHGVSOneHot) MaxGoroutines() int                            { return 0 }
693 func (formatHGVSOneHot) Filename() string                              { return "out.tsv" }
694 func (formatHGVSOneHot) PadLeft() bool                                 { return false }
695 func (formatHGVSOneHot) Head(out io.Writer, cgs []CompactGenome) error { return nil }
696 func (formatHGVSOneHot) Finish(string, io.Writer, string) error        { return nil }
697 func (formatHGVSOneHot) Print(out io.Writer, seqname string, varslice []tvVariant) error {
698         vars := map[hgvs.Variant]bool{}
699         for _, v := range varslice {
700                 if v.Ref != v.New {
701                         vars[v.Variant] = true
702                 }
703         }
704
705         // sort variants to ensure output is deterministic
706         sorted := make([]hgvs.Variant, 0, len(vars))
707         for v := range vars {
708                 sorted = append(sorted, v)
709         }
710         sort.Slice(sorted, func(a, b int) bool { return hgvs.Less(sorted[a], sorted[b]) })
711
712         for _, v := range sorted {
713                 if v.New == "-" {
714                         continue
715                 }
716                 fmt.Fprintf(out, "%s.%s", seqname, v.String())
717                 for i := 0; i < len(varslice); i += 2 {
718                         if varslice[i].Variant == v || varslice[i+1].Variant == v {
719                                 out.Write([]byte("\t1"))
720                         } else {
721                                 out.Write([]byte("\t0"))
722                         }
723                 }
724                 _, err := out.Write([]byte{'\n'})
725                 if err != nil {
726                         return err
727                 }
728         }
729         return nil
730 }
731
732 type formatHGVSNumpy struct {
733         sync.Mutex
734         writelock sync.Mutex
735         alleles   map[string][][]int8 // alleles[seqname][variantidx][genomeidx*2+phase]
736 }
737
738 func (*formatHGVSNumpy) MaxGoroutines() int                            { return 4 }
739 func (*formatHGVSNumpy) Filename() string                              { return "annotations.csv" }
740 func (*formatHGVSNumpy) PadLeft() bool                                 { return false }
741 func (*formatHGVSNumpy) Head(out io.Writer, cgs []CompactGenome) error { return nil }
742 func (f *formatHGVSNumpy) Print(outw io.Writer, seqname string, varslice []tvVariant) error {
743         // sort variants to ensure output is deterministic
744         sorted := make([]hgvs.Variant, 0, len(varslice))
745         for _, v := range varslice {
746                 sorted = append(sorted, v.Variant)
747         }
748         sort.Slice(sorted, func(a, b int) bool { return hgvs.Less(sorted[a], sorted[b]) })
749
750         f.Lock()
751         seqalleles := f.alleles[seqname]
752         f.Unlock()
753
754         // append a row to seqalleles for each unique non-ref variant
755         // in varslice.
756         var previous hgvs.Variant
757         for _, v := range sorted {
758                 if previous == v || v.Ref == v.New || v.New == "-" {
759                         continue
760                 }
761                 previous = v
762                 newrow := make([]int8, len(varslice))
763                 for i, allele := range varslice {
764                         if allele.Variant == v {
765                                 newrow[i] = 1
766                         } else if allele.Variant.New == "-" {
767                                 newrow[i] = -1
768                         }
769                 }
770                 seqalleles = append(seqalleles, newrow)
771                 _, err := fmt.Fprintf(outw, "%d,%q\n", len(seqalleles)-1, seqname+"."+v.String())
772                 if err != nil {
773                         return err
774                 }
775         }
776
777         f.Lock()
778         f.alleles[seqname] = seqalleles
779         f.Unlock()
780         return nil
781 }
782 func (f *formatHGVSNumpy) Finish(outdir string, _ io.Writer, seqname string) error {
783         // Write seqname's data to a .npy matrix with one row per
784         // genome and 2 columns per variant.
785         f.Lock()
786         seqalleles := f.alleles[seqname]
787         delete(f.alleles, seqname)
788         f.Unlock()
789         if len(seqalleles) == 0 {
790                 return nil
791         }
792         out := make([]int8, len(seqalleles)*len(seqalleles[0]))
793         rows := len(seqalleles[0]) / 2
794         cols := len(seqalleles) * 2
795         // copy seqalleles[varidx][genome*2+phase] to
796         // out[genome*nvars*2 + varidx*2 + phase]
797         for varidx, alleles := range seqalleles {
798                 for g := 0; g < len(alleles)/2; g++ {
799                         aa, ab := alleles[g*2], alleles[g*2+1]
800                         if aa < 0 || ab < 0 {
801                                 // no-call
802                                 out[g*cols+varidx*2] = -1
803                                 out[g*cols+varidx*2+1] = -1
804                         } else if aa > 0 && ab > 0 {
805                                 // hom
806                                 out[g*cols+varidx*2] = 1
807                         } else if aa > 0 || ab > 0 {
808                                 // het
809                                 out[g*cols+varidx*2+1] = 1
810                         }
811                 }
812         }
813         outf, err := os.OpenFile(outdir+"/matrix."+seqname+".npy", os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0777)
814         if err != nil {
815                 return err
816         }
817         defer outf.Close()
818         bufw := bufio.NewWriter(outf)
819         npw, err := gonpy.NewWriter(nopCloser{bufw})
820         if err != nil {
821                 return err
822         }
823         log.WithFields(logrus.Fields{
824                 "seqname": seqname,
825                 "rows":    rows,
826                 "cols":    cols,
827         }).Info("writing numpy")
828         npw.Shape = []int{rows, cols}
829         f.writelock.Lock() // serialize because WriteInt8 uses lots of memory
830         npw.WriteInt8(out)
831         f.writelock.Unlock()
832         err = bufw.Flush()
833         if err != nil {
834                 return err
835         }
836         err = outf.Close()
837         if err != nil {
838                 return err
839         }
840         return nil
841 }