Add an option to run PBC in single system mode

[misko]

There is an issue when running PBC across multiple systems. The current code takes the number of expansions in each direction for the worst case system in each batch. This can result in a very large amount of GPU memory required for expansion 50GB+

This PR adds an option "use_pbc_single" that runs PBC expansion for each system, one at a time.
The runtime is approximately 10% slower for PBC construction, but does not seem significant for the overall runtime.

[codecov]

Codecov Report

Attention: Patch coverage is 95.23810% with 1 line in your changes missing coverage. Please review.

Files	Patch %	Lines
src/fairchem/core/models/gemnet_gp/gemnet.py	0.00%	1 Missing ⚠️

Files	Coverage Δ
src/fairchem/core/models/base.py	87.96% <100.00%> (+0.96%)	⬆️
src/fairchem/core/models/dimenet_plus_plus.py	97.33% <100.00%> (+0.01%)	⬆️
...airchem/core/models/equiformer_v2/equiformer_v2.py	88.51% <100.00%> (+0.03%)	⬆️
src/fairchem/core/models/escn/escn.py	96.77% <100.00%> (+<0.01%)	⬆️
src/fairchem/core/models/gemnet/gemnet.py	93.95% <100.00%> (+0.02%)	⬆️
src/fairchem/core/models/gemnet_oc/gemnet_oc.py	90.04% <100.00%> (+0.02%)	⬆️
src/fairchem/core/models/painn/painn.py	90.78% <100.00%> (+0.03%)	⬆️
src/fairchem/core/models/schnet.py	98.03% <100.00%> (+0.03%)	⬆️
src/fairchem/core/models/scn/scn.py	93.89% <100.00%> (+0.01%)	⬆️
src/fairchem/core/models/gemnet_gp/gemnet.py	16.93% <0.00%> (-0.10%)	⬇️

[kyonofx]

LGTM!

[lbluque]

lgtm. Question though, how do you use pbc_single? Do I need to specify use_pbc_single: True to the model parameters in the config?

[misko]

lgtm. Question though, how do you use pbc_single? Do I need to specify use_pbc_single: True to the model parameters in the config?

I just pushed those changes! Realized it wasn't connected. The way we connect parameters to generate_graph is a bit interesting!