Collect#

class Collect(keys, offset_keys_dict=None, **kwargs)[source]#

Bases: object

Final projection of a sample dict into the model-facing batch contract.

Keeps only the requested keys, builds per-sample offset count tensors, and concatenates feature groups (e.g. feat_keys -> feat) into the fused input tensors a backbone consumes. This is normally the last transform in a pipeline, run after ToTensor. Reads the keys named in keys and in every *_keys group; the length used for offset is read from the source key’s first dimension (coord by default). Returns a new dict containing only the collected outputs. Registered as Collect – use this string as the type in a transform=[...] config list.

Each *_keys keyword argument names a list of source keys; the suffix _keys is stripped and the named arrays are concatenated along dim=1 (cast to float) under the resulting key. For example feat_keys produces feat – whose total channel width sets the backbone in_channels – and offset_keys_dict maps each output key to the source key whose row count becomes its [N] offset tensor.

Parameters:

  • keys (str | Sequence[str]) – keys to copy through verbatim into the output dict. A bare string is wrapped into a single-element list.

  • offset_keys_dict (dict, optional) – mapping of output offset key to the source key whose number of rows is recorded as a length-1 tensor. Defaults to dict(offset="coord").

  • **kwargs – feature-group specifications. Every keyword ending in _keys (e.g. feat_keys=["coord", "energy"]) names a sequence of keys that are concatenated along the channel dimension into a key with the _keys suffix removed (e.g. feat).

Note

Requires that the source key behind each offset entry (coord by default) and every key referenced in keys/*_keys is present. The concatenated feature tensor’s channel width must match the model’s in_channels.

Example

>>> import numpy as np
>>> from pimm.datasets.transform import Collect, ToTensor
>>> data = {"coord": np.array([[1., 2., 3.], [4., 5., 6.]], dtype="f4"),
...         "energy": np.array([[10.], [20.]], dtype="f4")}
>>> data = ToTensor()(data)  # Collect concatenates tensors, so run after ToTensor
>>> out = Collect(keys=("coord",), feat_keys=("coord", "energy"))(data)
>>> sorted(out)  # only kept keys, plus the built feat / offset
['coord', 'feat', 'offset']
>>> out["feat"]  # coord (3 ch) and energy (1 ch) fused -> 4 channels
tensor([[ 1.,  2.,  3., 10.],
        [ 4.,  5.,  6., 20.]])
>>> out["offset"]  # number of points in this sample
tensor([2])
On this page