From 1da795789e97c44799186158d651b3bfce6efe09 Mon Sep 17 00:00:00 2001 From: btracey Date: Tue, 16 Sep 2014 23:46:39 -0700 Subject: [PATCH] Fixed check for zero probability value in JensenShannon and updated test. --- stat.go | 4 +++- stat_test.go | 4 ++++ 2 files changed, 7 insertions(+), 1 deletion(-) diff --git a/stat.go b/stat.go index e60e186f..936c79ba 100644 --- a/stat.go +++ b/stat.go @@ -402,9 +402,11 @@ func JensenShannon(p, q []float64) float64 { for i, v := range p { qi := q[i] m := 0.5 * (v + qi) - if m != 0 { + if v != 0 { // add kl from p to m js += 0.5 * v * (math.Log(v) - math.Log(m)) + } + if qi != 0 { // add kl from q to m js += 0.5 * qi * (math.Log(qi) - math.Log(m)) } diff --git a/stat_test.go b/stat_test.go index 12d5de91..13b13961 100644 --- a/stat_test.go +++ b/stat_test.go @@ -357,6 +357,10 @@ func TestJensenShannon(t *testing.T) { js1 := 0.5*KullbackLeibler(p, m) + 0.5*KullbackLeibler(q, m) js2 := JensenShannon(p, q) + if math.IsNaN(js2) { + t.Errorf("In case %v, JS distance is NaN", i) + } + if math.Abs(js1-js2) > 1e-14 { t.Errorf("JS mismatch case %v. Expected %v, found %v.", i, js1, js2) }