def assign_pml_elems(sorted_elems, pml_elems, pml_partID='2'):
    """assign PML elements in the sorted element matrix

    Args:
        sorted_elems: sorted element matrix
        pml_elems: list of tuples of # PML elems on each axis edge
        ([[xmin, max], [ymin, ymax], ...)
        pml_partID: default = 2

    Returns:
        sorted_elems (to be written to new file)

    """
    sorted_elems['pid'][0:pml_elems[0][0], :, :] = pml_partID
    sorted_elems['pid'][-1:-pml_elems[0][1] - 1:-1, :, :] = pml_partID
    sorted_elems['pid'][:, 0:pml_elems[1][0], :] = pml_partID
    sorted_elems['pid'][:, -1:-pml_elems[1][1] - 1:-1, :] = pml_partID
    sorted_elems['pid'][:, :, 0:pml_elems[2][0]] = pml_partID
    sorted_elems['pid'][:, :, -1:-pml_elems[2][1] - 1:-1] = pml_partID

    return sorted_elems