go-hep.org/x/hep@v0.38.1/fads/momentum_smearing.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"
     9  	"math/rand/v2"
    10  	"reflect"
    11  	"sync"
    12  
    13  	"go-hep.org/x/hep/fmom"
    14  	"go-hep.org/x/hep/fwk"
    15  	"gonum.org/v1/gonum/stat/distuv"
    16  )
    17  
    18  type MomentumSmearing struct {
    19  	fwk.TaskBase
    20  
    21  	input  string
    22  	output string
    23  
    24  	smear func(x, y float64) float64
    25  	seed  uint64
    26  	src   *rand.Rand
    27  	srcmu sync.Mutex
    28  }
    29  
    30  func (tsk *MomentumSmearing) Configure(ctx fwk.Context) error {
    31  	var err error
    32  
    33  	err = tsk.DeclInPort(tsk.input, reflect.TypeOf([]Candidate{}))
    34  	if err != nil {
    35  		return err
    36  	}
    37  
    38  	err = tsk.DeclOutPort(tsk.output, reflect.TypeOf([]Candidate{}))
    39  	if err != nil {
    40  		return err
    41  	}
    42  
    43  	return err
    44  }
    45  
    46  func (tsk *MomentumSmearing) StartTask(ctx fwk.Context) error {
    47  	var err error
    48  	tsk.src = rand.New(rand.NewPCG(tsk.seed, tsk.seed))
    49  	return err
    50  }
    51  
    52  func (tsk *MomentumSmearing) StopTask(ctx fwk.Context) error {
    53  	var err error
    54  
    55  	return err
    56  }
    57  
    58  func (tsk *MomentumSmearing) Process(ctx fwk.Context) error {
    59  	var err error
    60  	store := ctx.Store()
    61  	msg := ctx.Msg()
    62  
    63  	v, err := store.Get(tsk.input)
    64  	if err != nil {
    65  		return err
    66  	}
    67  
    68  	input := v.([]Candidate)
    69  	msg.Debugf(">>> input: %v\n", len(input))
    70  
    71  	output := make([]Candidate, 0, len(input))
    72  	defer func() {
    73  		err = store.Put(tsk.output, output)
    74  	}()
    75  
    76  	for i := range input {
    77  		cand := &input[i]
    78  		eta := cand.Pos.Eta()
    79  		pt := cand.Mom.Pt()
    80  
    81  		// apply smearing
    82  		tsk.srcmu.Lock()
    83  		smearPt := distuv.Normal{Mu: pt, Sigma: tsk.smear(pt, eta) * pt, Src: tsk.src}
    84  		pt = smearPt.Rand()
    85  		tsk.srcmu.Unlock()
    86  
    87  		if pt <= 0 {
    88  			continue
    89  		}
    90  
    91  		mother := cand
    92  		c := cand.Clone()
    93  		eta = cand.Mom.Eta()
    94  		phi := cand.Mom.Phi()
    95  
    96  		pxs := pt * math.Cos(phi)
    97  		pys := pt * math.Sin(phi)
    98  		pzs := pt * math.Sinh(eta)
    99  		es := pt * math.Cosh(eta)
   100  		c.Mom = fmom.NewPxPyPzE(pxs, pys, pzs, es)
   101  		c.Add(mother)
   102  
   103  		output = append(output, *c)
   104  	}
   105  
   106  	msg.Debugf(">>> smeared: %v\n", len(output))
   107  
   108  	return err
   109  }
   110  
   111  func init() {
   112  	fwk.Register(reflect.TypeOf(MomentumSmearing{}),
   113  		func(typ, name string, mgr fwk.App) (fwk.Component, error) {
   114  			var err error
   115  			tsk := &MomentumSmearing{
   116  				TaskBase: fwk.NewTask(typ, name, mgr),
   117  				input:    "InputParticles",
   118  				output:   "OutputParticles",
   119  				smear:    func(x, y float64) float64 { return 0 },
   120  				seed:     1234,
   121  			}
   122  
   123  			err = tsk.DeclProp("Input", &tsk.input)
   124  			if err != nil {
   125  				return nil, err
   126  			}
   127  
   128  			err = tsk.DeclProp("Output", &tsk.output)
   129  			if err != nil {
   130  				return nil, err
   131  			}
   132  
   133  			err = tsk.DeclProp("Resolution", &tsk.smear)
   134  			if err != nil {
   135  				return nil, err
   136  			}
   137  
   138  			err = tsk.DeclProp("Seed", &tsk.seed)
   139  			if err != nil {
   140  				return nil, err
   141  			}
   142  
   143  			return tsk, err
   144  		},
   145  	)
   146  }