Skip to content

EAMxx: Change DPxx large scale subsidence routine to semi Lagrangian formulation (draft) #7413

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Draft
wants to merge 5 commits into
base: master
Choose a base branch
from
Draft

EAMxx: Change DPxx large scale subsidence routine to semi Lagrangian formulation (draft) #7413

Changes from all commits
Commits
Show all changes
5 commits
Select commit Hold shift + click to select a range
b2e903d
change IOP subsidence calculation from finite centered to semi lagran…
May 30, 2025
7294900
spacing issue
May 30, 2025
5137089
move where lininterp is established
May 30, 2025
14b3396
sl subsidence with inline linear interpolation routine
May 31, 2025
c9c3357
update comments for SL iop routine
Jun 3, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
177 changes: 58 additions & 119 deletions components/eamxx/src/physics/iop_forcing/eamxx_iop_forcing_process_interface.cpp
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -2,6 +2,7 @@

#include "share/field/field_utils.hpp"
#include "share/property_checks/field_within_interval_check.hpp"
#include "ekat/util/ekat_lin_interp.hpp"

namespace scream
{
Expand Down Expand Up @@ -138,6 +139,25 @@ void IOPForcing::initialize_impl (const RunType run_type)
if (iop_nudge_tq or iop_nudge_uv) m_helper_fields.at("horiz_mean_weights").deep_copy(one_over_num_dofs);
}
// =========================================================================================
// =========================================================================================
// =========================================================================================
// Inline helper for scalar 1D linear interpolation
KOKKOS_FUNCTION
inline Real linear_interp_1d(const Real* x, const Real* f, const int n, const Real x_interp) {
if (x_interp <= x[0]) return f[0];
if (x_interp >= x[n-1]) return f[n-1];

for (int i = 0; i < n-1; ++i) {
if (x_interp >= x[i] && x_interp <= x[i+1]) {
Real t = (x_interp - x[i]) / (x[i+1] - x[i]);
return (1.0 - t) * f[i] + t * f[i+1];
}
}
return f[n-1]; // fallback
}

// =========================================================================================
// Main semi-Lagrangian subsidence routine
KOKKOS_FUNCTION
void IOPForcing::
advance_iop_subsidence(const MemberType& team,
Expand All @@ -157,131 +177,50 @@ advance_iop_subsidence(const MemberType& team,
constexpr Real Rair = C::Rair;
constexpr Real Cpair = C::Cpair;

const auto n_q_tracers = Q.extent_int(0);
const auto nlev_packs = ekat::npack<Pack>(nlevs);

// Get some temporary views from WS
uview_1d<Pack> omega_int, delta_u, delta_v, delta_T, tmp;
workspace.take_many_contiguous_unsafe<4>({"omega_int", "delta_u", "delta_v", "delta_T"},
{&omega_int, &delta_u, &delta_v, &delta_T});
const auto delta_Q_slot = workspace.take_macro_block("delta_Q", n_q_tracers);
uview_2d<Pack> delta_Q(delta_Q_slot.data(), n_q_tracers, nlev_packs);
const int n_q_tracers = Q.extent_int(0);

// Scalar views for pack-based inputs
auto s_ref_p_mid = ekat::scalarize(ref_p_mid);
auto s_omega = ekat::scalarize(omega);
auto s_delta_u = ekat::scalarize(delta_u);
auto s_delta_v = ekat::scalarize(delta_v);
auto s_delta_T = ekat::scalarize(delta_T);
auto s_delta_Q = ekat::scalarize(delta_Q);
auto s_omega_int = ekat::scalarize(omega_int);

// Compute omega on the interface grid by using a weighted average in pressure
const int pack_begin = 1/Pack::n, pack_end = (nlevs-1)/Pack::n;
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, pack_begin, pack_end+1), [&] (const int k){
auto range_pack = ekat::range<IntPack>(k*Pack::n);
range_pack.set(range_pack<1, 1);
Pack ref_p_mid_k, ref_p_mid_km1, omega_k, omega_km1;
ekat::index_and_shift<-1>(s_ref_p_mid, range_pack, ref_p_mid_k, ref_p_mid_km1);
ekat::index_and_shift<-1>(s_omega, range_pack, omega_k, omega_km1);

const auto weight = (ref_p_int(k) - ref_p_mid_km1)/(ref_p_mid_k - ref_p_mid_km1);
omega_int(k).set(range_pack>=1 and range_pack<=nlevs-1,
weight*omega_k + (1-weight)*omega_km1);
auto s_omega = ekat::scalarize(omega);
auto s_u = ekat::scalarize(u);
auto s_v = ekat::scalarize(v);
auto s_T = ekat::scalarize(T);
auto s_Q = ekat::scalarize(Q);

const Real* x_ptr = s_ref_p_mid.data();
const Real* omega_ptr = s_omega.data();
Real* u_ptr = s_u.data();
Real* v_ptr = s_v.data();
Real* T_ptr = s_T.data();
Real* Q_ptr = s_Q.data();

Kokkos::parallel_for(Kokkos::TeamThreadRange(team, nlevs), [&](const int k) {
const Real p_ref = x_ptr[k];
const Real omega_k = omega_ptr[k];
const Real p_dep = p_ref - dt * omega_k;

// Note that I know I should probably be using ekat's linear interp
// routine but I couldn't figure out for the life of me how to interface
// with this without getting an onslaught of compile issues. HELP!?

// Interpolate u, v, T at departure level
u_ptr[k] = linear_interp_1d(x_ptr, u_ptr, nlevs, p_dep);
v_ptr[k] = linear_interp_1d(x_ptr, v_ptr, nlevs, p_dep);
T_ptr[k] = linear_interp_1d(x_ptr, T_ptr, nlevs, p_dep);

// Add thermal expansion correction
T_ptr[k] *= 1.0 + (dt * Rair / Cpair) * omega_k / p_ref;

// Interpolate each tracer
for (int m = 0; m < n_q_tracers; ++m) {
Real* tracer_ptr = &Q_ptr[m * nlevs];
tracer_ptr[k] = linear_interp_1d(x_ptr, tracer_ptr, nlevs, p_dep);
}
});
omega_int(0)[0] = 0;
omega_int(nlevs/Pack::n)[nlevs%Pack::n] = 0;

// Compute delta views for u, v, T, and Q (e.g., u(k+1) - u(k), k=0,...,nlevs-2)
ColOps::compute_midpoint_delta(team, nlevs-1, u, delta_u);
ColOps::compute_midpoint_delta(team, nlevs-1, v, delta_v);
ColOps::compute_midpoint_delta(team, nlevs-1, T, delta_T);
for (int iq=0; iq<n_q_tracers; ++iq) {
auto tracer = Kokkos::subview(Q, iq, Kokkos::ALL());
auto delta_tracer = Kokkos::subview(delta_Q, iq, Kokkos::ALL());
ColOps::compute_midpoint_delta(team, nlevs-1, tracer, delta_tracer);
}
team.team_barrier();
}

// Compute updated temperature, horizontal winds, and tracers
Kokkos::parallel_for(Kokkos::TeamVectorRange(team, nlev_packs), [&] (const int k) {
auto range_pack = ekat::range<IntPack>(k*Pack::n);
const auto at_top = range_pack==0;
const auto not_at_top = not at_top;
const auto at_bot = range_pack==nlevs-1;
const auto not_at_bot = not at_bot;
const bool any_at_top = at_top.any();
const bool any_at_bot = at_bot.any();

// Get delta(k-1) packs. The range pack should not
// contain index 0 (so that we don't attempt to access
// k=-1 index) or index > nlevs-2 (since delta_* views
// are size nlevs-1).
auto range_pack_for_m1_shift = range_pack;
range_pack_for_m1_shift.set(range_pack<1, 1);
range_pack_for_m1_shift.set(range_pack>nlevs-2, nlevs-2);
Pack delta_u_k, delta_u_km1,
delta_v_k, delta_v_km1,
delta_T_k, delta_T_km1;
ekat::index_and_shift<-1>(s_delta_u, range_pack_for_m1_shift, delta_u_k, delta_u_km1);
ekat::index_and_shift<-1>(s_delta_v, range_pack_for_m1_shift, delta_v_k, delta_v_km1);
ekat::index_and_shift<-1>(s_delta_T, range_pack_for_m1_shift, delta_T_k, delta_T_km1);

// At the top and bottom of the model, set the end points for
// delta_*_k and delta_*_km1 to be the first and last entries
// of delta_*, respectively.
if (any_at_top) {
delta_u_k.set(at_top, s_delta_u(0));
delta_v_k.set(at_top, s_delta_v(0));
delta_T_k.set(at_top, s_delta_T(0));
}
if (any_at_bot) {
delta_u_km1.set(at_bot, s_delta_u(nlevs-2));
delta_v_km1.set(at_bot, s_delta_v(nlevs-2));
delta_T_km1.set(at_bot, s_delta_T(nlevs-2));
}

// Get omega_int(k+1) pack. The range pack should not
// contain index > nlevs-1 (since omega_int is size nlevs+1).
auto range_pack_for_p1_shift = range_pack;
range_pack_for_p1_shift.set(range_pack>nlevs-1, nlevs-1);
Pack omega_int_k, omega_int_kp1;
ekat::index_and_shift<1>(s_omega_int, range_pack, omega_int_k, omega_int_kp1);

const auto fac = (dt/2)/ref_p_del(k);

// Update u
u(k).update(not_at_bot, fac*omega_int_kp1*delta_u_k, -1, 1);
u(k).update(not_at_top, fac*omega_int_k*delta_u_km1, -1, 1);

// Update v
v(k).update(not_at_bot, fac*omega_int_kp1*delta_v_k, -1, 1);
v(k).update(not_at_top, fac*omega_int_k*delta_v_km1, -1, 1);

// Before updating T, first scale using thermal
// expansion term due to LS vertical advection
T(k) *= 1 + (dt*Rair/Cpair)*omega(k)/ref_p_mid(k);

// Update T
T(k).update(not_at_bot, fac*omega_int_kp1*delta_T_k, -1, 1);
T(k).update(not_at_top, fac*omega_int_k*delta_T_km1, -1, 1);

// Update Q
Pack delta_tracer_k, delta_tracer_km1;
for (int iq=0; iq<n_q_tracers; ++iq) {
auto s_delta_tracer = Kokkos::subview(s_delta_Q, iq, Kokkos::ALL());
ekat::index_and_shift<-1>(s_delta_tracer, range_pack_for_m1_shift, delta_tracer_k, delta_tracer_km1);
if (any_at_top) delta_tracer_k.set(at_top, s_delta_tracer(0));
if (any_at_bot) delta_tracer_km1.set(at_bot, s_delta_tracer(nlevs-2));

Q(iq, k).update(not_at_bot, fac*omega_int_kp1*delta_tracer_k, -1, 1);
Q(iq, k).update(not_at_top, fac*omega_int_k*delta_tracer_km1, -1, 1);
}
});

// Release WS views
workspace.release_macro_block(delta_Q_slot, n_q_tracers);
workspace.release_many_contiguous<4>({&omega_int, &delta_u, &delta_v, &delta_T});
}
// =========================================================================================
KOKKOS_FUNCTION
void IOPForcing::
Expand Down
Loading