update scatter plot to support metadata coloring with variables (#152)

* update scatter plot to support metadata coloring with variables * update release notes * add test
IAMconsortium · Dec 6, 2018 · 58d489c · 58d489c
1 parent d76a3fa
commit 58d489c
Show file tree

Hide file tree

Showing 4 changed files with 25 additions and 5 deletions.
diff --git a/RELEASE_NOTES.md b/RELEASE_NOTES.md
@@ -1,6 +1,7 @@
 
 # Next Release
 
+- [#152](https://github.com/IAMconsortium/pyam/pull/152) Fix bug where scatter plots did not work with property metadata when using two variables (#136, #152)
 - [#151](https://github.com/IAMconsortium/pyam/pull/151) Fix bug where excel files were not being written on Windows and MacOSX (#149)
 - [#145](https://github.com/IAMconsortium/pyam/pull/145) Support full semantic and VCS-style versioning with `versioneer`
 

diff --git a/pyam/core.py b/pyam/core.py
@@ -893,7 +893,8 @@ def line_plot(self, x='year', y='value', **kwargs):
             if x != 'year' and y != 'year':
                 df = df.drop('year', axis=1)  # years causes NaNs
 
-        ax, handles, labels = plotting.line_plot(df, x=x, y=y, **kwargs)
+        ax, handles, labels = plotting.line_plot(
+            df.dropna(), x=x, y=y, **kwargs)
         return ax
 
     def stack_plot(self, *args, **kwargs):
@@ -937,20 +938,20 @@ def scatter(self, x, y, **kwargs):
             dfx = (
                 self
                 .filter(variable=x)
-                .as_pandas()
+                .as_pandas(with_metadata=True)
                 .rename(columns={'value': x, 'unit': 'xunit'})
                 .set_index(YEAR_IDX)
                 .drop('variable', axis=1)
             )
             dfy = (
                 self
                 .filter(variable=y)
-                .as_pandas()
+                .as_pandas(with_metadata=True)
                 .rename(columns={'value': y, 'unit': 'yunit'})
                 .set_index(YEAR_IDX)
                 .drop('variable', axis=1)
             )
-            df = dfx.join(dfy).reset_index()
+            df = dfx.join(dfy, lsuffix='_left', rsuffix='').reset_index()
         else:
             # filter, merge with meta, and rename value column to match var
             var = x if xisvar else y
@@ -960,7 +961,7 @@ def scatter(self, x, y, **kwargs):
                 .as_pandas(with_metadata=True)
                 .rename(columns={'value': var})
             )
-        ax = plotting.scatter(df, x, y, **kwargs)
+        ax = plotting.scatter(df.dropna(), x, y, **kwargs)
         return ax
 
     def map_regions(self, map_col, agg=None, copy_col=None, fname=None,

diff --git a/tests/expected_figs/test_scatter_variables_with_meta_color.png b/tests/expected_figs/test_scatter_variables_with_meta_color.png
diff --git a/tests/test_plotting.py b/tests/test_plotting.py
@@ -432,6 +432,24 @@ def test_scatter(plot_df):
     return fig
 
 
+@pytest.mark.mpl_image_compare(**MPL_KWARGS)
+def test_scatter_variables_with_meta_color(plot_df):
+    fig, ax = plt.subplots(figsize=(8, 8))
+    plot_df.categorize(
+        'foo', 'a',
+        criteria={'Primary Energy': {'up': 5, 'year': 2010}},
+        color='blue'
+    )
+    plot_df.categorize(
+        'foo', 'b',
+        criteria={'Primary Energy': {'lo': 5, 'year': 2010}},
+        color='red'
+    )
+    plot_df.scatter(ax=ax, x='Primary Energy', y='Primary Energy|Coal',
+                    color='foo')
+    return fig
+
+
 @pytest.mark.mpl_image_compare(**MPL_KWARGS)
 def test_scatter_with_lines(plot_df):
     fig, ax = plt.subplots(figsize=(8, 8))