def add_work_relations()

in clutrr/actors/ancestry.py [0:0]


    def add_work_relations(self, w=0.3):
        """
        Policy of adding working relations:
        - Add w work locations
        - Divide the population into these w bins
        - Add works_at relation
        - Within each bin:
            - Assign m managers
        :return:
        """
        num_pop = len(self.family_data)
        pop_ids = self.family_data.keys()
        work_locations = random.sample(self.store.attribute_store['work']['options'], int(num_pop * w))
        node_ct = self.node_ct
        work_bins = {}
        pop_per_loc = num_pop // len(work_locations)
        for wl in work_locations:
            self.work_data[node_ct] = Entity(name=wl, etype='work')
            w = random.sample(pop_ids, pop_per_loc)
            pop_ids = list(set(pop_ids) - set(w))
            work_bins[wl] = {"id": node_ct, "w": w}
            node_ct+=1
        if len(pop_ids) > 0:
            work_bins[work_locations[-1]]["w"].extend(pop_ids)
        self.node_ct = node_ct
        for wl in work_locations:
            e_id = work_bins[wl]["id"]
            pops = work_bins[wl]["w"]
            for p in pops:
                edge = (e_id, p)
                if edge not in self.family:
                    self.family[edge] = {'family':'', 'work': []}
                if 'work' not in self.family[edge]:
                    self.family[edge]['work'] = []
                self.family[edge]['work'].append('works_at')
            # select manager
            manager = random.choice(pops)
            for p in pops:
                edge = (p, manager)
                if edge not in self.family:
                    self.family[edge] = {'family':'', 'work': []}
                if 'work' not in self.family[edge]:
                    self.family[edge]['work'] = []
                self.family[edge]['work'].append('manager')