Utility functions for atom_diag