Extended get unit values to positions as well

pycroscopy/io/hdf_utils.py

@@ -1598,18 +1598,20 @@ def create_spec_inds_from_vals(ds_spec_val_mat):
     return ds_spec_inds_mat
 
 
-def get_unit_values(h5_spec_ind, h5_spec_val, dim_names=None):
+def get_unit_values(h5_inds, h5_vals, is_spec=True, dim_names=None):
     """
     Gets the unit arrays of values that describe the spectroscopic dimensions
 
     Parameters
     ----------
-    h5_spec_ind : h5py.Dataset
-        Spectroscopic Indices dataset
-    h5_spec_val : h5py.Dataset
-        Spectroscopic Values dataset
+    h5_inds : h5py.Dataset
+        Spectroscopic or Position Indices dataset
+    h5_vals : h5py.Dataset
+        Spectroscopic or Position Values dataset
+    is_spec : bool, recommended
+        Are the provided datasets spectral. Default = True
     dim_names : str, or list of str, Optional
-        Names of the dimensions of interest
+        Names of the dimensions of interest. Default = all
 
     Note - this function can be extended / modified for ancillary position dimensions as well
 
@@ -1619,33 +1621,41 @@ def get_unit_values(h5_spec_ind, h5_spec_val, dim_names=None):
         Dictionary containing the unit array for each dimension. The name of the dimensions are the keys.
 
     """
+    # First load to memory
+    inds_mat = h5_inds[()]
+    vals_mat = h5_vals[()]
+    if not is_spec:
+        # Convert to spectral shape
+        inds_mat = np.transpose(inds_mat)
+        vals_mat = np.transpose(vals_mat)
+
     # For all dimensions, find where the index = 0
     # basically, we are indexing all dimensions to 0
     first_indices = []
-    for dim_ind in range(h5_spec_ind.shape[0]):
-        first_indices.append(h5_spec_ind[dim_ind] == 0)
+    for dim_ind in range(inds_mat.shape[0]):
+        first_indices.append(inds_mat[dim_ind] == 0)
     first_indices = np.vstack(first_indices)
 
-    spec_dim_names = get_attr(h5_spec_ind, 'labels')
+    full_dim_names = get_attr(h5_inds, 'labels')
     if dim_names is None:
-        dim_names = spec_dim_names
+        dim_names = full_dim_names
     elif not isinstance(dim_names, list):
         dim_names = [dim_names]
 
     unit_values = dict()
     for dim_name in dim_names:
         # Find the row in the spectroscopic indices that corresponds to the dimensions we want to slice:
-        desired_row_ind = np.where(spec_dim_names == dim_name)[0][0]
+        desired_row_ind = np.where(full_dim_names == dim_name)[0][0]
 
         # Find indices of all other dimensions
-        remaining_dims = list(range(h5_spec_ind.shape[0]))
+        remaining_dims = list(range(inds_mat.shape[0]))
         remaining_dims.remove(desired_row_ind)
 
         # The intersection of all these indices should give the desired index for the desired row
         intersections = np.all(first_indices[remaining_dims, :], axis=0)
 
         # apply this slicing to the values dataset:
-        unit_values[dim_name] = h5_spec_val[desired_row_ind, intersections]
+        unit_values[dim_name] = vals_mat[desired_row_ind, intersections]
 
     return unit_values