Tensornet-q

[AntonioMirarchi]

This PR introduces updates to enable TorchMD-Net to handle additional properties, collectively referred to as extra_args, such as charges, partial charges, or spin. The implementation allows models to utilize these extra_args, though currently, only TensorNet supports these additional labels. At this stage, only one property (charge, partial charges, or spin) can be used at a time. Future updates may expand support to handle combinations of multiple extra_args, potentially incorporating learnable parameters to optimize their integration.

Main changes:

• Added support for extra_args in the forward pass of the lightning module (module.py), model (model.py), and model architecture (tensornet.py, etc.) ensuring backward compatibility by limiting changes to the forward pass.
• 'partial_charges' added as arg to the parser.
• Updated the HDF5 dataset to include total_charge and spin in addition to partial_charges.
• Standardized the memmap dataset nomenclature for partial_charges.
• Implemented a function in TensorNet to process extra_args.
• Updated and expanded PyTests to include cases for extra_args.
• Fixed test_calculator.py.

[peastman]

I added support for arbitrary extra properties in #289. My implementation was fully general, maintained perfect backward compatibility, and introduced them into the model in a way that works much better than what's currently done for q. But it was never merged.

[AntonioMirarchi]

Yes, I was also working on the inclusion of additional labels in PR #306, but this PR focuses on a simpler approach to ensure the reproducibility of a paper currently under revision. That said, if we aim for a more general solution or wish to incorporate your implementation, I’m completely open to it.

[peastman]

The approach in #289 was used in https://doi.org/10.1021/acs.jctc.4c00794. I did that work with PhysicsML—without this feature, I can't implement my models in TorchMD-Net.

I can tell you that based on extensive testing, that approach works much better than the one currently used for q.