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