diff --git a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_onedft.hpp b/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_onedft.hpp index a22321e6..6783d017 100644 --- a/src/xc_integrator/replicated/device/incore_replicated_xc_device_integrator_onedft.hpp +++ b/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 recv_raw(world_rank == 0 ? total_npts_sum : 0); + MPI_Gatherv(raw_grid_weights.data(), static_cast(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 perm(total_npts_sum); + for (int64_t j = 0; j < total_npts_sum; ++j) { + perm[atom_reorder_inv_perm[j]] = j; + } + std::vector 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 eps_on_grid_local; + std::vector 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 recvcounts3(recvcounts.size()); + std::vector 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(); } diff --git a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_onedft.hpp b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_onedft.hpp index e00bb9be..bdeba413 100644 --- a/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_onedft.hpp +++ b/src/xc_integrator/replicated/host/reference_replicated_xc_host_integrator_onedft.hpp @@ -1027,33 +1027,58 @@ void ReferenceReplicatedXCHostIntegrator:: // 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 eps_on_grid_local; + std::vector 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 sendcounts3(sendcounts.size()); + std::vector 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:: 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 diff --git a/tests/onedft_test.cxx b/tests/onedft_test.cxx index 88092b8c..5be0f55c 100644 --- a/tests/onedft_test.cxx +++ b/tests/onedft_test.cxx @@ -17,6 +17,7 @@ #include #include #include +#include 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 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 displaced_basis = basis; + const std::array 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 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{} ); + }; + + 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])); } } -} \ No newline at end of file +} diff --git a/tests/ref_data/h2o_qzvpp_high_l.hdf5 b/tests/ref_data/h2o_qzvpp_high_l.hdf5 new file mode 100644 index 00000000..6b001da2 Binary files /dev/null and b/tests/ref_data/h2o_qzvpp_high_l.hdf5 differ