Gtm/time series at grid cell

.gitignore

@@ -69,6 +69,7 @@ benchmarks/__pycache__/*
 grids/utilities/jigsaw/__pycache__/*
 post_proc/py/geometry_lat_lon_2d_plot/__pycache__/*
 post_proc/py/__pycache__/
+post_proc/py/time_series_compare/__pycache__
 
 src/core_*/inc
 .DS_Store

post_proc/py/time_series_compare/mpas_time_series_at_point.py

@@ -0,0 +1,107 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+"""
+Script to plot time series for a list of variables at a 
+particular point (lon,lat,vertlevel) and time period.
+
+Usage (optional input arguments may be given, see below):
+python3 mpas_time_series_at_point.py --dir /path/to/data
+--stream history --lat -2.12 --lon -59.00 --vertlevel 0 --t0 '2021-01-01 00:00:00' --tf '2021-01-02 18:00:00' 
+--dt 10800 --var swdnb,swupb
+
+Guilherme Torres Mendonça ([email protected])
+March 2024
+"""
+import numpy as np
+import matplotlib.pyplot as plt
+import argparse
+from utils import ( 
+    create_datetime_range,
+    concat_mpas_output,
+    cs_string_to_list,
+    find_nCells_from_latlon
+)
+import os
+
+# Create ArgumentParser object
+parser = argparse.ArgumentParser(description=
+                                 'plotting time series for a list of variables at a' 
+                                +'particular point and time period.')
+# Input arguments
+parser.add_argument('--dir', default=os.getcwd(), type=str, 
+                    help='data directory')
+parser.add_argument('--stream', default='history', type=str, 
+                    help='stream')
+parser.add_argument('--var', default='swdnb, swupb', type=str, 
+                    help='list of variables to plot (var1,var2,...)')
+parser.add_argument('--lat', default='-2.124375557901979', type=str, 
+                    help='latitude')
+parser.add_argument('--lon', default='-59.003420487096946', type=str, 
+                    help='longitude')
+parser.add_argument('--vertlevel', default=0, type=str, 
+                    help='vertical level')
+parser.add_argument('--t0', type=str, 
+                    help='initial datetime (YYYY-mm-dd HH:MM:SS)')
+parser.add_argument('--tf', type=str, 
+                    help='final datetime (YYYY-mm-dd HH:MM:SS)')
+parser.add_argument('--dt', type=str, 
+                    help='datetime step (seconds)')
+parser.add_argument('--closest_value', type=str, default='haversine',
+                    help='method to find grid cell from (lat,lon)')
+
+# Parse the command line arguments
+args = parser.parse_args()
+
+# Create datetime range and transform it into a list of datetime strings
+datetime_list = create_datetime_range(
+    t0=args.t0,
+    tf=args.tf,
+    dt=args.dt
+)
+
+# Select variable and concatenate files between t0 and tf
+vars_list = cs_string_to_list(args.var)
+cat_file = concat_mpas_output(
+    stream=args.stream,
+    datetime_list=datetime_list,
+    data_dir=args.dir,
+    vars_list=vars_list,
+    )
+
+# Get lat, lon, vertlevel
+lat = float(args.lat)
+lon = float(args.lon)
+vertlevel = int(args.vertlevel)
+
+# Find grid cell
+nCells, ds = find_nCells_from_latlon(cat_file,lon=lon,lat=lat,
+                                     method=args.closest_value)
+
+# Plot
+## Get again list of variables for plotting
+vars_list = cs_string_to_list(args.var)
+## Get again list of datetimes, but abbreviated
+datetime_list = create_datetime_range(
+    t0=args.t0,
+    tf=args.tf,
+    dt=args.dt,
+    short_strings='y'
+)
+
+plt.figure()
+for var in vars_list:
+    # Try choosing nVertLevels
+    try:
+        y_series = ds[var].sel(nCells=nCells,nVertLevels=vertlevel).values
+    except: 
+        y_series = ds[var].sel(nCells=nCells).values
+    plt.plot(datetime_list,y_series,linestyle='-',label=var)
+plt.tick_params(axis='both', which='major', labelsize=12)
+plt.xticks(np.arange(min(ds.Time), max(ds.Time)+1,1),
+           rotation=45, ha='right')
+plt.grid(True)
+plt.legend(loc='center left', bbox_to_anchor=(1, 0.5))
+plt.title(f"{args.stream} from {args.t0} to {args.tf}")
+plt.tight_layout()
+plt.savefig("time_series.png")
+plt.show()

post_proc/py/time_series_compare/utils.py

@@ -0,0 +1,173 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+import xarray as xr
+import numpy as np
+import math
+import pandas as pd
+
+#TODO: utils.py should contain generally useful functions for postprocessing.
+#      Therefore, it should live under MPAS-BR/post_proc/py/.
+
+derived_variables = {
+    'latCell': ['latitude'],
+    'lonCell': ['longitude'],
+    'latVertex': ['latitudeVertex'],
+    'lonVertex': ['longitudeVertex'],
+    'areaCell': ['area', 'resolution'],
+}
+
+def add_mpas_mesh_variables(ds, full=True, **kwargs):
+    for v in ds.data_vars:
+        if v not in derived_variables:
+            continue
+
+        newvs = derived_variables[v]
+
+        for newv in newvs:
+            if newv in ds:
+                #print(newv + ' already here')
+                continue
+
+            if 'lat' in v or 'lon' in v:
+                ds[newv] = xr.apply_ufunc(np.rad2deg, ds[v])
+                ds[newv] = ds[newv].where(ds[newv] <= 180.0, ds[newv] - 360.0)
+                ds[newv].attrs['units'] = 'degrees'
+
+            elif newv == 'area':
+                radius_circle = ds.attrs.get('sphere_radius', 1.0)
+                if radius_circle == 1:
+                    # need to correct to earth radius
+                    correction_rad_earth = 6371220.0
+                else:
+                    correction_rad_earth = 1
+
+                ds[newv] = (ds[v] / 10 ** 6) * correction_rad_earth**2
+                ds[newv].attrs['units'] = 'km^2 (assuming areaCell in m^2)'
+                ds[newv].attrs['long_name'] = 'Area of the cell in km^2'
+
+            elif newv == 'resolution':
+                radius_circle = ds.attrs.get('sphere_radius', 1.0)
+                if radius_circle == 1.0:
+                    #print('need to correct to earth radius!!')
+                    correction_rad_earth = 6371220.0
+                else:
+                    correction_rad_earth = 1
+
+                # km^2 (assuming areaCell in m^2)
+                area = (ds[v] / 10 ** 6) * correction_rad_earth**2
+
+                ds[newv] = 2 * (xr.apply_ufunc(np.sqrt, area / math.pi))
+                ds[newv].attrs['units'] = 'km'
+                ds[newv].attrs['long_name'] = 'Resolution of the cell (approx)'
+
+    return ds
+
+def bash_array_to_list(bash_array):
+    return bash_array.split("|")[:-1]
+
+def cs_string_to_list(cs_string):
+    return cs_string.split(",")
+
+def create_datetime_range(t0,tf,dt,short_strings='n'):
+    '''
+    Creates a list of datetime strings between t0 and tf, with 
+    time step dt.
+
+    INPUT: t0 (str) - initial datetime in format 'YYYY.mm.dd HH:MM:SS'
+           tf (str) - final datetime in format 'YYYY.mm.dd HH:MM:SS'
+           dt (str) - time step in seconds
+
+    OUTPUT: datetime_str_list (list) - list of datetime strings
+    '''
+    t0 = pd.Timestamp(t0)
+    tf = pd.Timestamp(tf)
+    datetime_range = pd.date_range(start=t0,end=tf,freq=f"{dt}S")
+    if short_strings == 'y':
+        datetime_str_list = [dt.strftime('%d/%m %Hh') for dt in datetime_range]
+    else:
+        datetime_str_list = [dt.strftime('%Y-%m-%d_%H.%M.%S') for dt in datetime_range]
+    return datetime_str_list
+
+def concat_mpas_output(stream,datetime_list,data_dir,vars_list):
+    '''
+    Concatenates in time, for variables vars_list,  
+    MPAS output files named as follows:
+
+    stream.YYYY.mm.dd_HH.MM.SS
+
+    INPUT: stream (str) - name of the stream
+           datetime_list (list) - list of datetime strings in format YYYY.mm.dd_HH.MM.SS
+           data_dir (str) - directory where output files are stored
+           vars_list (list) - list of variable names (strings)
+
+    OUTPUT: cat_file (Xarray dataset) - concatenated dataset
+    '''
+    # Set variable list to keep in file
+    vars_list.extend(['latCell','lonCell','latVertex','lonVertex','areaCell'])
+    # Read first file
+    filename1 = stream + '.' + datetime_list[0] + '.nc'
+    cat_file = xr.open_dataset(f'{data_dir}/{filename1}', engine='netcdf4')[vars_list]
+    # Read and concatenate that file to remaining files
+    for datetime in datetime_list[1:]:
+        filename = stream + '.' + datetime + '.nc'
+        ds_temp = xr.open_dataset(f'{data_dir}/{filename}', engine='netcdf4')[vars_list]
+        cat_file = xr.concat([cat_file, ds_temp], dim='Time')
+    return cat_file
+
+def get_distance_haversine(lats, lons, lat_ref, lon_ref):
+    '''
+    Using the haversine formula (https://en.wikipedia.org/wiki/Haversine_formula), 
+    returns the distance in km of each lat, lon point to (lat_ref,lon_ref).
+
+    from Marta G. Badarjí, slightly modified (based on 
+    pypi haversine package: https://pypi.org/project/haversine/) 
+    by G. Torres Mendonça
+    '''
+
+    radius = 6371.0088  # km
+    d_lat = np.radians(lats - lat_ref)  # lat distance in radians
+    d_lon = np.radians(lons - lon_ref)  # lon distance in radians
+
+    a = (np.sin(d_lat / 2.) * np.sin(d_lat / 2.) +
+         np.cos(np.radians(lat_ref)) * np.cos(np.radians(lats)) *
+         np.sin(d_lon / 2.) * np.sin(d_lon / 2.))
+    c = np.arcsin(np.sqrt(a))
+    d = 2 * radius * c
+
+    return d
+
+def closest_value_haversine(ds,lon,lat):
+    d = get_distance_haversine(ds['latitude'].values, ds['longitude'].values,
+                                            lat_ref=lat, lon_ref=lon)
+    ds['distance'] = xr.DataArray(d, dims=['nCells'])
+    #print ('distance data array')
+    #print (ds[['distance','longitude','latitude']])
+    nCells_index = ds['distance'].argmin().item()
+    #print ('nCells_index')
+    #print (nCells_index)
+    #print ('distance min')
+    #print (ds['distance'].min())
+    #print ('ds[nCells]:')
+    #print (ds['nCells'])
+    nCells_value = ds['nCells'].isel(nCells=nCells_index)
+    distance_value = ds['distance'].isel(nCells=nCells_index)
+    return nCells_value, distance_value
+
+def find_nCells_from_latlon(ds,lon,lat,method='haversine',verbose='y'):
+    ds = add_mpas_mesh_variables(ds)
+    if method == 'haversine':
+        index_closest, distance_value  = closest_value_haversine(ds=ds.sel(Time=0),lon=lon,lat=lat)
+    else:
+        print (f"{method} method not supported.")
+        exit (-1)
+    if verbose == 'y':
+        # Print information on (lon,lat) point
+        closest_lon = ds['longitude'].sel(Time=0,nCells=index_closest)
+        closest_lat = ds['latitude'].sel(Time=0,nCells=index_closest)
+        print ("input (lon,lat):", (lon,lat))
+        print ("closest (lon,lat):", (float(closest_lon),float(closest_lat)))
+        print ("distance to input point (km):", distance_value.values)
+        print ("corresponding nCells value:",index_closest.values)
+        print ('max/min lat:', ds['latitude'].values.max(),ds['latitude'].values.min())
+        print ('max/min lon:', ds['longitude'].values.max(),ds['longitude'].values.min())
+    return index_closest, ds