Fix OneDFT gradient task ordering and MPI distribution

src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_onedft.hpp

@@ -859,15 +859,35 @@ eval_exc_grad_onedft_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
     auto reorder_result = mpi_gather_and_reorder_gpu(
       den_eval, dden_eval, tau, grid_coords, grid_weights,
       atomic_grid_sizes_vec, total_npts, natoms, rt, recvcounts, displs);
+    const int total_npts_sum = reorder_result.total_npts;
     atom_reorder_inv_perm = std::move(reorder_result.inv_perm);
 
+    GAUXC_MPI_CODE(
+      std::vector<double> recv_raw(world_rank == 0 ? total_npts_sum : 0);
+      MPI_Gatherv(raw_grid_weights.data(), static_cast<int>(total_npts), MPI_DOUBLE,
+                  recv_raw.data(), recvcounts.data(), displs.data(),
+                  MPI_DOUBLE, 0, rt.comm());
+      if (world_rank == 0) {
+        raw_grid_weights = std::move(recv_raw);
+        std::vector<int64_t> perm(total_npts_sum);
+        for (int64_t j = 0; j < total_npts_sum; ++j) {
+          perm[atom_reorder_inv_perm[j]] = j;
+        }
+        std::vector<double> reordered_raw(total_npts_sum);
+        for (int64_t i = 0; i < total_npts_sum; ++i) {
+          reordered_raw[perm[i]] = raw_grid_weights[i];
+        }
+        raw_grid_weights = std::move(reordered_raw);
+      }
+    )
+
     if (world_rank == 0) {
       int64_t max_grid_size = *std::max_element(
         reorder_result.global_atomic_grid_sizes.begin(),
         reorder_result.global_atomic_grid_sizes.end());
       auto options = torch::TensorOptions().dtype(torch::kFloat64).device(torch::kCPU);
       features_dict = prepare_onedft_features(
-        natoms, reorder_result.total_npts, ndm, options, feature_keys,
+        natoms, total_npts_sum, ndm, options, feature_keys,
         den_eval.data(), dden_eval.data(), tau.data(),
         grid_coords.data(), grid_weights.data(), host_coords.data(),
         reorder_result.global_atomic_grid_sizes, max_grid_size,
@@ -969,17 +989,44 @@ eval_exc_grad_onedft_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
       den_eval.data(), dden_eval.data(), tau.data());
   }
 
-  // Scatter eps_on_grid to local tasks
+  // Scatter rank-ordered point values back to the local atom-ordered tasks.
   std::vector<double> eps_on_grid_local;
+  std::vector<double> points_grad_local;
   if (world_size == 1) {
     eps_on_grid_local = std::move(eps_on_grid_global);
+    points_grad_local = std::move(points_grad_global);
   } else {
-    GAUXC_GENERIC_EXCEPTION("OneDFT gradient with MPI not yet implemented");
+    eps_on_grid_local.resize(total_npts);
+    GAUXC_MPI_CODE(
+      MPI_Scatterv(eps_on_grid_global.data(), recvcounts.data(), displs.data(),
+                   MPI_DOUBLE, eps_on_grid_local.data(), total_npts, MPI_DOUBLE,
+                   0, rt.comm());
+
+      int has_points_grad = world_rank == 0 && !points_grad_global.empty();
+      MPI_Bcast(&has_points_grad, 1, MPI_INT, 0, rt.comm());
+      if (has_points_grad) {
+        std::vector<int> recvcounts3(recvcounts.size());
+        std::vector<int> displs3(displs.size());
+        for (size_t i = 0; i < recvcounts.size(); ++i) {
+          recvcounts3[i] = 3 * recvcounts[i];
+          displs3[i] = 3 * displs[i];
+        }
+        points_grad_local.resize(3 * total_npts);
+        MPI_Scatterv(points_grad_global.data(), recvcounts3.data(), displs3.data(),
+                     MPI_DOUBLE, points_grad_local.data(), 3 * total_npts,
+                     MPI_DOUBLE, 0, rt.comm());
+      }
+    )
   }
 
   // Zero out EXC_GRAD on host
   for (int i = 0; i < 3*natoms; ++i) EXC_GRAD[i] = 0.0;
 
+  if (!points_grad_local.empty() &&
+      points_grad_local.size() != 3 * eps_on_grid_local.size()) {
+    GAUXC_GENERIC_EXCEPTION("Inconsistent OneDFT point-gradient layout");
+  }
+
   // Phase 5: Pulay gradient + weight derivative on device
   // Use the standard gradient flow: for each batch, load OneDFT Vxc,
   // evaluate collocation (hessian), X-matrix, transform, and inc_exc_grad
@@ -1005,12 +1052,18 @@ eval_exc_grad_onedft_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
     size_t offset = 0;
     for (auto& task : tasks) {
       int64_t npts = task.points.size();
+      if (offset + npts > eps_on_grid_local.size()) {
+        GAUXC_GENERIC_EXCEPTION("Inconsistent OneDFT energy-density layout");
+      }
       task.feat.eps.resize(npts);
       std::copy(eps_on_grid_local.data() + offset,
                 eps_on_grid_local.data() + offset + npts,
                 task.feat.eps.begin());
       offset += npts;
     }
+    if (offset != eps_on_grid_local.size()) {
+      GAUXC_GENERIC_EXCEPTION("Inconsistent OneDFT energy-density layout");
+    }
   }
 
   // Normally task_comparator is used for sorting tasks, but for OneDFT
@@ -1121,17 +1174,17 @@ eval_exc_grad_onedft_( int64_t m, int64_t n, const value_type* Ps, int64_t ldps,
   rt.device_backend()->master_queue_synchronize();
 
   // Phase 6: Add autograd forces (points -> parent atoms, coords -> direct)
-  if (!points_grad_global.empty() && world_rank == 0) {
+  if (!points_grad_local.empty()) {
     size_t pg_offset = 0;
-    // Iterate in iParent order — re-sort tasks
+    // Iterate in the atom order used to assemble the local point arrays.
     std::stable_sort(tasks.begin(), tasks.end(),
       [](const auto& a, const auto& b) { return a.iParent < b.iParent; });
     for (const auto& task : tasks) {
       int iParent = task.iParent;
       for (size_t ipt = 0; ipt < task.points.size(); ++ipt) {
-        EXC_GRAD[3*iParent + 0] += points_grad_global[(pg_offset + ipt)*3 + 0];
-        EXC_GRAD[3*iParent + 1] += points_grad_global[(pg_offset + ipt)*3 + 1];
-        EXC_GRAD[3*iParent + 2] += points_grad_global[(pg_offset + ipt)*3 + 2];
+        EXC_GRAD[3*iParent + 0] += points_grad_local[(pg_offset + ipt)*3 + 0];
+        EXC_GRAD[3*iParent + 1] += points_grad_local[(pg_offset + ipt)*3 + 1];
+        EXC_GRAD[3*iParent + 2] += points_grad_local[(pg_offset + ipt)*3 + 2];
       }
       pg_offset += task.points.size();
     }

src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_onedft.hpp

@@ -1027,33 +1027,58 @@ void ReferenceReplicatedXCHostIntegrator<ValueType>::
   // Step 4: Scatter Vxc back to tasks
   send_buffer_onedft_outputs(2/*ndm*/, features_dict, tasks, rt, sendcounts, displs, atom_reorder_inv_perm);
 
-  // Scatter eps_on_grid to local tasks (for single rank, just distribute)
-  // For MPI, would need MPI_Scatterv — for now handle single rank
+  // Scatter rank-ordered point values back to the local atom-ordered tasks.
   std::vector<double> eps_on_grid_local;
+  std::vector<double> points_grad_local;
   if (rt.comm_size() == 1) {
     eps_on_grid_local = std::move(eps_on_grid_global);
+    points_grad_local = std::move(points_grad_global);
   } else {
-    // TODO: MPI scatter of eps_on_grid
-    GAUXC_GENERIC_EXCEPTION("OneDFT gradient with MPI not yet implemented");
+    const int local_npts = std::accumulate(
+      tasks.begin(), tasks.end(), 0,
+      [](const auto npts, const auto& task) { return npts + task.npts; });
+    eps_on_grid_local.resize(local_npts);
+    GAUXC_MPI_CODE(
+      MPI_Scatterv(eps_on_grid_global.data(), sendcounts.data(), displs.data(),
+                   MPI_DOUBLE, eps_on_grid_local.data(), local_npts, MPI_DOUBLE,
+                   0, rt.comm());
+
+      int has_points_grad = world_rank == 0 && !points_grad_global.empty();
+      MPI_Bcast(&has_points_grad, 1, MPI_INT, 0, rt.comm());
+      if (has_points_grad) {
+        std::vector<int> sendcounts3(sendcounts.size());
+        std::vector<int> displs3(displs.size());
+        for (size_t i = 0; i < sendcounts.size(); ++i) {
+          sendcounts3[i] = 3 * sendcounts[i];
+          displs3[i] = 3 * displs[i];
+        }
+        points_grad_local.resize(3 * local_npts);
+        MPI_Scatterv(points_grad_global.data(), sendcounts3.data(), displs3.data(),
+                     MPI_DOUBLE, points_grad_local.data(), 3 * local_npts,
+                     MPI_DOUBLE, 0, rt.comm());
+      }
+    )
   }
 
   // Zero out EXC_GRAD
   for (int i = 0; i < 3*natoms; ++i) EXC_GRAD[i] = 0.0;
 
-  // Step 6: Add autograd forces BEFORE Pulay (which re-sorts tasks!)
+  if (!points_grad_local.empty() &&
+      points_grad_local.size() != 3 * eps_on_grid_local.size()) {
+    GAUXC_GENERIC_EXCEPTION("Inconsistent OneDFT point-gradient layout");
+  }
+
+  // Step 6: Add autograd forces before the Pulay work reorders the tasks.
   // points.grad gives ∂E/∂r_g. Since grid points move with their parent atom,
-  // the force on atom A = Σ_{g∈A} points_grad[g].
-  // NOTE: Must be done while tasks are still in iParent-sorted order
-  //       (matching points_grad_global layout). exc_grad_local_work_onedft_
-  //       re-sorts tasks by workload, breaking the correspondence.
-  if (!points_grad_global.empty() && world_rank == 0) {
+  // the derivative for atom A contains Σ_{g∈A} points_grad[g].
+  if (!points_grad_local.empty()) {
     size_t offset = 0;
     for (const auto& task : tasks) {
       int iParent = task.iParent;
       for (size_t ipt = 0; ipt < task.points.size(); ++ipt) {
-        EXC_GRAD[3*iParent + 0] += points_grad_global[(offset + ipt)*3 + 0];
-        EXC_GRAD[3*iParent + 1] += points_grad_global[(offset + ipt)*3 + 1];
-        EXC_GRAD[3*iParent + 2] += points_grad_global[(offset + ipt)*3 + 2];
+        EXC_GRAD[3*iParent + 0] += points_grad_local[(offset + ipt)*3 + 0];
+        EXC_GRAD[3*iParent + 1] += points_grad_local[(offset + ipt)*3 + 1];
+        EXC_GRAD[3*iParent + 2] += points_grad_local[(offset + ipt)*3 + 2];
       }
       offset += task.points.size();
     }
@@ -1114,41 +1139,29 @@ void ReferenceReplicatedXCHostIntegrator<ValueType>::
   auto& tasks = this->load_balancer_->get_tasks();
   const size_t ntasks = tasks.size();
 
-  // Sort tasks for load balancing
+  // The tasks still have the exact order used to assemble eps_on_grid.
+  // Attach each point value to its task before reordering for load balancing.
   auto task_comparator = []( const XCTask& a, const XCTask& b ) {
     return (a.points.size() * a.bfn_screening.nbe) > (b.points.size() * b.bfn_screening.nbe);
   };
-  std::sort( tasks.begin(), tasks.end(), task_comparator );
-
-  // Build global eps_on_grid offset map: since tasks may be re-sorted,
-  // we need to distribute eps_on_grid to tasks. We use the task ordering
-  // from send_buffer_onedft_outputs (which was sorted by iParent).
-  // After re-sorting by task_comparator, we need a different approach.
-  // Actually, eps_on_grid_local was already in the send_buffer_onedft_outputs
-  // task order (sorted by iParent). The tasks are now re-sorted.
-  // We need to store eps per-task before re-sorting.
-  //
-  // WORKAROUND: Store per-task eps in task.feat before sorting.
-  // Actually, the simpler approach: don't re-sort. Use the current task order.
-  // The tasks were already sorted by iParent from prepare_onedft_features.
-  // Let's just rebuild the eps_per_task mapping.
-
-  // Build task -> eps mapping from the eps_on_grid vector (in iParent-sorted order)
-  // First, re-sort back to iParent order to match eps_on_grid
-  std::stable_sort( tasks.begin(), tasks.end(),
-    [](const auto& a, const auto& b) { return a.iParent < b.iParent; });
 
   // Distribute eps_on_grid to per-task storage
   {
     size_t offset = 0;
     for (auto& task : tasks) {
       int64_t npts = task.points.size();
+      if (offset + npts > eps_on_grid.size()) {
+        GAUXC_GENERIC_EXCEPTION("Inconsistent OneDFT energy-density layout");
+      }
       task.feat.eps.resize(npts);
       std::copy(eps_on_grid.data() + offset,
                 eps_on_grid.data() + offset + npts,
                 task.feat.eps.begin());
       offset += npts;
     }
+    if (offset != eps_on_grid.size()) {
+      GAUXC_GENERIC_EXCEPTION("Inconsistent OneDFT energy-density layout");
+    }
   }
 
   // Now sort by workload for the Pulay loop

tests/onedft_test.cxx

@@ -17,6 +17,7 @@
 #include <gauxc/util/mpi.hpp>
 #include <highfive/H5File.hpp>
 #include <Eigen/Core>
+#include <array>
 
 using namespace GauXC;
 
@@ -208,9 +209,96 @@ TEST_CASE( "OneDFT MPI Subgroup", "[onedft][mpi][subcomm]" ) {
 #endif
 }
 
+#ifdef GAUXC_HAS_HOST
+TEST_CASE( "OneDFT EXC Gradient Finite Difference", "[onedft][grad][fd]" ) {
+    using matrix_type = Eigen::MatrixXd;
+
+    auto rt = RuntimeEnvironment( GAUXC_MPI_CODE(MPI_COMM_WORLD) );
+
+    Molecule mol;
+    BasisSet<double> basis;
+    const std::string reference_file =
+        GAUXC_REF_DATA_PATH "/h2o_qzvpp_high_l.hdf5";
+    read_hdf5_record( mol, reference_file, "/MOLECULE" );
+    read_hdf5_record( basis, reference_file, "/BASIS" );
+
+    HighFive::File file( reference_file, HighFive::File::ReadOnly );
+    auto dset = file.getDataSet( "/DENSITY" );
+    auto dims = dset.getDimensions();
+    matrix_type P( dims[0], dims[1] );
+    dset.read( P.data() );
+
+    matrix_type Ps = P;
+    matrix_type Pz = matrix_type::Zero( dims[0], dims[1] );
+
+    constexpr size_t atom_index = 1;
+    constexpr size_t coordinate = 1;
+    constexpr double step = 1e-4;
+
+    auto evaluate = [&]( double displacement, bool gradient ) {
+        Molecule displaced_mol = mol;
+        BasisSet<double> displaced_basis = basis;
+        const std::array<double, 3> origin = {
+            mol[atom_index].x, mol[atom_index].y, mol[atom_index].z };
+
+        if( coordinate == 0 ) displaced_mol[atom_index].x += displacement;
+        if( coordinate == 1 ) displaced_mol[atom_index].y += displacement;
+        if( coordinate == 2 ) displaced_mol[atom_index].z += displacement;
+
+        for( auto& shell : displaced_basis ) {
+            const auto dx = shell.O()[0] - origin[0];
+            const auto dy = shell.O()[1] - origin[1];
+            const auto dz = shell.O()[2] - origin[2];
+            if( std::sqrt( dx*dx + dy*dy + dz*dz ) < 1e-8 ) {
+                shell.O()[coordinate] += displacement;
+            }
+            shell.set_shell_tolerance( 1e-10 );
+        }
+
+        auto mg = MolGridFactory::create_default_molgrid(
+            displaced_mol, PruningScheme::Robust, BatchSize(512),
+            RadialQuad::MuraKnowles, AtomicGridSizeDefault::FineGrid );
+        LoadBalancerFactory lb_factory( ExecutionSpace::Host, "Default" );
+        auto lb = lb_factory.get_instance(
+            rt, displaced_mol, mg, displaced_basis );
+        MolecularWeightsFactory mw_factory(
+            ExecutionSpace::Host, "Default", MolecularWeightsSettings{} );
+        auto mw = mw_factory.get_instance();
+        mw.modify_weights( lb );
+
+        functional_type func = functional_type(
+            ExchCXX::Backend::builtin, ExchCXX::Functional::PBE0,
+            ExchCXX::Spin::Unpolarized );
+        XCIntegratorFactory<matrix_type> integrator_factory(
+            ExecutionSpace::Host, "Replicated", "Default", "Default", "Default" );
+        auto integrator = integrator_factory.get_instance( func, lb );
+        OneDFTSettings settings;
+        settings.model = GAUXC_ONEDFT_MODEL_PATH "/pbe.fun";
+
+        if( gradient ) {
+            auto grad = integrator.eval_exc_grad_onedft( Ps, Pz, settings );
+            return std::make_pair( 0.0, std::move( grad ) );
+        }
+        auto [exc, vxc, vxcz] =
+            integrator.eval_exc_vxc_onedft( Ps, Pz, settings );
+        return std::make_pair( exc, std::vector<double>{} );
+    };
+
+    const auto analytical = evaluate( 0.0, true ).second;
+    const auto exc_plus = evaluate( step, false ).first;
+    const auto exc_minus = evaluate( -step, false ).first;
+    const auto numerical = (exc_plus - exc_minus) / (2.0 * step);
+
+    REQUIRE( analytical.size() == 3 * mol.size() );
+    CHECK( analytical[3*atom_index + coordinate] ==
+        Approx( numerical ).margin( 5e-6 ) );
+}
+#endif
+
 #if defined(GAUXC_HAS_HOST) && defined(GAUXC_HAS_DEVICE)
 void test_onedft_grad_host_device( std::string reference_file,
-    std::string onedft_model_path ) {
+    std::string onedft_model_path,
+    AtomicGridSizeDefault grid_size = AtomicGridSizeDefault::UltraFineGrid ) {
 
     using matrix_type = Eigen::MatrixXd;
     Molecule mol;
@@ -230,7 +318,7 @@ void test_onedft_grad_host_device( std::string reference_file,
     matrix_type Pz = matrix_type::Zero( dims[0], dims[1] );
 
     auto mg = MolGridFactory::create_default_molgrid( mol, PruningScheme::Unpruned,
-        BatchSize(512), RadialQuad::MuraKnowles, AtomicGridSizeDefault::UltraFineGrid );
+        BatchSize(512), RadialQuad::MuraKnowles, grid_size );
 
     functional_type func = functional_type( ExchCXX::Backend::builtin,
         ExchCXX::Functional::PBE0, ExchCXX::Spin::Unpolarized );
@@ -239,7 +327,7 @@ void test_onedft_grad_host_device( std::string reference_file,
     onedft_settings.model = onedft_model_path;
 
 #ifdef GAUXC_HAS_DEVICE
-    auto rt = DeviceRuntimeEnvironment( GAUXC_MPI_CODE(MPI_COMM_WORLD,) 0.9 );
+    auto rt = DeviceRuntimeEnvironment( GAUXC_MPI_CODE(MPI_COMM_WORLD,) 0.4 );
 #else
     auto rt = RuntimeEnvironment( GAUXC_MPI_CODE(MPI_COMM_WORLD) );
 #endif
@@ -280,6 +368,12 @@ TEST_CASE( "OneDFT EXC Gradient", "[onedft][grad]" ) {
             GAUXC_REF_DATA_PATH "/h2o2_def2-tzvp.hdf5",
             GAUXC_ONEDFT_MODEL_PATH "/tpss.fun" );
     }
+    SECTION( "H2O / QZVPP / pbe.fun (high-l)" ) {
+        test_onedft_grad_host_device(
+            GAUXC_REF_DATA_PATH "/h2o_qzvpp_high_l.hdf5",
+            GAUXC_ONEDFT_MODEL_PATH "/pbe.fun",
+            AtomicGridSizeDefault::FineGrid );
+    }
 }
 #endif
 
@@ -454,4 +548,4 @@ TEST_CASE( "Atom Reorder Permutation", "[onedft][reorder]" ) {
         CHECK(grad_dden[c*npts+i] == Approx(orig_dden[i*6+c]));
     }
   }
-}
+}