From ea8664f09876c683bc174a8aae42b256e29eb489 Mon Sep 17 00:00:00 2001 From: Guillem Simeon <55756547+guillemsimeon@users.noreply.github.com> Date: Fri, 23 Feb 2024 12:16:11 +0100 Subject: [PATCH] update hdf5 dataset (#292) update hdf5 dataset such that now charges, spin and partial charges can be included at the same time. also, renamed partial charges to 'pq' to be consistent with other datasets such as Ace --- torchmdnet/datasets/hdf.py | 6 +++++- 1 file changed, 5 insertions(+), 1 deletion(-) diff --git a/torchmdnet/datasets/hdf.py b/torchmdnet/datasets/hdf.py index 3d817a50c..ee6a5e0f0 100644 --- a/torchmdnet/datasets/hdf.py +++ b/torchmdnet/datasets/hdf.py @@ -53,9 +53,13 @@ def __init__(self, filename, dataset_preload_limit=1024, **kwargs): self.fields.append(("y", "energy", torch.float32)) if "forces" in group: self.fields.append(("neg_dy", "forces", torch.float32)) + if "charge" in group: + self.fields.append(("q", "charge", torch.float32)) + if "spin" in group: + self.fields.append(("s", "spin", torch.float32)) if "partial_charges" in group: self.fields.append( - ("partial_charges", "partial_charges", torch.float32) + ("pq", "partial_charges", torch.float32) ) assert ("energy" in group) or ( "forces" in group