wfc extrapolation / MD

ci/baseimage.cuda.Dockerfile

@@ -67,7 +67,7 @@ RUN spack install openblas %gcc +fortran
 
 RUN spack install magma %gcc +cuda +fortran ^openblas
 
-RUN spack install nlcglib@master %gcc +cuda
+RUN spack install nlcglib@develop %gcc +cuda
 
 # for the MPI hook
 RUN echo $(spack find --format='{prefix.lib}' mpich) > /etc/ld.so.conf.d/mpich.conf

ci/cscs-daint.yml

@@ -31,6 +31,7 @@ build cuda image mkl:
   extends: .container-builder
   needs: ["build base cuda image"]
   stage: build
+  timeout: 4h
   variables:
     DOCKERFILE: ci/build.Dockerfile
     PERSIST_IMAGE_NAME: discard
@@ -41,6 +42,7 @@ build cuda image:
   extends: .container-builder
   needs: ["build base cuda image"]
   stage: build
+  timeout: 4h
   variables:
     CSCS_REBUILD_POLICY: always
     DOCKERFILE: ci/build.Dockerfile
@@ -52,6 +54,7 @@ build rocm image:
   extends: .container-builder
   needs: ["build base rocm image"]
   stage: build
+  timeout: 4h
   variables:
     DOCKERFILE: ci/build.Dockerfile
     PERSIST_IMAGE_NAME: discard
@@ -62,6 +65,7 @@ build elpa cuda image:
   extends: .container-builder
   needs: ["build base cuda image"]
   stage: build
+  timeout: 4h
   variables:
     DOCKERFILE: ci/build.Dockerfile
     PERSIST_IMAGE_NAME: discard
@@ -72,6 +76,7 @@ build sequential eigen-solver cuda image:
   extends: .container-builder
   needs: ["build base cuda image"]
   stage: build
+  timeout: 4h
   variables:
     DOCKERFILE: ci/build.Dockerfile
     PERSIST_IMAGE_NAME: discard
@@ -82,6 +87,7 @@ build fp32 cuda image:
   extends: .container-builder
   needs: ["build base cuda image"]
   stage: build
+  timeout: 4h
   variables:
     DOCKERFILE: ci/build.Dockerfile
     PERSIST_IMAGE_NAME: discard
@@ -92,6 +98,7 @@ build vdwxc cuda image:
   extends: .container-builder
   needs: ["build base cuda image"]
   stage: build
+  timeout: 4h
   variables:
     DOCKERFILE: ci/build.Dockerfile
     PERSIST_IMAGE_NAME: discard
@@ -103,6 +110,7 @@ build nlcg image:
   extends: .container-builder
   needs: ["build base cuda image"]
   stage: build
+  timeout: 4h
   variables:
     DOCKERFILE: ci/build.Dockerfile
     PERSIST_IMAGE_NAME: discard
@@ -113,6 +121,7 @@ build clang image:
   extends: .container-builder
   needs: ["build base cuda image"]
   stage: build
+  timeout: 4h
   variables:
     DOCKERFILE: ci/build.Dockerfile
     PERSIST_IMAGE_NAME: discard
@@ -124,6 +133,7 @@ build openmpi image:
   extends: .container-builder
   needs: ["build base cuda image"]
   stage: build
+  timeout: 4h
   variables:
     DOCKERFILE: ci/build.Dockerfile
     PERSIST_IMAGE_NAME: discard

python_module/py_sirius.cpp

@@ -6,7 +6,8 @@
  * SPDX-License-Identifier: BSD-3-Clause
  */
 
-#include "k_point/k_point_set.hpp"
+#include "context/simulation_context.hpp"
+#include "k_point/k_point.hpp"
 #include "python_module_includes.hpp"
 #include <mpi.h>
 
@@ -216,6 +217,7 @@ PYBIND11_MODULE(py_sirius, m)
             .def("update", &Simulation_context::update)
             .def("use_symmetry", py::overload_cast<>(&Simulation_context::use_symmetry, py::const_))
             .def("processing_unit_memory_t", &Simulation_context::processing_unit_memory_t)
+            .def("spla_context", &Simulation_context::spla_context_ptr)
             .def(
                     "comm", [](Simulation_context& obj) { return make_pycomm(obj.comm()); },
                     py::return_value_policy::reference_internal)
@@ -403,8 +405,7 @@ PYBIND11_MODULE(py_sirius, m)
             .def("fv_states", &K_point<double>::fv_states, py::return_value_policy::reference_internal)
             .def("ctx", &K_point<double>::ctx, py::return_value_policy::reference_internal)
             .def("weight", &K_point<double>::weight)
-            .def("beta_projectors", py::overload_cast<>(&K_point<double>::beta_projectors, py::const_),
-                 py::return_value_policy::reference_internal)
+            .def("beta_projectors", py::overload_cast<>(&K_point<double>::beta_projectors_ptr, py::const_))
             .def("spinor_wave_functions", py::overload_cast<>(&K_point<double>::spinor_wave_functions),
                  py::return_value_policy::reference_internal);
 
@@ -415,7 +416,8 @@ PYBIND11_MODULE(py_sirius, m)
             .def(py::init<Simulation_context&, r3::vector<int>, r3::vector<int>, bool>(), py::keep_alive<1, 2>())
             .def(py::init<Simulation_context&, std::vector<int>, std::vector<int>, bool>(), py::keep_alive<1, 2>())
             .def("initialize", &K_point_set::initialize, py::arg("counts") = std::vector<int>{})
-            .def("ctx", &K_point_set::ctx, py::return_value_policy::reference_internal)
+            .def("ctx", static_cast<Simulation_context const& (K_point_set::*)() const>(&K_point_set::ctx),
+                 py::return_value_policy::reference_internal)
             .def("unit_cell", &K_point_set::unit_cell, py::return_value_policy::reference_internal)
             .def("_num_kpoints", &K_point_set::num_kpoints)
             .def("size", [](K_point_set& ks) -> int { return ks.spl_num_kpoints().local_size(); })

python_module/py_sirius_operators.cpp

@@ -7,6 +7,7 @@
  */
 
 #include "python_module_includes.hpp"
+#include <memory>
 #include <string>
 #include <iomanip>
 #include <complex>
@@ -25,16 +26,16 @@ using PT = double;
 void
 init_operators(py::module& m)
 {
-    py::class_<Non_local_operator<PT>>(m, "Non_local_operator")
+    py::class_<Non_local_operator<PT>, std::shared_ptr<Non_local_operator<PT>>>(m, "Non_local_operator")
             .def("get_matrix", &Non_local_operator<PT>::get_matrix<std::complex<double>>);
-    py::class_<D_operator<PT>, Non_local_operator<PT>>(m, "D_operator");
-    py::class_<Q_operator<PT>, Non_local_operator<PT>>(m, "Q_operator");
+    py::class_<D_operator<PT>, Non_local_operator<PT>, std::shared_ptr<D_operator<PT>>>(m, "D_operator");
+    py::class_<Q_operator<PT>, Non_local_operator<PT>, std::shared_ptr<Q_operator<PT>>>(m, "Q_operator");
 
     py::class_<Hamiltonian0<PT>>(m, "Hamiltonian0")
             .def(py::init<Potential&, bool>(), py::keep_alive<1, 2>(), "Potential"_a,
                  py::arg("precompute_lapw") = false)
-            .def("Q", &Hamiltonian0<PT>::Q, py::return_value_policy::reference_internal)
-            .def("D", &Hamiltonian0<PT>::D, py::return_value_policy::reference_internal)
+            .def("Q", &Hamiltonian0<PT>::Q_ptr)
+            .def("D", &Hamiltonian0<PT>::D_ptr)
             .def("Hk", &Hamiltonian0<PT>::operator(), py::keep_alive<0, 1>())
             .def("potential", &Hamiltonian0<PT>::potential, py::return_value_policy::reference_internal);
 
@@ -46,8 +47,9 @@ init_operators(py::module& m)
             .def_property_readonly("U", &Hamiltonian_k<PT>::U, py::return_value_policy::reference_internal);
 
     py::class_<S_k<complex_double>>(m, "S_k")
-            .def(py::init<Simulation_context&, const Q_operator<PT>&, const Beta_projectors_base<PT>&, int>(),
-                 py::keep_alive<1, 2>(), py::keep_alive<1, 3>(), py::keep_alive<1, 4>())
+            .def(py::init<memory_t, std::shared_ptr<spla::Context>, std::shared_ptr<Q_operator<PT> const>,
+                          std::shared_ptr<Beta_projectors_base<PT> const>, int>(),
+                 py::keep_alive<1, 2>())
             .def_property_readonly("size", &S_k<complex_double>::size)
             .def("apply", [](py::object& obj, py::array_t<complex_double>& X) {
                 using class_t = S_k<complex_double>;
@@ -71,9 +73,11 @@ init_operators(py::module& m)
                 return sk.apply(array, memory_t::host);
             });
 
+    py::class_<spla::Context, std::shared_ptr<spla::Context>>(m, "splaContext");
+
     py::class_<InverseS_k<complex_double>>(m, "InverseS_k")
-            .def(py::init<Simulation_context&, const Q_operator<PT>&, const Beta_projectors_base<PT>&, int>(),
-                 py::keep_alive<1, 2>(), py::keep_alive<1, 3>(), py::keep_alive<1, 4>())
+            .def(py::init<memory_t, std::shared_ptr<spla::Context>, std::shared_ptr<Q_operator<PT> const>,
+                          std::shared_ptr<Beta_projectors_base<PT> const>, int>())
             .def_property_readonly("size", &InverseS_k<complex_double>::size)
             .def("apply", [](py::object& obj, py::array_t<complex_double>& X) {
                 using class_t       = InverseS_k<complex_double>;
@@ -98,9 +102,9 @@ init_operators(py::module& m)
             });
 
     py::class_<Ultrasoft_preconditioner<complex_double>>(m, "Precond_us")
-            .def(py::init<Simulation_context&, const Q_operator<PT>&, int, const Beta_projectors_base<PT>&,
-                          const fft::Gvec&>(),
-                 py::keep_alive<1, 2>(), py::keep_alive<1, 3>(), py::keep_alive<1, 5>(), py::keep_alive<1, 6>())
+            .def(py::init<Simulation_context&, std::shared_ptr<Q_operator<PT> const>, int,
+                          std::shared_ptr<Beta_projectors_base<PT> const>, const fft::Gvec&>(),
+                 py::keep_alive<1, 2>())
             .def_property_readonly("size", &Ultrasoft_preconditioner<complex_double>::size)
             .def("apply", [](py::object& obj, py::array_t<complex_double>& X) {
                 using class_t    = Ultrasoft_preconditioner<complex_double>;
@@ -163,12 +167,13 @@ init_operators(py::module& m)
             .def("generate_j", py::overload_cast<beta_projectors_coeffs_t<PT>&, int, int>(
                                        &Beta_projector_generator<PT>::generate, py::const_));
 
-    py::class_<Beta_projectors_base<PT>>(m, "Beta_projectors_base")
+    py::class_<Beta_projectors_base<PT>, std::shared_ptr<Beta_projectors_base<PT>>>(m, "Beta_projectors_base")
             .def_property_readonly("num_chunks", &Beta_projectors_base<PT>::num_chunks)
             .def("make_generator", py::overload_cast<device_t>(&Beta_projectors_base<PT>::make_generator, py::const_),
                  py::keep_alive<1, 0>());
 
-    py::class_<Beta_projectors<PT>, Beta_projectors_base<PT>>(m, "Beta_projectors");
+    py::class_<Beta_projectors<PT>, Beta_projectors_base<PT>, std::shared_ptr<Beta_projectors<PT>>>(m,
+                                                                                                    "Beta_projectors");
     m.def("apply_U_operator", &apply_U_operator<double>, py::arg("ctx"), py::arg("spin_range"), py::arg("band_range"),
           py::arg("hub_wf"), py::arg("phi"), py::arg("u_op"), py::arg("hphi"));
 

python_module/sirius/coefficient_array/coefficient_array.py

@@ -231,7 +231,9 @@ def sum(self, **kwargs):
         """ """
         loc_sum = np.array(sum([np.sum(v, **kwargs) for _, v in self.items()]))
         reduced = MPI.COMM_WORLD.allreduce(loc_sum, op=MPI.SUM)
-        return np.ndarray.item(reduced)
+        if isinstance(reduced, np.ndarray):
+            return np.ndarray.item(reduced)
+        return reduced
 
     def log(self, **kwargs):
         out = type(self)()

python_module/sirius/edft/neugebauer.py

@@ -6,7 +6,7 @@
 import numpy as np
 from scipy.constants import physical_constants
 
-from ..coefficient_array import diag, inner, l2norm
+from ..coefficient_array import diag, inner, l2norm, identity_like
 from ..operators import US_Precond, Sinv_operator, S_operator
 from ..logger import Logger
 from ..py_sirius import sprint_magnetization
@@ -229,7 +229,6 @@ def run(
         tol=1e-10,
         kappa=0.3,
         tau=0.5,
-        cgtype="FR",
         K=IdentityPreconditioner(),
         callback=lambda *args, **kwargs: None,
         error_callback=lambda *args, **kwargs: None,
@@ -242,15 +241,6 @@ def run(
         is_converged -- bool
         """
 
-        if cgtype == "PR":
-            cg_update = polak_ribiere
-        elif cgtype == "FR":
-            cg_update = fletcher_reeves
-        elif cgtype == "SD":
-            cg_update = steepest_descent
-        else:
-            raise ValueError("wrong type")
-
         if self.is_ultrasoft:
             K = self.K
             # TODO cleanup signature of run and don't pass the preconditioner anymore
@@ -265,6 +255,7 @@ def run(
         X = X @ U
         # set occupation numbers from band energies
         fn, mu = self.M.smearing.fn(ek)
+
         # compute initial free energy
         FE, Hx = M(X, fn=fn, mu=mu, ek=ek)
         logger("initial F: %.13f" % FE)
@@ -278,6 +269,7 @@ def run(
             SX = self.S @ X
         else:
             SX = X
+
         XhKHX = SX.H @ (K @ HX)
         XhKX = SX.H @ (K @ SX)
         LL = _solve(XhKX, XhKHX)

python_module/sirius/edft/pseudo_hamiltonian.py

@@ -83,17 +83,15 @@ def __call__(self, t):
         X_new = self.X + t * self.G_X
         eta_new = self.eta + t * self.G_eta
         ek, Ul = eta_new.eigh()
-        X = loewdin(X_new, self.overlap) @ Ul
+        X_new = loewdin(X_new, self.overlap) @ Ul
         fn, mu = self.M.smearing.fn(ek)
 
-        # check fn
-
-        FE, Hx = self.M(X, fn=fn, mu=mu, ek=ek)
+        FE, Hx = self.M(X_new, fn=fn, mu=mu, ek=ek)
         return SimpleNamespace(
             **{
                 "free_energy": FE,
                 "Hx": Hx,
-                "X": X,
+                "X": X_new,
                 "fn": fn,
                 "mu": mu,
                 "ek": ek,

python_module/sirius/operators.py

@@ -1,4 +1,5 @@
-from .py_sirius import InverseS_k, S_k, Precond_us, Hamiltonian0
+from .py_sirius import InverseS_k, S_k, Precond_us, Hamiltonian0, MemoryEnum
+
 # from .coefficient_array import threaded
 from scipy.sparse.linalg import LinearOperator
 import numpy as np
@@ -24,7 +25,11 @@ def __init__(self, ctx, potential, kpointset):
         for ik, kp in enumerate(kpointset):
             for ispn in range(ctx.num_spins()):
                 self.ops[ik, ispn] = S_k(
-                    ctx, hamiltonian0.Q(), kp.beta_projectors(), ispn
+                    MemoryEnum.host,
+                    ctx.spla_context(),
+                    hamiltonian0.Q(),
+                    kp.beta_projectors(),
+                    ispn,
                 )
 
     def apply(self, cn):
@@ -37,8 +42,11 @@ def apply(self, cn):
     def __getitem__(self, key):
         def _matvec(X):
             return np.array(self.ops[key].apply(np.asfortranarray(X)))
+
         n = self.ops[key].size
-        return LinearOperator(dtype=np.complex128, shape=(n, n), matvec=_matvec, rmatvec=_matvec)
+        return LinearOperator(
+            dtype=np.complex128, shape=(n, n), matvec=_matvec, rmatvec=_matvec
+        )
 
     def __matmul__(self, cn):
         return self.apply(cn)
@@ -53,7 +61,11 @@ def __init__(self, ctx, potential, kpointset):
         for ik, kp in enumerate(kpointset):
             for ispn in range(ctx.num_spins()):
                 self.ops[ik, ispn] = InverseS_k(
-                    ctx, hamiltonian0.Q(), kp.beta_projectors(), ispn
+                    MemoryEnum.host,
+                    ctx.spla_context(),
+                    hamiltonian0.Q(),
+                    kp.beta_projectors(),
+                    ispn,
                 )
 
     def apply(self, cn):
@@ -65,8 +77,11 @@ def apply(self, cn):
     def __getitem__(self, key):
         def _matvec(X):
             return np.array(self.ops[key].apply(np.asfortranarray(X)))
+
         n = self.ops[key].size
-        return LinearOperator(dtype=np.complex128, shape=(n, n), matvec=_matvec, rmatvec=_matvec)
+        return LinearOperator(
+            dtype=np.complex128, shape=(n, n), matvec=_matvec, rmatvec=_matvec
+        )
 
     def __matmul__(self, cn):
         return self.apply(cn)
@@ -99,8 +114,11 @@ def apply(self, cn):
     def __getitem__(self, key):
         def _matvec(X):
             return np.array(self.ops[key].apply(np.asfortranarray(X)))
+
         n = self.ops[key].size
-        return LinearOperator(dtype=np.complex128, shape=(n, n), matvec=_matvec, rmatvec=_matvec)
+        return LinearOperator(
+            dtype=np.complex128, shape=(n, n), matvec=_matvec, rmatvec=_matvec
+        )
 
     def __matmul__(self, cn):
         return self.apply(cn)

python_module/test/test_mvp2_gradient.py

@@ -26,7 +26,7 @@ def make_new_ctx(pw_cutoff, gk_cutoff):
                 "pw_cutoff": pw_cutoff,
                 "gk_cutoff": gk_cutoff,
             },
-            "control": {"verbosity": 0},
+            "control": {"verbosity": 0, "processing_unit": "cpu"},
         }
         # create simulation context
         ctx = Simulation_context(json.dumps(inp))