gmxapi/src/gmx/core/export_context.cpp at 7fedeb36c997f90a5906af1f963db3c00a6b56a7 · kassonlab/gmxapi · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
//
// Created by Eric Irrgang on 3/18/18.
//

#include <iostream>
#include <zconf.h>
#include "core.h"

#include "gmxapi/context.h"

#include "gromacs/utility/init.h"

namespace gmxpy
{

namespace detail
{

namespace py = pybind11;


//-dd <vector> (0 0 0)
//Domain decomposition grid, 0 is optimize
//-npme <int> (-1)
//Number of separate ranks to be used for PME, -1 is guess
//-nt <int> (0)
//Total number of threads to start (0 is guess)
//-ntmpi <int> (0)
//Number of thread-MPI ranks to start (0 is guess)
//-ntomp <int> (0)
//Number of OpenMP threads per MPI rank to start (0 is guess)
//-ntomp_pme <int> (0)
//Number of OpenMP threads per MPI rank to start (0 is -ntomp)
//-nsteps <int> (-2)
//Run this number of steps, overrides .mdp file option (-1 means infinite, -2 means use mdp option, smaller is invalid)
//-maxh <real> (-1)
//Terminate after 0.99 times this time (hours)

void setMDArgs(std::vector<std::string>* mdargs, py::dict params)
{
    mdargs->clear();
    if (params.contains("grid"))
    {
        assert(py::len(params["grid"]) > 0);
        std::vector<std::string> vals;
        for (auto&& val : params["grid"])
        {
            vals.emplace_back(py::cast<std::string>(py::str(val)));
        }
        mdargs->emplace_back("-dd");
        for(auto&& val : vals)
        {
            mdargs->emplace_back(val);
        }
    }
    if (params.contains("pme_ranks"))
    {
        auto val = py::cast<std::string>(py::str(params["pme_ranks"]));
        mdargs->emplace_back("-npme");
        mdargs->emplace_back(val);
    }
    if (params.contains("threads"))
    {
        auto val = py::cast<std::string>(py::str(params["threads"]));
        mdargs->emplace_back("-nt");
        mdargs->emplace_back(val);
    }
    if (params.contains("tmpi"))
    {
        auto val = py::cast<std::string>(py::str(params["tmpi"]));
        mdargs->emplace_back("-ntmpi");
        mdargs->emplace_back(val);
    }
    if (params.contains("threads_per_rank"))
    {
        auto val = py::cast<std::string>(py::str(params["threads_per_rank"]));
        mdargs->emplace_back("-ntomp");
        mdargs->emplace_back(val);
    }
    if (params.contains("pme_threads_per_rank"))
    {
        auto val = py::cast<std::string>(py::str(params["pme_threads_per_rank"]));
        mdargs->emplace_back("-ntomp_pme");
        mdargs->emplace_back(val);
    }
    if (params.contains("steps"))
    {
        auto val = py::cast<std::string>(py::str(params["steps"]));
        mdargs->emplace_back("-nsteps");
        mdargs->emplace_back(val);
    }
    if (params.contains("max_hours"))
    {
        auto val = py::cast<std::string>(py::str(params["max_hours"]));
        mdargs->emplace_back("-maxh");
        mdargs->emplace_back(val);
    }
}

void export_context(py::module &m)
{
    using ::gmxapi::Context;

    using MDArgs = std::vector<std::string>;
    py::class_<MDArgs, std::unique_ptr<MDArgs>> mdargs(m, "MDArgs");
    mdargs.def(py::init(), "Create an empty MDArgs object.");
    mdargs.def("set", &setMDArgs, "Assign parameters in MDArgs from Python dict.");

    // Export execution context class
    py::class_<Context, std::shared_ptr<Context>> context(m, "Context");
    context.def(py::init(), "Create a default execution context.");
    context.def("setMDArgs", &Context::setMDArgs, "Set MD runtime parameters.");

    // During the registration of the gmx.core.Context Python type, perform appropriate environment initialization
    // and deinitialize at module destruction.
    gmx::init(nullptr, nullptr);
    m.add_object("_current_context", py::capsule([](){ gmx::finalize(); std::cout << "Shutting down GROMACS" << std::endl; }));
}

} // end namespace gmxpy::detail

} // end namespace gmxpy