Repository: gonum/gonum
Pull request #1867 (2023-06-13 19:35:36)

Fix Weighted panic
Please take a look.

Take is divided into 2 steps index search and after that Reweight with 0.
1. I added checks for heap over bounds during tree traversal.
2. I removed last and left because in each iteration of the loop we always make a move to either the left or right subtree and this branch never reached.
3. Instead of two implementations of reweight in the `Take` and `Reweight` functions, I left only one.
4. Instead of calling reset - O(n) when you check a negative weight, I do a weight calculation as in reset, but through sift up only the desired nodes for O(log(n)). It also removed the `scalar.EqualWithinAbsOrRel` check, since we are now summing up zeros instead of accumulating a float error.

https://github.com/gonum/gonum/issues/1866

codecov-commenter: ## [Codecov](https://app.codecov.io/gh/gonum/gonum/pull/1867?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gonum) Report
Patch coverage: **`88.23`**% and project coverage change: **`-0.53`** :warning:
> Comparison is base [(`258a51e`)](https://app.codecov.io/gh/gonum/gonum/commit/258a51e069623132f86d7efd2188b05d2578a92e?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gonum) 73.93% compared to head [(`891a633`)](https://app.codecov.io/gh/gonum/gonum/pull/1867?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gonum) 73.40%.

:exclamation: Your organization is not using the GitHub App Integration. As a result you may experience degraded service beginning May 15th. Please [install the Github App Integration](https://github.com/apps/codecov) for your organization. [Read more](https://about.codecov.io/blog/codecov-is-updating-its-github-integration/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gonum).

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #1867      +/-   ##
==========================================
- Coverage   73.93%   73.40%   -0.53%     
==========================================
  Files         526      510      -16     
  Lines       63112    61651    -1461     
==========================================
- Hits        46659    45257    -1402     
+ Misses      13884    13844      -40     
+ Partials     2569     2550      -19     
```


| [Impacted Files](https://app.codecov.io/gh/gonum/gonum/pull/1867?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gonum) | Coverage Δ | |
|---|---|---|
| [stat/sampleuv/weighted.go](https://app.codecov.io/gh/gonum/gonum/pull/1867?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gonum#diff-c3RhdC9zYW1wbGV1di93ZWlnaHRlZC5nbw==) | `88.46% <88.23%> (+5.48%)` | :arrow_up: |

... and [24 files with indirect coverage changes](https://app.codecov.io/gh/gonum/gonum/pull/1867/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gonum)


</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/gonum/gonum/pull/1867?src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gonum).   
:loudspeaker: Do you have feedback about the report comment? [Let us know in this issue](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=gonum).

kortschak: I'll take a look on the weekend. From an initial look, I'd like more detail on the fix in the description and less code movement.
kortschak: > can I change the golang version for staticcheck to 1.18?

Where do you want to do this?
khasanovbi: > Where do you want to do this?

[Here](https://github.com/gonum/gonum/blob/master/.github/workflows/staticcheck.yml#L16). I want to set it to v1 to use v1.3.0

Patch:
diff --git a/stat/sampleuv/weighted.go b/stat/sampleuv/weighted.go
index f08d5453..9cf04dc0 100644
--- a/stat/sampleuv/weighted.go
+++ b/stat/sampleuv/weighted.go
@@ -4,11 +4,7 @@
 
 package sampleuv
 
-import (
-	"golang.org/x/exp/rand"
-
-	"gonum.org/v1/gonum/floats/scalar"
-)
+import "golang.org/x/exp/rand"
 
 // Weighted provides sampling without replacement from a collection of items with
 // non-uniform probability.
@@ -58,66 +54,83 @@ func (s Weighted) Len() int { return len(s.weights) }
 // to the weight of the item. The weight of the item is then set to zero.
 // Take returns false if there are no items remaining.
 func (s Weighted) Take() (idx int, ok bool) {
-	const small = 1e-12
-	if scalar.EqualWithinAbsOrRel(s.heap[0], 0, small, small) {
+	if s.heap[0] == 0 {
 		return -1, false
 	}
 
 	var r float64
 	if s.rnd == nil {
-		r = s.heap[0] * rand.Float64()
+		r = rand.Float64()
 	} else {
-		r = s.heap[0] * s.rnd.Float64()
+		r = s.rnd.Float64()
 	}
-	i := 1
-	last := -1
-	left := len(s.weights)
+
+	r *= s.heap[0]
+	i := 0
 	for {
-		if r -= s.weights[i-1]; r <= 0 {
-			break // Fall within item i-1.
-		}
-		i <<= 1 // Move to left child.
-		if d := s.heap[i-1]; r > d {
-			r -= d
-			// If enough r to pass left child
-			// move to right child state will
-			// be caught at break above.
-			i++
+		r -= s.weights[i]
+		if r < 0 {
+			break // Fall within item i.
 		}
-		if i == last || left < 0 {
-			// No progression.
-			return -1, false
+
+		li := i*2 + 1 // Move to left child.
+		// Left node should exist, because r is non-negative,
+		// but there could be floating point errors, so we
+		// check index explicitly.
+		if li >= len(s.heap) {
+			break
 		}
-		last = i
-		left--
-	}
 
-	w, idx := s.weights[i-1], i-1
-
-	s.weights[i-1] = 0
-	for i > 0 {
-		s.heap[i-1] -= w
-		// The following condition is necessary to
-		// handle floating point error. If we see
-		// a heap value below zero, we know we need
-		// to rebuild it.
-		if s.heap[i-1] < 0 {
-			s.reset()
-			return idx, true
+		i = li
+
+		d := s.heap[i]
+		if r >= d {
+			// If there is enough r to pass left child try to
+			// move to the right child.
+			r -= d
+			ri := i + 1
+
+			if ri >= len(s.heap) {
+				break
+			}
+
+			i = ri
 		}
-		i >>= 1
 	}
 
-	return idx, true
+	s.Reweight(i, 0)
+
+	return i, true
 }
 
 // Reweight sets the weight of item idx to w.
 func (s Weighted) Reweight(idx int, w float64) {
-	w, s.weights[idx] = s.weights[idx]-w, w
-	idx++
-	for idx > 0 {
-		s.heap[idx-1] -= w
-		idx >>= 1
+	s.weights[idx] = w
+
+	// We want to keep the heap state here consistent
+	// with the result of a reset call. So we sum
+	// weights in the same order, since floating point
+	// addition is not associative.
+	for {
+		w = s.weights[idx]
+
+		ri := idx*2 + 2
+		if ri < len(s.heap) {
+			w += s.heap[ri]
+		}
+
+		li := ri - 1
+		if li < len(s.heap) {
+			w += s.heap[li]
+		}
+
+		s.heap[idx] = w
+
+		if idx == 0 {
+			break
+		}
+
+		idx = (idx - 1) / 2
 	}
 }
 


Test patch:
diff --git a/stat/sampleuv/weighted_test.go b/stat/sampleuv/weighted_test.go
index 3f6f58ee..a3d98a12 100644
--- a/stat/sampleuv/weighted_test.go
+++ b/stat/sampleuv/weighted_test.go
@@ -266,3 +266,72 @@ func TestWeightedNoResample(t *testing.T) {
 		}
 	}
 }
+
+var issue1866Tests = []struct {
+	weights []float64
+	seed    uint64
+}{
+	{weights: []float64{72000, 166.66666666666666}, seed: 202},
+	{weights: []float64{0.6666666666666666, 20412, 0}, seed: 0},
+}
+
+// See https://github.com/gonum/gonum/issues/1866
+func TestIssue1866(t *testing.T) {
+	for i, test := range issue1866Tests {
+		w := NewWeighted(test.weights, rand.NewSource(test.seed))
+		for j := 0; j < len(test.weights)+1; j++ {
+			if panicked(func() { w.Take() }) {
+				t.Errorf("unexpected panic for test %d iteration %d", i, j)
+			}
+		}
+	}
+}
+
+func panicked(fn func()) (yes bool) {
+	defer func() {
+		if recover() != nil {
+			yes = true
+		}
+	}()
+	fn()
+	return false
+}
+
+func FuzzWeightedTake(f *testing.F) {
+	f.Add(0.000001, 0.000000012)
+
+	f.Fuzz(func(t *testing.T, w1 float64, w2 float64) {
+		if w1 < 0 || w2 < 0 {
+			t.Skip()
+		}
+
+		weights := []float64{w1, w2}
+
+		weighted := NewWeighted(weights, rand.NewSource(0))
+
+		want := calculateNonZeroWeights(weights)
+		var got int
+		for {
+			_, ok := weighted.Take()
+			if !ok {
+				break
+			}
+			got++
+		}
+
+		if got != want {
+			t.Errorf("unexpected taken set size: got=%d, want=%d", got, want)
+		}
+	})
+}
+
+func calculateNonZeroWeights(weights []float64) int {
+	count := 0
+	for _, w := range weights {
+		if w != 0 {
+			count++
+		}
+	}
+
+	return count
+}