Possible V2 enhancement: correct4sensor_orientation does not broadcast across all frequencies

[k-a-mendoza]

The correct4sensor_orientation(Z_prime, Bx=0, By=90, Ex=0, Ey=90, Z_prime_error=None) method in mtpy.core.z only accepts 2x2 matrices for Z_prime and Z_prime_error. Given that there is a frequency dimension for every tensor, this seems a bit limiting.

Here's a conversion of the code into something that uses numpy only and allows for (2x2xn) inputs. I also cleaned up the function to be a little more flexible and expressive

def _rotate_tensor(array,ex=0,ey=90,bx=0,by=90):
    T = np.matrix(np.zeros((2, 2)))
    U = np.matrix(np.zeros((2, 2)))

    ex_complex_projection = np.exp(np.deg2rad(ex)*1.0j)
    ey_complex_projection = np.exp(np.deg2rad(ey)*1.0j)
    bx_complex_projection = np.exp(np.deg2rad(bx)*1.0j)
    by_complex_projection = np.exp(np.deg2rad(by)*1.0j)

    T[0, 0] = np.real(ex_complex_projection)
    T[1, 0] = np.imag(ex_complex_projection)

    T[0, 1] = np.real(ey_complex_projection)
    T[1, 1] = np.imag(ey_complex_projection)

    U[0, 0] = np.real(bx_complex_projection)
    U[1, 0] = np.imag(bx_complex_projection)

    U[0, 1] = np.real(by_complex_projection)
    U[1, 1] = np.imag(by_complex_projection)

    first_product  = np.einsum('...ij,jk->...ik',array,U.getI())
    second_product = np.einsum('ij,...jk->...ik',T,first_product)
    
    return second_product

def correct4sensor_orientation(z,z_err=None,**kwargs):
    
    assert np.iscomplexobj(z), f"z tensor is not of expected variable type. should be complex, is {z.dtype}"
    assert (len(z.shape)==3) | (len(z.shape)==2), \
            f"z tensor is not of expected shape. Should be nx2x2. is {z.shape}"
    assert (z_err is not None) & (np.isrealobj(z_err)),\
            f"z_error tensor is of incorrect variable type. Is {z_err.dtype}, should be float"
    assert (z_err is not None) & ((len(z_err.shape)==3) | (len(z_err.shape)==2)),\
            f"z_error tensor is of incorrect shape. Is {z_err.shape}, should be 2x2 or nx2x2"
    
    if len(z.shape)==3:
        assert (z.shape[1]==2) & (z.shape[1]==2), \
            f"z tensor ij dimensions not of expected shape. Should be nx2x2. is {z.shape}"
    if len(z.shape)==2:
        assert (z.shape[0]==2) & (z.shape[0]==2), \
            f"z tensor ij dimensions not of expected shape. Should be nx2x2. is {z.shape}"
    if len(z_err.shape)==3:
        assert (z_err.shape[1]==2) & (z_err.shape[1]==2), \
            f"z_error tensor ij dimensions not of expected shape. Should be nx2x2. is {z.shape}"
    if len(z_err.shape)==2:
        assert (z_err.shape[0]==2) & (z_err.shape[0]==2), \
            f"z_error tensor ij dimensions not of expected shape. Should be 2x2. is {z.shape}"

   
    # checks
    z_prime = _rotate_tensor(z,**kwargs)
    if z_err is not None:
        z_err = _rotate_tensor(z_err,**kwargs)
        
    return z_prime, z_err