rebalance: add ValueMap

pkg/fixedpoint/value_map.go

@@ -0,0 +1,135 @@
+package fixedpoint
+
+type ValueMap map[string]Value
+
+func (m ValueMap) Eq(n ValueMap) bool {
+	if len(m) != len(n) {
+		return false
+	}
+
+	for m_k, m_v := range m {
+		n_v, ok := n[m_k]
+		if !ok {
+			return false
+		}
+
+		if !m_v.Eq(n_v) {
+			return false
+		}
+	}
+
+	return true
+}
+
+func (m ValueMap) Add(n ValueMap) ValueMap {
+	if len(m) != len(n) {
+		panic("unequal length")
+	}
+
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Add(n[k])
+	}
+
+	return o
+}
+
+func (m ValueMap) Sub(n ValueMap) ValueMap {
+	if len(m) != len(n) {
+		panic("unequal length")
+	}
+
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Sub(n[k])
+	}
+
+	return o
+}
+
+func (m ValueMap) Mul(n ValueMap) ValueMap {
+	if len(m) != len(n) {
+		panic("unequal length")
+	}
+
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Mul(n[k])
+	}
+
+	return o
+}
+
+func (m ValueMap) Div(n ValueMap) ValueMap {
+	if len(m) != len(n) {
+		panic("unequal length")
+	}
+
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Div(n[k])
+	}
+
+	return o
+}
+
+func (m ValueMap) AddScalar(x Value) ValueMap {
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Add(x)
+	}
+
+	return o
+}
+
+func (m ValueMap) SubScalar(x Value) ValueMap {
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Sub(x)
+	}
+
+	return o
+}
+
+func (m ValueMap) MulScalar(x Value) ValueMap {
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Mul(x)
+	}
+
+	return o
+}
+
+func (m ValueMap) DivScalar(x Value) ValueMap {
+	o := ValueMap{}
+
+	for k, v := range m {
+		o[k] = v.Div(x)
+	}
+
+	return o
+}
+
+func (m ValueMap) Sum() Value {
+	var sum Value
+	for _, v := range m {
+		sum = sum.Add(v)
+	}
+	return sum
+}
+
+func (m ValueMap) Normalize() ValueMap {
+	sum := m.Sum()
+	if sum.Eq(Zero) {
+		panic("zero sum")
+	}
+
+	return m.DivScalar(sum)
+}

pkg/fixedpoint/value_map_test.go

@@ -0,0 +1,123 @@
+package fixedpoint
+
+import (
+	"testing"
+
+	"github.com/stretchr/testify/assert"
+)
+
+func Test_ValueMap_Eq(t *testing.T) {
+	m1 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	m2 := ValueMap{}
+
+	m3 := ValueMap{"A": NewFromFloat(5.0)}
+
+	m4 := ValueMap{
+		"A": NewFromFloat(6.0),
+		"B": NewFromFloat(7.0),
+	}
+
+	m5 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	assert.True(t, m1.Eq(m1))
+	assert.False(t, m1.Eq(m2))
+	assert.False(t, m1.Eq(m3))
+	assert.False(t, m1.Eq(m4))
+	assert.True(t, m1.Eq(m5))
+}
+
+func Test_ValueMap_Add(t *testing.T) {
+	m1 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	m2 := ValueMap{
+		"A": NewFromFloat(5.0),
+		"B": NewFromFloat(6.0),
+	}
+
+	m3 := ValueMap{
+		"A": NewFromFloat(8.0),
+		"B": NewFromFloat(10.0),
+	}
+
+	m4 := ValueMap{"A": NewFromFloat(8.0)}
+
+	assert.Equal(t, m3, m1.Add(m2))
+	assert.Panics(t, func() { m1.Add(m4) })
+}
+
+func Test_ValueMap_AddScalar(t *testing.T) {
+	x := NewFromFloat(5.0)
+
+	m1 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	m2 := ValueMap{
+		"A": NewFromFloat(3.0).Add(x),
+		"B": NewFromFloat(4.0).Add(x),
+	}
+
+	assert.Equal(t, m2, m1.AddScalar(x))
+}
+
+func Test_ValueMap_DivScalar(t *testing.T) {
+	x := NewFromFloat(5.0)
+
+	m1 := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	m2 := ValueMap{
+		"A": NewFromFloat(3.0).Div(x),
+		"B": NewFromFloat(4.0).Div(x),
+	}
+
+	assert.Equal(t, m2, m1.DivScalar(x))
+}
+
+func Test_ValueMap_Sum(t *testing.T) {
+	m := ValueMap{
+		"A": NewFromFloat(3.0),
+		"B": NewFromFloat(4.0),
+	}
+
+	assert.Equal(t, NewFromFloat(7.0), m.Sum())
+}
+
+func Test_ValueMap_Normalize(t *testing.T) {
+	a := NewFromFloat(3.0)
+	b := NewFromFloat(4.0)
+
+	m := ValueMap{
+		"A": a,
+		"B": b,
+	}
+
+	n := ValueMap{
+		"A": a.Div(a.Add(b)),
+		"B": b.Div(a.Add(b)),
+	}
+
+	assert.True(t, m.Normalize().Eq(n))
+}
+
+func Test_ValueMap_Normalize_zero_sum(t *testing.T) {
+	m := ValueMap{
+		"A": Zero,
+		"B": Zero,
+	}
+
+	assert.Panics(t, func() { m.Normalize() })
+}