Revisit shape of per-atom/molecule tensors

[wiederm]

Description

While writing the style guide in the developer section of the documentation I noticed that we haven't been consistent with the shape of tensors. I recognize that it is difficult to stay consistent within a network implementation, but I think we can try to recommend a few general rules and note whenever modelforge deviates from them in specific places for performance reasons.

• I believe tensors that describe per-ato, per-molecule or per-interaction properties/features should have the shape (nr_of_instances, nr_of_features). That means that the shape of the tensor that contains per-atom charge has shape (nr_of_atoms, 1), the tensor that defines the species shape (nr_of_atoms, 1), positions (nr_of_atoms, 3) and so on.

• In contrast, index tensors should be arrays. So, e.g., atomic_subsystem_indices, which we use to reduce the per-atom energy (nr_of_atoms, 1) to per-molecule energy (nr_of_molecule, 1) has shape (nr_of_atoms). We are always indexing in the first dimension.

• Todos

Notable points that this PR has either accomplished or will accomplish.

• update dataloader
• update dataclasses
• update networks
• update documentation

Status

• Ready to go