Add read_vars_in_frame and read_vars_in_frame_as_3d to slf.Serafin.Read

cli/slf_bottom_zones.py

@@ -138,9 +138,7 @@ def bottom(args):
             pos_B = output_header.var_IDs.index('B')
 
             for time_index, time in enumerate(resin.time):
-                var = np.empty((output_header.nb_var, output_header.nb_nodes), dtype=output_header.np_float_type)
-                for i, var_ID in enumerate(output_header.var_IDs):
-                    var[i, :] = resin.read_var_in_frame(time_index, var_ID)
+                var = resin.read_vars_in_frame(time_index)
 
                 # Replace bottom locally
                 nmodif = 0

cli/slf_last.py

@@ -33,16 +33,13 @@ def slf_last(args):
             else:
                 logger.warn('Input file is already single precision! Argument `--to_single_precision` is ignored')
 
-        values = np.empty((output_header.nb_var, output_header.nb_nodes), dtype=output_header.np_float_type)
         with Serafin.Write(args.out_slf, args.lang, overwrite=args.force) as resout:
             resout.write_header(output_header)
 
             time_index = len(resin.time) - 1
             time = resin.time[-1] if args.time is None else args.time
 
-            for i, var_ID in enumerate(output_header.var_IDs):
-                values[i, :] = resin.read_var_in_frame(time_index, var_ID)
-
+            values = resin.read_vars_in_frame(time_index)
             resout.write_entire_frame(output_header, time, values)
 
 

pyteltools/slf/Serafin.py

@@ -935,8 +935,35 @@ def read_var_in_frame(self, time_index, var_ID):
         self.file.seek(self.header.header_size + time_index * self.header.frame_size + 8 +
                        self.header.float_size + pos_var * (8 + self.header.float_size * self.header.nb_nodes), 0)
         self.file.read(4)
-        return np.array(self.header.unpack_float(self.file.read(self.header.float_size * self.header.nb_nodes),
-                                                 self.header.nb_nodes), dtype=self.header.np_float_type)
+        return np.array(self.header.unpack_float(
+            self.file.read(self.header.float_size * self.header.nb_nodes), self.header.nb_nodes),
+            dtype=self.header.np_float_type
+        )
+
+    def read_vars_in_frame(self, time_index, var_IDs=None):
+        """!
+        @brief Read multiple variables in a frame
+        @param time_index <int>: the index of the frame (0-based)
+        @param var_IDs <[str]>: list of variable IDs (if not present, all variables are considered)
+        @return <numpy 2D-array>: values of the variables with shape (number of variables, number of 2D nodes)
+        """
+        if var_IDs is None:
+            var_IDs = self.header.var_IDs
+        if time_index < 0:
+            raise SerafinRequestError('Impossible to read a negative time index!')
+        logger.debug('Reading variables %s at frame %i' % (var_IDs, time_index))
+
+        res = np.empty((len(var_IDs), self.header.nb_nodes))
+        for i, var_ID in enumerate(var_IDs):
+            pos_var = self._get_var_index(var_ID)
+            self.file.seek(self.header.header_size + time_index * self.header.frame_size + 8 +
+                           self.header.float_size + pos_var * (8 + self.header.float_size * self.header.nb_nodes), 0)
+            self.file.read(4)
+            res[i, :] = np.array(self.header.unpack_float(
+                self.file.read(self.header.float_size * self.header.nb_nodes), self.header.nb_nodes),
+                dtype=self.header.np_float_type
+            )
+        return res
 
     def read_var_in_frame_as_3d(self, time_index, var_ID):
         """!
@@ -950,6 +977,21 @@ def read_var_in_frame_as_3d(self, time_index, var_ID):
         new_shape = (self.header.nb_planes, self.header.nb_nodes_2d)
         return self.read_var_in_frame(time_index, var_ID).reshape(new_shape)
 
+    def read_vars_in_frame_as_3d(self, time_index, var_IDs=None):
+        """!
+        @brief Read multiple variables in a 3D frame
+        @param time_index <int>: the index of the frame (0-based)
+        @param var_IDs <[str]>: list of variable IDs (if not present, all variables are considered)
+        @return <numpy 3D-array>: values of the variables with shape
+            (number of variables, planes number, number of 2D nodes)
+        """
+        if self.header.is_2d:
+            raise SerafinRequestError('Reading values as 3D is only possible in 3D!')
+        if var_IDs is None:
+            var_IDs = self.header.var_IDs
+        new_shape = (len(var_IDs), self.header.nb_planes, self.header.nb_nodes_2d)
+        return self.read_vars_in_frame(time_index, var_IDs).reshape(new_shape)
+
     def read_var_in_frame_at_layer(self, time_index, var_ID, iplan):
         """!
         @brief Read a single variable in a frame at specific layer