Do the soft iterations in pssm_semigreedy obey the bias matrix

[amin-sagar]

Hello.
This is somewhat related to #107
I am fixing the positions of some of the residues using the bias matrix as follows.

bias = np.zeros((af_model._binder_len,20))
fixpos = [1,11,15]
nonfixpos = [0,2,3,4,5,6,7,8,9,10,12,13,14,16]
aa_order = residue_constants.restype_order
bias[fixpos,residue_constants.restype_order["E"]] = 1e9
bias[nonfixpos,residue_constants.restype_order["E"]] = -100

This ensures that I have E at the fixed positions in the designed sequences.
However, now if I try to add a term to the loss function to restrict the distances or angles between these residues using

def dist_loss(inputs, outputs):
  positions = outputs["structure_module"]["final_atom_positions"]
  D1 = positions[215,residue_constants.atom_order["OE1"]]
  D2 = positions[225,residue_constants.atom_order["OE1"]]
  squared_dist1 = jnp.sum((D1-D2)**2,axis=0)
  dist1 = jnp.sqrt(squared_dist1)
  dist_desired = 12
  dist = jax.nn.elu(dist1-dist_desired)
  return {"dist":dist}

I get an error

ValueError: cannot convert float NaN to integer

However, if I use CA, I don't get this error.
Interestingly, if I use a very small number of soft iterations (e.g. 2), the script runs but optimization is not optimal as expected.
It seems to me that this can happen if during the soft iterations, especially when the number of soft iterations is not too small, the residue at the fixed position is not actually "E".
Is this true?
If so, if there a way to ensure that the positions are fixed even during soft iterations?

I would be really grateful for any suggestions.
Best,
Amin.