From 07f867fdacee44a7773a57a36340adb95b3d6560 Mon Sep 17 00:00:00 2001 From: Elynn Wu Date: Tue, 4 Oct 2022 11:56:05 -0700 Subject: [PATCH 1/3] use quantity np sum instead --- driver/pace/driver/diagnostics.py | 2 +- 1 file changed, 1 insertion(+), 1 deletion(-) diff --git a/driver/pace/driver/diagnostics.py b/driver/pace/driver/diagnostics.py index 51e38f17e..507c8e85f 100644 --- a/driver/pace/driver/diagnostics.py +++ b/driver/pace/driver/diagnostics.py @@ -177,7 +177,7 @@ def _compute_column_integral( if len(q_in.shape) < 3: assert ShapeError(f"{name} does not have vertical levels.") column_integral = pace.util.Quantity( - sum( + q_in.np.sum( q_in.data[:, :, k] * delp.data[:, :, k] for k in range(q_in.metadata.extent[2]) ), From c476826d3a46b6361c9d3eac880e2a19802c2481 Mon Sep 17 00:00:00 2001 From: Elynn Wu Date: Fri, 7 Oct 2022 13:46:19 -0700 Subject: [PATCH 2/3] address PR comment --- driver/pace/driver/diagnostics.py | 14 ++++++-------- 1 file changed, 6 insertions(+), 8 deletions(-) diff --git a/driver/pace/driver/diagnostics.py b/driver/pace/driver/diagnostics.py index 507c8e85f..e0f2d0f6d 100644 --- a/driver/pace/driver/diagnostics.py +++ b/driver/pace/driver/diagnostics.py @@ -4,8 +4,6 @@ from datetime import datetime, timedelta from typing import List, Optional, Union -from sympy import ShapeError - import pace.driver import pace.dsl import pace.stencils @@ -174,13 +172,13 @@ def _compute_column_integral( q_in: tracer mixing ratio delp: pressure thickness of atmospheric layer """ - if len(q_in.shape) < 3: - assert ShapeError(f"{name} does not have vertical levels.") + if q_in.dims[2] != pace.util.Z_DIM: + raise NotImplementedError( + "this function assumes the z-dimension is the third dimension" + ) + k_slice = slice(q_in.origin[2], q_in.origin[2] + q_in.extent[2]) column_integral = pace.util.Quantity( - q_in.np.sum( - q_in.data[:, :, k] * delp.data[:, :, k] - for k in range(q_in.metadata.extent[2]) - ), + q_in.np.sum(q_in.data[:, :, k_slice] * delp.data[:, :, k_slice], axis=2), dims=("x", "y"), origin=q_in.metadata.origin[0:2], extent=(q_in.metadata.extent[0], q_in.metadata.extent[1]), From 3653a0525276f6b0087106f7c04616b8c0811e4e Mon Sep 17 00:00:00 2001 From: Elynn Wu Date: Wed, 12 Oct 2022 11:53:18 -0700 Subject: [PATCH 3/3] add assert if dim less than 3 --- driver/pace/driver/diagnostics.py | 1 + 1 file changed, 1 insertion(+) diff --git a/driver/pace/driver/diagnostics.py b/driver/pace/driver/diagnostics.py index e0f2d0f6d..1e91508e6 100644 --- a/driver/pace/driver/diagnostics.py +++ b/driver/pace/driver/diagnostics.py @@ -172,6 +172,7 @@ def _compute_column_integral( q_in: tracer mixing ratio delp: pressure thickness of atmospheric layer """ + assert len(q_in.dims) > 2 if q_in.dims[2] != pace.util.Z_DIM: raise NotImplementedError( "this function assumes the z-dimension is the third dimension"