The SWRegressor Class and Solar Wind Prediction

The primesw package is based around the SWRegressor class. This class subclasses the Pytorch Lightning LightningModule, and is flexible enough to load different model configurations that were trained to predict different plasma environments. For instance, you can load a SWRegressor wrapping the solar wind prediction model PRIME by specifying the keyword 'PRIME' to primesw.load():

import primesw as psw
prime = psw.load('PRIME')
primesw.prime.load(modelname=None, checkpoint=None)[source]

Loads a pretrained PRIME model, either one of the included models (PRIME, PRIME-SH, PRIME-PS, PRIME-PS-GEO, PRIME-PS-MMS) or a user supplied model checkpoint.

Parameters:

  • modelname ([str]) – Name of included pretrained model to load. Options include ‘PRIME’ (solar wind model), ‘PRIME-SH’ (magnetosheath model), or ‘PRIME-PS’ (plasmasheet model). For ‘PRIME-PS’, can specify Geotail-trained (‘PRIME-PS-GEO’) or MMS-trained (‘PRIME-PS-MMS’) models (if ‘PRIME-PS’ specified, MMS model loaded by default).

  • checkpoint ([str or Path]) – Path to user-supplied torch checkpoint .ckpt file. If specified, must also specify configuration path. Overrides supplied modelname.

Returns [SWRegressor] model:

Pretrained PRIME-like model.

The class method SWRegressor.predict_ts() is the way that most users will interface with the model. To generate predictions from Wind spacecraft data, specify start and stop times for the desired prediction. start and stop are strings with format 'YYYY-MM-DD HH:MM:SS'.

import primesw as psw
prime = psw.load('PRIME')
prime.predict_ts(start = '2020-01-01 00:00:00', stop = '2020-01-02 00:00:00')

If using data from an L1 monitor to make predictions, pass the input data using the in_data argument. If in_data is specified, start and stop should not be (and vice versa). in_data is also useful for making predicitons from synthetic solar wind data (see SWRegressor.build_synth_input()). For instance, one can predict what the solar wind at the bow shock nose would be if the solar wind flow at L1 was 700km/s:

import primesw as psw
prime = psw.load('PRIME')
prime.predict_ts(in_data = prime.build_synth_input(vx=-700))

By default for PRIME, predictions are made at the average location of the nose of Earth’s bow shock 13.25 Earth Radii upstream on the Geocentric Solar Ecliptic (GSE) x-axis. One can also specify a position to propagate to besides the default by specifying pos:

import primesw as psw
prime = psw.load('PRIME')
prime.predict_ts(start = '2020-01-01 00:00:00', stop = '2020-01-02 00:00:00', pos = [13.25, 5, 0])

All positions are in GSE coordinates with units of Earth Radii. It is not recommended to make predictions outside of the region any given model was trained on. For PRIME, that’s within 30 Earth radii of the Earth on the dayside.

class primesw.prime.SWRegressor(optimizer='adam', lr=0.001, lr_scheduler=None, patience=3, factor=0.5, weight_decay=0, total_iters=40, in_dim=14, tar_dim=1, pos_dim=3, in_norm=None, tar_norm=None, pos_norm=None, window=1, stride=1, interp_frac=1, decoder_type='linear', encoder_type='rnn', decoder_hidden_layers=[128], encoder_hidden_dim=128, encoder_num_layers=1, p_drop=0.1, pos_encoding_size=None, loss='mae', save_debug_ckpt=False, *args, **kwargs)[source]

This class wraps instances of the PRIME architecutre trained on different plasma regions throughout perigeospace. It is recommended to instantiate SWRegressor objects using primesw.load() and specifying the desired model (PRIME, PRIME-SH, PRIME-PS) rather than calling this class directly.

Parameters:

  • optimizer ([str]) – Optimization algorithm used to update model weights. Accepts any Pytorch optimizer alias.

  • lr ([float]) – Optimization algorithm learning rate.

  • lr_scheduler ([str]) – Optimization algorithm learning rate scheduler. Options are ‘cosine’, ‘cosine_warm’, ‘plateau’, ‘linear’, or ‘const’

  • patience ([int]) – Learning rate scheduler patience.

  • factor ([float]) – Learning rate scheduler factor.

  • weight_decay ([float]) – Optimization algorithm weight decay factor.

  • total_iters ([int]) – Total training epochs, used by learning rate scheduler.

  • in_dim ([int]) – Input timeseries dimensions (number of features).

  • tar_dim ([int]) – Target dimensions (number of features).

  • pos_dim ([stintr]) – Number of position dimensions (number of coordinates, generally 3).

  • in_norm ([dict]) – Dictionary of input features and their normalization factors.

  • tar_norm ([dict]) – Dictionary of target features and their normalization factors.

  • pos_norm ([dict]) – Dictionary of position features and their normalization factors.

  • window ([int]) – Timeseries window size.

  • stride ([int]) – Prediction lead time.

  • interp_frac ([float]) – Percent permissible interpolated data in each input timeseries.

  • decoder_type ([str]) – Decoder architecture. Options are ‘linear’ or ‘prob_linear’.

  • encoder_type ([str]) – Encoder architecture. Options are ‘linear’ or ‘rnn’.

  • decoder_hidden_layers ([list of int]) – Size of each of the hidden layers in decoder.

  • encoder_hidden_dim ([int]) – Dimension of hidden layers in encoder.

  • encoder_num_layers ([int]) – Number of layers in the encoder.

  • p_drop ([float]) – Dropout rate for model during training.

  • pos_encoding_size ([int]) – Order of random Fourier features applied to position data.

  • loss ([str]) – Loss function for model, options are ‘mae’ or ‘crps’. ‘crps’ only usable for ‘prob_linear’ decoders.

  • save_debug_ckpt ([bool]) – Whether to dump a debug packet on each validation epoch end.

build_real_input(start, stop)[source]

Load Wind spacecraft input data in between specified date strings.

Parameters:

  • start ([str]) – The start date of the data to load (‘YYYY-MM-DD’)

  • stop ([str]) – The end date of the data to load (‘YYYY-MM-DD’)

Returns [DataFrame] in_df:

Input dataframe suitable to predict from with self.predict().

build_synth_input(bx=0, by=0, bz=-5, vx=-400, vy=0, vz=0, ni=5, vt=30, rx=200, ry=0, rz=0, sme=None, smr=None, tilt=None)[source]

Builds a synthetic input array from user-specified quantities at L1. For input arrays made from measured data at L1, see SWRegressor.build_real_input.

Parameters:

  • x ([float, array-like]) – IMF Bx value (nT).

  • by ([float, array-like]) – IMF By value (nT).

  • bz ([float, array-like]) – IMF Bz value (nT).

  • vx ([float, array-like]) – Solar wind Vx value ().

  • vy ([float, array-like]) – Solar wind Vy value.

  • vz ([float, array-like]) – Solar wind Vz value.

  • ni ([float, array-like]) – Solar wind ion density value.

  • vt ([float, array-like]) – Solar wind ion thermal speed value.

  • rx ([float, array-like]) – Wind spacecraft position x value.

  • ry ([float, array-like]) – Wind spacecraft position y value.

  • rz ([float, array-like]) – Wind spacecraft position z value.

  • sme ([float, array-like]) – SuperMAG SME index (nT). Only used for plasmasheet model.

  • smr ([float, array-like]) – SuperMAG SMR index (nT). Only used for plasmasheet model.

  • tilt ([float, array-like]) – Earth dipole tilt angle (degrees). Only used for plasmasheet model.

Returns [Dataframe] in_df:

Input dataframe suitable to predict from with self.predict_ts().

predict_df(timeseries, position)[source]

Generate predictions from the model using DataFrames. In general it is recommended to use predict_ts() which is a more flexible wrapper of this method.

Parameters:

  • timeseries ([DataFrame]) – Input data from L1 monitor. Must contain the keys the model expects (check model.in_norm.keys()).

  • position ([DataFrame]) – Position(s) of desired prediction. Must contain the keys the model expects (check model.pos_norm.keys()).

Returns [DataFrame] tar_scaled:

Model output for the given timerange or input data. Includes means and standard deviations at each timestep.

predict_grid(ts, gridsize, x_extent, y_extent=None, z_extent=None, y=0, z=0, loc_mask=None, subtract_ecliptic=False)[source]

Generate predictions efficiently on a grid of points. Timeseries data can be generated with methods build_synth_input() or build_real_input(), depending on desired data source.

Parameters:

  • ts ([DataFrame]) – Timeseries of data at L1 (Can be synthetic)

  • gridsize ([float]) – Spacing of grid points

  • x_extent ([list]) – Range of x values to calculate on

  • y_extent ([list]) – Range of y values to calculate on. If None, z_extent must be specified.

  • z_extent ([list]) – Range of z values to calculate on. If None, y_extent must be specified.

  • y ([float, array-like]) – Y position that is held constant if y_extent is not specified. Default 0.

  • z ([float, array-like]) – Z position that is held constant if z_extent is not specified. Default 0.

  • loc_mask ([float, optional]) – RE from Earth to occlude (masking the magnetopause/magnetosheath)

Returns [ndarray] output_grid:

Array of predicted values on the grid. Shape (timestamps, x_extent/gridsize, y_extent/gridsize, z_extent/gridsize, features * 2)

Returns [Series] timestamps:

Series of datetimes corresponding to each grid’s time

predict_ts(start=None, stop=None, in_data=None, pos=array([13.25, 0., 0.]))[source]

Generate predictions from the model. Specify either a start and stop time or supply input data (in_data) from an L1 monitor. Predictions are made at a static position (13.25 RE upstream on GSE X axis), but a different or moving position can be specified with pos.

Parameters:

  • start ([str]) – Start time of desired prediction in format ‘YYYY-MM-DD HH:MM:SS’.

  • stop ([str]) – Stop time of desired prediction in format ‘YYYY-MM-DD HH:MM:SS’.

  • in_data ([DataFrame or ndarray]) – Input data from L1 monitor. If DataFrame, must contain the keys the model expects (check model.in_norm.keys()). If Numpy array, must be in the order expected by the model. Overrides start and stop if specified.

  • pos ([DataFrame or ndarray]) – Position(s) of desired prediction. For a static position, specify a 1D vector with three inputs corresponding to GSE X/Y/Z position. For a moving target, specify a 2D array or DataFrame with one entry per timestep in the timeseries. If DataFrame, must contain the keys the model expects (check model.pos_norm.keys()).

Returns:

[DataFrame] output: Model output for the given timerange or input data. Includes means and standard deviations at each timestep.