<gh_stars>0
# File: coordinate.py
# Author: <NAME> <EMAIL> Tsinghua Univ. & EAST
# Usage: Transform between various coordinates.
# Output: 

from numpy import ndarray as ndarray
import numpy as np
def rzphi2xyz(rzphi:ndarray, category:str='coord', merge_return:bool=True):
    # import the modules in the function scope in case of polluting namespace. 
    if category=='coord':
        x = rzphi[..., 0] * np.cos(rzphi[..., 2])
        y = rzphi[..., 0] * np.sin(rzphi[..., 2])
        z = rzphi[..., 1]
    elif category=='field':
        nPhi = rzphi.shape[2]
        Phi = np.arange(0, 2*np.pi, 2*np.pi/nPhi)
        x = rzphi[..., 0] * np.cos(Phi[None,None, :]) - rzphi[...,2] * np.sin(Phi[None,None,:])
        y = rzphi[..., 0] * np.sin(Phi[None,None, :]) + rzphi[...,2] * np.cos(Phi[None,None,:])
        z = rzphi[..., 1]
    else: raise ValueError('Choose appropriate category, either coord or field.')

    if merge_return:
        return np.stack((x, y, z), axis=-1)
    else:
        return x, y, z

def xyz2rzphi(xyz:ndarray, category:str='coord', merge_return:bool=True):
    # import the modules in the function scope in case of polluting namespace. 
    r = np.sqrt(np.square(xyz[..., 0]) + np.square(xyz[..., 1]))
    z = xyz[..., 2]
    # TODO: how to avoid zero division, I have not yet figured out
    if category=='coord':
        phi = np.arctan(xyz[..., 1]/xyz[..., 0]) 
    elif category=='field':
        nPhi = xyz.shape[2]
        Phi = np.arange(0, 2*np.pi, 2*np.pi/nPhi)
        phi = xyz[...,1] * np.cos(Phi[None,None,:]) + xyz[...,0] * np.sin(Phi[None,None,:])
    else: raise ValueError('Choose appropriate category, either coord or field.')

    if merge_return: return np.stack((r, z, phi), axis=-1)
    else: return r, z, phi

def coord_system_change(coord_from:str, coord_to:str, r:ndarray, merge_return:bool=True):
    if coord_from==coord_to: 
        if merge_return: return r
        else: return r[...,0], r[...,1], r[...,2]

    if coord_from=='XYZ':
        if coord_to=='RZPhi': return xyz2rzphi(r, merge_return=merge_return)
        else: raise ValueError(f"From {coord_from} to {coord_to} transfrom has not yet been prepared.")
    elif coord_from=='RZPhi':
        if coord_to=='XYZ': return rzphi2xyz(r, merge_return=merge_return)
        else: raise ValueError(f"From {coord_from} to {coord_to} transfrom has not yet been prepared.")
    else: raise ValueError(f"From {coord_from} to {coord_to} transfrom has not yet been prepared.")

def coord_mirror(coord:str, r:ndarray, plane:str):
    r_new = r.copy()
    # let the z components be inverse
    if plane=='xy':
        if coord=='XYZ': r_new[..., 2] *= -1.0 
        elif coord=='RZPhi': r_new[..., 1] *= -1.0
    else:
        raise ValueError(f"Not yet prepared for mirror around the plane {plane}.")
    return r_new
    
# - ds : $S$ neighboring difference
# - dsR : $dR/ds(s,\theta)$ 
# - dsZ : $dZ/ds(s,\theta)$ 
# - dtheta : $\theta$ neighboring difference 
# - dthetaR : $dR/d\theta(s,\theta)$ 
# - dthetaZ : $dZ/d\theta(s,\theta)$
def Jac_rz2stheta(S, TET, r_mesh, z_mesh):
    assert(S.ndim==1 and TET.ndim==1)
    assert(r_mesh.ndim==2 and z_mesh.ndim==2)

    ds = np.roll(S, 1) - np.roll(S,-1)
    ds[0], ds[-1] = S[1]-S[0], S[-1]-S[-2]
    dsR = np.roll(r_mesh, 1, axis=0) - np.roll(r_mesh, -1, axis=0)
    dsZ = np.roll(z_mesh, 1, axis=0) - np.roll(z_mesh, -1, axis=0)
    dsR[ 0,:], dsZ[ 0,:] = r_mesh[ 1,:]-r_mesh[ 0,:], z_mesh[ 1,:]-z_mesh[ 0,:]
    dsR[-1,:], dsZ[-1,:] = r_mesh[-1,:]-r_mesh[-2,:], z_mesh[-1,:]-z_mesh[-2,:]
    dsR = dsR / ds[:, None]
    dsZ = dsZ / ds[:, None]

    # TODO: Add assert that the theta* is equally-spaced.
    dtheta = np.roll(TET, 1) - np.roll(TET,-1)
    dtheta[0], dtheta[-1] = TET[1]-TET[0], TET[-1]-TET[-2]
    dthetaR = np.roll(r_mesh, 1, axis=1) - np.roll(r_mesh, -1, axis=1)
    dthetaZ = np.roll(z_mesh, 1, axis=1) - np.roll(z_mesh, -1, axis=1)
    dthetaR[:, 0], dthetaZ[:, 0] = r_mesh[:, 1]-r_mesh[:, 0], z_mesh[:, 1]-z_mesh[:, 0]
    dthetaR[:,-1], dthetaZ[:,-1] = r_mesh[:,-1]-r_mesh[:,-2], z_mesh[:,-1]-z_mesh[:,-2]
    dthetaR = dthetaR / dtheta[None,:]
    dthetaZ = dthetaZ / dtheta[None,:]

    det_kl = dsR*dthetaZ - dthetaR*dsZ
    # partial derivations in the "easy" direction: dsr=dr/ds, dsz=dz/ds...

    # loops on the (s,theta) grid indexes k and l
    # for k in range(ns):
    # 	for l in range(ntheta):
        
        # computation of the partial derivated in the interesting direction : drs=ds/dr...
        # A=[dsR dthetaR; dsZ dthetaZ]
        # check the determinant value not to be zero
        # test2(k,l)=det(A)
                            # B=A^(-1)
    dRs    = dthetaZ / det_kl # B(1,1)
    dZs    =-dthetaR / det_kl # B(1,2)
    dRtheta=-dsZ     / det_kl # B(2,1)
    dZtheta= dsR     / det_kl # B(2,2) 
    return dRs, dZs, dRtheta, dZtheta

def RZ2STET(RZ:ndarray, S:ndarray, TET:ndarray, r_mesh:ndarray, z_mesh:ndarray, dRZdSTET_mesh=None):
    
    ###### Unsatisfactory Version ######
    # from scipy.interpolate import griddata
    # S_mesh = np.ones_like(r_mesh) * S[:,None]
    # TET_mesh=np.ones_like(r_mesh)*TET[None,:]
    # points = np.stack((r_mesh.flatten(), z_mesh.flatten()), axis=-1)
    # points_pos = np.stack((r_mesh[:,:len(TET)//2+1].flatten(), z_mesh[:,:len(TET)//2+1].flatten()), axis=-1)
    # points_neg = np.stack((r_mesh[:,len(TET)//2-1:].flatten(), z_mesh[:,len(TET)//2-1:].flatten()), axis=-1)

    # grid_S = griddata(points, S_mesh.flatten(), (RZ[...,0], RZ[...,1]), method='linear')
    # RZ_pos = RZ[RZ[...,1]>=0,:]
    # RZ_neg = RZ[RZ[...,1]< 0,:]
    # # grid_TET = griddata(points,TET_mesh.flatten(),(RZ[...,0], RZ[...,1]), method='cubic') # Such interpolate would cause severe jump interpolated values at the infinitesimal (2pi,0) angle range
    # grid_TET_pos=griddata(points_pos,TET_mesh[:,:len(TET)//2+1].flatten(),(RZ_pos[...,0], RZ_pos[...,1]), method='linear')
    # grid_TET_neg=griddata(points_neg,TET_mesh[:,len(TET)//2-1:].flatten(),(RZ_neg[...,0], RZ_neg[...,1]), method='linear')
    # grid_TET = np.empty((RZ.shape[:-1]))
    # grid_TET[RZ[...,1]>=0], grid_TET[RZ[...,1]< 0] = grid_TET_pos, grid_TET_neg
    # return np.stack((grid_S, grid_TET), axis=-1)

    ###### Version 1 ######
    from scipy.linalg import solve

    RZ_mesh = np.stack((r_mesh,z_mesh), axis=-1)
    STET = np.empty_like(RZ)
    dRdS_mesh, dRdTET_mesh, dZdS_mesh, dZdTET_mesh \
        = calc_dRZdSTET_mesh(S, TET, r_mesh, z_mesh) if dRZdSTET_mesh is None \
        else dRZdSTET_mesh
        
    for x in np.ndindex(RZ.shape[:-1]):
        RZ_mesh_bias = np.linalg.norm(RZ_mesh - RZ[x][None,None,:], axis=-1)
        rz0_ind = np.unravel_index(np.argmin(RZ_mesh_bias, axis=None), RZ_mesh_bias.shape)
        r0,z0 = r_mesh[rz0_ind], z_mesh[rz0_ind]
        dRdTET = dRdTET_mesh[rz0_ind]
        dZdTET = dZdTET_mesh[rz0_ind]
        dRdS = dRdS_mesh[rz0_ind]
        dZdS = dZdS_mesh[rz0_ind]
        ds,dtheta = solve(
            a = np.array(
            [[dRdS, dRdTET], 
            [dZdS, dZdTET]]),
            b = np.array(RZ[x] - [r0,z0]))
        STET[x] = [S[rz0_ind[0]]+ds, TET[rz0_ind[1]]+dtheta]

    return STET

def calc_dRZdSTET_mesh(S:ndarray, TET:ndarray, r_mesh:ndarray, z_mesh:ndarray):
    dTET = TET[1]-TET[0] # The code requires that theta is equally spaced. One may cancel the requirement by using the numpy.gradient function.
    dS = S[1]-S[0]
    # 5 points center difference method to calculate the first order derivative 
    dRdTET_mesh = (
        -  np.roll(r_mesh, -2, axis=1)
        +8*np.roll(r_mesh, -1, axis=1)
        -8*np.roll(r_mesh,  1, axis=1)
        +  np.roll(r_mesh,  2, axis=1)
    ) / (12*dTET)
    dZdTET_mesh = (
        -  np.roll(z_mesh, -2, axis=1)
        +8*np.roll(z_mesh, -1, axis=1)
        -8*np.roll(z_mesh,  1, axis=1)
        +  np.roll(z_mesh,  2, axis=1)
    ) / (12*dTET)
    dRdS_mesh = (
        -  np.roll(r_mesh, -2, axis=0)
        +8*np.roll(r_mesh, -1, axis=0)
        -8*np.roll(r_mesh,  1, axis=0)
        +  np.roll(r_mesh,  2, axis=0)
    ) / (12*dS)
    dZdS_mesh = (
        -  np.roll(z_mesh, -2, axis=0)
        +8*np.roll(z_mesh, -1, axis=0)
        -8*np.roll(z_mesh,  1, axis=0)
        +  np.roll(z_mesh,  2, axis=0)
    ) / (12*dS)
    dRdS_mesh[1, :] = (r_mesh[2, :] - r_mesh[0, :]) / (2*dS)
    dRdS_mesh[-2,:] = (r_mesh[-1,:] - r_mesh[-3,:]) / (2*dS)
    dZdS_mesh[0, :] = (z_mesh[1, :] - z_mesh[0, :]) / (2*dS)
    dZdS_mesh[-1,:] = (z_mesh[1, :] - z_mesh[0, :]) / (2*dS)
    return dRdS_mesh, dRdTET_mesh, dZdS_mesh, dZdTET_mesh

def STET2RZ(STET:ndarray, S:ndarray, TET:ndarray, r_mesh:ndarray, z_mesh:ndarray):
    from scipy.interpolate import RegularGridInterpolator
    R_interp = RegularGridInterpolator((S,TET), r_mesh)
    Z_interp = RegularGridInterpolator((S,TET), z_mesh)
    grid_R = R_interp(STET)
    grid_Z = Z_interp(STET)
    # S_mesh = np.ones_like(r_mesh) * S[:,None]
    # TET_mesh=np.ones_like(r_mesh)*TET[None,:]
    # points = np.stack((S_mesh.flatten(), TET_mesh.flatten()), axis=-1)

    # grid_R = griddata(points, r_mesh.flatten(), (STET[...,0], STET[...,1]), method='linear')
    # grid_Z = griddata(points, z_mesh.flatten(), (STET[...,0], STET[...,1]), method='linear')
    return np.stack((grid_R, grid_Z), axis=-1)

if __name__ == "__main__":
    import file
    import matplotlib 
    matplotlib.use('tkagg') # Agg backend runs without a display
    import matplotlib.pyplot as plt
    equili_dict = file.EquiliDictRead('EAST', '73999_030400ms.mat')
    S, TET = equili_dict['S'][:,0], equili_dict['TET'][:,0]
    r_mesh,z_mesh = equili_dict['r_mesh'], equili_dict['z_mesh']

    STET = np.stack((np.ones((200))*0.7, np.linspace(0,2*np.pi,200, endpoint=True)), axis=-1)
    RZ = STET2RZ(STET, S, TET, r_mesh, z_mesh)
    STET = RZ2STET(RZ, S, TET, r_mesh, z_mesh)
    plt.scatter(STET[:,0], STET[:,1], s=0.05)

    STET = np.stack((np.ones((200))*0.5, np.linspace(0,2*np.pi,200, endpoint=True)), axis=-1)
    RZ = STET2RZ(STET, S, TET, r_mesh, z_mesh)
    STET = RZ2STET(RZ, S, TET, r_mesh, z_mesh)
    plt.scatter(STET[:,0], STET[:,1], s=0.05)

    STET = np.stack((np.ones((200))*0.004, np.linspace(0,2*np.pi,200, endpoint=True)), axis=-1)
    RZ = STET2RZ(STET, S, TET, r_mesh, z_mesh)
    STET = RZ2STET(RZ, S, TET, r_mesh, z_mesh)
    plt.scatter(STET[:,0], STET[:,1], s=0.05)

    plt.show()
# File: coordinate.py
# Author: <NAME> <EMAIL> Tsinghua Univ. & EAST
# Usage: Transform between various coordinates.
# Output: 

from numpy import ndarray as ndarray
import numpy as np
def rzphi2xyz(rzphi:ndarray, category:str='coord', merge_return:bool=True):
    # import the modules in the function scope in case of polluting namespace. 
    if category=='coord':
        x = rzphi[..., 0] * np.cos(rzphi[..., 2])
        y = rzphi[..., 0] * np.sin(rzphi[..., 2])
        z = rzphi[..., 1]
    elif category=='field':
        nPhi = rzphi.shape[2]
        Phi = np.arange(0, 2*np.pi, 2*np.pi/nPhi)
        x = rzphi[..., 0] * np.cos(Phi[None,None, :]) - rzphi[...,2] * np.sin(Phi[None,None,:])
        y = rzphi[..., 0] * np.sin(Phi[None,None, :]) + rzphi[...,2] * np.cos(Phi[None,None,:])
        z = rzphi[..., 1]
    else: raise ValueError('Choose appropriate category, either coord or field.')

    if merge_return:
        return np.stack((x, y, z), axis=-1)
    else:
        return x, y, z

def xyz2rzphi(xyz:ndarray, category:str='coord', merge_return:bool=True):
    # import the modules in the function scope in case of polluting namespace. 
    r = np.sqrt(np.square(xyz[..., 0]) + np.square(xyz[..., 1]))
    z = xyz[..., 2]
    # TODO: how to avoid zero division, I have not yet figured out
    if category=='coord':
        phi = np.arctan(xyz[..., 1]/xyz[..., 0]) 
    elif category=='field':
        nPhi = xyz.shape[2]
        Phi = np.arange(0, 2*np.pi, 2*np.pi/nPhi)
        phi = xyz[...,1] * np.cos(Phi[None,None,:]) + xyz[...,0] * np.sin(Phi[None,None,:])
    else: raise ValueError('Choose appropriate category, either coord or field.')

    if merge_return: return np.stack((r, z, phi), axis=-1)
    else: return r, z, phi

def coord_system_change(coord_from:str, coord_to:str, r:ndarray, merge_return:bool=True):
    if coord_from==coord_to: 
        if merge_return: return r
        else: return r[...,0], r[...,1], r[...,2]

    if coord_from=='XYZ':
        if coord_to=='RZPhi': return xyz2rzphi(r, merge_return=merge_return)
        else: raise ValueError(f"From {coord_from} to {coord_to} transfrom has not yet been prepared.")
    elif coord_from=='RZPhi':
        if coord_to=='XYZ': return rzphi2xyz(r, merge_return=merge_return)
        else: raise ValueError(f"From {coord_from} to {coord_to} transfrom has not yet been prepared.")
    else: raise ValueError(f"From {coord_from} to {coord_to} transfrom has not yet been prepared.")

def coord_mirror(coord:str, r:ndarray, plane:str):
    r_new = r.copy()
    # let the z components be inverse
    if plane=='xy':
        if coord=='XYZ': r_new[..., 2] *= -1.0 
        elif coord=='RZPhi': r_new[..., 1] *= -1.0
    else:
        raise ValueError(f"Not yet prepared for mirror around the plane {plane}.")
    return r_new
    
# - ds : $S$ neighboring difference
# - dsR : $dR/ds(s,\theta)$ 
# - dsZ : $dZ/ds(s,\theta)$ 
# - dtheta : $\theta$ neighboring difference 
# - dthetaR : $dR/d\theta(s,\theta)$ 
# - dthetaZ : $dZ/d\theta(s,\theta)$
def Jac_rz2stheta(S, TET, r_mesh, z_mesh):
    assert(S.ndim==1 and TET.ndim==1)
    assert(r_mesh.ndim==2 and z_mesh.ndim==2)

    ds = np.roll(S, 1) - np.roll(S,-1)
    ds[0], ds[-1] = S[1]-S[0], S[-1]-S[-2]
    dsR = np.roll(r_mesh, 1, axis=0) - np.roll(r_mesh, -1, axis=0)
    dsZ = np.roll(z_mesh, 1, axis=0) - np.roll(z_mesh, -1, axis=0)
    dsR[ 0,:], dsZ[ 0,:] = r_mesh[ 1,:]-r_mesh[ 0,:], z_mesh[ 1,:]-z_mesh[ 0,:]
    dsR[-1,:], dsZ[-1,:] = r_mesh[-1,:]-r_mesh[-2,:], z_mesh[-1,:]-z_mesh[-2,:]
    dsR = dsR / ds[:, None]
    dsZ = dsZ / ds[:, None]

    # TODO: Add assert that the theta* is equally-spaced.
    dtheta = np.roll(TET, 1) - np.roll(TET,-1)
    dtheta[0], dtheta[-1] = TET[1]-TET[0], TET[-1]-TET[-2]
    dthetaR = np.roll(r_mesh, 1, axis=1) - np.roll(r_mesh, -1, axis=1)
    dthetaZ = np.roll(z_mesh, 1, axis=1) - np.roll(z_mesh, -1, axis=1)
    dthetaR[:, 0], dthetaZ[:, 0] = r_mesh[:, 1]-r_mesh[:, 0], z_mesh[:, 1]-z_mesh[:, 0]
    dthetaR[:,-1], dthetaZ[:,-1] = r_mesh[:,-1]-r_mesh[:,-2], z_mesh[:,-1]-z_mesh[:,-2]
    dthetaR = dthetaR / dtheta[None,:]
    dthetaZ = dthetaZ / dtheta[None,:]

    det_kl = dsR*dthetaZ - dthetaR*dsZ
    # partial derivations in the "easy" direction: dsr=dr/ds, dsz=dz/ds...

    # loops on the (s,theta) grid indexes k and l
    # for k in range(ns):
    # 	for l in range(ntheta):
        
        # computation of the partial derivated in the interesting direction : drs=ds/dr...
        # A=[dsR dthetaR; dsZ dthetaZ]
        # check the determinant value not to be zero
        # test2(k,l)=det(A)
                            # B=A^(-1)
    dRs    = dthetaZ / det_kl # B(1,1)
    dZs    =-dthetaR / det_kl # B(1,2)
    dRtheta=-dsZ     / det_kl # B(2,1)
    dZtheta= dsR     / det_kl # B(2,2) 
    return dRs, dZs, dRtheta, dZtheta

def RZ2STET(RZ:ndarray, S:ndarray, TET:ndarray, r_mesh:ndarray, z_mesh:ndarray, dRZdSTET_mesh=None):
    
    ###### Unsatisfactory Version ######
    # from scipy.interpolate import griddata
    # S_mesh = np.ones_like(r_mesh) * S[:,None]
    # TET_mesh=np.ones_like(r_mesh)*TET[None,:]
    # points = np.stack((r_mesh.flatten(), z_mesh.flatten()), axis=-1)
    # points_pos = np.stack((r_mesh[:,:len(TET)//2+1].flatten(), z_mesh[:,:len(TET)//2+1].flatten()), axis=-1)
    # points_neg = np.stack((r_mesh[:,len(TET)//2-1:].flatten(), z_mesh[:,len(TET)//2-1:].flatten()), axis=-1)

    # grid_S = griddata(points, S_mesh.flatten(), (RZ[...,0], RZ[...,1]), method='linear')
    # RZ_pos = RZ[RZ[...,1]>=0,:]
    # RZ_neg = RZ[RZ[...,1]< 0,:]
    # # grid_TET = griddata(points,TET_mesh.flatten(),(RZ[...,0], RZ[...,1]), method='cubic') # Such interpolate would cause severe jump interpolated values at the infinitesimal (2pi,0) angle range
    # grid_TET_pos=griddata(points_pos,TET_mesh[:,:len(TET)//2+1].flatten(),(RZ_pos[...,0], RZ_pos[...,1]), method='linear')
    # grid_TET_neg=griddata(points_neg,TET_mesh[:,len(TET)//2-1:].flatten(),(RZ_neg[...,0], RZ_neg[...,1]), method='linear')
    # grid_TET = np.empty((RZ.shape[:-1]))
    # grid_TET[RZ[...,1]>=0], grid_TET[RZ[...,1]< 0] = grid_TET_pos, grid_TET_neg
    # return np.stack((grid_S, grid_TET), axis=-1)

    ###### Version 1 ######
    from scipy.linalg import solve

    RZ_mesh = np.stack((r_mesh,z_mesh), axis=-1)
    STET = np.empty_like(RZ)
    dRdS_mesh, dRdTET_mesh, dZdS_mesh, dZdTET_mesh \
        = calc_dRZdSTET_mesh(S, TET, r_mesh, z_mesh) if dRZdSTET_mesh is None \
        else dRZdSTET_mesh
        
    for x in np.ndindex(RZ.shape[:-1]):
        RZ_mesh_bias = np.linalg.norm(RZ_mesh - RZ[x][None,None,:], axis=-1)
        rz0_ind = np.unravel_index(np.argmin(RZ_mesh_bias, axis=None), RZ_mesh_bias.shape)
        r0,z0 = r_mesh[rz0_ind], z_mesh[rz0_ind]
        dRdTET = dRdTET_mesh[rz0_ind]
        dZdTET = dZdTET_mesh[rz0_ind]
        dRdS = dRdS_mesh[rz0_ind]
        dZdS = dZdS_mesh[rz0_ind]
        ds,dtheta = solve(
            a = np.array(
            [[dRdS, dRdTET], 
            [dZdS, dZdTET]]),
            b = np.array(RZ[x] - [r0,z0]))
        STET[x] = [S[rz0_ind[0]]+ds, TET[rz0_ind[1]]+dtheta]

    return STET

def calc_dRZdSTET_mesh(S:ndarray, TET:ndarray, r_mesh:ndarray, z_mesh:ndarray):
    dTET = TET[1]-TET[0] # The code requires that theta is equally spaced. One may cancel the requirement by using the numpy.gradient function.
    dS = S[1]-S[0]
    # 5 points center difference method to calculate the first order derivative 
    dRdTET_mesh = (
        -  np.roll(r_mesh, -2, axis=1)
        +8*np.roll(r_mesh, -1, axis=1)
        -8*np.roll(r_mesh,  1, axis=1)
        +  np.roll(r_mesh,  2, axis=1)
    ) / (12*dTET)
    dZdTET_mesh = (
        -  np.roll(z_mesh, -2, axis=1)
        +8*np.roll(z_mesh, -1, axis=1)
        -8*np.roll(z_mesh,  1, axis=1)
        +  np.roll(z_mesh,  2, axis=1)
    ) / (12*dTET)
    dRdS_mesh = (
        -  np.roll(r_mesh, -2, axis=0)
        +8*np.roll(r_mesh, -1, axis=0)
        -8*np.roll(r_mesh,  1, axis=0)
        +  np.roll(r_mesh,  2, axis=0)
    ) / (12*dS)
    dZdS_mesh = (
        -  np.roll(z_mesh, -2, axis=0)
        +8*np.roll(z_mesh, -1, axis=0)
        -8*np.roll(z_mesh,  1, axis=0)
        +  np.roll(z_mesh,  2, axis=0)
    ) / (12*dS)
    dRdS_mesh[1, :] = (r_mesh[2, :] - r_mesh[0, :]) / (2*dS)
    dRdS_mesh[-2,:] = (r_mesh[-1,:] - r_mesh[-3,:]) / (2*dS)
    dZdS_mesh[0, :] = (z_mesh[1, :] - z_mesh[0, :]) / (2*dS)
    dZdS_mesh[-1,:] = (z_mesh[1, :] - z_mesh[0, :]) / (2*dS)
    return dRdS_mesh, dRdTET_mesh, dZdS_mesh, dZdTET_mesh

def STET2RZ(STET:ndarray, S:ndarray, TET:ndarray, r_mesh:ndarray, z_mesh:ndarray):
    from scipy.interpolate import RegularGridInterpolator
    R_interp = RegularGridInterpolator((S,TET), r_mesh)
    Z_interp = RegularGridInterpolator((S,TET), z_mesh)
    grid_R = R_interp(STET)
    grid_Z = Z_interp(STET)
    # S_mesh = np.ones_like(r_mesh) * S[:,None]
    # TET_mesh=np.ones_like(r_mesh)*TET[None,:]
    # points = np.stack((S_mesh.flatten(), TET_mesh.flatten()), axis=-1)

    # grid_R = griddata(points, r_mesh.flatten(), (STET[...,0], STET[...,1]), method='linear')
    # grid_Z = griddata(points, z_mesh.flatten(), (STET[...,0], STET[...,1]), method='linear')
    return np.stack((grid_R, grid_Z), axis=-1)

if __name__ == "__main__":
    import file
    import matplotlib 
    matplotlib.use('tkagg') # Agg backend runs without a display
    import matplotlib.pyplot as plt
    equili_dict = file.EquiliDictRead('EAST', '73999_030400ms.mat')
    S, TET = equili_dict['S'][:,0], equili_dict['TET'][:,0]
    r_mesh,z_mesh = equili_dict['r_mesh'], equili_dict['z_mesh']

    STET = np.stack((np.ones((200))*0.7, np.linspace(0,2*np.pi,200, endpoint=True)), axis=-1)
    RZ = STET2RZ(STET, S, TET, r_mesh, z_mesh)
    STET = RZ2STET(RZ, S, TET, r_mesh, z_mesh)
    plt.scatter(STET[:,0], STET[:,1], s=0.05)

    STET = np.stack((np.ones((200))*0.5, np.linspace(0,2*np.pi,200, endpoint=True)), axis=-1)
    RZ = STET2RZ(STET, S, TET, r_mesh, z_mesh)
    STET = RZ2STET(RZ, S, TET, r_mesh, z_mesh)
    plt.scatter(STET[:,0], STET[:,1], s=0.05)

    STET = np.stack((np.ones((200))*0.004, np.linspace(0,2*np.pi,200, endpoint=True)), axis=-1)
    RZ = STET2RZ(STET, S, TET, r_mesh, z_mesh)
    STET = RZ2STET(RZ, S, TET, r_mesh, z_mesh)
    plt.scatter(STET[:,0], STET[:,1], s=0.05)

    plt.show()
# File: coordinate.py # Author: <NAME> <EMAIL> Tsinghua Univ. & EAST # Usage: Transform between various coordinates. # Output: # import the modules in the function scope in case of polluting namespace. # import the modules in the function scope in case of polluting namespace. # TODO: how to avoid zero division, I have not yet figured out # let the z components be inverse # - ds : $S$ neighboring difference # - dsR : $dR/ds(s,\theta)$ # - dsZ : $dZ/ds(s,\theta)$ # - dtheta : $\theta$ neighboring difference # - dthetaR : $dR/d\theta(s,\theta)$ # - dthetaZ : $dZ/d\theta(s,\theta)$ # TODO: Add assert that the theta* is equally-spaced. # partial derivations in the "easy" direction: dsr=dr/ds, dsz=dz/ds... # loops on the (s,theta) grid indexes k and l # for k in range(ns): # 	for l in range(ntheta): # computation of the partial derivated in the interesting direction : drs=ds/dr... # A=[dsR dthetaR; dsZ dthetaZ] # check the determinant value not to be zero # test2(k,l)=det(A) # B=A^(-1) # B(1,1) # B(1,2) # B(2,1) # B(2,2) ###### Unsatisfactory Version ###### # from scipy.interpolate import griddata # S_mesh = np.ones_like(r_mesh) * S[:,None] # TET_mesh=np.ones_like(r_mesh)*TET[None,:] # points = np.stack((r_mesh.flatten(), z_mesh.flatten()), axis=-1) # points_pos = np.stack((r_mesh[:,:len(TET)//2+1].flatten(), z_mesh[:,:len(TET)//2+1].flatten()), axis=-1) # points_neg = np.stack((r_mesh[:,len(TET)//2-1:].flatten(), z_mesh[:,len(TET)//2-1:].flatten()), axis=-1) # grid_S = griddata(points, S_mesh.flatten(), (RZ[...,0], RZ[...,1]), method='linear') # RZ_pos = RZ[RZ[...,1]>=0,:] # RZ_neg = RZ[RZ[...,1]< 0,:] # # grid_TET = griddata(points,TET_mesh.flatten(),(RZ[...,0], RZ[...,1]), method='cubic') # Such interpolate would cause severe jump interpolated values at the infinitesimal (2pi,0) angle range # grid_TET_pos=griddata(points_pos,TET_mesh[:,:len(TET)//2+1].flatten(),(RZ_pos[...,0], RZ_pos[...,1]), method='linear') # grid_TET_neg=griddata(points_neg,TET_mesh[:,len(TET)//2-1:].flatten(),(RZ_neg[...,0], RZ_neg[...,1]), method='linear') # grid_TET = np.empty((RZ.shape[:-1])) # grid_TET[RZ[...,1]>=0], grid_TET[RZ[...,1]< 0] = grid_TET_pos, grid_TET_neg # return np.stack((grid_S, grid_TET), axis=-1) ###### Version 1 ###### # The code requires that theta is equally spaced. One may cancel the requirement by using the numpy.gradient function. # 5 points center difference method to calculate the first order derivative # S_mesh = np.ones_like(r_mesh) * S[:,None] # TET_mesh=np.ones_like(r_mesh)*TET[None,:] # points = np.stack((S_mesh.flatten(), TET_mesh.flatten()), axis=-1) # grid_R = griddata(points, r_mesh.flatten(), (STET[...,0], STET[...,1]), method='linear') # grid_Z = griddata(points, z_mesh.flatten(), (STET[...,0], STET[...,1]), method='linear') # Agg backend runs without a display