Change loss chart visual to darker line (#1917)

sdv-dev · Apr 18, 2024 · 7732f0d · 7732f0d
1 parent 938bea0
commit 7732f0d
Show file tree

Hide file tree

Showing 2 changed files with 2 additions and 4 deletions.
diff --git a/sdv/single_table/ctgan.py b/sdv/single_table/ctgan.py
@@ -71,15 +71,13 @@ def get_loss_values_plot(self, title='CTGAN loss function'):
 
         # Tidy up the loss values data
         loss_df = self._model.loss_values.copy()
-        loss_df['Generator Loss'] = loss_df['Generator Loss'].apply(lambda x: x.item())
-        loss_df['Discriminator Loss'] = loss_df['Discriminator Loss'].apply(lambda x: x.item())
 
         # Create a pretty chart using Plotly Express
         fig = px.line(
             loss_df, x='Epoch',
             y=['Generator Loss', 'Discriminator Loss'],
             color_discrete_map={
-                'Generator Loss': visualization.PlotConfig.DATACEBO_BLUE,
+                'Generator Loss': visualization.PlotConfig.DATACEBO_DARK,
                 'Discriminator Loss': visualization.PlotConfig.DATACEBO_GREEN
             },
         )

diff --git a/tests/unit/single_table/test_ctgan.py b/tests/unit/single_table/test_ctgan.py
@@ -330,7 +330,7 @@ def test_get_loss_values_plot(self, mock_line_plot):
         assert (fig['x'] == 'Epoch')
         assert (fig['y'] == ['Generator Loss', 'Discriminator Loss'])
         assert (fig['color_discrete_map'] == {
-            'Generator Loss': visualization.PlotConfig.DATACEBO_BLUE,
+            'Generator Loss': visualization.PlotConfig.DATACEBO_DARK,
             'Discriminator Loss': visualization.PlotConfig.DATACEBO_GREEN
         })