ICT_IWT/run_tool.py at main · CW3E/ICT_IWT · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
import time
start_time = time.time()  # Record the start time

import sys
import os
import glob
import re
import xarray as xr
import numpy as np
import pandas as pd
import gc
import cartopy.crs as ccrs
import multiprocessing as mp
import gc
import warnings
warnings.filterwarnings("ignore", message="Engine.*loading failed")

sys.path.append('/data/projects/operations/ICT_IWT/modules')
from read_deterministic_data import calc_gfs_data
from plotter import plot_fields
from calc_funcs import format_timedelta_to_HHMMSS

model_name = sys.argv[1]
fdate = sys.argv[2]
F_lst = np.arange(0, 168+3, 3)

def plot_ICT_IWV(F, fdate):
    repo_path = "/data/projects/operations/ICT_IWT/"
    print('Reading {0} ...'.format(F))
    ds = calc_gfs_data(F, fdate)

    # Plot - use pconfig to change vars plotted
    outpath=f"{repo_path}figs/"
    pconfig = {
        'IVT': {
            'cfkey': 'ivt',
            'ckey': None,
            'ukey': 'ivtu',
            'vkey': 'ivtv',
            'title': "NCEP GFS IVT (kg m-1 s-1; shaded and vectors)"
        },
        'IWT': {
            'cfkey': 'iwt',
            'ckey': None,
            'ukey': 'ivtu',
            'vkey': 'ivtv',
            'title': "NCEP GFS IWT (kg m-1 s-1; shaded and vectors)"
        },
        'ICT': {
            'cfkey': 'ict',
            'ckey': None,
            'ukey': None,
            'vkey': None,
            'title': "NCEP GFS ICT (kg m-1 s-1; shaded)"
        },
        'RATIO': {
            'cfkey': 'ratio',
            'ckey': None,
            'ukey': None,
            'vkey': None,
            'title': "NCEP GFS ICT/IVT Ratio (%; shaded)"
        }
    }

    for varname in ['IVT', 'IWT', 'ICT', 'RATIO']:
        print(f'Plotting {varname}')

        config = {
                "model_name":"GFS",
                "output_fname" : outpath+'{model_name}_'+pconfig[varname]['cfkey']+'_{domain_name}_latest_F{lead_time}.png',
                "filled_var" : pconfig[varname]['cfkey'],
                "contour_var" : pconfig[varname]['ckey'],
                "vector_vars" : [pconfig[varname]['ukey'],pconfig[varname]['vkey']],
                "domain_dash_line" : False,
                "required_domains" :  ['IntWest'],
                "ulc_title":pconfig[varname]['title'],
                "llc_title":"Initialized: {init_time}",
                "lrc_title":"F-{lead_time:03d}  Valid: {valid_time}",
                "logo_file":f"/data/projects/operations/data/CW3E_Logo_Suite/5-Vertical-Acronym_Only/Digital/PNG/CW3E-Logo-Vertical-Acronym-FullColor.png",
                "font_file":f"{repo_path}/modules/helvetica.ttc"
            }


        plot_fields(ds, config)

    ds.close()
    del ds
    gc.collect()

def multiP_create_plots(F):
    plot_ICT_IWV(F, fdate)

if __name__ == '__main__':
    # sys.argv.append(None) ## add this in case date not specified in command line
    # fdate = sys.argv[1] ## set this to None to get most recently downloaded data

    start_time = pd.Timestamp.today()
    print(f'Creating ICT plots for {model_name}')

    ###################################
    ### CREATE PLOTS USING POOL.MAP ###
    ###################################

    print('...create plots ...')

    with mp.get_context('spawn').Pool(processes=16) as pool:
        # debug_print("Via map with exception")
        # debug_print("\tKicking off pool via map with exception")
        pool.map(multiP_create_plots,F_lst)
        # debug_print("\tKicked off pool via map with exception")
        pool.close()
        pool.join()


    end_time = pd.Timestamp.today()
    td = end_time - start_time
    td = format_timedelta_to_HHMMSS(td)
    print(f'Plots for {model_name} took {td} to run')