Represents a tensor of Geometries. Instance dimensions represent geometry unions and are reduced.
| 22 | @sliceable(keepdims='vector') |
| 23 | @dataclass(frozen=True, eq=False) |
| 24 | class GeometryStack(Geometry): |
| 25 | """ |
| 26 | Represents a tensor of Geometries. |
| 27 | Instance dimensions represent geometry unions and are reduced. |
| 28 | """ |
| 29 | geometries: Tensor |
| 30 | |
| 31 | def __post_init__(self): |
| 32 | assert isinstance(self.geometries, Tensor) and self.geometries.dtype.kind == object, f"geometries must be a Tensor of geometries but got {self.geometries}" |
| 33 | ranks = wrap(math.map(lambda g: g.spatial_rank, self.geometries, dims=object)) |
| 34 | assert ranks.min == ranks.max, f"Can only stack geometries of the same spatial rank but got ranks {ranks}" |
| 35 | |
| 36 | @cached_property |
| 37 | def shape(self) -> Shape: |
| 38 | return self.geometries.shape |
| 39 | |
| 40 | @property |
| 41 | def sets(self) -> Dict[str, Shape]: |
| 42 | all_names = set() |
| 43 | for g in self.geometries: |
| 44 | all_names.update(g.sets) |
| 45 | result = {} |
| 46 | for name in all_names: |
| 47 | all_dims = [] |
| 48 | for g in self.geometries: |
| 49 | all_dims.append(g.sets.get(name, EMPTY_SHAPE)) |
| 50 | all_dims = merge_shapes(all_dims, allow_varying_sizes=True) |
| 51 | zeros = all_dims.with_sizes(0) |
| 52 | set_shapes = [] |
| 53 | for g in self.geometries: |
| 54 | set_shapes.append(g.sets.get(name, zeros)) |
| 55 | set_shape = shape_stack(object_dims(self.geometries), *set_shapes) |
| 56 | result[name] = set_shape |
| 57 | return result |
| 58 | |
| 59 | @property |
| 60 | def object_dims(self): |
| 61 | return object_dims(self.geometries) |
| 62 | |
| 63 | def unstack(self, dimension) -> tuple: |
| 64 | if dimension == self.geometries.shape.name: |
| 65 | return tuple(self.geometries) |
| 66 | raise NotImplementedError() |
| 67 | |
| 68 | def _bounding_box(self): |
| 69 | boxes = math.map(bounding_box, self.geometries, dims=object) |
| 70 | return boxes.largest(instance) |
| 71 | |
| 72 | @property |
| 73 | def center(self) -> Tensor: |
| 74 | return self._bounding_box().center |
| 75 | |
| 76 | @property |
| 77 | def weighted_center(self): |
| 78 | centers = math.map(lambda g: g.center, self.geometries, dims=object) |
| 79 | if not instance(centers): |
| 80 | return centers |
| 81 | # --- volume-weighted mean over instance dimensions --- |
no outgoing calls
no test coverage detected