Line data Source code
1 : // SPDX-FileCopyrightText: 2024 PairInteraction Developers
2 : // SPDX-License-Identifier: LGPL-3.0-or-later
3 :
4 : #include "pairinteraction/database/Database.hpp"
5 :
6 : #include "pairinteraction/basis/BasisAtom.hpp"
7 : #include "pairinteraction/database/AtomDescriptionByParameters.hpp"
8 : #include "pairinteraction/database/AtomDescriptionByRanges.hpp"
9 : #include "pairinteraction/database/GitHubDownloader.hpp"
10 : #include "pairinteraction/database/ParquetManager.hpp"
11 : #include "pairinteraction/enums/OperatorType.hpp"
12 : #include "pairinteraction/ket/KetAtom.hpp"
13 : #include "pairinteraction/ket/KetNotUniqueError.hpp"
14 : #include "pairinteraction/utils/TaskControl.hpp"
15 : #include "pairinteraction/utils/hash.hpp"
16 : #include "pairinteraction/utils/ket_id.hpp"
17 : #include "pairinteraction/utils/paths.hpp"
18 : #include "pairinteraction/utils/streamed.hpp"
19 :
20 : #include <algorithm>
21 : #include <cpptrace/cpptrace.hpp>
22 : #include <duckdb.hpp>
23 : #include <fmt/core.h>
24 : #include <fmt/format.h>
25 : #include <fmt/ranges.h>
26 : #include <fstream>
27 : #include <iterator>
28 : #include <nlohmann/json.hpp>
29 : #include <oneapi/tbb.h>
30 : #include <spdlog/spdlog.h>
31 : #include <string>
32 : #include <system_error>
33 : #include <unordered_map>
34 : #include <unordered_set>
35 :
36 : namespace pairinteraction {
37 :
38 : namespace {
39 :
40 353 : std::string format_expectation_value_range(const std::string &value_column,
41 : const std::string &std_column,
42 : const Range<double> &range,
43 : double standard_deviation_factor) {
44 0 : return fmt::format("{} BETWEEN {}-{}*{} AND {}+{}*{}", value_column, range.min(),
45 353 : standard_deviation_factor, std_column, range.max(),
46 706 : standard_deviation_factor, std_column);
47 : }
48 :
49 : // Find the index of the result column with the given name.
50 4067 : size_t get_column_index(const std::vector<std::string> &names, const std::string &name) {
51 4067 : auto it = std::find(names.begin(), names.end(), name);
52 4067 : if (it == names.end()) {
53 0 : throw std::runtime_error("Missing database column '" + name + "'.");
54 : }
55 4067 : return static_cast<size_t>(std::distance(names.begin(), it));
56 : }
57 :
58 : // Read a single value of a duckdb result column as a double, regardless of its logical type.
59 910233 : double get_entry_as_double(duckdb::Vector &vector, const duckdb::LogicalType &type, size_t row) {
60 910233 : switch (type.id()) {
61 736406 : case duckdb::LogicalTypeId::DOUBLE:
62 736406 : return duckdb::FlatVector::GetData<double>(vector)[row];
63 87264 : case duckdb::LogicalTypeId::BIGINT:
64 87264 : return static_cast<double>(duckdb::FlatVector::GetData<int64_t>(vector)[row]);
65 86566 : case duckdb::LogicalTypeId::BOOLEAN:
66 86566 : return duckdb::FlatVector::GetData<bool>(vector)[row] ? 1.0 : 0.0;
67 0 : default:
68 0 : throw std::runtime_error("Cannot read database column of type " + type.ToString() +
69 0 : " as a quantum number.");
70 : }
71 : }
72 :
73 : struct QuantumNumbers {
74 : std::unordered_map<std::string, double> values;
75 : std::unordered_map<std::string, double> stds;
76 : };
77 :
78 43283 : QuantumNumbers get_quantum_numbers_from_row(duckdb::DataChunk &chunk,
79 : const std::vector<duckdb::LogicalType> &types,
80 : const std::vector<std::string> &names,
81 : const std::unordered_set<std::string> &excluded_columns,
82 : size_t row) {
83 43283 : QuantumNumbers quantum_numbers;
84 1038300 : for (size_t col = 0; col < names.size(); ++col) {
85 995015 : const std::string &name = names[col];
86 995016 : if (excluded_columns.contains(name)) {
87 129847 : continue;
88 : }
89 865166 : double value = get_entry_as_double(chunk.data[col], types[col], row);
90 865167 : if (name.starts_with("std_")) {
91 259698 : quantum_numbers.stds[name.substr(4)] = value;
92 605464 : } else if (name.starts_with("exp_")) {
93 259697 : quantum_numbers.values[name.substr(4)] = value;
94 : } else {
95 345772 : quantum_numbers.values[name] = value;
96 : }
97 : }
98 43283 : return quantum_numbers;
99 0 : }
100 : } // namespace
101 :
102 43281 : void ensure_consistent_quantum_numbers(double quantum_number_f, double quantum_number_m) {
103 43281 : if (2 * quantum_number_m != std::rint(2 * quantum_number_m)) {
104 0 : throw std::runtime_error("The quantum number m must be an integer or half-integer.");
105 : }
106 43281 : if (2 * quantum_number_f != std::rint(2 * quantum_number_f)) {
107 0 : throw std::runtime_error("The quantum number f must be an integer or half-integer.");
108 : }
109 43281 : if (quantum_number_f + quantum_number_m != std::rint(quantum_number_f + quantum_number_m)) {
110 0 : throw std::invalid_argument(
111 0 : "The quantum numbers f and m must be both either integers or half-integers.");
112 : }
113 43281 : if (std::abs(quantum_number_m) > quantum_number_f) {
114 1 : throw std::invalid_argument(
115 2 : "The absolute value of the quantum number m must be less than or equal to f.");
116 : }
117 43280 : }
118 :
119 0 : Database::Database() : Database(default_download_missing) {}
120 :
121 0 : Database::Database(bool download_missing)
122 0 : : Database(download_missing, default_use_cache, default_database_dir) {}
123 :
124 0 : Database::Database(std::filesystem::path database_dir)
125 0 : : Database(default_download_missing, default_use_cache, std::move(database_dir)) {}
126 :
127 3 : Database::Database(bool download_missing, bool use_cache, std::filesystem::path database_dir)
128 3 : : download_missing_(download_missing), use_cache_(use_cache),
129 3 : database_dir_(std::move(database_dir)), db(std::make_unique<duckdb::DuckDB>(nullptr)),
130 15 : con(std::make_unique<duckdb::Connection>(*db)) {
131 :
132 3 : if (database_dir_.empty()) {
133 0 : database_dir_ = default_database_dir;
134 : }
135 :
136 : // Ensure the database directory exists
137 3 : if (!std::filesystem::exists(database_dir_)) {
138 0 : std::filesystem::create_directories(database_dir_);
139 : }
140 3 : database_dir_ = std::filesystem::canonical(database_dir_);
141 3 : if (!std::filesystem::is_directory(database_dir_)) {
142 0 : throw std::filesystem::filesystem_error("Cannot access database", database_dir_.string(),
143 0 : std::make_error_code(std::errc::not_a_directory));
144 : }
145 3 : SPDLOG_INFO("Using database directory: {}", database_dir_.string());
146 :
147 : // Ensure that the config directory exists
148 3 : std::filesystem::path configdir = paths::get_config_directory();
149 3 : if (!std::filesystem::exists(configdir)) {
150 1 : std::filesystem::create_directories(configdir);
151 2 : } else if (!std::filesystem::is_directory(configdir)) {
152 0 : throw std::filesystem::filesystem_error("Cannot access config directory ",
153 0 : configdir.string(),
154 0 : std::make_error_code(std::errc::not_a_directory));
155 : }
156 :
157 : // Read in the database_repo_paths if a config file exists, otherwise use the default and
158 : // write it to the config file
159 3 : std::filesystem::path configfile = configdir / "database.json";
160 3 : std::string database_repo_host;
161 3 : std::vector<std::string> database_repo_paths;
162 3 : if (std::filesystem::exists(configfile)) {
163 2 : std::ifstream file(configfile);
164 2 : nlohmann::json doc = nlohmann::json::parse(file, nullptr, false);
165 :
166 4 : if (!doc.is_discarded() && doc.contains("hash") && doc.contains("database_repo_host") &&
167 2 : doc.contains("database_repo_paths")) {
168 2 : database_repo_host = doc["database_repo_host"].get<std::string>();
169 2 : database_repo_paths = doc["database_repo_paths"].get<std::vector<std::string>>();
170 :
171 : // If the values are not equal to the default values but the hash is consistent (i.e.,
172 : // the user has not changed anything manually), clear the values so that they can be
173 : // updated
174 4 : if (database_repo_host != default_database_repo_host ||
175 2 : database_repo_paths != default_database_repo_paths) {
176 0 : std::size_t seed = 0;
177 0 : utils::hash_combine(seed, database_repo_paths);
178 0 : utils::hash_combine(seed, database_repo_host);
179 0 : if (seed == doc["hash"].get<std::size_t>()) {
180 0 : database_repo_host.clear();
181 0 : database_repo_paths.clear();
182 : } else {
183 0 : SPDLOG_INFO("The database repository host and paths have been changed "
184 : "manually. Thus, they will not be updated automatically. To reset "
185 : "them, delete the file '{}'.",
186 : configfile.string());
187 : }
188 : }
189 : }
190 2 : }
191 :
192 : // Read in and store the default values if necessary
193 3 : if (database_repo_host.empty() || database_repo_paths.empty()) {
194 2 : SPDLOG_INFO("Updating the database repository host and paths:");
195 :
196 1 : database_repo_host = default_database_repo_host;
197 1 : database_repo_paths = default_database_repo_paths;
198 1 : std::ofstream file(configfile);
199 1 : nlohmann::json doc;
200 :
201 1 : SPDLOG_INFO("* New host: {}", default_database_repo_host);
202 2 : SPDLOG_INFO("* New paths: {}", fmt::join(default_database_repo_paths, ", "));
203 :
204 1 : doc["database_repo_host"] = default_database_repo_host;
205 1 : doc["database_repo_paths"] = database_repo_paths;
206 :
207 1 : std::size_t seed = 0;
208 1 : utils::hash_combine(seed, default_database_repo_paths);
209 1 : utils::hash_combine(seed, default_database_repo_host);
210 1 : doc["hash"] = seed;
211 :
212 1 : file << doc.dump(4);
213 1 : }
214 :
215 : // Limit the memory usage of duckdb's buffer manager
216 : {
217 3 : auto result = con->Query("PRAGMA max_memory = '8GB';");
218 3 : if (result->HasError()) {
219 0 : throw cpptrace::runtime_error("Error setting the memory limit: " + result->GetError());
220 : }
221 3 : }
222 :
223 : // Instantiate a database manager that provides access to database tables. If a table
224 : // is outdated/not available locally, it will be downloaded if download_missing_ is true.
225 3 : if (!download_missing_) {
226 3 : database_repo_paths.clear();
227 : }
228 3 : downloader = std::make_unique<GitHubDownloader>();
229 6 : manager = std::make_unique<ParquetManager>(database_dir_, *downloader, database_repo_paths,
230 6 : *con, use_cache_);
231 3 : manager->scan_local();
232 3 : manager->scan_remote();
233 :
234 : // Print versions of tables
235 3 : std::istringstream iss(manager->get_versions_info());
236 33 : for (std::string line; std::getline(iss, line);) {
237 30 : SPDLOG_INFO(line);
238 3 : }
239 9 : }
240 :
241 2 : Database::~Database() = default;
242 :
243 2037 : const std::unordered_set<std::string> &Database::get_column_names(const std::string &table_path) {
244 2037 : if (auto it = column_names_cache.find(table_path); it != column_names_cache.end()) {
245 2029 : return it->second;
246 : }
247 16 : auto result = con->Query(fmt::format(R"(SELECT * FROM '{}' LIMIT 0)", table_path));
248 8 : if (result->HasError()) {
249 0 : throw cpptrace::runtime_error("Error querying the database columns: " + result->GetError());
250 : }
251 8 : std::unordered_set<std::string> names(result->names.begin(), result->names.end());
252 :
253 : // Every states table must provide the columns required for constructing kets and must not
254 : // contain an 'm' column, since m added separately below.
255 32 : for (const auto &required : {"id", "f", "energy"}) {
256 24 : if (!names.contains(required)) {
257 0 : throw std::runtime_error(
258 0 : fmt::format("The database table '{}' is missing the required column '{}'.",
259 0 : table_path, required));
260 : }
261 : }
262 8 : if (names.contains("m")) {
263 0 : throw std::runtime_error(
264 0 : fmt::format("The database table '{}' must not contain a column 'm'.", table_path));
265 : }
266 :
267 : // If another thread inserted the same entry concurrently, insert keeps the existing value and
268 : // returns an iterator to it, so the reference stays valid (elements are never erased).
269 8 : return column_names_cache.insert({table_path, std::move(names)}).first->second;
270 8 : }
271 :
272 495 : std::shared_ptr<const KetAtom> Database::get_ket(const std::string &species,
273 : const AtomDescriptionByParameters &description) {
274 : // Check that the specifications are valid
275 495 : if (!description.quantum_numbers.contains("m")) {
276 0 : throw std::invalid_argument("The quantum number m must be specified.");
277 : }
278 2455 : for (const auto &[name, value] : description.quantum_numbers) {
279 1960 : if ((name == "f" || name == "m") && 2 * value != std::rint(2 * value)) {
280 0 : throw std::invalid_argument("The quantum number " + name +
281 0 : " must be an integer or half-integer.");
282 : }
283 1960 : if (name == "f" && value < 0) {
284 0 : throw std::invalid_argument("The quantum number " + name + " must be positive.");
285 : }
286 : }
287 :
288 495 : const auto &columns = get_column_names(manager->get_path(species, "states"));
289 :
290 : // Describe the state. The quantum numbers n, f and parity are matched exactly, while all other
291 : // quantum numbers are matched within a +-0.5 window (they can deviate from the requested value,
292 : // e.g. expectation values in MQDT). The result is ordered by the distance to the requested
293 : // values, so the nearest state is returned.
294 495 : std::string where;
295 495 : std::string where_separator;
296 495 : std::string orderby;
297 495 : std::string orderby_separator;
298 495 : if (description.energy.has_value()) {
299 : // The following condition derives from demanding that quantum number n that corresponds to
300 : // the energy "E_n = -1/(2*n^2)" is not off by more than 1 from the actual quantum number n,
301 : // i.e., "sqrt(-1/(2*E_n)) - sqrt(-1/(2*E_{n-1})) = 1"
302 0 : double n_from_energy = std::sqrt(-1 / (2 * description.energy.value()));
303 0 : where += where_separator +
304 0 : fmt::format("SQRT(-1/(2*energy)) BETWEEN {} AND {}", n_from_energy - 0.5,
305 0 : n_from_energy + 0.5);
306 0 : where_separator = " AND ";
307 0 : orderby += orderby_separator + fmt::format("(SQRT(-1/(2*energy)) - {})^2", n_from_energy);
308 0 : orderby_separator = " + ";
309 : }
310 2453 : for (const auto &[name, value] : description.quantum_numbers) {
311 1959 : if (name == "m") {
312 495 : continue; // m is encoded into the id, not stored as a queryable column
313 : }
314 1464 : std::string column = columns.contains("exp_" + name) ? "exp_" + name : name;
315 1464 : if (!columns.contains(column)) {
316 2 : throw std::invalid_argument(
317 2 : fmt::format("The quantum number '{}' is not stored in the database table for "
318 : "species '{}'.",
319 2 : name, species));
320 : }
321 1463 : double tolerance = (name == "n" || name == "f" || name == "parity") ? 0.0 : 0.5;
322 1463 : where += where_separator +
323 4389 : fmt::format("{} BETWEEN {} AND {}", column, value - tolerance, value + tolerance);
324 1463 : where_separator = " AND ";
325 2926 : orderby += orderby_separator + fmt::format("({} - {})^2", column, value);
326 1463 : orderby_separator = " + ";
327 1464 : }
328 494 : if (where_separator.empty()) {
329 0 : where += "FALSE";
330 : }
331 494 : if (orderby_separator.empty()) {
332 0 : orderby += "id";
333 : }
334 :
335 : // Ask the database for the described state
336 494 : set_task_status("Loading atomic ket from database...");
337 988 : auto result = con->Query(fmt::format(
338 : R"(SELECT *, {} AS order_val FROM '{}' WHERE {} ORDER BY order_val ASC LIMIT 2)", orderby,
339 1482 : manager->get_path(species, "states"), where));
340 :
341 494 : if (result->HasError()) {
342 0 : throw cpptrace::runtime_error("Error querying the database: " + result->GetError());
343 : }
344 :
345 494 : if (result->RowCount() == 0) {
346 5 : throw std::invalid_argument("No state found.");
347 : }
348 :
349 : // Get the first chunk of the results (the first chunk is sufficient as we need two rows at
350 : // most). Every column except energy, id and the synthetic order_val is treated as a quantum
351 : // number; m is not a database column and is injected from the description.
352 489 : const auto &types = result->types;
353 489 : const auto &names = result->names;
354 1956 : const std::unordered_set<std::string> excluded_columns = {"energy", "id", "order_val"};
355 489 : auto chunk = result->Fetch();
356 :
357 489 : size_t energy_column = get_column_index(names, "energy");
358 489 : size_t id_column = get_column_index(names, "id");
359 489 : double quantum_number_m = description.quantum_numbers.at("m");
360 :
361 490 : auto make_ket = [&](size_t row) {
362 : auto quantum_numbers =
363 490 : get_quantum_numbers_from_row(*chunk, types, names, excluded_columns, row);
364 490 : quantum_numbers.values["m"] = quantum_number_m;
365 490 : double energy = get_entry_as_double(chunk->data[energy_column], types[energy_column], row);
366 : auto id =
367 980 : utils::encode_as_ket_id({.id = static_cast<size_t>(duckdb::FlatVector::GetData<int64_t>(
368 490 : chunk->data[id_column])[row]),
369 490 : .m = quantum_number_m});
370 490 : return KetAtom(typename KetAtom::Private(), energy, species,
371 980 : std::move(quantum_numbers.values), std::move(quantum_numbers.stds), *this,
372 1470 : id);
373 490 : };
374 :
375 : // Check that the ket is uniquely specified. If not, throw a KetNotUniqueError.
376 489 : if (chunk->size() > 1) {
377 7 : size_t order_val_column = get_column_index(names, "order_val");
378 : auto order_val_0 =
379 7 : get_entry_as_double(chunk->data[order_val_column], types[order_val_column], 0);
380 : auto order_val_1 =
381 7 : get_entry_as_double(chunk->data[order_val_column], types[order_val_column], 1);
382 :
383 7 : if (order_val_1 - order_val_0 <= order_val_0) {
384 4 : throw KetNotUniqueError({
385 0 : std::make_shared<const KetAtom>(make_ket(0)),
386 1 : std::make_shared<const KetAtom>(make_ket(1)),
387 5 : });
388 : }
389 : }
390 :
391 : // Construct the state
392 488 : auto ket = make_ket(0);
393 :
394 : // Check database consistency
395 490 : ensure_consistent_quantum_numbers(ket.get_quantum_number("f"), quantum_number_m);
396 :
397 974 : return std::make_shared<const KetAtom>(std::move(ket));
398 1510 : }
399 :
400 : template <typename Scalar>
401 : std::shared_ptr<const BasisAtom<Scalar>>
402 1542 : Database::get_basis(const std::string &species, const AtomDescriptionByRanges &description,
403 : const std::vector<size_t> &additional_ket_ids) {
404 : // The quantum number m is restricted separately because it is generated by UNNEST below.
405 3084 : auto range_quantum_number_m = [&description]() {
406 1542 : auto it = description.quantum_number_ranges.find("m");
407 3084 : return it != description.quantum_number_ranges.end() ? it->second : Range<double>{};
408 1542 : }();
409 :
410 1542 : const auto &columns = get_column_names(manager->get_path(species, "states"));
411 :
412 : // Describe the states by all restrictions that do not involve the quantum number m
413 1542 : std::string where = "(";
414 1542 : std::string separator;
415 1542 : if (description.range_energy.is_finite()) {
416 57 : where += separator +
417 57 : fmt::format("energy BETWEEN {} AND {}", description.range_energy.min(),
418 57 : description.range_energy.max());
419 57 : separator = " AND ";
420 : }
421 3089 : for (const auto &[name, range] : description.quantum_number_ranges) {
422 829 : if (name == "m" || !range.is_finite()) {
423 111 : continue;
424 : }
425 718 : std::string exp_column = "exp_" + name;
426 718 : std::string std_column = "std_" + name;
427 718 : if (columns.contains(exp_column) && columns.contains(std_column)) {
428 353 : where += separator +
429 : format_expectation_value_range(
430 : exp_column, std_column, range,
431 353 : description.quantum_number_standard_deviation_factor);
432 365 : } else if (columns.contains(name)) {
433 364 : where +=
434 728 : separator + fmt::format("{} BETWEEN {} AND {}", name, range.min(), range.max());
435 : } else {
436 2 : throw std::invalid_argument(
437 : fmt::format("The quantum number '{}' is not stored in the database table for "
438 : "species '{}'.",
439 : name, species));
440 : }
441 717 : separator = " AND ";
442 : }
443 1541 : if (separator.empty()) {
444 : // If the description contains no restrictions at all, it describes no states
445 1181 : where += range_quantum_number_m.is_finite() ? "TRUE" : "FALSE";
446 : }
447 1541 : where += ")";
448 :
449 : // Describe the restriction of the quantum number m
450 1541 : std::string where_m = "(";
451 1541 : if (range_quantum_number_m.is_finite()) {
452 222 : where_m += fmt::format("m BETWEEN {} AND {}", range_quantum_number_m.min(),
453 222 : range_quantum_number_m.max());
454 : } else {
455 1430 : where_m += "TRUE";
456 : }
457 1541 : where_m += ")";
458 :
459 : // Create a table containing the described states
460 1541 : std::string canonical_basis_id;
461 : {
462 1541 : auto result = con->Query(R"(SELECT UUID()::varchar)");
463 1541 : if (result->HasError()) {
464 0 : throw cpptrace::runtime_error("Error selecting canonical_basis_id: " +
465 0 : result->GetError());
466 : }
467 1541 : canonical_basis_id =
468 3082 : duckdb::FlatVector::GetData<duckdb::string_t>(result->Fetch()->data[0])[0].GetString();
469 1541 : }
470 : {
471 1541 : set_task_status("Selecting atomic basis states...");
472 4623 : auto result = con->Query(fmt::format(
473 : R"(CREATE TEMP TABLE '{}' AS SELECT *, id*{}+(2*m+{})::bigint AS ketid FROM (
474 : SELECT *,
475 : UNNEST(list_transform(generate_series(0,(2*f)::bigint),
476 : x -> x::double-f)) AS m FROM (
477 : SELECT * FROM '{}' WHERE {}
478 : )
479 : ) WHERE {})",
480 : canonical_basis_id, utils::KET_ID_STRIDE, utils::M_OFFSET,
481 3082 : manager->get_path(species, "states"), where, where_m));
482 :
483 1541 : if (result->HasError()) {
484 0 : throw cpptrace::runtime_error("Error creating table: " + result->GetError());
485 : }
486 1541 : }
487 : {
488 3082 : auto result = con->Query(
489 : fmt::format(R"(ALTER TABLE '{}' ADD PRIMARY KEY (ketid))", canonical_basis_id));
490 :
491 1541 : if (result->HasError()) {
492 0 : throw cpptrace::runtime_error("Error adding primary key: " + result->GetError());
493 : }
494 1541 : }
495 :
496 : // Add the additional kets to the table if they are not already contained in it
497 1541 : if (!additional_ket_ids.empty()) {
498 1243 : set_task_status("Selecting additional kets...");
499 1243 : std::vector<std::string> values;
500 1243 : values.reserve(additional_ket_ids.size());
501 3371 : for (size_t ket_id : additional_ket_ids) {
502 2128 : auto [id, m] = utils::decode_from_ket_id(ket_id);
503 4256 : values.push_back(fmt::format("({}, {}::double, {})", id, m, ket_id));
504 : }
505 4972 : auto result = con->Query(fmt::format(
506 : R"(INSERT OR IGNORE INTO '{0}'
507 : SELECT s.*, v.m, v.ketid FROM '{1}' AS s
508 : JOIN (VALUES {2}) AS v(id, m, ketid) ON s.id = v.id)",
509 3729 : canonical_basis_id, manager->get_path(species, "states"), fmt::join(values, ",")));
510 :
511 1243 : if (result->HasError()) {
512 0 : throw cpptrace::runtime_error("Error adding additional kets: " + result->GetError());
513 : }
514 1243 : }
515 :
516 : // Ask the table for the extreme values of the quantum numbers
517 : {
518 1541 : set_task_status("Validating atomic basis coverage...");
519 :
520 : // Collect the finite quantum-number ranges to validate against the loaded basis. Energy is
521 : // handled separately because it is compared via its corresponding effective quantum number.
522 : struct CoverageCheck {
523 : std::string name;
524 : std::string column;
525 : double tolerance{}; // allowed slack between the requested and the available range
526 : Range<double> range;
527 : };
528 1541 : std::vector<CoverageCheck> checks;
529 2369 : for (const auto &[name, range] : description.quantum_number_ranges) {
530 828 : if (!range.is_finite()) {
531 0 : continue;
532 : }
533 828 : bool is_expectation_value =
534 828 : columns.contains("exp_" + name) && columns.contains("std_" + name);
535 828 : std::string column = is_expectation_value ? "exp_" + name : name;
536 828 : double tolerance =
537 828 : (name == "n" || name == "f" || name == "m" || name == "parity") ? 0.0 : 1.0;
538 828 : checks.push_back({name, column, tolerance, range});
539 : }
540 :
541 1541 : std::string select;
542 1541 : std::string separator;
543 1541 : if (description.range_energy.is_finite()) {
544 57 : select += separator + "MIN(energy) AS min_energy, MAX(energy) AS max_energy";
545 57 : separator = ", ";
546 : }
547 2369 : for (const auto &check : checks) {
548 828 : select += separator +
549 828 : fmt::format("MIN({0}) AS min_{1}, MAX({0}) AS max_{1}", check.column, check.name);
550 828 : separator = ", ";
551 : }
552 :
553 1541 : if (!separator.empty()) {
554 720 : auto result =
555 : con->Query(fmt::format(R"(SELECT {} FROM '{}')", select, canonical_basis_id));
556 :
557 360 : if (result->HasError()) {
558 0 : throw cpptrace::runtime_error("Error querying the database: " + result->GetError());
559 : }
560 :
561 360 : auto chunk = result->Fetch();
562 360 : const auto &types = result->types;
563 :
564 2130 : for (size_t i = 0; i < chunk->ColumnCount(); i++) {
565 1770 : if (duckdb::FlatVector::IsNull(chunk->data[i], 0)) {
566 0 : throw std::invalid_argument("No state found.");
567 : }
568 : }
569 :
570 360 : size_t idx = 0;
571 360 : if (description.range_energy.is_finite()) {
572 57 : auto min_energy = get_entry_as_double(chunk->data[idx], types[idx], 0);
573 57 : idx++;
574 114 : if (std::sqrt(-1 / (2 * min_energy)) - 1 >
575 57 : std::sqrt(-1 / (2 * description.range_energy.min()))) {
576 0 : SPDLOG_DEBUG("No state found with the requested minimum energy. Requested: {}, "
577 : "found: {}.",
578 : description.range_energy.min(), min_energy);
579 : }
580 57 : auto max_energy = get_entry_as_double(chunk->data[idx], types[idx], 0);
581 57 : idx++;
582 114 : if (std::sqrt(-1 / (2 * max_energy)) + 1 <
583 57 : std::sqrt(-1 / (2 * description.range_energy.max()))) {
584 0 : SPDLOG_DEBUG("No state found with the requested maximum energy. Requested: {}, "
585 : "found: {}.",
586 : description.range_energy.max(), max_energy);
587 : }
588 : }
589 1188 : for (const auto &check : checks) {
590 828 : auto min_value = get_entry_as_double(chunk->data[idx], types[idx], 0);
591 828 : idx++;
592 828 : if (min_value - check.tolerance > check.range.min()) {
593 122 : SPDLOG_DEBUG("No state found with the requested minimum quantum number {}. "
594 : "Requested: {}, found: {}.",
595 : check.name, check.range.min(), min_value);
596 : }
597 828 : auto max_value = get_entry_as_double(chunk->data[idx], types[idx], 0);
598 828 : idx++;
599 828 : if (max_value + check.tolerance < check.range.max()) {
600 54 : SPDLOG_DEBUG("No state found with the requested maximum quantum number {}. "
601 : "Requested: {}, found: {}.",
602 : check.name, check.range.max(), max_value);
603 : }
604 : }
605 360 : }
606 1541 : }
607 :
608 : // Ask the table for the described states
609 1541 : set_task_status("Loading atomic basis states...");
610 3082 : auto result =
611 : con->Query(fmt::format(R"(SELECT * FROM '{}' ORDER BY ketid ASC)", canonical_basis_id));
612 :
613 1541 : if (result->HasError()) {
614 0 : throw cpptrace::runtime_error("Error querying the database: " + result->GetError());
615 : }
616 :
617 1541 : if (result->RowCount() == 0) {
618 0 : throw std::invalid_argument("No state found.");
619 : }
620 :
621 : // Construct the states. Every column except energy and the raw/encoded id is treated as a
622 : // quantum number ("id" is the raw states-table id, "ketid" the m-encoded id of the basis
623 : // state).
624 1541 : const auto &types = result->types;
625 1541 : const auto &names = result->names;
626 6164 : const std::unordered_set<std::string> excluded_columns = {"energy", "id", "ketid"};
627 1541 : size_t energy_column = get_column_index(names, "energy");
628 1541 : size_t ketid_column = get_column_index(names, "ketid");
629 :
630 1541 : std::vector<std::shared_ptr<const KetAtom>> kets;
631 1541 : kets.reserve(result->RowCount());
632 1541 : double last_energy = std::numeric_limits<double>::lowest();
633 1541 : double min_quantum_number_nu = std::numeric_limits<double>::max();
634 :
635 3082 : for (auto chunk = result->Fetch(); chunk; chunk = result->Fetch()) {
636 1541 : set_task_status("Constructing atomic basis...");
637 :
638 44334 : for (size_t i = 0; i < chunk->size(); i++) {
639 85586 : auto quantum_numbers =
640 42793 : get_quantum_numbers_from_row(*chunk, types, names, excluded_columns, i);
641 42793 : double energy =
642 42793 : get_entry_as_double(chunk->data[energy_column], types[energy_column], i);
643 42793 : auto id = static_cast<size_t>(
644 42793 : duckdb::FlatVector::GetData<int64_t>(chunk->data[ketid_column])[i]);
645 :
646 : // Check database consistency
647 42793 : ensure_consistent_quantum_numbers(quantum_numbers.values.at("f"),
648 85586 : quantum_numbers.values.at("m"));
649 42793 : if (energy < last_energy) {
650 0 : throw std::runtime_error("The states are not sorted by energy.");
651 : }
652 42793 : last_energy = energy;
653 :
654 42793 : if (auto it = quantum_numbers.values.find("nu"); it != quantum_numbers.values.end()) {
655 42793 : min_quantum_number_nu = std::min(min_quantum_number_nu, it->second);
656 : }
657 :
658 : // Append a new state
659 42793 : kets.push_back(std::make_shared<const KetAtom>(
660 85586 : typename KetAtom::Private(), energy, species, std::move(quantum_numbers.values),
661 42793 : std::move(quantum_numbers.stds), *this, id));
662 : }
663 : }
664 :
665 : // Show a warning for low-lying states
666 1541 : if (min_quantum_number_nu < 25) {
667 50 : if (species.ends_with("_mqdt")) {
668 44 : SPDLOG_WARN("The multi-channel quantum defect theory might produce inaccurate results "
669 : "for effective principal quantum numbers < 25. The models get increasingly "
670 : "unreliable for small principal quantum numbers, leading to inaccurate "
671 : "matrix elements and energies. Due to missing data, even some states might "
672 : "not be present.");
673 : } else {
674 56 : SPDLOG_WARN(
675 : "The single-channel quantum defect theory can be inaccurate for effective "
676 : "principal quantum numbers < 25. This can lead to inaccurate matrix elements.");
677 : }
678 : }
679 :
680 0 : return std::make_shared<const BasisAtom<Scalar>>(typename BasisAtom<Scalar>::Private(),
681 1541 : std::move(kets), std::move(canonical_basis_id),
682 3082 : *this);
683 5453 : }
684 :
685 : template <typename Scalar>
686 5208 : Eigen::SparseMatrix<Scalar, Eigen::RowMajor> Database::get_matrix_elements_in_canonical_basis(
687 : std::shared_ptr<const BasisAtom<Scalar>> initial_basis,
688 : std::shared_ptr<const BasisAtom<Scalar>> final_basis, OperatorType type, int q) {
689 : using real_t = typename traits::NumTraits<Scalar>::real_t;
690 : using cached_matrix_ptr_t = std::shared_ptr<const cached_matrix_t>;
691 :
692 5208 : if (&initial_basis->get_database() != this || &final_basis->get_database() != this) {
693 1 : throw std::invalid_argument(
694 : "The initial and final bases must belong to the Database instance used for the "
695 : "matrix element calculation.");
696 : }
697 5214 : if (initial_basis->get_species() != final_basis->get_species()) {
698 0 : throw std::invalid_argument(fmt::format(
699 : "The initial and final bases must have the same species, but got '{}' and '{}'.",
700 0 : initial_basis->get_species(), final_basis->get_species()));
701 : }
702 :
703 5214 : std::string specifier;
704 5232 : int kappa{};
705 5232 : switch (type) {
706 3747 : case OperatorType::ELECTRIC_DIPOLE:
707 3747 : specifier = "matrix_elements_d";
708 3749 : kappa = 1;
709 3749 : break;
710 257 : case OperatorType::ELECTRIC_QUADRUPOLE:
711 257 : specifier = "matrix_elements_q";
712 257 : kappa = 2;
713 257 : break;
714 49 : case OperatorType::ELECTRIC_QUADRUPOLE_ZERO:
715 49 : specifier = "matrix_elements_q0";
716 49 : kappa = 0;
717 49 : break;
718 0 : case OperatorType::ELECTRIC_OCTUPOLE:
719 0 : specifier = "matrix_elements_o";
720 0 : kappa = 3;
721 0 : break;
722 56 : case OperatorType::MAGNETIC_DIPOLE:
723 56 : specifier = "matrix_elements_mu";
724 58 : kappa = 1;
725 58 : break;
726 4 : case OperatorType::ENERGY:
727 4 : specifier = "energy";
728 4 : kappa = 0;
729 4 : break;
730 1121 : case OperatorType::IDENTITY:
731 1121 : specifier = "identity";
732 1121 : kappa = 0;
733 1121 : break;
734 0 : default:
735 0 : throw std::invalid_argument("Unknown operator type.");
736 : }
737 :
738 5238 : std::string canonical_basis_id_initial = initial_basis->get_canonical_basis_id();
739 5224 : std::string canonical_basis_id_final = final_basis->get_canonical_basis_id();
740 5219 : std::string cache_key = fmt::format("{}_{}_{}_{}", specifier, q, canonical_basis_id_initial,
741 : canonical_basis_id_final);
742 5242 : auto &matrix_elements_cache = get_matrix_elements_cache();
743 5183 : std::promise<cached_matrix_ptr_t> matrix_promise;
744 5215 : auto [cache_it, inserted] =
745 : matrix_elements_cache.insert({cache_key, matrix_promise.get_future().share()});
746 :
747 5193 : if (inserted) {
748 : try {
749 1594 : Eigen::Index num_rows = final_basis->get_number_of_kets();
750 1594 : Eigen::Index num_cols = initial_basis->get_number_of_kets();
751 :
752 1595 : std::vector<int> outerIndexPtr;
753 1595 : std::vector<int> innerIndices;
754 1595 : std::vector<real_t> values;
755 :
756 : // Check that the specifications are valid
757 1595 : if (std::abs(q) > kappa) {
758 0 : throw std::invalid_argument("Invalid q.");
759 : }
760 :
761 : // Ask the database for the operator
762 1595 : set_task_status("Loading matrix elements from database...");
763 1595 : std::string species = initial_basis->get_species();
764 1595 : duckdb::unique_ptr<duckdb::MaterializedQueryResult> result;
765 1595 : if (specifier == "identity") {
766 1948 : result = con->Query(fmt::format(
767 : R"(SELECT s2.ketid AS row, s1.ketid AS col, 1.0::DOUBLE AS val
768 : FROM '{}' AS s1
769 : INNER JOIN '{}' AS s2 ON s1.ketid = s2.ketid
770 : ORDER BY row ASC)",
771 : canonical_basis_id_initial, canonical_basis_id_final));
772 621 : } else if (specifier == "energy") {
773 6 : result = con->Query(fmt::format(
774 : R"(SELECT s2.ketid AS row, s1.ketid AS col, s1.energy AS val
775 : FROM '{}' AS s1
776 : INNER JOIN '{}' AS s2 ON s1.ketid = s2.ketid
777 : ORDER BY row ASC)",
778 : canonical_basis_id_initial, canonical_basis_id_final));
779 : } else {
780 2471 : result = con->Query(fmt::format(
781 : R"(WITH s1 AS (
782 : SELECT id, f, m, ketid FROM '{}'
783 : ),
784 : s2 AS (
785 : SELECT id, f, m, ketid FROM '{}'
786 : ),
787 : b AS (
788 : SELECT MIN(f) AS min_f, MAX(f) AS max_f,
789 : MIN(id) AS min_id, MAX(id) AS max_id
790 : FROM (SELECT f, id FROM s1 UNION ALL SELECT f, id FROM s2)
791 : ),
792 : w_filtered AS (
793 : SELECT *
794 : FROM '{}'
795 : WHERE kappa = {} AND q = {} AND
796 : f_initial BETWEEN (SELECT min_f FROM b) AND (SELECT max_f FROM b) AND
797 : f_final BETWEEN (SELECT min_f FROM b) AND (SELECT max_f FROM b)
798 : ),
799 : e_filtered AS (
800 : SELECT *
801 : FROM '{}'
802 : WHERE
803 : id_initial BETWEEN (SELECT min_id FROM b) AND (SELECT max_id FROM b) AND
804 : id_final BETWEEN (SELECT min_id FROM b) AND (SELECT max_id FROM b)
805 : )
806 : SELECT
807 : s2.ketid AS row,
808 : s1.ketid AS col,
809 : e.val*w.val AS val
810 : FROM e_filtered AS e
811 : JOIN s1 ON e.id_initial = s1.id
812 : JOIN s2 ON e.id_final = s2.id
813 : JOIN w_filtered AS w ON
814 : w.f_initial = s1.f AND w.m_initial = s1.m AND
815 : w.f_final = s2.f AND w.m_final = s2.m
816 : ORDER BY row ASC, col ASC)",
817 : canonical_basis_id_initial, canonical_basis_id_final,
818 1235 : manager->get_path("misc", "wigner"), kappa, q,
819 : manager->get_path(species, specifier)));
820 : }
821 :
822 1595 : if (result->HasError()) {
823 0 : throw cpptrace::runtime_error("Error querying the database: " + result->GetError());
824 : }
825 :
826 : // Check the types of the columns
827 1595 : const auto &types = result->types;
828 1595 : const auto &labels = result->names;
829 6380 : const std::vector<duckdb::LogicalType> ref_types = {duckdb::LogicalType::BIGINT,
830 : duckdb::LogicalType::BIGINT,
831 : duckdb::LogicalType::DOUBLE};
832 6380 : for (size_t i = 0; i < types.size(); i++) {
833 4785 : if (types[i] != ref_types[i]) {
834 0 : throw std::runtime_error("Wrong type for '" + labels[i] + "'.");
835 : }
836 : }
837 :
838 1595 : set_task_status("Constructing matrix elements...");
839 :
840 : // Construct the matrix
841 1595 : int num_entries = static_cast<int>(result->RowCount());
842 1595 : outerIndexPtr.reserve(num_rows + 1);
843 1595 : innerIndices.reserve(num_entries);
844 1595 : values.reserve(num_entries);
845 :
846 1595 : int last_row = -1;
847 :
848 3399 : for (auto chunk = result->Fetch(); chunk; chunk = result->Fetch()) {
849 1804 : auto *chunk_row = duckdb::FlatVector::GetData<int64_t>(chunk->data[0]);
850 1804 : auto *chunk_col = duckdb::FlatVector::GetData<int64_t>(chunk->data[1]);
851 1804 : auto *chunk_val = duckdb::FlatVector::GetData<double>(chunk->data[2]);
852 :
853 782368 : for (size_t i = 0; i < chunk->size(); i++) {
854 780564 : int row = final_basis->get_ket_index_from_id(chunk_row[i]);
855 780564 : if (row != last_row) {
856 51156 : if (row < last_row) {
857 0 : throw std::runtime_error("The rows are not sorted.");
858 : }
859 107788 : for (; last_row < row; last_row++) {
860 56632 : outerIndexPtr.push_back(static_cast<int>(innerIndices.size()));
861 : }
862 : }
863 780564 : innerIndices.push_back(initial_basis->get_ket_index_from_id(chunk_col[i]));
864 780564 : values.push_back(chunk_val[i]);
865 : }
866 : }
867 :
868 5294 : for (; last_row < num_rows + 1; last_row++) {
869 3699 : outerIndexPtr.push_back(static_cast<int>(innerIndices.size()));
870 : }
871 :
872 1595 : Eigen::Map<const cached_matrix_t> matrix_map(num_rows, num_cols, values.size(),
873 1595 : outerIndexPtr.data(), innerIndices.data(),
874 1595 : values.data());
875 :
876 1595 : auto cached_matrix = std::make_shared<const cached_matrix_t>(matrix_map);
877 1595 : matrix_promise.set_value(std::move(cached_matrix));
878 1595 : } catch (...) {
879 0 : matrix_promise.set_exception(std::current_exception());
880 0 : throw;
881 : }
882 : }
883 :
884 5194 : set_task_status("Returning matrix elements in canonical basis...");
885 :
886 10490 : return cache_it->second.get()->template cast<Scalar>();
887 6824 : }
888 :
889 4 : bool Database::get_download_missing() const { return download_missing_; }
890 :
891 2 : bool Database::get_use_cache() const { return use_cache_; }
892 :
893 10 : std::filesystem::path Database::get_database_dir() const { return database_dir_; }
894 :
895 0 : std::string Database::get_versions_info() const { return manager->get_versions_info(); }
896 :
897 5237 : Database::matrix_elements_cache_t &Database::get_matrix_elements_cache() {
898 5237 : static matrix_elements_cache_t matrix_elements_cache;
899 5197 : return matrix_elements_cache;
900 : }
901 :
902 45 : Database &Database::get_global_instance() {
903 90 : return get_global_instance_without_checks(default_download_missing, default_use_cache,
904 90 : default_database_dir);
905 : }
906 :
907 0 : Database &Database::get_global_instance(bool download_missing) {
908 0 : Database &database = get_global_instance_without_checks(download_missing, default_use_cache,
909 : default_database_dir);
910 0 : if (download_missing != database.download_missing_) {
911 0 : throw std::invalid_argument(
912 0 : "The 'download_missing' argument must not change between calls to the method.");
913 : }
914 0 : return database;
915 : }
916 :
917 0 : Database &Database::get_global_instance(std::filesystem::path database_dir) {
918 0 : if (database_dir.empty()) {
919 0 : database_dir = default_database_dir;
920 : }
921 0 : Database &database = get_global_instance_without_checks(default_download_missing,
922 : default_use_cache, database_dir);
923 0 : if (!std::filesystem::exists(database_dir) ||
924 0 : std::filesystem::canonical(database_dir) != database.database_dir_) {
925 0 : throw std::invalid_argument(
926 0 : "The 'database_dir' argument must not change between calls to the method.");
927 : }
928 0 : return database;
929 : }
930 :
931 1 : Database &Database::get_global_instance(bool download_missing, bool use_cache,
932 : std::filesystem::path database_dir) {
933 1 : if (database_dir.empty()) {
934 0 : database_dir = default_database_dir;
935 : }
936 : Database &database =
937 1 : get_global_instance_without_checks(download_missing, use_cache, database_dir);
938 1 : if (download_missing != database.download_missing_ || use_cache != database.use_cache_ ||
939 3 : !std::filesystem::exists(database_dir) ||
940 2 : std::filesystem::canonical(database_dir) != database.database_dir_) {
941 0 : throw std::invalid_argument(
942 : "The 'download_missing', 'use_cache' and 'database_dir' arguments must not "
943 0 : "change between calls to the method.");
944 : }
945 1 : return database;
946 : }
947 :
948 46 : Database &Database::get_global_instance_without_checks(bool download_missing, bool use_cache,
949 : std::filesystem::path database_dir) {
950 46 : static Database database(download_missing, use_cache, std::move(database_dir));
951 46 : return database;
952 : }
953 :
954 : struct database_dir_noexcept : std::filesystem::path {
955 2 : explicit database_dir_noexcept() noexcept try : std
956 2 : ::filesystem::path(paths::get_cache_directory() / "database") {}
957 0 : catch (...) {
958 0 : SPDLOG_ERROR("Error getting the PairInteraction cache directory.");
959 0 : std::terminate();
960 2 : }
961 : };
962 :
963 : const std::filesystem::path Database::default_database_dir = database_dir_noexcept();
964 :
965 : // Explicit instantiations
966 : // NOLINTBEGIN(bugprone-macro-parentheses, cppcoreguidelines-macro-usage)
967 : #define INSTANTIATE_GETTERS(SCALAR) \
968 : template std::shared_ptr<const BasisAtom<SCALAR>> Database::get_basis<SCALAR>( \
969 : const std::string &species, const AtomDescriptionByRanges &description, \
970 : const std::vector<size_t> &additional_ket_ids); \
971 : template Eigen::SparseMatrix<SCALAR, Eigen::RowMajor> \
972 : Database::get_matrix_elements_in_canonical_basis<SCALAR>( \
973 : std::shared_ptr<const BasisAtom<SCALAR>> initial_basis, \
974 : std::shared_ptr<const BasisAtom<SCALAR>> final_basis, OperatorType type, int q);
975 : // NOLINTEND(bugprone-macro-parentheses, cppcoreguidelines-macro-usage)
976 :
977 : INSTANTIATE_GETTERS(double)
978 : INSTANTIATE_GETTERS(std::complex<double>)
979 :
980 : #undef INSTANTIATE_GETTERS
981 : } // namespace pairinteraction
|
