Potential fails to predict structure without threebody interaction in threebody cutoff of 4 Å

[JiQi535]

Hi, I'm reporting that the current M3GNet Potential cannot predict structures without threebody interactions within the threebody cutoff of 4 Å. I believe this is not a desired behavior. Part of the error is listed here:

The error comes from the below lines:

m3gnet/m3gnet/models/_base.py

Lines 253 to 261 in b2a92f7

	@tf.function(experimental_relax_shapes=True)
	def get_energies(self, graph: List) -> tf.Tensor:
	"""
	get energies from a list repr of a graph
	Args:
	graph (List): list repr of a graph
	Returns:
	"""
	return self.model(graph)

As designed, self.model(graph) actually works fine to get energy of structures without threebody interactions, but the line of @tf.function(experimental_relax_shapes=True) leads to the bug. When commenting out these lines of @tf.function from the get_energies and get_efs_tensor functions, this bug no longer exists.

@chc273 When available, would you please comment on whether or not we should keep the @tf.function, and if we should keep them, how we can make the Potential behave as designed?

[shyuep]

I think the most robust way would be to skip the 3-body update to bond, i.e., treat it as a simple identity pass through, when there are no 3-body terms at all?

[chc273]

@shyuep @JiQi535

Run the model on a "good" structure first to compile static graph and then subsequent runs will be ok

[chc273]

Also see my previous PR, I have added a test case that just has one atom with no bond nor three body. The model runs fine

[JiQi535]

@chc273 Thanks for the comments. m3gnet fails to run the "test_single_atoms" test function with the latest tensorflow but works with tensorflow=2.9.1. So it has to be this version of tensorflow.

I don't know if it is of interest, but to make m3gnet work with the latest tensorflow, I'm trying to understand why the error in the snapshot emerged.

m3gnet/m3gnet/utils/_math.py

Lines 288 to 305 in b2a92f7

	if tf.shape(sbf)[0] == 0:
	return sbf
	if not use_phi:
	repeats_sbf = [1] * max_l * max_n
	block_size = [1] * max_l
	else:
	# [1, 1, 1, ..., 1, 3, 3, 3, ..., 3, ...]
	repeats_sbf = np.repeat(2 * np.arange(max_l) + 1, repeats=max_n)
	# tf.repeat(2 * tf.range(max_l) + 1, repeats=max_n)
	block_size = 2 * np.arange(max_l) + 1
	# 2 * tf.range(max_l) + 1
	expanded_sbf = tf.repeat(sbf, repeats=repeats_sbf, axis=1)
	expanded_shf = _block_repeat(shf, block_size=block_size, repeats=[max_n] * max_l)
	shape = max_n * max_l
	if use_phi:
	shape *= max_l
	return tf.reshape(expanded_sbf * expanded_shf, [-1, shape])

From my understanding, the above code for basis_expansion tries to return sbf if it is empty, i.e., tf.shape(sbf)[0] =0 when there is no 3-body interactions. It works as designed for M3GNet object but fails for Potential objects. I found that the @tf.function added to functions of Potential leads to this different behavior of M3GNet and Potential. . Would you give some suggestions on this? Does the @tf.function provide a significant improvement in efficiency or other factors so that we have to keep it?

[chc273]

Please don't touch it. I added there for that reason you described