diff --git a/scripts/realtime.py b/scripts/realtime.py new file mode 100644 index 0000000..db94d4e --- /dev/null +++ b/scripts/realtime.py @@ -0,0 +1,453 @@ +import argparse +#----------------------------- +# Command Line Options +#----------------------------- + +description = "Create real-time maps of the magnetic field across the US." +parser = argparse.ArgumentParser(description=description) + +parser.add_argument('--outdir', action='store', + required=False, default='.', + help='Where the output files will be written to. (default: "."') + +parser.add_argument('--timeseries', action='store_true', + default=False, required=False, + help='Make a plot of time series') + +parser.add_argument('--secmap', action='store_true', + default=False, required=False, + help='Make the most recent map of SEC predicted magnetic fields') + +parser.add_argument('--movie', action='store_true', + default=False, required=False, + help='Make a movie of the magnetic field changing over the past day') + +parser.add_argument('--server', action='store', + default='https://geomag.usgs.gov', + required=False, + help='Server url to query for data (default: %(default)s)') + +parser.add_argument('--verbose', action='store_true', + default=False, required=False, + help='Print out timing information') + +args = parser.parse_args() + +SERVER_URL = args.server + +# server requests +import requests +import datetime +import time + +# Numerical +import numpy as np +import pandas as pd +from secs import SECS +import bezpy + +# Plotting libraries +import matplotlib as mpl +import matplotlib.pyplot as plt + +#------------------ +# Defaults +#------------------ +t0 = time.time() + +R_earth = 6378000 +Bscale = 12 +arrow_width = 0.004 +cmap = mpl.cm.plasma +norm = mpl.colors.LogNorm(vmin=1, vmax=1000) + +# Mapping library +import cartopy +import cartopy.crs as ccrs +import cartopy.feature as cfeature + +bbox = [-135, -65, 20, 75] +projection = ccrs.LambertAzimuthalEqualArea(central_latitude=42.5, central_longitude=-95) +proj_data = ccrs.PlateCarree() + +def scale_vectors(x, y, scale=np.log10): + """Scale vectors while preserving the angle. + + Parameters + ========== + x: Cartesian x coordinate(s) + y: Cartesian y coordinate(s) + + scale: function to scale the magnitude by (Default: log10) + """ + mag = np.sqrt(x**2 + y**2) + # Prevent logs of values <1, which could be large negatives and flip the true sign + with np.errstate(invalid='ignore'): + mag[mag<1] = 1. + angle = np.arctan2(y, x) + newx = scale(mag)*np.cos(angle) + newy = scale(mag)*np.sin(angle) + + return (newx, newy) + +#------------------------------------------------ +# Routines for downloading geomag data from USGS +#------------------------------------------------ +USGS_CODES = ['BRW', 'DED', 'CMO', 'SIT', 'SHU', + 'NEW', 'BOU', 'FRD', 'FRN', 'TUC', + 'BSL', 'HON', 'SJG', 'GUA'] + +# Actual sites +NRCAN_CODES = ['ALE', 'BLC', 'BRD', 'CBB', 'EUA', + 'FCC', 'IQA', 'MEA', 'OTT', 'RES', + 'STJ', 'SNK', 'VIC', 'YKC'] + +# Removed ALE, EUA, SNK +NRCAN_CODES = ['BLC', 'BRD', 'CBB', + 'FCC', 'IQA', 'MEA', 'OTT', 'RES', + 'STJ', 'VIC', 'YKC'] + +class GeomagObservatory: + def __init__(self, name='BOU'): + """Object to download and store geomagnetic observatory data""" + + self.name = name + # Location code on USGS servers + self.loc_code = 'R0' + # Canadian data is in R1 (GOES) + if name in NRCAN_CODES: + self.loc_code = 'R1' + + self.dt = datetime.timedelta(days=1) + self._update_times() + self.df = None + self.update_data() + + def _update_times(self): + self.end_time = datetime.datetime.utcnow() + self.start_time = self.end_time - self.dt + + def update_data(self): + self._update_times() + + url = ("{server_url}/ws/edge/?format=json".format(server_url=SERVER_URL) + + "&id={name}".format(name=self.name) + + "&type={loc_code}".format(loc_code=self.loc_code) + + "&starttime={starttime}".format(starttime=self.start_time.strftime("%Y-%m-%dT%H:%M:%SZ")) + + "&endtime={endtime}".format(endtime=self.end_time.strftime("%Y-%m-%dT%H:%M:%SZ")) + + "&sampling_period=60" + + "&elements=X,Y,Z") # other codes: ,F,SV,SQ,Dist + + try: + r = requests.get(url) + data = r.json() + self.df = pd.DataFrame(index=pd.to_datetime(data['times']), + data={'X': data['values'][0]['values'], + 'Y': data['values'][1]['values'], + 'Z': data['values'][2]['values']}, + dtype=np.float).dropna() + + self.df['Direction'] = np.rad2deg(np.arctan2(self.df['X'], self.df['Y'])) + except: + print(f"Failed to update {self.name}") + +observatories = [GeomagObservatory(obs) for obs in USGS_CODES] + \ + [GeomagObservatory(obs) for obs in NRCAN_CODES] + +good_observatories = [obs for obs in observatories + if ((obs.df is not None) and (not obs.df.empty))] +missing_observatories = [obs for obs in observatories if obs not in good_observatories] + +good_xy = [] +missing_xy = [] + +for obs in observatories: + site = bezpy.mag.get_iaga_observatory(obs.name) + lonlat = (site['longitude'], site['latitude']) + if obs in good_observatories: + good_xy.append(lonlat) + else: + missing_xy.append(lonlat) + +good_xy = np.array(good_xy) +missing_xy = np.array(missing_xy) + +nobs = len(good_observatories) +B = pd.concat([obs.df[['X', 'Y', 'Z']] for obs in good_observatories], axis=1) +# Resample, then interpolate +B = B.resample('15Min').median().interpolate(limit_direction='both') +# Detrend +B = B - B.median() + +# Create numpy array +B_obs = np.zeros((len(B), nobs, 3)) +B_obs[:,:,0] = B['X'] +B_obs[:,:,1] = B['Y'] +B_obs[:,:,2] = B['Z'] + +# Save the most recent time string for file naming +time_string = B.index[-1].strftime("%Y%m%d-%H%M") + +if args.verbose: + print("Downloading data: {:.2f} s".format(time.time()-t0)) + +#----------------------- +# Set up regular grids +#----------------------- +delta_x = delta_y = 2 + +xs = np.arange(-180, -20, delta_x) +ys = np.arange(0, 89, delta_y) +yy, xx = np.meshgrid(ys, xs) +# Change it to an N x 2 matrix +gridded_xys = np.vstack([np.ravel(xx), np.ravel(yy)]).T + +# Plot grid +# NOTE: This specifies box edges for pcolormesh, which must be +# one longer than the 'data points' +grid_shape = xx.shape + +xs = np.append((xs - delta_x/2), xs[-1] + delta_x/2) +ys = np.append((ys - delta_y/2), ys[-1] + delta_y/2) +yy_edges, xx_edges = np.meshgrid(ys, xs) + +obs_lat_lon_r = np.zeros((len(good_xy), 3)) +obs_lat_lon_r[:,0] = good_xy[:,1] +obs_lat_lon_r[:,1] = good_xy[:,0] +obs_lat_lon_r[:,2] = R_earth + +obs_var = np.ones(obs_lat_lon_r.shape) +# Don't include Z component in fits for now +obs_var[:,2] = np.inf + +gridded_lat_lon_r = np.zeros((len(gridded_xys), 3)) +gridded_lat_lon_r[:,0] = gridded_xys[:,1] +gridded_lat_lon_r[:,1] = gridded_xys[:,0] +gridded_lat_lon_r[:,2] = R_earth + +# specify the SECS grid +lat, lon, r = np.meshgrid(np.linspace(15,85,36), + np.linspace(-175,-25,76), + R_earth+110000, indexing='ij') +secs_lat_lon_r = np.hstack((lat.reshape(-1,1), + lon.reshape(-1,1), + r.reshape(-1,1))) + +#-------------- +# SECS +#-------------- +secs = SECS(sec_df_loc=secs_lat_lon_r) +secs.fit(obs_loc=obs_lat_lon_r, obs_B=B_obs, + obs_var=obs_var, epsilon=0.05) + +B_gridded = secs.predict_B(gridded_lat_lon_r) + +if args.verbose: + print("SECS Interpolation: {:.2f} s".format(time.time()-t0)) + +#--------------- +# Plotting +#--------------- +def setup_map(): + scale = '10m' + coast = cfeature.NaturalEarthFeature(category='physical', scale=scale, + edgecolor='k', + facecolor='none', name='coastline') + states = cfeature.NaturalEarthFeature( + category='cultural', + name='admin_1_states_provinces_lines', + scale=scale, + facecolor='none', + edgecolor='k') + countries = cfeature.NaturalEarthFeature( + category='cultural', + name='admin_0_countries', + scale=scale, + facecolor='none', + edgecolor='k') + + fig, ax = plt.subplots(figsize=(8, 8), + subplot_kw=dict(projection=projection)) + + ax.set_extent(bbox, proj_data) + ax.add_feature(coast) + ax.add_feature(states, alpha=0.8) + ax.add_feature(countries) + + return (fig, ax) + + +if args.timeseries: + fig, (ax1, ax2) = plt.subplots(nrows=2, figsize=(8, 8)) + + for obs in good_observatories: + ax1.plot(obs.df.index, obs.df['X']-obs.df['X'].median(), label=obs.name) + ax2.plot(obs.df.index, obs.df['Y']-obs.df['Y'].median(), label=obs.name) + + ax1.legend(loc='upper center', bbox_to_anchor=(0.5, 1.15), + ncol=5, fancybox=True, shadow=True) + ax1.set_ylabel('B$_x$ (nT)') + ax2.set_ylabel('B$_y$ (nT)') + + ylim = ax1.get_ylim() + max_y = np.max(np.abs(ylim)) + ax1.set_ylim(-max_y, max_y) + + ylim = ax2.get_ylim() + max_y = np.max(np.abs(ylim)) + ax2.set_ylim(-max_y, max_y) + + fig.savefig(args.outdir + '/timeseries_{}.png'.format(time_string), bbox_inches='tight') + + if args.verbose: + print("Time series figure: {:.2f} s".format(time.time()-t0)) + +if args.secmap: + # Make a single map and save it + t = -1 + Bx = B_gridded[t,:,0] + By = B_gridded[t,:,1] + B_mag = np.sqrt(Bx**2 + By**2) + By_log, Bx_log = scale_vectors(By, Bx) + + fig, ax = setup_map() + + cax = ax.pcolormesh(xx_edges, yy_edges, B_mag.reshape(grid_shape), + cmap=cmap, norm=norm, + transform=proj_data) + + cbar = plt.colorbar(cax, orientation='horizontal', label='B (nT)') + + Q = ax.quiver(gridded_xys[:,0], gridded_xys[:,1], + By_log, Bx_log, color='k', + transform=proj_data, + scale_units='inches', scale=Bscale, width=arrow_width, + alpha=1., regrid_shape=15) + + Bx = B_obs[t,:,0] + By = B_obs[t,:,1] + B_mag = np.sqrt(Bx**2 + By**2) + + By_log, Bx_log = scale_vectors(By, Bx) + + #obs_scatter = ax.scatter(good_xy[:,0], good_xy[:,1], + # c='g', s=50, alpha=0.75, + # zorder=9, transform=proj_data) + # Plot the missing observatories in Red + missing_scatter = ax.scatter(missing_xy[:,0], missing_xy[:,1], + c='r', s=50, alpha=0.75, + zorder=9, transform=proj_data) + + Q_obs = ax.quiver(good_xy[:,0], good_xy[:,1], + By_log, Bx_log, color='g', + transform=proj_data, + scale_units='inches', scale=Bscale, width=arrow_width, + alpha=1., zorder=10) + + title = ax.set_title("SECS Magnetic Field\n{0}".format( + B.index[t].strftime("%Y/%m/%d %H:%M"))) + + fig.savefig(args.outdir + "/secs_map_{}.png".format(time_string), bbox_inches='tight') + + if args.verbose: + print("SECS map: {:.2f} s".format(time.time()-t0)) + +if args.movie: + # Animations + from matplotlib import animation + + # Make a single map and save it + fig, ax = setup_map() + + t = -1 + Bx = B_gridded[t,:,0] + By = B_gridded[t,:,1] + B_mag = np.sqrt(Bx**2 + By**2) + By_log, Bx_log = scale_vectors(By, Bx) + + cax = ax.pcolormesh(xx_edges, yy_edges, B_mag.reshape(grid_shape), + cmap=cmap, norm=norm, + transform=proj_data) + + cbar = plt.colorbar(cax, orientation='horizontal', label='B (nT)') + + Q = ax.quiver(gridded_xys[:,0], gridded_xys[:,1], + By_log, Bx_log, color='k', + transform=proj_data, + scale_units='inches', scale=Bscale, width=arrow_width, + alpha=1., regrid_shape=15) + + #---------------------------------------- + # Regrid in map space and + # calculate B-fields on this new grid. + #---------------------------------------- + # Get the uniform xy coordinates of the quiver grid and then delete that quiver + xy = Q.get_offsets() + Q.remove() + + uniform_xy = proj_data.transform_points(projection, xy[:,0], xy[:,1]) + uniform_xy[:,2] = R_earth + # Need to send lat/lon/r, so swap axis 0/1 + B_uniform = secs.predict_B(uniform_xy[:,[1,0,2]]) + + Bx = B_uniform[t,:,0] + By = B_uniform[t,:,1] + + By_log, Bx_log = scale_vectors(By, Bx) + + Q = ax.quiver(uniform_xy[:,0], uniform_xy[:,1], + By_log, Bx_log, color='k', + transform=proj_data, + scale_units='inches', scale=Bscale, width=arrow_width, + alpha=1.) + + Bx = B_obs[t,:,0] + By = B_obs[t,:,1] + + By_log, Bx_log = scale_vectors(By, Bx) + + #obs_scatter = ax.scatter(good_xy[:,0], good_xy[:,1], + # c='g', s=50, alpha=0.75, + # zorder=9, transform=proj_data) + # Plot the missing observatories in Red + missing_scatter = ax.scatter(missing_xy[:,0], missing_xy[:,1], + c='r', s=50, alpha=0.75, + zorder=9, transform=proj_data) + + Q_obs = ax.quiver(good_xy[:,0], good_xy[:,1], + By_log, Bx_log, color='g', + transform=proj_data, + scale_units='inches', scale=Bscale, width=arrow_width, + alpha=1., zorder=10) + + title = ax.set_title("SECS Magnetic Field\n{0}".format( + B.index[t].strftime("%Y/%m/%d %H:%M"))) + + def animate(t): + title.set_text("{0}".format(B.index[t].strftime("%Y/%m/%d %H:%M"))) + + # pcolormesh + Bx = B_gridded[t,:,0] + By = B_gridded[t,:,1] + B_mag = np.sqrt(Bx**2 + By**2) + + cax.set_array(B_mag) + + # Uniform grid vectors + Bx = B_uniform[t,:,0] + By = B_uniform[t,:,1] + By_log, Bx_log = scale_vectors(By, Bx) + Q.set_UVC(By_log, Bx_log) + + # Observations + Bx = B_obs[t,:,0] + By = B_obs[t,:,1] + + By_log, Bx_log = scale_vectors(By, Bx) + Q_obs.set_UVC(By_log, Bx_log) + + anim = animation.FuncAnimation(fig, animate, frames=range(len(B)), interval=150) + anim.save('daily_movie_{}.mp4'.format(time_string)) + + if args.verbose: + print("SECS animation: {:.2f} s".format(time.time()-t0)) diff --git a/setup.py b/setup.py index 44295c3..6845413 100644 --- a/setup.py +++ b/setup.py @@ -99,9 +99,6 @@ def run(self): 'License :: OSI Approved :: MIT License', 'Programming Language :: Python', 'Programming Language :: Python :: 3', - 'Programming Language :: Python :: 3.3', - 'Programming Language :: Python :: 3.4', - 'Programming Language :: Python :: 3.5', 'Programming Language :: Python :: 3.6', 'Programming Language :: Python :: Implementation :: CPython', 'Programming Language :: Python :: Implementation :: PyPy'