Batched calculator over structures

[frostedoyster]

Especially during training (but also sometimes at evaluation time), it is useful to calculate properties for batches of structures

[E-Rum]

There is no way to properly batch the current implementation. Since the best we can do is loop over systems, we decided not to include this inside the machinery, as one can always loop externally in their own code.

[ericyuan00000]

What's the major issue for batching the current implementation? I would imagine a proper combination of scatter and gather operations (from pytorch_geometric for example) could be used for summing this, if that's the problem.

[ericyuan00000]

To my understanding, the batching problem would only show up when summing over i. One should be careful not to sum the trigonometric factors and the charges that do not belong to the same graph, and a correct cell volume should be grabbed. These two are the only issues. Is there anything else that I haven't thought through?

[E-Rum]

@ericyuan00000 thank you for your comment. I think it makes sense, and indeed, there are cases where batching is possible. The problem arises when you have a batch of systems with different cell sizes, because in that case, the k-vectors will differ between systems — both in values and in number.

Of course, one could try to precompute all these vectors before training, based on the entire dataset, but that seems like a very costly procedure and also reduces pipeline flexibility.

If you think there’s another way to overcome this problem, we’d be very glad to hear your thoughts!

[ericyuan00000]

That makes a lot of sense — definitely a challenging setup. In this case, a bipartite graph from i to k seems necessary. Each system (with potentially different cell sizes and numbers of k-vectors) would have its own set of i-nodes and k-nodes, connected by i*k edges. A batch of systems could just be formed by concatenating nodes, with edge indices properly offset. The inner product can still be handled in a graph-friendly way. I believe the time and memory complexity would still be O(N^2).

I can try putting together a simple implementation — hopefully it won’t come with too much of an efficiency hit.

[E-Rum]

Sounds reasonable! We gladly welcome any efforts in this regard, so if you need any help — let us know. When you have something ready to show us — open up a PR and we will definitely review it!