parallel.py

import time
import bisect
import copy
import numpy as np
import pandas as pd
import networkx as nx
import scipy
import scipy.optimize
import scipy as sp
import os, math
import pickle
import matplotlib.pyplot as plt
from joblib import Parallel, delayed
from pathos.multiprocessing import ProcessingPool as Pool

from lib.dynamics import DiseaseModel
from lib.priorityqueue import PriorityQueue
from lib.measures import (MeasureList, BetaMultiplierMeasure,
                      SocialDistancingForAllMeasure, BetaMultiplierMeasureByType,
                      SocialDistancingForPositiveMeasure, SocialDistancingByAgeMeasure, SocialDistancingForSmartTracing, ComplianceForAllMeasure)

from concurrent.futures import ProcessPoolExecutor, ThreadPoolExecutor
from lib.mobilitysim import MobilitySimulator

# Comment this in if you want to do map plots
STORE_MOB = False

pp_legal_states = ['susc', 'expo', 'ipre', 'isym', 'iasy', 'posi', 'nega', 'resi', 'dead', 'hosp']


class ParallelSummary(object):
    """
    Summary class of several restarts
    """

    def __init__(self, max_time, repeats, n_people, n_sites, site_loc, home_loc):

        self.max_time = max_time
        self.random_repeats = repeats
        self.n_people = n_people
        self.n_sites = n_sites
        self.site_loc = site_loc
        self.home_loc = home_loc
       
        self.state = {
            'susc': np.ones((repeats, n_people), dtype='bool'),
            'expo': np.zeros((repeats, n_people), dtype='bool'),
            'ipre': np.zeros((repeats, n_people), dtype='bool'),
            'isym': np.zeros((repeats, n_people), dtype='bool'),
            'iasy': np.zeros((repeats, n_people), dtype='bool'),
            'posi': np.zeros((repeats, n_people), dtype='bool'),
            'nega': np.zeros((repeats, n_people), dtype='bool'),
            'resi': np.zeros((repeats, n_people), dtype='bool'),
            'dead': np.zeros((repeats, n_people), dtype='bool'),
            'hosp': np.zeros((repeats, n_people), dtype='bool'),
        }

        self.state_started_at = {
            'susc': - np.inf * np.ones((repeats, n_people), dtype='float'),
            'expo': np.inf * np.ones((repeats, n_people), dtype='float'),
            'ipre': np.inf * np.ones((repeats, n_people), dtype='float'),
            'isym': np.inf * np.ones((repeats, n_people), dtype='float'),
            'iasy': np.inf * np.ones((repeats, n_people), dtype='float'),
            'posi': np.inf * np.ones((repeats, n_people), dtype='float'),
            'nega': np.inf * np.ones((repeats, n_people), dtype='float'),
            'resi': np.inf * np.ones((repeats, n_people), dtype='float'),
            'dead': np.inf * np.ones((repeats, n_people), dtype='float'),
            'hosp': np.inf * np.ones((repeats, n_people), dtype='float'),
        }
        self.state_ended_at = {
            'susc': np.inf * np.ones((repeats, n_people), dtype='float'),
            'expo': np.inf * np.ones((repeats, n_people), dtype='float'),
            'ipre': np.inf * np.ones((repeats, n_people), dtype='float'),
            'isym': np.inf * np.ones((repeats, n_people), dtype='float'),
            'iasy': np.inf * np.ones((repeats, n_people), dtype='float'),
            'posi': np.inf * np.ones((repeats, n_people), dtype='float'),
            'nega': np.inf * np.ones((repeats, n_people), dtype='float'),
            'resi': np.inf * np.ones((repeats, n_people), dtype='float'),
            'dead': np.inf * np.ones((repeats, n_people), dtype='float'),
            'hosp': np.inf * np.ones((repeats, n_people), dtype='float'),
        }
        
        self.measure_list = []
        self.mob = []
        
        self.people_age = np.zeros((repeats, n_people), dtype='int')

        self.children_count_iasy = np.zeros((repeats, n_people), dtype='int')
        self.children_count_ipre = np.zeros((repeats, n_people), dtype='int')
        self.children_count_isym = np.zeros((repeats, n_people), dtype='int')


def create_ParallelSummary_from_DiseaseModel(sim):

    summary = ParallelSummary(sim.max_time, 1, sim.n_people, sim.mob.num_sites, sim.mob.site_loc, sim.mob.home_loc)

    for code in pp_legal_states:
        summary.state[code][0, :] = sim.state[code]
        summary.state_started_at[code][0, :] = sim.state_started_at[code]
        summary.state_ended_at[code][0, :] = sim.state_ended_at[code]

    summary.measure_list.append(sim.measure_list)
    if STORE_MOB:
        summary.mob.append(sim.mob)
    
    summary.people_age[0, :] = sim.mob.people_age
        
    summary.children_count_iasy[0, :] = sim.children_count_iasy
    summary.children_count_ipre[0, :] = sim.children_count_ipre
    summary.children_count_isym[0, :] = sim.children_count_isym
    return summary


def pp_launch(r, kwargs, distributions, params, initial_counts, testing_params, measure_list, max_time):

    mob = MobilitySimulator(**kwargs)
    mob.simulate(max_time=max_time)

    sim = DiseaseModel(mob, distributions)

    sim.launch_epidemic(
        params=params,
        initial_counts=initial_counts,
        testing_params=testing_params,
        measure_list=measure_list,
        verbose=False)

    result = {
        'state' : sim.state,
        'state_started_at': sim.state_started_at,
        'state_ended_at': sim.state_ended_at,
        'measure_list' : copy.deepcopy(sim.measure_list),
        'people_age' : sim.mob.people_age,
        'children_count_iasy': sim.children_count_iasy,
        'children_count_ipre': sim.children_count_ipre,
        'children_count_isym': sim.children_count_isym,
    }
    if STORE_MOB:
        result['mob'] = sim.mob

    return result


def launch_parallel_simulations(mob_settings, distributions, random_repeats, cpu_count, params, 
    initial_seeds, testing_params, measure_list, max_time, num_people, num_sites, site_loc, home_loc, 
    verbose=True, synthetic=False):
    

    with open(mob_settings, 'rb') as fp:
        kwargs = pickle.load(fp)

    mob_setting_list = [copy.deepcopy(kwargs) for _ in range(random_repeats)]
    distributions_list = [copy.deepcopy(distributions) for _ in range(random_repeats)]
    measure_list_list = [copy.deepcopy(measure_list) for _ in range(random_repeats)]
    params_list = [copy.deepcopy(params) for _ in range(random_repeats)]
    initial_seeds_list = [copy.deepcopy(initial_seeds) for _ in range(random_repeats)]
    testing_params_list = [copy.deepcopy(testing_params) for _ in range(random_repeats)]
    max_time_list = [copy.deepcopy(max_time) for _ in range(random_repeats)]
    repeat_ids = list(range(random_repeats))

    if verbose:
        print('Launching simulations...')

    with ProcessPoolExecutor(cpu_count) as ex:
        res = ex.map(pp_launch, repeat_ids, mob_setting_list, distributions_list, params_list,
                     initial_seeds_list, testing_params_list, measure_list_list, max_time_list)
    
    # collect all result (the fact that mob is still available here is due to the for loop)
    summary = ParallelSummary(max_time, random_repeats, num_people, num_sites, site_loc, home_loc)
    
    for r, result in enumerate(res):

        for code in pp_legal_states:
            summary.state[code][r, :] = result['state'][code]
            summary.state_started_at[code][r, :] = result['state_started_at'][code]
            summary.state_ended_at[code][r, :] = result['state_ended_at'][code]
        
        summary.measure_list.append(result['measure_list'])

        if STORE_MOB:
            summary.mob.append(result['mob']) 

        summary.people_age[r, :] = result['people_age']
        
        summary.children_count_iasy[r, :] = result['children_count_iasy']
        summary.children_count_ipre[r, :] = result['children_count_ipre']
        summary.children_count_isym[r, :] = result['children_count_isym']

    return summary