Fix laser_nll going inf when residues are fixed

[ssiddhantsharma]

What

Fixes laser_nll / laser_bs_nll becoming inf whenever --fixed-identity-residue-indices is used.

Why

In sample_sequences, the per-design NLL is a mean over every residue:

nll = (-1 * torch.log10(curr_probs)).cpu().numpy().mean()

When residues are fixed (fix_beta=True -> chain_mask = self.fixed_rotamers), the clamped positions carry a recorded probability of 0, so log10(0) sends the mean to -inf and nll to inf. Every design with at least one fixed residue reports inf for both laser_nll and laser_bs_nll, and mean_laser_score / mean_laser_bs_score follow. The scores are meant to measure the quality of the designed sequence, so the clamped positions should not be in the average anyway.

Fix

Average the NLL over the designed positions only (exclude chain_mask), and intersect the binding-site mask with the designed set:

designed = ~batch_data.chain_mask[curr_batch_mask].bool()
bs_mask = batch_data.first_shell_ligand_contact_mask[curr_batch_mask] & designed
nll = (-1 * torch.log10(curr_probs[designed])).cpu().numpy().mean()
bs_nll = (-1 * torch.log10(curr_probs[bs_mask])).cpu().numpy().mean()

With no fixed residues, designed is all-True and behavior is unchanged.

Testing

One-iteration run with four fixed residues, before vs after:

• before: laser_nll = [inf, inf, ...], mean_laser_score = inf
• after: laser_nll = [0.365, 0.367], laser_bs_nll = [0.143, 0.126], mean_laser_score = 0.366

The binding-site NLL being lower than the full-sequence NLL is the expected direction (interface positions are higher-confidence).

Scope

Only run_nise_boltz2x.py (the LASErMPNN path). run_nise_boltz1x.py does not use this pattern and does not expose fixed-identity residues; the LigandMPNN variant computes its own lmpnn_nll on a separate path and is out of scope here.