Skip to content

Commit

Permalink
partially moves some math functions to kokkos, adds some KOKKOS_FUNCT…
Browse files Browse the repository at this point in the history 
…ION to functors that will be used on device
  • Loading branch information
@ecoon
ecoon committed Jun 19, 2024
1 parent 069f117 commit 2d78686
Show file tree
Hide file tree
Showing 4 changed files with 38 additions and 28 deletions.
4 changes: 2 additions & 2 deletions src/pks/flow/constitutive_relations/sources/soil_resistance_sakagucki_zeng_model.hh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -106,8 +106,8 @@ class SoilResistanceSakaguckiZengModel {
if (sat_gas_(c, 0) == 0.) {
r_soil = 0.; // ponded water
} else {
double vp_diffusion = 2.2e-5 * pow(poro_(c, 0), 2) * pow(1 - sr_, 2 + 3 * b_);
double L_Rsoil = d_ * (exp(pow(sat_gas_(c, 0), 5)) - 1) / (exp(1) - 1);
double vp_diffusion = 2.2e-5 * Kokkos::pow(poro_(c, 0), 2) * Kokkos::pow(1 - sr_, 2 + 3 * b_);
double L_Rsoil = d_ * (Kokkos::exp(Kokkos::pow(sat_gas_(c, 0), 5)) - 1) / (Kokkos::exp(1) - 1);
r_soil = L_Rsoil / vp_diffusion;
}
res_(sc, 0) = r_soil;
Expand Down
42 changes: 25 additions & 17 deletions src/pks/flow/constitutive_relations/wrm/wrm_van_genuchten.cc
View file
Original file line number Diff line number Diff line change
Expand Up @@ -9,6 +9,8 @@
*/

#include <cmath>
#include "Kokkos_MathematicalFunctions.hpp"

#include "dbc.hh"
#include "errors.hh"
#include "Spline.hh"
Expand Down Expand Up @@ -106,15 +108,16 @@ WRMVanGenuchten::WRMVanGenuchten(Teuchos::ParameterList& plist)
* The original curve is regulized on interval (s0, 1) using the
* Hermite interpolant of order 3. Formulas (3.11)-(3.12).
****************************************************************** */
KOKKOS_FUNCTION
double
WRMVanGenuchten::k_relative(double s) const
{
if (s <= s0_) {
double se = (s - sr_) / (1 - sr_);
if (function_ == RelPermFunction_kind::MUALEM) {
return pow(se, l_) * pow(1.0 - pow(1.0 - pow(se, 1.0 / m_), m_), 2.0);
return Kokkos::pow(se, l_) * Kokkos::pow(1.0 - Kokkos::pow(1.0 - Kokkos::pow(se, 1.0 / m_), m_), 2.0);
} else {
return se * se * (1.0 - pow(1.0 - pow(se, 1.0 / m_), m_));
return se * se * (1.0 - Kokkos::pow(1.0 - Kokkos::pow(se, 1.0 / m_), m_));
}
} else if (s == 1.0) {
return 1.0;
Expand All @@ -127,20 +130,21 @@ WRMVanGenuchten::k_relative(double s) const
/* ******************************************************************
* D Relative permeability / D capillary pressure pc.
****************************************************************** */
KOKKOS_FUNCTION
double
WRMVanGenuchten::d_k_relative(double s) const
{
if (s <= s0_) {
double se = (s - sr_) / (1 - sr_);

double x = pow(se, 1.0 / m_);
double x = Kokkos::pow(se, 1.0 / m_);
if (fabs(1.0 - x) < FLOW_WRM_TOLERANCE) return 0.0;
if (fabs(x) < FLOW_WRM_TOLERANCE) return 0.0;

double y = pow(1.0 - x, m_);
double y = Kokkos::pow(1.0 - x, m_);
double dkdse;
if (function_ == RelPermFunction_kind::MUALEM)
dkdse = (1.0 - y) * (l_ * (1.0 - y) + 2 * x * y / (1.0 - x)) * pow(se, l_ - 1.0);
dkdse = (1.0 - y) * (l_ * (1.0 - y) + 2 * x * y / (1.0 - x)) * Kokkos::pow(se, l_ - 1.0);
else
dkdse = (2 * (1.0 - y) + x / (1.0 - x)) * se;

Expand All @@ -157,11 +161,12 @@ WRMVanGenuchten::d_k_relative(double s) const
/* ******************************************************************
* Saturation formula (3.5)-(3.6).
****************************************************************** */
KOKKOS_FUNCTION
double
WRMVanGenuchten::saturation(double pc) const
{
if (pc > pc0_) {
return std::pow(1.0 + std::pow(alpha_ * pc, n_), -m_) * (1.0 - sr_) + sr_;
return Kokkos::pow(1.0 + Kokkos::pow(alpha_ * pc, n_), -m_) * (1.0 - sr_) + sr_;
} else if (pc <= 0.) {
return 1.0;
} else {
Expand All @@ -173,12 +178,13 @@ WRMVanGenuchten::saturation(double pc) const
/* ******************************************************************
* Derivative of the saturation formula w.r.t. capillary pressure.
****************************************************************** */
KOKKOS_FUNCTION
double
WRMVanGenuchten::d_saturation(double pc) const
{
if (pc > pc0_) {
return -m_ * n_ * std::pow(1.0 + std::pow(alpha_ * pc, n_), -m_ - 1.0) *
std::pow(alpha_ * pc, n_ - 1) * alpha_ * (1.0 - sr_);
return -m_ * n_ * Kokkos::pow(1.0 + Kokkos::pow(alpha_ * pc, n_), -m_ - 1.0) *
Kokkos::pow(alpha_ * pc, n_ - 1) * alpha_ * (1.0 - sr_);
} else if (pc <= 0.) {
return 0.0;
} else {
Expand All @@ -189,34 +195,36 @@ WRMVanGenuchten::d_saturation(double pc) const
/* ******************************************************************
* Pressure as a function of saturation.
****************************************************************** */
KOKKOS_FUNCTION
double
WRMVanGenuchten::capillaryPressure(double s) const
{
double se = (s - sr_) / (1.0 - sr_);
se = std::min<double>(se, 1.0);
se = std::max<double>(se, 1.e-40);
se = Kokkos::min(se, 1.0);
se = Kokkos::max(se, 1.e-40);
if (se < 1.e-8) {
return std::pow(se, -1.0 / (m_ * n_)) / alpha_;
return Kokkos::pow(se, -1.0 / (m_ * n_)) / alpha_;
} else {
return (std::pow(std::pow(se, -1.0 / m_) - 1.0, 1 / n_)) / alpha_;
return (Kokkos::pow(Kokkos::pow(se, -1.0 / m_) - 1.0, 1 / n_)) / alpha_;
}
}


/* ******************************************************************
* Derivative of pressure formulat w.r.t. saturation.
****************************************************************** */
KOKKOS_FUNCTION
double
WRMVanGenuchten::d_capillaryPressure(double s) const
{
double se = (s - sr_) / (1.0 - sr_);
se = std::min<double>(se, 1.0);
se = std::max<double>(se, 1.e-40);
se = Kokkos::min(se, 1.0);
se = Kokkos::max(se, 1.e-40);
if (se < 1.e-8) {
return -1.0 / (m_ * n_ * alpha_) * std::pow(se, -1.0 / (m_ * n_) - 1.) / (1.0 - sr_);
return -1.0 / (m_ * n_ * alpha_) * Kokkos::pow(se, -1.0 / (m_ * n_) - 1.) / (1.0 - sr_);
} else {
return -1.0 / (m_ * n_ * alpha_) * std::pow(std::pow(se, -1.0 / m_) - 1.0, 1 / n_ - 1.0) *
std::pow(se, -1.0 / m_ - 1.0) / (1.0 - sr_);
return -1.0 / (m_ * n_ * alpha_) * Kokkos::pow(Kokkos::pow(se, -1.0 / m_) - 1.0, 1 / n_ - 1.0) *
Kokkos::pow(se, -1.0 / m_ - 1.0) / (1.0 - sr_);
}
}

Expand Down
16 changes: 9 additions & 7 deletions src/pks/flow/constitutive_relations/wrm/wrm_van_genuchten.hh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -63,15 +63,17 @@ class WRMVanGenuchten {
static const std::string eval_type;

explicit WRMVanGenuchten(Teuchos::ParameterList& plist);
KOKKOS_FUNCTION WRMVanGenuchten(const WRMVanGenuchten& other) = default;
KOKKOS_INLINE_FUNCTION ~WRMVanGenuchten() = default;

// required methods from the base class
double k_relative(double saturation) const;
double d_k_relative(double saturation) const;
double saturation(double pc) const;
double d_saturation(double pc) const;
double capillaryPressure(double saturation) const;
double d_capillaryPressure(double saturation) const;
double residualSaturation() const { return sr_; }
KOKKOS_FUNCTION double k_relative(double saturation) const;
KOKKOS_FUNCTION double d_k_relative(double saturation) const;
KOKKOS_FUNCTION double saturation(double pc) const;
KOKKOS_FUNCTION double d_saturation(double pc) const;
KOKKOS_FUNCTION double capillaryPressure(double saturation) const;
KOKKOS_FUNCTION double d_capillaryPressure(double saturation) const;
KOKKOS_FUNCTION double residualSaturation() const { return sr_; }

private:
double m_; // van Genuchten parameters: m, n, alpha
Expand Down
4 changes: 2 additions & 2 deletions src/pks/surface_balance/constitutive_relations/land_cover/snow_meltrate_model.hh
View file
Original file line number Diff line number Diff line change
Expand Up @@ -118,7 +118,7 @@ class SnowMeltRateModel {
res_(i, 0) = melt_rate_ * (exp_temp_(i, 0) - 273.15);

if (swe_(i, 0) < land_cover_.snow_transition_depth) {
res_(i, 0) *= std::max(0., swe_(i, 0) / land_cover_.snow_transition_depth);
res_(i, 0) *= Kokkos::max(0., swe_(i, 0) / land_cover_.snow_transition_depth);
}

} else {
Expand All @@ -132,7 +132,7 @@ class SnowMeltRateModel {
if (exp_temp_(i, 0) > 273.15) {
res_(i, 0) = melt_rate_;
if (swe_(i, 0) < land_cover_.snow_transition_depth) {
res_(i, 0) *= std::max(0., swe_(i, 0) / land_cover_.snow_transition_depth);
res_(i, 0) *= Kokkos::max(0., swe_(i, 0) / land_cover_.snow_transition_depth);
}
} else {
res_(i, 0) = 0.0;
Expand Down

0 comments on commit 2d78686

Please sign in to comment.