Add filter options to export cmd.
[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][][]bool{}}
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                                 if len(cg.Variants) <= int(libref.Tag)*2+phase {
404                                         continue
405                                 }
406                                 variant := cg.Variants[int(libref.Tag)*2+phase]
407                                 if variant == 0 {
408                                         continue
409                                 }
410                                 tagcoverage++
411                                 if variant == libref.Variant {
412                                         continue
413                                 }
414                                 glibref := tileLibRef{Tag: libref.Tag, Variant: variant}
415                                 vars, ok := diffs[glibref]
416                                 if !ok {
417                                         genomeseq := tilelib.TileVariantSequence(glibref)
418                                         if len(genomeseq) == 0 {
419                                                 // Hash is known but sequence
420                                                 // is not, e.g., retainNoCalls
421                                                 // was false during import
422                                                 continue
423                                         }
424                                         if len(genomeseq) > maxTileSize {
425                                                 continue
426                                         }
427                                         refSequence := refseq
428                                         // If needed, extend the
429                                         // reference sequence up to
430                                         // the tag at the end of the
431                                         // genomeseq sequence.
432                                         refstepend := refstep + 1
433                                         for refstepend < len(reftiles) && len(refSequence) >= taglen && !bytes.EqualFold(refSequence[len(refSequence)-taglen:], genomeseq[len(genomeseq)-taglen:]) && len(refSequence) <= maxTileSize {
434                                                 if &refSequence[0] == &refseq[0] {
435                                                         refSequence = append([]byte(nil), refSequence...)
436                                                 }
437                                                 refSequence = append(refSequence, tilelib.TileVariantSequence(reftiles[refstepend])...)
438                                                 refstepend++
439                                         }
440                                         // (TODO: handle no-calls)
441                                         if len(refSequence) <= maxTileSize {
442                                                 refstr := strings.ToUpper(string(refSequence))
443                                                 genomestr := strings.ToUpper(string(genomeseq))
444                                                 vars, _ = hgvs.Diff(refstr, genomestr, time.Second)
445                                         }
446                                         diffs[glibref] = vars
447                                 }
448                                 for _, v := range vars {
449                                         if padLeft {
450                                                 v = v.PadLeft()
451                                         }
452                                         v.Position += refpos
453                                         varslice := variantAt[v.Position]
454                                         if varslice == nil {
455                                                 varslice = make([]tvVariant, len(cgs)*2)
456                                                 variantAt[v.Position] = varslice
457                                         }
458                                         varslice[cgidx*2+phase].Variant = v
459                                         if varslice[cgidx*2+phase].librefs == nil {
460                                                 varslice[cgidx*2+phase].librefs = map[tileLibRef]bool{glibref: true}
461                                         } else {
462                                                 varslice[cgidx*2+phase].librefs[glibref] = true
463                                         }
464                                 }
465                         }
466                 }
467                 refpos += len(refseq) - taglen
468
469                 // Flush entries from variantAt that are behind
470                 // refpos. Flush all entries if this is the last
471                 // reftile of the path/chromosome.
472                 flushpos := make([]int, 0, len(variantAt))
473                 lastrefstep := refstep == len(reftiles)-1
474                 for pos := range variantAt {
475                         if lastrefstep || pos <= refpos {
476                                 flushpos = append(flushpos, pos)
477                         }
478                 }
479                 sort.Slice(flushpos, func(i, j int) bool { return flushpos[i] < flushpos[j] })
480                 flushvariants := make([][]tvVariant, len(flushpos))
481                 for i, pos := range flushpos {
482                         varslice := variantAt[pos]
483                         delete(variantAt, pos)
484                         for i := range varslice {
485                                 if varslice[i].Position == 0 {
486                                         varslice[i].Position = pos
487                                 }
488                         }
489                         flushvariants[i] = varslice
490                 }
491                 outmtx.Lock()
492                 go func() {
493                         defer outmtx.Unlock()
494                         for _, varslice := range flushvariants {
495                                 callback(varslice)
496                         }
497                 }()
498                 if bedw != nil && len(refseq) > 0 {
499                         tilestart := refpos - len(refseq) + taglen
500                         tileend := refpos
501                         if !lastrefstep {
502                                 tileend += taglen
503                         }
504                         thickstart := tilestart + taglen
505                         if refstep == 0 {
506                                 thickstart = 0
507                         }
508                         thickend := refpos
509
510                         // coverage score, 0 to 1000
511                         score := 1000
512                         if len(cgs) > 0 {
513                                 score = 1000 * tagcoverage / len(cgs) / 2
514                         }
515
516                         fmt.Fprintf(bedw, "%s %d %d %d %d . %d %d\n",
517                                 seqname, tilestart, tileend,
518                                 libref.Tag,
519                                 score,
520                                 thickstart, thickend)
521                 }
522         }
523 }
524
525 func bucketVarsliceByRef(varslice []tvVariant) map[string]map[string]int {
526         byref := map[string]map[string]int{}
527         for _, v := range varslice {
528                 if v.Ref == "" && v.New == "" {
529                         continue
530                 }
531                 alts := byref[v.Ref]
532                 if alts == nil {
533                         alts = map[string]int{}
534                         byref[v.Ref] = alts
535                 }
536                 alts[v.New]++
537         }
538         return byref
539 }
540
541 type formatVCF struct{}
542
543 func (formatVCF) MaxGoroutines() int                     { return 0 }
544 func (formatVCF) Filename() string                       { return "out.vcf" }
545 func (formatVCF) PadLeft() bool                          { return true }
546 func (formatVCF) Finish(string, io.Writer, string) error { return nil }
547 func (formatVCF) Head(out io.Writer, cgs []CompactGenome) error {
548         _, err := fmt.Fprint(out, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\n")
549         return err
550 }
551 func (formatVCF) Print(out io.Writer, seqname string, varslice []tvVariant) error {
552         for ref, alts := range bucketVarsliceByRef(varslice) {
553                 altslice := make([]string, 0, len(alts))
554                 for alt := range alts {
555                         altslice = append(altslice, alt)
556                 }
557                 sort.Strings(altslice)
558
559                 info := "AC="
560                 for i, a := range altslice {
561                         if i > 0 {
562                                 info += ","
563                         }
564                         info += strconv.Itoa(alts[a])
565                 }
566                 _, 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)
567                 if err != nil {
568                         return err
569                 }
570         }
571         return nil
572 }
573
574 type formatPVCF struct{}
575
576 func (formatPVCF) MaxGoroutines() int                     { return 0 }
577 func (formatPVCF) Filename() string                       { return "out.vcf" }
578 func (formatPVCF) PadLeft() bool                          { return true }
579 func (formatPVCF) Finish(string, io.Writer, string) error { return nil }
580 func (formatPVCF) Head(out io.Writer, cgs []CompactGenome) error {
581         fmt.Fprintln(out, `##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">`)
582         fmt.Fprintf(out, "#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT")
583         for _, cg := range cgs {
584                 fmt.Fprintf(out, "\t%s", cg.Name)
585         }
586         _, err := fmt.Fprintf(out, "\n")
587         return err
588 }
589
590 func (formatPVCF) Print(out io.Writer, seqname string, varslice []tvVariant) error {
591         for ref, alts := range bucketVarsliceByRef(varslice) {
592                 altslice := make([]string, 0, len(alts))
593                 for alt := range alts {
594                         altslice = append(altslice, alt)
595                 }
596                 sort.Strings(altslice)
597                 for i, a := range altslice {
598                         alts[a] = i + 1
599                 }
600                 _, err := fmt.Fprintf(out, "%s\t%d\t.\t%s\t%s\t.\t.\t.\tGT", seqname, varslice[0].Position, ref, strings.Join(altslice, ","))
601                 if err != nil {
602                         return err
603                 }
604                 for i := 0; i < len(varslice); i += 2 {
605                         v1, v2 := varslice[i], varslice[i+1]
606                         a1, a2 := alts[v1.New], alts[v2.New]
607                         if v1.Ref != ref {
608                                 // variant on allele 0 belongs on a
609                                 // different output line -- same
610                                 // chr,pos but different "ref" length
611                                 a1 = 0
612                         }
613                         if v2.Ref != ref {
614                                 a2 = 0
615                         }
616                         _, err := fmt.Fprintf(out, "\t%d/%d", a1, a2)
617                         if err != nil {
618                                 return err
619                         }
620                 }
621                 _, err = out.Write([]byte{'\n'})
622                 if err != nil {
623                         return err
624                 }
625         }
626         return nil
627 }
628
629 type formatHGVS struct{}
630
631 func (formatHGVS) MaxGoroutines() int                            { return 0 }
632 func (formatHGVS) Filename() string                              { return "out.tsv" }
633 func (formatHGVS) PadLeft() bool                                 { return false }
634 func (formatHGVS) Head(out io.Writer, cgs []CompactGenome) error { return nil }
635 func (formatHGVS) Finish(string, io.Writer, string) error        { return nil }
636 func (formatHGVS) Print(out io.Writer, seqname string, varslice []tvVariant) error {
637         for i := 0; i < len(varslice)/2; i++ {
638                 if i > 0 {
639                         out.Write([]byte{'\t'})
640                 }
641                 var1, var2 := varslice[i*2], varslice[i*2+1]
642                 if var1.Variant == var2.Variant {
643                         if var1.Ref == var1.New {
644                                 _, err := out.Write([]byte{'.'})
645                                 if err != nil {
646                                         return err
647                                 }
648                         } else {
649                                 _, err := fmt.Fprintf(out, "%s:g.%s", seqname, var1.String())
650                                 if err != nil {
651                                         return err
652                                 }
653                         }
654                 } else {
655                         _, err := fmt.Fprintf(out, "%s:g.[%s];[%s]", seqname, var1.String(), var2.String())
656                         if err != nil {
657                                 return err
658                         }
659                 }
660         }
661         _, err := out.Write([]byte{'\n'})
662         return err
663 }
664
665 type formatHGVSOneHot struct{}
666
667 func (formatHGVSOneHot) MaxGoroutines() int                            { return 0 }
668 func (formatHGVSOneHot) Filename() string                              { return "out.tsv" }
669 func (formatHGVSOneHot) PadLeft() bool                                 { return false }
670 func (formatHGVSOneHot) Head(out io.Writer, cgs []CompactGenome) error { return nil }
671 func (formatHGVSOneHot) Finish(string, io.Writer, string) error        { return nil }
672 func (formatHGVSOneHot) Print(out io.Writer, seqname string, varslice []tvVariant) error {
673         vars := map[hgvs.Variant]bool{}
674         for _, v := range varslice {
675                 if v.Ref != v.New {
676                         vars[v.Variant] = true
677                 }
678         }
679
680         // sort variants to ensure output is deterministic
681         sorted := make([]hgvs.Variant, 0, len(vars))
682         for v := range vars {
683                 sorted = append(sorted, v)
684         }
685         sort.Slice(sorted, func(a, b int) bool { return hgvs.Less(sorted[a], sorted[b]) })
686
687         for _, v := range sorted {
688                 fmt.Fprintf(out, "%s.%s", seqname, v.String())
689                 for i := 0; i < len(varslice); i += 2 {
690                         if varslice[i].Variant == v || varslice[i+1].Variant == v {
691                                 out.Write([]byte("\t1"))
692                         } else {
693                                 out.Write([]byte("\t0"))
694                         }
695                 }
696                 _, err := out.Write([]byte{'\n'})
697                 if err != nil {
698                         return err
699                 }
700         }
701         return nil
702 }
703
704 type formatHGVSNumpy struct {
705         sync.Mutex
706         writelock sync.Mutex
707         alleles   map[string][][]bool // alleles[seqname][variantidx][genomeidx*2+phase]
708 }
709
710 func (*formatHGVSNumpy) MaxGoroutines() int                            { return 4 }
711 func (*formatHGVSNumpy) Filename() string                              { return "annotations.csv" }
712 func (*formatHGVSNumpy) PadLeft() bool                                 { return false }
713 func (*formatHGVSNumpy) Head(out io.Writer, cgs []CompactGenome) error { return nil }
714 func (f *formatHGVSNumpy) Print(outw io.Writer, seqname string, varslice []tvVariant) error {
715         // sort variants to ensure output is deterministic
716         sorted := make([]hgvs.Variant, 0, len(varslice))
717         for _, v := range varslice {
718                 sorted = append(sorted, v.Variant)
719         }
720         sort.Slice(sorted, func(a, b int) bool { return hgvs.Less(sorted[a], sorted[b]) })
721
722         f.Lock()
723         seqalleles := f.alleles[seqname]
724         f.Unlock()
725
726         // append a row to seqvariants and seqalleles for each unique
727         // non-ref variant in varslice.
728         var previous hgvs.Variant
729         for _, v := range sorted {
730                 if previous == v || v.Ref == v.New {
731                         continue
732                 }
733                 previous = v
734                 newrow := make([]bool, len(varslice))
735                 for i, allele := range varslice {
736                         if allele.Variant == v {
737                                 newrow[i] = true
738                         }
739                 }
740                 seqalleles = append(seqalleles, newrow)
741                 _, err := fmt.Fprintf(outw, "%d,%q\n", len(seqalleles)-1, seqname+"."+v.String())
742                 if err != nil {
743                         return err
744                 }
745         }
746
747         f.Lock()
748         f.alleles[seqname] = seqalleles
749         f.Unlock()
750         return nil
751 }
752 func (f *formatHGVSNumpy) Finish(outdir string, _ io.Writer, seqname string) error {
753         // Write seqname's data to a .npy matrix with one row per
754         // genome and 2 columns per variant.
755         seqalleles := f.alleles[seqname]
756         f.Lock()
757         delete(f.alleles, seqname)
758         f.Unlock()
759         if len(seqalleles) == 0 {
760                 return nil
761         }
762         out := make([]int8, len(seqalleles)*len(seqalleles[0]))
763         rows := len(seqalleles[0]) / 2
764         cols := len(seqalleles) * 2
765         // copy seqalleles[varidx][genome*2+phase] to
766         // out[genome*nvars*2 + varidx*2 + phase]
767         for varidx, alleles := range seqalleles {
768                 for g := 0; g < len(alleles)/2; g++ {
769                         aa, ab := alleles[g*2], alleles[g*2+1]
770                         if aa && ab {
771                                 // hom
772                                 out[g*cols+varidx*2] = 1
773                         } else if aa || ab {
774                                 // het
775                                 out[g*cols+varidx*2+1] = 1
776                         }
777                 }
778         }
779         outf, err := os.OpenFile(outdir+"/matrix."+seqname+".npy", os.O_CREATE|os.O_EXCL|os.O_WRONLY, 0777)
780         if err != nil {
781                 return err
782         }
783         defer outf.Close()
784         bufw := bufio.NewWriter(outf)
785         npw, err := gonpy.NewWriter(nopCloser{bufw})
786         if err != nil {
787                 return err
788         }
789         log.WithFields(logrus.Fields{
790                 "seqname": seqname,
791                 "rows":    rows,
792                 "cols":    cols,
793         }).Info("writing numpy")
794         npw.Shape = []int{rows, cols}
795         f.writelock.Lock() // serialize because WriteInt8 uses lots of memory
796         npw.WriteInt8(out)
797         f.writelock.Unlock()
798         err = bufw.Flush()
799         if err != nil {
800                 return err
801         }
802         err = outf.Close()
803         if err != nil {
804                 return err
805         }
806         return nil
807 }