beam_manipulation.py 文件源码

def rotateVec(vec, rotaxis, theta):
    """
    Given a 3-vector vec, rotate about rotaxis by $\theta$

    Also accepts iterable input for vec and rotaxis if the arguments are
    compatible lengths.
    """
    assert not (np.any(np.isnan(vec)) or
                np.any(np.isnan(rotaxis)) or
                np.any(np.isnan(theta))), "Inputs must not be NaN"
    if np.shape(vec) == (3, ):
        R = rotationMatrix3D(rotaxis, theta)
        norm = np.linalg.norm(vec)
        res = np.dot(R, vec)
        assert np.isclose(np.linalg.norm(res), norm), "Rotation changed vector norm"
        return np.dot(R, vec)
    else:
        assert np.shape(vec)[0] == np.shape(rotaxis)[0] == np.shape(theta)[0], "Dimension mismatch in rotateVec()"
        # Unfortunately, seems that np.dot can't be coerced into doing this operation all at once
        # Tried to build a tensor of rotation matrices and use np.einsum, but couldn't get good reuslts.
        # If this becomes slow at any point, it's a good target for optimization.
        res = np.zeros(shape=(np.shape(vec)[0], 3))
        for i, (v, r, t) in enumerate(zip(vec, rotaxis, theta)):
            # print("In rotateVec(): r={}, t={}".format(r, t))
            norm = np.linalg.norm(v)
            R = rotationMatrix3D(r, t)
            res[i] = np.dot(R, v)
            assert np.isclose(np.linalg.norm(res[i]), norm), "Rotation changed vector norm: v={}, r={}, t={}, R={}".format(v, r, t, R)
        return np.hsplit(res, 3)