Data

Dataset handling and data loading utilities.

DataGroup

DataGroup(data=None, keys=None, shard=None)

Dict-like container for data arrays with consistent shape.

Args:

data (str | Dict[str, np.ndarray] | None): The data to be used in the dataset. It can be a string representing the path to the data file in NPZ format or a dictionary where keys are strings and values are numpy arrays.

keys (List[str] | None): A list of keys to be used from the data dictionary.

shard (Tuple[int, int] | None): A tuple representing the shard index and total shards.

Source code in aimnet/data/sgdataset.py

def __init__(
    self,
    data: str | dict[str, np.ndarray] | h5py.Group | None = None,
    keys: list[str] | None = None,
    shard: tuple[int, int] | None = None,
):
    # main container for data
    self._data: dict[str, np.ndarray] = {}

    if data is None:
        data = {}

    s = slice(shard[0], None, shard[1]) if shard is not None else slice(None)

    # load  to dict
    if isinstance(data, str):
        if not os.path.isfile(data):
            raise FileNotFoundError(f"{data} does not exist or not a file.")
        data = np.load(data, mmap_mode="r")
        if not hasattr(data, "files"):
            raise ValueError(f"Data file {data} does not contain named arrays.")

    # take only keys
    selected_keys = set(keys) if keys is not None else set(data.keys())  # type: ignore[union-attr]
    data = {k: v[s] for k, v in data.items() if k in selected_keys}  # type: ignore[union-attr]

    # check data
    _n = None
    for k, v in data.items():
        if not isinstance(k, str):
            raise TypeError(f"Expected key to be of type str, but got {type(k).__name__}")
        if keys is not None and k not in keys:
            continue
        if _n is None:
            _n = len(v)
        if len(v) != _n:
            raise ValueError(f"Inconsistent data shape for key {k}. Expected first dimension {_n}, got {len(v)}.")
        self._data[k] = v

cv_split(cv=5, seed=None)

Return list of cv tuples containing train and val DataGroups

Source code in aimnet/data/sgdataset.py

def cv_split(self, cv: int = 5, seed=None):
    """Return list of `cv` tuples containing train and val `DataGroup`s"""
    fractions = [1 / cv] * cv
    parts = self.random_split(*fractions, seed=seed)
    splits = []
    for icv in range(cv):
        val = parts[icv]
        train_parts = [parts[_i] for _i in range(cv) if _i != icv]
        # Build a fresh DataGroup: cat() mutates its receiver in place, which
        # would corrupt the shared parts for the remaining folds.
        train = self.__class__({k: np.concatenate([p[k] for p in train_parts], axis=0) for k in self.keys()})
        splits.append((train, val))
    return splits

merge(other, strict=True)

Concatenate other into this group along the sample axis.

With strict=True both groups must carry exactly the same data keys. With strict=False the result keeps only the keys they share: keys present here but missing from other are dropped. Keys present in other but missing here are ignored in both modes -- this only ever grows the sample axis, never the key set. other is not modified.

Source code in aimnet/data/sgdataset.py

def merge(self, other, strict=True):
    """Concatenate `other` into this group along the sample axis.

    With `strict=True` both groups must carry exactly the same data keys.
    With `strict=False` the result keeps only the keys they share: keys
    present here but missing from `other` are dropped. Keys present in
    `other` but missing here are ignored in both modes -- this only ever
    grows the sample axis, never the key set. `other` is not modified.
    """
    if strict:
        if set(self.keys()) != set(other.keys()):
            raise ValueError("Data keys do not match between the datasets.")
        keys = self.keys()
    else:
        keys = set(self.keys()) & set(other.keys())
    for k in list(self.keys()):
        if k in keys:
            self._data[k] = np.concatenate([self[k], other[k]], axis=0)
        else:
            self.pop(k)

rename_key(old, new)

Rename data key old to new in place.

Source code in aimnet/data/sgdataset.py

def rename_key(self, old, new):
    """Rename data key `old` to `new` in place."""
    self[new] = self.pop(old)

sample(idx, keys=None)

Return a new DataGroup with the data indexed by idx.

Source code in aimnet/data/sgdataset.py

def sample(self, idx, keys=None) -> "DataGroup":
    """Return a new `DataGroup` with the data indexed by `idx`."""
    if keys is None:
        keys = self.keys()
    if isinstance(idx, int):
        idx = slice(idx, idx + 1)
    return self.__class__({k: self[k][idx] for k in keys})

to_dict()

Return the live {key: array} mapping backing this group.

This is the internal dict, not a copy: mutating it mutates the group.

Source code in aimnet/data/sgdataset.py

def to_dict(self):
    """Return the live `{key: array}` mapping backing this group.

    This is the internal dict, not a copy: mutating it mutates the group.
    """
    return self._data

SizeGroupedDataset

SizeGroupedDataset(data=None, keys=None, shard=None)

Source code in aimnet/data/sgdataset.py

def __init__(
    self,
    data: str | list[str] | dict[int, dict[str, np.ndarray]] | dict[int, DataGroup] | None = None,
    keys: list[str] | None = None,
    shard: tuple[int, int] | None = None,
):
    # main containers
    self._data: dict[int, DataGroup] = {}
    self._meta: dict[str, str] = {}

    # load data
    if isinstance(data, str):
        if os.path.isdir(data):
            self.load_datadir(data, keys=keys, shard=shard)
        else:
            self.load_h5(data, keys=keys, shard=shard)
    elif isinstance(data, (list, tuple)):
        self.load_files(data, shard=shard)
    elif isinstance(data, dict):
        self.load_dict(data)
    self.loader_mode = False
    self.x: list[str] = []
    self.y: list[str] = []

merge(other, strict=True)

Merge other into this dataset in place.

Size groups present in both are concatenated; groups only in other are adopted as-is. With strict=True both datasets must carry exactly the same data keys.

Two side effects to be aware of, both long-standing:

• With strict=False, other is also reduced in place to the shared key set. Pass a copy if you still need it intact.
• An adopted group is taken by reference, not copied, so that group is then shared between the two datasets.

Source code in aimnet/data/sgdataset.py

def merge(self, other, strict=True):
    """Merge `other` into this dataset in place.

    Size groups present in both are concatenated; groups only in `other`
    are adopted as-is. With `strict=True` both datasets must carry exactly
    the same data keys.

    Two side effects to be aware of, both long-standing:

    - With `strict=False`, `other` is also reduced in place to the shared
      key set. Pass a copy if you still need it intact.
    - An adopted group is taken by reference, not copied, so that group is
      then shared between the two datasets.
    """
    if not isinstance(other, self.__class__):
        other = self.__class__(other)
    if strict:
        if set(other.datakeys()) != set(self.datakeys()):
            raise ValueError("Data keys do not match between the datasets.")
    else:
        keys = set(other.datakeys()) & set(self.datakeys())
        for k in list(self.datakeys()):
            if k not in keys:
                for g in self.groups:
                    g.pop(k)
        for k in list(other.datakeys()):
            if k not in keys:
                for g in other.groups:
                    g.pop(k)
    # NB: iterate `other.keys()` explicitly. `SizeGroupedDataset` defines
    # `__getitem__` but no `__iter__`, so a bare `for k in other` falls back
    # to the legacy integer sequence protocol: it probes `other[0]`,
    # `other[1]`, ... and stops only on `IndexError`. Keys here are group
    # sizes, so the first size missing from a gap-free run starting at 0
    # raises `KeyError` instead -- i.e. on every realistic dataset.
    for k in list(other.keys()):
        if k in self:
            self[k].cat(other[k])  # type: ignore[attr-defined]
        else:
            self[k] = other[k]  # type: ignore[attr-defined]

rename_datakey(old, new)

Rename data key old to new across every size group.

Source code in aimnet/data/sgdataset.py

def rename_datakey(self, old, new):
    """Rename data key `old` to `new` across every size group."""
    for g in self.groups:
        g.rename_key(old, new)

SizeGroupedSampler

SizeGroupedSampler(ds, batch_size, batch_mode='molecules', shuffle=False, batches_per_epoch=-1, seed=None)

Source code in aimnet/data/sgdataset.py

def __init__(
    self,
    ds: SizeGroupedDataset,
    batch_size: int,
    batch_mode: str = "molecules",
    shuffle: bool = False,
    batches_per_epoch: int = -1,
    seed: int | None = None,
):
    self.ds = ds
    self.batch_size = batch_size
    if batch_mode not in ["molecules", "atoms"]:
        raise ValueError(f"Unknown batch_mode {batch_mode}")
    self.batch_mode = batch_mode
    self.shuffle = shuffle
    self.batches_per_epoch = batches_per_epoch
    self.seed = seed