[question]: Around training the model.

[alexanderbonnet]

[question]: Around training the model.

[patrickbryant1]

Hi,

Great 👍

The distogram is predicted in bins mapped from the pair representation. Therefore, the amino acid type doesn't matter as long as the ground truth coordinates (CB for protein) is provided for that loss.

Hope this helps.

Best,

Patrick

[patrickbryant1]

Hi,

I would like to explain this in detail, but don't have time right now. I will provide a better answer in the first weeks of July.

Right now what I can say is that comment is probably wrong. What matters is that each 'ligand atom token' can be mapped to a CA and that the masks use only CA for the losses. Any amino acid selection will take care of this.

Hope this helps (somewhat).

Best,

Patrick

[alexanderbonnet]

Thanks for the quick answer!

You may disregard the second portion of the question, my issues were due to poor indexing on my part in one of the frame aligned point error losses. All looks good now and I am getting expected behavior during training.

For the distogram loss, it still looks to me like using any other amino acid than glycine would essentially remove the ligand from the distogram loss, as the distances considered are between CBs (except for glycine, that uses CAs, and would be compatible with setting the ligand heavy atoms to CAs).

Umol/src/net/model/tf/data_transforms.py

Lines 308 to 324 in f7cd2b4

	def pseudo_beta_fn(aatype, all_atom_positions, all_atom_masks):
	"""Create pseudo beta features."""
	is_gly = tf.equal(aatype, residue_constants.restype_order['G'])
	ca_idx = residue_constants.atom_order['CA']
	cb_idx = residue_constants.atom_order['CB']
	pseudo_beta = tf.where(
	tf.tile(is_gly[..., None], [1] * len(is_gly.shape) + [3]),
	all_atom_positions[..., ca_idx, :],
	all_atom_positions[..., cb_idx, :])
	if all_atom_masks is not None:
	pseudo_beta_mask = tf.where(
	is_gly, all_atom_masks[..., ca_idx], all_atom_masks[..., cb_idx])
	pseudo_beta_mask = tf.cast(pseudo_beta_mask, tf.float32)
	return pseudo_beta, pseudo_beta_mask
	else:
	return pseudo_beta

Umol/src/net/model/modules.py

Lines 1244 to 1273 in f7cd2b4

	def _distogram_log_loss(logits, bin_edges, batch, num_bins):
	"""Log loss of a distogram."""
	assert len(logits.shape) == 3
	positions = batch['pseudo_beta']
	mask = batch['pseudo_beta_mask']
	assert positions.shape[-1] == 3
	sq_breaks = jnp.square(bin_edges)
	dist2 = jnp.sum(
	jnp.square(
	jnp.expand_dims(positions, axis=-2) -
	jnp.expand_dims(positions, axis=-3)),
	axis=-1,
	keepdims=True)
	true_bins = jnp.sum(dist2 > sq_breaks, axis=-1)
	errors = softmax_cross_entropy(
	labels=jax.nn.one_hot(true_bins, num_bins), logits=logits)
	square_mask = jnp.expand_dims(mask, axis=-2) * jnp.expand_dims(mask, axis=-1)
	avg_error = (
	jnp.sum(errors * square_mask, axis=(-2, -1)) /
	(1e-6 + jnp.sum(square_mask, axis=(-2, -1))))
	dist2 = dist2[..., 0]
	return dict(loss=avg_error, true_dist=jnp.sqrt(1e-6 + dist2))

I think I should be fine for the most part, but would love to have detailed explanations regardless if you find the time.

Thanks again,
Alexander

[patrickbryant1]

I will close this issue now, please email me if you have any further more detailed questions and I will try my best to answer.

I realize the code is not that neat, but I partly blame DeepMind for this and all their nested calls to all the different files.

Best,

Patrick

[alexanderbonnet]

I will close this issue now, please email me if you have any further more detailed questions and I will try my best to answer.

I've been able to get some results I'm pretty happy with so I should be all set, thanks! Will definitely shoot you an e-mail if I need to.

I realize the code is not that neat, but I partly blame DeepMind for this and all their nested calls to all the different files.

Yup, not the easiest piece of code to work with for sure...

Thanks again,
Alexander