Merge branch 'dev' into infe_prof_and_same_diff

SABS-R3-Epidemiology · Mar 4, 2024 · da50cc1 · da50cc1
2 parents db2a1a3 + 8a9434a
commit da50cc1
Show file tree

Hide file tree

Showing 3 changed files with 16 additions and 8 deletions.
diff --git a/epios/post_process.py b/epios/post_process.py
@@ -444,6 +444,7 @@ def _compare(self, time_sample, gen_plot=False, scale_method: str = 'proportiona
             # Find the difference between estimated infection level and the real one
             diff = np.array(true_result) - result_scaled
             if gen_plot:
+                plt.figure()
                 plt.plot(time_sample, result_scaled, label='Predicted result', linestyle='--')
                 plt.plot(time_sample, true_result, label='True result')
                 plt.plot(time_sample, np.abs(diff), label='Absolute difference')
@@ -488,10 +489,11 @@ def _wrapper_Region_AgeRegion(self, sampling_method, sample_size, time_sample, n
                     else:
                         if sampling_method == 'AgeRegion':
                             sampler_class = SamplerAgeRegion(data=self.demo_data, data_store_path=data_store_path,
-                                                             num_age_group=num_age_group,
+                                                             num_age_group=num_age_group, pre_process=False,
                                                              age_group_width=age_group_width)
                         else:
-                            sampler_class = SamplerRegion(data=self.demo_data, data_store_path=data_store_path)
+                            sampler_class = SamplerRegion(data=self.demo_data, data_store_path=data_store_path,
+                                                          pre_process=False)
                     try:
                         people = sampler_class.sample(sample_size=sample_size, additional_sample=additional_sample)
                     except NameError:
@@ -658,7 +660,9 @@ def _wrapper_Region_AgeRegion(self, sampling_method, sample_size, time_sample, n
 
             # Plot the figure
             if gen_plot:
-                plt.plot(time_sample, infected_rate)
+                plt.figure()
+                infected_population = np.array(infected_rate) * len(self.demo_data)
+                plt.plot(time_sample, infected_population)
                 plt.xlabel('Time')
                 plt.ylabel('Population')
                 plt.xlim(0, max(time_sample))
@@ -728,9 +732,11 @@ def _wrapper_Age_Base(self, sampling_method, sample_size, time_sample,
                     else:  # After the data process, we can directly read files processed at the first time
                         if sampling_method == 'Age':
                             sampler_class = SamplerAge(data=self.demo_data, data_store_path=data_store_path,
-                                                       num_age_group=num_age_group, age_group_width=age_group_width)
+                                                       num_age_group=num_age_group, age_group_width=age_group_width,
+                                                       pre_process=False)
                         else:
-                            sampler_class = Sampler(data=self.demo_data, data_store_path=data_store_path)
+                            sampler_class = Sampler(data=self.demo_data, data_store_path=data_store_path,
+                                                    pre_process=False)
                     people = sampler_class.sample(sample_size=sample_size)
 
                     X = SamplingMaker(non_resp_rate=0, data=self.time_data,
@@ -742,7 +748,9 @@ def _wrapper_Age_Base(self, sampling_method, sample_size, time_sample,
 
             # Plot the figure
             if gen_plot:
-                plt.plot(time_sample, infected_rate)
+                plt.figure()
+                infected_population = np.array(infected_rate) * len(self.demo_data)
+                plt.plot(time_sample, infected_population)
                 plt.xlabel('Time')
                 plt.ylabel('Population')
                 plt.xlim(0, max(time_sample))

diff --git a/epios/sampler_region.py b/epios/sampler_region.py
@@ -99,7 +99,7 @@ def multinomial_draw(self, n: int, prob: list):
         cap_region = []
         record_cap_region = []
         for i in range(len(prob)):
-            cap_region.append(min(max(n * prob[i] + 0.005 * n, 1),
+            cap_region.append(min(n * prob[i] + 0.005 * n + 1,
                                   self.geoinfo[self.geoinfo['cell'] == i]['Susceptible'].sum()))
             record_cap_region.append(self.geoinfo[self.geoinfo['cell'] == i]['Susceptible'].sum())
         cap_region = [cap_region, list(np.arange(len(cap_region)))]

diff --git a/epios/tests/test_re_scaler.py b/epios/tests/test_re_scaler.py
@@ -1,7 +1,7 @@
 from unittest import TestCase
 from numpy.random import rand
 from numpy import array
-from epios.re_scaler import ReScaler
+from re_scaler import ReScaler
 
 
 class TestRS(TestCase):