Rewrite the horrible Graph API

epfldata · Feb 13, 2024 · 051a2c7 · 051a2c7
1 parent d1b3ad4
commit 051a2c7
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 25 deletions.
diff --git a/example/src/main/scala/example/cellularAutomata/gameOfLifeRPC/Example.scala b/example/src/main/scala/example/cellularAutomata/gameOfLifeRPC/Example.scala
@@ -3,7 +3,7 @@ package gameOfLifeRPC
 
 import scala.collection.mutable.{Map => MutMap}
 
-import cloudcity.lib.Graph.GenerateGraph.Torus2DGraph
+import cloudcity.lib.Graph.Torus2DGraph
 
 object MainInit {
     val liftedMain = meta.classLifting.liteLift {

diff --git a/example/src/main/scala/example/epidemic/Example.scala b/example/src/main/scala/example/epidemic/Example.scala
@@ -3,8 +3,7 @@ package epidemic
 
 object MainInit {
     val liftedMain = meta.classLifting.liteLift {
-        def apply(population: Int, p: Double): IndexedSeq[Actor] = {
-            val socialGraph = cloudcity.lib.Graph.GenerateGraph.ErdosRenyiGraph(population, p)
+        def apply(socialGraph: cloudcity.lib.Graph.Graph): IndexedSeq[Actor] = {
             val citizens = socialGraph.map(i => {
                 val person = new Person(Random.nextInt(90) + 10)
                 person.id = i._1

diff --git a/library/src/main/scala/Graph/GenerateGraph.scala b/library/src/main/scala/Graph/GenerateGraph.scala
@@ -1,13 +1,19 @@
 package cloudcity.lib
-package Graph.GenerateGraph
+package Graph
 
 import scala.util.Random
 import scala.math.floor
 
-trait GenerateGraph
+trait Graph {
+    val g: Map[Long, Iterable[Long]]
+}
+
+object Graph{
+    implicit def toGraph(g: Graph): Map[Long, Iterable[Long]] = g.g
+}
 
-object ErdosRenyiGraph extends GenerateGraph {
-    def apply(totalVertices: Int, edgeProb: Double, startingIndex: Int = 0): Map[Long, Iterable[Long]] = {
+case class ErdosRenyiGraph(totalVertices: Int, edgeProb: Double, startingIndex: Int = 0) extends Graph {
+    val g: Map[Long, Iterable[Long]] = {
         val nodes = Range(startingIndex, totalVertices+startingIndex)
         Range(0, totalVertices).map(i => {
             (i.toLong + startingIndex, nodes.filter(n => {
@@ -17,8 +23,8 @@ object ErdosRenyiGraph extends GenerateGraph {
     }
 }
 
-object SBMGraph extends GenerateGraph {
-    def apply(totalVertices: Int, p: Double, q: Double, blocks: Int, startingIndex: Int = 0): Map[Long, Iterable[Long]] = {
+case class SBMGraph(totalVertices: Int, p: Double, q: Double, blocks: Int, startingIndex: Int = 0) extends Graph {
+    val g: Map[Long, Iterable[Long]] = {
         val verticesPerBlock: Int = floor(totalVertices / blocks).toInt
         var graph: Map[Long, Iterable[Long]] = Map[Long, Iterable[Long]]()
         // The edge probability between two vertices in the same block is p
@@ -46,8 +52,8 @@ object SBMGraph extends GenerateGraph {
     }
 }
 
-object Torus2DGraph extends GenerateGraph {
-    def apply(width: Int, height: Int, startingIndex: Int = 0): Map[Long, IndexedSeq[Long]] = {
+case class Torus2DGraph(width: Int, height: Int, startingIndex: Int = 0) extends Graph {
+    val g: Map[Long, IndexedSeq[Long]] = {
         Range(0, width * height).map(index => {
             val x = index % width
             val y = index / width
@@ -64,8 +70,8 @@ object Torus2DGraph extends GenerateGraph {
     }
 }
 
-object BipartiteGraph extends GenerateGraph {
-    def apply(set1Size: Int, set2Size: Int, startingIndex: Int = 0): Map[Long, Iterable[Long]] = {        
+case class BipartiteGraph(set1Size: Int, set2Size: Int, startingIndex: Int = 0) extends Graph {
+    val g: Map[Long, Iterable[Long]] = {        
         (Range(startingIndex, set1Size + startingIndex).map(i => {
             (i.toLong, Range(set1Size + startingIndex, set2Size).map(_.toLong))
         }) ++ Range(startingIndex + set1Size, startingIndex + set1Size + set2Size).map(i => {

diff --git a/library/src/test/scala/GraphSpec.scala b/library/src/test/scala/GraphSpec.scala
@@ -7,59 +7,59 @@ import org.scalatest.FlatSpec
 class GraphSpec extends FlatSpec {
 
   "Each cell in a 2D grid" should "have eight neighbors" in {
-    val graph1 = GenerateGraph.Torus2DGraph(3, 3)
+    val graph1 = Torus2DGraph(3, 3)
     graph1.foreach(i => assert(i._2.size == 8))
     assert(graph1.size == 9)
 
-    val graph2 = GenerateGraph.Torus2DGraph(100, 100)
+    val graph2 = Torus2DGraph(100, 100)
     graph2.foreach(i => assert(i._2.size == 8))
     assert(graph2.size == 10000)
   }
 
   "Each cell in a 2D grid with offsets" should "have eight neighbors" in {
-    val graph1 = GenerateGraph.Torus2DGraph(3, 3, 13)
+    val graph1 = Torus2DGraph(3, 3, 13)
     graph1.foreach(i => assert(i._2.size == 8))
     assert(graph1.size == 9)
 
-    val graph2 = GenerateGraph.Torus2DGraph(100, 100, 11)
+    val graph2 = Torus2DGraph(100, 100, 11)
     graph2.foreach(i => assert(i._2.size == 8))
     assert(graph2.size == 10000)
   }
 
   "SBM graph" should "generate a line for each node" in {
     val totalNodes = 36
-    val graph = GenerateGraph.SBMGraph(totalNodes, 0.1, 0, 5)
+    val graph = SBMGraph(totalNodes, 0.1, 0, 5)
     assert(graph.size == totalNodes)
     val totalNodes2 = 1000
-    val graph2 = GenerateGraph.SBMGraph(totalNodes2, 0.1, 0, 5)
+    val graph2 = SBMGraph(totalNodes2, 0.1, 0, 5)
     assert(graph2.size == totalNodes2)
   }
 
   "SBM graph with offsets" should "generate a line for each node" in {
     val totalNodes = 36
-    val graph = GenerateGraph.SBMGraph(totalNodes, 0.1, 0, 5, 10)
+    val graph = SBMGraph(totalNodes, 0.1, 0, 5, 10)
     assert(graph.size == totalNodes)
     val totalNodes2 = 1000
-    val graph2 = GenerateGraph.SBMGraph(totalNodes2, 0.1, 0, 5, 1)
+    val graph2 = SBMGraph(totalNodes2, 0.1, 0, 5, 1)
     assert(graph2.size == totalNodes2)
   }
 
 
   "ERM graph" should "generate a line for each node" in {
     val totalNodes = 71
-    val graph = GenerateGraph.ErdosRenyiGraph(totalNodes, 0.1)
+    val graph = ErdosRenyiGraph(totalNodes, 0.1)
     assert(graph.size == totalNodes)
     val totalNodes2 = 1000
-    val graph2 = GenerateGraph.ErdosRenyiGraph(totalNodes2, 0.1)
+    val graph2 = ErdosRenyiGraph(totalNodes2, 0.1)
     assert(graph2.size == totalNodes2)
   }
 
   "ERM graph with offsets" should "generate a line for each node" in {
     val totalNodes = 71
-    val graph = GenerateGraph.ErdosRenyiGraph(totalNodes, 0.1, 1)
+    val graph = ErdosRenyiGraph(totalNodes, 0.1, 1)
     assert(graph.size == totalNodes)
     val totalNodes2 = 1000
-    val graph2 = GenerateGraph.ErdosRenyiGraph(totalNodes2, 0.1, 11)
+    val graph2 = ErdosRenyiGraph(totalNodes2, 0.1, 11)
     assert(graph2.size == totalNodes2)
   }
 }