From 511ba3144d632a37805458190dade3b91a8f4814 Mon Sep 17 00:00:00 2001 From: Guillem Simeon <55756547+guillemsimeon@users.noreply.github.com> Date: Fri, 23 Feb 2024 12:10:09 +0100 Subject: [PATCH] update hdf5 dataset 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