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            assert len(part_result) == part.shape[0], f"Result: {part_result.shape} vs {part.shape}"
 45            results.append(part_result)
 46        return QRField(results)
 47
 48    @staticmethod
 49    def _selects_axis0(key) -> bool:
 50        """Does this key select on the entity axis (axis 0)? Only the leading index is examined; anything
 51        after it acts on axes >=1 and is passed straight through to each pool. Axis 0 is the one that
 52        matters because a QRField spans several pools (acts contiguous, isn't), so a key that picks,
 53        reorders or strides entities has no per-pool meaning.
 54
 55        raises:  [0]   [-1]   [0:2]   [::2]   [mask]   [[0, 2]]   [0, 1]   [None]   [()]
 56        ok:      [:]   [...]   [:, 0]   [:, 1:3]   [..., None]   [:, 0, 2]
 57
 58        - `[0:2]` / `[::2]` are slices but drop entities. Only `slice(None)` is whole.
 59        - `[None]` prepends an axis *before* the entity axis; `[()]` is rejected conservatively.
 60        - `[0, 1]` leads with an int, so it is one entity's component, not all entities'.
 61        """
 62        if isinstance(key, tuple) and key:
 63            key = key[0]
 64        return not (key is Ellipsis or (isinstance(key, slice) and key == slice(None)))
 65
 66    def __array_ufunc__(self, ufunc, method, *inputs, out=None, **kwargs):
 67        """wrapper for elementwise (python) primitives, e.g. qr.position += 1"""
 68        if method != "__call__":
 69            return NotImplemented
 70        if out is None:
 71            return self._apply_fn_on_parts(ufunc, inputs, **kwargs)
 72
 73        assert len(out) == 1 and isinstance(out[0], QRField), out
 74        for i in range(len(self.parts)):
 75            pool_args = [self._chunk(x, i) for x in inputs]
 76            ufunc(*pool_args, out=out[0].parts[i], **kwargs)
 77        return out[0]
 78
 79    def __array_function__(self, func: Callable, _types, args: list, kwargs: dict):
 80        """wrapper for elementwise numpy functions, e.g. qr.velocity = np.where(mask, -qr.velocity, qr.velocity)"""
 81        return self._apply_fn_on_parts(func, args, **kwargs)
 82
 83    # qr.position = <field | scalar | per-entity broadcast>   -> scatter through the views
 84    def __setitem__(self, key, value):
 85        if QRField._selects_axis0(key):
 86            raise TypeError(f"Only batch writes are supported, e.g. `qr.attr = xxx` or `qr.attr[:, k] = xxx` "
 87                            f"({key=} selects entities, which crosses pools). "
 88                            f"For one entity use `world.get_entity(qr.entity_ids[i]).attr = xxx`")
 89
 90        if isinstance(value, QRField):
 91            for i, part in enumerate(self.parts):
 92                part[key] = value.parts[i]
 93            return
 94
 95        # follow numpy's rules for broadcasting
 96        views = [part[key] for part in self.parts] # per-pool destinations (views)
 97        logical = (self.len, *views[0].shape[1:]) # the (N, *e) the user "sees"
 98        full = np.broadcast_to(value, logical) # numpy rules: (*e,)/scalar fill, (N,*e) positional; raises otherwise
 99        for v, chunk in zip(views, np.split(full, np.cumsum(self._lens)[:-1])):
100            v[:] = chunk
101
102    def __getitem__(self, key):
103        if QRField._selects_axis0(key):
104            raise TypeError(("Only batch updates are supported, e.g. `qr.attr=xxx` or `qr.attr[:, k]=xxx`. Use "
105                            ".numpy() for a proper array. For entity-level ops use `world.get_entity(eid).attr=xxx`"))
106
107        return QRField([part[key] for part in self.parts])
108
109    def __iter__(self):
110        for part in self.parts:
111            yield from part
112
113    def __len__(self):
114        return self.len
115
116    def __repr__(self):
117        return f"[Field] Shape: {self.shape} (across {len(self.parts)} pools)"