def atom_extract(Mask_3D, atom):
    """
    Given (a list of) 3D Mask, extract the coord's of needed atom.
    """
    xyz_list = []
    if type(Mask_3D) is list:
        for i in Mask_3D:
            for x, y in enumerate(i.des):
                if y == atom:
                    if i.values.size == 3:
                        i.values = i.values.reshape(1, 3)
                    xyz_list.append(i.values[x])
    else:
        for x, y in enumerate(Mask_3D.des):
            if y == atom:
                if Mask_3D.values.size == 3:
                    Mask_3D.values = Mask_3D.values.reshape(1, 3)
                xyz_list.append(Mask_3D.values[x])
    return xyz_list
