Code with indicators.

Geographie-cites · Apr 19, 2016 · 2ba2380 · 2ba2380
1 parent f557217
commit 2ba2380
Show file tree

Hide file tree

Showing 7 changed files with 106 additions and 58 deletions.
diff --git a/src/main/scala/fr/geocites/micmac/Test.scala b/src/main/scala/fr/geocites/micmac/Test.scala
@@ -43,8 +43,8 @@ object Test extends App {
   val planeCapacity = 80
   val planeSpeed = 0.5
 
-  val airportIntegrator = integrator(1, 0.01)
-  val planeIntegrator = integrator(1, 0.01)
+  val airportIntegrator = integrator(beta / populationByNode, 0.01)
+  val planeIntegrator = integrator(beta / populationByNode, 0.01)
 
   def sir(s: Double, i: Double, r: Double) =
     SIR(s = s, i = i, r = r,  alpha = alpha, beta = beta)
@@ -62,7 +62,14 @@ object Test extends App {
   val airports =
     randomAirports[Context](
       territory,
-      Airport(_,_,_,  sir(s = populationByNode - 1, i = 1, r = 0), populationToFly / nodes),
+      (index, x, y, infected) =>
+        Airport(
+          index,
+          x,
+          y,
+          sir(s = populationByNode - infected, infected, r = 0),
+          populationToFly
+        ),
       nodes
     )
 
@@ -86,42 +93,24 @@ object Test extends App {
     updateAirportSIR andThen updatePlaneSIR andThen
      dynamic.planeDepartures[M](
        planeCapacity = planeCapacity,
+       populationToFly = populationToFly,
        destination = dynamic.randomDestination[M],
        buildSIR = sir) andThen
       dynamic.planeArrivals[M](planeSpeed) andThen
-      updateMaxStepI[M] andThen updateStep[M]
+      updateMaxStepI[M] andThen
+      updateStep[M] andThen
+      updateInfectedNodes[M]
   }
 
-
   val initialise = airports >>= world
 
   val simulation =
     initialise >>=
       context.run(
         evolve,
-        or(endOfEpidemy[Context], stopAfter[Context](10000))
-      )
-
-  val initialState = SimulationState(0, new Random(42), None)
-
-//  println(populationToFly)
-//  println((initialise >>= simulation).eval(initialState))
-//
-  println(simulation.run(initialState))
-//
-//
-//  val network  = (Kleisli { _: Any => airports } andThen world).apply().eval(initialState)
-//
-//  println(evolve(network))
-
-
-
-
-//
-//  println(.map(evolution))
-
-  //println(airports.map { world.run }.eval(initialState))
-
+        or(endOfEpidemy[Context], stopAfter[Context](20000))
+      ) andThen indicators[Context]
 
+  println(simulation.run(initialState(nodes, new Random(42))))
 
 }
diff --git a/src/main/scala/fr/geocites/micmac/context.scala b/src/main/scala/fr/geocites/micmac/context.scala
@@ -27,8 +27,21 @@ import monocle.state.all._
 
 object context {
 
-  @Lenses case class MicMacState(network: Network, flyingPlanes: Vector[Plane])
-  @Lenses case class SimulationState(step: Long, rng: Random, maxIStep: Option[MaxIStep])
+  def initialState(airports: Int, rng: Random) = {
+    SimulationState(
+      step = 0,
+      rng = rng,
+      maxIStep = None,
+      infectionStep = Vector.fill(airports)(None)
+    )
+  }
+
+  @Lenses case class SimulationState(
+    step: Long,
+    rng: Random,
+    maxIStep: Option[MaxIStep] = None,
+    infectionStep: Vector[Option[Long]]
+  )
 
   type Context[X] = State[SimulationState, X]
 
@@ -39,6 +52,7 @@ object context {
 
   implicit def sObservable = new Observable[Context] {
     override def maxIStep = SimulationState.maxIStep.mod
+    override def infectionStep = SimulationState.infectionStep.mod
   }
 
   implicit def sRNG = new RNG[Context] {

diff --git a/src/main/scala/fr/geocites/micmac/dynamic.scala b/src/main/scala/fr/geocites/micmac/dynamic.scala
@@ -17,9 +17,7 @@
   */
 package fr.geocites.micmac
 
-import fr.geocites.micmac.context.MicMacState
 import fr.geocites.micmac.sir.Integrator
-
 import scala.concurrent.duration.Duration
 import scalaz._
 import Scalaz._
@@ -30,17 +28,25 @@ import sir._
 
 object dynamic {
 
-  def populationToFly(sir: SIR, integrator: Integrator, epidemyDuration: Duration, mobilityRate: Double, epsilon: Double, nbAirports: Int): Long = {
+  def populationToFly(
+    sir: SIR,
+    integrator: Integrator,
+    epidemyDuration: Duration,
+    mobilityRate: Double,
+    epsilon: Double,
+    nbAirports: Int): Double = {
+
     def steps(sir: SIR, step: Int = 0): Int = {
       val newSir = integrator(sir)
       if (newSir.i < epsilon) step
       else steps(newSir, step + 1)
     }
 
     val nbSteps = steps(sir)
-    def dt = epidemyDuration.toHours / nbSteps
-    def totalPopulationToFly = (sir.s + sir.i + sir.r) * mobilityRate * dt
-    (totalPopulationToFly / nbSteps).toLong
+    def dt = epidemyDuration.toHours.toDouble
+    def totalPopulationToFly = sir.total * mobilityRate * dt
+
+    (totalPopulationToFly / nbSteps / nbAirports)
   }
 
   def updateSIRs[T](integrator: Integrator, sir: monocle.Traversal[T, SIR])(t: T) =
@@ -57,6 +63,7 @@ object dynamic {
 
   def planeDepartures[M[_]: Monad: RNG: Step](
     planeCapacity: Int,
+    populationToFly: Double,
     destination: (Airport, Network) => M[Airport],
     buildSIR: (Double, Double, Double) => SIR) = Kleisli[M, MicMacState, MicMacState] { modelState =>
     def departures(state: MicMacState): M[Vector[(Airport, List[Plane])]] =
@@ -82,8 +89,11 @@ object dynamic {
     } yield {
       val (newAirports, departedPlanes) = v.unzip
 
-      ((MicMacState.network composeLens Network.airports).set(newAirports) andThen
-        MicMacState.flyingPlanes.modify(_ ++ departedPlanes.flatten)) (modelState)
+      def setNewAirports = (MicMacState.network composeLens Network.airports) set newAirports
+      def updatePopulationToFly = (MicMacState.network composeTraversal Network.airportsTraversal composeLens Airport.populationToFly) modify (_ + populationToFly)
+      def addDepartingPlanes = MicMacState.flyingPlanes modify(_ ++ departedPlanes.flatten)
+
+      (setNewAirports andThen updatePopulationToFly andThen addDepartingPlanes) (modelState)
     }
   }
 
@@ -141,10 +151,10 @@ object dynamic {
 
       def fillPlanes(airport: Airport, planes: List[Plane] = List.empty): M[(Airport, List[Plane])] =
         // Plane should be full to leave
-        if(Airport.populationToFly.get(airport) < planeCapacity) (airport, planes).point[M]
+        if (Airport.populationToFly.get(airport) < planeCapacity) (airport, planes).point[M]
         else
           for {
-            plane <-  buildPlane(0, 0, 0)
+            plane <- buildPlane(0, 0, 0)
             filled <- fillPlane(airport, plane)
             (newAirport, newPlane) = filled
             res <- fillPlanes(newAirport, newPlane :: planes)

diff --git a/src/main/scala/fr/geocites/micmac/network.scala b/src/main/scala/fr/geocites/micmac/network.scala
@@ -23,21 +23,26 @@ import Scalaz._
 
 object network {
 
-  def randomAirports[M[_]: Monad: RNG](
+  def randomAirports[M[_]: Monad](
     territory: Territory,
-    buildAirport: (Int, Double, Double) => Airport,
-    number: Int) = {
+    buildAirport: (Int, Double, Double, Int) => Airport,
+    number: Int)(implicit rng: RNG[M]) = {
 
-    def randomAirport(i: Int): M[Airport] =
+
+    def randomAirport(i: Int, infectedIndex: Int): M[Airport] =
       for {
         rng <- implicitly[RNG[M]].rng
       } yield {
         val x = (rng.nextDouble * territory.length)
         val y = (rng.nextDouble * territory.width)
-        buildAirport(i, x, y)
+        def infected = if(infectedIndex == i) 1 else 0
+        buildAirport(i, x, y, infected)
       }
 
-    (0 until number).toVector.traverseU(randomAirport)
+    for {
+      infectedIndex <- rng.rng.map(_.nextInt(number))
+      airports <- (0 until number).toVector.traverseU(randomAirport(_, infectedIndex))
+    } yield airports
   }
 
   def randomNetwork[M[_]: Monad](edges: Int)(implicit mRNG: RNG[M]) = Kleisli[M, Vector[Airport], Network] {  airports: Vector[Airport] =>

diff --git a/src/main/scala/fr/geocites/micmac/observable.scala b/src/main/scala/fr/geocites/micmac/observable.scala
@@ -22,14 +22,14 @@ import fr.geocites.micmac.context._
 import scalaz._
 import Scalaz._
 
+
 object observable {
 
-  def totalI(state: MicMacState) =
-    state.flyingPlanes.map(_.sir.i).sum +
-      state.network.airports.map(_.sir.i).sum
+  def total(on: SIR => Double)(state: MicMacState) =
+    state.flyingPlanes.map(p => on(p.sir)).sum + state.network.airports.map(a => on(a.sir)).sum
 
   def updateMaxStepI[M[_]: Monad](implicit obs: Observable[M], step: Step[M]) = Kleisli[M, MicMacState, MicMacState] { state: MicMacState =>
-    val totalIValue = totalI(state)
+    val totalIValue = total(SIR.i.get)(state)
 
     def update(step: Long) =
       obs.maxIStep.modify {
@@ -44,4 +44,34 @@ object observable {
     } yield state
   }
 
+  case class Indicators(maxIStep: MaxIStep, infectedRatio: Double)
+
+  def indicators[M[_]: Monad](implicit obs: Observable[M]) = Kleisli[M, MicMacState, Option[Indicators]] { state =>
+    def infectedRatio = {
+      def totalPopulation = total(_.total)(state)
+      def recovered = total(_.r)(state)
+      recovered / totalPopulation
+    }
+
+    for {
+      maxIStep <- obs.maxIStep.get
+    } yield maxIStep.map(mis => Indicators(mis, infectedRatio))
+
+  }
+
+  def updateInfectedNodes[M[_]: Monad](implicit obs: Observable[M], step: Step[M]) = Kleisli[M, MicMacState, MicMacState] { state =>
+    def update(infectionSteps: Vector[Option[Long]], currentStep: Long) =
+      infectionSteps zip (MicMacState.network composeLens Network.airports).get(state) map {
+        case (is, a) =>
+          def infected = a.sir.i > 0
+          is orElse (if(infected) Some(currentStep) else None)
+      }
+
+    for {
+      s <- step.step.get
+      _ <- obs.infectionStep.modify(update(_, s))
+    } yield state
+  }
+
+
 }
diff --git a/src/main/scala/fr/geocites/micmac/package.scala b/src/main/scala/fr/geocites/micmac/package.scala
@@ -49,13 +49,9 @@ package object micmac {
 
   @typeclass trait Observable[M[_]] {
     def maxIStep: Field[M, Option[MaxIStep]]
+    def infectionStep: Field[M, Vector[Option[Long]]]
   }
 
-  /*trait ModelState[M[_], S] {
-    def get: M[S]
-    def set(s: S): M[Unit]
-  }*/
-
   case class Territory(length: Int, width: Int) {
     def area = length * width
   }
@@ -96,4 +92,10 @@ package object micmac {
 
   @Lenses case class Plane(capacity: Int, sir: SIR, origin: Int, destination: Int, departureStep: Long)
 
+  object MicMacState {
+    def flyingPlaneTraversal = (MicMacState.flyingPlanes composeTraversal Each.each)
+  }
+
+  @Lenses case class MicMacState(network: Network, flyingPlanes: Vector[Plane])
+
 }
diff --git a/src/main/scala/fr/geocites/micmac/stop.scala b/src/main/scala/fr/geocites/micmac/stop.scala
@@ -17,8 +17,6 @@
   */
 package fr.geocites.micmac
 
-import fr.geocites.micmac.context.MicMacState
-
 import scalaz._
 import Scalaz._
 
@@ -32,7 +30,7 @@ object stop {
     def finished(maxIStep: Option[MaxIStep], step: Long) =
       maxIStep match {
         case Some(maxIStep) =>
-          if(step >= maxIStep.step + 1) totalI(state) < 1.0
+          if(step >= maxIStep.step + 1) total(_.i)(state) < 1.0
           else false
         case None => false
       }