Add DerivedTCWV diagnostic for TCWV calculation

[wm-ytakano]

Earth2Studio Pull Request

Description

This PR adds a new diagnostic model DerivedTCWV that calculates Total Column Water Vapor (TCWV) from specific humidity at pressure levels and surface pressure.

Motivation

Some prognostic models (e.g., Pangu-Weather) do not output TCWV directly in their predictions. However, TCWV is a required input variable for diagnostic models like PrecipitationAFNO. This diagnostic model bridges that gap by computing TCWV from the commonly available specific humidity fields at pressure levels and surface pressure.

Implementation Details

• DerivedTCWV: Implements vertical integration of specific humidity using the trapezoidal rule:
  • TCWV = (1/g) * ∫(p_surface to p_top) q dp
  • Properly handles the layer between surface pressure and the lowest pressure level
  • Configurable pressure levels (default: [1000, 850, 700, 500, 300, 200, 100] hPa)
  • Input: specific humidity at pressure levels (q{level}) + surface pressure (sp)
  • Output: total column water vapor (tcwv) in kg/m²

Changes

• Added DerivedTCWV class in earth2studio/models/dx/derived.py
• Added unit tests in test/models/dx/test_derived.py
• Updated documentation in docs/modules/models.rst
• Updated CHANGELOG.md

Checklist

• I am familiar with the Contributing Guidelines.
• New or existing tests cover these changes.
• The documentation is up to date with these changes.
• The CHANGELOG.md is up to date with these changes.
• An issue is linked to this pull request.
• Assess and address Greptile feedback (AI code review bot for guidance; use discretion, addressing all feedback is not required).

Dependencies

[greptile-apps]

Disclaimer: This is AI-generated, please review response for accuracy

Greptile Summary

This PR adds DerivedTCWV, a diagnostic model that calculates Total Column Water Vapor from specific humidity at pressure levels and surface pressure, enabling prognostic models like Pangu-Weather (which don't output TCWV directly) to work with diagnostic models like PrecipitationAFNO that require TCWV as input.

The implementation uses vertical integration with the trapezoidal rule between pressure levels and handles the surface layer appropriately by using a rectangular approximation from surface pressure to the lowest valid pressure level. The code properly masks pressure levels below the surface and handles edge cases where the surface pressure falls between defined pressure levels.

Key changes:

• Implemented TCWV calculation using the formula: TCWV = (1/g) × ∫ q dp
• Properly handles variable surface pressure with masking logic
• Uses configurable pressure levels (default: [1000, 850, 700, 500, 300, 200, 100] hPa)
• Updated documentation and exports
• Added basic test coverage (shape validation, zero input, invalid coordinates)

Confidence Score: 4/5

• This PR is safe to merge with minimal risk - the implementation is well-structured and the integration logic is sound
• The implementation correctly handles the vertical integration with proper masking for surface pressure boundaries. The code follows established patterns in the codebase and includes appropriate documentation. Minor improvements could be made to test coverage (numerical accuracy validation), but the core functionality is solid and the changes are isolated to a new diagnostic model
• Pay attention to test/models/dx/test_derived.py - consider adding numerical validation tests to verify integration accuracy

Important Files Changed

Filename	Overview
earth2studio/models/dx/derived.py	Added DerivedTCWV class implementing vertical integration of specific humidity to compute TCWV using trapezoidal rule with proper surface pressure handling
test/models/dx/test_derived.py	Added tests for DerivedTCWV covering basic functionality, zero input, and invalid coordinates, but lacks numerical validation of integration accuracy

[greptile-apps]

Additional Comments (1)

1. test/models/dx/test_derived.py, line 472-489 (link)

   style: tests validate shapes and zero input but lack numerical accuracy checks for the integration

   consider adding a test case with known values to verify the trapezoidal integration. for example, with constant q and uniform pressure spacing, the analytical TCWV can be calculated and compared

   # example: test with constant q and known pressure range
   q_const = 0.01  # kg/kg
   p_surface = 101325  # pa
   p_top = 10000  # pa
   expected_tcwv = q_const * (p_surface - p_top) / model.g

   _{Note: If this suggestion doesn't match your team's coding style, reply to this and let me know. I'll remember it for next time!}

_{5 files reviewed, 1 comment}

_{Edit Code Review Agent Settings | Greptile}