datasets.ANI1

class datasets.ANI1(*args: Any, **kwargs: Any)[source]

ANI1 benchmark database. This class adds convenience functions to download ANI1 from figshare and load the data into pytorch.

References

Parameters:

  • datapath – path to dataset

  • num_heavy_atoms – number of heavy atoms. (See ‘Table 1’ in Ref. [1])

  • high_energies – add high energy conformations. (See ‘Technical Validation’ of Ref. [1])

  • batch_size – (train) batch size

  • num_train – number of training examples

  • num_val – number of validation examples

  • num_test – number of test examples

  • split_file – path to npz file with data partitions

  • format – dataset format

  • load_properties – subset of properties to load

  • val_batch_size – validation batch size. If None, use test_batch_size, then batch_size.

  • test_batch_size – test batch size. If None, use val_batch_size, then batch_size.

  • transforms – Transform applied to each system separately before batching.

  • train_transforms – Overrides transform_fn for training.

  • val_transforms – Overrides transform_fn for validation.

  • test_transforms – Overrides transform_fn for testing.

  • num_workers – Number of data loader workers.

  • num_val_workers – Number of validation data loader workers (overrides num_workers).

  • num_test_workers – Number of test data loader workers (overrides num_workers).

  • property_units – Dictionary from property to corresponding unit as a string (eV, kcal/mol, …).

  • distance_unit – Unit of the atom positions and cell as a string (Ang, Bohr, …).