From b31f5885df318bf3d7c125b13a187d4ac6fb8159 Mon Sep 17 00:00:00 2001 From: Mark Piper Date: Fri, 28 Feb 2020 16:48:24 -0700 Subject: [PATCH] Change behavior of grid_x, y, and z for rectilinear grids For a rectilinear grid, these functions should return, respectively, a vector of length ncols, nrows, and nzlevels. --- pymt/framework/bmi_bridge.py | 21 ++++++++++++++++++--- 1 file changed, 18 insertions(+), 3 deletions(-) diff --git a/pymt/framework/bmi_bridge.py b/pymt/framework/bmi_bridge.py index 1cd726c6..2307ba80 100644 --- a/pymt/framework/bmi_bridge.py +++ b/pymt/framework/bmi_bridge.py @@ -441,19 +441,34 @@ def grid_nodes_per_face(self, grid, out=None): def grid_x(self, grid, out=None): if out is None: - out = np.empty(self.grid_node_count(grid), dtype=float) + if self.grid_type(grid) == 'rectilinear': + out = np.empty(self.grid_shape(grid)[1], dtype=float) + else: + out = np.empty(self.grid_node_count(grid), dtype=float) self.bmi.get_grid_x(grid, out) return out def grid_y(self, grid, out=None): if out is None: - out = np.empty(self.grid_node_count(grid), dtype=float) + if self.grid_type(grid) == 'rectilinear': + out = np.empty(self.grid_shape(grid)[0], dtype=float) + else: + out = np.empty(self.grid_node_count(grid), dtype=float) self.bmi.get_grid_y(grid, out) return out def grid_z(self, grid, out=None): if out is None: - out = np.empty(self.node_count(grid), dtype=float) + if self.grid_type(grid) == 'rectilinear': + shape = self.grid_shape(grid) + try: + zdim = shape[2] + except IndexError: + zdim = 1 + out = np.empty(zdim, dtype=float) + else: + out = np.empty(self.grid_node_count(grid), dtype=float) + self.bmi.get_grid_z(grid, out) return out