Merge pull request #91 from IMMM-SFA/hotfix/plotting

fix plotting bug
IMMM-SFA · Aug 2, 2022 · aa06e07 · aa06e07
2 parents dd30856 + eb16d0d
commit aa06e07
Show file tree

Hide file tree

Showing 4 changed files with 16 additions and 9 deletions.
diff --git a/mosartwmpy/_version.py b/mosartwmpy/_version.py
@@ -1 +1 @@
-__version__ = "0.3.0"
+__version__ = "0.3.1"
diff --git a/mosartwmpy/model.py b/mosartwmpy/model.py
@@ -264,9 +264,10 @@ def plot_variable(
             self,
             variable: str,
             log_scale: bool = False,
+            show: bool = True,
     ):
         """Display a colormap of a spatial variable at the current timestep."""
-        data = self.unmask(self.get_value_ptr(variable)).reshape(self.get_grid_shape())
+        data = self.get_value_ptr(variable).reshape(self.get_grid_shape())
         if log_scale:
             data = np.where(data > 0, data, np.nan)
         xr.DataArray(
@@ -281,7 +282,9 @@ def plot_variable(
             cmap='winter_r',
             norm=colors.LogNorm() if log_scale else None,
         )
-        plt.show()
+        if show:
+            plt.show()
+        return plt
 
     def unmask(self, vector: np.ndarray) -> np.ndarray:
         unmasked = np.empty_like(self.mask, dtype=vector.dtype)

diff --git a/mosartwmpy/plotting/plot.py b/mosartwmpy/plotting/plot.py
@@ -90,12 +90,12 @@ def plot_reservoir(
     offset_x = spacing[0] / 2
     offset_y = spacing[1] / 2
 
-    dam_index = model.grid.reservoir_id.tolist().index(grand_id)
+    dam_index = model.unmask(model.grid.reservoir_id).tolist().index(grand_id)
     ilat = dam_index // model.get_grid_shape()[1]
     ilon = dam_index % model.get_grid_shape()[1]
     dependent_cell_indices = model.grid.reservoir_dependency_database.query(
         f'reservoir_id == {grand_id}').index.values
-    basin_cell_indices = np.nonzero(model.grid.outlet_id == model.grid.outlet_id[dam_index])[0]
+    basin_cell_indices = np.nonzero(model.unmask(model.grid.outlet_id) == model.unmask(model.grid.outlet_id)[dam_index])[0]
 
     if grand_file_path is not None:
         grand = gpd.read_file(grand_file_path)
@@ -109,7 +109,7 @@ def plot_reservoir(
     stream_indices = [dam_index]
     j = dam_index
     while True:
-        i = model.grid.downstream_id[j]
+        i = model.unmask(model.grid.downstream_id)[j]
         if (i >= 0) and (i != j):
             stream_indices.append(i)
             j = i
@@ -118,7 +118,7 @@ def plot_reservoir(
     stream_indices = np.array(stream_indices)
 
     to_plot = [{
-        'geometry': Point(model.grid.longitude[dam_index], model.grid.latitude[dam_index]),
+        'geometry': Point(model.unmask(model.grid.longitude)[dam_index], model.unmask(model.grid.latitude)[dam_index]),
         'type': 'dam',
     }]
     if grand_point is not None:
@@ -133,8 +133,8 @@ def plot_reservoir(
         })
     for t, array in dict(dependency=dependent_cell_indices, basin=basin_cell_indices, stream=stream_indices).items():
         for i in array:
-            x = model.grid.longitude[i]
-            y = model.grid.latitude[i]
+            x = model.unmask(model.grid.longitude)[i]
+            y = model.unmask(model.grid.latitude)[i]
             to_plot.append({
                 'geometry': Polygon([
                     (x - offset_x, y - offset_y),

diff --git a/mosartwmpy/tests/test_model.py b/mosartwmpy/tests/test_model.py
@@ -36,6 +36,10 @@ def test_can_run(self):
         self.model.update()
         self.assertTrue(True, "model initializes and updates")
 
+    def test_can_plot(self):
+        self.model.update()
+        self.model.plot_variable('surface_water_amount', log_scale=True, show=False)
+
     def test_input_io(self):
         self.assertGreater(self.model.get_input_item_count(), 0, "model can count input")
         input_vars = self.model.get_input_var_names()