<commit_before>#include <sirius.h>
#include <json.hpp>

using namespace sirius;
using json = nlohmann::json;

const std::string aiida_output_file = "output_aiida.json";

enum class task_t : int
{
    ground_state_new     = 0,
    ground_state_restart = 1,
    k_point_path         = 2
};

const double au2angs = 0.5291772108;

void json_output_common(json& dict__)
{
    dict__["git_hash"] = git_hash;
    dict__["build_date"] = build_date;
    dict__["comm_world_size"] = mpi_comm_world().size();
    dict__["threads_per_rank"] = omp_get_max_threads();
}

std::unique_ptr<Simulation_context> create_sim_ctx(std::string     fname__,
                                                   cmd_args const& args__)
{
    auto ctx_ptr = std::unique_ptr<Simulation_context>(new Simulation_context(fname__, mpi_comm_world()));
    Simulation_context& ctx = *ctx_ptr;

    auto& inp = ctx.parameters_input();
    if (inp.gamma_point_ && !(inp.ngridk_[0] * inp.ngridk_[1] * inp.ngridk_[2] == 1)) {
        TERMINATE("this is not a Gamma-point calculation")
    }

    auto mpi_grid_dims = args__.value<std::vector<int>>("mpi_grid", ctx.mpi_grid_dims());
    ctx.set_mpi_grid_dims(mpi_grid_dims);

    auto std_evp_solver_name = args__.value<std::string>("std_evp_solver_name", ctx.control().std_evp_solver_name_);
    ctx.set_std_evp_solver_name(std_evp_solver_name);

    auto gen_evp_solver_name = args__.value<std::string>("gen_evp_solver_name", ctx.control().gen_evp_solver_name_);
    ctx.set_gen_evp_solver_name(gen_evp_solver_name);

    auto pu = args__.value<std::string>("processing_unit", ctx.control().processing_unit_);
    if (pu == "") {
        #ifdef __GPU
        pu = "gpu";
        #else
        pu = "cpu";
        #endif
    }
    ctx.set_processing_unit(pu);

    return std::move(ctx_ptr);
}


double ground_state(Simulation_context& ctx,
                    task_t              task,
                    cmd_args const&     args,
                    int                 write_output)
{
    if (ctx.comm().rank() == 0 && ctx.control().print_memory_usage_) {
        MEMORY_USAGE_INFO();
    }

    Potential potential(ctx);
    potential.allocate();

    Density density(ctx);
    density.allocate();

    Hamiltonian H(ctx, potential);

    if (ctx.comm().rank() == 0 && ctx.control().print_memory_usage_) {
        MEMORY_USAGE_INFO();
    }

    auto& inp = ctx.parameters_input();

    K_point_set ks(ctx, inp.ngridk_, inp.shiftk_, ctx.use_symmetry());
    ks.initialize();

    if (ctx.comm().rank() == 0 && ctx.control().print_memory_usage_) {
        MEMORY_USAGE_INFO();
    }

    std::string ref_file = args.value<std::string>("test_against", "");
    /* don't write output if we compare against the reference calculation */
    bool write_state = (ref_file.size() == 0);

    DFT_ground_state dft(ctx, H, density, ks);

    if (task == task_t::ground_state_restart) {
        if (!Utils::file_exists(storage_file_name)) {
            TERMINATE("storage file is not found");
        }
        density.load();
        potential.load();
    } else {
        density.initial_density();
        potential.generate(density);
        if (!ctx.full_potential()) {
            dft.band().initialize_subspace(ks, H);
        }
    }

    /* launch the calculation */
    int result = dft.find(inp.potential_tol_, inp.energy_tol_, inp.num_dft_iter_, write_state);

    dft.print_magnetic_moment();

    if (ref_file.size() != 0) {
        auto dict = dft.serialize();
        json dict_ref;
        std::ifstream(ref_file) >> dict_ref;

        double e1 = dict["energy"]["total"];
        double e2 = dict_ref["ground_state"]["energy"]["total"];

        if (std::abs(e1 - e2) > 1e-6) {
            printf("total energy is different: %18.7f computed vs. %18.7f reference\n", e1, e2);
            sirius::terminate(1);
        }
    }

    if (!ctx.full_potential()) {
        if (ctx.control().print_stress_) {
            Stress s(ctx, ks, density, potential);
            s.print_info();
        }
        if (ctx.control().print_forces_) {
            Forces_PS f(ctx, density, potential, ks);
            f.print_info();
        }
    }

    if (write_state && write_output) {
        json dict;
        json_output_common(dict);

        dict["task"] = static_cast<int>(task);
        dict["ground_state"] = dft.serialize();
        dict["timers"] = sddk::timer::serialize_timers();
        if (ctx.comm().rank() == 0) {
            std::ofstream ofs(std::string("output_") + ctx.start_time_tag() + std::string(".json"),
                              std::ofstream::out | std::ofstream::trunc);
            ofs << dict.dump(4);
        }

        if (args.exist("aiida_output")) {
            json dict;
            json_output_common(dict);
            dict["task"] = static_cast<int>(task);
            if (result >= 0) {
                dict["task_status"] = "converged";
                dict["num_scf_iterations"] =  result;
            } else {
                dict["task_status"] = "unconverged";
            }
            dict["volume"] = ctx.unit_cell().omega() * std::pow(au2angs, 3);
            dict["volume_units"] = "angstrom^3";
            dict["energy"] = dft.total_energy() * ha2ev;
            dict["energy_units"] = "eV";
            if (ctx.comm().rank() == 0) {
                std::ofstream ofs(aiida_output_file, std::ofstream::out | std::ofstream::trunc);
                ofs << dict.dump(4);
            }
        }
    }

    /* wait for all */
    ctx.comm().barrier();

    if (ctx.control().print_timers_)  {
        sddk::timer::print();
    }

    return dft.total_energy();
}

/// Run a task based on a command line input.
void run_tasks(cmd_args const& args)
{
    /* get the task id */
    task_t task = static_cast<task_t>(args.value<int>("task", 0));
    /* get the input file name */
    std::string fname = args.value<std::string>("input", "sirius.json");
    if (!Utils::file_exists(fname)) {
        if (mpi_comm_world().rank() == 0) {
            printf("input file does not exist\n");
        }
        return;
    }

    if (task == task_t::ground_state_new || task == task_t::ground_state_restart) {
        auto ctx = create_sim_ctx(fname, args);<commit_msg>only first ranks writes<commit_after>#include <sirius.h>
#include <json.hpp>

using namespace sirius;
using json = nlohmann::json;

const std::string aiida_output_file = "output_aiida.json";

enum class task_t : int
{
    ground_state_new     = 0,
    ground_state_restart = 1,
    k_point_path         = 2
};

const double au2angs = 0.5291772108;

void json_output_common(json& dict__)
{
    dict__["git_hash"] = git_hash;
    dict__["build_date"] = build_date;
    dict__["comm_world_size"] = mpi_comm_world().size();
    dict__["threads_per_rank"] = omp_get_max_threads();
}

std::unique_ptr<Simulation_context> create_sim_ctx(std::string     fname__,
                                                   cmd_args const& args__)
{
    auto ctx_ptr = std::unique_ptr<Simulation_context>(new Simulation_context(fname__, mpi_comm_world()));
    Simulation_context& ctx = *ctx_ptr;

    auto& inp = ctx.parameters_input();
    if (inp.gamma_point_ && !(inp.ngridk_[0] * inp.ngridk_[1] * inp.ngridk_[2] == 1)) {
        TERMINATE("this is not a Gamma-point calculation")
    }

    auto mpi_grid_dims = args__.value<std::vector<int>>("mpi_grid", ctx.mpi_grid_dims());
    ctx.set_mpi_grid_dims(mpi_grid_dims);

    auto std_evp_solver_name = args__.value<std::string>("std_evp_solver_name", ctx.control().std_evp_solver_name_);
    ctx.set_std_evp_solver_name(std_evp_solver_name);

    auto gen_evp_solver_name = args__.value<std::string>("gen_evp_solver_name", ctx.control().gen_evp_solver_name_);
    ctx.set_gen_evp_solver_name(gen_evp_solver_name);

    auto pu = args__.value<std::string>("processing_unit", ctx.control().processing_unit_);
    if (pu == "") {
        #ifdef __GPU
        pu = "gpu";
        #else
        pu = "cpu";
        #endif
    }
    ctx.set_processing_unit(pu);

    return std::move(ctx_ptr);
}


double ground_state(Simulation_context& ctx,
                    task_t              task,
                    cmd_args const&     args,
                    int                 write_output)
{
    if (ctx.comm().rank() == 0 && ctx.control().print_memory_usage_) {
        MEMORY_USAGE_INFO();
    }

    Potential potential(ctx);
    potential.allocate();

    Density density(ctx);
    density.allocate();

    Hamiltonian H(ctx, potential);

    if (ctx.comm().rank() == 0 && ctx.control().print_memory_usage_) {
        MEMORY_USAGE_INFO();
    }

    auto& inp = ctx.parameters_input();

    K_point_set ks(ctx, inp.ngridk_, inp.shiftk_, ctx.use_symmetry());
    ks.initialize();

    if (ctx.comm().rank() == 0 && ctx.control().print_memory_usage_) {
        MEMORY_USAGE_INFO();
    }

    std::string ref_file = args.value<std::string>("test_against", "");
    /* don't write output if we compare against the reference calculation */
    bool write_state = (ref_file.size() == 0);

    DFT_ground_state dft(ctx, H, density, ks);

    if (task == task_t::ground_state_restart) {
        if (!Utils::file_exists(storage_file_name)) {
            TERMINATE("storage file is not found");
        }
        density.load();
        potential.load();
    } else {
        density.initial_density();
        potential.generate(density);
        if (!ctx.full_potential()) {
            dft.band().initialize_subspace(ks, H);
        }
    }

    /* launch the calculation */
    int result = dft.find(inp.potential_tol_, inp.energy_tol_, inp.num_dft_iter_, write_state);

    dft.print_magnetic_moment();

    if (ref_file.size() != 0) {
        auto dict = dft.serialize();
        json dict_ref;
        std::ifstream(ref_file) >> dict_ref;

        double e1 = dict["energy"]["total"];
        double e2 = dict_ref["ground_state"]["energy"]["total"];

        if (std::abs(e1 - e2) > 1e-6) {
            printf("total energy is different: %18.7f computed vs. %18.7f reference\n", e1, e2);
            sirius::terminate(1);
        }
    }

    if (!ctx.full_potential()) {
        if (ctx.control().print_stress_) {
            Stress s(ctx, ks, density, potential);
            s.print_info();
        }
        if (ctx.control().print_forces_) {
            Forces_PS f(ctx, density, potential, ks);
            f.print_info();
        }
    }

    if (write_state && write_output) {
        json dict;
        json_output_common(dict);

        dict["task"] = static_cast<int>(task);
        dict["ground_state"] = dft.serialize();
        dict["timers"] = sddk::timer::serialize_timers();
        if (ctx.comm().rank() == 0) {
            std::ofstream ofs(std::string("output_") + ctx.start_time_tag() + std::string(".json"),
                              std::ofstream::out | std::ofstream::trunc);
            ofs << dict.dump(4);
        }

        if (args.exist("aiida_output")) {
            json dict;
            json_output_common(dict);
            dict["task"] = static_cast<int>(task);
            if (result >= 0) {
                dict["task_status"] = "converged";
                dict["num_scf_iterations"] =  result;
            } else {
                dict["task_status"] = "unconverged";
            }
            dict["volume"] = ctx.unit_cell().omega() * std::pow(au2angs, 3);
            dict["volume_units"] = "angstrom^3";
            dict["energy"] = dft.total_energy() * ha2ev;
            dict["energy_units"] = "eV";
            if (ctx.comm().rank() == 0) {
                std::ofstream ofs(aiida_output_file, std::ofstream::out | std::ofstream::trunc);
                ofs << dict.dump(4);
            }
        }
    }

    /* wait for all */
    ctx.comm().barrier();

    if (ctx.control().print_timers_ && ctx.comm().rank() == 0)  {
        sddk::timer::print();
    }

    return dft.total_energy();
}

/// Run a task based on a command line input.
void run_tasks(cmd_args const& args)
{
    /* get the task id */
    task_t task = static_cast<task_t>(args.value<int>("task", 0));
    /* get the input file name */
    std::string fname = args.value<std::string>("input", "sirius.json");
    if (!Utils::file_exists(fname)) {
        if (mpi_comm_world().rank() == 0) {
            printf("input file does not exist\n");
        }
        return;
    }

    if (task == task_t::ground_state_new || task == task_t::ground_state_restart) {
        auto ctx = create_sim_ctx(fname, args);
{'fraction_non_alphanumeric': 0.14922263183202436, 'fraction_numerical': 0.008094245872736015, 'mean_word_length': 2.4606211869107044, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 41, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 6, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '5188769', 'n_tokens_mistral': 4233, 'n_tokens_neox': 4099, 'n_words': 730}