Adds Generator to start without dataset, including yaml and batch runner options; Starts adding a priority based Queue reordering.

sabvdf · sabvdf · commit cbaeaeb7bdcc · 2025-06-26T18:16:13.000+02:00
diff --git a/.vscode/extensions.json b/.vscode/extensions.json
@@ -5,6 +5,7 @@
         "github.vscode-github-actions",
         "fnando.linter",
         "ms-python.black-formatter",
-        "virtualplay.vibrant-semantics-dark"
+        "virtualplay.vibrant-semantics-dark",
+        "ms-toolsai.datawrangler"
     ]
-}
+}
diff --git a/datasim/generator.py b/datasim/generator.py
@@ -1,3 +1,176 @@
+from typing import Any, Dict, Final, Optional, Tuple, Type
+from typing_extensions import Literal
+from git import List
+import numpy as np
+
+from .types import Number, Value
+
+
+class Sampler:
+    property: str
+
+    def __init__(self, property: str):
+        self.property = property
+
+    @staticmethod
+    def _from_yaml(property: str, params: Dict) -> "Sampler":
+        if isinstance(params, Value):
+            return StaticSampler(property, params, start=params)
+        elif isinstance(params, Dict):
+            if "value" in params:
+                return StaticSampler(
+                    property,
+                    params["value"],
+                    params.get("sample", "independent") == "accumulate",
+                    params.get("start", None),
+                )
+
+            if "distribution" in params:
+                return DistributionSampler(
+                    property,
+                    params["distribution"],
+                    params.get("parameters", {}),
+                    params.get("sample", "independent"),
+                )
+
+            return StaticSampler(property, None)
+
+    def next(self) -> Value:
+        return None
+
+
+class StaticSampler(Sampler):
+    value: Value
+    accumulate: bool
+    step: Value
+
+    def __init__(
+        self, property: str, value: Value, accumulate: bool = False, start: Value = None
+    ):
+        super().__init__(property)
+        self.accumulate = accumulate
+        self.value = start if start else 0.0 if isinstance(value, float) else 0
+        self.step = value
+
+    def next(self) -> Value:
+        if (
+            self.accumulate
+            and isinstance(self.value, int | float)
+            and isinstance(self.step, int | float)
+        ):
+            self.value += self.step
+        return self.value
+
+
+class DistributionSampler(Sampler):
+    value: float
+    accumulate: bool
+    rng: np.random.Generator
+    np_function: Any
+    parameters: Dict
+
+    def __init__(
+        self,
+        property: str,
+        np_generator: str,
+        parameters: Dict,
+        accumulation: Literal["independent", "accumulate"],
+        start: Optional[float] = None,
+    ):
+        super().__init__(property)
+        self.value = start if start else 0.0
+        self.accumulate = accumulation == "accumulate"
+        self.rng = np.random.default_rng()
+        self.np_function = getattr(self.rng, np_generator)
+        self.parameters = parameters
+
+    def next(self) -> Value:
+        sample = self.np_function(**self.parameters)
+        if not self.accumulate:
+            return sample
+
+        self.value += sample
+        return self.value
+
+
 class Generator:
-    def __init__(self):
-        pass
+    id: Final[str]
+    data_class: Final[str]
+    subset_key: Final[str]
+    subsets: Final[Dict]
+
+    def __init__(
+        self, world, id: str, data_class: str, subset_key: str, subsets: List[Dict]
+    ):
+        self.id = id
+        self.data_class = data_class
+        self.subset_key = subset_key
+        self.subsets = {}
+        for subset in subsets:
+            self.subsets[subset[subset_key]] = subset
+
+        world.add(self)
+
+    @staticmethod
+    def _from_yaml(world, params: Dict) -> "Generator":
+        id = list(params.keys())
+        if len(id) > 1:
+            raise ValueError(f"Unable to parse yaml: Multiple keys found in {params}")
+
+        id = id[0]
+        params = params[id]
+
+        return Generator(world, id, params["class"], params["key"], params["subsets"])
+
+    def generate(
+        self,
+        type: Type,
+        limits: Dict[str, Tuple[Literal["<", ">"], Number]] = {},
+        counts: Dict[Tuple[str, Value], int] = {},
+        sort: Optional[str] = None,
+        sort_direction: Literal["asc", "desc"] = "asc",
+    ) -> List:
+        data = []
+        id = 0
+
+        for subset in self.subsets.values():
+            samplers: Dict[str, Sampler] = {}
+            for property, parameters in subset.items():
+                samplers[property] = Sampler._from_yaml(property, parameters)
+
+            count = 0
+            for (key, value), max in counts.items():
+                if subset.get(key, None) == value:
+                    count = max
+                    break
+
+            while True:
+                id += 1
+                next = type()
+                next.id = f"{id}"
+                for property, sampler in samplers.items():
+                    setattr(next, property, sampler.next())
+
+                for key, (compare, limit) in limits.items():
+                    value = getattr(next, key, None)
+                    if value and (
+                        (compare == ">" and value > limit)
+                        or (compare == "<" and value < limit)
+                    ):
+                        break
+
+                else:
+                    data.append(next)
+
+                    if count > 0:
+                        count -= 1
+                        if count == 0:
+                            break
+                    continue
+
+                break
+
+        if sort:
+            data.sort(key=lambda o: getattr(o, sort), reverse=sort_direction == "desc")
+
+        return data
diff --git a/datasim/queue.py b/datasim/queue.py
@@ -195,18 +195,17 @@ def peek_with_amount(self) -> Tuple[EntityType, Number] | None:
 
         return self.queue[-1]
 
-    def prioritize(self, entity: EntityType) -> bool:
-        """Pushes an entity to the front of the list.
-
-        If the entity was not in the list, it will not be added;
-        If the entity is in the list more than once, the copy furthest to the back will be put at the front.
+    def queue_prioritized(self, entity: EntityType, sort_function) -> bool:
+        """Pushes an entity to a sorted place in the list
 
         Args:
             entity (Entity): _description_
+            sort_function (lambda): _function that evaluates to a __gt__ comparable type
         """
         (_, entry) = [
             (i, (e, a)) for i, (e, a) in enumerate(self.queue) if e is entity
         ][0]
+        # TODO: change
         if self.queue.remove(entry):
             self.queue.append(entry)
             self.changed_tick = self.world.ticks
diff --git a/datasim/world.py b/datasim/world.py
@@ -8,10 +8,11 @@
 from typing import Any, Dict, Final, List, Optional, Tuple
 
 from .constant import Constant
-from .output import Output
+from .dataset import DataFrameData, Dataset, DataSource
 from .entity import Entity
+from .generator import Generator
 from .logging import log, LogLevel
-from .dataset import DataFrameData, Dataset, DataSource
+from .output import Output
 from .quantity import Quantity
 from .queue import Queue
 from .resource import Resource
@@ -36,6 +37,7 @@ class World(ABC):
     _entity_registry: Final[dict[type, int]] = {}
     datasets: Final[Dict[str, Dataset]]
     constants: Final[Dict[str, Any]]
+    generators: Final[Dict[str, Generator]]
     resources: Final[Dict[str, Resource]]
     queues: Final[Dict[str, Queue]]
     quantities: Final[Dict[str, Quantity]]
@@ -119,6 +121,7 @@ def __init__(
         self._entity_dict = {}
         self.datasets = {}
         self.constants = {}
+        self.generators = {}
         self.resources = {}
         self.queues = {}
         self.quantities = {}
@@ -137,6 +140,10 @@ def __init__(
                 for constant in definition["constants"]:
                     Constant._from_yaml(self, constant)
 
+            if "generators" in definition:
+                for generator in definition["generators"]:
+                    Generator._from_yaml(self, generator)
+
             if "resources" in definition:
                 for resource in definition["resources"]:
                     Resource._from_yaml(self, resource)
@@ -158,7 +165,7 @@ def reset(self):
         """Reset the World so you can start a different simulation."""
         self.active = False
 
-    def add(self, obj: Constant | Entity | Resource | Queue | Quantity):
+    def add(self, obj: Constant | Generator | Entity | Resource | Queue | Quantity):
         """Add an entity to this :class:`World`.
 
         Args:
@@ -176,6 +183,8 @@ def add(self, obj: Constant | Entity | Resource | Queue | Quantity):
 
         if isinstance(obj, Constant):
             self.constants[obj.id] = obj
+        elif isinstance(obj, Generator):
+            self.generators[obj.id] = obj
         elif isinstance(obj, Entity):
             self.entities.append(obj)
             self._entity_dict[obj.id] = obj
@@ -186,7 +195,9 @@ def add(self, obj: Constant | Entity | Resource | Queue | Quantity):
         elif isinstance(obj, Quantity):
             self.quantities[obj.id] = obj
 
-    def remove(self, obj: Constant | Entity | Resource | Queue | Quantity) -> bool:
+    def remove(
+        self, obj: Constant | Generator | Entity | Resource | Queue | Quantity
+    ) -> bool:
         """Remove an entity from this :class:`World`.
 
         Args:
@@ -200,6 +211,8 @@ def remove(self, obj: Constant | Entity | Resource | Queue | Quantity) -> bool:
 
             if isinstance(obj, Constant):
                 self.constants.pop(obj.id)
+            elif isinstance(obj, Generator):
+                self.generators.pop(obj.id)
             elif isinstance(obj, Entity):
                 self.entities.remove(obj)
                 self._entity_dict.pop(obj.id)
@@ -223,17 +236,28 @@ def _set_variation(self, selector: str, value: Value):
         obj_path = selector.split(".")
         current = self
         for path_part in obj_path[:-1]:
-            current = getattr(current, path_part)
+            current = (
+                current.get(path_part)
+                if isinstance(current, dict)
+                else getattr(current, path_part)
+            )
 
-        destination = getattr(current, obj_path[-1])
+        destination = (
+            current.get(obj_path[-1])
+            if isinstance(current, dict)
+            else getattr(current, obj_path[-1])
+        )
 
         if (
             destination is None
             or isinstance(destination, int)
             or isinstance(destination, float)
             or isinstance(destination, str)
         ):
-            setattr(current, obj_path[-1], value)
+            if isinstance(current, dict):
+                current[obj_path[-1]] = value
+            else:
+                setattr(current, obj_path[-1], value)
         elif isinstance(destination, Constant):
             destination.value = value
 
diff --git a/examples/icu/icu.py b/examples/icu/icu.py
@@ -3,7 +3,7 @@
 
 import pandas as pd
 
-from datasim import Quantity, Queue, Resource, StateData, World
+from datasim import log, LogLevel, Quantity, Queue, Resource, StateData, World
 from .patient import DiedPatientState, PatientData, Patient, TreatedPatientState
 
 
@@ -33,12 +33,21 @@ def __init__(
             variation_dict=variation_dict,
         )
 
-        self.load_patient_data("examples/icu/simulatiedata.csv")
+        # Option: load from CSV or simulate input using Generator
+        # self.load_patient_data("examples/icu/simulatiedata.csv")
+        self.generate_patient_data(500)
 
     def load_patient_data(self, filename: str):
         self.patients = []
         for row in list(reader(open(filename)))[1:]:
             self.patients.append(PatientData(row))
+        log(f"Loaded data for {len(self.patients)} patients", LogLevel.debug)
+
+    def generate_patient_data(self, end_enter_time: float):
+        self.patients = self.generators["patient_generator"].generate(
+            PatientData, limits={"enter_time": (">", end_enter_time)}, sort="enter_time"
+        )
+        log(f"Generated data for {len(self.patients)} patients", LogLevel.debug)
 
     def remove(self, obj):
         if isinstance(obj, Patient):
@@ -69,7 +78,11 @@ def before_entities_update(self):
             next = self.patients_waiting.peek()
 
     def after_entities_update(self):
-        if len(self.patients_waiting) == 0 and len(self.entities) == 0:
+        if (
+            len(self.patients_waiting) == 0
+            and len(self.entities) == 0
+            and len(self.patients) == 0
+        ):
             self.stopped = True
 
     def aggregate_data(self) -> Dict[str, pd.DataFrame]:
diff --git a/examples/icu/icu.yaml b/examples/icu/icu.yaml
diff --git a/examples/icu/patient.py b/examples/icu/patient.py
