go-hep.org/x/hep@v0.38.1/fads/energy_scale.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  	"reflect"
     9  
    10  	"go-hep.org/x/hep/fmom"
    11  	"go-hep.org/x/hep/fwk"
    12  )
    13  
    14  type EnergyScale struct {
    15  	fwk.TaskBase
    16  
    17  	input  string
    18  	output string
    19  
    20  	scale func(pt, eta float64) float64
    21  }
    22  
    23  func (tsk *EnergyScale) Configure(ctx fwk.Context) error {
    24  	var err error
    25  
    26  	err = tsk.DeclInPort(tsk.input, reflect.TypeOf([]Candidate{}))
    27  	if err != nil {
    28  		return err
    29  	}
    30  
    31  	err = tsk.DeclOutPort(tsk.output, reflect.TypeOf([]Candidate{}))
    32  	if err != nil {
    33  		return err
    34  	}
    35  
    36  	return err
    37  }
    38  
    39  func (tsk *EnergyScale) StartTask(ctx fwk.Context) error {
    40  	var err error
    41  
    42  	return err
    43  }
    44  
    45  func (tsk *EnergyScale) StopTask(ctx fwk.Context) error {
    46  	var err error
    47  
    48  	return err
    49  }
    50  
    51  func (tsk *EnergyScale) Process(ctx fwk.Context) error {
    52  	var err error
    53  
    54  	store := ctx.Store()
    55  	msg := ctx.Msg()
    56  
    57  	v, err := store.Get(tsk.input)
    58  	if err != nil {
    59  		return err
    60  	}
    61  
    62  	input := v.([]Candidate)
    63  	msg.Debugf(">>> input: %v\n", len(input))
    64  
    65  	output := make([]Candidate, 0, len(input))
    66  	defer func() {
    67  		err = store.Put(tsk.output, output)
    68  	}()
    69  
    70  	for i := range input {
    71  		cand := input[i].Clone()
    72  
    73  		eta := cand.Mom.Eta()
    74  		pt := cand.Mom.Pt()
    75  
    76  		// get new scale
    77  		scale := tsk.scale(pt, eta)
    78  		if scale > 0 {
    79  			mom := fmom.Scale(scale, &cand.Mom)
    80  			cand.Mom.Set(mom)
    81  		}
    82  
    83  		output = append(output, *cand)
    84  	}
    85  
    86  	msg.Debugf(">>> scaled: %v\n", len(output))
    87  
    88  	return err
    89  }
    90  
    91  func newEnergyScale(typ, name string, mgr fwk.App) (fwk.Component, error) {
    92  	var err error
    93  
    94  	tsk := &EnergyScale{
    95  		TaskBase: fwk.NewTask(typ, name, mgr),
    96  		input:    "InputParticles",
    97  		output:   "OutputParticles",
    98  		scale:    func(pt, eta float64) float64 { return 0.0 },
    99  	}
   100  
   101  	err = tsk.DeclProp("Input", &tsk.input)
   102  	if err != nil {
   103  		return nil, err
   104  	}
   105  
   106  	err = tsk.DeclProp("Output", &tsk.output)
   107  	if err != nil {
   108  		return nil, err
   109  	}
   110  
   111  	err = tsk.DeclProp("Scale", &tsk.scale)
   112  	if err != nil {
   113  		return nil, err
   114  	}
   115  
   116  	return tsk, err
   117  }
   118  
   119  func init() {
   120  	fwk.Register(reflect.TypeOf(EnergyScale{}), newEnergyScale)
   121  }