Fixes (#704)

sdv-dev · Sep 7, 2023 · 8fc2e5e · 8fc2e5e
1 parent 4f75aae
commit 8fc2e5e
Show file tree

Hide file tree

Showing 2 changed files with 8 additions and 9 deletions.
diff --git a/rdt/transformers/numerical.py b/rdt/transformers/numerical.py
@@ -463,7 +463,6 @@ def _fit(self, data):
  n_components=self.max_clusters,
  weight_concentration_prior_type='dirichlet_process',
  weight_concentration_prior=0.001,
- n_init=1,
  random_state=self._get_current_random_seed()
  )
 
@@ -494,10 +493,12 @@ def _transform(self, data):
 
  data = data.reshape((len(data), 1))
  means = self._bgm_transformer.means_.reshape((1, self.max_clusters))
-
+ means = means[:, self.valid_component_indicator]
  stds = np.sqrt(self._bgm_transformer.covariances_).reshape((1, self.max_clusters))
+ stds = stds[:, self.valid_component_indicator]
+
+ # Multiply stds by 4 so that a value will be in the range [-1,1] with 99.99% probability
  normalized_values = (data - means) / (self.STD_MULTIPLIER * stds)
- normalized_values = normalized_values[:, self.valid_component_indicator]
  component_probs = self._bgm_transformer.predict_proba(data)
  component_probs = component_probs[:, self.valid_component_indicator]
 
@@ -524,7 +525,8 @@ def _reverse_transform_helper(self, data):
  normalized = np.clip(data[:, 0], -1, 1)
  means = self._bgm_transformer.means_.reshape([-1])
  stds = np.sqrt(self._bgm_transformer.covariances_).reshape([-1])
- selected_component = data[:, 1].astype(int) # maybe round instead?
+ selected_component = data[:, 1].round().astype(int)
+ selected_component = selected_component.clip(0, self.valid_component_indicator.sum() - 1)
  std_t = stds[self.valid_component_indicator][selected_component]
  mean_t = means[self.valid_component_indicator][selected_component]
  reversed_data = normalized * self.STD_MULTIPLIER * std_t + mean_t
@@ -546,8 +548,6 @@ def _reverse_transform(self, data):
 
  recovered_data = self._reverse_transform_helper(data)
  if self.null_transformer and self.null_transformer.models_missing_values():
- data = np.stack([recovered_data, data[:, -1]], axis=1) # noqa: PD013
- else:
- data = recovered_data
+ recovered_data = np.stack([recovered_data, data[:, -1]], axis=1) # noqa: PD013
 
- return super()._reverse_transform(data)
+ return super()._reverse_transform(recovered_data)
diff --git a/tests/unit/transformers/test_numerical.py b/tests/unit/transformers/test_numerical.py
@@ -1213,7 +1213,6 @@ def test__fit(self, mock_bgm):
  n_components=10,
  weight_concentration_prior_type='dirichlet_process',
  weight_concentration_prior=0.001,
- n_init=1,
  random_state=0
  )