major rework of FitH5 and writer

xsocs/io/FitH5.py

@@ -90,13 +90,14 @@ class FitH5(XsocsH5Base):
     - all arrays are 1D.
     """
 
-    _QSPACE_AXIS_PATH = '{entry}/qspace_axis/{axis}'
+    _QSPACE_AXIS_PATH = '{entry}/qspace_axis'
     _STATUS_PATH = '{entry}/status/{axis}'
     _RESULT_GRP_PATH = '{entry}/{process}/results'
     _RESULT_PATH = '{entry}/{process}/results/{result}/{axis}'
     _SCAN_X_PATH = '{entry}/sample/x_pos'
     _SCAN_Y_PATH = '{entry}/sample/y_pos'
     _BACKGROUND_MODE_PATH = '{entry}/background_mode'
+    _AXIS_NAMES_PATH = '{entry}/axes'
 
     def entries(self):
         """Return the entry names.
@@ -136,13 +137,11 @@ class FitH5(XsocsH5Base):
         """Returns the fit status for the given entry/process/axis
 
         :param str entry:
-        :param axis: FitH5QAxis.qx_axis, FitH5QAxis.qy_axis
-         or FitH5QAxis.qz_axis
+        :param int axis:
         :rtype: numpy.ndarray
         """
-        axis_name = FitH5QAxis.axis_name(axis)
-        status_path = self._STATUS_PATH.format(entry=entry,
-                                               axis=axis_name)
+        axis_name = self.get_qspace_dimension_names(entry)[axis]
+        status_path = self._STATUS_PATH.format(entry=entry, axis=axis_name)
         status = self._get_array_data(status_path)
 
         # 30 Jan 2017.
@@ -170,41 +169,15 @@ class FitH5(XsocsH5Base):
         return (self._get_array_data(self._SCAN_X_PATH.format(entry=entry)),
                 self._get_array_data(self._SCAN_Y_PATH.format(entry=entry)))
 
-    def get_qx(self, entry):
-        """Returns the axis values for qx for the given entry.
-
-        :param str entry:
-        :rtype: numpy.ndarray
-        """
-        return self.__get_axis_values(entry, FitH5QAxis.qx_axis)
-
-    def get_qy(self, entry):
-        """Returns the axis values for qy for the given entry.
-
-        :param str entry:
-        :rtype: numpy.ndarray
-        """
-        return self.__get_axis_values(entry, FitH5QAxis.qy_axis)
-
-    def get_qz(self, entry):
-        """Returns the axis values for qz for the given entry.
-
-        :param str entry:
-        :rtype: numpy.ndarray
-        """
-        return self.__get_axis_values(entry, FitH5QAxis.qz_axis)
-
-    def __get_axis_values(self, entry, axis):
+    def get_qspace_dimension_values(self, entry):
         """Returns the axis values.
 
         :param str entry:
-        :param axis: FitH5QAxis.qx_axis, FitH5QAxis.qy_axis
-         or FitH5QAxis.qz_axis
-        :rtype: numpy.ndarray
+        :rtype: List[numpy.ndarray]
         """
-        axis_name = FitH5QAxis.axis_name(axis)
-        return self._get_array_data(self._QSPACE_AXIS_PATH.format(
-            entry=entry, axis=axis_name))
+        base_path = self._QSPACE_AXIS_PATH.format(entry=entry)
+        return [self._get_array_data('/'.join((base_path, name)))
+                for name in self.get_qspace_dimension_names(entry)]
 
     def get_axis_result(self, entry, process, result, axis):
         """Returns the results for the given entry/process/result name/axis.
@@ -212,12 +185,10 @@ class FitH5(XsocsH5Base):
         :param str entry:
         :param str process:
         :param str result:
-        :param axis: FitH5QAxis.qx_axis, FitH5QAxis.qy_axis or
-         FitH5QAxis.qz_axis
+        :param int axis:
         :rtype: numpy.ndarray
         """
-        assert axis in FitH5QAxis.axis_values
-        axis_name = FitH5QAxis.axis_name(axis)
+        axis_name = self.get_qspace_dimension_names(entry)[axis]
         result_path = self._RESULT_PATH.format(entry=entry,
                                                process=process,
                                                result=result,
@@ -234,6 +205,19 @@ class FitH5(XsocsH5Base):
             self._BACKGROUND_MODE_PATH.format(entry=entry))
         return mode if mode is not None else BackgroundTypes.NONE
 
+    def get_qspace_dimension_names(self, entry):
+        """Returns names of QSpace axes
+
+        :param str entry: HDF5 entry name for which to retrieve information
+        :rtype: List[str]
+        """
+        axes = self._get_array_data(
+            self._AXIS_NAMES_PATH.format(entry=entry))
+        if axes is None:  # backward compatibility
+            return QSpaceCoordinates.axes_names(QSpaceCoordinates.CARTESIAN)
+        else:
+            return axes
+
     def export_csv(self, entry, filename):
         """Exports an entry results as csv.
 
@@ -254,7 +238,7 @@ class FitH5(XsocsH5Base):
             header_list = ['X', 'Y']
             for process in processes:
                 result_names = self.get_result_names(entry, process)
-                for axis in FitH5QAxis.axis_names:
+                for axis in self.get_qspace_dimension_names(entry):
                     for result_name in result_names:
                         header_process.append(process)
                         header_list.append(result_name + '_' + axis)
@@ -271,16 +255,12 @@ class FitH5(XsocsH5Base):
             col_idx = 2
             for process in processes:
                 result_names = self.get_result_names(entry, process)
-                for axis in FitH5QAxis.axis_values:
+                for axis in range(self.get_qspace_dimension_names(entry)):
                     for result_name in result_names:
-                        result = self.get_axis_result(entry,
-                                                      process,
-                                                      result_name,
-                                                      axis)
-                        results[:, col_idx] = result
+                        results[:, col_idx] = self.get_axis_result(
+                            entry, process, result_name, axis)
                         col_idx += 1
-                    results[:, col_idx] = self.get_status(entry,
-                                                          axis)
+                    results[:, col_idx] = self.get_status(entry, axis)
                     col_idx += 1
 
             numpy.savetxt(filename,
@@ -295,26 +275,37 @@ class FitH5Writer(FitH5):
     """Class to write fit/COM results in a HDF5 file
 
     :param str h5_f: Filename where to write
-    :param QSpaceCoordinates coordinates: Kind of QSpace coordinates
+    :param str entry:
+    :param List[str] axis_names: Names of QSpace axes
+    :param BackgroundTypes background_mode: Used background subtraction
     :param str mode: File opening mode
     """
 
     def __init__(self, h5_f,
-                 coordinates=QSpaceCoordinates.CARTESIAN,
+                 entry,
+                 axis_names,
+                 background_mode=BackgroundTypes.NONE,
                  mode='w'):
-        assert coordinates in QSpaceCoordinates.ALLOWED
-        self.__coordinates = coordinates
+        assert background_mode in BackgroundTypes.ALLOWED
+
         super(FitH5Writer, self).__init__(h5_f, mode)
 
+        self.__entry = entry
+        self.__create_entry(self.__entry)
+        self._create_dataset(self._AXIS_NAMES_PATH.format(entry=entry),
+                             data=str_to_h5_utf8([name for name in axis_names]))
+        self._create_dataset(self._BACKGROUND_MODE_PATH.format(entry=entry),
+                             data=text_type(background_mode))
+
     def __axis_name(self, dimension):
-        """Returns the names of the QSpace axis
+        """Helper returning the names of the QSpace axis
 
         :param int dimension: Dimension index of the axis
         :rtype: str
         """
-        return QSpaceCoordinates.axes_names(self.__coordinates)[dimension]
+        return self.get_qspace_dimension_names(self.__entry)[dimension]
 
-    def create_entry(self, entry):
+    def __create_entry(self, entry):
         """Create group to store result for entry
 
         :param str entry:
@@ -329,23 +320,18 @@ class FitH5Writer(FitH5):
             entry_grp = h5_file.require_group(entry)
             entry_grp.attrs['NX_class'] = str_to_h5_utf8('NXentry')
 
-    def create_process(self, entry, process):
+    def create_process(self,process):
         """Create group to store a process in entry
 
-        :param str entry:
         :param str process:
         """
         # TODO : check that there isn't already an existing process
         with self._get_file() as h5_file:
-
-            processes = self.processes(entry)
-            if process in processes:
+            if process in self.processes(self.__entry):
                 raise ValueError('Process <{0}> already exists.'
                                  ''.format(process))
 
-            if entry not in h5_file:
-                self.create_entry(entry)
-            entry_grp = h5_file[entry]
+            entry_grp = h5_file[self.__entry]
 
             # TODO : check if it exists
             process_grp = entry_grp.require_group(process)
@@ -353,64 +339,48 @@ class FitH5Writer(FitH5):
             results_grp = process_grp.require_group('results')
             results_grp.attrs['NX_class'] = str_to_h5_utf8('NXcollection')
 
-    def set_sample_positions(self, entry, x, y):
+    def set_sample_positions(self, x, y):
         """Write sample positions (x, y) in file
 
-        :param str entry:
         :param numpy.ndarray x:
         :param numpy.ndarray y:
         """
-        self._set_array_data(self._SCAN_X_PATH.format(entry=entry), x)
-        self._set_array_data(self._SCAN_Y_PATH.format(entry=entry), y)
+        self._set_array_data(self._SCAN_X_PATH.format(entry=self.__entry), x)
+        self._set_array_data(self._SCAN_Y_PATH.format(entry=self.__entry), y)
 
-    def set_status(self, entry, dimension, data):
+    def set_status(self, dimension, data):
         """Write fit/COM status in the file
 
-        :param str entry:
         :param int dimension:
         :param numpy.ndarray data:
         """
         status_path = self._STATUS_PATH.format(
-            entry=entry, axis=self.__axis_name(dimension))
+            entry=self.__entry, axis=self.__axis_name(dimension))
         self._set_array_data(status_path, data)
 
-    def set_result(self, entry, process, dimension, name, data):
+    def set_result(self, process, dimension, name, data):
         """Write a fit/COM result parameter to the HDF5 file
 
-        :param str entry:
         :param str process:
         :param int dimension:
         :param str name:
         :param numpy.ndarray data:
         """
-        axis_name = self.__axis_name(dimension)
-        result_path = self._RESULT_PATH.format(entry=entry,
-                                               process=process,
-                                               result=name,
-                                               axis=axis_name)
+        result_path = self._RESULT_PATH.format(
+            entry=self.__entry,
+            process=process,
+            result=name,
+            axis=self.__axis_name(dimension))
         self._set_array_data(result_path, data)
 
-    def set_qspace_dimension_values(self, entry, dim0, dim1, dim2):
+    def set_qspace_dimension_values(self, dim0, dim1, dim2):
         """Write qspace axes coordinates for each dimension
 
-        :param str entry:
         :param numpy.ndarray dim0:
         :param numpy.ndarray dim1:
         :param numpy.ndarray dim2:
         """
+        base_path = self._QSPACE_AXIS_PATH.format(entry=self.__entry)
         for index, values in enumerate((dim0, dim1, dim2)):
             self._set_array_data(
-                self._QSPACE_AXIS_PATH.format(entry=entry,
-                                              axis=self.__axis_name(index)),
-                values)
-
-    def set_background_mode(self, entry, mode):
-        """Returns the background subtraction mode used
-
-        :param str entry: HDF5 entry name for which to retrieve information
-        :param str mode: Background mode description
-        """
-        assert mode in BackgroundTypes.ALLOWED
-        self._create_dataset(
-            self._BACKGROUND_MODE_PATH.format(entry=entry),
-            data=text_type(mode))
+                '/'.join((base_path, self.__axis_name(index))), values)

xsocs/process/fit/peak_fit.py

@@ -155,23 +155,22 @@ class FitResult(object):
         else:
             raise RuntimeError('Unknown Fit Type')
 
-        with FitH5Writer(fit_h5, mode=mode) as fitH5:
-            fitH5.create_entry(fit_name)
-            fitH5.create_process(fit_name, result_name)
-
-            fitH5.set_sample_positions(fit_name, self.sample_x, self.sample_y)
-            fitH5.set_qspace_dimension_values(
-                fit_name, *self.qspace_dimension_values)
-            fitH5.set_background_mode(fit_name, self.background_mode)
+        with FitH5Writer(fit_h5,
+                         entry=fit_name,
+                         axis_names=self.qspace_dimension_names,
+                         background_mode=self.background_mode,
+                         mode=mode) as fitH5:
+            fitH5.create_process(result_name)
+            fitH5.set_sample_positions(self.sample_x, self.sample_y)
+            fitH5.set_qspace_dimension_values(*self.qspace_dimension_values)
 
             for dimension, array in enumerate(self._fit_results):
                 for name in self.available_result_names:
                     results = self.get_results(dimension, name, copy=False)
                     if name == 'Status':
-                        fitH5.set_status(fit_name, dimension, results)
+                        fitH5.set_status(dimension, results)
                     else:
-                        fitH5.set_result(
-                            fit_name, result_name, dimension, name, results)
+                        fitH5.set_result(result_name, dimension, name, results)
 
 
 class PeakFitter(object):