Ejemplo n.º 1
0
class Simulation(object):
    def __init__(self, params, ind_type=Individual, create_pop=True):
        # convert ConfigObj to dictionary to store params
        self.params = dict(params)
        self.params_adj = {}
        self.rng = Random(self.params['seed'])
        self.load_demographic_data()
        if create_pop:
            self.create_population(ind_type, params['logging'])

        # for storing general data
        self.max_hh = 50  # maximum possible hh size
        self.pop_size = []
        self.age_dist = []
        self.hh_size_dist = []
        self.hh_size_dist_counts = []
        self.hh_size_avg = []
        self.hh_count = []
        self.fam_types = []
        self.start_time = None
        self.end_time = None

    def create_population(self, ind_type, logging=True):
        """
        Create a population according to specified age and household size
        distributions.
        """

        self.P = Pop_HH(ind_type, logging)
        self.P.gen_hh_age_structured_pop(self.params['pop_size'], self.hh_comp,
                                         self.age_dist,
                                         self.params['age_cutoffs'], self.rng)
        self.P.allocate_couples()
        self.P.print_population_summary()

    def parse_age_rates(self, filename, factor, final):
        """
        Parse an age-year-rate table to produce a dictionary, keyed by age, 
        with each entry being a list of annual rates (by year).
        
        Setting final to 'True' appends an age 100 rate of >1 (e.g., to 
        ensure everyone dies!
        """

        dat = load_age_rates(filename)
        rates = {}
        for line in dat:
            rates[line[0]] = [x * factor for x in line[1:]]
        if final:
            rates[101] = [100 for x in dat[0][1:]]  # everybody dies...
        return rates

    def load_demographic_data(self):
        """
        Load data on age-specific demographic processes (mortality/fertility)
        and adjust event probabilities according to time-step.
        """

        # load household size distribution and age distribution
        self.hh_comp = load_probs(
            os.path.join(self.params['resource_prefix'],
                         self.params['hh_composition']), False)
        self.params['age_cutoffs'] = [int(x)
                                      for x in self.hh_comp[0][1:][0]]  # yuk!
        self.age_dist = load_probs(
            os.path.join(self.params['resource_prefix'],
                         self.params['age_distribution']))

        annual_factor = self.params['t_dur'] / 365.0

        # load and scale MORTALITY rates
        self.death_rates = {}
        self.death_rates[0] = self.parse_age_rates(
            os.path.join(self.params['resource_prefix'],
                         self.params['death_rates_m']), annual_factor, True)
        self.death_rates[1] = self.parse_age_rates(
            os.path.join(self.params['resource_prefix'],
                         self.params['death_rates_f']), annual_factor, True)

        ### load FERTILITY age probs (don't require scaling) for closed pops
        self.fertility_age_probs = load_prob_tables(
            os.path.join(self.params['resource_prefix'],
                         self.params['fertility_age_probs']))
        self.fertility_parity_probs = load_probs_new(
            os.path.join(self.params['resource_prefix'],
                         self.params['fertility_parity_probs']))

        ### load and scale leav/couple/divorce and growth rates
        if self.params['dyn_rates']:
            # rates will be a list of annual values
            self.params['leaving_probs'] = load_prob_list(
                os.path.join(self.params['resource_prefix'],
                             self.params['leaving_prob_file']))
            self.params['couple_probs'] = load_prob_list(
                os.path.join(self.params['resource_prefix'],
                             self.params['couple_prob_file']))
            self.params['divorce_probs'] = load_prob_list(
                os.path.join(self.params['resource_prefix'],
                             self.params['divorce_prob_file']))
            self.params['growth_rates'] = load_prob_list(
                os.path.join(self.params['resource_prefix'],
                             self.params['growth_rate_file']))
            self.params['imm_rates'] = load_prob_list(
                os.path.join(self.params['resource_prefix'],
                             self.params['imm_rate_file']))

            self.params_adj['leaving_probs'] = [
                adjust_prob(x, self.params['t_dur'])
                for x in self.params['leaving_probs']
            ]
            self.params_adj['couple_probs'] = [
                adjust_prob(x, self.params['t_dur'])
                for x in self.params['couple_probs']
            ]
            self.params_adj['divorce_probs'] = [
                adjust_prob(x, self.params['t_dur'])
                for x in self.params['divorce_probs']
            ]
            self.params_adj['growth_rates'] = [
                adjust_prob(x, self.params['t_dur'])
                for x in self.params['growth_rates']
            ]
            self.params_adj['imm_rates'] = [
                adjust_prob(x, self.params['t_dur'])
                for x in self.params['imm_rates']
            ]

            self.dyn_years = min(
                len(self.death_rates[0][0]) - 1,
                len(self.fertility_age_probs) - 1,
                len(self.params_adj['leaving_probs']) - 1,
                len(self.params_adj['couple_probs']) - 1,
                len(self.params_adj['divorce_probs']) - 1,
                len(self.params_adj['growth_rates']) - 1)

        else:
            # adjust demographic event probabilities according to time step
            self.params_adj['couple_probs'] = [
                adjust_prob(self.params['couple_prob'], self.params['t_dur'])
            ]
            self.params_adj['leaving_probs'] = [
                adjust_prob(self.params['leaving_prob'], self.params['t_dur'])
            ]
            self.params_adj['divorce_probs'] = [
                adjust_prob(self.params['divorce_prob'], self.params['t_dur'])
            ]
            self.params_adj['growth_rates'] = [
                adjust_prob(self.params['growth_rate'], self.params['t_dur'])
            ]
            self.params_adj['imm_rates'] = [
                adjust_prob(self.params['imm_rate'], self.params['t_dur'])
            ]

    def update_individual_demo(self, t, ind, index=0):
        """
        Update individual ind; check for death, couple formation, leaving home
        or divorce, as possible and appropriate.
        """

        death = None
        birth = None

        couple_prob = self.params_adj['couple_probs'][index]
        #        if ind.divorced and ind.deps: couple_prob *= 0.5

        # DEATH / BIRTH:
        if self.rng.random() > exp(-self.death_rates[ind.sex][ind.age][index]):
            death = ind
            mother = ind
            while mother is ind:  # make sure dead individual isn't selected as mother!
                mother = self.choose_mother(index)
            birth = self.update_death_birth(t, ind, mother)

        # COUPLE FORMATION:
        elif self.params['couple_age'] < ind.age < 60 \
                and not ind.partner \
                and self.rng.random() < couple_prob:
            partner = self.choose_partner(ind)
            if partner:
                self.P.form_couple(t, ind, partner)

        # LEAVING HOME:
        elif ind.age > self.params['leaving_age'] \
                and ind.with_parents \
                and not ind.partner \
                and self.rng.random() < self.params_adj['leaving_probs'][index]:
            self.P.leave_home(t, ind)

        # DIVORCE:
        elif self.params['divorce_age'] < ind.age < 50 \
                and ind.partner \
                and self.rng.random() < self.params_adj['divorce_probs'][index]:
            self.P.separate_couple(t, ind)

        # ELSE: individual has a quiet year...
        return death, birth

    def choose_mother(self, index):
        """
        Choose a new mother on the basis of fertility rates.

        NOTE: there is still a very small possibility (in *very* small
        populations) that this will hang due to candidates remaining
        forever empty.  Should add a check to prevent this and exit gracefully.
        """

        candidates = []
        while not candidates:
            tgt_age = int(
                sample_table(self.fertility_age_probs[index], self.rng)[0])
            tgt_prev_min = 0
            tgt_prev_max = 100
            if self.params['use_parity']:
                tgt_prev_min = int(
                    sample_table(
                        self.fertility_parity_probs[(tgt_age - 15) / 5],
                        self.rng)[0])
                # effectively transform 5 into 5+
                tgt_prev_max = tgt_prev_min if tgt_prev_min < 5 else 20
            tgt_set = self.P.individuals_by_age(tgt_age, tgt_age)
            candidates = [x
                    for x in tgt_set \
                    if x.sex == 1 \
                    and x.can_birth() \
                    and not x.with_parents \
                    and tgt_prev_min <= len(x.children) <= tgt_prev_max
                    ]
        return self.rng.choice(candidates)

    def choose_partner(self, ind):
        """
        Choose a partner for i_id, subject to parameter constraints.

        :param ind: the first partner in the couple.
        :type ind: Individual
        :returns: partner if successful, otherwise None.
        """

        mean_age = ind.age+self.params['partner_age_diff'] \
                if ind.sex == 0 else ind.age-self.params['partner_age_diff']
        tgt_age = 0
        while tgt_age < self.params['min_partner_age']:
            tgt_age = int(
                self.rng.gauss(mean_age, self.params['partner_age_sd']))
            tgt_set = self.P.individuals_by_age(tgt_age, tgt_age)
            candidates = [x \
                for x in tgt_set \
                if not x.partner \
                and x.sex != ind.sex \
                and x not in self.P.hh_members(ind)
                ]

        # abort if no eligible partner exists
        return None if candidates == [] else self.rng.choice(candidates)

    def update_death_birth(self, t, ind, mother):
        """
        Replace a dying individual with a newborn.  If no individual to die
        is passed, only a birth occurs; if no mother is passed, only a death
        occurs.

        :param t: the current time step.
        :type t: int
        :param ind: the individual to die.
        :type ind: Individual
        :returns: the new individual.
        """

        if ind:
            orphans = self.P.death(t, ind)
            self.P.process_orphans(t, orphans, self.params['age_cutoffs'][-2],
                                   self.rng)

        if mother:
            sex = self.rng.randint(0, 1)
            new_ind = self.P.birth(t, self.rng, mother, mother.partner, sex)
            return new_ind

        return None

    def update_all_demo(self, t):
        """
        Update population over period of t days.

        Returns list of births, deaths, immigrants and birthdays.
        """

        birthdays = self.P.age_population(self.params['t_dur'])

        deaths = []
        births = []

        # calculate index for fertility and mortality rates
        # basically: use first entry for burn-in, then one entry every
        # 'period' years, then use the final entry for any remaining years.
        index = min(max(0, (t - (self.params['demo_burn']*365)) /
                (self.params['demo_period']*365)), self.dyn_years) \
                        if self.params['dyn_rates'] else 0

        cur_inds = self.P.I.values()
        for ind in cur_inds:
            death, birth = self.update_individual_demo(t, ind, index)
            if death: deaths.append(death)
            if birth: births.append(birth)

        #population growth
        for x in xrange(
                int(len(self.P.I) * self.params_adj['growth_rates'][index])):
            mother = self.choose_mother(index)
            births.append(self.update_death_birth(t, None, mother))

        #immigration
        imm_count = 0
        imm_tgt = int(len(self.P.I) * self.params_adj['imm_rates'][index])
        source_hh_ids = []
        immigrants = []
        while imm_count < imm_tgt:
            hh_id = self.rng.choice(self.P.groups['household'].keys())
            imm_count += len(self.P.groups['household'][hh_id])
            source_hh_ids.append(hh_id)
        for hh_id in source_hh_ids:
            new_hh_id = self.P.duplicate_household(t, hh_id)
            immigrants.extend(self.P.groups['household'][new_hh_id])

        return births, deaths, immigrants, birthdays

    def record_stats_demo(self, t):
        if self.params['record_interval'] <= 0 or t % self.params[
                'record_interval'] is not 0:
            return

        self.pop_size.append(len(self.P.I))
        self.age_dist.append(self.P.age_dist(101, 101)[0])
        self.hh_size_dist.append(
            self.P.group_size_dist('household', self.max_hh)[0])
        self.hh_size_dist_counts.append(
            self.P.group_size_dist('household', self.max_hh, False)[0])
        self.hh_size_avg.append(self.P.group_size_avg('household'))
        self.hh_count.append(len(self.P.groups['household']))

        self.fam_types.append(self.P.sum_hh_stats_group())

    def run(self):
        """
        Run simulated population (NB: this probably won't be used, as
        external running object may wish to just call update_all itself...)
        """

        timesteps = self.params['years'] * (365 / self.params['t_dur'])
        i = 0
        while i < timesteps:
            t = i * self.params['t_dur']
            print i, t
            self.update_all_demo(t)
            i += 1
Ejemplo n.º 2
0
class Simulation(object):

    def __init__(self, params, ind_type=Individual, create_pop=True):
        # convert ConfigObj to dictionary to store params
        self.params = dict(params)
        self.params_adj={}
        self.rng = Random(self.params['seed'])
        self.load_demographic_data()
        if create_pop:
            self.create_population(ind_type, params['logging'])

        # for storing general data
        self.max_hh = 50 # maximum possible hh size
        self.pop_size = []
        self.age_dist = []
        self.hh_size_dist = []
        self.hh_size_dist_counts = []
        self.hh_size_avg = []
        self.hh_count = []
        self.fam_types = []
        self.start_time = None
        self.end_time = None



    def create_population(self, ind_type, logging=True):
        """
        Create a population according to specified age and household size
        distributions.
        """

        self.P = Pop_HH(ind_type, logging)
        self.P.gen_hh_age_structured_pop(self.params['pop_size'], self.hh_comp, 
                self.age_dist, self.params['age_cutoffs'], self.rng)
        self.P.allocate_couples()
        self.P.print_population_summary()


    def parse_age_rates(self, filename, factor, final):
        """
        Parse an age-year-rate table to produce a dictionary, keyed by age, 
        with each entry being a list of annual rates (by year).
        
        Setting final to 'True' appends an age 100 rate of >1 (e.g., to 
        ensure everyone dies!
        """

        dat = load_age_rates(filename)
        rates = {}
        for line in dat:
            rates[line[0]] = [x * factor for x in line[1:]]
        if final:
            rates[101] = [100 for x in dat[0][1:]] # everybody dies...
        return rates


    def load_demographic_data(self):
        """
        Load data on age-specific demographic processes (mortality/fertility)
        and adjust event probabilities according to time-step.
        """

        # load household size distribution and age distribution
        self.hh_comp = load_probs(os.path.join(self.params['resource_prefix'], 
                    self.params['hh_composition']), False)
        self.params['age_cutoffs'] = [int(x) for x in self.hh_comp[0][1:][0]]  # yuk!
        self.age_dist = load_probs(os.path.join(self.params['resource_prefix'], 
                    self.params['age_distribution']))

        annual_factor = self.params['t_dur']/365.0

        # load and scale MORTALITY rates
        self.death_rates = {}
        self.death_rates[0] = self.parse_age_rates(os.path.join(
            self.params['resource_prefix'], 
            self.params['death_rates_m']), annual_factor, True)
        self.death_rates[1] = self.parse_age_rates(os.path.join(
            self.params['resource_prefix'], 
            self.params['death_rates_f']), annual_factor, True)

        ### load FERTILITY age probs (don't require scaling) for closed pops
        self.fertility_age_probs = load_prob_tables(os.path.join(
            self.params['resource_prefix'], 
            self.params['fertility_age_probs']))
        self.fertility_parity_probs = load_probs_new(os.path.join(
            self.params['resource_prefix'],
            self.params['fertility_parity_probs']))

        ### load and scale leav/couple/divorce and growth rates
        if self.params['dyn_rates']:
            # rates will be a list of annual values
            self.params['leaving_probs'] = load_prob_list(os.path.join(
                self.params['resource_prefix'], self.params['leaving_prob_file']))
            self.params['couple_probs'] = load_prob_list(os.path.join(
                self.params['resource_prefix'], self.params['couple_prob_file']))
            self.params['divorce_probs'] = load_prob_list(os.path.join(
                self.params['resource_prefix'], self.params['divorce_prob_file']))
            self.params['growth_rates'] = load_prob_list(os.path.join(
                self.params['resource_prefix'], self.params['growth_rate_file']))
            self.params['imm_rates'] = load_prob_list(os.path.join(
                self.params['resource_prefix'], self.params['imm_rate_file']))

            self.params_adj['leaving_probs'] = [adjust_prob(x, self.params['t_dur']) 
                for x in self.params['leaving_probs']]
            self.params_adj['couple_probs'] = [adjust_prob(x, self.params['t_dur']) 
                for x in self.params['couple_probs']]
            self.params_adj['divorce_probs'] = [adjust_prob(x, self.params['t_dur']) 
                for x in self.params['divorce_probs']]
            self.params_adj['growth_rates'] = [adjust_prob(x, self.params['t_dur']) 
                for x in self.params['growth_rates']]
            self.params_adj['imm_rates'] = [adjust_prob(x, self.params['t_dur']) 
                for x in self.params['imm_rates']]

            self.dyn_years = min(len(self.death_rates[0][0])-1, len(self.fertility_age_probs)-1,
                    len(self.params_adj['leaving_probs'])-1, len(self.params_adj['couple_probs'])-1,
                    len(self.params_adj['divorce_probs'])-1, len(self.params_adj['growth_rates'])-1)

        else:
            # adjust demographic event probabilities according to time step
            self.params_adj['couple_probs'] = [adjust_prob(
                self.params['couple_prob'], self.params['t_dur'])]
            self.params_adj['leaving_probs'] = [adjust_prob(
                self.params['leaving_prob'], self.params['t_dur'])]
            self.params_adj['divorce_probs'] = [adjust_prob(
                self.params['divorce_prob'], self.params['t_dur'])]
            self.params_adj['growth_rates'] = [adjust_prob(
                self.params['growth_rate'], self.params['t_dur'])]
            self.params_adj['imm_rates'] = [adjust_prob(
                self.params['imm_rate'], self.params['t_dur'])]


    def update_individual_demo(self, t, ind, index=0):
        """
        Update individual ind; check for death, couple formation, leaving home
        or divorce, as possible and appropriate.
        """

        death = None; birth = None

        couple_prob = self.params_adj['couple_probs'][index]
#        if ind.divorced and ind.deps: couple_prob *= 0.5

        # DEATH / BIRTH: 
        if self.rng.random() > exp(-self.death_rates[ind.sex][ind.age][index]):
            death = ind
            mother = ind
            while mother is ind:    # make sure dead individual isn't selected as mother!
                mother = self.choose_mother(index)
            birth = self.update_death_birth(t, ind, mother)

        # COUPLE FORMATION:
        elif self.params['couple_age'] < ind.age < 60 \
                and not ind.partner \
                and self.rng.random() < couple_prob:
            partner = self.choose_partner(ind)
            if partner:
                self.P.form_couple(t, ind, partner)

        # LEAVING HOME:
        elif ind.age > self.params['leaving_age'] \
                and ind.with_parents \
                and not ind.partner \
                and self.rng.random() < self.params_adj['leaving_probs'][index]:
            self.P.leave_home(t, ind)

        # DIVORCE:
        elif self.params['divorce_age'] < ind.age < 50 \
                and ind.partner \
                and self.rng.random() < self.params_adj['divorce_probs'][index]:
            self.P.separate_couple(t, ind)

        # ELSE: individual has a quiet year...
        return death, birth 


    def choose_mother(self, index):
        """
        Choose a new mother on the basis of fertility rates.

        NOTE: there is still a very small possibility (in *very* small
        populations) that this will hang due to candidates remaining
        forever empty.  Should add a check to prevent this and exit gracefully.
        """

        candidates = []
        while not candidates:
            tgt_age = int(sample_table(self.fertility_age_probs[index], self.rng)[0])
            tgt_prev_min = 0; tgt_prev_max = 100
            if self.params['use_parity']:
                tgt_prev_min = int(sample_table(
                    self.fertility_parity_probs[(tgt_age-15)/5], self.rng)[0])
                # effectively transform 5 into 5+
                tgt_prev_max = tgt_prev_min if tgt_prev_min < 5 else 20
            tgt_set = self.P.individuals_by_age(tgt_age, tgt_age)
            candidates = [x
                    for x in tgt_set \
                    if x.sex == 1 \
                    and x.can_birth() \
                    and not x.with_parents \
                    and tgt_prev_min <= len(x.children) <= tgt_prev_max
                    ]
        return self.rng.choice(candidates)


    def choose_partner(self, ind):
        """
        Choose a partner for i_id, subject to parameter constraints.

        :param ind: the first partner in the couple.
        :type ind: Individual
        :returns: partner if successful, otherwise None.
        """

        mean_age = ind.age+self.params['partner_age_diff'] \
                if ind.sex == 0 else ind.age-self.params['partner_age_diff']
        tgt_age = 0
        while tgt_age < self.params['min_partner_age']:
            tgt_age = int(self.rng.gauss(mean_age, self.params['partner_age_sd']))
            tgt_set = self.P.individuals_by_age(tgt_age, tgt_age)
            candidates = [x \
                for x in tgt_set \
                if not x.partner \
                and x.sex != ind.sex \
                and x not in self.P.hh_members(ind)
                ]

        # abort if no eligible partner exists
        return None if candidates == [] else self.rng.choice(candidates)



    def update_death_birth(self, t, ind, mother):
        """
        Replace a dying individual with a newborn.  If no individual to die
        is passed, only a birth occurs; if no mother is passed, only a death
        occurs.

        :param t: the current time step.
        :type t: int
        :param ind: the individual to die.
        :type ind: Individual
        :returns: the new individual.
        """

        if ind:
            orphans = self.P.death(t, ind)
            self.P.process_orphans(t, orphans, self.params['age_cutoffs'][-2], self.rng)
        
        if mother:
            sex = self.rng.randint(0, 1)
            new_ind = self.P.birth(t, self.rng, mother, mother.partner, sex)
            return new_ind
        
        return None



    def update_all_demo(self, t):
        """
        Update population over period of t days.

        Returns list of births, deaths, immigrants and birthdays.
        """

        birthdays = self.P.age_population(self.params['t_dur'])

        deaths = []; births = []

        # calculate index for fertility and mortality rates
        # basically: use first entry for burn-in, then one entry every 
        # 'period' years, then use the final entry for any remaining years.
        index = min(max(0, (t - (self.params['demo_burn']*365)) / 
                (self.params['demo_period']*365)), self.dyn_years) \
                        if self.params['dyn_rates'] else 0

        cur_inds = self.P.I.values()
        for ind in cur_inds: 
            death, birth = self.update_individual_demo(t, ind, index)
            if death: deaths.append(death)
            if birth: births.append(birth)

        #population growth
        for x in xrange(int(len(self.P.I) * self.params_adj['growth_rates'][index])):
            mother = self.choose_mother(index)
            births.append(self.update_death_birth(t, None, mother))

        #immigration
        imm_count = 0
        imm_tgt = int(len(self.P.I) * self.params_adj['imm_rates'][index])
        source_hh_ids = []
        immigrants = []
        while imm_count < imm_tgt:
            hh_id = self.rng.choice(self.P.groups['household'].keys())
            imm_count += len(self.P.groups['household'][hh_id])
            source_hh_ids.append(hh_id)
        for hh_id in source_hh_ids:
            new_hh_id = self.P.duplicate_household(t, hh_id)
            immigrants.extend(self.P.groups['household'][new_hh_id])
        
        return births, deaths, immigrants, birthdays



    def record_stats_demo(self, t):
        if self.params['record_interval'] <= 0 or t%self.params['record_interval'] is not 0:
            return
        
        self.pop_size.append(len(self.P.I))
        self.age_dist.append(self.P.age_dist(101,101)[0])
        self.hh_size_dist.append(self.P.group_size_dist('household', self.max_hh)[0])
        self.hh_size_dist_counts.append(self.P.group_size_dist('household',self.max_hh,False)[0])
        self.hh_size_avg.append(self.P.group_size_avg('household'))
        self.hh_count.append(len(self.P.groups['household']))

        self.fam_types.append(self.P.sum_hh_stats_group())



    def run(self):
        """
        Run simulated population (NB: this probably won't be used, as
        external running object may wish to just call update_all itself...)
        """

        timesteps = self.params['years'] * (365/self.params['t_dur'])
        i = 0
        while i < timesteps:
            t = i*self.params['t_dur']
            print i, t
            self.update_all_demo(t)
            i += 1