Modularize GaussianCopulaSynthesizer._fit (#2273)

Author: pvk-developer

sdv/single_table/copulas.py

@@ -135,16 +135,28 @@ def _fit(self, processed_data):
         log_numerical_distributions_error(
             self.numerical_distributions, processed_data.columns, LOGGER
         )
-        self._num_rows = len(processed_data)
+        self._num_rows = self._learn_num_rows(processed_data)
+        numerical_distributions = self._get_numerical_distributions(processed_data)
+        self._model = self._initialize_model(numerical_distributions)
+        self._fit_model(processed_data)
 
+    def _learn_num_rows(self, processed_data):
+        return len(processed_data)
+
+    def _get_numerical_distributions(self, processed_data):
         numerical_distributions = deepcopy(self._numerical_distributions)
         for column in processed_data.columns:
             if column not in numerical_distributions:
                 numerical_distributions[column] = self._numerical_distributions.get(
                     column, self._default_distribution
                 )
-        self._model = multivariate.GaussianMultivariate(distribution=numerical_distributions)
 
+        return numerical_distributions
+
+    def _initialize_model(self, numerical_distributions):
+        return multivariate.GaussianMultivariate(distribution=numerical_distributions)
+
+    def _fit_model(self, processed_data):
         with warnings.catch_warnings():
             warnings.filterwarnings('ignore', module='scipy')
             self._model.fit(processed_data)

tests/unit/single_table/test_copulas.py

@@ -228,6 +228,83 @@ def test__fit(self, mock_multivariate, mock_warnings):
         mock_warnings.catch_warnings.assert_called_once()
         instance._num_rows == 10
 
+    def test__fit_mocked_instance(self):
+        """Test that the `_fit` method calls the modularized functions."""
+        # Setup
+        instance = Mock(numerical_distributions={})
+        processed_data = Mock(columns=[])
+        numerical_distributions = Mock()
+        instance._get_numerical_distributions.return_value = numerical_distributions
+
+        # Run
+        GaussianCopulaSynthesizer._fit(instance, processed_data)
+
+        # Assert
+        instance._learn_num_rows.assert_called_once_with(processed_data)
+        instance._get_numerical_distributions.assert_called_once_with(processed_data)
+        instance._initialize_model.assert_called_once_with(numerical_distributions)
+        instance._fit_model.assert_called_once_with(processed_data)
+
+    def test__learn_num_rows(self):
+        """Test that the `_learn_num_rows` method returns the correct number of rows."""
+        # Setup
+        metadata = Metadata()
+        instance = GaussianCopulaSynthesizer(metadata)
+        processed_data = pd.DataFrame({'a': range(5), 'b': range(5)})
+
+        # Run
+        result = instance._learn_num_rows(processed_data)
+
+        # Assert
+        assert result == 5
+
+    def test__get_numerical_distributions_with_existing_columns(self):
+        """Test that `_get_numerical_distributions` returns correct distributions."""
+        # Setup
+        metadata = Metadata()
+        instance = GaussianCopulaSynthesizer(metadata)
+        instance._numerical_distributions = {'a': 'dist_a', 'b': 'dist_b'}
+        instance._default_distribution = 'default_dist'
+
+        processed_data = Mock()
+        processed_data.columns = ['a', 'b', 'c']
+
+        # Run
+        result = instance._get_numerical_distributions(processed_data)
+
+        # Assert
+        expected_result = {'a': 'dist_a', 'b': 'dist_b', 'c': 'default_dist'}
+        assert result == expected_result
+
+    @patch('sdv.single_table.copulas.multivariate.GaussianMultivariate')
+    def test__initialize_model(self, mock_gaussian_multivariate):
+        """Test that `_initialize_model` calls the GaussianMultivariate with correct parameters."""
+        # Setup
+        metadata = Metadata()
+        instance = GaussianCopulaSynthesizer(metadata)
+        numerical_distributions = {'a': 'dist_a', 'b': 'dist_b'}
+
+        # Run
+        model = instance._initialize_model(numerical_distributions)
+
+        # Assert
+        mock_gaussian_multivariate.assert_called_once_with(distribution=numerical_distributions)
+        assert model == mock_gaussian_multivariate.return_value
+
+    def test__fit_model(self):
+        """Test that `_fit_model` fits the model correctly."""
+        # Setup
+        metadata = Metadata()
+        instance = GaussianCopulaSynthesizer(metadata)
+        instance._model = Mock()
+        processed_data = Mock()
+
+        # Run
+        instance._fit_model(processed_data)
+
+        # Assert
+        instance._model.fit.assert_called_once_with(processed_data)
+
     def test__get_nearest_correlation_matrix_valid(self):
         """Test ``_get_nearest_correlation_matrix`` with a psd input.