From 6cee873b5f3108ecf8dfc39bb353d02af7a09e08 Mon Sep 17 00:00:00 2001 From: Jonathan Bloedow Date: Wed, 7 Aug 2024 13:31:05 -0700 Subject: [PATCH] 1) ages_in_days was using entire array including expansion slots; 2) Removed unneeded split_index 'min' calc; 3) Added expected_new_deaths_per_year since probably more useful to the way FUSION calculates population. --- src/idmlaser/numpynumba/population.py | 11 +++++++---- 1 file changed, 7 insertions(+), 4 deletions(-) diff --git a/src/idmlaser/numpynumba/population.py b/src/idmlaser/numpynumba/population.py index f37c8ce..cae2236 100644 --- a/src/idmlaser/numpynumba/population.py +++ b/src/idmlaser/numpynumba/population.py @@ -197,7 +197,7 @@ def eliminate_eulas(self, eula_age_in_years: float): # Calculate the age of each individual in days current_day = 0 # Adjust this if you have a simulation day tracker - ages_in_days = current_day - self.__dict__['dob'] + ages_in_days = current_day - self.__dict__['dob'][:self.count] # Calculate number of expansion slots birth_cap = (self.capacity-self.count) * 4 # hack coz right now seem to have "too many" births @@ -210,9 +210,6 @@ def eliminate_eulas(self, eula_age_in_years: float): split_index = np.searchsorted(ages_in_days[sorted_indices], age_threshold_in_days) print( f"split_index = {split_index}" ) - # Ensure split_index does not exceed the size of the sorted_indices; probably unnecessary - split_index = min(split_index, len(sorted_indices)) - # Keep only the individuals below the age threshold for key, value in self.__dict__.items(): if isinstance(value, np.ndarray) and value.size == self._capacity: @@ -254,6 +251,7 @@ def expected_pops_over_years(self, eula_age_in_years=5): unique_nodeids = np.unique(nodeid_filtered) nodeid_indices = {nodeid: i for i, nodeid in enumerate(unique_nodeids)} self.total_population_per_year = np.zeros((len(unique_nodeids), 20), dtype=int) + self.expected_new_deaths_per_year = np.zeros((len(unique_nodeids), 20), dtype=int) # Accumulate deaths by year and node for i in tqdm(range(len(death_year))): @@ -263,6 +261,11 @@ def expected_pops_over_years(self, eula_age_in_years=5): # Convert deaths to populations by subtracting cumulative deaths from the initial population initial_population_counts = np.bincount(nodeid_filtered, minlength=len(unique_nodeids)) cumulative_deaths = np.cumsum(self.total_population_per_year, axis=1) + + # Calculate new deaths per year + self.expected_new_deaths_per_year[:, 0] = self.total_population_per_year[:, 0] + self.expected_new_deaths_per_year[:, 1:] = np.diff(cumulative_deaths, axis=1) + self.total_population_per_year = initial_population_counts[:, None] - cumulative_deaths # Optional: print the resulting populations