models

Author: bmtgoncalves

models/SEIIR.yaml

@@ -0,0 +1,20 @@
+# Epidemic Model with 5 compartments and 6 transitions:
+Name: SEIIR
+
+Parameters:
+  rate : 0.166667
+  beta : 0.222222
+  epsilon_a : 0.160000
+  epsilon_s : 0.240000
+  mu : 0.100000
+
+
+Transitions:
+  - S + Ia = E rate
+  - S + Is = E beta
+  - E -> Ia epsilon_a
+  - E -> Is epsilon_s
+  - Ia -> R mu
+  - Is -> R mu
+
+# R0=2.00

models/SEIIRD.yaml

@@ -0,0 +1,23 @@
+# Epidemic Model with 6 compartments and 7 transitions:
+Name: SEIIRD
+
+Parameters:
+  rbeta : 0.166667
+  beta : 0.222222
+  epsilon_a : 0.160000
+  epsilon_s : 0.240000
+  mu : 0.100000
+  mu_nd : 0.090000
+  mu_d : 0.010000
+
+
+Transitions:
+  - S + Ia = E rbeta
+  - S + Is = E beta
+  - E -> Ia epsilon_a 
+  - E -> Is epsilon_s
+  - Ia -> R mu
+  - Is -> R mu_nd
+  - Is -> D mu_d
+
+# R0=2.00

models/SEIR.yaml

@@ -0,0 +1,15 @@
+# Epidemic Model with 4 compartments and 3 transitions:
+Name: SEIR
+
+Parameters:
+  beta : 0.200000
+  epsilon : 0.400000
+  mu : 0.100000
+
+
+Transitions:
+  - S + I = E beta
+  - E -> I epsilon 
+  - I -> R mu
+
+# R0=2.00

models/SEIRS.yaml

@@ -0,0 +1,15 @@
+# Epidemic Model with 4 compartments and 4 transitions:
+Name: SEIRS
+
+Parameters:
+  beta : 0.200000
+  epsilon : 0.400000
+  mu : 0.100000
+  rho : 0.300000
+
+
+Transitions:
+  - S + I = E beta
+  - E -> I epsilon 
+  - I -> R mu
+  - R -> S rho

models/SI.yaml

@@ -0,0 +1,8 @@
+# Epidemic Model with 2 compartments and 1 transitions:
+
+Parameters:
+  beta : 0.200000
+
+
+Transitions:
+  - S + I = I beta

models/model_list.txt

@@ -0,0 +1,6 @@
+SEIIR.yaml
+SEIIRD.yaml
+SEIR.yaml
+SEIRS.yaml
+SI.yaml
+SIR.yaml

poetry.lock

@@ -21,6 +21,7 @@ dependencies = [
  "scipy>=1.7",
  "tqdm>=4",
  "pyyaml (>=6.0.2,<7.0.0)",
+ "seaborn (>=0.13.2,<0.14.0)",
 ]
 [project.urls]
 Homepage = "https://github.com/DataForScience/epidemik"

pyproject.toml

@@ -9,6 +9,8 @@
 import time
 import os
 import re
+from urllib.parse import urlparse
+from urllib.request import urlretrieve
 
 import networkx as nx
 import numpy as np
@@ -21,7 +23,7 @@
 
 from typing import Union
 
-from .utils import *
+from . import utils
 
 
 class EpiModel(object):
@@ -325,7 +327,7 @@ def plot(
                 ax = plt.gca()
 
             for comp in self.values_.columns:
-                (self.values_[comp] / N).plot(c=epi_colors[comp[0]], **kwargs)
+                (self.values_[comp] / N).plot(c=utils.EPI_COLORS[comp[0]], **kwargs)
 
             ax.legend(self.values_.columns)
             ax.set_xlabel("Time")
@@ -340,7 +342,7 @@ def plot(
             return ax
         except Exception as e:
             print(e)
-            raise NotInitialized("You must call integrate() or simulate() first")
+            raise utils.NotInitialized("You must call integrate() or simulate() first")
 
     def __getattr__(self, name: str) -> pd.Series:
         """
@@ -670,7 +672,7 @@ def draw_model(self, ax: Union[plt.Axes, None] = None, show: bool = True) -> Non
                 orig_pos = pos[node[3:]]
                 pos[node] = [orig_pos[0], orig_pos[1] - 1]
             else:
-                node_colors.append(epi_colors[node[0]])
+                node_colors.append(utils.EPI_COLORS[node[0]])
 
         edge_labels = {}
 
@@ -821,3 +823,27 @@ def load_model(filename: str) -> None:
                 model.name = data[key]
 
         return model
+
+
+    def download_model(filename: str, repo: Union[str, None] = None) -> None:
+        """
+        Download model from offical repository
+        """
+        if repo is None:
+            repo = utils.OFFICIAL_REPO
+
+        parsed_repo = urlparse(repo)
+
+        if parsed_repo.netloc == 'github.com':
+            repo = repo.replace('github.com', 'raw.githubusercontent.com')
+            remote_path = repo + os.path.join('refs/heads/models/models/', filename)
+
+        if not os.path.exists(utils.LOCAL_DIRECTORY):
+            os.makedirs(utils.LOCAL_DIRECTORY)
+
+        local_path = os.path.join(utils.LOCAL_DIRECTORY, filename)
+
+        if not os.path.exists(local_path):
+            urlretrieve(remote_path, local_path)
+
+        return EpiModel.load_model(local_path)

src/epidemik/EpiModel.py

@@ -1,13 +1,37 @@
 from collections import defaultdict
+import os
+import sys
 
+__all__ = [
+    "OFFICIAL_REPO",
+    "EPI_COLORS",
+    "NotInitialized",
+    "get_cache_directory",
+]
 
 class NotInitialized(Exception):
     pass
 
+APPNAME = os.path.basename(os.path.realpath('.'))
+OFFICIAL_REPO = "https://github.com/DataForScience/epidemik/"
 
-epi_colors = defaultdict(lambda: "#f39019")
-epi_colors["S"] = "#51a7f9"
-epi_colors["E"] = "#f9e351"
-epi_colors["I"] = "#cf51f9"
-epi_colors["R"] = "#70bf41"
-epi_colors["D"] = "#8b8b8b"
+EPI_COLORS = defaultdict(lambda: "#f39019")
+EPI_COLORS["S"] = "#51a7f9"
+EPI_COLORS["E"] = "#f9e351"
+EPI_COLORS["I"] = "#cf51f9"
+EPI_COLORS["R"] = "#70bf41"
+EPI_COLORS["D"] = "#8b8b8b"
+
+def get_cache_directory():
+    """
+    Return the location of the cache directory for the current platform.
+    """
+    system = sys.platform
+
+    if system == 'darwin':
+        path = os.path.join(os.path.expanduser('~/Library/Caches'), APPNAME)
+    else:
+        path = os.getenv('XDG_CACHE_HOME', os.path.expanduser('~/.cache'))
+        path = os.path.join(path, APPNAME)
+
+    return path