diff --git a/srcPython/run_plot_model_results.py b/srcPython/aetherpy/run_plot_model_results.py similarity index 100% rename from srcPython/run_plot_model_results.py rename to srcPython/aetherpy/run_plot_model_results.py diff --git a/srcPython/thermo_histograms.py b/srcPython/aetherpy/thermo_histograms.py similarity index 100% rename from srcPython/thermo_histograms.py rename to srcPython/aetherpy/thermo_histograms.py diff --git a/srcPython/thermo_plot.py b/srcPython/aetherpy/thermo_plot.py similarity index 100% rename from srcPython/thermo_plot.py rename to srcPython/aetherpy/thermo_plot.py diff --git a/srcPython/gitm_3d_test.py b/srcPython/deprecated/gitm_3d_test.py similarity index 100% rename from srcPython/gitm_3d_test.py rename to srcPython/deprecated/gitm_3d_test.py diff --git a/srcPython/useful_functions.py b/srcPython/deprecated/useful_functions.py similarity index 100% rename from srcPython/useful_functions.py rename to srcPython/deprecated/useful_functions.py diff --git a/srcPython/pGITM.py b/srcPython/pGITM.py index 4851760c..2d7f7ea1 100755 --- a/srcPython/pGITM.py +++ b/srcPython/pGITM.py @@ -19,6 +19,14 @@ from scipy.io import FortranFile from struct import unpack +# check if netcdf can be imported. if not, don't error unless writing netcdf files +canWriteCDF = True +try: + from netCDF4 import Dataset + from pyitm.fileio import variables +except: + canWriteCDF = False + # This should work on most systems: endChar='<' @@ -28,7 +36,8 @@ def get_args_pgitm(): parser = argparse.ArgumentParser( - description = 'Post-process GITM model results') + description = "Post-process GITM model results." + " (note: NetCDF functionality is only available through post_process.py)") parser.add_argument('-v', \ action='store_true', default = False, \ help = 'set verbose to true') @@ -86,6 +95,132 @@ def write_gitm_file(file, data, isVerbose = False): fp.close() +def create_netcdf(filename, data, isVerbose=False): + + if isVerbose: + print(" --> Creating netCDF file:", filename) + + # get dimensions: + nz, ny, nx = data["Longitude"].shape + + with Dataset(filename, mode="w", format="NETCDF4") as ncfile: + # Dimensions + t = ncfile.createDimension('time', None) + xdim = ncfile.createDimension('lon', nx) + ydim = ncfile.createDimension('lat', ny) + zdim = ncfile.createDimension('z', nz) + + # time! + reftime = data['time'].date() + time = ncfile.createVariable('time', np.float64, ('time',)) + time.units = 'seconds since ' + str(reftime) + time.long_name = 'time' + time[:] = [(data['time'] + - datetime.datetime(reftime.year, reftime.month, reftime.day) + ).total_seconds()] + + if isVerbose: + print(" --> Created dataset with dimensions:") + [print(dim) for dim in ncfile.dimensions.items()] + + # Add in attributes. Version, etc. + ncfile.title = f"{filename[:5]} GITM Outputs" + ncfile.model = "GITM" + ncfile.version = data['version'] + + if isVerbose: + print(" --> Attributes added. Current ncfile:") + print(ncfile) + + # We'll try to add the coordinates, if they can be added cleanly... + if np.all(data['Longitude'][0, 0, :] == data['Longitude']): + lon = ncfile.createVariable('lon', np.float64, ('lon')) + lon[:] = np.rad2deg(data['Longitude'][0, 0, :]) + lon.units = 'degrees_east' + lon.long_name = 'Longitude' + + if np.all([data['Latitude'][0, :, 0] == data['Latitude'][i, :, :].T for i in range(nx)]):#, axis=(0, 2)): + lat = ncfile.createVariable('lat', np.float64, ('lat')) + lat[:] = np.rad2deg(data['Latitude'][0, :, 0]) + lat.units = 'degrees_north' + lat.long_name = 'Latitude' + + if np.all(data['Altitude'][:, 0, 0] == data['Altitude'].T):#, axis=(0, 1)): + alt = ncfile.createVariable('z', np.float64, ('z')) + alt[:] = data['Altitude'][:, 0, 0] / 1000 + alt.units = 'km' + alt.long_name = 'Altitude' + + # Then put the variables in + for varname in data['vars']: + unit = None + newname = variables.get_short_names([varname])[0] + thisvar = ncfile.createVariable(newname, np.float64, + ('time', 'lon', 'lat', 'z')) + thisvar[:] = data[varname].T + if 'V' in newname: + unit = 'm/s' + elif '[' in varname: + unit='/m3' + elif 'Temp' in varname: + unit = 'K' + if unit: + ncfile[newname].units = unit + + ncfile[newname].long_name = variables.get_long_names([varname])[0] + + return + +def append_netcdf(filename, data, isVerbose): + + if isVerbose: + print(" --> file already found. appending...") + + with Dataset(filename, mode='a', format='NETCDF4') as ncfile: + time = ncfile['time'] + + reftime = datetime.datetime.strptime(time.units.split("since ")[-1], "%Y-%m-%d") + time[len(time)] = (data["time"] - reftime).total_seconds() + + for varname in data['vars']: + newname = variables.get_short_names([varname])[0] + ncfile[newname][-1, ...] = data[varname].T + + return + + +def write_to_netcdf(file, data, combine=False, runname='', isVerbose=False): + """ + Dispatcher for writing/appending netCDFs + + inputs: + file (str): file name (that the bin would have been written to). Used to + determine if the data is 2D/3D & whether we need to drop ghost cells + data (dict): The gitm outputs. + runname (str, optional): name to add to processed files. output format is: + 'runname_3DALL.nc', or if no runname is given its just '3DALL.nc' + isVerbose (bool): whether to print extra info. + + """ + + if combine: + file = file.split('_')[0] + '.nc' + if runname: + file = runname + '_' + file + else: + file = file.replace('bin', 'nc') + + if file.startswith('1D'): + raise ValueError("Cannot write netcdf files for 1D outputs yet!") + + if (isVerbose): + print(' -> Writing NetCDF for: ', file, " to:", runname) + + if not os.path.exists(file): + create_netcdf(file, data, isVerbose) + else: + append_netcdf(file, data, isVerbose) + # ---------------------------------------------------------------------------- # # ---------------------------------------------------------------------------- @@ -298,7 +433,8 @@ def remove_files(header, isVerbose = False): # # ---------------------------------------------------------------------------- -def process_one_file(header, isVerbose = False): +def process_one_file(header, isVerbose = False, dowrite=True, + write_nc=False, combine=False, runname=''): fileOut = header.replace('.header','.bin') @@ -362,11 +498,23 @@ def process_one_file(header, isVerbose = False): iLonStart:iLonEnd] iBlock += 1 - write_gitm_file(fileOut, data, isVerbose = isVerbose) + if dowrite: + if write_nc: + if canWriteCDF: + write_to_netcdf(fileOut, data, runname=runname, combine=combine, + isVerbose=isVerbose) + else: + raise ImportError("Cant import netCDF4") + else: # default !! + write_gitm_file(fileOut, data, isVerbose = isVerbose) + return + + # return data if dowrite is false (for debugging) return data -def post_process_gitm(dir, doRemove, isVerbose = False): +def post_process_gitm(dir, doRemove, isVerbose = False, + write_nc=False, combine=False, runname=''): # Get into the correct directory for processing processDir = dir @@ -384,6 +532,8 @@ def post_process_gitm(dir, doRemove, isVerbose = False): headers = sorted(glob('*.header')) if (isVerbose): print('Found %i files to process' % len(headers)) + elif len(headers) == 0: + print(" -> Did not find any files to process.") # Process each single header: for head in headers: @@ -392,11 +542,16 @@ def post_process_gitm(dir, doRemove, isVerbose = False): print(f' --> Processing {head}') if os.path.exists("../../PostGITM.exe"): + if write_nc: + raise ValueError( + "Cannot use PostGITM and netCDF at the same time, yet." + "\n >> either remove PostGITM.exe or disable -nc option") # Check if we can use the old postprocessor. it's faster. # Pipe the header filename into PostGITM.exe run(f"echo {head} | ../../PostGITM.exe ", check=True, shell=True) else: - data = process_one_file(head, isVerbose = isVerbose) + process_one_file(head, isVerbose=isVerbose, + write_nc=write_nc, combine=combine, runname=runname) if (doRemove): remove_files(head, isVerbose = isVerbose) diff --git a/srcPython/plot_drivers.py b/srcPython/plot_drivers.py index f1983101..c836008b 100755 --- a/srcPython/plot_drivers.py +++ b/srcPython/plot_drivers.py @@ -5,6 +5,7 @@ import numpy as np import matplotlib.pyplot as plt import matplotlib.dates as dates +import argparse, os from swmf_imf import * from read_ae import * @@ -47,72 +48,127 @@ def plot_day_boundaries(ax, times): ax.xaxis.set_major_formatter(dates.DateFormatter('%d-%H')) return -imfFile = 'imf.txt' -aeFile = 'ae.txt' - -imf = read_swmf_file(imfFile) -ae = read_ae(aeFile) - -fig = plt.figure(figsize = (7,10)) -plt.rcParams.update({'font.size': 14}) - -zeros = np.array(imf["bz"]) * 0.0 - -nY = 4 -ax = [] - -xStart = 0.15 -xSize = 0.81 -ySize = 0.19 -yGap = 0.04 -yEnd = 0.95 - -iY = 1 -ax.append(fig.add_subplot(nY*100 + 10 + iY)) -yStart = yEnd - ySize * iY - yGap * (iY-1) -ax[-1].set_position([xStart, yStart, xSize, ySize]) - -ax[-1].plot(imf["times"], imf["bz"]) -ax[-1].plot(imf["times"], zeros, 'k:') -ax[-1].set_ylabel('(a) IMF Bz (nT)') -ax[-1].set_title('IMF, Solar Wind, and AE Drivers for GITM') -plot_day_boundaries(ax[-1], imf["times"]) - -iY = 2 -ax.append(fig.add_subplot(nY*100 + 10 + iY)) -yStart = yEnd - ySize * iY - yGap * (iY-1) -ax[-1].set_position([xStart, yStart, xSize, ySize]) -ax[-1].plot(imf["times"], imf["by"]) -ax[-1].plot(imf["times"], zeros, 'k:') -ax[-1].set_ylabel('(b) IMF By (nT)') -plot_day_boundaries(ax[-1], imf["times"]) - -iY = 3 -ax.append(fig.add_subplot(nY*100 + 10 + iY)) -yStart = yEnd - ySize * iY - yGap * (iY-1) -ax[-1].set_position([xStart, yStart, xSize, ySize]) -ax[-1].plot(imf["times"], imf["vx"]) -ax[-1].set_ylabel('(c) SW Vx (km/s)') -plot_day_boundaries(ax[-1], imf["times"]) - -iY = 4 -ax.append(fig.add_subplot(nY*100 + 10 + iY)) -yStart = yEnd - ySize * iY - yGap * (iY-1) -ax[-1].set_position([xStart, yStart, xSize, ySize]) -ax[-1].plot(ae["time"], ae["ae"]) -ax[-1].set_ylabel('(d) AE (nT)') -ax[-1].set_ylim(0.0, np.max(ae["ae"])*1.05) - -plot_day_boundaries(ax[-1], imf["times"]) - -mint = np.min(imf["times"]) -maxt = np.max(imf["times"]) - -sTime = mint.strftime('%b %d, %Y %H:%M UT') + ' - ' + \ - maxt.strftime('%b %d, %Y %H:%M UT') -ax[-1].set_xlabel(sTime) - -plotfile = imf["times"][0].strftime('drivers%Y%m%d.png') -fig.savefig(plotfile) -plt.close() + +def main(imf=None, ae=None): + """ + Makes plots of IMF and/or AE + + inputs: + + imf: returned from swmf_imf.read_swmf_file() + ae: returned from read_ae.readae() + + Outputs: + nothing + + A plot is created called 'driversYYYYMMDD.png' (min date given) + + """ + + + fig = plt.figure(figsize = (7,10)) + nY = 4 + if imf is not None: + zeros = np.array(imf["bz"]) * 0.0 + elif ae is not None: + zeros = np.array(ae["ae"]) * 0.0 + else: + raise ValueError("Neither AE nor IMF were provided!") + + plt.rcParams.update({'font.size': 14}) + + ax = [] + + xStart = 0.15 + xSize = 0.81 + ySize = 0.19 + yGap = 0.04 + yEnd = 0.95 + + iY = 0 + + if imf is not None: + iY += 1 + ax.append(fig.add_subplot(nY*100 + 10 + iY)) + yStart = yEnd - ySize * iY - yGap * (iY-1) + ax[-1].set_position([xStart, yStart, xSize, ySize]) + + ax[-1].plot(imf["times"], imf["bz"]) + ax[-1].plot(imf["times"], zeros, 'k:') + ax[-1].set_ylabel('(a) IMF Bz (nT)') + ax[-1].set_title('IMF, Solar Wind, and AE Drivers for GITM') + plot_day_boundaries(ax[-1], imf["times"]) + + iY += 1 + ax.append(fig.add_subplot(nY*100 + 10 + iY)) + yStart = yEnd - ySize * iY - yGap * (iY-1) + ax[-1].set_position([xStart, yStart, xSize, ySize]) + ax[-1].plot(imf["times"], imf["by"]) + ax[-1].plot(imf["times"], zeros, 'k:') + ax[-1].set_ylabel('(b) IMF By (nT)') + plot_day_boundaries(ax[-1], imf["times"]) + + iY += 1 + ax.append(fig.add_subplot(nY*100 + 10 + iY)) + yStart = yEnd - ySize * iY - yGap * (iY-1) + ax[-1].set_position([xStart, yStart, xSize, ySize]) + ax[-1].plot(imf["times"], imf["vx"]) + ax[-1].set_ylabel('(c) SW Vx (km/s)') + plot_day_boundaries(ax[-1], imf["times"]) + mint = np.min(imf["times"]) + maxt = np.max(imf["times"]) + + if ae is not None: + iY += 1 + ax.append(fig.add_subplot(nY*100 + 10 + iY)) + yStart = yEnd - ySize * iY - yGap * (iY-1) + ax[-1].set_position([xStart, yStart, xSize, ySize]) + ax[-1].plot(ae["time"], ae["ae"]) + ax[-1].set_ylabel('(d) AE (nT)') + ax[-1].set_ylim(0.0, np.max(ae["ae"])*1.05) + + plot_day_boundaries(ax[-1], ae["time"]) + mint = np.min(ae["time"]) + maxt = np.max(ae["time"]) + + sTime = mint.strftime('%b %d, %Y %H:%M UT') + ' - ' + \ + maxt.strftime('%b %d, %Y %H:%M UT') + ax[-1].set_xlabel(sTime) + + plotfile = mint.strftime('drivers%Y%m%d.png') + fig.savefig(plotfile) + plt.close() + + + +def parse_args(): + parser = argparse.ArgumentParser( + description = 'Plot AE and IMF data') + + parser.add_argument('-imf', type=str, default='imf.txt', + help='Path to IMF file') + + parser.add_argument('-ae', type=str, default='ae.txt', + help='Path to AE file') + + args = parser.parse_args() + + return args + + +if __name__ == '__main__': + + args = parse_args() + + if os.path.exists(args.imf): + imf = read_swmf_file(args.imf) + else: + imf = None + + if os.path.exists(args.ae): + ae = read_ae(args.ae) + else: + ae=None + + main(imf, ae) diff --git a/srcPython/post_process.py b/srcPython/post_process.py index 755f7455..4e7246b3 100755 --- a/srcPython/post_process.py +++ b/srcPython/post_process.py @@ -44,12 +44,8 @@ def parse_args_post(): parser.add_argument('-totaltime', help = 'specify how long to run in total in hours, (default 0 - only run once)', default = 0, type = int) - - parser.add_argument('-q', - help = 'Run with verbose turned off', - action = 'store_true') - parser.add_argument('-v', + parser.add_argument('-v', '--verbose', help = 'Run with verbose', action = 'store_true') @@ -60,7 +56,20 @@ def parse_args_post(): parser.add_argument('-tgz', help = "tar and zip raw GITM file instead of process", action = 'store_true') + + parser.add_argument('-nc', action = 'store_true', + help = "Postprocess to netCDF files instead on '.bin'?") + parser.add_argument('--combine', action='store_true', + help="If processing to netCDF, we can combine each timestep to a" + " single file per output type. (ex: 3DALL.nc, etc.). " + "Will not work without -nc or if using remote.") + + parser.add_argument('--runname', type=str, default='', help= + "When combining, this is prepended to the output type in the " + "resulting files: '[runname]_3DALL.nc'. Default is no descriptor.") + + args = parser.parse_args() return args @@ -403,7 +412,9 @@ def transfer_model_output_files(user, server, dir): # Post process and then transfer files once: # ---------------------------------------------------------------------- -def do_loop(doTarZip, user, server, dir, IsRemote): +def do_loop(doTarZip, user, server, dir, IsRemote, + write_nc=False, combine=False, runname='', # For writing outputs to netCDF + ): DidWork = True @@ -437,9 +448,9 @@ def do_loop(doTarZip, user, server, dir, IsRemote): if (not os.path.exists(processDir)): # Maybe we are already in the data directory??? processDir = '.' - DidWork = post_process_gitm(processDir, DoRm, isVerbose = IsVerbose) - #DidWork = post_process_gitm() - + DidWork = post_process_gitm(processDir, DoRm, isVerbose = IsVerbose, + write_nc=write_nc, combine=combine, runname=runname) + # 4 - Check if remote data directory exists, make it if it doesn't: data_remote = '/data' if (IsRemote and DidWork): @@ -491,10 +502,24 @@ def load_remote_file(args): # make local variables for arguments: doTarZip = args.tgz - IsVerbose = args.v + IsVerbose = args.verbose if (args.norm): DoRm = False + # Sanity check netCDF-related arguments... + if args.nc or args.combine or (args.runname != ''): + if not args.nc: + raise ValueError(f"Combine and Runname can only be used with NetCDF files!") + if (args.runname != '') and not args.combine: + print(f"Using runname '{args.runname}' and combining") + args.combine = True + # For compatibility, make sure we can import PyITM if we're using netcdf files + try: + import pyitm + except ImportError: + raise ImportError("\n>> PyITM is required for NetCDF post-processing!\n" + " It is available at https://github.com/GITMCode/PyITM.git") + # Check if stop file exists: check_for_stop_file(delete = True) @@ -505,8 +530,13 @@ def load_remote_file(args): # load remote file every iteration so we can change it if needed: IsRemote, user, server, dir = load_remote_file(args) print('Move files to remote system? ', IsRemote) - - DidWork = do_loop(doTarZip, user, server, dir, IsRemote) + print('Process into netCDF files? ', args.nc) + if IsRemote and args.nc and args.combine: + raise ValueError( + "The remote transfer & netcdf combine options cannot be used together") + + DidWork = do_loop(doTarZip, user, server, dir, IsRemote, + args.nc, args.combine, args.runname) if (DidWork): # Check if stop file exists: stopCheck = check_for_stop_file() diff --git a/srcPython/example_gitm_cindi_plot_script.py b/srcPython/python2/example_gitm_cindi_plot_script.py similarity index 100% rename from srcPython/example_gitm_cindi_plot_script.py rename to srcPython/python2/example_gitm_cindi_plot_script.py diff --git a/srcPython/gitm.py b/srcPython/python2/gitm.py similarity index 100% rename from srcPython/gitm.py rename to srcPython/python2/gitm.py diff --git a/srcPython/gitm_3D_global_plots.py b/srcPython/python2/gitm_3D_global_plots.py similarity index 100% rename from srcPython/gitm_3D_global_plots.py rename to srcPython/python2/gitm_3D_global_plots.py diff --git a/srcPython/gitm_alt_plots.py b/srcPython/python2/gitm_alt_plots.py similarity index 100% rename from srcPython/gitm_alt_plots.py rename to srcPython/python2/gitm_alt_plots.py diff --git a/srcPython/gitm_comparison_plots.py b/srcPython/python2/gitm_comparison_plots.py similarity index 100% rename from srcPython/gitm_comparison_plots.py rename to srcPython/python2/gitm_comparison_plots.py diff --git a/srcPython/gitm_diff_images.py b/srcPython/python2/gitm_diff_images.py similarity index 100% rename from srcPython/gitm_diff_images.py rename to srcPython/python2/gitm_diff_images.py diff --git a/srcPython/gitm_loc_rout.py b/srcPython/python2/gitm_loc_rout.py similarity index 100% rename from srcPython/gitm_loc_rout.py rename to srcPython/python2/gitm_loc_rout.py diff --git a/srcPython/gitm_movie_script.py b/srcPython/python2/gitm_movie_script.py similarity index 100% rename from srcPython/gitm_movie_script.py rename to srcPython/python2/gitm_movie_script.py diff --git a/srcPython/gitm_plot_rout.py b/srcPython/python2/gitm_plot_rout.py similarity index 100% rename from srcPython/gitm_plot_rout.py rename to srcPython/python2/gitm_plot_rout.py diff --git a/srcPython/gitm_sat_plots.py b/srcPython/python2/gitm_sat_plots.py similarity index 100% rename from srcPython/gitm_sat_plots.py rename to srcPython/python2/gitm_sat_plots.py diff --git a/srcPython/gitm_time.py b/srcPython/python2/gitm_time.py similarity index 100% rename from srcPython/gitm_time.py rename to srcPython/python2/gitm_time.py diff --git a/srcPython/load_files.py b/srcPython/python2/load_files.py similarity index 100% rename from srcPython/load_files.py rename to srcPython/python2/load_files.py diff --git a/srcPython/plot_3D_global.py b/srcPython/python2/plot_3D_global.py similarity index 100% rename from srcPython/plot_3D_global.py rename to srcPython/python2/plot_3D_global.py diff --git a/srcPython/plot_alt_profiles.py b/srcPython/python2/plot_alt_profiles.py similarity index 100% rename from srcPython/plot_alt_profiles.py rename to srcPython/python2/plot_alt_profiles.py diff --git a/srcPython/plot_stats.py b/srcPython/python2/plot_stats.py similarity index 100% rename from srcPython/plot_stats.py rename to srcPython/python2/plot_stats.py diff --git a/srcPython/read_files.py b/srcPython/python2/read_files.py similarity index 100% rename from srcPython/read_files.py rename to srcPython/python2/read_files.py diff --git a/srcPython/solar_rout.py b/srcPython/python2/solar_rout.py similarity index 100% rename from srcPython/solar_rout.py rename to srcPython/python2/solar_rout.py diff --git a/srcPython/write_fftcoeff_wpgitm.py b/srcPython/python2/write_fftcoeff_wpgitm.py similarity index 100% rename from srcPython/write_fftcoeff_wpgitm.py rename to srcPython/python2/write_fftcoeff_wpgitm.py diff --git a/srcPython/write_files.py b/srcPython/python2/write_files.py similarity index 100% rename from srcPython/write_files.py rename to srcPython/python2/write_files.py