Add test and improve variables

NNPDF · May 26, 2023 · bd676ee · bd676ee
1 parent 97dffd9
commit bd676ee
Show file tree

Hide file tree

Showing 4 changed files with 30 additions and 14 deletions.
diff --git a/validphys2/src/validphys/dataplots.py b/validphys2/src/validphys/dataplots.py
@@ -895,20 +895,20 @@ def plot_smpdf(pdf, dataset, obs_pdf_correlations, mark_threshold:float=0.9):
     categorical = not np.issubdtype(plotting_var.dtype, np.number)
     if categorical:
         # Categorical
-        keys, values = np.unique(plotting_var, return_inverse=True)
+        categorical_keys, values = np.unique(plotting_var, return_inverse=True)
         plotting_var = values
-        nunique = len(keys)
-        if nunique <= len(cm.Set2.colors):
-            cmap = mcolors.ListedColormap(cm.Set2.colors[:nunique])
+        num_categories = len(categorical_keys)
+        if num_categories <= len(cm.Set2.colors):
+            cmap = mcolors.ListedColormap(cm.Set2.colors[:num_categories])
         else:
-            cmap = cm.viridis.resample(nunique)
-        bins = np.linspace(0, nunique, nunique + 1)
-        norm = mcolors.BoundaryNorm(bins, nunique)
+            cmap = cm.viridis.resample(num_categories)
+        bins = np.linspace(0, num_categories, num_categories + 1)
+        norm = mcolors.BoundaryNorm(bins, num_categories)
 
     else:
         cmap = cm.viridis
         #TODO: vmin vmax should be global or by figure?
-        vmin,vmax = min(plotting_var), max(plotting_var)
+        vmin, vmax = min(plotting_var), max(plotting_var)
         if info.x_scale == 'log':
             norm = mcolors.LogNorm(vmin, vmax)
         else:
@@ -931,12 +931,12 @@ def plot_smpdf(pdf, dataset, obs_pdf_correlations, mark_threshold:float=0.9):
         label = info.group_label(same_vals, info.figure_by)
         #TODO: PY36ScalarMappable
         #TODO Improve title?
-        title = "%s %s\n[%s]" % (info.dataset_label, '(%s)'%label if label else '' ,pdf.label)
+        title = f"{info.dataset_label} {label if label else ''}\n[{pdf.label}]"
 
         #Start plotting
-        w,h = plt.rcParams["figure.figsize"]
+        w,h = mpl.rcParams["figure.figsize"]
         h*=2.5
-        fig,axes = plt.subplots(nrows=nf ,sharex=True, figsize=(w,h), sharey=True)
+        fig, axes = plotutils.subplots(nrows=nf, sharex=True, figsize=(w,h), sharey=True)
         fig.suptitle(title)
         colors = sm.to_rgba(fb["__plotting_var"])
         for flindex, (ax, fl) in enumerate(zip(axes, fls)):
@@ -954,7 +954,6 @@ def plot_smpdf(pdf, dataset, obs_pdf_correlations, mark_threshold:float=0.9):
             ax.set_ylim(-1,1)
             ax.set_xlim(x[0], x[-1])
         ax.set_xlabel('$x$')
-        #fig.subplots_adjust(hspace=0)
 
         cbar = fig.colorbar(
             sm,
@@ -963,8 +962,8 @@ def plot_smpdf(pdf, dataset, obs_pdf_correlations, mark_threshold:float=0.9):
             aspect=100,
         )
         if categorical:
-            cbar.set_ticks(np.linspace(0.5, nunique - 0.5, nunique))
-            cbar.ax.set_yticklabels(keys)
+            cbar.set_ticks(np.linspace(0.5, num_categories - 0.5, num_categories))
+            cbar.ax.set_yticklabels(categorical_keys)
 
         #TODO: Fix title for this
         #fig.tight_layout()

diff --git a/validphys2/src/validphys/tests/baseline/test_plot_smpdf_categorical.png b/validphys2/src/validphys/tests/baseline/test_plot_smpdf_categorical.png
diff --git a/validphys2/src/validphys/tests/conftest.py b/validphys2/src/validphys/tests/conftest.py
@@ -27,6 +27,8 @@ def tmp(tmpdir):
 
 SINGLE_DATASET = {'dataset': 'NMC'}
 
+SINGLE_CATEGICAL = {"dataset": "ATLAS_WZ_TOT_13TEV", 'cfac': ["QCD"]}
+
 DATA = [
     {'dataset': 'NMC'},
     {'dataset': 'ATLASTTBARTOT', 'cfac':['QCD']},
@@ -95,6 +97,16 @@ def single_data_internal_cuts_config(data_internal_cuts_config):
     config_dict.update(dataset_input=DATA[0])
     return config_dict
 
+@pytest.fixture(scope='module')
+def single_data_categorical_internal_cuts_config(data_internal_cuts_config):
+    """Test dataset with categorical plotting variables"""
+    return {
+        **data_internal_cuts_config,
+        'dataset_input': SINGLE_CATEGICAL,
+        # NOTE: The old theory is currently bugged for this dataset
+        'theoryid': THEORYID_NEW,
+    }
+
 @pytest.fixture(scope='module')
 def single_data_single_point_internal_cuts_config(single_data_internal_cuts_config):
     config_dict = dict(single_data_internal_cuts_config)

diff --git a/validphys2/src/validphys/tests/test_plots.py b/validphys2/src/validphys/tests/test_plots.py
@@ -40,6 +40,11 @@ def test_dataspecschi2():
 def test_plot_smpdf(single_data_internal_cuts_config):
     return next(API.plot_smpdf(**single_data_internal_cuts_config))
 
+@pytest.mark.linux
+@pytest.mark.mpl_image_compare
+def test_plot_smpdf_categorical(single_data_categorical_internal_cuts_config):
+    return next(API.plot_smpdf(**single_data_categorical_internal_cuts_config))
+
 @pytest.mark.linux
 @pytest.mark.mpl_image_compare
 def test_plot_obscorrs(single_data_internal_cuts_config):