coala_interface.hpp

Go to the documentation of this file.

// -------------------------------------------------------//
//
// SHAMROCK code for hydrodynamics
// Copyright (c) 2021-2026 Timothée David--Cléris <tim.shamrock@proton.me>
// SPDX-License-Identifier: CeCILL Free Software License Agreement v2.1
// Shamrock is licensed under the CeCILL 2.1 License, see LICENSE for more information
//
// -------------------------------------------------------//
 
#pragma once

 
#include "shambase/assert.hpp"
#include "shambase/mdspan_concepts.hpp"
#include <experimental/mdspan>
#include <concepts>
 
namespace shamphys {

    template<class T>
    inline void compute_gij_k0(
        auto &&rho_dust,
        T rho_eps,
        shambase::is_mdspan_rank<1> auto massgrid,
        shambase::is_mdspan_rank<1> auto gij)
        requires requires(decltype(rho_dust) rd, int j) {
            { rd(j) } -> std::same_as<T>;
        }
    {
 
        SHAM_ASSERT(massgrid.extent(0) == gij.extent(0) + 1);
 
        for (std::size_t j = 0; j < gij.extent(0); ++j) {
            T rho_d = rho_dust(j);
            gij(j)  = (rho_d > rho_eps) ? rho_d / (massgrid[j + 1] - massgrid[j]) : 0;
        }
    }

    template<class Func>
        requires requires(Func f, int a, int b) {
            { f(a, b) };
        }
    inline void compute_flux_coag_k0_kdv(
        int nbins,
        shambase::is_mdspan_rank<1> auto gij,
        shambase::is_mdspan_rank<3> auto tensor_tabflux_coag,
        Func &&dv,
        shambase::is_mdspan_rank<1> auto flux) {
 
        SHAM_ASSERT(gij.extent(0) == nbins);
        SHAM_ASSERT(flux.extent(0) == nbins);
        SHAM_ASSERT(tensor_tabflux_coag.extent(0) == nbins);
        SHAM_ASSERT(tensor_tabflux_coag.extent(1) == nbins);
        SHAM_ASSERT(tensor_tabflux_coag.extent(2) == nbins);
 
        /*
         * Python version:
         * flux = np.einsum("jlm,lm,l,m->j", tensor_tabflux_coag, dv, gij, gij)
         */
 
        for (int j = 0; j < nbins; ++j) {
            double sum = 0.0;
            for (int l = 0; l < nbins; ++l) {
                for (int m = 0; m < nbins; ++m) {
                    sum += tensor_tabflux_coag(j, l, m) * dv(l, m) * gij[l] * gij[m];
                }
            }
            flux[j] = sum;
        }
    }

    void coala_flux_diff(
        shambase::is_mdspan_rank<1> auto flux, shambase::is_mdspan_rank<1> auto S_coag) {
 
        SHAM_ASSERT(flux.extent(0) == S_coag.extent(0));
 
        S_coag(0) = -flux(0);
        for (int j = 1; j < flux.extent(0); ++j) {
            S_coag(j) = flux(j - 1) - flux(j);
        }
    }
 
    template<class T, class FuncDv, class FuncRhoDust>
    void coala_k0_source_term(
        int nbins,
        /* inputs */
        FuncDv &&dv,
        FuncRhoDust &&rho_dust,
        T rho_eps,
        shambase::is_mdspan_rank<1> auto massgrid,
        /* COALA inputs */
        shambase::is_mdspan_rank<3> auto tabflux_coag,
        /* internal */
        shambase::is_mdspan_rank<1> auto gij,
        shambase::is_mdspan_rank<1> auto flux,
        /* output */
        shambase::is_mdspan_rank<1> auto S_coag) {
 
        // init the gij coefficients
        shamphys::compute_gij_k0(rho_dust, rho_eps, massgrid, gij);
 
        // compute flux for all dust bins
        shamphys::compute_flux_coag_k0_kdv(nbins, gij, tabflux_coag, dv, flux);
 
        // compute flux diff and store result
        shamphys::coala_flux_diff(flux, S_coag);
    }
 
} // namespace shamphys