Implement xarray as containers for data from MDSPlus

mamba_environment.yml

@@ -14,3 +14,4 @@ dependencies:
   - configobj
   - pyyaml
   - sphinx
+  - xarray

src/wipplpy/__init__.py

@@ -8,3 +8,4 @@
 ]
 
 from wipplpy import modules
+from .mst.diagnostics.summary import PlasmaCurrent

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.

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"
 }

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()