RFDiffusion3 to ProteinMPNN API via AtomArray objects fails with duplicate atom names

[kjczarne]

Here's another fun issue.

I've attempted passing the AtomArray outputs from RFdiffusion3 directly to ProteinMPNN:

            engine = MPNNInferenceEngine(
                model_type="protein_mpnn",
                checkpoint_path=str(weights_path),
                is_legacy_weights=True,
                out_directory=None,
                write_fasta=False,
                write_structures=False
            )
            # Fix some annotations
            atom_array._annot['coord_to_be_noised'] = np.stack(atom_array._annot['coord_to_be_noised'])  # type: ignore
            outputs = engine.run(atom_arrays=[atom_array], input_dicts=None)
            print(outputs)

The atom_array object above has been directly fed from RFD3InferenceEngine. So I would expect this to work without issues. However I get:

Duplicate atom names detected in the same residue -- cannot infer struct_conn. This may happen when a non-polymer is loaded from a CIF file without using `atomworks.io.parser.parse`. 

For this issue I have no proposed fix at the moment as I'm unsure whether it's RFD3 not dumping correct .atom_name annotations in the AtomArray or if I'm just doing something fundamentally wrong.

[kjczarne]

A stop-gap measure was to fix up atom names:

def fix_atom_names(atom_array: AtomArray) -> AtomArray:
    """
    Assign proper PDB-style atom names to atoms in an AtomArray.
    
    Args:
        atom_array: biotite.structure.AtomArray
            Input structure with potentially duplicate atom names
        
    Returns:
        fixed_array: biotite.structure.AtomArray
            Structure with corrected atom names
    """
    # Create a copy to avoid modifying the original
    fixed_array = atom_array.copy()
    
    # Group atoms by residue
    residue_starts = struc.get_residue_starts(atom_array)
    
    # Counter for each element type within each residue
    for i in range(len(residue_starts)):
        start = residue_starts[i]
        end = residue_starts[i + 1] if i + 1 < len(residue_starts) else len(atom_array)
        
        # Get atoms in this residue
        residue_atoms = atom_array[start:end]
        
        # Count occurrences of each element
        element_counts = {}
        
        for j, atom_idx in enumerate(range(start, end)):
            element = residue_atoms.element[j]
            
            # Increment counter for this element
            if element not in element_counts:
                element_counts[element] = 1
            else:
                element_counts[element] += 1
            
            # Create atom name: element + count
            # For single atoms of a type, often no number is added
            count = element_counts[element]
            if count == 1:
                # Check if there will be more of this element
                total_of_element = np.sum(residue_atoms.element == element)
                if total_of_element == 1:
                    atom_name = element
                else:
                    atom_name = f"{element}{count}"
            else:
                atom_name = f"{element}{count}"
            
            # PDB atom names are 4 characters, left-justified for 2-char elements
            # and space-padded on the left for 1-char elements
            if len(element) == 1:
                atom_name = f" {atom_name:<3}"
            else:
                atom_name = f"{atom_name:<4}"
            
            fixed_array.atom_name[atom_idx] = atom_name
    
    return fixed_array

However when the MPNNInferenceInput object is instantiated, the atom_array loses all the coordinate data for some reason (each of the X Y Z coordinates becomes np.nan). This can be observed by setting a breakpoint at L713 in mpnn.utils.inference. Before parse_atom_array is called, the atom array contains correct coordinates for each atom. After the function is called, all coordinates become np.nan and ProteinMPNN crashes.

[Ubiquinone-dot]

Interesting, do you have the original inputs to rfd3?

[kjczarne]

I will prepare a small example to repro this today and get back to you, need to finish some urgent work stuff first 😄