Density on grid (#49)

* g2g seems to work

* backup

* density2grid works (I think)

* fix warnings

Author: ryanmrichard

src/scf/xc/density2grid.cpp

@@ -0,0 +1,96 @@
+/*
+ * Copyright 2025 NWChemEx-Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "xc.hpp"
+#include <simde/integration_grids/collocation_matrix.hpp>
+
+namespace scf::xc {
+namespace {
+const auto desc = R"(
+
+DensityCollocationMatrix
+-----------------
+)";
+
+struct Kernel {
+    using buffer_base = tensorwrapper::buffer::BufferBase;
+
+    template<typename FloatType>
+    auto run(const buffer_base& aos_on_grid, const buffer_base& X,
+             parallelzone::runtime::RuntimeView& rv) {
+        tensorwrapper::allocator::Eigen<FloatType> allocator(rv);
+
+        const auto& eigen_aos_on_grid = allocator.rebind(aos_on_grid);
+        const auto* paos_on_grid      = eigen_aos_on_grid.get_immutable_data();
+        const auto& eigen_X           = allocator.rebind(X);
+        const auto* pX                = eigen_X.get_immutable_data();
+        const auto& shape_X           = eigen_X.layout().shape().as_smooth();
+        auto n_aos                    = shape_X.extent(0);
+        auto n_grid                   = shape_X.extent(1);
+
+        tensorwrapper::shape::Smooth rv_shape{n_grid};
+        tensorwrapper::layout::Physical rv_layout(rv_shape);
+        auto rv_buffer = allocator.allocate(rv_layout);
+
+        // AOs on rows, grid points on columns
+        for(std::size_t grid_i = 0; grid_i < n_grid; ++grid_i) {
+            FloatType sum = 0;
+            for(std::size_t ao_i = 0; ao_i < n_aos; ++ao_i) {
+                const auto idx = ao_i * n_grid + grid_i;
+                sum += paos_on_grid[idx] * pX[idx];
+            }
+            rv_buffer->set_elem(std::vector{grid_i}, sum);
+        }
+        return simde::type::tensor(rv_shape, std::move(rv_buffer));
+    }
+};
+} // namespace
+
+using pt         = simde::EDensityCollocationMatrix;
+using ao2grid_pt = simde::AOCollocationMatrix;
+
+MODULE_CTOR(Density2Grid) {
+    satisfies_property_type<pt>();
+    description(desc);
+
+    add_submodule<ao2grid_pt>("AOs on a grid");
+}
+
+MODULE_RUN(Density2Grid) {
+    const auto& [grid, density] = pt::unwrap_inputs(inputs);
+
+    const auto& rho = density.value();
+    const auto& aos = density.basis_set().ao_basis_set();
+
+    auto& ao2grid_mod = submods.at("AOs on a grid");
+    auto aos_on_grid  = ao2grid_mod.run_as<ao2grid_pt>(grid, aos);
+
+    simde::type::tensor X;
+    X("m,i") = rho("m,n") * aos_on_grid("n,i");
+
+    using tensorwrapper::utilities::floating_point_dispatch;
+    Kernel k;
+    auto& runtime          = get_runtime();
+    const auto& aos_buffer = aos_on_grid.buffer();
+    const auto& X_buffer   = X.buffer();
+    auto rho_on_grid =
+      floating_point_dispatch(k, aos_buffer, X_buffer, runtime);
+
+    auto rv = results();
+    return pt::wrap_results(rv, std::move(rho_on_grid));
+}
+
+} // namespace scf::xc

src/scf/xc/gau2grid.cpp

@@ -40,7 +40,7 @@ std::vector<T> flatten_grid(const chemist::Grid& grid) {
 
 } // namespace
 
-using pt = simde::CollocationMatrix;
+using pt = simde::AOCollocationMatrix;
 
 MODULE_CTOR(Gau2Grid) {
     satisfies_property_type<pt>();

src/scf/xc/xc.cpp

@@ -22,6 +22,7 @@ void load_modules(pluginplay::ModuleManager& mm) {
     gauxc::load_modules(mm);
     mm.add_module<Gau2Grid>("Gau2Grid");
     mm.add_module<GridFromFile>("Grid From File");
+    mm.add_module<Density2Grid>("Density2Grid");
 }
 
 void set_defaults(pluginplay::ModuleManager& mm) { gauxc::set_defaults(mm); }

src/scf/xc/xc.hpp

@@ -19,6 +19,7 @@
 
 namespace scf::xc {
 
+DECLARE_MODULE(Density2Grid);
 DECLARE_MODULE(Gau2Grid);
 DECLARE_MODULE(GridFromFile);
 

tests/cxx/unit_tests/xc/density2grid.cpp

@@ -0,0 +1,48 @@
+/*
+ * Copyright 2025 NWChemEx-Project
+ *
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ */
+
+#include "../../test_scf.hpp"
+#include <iostream>
+#include <pluginplay/pluginplay.hpp>
+
+using namespace scf;
+
+using pt    = simde::EDensityCollocationMatrix;
+using ao_pt = simde::AOCollocationMatrix;
+using tensorwrapper::operations::approximately_equal;
+
+TEST_CASE("Density2Grid") {
+    pluginplay::ModuleManager mm;
+    scf::load_modules(mm);
+    auto& mod = mm.at("Density2Grid");
+
+    chemist::Grid grid; // Value doesn't matter for this test b/c of submodule
+
+    SECTION("H2") {
+        auto rho = test_scf::h2_density<double>();
+        simde::type::tensor ao2grid{{1.0, 2.0, 3.0, 4.0}, {5.0, 6.0, 7.0, 8.0}};
+        auto ao_mod =
+          pluginplay::make_lambda<ao_pt>([&](auto&& grid_in, auto&& aos) {
+              REQUIRE(grid_in == grid);
+              REQUIRE(aos == rho.basis_set().ao_basis_set());
+              return ao2grid;
+          });
+        mod.change_submod("AOs on a grid", ao_mod);
+        auto rv = mod.run_as<pt>(grid, rho);
+        simde::type::tensor corr{11.513088, 20.467712, 31.9808, 46.052352};
+        REQUIRE(approximately_equal(rv, corr, 1e-4));
+    }
+}

tests/cxx/unit_tests/xc/gau2grid.cpp

@@ -20,7 +20,7 @@
 
 using namespace scf;
 
-using pt = simde::CollocationMatrix;
+using pt = simde::AOCollocationMatrix;
 
 using tensorwrapper::operations::approximately_equal;