// Get atoms by degree
atoms_by_deg = [
mol.get_atoms_with_deg(deg)
for deg in range(min_deg, max_deg + 1)
for mol in mols
]
// stack the atoms
all_atoms = np.vstack(atoms_by_deg)
// Sort all atoms by degree.
// Get the size of each atom list separated by molecule id, then by degree
mol_deg_sz = [[mol.get_num_atoms_with_deg(deg)
After Change
// Get atoms by degree
atoms_by_deg = np.concatenate([x.atom_features for x in mols])
degree_vector = np.concatenate([x.degree_list for x in mols], axis=0)
all_atoms = atoms_by_deg[degree_vector.argsort(kind="mergesort")]
// Sort all atoms by degree.
