use fill_value

sdv-dev · Aug 11, 2023 · 712fe62 · 712fe62
1 parent 754b8b2
commit 712fe62
Showing 1 changed file with 2 additions and 4 deletions.
diff --git a/rdt/transformers/categorical.py b/rdt/transformers/categorical.py
@@ -128,8 +128,7 @@ def _fit(self, data):
         labels = self._order_categories(labels)
         freq = data.value_counts(normalize=True, dropna=False)
         nan_value = freq[np.nan] if np.nan in freq.index else None
-        freq = freq.reindex(labels).array
-        freq[np.isnan(freq)] = nan_value
+        freq = freq.reindex(labels, fill_value=nan_value).array
 
         self.frequencies, self.intervals = self._compute_frequencies_intervals(labels, freq)
 
@@ -269,8 +268,7 @@ def _fit(self, data):
             freq = data.value_counts(normalize=True, dropna=False)
 
         nan_value = freq[np.nan] if np.nan in freq.index else None
-        freq = freq.reindex(self.order).array
-        freq[np.isnan(freq)] = nan_value
+        freq = freq.reindex(self.order, fill_value=nan_value).array
 
         self.frequencies, self.intervals = self._compute_frequencies_intervals(self.order, freq)