microecs.pool

 1"""pool.py - A pool of entities of the same type (same list of components). Basically a dynamic array with numpy"""
 2import numpy as np
 3from .utils import Shape, logger
 4
 5# Note: if Pool gets new fields, add them here! Otherwise the user code may overwrite them.
 6_POOL_INTERNAL_ATTRS = {"size", "capacity", "fields", "shapes", "dtypes", "data", "fields_set"}
 7
 8class Pool:
 9    """
10    Pool is a dynamic array of entities data given a list of fields, shapes and dtypes (from traits).
11    Pool has no concept of entity ids.
12    """
13    INITIAL_CAPACITY = 100
14
15    def __init__(self, fields: list[str], shapes: list[Shape], dtypes: list[np.dtype]):
16        if not len(fields) == len(shapes) == len(dtypes):
17            raise ValueError(f"Lens not match: {len(fields)} - {len(shapes)} - {len(dtypes)}")
18        if (fields_set := set(fields)) & POOL_RESERVED_NAMES:
19            raise ValueError(f"One of {fields=} in {POOL_RESERVED_NAMES}")
20
21        self.fields = fields
22        self.shapes = shapes
23        self.dtypes = dtypes
24        self.fields_set = fields_set # useful for fast checking (e.g. in entity)
25
26        self.data: dict[str, np.ndarray] = {} # the actual data is stored in dict of dynamic arrrays, one per field
27        self.size = 0
28        self.capacity = Pool.INITIAL_CAPACITY
29        for _field, shape, dtype in zip(fields, shapes, dtypes):
30            self.data[_field] = np.empty(shape=(self.capacity, *shape), dtype=dtype)
31
32    def add_entity(self, entity_data: dict[str, np.ndarray]) -> int:
33        """Adds an entity to the pool. All the fields required by this pool must be provided in entity_data"""
34        if self.size == self.capacity:
35            self._realloc(self.capacity * 2)
36            logger.debug(f"Capacity extended from {self.capacity // 2} to {self.capacity}")
37
38        for _field, field_shape, field_dtype in zip(self.fields, self.shapes, self.dtypes):
39            new_item = entity_data[_field] # checked in World._get_entity_pool(entity).
40            if (dtp := new_item.dtype) != field_dtype or new_item.shape != field_shape:
41                raise ValueError(f"Field {_field}. Dtype: {dtp} {field_dtype=}. Shape: {new_item.shape} {field_shape=}")
42            self.data[_field][self.size] = new_item
43        self.size += 1
44        return self.size - 1
45
46    def remove_entity(self, entity_index: int):
47        """removes an entity given an index (NOT ID) inside this pool"""
48        if not 0 <= entity_index < self.size:
49            raise IndexError(f"OOB: {entity_index=}, {self.size=}")
50        for _field in self.fields:
51            self.data[_field][entity_index] = self.data[_field][self.size - 1]
52        self.size -= 1
53
54        if self.size < self.capacity / 4 and self.capacity > Pool.INITIAL_CAPACITY:
55            self._realloc(self.capacity // 2)
56
57    def pop_entity(self, entity_index: int) -> dict[str, np.ndarray]:
58        """pops an entity given an index (NOT ID) inside this pool and returns the data"""
59        res = {_field: self.data[_field][entity_index].copy() for _field in self.fields}
60        self.remove_entity(entity_index)
61        return res
62
63    def _realloc(self, new_capacity: int):
64        for _field, shape, dtype in zip(self.fields, self.shapes, self.dtypes):
65            old_data = self.data[_field]
66            self.data[_field] = np.empty(shape=(new_capacity, *shape), dtype=dtype)
67            self.data[_field][0:self.size] = old_data[0:self.size]
68        self.capacity = new_capacity
69
70    def __getattr__(self, name):
71        if (data := self.__dict__.get("data")) is not None and name in data:
72            return data[name][0: self.size]
73        raise AttributeError(name)
74
75    def __setattr__(self, name, value):
76        if (data := self.__dict__.get("data")) is not None and name in data:
77            raise ValueError(f"Cannot explicitly set anything to Pool. Use `pool.component[:] = ...` ({name=})")
78        super().__setattr__(name, value)
79
80    def __len__(self):
81        return self.size
82
83POOL_RESERVED_NAMES = _POOL_INTERNAL_ATTRS | {n for n in vars(Pool) if not n.startswith("__")}