% This script plots the spectra (including interpolations) in
% Figure 2 of 
%
% T. Schneider, L. R. Leung, and R. C. J. Wills. Opinion: Optimizing 
% climate models with process- knowledge, resolution, and AI. 
% Atmos. Chem. Phys., 2024.

% It uses MOZAIC spectra provided by Joern Callies, from
%
% J. Callies, R. Ferrari, and O. Buhler. Transition from
% geostrophic turbulence to inertia–gravity waves in the
% atmospheric energy spectrum. Proc. Natl. Acad. Sci.,
% 111:17033–17038, 2014
%
% and NAWDEX spectra provided by Ulrich Schumann from
%
% U. Schumann. The horizontal spectrum of vertical velocities near
% the tropopause from global to gravity wave
% scales. J. Atmos. Sci., 76:3847–3862, 2019.
%
% Please reference these primary sources when using the data or the
% code provided here. 
%
% This script requires, in addition to Matlab base, the Curvefitting and 
% Statistics and Machine Learning toolboxes.
%
% The readNPY functions are provided here for completeness; they are
% originally from https://github.com/kwikteam/npy-matlab/blob/master/npy-matlab/readNPY.m
%
% (c) Tapio Schneider, Caltech


% Minimum and maximum wavelength to include in plots and integrals
l_min        = 0.01;   % minimum scale to extrapolate to [km]
l_max        = 5e3;    % maximum scale to extrapolate to [km]

% Empirical conversion factor between wavelength and model
% resolution needed to resolve a wavelength 
delta_x_conv = 7;     

% Load (larger-scale)  MOZAIC spectra from Callies et al. (2014)
ku          = readNPY('ku.npy'); % inverse wavelength [1/m]
kv          = readNPY('kv.npy');
Pu          = readNPY('Pu.npy'); % Power spectral density u [m^3/s^2]
Pv          = readNPY('Pv.npy'); % Power spectral density v [m^3/s^2]
h_ls        = (Pu + Pv)/2;

% Assert that wavenumber axes for u and v are identical
assert(max(abs(ku-kv)) < eps, 'Incompatible axes.')

% Compute wavelength and wavenumber axes for larger-scale spectra
lambda_ls   = 1e-3 * 1./ku;     % Convert to wavelength in units of km, to be
                                % consistent with NAWDEX 
k_ls        = 2*pi*1e3*ku;      % Wavenumber [1/km] analogous to NAWDEX

% Truncate larger-scale data to wavelength above cutoff in Callies et al.
% (2014) [because spectra at smaller scales may be inaccurate]
lambda_lsmin = 20;              % Large-scale cutoff [km]
ls_idx      = find(lambda_ls >= lambda_lsmin);
k_ls        = k_ls(ls_idx);
lambda_ls   = lambda_ls(ls_idx);
h_ls        = h_ls(ls_idx);

% Load (smaller-scale) NAWDEX spectra from Schumann (2019)
load('Schumann_table.csv');
k_ss        = Schumann_table(:, 1); % wavenumber [1/km]
lambda_ss   = Schumann_table(:, 2); % wavelength [km] = 2*pi/k
h_ss        = Schumann_table(:, 3); % spectral power horizontal [m^3/s^2]
v_ss        = Schumann_table(:, 4); % spectral power vertical [m^3/s^2]

% Truncate smaller-scale data to wavelength above cutoff in Schumann (2019) 
lambda_ssmin = 0.1;             % Small-scale cutoff [km]
ss_idx      = find(lambda_ss >= lambda_ssmin);
k_ss        = k_ss(ss_idx);
lambda_ss   = lambda_ss(ss_idx);
h_ss        = h_ss(ss_idx);
v_ss        = v_ss(ss_idx);

% Find range of overlap of wavelengths for larger-scale and 
% smaller-scale horizontal spectra
lmax_ovr    = max(lambda_ss);
lmin_ovr    = min(lambda_ls);
ils_ovr     = find(lambda_ls >= lmin_ovr & lambda_ls <= lmax_ovr);
iss_ovr     = find(lambda_ss >= lmin_ovr & lambda_ss <= lmax_ovr);

% Interpolate to common axis in overlap range, with equally spaced points
% in log space
lambda_ovr  = exp(linspace(log(lmin_ovr), log(lmax_ovr), 100));  
hss_ovr     = interp1(lambda_ss(iss_ovr), h_ss(iss_ovr), lambda_ovr, 'spline');
hls_ovr     = interp1(lambda_ls(ils_ovr), h_ls(ils_ovr), lambda_ovr, 'spline');

% Compute scale factor between spectra from mean ratio in overlap region
ss_scale    = mean(hls_ovr./hss_ovr);

% Common wavenumber axis (equally spaced in log space) for interpolation
% and extrapolation of spectra and computation of integrals
k_max       = 2*pi/l_min;
k_min       = 2*pi/l_max;
k_com       = exp(linspace(log(k_min), log(k_max), 500));
lambda_com  = 2*pi./k_com;
deltax_com  = lambda_com/delta_x_conv;

% Indices of common axis with wavelength within those measured and outside
in_idx      = find(lambda_com >= min(lambda_ss));
out_idx     = find(lambda_com < min(lambda_ss));

% Splice re-scaled NAWDEX horizontal spectrum to larger-scale spectrum 
% where the larger-scale spectrum ends; interpolate with monotonicity 
% preserving cubicspline in log-log space, including extrapolation to 
% larger scales than measured 
h_com       = zeros(size(k_com));
h_com(in_idx) = exp(interp1(log([k_ls; k_ss]), ...
    log([h_ls; ss_scale * h_ss]), ...
    log(k_com(in_idx)), 'pchip', 'extrap'));

% Same splicing, interpolation, and extrapolation for vertical spectra 
v_com       = zeros(size(k_com));
v_com(in_idx) = exp(interp1(log(k_ss(ss_idx)), log(v_ss(ss_idx)), ...
    log(k_com(in_idx)), 'pchip', 'extrap'));

% Compute linear fit in log-log space in wavelength range near bottom of spectrum
lrange_min = min(lambda_ss);  % minimum wavelength for fit [km]
lrange_max = 1;               % maximum wavelength for fit [km]
fit_idx    = find(lambda_com >= lrange_min & lambda_com <= lrange_max);
Bh         = robustfit(log(k_com(fit_idx)), log(h_com(fit_idx)));
Bv         = robustfit(log(k_com(fit_idx)), log(v_com(fit_idx)));

% Extrapolate to smaller scales using linear fit in log-log space
h_com(out_idx) = exp(Bh(1) + Bh(2) * log(k_com(out_idx)));
v_com(out_idx) = exp(Bv(1) + Bv(2) * log(k_com(out_idx)));

% Compute normalized integrals over commong wavenumber range
hint = cumsum(h_com .* gradient(k_com)); % Horizontal spectra
hint = hint/hint(end);
vint = cumsum(v_com .* gradient(k_com)); % Vertical spectra
vint = vint/vint(end);

% Colors for plot
light_blue = [158,202,225]/255;
blue       = [49,130,189]/255;
orange     = [230,85,13]/255;

% Plot results
clf
init_figure

% Spectra, upper panel
subplot(2,1,1)
loglog(1./lambda_ls, h_ls, 'o-', 'Color', blue)
hold on
loglog(1./lambda_ss, h_ss, 'o-', 'Color', light_blue)
loglog(1./lambda_ss, v_ss, 'o-', 'Color', orange)
loglog(1./lambda_com, h_com, '-.', 'Color', blue)
loglog(1./lambda_com, v_com, '-.', 'Color', orange)

% Reference lines
% k^{-3} spectrum
invl3 = exp(linspace(log(1/4000), log(1/500), 20));
E3    = 0.08 * invl3.^(-3);
loglog(invl3, E3, 'k-')

% k^{-5/3} spectrum
invl53 = exp(linspace(log(1/2), log(1/.20), 20));
E53    = 4e1 * invl53.^(-5/3);
loglog(invl53, E53, 'k-')

inv_lambda_lim = [1/l_max, 1/l_min];
inv_deltax_lim = delta_x_conv*inv_lambda_lim;
set(gca, 'Box', 'off', ...
    'TickDir', 'out', ...
    'XMinorTick', 'off', ...
    'YMinorTick', 'off', ...
    'Ylim', [1e-3, 1e9], ...
    'xlim', inv_lambda_lim)
xticks([1e-3, 1e-2, 1e-1, 1e0, 1e1, 1e2])
yticks([1e-3 1e0, 1e3, 1e6, 1e9])
ylabel('Energy density (m^3 s^{-2})')
xlabel('Inverse Wavelength (km^{-1})')
lgd1 = legend('Horizontal (MOZAIC)', 'Horizontal (NAWDEX)', 'Vertical (NAWDEX)');
set(lgd1, 'Box', 'off', 'Location', 'northeast')

% Cumulative spectra, lower panel
subplot(2,1,2);
semilogx(1./deltax_com, hint, '-', 'Color', blue)
hold on
semilogx(1./deltax_com, vint, '-', 'Color', orange)
semilogx([inv_deltax_lim(1), inv_deltax_lim(2)], [0.5 0.5], 'k:') % reference line

% fit and plot reference line on cumulative spectrum
iref       = find(lambda_com <= 25 & lambda_com >= 1e-2);
deltax_ref = deltax_com(iref); 

fo = fitoptions('Method','NonlinearLeastSquares',...
               'StartPoint',[0.1 2/3]);
ft = fittype('1-a*(x^b-x_min^b)','problem','x_min','options',fo);
[curve2,gof2] = fit(deltax_ref', vint(iref)', ft, 'problem', min(deltax_com));
a = curve2.a;
b = curve2.b;
fprintf('Cumulative spectrum reference exponent = %1.2f\n', b)
fprintf('\n')

CE_ref     = 1 - a * (deltax_ref.^b - min(deltax_com)^b );
semilogx(1./deltax_ref, CE_ref, '--', 'Color', orange) % reference line

% labels

%xlabel('Inverse Wavelength (km^{-1})')
xlabel('Horizontal resolution (km)')
ylabel('Cumulative Energy')
xticks([1e-2, 1e-1, 1e0, 1e1, 1e2])
yticks([0:0.25:1])

set(gca, 'Box', 'off', ...
    'TickDir', 'out', ...
    'XMinorTick', 'off', ...
    'YMinorTick', 'off', ...
    'Ylim', [0, 1.05], ...
    'Xlim', inv_deltax_lim, ...
    'XTickLabel', {'100', '10', '1', '0.1', '0.01'}, ...
    'YTickLabel', {'0', '', '0.5', '', '1'})

lgd2 = legend('Horizontal', 'Vertical');
set(lgd2, 'Box', 'off', 'Location', 'southeast')

% add inverse wavelength axis on top
hAx1=gca;                % get the first axis handle so can address each desired
hAx2=axes('Position',hAx1.Position, ...
          'XAxisLocation','top', ...
          'YAxisLocation','left', ...
          'XScale', 'log', ...
          'xlim', inv_lambda_lim, ...
          'Color','none', ...
          'TickDir', 'out', ...
          'xtick', [1e-3, 1e-2, 1e-1, 1e0, 1e1, 1e2], ...
          'XMinorTick', 'off', ...
          'YTick',[], ...
          'Box', 'off');

print('Atmos-Spectrum.pdf', '-dpdf')

% get some representative values
disp('Fractions of vertical and horizontal kinetic energy captured at a given minimum wavelength:')
lambda_ref = [1000, 100, 10, 1, .1];
for j=1:length(lambda_ref)
    [m, idx] = min(abs(lambda_ref(j) - lambda_com));
    fprintf('lambda_min = %5.1f; vint = %2.2f; hint = %2.2f\n', lambda_com(idx), 100*vint(idx), 100*hint(idx));
end