Merge pull request #557 from OpenCOMPES/fix_normed_dtype

Fix normed dtype

Author: rettigl

src/sed/binning/binning.py

@@ -465,6 +465,7 @@ def normalization_histogram_from_timed_dataframe(
     axis: str,
     bin_centers: np.ndarray,
     time_unit: float,
+    **kwds,
 ) -> xr.DataArray:
     """Get a normalization histogram from a timed dataframe.
 
@@ -475,17 +476,12 @@ def normalization_histogram_from_timed_dataframe(
             histogram.
         bin_centers (np.ndarray): Bin centers used for binning of the axis.
         time_unit (float): Time unit the data frame entries are based on.
+        **kwds: Additional keyword arguments passed to the bin_dataframe function.
 
     Returns:
         xr.DataArray: Calculated normalization histogram.
     """
-    bins = df[axis].map_partitions(
-        pd.cut,
-        bins=bin_centers_to_bin_edges(bin_centers),
-    )
-
-    histogram = df[axis].groupby([bins]).count().compute().values * time_unit
-    # histogram = bin_dataframe(df, axes=[axis], bins=[bin_centers]) * time_unit
+    histogram = bin_dataframe(df, axes=[axis], bins=[bin_centers], **kwds) * time_unit
 
     data_array = xr.DataArray(
         data=histogram,

src/sed/core/processor.py

@@ -2283,6 +2283,8 @@ def compute(
                 )
                 # if the axes are named correctly, xarray figures out the normalization correctly
                 self._normalized = self._binned / self._normalization_histogram
+                # Set datatype of binned data
+                self._normalized.data = self._normalized.data.astype(self._binned.data.dtype)
                 self._attributes.add(
                     self._normalization_histogram.values,
                     name="normalization_histogram",
@@ -2353,36 +2355,35 @@ def get_normalization_histogram(
 
         if isinstance(df_partitions, int):
             df_partitions = list(range(0, min(df_partitions, self._dataframe.npartitions)))
+
         if use_time_stamps or self._timed_dataframe is None:
             if df_partitions is not None:
-                self._normalization_histogram = normalization_histogram_from_timestamps(
-                    self._dataframe.partitions[df_partitions],
-                    axis,
-                    self._binned.coords[axis].values,
-                    self._config["dataframe"]["columns"]["timestamp"],
-                )
+                dataframe = self._dataframe.partitions[df_partitions]
             else:
-                self._normalization_histogram = normalization_histogram_from_timestamps(
-                    self._dataframe,
-                    axis,
-                    self._binned.coords[axis].values,
-                    self._config["dataframe"]["columns"]["timestamp"],
-                )
+                dataframe = self._dataframe
+            self._normalization_histogram = normalization_histogram_from_timestamps(
+                df=dataframe,
+                axis=axis,
+                bin_centers=self._binned.coords[axis].values,
+                time_stamp_column=self._config["dataframe"]["columns"]["timestamp"],
+            )
         else:
             if df_partitions is not None:
-                self._normalization_histogram = normalization_histogram_from_timed_dataframe(
-                    self._timed_dataframe.partitions[df_partitions],
-                    axis,
-                    self._binned.coords[axis].values,
-                    self._config["dataframe"]["timed_dataframe_unit_time"],
-                )
+                timed_dataframe = self._timed_dataframe.partitions[df_partitions]
             else:
-                self._normalization_histogram = normalization_histogram_from_timed_dataframe(
-                    self._timed_dataframe,
-                    axis,
-                    self._binned.coords[axis].values,
-                    self._config["dataframe"]["timed_dataframe_unit_time"],
-                )
+                timed_dataframe = self._timed_dataframe
+            self._normalization_histogram = normalization_histogram_from_timed_dataframe(
+                df=timed_dataframe,
+                axis=axis,
+                bin_centers=self._binned.coords[axis].values,
+                time_unit=self._config["dataframe"]["timed_dataframe_unit_time"],
+                hist_mode=self.config["binning"]["hist_mode"],
+                mode=self.config["binning"]["mode"],
+                pbar=self.config["binning"]["pbar"],
+                n_cores=self.config["core"]["num_cores"],
+                threads_per_worker=self.config["binning"]["threads_per_worker"],
+                threadpool_api=self.config["binning"]["threadpool_API"],
+            )
 
         return self._normalization_histogram
 

tests/test_processor.py

@@ -1008,6 +1008,8 @@ def test_compute_with_normalization() -> None:
         processor.binned.data,
         (processor.normalized * processor.normalization_histogram).data,
     )
+    # check dtype
+    assert processor.normalized.dtype == processor.binned.dtype
     # bin only second dataframe partition
     result2 = processor.compute(
         bins=bins,