Cleaned up a sped up the eula calculation code.

src/idmlaser/numpynumba/population.py

@@ -4,9 +4,16 @@
 
 from tqdm import tqdm
 import numpy as np
+import numba as nb
 import h5py
 import pdb
 
+@nb.njit(parallel=True)
+def accumulate_deaths_parallel(nodeid_filtered, death_year, nodeid_indices_array, total_population_per_year):
+ for i in nb.prange(len(death_year)):
+ node_index = nodeid_indices_array[nodeid_filtered[i]]
+ total_population_per_year[node_index, death_year[i]] += 1
+
 class Population:
  """Array-based Agent Based Population Class"""
 
@@ -246,27 +253,31 @@ def expected_pops_over_years(self, eula_age_in_years=5, years=10):
 
  # Calculate the year of death for each individual (0-based index for years 1-years)
  death_year = dod_filtered // 365
- death_year[death_year >= years] = years-1 # Cap deaths at year years
+ # Let's initially count for all the years, even if we only keep first 10.
+ # No, too slow. Let's count 0 through years-1, and put 'everything else' in 'years'
+ death_year[death_year >= years] = years-1 # Cap deaths at years-1
 
- # Initialize the total_population_per_year array
+ # Initialize the expected_new_deaths array
  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), years), dtype=int)
  self.expected_new_deaths_per_year = np.zeros((len(unique_nodeids), years), dtype=int)
 
- # Accumulate deaths by year and node
- for i in tqdm(range(len(death_year))):
- node_index = nodeid_indices[nodeid_filtered[i]]
-  self.total_population_per_year[node_index, death_year[i]] += 1
+ nodeid_indices_array = np.zeros(unique_nodeids.max() + 1, dtype=np.int32)
+ for i, nodeid in enumerate(unique_nodeids):
+ nodeid_indices_array[nodeid] = i
+ accumulate_deaths_parallel(nodeid_filtered, death_year, nodeid_indices_array, self.expected_new_deaths_per_year)
 
  # 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)
+ cumulative_deaths = np.cumsum(self.expected_new_deaths_per_year, axis=1) # xtra dupe element at end I don't understand yet
 
  # Calculate new deaths per year
- self.expected_new_deaths_per_year[:, 0] = self.total_population_per_year[:, 0]
+ self.expected_new_deaths_per_year[:, 0] = self.expected_new_deaths_per_year[:, 0]
  self.expected_new_deaths_per_year[:, 1:] = np.diff(cumulative_deaths, axis=1)
+ self.expected_new_deaths_per_year = self.expected_new_deaths_per_year[:, 0:10] # discard extra killall column
 
+ # Calculate the initial population counts, though not currently used
+ initial_population_counts = np.bincount(nodeid_filtered, minlength=len(unique_nodeids))
+ # Calculate total_population_per_year at the end, using cumulative deaths
  self.total_population_per_year = initial_population_counts[:, None] - cumulative_deaths
 
  # Optional: print the resulting populations