Skip to content

Commit

Permalink
update hdf5 dataset (torchmd#292)
Browse files Browse the repository at this point in the history
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
  • Loading branch information
guillemsimeon authored Feb 23, 2024
1 parent 166b7db commit ea8664f
Showing 1 changed file with 5 additions and 1 deletion.
6 changes: 5 additions & 1 deletion torchmdnet/datasets/hdf.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down

0 comments on commit ea8664f

Please sign in to comment.