wrapper.cpp

Go to the documentation of this file.

/*
 * Copyright 2023 Blue Brain Project, EPFL.
 * See the top-level LICENSE file for details.
 *
 * SPDX-License-Identifier: Apache-2.0
 */
#include <filesystem>
#include <optional>
#include <set>
#include <vector>
 
// 3rd party headers
#include <fmt/format.h>
#include <pybind11/embed.h>
#include <pybind11/stl.h>
 
// NMODL headers
#include "codegen/codegen_naming.hpp"
#include "pybind/ode_py.hpp"
#include "pybind/wrapper.hpp"
#include "pybind/pyembed.hpp"
#include "utils/common_utils.hpp"
 
 
namespace fs = std::filesystem;
namespace py = pybind11;
using namespace py::literals;
 
namespace nmodl {
namespace pybind_wrappers {
 
// This wrapper is used for obtaining better coverage in `ode.py`.
// Since we embed the `ode.py` as a string, there is no way to check what was covered via running
// pytest or similar. Instead, we use the coverage.py API directly.
static void run_python_script(const std::string& script, const py::dict& locals) {
#ifdef NRN_ENABLE_COVERAGE
    // to prevent race conditions during testing, we generate a random suffix
    const auto& suffix =
        nmodl::utils::generate_random_string(20, nmodl::utils::UseNumbersInString::WithoutNumbers);
 
    py::exec(fmt::format(R"(
import coverage
cov = coverage.Coverage(data_suffix='{}')
cov.start()
)",
                         suffix),
             locals);
    const auto& code_with_mapping = std::string("exec(compile(r'''" + ode_py + script + "''', '" +
                                                ode_py_path + "', 'exec'))");
    py::exec(code_with_mapping, locals);
#else
    py::exec(ode_py + script, locals);
#endif
 
#ifdef NRN_ENABLE_COVERAGE
    const auto& path = fs::current_path() / fmt::format("coverage_{}.xml", suffix);
    py::exec(fmt::format(R"(
cov.stop()
cov.save()
# Check if we have any coverage data
data = cov.get_data()
if data.measured_files():
    cov.xml_report(outfile='{}')
)",
                         path.string()),
             locals);
#endif
}
 
std::tuple<std::vector<std::string>, std::vector<std::string>, std::string>
call_solve_linear_system(const std::vector<std::string>& eq_system,
                         const std::vector<std::string>& state_vars,
                         const std::set<std::string>& vars,
                         bool small_system,
                         bool elimination,
                         const std::string& tmp_unique_prefix,
                         const std::set<std::string>& function_calls) {
    const auto locals = py::dict("eq_strings"_a = eq_system,
                                 "state_vars"_a = state_vars,
                                 "vars"_a = vars,
                                 "small_system"_a = small_system,
                                 "do_cse"_a = elimination,
                                 "function_calls"_a = function_calls,
                                 "tmp_unique_prefix"_a = tmp_unique_prefix);
    std::string script = R"(
exception_message = ""
try:
    solutions, new_local_vars = solve_lin_system(eq_strings,
                                                 state_vars,
                                                 vars,
                                                 function_calls,
                                                 tmp_unique_prefix,
                                                 small_system,
                                                 do_cse)
except Exception as e:
    # if we fail, fail silently and return empty string
    import traceback
    solutions = [""]
    new_local_vars = [""]
    exception_message = traceback.format_exc()
)";
    run_python_script(script, locals);
    // returns a vector of solutions, i.e. new statements to add to block:
    auto solutions = locals["solutions"].cast<std::vector<std::string>>();
    // and a vector of new local variables that need to be declared in the block:
    auto new_local_vars = locals["new_local_vars"].cast<std::vector<std::string>>();
    // may also return a python exception message:
    auto exception_message = locals["exception_message"].cast<std::string>();
 
    return {std::move(solutions), std::move(new_local_vars), std::move(exception_message)};
}
 
 
std::tuple<std::vector<std::string>, std::string> call_solve_nonlinear_system(
    const std::vector<std::string>& eq_system,
    const std::vector<std::string>& state_vars,
    const std::set<std::string>& vars,
    const std::set<std::string>& function_calls) {
    const auto locals = py::dict("equation_strings"_a = eq_system,
                                 "state_vars"_a = state_vars,
                                 "vars"_a = vars,
                                 "function_calls"_a = function_calls);
    std::string script = R"(
exception_message = ""
try:
    solutions = solve_non_lin_system(equation_strings,
                                               state_vars,
                                               vars,
                                               function_calls)
except Exception as e:
    # if we fail, fail silently and return empty string
    import traceback
    solutions = [""]
    exception_message = traceback.format_exc()
)";
 
    run_python_script(script, locals);
    // returns a vector of solutions, i.e. new statements to add to block:
    auto solutions = locals["solutions"].cast<std::vector<std::string>>();
    // may also return a python exception message:
    auto exception_message = locals["exception_message"].cast<std::string>();
 
    return {std::move(solutions), std::move(exception_message)};
}
 
 
std::tuple<std::string, std::string> call_diffeq_solver(const std::string& node_as_nmodl,
                                                        const std::string& dt_var,
                                                        const std::set<std::string>& vars,
                                                        bool use_pade_approx,
                                                        const std::set<std::string>& function_calls,
                                                        const std::string& method) {
    const auto locals = py::dict("equation_string"_a = node_as_nmodl,
                                 "dt_var"_a = dt_var,
                                 "vars"_a = vars,
                                 "use_pade_approx"_a = use_pade_approx,
                                 "function_calls"_a = function_calls);
 
    if (method == codegen::naming::EULER_METHOD) {
        // replace x' = f(x) differential equation
        // with forwards Euler timestep:
        // x = x + f(x) * dt
        std::string script = R"(
exception_message = ""
try:
    solution = forwards_euler2c(equation_string, dt_var, vars, function_calls)
except Exception as e:
    # if we fail, fail silently and return empty string
    import traceback
    solution = ""
    exception_message = traceback.format_exc()
)";
 
        run_python_script(script, locals);
    } else if (method == codegen::naming::CNEXP_METHOD) {
        // replace x' = f(x) differential equation
        // with analytic solution for x(t+dt) in terms of x(t)
        // x = ...
        std::string script = R"(
exception_message = ""
try:
    solution = integrate2c(equation_string, dt_var, vars,
                           use_pade_approx)
except Exception as e:
    # if we fail, fail silently and return empty string
    import traceback
    solution = ""
    exception_message = traceback.format_exc()
)";
 
        run_python_script(script, locals);
    } else {
        // nothing to do, but the caller should know.
        return {};
    }
    auto solution = locals["solution"].cast<std::string>();
    auto exception_message = locals["exception_message"].cast<std::string>();
 
    return {std::move(solution), std::move(exception_message)};
}
 
 
std::tuple<std::string, std::string> call_analytic_diff(
    const std::vector<std::string>& expressions,
    const std::set<std::string>& used_names_in_block) {
    auto locals = py::dict("expressions"_a = expressions, "vars"_a = used_names_in_block);
    std::string script = R"(
exception_message = ""
try:
    rhs = expressions[-1].split("=", 1)[1]
    solution = differentiate2c(rhs,
                               "v",
                               vars,
                               expressions[:-1]
               )
except Exception as e:
    # if we fail, fail silently and return empty string
    import traceback
    solution = ""
    exception_message = traceback.format_exc()
)";
 
    run_python_script(script, locals);
 
    auto solution = locals["solution"].cast<std::string>();
    auto exception_message = locals["exception_message"].cast<std::string>();
 
    return {std::move(solution), std::move(exception_message)};
}
 
std::tuple<std::string, std::string> call_diff2c(
    const std::string& expression,
    const std::pair<std::string, std::optional<int>>& variable,
    const std::unordered_set<std::string>& indexed_vars) {
    std::string statements;
    // only indexed variables require special treatment
    for (const auto& var: indexed_vars) {
        statements += fmt::format("_allvars.append(sp.IndexedBase('{}', shape=[1]))\n", var);
    }
    auto [name, property] = variable;
    if (property.has_value()) {
        name = fmt::format("sp.IndexedBase('{}', shape=[1])", name);
        statements += fmt::format("_allvars.append({})", name);
    } else {
        name = fmt::format("'{}'", name);
    }
    auto locals = py::dict("expression"_a = expression);
    std::string script =
        fmt::format(R"(
_allvars = []
{}
variable = {}
exception_message = ""
try:
    solution = differentiate2c(expression,
                               variable,
                               _allvars,
               )
except Exception as e:
    # if we fail, fail silently and return empty string
    solution = ""
    exception_message = str(e)
)",
                    statements,
                    property.has_value() ? fmt::format("{}[{}]", name, property.value()) : name);
 
    run_python_script(script, locals);
 
    auto solution = locals["solution"].cast<std::string>();
    auto exception_message = locals["exception_message"].cast<std::string>();
 
    return {std::move(solution), std::move(exception_message)};
}
 
void initialize_interpreter_func() {
    pybind11::initialize_interpreter(true);
}
 
void finalize_interpreter_func() {
    pybind11::finalize_interpreter();
}
 
// Prevent mangling for easier `dlsym`.
extern "C" {
NMODL_EXPORT pybind_wrap_api nmodl_init_pybind_wrapper_api() noexcept {
    return {&nmodl::pybind_wrappers::initialize_interpreter_func,
            &nmodl::pybind_wrappers::finalize_interpreter_func,
            &call_solve_nonlinear_system,
            &call_solve_linear_system,
            &call_diffeq_solver,
            &call_analytic_diff,
            &call_diff2c};
}
}
 
}  // namespace pybind_wrappers
}  // namespace nmodl