71 lines
1.9 KiB
Go
71 lines
1.9 KiB
Go
package simulation
|
|
|
|
import (
|
|
"fmt"
|
|
"math/rand"
|
|
)
|
|
|
|
// TransitionMatrix is _almost_ a left stochastic matrix.
|
|
// It is technically not one due to not normalizing the column values.
|
|
// In the future, if we want to find the steady state distribution,
|
|
// it will be quite easy to normalize these values to get a stochastic matrix.
|
|
// Floats aren't currently used as the default due to non-determinism across
|
|
// architectures
|
|
type TransitionMatrix struct {
|
|
weights [][]int
|
|
// total in each column
|
|
totals []int
|
|
n int
|
|
}
|
|
|
|
// CreateTransitionMatrix creates a transition matrix from the provided weights.
|
|
// TODO: Provide example usage
|
|
func CreateTransitionMatrix(weights [][]int) (TransitionMatrix, error) {
|
|
n := len(weights)
|
|
for i := 0; i < n; i++ {
|
|
if len(weights[i]) != n {
|
|
return TransitionMatrix{}, fmt.Errorf("Transition Matrix: Non-square matrix provided, error on row %d", i)
|
|
}
|
|
}
|
|
totals := make([]int, n)
|
|
for row := 0; row < n; row++ {
|
|
for col := 0; col < n; col++ {
|
|
totals[col] += weights[row][col]
|
|
}
|
|
}
|
|
return TransitionMatrix{weights, totals, n}, nil
|
|
}
|
|
|
|
// NextState returns the next state randomly chosen using r, and the weightings provided
|
|
// in the transition matrix.
|
|
func (t TransitionMatrix) NextState(r *rand.Rand, i int) int {
|
|
randNum := r.Intn(t.totals[i])
|
|
for row := 0; row < t.n; row++ {
|
|
if randNum < t.weights[row][i] {
|
|
return row
|
|
}
|
|
randNum -= t.weights[row][i]
|
|
}
|
|
// This line should never get executed
|
|
return -1
|
|
}
|
|
|
|
// GetMemberOfInitialState takes an initial array of weights, of size n.
|
|
// It returns a weighted random number in [0,n).
|
|
func GetMemberOfInitialState(r *rand.Rand, weights []int) int {
|
|
n := len(weights)
|
|
total := 0
|
|
for i := 0; i < n; i++ {
|
|
total += weights[i]
|
|
}
|
|
randNum := r.Intn(total)
|
|
for state := 0; state < n; state++ {
|
|
if randNum < weights[state] {
|
|
return state
|
|
}
|
|
randNum -= weights[state]
|
|
}
|
|
// This line should never get executed
|
|
return -1
|
|
}
|