go-hep.org/x/hep@v0.38.1/groot/rtree/copy_test.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 package rtree_test 6 7 import ( 8 "bytes" 9 "fmt" 10 "os" 11 "path/filepath" 12 "testing" 13 14 "go-hep.org/x/hep/groot" 15 "go-hep.org/x/hep/groot/rcmd" 16 "go-hep.org/x/hep/groot/riofs" 17 "go-hep.org/x/hep/groot/rtree" 18 ) 19 20 func TestCopyTree(t *testing.T) { 21 const deep = true 22 tmp, err := os.MkdirTemp("", "groot-rtree-copy-") 23 if err != nil { 24 t.Fatalf("could not create tmpdir: %+v", err) 25 } 26 defer os.RemoveAll(tmp) 27 28 for _, tc := range []struct { 29 file string 30 tree string 31 branches map[string]int 32 nevts int64 33 want string 34 }{ 35 { 36 file: "../testdata/simple.root", 37 tree: "tree", 38 nevts: 4, 39 want: `key[000]: tree;1 "fake data" (TTree) 40 [000][one]: 1 41 [000][two]: 1.1 42 [000][three]: uno 43 [001][one]: 2 44 [001][two]: 2.2 45 [001][three]: dos 46 [002][one]: 3 47 [002][two]: 3.3 48 [002][three]: tres 49 [003][one]: 4 50 [003][two]: 4.4 51 [003][three]: quatro 52 `, 53 }, 54 { 55 file: "../testdata/simple.root", 56 tree: "tree", 57 branches: map[string]int{ 58 "one": 1, 59 "two": 1, 60 "three": 1, 61 }, 62 nevts: 4, 63 want: `key[000]: tree;1 "fake data" (TTree) 64 [000][one]: 1 65 [000][two]: 1.1 66 [000][three]: uno 67 [001][one]: 2 68 [001][two]: 2.2 69 [001][three]: dos 70 [002][one]: 3 71 [002][two]: 3.3 72 [002][three]: tres 73 [003][one]: 4 74 [003][two]: 4.4 75 [003][three]: quatro 76 `, 77 }, 78 { 79 file: "../testdata/simple.root", 80 tree: "tree", 81 branches: map[string]int{ 82 "one": 1, 83 "three": 1, 84 }, 85 nevts: 4, 86 want: `key[000]: tree;1 "fake data" (TTree) 87 [000][one]: 1 88 [000][three]: uno 89 [001][one]: 2 90 [001][three]: dos 91 [002][one]: 3 92 [002][three]: tres 93 [003][one]: 4 94 [003][three]: quatro 95 `, 96 }, 97 { 98 file: "../testdata/simple.root", 99 tree: "tree", 100 branches: map[string]int{ 101 "one": 1, 102 "two": 1, 103 "three": 1, 104 }, 105 nevts: 0, 106 want: "key[000]: tree;1 \"fake data\" (TTree)\n", 107 }, 108 { 109 file: "../testdata/simple.root", 110 tree: "tree", 111 branches: map[string]int{ 112 "one": 1, 113 "two": 1, 114 "three": 1, 115 }, 116 nevts: -1, 117 want: `key[000]: tree;1 "fake data" (TTree) 118 [000][one]: 1 119 [000][two]: 1.1 120 [000][three]: uno 121 [001][one]: 2 122 [001][two]: 2.2 123 [001][three]: dos 124 [002][one]: 3 125 [002][two]: 3.3 126 [002][three]: tres 127 [003][one]: 4 128 [003][two]: 4.4 129 [003][three]: quatro 130 `, 131 }, 132 { 133 file: "../testdata/leaves.root", 134 tree: "tree", 135 nevts: 4, 136 want: `key[000]: tree;1 "my tree title" (TTree) 137 [000][B]: true 138 [000][Str]: str-0 139 [000][I8]: 0 140 [000][I16]: 0 141 [000][I32]: 0 142 [000][I64]: 0 143 [000][G64]: 0 144 [000][U8]: 0 145 [000][U16]: 0 146 [000][U32]: 0 147 [000][U64]: 0 148 [000][UGG]: 0 149 [000][F32]: 0 150 [000][F64]: 0 151 [000][D16]: 0 152 [000][D32]: 0 153 [000][ArrBs]: [true false false false false false false false false false] 154 [000][ArrI8]: [0 0 0 0 0 0 0 0 0 0] 155 [000][ArrI16]: [0 0 0 0 0 0 0 0 0 0] 156 [000][ArrI32]: [0 0 0 0 0 0 0 0 0 0] 157 [000][ArrI64]: [0 0 0 0 0 0 0 0 0 0] 158 [000][ArrG64]: [0 0 0 0 0 0 0 0 0 0] 159 [000][ArrU8]: [0 0 0 0 0 0 0 0 0 0] 160 [000][ArrU16]: [0 0 0 0 0 0 0 0 0 0] 161 [000][ArrU32]: [0 0 0 0 0 0 0 0 0 0] 162 [000][ArrU64]: [0 0 0 0 0 0 0 0 0 0] 163 [000][ArrUGG]: [0 0 0 0 0 0 0 0 0 0] 164 [000][ArrF32]: [0 0 0 0 0 0 0 0 0 0] 165 [000][ArrF64]: [0 0 0 0 0 0 0 0 0 0] 166 [000][ArrD16]: [0 0 0 0 0 0 0 0 0 0] 167 [000][ArrD32]: [0 0 0 0 0 0 0 0 0 0] 168 [000][N]: 0 169 [000][SliBs]: [] 170 [000][SliI8]: [] 171 [000][SliI16]: [] 172 [000][SliI32]: [] 173 [000][SliI64]: [] 174 [000][SliG64]: [] 175 [000][SliU8]: [] 176 [000][SliU16]: [] 177 [000][SliU32]: [] 178 [000][SliU64]: [] 179 [000][SliUGG]: [] 180 [000][SliF32]: [] 181 [000][SliF64]: [] 182 [000][SliD16]: [] 183 [000][SliD32]: [] 184 [001][B]: false 185 [001][Str]: str-1 186 [001][I8]: -1 187 [001][I16]: -1 188 [001][I32]: -1 189 [001][I64]: -1 190 [001][G64]: -1 191 [001][U8]: 1 192 [001][U16]: 1 193 [001][U32]: 1 194 [001][U64]: 1 195 [001][UGG]: 1 196 [001][F32]: 1 197 [001][F64]: 1 198 [001][D16]: 1 199 [001][D32]: 1 200 [001][ArrBs]: [false true false false false false false false false false] 201 [001][ArrI8]: [-1 -1 -1 -1 -1 -1 -1 -1 -1 -1] 202 [001][ArrI16]: [-1 -1 -1 -1 -1 -1 -1 -1 -1 -1] 203 [001][ArrI32]: [-1 -1 -1 -1 -1 -1 -1 -1 -1 -1] 204 [001][ArrI64]: [-1 -1 -1 -1 -1 -1 -1 -1 -1 -1] 205 [001][ArrG64]: [-1 -1 -1 -1 -1 -1 -1 -1 -1 -1] 206 [001][ArrU8]: [1 1 1 1 1 1 1 1 1 1] 207 [001][ArrU16]: [1 1 1 1 1 1 1 1 1 1] 208 [001][ArrU32]: [1 1 1 1 1 1 1 1 1 1] 209 [001][ArrU64]: [1 1 1 1 1 1 1 1 1 1] 210 [001][ArrUGG]: [1 1 1 1 1 1 1 1 1 1] 211 [001][ArrF32]: [1 1 1 1 1 1 1 1 1 1] 212 [001][ArrF64]: [1 1 1 1 1 1 1 1 1 1] 213 [001][ArrD16]: [1 1 1 1 1 1 1 1 1 1] 214 [001][ArrD32]: [1 1 1 1 1 1 1 1 1 1] 215 [001][N]: 1 216 [001][SliBs]: [true] 217 [001][SliI8]: [-1] 218 [001][SliI16]: [-1] 219 [001][SliI32]: [-1] 220 [001][SliI64]: [-1] 221 [001][SliG64]: [-1] 222 [001][SliU8]: [1] 223 [001][SliU16]: [1] 224 [001][SliU32]: [1] 225 [001][SliU64]: [1] 226 [001][SliUGG]: [1] 227 [001][SliF32]: [1] 228 [001][SliF64]: [1] 229 [001][SliD16]: [1] 230 [001][SliD32]: [1] 231 [002][B]: true 232 [002][Str]: str-2 233 [002][I8]: -2 234 [002][I16]: -2 235 [002][I32]: -2 236 [002][I64]: -2 237 [002][G64]: -2 238 [002][U8]: 2 239 [002][U16]: 2 240 [002][U32]: 2 241 [002][U64]: 2 242 [002][UGG]: 2 243 [002][F32]: 2 244 [002][F64]: 2 245 [002][D16]: 2 246 [002][D32]: 2 247 [002][ArrBs]: [false false true false false false false false false false] 248 [002][ArrI8]: [-2 -2 -2 -2 -2 -2 -2 -2 -2 -2] 249 [002][ArrI16]: [-2 -2 -2 -2 -2 -2 -2 -2 -2 -2] 250 [002][ArrI32]: [-2 -2 -2 -2 -2 -2 -2 -2 -2 -2] 251 [002][ArrI64]: [-2 -2 -2 -2 -2 -2 -2 -2 -2 -2] 252 [002][ArrG64]: [-2 -2 -2 -2 -2 -2 -2 -2 -2 -2] 253 [002][ArrU8]: [2 2 2 2 2 2 2 2 2 2] 254 [002][ArrU16]: [2 2 2 2 2 2 2 2 2 2] 255 [002][ArrU32]: [2 2 2 2 2 2 2 2 2 2] 256 [002][ArrU64]: [2 2 2 2 2 2 2 2 2 2] 257 [002][ArrUGG]: [2 2 2 2 2 2 2 2 2 2] 258 [002][ArrF32]: [2 2 2 2 2 2 2 2 2 2] 259 [002][ArrF64]: [2 2 2 2 2 2 2 2 2 2] 260 [002][ArrD16]: [2 2 2 2 2 2 2 2 2 2] 261 [002][ArrD32]: [2 2 2 2 2 2 2 2 2 2] 262 [002][N]: 2 263 [002][SliBs]: [false true] 264 [002][SliI8]: [-2 -2] 265 [002][SliI16]: [-2 -2] 266 [002][SliI32]: [-2 -2] 267 [002][SliI64]: [-2 -2] 268 [002][SliG64]: [-2 -2] 269 [002][SliU8]: [2 2] 270 [002][SliU16]: [2 2] 271 [002][SliU32]: [2 2] 272 [002][SliU64]: [2 2] 273 [002][SliUGG]: [2 2] 274 [002][SliF32]: [2 2] 275 [002][SliF64]: [2 2] 276 [002][SliD16]: [2 2] 277 [002][SliD32]: [2 2] 278 [003][B]: false 279 [003][Str]: str-3 280 [003][I8]: -3 281 [003][I16]: -3 282 [003][I32]: -3 283 [003][I64]: -3 284 [003][G64]: -3 285 [003][U8]: 3 286 [003][U16]: 3 287 [003][U32]: 3 288 [003][U64]: 3 289 [003][UGG]: 3 290 [003][F32]: 3 291 [003][F64]: 3 292 [003][D16]: 3 293 [003][D32]: 3 294 [003][ArrBs]: [false false false true false false false false false false] 295 [003][ArrI8]: [-3 -3 -3 -3 -3 -3 -3 -3 -3 -3] 296 [003][ArrI16]: [-3 -3 -3 -3 -3 -3 -3 -3 -3 -3] 297 [003][ArrI32]: [-3 -3 -3 -3 -3 -3 -3 -3 -3 -3] 298 [003][ArrI64]: [-3 -3 -3 -3 -3 -3 -3 -3 -3 -3] 299 [003][ArrG64]: [-3 -3 -3 -3 -3 -3 -3 -3 -3 -3] 300 [003][ArrU8]: [3 3 3 3 3 3 3 3 3 3] 301 [003][ArrU16]: [3 3 3 3 3 3 3 3 3 3] 302 [003][ArrU32]: [3 3 3 3 3 3 3 3 3 3] 303 [003][ArrU64]: [3 3 3 3 3 3 3 3 3 3] 304 [003][ArrUGG]: [3 3 3 3 3 3 3 3 3 3] 305 [003][ArrF32]: [3 3 3 3 3 3 3 3 3 3] 306 [003][ArrF64]: [3 3 3 3 3 3 3 3 3 3] 307 [003][ArrD16]: [3 3 3 3 3 3 3 3 3 3] 308 [003][ArrD32]: [3 3 3 3 3 3 3 3 3 3] 309 [003][N]: 3 310 [003][SliBs]: [false false true] 311 [003][SliI8]: [-3 -3 -3] 312 [003][SliI16]: [-3 -3 -3] 313 [003][SliI32]: [-3 -3 -3] 314 [003][SliI64]: [-3 -3 -3] 315 [003][SliG64]: [-3 -3 -3] 316 [003][SliU8]: [3 3 3] 317 [003][SliU16]: [3 3 3] 318 [003][SliU32]: [3 3 3] 319 [003][SliU64]: [3 3 3] 320 [003][SliUGG]: [3 3 3] 321 [003][SliF32]: [3 3 3] 322 [003][SliF64]: [3 3 3] 323 [003][SliD16]: [3 3 3] 324 [003][SliD32]: [3 3 3] 325 `, 326 }, 327 } { 328 t.Run(tc.file, func(t *testing.T) { 329 f, err := groot.Open(tc.file) 330 if err != nil { 331 t.Fatalf("%+v", err) 332 } 333 defer f.Close() 334 335 obj, err := riofs.Dir(f).Get(tc.tree) 336 if err != nil { 337 t.Fatalf("could not get input tree: %+v", err) 338 } 339 340 src := obj.(rtree.Tree) 341 342 wvars := rtree.WriteVarsFromTree(src) 343 if len(tc.branches) > 0 { 344 all := wvars 345 wvars = make([]rtree.WriteVar, 0, len(tc.branches)) 346 for _, wvar := range all { 347 if _, ok := tc.branches[wvar.Name]; !ok { 348 continue 349 } 350 wvars = append(wvars, wvar) 351 } 352 } 353 354 oname := filepath.Join(tmp, fmt.Sprintf("copy-%s.root", filepath.Base(tc.file))) 355 o, err := groot.Create(oname) 356 if err != nil { 357 t.Fatalf("%+v", err) 358 } 359 defer o.Close() 360 361 dst, err := rtree.NewWriter(o, src.Name(), wvars, rtree.WithTitle(src.Title())) 362 if err != nil { 363 t.Fatalf("could not create tree writer: %+v", err) 364 } 365 366 r, err := rtree.NewReader(src, nil, rtree.WithRange(0, tc.nevts)) 367 if err != nil { 368 t.Fatalf("could not create tree reader: %+v", err) 369 } 370 defer r.Close() 371 372 _, err = rtree.Copy(dst, r) 373 if err != nil { 374 t.Fatalf("could not copy tree: %+v", err) 375 } 376 377 err = dst.Close() 378 if err != nil { 379 t.Fatalf("could not close tree: %+v", err) 380 } 381 382 err = o.Close() 383 if err != nil { 384 t.Fatalf("could not close file: %+v", err) 385 } 386 387 got := new(bytes.Buffer) 388 err = rcmd.Dump(got, oname, deep, nil) 389 if err != nil { 390 t.Errorf("could not dump output file: %+v", err) 391 } 392 393 if got, want := got.String(), tc.want; got != want { 394 t.Fatalf("invalid root-dump output:\ngot:\n%s\nwant:\n%s\n", got, want) 395 } 396 }) 397 } 398 }