# Copyright 2018 Regents of the University of Colorado. All Rights Reserved.
# Released under the MIT license.
# This software was developed at the University of Colorado's Laboratory for Atmospheric and Space Physics.
# Verify current version before use at: https://github.com/MAVENSDC/PyTplot
import os
import pickle
import numpy as np
import pyspedas
from pyspedas.tplot_tools import options, store_data
from pyspedas.tplot_tools.tplot_options import tplot_options
from scipy.io import readsav
import logging
[docs]
def tplot_restore(filename):
"""
Restore tplot variables that have been saved to a file.
If the filename has a suffix ".tplot", it is assumed to be in the IDL .sav format as written by the IDL tplot_save routine.
In this case, it is read using the scipy.io.readsav() routine.
For any other filename suffix, the file is assumed to be a Python pickle file, as written by the PySPEDAS tplot_save routine.
Most of the metadata (units, coordinate systems, CDF attributes etc.) should be correctly represented in the
tplot variables created by tplot_restore.
Some IDL plot options may not have direct counterparts in
the PySPEDAS matplotlib-based plot options.
For multidimensional variables (spectra, particle distributions, etc.), the IDL v, v1, v2 (etc.) attributes should also be correctly
represented in the output tplot variables.
Parameters
-----------
filename : str
The path to the ".tplot" file created by IDL tplot_save, or for any other suffix, the pickle file created by PySPEDAS tplot_save.
Returns
--------
None
Examples
--------
>>> # Restore the saved data from the tplot_save example
>>> import pyspedas
>>> pyspedas.tplot_restore('C:/temp/variable1.pyspedas')
"""
#Error check
if not (os.path.isfile(filename)):
logging.error("tplot_restore: %s is not a valid file name",filename)
return
#Check if the restored file was an IDL file
if filename.endswith('.tplot'):
temp_tplot = readsav(filename)
for i in range(len(temp_tplot['dq'])):
if isinstance(temp_tplot['dq'][i][0], str):
logging.warning("tplot_restore: Error reading variable; this error occurs when the variable wasn't loaded in IDL when the SAV file was created.")
continue
data_name = temp_tplot['dq'][i][0].decode("utf-8")
try:
temp_x_data = temp_tplot['dq'][i][1][0][0].squeeze()
except AttributeError as err:
logging.warning("tplot_restore: Attribute error squeezing temp_x_data, index %d, name %s. (Value is scalar rather than array?)",i,data_name)
continue
#Pandas reads in data the other way I guess
if len(temp_tplot['dq'][i][1][0][2].shape) == 4:
temp_y_data = np.transpose(temp_tplot['dq'][i][1][0][2], axes=(3, 2, 1, 0))
elif len(temp_tplot['dq'][i][1][0][2].shape) == 3:
temp_y_data = np.transpose(temp_tplot['dq'][i][1][0][2], axes=(2, 1, 0))
elif len(temp_tplot['dq'][i][1][0][2].shape) == 2:
temp_y_data = np.transpose(temp_tplot['dq'][i][1][0][2])
else:
temp_y_data = temp_tplot['dq'][i][1][0][2]
# Check for loss of leading dimension if only 1 timestamp is present
shape = temp_x_data.shape
if len(shape) == 0:
temp_x_data = np.atleast_1d(temp_x_data)
# variable contains V1, V2 and V3 (e.g., DF as a function of energy, theta, phi)
if len(temp_tplot['dq'][i][1][0]) == 10:
temp_v1_data = temp_tplot['dq'][i][1][0][4]
temp_v2_data = temp_tplot['dq'][i][1][0][6]
temp_v3_data = temp_tplot['dq'][i][1][0][8]
# Data array gets transposed...so we need to swap v1 and v3, while v2 stays the same
temp_v1_data, temp_v3_data = temp_v3_data, temp_v1_data
# Change from little endian to big endian, since pandas apparently hates little endian
# We might want to move this into the store_data procedure eventually
if temp_x_data.dtype.byteorder == '>':
temp_x_data = temp_x_data.byteswap().view(temp_x_data.dtype.newbyteorder())
if temp_y_data.dtype.byteorder == '>':
temp_y_data = temp_y_data.byteswap().view(temp_y_data.dtype.newbyteorder())
if temp_v1_data.dtype.byteorder == '>':
temp_v1_data = temp_v1_data.byteswap().view(temp_v1_data.dtype.newbyteorder())
if temp_v2_data.dtype.byteorder == '>':
temp_v2_data = temp_v2_data.byteswap().view(temp_v2_data.dtype.newbyteorder())
if temp_v3_data.dtype.byteorder == '>':
temp_v3_data = temp_v3_data.byteswap().view(temp_v3_data.dtype.newbyteorder())
# support time-varying depends
if len(temp_v1_data.shape) == 2:
temp_v1_data = np.transpose(temp_v1_data)
if len(temp_v2_data.shape) == 2:
temp_v2_data = np.transpose(temp_v2_data)
if len(temp_v3_data.shape) == 2:
temp_v3_data = np.transpose(temp_v3_data)
store_data(data_name, data={'x': temp_x_data, 'y': temp_y_data, 'v1': temp_v1_data, 'v2': temp_v2_data, 'v3': temp_v3_data})
# variable contains V1, V2 (e.g., DF as a function of energy, angle)
elif len(temp_tplot['dq'][i][1][0]) == 8:
temp_v1_data = temp_tplot['dq'][i][1][0][4]
temp_v2_data = temp_tplot['dq'][i][1][0][6]
# Data array gets transposed, so we have to swap v1 and v2
temp_v1_data, temp_v2_data = temp_v2_data, temp_v1_data
#Change from little endian to big endian, since pandas apparently hates little endian
#We might want to move this into the store_data procedure eventually
if temp_x_data.dtype.byteorder == '>':
temp_x_data = temp_x_data.byteswap().view(temp_x_data.dtype.newbyteorder())
if temp_y_data.dtype.byteorder == '>':
temp_y_data = temp_y_data.byteswap().view(temp_y_data.dtype.newbyteorder())
if temp_v1_data.dtype.byteorder == '>':
temp_v1_data = temp_v1_data.byteswap().view(temp_v1_data.dtype.newbyteorder())
if temp_v2_data.dtype.byteorder == '>':
temp_v2_data = temp_v2_data.byteswap().view(temp_v2_data.dtype.newbyteorder())
# support time-varying depends
if len(temp_v1_data.shape) == 2:
temp_v1_data = np.transpose(temp_v1_data)
if len(temp_v2_data.shape) == 2:
temp_v2_data = np.transpose(temp_v2_data)
store_data(data_name, data={'x': temp_x_data, 'y': temp_y_data, 'v1': temp_v1_data, 'v2': temp_v2_data})
# If there are 4 fields, that means it is a spectrogram
# 6 fields is a spectrogram with a time varying Y axis
elif len(temp_tplot['dq'][i][1][0]) == 5 or len(temp_tplot['dq'][i][1][0]) == 6:
temp_v_data = temp_tplot['dq'][i][1][0][4]
# Change from little endian to big endian, since pandas apparently hates little endian
# We might want to move this into the store_data procedure eventually
if temp_x_data.dtype.byteorder == '>':
temp_x_data = temp_x_data.byteswap().view(temp_x_data.dtype.newbyteorder())
if temp_y_data.dtype.byteorder == '>':
temp_y_data = temp_y_data.byteswap().view(temp_y_data.dtype.newbyteorder())
if temp_v_data.dtype.byteorder == '>':
temp_v_data = temp_v_data.byteswap().view(temp_v_data.dtype.newbyteorder())
# support time-varying depends
if len(temp_v_data.shape) == 2:
temp_v_data = np.transpose(temp_v_data)
store_data(data_name, data={'x':temp_x_data, 'y':temp_y_data, 'v':temp_v_data})
else:
# Change from little endian to big endian, since pandas apparently hates little endian
# We might want to move this into the store_data procedure eventually
if temp_x_data.dtype.byteorder == '>':
temp_x_data = temp_x_data.byteswap().view(temp_x_data.dtype.newbyteorder())
if temp_y_data.dtype.byteorder == '>':
temp_y_data = temp_y_data.byteswap().view(temp_y_data.dtype.newbyteorder())
store_data(data_name, data={'x':temp_x_data, 'y':temp_y_data})
if temp_tplot['dq'][i][3].dtype.names is not None:
for option_name in temp_tplot['dq'][i][3].dtype.names:
if option_name.lower() == 'data_att':
arr = temp_tplot['dq'][i][3][option_name][0]
att_names = arr.dtype.names
# extract the values associated with the field names
att_values = arr.item()
# ensure the values are decoded to strings
att_values = [value.decode('utf-8') for value in att_values if isinstance(value, bytes)]
# create a dictionary with the desired mappings
data_att = {name.lower(): value for name, value in zip(att_names, att_values)}
pyspedas.tplot_tools.data_quants[data_name].attrs['data_att'] = data_att
if option_name.lower() not in ['color', 'colors']:
options(data_name, option_name, temp_tplot['dq'][i][3][option_name][0], quiet=True)
pyspedas.tplot_tools.data_quants[data_name].attrs['plot_options']['trange'] = temp_tplot['dq'][i][4].tolist()
pyspedas.tplot_tools.data_quants[data_name].attrs['plot_options']['create_time'] = temp_tplot['dq'][i][6]
if not np.isscalar(temp_tplot['tv']): # Skip if it is a scalar
for option_name in temp_tplot['tv'][0][0].dtype.names:
# the following should be set on the tplot variable, not for the entire session
#if option_name == 'TRANGE':
# # x_range of [0, 0] causes tplot to create an empty figure
# if temp_tplot['tv'][0][0][option_name][0][0] != 0 or temp_tplot['tv'][0][0][option_name][0][1] != 0:
# tplot_options('x_range', temp_tplot['tv'][0][0][option_name][0])
if option_name == 'WSIZE':
tplot_options('wsize', temp_tplot['tv'][0][0][option_name][0])
if option_name == 'VAR_LABEL':
tplot_options('var_label', temp_tplot['tv'][0][0][option_name][0])
if 'P' in temp_tplot['tv'][0][1].tolist():
for option_name in temp_tplot['tv'][0][1]['P'][0].dtype.names:
if option_name == 'TITLE':
tplot_options('title', temp_tplot['tv'][0][1]['P'][0][option_name][0])
# correct legend_names array
plt_options = pyspedas.tplot_tools.data_quants[data_name].attrs['plot_options']
yaxis_opts = plt_options.get('yaxis_opt')
if yaxis_opts is not None:
yaxis_opts = plt_options.get('yaxis_opt')
if yaxis_opts.get('legend_names') is not None:
lnames = pyspedas.tplot_tools.data_quants[data_name].attrs['plot_options']['yaxis_opt']['legend_names'][0]
if isinstance(lnames, list) or isinstance(lnames, np.ndarray):
pyspedas.tplot_tools.data_quants[data_name].attrs['plot_options']['yaxis_opt']['legend_names'] = [lname.decode('utf-8') for lname in lnames]
else:
pyspedas.tplot_tools.data_quants[data_name].attrs['plot_options']['yaxis_opt']['legend_names'] = [lnames.decode('utf-8')]
# decode any other string options
for y_key in yaxis_opts.keys():
if isinstance(yaxis_opts[y_key], bytes):
yaxis_opts[y_key] = yaxis_opts[y_key].decode("utf-8")
#temp_tplot['tv'][0][1] is all of the "settings" variables
#temp_tplot['tv'][0][1]['D'][0] is "device" options
#temp_tplot['tv'][0][1]['P'][0] is "plot" options
#temp_tplot['tv'][0][1]['X'][0] is x axis options
#temp_tplot['tv'][0][1]['Y'][0] is y axis options
####################################################################
else:
in_file = open(filename,"rb")
temp = pickle.load(in_file)
num_data_quants = temp[0]
for i in range(0, num_data_quants):
if isinstance(temp[i+1], dict):
# NRV variable
pyspedas.tplot_tools.data_quants[temp[i+1]['name']] = temp[i+1]
else:
pyspedas.tplot_tools.data_quants[temp[i+1].name] = temp[i+1]
pyspedas.tplot_tools.tplot_opt_glob = temp[num_data_quants+1]
in_file.close()
return