diff --git a/World.py b/World.py index 0606711..091cdff 100644 --- a/World.py +++ b/World.py @@ -3,7 +3,7 @@ from simulation.workplace import Workplace, WorkplaceSize # from data_distribution import * from timeit import default_timer as timer -from typing import List +from typing import List, Literal, Any import numpy as np from multiprocessing import Pool @@ -257,21 +257,19 @@ def generate_province(self, province: str, n_processes: int = 12, verbose=3): # test this code later workplaces = [] while total_workers < len(people_that_work): - temp = np.random.choice([0, 1, 2, 3]) - match temp: - case 0: size = WorkplaceSize.SMALL - case 1: size = WorkplaceSize.MEDIUM - case 2: size = WorkplaceSize.LARGE - case 3: size = WorkplaceSize.EXTRA_LARGE + temp: Literal[0,1,2,3] = np.random.choice([0, 1, 2, 3]) + wp = Workplace(province=province, size=WorkplaceSize(temp)) assigned_people = [] for i in range(len(works_by_people)): if people_mask[i] == 0 and works_by_people[i] == temp: assigned_people.append(people_that_work[i]) people_mask[i] = 1 + if wp.number_of_people == len(assigned_people): + # salir del ciclo si ya se lleno el centro de trabajo + break + wp.fill_with_people(assigned_people) - wp = Workplace(province=province, size=size, - people_ids=assigned_people) wp_id = self.db.insert_one('Workplace', wp.serialize()).inserted_id total_workers += len(assigned_people) workplaces.append(wp_id) @@ -352,8 +350,20 @@ def province_execution(self, population_name: str, day: int, p: str): province = self.db.get_one("Province", {"_id": p}) if 'workplaces' in province: - for workplace in province['workplaces']: - pass + for workplace_id in province['workplaces']: + workplace_dict: dict[str, Any] = self.db.get_one('Workplace', + {'_id': workplace_id}) + workplace = Workplace.load_serialized(workplace_dict) + + pairs=self.generate_contacts( + workplace.workers_ids, + np.random.randint(0,workplace.number_of_people)) + + self.insert_pairs(day, + pairs, + population_name, + str(workplace_id)+' workplace', + province) if 'schools' in province: for sc_tp in province['schools']: @@ -378,7 +388,7 @@ def province_execution(self, population_name: str, day: int, p: str): for house in neighborhood: # print(house) - if not 'persons' in house: + if 'persons' not in house: continue persons.append(house['persons']) pairs = self.generate_contacts( diff --git a/database/mongodb_client.py b/database/mongodb_client.py index 8a043c7..ff5706c 100644 --- a/database/mongodb_client.py +++ b/database/mongodb_client.py @@ -1,5 +1,5 @@ from pymongo import MongoClient - +from typing import Any class MongoCRUD: def __init__(self, database_name): @@ -53,7 +53,7 @@ def get_data(self, collection_name, filter_query={}, projection_fields=None): # yield doc return list(collection.find(filter_query, projection_fields)) - def get_one(self, collection_name, filter_query={}, projection_fields=None): + def get_one(self, collection_name, filter_query={}, projection_fields=None)-> Any: """ Retrieve documents from a MongoDB collection based on a filter query and projection fields. diff --git a/simulation/workplace.py b/simulation/workplace.py index 880d6ba..909cc24 100644 --- a/simulation/workplace.py +++ b/simulation/workplace.py @@ -1,7 +1,7 @@ from enum import Enum -from typing import List +from typing import List, Dict, Any import numpy as np @@ -13,10 +13,19 @@ class WorkplaceSize(Enum): class Workplace: - def __init__(self, province: str, size: WorkplaceSize, people_ids: List): + def __init__(self, province: str, size: WorkplaceSize,): self.province = province - self.size = size - self.people_ids = people_ids + self.size_type = size + + self.workers_ids = [] + self.number_of_people=0 + self.__assign_people_by_size() + + def fill_with_people(self, people_ids: List): + assert len(self.workers_ids)+len(people_ids)>self.number_of_people,\ + f'se esta asignando demasiada gente a un centro de trabajo de tipo {self.size_type}' + + self.workers_ids += people_ids def serialize(self): """serialize object @@ -26,12 +35,21 @@ def serialize(self): """ return { 'province': self.province, - 'size': str(self.size), - 'people_ids': self.people_ids + 'size_type': self.size_type.value, + 'workers_ids': self.workers_ids } + @classmethod + def load_serialized(cls, data: Dict[str, Any]) -> "Workplace": + province: str = data["province"] + size_type: int = data["school_type"] + workplace: Workplace = cls(province, WorkplaceSize(size_type)) + workplace.workers_ids = data['workers_ids'] + workplace.number_of_people = len(workplace.workers_ids) + return workplace + def __assign_people_by_size(self): - match self.size: + match self.size_type: case WorkplaceSize.SMALL: self.number_of_people = np.random.randint(1, 10) case WorkplaceSize.MEDIUM: