Atomic encodings

[KnathanM]

Description

#722 asked about getting an interface for atomic encodings. This is simple to implement and may not be entirely necessary because users can do it themselves as @davidegraff showed in a comment. But in v2.2 we plan to add support for atom and bond targets which would probably benefit from a function at the model level to get the atom encodings. We could add it now, given that it would be simple.

Example

atom_encodings = []
for batch in dataloader:
    bmg, V_d, *_ = batch
    atom_encodings.extend(model.atomic_encodings(bmg, V_d)

Each element in atom_encodings is the atomic fingerprints for a single molecule.

Questions

What is a good name for this function? I brainstormed atom_encodings, atomic_encoding, atoms_fingerprint.

I still have a couple things to do before merging this if the reviews are positive.

Checklist

• Add unit tests
• Add example in notebooks

[davidegraff]

Suggested change

	return H_v.split(torch.bincount(bmg.batch).tolist())
	sizes = torch.bincount(bmg.batch, minlength=len(bmg)).tolist()
	return H_v.split(sizes)

[davidegraff]

same as above

[davidegraff]

FWIW I prefer atom_encodings. I prefer to avoid the term "fingerprint" as much as possible as it's not defined in the CS community, but the chemists occasionally like to use it so that's why it's been kept around. You'll also want to add a unit test or two based on the output shape. There are two strategies here:

1. keep the code as-is and unit test that the function, given an input batch with some number of molecules produces a list[Tensor] of shape $(b, n_{a,i}, d)$, where $n_{a,i}$ is the number atoms in molecule $i$
2. refactor this function to call some other utility function split_into_tensors() (or some other descriptive name) that handles the operation of taking a tensor of shape $(n, d)$ a batch index $\mathbf i \in [0 .. b]^{n}$ and splits it into a list[Tensor] of shape $(b, n_{a,i}, d)$, where "..."

The advantage of (2) is that it's simpler and more isolated. The unit test only needs to cover the split_into_tensors() function (as testing atom_encodings() would be an "integration" test now). In comparison, (1) can result in failures for a number of reasons, with only one of them being the fault of the atom_encodings() function (because it relies featurization and message-passing to work correctly.) The fix to this is to "mock" the output of message_passing to be correct by construction. I'm ambivalent between the two approaches: mocking (1) and refactoring (2).