def rotate_x(vector, theta):
    """ Rotate a 3D vector about the +z axis
        Inputs:
            vector: 3d vector array
            theta [rad]: angle to rotate the vector about
        Returns: rotated vector
    """
    import numpy as np
    if np.ndim(vector) == 1:
        R = np.array([[1.0, 0.0, 0.0], [0.0, np.cos(theta), -np.sin(theta)],
            [0.0, np.sin(theta), np.cos(theta)]])
        out = np.zeros(3)
        out[0] = R[0, 0] * vector[0] + R[0, 1] * vector[1] + R[0, 2] * vector[2
            ]
        out[1] = R[1, 0] * vector[0] + R[1, 1] * vector[1] + R[1, 2] * vector[2
            ]
        out[2] = R[2, 0] * vector[0] + R[2, 1] * vector[1] + R[2, 2] * vector[2
            ]
    else:
        R = np.array([[[1.0] * len(theta), 0.0, 0.0], [0.0, np.cos(theta), 
            -np.sin(theta)], [0.0, np.sin(theta), np.cos(theta)]])
        out = np.zeros((3, vector.shape[1]))
        out[0] = R[0, 0] * vector[(0), :] + R[0, 1] * vector[(1), :] + R[0, 2
            ] * vector[(2), :]
        out[1] = R[1, 0] * vector[(0), :] + R[1, 1] * vector[(1), :] + R[1, 2
            ] * vector[(2), :]
        out[2] = R[2, 0] * vector[(0), :] + R[2, 1] * vector[(1), :] + R[2, 2
            ] * vector[(2), :]
    return out
