diff --git a/mamba_environment.yml b/mamba_environment.yml index 62df5eb..b596316 100644 --- a/mamba_environment.yml +++ b/mamba_environment.yml @@ -14,3 +14,4 @@ dependencies: - configobj - pyyaml - sphinx + - xarray \ No newline at end of file diff --git a/mamba_environment_linux.yml b/mamba_environment_linux.yml deleted file mode 100644 index 90ff2a2..0000000 --- a/mamba_environment_linux.yml +++ /dev/null @@ -1,15 +0,0 @@ -name: WiPPLPy -channels: - - defaults - - anaconda - - conda-forge -dependencies: - - python=3.9 - - jupyter - - numpy - - scipy - - matplotlib - - tqdm - - configobj - - pyyaml - - sphinx diff --git a/src/wipplpy/__init__.py b/src/wipplpy/__init__.py index 37579ee..66a65b3 100644 --- a/src/wipplpy/__init__.py +++ b/src/wipplpy/__init__.py @@ -8,3 +8,4 @@ ] from wipplpy import modules +from .mst.diagnostics.summary import PlasmaCurrent \ No newline at end of file diff --git a/src/wipplpy/modules/generic_get_data.py b/src/wipplpy/modules/generic_get_data.py index 1d87876..f8a8783 100644 --- a/src/wipplpy/modules/generic_get_data.py +++ b/src/wipplpy/modules/generic_get_data.py @@ -2,6 +2,7 @@ import re import numpy as np +import xarray as xr from MDSplus.connection import Connection, MdsIpException from MDSplus.mdsExceptions import MDSplusException, SsSUCCESS from scipy.io import loadmat, savemat @@ -67,11 +68,13 @@ def __init__(self, call_string, name=None, signal=True): self.name = name self.signal = signal + def __str__(self) -> str: if self.signal: return f"Get({self.call_string}, is_signal)" else: return f"Get({self.call_string})" + def full_str(self, index_range=None, sample_period=1): """ @@ -164,7 +167,7 @@ def __init__( # noqa: PLR0913 ): """ Generic class for dealing with data from a remote MDSplus database. - + Parameters ---------- tree : Connection or int @@ -205,13 +208,19 @@ def __init__( # noqa: PLR0913 try: self.shot_number = int(tree) except ValueError: - self.shot_number = tree.shot_number + self.shot_number = tree.shot + except TypeError: + self.shot_number = tree.shot self.ignore_errors = ignore_errors self.silence_error_logging = silence_error_logging self._tree = None + #Make attributes for dimension_names and dimension_units if specified + self.units=[] + self.names=[] + # Initialize the time index range as empty and then try to get something for it. We do this because some code in _to_time_index_range requires it. self.time_index_range = None self.sample_period = sample_period @@ -228,6 +237,8 @@ def __init__( # noqa: PLR0913 elif len(get_calls) == 0: logging.info("Did not do any get calls for preloading data into object.") + self.get_calls = get_calls + # Hold all the calls and call data gotten from MDSplus. # TODO: Change how calls are saved so that we can change loaded attributes and save the update. self.saved_calls = {} @@ -310,6 +321,7 @@ def _to_time_index(self, time): index : int Data index from start of digitizer recording. """ + try: times = self.time except AttributeError: @@ -420,6 +432,8 @@ def get( # noqa: PLR0912, PLR0915 Whether to change the data type of what MDSplus returns. load_from_saved : bool, default=True Whether to try to load the data from the saved calls. + return_node : bool, default = False + Whether to return the data node instead of returning the data itself. Used for gathering xarray. Returns ------- @@ -488,6 +502,7 @@ def get( # noqa: PLR0912, PLR0915 try_loading = True self.tree = get_remote_shot_tree(self.shot_number, reconnect=True) + # Might be good to put the rest of these checks into a separate function, so it can be called here and in the gather_data_array if self.ignore_errors: try: data = node.data() @@ -512,6 +527,68 @@ def get( # noqa: PLR0912, PLR0915 return data + def get_xarray(self, np_data_type=np.float64, change_data=True, load_from_saved=True): + """ + Get data and dimensions from MDSplus tree, then store in an xarray along with metadata and proper dimensions as defined in the class + + Parameters + ---------- + get_call : str + A string that defines the call to use on the tree. + np_data_type : data-type, default=np.float64 + The numpy data type to change the data to. + change_data : bool, default=True + Whether to change the data type of what MDSplus returns. + load_from_saved : bool, default=True + Whether to try to load the data from the saved calls. + + Returns + ------- + data : xarray DataArray + The data from the tree packaged into an xarray + """ + + units = self.units + names = self.names + get_call = self.get_calls[0] + print(get_call) + + #Needs to get expanded for multiple get_calls passed!!!!!!!!!!!!!&&&&&&&&&&&&&&&&&&&)))))))()()() + + data = self.get(get_call, np_data_type=np_data_type, change_data=change_data, load_from_saved=load_from_saved) + + num_dimensions = data.ndim + if isinstance(units, list) and isinstance(names, list): + if num_dimensions + 1 != len(units): + raise Exception( + "*** Number of dimension names and number of dimension units does not match" + ) + + if num_dimensions + 1 != len(names): + raise Exception( + "*** Number of dimension names and number of data dimensions does not match" + ) + + coords_dict = {} + units_dict = {} + + for dim_num, dim_name in enumerate(names[1:]): + coords_dict[dim_name] = self.get('DIM_OF('+ get_call + ', {:d})'.format(dim_num), np_data_type=np_data_type, change_data=change_data, load_from_saved=load_from_saved) + units_dict[dim_name] = units[dim_num+1] + + units_dict[names[0]] = units[0] + + coords_dict["shot_num"] = int(self.get('$SHOT')) + + # get the rest of the attributes together into + attributes = {"units": units_dict, "call_string": get_call, "version": 1.0, 'tree': self.get('$EXPT', change_data = False)} + + data_xarray = xr.DataArray( + data, dims=names[1:], coords=coords_dict, attrs=attributes + ) + + self.xarray = data_xarray + def to_raw_index(self, time_index): """ Convert an index that works for the entire time range and change it to work for this objects arrays. diff --git a/src/wipplpy/modules/shot_loading_config.json b/src/wipplpy/modules/shot_loading_config.json index 0afc389..50a53f0 100644 --- a/src/wipplpy/modules/shot_loading_config.json +++ b/src/wipplpy/modules/shot_loading_config.json @@ -1,4 +1,4 @@ { - "tree_name": "wipal", - "server_name": "skywalker.physics.wisc.edu" + "tree_name": "mst", + "server_name": "toni.physics.wisc.edu" } diff --git a/src/wipplpy/mst/diagnostics/summary.py b/src/wipplpy/mst/diagnostics/summary.py new file mode 100644 index 0000000..8fd9aa7 --- /dev/null +++ b/src/wipplpy/mst/diagnostics/summary.py @@ -0,0 +1,24 @@ +import matplotlib.pyplot as plt +from wipplpy.modules.generic_get_data import Data as Data +import numpy as np +import xarray as xr +import MDSplus as mds + +class PlasmaCurrent(Data): + def __init__(self, tree): + #Ordering does matter. Define calls, super init, then dimension specifications. Figuring out how to make it not matter, or streamline whats needed would be cool + calls = ['\MST_OPS::IP_F'] #want to add ability to have multiple in this list + super().__init__(tree, [True], calls) # could augment second argument as a keyword with Default None in which case it does all of calls + self.units = ['kA', 's'] + self.names = ['plasma_current', 'shot_time'] + + def plot(self): + if not hasattr(self, "xarray"): + self.get_xarray() + print("Put together xarray") + + fig, ax = plt.subplots(figsize=(8,4)) + plt.plot(self.xarray.shot_time * 1000, self.xarray.data) + ax.set_xlabel('Time (ms)') + ax.set_ylabel('Plasma Current (kA)') + plt.show() \ No newline at end of file