go-hep.org/x/hep@v0.38.1/fads/efficiency.go (about)

     1  // Copyright ©2017 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 fads
     6  
     7  import (
     8  	"math/rand/v2"
     9  	"reflect"
    10  	"sync"
    11  
    12  	"go-hep.org/x/hep/fwk"
    13  	"gonum.org/v1/gonum/stat/distuv"
    14  )
    15  
    16  type Efficiency struct {
    17  	fwk.TaskBase
    18  
    19  	input  string
    20  	output string
    21  
    22  	eff  func(pt, eta float64) float64
    23  	seed uint64
    24  	dist distuv.Uniform
    25  	dmu  sync.Mutex
    26  }
    27  
    28  func (tsk *Efficiency) Configure(ctx fwk.Context) error {
    29  	var err error
    30  	err = tsk.DeclInPort(tsk.input, reflect.TypeOf([]Candidate{}))
    31  	if err != nil {
    32  		return err
    33  	}
    34  
    35  	err = tsk.DeclOutPort(tsk.output, reflect.TypeOf([]Candidate{}))
    36  	if err != nil {
    37  		return err
    38  	}
    39  
    40  	return err
    41  }
    42  
    43  func (tsk *Efficiency) StartTask(ctx fwk.Context) error {
    44  	var err error
    45  	src := rand.New(rand.NewPCG(tsk.seed, tsk.seed))
    46  	tsk.dist = distuv.Uniform{Min: 0, Max: 1, Src: src}
    47  	return err
    48  }
    49  
    50  func (tsk *Efficiency) StopTask(ctx fwk.Context) error {
    51  	var err error
    52  
    53  	return err
    54  }
    55  
    56  func (tsk *Efficiency) Process(ctx fwk.Context) error {
    57  	var err error
    58  	store := ctx.Store()
    59  	msg := ctx.Msg()
    60  
    61  	v, err := store.Get(tsk.input)
    62  	if err != nil {
    63  		return err
    64  	}
    65  
    66  	input := v.([]Candidate)
    67  	msg.Debugf(">>> input: %v\n", len(input))
    68  
    69  	output := make([]Candidate, 0, len(input))
    70  	defer func() {
    71  		err = store.Put(tsk.output, output)
    72  	}()
    73  
    74  	for i := range input {
    75  		cand := &input[i]
    76  		eta := cand.Pos.Eta()
    77  		pt := cand.Mom.Pt()
    78  
    79  		// apply efficiency
    80  		tsk.dmu.Lock()
    81  		eff := tsk.dist.Rand()
    82  		tsk.dmu.Unlock()
    83  		max := tsk.eff(pt, eta)
    84  		if eff > max {
    85  			continue
    86  		}
    87  
    88  		output = append(output, *cand)
    89  	}
    90  
    91  	msg.Debugf(">>> output: %v\n", len(output))
    92  
    93  	return err
    94  }
    95  
    96  func newEfficiency(typ, name string, mgr fwk.App) (fwk.Component, error) {
    97  	var err error
    98  	tsk := &Efficiency{
    99  		TaskBase: fwk.NewTask(typ, name, mgr),
   100  		input:    "InputParticles",
   101  		output:   "OutputParticles",
   102  		eff:      func(x, y float64) float64 { return 1 },
   103  		seed:     1234,
   104  	}
   105  	err = tsk.DeclProp("Input", &tsk.input)
   106  	if err != nil {
   107  		return nil, err
   108  	}
   109  
   110  	err = tsk.DeclProp("Output", &tsk.output)
   111  	if err != nil {
   112  		return nil, err
   113  	}
   114  
   115  	err = tsk.DeclProp("Eff", &tsk.eff)
   116  	if err != nil {
   117  		return nil, err
   118  	}
   119  
   120  	err = tsk.DeclProp("Seed", &tsk.seed)
   121  	if err != nil {
   122  		return nil, err
   123  	}
   124  
   125  	return tsk, err
   126  }
   127  
   128  func init() {
   129  	fwk.Register(reflect.TypeOf(Efficiency{}), newEfficiency)
   130  }