Merge pull request #341 from NWChemEx/py_module_added_inputs

Additional Inputs in Python modules

docs/source/tutorials/modules/basics.rst

@@ -76,6 +76,18 @@ module to aid in defining the module's constructor. At a minimum the constructor
 must set the property type(s) that the module satisfies. It is also a good idea
 to provide a short description about the algorithm implemented in the module.
 
+.. hint::
+
+   The property type will be needed in a few places so it is a good idea to set
+   a typedef, this way if you need to change the property type for any reason
+   you only need to change it in the typedef.
+
+.. hint::
+
+   The raw string literal introduced in C++11 denoted by
+   ``auto str = R"(...)";`` is a great way to write a lengthy description
+   without having to use escapes or other workarounds for line endings.
+
 For our ``CoulombsLaw`` module a basic constructor looks like:
 
 .. tabs::
@@ -84,7 +96,7 @@ For our ``CoulombsLaw`` module a basic constructor looks like:
 
       .. literalinclude:: ../../../../tests/cxx/doc_snippets/coulombs_law.cpp
          :language: c++
-         :lines: 36-39
+         :lines: 21-33, 36-39
 
    .. tab:: Python
 
@@ -105,17 +117,28 @@ version. A full list of meta-data is available (TODO: Add link).
    PluginPlay automatically documents input parameters and submodules so there
    is no need to include these in your description.
 
-.. hint::
 
-   The property type will be needed in a few places so it is a good idea to set
-   a typedef, this way if you need to change the property type for any reason
-   you only need to change it in the typedef.
+For some modules, inputs that are not defined by the property type(s)
+may be needed. Additional inputs can be specified in the constructor with the
+``add_inputs`` method. A similar option for added results exist, though it is
+less common. Additional inputs have to be set before the module is run, either
+with a default value or by calling the ``change_input`` method. Here are 
+examples for adding a screening threshold to our ``ScreenedCoulombsLaw`` module:
 
-.. hint::
+.. tabs::
+
+   .. tab:: C++
+
+      .. literalinclude:: ../../../../tests/cxx/doc_snippets/screened_coulombs_law.cpp
+         :language: c++
+         :lines: 39-46
+
+   .. tab:: Python
+
+      .. literalinclude:: ../../../../tests/python/doc_snippets/coulombslaw_force.py
+         :language: python
+         :lines: 47-51
 
-   The raw string literal introduced in C++11 denoted by
-   ``auto str = R"(...)";`` is a great way to write a lengthy description
-   without having to use escapes or other workarounds for line endings.
 
 
 Defining the Run Member
@@ -155,8 +178,29 @@ vein the results of every module are also type-erased and every property type
 defines a static function ``wrap_results`` which takes the typed results and
 returns a result map with the type-erased results.
 
-The full definition of the ``run_`` member (including the source code for
-computing the electric field) is:
+.. note::
+
+   If a module has additional inputs beyond those specified by the property
+   type, they will have to be specifically unpacked from the ``inputs``. In the
+   case of the ``ScreenedCoulombsLaw`` module, the screening threshold is 
+   acquired as:
+
+   .. tabs::
+
+    .. tab:: C++
+
+        .. literalinclude:: ../../../../tests/cxx/doc_snippets/screened_coulombs_law.cpp
+            :language: c++
+            :lines: 48-50
+
+    .. tab:: Python
+
+        .. literalinclude:: ../../../../tests/python/doc_snippets/coulombslaw_force.py
+            :language: python
+            :lines: 53-56
+
+The full definition of the ``run_`` member for the ``CoulombsLaw`` module 
+(including the source code for computing the electric field) is:
 
 .. tabs::
 

src/python/export_module_base.cpp

@@ -45,27 +45,33 @@ void export_module_base(py_module_reference m) {
       //   .def("reset_internal_cache", &ModuleBase::reset_internal_cache)
       .def("description", &Publicist::description)
       .def("citation", &Publicist::citation)
-      .def("add_input",
-           [](ModuleBase& self, std::string key) {
-               auto& new_i = self.inputs()[key];
-               new_i.set_type<python::PythonWrapper>();
-               return new_i;
-           })
-      .def("add_result",
-           [](ModuleBase& self, std::string key) {
-               auto& new_r = self.results()[key];
-               new_r.set_type<python::PythonWrapper>();
-               return new_r;
-           })
-      .def("add_submodule",
-           [](ModuleBase& self, pybind11::object pt, std::string key) {
-               auto py_sr = pybind11::cast(SubmoduleRequest());
-               py_sr.attr("set_type")(pt);
-               auto sr     = py_sr.cast<SubmoduleRequest>();
-               auto& smods = self.submods();
-               smods.emplace(key, std::move(sr));
-               return smods.at(key);
-           })
+      .def(
+        "add_input",
+        [](ModuleBase& self, std::string key) {
+            auto& new_i = self.inputs()[key];
+            new_i.set_type<python::PythonWrapper>();
+            return &new_i;
+        },
+        pybind11::return_value_policy::reference)
+      .def(
+        "add_result",
+        [](ModuleBase& self, std::string key) {
+            auto& new_r = self.results()[key];
+            new_r.set_type<python::PythonWrapper>();
+            return &new_r;
+        },
+        pybind11::return_value_policy::reference)
+      .def(
+        "add_submodule",
+        [](ModuleBase& self, pybind11::object pt, std::string key) {
+            auto py_sr = pybind11::cast(SubmoduleRequest());
+            py_sr.attr("set_type")(pt);
+            auto sr     = py_sr.cast<SubmoduleRequest>();
+            auto& smods = self.submods();
+            smods.emplace(key, std::move(sr));
+            return &smods.at(key);
+        },
+        pybind11::return_value_policy::reference)
       .def("satisfies_property_type",
            [](ModuleBase& self, pybind11::object pt) {
                auto info    = pt.attr("type")().cast<python::PyTypeInfo>();

src/python/fields/export_module_input.cpp

@@ -35,15 +35,20 @@ void export_module_input(py_module_reference m) {
            [](ModuleInput& i, pybind11::object o) {
                return i.is_valid(any::make_any_field<PythonWrapper>(o));
            })
-      .def("change",
-           [](ModuleInput& i, pybind11::object o) {
-               i.change(any::make_any_field<PythonWrapper>(o));
-               return i;
-           })
-      .def("set_default",
-           [](ModuleInput& i, pybind11::object o) {
-               return i.set_default(any::make_any_field<PythonWrapper>(o));
-           })
+      .def(
+        "change",
+        [](ModuleInput& i, pybind11::object o) {
+            i.change(any::make_any_field<PythonWrapper>(o));
+            return &i;
+        },
+        pybind11::return_value_policy::reference)
+      .def(
+        "set_default",
+        [](ModuleInput& i, pybind11::object o) {
+            i.set_default(any::make_any_field<PythonWrapper>(o));
+            return &i;
+        },
+        pybind11::return_value_policy::reference)
       .def("set_description", &ModuleInput::set_description)
       //.def("add_check")
       .def("make_optional", &ModuleInput::make_optional)

tests/python/doc_snippets/coulombslaw_force.py

@@ -14,7 +14,7 @@
 
 import pluginplay as pp
 import pluginplay_examples as ppe
-from math import sqrt
+from math import sqrt, inf
 
 
 class CoulombsLaw(pp.ModuleBase):
@@ -43,6 +43,36 @@ def run_(self, inputs, submods):
         return pt.wrap_results(rv, E)
 
 
+class ScreenedCoulombsLaw(pp.ModuleBase):
+
+    def __init__(self):
+        pp.ModuleBase.__init__(self)
+        self.description("Screened Electric Field From Coulomb's Law")
+        self.satisfies_property_type(ppe.ElectricField())
+        self.add_input("threshold").set_default(inf)
+
+    def run_(self, inputs, submods):
+        pt = ppe.ElectricField()
+        [r, charges] = pt.unwrap_inputs(inputs)
+        thresh = inputs["threshold"].value()
+        E = [0.0, 0.0, 0.0]
+        for charge in charges:
+            q = charge.m_charge
+            ri = charge.m_r
+            rij = [i for i in r]
+            for i in range(3):
+                rij[i] -= charge.m_r[i]
+            rij2 = rij[0]**2 + rij[1]**2 + rij[2]**2
+            if sqrt(rij2) >= thresh:
+                continue
+            ri2 = ri[0]**2 + ri[1]**2 + ri[2]**2
+            mag_ri = sqrt(ri2)
+            for i in range(3):
+                E[i] += (q * ri[i] / (mag_ri * rij2))
+        rv = self.results()
+        return pt.wrap_results(rv, E)
+
+
 class ClassicalForce(pp.ModuleBase):
 
     def __init__(self):

tests/python/doc_snippets/test_python_modules.py

@@ -32,10 +32,22 @@ def test_modules(self):
         mm = pp.ModuleManager()
 
         mm.add_module("My Coulomb's Law", clf.CoulombsLaw())
+        mm.add_module("My Screened Coulomb's Law 1", clf.ScreenedCoulombsLaw())
+        mm.add_module("My Screened Coulomb's Law 2", clf.ScreenedCoulombsLaw())
         mm.add_module("My Force", clf.ClassicalForce())
         mm.change_submod("My Force", "electric field", "My Coulomb's Law")
 
-        field = mm.at("My Coulomb's Law").run_as(ppe.ElectricField(), r, pvc)
-        self.assertTrue(field == [1.5, 0.0, 0.0])
+        efield = ppe.ElectricField()
+
+        field0 = mm.at("My Coulomb's Law").run_as(efield, r, pvc)
+        self.assertTrue(field0 == [1.5, 0.0, 0.0])
+
+        field1 = mm.at("My Screened Coulomb's Law 1").run_as(efield, r, pvc)
+        self.assertTrue(field1 == [1.5, 0.0, 0.0])
+
+        mm.change_input("My Screened Coulomb's Law 2", "threshold", 1.0)
+        field2 = mm.at("My Screened Coulomb's Law 2").run_as(efield, r, pvc)
+        self.assertTrue(field2 == [0.0, 0.0, 0.0])
+
         cforce = mm.at("My Force").run_as(ppe.Force(), q, m, a, pvc)
         self.assertTrue(cforce == [5.5, 0.0, 0.0])