def draw2d(atom, **kwargs):
    """ Draw a 2D viewer with this atom highlighted (Jupyter only).
    In biomolecules, only draws the atom's residue.

    Args:
        width (int): width of viewer in pixels
        height (int): height of viewer in pixels

    Returns:
        mdt.ChemicalGraphViewer: viewer object
    """
    if atom.molecule:
        if atom.molecule.is_small_molecule:
            return atom.molecule.draw2d(highlight_atoms=[atom], **kwargs)
        else:
            return atom.residue.draw2d(highlight_atoms=[atom], **kwargs)
    else:
        raise ValueError("Cant draw an atom that's not part of a molecule.")
