gninatorch.dataloaders module

class gninatorch.dataloaders.GriddedExamplesLoader(example_provider, grid_maker, label_pos: int = 0, affinity_pos: int | None = None, flexlabel_pos: int | None = None, random_translation: float = 0.0, random_rotation: bool = False, device: device = device(type='cpu'), grids_only: bool = False)[source]

Bases: object

Load example and compute atomic density on a grid.

Parameters:
  • example_provider (molgrid.ExampleProvider) – :package:`molgrid` example provider

  • grid_maker (molgrid.GridMaker) – :package:`molgrid` grid maker

  • label_pos (int) – Ligand pose annotation label position

  • affinity_pos (Optional[int]) – Affinity annotation position

  • flexlabel_pos (Optional[int]) – Receptor (side chains) pose annotation label position

  • random_translation (float) – Random translation applied to each example on each cartesian axis

  • random_rotation (bool) – Uniform random rotation applied to each example

  • device (torch.device) – Device

  • grid_only (bool) – If True, return only the grid, otherwise return grid and labels

Notes

The batch size is defined in the example_provider as default_batch_size. The number of batches actually depend on the molgrid.IterationScheme used, also defined in the example_provider.

If molgrid.IterationScheme.SmallEpoch is used, examples are seen at most once. If molgrid.IterationScheme.LargeEpoch is used, examples are seen at least once.

The last batch is not padded with examples of the next epoch, in contrast with molgrid.ExampleProvider default behaviour.