go-hep.org/x/hep@v0.38.1/fads/merger.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  	"reflect"
    10  
    11  	"go-hep.org/x/hep/fmom"
    12  	"go-hep.org/x/hep/fwk"
    13  )
    14  
    15  type Merger struct {
    16  	fwk.TaskBase
    17  
    18  	inputs []string
    19  	output string
    20  	outene string
    21  	outmom string
    22  }
    23  
    24  func (tsk *Merger) Configure(ctx fwk.Context) error {
    25  	var err error
    26  
    27  	for _, input := range tsk.inputs {
    28  		err = tsk.DeclInPort(input, reflect.TypeOf([]Candidate{}))
    29  		if err != nil {
    30  			return err
    31  		}
    32  	}
    33  
    34  	err = tsk.DeclOutPort(tsk.output, reflect.TypeOf([]Candidate{}))
    35  	if err != nil {
    36  		return err
    37  	}
    38  
    39  	err = tsk.DeclOutPort(tsk.outene, reflect.TypeOf(Candidate{}))
    40  	if err != nil {
    41  		return err
    42  	}
    43  
    44  	err = tsk.DeclOutPort(tsk.outmom, reflect.TypeOf(Candidate{}))
    45  	if err != nil {
    46  		return err
    47  	}
    48  
    49  	return err
    50  }
    51  
    52  func (tsk *Merger) StartTask(ctx fwk.Context) error {
    53  	var err error
    54  
    55  	return err
    56  }
    57  
    58  func (tsk *Merger) StopTask(ctx fwk.Context) error {
    59  	var err error
    60  
    61  	return err
    62  }
    63  
    64  func (tsk *Merger) Process(ctx fwk.Context) error {
    65  	var err error
    66  
    67  	store := ctx.Store()
    68  	msg := ctx.Msg()
    69  
    70  	output := make([]Candidate, 0)
    71  	defer func() {
    72  		err = store.Put(tsk.output, output)
    73  	}()
    74  
    75  	sumpt := 0.0
    76  	sumene := 0.0
    77  	p4 := fmom.NewPxPyPzE(0, 0, 0, 0)
    78  	var mom fmom.P4 = &p4
    79  
    80  	for _, k := range tsk.inputs {
    81  		v, err := store.Get(k)
    82  		if err != nil {
    83  			return err
    84  		}
    85  		input := v.([]Candidate)
    86  		msg.Debugf(">>> input[%s]: %v\n", k, len(input))
    87  
    88  		for i := range input {
    89  			cand := &input[i]
    90  			cmom := cand.Mom
    91  			mom = fmom.IAdd(mom, &cmom)
    92  			sumpt += cmom.Pt()
    93  			sumene += cmom.E()
    94  
    95  			output = append(output, *cand)
    96  		}
    97  	}
    98  
    99  	var cmom Candidate
   100  	cmom.Mom.Set(mom)
   101  
   102  	err = store.Put(tsk.outmom, cmom)
   103  	if err != nil {
   104  		return err
   105  	}
   106  
   107  	var cene Candidate
   108  
   109  	eta := 0.0
   110  	phi := 0.0
   111  	px := sumpt * math.Cos(phi)
   112  	py := sumpt * math.Sin(phi)
   113  	pz := sumpt * math.Sinh(eta)
   114  
   115  	cene.Mom = fmom.NewPxPyPzE(px, py, pz, sumene)
   116  	err = store.Put(tsk.outene, cene)
   117  	if err != nil {
   118  		return err
   119  	}
   120  
   121  	msg.Debugf(">>> output: %v\n", len(output))
   122  	return err
   123  }
   124  
   125  func newMerger(typ, name string, mgr fwk.App) (fwk.Component, error) {
   126  	var err error
   127  
   128  	tsk := &Merger{
   129  		TaskBase: fwk.NewTask(typ, name, mgr),
   130  		inputs:   []string{},
   131  		output:   "candidates",
   132  		outene:   "energy",
   133  		outmom:   "momentum",
   134  	}
   135  
   136  	err = tsk.DeclProp("Inputs", &tsk.inputs)
   137  	if err != nil {
   138  		return nil, err
   139  	}
   140  
   141  	err = tsk.DeclProp("Output", &tsk.output)
   142  	if err != nil {
   143  		return nil, err
   144  	}
   145  
   146  	err = tsk.DeclProp("MomentumOutput", &tsk.outmom)
   147  	if err != nil {
   148  		return nil, err
   149  	}
   150  
   151  	err = tsk.DeclProp("EnergyOutput", &tsk.outene)
   152  	if err != nil {
   153  		return nil, err
   154  	}
   155  
   156  	return tsk, err
   157  }
   158  
   159  func init() {
   160  	fwk.Register(reflect.TypeOf(Merger{}), newMerger)
   161  }