1 // Copyright (C) The Lightning Authors. All rights reserved.
3 // SPDX-License-Identifier: AGPL-3.0
25 "git.arvados.org/arvados.git/sdk/go/arvados"
26 "github.com/arvados/lightning/hgvs"
27 "github.com/kshedden/gonpy"
28 "github.com/sirupsen/logrus"
29 log "github.com/sirupsen/logrus"
32 type exportNumpy struct {
36 func (cmd *exportNumpy) RunCommand(prog string, args []string, stdin io.Reader, stdout, stderr io.Writer) int {
40 fmt.Fprintf(stderr, "%s\n", err)
43 flags := flag.NewFlagSet("", flag.ContinueOnError)
44 flags.SetOutput(stderr)
45 pprof := flags.String("pprof", "", "serve Go profile data at http://`[addr]:port`")
46 runlocal := flags.Bool("local", false, "run on local host (default: run in an arvados container)")
47 projectUUID := flags.String("project", "", "project `UUID` for output data")
48 priority := flags.Int("priority", 500, "container request priority")
49 inputDir := flags.String("input-dir", "./in", "input `directory`")
50 outputDir := flags.String("output-dir", "./out", "output `directory`")
51 annotationsFilename := flags.String("output-annotations", "", "output `file` for tile variant annotations csv")
52 librefsFilename := flags.String("output-onehot2tilevar", "", "when using -one-hot, create csv `file` mapping column# to tag# and variant#")
53 labelsFilename := flags.String("output-labels", "", "output `file` for genome labels csv")
54 regionsFilename := flags.String("regions", "", "only output columns/annotations that intersect regions in specified bed `file`")
55 expandRegions := flags.Int("expand-regions", 0, "expand specified regions by `N` base pairs on each side`")
56 onehot := flags.Bool("one-hot", false, "recode tile variants as one-hot")
57 chunks := flags.Int("chunks", 1, "split output into `N` numpy files")
58 cmd.filter.Flags(flags)
59 err = flags.Parse(args)
60 if err == flag.ErrHelp {
63 } else if err != nil {
69 log.Println(http.ListenAndServe(*pprof, nil))
74 runner := arvadosContainerRunner{
75 Name: "lightning export-numpy",
76 Client: arvados.NewClientFromEnv(),
77 ProjectUUID: *projectUUID,
84 err = runner.TranslatePaths(inputDir, regionsFilename)
88 runner.Args = []string{"export-numpy", "-local=true",
90 fmt.Sprintf("-one-hot=%v", *onehot),
91 "-input-dir", *inputDir,
92 "-output-dir", "/mnt/output",
93 "-output-annotations", "/mnt/output/annotations.csv",
94 "-output-onehot2tilevar", "/mnt/output/onehot2tilevar.csv",
95 "-output-labels", "/mnt/output/labels.csv",
96 "-regions", *regionsFilename,
97 "-expand-regions", fmt.Sprintf("%d", *expandRegions),
98 "-chunks", fmt.Sprintf("%d", *chunks),
100 runner.Args = append(runner.Args, cmd.filter.Args()...)
102 output, err = runner.Run()
106 fmt.Fprintln(stdout, output+"/matrix.npy")
110 tilelib := &tileLibrary{
112 retainTileSequences: true,
113 compactGenomes: map[string][]tileVariantID{},
115 err = tilelib.LoadDir(context.Background(), *inputDir)
120 log.Info("filtering")
121 cmd.filter.Apply(tilelib)
125 log.Info("building lowqual map")
126 lowqual := lowqual(tilelib)
127 names := cgnames(tilelib)
129 if *labelsFilename != "" {
130 log.Infof("writing labels to %s", *labelsFilename)
132 f, err = os.OpenFile(*labelsFilename, os.O_CREATE|os.O_WRONLY, 0777)
137 outBasename := "matrix.npy"
138 for i, name := range names {
139 _, err = fmt.Fprintf(f, "%d,%q,%q\n", i, trimFilenameForLabel(name), outBasename)
141 err = fmt.Errorf("write %s: %w", *labelsFilename, err)
147 err = fmt.Errorf("close %s: %w", *labelsFilename, err)
152 log.Info("determining which tiles intersect given regions")
153 dropTiles, err := chooseTiles(tilelib, *regionsFilename, *expandRegions)
158 annotation2tvs := map[string]map[hgvs.Variant][]tileLibRef{}
159 if *annotationsFilename != "" {
160 log.Info("writing annotations")
161 var annow io.WriteCloser
162 annow, err = os.OpenFile(*annotationsFilename, os.O_CREATE|os.O_WRONLY, 0666)
170 dropTiles: dropTiles,
171 reportAnnotation: func(tag tagID, _ int, variant tileVariantID, refname string, seqname string, pdi hgvs.Variant) {
174 if annotation2tvs[seqname] == nil {
175 annotation2tvs[seqname] = map[hgvs.Variant][]tileLibRef{}
177 annotation2tvs[seqname][pdi] = append(annotation2tvs[seqname][pdi], tileLibRef{Tag: tag, Variant: variant})
179 }).exportTileDiffs(annow, tilelib)
189 var lastErr atomic.Value
190 var wg sync.WaitGroup
191 for seqname, pdivars := range annotation2tvs {
192 seqname, pdivars := seqname, pdivars
196 log.Infof("choosing hgvs columns for seq %s", seqname)
197 var pdis []hgvs.Variant
198 for pdi, librefs := range pdivars {
199 // Include this HGVS column if it was
200 // seen in a variant of any
202 for _, libref := range librefs {
203 if int(libref.Tag) >= len(dropTiles) || !dropTiles[libref.Tag] {
204 pdis = append(pdis, pdi)
209 sort.Slice(pdis, func(i, j int) bool {
210 if cmp := pdis[i].Position - pdis[j].Position; cmp != 0 {
212 } else if pdis[i].Ref != pdis[j].Ref {
213 return pdis[i].Ref < pdis[j].Ref
215 return pdis[i].New < pdis[j].New
218 log.Infof("writing column labels for seq %s", seqname)
220 for _, pdi := range pdis {
221 fmt.Fprintf(&buf, "%s:g.%s\n", seqname, pdi.String())
223 err := ioutil.WriteFile(*outputDir+"/"+seqname+".columns.csv", buf.Bytes(), 0777)
229 log.Infof("building hgvs matrix for seq %s", seqname)
230 data := make([]int8, len(names)*len(pdis)*2)
231 for row, name := range names {
232 cg := tilelib.compactGenomes[name]
233 rowstart := row * len(pdis) * 2
234 for col, pdi := range pdis {
235 for _, libref := range pdivars[pdi] {
236 if len(cg) <= int(libref.Tag)*2+1 {
239 for phase := 0; phase < 2; phase++ {
240 if cg[int(libref.Tag)*2+phase] == libref.Variant {
241 data[rowstart+col*2+phase] = 1
247 log.Infof("writing hgvs numpy for seq %s", seqname)
248 f, err := os.OpenFile(*outputDir+"/"+seqname+".npy", os.O_CREATE|os.O_WRONLY, 0777)
254 // gonpy closes our writer and ignores errors. Give it a nopCloser so we can close f properly.
255 npw, err := gonpy.NewWriter(nopCloser{f})
260 npw.Shape = []int{len(names), len(pdis) * 2}
261 err = npw.WriteInt8(data)
274 if e, ok := lastErr.Load().(error); ok {
279 chunksize := (len(tilelib.variant) + *chunks - 1) / *chunks
280 for chunk := 0; chunk < *chunks; chunk++ {
281 log.Infof("preparing chunk %d of %d", chunk+1, *chunks)
282 tagstart := chunk * chunksize
283 tagend := tagstart + chunksize
284 if tagend > len(tilelib.variant) {
285 tagend = len(tilelib.variant)
287 out, rows, cols := cgs2array(tilelib, names, lowqual, dropTiles, tagstart, tagend)
289 var npw *gonpy.NpyWriter
290 var output io.WriteCloser
291 fnm := *outputDir + "/matrix.npy"
293 fnm = fmt.Sprintf("%s/matrix.%d.npy", *outputDir, chunk)
295 output, err = os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY, 0777)
300 bufw := bufio.NewWriter(output)
301 npw, err = gonpy.NewWriter(nopCloser{bufw})
306 log.Info("recoding to onehot")
307 recoded, librefs, recodedcols := recodeOnehot(out, cols)
308 out, cols = recoded, recodedcols
309 if *librefsFilename != "" {
310 log.Infof("writing onehot column mapping")
311 err = cmd.writeLibRefs(*librefsFilename, tilelib, librefs)
317 log.WithFields(logrus.Fields{
321 }).Info("writing numpy")
322 npw.Shape = []int{rows, cols}
336 func (*exportNumpy) writeLibRefs(fnm string, tilelib *tileLibrary, librefs []tileLibRef) error {
337 f, err := os.OpenFile(fnm, os.O_CREATE|os.O_WRONLY, 0666)
342 for i, libref := range librefs {
343 _, err = fmt.Fprintf(f, "%d,%d,%d\n", i, libref.Tag, libref.Variant)
351 func cgnames(tilelib *tileLibrary) (cgnames []string) {
352 for name := range tilelib.compactGenomes {
353 cgnames = append(cgnames, name)
355 sort.Slice(cgnames, func(i, j int) bool {
356 return trimFilenameForLabel(cgnames[i]) < trimFilenameForLabel(cgnames[j])
361 func lowqual(tilelib *tileLibrary) (lowqual []map[tileVariantID]bool) {
362 lowqual = make([]map[tileVariantID]bool, len(tilelib.variant))
363 for tag, variants := range tilelib.variant {
365 for varidx, hash := range variants {
366 if len(tilelib.hashSequence(hash)) == 0 {
368 lq = map[tileVariantID]bool{}
371 lq[tileVariantID(varidx+1)] = true
378 func cgs2array(tilelib *tileLibrary, names []string, lowqual []map[tileVariantID]bool, dropTiles []bool, tagstart, tagend int) (data []int16, rows, cols int) {
379 rows = len(tilelib.compactGenomes)
380 for tag := tagstart; tag < tagend; tag++ {
381 if len(dropTiles) <= tag || !dropTiles[tag] {
385 data = make([]int16, rows*cols)
386 for row, name := range names {
387 cg := tilelib.compactGenomes[name]
389 for tag := tagstart; tag < tagend && tag*2+1 < len(cg); tag++ {
390 if len(dropTiles) > tag && dropTiles[tag] {
393 for phase := 0; phase < 2; phase++ {
395 if v > 0 && lowqual[tag][v] {
396 data[row*cols+outidx] = -1
398 data[row*cols+outidx] = int16(v)
407 func makeMask(regionsFilename string, expandRegions int) (*mask, error) {
408 log.Printf("makeMask: reading %s", regionsFilename)
409 rfile, err := zopen(regionsFilename)
414 regions, err := io.ReadAll(rfile)
419 log.Print("makeMask: building mask")
421 for _, line := range bytes.Split(regions, []byte{'\n'}) {
422 if bytes.HasPrefix(line, []byte{'#'}) {
425 fields := bytes.Split(line, []byte{'\t'})
429 refseqname := string(fields[0])
430 if strings.HasPrefix(refseqname, "chr") {
431 refseqname = refseqname[3:]
433 start, err1 := strconv.Atoi(string(fields[1]))
434 end, err2 := strconv.Atoi(string(fields[2]))
435 if err1 == nil && err2 == nil {
438 start, err1 = strconv.Atoi(string(fields[3]))
439 end, err2 = strconv.Atoi(string(fields[4]))
440 if err1 == nil && err2 == nil {
444 return nil, fmt.Errorf("cannot parse input line as BED or GFF/GTF: %q", line)
447 mask.Add(refseqname, start-expandRegions, end+expandRegions)
449 log.Print("makeMask: mask.Freeze")
454 func chooseTiles(tilelib *tileLibrary, regionsFilename string, expandRegions int) (drop []bool, err error) {
455 if regionsFilename == "" {
458 mask, err := makeMask(regionsFilename, expandRegions)
463 tagset := tilelib.taglib.Tags()
464 if len(tagset) == 0 {
465 err = errors.New("cannot choose tiles by region in a library without tags")
468 taglen := len(tagset[0])
470 log.Print("chooseTiles: check ref tiles")
471 // Find position+size of each reference tile, and if it
472 // intersects any of the desired regions, set drop[tag]=false.
474 // (Note it doesn't quite work to do the more obvious thing --
475 // start with drop=false and change to true when ref tiles
476 // intersect target regions -- because that would give us
477 // drop=false for tiles that don't appear at all in the
480 // TODO: (optionally?) don't drop tags for which some tile
481 // variants are spanning tiles, i.e., where the reference tile
482 // does not intersect the desired regions, but a spanning tile
483 // from a genome does.
484 drop = make([]bool, len(tilelib.variant))
485 for i := range drop {
488 for refname, refseqs := range tilelib.refseqs {
489 for refseqname, reftiles := range refseqs {
490 if strings.HasPrefix(refseqname, "chr") {
491 refseqname = refseqname[3:]
494 for _, libref := range reftiles {
495 if libref.Variant < 1 {
496 err = fmt.Errorf("reference %q seq %q uses variant zero at tag %d", refname, refseqname, libref.Tag)
499 seq := tilelib.TileVariantSequence(libref)
500 if len(seq) < taglen {
501 err = fmt.Errorf("reference %q seq %q uses tile %d variant %d with sequence len %d < taglen %d", refname, refseqname, libref.Tag, libref.Variant, len(seq), taglen)
505 tileend = tilestart + len(seq) - taglen
506 if mask.Check(refseqname, tilestart, tileend) {
507 drop[libref.Tag] = false
513 log.Print("chooseTiles: done")
517 func recodeOnehot(in []int16, incols int) (out []int16, librefs []tileLibRef, outcols int) {
518 rows := len(in) / incols
519 maxvalue := make([]int16, incols)
520 for row := 0; row < rows; row++ {
521 for col := 0; col < incols; col++ {
522 if v := in[row*incols+col]; maxvalue[col] < v {
527 outcol := make([]int, incols)
529 for incol, maxv := range maxvalue {
530 outcol[incol] = outcols
534 for v := 1; v <= int(maxv); v++ {
535 librefs = append(librefs, tileLibRef{Tag: tagID(incol), Variant: tileVariantID(v)})
539 log.Printf("recodeOnehot: dropped %d input cols with zero maxvalue", dropped)
541 out = make([]int16, rows*outcols)
542 for inidx, row := 0, 0; row < rows; row++ {
543 outrow := out[row*outcols:]
544 for col := 0; col < incols; col++ {
545 if v := in[inidx]; v > 0 {
546 outrow[outcol[col]+int(v)-1] = 1
554 type nopCloser struct {
558 func (nopCloser) Close() error { return nil }
560 func trimFilenameForLabel(s string) string {
561 if i := strings.LastIndex(s, "/"); i >= 0 {
564 s = strings.TrimSuffix(s, ".gz")
565 s = strings.TrimSuffix(s, ".fa")
566 s = strings.TrimSuffix(s, ".fasta")
567 s = strings.TrimSuffix(s, ".1")
568 s = strings.TrimSuffix(s, ".2")
569 s = strings.TrimSuffix(s, ".gz")
570 s = strings.TrimSuffix(s, ".vcf")
571 s = strings.Replace(s, ",", "-", -1)