export.go

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