added files

poetry run python plot_fastran.py --command=check --ncfile=f000001.00001 --input=plot.json --pdf=fastran.pdf

{

"yrange": {

    "density": [0, 30.0],

    "temperature": [0, 30.0],

    "rotation": [0, 10.0e3],

    "heating": [0, 1.0],

    "current": [0, 2.5],

    "pressure": [0, 300.0e3],

    "q": [0, 10.0],

    "fluxe": [0, 300.0],

    "fluxi": [0, 300.0]

},

"pannel": {

    "nrows": 2,

    "ncols": 4

},

"define": {

    "density": [ 

    	["ne", 0, "k--"],

    	["ne", -1, "k"]

    ],

    "temperature": [

    	["te", -1, "r"  ],

    	["ti", -1, "b"  ],

    	["te",  0, "r--"],

    	["ti",  0, "b--"]

    ], 

    "rotation": [

    	["omega", -1, "k"  ],

    	["omega",  0, "k--"]

    ],

    "heating": [ 

    	["pi_nb", -1, "k"],

    	["pi_rf", -1, "b"],

    	["pe_rf", -1, "r"]

    ],

    "current": [ 

    	["j_tot", -1, "k"],

    	["j_rf", -1, "r"],

    	["j_nb", -1, "g"]

    ],

    "q": [

    	["q", -1, "k"]

    ],

    "fluxe": [

    	["fluxe_exp", -1, "k"]

    ],

    "fluxi": [

    	["fluxi_exp", -1, "k"]

    ]

}

}

import os

import sys

import glob

import json

import argparse

from plotbase import PlotFastran

def check(input_params, fname_fastran='fastran.nc', fname_pdf='fastran.pdf'):

    plt = PlotFastran(input_params, fname_pdf)

    plt.load_fastran(fname_fastran)

    plt.set_page()

    plt.close()  

def scan(input_params, fname_fastran='fastran.nc', fname_pdf='fastran.pdf'):

    plt = PlotFastran(input_params, fname_pdf)

    files = sorted(glob.glob(fname_fastran))

    for f in files:

       ishot = int(os.path.basename(f).split('.')[0][1:])

       itime = int(os.path.basename(f).split('.')[1])

       print(f, ishot, itime)

       txt = f'ishot = {ishot}, itime = {itime}'

       # txt = 0

       plt.load_fastran(f)

       plt.set_page(txt=txt)

    plt.close()

def main(argv=None):

    parser = argparse.ArgumentParser()

    parser.add_argument('-c', '--command', default='check')    

    parser.add_argument('-n', '--ncfile', default='fastran.nc')    

    parser.add_argument('-i', '--input', default='plot.json')    

    parser.add_argument('-p', '--pdf', default='fastran.pdf')    

    args = parser.parse_args(sys.argv[1:])

    fname_pdf = args.pdf

    fname_fastran = args.ncfile

    fname_input = args.input

    with open(fname_input, 'r') as f:

        input_params =  json.load(f)

    if args.command == 'check':

        check(input_params, fname_fastran, fname_pdf)

    elif args.command == 'scan':

        print(fname_fastran)

        scan(input_params, fname_fastran + '*', fname_pdf)

if __name__=="__main__":

    sys.exit(main())

#!/usr/bin/env python

import sys

import os

import glob

import json

import numpy as np

import netCDF4

from Namelist import Namelist

# default plot preferences

import matplotlib

matplotlib.rcParams['axes.linewidth']=0.5

matplotlib.rc('pdf', fonttype=42)

matplotlib.rcParams['figure.max_open_warning'] = 0

from matplotlib.backends.backend_pdf import PdfPages

import matplotlib.pyplot as plt

class PlotFastran():

    def __init__(self, input_params, fname_pdf='fastran.pdf', page_size_x=14.5, page_size_y=8.5, verbose=False):

        self.pdf = PdfPages(fname_pdf)

        self.input_params = input_params

        self.page_size_x = page_size_x

        self.page_size_y = page_size_y

        self.verbose = verbose

    def close(self):

        self.pdf.close()

    def load_fastran(self, fn_ncfile):

        self.fastran = netCDF4.Dataset(fn_ncfile, 'r', format='NETCDF4')

        self.basename = os.path.basename(fn_ncfile)

    def load_instate(self, fn_instate):

        print(fn_instate)

        self.instate = Namelist(fn_instate)

    def get1d(self, key):

        return self.fastran.variables[key][:]

    def get2d(self, key):

        return self.fastran.variables[key][:,:]

    def set_page(self, txt=''):

        nrows = self.input_params['pannel']['nrows']

        ncols = self.input_params['pannel']['ncols']

        self.fig, self.ax = plt.subplots(

            nrows=nrows, 

            ncols=ncols, 

            figsize=(self.page_size_x, self.page_size_y)

            )

        plt.subplots_adjust(top=0.9, bottom=0.1, left=0.05, right=0.95)

        if txt:

            plt.figtext(0.5, 0.975, txt, horizontalalignment='center', verticalalignment='top')

        plot_list = [key for key in self.input_params['define']]

        print(plot_list)

        if len(plot_list) is not nrows * ncols:

            raise Exception('Number of pannel != nrows * ncols')

        k_plot = 0

        for i_plot in range(nrows):

            for j_plot in range(ncols):

                pannel = plot_list[k_plot]

                define = self.input_params['define'][pannel]

                p = self.ax[i_plot, j_plot]

                for key in define:

                    print(pannel, key)

                    if self.verbose: print(pannel, key)

                    x = self.fastran.variables['rho'][:]

                    y = self.fastran.variables[key[0]][key[1], :]

                    p.plot(x, y, key[2])

                p.set_title(pannel)

                p.set_xlim(0, 1)

                if pannel in self.input_params['yrange']:

                    yrange = self.input_params['yrange'][pannel]

                    if self.verbose: print('set plot range', pannel, yrange)

                    p.set_ylim(yrange[0], yrange[1])

                else:

                    if self.verbose: print('default plot range', pannel)

                k_plot += 1

        self.pdf.savefig()

        plt.clf()