# ChimeraX script to compute the per-atom RMSD for the open models.
# It assumes the models all have the same residues and atoms and in the same order in the input files,
# for example, computed by AlphaFold.
# It assigns the atom attribute "rmsd" which can be used to color the structure with Render by Attribute.

def atom_rmsd(session, atoms):
    coords = []
    for structure, satoms in atoms.by_structure:
        coords.append(satoms.scene_coords)
    from numpy import var, sqrt
    atom_rmsds = sqrt(var(coords, axis=0).sum(axis=1))
    from chimerax.atomic import Atom
    Atom.register_attr(session, 'rmsd', 'Per-atom RMSD script', attr_type = float)
    for structure, satoms in atoms.by_structure:
        for a, rmsd in zip(satoms, atom_rmsds):
            a.rmsd = rmsd

def register_command(logger):
    from chimerax.core.commands import register, CmdDesc
    from chimerax.atomic import AtomsArg
    desc = CmdDesc(required = [('atoms', AtomsArg)],
                   synopsis='Compute per-atom RMSD for a set of structures that all have the same atoms')
    register('atomrmsd', desc, atom_rmsd, logger=logger)

register_command(session.logger)