From 59fcfce00d681a8992839ffdb7045c5d8848220b Mon Sep 17 00:00:00 2001
From: Giulio Rossetti <giulio.rossetti@gmail.com>
Date: Tue, 31 Mar 2020 20:58:08 +0200
Subject: [PATCH] :arrow_up: UTLDR model

---
 ndlib/models/epidemics/UTLDRModel.py | 33 ++++++++++++++++++++++++++--
 ndlib/test/test_ndlib.py             |  2 ++
 2 files changed, 33 insertions(+), 2 deletions(-)

diff --git a/ndlib/models/epidemics/UTLDRModel.py b/ndlib/models/epidemics/UTLDRModel.py
index 40a4143..556fcaa 100644
--- a/ndlib/models/epidemics/UTLDRModel.py
+++ b/ndlib/models/epidemics/UTLDRModel.py
@@ -141,6 +141,18 @@ def __init__(self, graph, seed=None):
                     "optional": True,
                     "default": 1
                 },
+                "z": {
+                    "descr": "Probability of infection from corpses",
+                    "range": [0, 1],
+                    "optional": True,
+                    "default": 0
+                },
+                "s": {
+                    "descr": "Probability of absent immunization",
+                    "range": [0, 1],
+                    "optional": True,
+                    "default": 0
+                },
             },
              "nodes": {
                 "activity": {
@@ -293,7 +305,12 @@ def iteration(self, node_status=True):
 
             ####################### Resolved Compartment ###########################
 
-            elif u_status == self.available_statuses['Dead'] or u_status == self.available_statuses['Recovered']:
+            elif u_status == self.available_statuses['Recovered']:
+                immunity = np.random.random_sample()
+                if immunity < self.params['model']['s']:
+                    actual_status[u] = self.available_statuses['Susceptible']
+
+            elif u_status == self.available_statuses['Dead']:
                 pass
 
         delta, node_count, status_delta = self.status_delta(actual_status)
@@ -400,7 +417,12 @@ def __Susceptible_to_Exposed(self, u, neighbors, lockdown=False):
         l_range = np.random.random_sample()  # long range interaction
         if l_range < l_range_proba:
             # filtering out quarantined and dead nodes
-            candidates = [n for n in self.graph.nodes if self.status[n] not in [self.available_statuses['Tested_E'], self.available_statuses['Tested_I'], self.available_statuses['Dead']]]
+            if self.params['model']['z'] == 0:
+                candidates = [n for n in self.graph.nodes if self.status[n] not in [self.available_statuses['Tested_E'], self.available_statuses['Tested_I'], self.available_statuses['Dead']]]
+            else:
+                candidates = [n for n in self.graph.nodes if self.status[n] not in [self.available_statuses['Tested_E'],
+                                                                                    self.available_statuses['Tested_I']]]
+
             interactions.extend(list(np.random.choice(a=candidates, size=int(social_interactions*self.params['model']['lsize']), replace=True)))
 
         for v in interactions:
@@ -411,6 +433,13 @@ def __Susceptible_to_Exposed(self, u, neighbors, lockdown=False):
                         return self.available_statuses['Lockdown_E']
                     return self.available_statuses['Exposed']
 
+            elif self.status[v] == self.available_statuses['Dead']:
+                zp = np.random.random_sample()
+                if zp < self.params['model']['z']:
+                    if lockdown:
+                        return self.available_statuses['Lockdown_E']
+                    return self.available_statuses['Exposed']
+
         if lockdown:
             return self.available_statuses['Lockdown_S']
         return self.available_statuses['Susceptible']
diff --git a/ndlib/test/test_ndlib.py b/ndlib/test/test_ndlib.py
index e3e40eb..f2e5015 100644
--- a/ndlib/test/test_ndlib.py
+++ b/ndlib/test/test_ndlib.py
@@ -66,6 +66,8 @@ def test_utldr(self):
             config.add_model_parameter("epsilon_e", 1)
             config.add_model_parameter("icu_b", 0.15)
             config.add_model_parameter("iota", 0.20)
+            config.add_model_parameter("z", 0.2)
+            config.add_model_parameter("s", 0.05)
 
             # Lockdown
             config.add_model_parameter("lambda", 0.8)