Workplace added to simulation.

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(

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.
 

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: