go-hep.org/x/hep@v0.38.1/cmd/npy2root/main.go (about)

     1  // Copyright ©2019 The go-hep Authors. All rights reserved.
     2  // Use of this source code is governed by a BSD-style
     3  // license that can be found in the LICENSE file.
     4  
     5  // Command npy2root converts the content of a NumPy data file to a ROOT file and tree.
     6  //
     7  // Usage: npy2root [OPTIONS] input.npy
     8  //
     9  // The NumPy data file format is described here:
    10  //
    11  //	https://numpy.org/neps/nep-0001-npy-format.html
    12  //
    13  // Example:
    14  //
    15  //	$> npyio-ls input.npy
    16  //	================================================================================
    17  //	file: input.npy
    18  //	npy-header: Header{Major:1, Minor:0, Descr:{Type:<f8, Fortran:false, Shape:[2 3]}}
    19  //	data = [0 1 2 3 4 5]
    20  //
    21  //	$> npy2root -o output.root -t mytree ./input.npy
    22  //	$> root-ls -t ./output.root
    23  //	=== [./output.root] ===
    24  //	version: 61804
    25  //	  TTree   mytree       mytree  (entries=2)
    26  //	    numpy "numpy[3]/D" TBranch
    27  //
    28  //	$> root-dump ./output.root
    29  //	>>> file[./output.root]
    30  //	key[000]: mytree;1 "mytree" (TTree)
    31  //	[000][numpy]: [0 1 2]
    32  //	[001][numpy]: [3 4 5]
    33  //
    34  // Options:
    35  //
    36  //	-o string
    37  //	  	path to output ROOT file (default "output.root")
    38  //	-t string
    39  //	  	name of the output ROOT tree (default "tree")
    40  package main
    41  
    42  import (
    43  	"flag"
    44  	"fmt"
    45  	"log"
    46  	"os"
    47  
    48  	"codeberg.org/sbinet/npyio/npy"
    49  	"git.sr.ht/~sbinet/go-arrow/arrio"
    50  	"go-hep.org/x/hep/groot"
    51  	"go-hep.org/x/hep/groot/rarrow"
    52  	"go-hep.org/x/hep/groot/rnpy"
    53  	"go-hep.org/x/hep/groot/rtree"
    54  )
    55  
    56  func main() {
    57  	log.SetPrefix("npy2root: ")
    58  	log.SetFlags(0)
    59  
    60  	oname := flag.String("o", "output.root", "path to output ROOT file")
    61  	tname := flag.String("t", "tree", "name of the output ROOT tree")
    62  
    63  	flag.Usage = func() {
    64  		fmt.Printf(`npy2root converts the content of a NumPy data file to a ROOT file and tree.
    65  
    66  Usage: npy2root [OPTIONS] input.npy
    67  
    68  The NumPy data file format is described here:
    69  
    70   https://numpy.org/neps/nep-0001-npy-format.html
    71  
    72  Example:
    73  
    74   $> npyio-ls input.npy
    75   ================================================================================
    76   file: input.npy
    77   npy-header: Header{Major:1, Minor:0, Descr:{Type:<f8, Fortran:false, Shape:[2 3]}}
    78   data = [0 1 2 3 4 5]
    79  
    80   $> npy2root -o output.root -t mytree ./input.npy
    81   $> root-ls -t ./output.root
    82   === [./output.root] ===
    83   version: 61804
    84     TTree   mytree       mytree  (entries=2)
    85       numpy "numpy[3]/D" TBranch
    86  
    87   $> root-dump ./output.root
    88   >>> file[./output.root]
    89   key[000]: mytree;1 "mytree" (TTree)
    90   [000][numpy]: [0 1 2]
    91   [001][numpy]: [3 4 5]
    92  
    93  Options:
    94  `)
    95  		flag.PrintDefaults()
    96  	}
    97  
    98  	flag.Parse()
    99  
   100  	if flag.NArg() != 1 {
   101  		flag.Usage()
   102  		log.Fatalf("missing input NumPy data file")
   103  	}
   104  
   105  	fname := flag.Arg(0)
   106  	err := process(*oname, *tname, fname)
   107  	if err != nil {
   108  		log.Fatalf("could not convert %q: %+v", fname, err)
   109  	}
   110  }
   111  
   112  func process(oname, tname, fname string) error {
   113  	src, err := os.Open(fname)
   114  	if err != nil {
   115  		return fmt.Errorf("could not open numpy file %q: %w", fname, err)
   116  	}
   117  	defer src.Close()
   118  
   119  	npy, err := npy.NewReader(src)
   120  	if err != nil {
   121  		return fmt.Errorf("could not create numpy file reader %q: %w", fname, err)
   122  	}
   123  
   124  	rec := rnpy.NewRecord(npy)
   125  	defer rec.Release()
   126  
   127  	dst, err := groot.Create(oname)
   128  	if err != nil {
   129  		return fmt.Errorf("could not create output ROOT file %q: %w", oname, err)
   130  	}
   131  	defer dst.Close()
   132  
   133  	t, err := rarrow.NewFlatTreeWriter(dst, tname, rec.Schema(), rtree.WithTitle(tname))
   134  	if err != nil {
   135  		return fmt.Errorf("could not create output ROOT tree %q: %w", tname, err)
   136  	}
   137  
   138  	_, err = arrio.Copy(t, rnpy.NewRecordReader(rec))
   139  	if err != nil {
   140  		return fmt.Errorf("could not copy numpy array to ROOT tree %q: %w", tname, err)
   141  	}
   142  
   143  	err = t.Close()
   144  	if err != nil {
   145  		return fmt.Errorf("could not close output ROOT tree %q: %w", tname, err)
   146  	}
   147  
   148  	err = dst.Close()
   149  	if err != nil {
   150  		return fmt.Errorf("could not close output ROOT file %q: %w", oname, err)
   151  	}
   152  
   153  	return nil
   154  }