microecs.qr_field

  1"""qr_field.py - Query Result Field. A single field that implements numpy's interface for interop"""
  2from typing import Callable, T
  3import numpy as np
  4from microecs.utils import Shape
  5
  6class QRField(np.lib.mixins.NDArrayOperatorsMixin):
  7    """
  8    Field is a single field (column) from a QueryResult object obtained from world.query(...).
  9    Parts may come from different Pools, so we try to make a contiguius-like view from discountinous arrays.
 10    """
 11    def __init__(self, parts: list[np.ndarray]):
 12        if len(parts) in (0, 1):
 13            raise ValueError("Cannot instantiate QRField with a single part. Use _QRArray for that optimized path.")
 14        self.parts = parts
 15        self._lens = [len(p) for p in self.parts]
 16        self.len = sum(self._lens)
 17        self.shape: Shape = (len(self), *self.parts[0].shape[1:])
 18        self.dtype = self.parts[0].dtype
 19        self._bounds: np.ndarray | None = None
 20
 21    def numpy(self) -> np.ndarray:
 22        """Creates a numpy array from the underlying pool parts. Note: guaranteed len(self.parts) >= 2 (init)"""
 23        return np.concatenate(self.parts)
 24
 25    def _chunk(self, x: T, i: int) -> T:
 26        if isinstance(x, QRField):
 27            return x.parts[i]
 28        if isinstance(x, np.ndarray) and x.ndim == len(self.shape) and x.shape[0] == self.len:
 29            if self._bounds is None: # lazy instantiate because it's expensive to do this unless needed.
 30                self._bounds = np.cumsum([0, *self._lens])
 31            return x[self._bounds[i]:self._bounds[i + 1]]
 32        return x
 33
 34    def _apply_fn_on_parts(self, fn: Callable, op_args: list, **kwargs):
 35        # op args can be 1 element (-qr.velocity), 2 elements (qr.position * 0.1), 3 elements (np.where(a, b, c)), etc.
 36        # all of them must be chunked based on how many we have in this Field so each subpart is called independently.
 37
 38        results = []
 39        for i, part in enumerate(self.parts):
 40            pool_args = [self._chunk(x, i) for x in op_args]
 41            part_result: QRField = fn(*pool_args, **kwargs)
 42            # we expect f(arr(N, ...)) -> arr(N, ...) where N = number of items in the pool
 43            # for e.g. np.linalg.norm(velocity, axis=1) should do (N, 2) -> (N, 1) so the first axis is preserved
 44            if len(part_result) != part.shape[0]:
 45                raise ValueError(f"Result: {part_result.shape} vs {part.shape}")
 46            results.append(part_result)
 47        return QRField(results)
 48
 49    def _row_level_mask(self, key) -> list[np.ndarray] | None:
 50        """Per-part row masks for `qr.f[mask]` (task 51), or None if `key` is not a boolean mask.
 51
 52        Accepted (all equal to numpy's entity-ROW selection, `qr.f.numpy()[mask]`):
 53        - a bool ndarray of shape `(n,)` or `(n, 1)` -- the `(n, 1)` form (natural for `qr.f == x`
 54          comparisons on a `(1,)` field) is squeezed to the entity axis. numpy's own `(n, 1)` boolean
 55          indexing is elementwise-flattening, so here the `(n, 1)` form is an ENTITY mask;
 56        - a bool QRField, e.g. `qr.velocity[:, 0] > 3` -- coerced via `.numpy()`.
 57
 58        Returns one bool mask slice per pool, aligned with `self.parts` (gather: `part[mask]`; scatter:
 59        `part[mask] = chunk` with `len(chunk) == mask.sum()`). Raises ValueError on a wrong length or
 60        ndim. Returns None for everything else -- including non-bool dtypes -- so a non-mask key falls
 61        through to the `_selects_axis0` guard and keeps its TypeError (a non-bool "mask" is user error,
 62        not a mask)."""
 63        if isinstance(key, QRField):
 64            key = key.numpy()
 65        if not isinstance(key, np.ndarray) or key.dtype != bool:
 66            return None
 67        if key.shape == (len(self), 1):
 68            key = key[:, 0]
 69        if key.shape != (len(self), ):
 70            raise ValueError(f"mask {key.shape} does not match {self.len} entities (field shape {self.shape})")
 71        res, i = [], 0
 72        for part in self.parts:
 73            res.append(key[i:i + len(part)])
 74            i += len(part)
 75        return res
 76
 77    @staticmethod
 78    def _selects_axis0(key) -> bool:
 79        """Does this key select on the entity axis (axis 0)? Only the leading index is examined; anything
 80        after it acts on axes >=1 and is passed straight through to each pool. Axis 0 is the one that
 81        matters because a QRField spans several pools (acts contiguous, isn't), so a key that picks,
 82        reorders or strides entities has no per-pool meaning.
 83
 84        raises:  [0]   [-1]   [0:2]   [::2]   [[0, 2]]   [0, 1]   [None]   [()]
 85        ok:      [:]   [...]   [:, 0]   [:, 1:3]   [..., None]   [:, 0, 2]   [mask] (via _row_level_mask)
 86
 87        - `[0:2]` / `[::2]` are slices but drop entities. Only `slice(None)` is whole.
 88        - `[None]` prepends an axis *before* the entity axis; `[()]` is rejected conservatively.
 89        - `[0, 1]` leads with an int, so it is one entity's component, not all entities'.
 90        - `[mask]`: a boolean entity-row mask HAS per-pool meaning, so both `__getitem__` and
 91          `__setitem__` intercept it via `_row_level_mask` BEFORE this guard (task 51) -- this guard
 92          itself never sees a mask.
 93        """
 94        if isinstance(key, tuple) and key:
 95            key = key[0]
 96        return not (key is Ellipsis or (isinstance(key, slice) and key == slice(None)))
 97
 98    def __array_ufunc__(self, ufunc, method, *inputs, out=None, **kwargs):
 99        """wrapper for elementwise (python) primitives, e.g. qr.position += 1"""
100        if method != "__call__":
101            return NotImplemented
102        if out is None:
103            return self._apply_fn_on_parts(ufunc, inputs, **kwargs)
104        if len(out) != 1 or not isinstance(out[0], QRField):
105            raise TypeError(out)
106
107        for i in range(len(self.parts)):
108            pool_args = [self._chunk(x, i) for x in inputs]
109            ufunc(*pool_args, out=out[0].parts[i], **kwargs)
110        return out[0]
111
112    def __array_function__(self, func: Callable, _types, args: list, kwargs: dict):
113        """wrapper for elementwise numpy functions, e.g. qr.velocity = np.where(mask, -qr.velocity, qr.velocity)"""
114        return self._apply_fn_on_parts(func, args, **kwargs)
115
116    # qr.position = <field | scalar | per-entity broadcast>   -> scatter through the views
117    def __setitem__(self, key, value):
118        # qr.f[mask] = v: numpy-parity scatter into the masked rows (task 51)
119        if (mask_parts := self._row_level_mask(key)) is not None:
120            if isinstance(value, QRField):
121                value = value[key]
122            items_per_part = [x.sum() for x in mask_parts] # e.g. [[False, True], [True, True, True]] -> [1, 3]
123            value_broadcasted = np.broadcast_to(value, (sum(items_per_part), *self.shape[1:])) # (sum([1,3]) = 4, *nf)
124            i = 0
125            # then, distribute parts of value_broadcasted to each of this field's parts
126            for part, len_mask, mask_part in zip(self.parts, items_per_part, mask_parts):
127                part[mask_part] = value_broadcasted[i:i+len_mask]
128                i += len_mask
129            return
130
131        if QRField._selects_axis0(key):
132            raise TypeError(f"Only batch writes are supported, e.g. `qr.attr = xxx` or `qr.attr[:, k] = xxx` "
133                            f"({key=} selects entities, which crosses pools). "
134                            f"For one entity use `world.get_entity(qr.entity_ids[i]).attr = xxx`")
135
136        if isinstance(value, QRField):
137            for i, part in enumerate(self.parts):
138                part[key] = value.parts[i]
139            return
140
141        # follow numpy's rules for broadcasting
142        views = [part[key] for part in self.parts] # per-pool destinations (views)
143        logical = (self.len, *views[0].shape[1:]) # the (N, *e) the user "sees"
144        full = np.broadcast_to(value, logical) # numpy rules: (*e,)/scalar fill, (N,*e) positional; raises otherwise
145        for v, chunk in zip(views, np.split(full, np.cumsum(self._lens)[:-1])):
146            v[:] = chunk
147
148    def __getitem__(self, key):
149        # qr.f[mask]: numpy-parity gather of the masked rows (task 51)
150        if (mask_parts := self._row_level_mask(key)) is not None:
151            # no dtype= kwarg: every part already shares the field's dtype (one component field across
152            # pools), so concatenate infers it -- object-dtype fields included
153            return np.concatenate([part[mask] for part, mask in zip(self.parts, mask_parts)])
154
155        if QRField._selects_axis0(key):
156            raise TypeError(("Only batch updates are supported, e.g. `qr.attr=xxx` or `qr.attr[:, k]=xxx`. Use "
157                            ".numpy() for a proper array. For entity-level ops use `world.get_entity(eid).attr=xxx`"))
158
159        return QRField([part[key] for part in self.parts])
160
161    def __iter__(self):
162        for part in self.parts:
163            yield from part
164
165    def __len__(self):
166        return self.len
167
168    def __repr__(self):
169        return f"[Field] Shape: {self.shape} (across {len(self.parts)} pools)"