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

     1  // Copyright ©2018 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  // yoda2root converts YODA files containing hbook-like values (H1D, H2D, P1D, ...)
     6  // into ROOT files.
     7  //
     8  // Example:
     9  //
    10  //	$> yoda2root rivet.yoda rivet.root
    11  //	$> yoda2root rivet.yoda.gz rivet.root
    12  package main // import "go-hep.org/x/hep/cmd/yoda2root"
    13  
    14  import (
    15  	"compress/gzip"
    16  	"flag"
    17  	"fmt"
    18  	"io"
    19  	"log"
    20  	"os"
    21  	"path/filepath"
    22  
    23  	"go-hep.org/x/hep/groot"
    24  	"go-hep.org/x/hep/groot/riofs"
    25  	_ "go-hep.org/x/hep/groot/riofs/plugin/http"
    26  	_ "go-hep.org/x/hep/groot/riofs/plugin/xrootd"
    27  	"go-hep.org/x/hep/groot/root"
    28  	"go-hep.org/x/hep/hbook"
    29  	"go-hep.org/x/hep/hbook/rootcnv"
    30  	"go-hep.org/x/hep/hbook/yodacnv"
    31  )
    32  
    33  func main() {
    34  	log.SetFlags(0)
    35  	log.SetPrefix("yoda2root: ")
    36  	log.SetOutput(os.Stderr)
    37  
    38  	flag.Usage = func() {
    39  		fmt.Fprintf(
    40  			os.Stderr,
    41  			`Usage: yoda2root [options] <file1.yoda> [<file2.yoda> [...]] file.root
    42  
    43  ex:
    44   $> yoda2root rivet.yoda rivet.root
    45   $> yoda2root rivet.yoda.gz rivet.root
    46  `)
    47  	}
    48  
    49  	flag.Parse()
    50  
    51  	if flag.NArg() < 2 {
    52  		log.Printf("missing input and/or output file name(s)")
    53  		flag.Usage()
    54  		flag.PrintDefaults()
    55  		os.Exit(1)
    56  	}
    57  
    58  	oname := flag.Arg(flag.NArg() - 1)
    59  	f, err := groot.Create(oname)
    60  	if err != nil {
    61  		log.Fatalf("could not create output ROOT file: %v", err)
    62  	}
    63  	defer f.Close()
    64  
    65  	for _, fname := range flag.Args()[:flag.NArg()-1] {
    66  		err = convert(f, fname)
    67  		if err != nil {
    68  			log.Fatal(err)
    69  		}
    70  	}
    71  
    72  	err = f.Close()
    73  	if err != nil {
    74  		log.Fatalf("could not close output ROOT file: %v", err)
    75  	}
    76  }
    77  
    78  func convert(w *riofs.File, fname string) error {
    79  	var r io.ReadCloser
    80  	r, err := os.Open(fname)
    81  	if err != nil {
    82  		return fmt.Errorf("error opening file [%s]: %w", fname, err)
    83  	}
    84  	defer r.Close()
    85  
    86  	if filepath.Ext(fname) == ".gz" {
    87  		rz, err := gzip.NewReader(r)
    88  		if err != nil {
    89  			return fmt.Errorf("could not open gzip file [%s]: %w", fname, err)
    90  		}
    91  		defer rz.Close()
    92  		r = rz
    93  	}
    94  
    95  	vs, err := yodacnv.Read(r)
    96  	if err != nil {
    97  		return fmt.Errorf("error decoding YODA file [%s]: %w", fname, err)
    98  	}
    99  
   100  	for i, v := range vs {
   101  		var (
   102  			obj root.Object
   103  			key string
   104  		)
   105  		switch v := v.(type) {
   106  		case *hbook.H1D:
   107  			key = "h1"
   108  			obj = rootcnv.FromH1D(v)
   109  		case *hbook.H2D:
   110  			key = "h2"
   111  			obj = rootcnv.FromH2D(v)
   112  		case *hbook.S2D:
   113  			key = "scatter"
   114  			obj = rootcnv.FromS2D(v)
   115  		default:
   116  			log.Printf("%s: no YODA -> ROOT conversion for %T", fname, v)
   117  			continue
   118  		}
   119  		switch v.Name() {
   120  		case "":
   121  			key = fmt.Sprintf("yoda-%s-%03d", key, i)
   122  		default:
   123  			key = v.Name()
   124  		}
   125  		err = riofs.Dir(w).Put(key, obj)
   126  		if err != nil {
   127  			return fmt.Errorf("error writing %q from YODA file [%s]: %w", v.Name(), fname, err)
   128  		}
   129  	}
   130  
   131  	return nil
   132  }