AtomsDataset and DataModule Refactor

.gitignore

@@ -128,4 +128,4 @@ interfaces/lammps/examples/*/*.dat
 interfaces/lammps/examples/*/deployed_model
 
 # batchwise optimizer examples
-examples/howtos/howto_batchwise_relaxations_outputs/*
+examples/howtos/howto_batchwise_relaxations_outputs/*

src/schnetpack/cli.py

@@ -17,7 +17,7 @@
 import schnetpack as spk
 from schnetpack.utils import str2class
 from schnetpack.utils.script import log_hyperparameters, print_config
-from schnetpack.data import BaseAtomsData, AtomsLoader
+from schnetpack.data import ASEAtomsData, AtomsLoader
 from schnetpack.train import PredictionWriter
 from schnetpack import properties
 from schnetpack.utils import load_model
@@ -178,14 +178,16 @@ def train(config: DictConfig):
 
     # Evaluate model on test set after training
     log.info("Starting testing.")
-    trainer.test(model=task, datamodule=datamodule, ckpt_path="best")
+    trainer.test(
+        model=task, datamodule=datamodule, ckpt_path="best", weights_only=False
+    )
 
     # Store best model
     best_path = trainer.checkpoint_callback.best_model_path
     log.info(f"Best checkpoint path:\n{best_path}")
 
     log.info(f"Store best model")
-    best_task = type(task).load_from_checkpoint(best_path)
+    best_task = type(task).load_from_checkpoint(best_path, weights_only=False)
     torch.save(best_task, config.globals.model_path + ".task")
 
     best_task.save_model(config.globals.model_path, do_postprocessing=True)
@@ -195,7 +197,7 @@ def train(config: DictConfig):
 @hydra.main(config_path="configs", config_name="predict", version_base="1.2")
 def predict(config: DictConfig):
     log.info(f"Load data from `{config.data.datapath}`")
-    dataset: BaseAtomsData = hydra.utils.instantiate(config.data)
+    dataset: ASEAtomsData = hydra.utils.instantiate(config.data)
     loader = AtomsLoader(dataset, batch_size=config.batch_size, num_workers=8)
 
     model = load_model("best_model")

src/schnetpack/configs/data/ani1.yaml

@@ -1,16 +1,18 @@
+# @package data
 defaults:
   - custom
 
-_target_: schnetpack.datasets.ANI1
+dataset:
+  _target_: schnetpack.datasets.ANI1
+  datapath: ${run.data_dir}/ani1.db # data_dir is specified in train.yaml
+  num_heavy_atoms: 8
+  high_energies: false
+  distance_unit: Ang
+  property_units:
+    energy: eV
+  transforms: ${data.transforms}
+
 
-datapath: ${run.data_dir}/ani1.db  # data_dir is specified in train.yaml
 batch_size: 32
 num_train: 10000000
 num_val: 100000
-num_heavy_atoms: 8
-high_energies: False
-
-# convert to typically used units
-distance_unit: Ang
-property_units:
-  energy: eV

src/schnetpack/configs/data/custom.yaml

@@ -1,12 +1,25 @@
-_target_: schnetpack.data.AtomsDataModule
+# @package data
+_target_: schnetpack.data.datamodule_v2.AtomsDataModuleV2
+
+dataset:
+  _target_: schnetpack.data.ASEAtomsData
+  datapath: ???
+  load_properties: null
+  distance_unit: Ang
+  property_units: {}
+  transforms: ${data.transforms}
+  train_transforms: null
+  val_transforms: null
+  test_transforms: null
 
-datapath: ???
-data_workdir: null
 batch_size: 10
 num_train: ???
 num_val: ???
 num_test: null
+split_file: ${run.data_dir}/split.npz
+splitting: null
 num_workers: 8
-num_val_workers: null
-num_test_workers: null
-train_sampler_cls: null
+train_sampler_cls: null
+train_sampler_args: {}
+pin_memory: false
+

src/schnetpack/configs/data/iso17.yaml

@@ -1,10 +1,12 @@
+# @package data
 defaults:
   - custom
 
-_target_: schnetpack.datasets.ISO17
+dataset:
+  _target_: schnetpack.datasets.ISO17
+  datapath: ${run.data_dir}/${data.folder}.db # data_dir is specified in train.yaml
+  folder: reference
 
-datapath: ${run.data_dir}/${data.folder}.db  # data_dir is specified in train.yaml
-folder: reference
 batch_size: 32
 num_train: 0.9
 num_val: 0.1

src/schnetpack/configs/data/materials_project.yaml

@@ -1,10 +1,12 @@
+# @package data
 defaults:
   - custom
 
-_target_: schnetpack.datasets.MaterialsProject
+dataset:
+  _target_: schnetpack.datasets.MaterialsProject
+  datapath: ${run.data_dir}/materials_project.db # data_dir is specified in train.yaml
+  apikey: ???
 
-datapath: ${run.data_dir}/materials_project.db  # data_dir is specified in train.yaml
 batch_size: 32
 num_train: 60000
-num_val: 2000
-apikey: ???
+num_val: 2000

src/schnetpack/configs/data/md17.yaml

@@ -1,10 +1,14 @@
+# @package data
 defaults:
   - custom
 
-_target_: schnetpack.datasets.MD17
-
-datapath: ${run.data_dir}/${data.molecule}.db  # data_dir is specified in train.yaml
 molecule: aspirin
+
+dataset:
+  _target_: schnetpack.datasets.MD17
+  datapath: ${run.data_dir}/${data.molecule}.db  # data_dir is specified in train.yaml
+  molecule: ${data.molecule}
+
 batch_size: 10
 num_train: 950
 num_val: 50

src/schnetpack/configs/data/md22.yaml

@@ -1,10 +1,14 @@
+# @package data
 defaults:
   - custom
 
-_target_: schnetpack.datasets.MD22
-
-datapath: ${run.data_dir}/${data.molecule}.db  # data_dir is specified in train.yaml
 molecule: Ac-Ala3-NHMe
+
+dataset:
+  _target_: schnetpack.datasets.MD22
+  datapath: ${run.data_dir}/${data.molecule}.db
+  molecule: ${data.molecule}
+
 batch_size: 10
 num_train: 5700
 num_val: 300

src/schnetpack/configs/data/omdb.yaml

@@ -1,10 +1,12 @@
+# @package data
 defaults:
   - custom
 
-_target_: schnetpack.datasets.OrganicMaterialsDatabase
+dataset:
+  _target_: schnetpack.datasets.OrganicMaterialsDatabase 
+  datapath: ${run.data_dir}/omdb.db # data_dir is specified in train.yaml
+  raw_path: null
 
-datapath: ${run.data_dir}/omdb.db  # data_dir is specified in train.yaml
 batch_size: 32
 num_train: 0.8
 num_val: 0.1
-raw_path: null

src/schnetpack/configs/data/qm7x.yaml

@@ -1,9 +1,14 @@
+# @package data
 defaults:
   - custom
 
-_target_: schnetpack.datasets.QM7X
+dataset:
+  _target_: schnetpack.datasets.QM7X
+  datapath: ${run.data_dir}/qm7x.db  # data_dir is specified in train.yaml
+  remove_duplicates: true
+  only_equilibrium: false
+  only_non_equilibrium: false
 
-datapath: ${run.data_dir}/qm7x.db  # data_dir is specified in train.yaml
 batch_size: 100
 num_train: 5550
 num_val: 700

src/schnetpack/configs/data/qm9.yaml

@@ -1,22 +1,26 @@
+# @package data
 defaults:
   - custom
 
-_target_: schnetpack.datasets.QM9
+dataset:
+  _target_: schnetpack.datasets.qm9.QM9
+  datapath: ${run.data_dir}/qm9.db
+  remove_uncharacterized: true
+  load_properties: null
+  distance_unit: Ang
+  property_units:
+    energy_U0: eV
+    energy_U: eV
+    enthalpy_H: eV
+    free_energy: eV
+    homo: eV
+    lumo: eV
+    gap: eV
+    zpve: eV
+  transforms: ${data.transforms}
 
-datapath: ${run.data_dir}/qm9.db  # data_dir is specified in train.yaml
 batch_size: 100
 num_train: 110000
 num_val: 10000
-remove_uncharacterized: True
-
-# convert to typically used units
-distance_unit: Ang
-property_units:
-  energy_U0: eV
-  energy_U: eV
-  enthalpy_H: eV
-  free_energy: eV
-  homo: eV
-  lumo: eV
-  gap: eV
-  zpve: eV
+num_test: 10000
+num_workers: 2

src/schnetpack/configs/data/rmd17.yaml

@@ -1,11 +1,13 @@
+# @package data
 defaults:
   - custom
+molecule: aspirin
 
-_target_: schnetpack.datasets.rMD17
+dataset:
+  _target_: schnetpack.datasets.rMD17
+  datapath: ${run.data_dir}/rmd17_${data.molecule}.db  # data_dir is specified in train.yaml
+  molecule: ${data.molecule}
 
-datapath: ${run.data_dir}/rmd17_${data.molecule}.db  # data_dir is specified in train.yaml
-molecule: aspirin
 batch_size: 10
 num_train: 950
 num_val: 50
-split_id: null

src/schnetpack/data/__init__.py

@@ -4,3 +4,5 @@
 from .splitting import *
 from .datamodule import *
 from .sampler import *
+from .datamodule_v2 import *
+from .provider import *