Revision/dataset to enable batch post processing

modelforge-curate/modelforge/curate/datasets/geom_qm9_curation.py

@@ -32,8 +32,7 @@ class GEOMQM9Curation(DatasetCuration):
     DOI for original qm9 dataset: 10.6084/m9.figshare.c.978904.v5
 
     The Geometric Ensemble Of Molecules (GEOM) dataset contains conformers for 133,000 species from the QM9 dataset.
-    Conformer structures were generated with CREST and GFN2-XTB.  Energies were evaluated using
-    DFT via ORCA 5.0.2  using the r2scan-3c functional and mTZVPP basis.
+    Conformer structures were generated with CREST and GFN2-XTB.  Energies were evaluated using xtb.
 
     Citation:   Axelrod, S., Gómez-Bombarelli, R.
                 GEOM, energy-annotated molecular conformations for property prediction and molecular generation.
@@ -151,7 +150,6 @@ def _process_downloaded(
             for conformer in all_records[data_key]["conformers"]:
                 xyz_temp = conformer["xyz"]
                 energy_temp = conformer["totalenergy"]
-
                 # parse xyz_temp to separate atomic numbers and positions
                 # stored in standard atomic_number, x, y, z format
                 atomic_numbers = np.array(xyz_temp)[:, 0]

modelforge-curate/modelforge/curate/datasets/scripts/curate_geom_qm9.py

@@ -17,8 +17,7 @@
 DOI for original qm9 dataset: 10.6084/m9.figshare.c.978904.v5
 
 The Geometric Ensemble Of Molecules (GEOM) dataset contains conformers for 133,000 species from the QM9 dataset.
-Conformer structures were generated with CREST and GFN2-XTB.  Energies were evaluated using
-DFT via ORCA 5.0.2  using the r2scan-3c functional and mTZVPP basis.
+Conformer structures were generated with CREST and GFN2-XTB.  Energies reported here correspond to xtb energy.
 
 Citation:   Axelrod, S., Gómez-Bombarelli, R.
             GEOM, energy-annotated molecular conformations for property prediction and molecular generation.
@@ -43,7 +42,7 @@ def main():
 
     # We'll want to provide some simple means of versioning
     # if we make updates to either the underlying dataset, curation modules, or parameters given to the code
-    version = "1.0"
+    version = "1.1"
     # version of the dataset to curate
     version_select = f"v_4"
     # Curate the test dataset with 1000 total conformers

modelforge-curate/modelforge/curate/examples/basic_usage.ipynb

@@ -1139,7 +1139,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.13.13"
+   "version": "3.14.5"
   }
  },
  "nbformat": 4,

modelforge-curate/modelforge/curate/examples/record_and_sourcedataset.ipynb

@@ -661,7 +661,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.11.8"
+   "version": "3.14.5"
   }
  },
  "nbformat": 4,

modelforge-curate/modelforge/curate/sourcedataset.py

@@ -51,9 +51,10 @@ def __init__(
             If True, use an existing database.
             If False, removes the existing database and creates a new one.
         """
+        from collections import OrderedDict
 
         self.name = name
-        self.records = {}
+        self.records = OrderedDict()
         self.append_property = append_property
         self.local_db_dir = local_db_dir
         os.makedirs(self.local_db_dir, exist_ok=True)

modelforge/potential/potential.py

@@ -1526,3 +1526,149 @@ def load_inference_model_from_checkpoint(
 
     # Return the model
     return potential
+
+
+def evaluate_batch(
+    dm: DataModule,
+    potential: Potential,
+    subset: Literal["train", "test", "val"],
+    include_metadata: bool = True,
+) -> dict:
+    """
+    This will evaluate the dataset subset using the specified potential.
+
+    This will perform batch processing (for efficiency). which requires we use the NeighborlistForTraining in the potential
+    by setting `use_training_mode_neighborlist` to True when initializing the potential.
+
+
+    Parameters
+    ----------
+    dm: DataModule
+        Instance of the DataModule that contains the dataset
+    potential: Potential
+        Instance of the NNP potential
+    subset: Literal["train", "test", "val"]
+        Which subset to evaluate.
+    include_metadata: bool, default True
+        This will also include metadata information in the output dictionary, providing the reference data used in training.
+
+    Returns
+    -------
+    dict -- this dictionary is keyed on the top most level with the original index in the torch_datset.
+            Each entry includes `name` which corresponds to the name of  record in the original dataset; note, since
+            records in the dataset may have multiple configurations, multiple entries may have the same name.
+            Additionally, in this dict, there is the `output` which stores the output of the potential, and `metadata`
+            which includes data which can be used to train against.
+
+    """
+
+    from modelforge.potential.neighbors import NeighborListForTraining
+
+    if not isinstance(potential.neighborlist, NeighborListForTraining):
+        raise Exception(
+            "For batch processing, must use the NeighborListForTraining; set use_training_mode_neighborlist=True when setting up the potential"
+        )
+
+    # we'll want to grab the batch and the associated indices so we can generate a more rich summary output
+    if subset == "train":
+        batches = dm.train_dataloader()
+        indices = dm.train_dataset.indices
+    elif subset == "test":
+        batches = dm.test_dataloader()
+        indices = dm.test_dataset.indices
+    elif subset == "val":
+        batches = dm.val_dataloader()
+        indices = dm.val_dataset.indices
+
+    datapoints = {}
+
+    # first loop over all the indices to get the name of the associated system so we can more easily associated a
+    # datapoint with the dataset
+    for i in indices:
+        name, datapoint = dm.torch_dataset._get_datapoint(i)
+        datapoints[i] = {"name": name, "index": i, "output": {}}
+        if include_metadata:
+            datapoints[i]["metadata"] = {}
+    count = 0
+    # loop over the batch to compute
+    per_atom_properties = []
+    per_system_properties = []
+
+    per_atom_metadata_properties = []
+    per_system_metadata_properties = []
+
+    for batch in batches:
+        output = potential(batch.nnp_input)
+
+        # we'll need to parse the atomic_subsystem_indices
+        atomic_subsystem_indices = output["atomic_subsystem_indices"]
+
+        # we need to the total number of unique systems and length of each
+        mol_ids, atomic_subsystem_counts = torch.unique(
+            atomic_subsystem_indices, return_counts=True
+        )
+        n_systems = len(mol_ids)
+
+        # cumulative sum gives us the indices where each system starts
+        # to be able to use tensor_split for per_atom properties.  these need to be on the CPU
+        atomic_subsystem_counts_sum = torch.cumsum(atomic_subsystem_counts, dim=0).to(
+            "cpu"
+        )
+
+        # if this is our first pass through, figure out what the keys are
+        if count == 0:
+            keys = output.keys()
+
+            for key in keys:
+                if key.startswith("per_atom"):
+                    per_atom_properties.append(key)
+                if key.startswith("per_system"):
+                    per_system_properties.append(key)
+
+        per_atom_properties.append("atomic_numbers")
+
+        for key in per_system_properties:
+            temp = output[key]
+            for i in range(0, n_systems):
+                ii = count + i
+                real_index = indices[ii]
+                datapoints[real_index]["output"][key] = temp[i]
+
+        for key in per_atom_properties:
+            temp = output[key]
+            temp_index = torch.tensor_split(temp, atomic_subsystem_counts_sum[:-1])
+
+            for i in range(0, n_systems):
+                ii = count + i
+                real_index = indices[ii]
+                datapoints[real_index]["output"][key] = temp_index[i]
+        if include_metadata:
+            # let us grab all the keys from metadata slots if the first pass through
+            if count == 0:
+                keys = batch.metadata.__slots__
+
+                for key in keys:
+                    if key.startswith("per_atom"):
+                        per_atom_metadata_properties.append(key)
+                    if key.startswith("per_system"):
+                        per_system_metadata_properties.append(key)
+
+            for key in per_system_metadata_properties:
+                temp = getattr(batch.metadata, key)
+                for i in range(0, n_systems):
+                    ii = count + i
+                    real_index = indices[ii]
+                    datapoints[real_index]["metadata"][key] = temp[i]
+
+            for key in per_atom_metadata_properties:
+                temp = getattr(batch.metadata, key)
+                temp_index = torch.tensor_split(temp, atomic_subsystem_counts_sum[:-1])
+
+                for i in range(0, n_systems):
+                    ii = count + i
+                    real_index = indices[ii]
+                    datapoints[real_index]["metadata"][key] = temp_index[i]
+
+        count = count + n_systems
+
+    return datapoints

modelforge/tests/test_dataset.py

@@ -34,6 +34,7 @@ def test_dataset_basic_operations():
     geom_shape = (total_atoms_series, 3)
     atomic_numbers_shape = (total_atoms_single, 1)
     internal_energy_shape = (total_confs, 1)
+    names = ["one", "two", "three"]
     input_data = {
         "geometry": np.arange(geom_shape[0] * geom_shape[1]).reshape(geom_shape),
         "atomic_numbers": np.arange(atomic_numbers_shape[0]).reshape(
@@ -44,6 +45,7 @@ def test_dataset_basic_operations():
         ),
         "atomic_subsystem_counts": atomic_subsystem_counts,
         "n_confs": n_confs,
+        "names": np.array(names),
     }
 
     property_names = PropertyNames(
@@ -93,6 +95,13 @@ def test_dataset_basic_operations():
         assert np.array_equal(
             conf_data.metadata.per_system_energy, energy_true[conf_idx]
         )
+        # test the _get_datapoint operation which will return the name and the output from _get_item
+        fetched_point = dataset._get_datapoint(conf_idx)
+        assert fetched_point[0] == names[conf_idx]
+        assert np.array_equal(
+            fetched_point[1].nnp_input.atomic_numbers,
+            conf_data.nnp_input.atomic_numbers,
+        )
 
     for rec_idx in range(dataset.record_len()):
         assert np.array_equal(
@@ -389,6 +398,7 @@ def test_metadata_validation(prep_temp_dir, load_test_dataset, dataset_temp_dir)
             "internal_energy_at_0K",
             "positions",
             "dipole_moment_per_system",
+            "names",
         ],
         "element_filter": str([[6, 1]]),
         "hdf5_checksum": "befb3ef66d74f436ef399bf68eda9b90",
@@ -419,6 +429,7 @@ def test_metadata_validation(prep_temp_dir, load_test_dataset, dataset_temp_dir)
             "internal_energy_at_0K",
             "not_a_property",
             "dipole_moment_per_system",
+            "names",
         ],
         "element_filter": str([[6, 1]]),
         "hdf5_checksum": "befb3ef66d74f436ef399bf68eda9b90",
@@ -440,6 +451,7 @@ def test_metadata_validation(prep_temp_dir, load_test_dataset, dataset_temp_dir)
             "internal_energy_at_0K",
             "not_a_property",
             "dipole_moment_per_system",
+            "names",
         ],
         "element_filter": str([[6, 1]]),
         "hdf5_checksum": "coiejfweoijfowklewke33883",
@@ -461,6 +473,7 @@ def test_metadata_validation(prep_temp_dir, load_test_dataset, dataset_temp_dir)
             "internal_energy_at_0K",
             "not_a_property",
             "dipole_moment_per_system",
+            "names",
         ],
         "element_filter": str([[6, 12]]),
         "hdf5_checksum": "coiejfweoijfowklewke33883",
@@ -1793,6 +1806,7 @@ def test_hdf5_to_torchdataset(prep_temp_dir):
     with open(yaml_file_path_out, "w") as yaml_file:
         yaml.safe_dump(yaml_content, yaml_file)
 
+    print(yaml_content)
     dm = initialize_datamodule(
         dataset_name=dataset_name,
         splitting_strategy=FirstComeFirstServeSplittingStrategy(),
@@ -1940,6 +1954,14 @@ def test_hdf5_to_torchdataset(prep_temp_dir):
         torch.tensor([[1], [1]]),
     )
 
+    # check that we can access the "name" of the system
+
+    data_point_names = ["mol1", "mol1", "mol2", "mol2", "mol3", "mol3"]
+    for i in range(0, len(dm.torch_dataset)):
+
+        datapoint_name, _ = dm.torch_dataset._get_datapoint(i)
+        assert datapoint_name == data_point_names[i]
+
 
 @pytest.mark.parametrize("grouped", [True, False])
 def test_local_dataset(grouped, prep_temp_dir):

modelforge/tests/test_potentials.py

@@ -883,6 +883,96 @@ def test_loading_from_checkpoint_file(
     assert torch.allclose(output["per_system_energy"], torch.tensor([[-1656.3508]]))
 
 
+def test_batch_evaluation(datamodule_factory, prep_temp_dir, dataset_temp_dir):
+
+    from modelforge.dataset.utils import RandomRecordSplittingStrategy
+
+    from modelforge.utils.io import get_path_string
+    from modelforge.tests import data
+
+    local_cache_dir = f"{str(prep_temp_dir)}/test_loading_from_checkpoint_file"
+    dataset_cache_dir = str(dataset_temp_dir)
+
+    from modelforge.potential.potential import load_inference_model_from_checkpoint
+
+    ckpt_file = get_path_string(data) + "/best_SchNet-qm9-epoch=433.ckpt"
+
+    potential = load_inference_model_from_checkpoint(ckpt_file)
+    assert potential is not None
+
+    # prepare reference value
+    dm_min = datamodule_factory(
+        dataset_name="QM9",
+        batch_size=512,
+        splitting_strategy=RandomRecordSplittingStrategy(),
+        remove_self_energies=True,
+        local_cache_dir=local_cache_dir,
+        dataset_cache_dir=dataset_cache_dir,
+    )
+
+    from modelforge.potential.potential import evaluate_batch
+
+    # this will raise an exception because we didn't specify to use the neighborlist for training
+    with pytest.raises(Exception):
+        output_batch = evaluate_batch(dm_min, potential, subset="test")
+
+    potential = load_inference_model_from_checkpoint(
+        ckpt_file, jit=False, use_training_mode_neighborlist=True
+    )
+    output_batch = evaluate_batch(dm_min, potential, subset="test")
+
+    temp = output_batch[256]
+
+    # just pick a few outputs to examine
+    assert temp["name"] == "dsgdb9nsd_000257"
+    assert temp["index"] == 256
+    assert torch.allclose(
+        temp["output"]["per_system_total_charge"],
+        torch.tensor([0.0], dtype=torch.float),
+    )
+    assert torch.allclose(
+        temp["output"]["per_system_energy"], torch.tensor([[-5469.7192]])
+    )
+    assert torch.allclose(
+        temp["output"]["per_atom_energy"],
+        torch.tensor(
+            [
+                [2.2452],
+                [2.2952],
+                [2.2451],
+                [-3.1492],
+                [-2.4401],
+                [-3.7545],
+                [2.4187],
+                [2.1685],
+                [2.1486],
+                [1.2201],
+                [2.1487],
+                [2.1684],
+                [2.4188],
+                [2.2497],
+            ]
+        ),
+        atol=1e-3,
+    )
+    assert torch.allclose(
+        temp["metadata"]["per_system_energy"],
+        torch.tensor([-5470.4511], dtype=torch.float64),
+    )
+
+    temp = output_batch[526]
+
+    assert temp["name"] == "dsgdb9nsd_000527"
+    assert temp["index"] == 526
+    assert torch.allclose(
+        temp["output"]["per_system_energy"], torch.tensor([-5405.9854])
+    )
+    assert torch.allclose(
+        temp["metadata"]["per_system_energy"],
+        torch.tensor([-5406.6088], dtype=torch.float64),
+    )
+
+
 @pytest.mark.parametrize(
     "potential_name", _Implemented_NNPs.get_all_neural_network_names()
 )