go-hep.org/x/hep@v0.38.1/groot/rtree/join_example_test.go (about)

     1  // Copyright ©2020 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  package rtree_test
     6  
     7  import (
     8  	"fmt"
     9  	"log"
    10  	"strings"
    11  
    12  	"go-hep.org/x/hep/groot"
    13  	"go-hep.org/x/hep/groot/rtree"
    14  )
    15  
    16  func ExampleJoin() {
    17  
    18  	get := func(fname, tname string) (rtree.Tree, func() error) {
    19  		f, err := groot.Open(fname)
    20  		if err != nil {
    21  			log.Fatal(err)
    22  		}
    23  		t, err := f.Get(tname)
    24  		if err != nil {
    25  			_ = f.Close()
    26  			log.Fatal(err)
    27  		}
    28  		return t.(rtree.Tree), f.Close
    29  	}
    30  	chk := func(f func() error) {
    31  		err := f()
    32  		if err != nil {
    33  			log.Fatal(err)
    34  		}
    35  	}
    36  
    37  	t1, f1 := get("../testdata/join1.root", "j1")
    38  	defer chk(f1)
    39  
    40  	t2, f2 := get("../testdata/join2.root", "j2")
    41  	defer chk(f2)
    42  
    43  	t3, f3 := get("../testdata/join3.root", "j3")
    44  	defer chk(f3)
    45  
    46  	join, err := rtree.Join(t1, t2, t3)
    47  	if err != nil {
    48  		log.Fatalf("could not join trees: %+v", err)
    49  	}
    50  
    51  	fmt.Printf("t1:   %s (nevts=%d)\n", t1.Name(), t1.Entries())
    52  	fmt.Printf("t2:   %s (nevts=%d)\n", t2.Name(), t2.Entries())
    53  	fmt.Printf("t3:   %s (nevts=%d)\n", t3.Name(), t3.Entries())
    54  	fmt.Printf("join: %s\n", join.Name())
    55  	fmt.Printf("entries: %d\n", join.Entries())
    56  
    57  	rvars := rtree.NewReadVars(join)
    58  	r, err := rtree.NewReader(join, rvars)
    59  	if err != nil {
    60  		log.Fatalf("could not create reader for joined trees: %+v", err)
    61  	}
    62  	defer r.Close()
    63  
    64  	rf1, err := r.FormulaFunc(
    65  		[]string{"b10", "b30", "b20"},
    66  		func(b1, b3, b2 float64) float64 {
    67  			return b1 + b2 + b3
    68  		},
    69  	)
    70  	if err != nil {
    71  		log.Fatalf("could not bind formula: %+v", err)
    72  	}
    73  	fct1 := rf1.Func().(func() float64)
    74  
    75  	err = r.Read(func(rctx rtree.RCtx) error {
    76  		for _, rv := range rvars {
    77  			fmt.Printf("join[%03d][%s]: %v\n", rctx.Entry, rv.Name, rv.Deref())
    78  		}
    79  		fmt.Printf("join[%03d][fun]: %v\n", rctx.Entry, fct1())
    80  		return nil
    81  	})
    82  
    83  	if err != nil {
    84  		log.Fatalf("could not process events: %+v", err)
    85  	}
    86  
    87  	// Output:
    88  	// t1:   j1 (nevts=10)
    89  	// t2:   j2 (nevts=10)
    90  	// t3:   j3 (nevts=10)
    91  	// join: join_j1_j2_j3
    92  	// entries: 10
    93  	// join[000][b10]: 101
    94  	// join[000][b11]: 101
    95  	// join[000][b12]: j1-101
    96  	// join[000][b20]: 201
    97  	// join[000][b21]: 201
    98  	// join[000][b22]: j2-201
    99  	// join[000][b30]: 301
   100  	// join[000][b31]: 301
   101  	// join[000][b32]: j3-301
   102  	// join[000][fun]: 603
   103  	// join[001][b10]: 102
   104  	// join[001][b11]: 102
   105  	// join[001][b12]: j1-102
   106  	// join[001][b20]: 202
   107  	// join[001][b21]: 202
   108  	// join[001][b22]: j2-202
   109  	// join[001][b30]: 302
   110  	// join[001][b31]: 302
   111  	// join[001][b32]: j3-302
   112  	// join[001][fun]: 606
   113  	// join[002][b10]: 103
   114  	// join[002][b11]: 103
   115  	// join[002][b12]: j1-103
   116  	// join[002][b20]: 203
   117  	// join[002][b21]: 203
   118  	// join[002][b22]: j2-203
   119  	// join[002][b30]: 303
   120  	// join[002][b31]: 303
   121  	// join[002][b32]: j3-303
   122  	// join[002][fun]: 609
   123  	// join[003][b10]: 104
   124  	// join[003][b11]: 104
   125  	// join[003][b12]: j1-104
   126  	// join[003][b20]: 204
   127  	// join[003][b21]: 204
   128  	// join[003][b22]: j2-204
   129  	// join[003][b30]: 304
   130  	// join[003][b31]: 304
   131  	// join[003][b32]: j3-304
   132  	// join[003][fun]: 612
   133  	// join[004][b10]: 105
   134  	// join[004][b11]: 105
   135  	// join[004][b12]: j1-105
   136  	// join[004][b20]: 205
   137  	// join[004][b21]: 205
   138  	// join[004][b22]: j2-205
   139  	// join[004][b30]: 305
   140  	// join[004][b31]: 305
   141  	// join[004][b32]: j3-305
   142  	// join[004][fun]: 615
   143  	// join[005][b10]: 106
   144  	// join[005][b11]: 106
   145  	// join[005][b12]: j1-106
   146  	// join[005][b20]: 206
   147  	// join[005][b21]: 206
   148  	// join[005][b22]: j2-206
   149  	// join[005][b30]: 306
   150  	// join[005][b31]: 306
   151  	// join[005][b32]: j3-306
   152  	// join[005][fun]: 618
   153  	// join[006][b10]: 107
   154  	// join[006][b11]: 107
   155  	// join[006][b12]: j1-107
   156  	// join[006][b20]: 207
   157  	// join[006][b21]: 207
   158  	// join[006][b22]: j2-207
   159  	// join[006][b30]: 307
   160  	// join[006][b31]: 307
   161  	// join[006][b32]: j3-307
   162  	// join[006][fun]: 621
   163  	// join[007][b10]: 108
   164  	// join[007][b11]: 108
   165  	// join[007][b12]: j1-108
   166  	// join[007][b20]: 208
   167  	// join[007][b21]: 208
   168  	// join[007][b22]: j2-208
   169  	// join[007][b30]: 308
   170  	// join[007][b31]: 308
   171  	// join[007][b32]: j3-308
   172  	// join[007][fun]: 624
   173  	// join[008][b10]: 109
   174  	// join[008][b11]: 109
   175  	// join[008][b12]: j1-109
   176  	// join[008][b20]: 209
   177  	// join[008][b21]: 209
   178  	// join[008][b22]: j2-209
   179  	// join[008][b30]: 309
   180  	// join[008][b31]: 309
   181  	// join[008][b32]: j3-309
   182  	// join[008][fun]: 627
   183  	// join[009][b10]: 110
   184  	// join[009][b11]: 110
   185  	// join[009][b12]: j1-110
   186  	// join[009][b20]: 210
   187  	// join[009][b21]: 210
   188  	// join[009][b22]: j2-210
   189  	// join[009][b30]: 310
   190  	// join[009][b31]: 310
   191  	// join[009][b32]: j3-310
   192  	// join[009][fun]: 630
   193  }
   194  
   195  func ExampleJoin_withReadVarSelection() {
   196  
   197  	get := func(fname, tname string) (rtree.Tree, func() error) {
   198  		f, err := groot.Open(fname)
   199  		if err != nil {
   200  			log.Fatal(err)
   201  		}
   202  		t, err := f.Get(tname)
   203  		if err != nil {
   204  			_ = f.Close()
   205  			log.Fatal(err)
   206  		}
   207  		return t.(rtree.Tree), f.Close
   208  	}
   209  	chk := func(f func() error) {
   210  		err := f()
   211  		if err != nil {
   212  			log.Fatal(err)
   213  		}
   214  	}
   215  
   216  	t1, f1 := get("../testdata/join1.root", "j1")
   217  	defer chk(f1)
   218  
   219  	t2, f2 := get("../testdata/join2.root", "j2")
   220  	defer chk(f2)
   221  
   222  	t3, f3 := get("../testdata/join3.root", "j3")
   223  	defer chk(f3)
   224  
   225  	join, err := rtree.Join(t1, t2, t3)
   226  	if err != nil {
   227  		log.Fatalf("could not join trees: %+v", err)
   228  	}
   229  
   230  	fmt.Printf("t1:   %s (nevts=%d)\n", t1.Name(), t1.Entries())
   231  	fmt.Printf("t2:   %s (nevts=%d)\n", t2.Name(), t2.Entries())
   232  	fmt.Printf("t3:   %s (nevts=%d)\n", t3.Name(), t3.Entries())
   233  	fmt.Printf("join: %s\n", join.Name())
   234  	fmt.Printf("entries: %d\n", join.Entries())
   235  
   236  	rvars := []rtree.ReadVar{
   237  		{Name: "b10", Value: new(float64)},
   238  		{Name: "b20", Value: new(float64)},
   239  	}
   240  
   241  	r, err := rtree.NewReader(join, rvars, rtree.WithRange(3, 8))
   242  	if err != nil {
   243  		log.Fatalf("could not create reader for joined trees: %+v", err)
   244  	}
   245  	defer r.Close()
   246  
   247  	rf1, err := r.FormulaFunc(
   248  		[]string{"b12", "b32", "b22"},
   249  		func(b1, b3, b2 string) string {
   250  			return strings.Join([]string{b1, b3, b2}, ", ")
   251  		},
   252  	)
   253  	if err != nil {
   254  		log.Fatalf("could not bind formula: %+v", err)
   255  	}
   256  	fct1 := rf1.Func().(func() string)
   257  
   258  	err = r.Read(func(rctx rtree.RCtx) error {
   259  		for _, rv := range rvars {
   260  			fmt.Printf("join[%03d][%s]: %v\n", rctx.Entry, rv.Name, rv.Deref())
   261  		}
   262  		fmt.Printf("join[%03d][fun]: %v\n", rctx.Entry, fct1())
   263  		return nil
   264  	})
   265  
   266  	if err != nil {
   267  		log.Fatalf("could not process events: %+v", err)
   268  	}
   269  
   270  	// Output:
   271  	// t1:   j1 (nevts=10)
   272  	// t2:   j2 (nevts=10)
   273  	// t3:   j3 (nevts=10)
   274  	// join: join_j1_j2_j3
   275  	// entries: 10
   276  	// join[003][b10]: 104
   277  	// join[003][b20]: 204
   278  	// join[003][fun]: j1-104, j3-304, j2-204
   279  	// join[004][b10]: 105
   280  	// join[004][b20]: 205
   281  	// join[004][fun]: j1-105, j3-305, j2-205
   282  	// join[005][b10]: 106
   283  	// join[005][b20]: 206
   284  	// join[005][fun]: j1-106, j3-306, j2-206
   285  	// join[006][b10]: 107
   286  	// join[006][b20]: 207
   287  	// join[006][fun]: j1-107, j3-307, j2-207
   288  	// join[007][b10]: 108
   289  	// join[007][b20]: 208
   290  	// join[007][fun]: j1-108, j3-308, j2-208
   291  }