Data Extraction

This notebook demonstrates gcmprocpy’s data extraction functions (arr_*) for slicing model output along various dimensions, as well as height interpolation utilities.

Note: This notebook requires TIE-GCM or WACCM-X model output files.

[ ]:

import warnings
warnings.filterwarnings('ignore')

import numpy as np
import gcmprocpy as gy

directory = '/glade/work/nikhilr/tiegcm3.0/benchmarks/2.5/seasons/decsol_smin/hist'
datasets = gy.load_datasets(directory, dataset_filter='sech')
times = gy.time_list(datasets)
t_val = times[0]
print(f'Using time: {t_val}')

Using time: 2002-12-21T01:00:00.000000000

arr_lat_lon — Latitude x Longitude Slice

Extract a 2D lat/lon array at a given time and pressure level. Use plot_mode=True to get a PlotData object with metadata.

[17]:

result = gy.arr_lat_lon(datasets, 'TN', t_val, selected_lev_ilev=5.0, plot_mode=True)
print(f'shape:  {result.values.shape}')
print(f'unit:   {result.variable_unit}')
print(f'level:  {result.selected_lev}')
print(f'lats:   {result.lats[0]:.1f} to {result.lats[-1]:.1f}')
print(f'lons:   {result.lons[0]:.1f} to {result.lons[-1]:.1f}')

shape:  (72, 144)
unit:   K
level:  5.0
lats:   -88.8 to 88.8
lons:   -180.0 to 177.5

arr_var — Full Variable at a Time

Extract the full 3D field (lev x lat x lon) at a single time step.

[18]:

result = gy.arr_var(datasets, 'TN', t_val, plot_mode=True)
print(f'shape: {result.values.shape}  (nlev x nlat x nlon)')
print(f'levs:  {result.levs.shape}')

shape: (57, 72, 144)  (nlev x nlat x nlon)
levs:  (57,)

arr_lev_var — Vertical Profile

Extract a 1D vertical profile at a specific latitude and longitude.

[19]:

result = gy.arr_lev_var(datasets, 'TN', t_val, selected_lat=0.0, selected_lon=0.0, plot_mode=True)
print(f'shape: {result.values.shape}')
print(f'levs:  {result.levs}')

shape: (57,)
levs:  [-6.875 -6.625 -6.375 -6.125 -5.875 -5.625 -5.375 -5.125 -4.875 -4.625
 -4.375 -4.125 -3.875 -3.625 -3.375 -3.125 -2.875 -2.625 -2.375 -2.125
 -1.875 -1.625 -1.375 -1.125 -0.875 -0.625 -0.375 -0.125  0.125  0.375
  0.625  0.875  1.125  1.375  1.625  1.875  2.125  2.375  2.625  2.875
  3.125  3.375  3.625  3.875  4.125  4.375  4.625  4.875  5.125  5.375
  5.625  5.875  6.125  6.375  6.625  6.875  7.125]

arr_lev_lon — Level x Longitude Cross-Section

Extract a 2D lev/lon cross-section at a given latitude.

[20]:

result = gy.arr_lev_lon(datasets, 'TN', t_val, selected_lat=0.0, plot_mode=True)
print(f'shape: {result.values.shape}  (nlev x nlon)')
print(f'levs:  {result.levs.shape}')
print(f'lons:  {result.lons.shape}')

shape: (57, 144)  (nlev x nlon)
levs:  (57,)
lons:  (144,)

arr_lev_lat — Level x Latitude Cross-Section

Extract a 2D lev/lat cross-section at a given longitude.

[21]:

result = gy.arr_lev_lat(datasets, 'TN', t_val, selected_lon=0.0, plot_mode=True)
print(f'shape: {result.values.shape}  (nlev x nlat)')
print(f'levs:  {result.levs.shape}')
print(f'lats:  {result.lats.shape}')

shape: (57, 72)  (nlev x nlat)
levs:  (57,)
lats:  (72,)

arr_lev_time — Level x Time Cross-Section

Extract a 2D lev/time cross-section at a specific latitude and longitude.

[22]:

result = gy.arr_lev_time(datasets, 'TN', selected_lat=0.0, selected_lon=0.0, plot_mode=True)
print(f'shape:  {result.values.shape}  (nlev x ntime)')
print(f'levs:   {result.levs.shape}')
print(f'mtimes: {len(result.mtime_values)}')

shape:  (57, 120)  (nlev x ntime)
levs:   (57,)
mtimes: 120

arr_lat_time — Latitude x Time Cross-Section

Extract a 2D lat/time cross-section at a given longitude and level.

[23]:

result = gy.arr_lat_time(datasets, 'TN', selected_lon=0.0, selected_lev_ilev=5.0, plot_mode=True)
print(f'shape:  {result.values.shape}  (nlat x ntime)')
print(f'lats:   {result.lats.shape}')
print(f'mtimes: {len(result.mtime_values)}')

shape:  (72, 120)  (nlat x ntime)
lats:   (72,)
mtimes: 120

arr_lon_time — Longitude x Time Cross-Section

[24]:

result = gy.arr_lon_time(datasets, 'TN', selected_lat=0.0, selected_lev_ilev=5.0, plot_mode=True)
print(f'shape:  {result.values.shape}  (nlon x ntime)')
print(f'lons:   {result.lons.shape}')
print(f'mtimes: {len(result.mtime_values)}')

shape:  (144, 120)  (nlon x ntime)
lons:   (144,)
mtimes: 120

arr_var_time — Variable vs Time Series

[25]:

result = gy.arr_var_time(datasets, 'TN', selected_lat=0.0, selected_lon=0.0, selected_lev_ilev=5.0, plot_mode=True)
print(f'shape:  {result.values.shape}  (ntime,)')
print(f'mtimes: {len(result.mtime_values)}')

shape:  (120,)  (ntime,)
mtimes: 120

batch_arr_lat_lon — Multiple Variables at Once

Extract several variables in a single pass over the datasets.

[26]:

results = gy.batch_arr_lat_lon(datasets, ['TN', 'O1', 'NO'], t_val, selected_lev_ilev=5.0, plot_mode=True)
for name, r in results.items():
    print(f'{name}: shape={r.values.shape}, unit={r.variable_unit}')

TN: shape=(72, 144), unit=K
O1: shape=(72, 144), unit=mmr
NO: shape=(72, 144), unit=mmr

arr_sat_track — Satellite Track Interpolation

Interpolate model data along a satellite trajectory using trilinear interpolation.

[27]:

# Simulate a LEO ascending pass
sat_time = np.array([times[0] + np.timedelta64(i * 6, 'm') for i in range(20)])
sat_lat = np.linspace(-60, 60, 20)
sat_lon = np.linspace(-120, 120, 20)

# 1D extraction at a fixed level
values = gy.arr_sat_track(datasets, 'NO', sat_time, sat_lat, sat_lon, selected_lev_ilev=5.0)
print(f'At fixed level: shape={values.shape}')

# 2D extraction across all levels
values_2d = gy.arr_sat_track(datasets, 'NO', sat_time, sat_lat, sat_lon)
print(f'All levels:     shape={values_2d.shape}')

At fixed level: shape=(20,)
All levels:     shape=(57, 20)

Height Interpolation

height_to_pres_level

Convert a target height in km to the nearest pressure level.

[28]:

pres_level = gy.height_to_pres_level(datasets, t_val, 300.0)
print(f'300 km -> pressure level {pres_level}')

# At a specific location
pres_level_ll = gy.height_to_pres_level(datasets, t_val, 300.0, latitude=0.0, longitude=0.0)
print(f'300 km at equator, 0E -> pressure level {pres_level_ll}')

# Different heights map to different levels
for h in [100.0, 200.0, 300.0, 400.0, 500.0]:
    pl = gy.height_to_pres_level(datasets, t_val, h)
    print(f'  {h:.0f} km -> {pl}')

300 km -> pressure level 3.0
300 km at equator, 0E -> pressure level 3.25
  100 km -> -6.25
  200 km -> 0.0
  300 km -> 3.0
  400 km -> 5.25
  500 km -> 7.0

interpolate_to_height

Interpolate a 2D field from pressure levels to constant height surfaces.

[29]:

# Extract a lev x lat cross-section on pressure levels
result_lev_lat = gy.arr_lev_lat(datasets, 'TN', t_val, selected_lon=0.0, plot_mode=True)

# Interpolate to 30 evenly spaced height levels
interp_vals, target_heights = gy.interpolate_to_height(
    datasets, result_lev_lat.values, result_lev_lat.levs, t_val, n_heights=30)
print(f'Input shape:  {result_lev_lat.values.shape} (nlev x nlat)')
print(f'Output shape: {interp_vals.shape} (n_heights x nlat)')
print(f'Height range: {target_heights[0]:.1f} to {target_heights[-1]:.1f} km')

# Interpolate to specific heights
custom_hts = np.array([100.0, 200.0, 300.0, 400.0, 500.0])
interp_custom, _ = gy.interpolate_to_height(
    datasets, result_lev_lat.values, result_lev_lat.levs, t_val, target_heights=custom_hts)
print(f'\nCustom heights shape: {interp_custom.shape}')

Input shape:  (57, 72) (nlev x nlat)
Output shape: (30, 72) (n_heights x nlat)
Height range: 96.4 to 482.7 km

Custom heights shape: (5, 72)

Cleanup

[30]:

gy.close_datasets(datasets)
print('Datasets closed.')

Datasets closed.