go-hep.org/x/hep@v0.38.1/fads/energy_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 EnergySmearing struct {
    19  	fwk.TaskBase
    20  
    21  	input  string
    22  	output string
    23  
    24  	smear func(eta, ene float64) float64
    25  	seed  uint64
    26  	src   *rand.Rand
    27  	srcmu sync.Mutex
    28  }
    29  
    30  func (tsk *EnergySmearing) 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 *EnergySmearing) 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 *EnergySmearing) StopTask(ctx fwk.Context) error {
    53  	var err error
    54  
    55  	return err
    56  }
    57  
    58  func (tsk *EnergySmearing) 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  		ene := cand.Mom.E()
    80  
    81  		// apply smearing
    82  		tsk.srcmu.Lock()
    83  		smearEne := distuv.Normal{Mu: ene, Sigma: tsk.smear(eta, ene), Src: tsk.src}
    84  		ene = smearEne.Rand()
    85  		tsk.srcmu.Unlock()
    86  
    87  		if ene <= 0 {
    88  			continue
    89  		}
    90  
    91  		mother := cand
    92  		c := cand.Clone()
    93  		eta = cand.Mom.Eta()
    94  		phi := cand.Mom.Phi()
    95  		pt := ene / math.Cosh(eta)
    96  
    97  		pxs := pt * math.Cos(phi)
    98  		pys := pt * math.Sin(phi)
    99  		pzs := pt * math.Sinh(eta)
   100  
   101  		c.Mom = fmom.NewPxPyPzE(pxs, pys, pzs, ene)
   102  		c.Add(mother)
   103  
   104  		output = append(output, *c)
   105  	}
   106  
   107  	msg.Debugf(">>> smeared: %v\n", len(output))
   108  
   109  	return err
   110  }
   111  
   112  func init() {
   113  	fwk.Register(reflect.TypeOf(EnergySmearing{}),
   114  		func(typ, name string, mgr fwk.App) (fwk.Component, error) {
   115  			var err error
   116  			tsk := &EnergySmearing{
   117  				TaskBase: fwk.NewTask(typ, name, mgr),
   118  				input:    "InputParticles",
   119  				output:   "OutputParticles",
   120  				smear:    func(x, y float64) float64 { return 0 },
   121  				seed:     1234,
   122  			}
   123  
   124  			err = tsk.DeclProp("Input", &tsk.input)
   125  			if err != nil {
   126  				return nil, err
   127  			}
   128  
   129  			err = tsk.DeclProp("Output", &tsk.output)
   130  			if err != nil {
   131  				return nil, err
   132  			}
   133  
   134  			err = tsk.DeclProp("Resolution", &tsk.smear)
   135  			if err != nil {
   136  				return nil, err
   137  			}
   138  
   139  			err = tsk.DeclProp("Seed", &tsk.seed)
   140  			if err != nil {
   141  				return nil, err
   142  			}
   143  
   144  			return tsk, err
   145  		},
   146  	)
   147  }