equivalent numpy.random.choice function in Golang

Question

Basically what the title says is, is there any function / library in Golang that does numpy.random.choice?

What I want to do is, I want to shuffle a slice, based on the probability for each element.

The closest thing I got is the rand.Shuffle, but AFAIK it is not considering the probabilities of each element, and I couldn't find the equivalent function in gonum package as well.

Referrence:

numpy.random.choice docs --> https://numpy.org/doc/stable/reference/random/generated/numpy.random.choice.html

How about this? https://go.dev/play/p/8YvCXQzQHzm – spike014 Jul 20 '23 at 09:36 — spike014, Jul 20 '23 at 09:36
https://go.dev/play/p/99AYtrvSL25 can try this? – Prashant Luhar Jul 21 '23 at 17:23 — Prashant Luhar, Jul 21 '23 at 17:23

fizzie · Answer 1 · 2023-07-29T19:26:32.170

The corresponding Gonum routine can be found in the sampling package, gonum.org/v1/gonum/stat/sampleuv. Specifically, the Weighted type provides sampling without replacement from with given (non-uniform) probabilities, equivalent to numpy.random.choice with replace=False.

Assuming that by "shuffle a slice, based on the probability for each element" you mean sampling without replacement using the given weights until no elements remain, and returning the items in the order they were selected, the following code accomplishes this:

package main

import (
    "fmt"

    "gonum.org/v1/gonum/stat/sampleuv"
)

func shuffleWithWeights[S ~[]E, E any](items S, w []float64) []E {
    samp := sampleuv.NewWeighted(w, nil)
    n := len(items)
    result := make([]E, n)
    for i := 0; i < n; i++ {
        idx, _ := samp.Take()
        result[i] = items[idx]
    }
    return result
}

func main() {
    items := []string{"a", "b", "c", "d"}
    w := []float64{1, 2, 4, 8}
    for i := 0; i < 10; i++ {
        fmt.Println(shuffleWithWeights(items, w))
    }
}

A sample run of the code produced:

[d c b a]
[b d c a]
[d b c a]
[d b c a]
[d c a b]
[c d b a]
[d c a b]
[a d c b]
[c d b a]
[d c a b]

As expected, d is selected first with a probability that's approximately 8/15. (Properly rigorous validation of the results left for the reader.)

For sampling with replacement (replace=True in numpy.random.choice), you can just draw samples from the gonum.org/v1/gonum/stat/distuv.Categorical distribution. For example:

package main

import (
    "fmt"

    "gonum.org/v1/gonum/stat/distuv"
)

func drawN[S ~[]E, E any](items S, w []float64, n int) []E {
    dist := distuv.NewCategorical(w, nil)
    result := make([]E, n)
    for i := 0; i < n; i++ {
        result[i] = items[int(dist.Rand())]
    }
    return result
}

func main() {
    items := []string{"a", "b", "c", "d"}
    w := []float64{1, 2, 4, 8}
    fmt.Println(drawN(items, w, 50))
}

This might produce:

[d c c d b d c b d d d c d b d d d c d d c d d c d d a d a d c c c d d b d d b c d d c d d c d a d d]

There may be more convenience methods in Gonum to accomplish these things; I'm not very familiar with its API.

equivalent numpy.random.choice function in Golang

1 Answers1