8000 Final round of improvements for this project · dtcenter/mpas_plot@7741428 · GitHub
[go: up one dir, main page]
Skip to content

Commit 7741428

Browse files
mkavulichmkavulich
committed
Final round of improvements for this project
 - Add parallelism with starmap: script will create plots in parallel with the provided number of tasks via the -p argument - Change example file to be much smaller - Fix problem with memory leakage when making a large number of plots Ref: matplotlib/matplotlib#20300
1 parent e956092 commit 7741428

File tree

3 files changed

+91
-27
lines changed

3 files changed

+91
-27
lines changed

‎README.md

Lines changed: 6 additions & 3 deletions
Original file line numberDiff line numberDiff line change
@@ -32,13 +32,17 @@ The plotting script is built with argparse, so you can see a summary of the argu
3232

3333
```
3434
$ python plot_mpas_netcdf.py -h
35-
usage: plot_mpas_netcdf.py [-h] [-c CONFIG] [-d]
35+
usage: plot_mpas_netcdf.py [-h] [-c CONFIG] [-d] [-p PROCS]
36+
3637
Script for plotting MPAS input and/or output in native NetCDF format
38+
3739
options:
3840
-h, --help show this help message and exit
3941
-c CONFIG, --config CONFIG
4042
File used to specify plotting options
4143
-d, --debug Script will be run in debug mode with more verbose output
44+
-p PROCS, --procs PROCS
45+
Number of processors for generating plots in parallel
4246
```
4347

4448
The config file is where you will specify all the various options for what you want to plot, including which files, variables, levels, etc you want to plot. To setup the script to use your specific options, you’ll need to create a configuration file (`config_plot.yaml`). An example file `config_plot.yaml.example` is provided for reference, and you can view all available options in the `default_options.yaml` file.
@@ -65,7 +69,6 @@ This plotting utility is in a very early form, and has several known limitations
6569

6670
1. The user must know the name of the variable they want to plot, as well as the number of vertical levels if the variable has multiple.
6771
2. Only the [PlateCarree](https://scitools.org.uk/cartopy/docs/latest/reference/projections.html#platecarree) projection is currently supported for output maps
68-
3. The plotting script runs serially, which means it can take a long time to create a lot of large-domain plots.
69-
4. Certain variables that have additional dimensions such as grid property values (e.g. kiteAreasOnVertex) may not work out-of-the-box.
72+
3. Certain variables that have additional dimensions such as grid property values (e.g. kiteAreasOnVertex) may not work out-of-the-box.
7073

7174

‎config_plot.yaml.example

Lines changed: 2 additions & 2 deletions
Original file line numberDiff line numberDiff line change
@@ -1,6 +1,6 @@
11
data:
2-
filename: /scratch1/BMC/hmtb/kavulich/MPAS/plotting_scripts/debby/history.2024-08-08_12.00.00.nc
3-
# gridfile: /scratch1/BMC/hmtb/kavulich/MPAS/plotting_scripts/debby/history.2024-08-08_12.00.00.nc
2+
filename: /scratch1/BMC/hmtb/kavulich/MPAS/plotting_scripts/DTC_fork/mpas_app/ush/plotting/history.2023-09-15_09.00.00.nc
3+
# gridfile: /scratch1/BMC/hmtb/kavulich/MPAS/plotting_scripts/DTC_fork/mpas_app/ush/plotting/history.2023-09-15_09.00.00.nc
44
var:
55
- qv
66
lev:

‎plot_mpas_netcdf.py

Lines changed: 83 additions & 22 deletions
Original file line numberDiff line numberDiff line change
@@ -9,15 +9,22 @@
99
import os
1010
import sys
1111
import time
12+
import traceback
13+
from multiprocessing import Pool
1214

1315
print("Importing uxarray; this may take a while...")
1416
import uxarray as ux
17+
import matplotlib as mpl
18+
#This is needed to solve memory leak with large numbers of plots
19+
#https://github.com/matplotlib/matplotlib/issues/20300
20+
mpl.use('agg')
1521
import matplotlib.pyplot as plt
1622
import cartopy.feature as cfeature
1723
import cartopy.crs as ccrs
1824

1925
import uwtools.api.config as uwconfig
2026

27+
2128
def load_dataset(fn: str, gf: str = "") -> tuple[ux.UxDataset,ux.Grid]:
2229
"""
2330
Program loads the dataset from the specified MPAS NetCDF data file and grid file and returns
@@ -32,9 +39,54 @@ def load_dataset(fn: str, gf: str = "") -> tuple[ux.UxDataset,ux.Grid]:
3239
gf=fn
3340
return ux.open_dataset(gf,fn),ux.open_grid(gf)
3441

35-
def plotit(config_d: dict,uxds: ux.UxDataset,grid: ux.Grid,filepath: str) -> None:
42+
43+
def plotitparallel(config_d: dict,uxds: ux.UxDataset,grid: ux.Grid,filepath: str,variable: list=[],level: list=[]) -> None:
44+
"""
45+
The a wrapper for plotit() used for calling it recursively and in parallel with use of starmap
46+
Args:
47+
config_d (dict): A dictionary containing experiment settings
48+
uxds (ux.UxDataset): A ux.UxDataset object containing the data to be plotted
49+
grid (ux.Grid): A ux.Grid object containing the unstructured grid information
50+
filepath (str): The filename of the input data that was read into the ux objects
51+
variable (list): A one-item list, the variable name to plot. Used to overwrite the
52+
value in data:var (used for recursion over all variables)
53+
level (list): A one-item list, the vertical level to plot. If provided, overwrites
54+
the value in data:lev (used for recursion over all levels)
55+
56+
Returns:
57+
None
58+
59+
"""
60+
61+
newconf = copy.deepcopy(config_d)
62+
if variable:
63+
newconf["data"]["var"]=variable
64+
if level:
65+
newconf["data"]["lev"]=level
66+
67+
logging.debug(f'Trying to plot level {newconf["data"]["lev"]} for variable {newconf["data"]["var"]}')
68+
try:
69+
plotit(newconf,uxds,grid,filepath,1)
70+
except Exception as e:
71+
logging.warning(f'Could not plot variable {newconf["data"]["var"]}, level {newconf["data"]["lev"]}')
72+
logging.warning(f"{traceback.print_tb(e.__traceback__)}:")
73+
logging.warning(f"{type(e).__name__}:")
74+
logging.warning(e)
75+
76+
77+
78+
def plotit(config_d: dict,uxds: ux.UxDataset,grid: ux.Grid,filepath: str,parproc: int) -> None:
3679
"""
3780
The main program that makes the plot(s)
81+
Args:
82+
config_d (dict): A dictionary containing experiment settings
83+
uxds (ux.UxDataset): A ux.UxDataset object containing the data to be plotted
84+
grid (ux.Grid): A ux.Grid object containing the unstructured grid information
85+
filepath (str): The filename of the input data that was read into the ux objects
86+
parproc (int): The number of processors available for generating plots in parallel
87+
88+
Returns:
89+
None
3890
"""
3991

4092
filename=os.path.basename(filepath)
@@ -44,31 +96,37 @@ def plotit(config_d: dict,uxds: ux.UxDataset,grid: ux.Grid,filepath: str) -> Non
4496
logging.debug(f"Available data variables:\n{list(uxds.data_vars.keys())}")
4597
# To plot all variables, call plotit() recursively, trapping errors
4698
if config_d["data"]["var"]=="all":
47-
newconf = copy.deepcopy(config_d)
99+
args = []
48100
for var in uxds:
49-
logging.debug(f"Trying to plot variable {var}")
50-
newconf["data"]["var"]=[var]
51-
try:
52-
plotit(newconf,uxds,grid,filepath)
53-
except Exception as e:
54-
logging.warning(f"Could not plot variable {var}")
55-
logging.warning(f"{type(e).__name__}:")
56-
logging.warning(e)
101+
# Create argument tuples for each call to plotit() for use with starmap
102+
args.append( (config_d,uxds,grid,filepath,[var]) )
103+
if parproc > 1:
104+
with Pool(processes=parproc) as pool:
105+
pool.starmap(plotitparallel, args)
106+
else:
107+
i=0
108+
for instance in args:
109+
# '*' operator unpacks tuple for use as args
110+
plotitparallel(*args[i])
111+
i+=1
57112
# To plot all levels, call plotit() recursively, trapping errors
58113
elif config_d["data"]["lev"]=="all":
59-
newconf = copy.deepcopy(config_d)
60-
if "nVertLevels" in uxds[newconf["data"]["var"]].dims:
61-
levs = range(0,len(uxds[newconf["data"]["var"]]["nVertLevels"]))
114+
args = []
115+
if "nVertLevels" in uxds[config_d["data"]["var"]].dims:
116+
levs = range(0,len(uxds[config_d["data"]["var"]]["nVertLevels"]))
62117
else:
63118
levs = [0]
64119
for lev in levs:
65-
logging.debug(f"Trying to plot level {lev} for variable {newconf['data']['var']}")
66-
newconf["data"]["lev"]=[lev]
67-
try:
68-
plotit(newconf,uxds,grid,filepath)
69-
except Exception as e:
70-
logging.warning(f"Could not plot variable {newconf['data']['var']}, level {lev}")
71-
logging.warning(e)
120+
# Create argument tuples for each call to plotit() for use with starmap
121+
args.append( (config_d,uxds,grid,filepath,config_d["data"]["var"],[lev]) )
122+
if parproc > 1:
123+
with Pool(processes=parproc) as pool:
124+
pool.starmap(plotitparallel, args)
125+
else:
126+
i=0
127+
for instance in args:
128+
plotitparallel(*args[i])
129+
i+=1
72130

73131
elif isinstance(config_d["data"]["var"], list):
74132
for var in config_d["data"]["var"]:
@@ -192,6 +250,7 @@ def plotit(config_d: dict,uxds: ux.UxDataset,grid: ux.Grid,filepath: str) -> Non
192250
"time": "no_Time_dimension"
193251
})
194252

253+
195254
# Check if the file already exists, if so act according to plot:exists setting
196255
outfnme=outfnme.format_map(patterns)
197256
outfile=f"{outfnme.format_map(patterns)}.{fmt}"
@@ -229,7 +288,7 @@ def plotit(config_d: dict,uxds: ux.UxDataset,grid: ux.Grid,filepath: str) -> Non
229288
os.makedirs(os.path.dirname(outfile),exist_ok=True)
230289
logging.debug(f"Saving plot {outfile}")
231290
plt.savefig(outfile,format=fmt)
232-
plt.close()
291+
plt.close(fig)
233292
logging.debug(f"Done. Plot generation {time.time()-plotstart} seconds")
234293

235294

@@ -324,6 +383,8 @@ def setup_config(config: str, default: str="default_options.yaml") -> dict:
324383
help='File used to specify plotting options')
325384
parser.add_argument('-d', '--debug', action='store_true',
326385
help='Script will be run in debug mode with more verbose output')
386+
parser.add_argument('-p', '--procs', type=int, default=1,
387+
help='Number of processors for generating plots in parallel')
327388

328389
args = parser.parse_args()
329390

@@ -355,4 +416,4 @@ def setup_config(config: str, default: str="default_options.yaml") -> dict:
355416
logging.debug(f"{dataset=}")
356417
logging.debug(f"{grid=}")
357418
# Make the plots!
358-
plotit(expt_config,dataset,grid,f)
419+
plotit(expt_config,dataset,grid,f,args.procs)

0 commit comments

Comments
 (0)
0