diff --git a/scripts/plotting/ensemble_track_error_ellipses.py b/scripts/plotting/ensemble_track_error_ellipses.py new file mode 100755 index 00000000..18475365 --- /dev/null +++ b/scripts/plotting/ensemble_track_error_ellipses.py @@ -0,0 +1,249 @@ +#!/usr/bin/env python3 + +"""Generates error ellipses from atcfunix files of an ensemble forecast. + +This is a temporary location for this script; the plot_ellipses() +function should be lifted out and put somewhere more sensible. + +Command-line arguments: +1. Basin (AL, EP, IO, etc.) +2. Numeric storm id (3, 10, 91, etc.) +3. Analysis time in YYYYMMDDHH format (20251028) +4. Model identifier (GFSO, HFSA, etc.) +5. List of atcfunix filenames + +Model identifiers with AVN in atcfunix files are renamed to GFS (ie. AVNO becomes GFSO). + +`./ellipses.py AL 13 20251028 HFSA H*/*2025102800*atcfunix` +""" + +import sys +import collections +import numpy +import datetime +import math + +import cartopy + +import matplotlib +from matplotlib import pyplot as plt +from cartopy import crs as ccrs + +class AtcfUnixLine: + """ + Utility class with track information. + """ + def __init__(self, line): + # Line looks like this: + # AL, 10, 2025100712, 03, HFSA, 000, 110N, 430W + # 0000000000111111111122222222223333333333444444 + # 0123456789012345678901234567890123456789012345 + + self.basin2 = line[0:2] + self.stormid = int(line[4:6].strip(), 10) + self.YYYYMMDDHH = line[8:18] + self.atime = datetime.datetime.strptime(self.YYYYMMDDHH, '%Y%m%d%H') + self.model = line[24:28].strip().replace('AVN', 'GFS') + self.lead_time_hour = int(line[30:33].strip(), 10) + + lat_int = int(line[35:38].strip(), 10) + latNS = line[38] + self.lat = lat_int / 10.0 + if latNS == 'S': + self.lat = -self.lat + + lon_int = int(line[41:45].strip(), 10) + lonEW = line[45] + self.lon = lon_int / 10.0 + if lonEW == 'W': + self.lon = -self.lon + + def set_ensemble_id(self, ensid): + self.ensemble_id = ensid + +def read_atcfunix(filename, filter_basin2, filter_stormid, filter_atime, filter_model): + """ + Reads one atcfunix file, scanning for lines that match the filters. + + @returns a list of AtcfUnixLine objects sorted by lead time. + """ + #print(f'Scan {filename} for basin {filter_basin2} stormid {filter_stormid} atime {filter_atime} model {filter_model}') + by_time = [] + with open(filename, 'rt') as fd: + for line in fd: + parsed = AtcfUnixLine(line) + + if filter_basin2 is not None and filter_basin2 != parsed.basin2: + #print(f'bad basin {parsed.basin2}: {line}') + continue + if filter_stormid is not None and filter_stormid != parsed.stormid: + #print(f'bad stormid {parsed.stormid}: {line}') + continue + if filter_atime is not None and filter_atime != parsed.atime: + #print(f'bad atime: {line}') + continue + if filter_model is not None and filter_model != parsed.model: + #print(f'bad model: {model}') + continue + + #print(f'accept line {line}') + by_time.append(parsed) + + by_time.sort(key=lambda x: x.lead_time_hour) + return by_time + +def recentered_longitudes(inlon): + """Ensures no break in longitudes when the ensemble crosses a dateline. + Longitudes are shifted to min(inlon) ... 360+min(inlon)""" + minlon = min(inlon) + outlon = [ minlon + (l - minlon) % 360 for l in inlon ] + return outlon + +def read_ensemble(filenames, filter_basin2, filter_stormid, filter_atime, filter_model): + """ + Reads a list of atcfunix files, one per ensemble member. + + @returns A list, sorted by ensemble id, of lists of AtcfUnixLine objects. The inner lists are the result of a call to read_atcfunix() on one file. + """ + by_id = [ read_atcfunix(filename, filter_basin2, filter_stormid, filter_atime, filter_model) + for filename in filenames] + ensemble_id = -1 + for lines in by_id: + ensemble_id += 1 + for line in lines: + line.set_ensemble_id(ensemble_id) + return by_id + +def fixes_by_lead_time(by_id): + """ + Obtains fixes for each lead time and ensemble member. + + @param by_id the return value of read_ensemble() + + @returns A three-dimensional array fixes[lead_time][ensemble_id][2]. Outer dimension is the lead time, sorted by increasing lead time. Middle dimension is the ensemble id, sorted numerically increasing. Inner dimension is the fix location as [ lon, lat ] + """ + + fixes = collections.defaultdict(dict) + for by_lead_time in by_id: + for line in by_lead_time: + fixes[line.lead_time_hour][line.ensemble_id] = [ line.lon, line.lat ] + return fixes + +def covariance_and_mean_by_lead_time(fixes): + """ + Calculates covariance matrixes and mean locations for each lead time. + + @param fixes from fixes_by_lead_time() + + @returns A list (one per lead time) of lists. Inner lists have five elements: covariance matrix, mean longitude, mean latitude, list of longitude fixes, and list of latitude fixes. + """ + cov_and_mean = {} + for lead_time_hour, fixes_by_id in fixes.items(): + lon = [] + lat = [] + for ensemble_id, fix in fixes_by_id.items(): + lon.append(fix[0]) + lat.append(fix[1]) + if not lon and not lat: + continue # no data at this time + + lon = recentered_longitudes(lon) + + meanlon = sum(lon) / len(lon) + meanlat = sum(lat) / len(lat) + if len(lon) < 2: + # Need at least 2 points for an error estimate + cov_and_mean.append([ None, meanlon, meanlat ]) + continue + normlon = [ x - meanlon for x in lon ] + normlat = [ y - meanlat for y in lat ] + cov = numpy.cov(numpy.asarray(normlon), numpy.asarray(normlat)) + cov_and_mean[lead_time_hour] = [ cov, meanlon, meanlat, lon, lat ] + return cov_and_mean + +def plot_ellipses(cov_and_mean, plot_hours, colors): + """ + Plots a track with error ellipses. + + @param cov_and_mean from covariance_and_mean_by_lead_time + @param plot_hours list of forecast lead times to plot + @param colors list of colors cycled through for each lead time + """ + i = -1 + center_x = [] + center_y = [] + fig = plt.figure(figsize=(9,6), dpi=168) + ax = fig.add_subplot(1,1,1, projection=ccrs.PlateCarree()) + legend_contents = [] + for lead_time_hour, cov_meanlon_meanlat_lon_lat in cov_and_mean.items(): + if lead_time_hour not in plot_hours: + #print(f'skip {lead_time_hour}hr lead time which was not requested') + continue + + i += 1 + + cov, meanlon, meanlat, lon, lat = cov_meanlon_meanlat_lon_lat + eig = numpy.linalg.eig(cov) + del cov + + center_x.append(meanlon) + center_y.append(meanlat) + + val0 = eig.eigenvalues[0] + val1 = eig.eigenvalues[1] + + vec0 = eig.eigenvectors[0] + vec1 = eig.eigenvectors[1] + + del eig + + if val0 > val1: + anglefrom='val0' + angle = math.atan2(vec0[1], vec0[0]) * 180 / math.pi + else: + anglefrom='val1' + angle = math.atan2(vec1[1], vec1[0]) * 180 / math.pi + + stdev0 = math.sqrt(val0) + stdev1 = math.sqrt(val1) + + del val0, val1, vec0, vec1 + + color = colors[i % len(colors)] + + stdfac = 2 # one standard deviation + + plt.scatter(lon, lat, 2, color=color, transform=ccrs.PlateCarree()) + ellipse = matplotlib.patches.Ellipse([meanlon, meanlat], width=stdev0*stdfac, height=stdev1*stdfac, angle=angle, edgecolor=color, fill=False, transform=ccrs.PlateCarree(), linewidth=1.5) + ax.add_patch(ellipse) + + legend_contents.append(matplotlib.patches.Patch(facecolor=color, alpha=0.7, edgecolor=color, label=f'{lead_time_hour} hrs')) + + + print(repr(legend_contents)) + gl = ax.gridlines(draw_labels=True, dms=True, x_inline=False, y_inline=False, alpha=0.5) + gl.top_labels=False + gl.right_labels = False + ax.coastlines() + plt.plot(center_x, center_y, 'kx-', transform=ccrs.PlateCarree()) + plt.legend(handles=legend_contents) + plt.show() + + +def main(): + """Main program. Arguments from sys.argv; see module description for details.""" + filter_basin2 = sys.argv[1] + filter_stormid = int(sys.argv[2], 10) + filter_atime = datetime.datetime.strptime(sys.argv[3], '%Y%m%d%H') + filter_model = sys.argv[4] + filenames = sys.argv[5:] + by_id = read_ensemble(filenames, filter_basin2, filter_stormid, filter_atime, filter_model) + fixes = fixes_by_lead_time(by_id) + cov_and_mean = covariance_and_mean_by_lead_time(fixes) + + colors = [ 'tab:blue', 'tab:orange', 'tab:green', 'tab:red', 'tab:purple', 'tab:cyan', 'tab:pink', 'tab:gray', 'tab:olive', 'tab:brown' ] + hours = [ 0, 24, 48, 72, 96, 120 ] + plot_ellipses(cov_and_mean, hours, colors) + +if __name__ == '__main__': + main()