Factors for Variable Elimination

ve/factor.go

@@ -0,0 +1,62 @@
+package ve
+
+import "fmt"
+
+type Factor struct {
+	id        int
+	data      []float64
+	variables variables
+}
+
+func (f *Factor) Index(indices []int) int {
+	return f.variables.Index(indices)
+}
+
+func (f *Factor) Outcomes(index int, indices []int) {
+	f.variables.Outcomes(index, indices)
+}
+
+func (f *Factor) RowIndex(indices []int) (int, int) {
+	if len(indices) != len(f.variables)-1 {
+		panic(fmt.Sprintf("factor with %d variables can't use %d row indices", len(f.variables), len(indices)))
+	}
+	cols := int(f.variables[len(f.variables)-1].outcomes)
+	idx := 0
+	stride := 1
+	curr := len(f.variables) - 2
+	for curr >= 0 {
+		stride *= int(f.variables[curr+1].outcomes)
+		idx += indices[curr] * stride
+		curr--
+	}
+	return idx, cols
+}
+
+func (f *Factor) Get(indices []int) float64 {
+	idx := f.Index(indices)
+	return f.data[idx]
+}
+
+func (f *Factor) Set(indices []int, value float64) {
+	idx := f.Index(indices)
+	f.data[idx] = value
+}
+
+func (f *Factor) GetRow(indices []int) []float64 {
+	idx, ln := f.RowIndex(indices)
+	return f.data[idx : idx+ln]
+}
+
+func (f *Factor) Normalize() {
+	sum := 0.0
+	for _, v := range f.data {
+		sum += v
+	}
+	if sum == 0 {
+		return
+	}
+
+	for i := range f.data {
+		f.data[i] /= sum
+	}
+}

ve/factor_test.go

@@ -0,0 +1,88 @@
+package ve
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func TestFactorIndex(t *testing.T) {
+	f := Factor{
+		variables: []Variable{
+			{id: 0, outcomes: 3},
+			{id: 1, outcomes: 2},
+		},
+		data: make([]float64, 6),
+	}
+
+	assert.Equal(t, 0, f.Index([]int{0, 0}))
+	assert.Equal(t, 1, f.Index([]int{0, 1}))
+	assert.Equal(t, 2, f.Index([]int{1, 0}))
+	assert.Equal(t, 5, f.Index([]int{2, 1}))
+}
+
+func TestFactorRowIndex(t *testing.T) {
+	f := Factor{
+		variables: []Variable{
+			{id: 0, outcomes: 3},
+			{id: 1, outcomes: 2},
+			{id: 2, outcomes: 4},
+		},
+		data: make([]float64, 24),
+	}
+
+	idx, ln := f.RowIndex([]int{0, 0})
+	assert.Equal(t, 0, idx)
+	assert.Equal(t, 4, ln)
+
+	idx, ln = f.RowIndex([]int{0, 1})
+	assert.Equal(t, 4, idx)
+	assert.Equal(t, 4, ln)
+
+	idx, ln = f.RowIndex([]int{1, 0})
+	assert.Equal(t, 8, idx)
+	assert.Equal(t, 4, ln)
+
+	idx, ln = f.RowIndex([]int{1, 1})
+	assert.Equal(t, 12, idx)
+	assert.Equal(t, 4, ln)
+}
+
+func TestFactorOutcomes(t *testing.T) {
+	f := Factor{
+		variables: []Variable{
+			{id: 0, outcomes: 3},
+			{id: 1, outcomes: 2},
+		},
+		data: make([]float64, 6),
+	}
+
+	result := []int{0, 0}
+
+	f.Outcomes(0, result)
+	assert.Equal(t, []int{0, 0}, result)
+
+	f.Outcomes(1, result)
+	assert.Equal(t, []int{0, 1}, result)
+
+	f.Outcomes(2, result)
+	assert.Equal(t, []int{1, 0}, result)
+
+	f.Outcomes(5, result)
+	assert.Equal(t, []int{2, 1}, result)
+}
+
+func TestFactorNormalize(t *testing.T) {
+	f := Factor{
+		variables: []Variable{
+			{id: 0, outcomes: 3},
+		},
+		data: []float64{
+			2, 1, 1,
+		},
+	}
+
+	f.Normalize()
+
+	assert.Equal(t, []float64{0.5, 0.25, 0.25}, f.data)
+}