Arbitrage with Bellman Ford

bellman_ford.go

package bellmanford

import (
	"math"
)

// Graph represents a graph consisting of edges and vertices
type Graph struct {
	edges    []*Edge
	vertices []uint
}

// Edge represents a weighted line between two nodes
type Edge struct {
	From, To uint
	Weight   float64
}

// NewEdge returns a pointer to a new Edge
func NewEdge(from, to uint, weight float64) *Edge {
	return &Edge{From: from, To: to, Weight: weight}
}

// NewGraph returns a graph consisting of given edges and vertices (vertices must count from 0 upwards)
func NewGraph(edges []*Edge, vertices []uint) *Graph {
	return &Graph{edges: edges, vertices: vertices}
}

// FindArbitrageLoop returns either an arbitrage loop or a nil map
func (g *Graph) FindArbitrageLoop(source uint) []uint {
	predecessors, distances := g.BellmanFord(source)
	return g.FindNegativeWeightCycle(predecessors, distances, source)
}

// BellmanFord determines the shortest path and returns the predecessors and distances
func (g *Graph) BellmanFord(source uint) ([]uint, []float64) {
	size := len(g.vertices)
	distances := make([]float64, size)
	predecessors := make([]uint, size)
	for _, v := range g.vertices {
		distances[v] = math.MaxFloat64
	}
	distances[source] = 0

	for i, changes := 0, 0; i < size-1; i, changes = i+1, 0 {
		for _, edge := range g.edges {
			if newDist := distances[edge.From] + edge.Weight; newDist < distances[edge.To] {
				distances[edge.To] = newDist
				predecessors[edge.To] = edge.From
				changes++
			}
		}
		if changes == 0 {
			break
		}
	}
	return predecessors, distances
}

// FindNegativeWeightCycle finds a negative weight cycle from predecessors and a source
func (g *Graph) FindNegativeWeightCycle(predecessors []uint, distances []float64, source uint) []uint {
	for _, edge := range g.edges {
		if distances[edge.From]+edge.Weight < distances[edge.To] {
			return arbitrageLoop(predecessors, source)
		}
	}
	return nil
}

func arbitrageLoop(predecessors []uint, source uint) []uint {
	size := len(predecessors)
	loop := make([]uint, size)
	loop[0] = source

	exists := make([]bool, size)
	exists[source] = true

	indices := make([]uint, size)

	var index, next uint
	for index, next = 1, source; ; index++ {
		next = predecessors[next]
		loop[index] = next
		if exists[next] {
			return loop[indices[next] : index+1]
		}
		indices[next] = index
		exists[next] = true
	}
}

bellman_ford_test.go

package bellmanford

import (
	"math"
	"testing"
)

func _newGraph() *Graph {
	return &Graph{
		vertices: []uint{0, 1, 2, 3, 4, 5},
		edges: []*Edge{
			&Edge{To: 1, From: 0, Weight: -math.Log(1.380)},
			&Edge{To: 2, From: 1, Weight: -math.Log(3.08)},
			&Edge{To: 3, From: 2, Weight: -math.Log(15.120)},
			&Edge{To: 4, From: 3, Weight: -math.Log(0.012)},
			&Edge{To: 0, From: 4, Weight: -math.Log(1.30)},
			&Edge{To: 5, From: 4, Weight: -math.Log(0.57)}},
	}
}

func BenchmarkNewGraph(b *testing.B) {
	for n := 0; n < b.N; n++ {
		_ = _newGraph()
	}
}

func BenchmarkBellmanFord(b *testing.B) {
	g := _newGraph()
	var source uint = 1
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		g.BellmanFord(source)
	}
}

func BenchmarkFindNegativeWeightCycle(b *testing.B) {
	g := _newGraph()
	var source uint = 1
	predecessors, distances := g.BellmanFord(source)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		g.FindNegativeWeightCycle(predecessors, distances, source)
	}
}

func BenchmarkArbitrageLoop(b *testing.B) {
	g := _newGraph()
	var source uint = 1
	predecessors, _ := g.BellmanFord(source)
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		arbitrageLoop(predecessors, source)
	}
}

func BenchmarkFindArbitrageLoop(b *testing.B) {
	g := _newGraph()
	var source uint = 1
	b.ResetTimer()
	for n := 0; n < b.N; n++ {
		g.FindArbitrageLoop(source)
	}
}

func TestFullSequence(t *testing.T) {
	results := map[uint][]uint{
		0: []uint{0, 4, 3, 2, 1, 0},
		1: []uint{1, 0, 4, 3, 2, 1},
		2: []uint{2, 1, 0, 4, 3, 2},
		3: []uint{3, 2, 1, 0, 4, 3},
		4: []uint{4, 3, 2, 1, 0, 4},
	}
	for source, res := range results {
		g := _newGraph()
		loop := g.FindArbitrageLoop(source)
		if len(loop) != len(res) {
			t.Fatalf("loops have different lengths (%d != %d)", loop, res)
		}
		for i, v := range loop {
			if res[i] != v {
				t.Fatalf("incorrect arbitrage loop (%v != %v; source is %d)\n", loop, res, source)
			}
		}
	}
}

func _getPythonGraph() *Graph {
	return &Graph{edges: []*Edge{
		&Edge{From: 0, To: 1, Weight: -4.582438665548869},
		&Edge{From: 0, To: 2, Weight: 0.2981813979749493},
		&Edge{From: 0, To: 3, Weight: 4.838300943835368},
		&Edge{From: 1, To: 0, Weight: 4.585249918552961},
		&Edge{From: 1, To: 2, Weight: 4.836396313495658},
		&Edge{From: 1, To: 3, Weight: 9.375215015166416},
		&Edge{From: 2, To: 0, Weight: -0.3751523503802663},
		&Edge{From: 2, To: 1, Weight: -5.004605689846387},
		&Edge{From: 2, To: 3, Weight: 4.362953685292599},
		&Edge{From: 3, To: 0, Weight: -4.6488526240960395},
		&Edge{From: 3, To: 1, Weight: -9.277409346383422},
		&Edge{From: 3, To: 2, Weight: -4.344533438603351},
	}, vertices: []uint{0, 1, 2, 3}}
}

func TestWithPythonDataSource(t *testing.T) {
	results := map[uint][]uint{
		0: []uint{2, 1, 2},
		1: []uint{1, 2, 1},
		2: []uint{2, 1, 2},
		3: []uint{2, 1, 2},
	}
	for source, res := range results {
		g := _getPythonGraph()
		loop := g.FindArbitrageLoop(source)
		if len(loop) != len(res) {
			t.Fatalf("loops have different lengths (%d != %d)", loop, res)
		}
		for i, v := range loop {
			if res[i] != v {
				t.Fatalf("incorrect arbitrage loop (%v != %v; source is %d)\n", loop, res, source)
			}
		}
	}
}