예제 #1
0
def plot_field(
    field,
    filename="C1.h5",
    points=200,
    slice=0,
    range=None,
    xrange=None,
    yrange=None,
    rrange=None,
    zrange=None,
    palette=None,
    lcfs=None,
    bound=None,
    linfac=1.0,
    phi=0.0,
    iabs=None,
    iphase=None,
    isum=None,
    iavg=None,
    idiff=None,
    ilinear=None,
    iequil=None,
    icomplex=None,
    ntor=None,
    title=None,
    fs=1.0,
    ax=None,
):

    if isinstance(field, basestring):
        # Read this field
        if title is None:
            title = field
        field = read_field(
            field,
            slice=slice,
            filename=filename,
            phi=phi,
            points=points,
            rrange=rrange,
            zrange=zrange,
            linfac=linfac,
            iabs=iabs,
            iphase=iphase,
            isum=isum,
            iavg=iavg,
            idiff=idiff,
            ilinear=ilinear,
            iequil=iequil,
            icomplex=icomplex,
            ntor=ntor,
        )

    data = np.real(field.data)

    if range is None:
        vmin = data.min()
        vmax = data.max()
    else:
        vmin = range[0]
        vmax = range[1]

    if palette is None:
        if vmin >= 0.0:
            palette = "inferno"
            col_lcfs = "w"
        else:
            palette = "RdBu_r"
            col_lcfs = "k"
    else:
        col_lcfs = "w"

    cmap = mpl.colors.ListedColormap(sns.color_palette(palette, 256))

    extent = [field.R.data[0], field.R.data[-1], field.Z.data[0], field.Z.data[-1]]

    if ax is None:
        f, ax = plt.subplots(figsize=[fs * 8, fs * 12])
    else:
        f = None

    im = ax.imshow(data.T, origin="lower", vmin=vmin, vmax=vmax, extent=extent, cmap=cmap)

    if xrange is None:
        xrange = [field.R.data[0], field.R.data[-1]]
    if yrange is None:
        yrange = [field.Z.data[0], field.Z.data[-1]]

    ax.set_xlim(xrange)
    ax.set_xlabel(r"$R$ (m)", fontsize=fs * 28)
    ax.set_ylim(yrange)
    ax.set_ylabel(r"$Z$ (m)", fontsize=fs * 28)
    if title is not None:
        ax.set_title(title, fontsize=fs * 32)
    ax.tick_params(labelsize=fs * 24)

    if f is not None:
        cb = f.colorbar(im, format="%1.3g")
        cb.ax.tick_params(labelsize=fs * 24)
    else:
        div = make_axes_locatable(ax)
        cax = div.append_axes("right", size="10%", pad=0.02)
        cb = plt.colorbar(im, cax=cax, format="%1.3g")
        cb.ax.tick_params(labelsize=fs * 24)

    if lcfs is not None:
        if not isinstance(filename, basestring):
            filename = filename[0]

        isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE, filename)
        fio_py.get_options(isrc)

        psi = read_field("psi", slice=-1, filename=filename, points=points, rrange=rrange, zrange=zrange, iequil=1)

        ipsi_lcfs = fio_py.get_series(isrc, fio_py.FIO_LCFS_PSI)

        psi_lcfs = fio_py.eval_series(ipsi_lcfs, 0.0)

        fio_py.close_series(ipsi_lcfs)

        ax.contour(psi.data.T, [psi_lcfs], hold="on", origin="lower", extent=extent, colors=col_lcfs, linewidths=1)

    if bound is not None:
        (Wi, Wo) = get_wall()
        ax.plot(Wi[:, 0], Wi[:, 1], "--", color=col_lcfs, linewidth=1)
        ax.plot(Wo[:, 0], Wo[:, 1], "--", color=col_lcfs, linewidth=1)

    return (f, ax, im, cb)
예제 #2
0
def plot_field(field,
               filename='C1.h5',
               points=200,
               slice=0,
               range=None,
               xrange=None,
               yrange=None,
               rrange=None,
               zrange=None,
               palette=None,
               lcfs=None,
               bound=None,
               fac=1.,
               linfac=1.,
               phi=0.,
               iabs=None,
               iphase=None,
               isum=None,
               iavg=None,
               idiff=None,
               ilinear=None,
               iequil=None,
               icomplex=None,
               ntor=None,
               title=None,
               fs=1.0,
               ax=None,
               symrange=False,
               cb_label=None,
               minval=None,
               maxval=None,
               nimrod=False,
               make_cb=True):

    if isinstance(field, str):
        # Read this field
        if title is None:
            title = field
        field = read_field(field,
                           slice=slice,
                           filename=filename,
                           phi=phi,
                           points=points,
                           rrange=rrange,
                           zrange=zrange,
                           linfac=linfac,
                           iabs=iabs,
                           iphase=iphase,
                           isum=isum,
                           iavg=iavg,
                           idiff=idiff,
                           ilinear=ilinear,
                           iequil=iequil,
                           icomplex=icomplex,
                           ntor=ntor,
                           nimrod=nimrod)

    data = np.nan_to_num(fac * np.real(field.data))

    if range is None:
        if symrange:
            vmin = -abs(data).max()
            vmax = abs(data).max()
        else:
            vmin = data.min()
            vmax = data.max()
    else:
        vmin = range[0]
        vmax = range[1]
    if minval is not None:
        vmin = minval
    if maxval is not None:
        vmax = maxval

    if palette is None:
        if vmin >= 0.:
            palette = 'inferno'
            col_lcfs = 'w'
        else:
            palette = 'RdBu_r'
            col_lcfs = 'k'
    else:
        col_lcfs = 'w'

    cmap = mpl.colors.ListedColormap(sns.color_palette(palette, 256))

    extent = [
        field.R.data[0], field.R.data[-1], field.Z.data[0], field.Z.data[-1]
    ]

    if ax is None:
        f, ax = plt.subplots(figsize=[fs * 9, fs * 12])
    else:
        f = None

    im = ax.imshow(data.T,
                   origin='lower',
                   vmin=vmin,
                   vmax=vmax,
                   extent=extent,
                   cmap=cmap)

    if xrange is None:
        xrange = [field.R.data[0], field.R.data[-1]]
    if yrange is None:
        yrange = [field.Z.data[0], field.Z.data[-1]]

    ax.set_xlim(xrange)
    ax.set_xlabel(r'$R$ (m)', fontsize=fs * 28)
    ax.set_ylim(yrange)
    ax.set_ylabel(r'$Z$ (m)', fontsize=fs * 28)
    if title is not None:
        ax.set_title(title, fontsize=fs * 32)
    ax.tick_params(labelsize=fs * 24)

    if make_cb:
        if f is not None:
            cb = f.colorbar(im, format='%1.3g')
            cb.ax.tick_params(labelsize=fs * 24)
        else:
            div = make_axes_locatable(ax)
            cax = div.append_axes("right", size="10%", pad=0.02)
            cb = plt.colorbar(im, cax=cax, format='%1.3g')
            cb.ax.tick_params(labelsize=fs * 24)
        if cb_label is not None:
            cb.ax.get_yaxis().labelpad = fs * 36
            cb.ax.set_ylabel(cb_label, rotation=270, fontsize=fs * 24)
    else:
        cb = None

    if lcfs is not None:
        if not isinstance(filename, str):
            filename = filename[0]

        isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE, filename)
        fio_py.get_options(isrc)

        psi = read_field('psi',
                         slice=-1,
                         filename=filename,
                         points=points,
                         rrange=rrange,
                         zrange=zrange,
                         iequil=1)

        ipsi_lcfs = fio_py.get_series(isrc, fio_py.FIO_LCFS_PSI)

        psi_lcfs = fio_py.eval_series(ipsi_lcfs, 0.)

        fio_py.close_series(ipsi_lcfs)

        ax.contour(psi.data.T, [psi_lcfs],
                   hold='on',
                   origin='lower',
                   extent=extent,
                   colors=col_lcfs,
                   linewidths=1)

    if isinstance(bound, tuple):
        R0, a, delta, Z0, b = bound
        theta = np.linspace(-np.pi / 2., 3 * np.pi / 2, 1000)
        R = R0 + a * np.cos(theta + delta * np.sin(theta))
        Z = Z0 + b * np.sin(theta)
        ax.plot(R, Z, '-', color=col_lcfs, linewidth=fs * 3)
    elif bound is not None:
        (Wi, Wo) = get_wall()
        ax.plot(Wi[:, 0], Wi[:, 1], '--', color=col_lcfs, linewidth=1)
        ax.plot(Wo[:, 0], Wo[:, 1], '--', color=col_lcfs, linewidth=1)

    if f is not None:
        f.tight_layout()
    else:
        plt.tight_layout()

    return (f, ax, im, cb)
예제 #3
0
def plot_shape(folder='./', rrange=None, zrange=None, bound=False, ax=None,
               legend=True, fs=1.0, Nlvl_in=10, Nlvl_out=1, linewidth=3, 
               title=None):
    
    if ax is None:
        sns.set_style('white')
        f, ax = plt.subplots(figsize=[fs*8,fs*12])
    else:
        f = None
    
    if isinstance(folder,basestring):
           folder = [folder]
        
    Nf = len(folder)
    if Nf <=6:
        cols = sns.color_palette(n_colors=Nf,palette='colorblind')
    else:
        cols = sns.color_palette(n_colors=Nf,palette='inferno')
    
    cols[-1] = (0.,0.,0.)    
    
    for i in range(Nf):
        
        fn = folder[i]+'/C1.h5'

        isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE,fn)
        fio_py.get_options(isrc)
        
        ipsi_axis = fio_py.get_series(isrc, fio_py.FIO_MAGAXIS_PSI)
        ipsi_lcfs = fio_py.get_series(isrc, fio_py.FIO_LCFS_PSI)
    
        psi_axis = fio_py.eval_series(ipsi_axis, 0.)
        psi_lcfs = fio_py.eval_series(ipsi_lcfs, 0.)
    
        fio_py.close_series(ipsi_axis)
        fio_py.close_series(ipsi_lcfs)

#        print([psi_axis,psi_lcfs])
        
        levels = (np.arange(Nlvl_in+Nlvl_out+1.)/Nlvl_in)
        levels = levels*(psi_lcfs-psi_axis)+psi_axis
        if psi_lcfs < psi_axis:
            levels = np.flipud(levels)
        
        psi = read_field('psi',slice=-1,filename=fn,rrange=rrange,
                         zrange=zrange)

        col = mpl.colors.rgb2hex(cols[i-1])
        if i==0:
            hold='off'
        else:
            hold='on'
            
        ax.contour(psi.R,psi.Z,psi.data.T,levels,hold=hold,colors=col,
                        linewidths=linewidth,linestyles='solid')
                        
        if legend:
            h, = ax.plot([np.inf,np.inf],[np.inf,np.inf],'-',color=col,
                         linewidth=linewidth)
            
            h.set_label(folder[i])
    
    if bound:
        (Wi,Wo) = get_wall()
        ax.plot(Wi[:,0],Wi[:,1],'r-',linewidth=1)
        ax.plot(Wo[:,0],Wo[:,1],'r-',linewidth=1)
        
    
    if rrange is None:
        rrange = [min(psi.R),max(psi.R)]
    if zrange is  None:
        zrange = [min(psi.Z),max(psi.Z)]
        
    ax.set_xlim(rrange)
    ax.set_ylim(zrange)
    ax.set_xlabel(r'$R$ (m)',fontsize=28*fs)
    ax.set_ylabel(r'$Z$ (m)',fontsize=28*fs)
    ax.tick_params(labelsize=24*fs)
    if title is not None:
        ax.set_title(title,fontsize=32*fs)
    if legend:
        ax.legend(fontsize=24*fs)
    if f is not None:
        plt.tight_layout()
    
    return
예제 #4
0
def read_field(name, slice=0, filename='C1.h5', points=200, phi=0.,
               rrange=None, zrange=None, linfac=1., iabs=None, iphase=None,
               isum=None, iavg=None, idiff=None,ilinear=None, iequil=None, 
               icomplex=None, ntor=None):
    
    # Redefine zero options to None for convenience
    if iabs==0:
        iabs = None
    if iphase==0:
        iphase = None
    if isum==0:
        isum = None
    if iavg==0:
        iavg = None
    if idiff==0:
        idiff = None
    if ilinear==0:
        ilinear = None
    if iequil==0:
        iequil = None
    if icomplex==0:
        icomplex = None
    
    if (isum is not None) or (iavg is not None) or (idiff is not None):
        
        if isinstance(filename,basestring):
            filename = [filename]
        if isinstance(slice,int):
            slice = [slice]
        
        N_files  = len(filename)
        N_slices = len(slice)
        N = N_files*N_slices
        
        if (N_files > 1) and (N_slices > 1):
            print("Error:  can't combine mutliple filenames and slices")
            return None
        elif N_files > 1:
            slice = slice*N
        elif N_slices > 1:
            filename = filename*N
            
        if not hasattr(linfac,"__len__"):
            # Duplicate constant linear factor                
            linfac = np.ones(N)*linfac
        elif len(linfac) != N:
            print("Error:  incorrect number of linfacs")
            return None
            
        if not hasattr(phi,"__len__"):
            # Duplicate constant linear factor                
            phi = np.ones(N)*phi
        elif len(phi) != N:
            print("Error:  incorrect number of phis")
            return None
            
        field = read_field(name, slice=slice[0], filename=filename[0], 
                           points=points, phi=phi[0], rrange=rrange, 
                           zrange=zrange, linfac=linfac[0], 
                           ilinear=ilinear, iequil=iequil, 
                           icomplex=icomplex, ntor=ntor)
                               
        if N==1:
            return field
                                 
        for i in [x+1 for x in range(N-1)]:
            field2 = read_field(name, slice=slice[i], filename=filename[i], 
                                points=points, phi=phi[i], rrange=rrange, 
                                zrange=zrange, linfac=linfac[i], 
                                ilinear=ilinear, iequil=iequil, 
                                icomplex=icomplex, ntor=ntor)
            if (isum is not None) or (iavg is not None):
                field = field + field2
            elif (idiff is not None):
                field = field + field2*((-1)**i)
                
            
        if iavg is not None:
            return field/N
        else:
            return field
            
    else:
        if not isinstance(filename,basestring):
            print('Warning:  Only considering first filename')
            filename = filename[0]
        if  hasattr(slice,"__len__"):
            print('Warning:  Only considering first slice')
            slice = slice[0]
        if hasattr(linfac,"__len__"):
            print('Warning:  Only considering first linfac')
            linfac = linfac[0]
        if hasattr(phi,"__len__"):
            print('Warning:  Only considering first phi')
            phi = phi[0]
            
        
    
    # initialize the M3D-C1 file
    isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE,filename)
    fio_py.get_options(isrc)
    
    # Set up options given as imput
    fio_py.set_int_option(fio_py.FIO_TIMESLICE,slice)
    
    if ilinear is not None:    
        fio_py.set_int_option(fio_py.FIO_PART,fio_py.FIO_PERTURBED_ONLY)
    elif iequil is not None:
        fio_py.set_int_option(fio_py.FIO_PART,fio_py.FIO_EQUILIBRIUM_ONLY)
    else:
        fio_py.set_int_option(fio_py.FIO_PART,fio_py.FIO_TOTAL)
    
    
    if rrange is None:
        rrange = [0.75,2.5]      # I want to read this from the file somehow
    if zrange is None:
        zrange = [-1.5,1.5]  # I want to read this from the file somehow

        
    R = np.linspace(rrange[0],rrange[1],points)
    Z = np.linspace(zrange[0],zrange[1],points)
    
    # Are we reading a real field?
    if icomplex is None:
        ntor = None
    elif ntor is None:
        print('Warning:  ntor must be assigned for complex evaluation')
        print('Reading real field')
        icomplex = None
    elif ilinear is None:
        print('Warning:  Trying to access complex values but ilinear=None')
        print('          This has not been properly implemented yet')
        print('Reading real field')
        icomplex = None
        ntor = None
    
    if icomplex is None:
        dtype = np.float64
    else:
        dtype = np.complex128
    data = np.zeros((points,points),dtype=dtype)   
    
    # Lists for scalar fields read from scalars
    total_pressure    = ['total pressure','p']
    ion_pressure      = ['ion pressure','pi']
    electron_pressure = ['electron pressure','pe']
    ion_density       = ['ion density', 'ni', 'den']
    electron_density  = ['electron density', 'ne']
    
    # Dictionaries for scalar fields read  from vectors
    B_comp =  {'bx':'R', 'by':'phi', 'bt':'phi', 'bphi':'phi', 'bz':'Z'}
    J_comp =  {'jx':'R', 'jy':'phi', 'jt':'phi', 'jphi':'phi', 'jz':'Z'}
    E_comp =  {'ex':'R', 'ey':'phi', 'et':'phi', 'ephi':'phi', 'ez':'Z'}
    v_comp =  {'vx':'R', 'vy':'phi', 'vt':'phi', 'vphi':'phi', 'vz':'Z'}
    A_comp =  {'ax':'R', 'ay':'phi', 'at':'phi', 'aphi':'phi', 'az':'Z'}
    comps  = ['R', 'phi', 'Z']
    
    # Lists for vector fields
    magnetic_field   = ['magnetic field', 'b field', 'bfield', 'b_field']
    current_density  = ['current density', 'j']
    electric_field   = ['electric field', 'e', 'efield', 'e_field']
    fluid_velocity   = ['fluid velocity', 'v']
    vector_potential = ['vector potential', 'a']
    
    # Lists for composite fields
    major_radius         = ['major radius', 'r', 'x']
    height               = ['height', 'z']
    ion_temperature      = ['ion_temperature', 'ti']
    electron_temperature = ['electron_temperature', 'te']

    name_lc = name.lower()

    # Primitive scalar fields
    if name_lc in total_pressure:

        field = xr.DataArray(data,[('R',R),('Z',Z)])
        field.attrs['type']    = fio_py.FIO_TOTAL_PRESSURE
        field.name = total_pressure[0]
    
    elif name_lc in ion_pressure:

        field = xr.DataArray(data,[('R',R),('Z',Z)])
        field.attrs['species'] = fio_py.FIO_MAIN_ION
        field.attrs['type']    = fio_py.FIO_PRESSURE
        field.name = ion_pressure[0]

    elif name_lc in ion_density:

        field = xr.DataArray(data,[('R',R),('Z',Z)])
        field.attrs['species'] = fio_py.FIO_MAIN_ION
        field.attrs['type']    = fio_py.FIO_DENSITY
        field.name = ion_density[0]

    elif name_lc in electron_pressure:

        field = xr.DataArray(data,[('R',R),('Z',Z)])
        field.attrs['species'] = fio_py.FIO_ELECTRON
        field.attrs['type']    = fio_py.FIO_PRESSURE
        field.name = electron_pressure[0]
        
    elif name_lc in electron_density:

        field = xr.DataArray(data,[('R',R),('Z',Z)])
        field.attrs['species'] = fio_py.FIO_ELECTRON
        field.attrs['type']    = fio_py.FIO_DENSITY
        field.name = electron_density[0]
        
    # Primitive vector fields
    elif name_lc in magnetic_field:

        data = np.array([data,data,data]) 
        
        field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])    
        field.attrs['type'] = fio_py.FIO_MAGNETIC_FIELD
        field.name = magnetic_field[0]

    elif name_lc in current_density:
        
        data = np.array([data,data,data]) 
        
        field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])  
        field.attrs['type'] = fio_py.FIO_CURRENT_DENSITY
        field.name = current_density[0]
    
    elif name_lc in electric_field:
        
        data = np.array([data,data,data])
        
        field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])    
        field.attrs['type'] = fio_py.FIO_ELECTRIC_FIELD
        field.name = electric_field[0]
        
    elif name_lc in fluid_velocity:
        
        data = np.array([data,data,data])
        
        field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])    
        field.attrs['type'] = fio_py.FIO_FLUID_VELOCITY
        field.name = fluid_velocity[0]
        
    elif name_lc in vector_potential:
        
        data = np.array([data,data,data])
        
        field = xr.DataArray(data,[('component',comps),('R',R),('Z',Z)])    
        field.attrs['type'] = fio_py.FIO_VECTOR_POTENTIAL
        field.name = vector_potential[0]
                
    #  Vector components
    elif name_lc in B_comp:

        B = read_field('bfield', slice=slice, filename=filename, points=points, 
                       phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
                       iequil=iequil, icomplex=icomplex, ntor=ntor)
        
        field = B.sel(component=B_comp[name_lc])
        field.attrs['type'] = 'component'

    elif name_lc in J_comp:

        J = read_field('j', slice=slice, filename=filename, points=points, 
                       phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
                       iequil=iequil, icomplex=icomplex, ntor=ntor)
        
        field = J.sel(component=J_comp[name_lc])
        field.attrs['type'] = 'component'
    
    elif name_lc in E_comp:

        E = read_field('efield', slice=slice, filename=filename, points=points, 
                       phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
                       iequil=iequil, icomplex=icomplex, ntor=ntor)
        
        field = E.sel(component=E_comp[name_lc])
        field.attrs['type'] = 'component'
        
    elif name_lc in v_comp:

        v = read_field('v', slice=slice, filename=filename, points=points, 
                       phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
                       iequil=iequil, icomplex=icomplex, ntor=ntor)
        
        field = v.sel(component=v_comp[name_lc])
        field.attrs['type'] = 'component'
        
    elif name_lc in A_comp:

        A = read_field('a', slice=slice, filename=filename, points=points, 
                       phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
                       iequil=iequil, icomplex=icomplex, ntor=ntor)
        
        field = A.sel(component=A_comp[name_lc])
        field.attrs['type'] = 'component'
                      
    # Composite fields
    elif name_lc == 'zero':
        field = xr.DataArray(data,[('R',R),('Z',Z)])
        field.attrs['type'] = 'composite'
        field.name = name_lc
        
    elif name_lc in major_radius:
        
        field = xr.DataArray(data,[('R',R),('Z',Z)])
        
        for j in range(points):

            field[dict(Z=j)]=field.coords['R']
        
        field.attrs['type'] = 'composite'
        field.name = major_radius[0]
    
    elif name_lc in height:
        
        for i in range(points):
            field[dict(R=i)]=field.coords['Z']
        
        field.attrs['type'] = 'composite'
        field.name = height[0]
        
    elif name_lc == 'beta':
        
        if ilinear is not None:
            
            print(name + ': Cannot calculate linear version')
            print('Running for total instead')   
        
        P = read_field('p', slice=slice, filename=filename, points=points,
                       phi=phi, rrange=rrange, zrange=zrange,iequil=iequil)
        B = read_field('b', slice=slice, filename=filename, points=points,
                       phi=phi, rrange=rrange, zrange=zrange,iequil=iequil)
                       
        field = 2.0*spc.mu_0*P/(B**2)
        
        field.attrs['type'] = 'composite'
        field.name = name_lc
        
    elif name_lc in ion_temperature:
        
        Pi = read_field('pi', slice=slice, filename=filename, points=points,
                        phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
                        iequil=iequil, icomplex=icomplex, ntor=ntor)
                        
        ni = read_field('ni', slice=slice, filename=filename, points=points,
                        phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
                        iequil=iequil, icomplex=icomplex, ntor=ntor)
                        
        if ilinear is None:
            Ti = Pi/ni            
        else:
            Pi0 = read_field('pi', slice=slice, filename=filename, points=points,
                             phi=phi, rrange=rrange, zrange=zrange, iequil=1)
                        
            ni0 = read_field('ni', slice=slice, filename=filename, points=points,
                             phi=phi, rrange=rrange, zrange=zrange, iequil=1)
            
            Ti = Pi/ni0 - (Pi0*ni)/(ni0**2)
            
        J_eV = spc.physical_constants['electron volt'][0]
        
        field = Ti/J_eV
        
        field.attrs['type'] = 'composite'
        field.name = ion_temperature[0]
        
    elif name_lc in electron_temperature:
        
        Pe = read_field('pe', slice=slice, filename=filename, points=points,
                        phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
                        iequil=iequil, icomplex=icomplex, ntor=ntor)
                        
        ne = read_field('ne', slice=slice, filename=filename, points=points,
                        phi=phi, rrange=rrange, zrange=zrange, ilinear=ilinear,
                        iequil=iequil, icomplex=icomplex, ntor=ntor)
                        
        if ilinear is not None:
            Pe0 = read_field('pe', slice=slice, filename=filename, points=points,
                             phi=phi, rrange=rrange, zrange=zrange, iequil=1)
                        
            ne0 = read_field('ne', slice=slice, filename=filename, points=points,
                             phi=phi, rrange=rrange, zrange=zrange, iequil=1)
            Te = Pe/ne0 - (Pe0*ne)/(ne0**2)
        else:
            Te = Pe/ne  
            
            
        J_eV = spc.physical_constants['electron volt'][0]
        
        field = Te/J_eV
        
        field.attrs['type'] = 'composite'
        field.name = electron_temperature[0]
        
    elif name_lc == 'b':
        
        B = read_field('bfield', slice=slice, filename=filename,
                       points=points, phi=phi, rrange=rrange, zrange=zrange, 
                       ilinear=ilinear, iequil=iequil, icomplex=icomplex,
                       ntor=ntor)
                      
        Bx = B.sel(component='R')
        By = B.sel(component='phi')
        Bz = B.sel(component='Z')
        
        if ilinear is not None:
            B0 = read_field('bfield', slice=slice, filename=filename, phi=phi, 
                            points=points, rrange=rrange, zrange=zrange, 
                            iequil=1)
            Bx0 = B0.sel(component='R')
            By0 = B0.sel(component='phi')
            Bz0 = B0.sel(component='Z')
        
            b0 = (Bx0**2 + By0**2 + Bz0**2)**0.5

            b = (Bx*Bx0 + By*By0 + Bz*Bz0)/b0
            
        else:
            b = (Bx**2 + By**2 + Bz**2)**0.5
            
        field = b.drop('component')
        field.attrs['type'] = 'composite'
        field.name = name_lc
        
    elif name_lc == 'b2':
        
        B = read_field('bfield', slice=slice, filename=filename, 
                       points=points, phi=phi, rrange=rrange, zrange=zrange, 
                       ilinear=ilinear, iequil=iequil, icomplex=icomplex,
                       ntor=ntor)
                      
        Bx = B.sel(component='R')
        By = B.sel(component='phi')
        Bz = B.sel(component='Z')
        
        if ilinear is not None:
            B0 = read_field('bfield', slice=slice, filename=filename, phi=phi, 
                            points=points, rrange=rrange, zrange=zrange, 
                            iequil=1)
            Bx0 = B0.sel(component='R')
            By0 = B0.sel(component='phi')
            Bz0 = B0.sel(component='Z')
        
            b2 = 2.0*(Bx*Bx0 + By*By0 + Bz*Bz0)
            
        else:
            b2 = Bx**2 + By**2 + Bz**2
            
        field = b2.drop('component')
        field.attrs['type'] = 'composite'
        field.name = name_lc
    
    elif name_lc == 'omega':

        vy = read_field('vy', slice=slice, filename=filename, phi=phi, 
                        points=points, rrange=rrange, zrange=zrange, 
                        ilinear=ilinear, iequil=iequil, icomplex=icomplex,
                       ntor=ntor)
        R = read_field('r', slice=slice, filename=filename, phi=phi, 
                       points=points, rrange=rrange, zrange=zrange, 
                       ilinear=ilinear, iequil=iequil, icomplex=icomplex,
                       ntor=ntor)
                        
        field = vy/R
        
        field.attrs['type'] = 'composite'
        field.name = name_lc
        
    elif name_lc == 'psi':
        
        Aphi = read_field('ay', slice=slice, filename=filename, phi=phi, 
                        points=points, rrange=rrange, zrange=zrange, 
                        ilinear=ilinear, iequil=iequil, icomplex=icomplex,
                       ntor=ntor)
        R = read_field('r', slice=slice, filename=filename, phi=phi, 
                       points=points, rrange=rrange, zrange=zrange, 
                       ilinear=ilinear, iequil=iequil, icomplex=icomplex,
                       ntor=ntor)

        field = R*Aphi
        field = field.drop('component')
        
        field.attrs['type'] = 'composite'
        field.name = name_lc
        
    else:
        print "Field '" + name_lc + "' is+- not defined"
        return None
        
    field.attrs['ntor'] = ntor
    
    #**************************************************************************
    
    if not isinstance(field.attrs['type'],basestring):
        
        if 'species' in field.attrs:
            fio_py.set_int_option(fio_py.FIO_SPECIES, field.attrs['species'])
    
        handle = fio_py.get_field(isrc,field.attrs['type'])
        
        if len(field.shape) == 2:
            enum = np.ndenumerate(field)
            eval_field = fio_py.eval_scalar_field
        elif len(field.shape) == 3:
            enum = np.ndenumerate(field.data[0,:,:])
            eval_field = fio_py.eval_vector_field
        
        for (i,j), dummy in enum:
    
            # Real part
    
            x = (R[i], phi*np.pi/180., Z[j])
            
            val = np.asarray(eval_field(handle, x),dtype=dtype)
            
            # Imaginary part
            if icomplex is not None:
                y = (x[0], x[1] + 3.0*np.pi/(2.0*ntor), x[2])
                val += 1j*np.asarray(eval_field(handle, y),dtype=dtype)
            
            field[dict(R=i,Z=j)] = val
    
    return linfac*field
예제 #5
0
def read_field(name,
               slice=0,
               filename='C1.h5',
               points=200,
               phi=0.,
               rrange=None,
               zrange=None,
               linfac=1.,
               iabs=None,
               iphase=None,
               isum=None,
               iavg=None,
               idiff=None,
               ilinear=None,
               iequil=None,
               icomplex=None,
               ntor=None,
               nimrod=False):

    # Redefine zero options to None for convenience
    if iabs == 0:
        iabs = None
    if iphase == 0:
        iphase = None
    if isum == 0:
        isum = None
    if iavg == 0:
        iavg = None
    if idiff == 0:
        idiff = None
    if ilinear == 0:
        ilinear = None
    if iequil == 0:
        iequil = None
    if icomplex == 0:
        icomplex = None

    if (isum is not None) or (iavg is not None) or (idiff is not None):

        if isinstance(filename, str):
            filename = [filename]
        if isinstance(slice, int):
            slice = [slice]

        N_files = len(filename)
        N_slices = len(slice)
        N = N_files * N_slices

        if (N_files > 1) and (N_slices > 1):
            print("Error:  can't combine mutliple filenames and slices")
            return None
        elif N_files > 1:
            slice = slice * N
        elif N_slices > 1:
            filename = filename * N

        if not hasattr(linfac, "__len__"):
            # Duplicate constant linear factor
            linfac = np.ones(N) * linfac
        elif len(linfac) != N:
            print("Error:  incorrect number of linfacs")
            return None

        if not hasattr(phi, "__len__"):
            # Duplicate constant linear factor
            phi = np.ones(N) * phi
        elif len(phi) != N:
            print("Error:  incorrect number of phis")
            return None

        field = read_field(name,
                           slice=slice[0],
                           filename=filename[0],
                           points=points,
                           phi=phi[0],
                           rrange=rrange,
                           zrange=zrange,
                           linfac=linfac[0],
                           ilinear=ilinear,
                           iequil=iequil,
                           icomplex=icomplex,
                           ntor=ntor)

        if N == 1:
            return field

        for i in [x + 1 for x in range(N - 1)]:
            field2 = read_field(name,
                                slice=slice[i],
                                filename=filename[i],
                                points=points,
                                phi=phi[i],
                                rrange=rrange,
                                zrange=zrange,
                                linfac=linfac[i],
                                ilinear=ilinear,
                                iequil=iequil,
                                icomplex=icomplex,
                                ntor=ntor)
            if (isum is not None) or (iavg is not None):
                field = field + field2
            elif (idiff is not None):
                field = field + field2 * ((-1)**i)

        if iavg is not None:
            return field / N
        else:
            return field

    else:
        if not isinstance(filename, str):
            print('Warning:  Only considering first filename')
            filename = filename[0]
        if hasattr(slice, "__len__"):
            print('Warning:  Only considering first slice')
            slice = slice[0]
        if hasattr(linfac, "__len__"):
            print('Warning:  Only considering first linfac')
            linfac = linfac[0]
        if hasattr(phi, "__len__"):
            print('Warning:  Only considering first phi')
            phi = phi[0]

    if nimrod:
        #filename is actually a folder
        R = np.loadtxt(filename + '/r.txt')
        Z = np.loadtxt(filename + '/z.txt')
        data = np.loadtxt(filename + '/' + name + ('%d' % slice).zfill(5) +
                          '.txt')
        data = np.transpose(np.nan_to_num(data))
        field = xr.DataArray(data, coords=[('Z', Z), ('R', R)])
        return field

    # initialize the M3D-C1 file
    isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE, filename)
    fio_py.get_options(isrc)

    # Set up options given as imput
    fio_py.set_int_option(fio_py.FIO_TIMESLICE, slice)

    if ilinear is not None:
        fio_py.set_int_option(fio_py.FIO_PART, fio_py.FIO_PERTURBED_ONLY)
    elif iequil is not None:
        fio_py.set_int_option(fio_py.FIO_PART, fio_py.FIO_EQUILIBRIUM_ONLY)
    else:
        fio_py.set_int_option(fio_py.FIO_PART, fio_py.FIO_TOTAL)

    if rrange is None:
        rrange = [0.75, 2.5]  # I want to read this from the file somehow
    if zrange is None:
        zrange = [-1.5, 1.5]  # I want to read this from the file somehow

    R = np.linspace(rrange[0], rrange[1], points)
    Z = np.linspace(zrange[0], zrange[1], points)

    # Are we reading a real field?
    if icomplex is None:
        ntor = None
    elif ntor is None:
        print('Warning:  ntor must be assigned for complex evaluation')
        print('Reading real field')
        icomplex = None
    elif ilinear is None:
        print('Warning:  Trying to access complex values but ilinear=None')
        print('          This has not been properly implemented yet')
        print('Reading real field')
        icomplex = None
        ntor = None

    if icomplex is None:
        dtype = np.float64
    else:
        dtype = np.complex128
    data = np.zeros((points, points), dtype=dtype)

    # Lists for scalar fields read from scalars
    total_pressure = ['total pressure', 'p']
    ion_pressure = ['ion pressure', 'pi']
    electron_pressure = ['electron pressure', 'pe']
    ion_density = ['ion density', 'ni', 'den']
    electron_density = ['electron density', 'ne']

    # Dictionaries for scalar fields read  from vectors
    B_comp = {'bx': 'R', 'by': 'phi', 'bt': 'phi', 'bphi': 'phi', 'bz': 'Z'}
    J_comp = {'jx': 'R', 'jy': 'phi', 'jt': 'phi', 'jphi': 'phi', 'jz': 'Z'}
    E_comp = {'ex': 'R', 'ey': 'phi', 'et': 'phi', 'ephi': 'phi', 'ez': 'Z'}
    v_comp = {'vx': 'R', 'vy': 'phi', 'vt': 'phi', 'vphi': 'phi', 'vz': 'Z'}
    A_comp = {'ax': 'R', 'ay': 'phi', 'at': 'phi', 'aphi': 'phi', 'az': 'Z'}
    comps = ['R', 'phi', 'Z']

    # Lists for vector fields
    magnetic_field = ['magnetic field', 'b field', 'bfield', 'b_field']
    current_density = ['current density', 'j']
    electric_field = ['electric field', 'e', 'efield', 'e_field']
    fluid_velocity = ['fluid velocity', 'v']
    vector_potential = ['vector potential', 'a']

    # Lists for composite fields
    major_radius = ['major radius', 'r', 'x']
    height = ['height', 'z']
    ion_temperature = ['ion_temperature', 'ti']
    electron_temperature = ['electron_temperature', 'te']
    Bpl_comp = {
        'bx_plasma': 'R',
        'by_plasma': 'phi',
        'bt_plasma': 'phi',
        'bphi_plasma': 'phi',
        'bz_plasma': 'Z'
    }
    Jpl_comp = {
        'jx_plasma': 'R',
        'jy_plasma': 'phi',
        'jt_plasma': 'phi',
        'jphi_plasma': 'phi',
        'jz_plasma': 'Z'
    }

    name_lc = name.lower()

    # Primitive scalar fields
    if name_lc in total_pressure:

        field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
        field.attrs['type'] = fio_py.FIO_TOTAL_PRESSURE
        field.name = total_pressure[0]

    elif name_lc in ion_pressure:

        field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
        field.attrs['species'] = fio_py.FIO_MAIN_ION
        field.attrs['type'] = fio_py.FIO_PRESSURE
        field.name = ion_pressure[0]

    elif name_lc in ion_density:

        field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
        field.attrs['species'] = fio_py.FIO_MAIN_ION
        field.attrs['type'] = fio_py.FIO_DENSITY
        field.name = ion_density[0]

    elif name_lc in electron_pressure:

        field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
        field.attrs['species'] = fio_py.FIO_ELECTRON
        field.attrs['type'] = fio_py.FIO_PRESSURE
        field.name = electron_pressure[0]

    elif name_lc in electron_density:

        field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
        field.attrs['species'] = fio_py.FIO_ELECTRON
        field.attrs['type'] = fio_py.FIO_DENSITY
        field.name = electron_density[0]

    # Primitive vector fields
    elif name_lc in magnetic_field:

        data = np.array([data, data, data])

        field = xr.DataArray(data,
                             coords=[('component', comps), ('R', R), ('Z', Z)])
        field.attrs['type'] = fio_py.FIO_MAGNETIC_FIELD
        field.name = magnetic_field[0]

    elif name_lc in current_density:

        data = np.array([data, data, data])

        field = xr.DataArray(data,
                             coords=[('component', comps), ('R', R), ('Z', Z)])
        field.attrs['type'] = fio_py.FIO_CURRENT_DENSITY
        field.name = current_density[0]

    elif name_lc in electric_field:

        data = np.array([data, data, data])

        field = xr.DataArray(data,
                             coords=[('component', comps), ('R', R), ('Z', Z)])
        field.attrs['type'] = fio_py.FIO_ELECTRIC_FIELD
        field.name = electric_field[0]

    elif name_lc in fluid_velocity:

        data = np.array([data, data, data])

        field = xr.DataArray(data,
                             coords=[('component', comps), ('R', R), ('Z', Z)])
        field.attrs['type'] = fio_py.FIO_FLUID_VELOCITY
        field.name = fluid_velocity[0]

    elif name_lc in vector_potential:

        data = np.array([data, data, data])

        field = xr.DataArray(data,
                             coords=[('component', comps), ('R', R), ('Z', Z)])
        field.attrs['type'] = fio_py.FIO_VECTOR_POTENTIAL
        field.name = vector_potential[0]

    #  Vector components
    elif name_lc in B_comp:

        B = read_field('bfield',
                       slice=slice,
                       filename=filename,
                       points=points,
                       phi=phi,
                       rrange=rrange,
                       zrange=zrange,
                       ilinear=ilinear,
                       iequil=iequil,
                       icomplex=icomplex,
                       ntor=ntor)

        field = B.sel(component=B_comp[name_lc])
        field.attrs['type'] = 'component'

    elif name_lc in J_comp:

        J = read_field('j',
                       slice=slice,
                       filename=filename,
                       points=points,
                       phi=phi,
                       rrange=rrange,
                       zrange=zrange,
                       ilinear=ilinear,
                       iequil=iequil,
                       icomplex=icomplex,
                       ntor=ntor)

        field = J.sel(component=J_comp[name_lc])
        field.attrs['type'] = 'component'

    elif name_lc in E_comp:

        E = read_field('efield',
                       slice=slice,
                       filename=filename,
                       points=points,
                       phi=phi,
                       rrange=rrange,
                       zrange=zrange,
                       ilinear=ilinear,
                       iequil=iequil,
                       icomplex=icomplex,
                       ntor=ntor)

        field = E.sel(component=E_comp[name_lc])
        field.attrs['type'] = 'component'

    elif name_lc in v_comp:

        v = read_field('v',
                       slice=slice,
                       filename=filename,
                       points=points,
                       phi=phi,
                       rrange=rrange,
                       zrange=zrange,
                       ilinear=ilinear,
                       iequil=iequil,
                       icomplex=icomplex,
                       ntor=ntor)

        field = v.sel(component=v_comp[name_lc])
        field.attrs['type'] = 'component'

    elif name_lc in A_comp:

        A = read_field('a',
                       slice=slice,
                       filename=filename,
                       points=points,
                       phi=phi,
                       rrange=rrange,
                       zrange=zrange,
                       ilinear=ilinear,
                       iequil=iequil,
                       icomplex=icomplex,
                       ntor=ntor)

        field = A.sel(component=A_comp[name_lc])
        field.attrs['type'] = 'component'

    # Composite fields
    elif name_lc == 'zero':
        field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])
        field.attrs['type'] = 'composite'
        field.name = name_lc

    elif name_lc in major_radius:

        field = xr.DataArray(data, coords=[('R', R), ('Z', Z)])

        for j in range(points):

            field[dict(Z=j)] = field.coords['R']

        field.attrs['type'] = 'composite'
        field.name = major_radius[0]

    elif name_lc in height:

        for i in range(points):
            field[dict(R=i)] = field.coords['Z']

        field.attrs['type'] = 'composite'
        field.name = height[0]

    elif name_lc == 'beta':

        if ilinear is not None:

            print(name + ': Cannot calculate linear version')
            print('Running for total instead')

        P = read_field('p',
                       slice=slice,
                       filename=filename,
                       points=points,
                       phi=phi,
                       rrange=rrange,
                       zrange=zrange,
                       iequil=iequil)
        B = read_field('b',
                       slice=slice,
                       filename=filename,
                       points=points,
                       phi=phi,
                       rrange=rrange,
                       zrange=zrange,
                       iequil=iequil)

        field = 2.0 * spc.mu_0 * P / (B**2)

        field.attrs['type'] = 'composite'
        field.name = name_lc

    elif name_lc in ion_temperature:

        Pi = read_field('pi',
                        slice=slice,
                        filename=filename,
                        points=points,
                        phi=phi,
                        rrange=rrange,
                        zrange=zrange,
                        ilinear=ilinear,
                        iequil=iequil,
                        icomplex=icomplex,
                        ntor=ntor)

        ni = read_field('ni',
                        slice=slice,
                        filename=filename,
                        points=points,
                        phi=phi,
                        rrange=rrange,
                        zrange=zrange,
                        ilinear=ilinear,
                        iequil=iequil,
                        icomplex=icomplex,
                        ntor=ntor)

        if ilinear is None:
            Ti = Pi / ni
        else:
            Pi0 = read_field('pi',
                             slice=slice,
                             filename=filename,
                             points=points,
                             phi=phi,
                             rrange=rrange,
                             zrange=zrange,
                             iequil=1)

            ni0 = read_field('ni',
                             slice=slice,
                             filename=filename,
                             points=points,
                             phi=phi,
                             rrange=rrange,
                             zrange=zrange,
                             iequil=1)

            Ti = Pi / ni0 - (Pi0 * ni) / (ni0**2)

        J_eV = spc.physical_constants['electron volt'][0]

        field = Ti / J_eV

        field.attrs['type'] = 'composite'
        field.name = ion_temperature[0]

    elif name_lc in electron_temperature:

        Pe = read_field('pe',
                        slice=slice,
                        filename=filename,
                        points=points,
                        phi=phi,
                        rrange=rrange,
                        zrange=zrange,
                        ilinear=ilinear,
                        iequil=iequil,
                        icomplex=icomplex,
                        ntor=ntor)

        ne = read_field('ne',
                        slice=slice,
                        filename=filename,
                        points=points,
                        phi=phi,
                        rrange=rrange,
                        zrange=zrange,
                        ilinear=ilinear,
                        iequil=iequil,
                        icomplex=icomplex,
                        ntor=ntor)

        if ilinear is not None:
            Pe0 = read_field('pe',
                             slice=slice,
                             filename=filename,
                             points=points,
                             phi=phi,
                             rrange=rrange,
                             zrange=zrange,
                             iequil=1)

            ne0 = read_field('ne',
                             slice=slice,
                             filename=filename,
                             points=points,
                             phi=phi,
                             rrange=rrange,
                             zrange=zrange,
                             iequil=1)
            Te = Pe / ne0 - (Pe0 * ne) / (ne0**2)
        else:
            Te = Pe / ne

        J_eV = spc.physical_constants['electron volt'][0]

        field = Te / J_eV

        field.attrs['type'] = 'composite'
        field.name = electron_temperature[0]

    elif name_lc == 'b':

        B = read_field('bfield',
                       slice=slice,
                       filename=filename,
                       points=points,
                       phi=phi,
                       rrange=rrange,
                       zrange=zrange,
                       ilinear=ilinear,
                       iequil=iequil,
                       icomplex=icomplex,
                       ntor=ntor)

        Bx = B.sel(component='R')
        By = B.sel(component='phi')
        Bz = B.sel(component='Z')

        if ilinear is not None:
            B0 = read_field('bfield',
                            slice=slice,
                            filename=filename,
                            phi=phi,
                            points=points,
                            rrange=rrange,
                            zrange=zrange,
                            iequil=1)
            Bx0 = B0.sel(component='R')
            By0 = B0.sel(component='phi')
            Bz0 = B0.sel(component='Z')

            b0 = (Bx0**2 + By0**2 + Bz0**2)**0.5

            b = (Bx * Bx0 + By * By0 + Bz * Bz0) / b0

        else:
            b = (Bx**2 + By**2 + Bz**2)**0.5

        field = b.drop('component')
        field.attrs['type'] = 'composite'
        field.name = name_lc

    elif name_lc == 'b2':

        B = read_field('bfield',
                       slice=slice,
                       filename=filename,
                       points=points,
                       phi=phi,
                       rrange=rrange,
                       zrange=zrange,
                       ilinear=ilinear,
                       iequil=iequil,
                       icomplex=icomplex,
                       ntor=ntor)

        Bx = B.sel(component='R')
        By = B.sel(component='phi')
        Bz = B.sel(component='Z')

        if ilinear is not None:
            B0 = read_field('bfield',
                            slice=slice,
                            filename=filename,
                            phi=phi,
                            points=points,
                            rrange=rrange,
                            zrange=zrange,
                            iequil=1)
            Bx0 = B0.sel(component='R')
            By0 = B0.sel(component='phi')
            Bz0 = B0.sel(component='Z')

            b2 = 2.0 * (Bx * Bx0 + By * By0 + Bz * Bz0)

        else:
            b2 = Bx**2 + By**2 + Bz**2

        field = b2.drop('component')
        field.attrs['type'] = 'composite'
        field.name = name_lc

    elif name_lc in Bpl_comp:

        B0 = read_field('bfield',
                        slice=0,
                        filename=filename,
                        points=points,
                        phi=phi,
                        rrange=rrange,
                        zrange=zrange,
                        ilinear=ilinear,
                        iequil=iequil,
                        icomplex=icomplex,
                        ntor=ntor)
        B1 = read_field('bfield',
                        slice=slice,
                        filename=filename,
                        points=points,
                        phi=phi,
                        rrange=rrange,
                        zrange=zrange,
                        ilinear=ilinear,
                        iequil=iequil,
                        icomplex=icomplex,
                        ntor=ntor)

        Bpl = B1 - B0
        field = Bpl.sel(component=Bpl_comp[name_lc])
        field.attrs['type'] = 'component'

    elif name_lc in Jpl_comp:

        J0 = read_field('j',
                        slice=0,
                        filename=filename,
                        points=points,
                        phi=phi,
                        rrange=rrange,
                        zrange=zrange,
                        ilinear=ilinear,
                        iequil=iequil,
                        icomplex=icomplex,
                        ntor=ntor)
        J1 = read_field('j',
                        slice=slice,
                        filename=filename,
                        points=points,
                        phi=phi,
                        rrange=rrange,
                        zrange=zrange,
                        ilinear=ilinear,
                        iequil=iequil,
                        icomplex=icomplex,
                        ntor=ntor)
        Jpl = J1 - J0
        field = Jpl.sel(component=Jpl_comp[name_lc])
        field.attrs['type'] = 'component'

    elif name_lc == 'omega':

        vy = read_field('vy',
                        slice=slice,
                        filename=filename,
                        phi=phi,
                        points=points,
                        rrange=rrange,
                        zrange=zrange,
                        ilinear=ilinear,
                        iequil=iequil,
                        icomplex=icomplex,
                        ntor=ntor)
        R = read_field('r',
                       slice=slice,
                       filename=filename,
                       phi=phi,
                       points=points,
                       rrange=rrange,
                       zrange=zrange,
                       ilinear=ilinear,
                       iequil=iequil,
                       icomplex=icomplex,
                       ntor=ntor)

        field = vy / R

        field.attrs['type'] = 'composite'
        field.name = name_lc

    elif name_lc == 'psi':

        Aphi = read_field('ay',
                          slice=slice,
                          filename=filename,
                          phi=phi,
                          points=points,
                          rrange=rrange,
                          zrange=zrange,
                          ilinear=ilinear,
                          iequil=iequil,
                          icomplex=icomplex,
                          ntor=ntor)
        R = read_field('r',
                       slice=slice,
                       filename=filename,
                       phi=phi,
                       points=points,
                       rrange=rrange,
                       zrange=zrange,
                       ilinear=ilinear,
                       iequil=iequil,
                       icomplex=icomplex,
                       ntor=ntor)

        field = R * Aphi
        field = field.drop('component')

        field.attrs['type'] = 'composite'
        field.name = name_lc

    else:
        print("Field '" + name_lc + "' is not defined")
        return None

    field.attrs['ntor'] = ntor

    #**************************************************************************

    if not isinstance(field.attrs['type'], str):

        if 'species' in field.attrs:
            fio_py.set_int_option(fio_py.FIO_SPECIES, field.attrs['species'])

        handle = fio_py.get_field(isrc, field.attrs['type'])
        hint = fio_py.allocate_hint(isrc)

        if len(field.shape) == 2:
            enum = np.ndenumerate(field)
            eval_field = fio_py.eval_scalar_field
        elif len(field.shape) == 3:
            enum = np.ndenumerate(field.data[0, :, :])
            eval_field = fio_py.eval_vector_field

        for (i, j), dummy in enum:

            # Real part

            x = (R[i], phi * np.pi / 180., Z[j])

            val = np.asarray(eval_field(handle, x, hint), dtype=dtype)

            # Imaginary part
            if icomplex is not None:
                y = (x[0], x[1] + 3.0 * np.pi / (2.0 * ntor), x[2])
                val += 1j * np.asarray(eval_field(handle, y, hint),
                                       dtype=dtype)

            field[dict(R=i, Z=j)] = val

    return linfac * field
예제 #6
0
import fio_py

filename = "C1.h5"

isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE, filename)

fio_py.get_options(isrc)
ia = fio_py.get_field(isrc, fio_py.FIO_VECTOR_POTENTIAL)
imag = fio_py.get_field(isrc, fio_py.FIO_MAGNETIC_FIELD)
ipres = fio_py.get_field(isrc, fio_py.FIO_TOTAL_PRESSURE)
ivel = fio_py.get_field(isrc, fio_py.FIO_FLUID_VELOCITY)
ij = fio_py.get_field(isrc, fio_py.FIO_CURRENT_DENSITY)

fio_py.set_int_option(fio_py.FIO_SPECIES, fio_py.FIO_ELECTRON)
idens = fio_py.get_field(isrc, fio_py.FIO_DENSITY)

fio_py.set_int_option(fio_py.FIO_SPECIES, fio_py.FIO_MAIN_ION)
ipi = fio_py.get_field(isrc, fio_py.FIO_PRESSURE)

list = fio_py.get_available_fields(isrc)

print 'Available fields:'
for field in list:
    print ' ', fio_py.get_field_name(field)

cs = fio_py.get_int_parameter(isrc, FIO_GEOMETRY)
if cs == fio_py.FIO_CYLINDRICAL:
    print 'Using CYLINDRICAL coordinate system'
else:
    print 'Using CARTESIAN coordinate system'
예제 #7
0
def plot_shape(folder='./',
               rrange=None,
               zrange=None,
               bound=False,
               ax=None,
               legend=True,
               fs=1.0,
               Nlvl_in=10,
               Nlvl_out=1,
               linewidth=3,
               title=None,
               labels=None):

    if ax is None:
        sns.set_style('white')
        f, ax = plt.subplots(figsize=[fs * 8, fs * 12])
    else:
        f = None

    if isinstance(folder, str):
        folder = [folder]

    Nf = len(folder)
    if Nf <= 6:
        cols = sns.color_palette(n_colors=6, palette='colorblind')
    else:
        cols = sns.color_palette(n_colors=Nf, palette='inferno')

    cols[-1] = (0., 0., 0.)

    for i in range(Nf):

        fn = folder[i] + '/C1.h5'

        isrc = fio_py.open_source(fio_py.FIO_M3DC1_SOURCE, fn)
        fio_py.get_options(isrc)

        ipsi_axis = fio_py.get_series(isrc, fio_py.FIO_MAGAXIS_PSI)
        ipsi_lcfs = fio_py.get_series(isrc, fio_py.FIO_LCFS_PSI)

        psi_axis = fio_py.eval_series(ipsi_axis, 0.)
        psi_lcfs = fio_py.eval_series(ipsi_lcfs, 0.)

        fio_py.close_series(ipsi_axis)
        fio_py.close_series(ipsi_lcfs)

        #        print([psi_axis,psi_lcfs])

        levels = (np.arange(Nlvl_in + Nlvl_out + 1.) / Nlvl_in)
        levels = levels * (psi_lcfs - psi_axis) + psi_axis
        if psi_lcfs < psi_axis:
            levels = np.flipud(levels)

        psi = read_field('psi',
                         slice=-1,
                         filename=fn,
                         rrange=rrange,
                         zrange=zrange)

        col = mpl.colors.rgb2hex(cols[i])
        if i == 0:
            hold = 'off'
        else:
            hold = 'on'

        ax.contour(psi.R,
                   psi.Z,
                   psi.data.T,
                   levels,
                   hold=hold,
                   colors=col,
                   linewidths=linewidth,
                   linestyles='solid')

        if legend:
            h, = ax.plot([np.inf, np.inf], [np.inf, np.inf],
                         '-',
                         color=col,
                         linewidth=linewidth)
            if labels is None:
                h.set_label(folder[i])
            else:
                h.set_label(labels[i])
    if bound:
        (Wi, Wo) = get_wall()
        ax.plot(Wi[:, 0], Wi[:, 1], 'k-', linewidth=2)
        ax.plot(Wo[:, 0], Wo[:, 1], 'k-', linewidth=2)

    if rrange is None:
        rrange = [min(psi.R), max(psi.R)]
    if zrange is None:
        zrange = [min(psi.Z), max(psi.Z)]

    ax.set_xlim(rrange)
    ax.set_ylim(zrange)
    ax.set_xlabel(r'$R$ (m)', fontsize=28 * fs)
    ax.set_ylabel(r'$Z$ (m)', fontsize=28 * fs)
    ax.tick_params(labelsize=24 * fs)
    if title is not None:
        ax.set_title(title, fontsize=32 * fs)
    if legend:
        ax.legend(fontsize=20 * fs, loc='upper right')
    if f is not None:
        plt.tight_layout()

    return