Commit
This commit does not belong to any branch on this repository, and may belong to a fork outside of the repository.
groot/rtree: implement a Formula interpreter
Fixes #634.
- Loading branch information
Showing
6 changed files
with
313 additions
and
0 deletions.
There are no files selected for viewing
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,86 @@ | ||
| // Copyright 2020 The go-hep Authors. All rights reserved. | ||
| // Use of this source code is governed by a BSD-style | ||
| // license that can be found in the LICENSE file. | ||
|
|
||
| package rtree | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "reflect" | ||
| "strings" | ||
|
|
||
| "github.com/containous/yaegi/interp" | ||
| "github.com/containous/yaegi/stdlib" | ||
| ) | ||
|
|
||
| // Formula is a mathematical formula bound to variables (branches) of | ||
| // a given ROOT tree. | ||
| // | ||
| // Formulae are attached to a rtree.Reader. | ||
| type Formula struct { | ||
| r *Reader | ||
| expr string | ||
| prog string | ||
| eval *interp.Interpreter | ||
| fct func() interface{} | ||
| } | ||
|
|
||
| func newFormula(r *Reader, expr string, imports []string) (Formula, error) { | ||
| var ( | ||
| eval = interp.New(interp.Options{}) | ||
| pkg = "groot_rtree" | ||
| uses = interp.Exports{ | ||
| pkg: make(map[string]reflect.Value), | ||
| } | ||
| prog = new(strings.Builder) | ||
| ) | ||
|
|
||
| for _, name := range imports { | ||
| if _, ok := stdlib.Symbols[name]; !ok { | ||
| return Formula{}, fmt.Errorf("rtree: no known stdlib import for %q", name) | ||
| } | ||
| fmt.Fprintf(prog, "import %q\n", name) | ||
| } | ||
|
|
||
| fmt.Fprintf(prog, "import %q\n", pkg) | ||
| fmt.Fprintf(prog, "func _groot_rtree_func_eval() interface{} {\n") | ||
|
|
||
| for _, rvar := range r.rvars { | ||
| name := "Var_" + rvar.Name | ||
| uses[pkg][name] = reflect.ValueOf(rvar.Value) | ||
| // FIXME(sbinet): only load rvars that are actually used. | ||
| fmt.Fprintf(prog, "\t%s := *%s.%s // %T\n", rvar.Name, pkg, name, rvar.Value) | ||
| } | ||
|
|
||
| eval.Use(stdlib.Symbols) | ||
| eval.Use(uses) | ||
|
|
||
| fmt.Fprintf(prog, | ||
| "\t_groot_return := %s\n\treturn &_groot_return\n}", | ||
| expr, | ||
| ) | ||
|
|
||
| _, err := eval.Eval(prog.String()) | ||
| if err != nil { | ||
| return Formula{}, fmt.Errorf("rtree: could not define formula eval-func: %w", err) | ||
| } | ||
|
|
||
| f, err := eval.Eval("_groot_rtree_func_eval") | ||
| if err != nil { | ||
| return Formula{}, fmt.Errorf("rtree: could not retrieve formula eval-func: %w", err) | ||
| } | ||
|
|
||
| form := Formula{ | ||
| r: r, | ||
| expr: expr, | ||
| prog: prog.String(), | ||
| eval: eval, | ||
| fct: f.Interface().(func() interface{}), | ||
| } | ||
|
|
||
| return form, nil | ||
| } | ||
|
|
||
| func (form *Formula) Eval() interface{} { | ||
| return reflect.ValueOf(form.fct()).Elem().Interface() | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
| Original file line number | Diff line number | Diff line change |
|---|---|---|
| @@ -0,0 +1,155 @@ | ||
| // Copyright 2020 The go-hep Authors. All rights reserved. | ||
| // Use of this source code is governed by a BSD-style | ||
| // license that can be found in the LICENSE file. | ||
|
|
||
| package rtree | ||
|
|
||
| import ( | ||
| "fmt" | ||
| "reflect" | ||
| "testing" | ||
|
|
||
| "go-hep.org/x/hep/groot/riofs" | ||
| "go-hep.org/x/hep/groot/root" | ||
| ) | ||
|
|
||
| func TestFormula(t *testing.T) { | ||
| for _, tc := range []struct { | ||
| fname string | ||
| tname string | ||
| expr string | ||
| imports []string | ||
| want []interface{} | ||
| err error | ||
| }{ | ||
| { | ||
| fname: "../testdata/simple.root", | ||
| tname: "tree", | ||
| expr: "one", | ||
| want: []interface{}{int32(1), int32(2)}, | ||
| }, | ||
| { | ||
| fname: "../testdata/simple.root", | ||
| tname: "tree", | ||
| expr: "one*one", | ||
| want: []interface{}{int32(1), int32(4)}, | ||
| }, | ||
| { | ||
| fname: "../testdata/simple.root", | ||
| tname: "tree", | ||
| expr: "math.Sqrt(float64(one*one))", | ||
| imports: []string{"math"}, | ||
| want: []interface{}{float64(1), float64(2)}, | ||
| }, | ||
| { | ||
| fname: "../testdata/simple.root", | ||
| tname: "tree", | ||
| expr: `fmt.Sprintf("%d", one)`, | ||
| imports: []string{"fmt"}, | ||
| want: []interface{}{"1", "2"}, | ||
| }, | ||
| { | ||
| fname: "../testdata/leaves.root", | ||
| tname: "tree", | ||
| expr: "ArrU64", | ||
| want: []interface{}{[10]uint64{}, [10]uint64{1, 1, 1, 1, 1, 1, 1, 1, 1, 1}}, | ||
| }, | ||
| { | ||
| fname: "../testdata/leaves.root", | ||
| tname: "tree", | ||
| expr: "ArrU64[0]", | ||
| want: []interface{}{uint64(0), uint64(1)}, | ||
| }, | ||
| { | ||
| fname: "../testdata/leaves.root", | ||
| tname: "tree", | ||
| expr: "D32", | ||
| want: []interface{}{root.Double32(0), root.Double32(1)}, | ||
| }, | ||
| { | ||
| fname: "../testdata/leaves.root", | ||
| tname: "tree", | ||
| expr: "float64(D32)+float64(len(SliI64))", | ||
| want: []interface{}{0.0, 2.0}, | ||
| }, | ||
| { | ||
| fname: "../testdata/simple.root", | ||
| tname: "tree", | ||
| expr: "ones", | ||
| err: fmt.Errorf("rtree: could not create Formula: rtree: could not define formula eval-func: 6:19: undefined: ones"), | ||
| }, | ||
| { | ||
| fname: "../testdata/simple.root", | ||
| tname: "tree", | ||
| expr: "one", | ||
| imports: []string{"go-hep.org/x/hep/groot"}, | ||
| err: fmt.Errorf(`rtree: could not create Formula: rtree: no known stdlib import for "go-hep.org/x/hep/groot"`), | ||
| }, | ||
| { | ||
| fname: "../testdata/simple.root", | ||
| tname: "tree", | ||
| expr: "one+three", | ||
| err: fmt.Errorf(`rtree: could not create Formula: rtree: could not define formula eval-func: 6:19: mismatched types .int32 and .string`), | ||
| }, | ||
| { | ||
| fname: "../testdata/simple.root", | ||
| tname: "tree", | ||
| expr: "math.Sqrt(float64(one))", | ||
| err: fmt.Errorf(`rtree: could not create Formula: rtree: could not define formula eval-func: 6:19: undefined: math`), | ||
| }, | ||
| } { | ||
| t.Run(tc.expr, func(t *testing.T) { | ||
| f, err := riofs.Open(tc.fname) | ||
| if err != nil { | ||
| t.Fatal(err) | ||
| } | ||
| defer f.Close() | ||
|
|
||
| o, err := riofs.Dir(f).Get(tc.tname) | ||
| if err != nil { | ||
| t.Fatal(err) | ||
| } | ||
|
|
||
| tree := o.(Tree) | ||
|
|
||
| r, err := NewReader(tree, NewReadVars(tree), WithRange(0, 2)) | ||
| if err != nil { | ||
| t.Fatal(err) | ||
| } | ||
| defer r.Close() | ||
|
|
||
| form, err := r.Formula(tc.expr, tc.imports) | ||
| switch { | ||
| case err != nil && tc.err != nil: | ||
| if got, want := err.Error(), tc.err.Error(); got != want { | ||
| t.Fatalf("invalid error.\ngot= %v\nwant=%v", got, want) | ||
| } | ||
| return | ||
| case err != nil && tc.err == nil: | ||
| t.Fatalf("unexpected error: %+v", err) | ||
| case err == nil && tc.err != nil: | ||
| t.Fatalf("expected an error: %v (got=nil)", tc.err) | ||
| case err == nil && tc.err == nil: | ||
| // ok. | ||
| } | ||
|
|
||
| defer func() { | ||
| e := recover() | ||
| if e != nil { | ||
| t.Fatalf("could not run form-eval:\n%s\n%+v", form.prog, e) | ||
| } | ||
| }() | ||
|
|
||
| err = r.Read(func(ctx RCtx) error { | ||
| got := form.Eval() | ||
| if got, want := got, tc.want[ctx.Entry]; !reflect.DeepEqual(got, want) { | ||
| return fmt.Errorf("entry[%d]: invalid form-eval:\ngot=%v (%T)\nwant=%v (%T)", ctx.Entry, got, got, want, want) | ||
| } | ||
| return nil | ||
| }) | ||
| if err != nil { | ||
| t.Fatalf("error: %+v", err) | ||
| } | ||
| }) | ||
| } | ||
| } |
This file contains bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters