<jupyter_start><jupyter_text>Offline validation tendency profiles<jupyter_code>import os
import xarray as xr
import numpy as np
from dask.diagnostics import ProgressBar
import fsspec
from matplotlib import pyplot as plt
import matplotlib as mpl
import string
import loaders
import yaml
from vcm.fv3.metadata import standardize_fv3_diagnostics
from vcm.catalog import catalog as CATALOG
from vcm.interpolate import interpolate_to_pressure_levels, PRESSURE_GRID
from diagnostics_utils import weighted_average
%run predictive_mapper.py
mpl.rcParams.update({'font.size': 8})
colors = plt.rcParams['axes.prop_cycle'].by_key()['color']
colors = [colors[ind] for ind in [0, 5, 6]]
mpl.rcParams['axes.prop_cycle'] = mpl.cycler(color=colors) 
N2F_MODEL_URLS = {
    '$TqR$-RF': 'gs://vcm-ml-experiments/2021-06-21-nudge-to-c3072-dq1-dq2-only/rf/trained_models/postphysics_ML_tendencies',
    '$TqR$-NN': 'gs://vcm-ml-experiments/2021-05-11-nudge-to-c3072-corrected-winds/nn-ensemble-model/trained_models/dq1-dq2',
    
}

N2F_TRAIN_TEST_DATA_URL = 'gs://vcm-ml-experiments/2021-04-13-n2f-c3072/3-hrly-ave-rad-precip-setting-30-min-rad-timestep-shifted-start-tke-edmf-3-hrly-ave-physics-tendencies'

N2F_TEST_CONFIG = 'gs://vcm-ml-experiments/2021-06-21-nudge-to-c3072-dq1-dq2-only/rf/offline_diags/postphysics_ML_tendencies/config.yaml'

OUTPUTDIR = 'figures'
# load the models for making ML predictions and make a mapper

nudging_variables = [
    "air_temperature",
    "specific_humidity",
    "x_wind",
    "y_wind",
    "pressure_thickness_of_atmospheric_layer"
]

models = {}
for ml_type, url in N2F_MODEL_URLS.items():
    models[ml_type] = fv3fit.load(url)
    
grid_c48 = standardize_fv3_diagnostics(CATALOG["grid/c48"].to_dask())

additional_vars = ['pressure_thickness_of_atmospheric_layer',]

predictive_mapper = PredictiveMapper(
    models,
    N2F_TRAIN_TEST_DATA_URL,
    loaders.open_nudge_to_fine,
    {'nudging_variables': nudging_variables},
    grid_c48,
    additional_vars=additional_vars
)
# evaluate over all of the test times
with fsspec.open(N2F_TEST_CONFIG, 'r') as f:
    config = yaml.safe_load(f)
test_times = config['batch_kwargs']['timesteps']
# augment the grid to make the plot less choppy
pressures = xr.DataArray(sorted(list(PRESSURE_GRID.values) + list(np.arange(27500., 77500., 5000.))), dims=PRESSURE_GRID.dims)

# load the predictions and targets, interpolate, and average
tendencies = xr.concat([predictive_mapper[time] for time in test_times], dim='time')
tendencies_pressure = interpolate_to_pressure_levels(
    tendencies, tendencies.pressure_thickness_of_atmospheric_layer, dim='z', levels=pressures
)
tendencies_pressure_profiles = weighted_average(
    tendencies_pressure, grid_c48['area']
)
with ProgressBar():
    tendencies_pressure_profiles = tendencies_pressure_profiles.mean(dim='time').load()
tendencies_pressure_profiles = tendencies_pressure_profiles.assign_coords(
    pressure=(tendencies_pressure_profiles.pressure/100).assign_attrs({'units': 'hPa'}))
def plot_dQ2_profile_case(da):
    fig, ax = plt.subplots(1, 1)
    ax.plot([0, 0], [1e3, 1], '-', color=[0.5, 0.5, 0.5], lw=1.5)
    hr, = da.sel(derivation='$TqR$-RF').plot(y='pressure', yincrease=False, ax=ax, lw=2, label='$TqR$-RF')
    hn, = da.sel(derivation='$TqR$-NN').plot(y='pressure', yincrease=False, ax=ax, lw=2, color=colors[1], label='$TqR$-NN')
    ht, = da.sel(derivation='target').plot(y='pressure', yincrease=False, ax=ax, lw=1, color='k', label='target')
    ax.legend()
    ax.set_xlim([-0.75e-1, 2.5e-2])
    ax.set_xticks(np.linspace(-5.0e-2, 0.0, 3))
    ax.set_ylim([1e3, 1])
    ax.set_ylabel('pressure [hPa]')
    ax.set_xlabel('g/kg/3-hr')
    ax.set_title('$\Delta Q_q$')
    ax.grid(axis='x')
    fig.set_size_inches([3.8, 4.5])
    fig.savefig(os.path.join(OUTPUTDIR, f"Figure_7_offline_moistening_profile.eps"), bbox_inches='tight')
plot_dQ2_profile_case(1000*3*3600*(tendencies_pressure_profiles.specific_humidity_tendency_due_to_ML))
# this converts matplotlib eps files to a more manageable size

!epstopdf figures/Figure_7_offline_moistening_profile.eps
!pdftops -eps figures/Figure_7_offline_moistening_profile.pdf
!rm figures/Figure_7_offline_moistening_profile.pdf<jupyter_output><empty_output>
{'fraction_non_alphanumeric': 0.13709296838131194, 'fraction_numerical': 0.03610193487494101, 'mean_word_length': 5.804173354735153, 'pattern_counts': {'":': 0, '<': 5, '<?xml version=': 0, '>': 5, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '1', 'dir': 'jupyter-structured-clean-dedup', 'id': '610876', 'license': 'MIT', 'n_tokens_mistral': 1779, 'n_tokens_neox': 1646, 'n_words': 254, 'path': 'notebooks/Figure-7-offline-validation-tendency-profiles.ipynb', 'repo_name': 'ai2cm/nudge-to-fine-3km-manuscript-workflow', 'size': '4238'}