Support periodic systems in the energy target

[JMorado]

Description

This PR adds support for using the energy target with periodic systems.

Status

• Ready to go

[lilyminium]

Thanks @JMorado! This seems like a good addition, but it'd be great to get more eyes on it -- @fjclark would you be up for reviewing this?

[fjclark]

Yes I'll take a look this week!

[fjclark]

Thanks, this looks good to me and I don't have any comments.

[JMorado]

Thanks, @fjclark! It looks like I need to run Ruff formatting on the files. I'll do that and push the changes.

[fjclark]

Great! If you run make env it should install the pre-commit hook which includes this.

[fjclark]

Thanks @JMorado! Again, this looks good to me. I know Lily is away at the moment -- are you happy for this to be merged @j-wags? It would be great to have a release of this on conda-forge with the latest updates soon (this would unblock me getting presto on conda-forge). Thanks.

[mattwthompson]

Given that

• @lilyminium asked @fjclark to be the reviewer
• @JMorado and @fjclark are happy with this in its current state
• All other checks are green

I'm moving ahead with merging this

[mattwthompson]

@JMorado could you have a look at the tests? This doesn't look concerning to me in substance, but I don't know what might have changed, if anything, in the past couple of weeks

>       assert entries == expected_entries
E       AssertionError: assert [{'box_vector..., 71.]), ...}] == [{'coords': a...[H:4])[H:5]'}]
E         
E         At index 0 diff: {'id': None, 'smiles': '[C:1]([O:2][H:6])([H:3])([H:4])[H:5]', 'coords': tensor([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13.,\n        14., 15., 16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26., 27.,\n        28., 29., 30., 31., 32., 33., 34., 35.]), 'box_vectors': None, 'energy': tensor([0., 3.]), 'forces': tensor([36., 37., 38., 39., 40., 41., 42., 43., 44., 45., 46., 47., 48., 49.,\n        50., 51., 52., 53., 54., 55., 56., 57., 58., 59., 60., 61., 62., 63.,\n        64., 65., 66., 67., 68., 69., 70., 71.])} != {'smiles': '[C:1]([O:2][H:6])([H:3])([H:4])[H:5]', 'coords': approx([0.0 ± 1.0e-12, 1.0 ± 1.0e-06, 2.0 ± 2.0e-06, 3.0 ± 3.0e-06, 4.0 ± 4.0e-06, 5.0 ± 5.0e-06, 6.0 ± 6.0e-06, 7.0 ± 7.0e-06, 8.0 ± 8.0e-06, 9.0 ± 9.0e-06, 10.0 ± 1.0e-05, 11.0 ± 1.1e-05, 12.0 ± 1.2e-05, 13.0 ± 1.3e-05, 14.0 ± 1.4e-05, 15.0 ± 1.5e-05, 16.0 ± 1.6e-05, 17.0 ± 1.7e-05, 18.0 ± 1.8e-05, 19.0 ± 1.9e-05, 20.0 ± 2.0e-05, 21.0 ± 2.1e-05, 22.0 ± 2.2e-05, 23.0 ± 2.3e-05, 24.0 ± 2.4e-05, 25.0 ± 2.5e-05, 26.0 ± 2.6e-05, 27.0 ± 2.7e-05, 28.0 ± 2.8e-05, 29.0 ± 2.9e-05, 30.0 ± 3.0e-05, 31.0 ± 3.1e-05, 32.0 ± 3.2e-05, 33.0 ± 3.3e-05, 34.0 ± 3.4e-05, 35.0 ± 3.5e-05]), 'energy': approx([0.0 ± 1.0e-12, 3.0 ± 3.0e-06]), 'forces': approx([36.0 ± 3.6e-05, 37.0 ± 3.7e-05, 38.0 ± 3.8e-05, 39.0 ± 3.9e-05, 40.0 ± 4.0e-05, 41.0 ± 4.1e-05, 42.0 ± 4.2e-05, 43.0 ± 4.3e-05, 44.0 ± 4.4e-05, 45.0 ± 4.5e-05, 46.0 ± 4.6e-05, 47.0 ± 4.7e-05, 48.0 ± 4.8e-05, 49.0 ± 4.9e-05, 50.0 ± 5.0e-05, 51.0 ± 5.1e-05, 52.0 ± 5.2e-05, 53.0 ± 5.3e-05, 54.0 ± 5.4e-05, 55.0 ± 5.5e-05, 56.0 ± 5.6e-05, 57.0 ± 5.7e-05, 58.0 ± 5.8e-05, 59.0 ± 5.9e-05, 60.0 ± 6.0e-05, 61.0 ± 6.1e-05, 62.0 ± 6.2e-05, 63.0 ± 6.3e-05, 64.0 ± 6.4e-05, 65.0 ± 6.5e-05, 66.0 ± 6.6e-05, 67.0 ± 6.7e-05, 68.0 ± 6.8e-05, 69.0 ± 6.9e-05, 70.0 ± 7.0e-05, 71.0 ± 7.1e-05])}
E         
E         Full diff:
E           [
E               {
E         -         'coords': approx([0.0 ± 1.0e-12, 1.0 ± 1.0e-06, 2.0 ± 2.0e-06, 3.0 ± 3.0e-06, 4.0 ± 4.0e-06, 5.0 ± 5.0e-06, 6.0 ± 6.0e-06, 7.0 ± 7.0e-06, 8.0 ± 8.0e-06, 9.0 ± 9.0e-06, 10.0 ± 1.0e-05, 11.0 ± 1.1e-05, 12.0 ± 1.2e-05, 13.0 ± 1.3e-05, 14.0 ± 1.4e-05, 15.0 ± 1.5e-05, 16.0 ± 1.6e-05, 17.0 ± 1.7e-05, 18.0 ± 1.8e-05, 19.0 ± 1.9e-05, 20.0 ± 2.0e-05, 21.0 ± 2.1e-05, 22.0 ± 2.2e-05, 23.0 ± 2.3e-05, 24.0 ± 2.4e-05, 25.0 ± 2.5e-05, 26.0 ± 2.6e-05, 27.0 ± 2.7e-05, 28.0 ± 2.8e-05, 29.0 ± 2.9e-05, 30.0 ± 3.0e-05, 31.0 ± 3.1e-05, 32.0 ± 3.2e-05, 33.0 ± 3.3e-05, 34.0 ± 3.4e-05, 35.0 ± 3.5e-05]),
E         -         'energy': approx([0.0 ± 1.0e-12, 3.0 ± 3.0e-06]),
E         -         'forces': approx([36.0 ± 3.6e-05, 37.0 ± 3.7e-05, 38.0 ± 3.8e-05, 39.0 ± 3.9e-05, 40.0 ± 4.0e-05, 41.0 ± 4.1e-05, 42.0 ± 4.2e-05, 43.0 ± 4.3e-05, 44.0 ± 4.4e-05, 45.0 ± 4.5e-05, 46.0 ± 4.6e-05, 47.0 ± 4.7e-05, 48.0 ± 4.8e-05, 49.0 ± 4.9e-05, 50.0 ± 5.0e-05, 51.0 ± 5.1e-05, 52.0 ± 5.2e-05, 53.0 ± 5.3e-05, 54.0 ± 5.4e-05, 55.0 ± 5.5e-05, 56.0 ± 5.6e-05, 57.0 ± 5.7e-05, 58.0 ± 5.8e-05, 59.0 ± 5.9e-05, 60.0 ± 6.0e-05, 61.0 ± 6.1e-05, 62.0 ± 6.2e-05, 63.0 ± 6.3e-05, 64.0 ± 6.4e-05, 65.0 ± 6.5e-05, 66.0 ± 6.6e-05, 67.0 ± 6.7e-05, 68.0 ± 6.8e-05, 69.0 ± 6.9e-05, 70.0 ± 7.0e-05, 71.0 ± 7.1e-05]),
E         +         'box_vectors': None,
E         +         'coords': tensor([ 0.,  1.,  2.,  3.,  4.,  5.,  6.,  7.,  8.,  9., 10., 11., 12., 13.,
E         +         14., 15., 16., 17., 18., 19., 20., 21., 22., 23., 24., 25., 26., 27.,
E         +         28., 29., 30., 31., 32., 33., 34., 35.]),
E         +         'energy': tensor([0., 3.]),
E         +         'forces': tensor([36., 37., 38., 39., 40., 41., 42., 43., 44., 45., 46., 47., 48., 49.,
E         +         50., 51., 52., 53., 54., 55., 56., 57., 58., 59., 60., 61., 62., 63.,
E         +         64., 65., 66., 67., 68., 69., 70., 71.]),
E         +         'id': None,
E                   'smiles': '[C:1]([O:2][H:6])([H:3])([H:4])[H:5]',
E               },
E           ]

[JMorado]

Thanks for flagging this @mattwthompson. It looks like #88 added create_dataset_from_generator functions and respective tests which were still using the old data schema.

[JMorado]

All tests are passing now -- it looks good to merge, unless there are objections to my last commit.