fmom: add SetPtEtaPhi{E,M} to PxPyPzE

Author: sbinet

fmom/pxpypze.go

@@ -200,6 +200,32 @@ func (p4 *PxPyPzE) Set(p P4) {
 	p4.P4.T = p.E()
 }
 
+func (p4 *PxPyPzE) SetPtEtaPhiM(pt, eta, phi, m float64) {
+	sin, cos := math.Sincos(phi)
+	pt = math.Abs(pt)
+	p4.P4.X = pt * cos
+	p4.P4.Y = pt * sin
+	p4.P4.Z = pt * math.Sinh(eta)
+
+	p2 := p4.P4.X*p4.P4.X + p4.P4.Y*p4.P4.Y + p4.P4.Z*p4.P4.Z
+	m2 := m * m
+	switch {
+	case m >= 0:
+		p4.P4.T = math.Sqrt(p2 + m2)
+	default:
+		p4.P4.T = math.Sqrt(math.Max(0, p2-m2))
+	}
+}
+
+func (p4 *PxPyPzE) SetPtEtaPhiE(pt, eta, phi, e float64) {
+	sin, cos := math.Sincos(phi)
+	pt = math.Abs(pt)
+	p4.P4.X = pt * cos
+	p4.P4.Y = pt * sin
+	p4.P4.Z = pt * math.Sinh(eta)
+	p4.P4.T = e
+}
+
 var (
 	_ P4           = (*PxPyPzE)(nil)
 	_ fmt.Stringer = (*PxPyPzE)(nil)

fmom/pxpypze_test.go

@@ -5,6 +5,7 @@
 package fmom_test
 
 import (
+	"math"
 	"testing"
 
 	"go-hep.org/x/hep/fmom"
@@ -139,4 +140,32 @@ func TestPxPyPzE(t *testing.T) {
 			}
 		})
 	}
+
+	t.Run("set-PtEtaPhiM", func(t *testing.T) {
+		p1 := fmom.NewPxPyPzE(10, 20, 30, 40)
+		p1.SetPtEtaPhiM(100, 1.5, 1/3.*math.Pi, 10)
+		want := fmom.NewPxPyPzE(
+			49.99999999999999,
+			86.60254037844388,
+			212.9279455094818,
+			235.45341360636257,
+		)
+		if got, want := p1, want; got != want {
+			t.Fatalf("invalid p4:\ngot= %v\nwant=%v", got, want)
+		}
+	})
+
+	t.Run("set-PtEtaPhiE", func(t *testing.T) {
+		p1 := fmom.NewPxPyPzE(10, 20, 30, 40)
+		p1.SetPtEtaPhiE(100, 1.5, 1/3.*math.Pi, 10)
+		want := fmom.NewPxPyPzE(
+			49.99999999999999,
+			86.60254037844388,
+			212.9279455094818,
+			10,
+		)
+		if got, want := p1, want; got != want {
+			t.Fatalf("invalid p4:\ngot= %v\nwant=%v", got, want)
+		}
+	})
 }