Implement golden tests for Poisson distribution

Signed-off-by: Jakub Sztandera <kubuxu@protocol.ai>

chain/types/electionproof.go

@@ -67,23 +67,25 @@ func init() {
 	expDenoCoef = parse(deno)
 }
 
-// expneg accepts x in Q.256 format, in range of 0 to 5 and computes e^-x,
-// output is in Q.256 format
+// expneg accepts x in Q.256 format and computes e^-x.
+// It is most precise within [0, 1.725) range, where error is less than 3.4e-30.
+// Over the [0, 5) range its error is less than 4.6e-15.
+// Output is in Q.256 format.
 func expneg(x *big.Int) *big.Int {
 	// exp is approximated by rational function
-	// polynomials of the rational function are evaluated using Horners method
-
+	// polynomials of the rational function are evaluated using Horner's method
 	num := polyval(expNumCoef, x)   // Q.256
 	deno := polyval(expDenoCoef, x) // Q.256
 
 	num = num.Lsh(num, precision) // Q.512
 	return num.Div(num, deno)     // Q.512 / Q.256 => Q.256
 }
 
-// polyval evaluates a polynomial given by coefficients `p` in Q.256 format,
-// at point `x` in Q.256 format, output is in Q.256
-
+// polyval evaluates a polynomial given by coefficients `p` in Q.256 format
+// at point `x` in Q.256 format. Output is in Q.256.
+// Coefficients should be ordered from the highest order coefficient to the lowest.
 func polyval(p []*big.Int, x *big.Int) *big.Int {
+	// evaluation using Horner's method
 	res := new(big.Int).Set(p[0]) // Q.256
 	tmp := new(big.Int)           // big.Int.Mul doesn't like when input is reused as output
 	for _, c := range p[1:] {
@@ -108,24 +110,22 @@ var MaxWinCount = 3 * build.BlocksPerEpoch
 type poiss struct {
 	lam  *big.Int
 	pmf  *big.Int
-	tmp  *big.Int
 	icdf *big.Int
 
-	k    uint64
-	kBig *big.Int
+	tmp *big.Int // temporary variable for optmization
+
+	k uint64
 }
 
 // newPoiss starts poisson inverted CDF
 // lambda is in Q.256 format
 // returns (instance, `1-poisscdf(0, lambda)`)
 // CDF value returend is reused when calling `next`
 func newPoiss(lambda *big.Int) (*poiss, *big.Int) {
-	// e^-lambda
-	elam := expneg(lambda) // Q.256
 
 	// pmf(k) = (lambda^k)*(e^lambda) / k!
-	// k = 0 here so it similifies to just e^labda
-	pmf := new(big.Int).Set(elam) // Q.256
+	// k = 0 here, so it simplifies to just e^-lambda
+	pmf := expneg(lambda) // Q.256
 
 	// icdf(k) = 1 - ∑ᵏᵢ₌₀ pmf(i)
 	// icdf(0) = 1 - pmf(0)
@@ -134,34 +134,34 @@ func newPoiss(lambda *big.Int) (*poiss, *big.Int) {
 	icdf = icdf.Sub(icdf, pmf)       // Q.256
 
 	k := uint64(0)
-	kBig := new(big.Int).SetUint64(k) // Q.0
 
 	p := &poiss{
 		lam: lambda,
 		pmf: pmf,
 
-		tmp:  elam,
+		tmp:  new(big.Int),
 		icdf: icdf,
 
-		k:    k,
-		kBig: kBig,
+		k: k,
 	}
 
 	return p, icdf
 }
 
-// next computes next `k++, 1-poisscdf(k, lam)`
+// next computes `k++, 1-poisscdf(k, lam)`
 // return is in Q.256 format
 func (p *poiss) next() *big.Int {
-	// incrementally compute next pfm and icdf
+	// incrementally compute next pmf and icdf
+
 	// pmf(k) = (lambda^k)*(e^lambda) / k!
 	// so pmf(k) = pmf(k-1) * lambda / k
-
 	p.k++
-	p.kBig = p.kBig.SetUint64(p.k) // Q.0
+	p.tmp.SetUint64(p.k) // Q.0
 
 	// calculate pmf for k
-	p.pmf = p.pmf.Div(p.pmf, p.kBig)    // Q.256 / Q.0 => Q.256
+	p.pmf = p.pmf.Div(p.pmf, p.tmp) // Q.256 / Q.0 => Q.256
+	// we are using `tmp` as target for multiplication as using an input as output
+	// for Int.Mul causes allocations
 	p.tmp = p.tmp.Mul(p.pmf, p.lam)     // Q.256 * Q.256 => Q.512
 	p.pmf = p.pmf.Rsh(p.tmp, precision) // Q.512 >> 256 => Q.256
 
@@ -171,11 +171,9 @@ func (p *poiss) next() *big.Int {
 	return p.icdf
 }
 
-// poissStep performs a step in evaluation of Poisson distribution
-// k should be incremented after each evaluation step
-// tmp is scratch space
-// ouput is (pmf, icdf)
-
+// ComputeWinCount uses VRFProof to compute number of wins
+// The algorithm is based on Algorand's Sortition with Binomial distribution
+// replaced by Poisson distribution.
 func (ep *ElectionProof) ComputeWinCount(power BigInt, totalPower BigInt) uint64 {
 	h := blake2b.Sum256(ep.VRFProof)
 

chain/types/electionproof_test.go

@@ -1,12 +1,52 @@
 package types
 
 import (
+	"bytes"
 	"fmt"
 	"math/big"
 	"os"
 	"testing"
+
+	"github.com/xorcare/golden"
 )
 
+func TestPoissonFunction(t *testing.T) {
+	tests := []struct {
+		lambdaBase  uint64
+		lambdaShift uint
+	}{
+		{10, 10},      // 0.0097
+		{209714, 20},  // 0.19999885
+		{1036915, 20}, // 0.9888792038
+		{1706, 10},    // 1.6660
+		{2, 0},        // 2
+		{5242879, 20}, //4.9999990
+		{5, 0},        // 5
+	}
+
+	for _, test := range tests {
+		t.Run(fmt.Sprintf("lam-%d-%d", test.lambdaBase, test.lambdaShift), func(t *testing.T) {
+			test := test
+
+			b := &bytes.Buffer{}
+			b.WriteString("icdf\n")
+
+			lam := new(big.Int).SetUint64(test.lambdaBase)
+			lam = lam.Lsh(lam, precision-test.lambdaShift)
+			p, icdf := newPoiss(lam)
+
+			b.WriteString(icdf.String())
+			b.WriteRune('\n')
+
+			for i := 0; i < 15; i++ {
+				b.WriteString(p.next().String())
+				b.WriteRune('\n')
+			}
+			golden.Assert(t, []byte(b.String()))
+		})
+	}
+}
+
 func q256ToF(x *big.Int) float64 {
 	deno := big.NewInt(1)
 	deno = deno.Lsh(deno, 256)

chain/types/testdata/TestPoissonFunction/lam-10-10.golden

@@ -0,0 +1,17 @@
+icdf
+1125278653883954157340515998824199281259686237023612910954531537010133365568
+5485581780123676224984074899749002236148166431650497590243915223971292577
+17842170241691453912141677461647358103546946338181118738604570718548080
+43538699106868068808561503680708109912117284430088557923220935824303
+85008817354919776986513548953593326094752560579206896166859206326
+138328508214407784218646352510285887677303021571444942144212932
+192946940473264032619492808002293590266469557064324916486706
+235498963868775663540157747991701194152938740691242898919
+255504995002156513848360474242833468958493714151012216
+249505772801580009049072856702435230216536441494110
+232834101038081471620316286291642591265760137159
+11533551765583311484083562200606025547302501012
+11353456917542541496993529059255898133794641608
+11353321629946357024346977051186975261639061786
+11353321535577219060847642123725989684858819334
+11353321535515780819985988910882590605707269697

chain/types/testdata/TestPoissonFunction/lam-1036915-20.golden

@@ -0,0 +1,17 @@
+icdf
+72718197862321603787957847055188249408647877796280116987586082825356498974798
+30123322455004902866096392147720313525979066051167804713805891738863284666878
+9062729215708086547397523150443475826195010851218953471088727298493148732956
+2120601657722952965249404258310617497861606438105371150448777320082300909521
+404370264674629622846679825118378024155401061236248117877372450081489268106
+64941157977031699837280219244596630364570458903895153606060064374226923144
+8998760514291795062091202215054973561658841304177095664084807386711484044
+1095864305518189121413406860322570887693509822868614985675874891701983877
+118988091690998292615838041961723734187855286032350465668217096980526413
+11653361409102593830837021393444274321721937984503223867724441309766994
+1039253147016500070761515827557061937775360855994320103804816785188376
+85064880905559492020902336338153555957806293407638915883403930221828
+6433469833589351070633969345898155559842007456615136652210720692299
+452164666340411064711833496915674079929092772106312520794348036629
+29679788379529243880483477334855509673275091060271619731000625321
+1827354397264548824373139406487609602015834650190678153972930556

chain/types/testdata/TestPoissonFunction/lam-1706-10.golden

@@ -0,0 +1,17 @@
+icdf
+93907545465879218275260347624273708225148723352890415890955745471328115138024
+57447553596668785643406883388130520172829512611140657354486862128150346836988
+27076095525930017054568011324233899656590951319429188573619716140132147265911
+10209654292635635800479757502291310268341281539592025246744927385067352842651
+3184715634432458451975226003210729824895496226017269232182332263939291493510
+843984120585852874524302031056145794325474791455055607017530061469667926785
+194034907919461416983404196340045475941285897027461784665454449911533518447
+39345544525367691179167072173518251727638261160626536209449847049064587000
+7131182470884793691126479665205819536661348710819293305892247868965466268
+1167890189531948301087715470416213490892402026859749426392544722128541359
+174396377814374645286335419995210764907848995332895729280582459579342738
+23925812935985027317838050852967276757134553590636105055350959232313899
+3035286920153916620063269622769735189335169019323041991528942663951986
+358060554242868329017312833503133182524340437692927389149107908421231
+39467628723598770945019147503633808666969235107758896427288845018926
+4082242594961149456000070139366495398696105445630156285138889148853

chain/types/testdata/TestPoissonFunction/lam-2-0.golden

@@ -0,0 +1,17 @@
+icdf
+100121334043824876020967110392587568107053060743383522309741056272383359786578
+68779823656842237215759361160386888614619212898869438850308000801323820079862
+37438313269859598410551611928186209122185365054355355390874945330264280373146
+16543973011871172540413112440052422793896133158012633084586241682891253902002
+6096802882876959605343862695985529629751517209841271931441889859204740666430
+1917934831279274431316162798358772364093670830572727470184149129730135372202
+524978814080046039973596165816519942207722037483212649764902219905266940794
+126991380594552213875719985090162107383165239457636986787974531383875960392
+27494522223178757351250939908572648677026039951243071043742609253528215292
+5384109251762433679146707645997213408995106727599978656135515446784271938
+962026657479168944725861193482126355388920082871360178614096685435483268
+158011640336757174831161838479383254733249783829793182701111456099339874
+24009137479688546515378612645592737957304733989532016715613917876649310
+3393367809370296005258116363471119991774726321799529640921988919312302
+448257856467688789526616894596603139556153797837745773108856211121302
+55576529414007827429083632080000892593677461309507924067105183362502

chain/types/testdata/TestPoissonFunction/lam-209714-20.golden

@@ -0,0 +1,17 @@
+icdf
+20989436322611254574979389012727393136091537312055593688180077279147576363096
+2029014232696520721523332453453803636238058967796600089005757967894658844577
+132982872945592560215194824292206599698851338836977427578620974467337246296
+6581505574095040906286115477587166321018484661524779791358182693980706360
+261473369241451189291715224208035992983406808190620756138506410631265448
+8673527587881831627652027122245179992541421812500366598395022980519170
+246914417172755020716947338861248744263385116558896234139925344965714
+6155418508705151241339695503211918565434455355693098789916851059631
+136477982954924943101944815709613669983383501630000081908122878386
+2724514397229876659600187515561464490237868059330199615762951760
+49460332945698577716770256622916953121860884014393828193908634
+823264627479642334265449493920662550930201158945634649498769
+12651460568281449375957051928671501418597070452648092732548
+180560129118157986659345179422379755776654969450793746138
+2405427973892505908363489341734991530618943133521671600
+30045550760659097055494870911555042207154242485930695

chain/types/testdata/TestPoissonFunction/lam-5-0.golden

@@ -0,0 +1,17 @@
+icdf
+115011888277122352219705460287186602222085188670665840381425306072442950421456
+111110883476153136200377836679680074066161208695792222091263916395092054329056
+101358371473730096152058777660913753676351258758608176365860442201714814098056
+85104184803025029404860345962969886360001342196634766823521318546086080379726
+64786451464643695970862306340540052214563946494168004895597413976550163231816
+44468718126262362536864266718110218069126550791701242967673509407014246083906
+27537273677611251341865900366085356281262054372978941361070255599067648460651
+15443384785717600488295638686067597861358842645320154499210788593391507301186
+7884704228284068704814225136056498848919335315533412710548621714843919076521
+3685437251932106602880106497161443842008497910096333939069640115650814507266
+1585803763756125551913047177713916338553079207377794553330149316054262222641
+631424905494315983291656577965040200618797978869367559812198952601283911451
+233767047885228663032743828069675143146180800324189645846386301162542948456
+80821718035579693702392770417611659502866500883736602013381435224565655001
+26198385946419347512981678399017558201682822512146229215879697389573764486
+7990608583365898783177981059486191101288263054949438283379118111243134316

chain/types/testdata/TestPoissonFunction/lam-5242879-20.golden

@@ -0,0 +1,17 @@
+icdf
+115011887533064380050215202002871794783671664388541274939039184893270600703151
+111110879755863630144777547269452207577572562543679248972859798819416242535921
+101358362173007217989878395557465593373392010546833825352856759707561945139525
+85104169301821710755220628061805234978705652326202881500299286939503312327242
+64786432088141395502693562453332343133008530920262028792733819517798429528997
+44468698749761909885846743680008378684637277300179598543979674797636862943384
+27537257530529080605729671834720742022613060678716434560927439906969908575619
+15443373252088578337713549627194971124753642243467082552611819728291522561946
+7884697019766617162458477655388623015603913245952633503615157778326861227793
+3685433247200570820185174890022164698361097802621279879954268046011715772231
+1585801761390548420193995297013958176769177427873274589439743405082427147382
+631423995328231141553650878972791975288890578033071631113620389859816342901
+233766668649395912353875000216328335294367462954196928187468994385755448892
+80821572175657684536655650469711580587115741420816601432036447172153463116
+26198333853594769407667441209847842642730676168975225661522405972966431948
+7990591219092428810606628986022697269060757693982546042792694706871429282

go.mod

@@ -111,6 +111,7 @@ require (
 	github.com/whyrusleeping/cbor-gen v0.0.0-20200504204219-64967432584d
 	github.com/whyrusleeping/multiaddr-filter v0.0.0-20160516205228-e903e4adabd7
 	github.com/whyrusleeping/pubsub v0.0.0-20131020042734-02de8aa2db3d
+	github.com/xorcare/golden v0.6.1-0.20191112154924-b87f686d7542
 	go.opencensus.io v0.22.3
 	go.uber.org/dig v1.8.0 // indirect
 	go.uber.org/fx v1.9.0

go.sum

@@ -1373,6 +1373,8 @@ github.com/whyrusleeping/timecache v0.0.0-20160911033111-cfcb2f1abfee/go.mod h1:
 github.com/whyrusleeping/yamux v1.1.5/go.mod h1:E8LnQQ8HKx5KD29HZFUwM1PxCOdPRzGwur1mcYhXcD8=
 github.com/x-cray/logrus-prefixed-formatter v0.5.2/go.mod h1:2duySbKsL6M18s5GU7VPsoEPHyzalCE06qoARUCeBBE=
 github.com/xiang90/probing v0.0.0-20190116061207-43a291ad63a2/go.mod h1:UETIi67q53MR2AWcXfiuqkDkRtnGDLqkBTpCHuJHxtU=
+github.com/xorcare/golden v0.6.1-0.20191112154924-b87f686d7542 h1:oWgZJmC1DorFZDpfMfWg7xk29yEOZiXmo/wZl+utTI8=
+github.com/xorcare/golden v0.6.1-0.20191112154924-b87f686d7542/go.mod h1:7T39/ZMvaSEZlBPoYfVFmsBLmUl3uz9IuzWj/U6FtvQ=
 github.com/xordataexchange/crypt v0.0.3-0.20170626215501-b2862e3d0a77/go.mod h1:aYKd//L2LvnjZzWKhF00oedf4jCCReLcmhLdhm1A27Q=
 github.com/yuin/goldmark v1.1.25/go.mod h1:3hX8gzYuyVAZsxl0MRgGTJEmQBFcNTphYh9decYSb74=
 go.dedis.ch/fixbuf v1.0.3 h1:hGcV9Cd/znUxlusJ64eAlExS+5cJDIyTyEG+otu5wQs=