From e97a4cc8057ee17e59b0693926f8ff2e9b55db00 Mon Sep 17 00:00:00 2001
From: Luis Aleixo <luis.aleixo@cern.ch>
Date: Tue, 7 May 2024 15:16:29 +0200
Subject: [PATCH] added new model input

---
 caimira/models.py | 13 +++++++++++--
 1 file changed, 11 insertions(+), 2 deletions(-)

diff --git a/caimira/models.py b/caimira/models.py
index 192a6afc..04098c44 100644
--- a/caimira/models.py
+++ b/caimira/models.py
@@ -1331,6 +1331,9 @@ class ShortRangeModel:
     #: Interpersonal distances
     distance: _VectorisedFloat
 
+    #: Total people with short-range interactions
+    total_people: int = 1
+
     def dilution_factor(self) -> _VectorisedFloat:
         '''
         The dilution factor for the respective expiratory activity type.
@@ -1816,10 +1819,16 @@ class ExposureModel:
         """
         The expect_new_cases should always take the long-range infection_probability and multiply by the occupants exposed to long-range.
         """
+        prob_inf: _VectorisedFloat = self.infection_probability()
         if self.short_range != ():
-            return nested_replace(self, {'short_range': ()}).infection_probability() * self.exposed.number / 100
+            # If short-range interaction are defined, the total expected number of new cases
+            # has to take into account both the long- and short-range probability of infection.
+            long_range_model = nested_replace(self, {'short_range': (),})
+            long_range_prob = long_range_model.infection_probability()
+            short_range_total_people = self.short_range[0].total_people
+            return (prob_inf * short_range_total_people + long_range_prob * self.exposed.number) / 100
 
-        return self.infection_probability() * self.exposed.number / 100
+        return prob_inf * self.exposed.number / 100
 
     def reproduction_number(self) -> _VectorisedFloat:
         """