polarisation.py 文件源码

def polmatrixmultiply(cm, vec, polaxis=1):
    """Matrix multiply of appropriate axis of vec [...,:] by cm

    For an image vec has axes [nchan, npol, ny, nx] and polaxis=1
    For visibility vec has axes [row, nchan, npol] and polaxis=2

    :param cm: matrix to apply
    :param vec: array to be multiplied [...,:]
    :param polaxis: which axis contains the polarisation
    :return: multiplied vec
    """
    if len(vec.shape) == 1:
        return numpy.dot(cm, vec)
    else:
        # This tensor swaps the first two axes so we need to tranpose back
        result = numpy.tensordot(cm, vec, axes=(1, polaxis))
        permut = list(range(len(result.shape)))
        permut[0], permut[polaxis] = permut[polaxis], permut[0]
        return numpy.transpose(result, axes=permut)