Incorrect behavior for atom_messages = True

[Varal7]

I think there is currently a bug where args.atom_messages = True.
BatchMolGraph uses the get_bond_dim() function without forwarding the argument which seems will lead to incorrect dimension of the bond features.

[fhvermei]

The bond feature vector in BatchMolGraph are indeed constructed using both the atom and bond features. When the bond feature vector is called for the MPNN (BatchMolGraph.get_components(atom_messages=True)) the size of the bond feature vector is corrected and the atom features are removed.

[Varal7]

I am talking about this line:

chemprop/chemprop/features/featurization.py

Line 206 in 8258fcb

self.bond_fdim = get_bond_fdim()

But I can run some tests tomorrow and submit a proper bug report if needs be

[hesther]

Thank you for your question! Just had a look - I agree with Florence. Upon initialising BatchMolGraph, the bond dimensions are set up alike no matter whether atom messages are used or not (line 206), and the bond features are set up with the full vector (bond + atom features). But when returning the atom and bond features in get_components, the bond features get truncated up to their correct length if atom messages are used (with the length determined by calling get_bond_fdim again with atom_messages=True):

chemprop/chemprop/features/featurization.py

Lines 246 to 272 in 8258fcb

	def get_components(self, atom_messages: bool = False) -> Tuple[torch.FloatTensor, torch.FloatTensor,
	torch.LongTensor, torch.LongTensor, torch.LongTensor,
	List[Tuple[int, int]], List[Tuple[int, int]]]:
	"""
	Returns the components of the :class:`BatchMolGraph`.
	The returned components are, in order:
	* :code:`f_atoms`
	* :code:`f_bonds`
	* :code:`a2b`
	* :code:`b2a`
	* :code:`b2revb`
	* :code:`a_scope`
	* :code:`b_scope`
	:param atom_messages: Whether to use atom messages instead of bond messages. This changes the bond feature
	vector to contain only bond features rather than both atom and bond features.
	:return: A tuple containing PyTorch tensors with the atom features, bond features, graph structure,
	and scope of the atoms and bonds (i.e., the indices of the molecules they belong to).
	"""
	if atom_messages:
	f_bonds = self.f_bonds[:, :get_bond_fdim(atom_messages=atom_messages)]
	else:
	f_bonds = self.f_bonds
	return self.f_atoms, f_bonds, self.a2b, self.b2a, self.b2revb, self.a_scope, self.b_scope

This may be a bit clumsy (and there might be faster ways to do this), but correct.

I will close this issue for now, but please feel free to reopen if you have further questions! Yours, Esther

[hesther]

While looking through the code more thoroughly today, I noticed that there actually IS a bug (although a different one than reported):

We set up the bond features in the order atom(length 133)-bond(length 14):

chemprop/chemprop/features/featurization.py

Lines 171 to 172 in 8258fcb

	self.f_bonds.append(self.f_atoms[a1] + f_bond)
	self.f_bonds.append(self.f_atoms[a2] + f_bond)

but then cut out the first (!) 14 values, instead of the last:

chemprop/chemprop/features/featurization.py

Lines 267 to 270 in 8258fcb

	if atom_messages:
	f_bonds = self.f_bonds[:, :get_bond_fdim(atom_messages=atom_messages)]
	else:
	f_bonds = self.f_bonds

Can someone else confirm this, so I can write a PR?

[cjmcgill]

@hesther It looks this way to me as well