phigrape/io.cpp

#ifdef HAS_HDF5
#include "hdf5.h"
#endif

#include "double3.h"
#include "io.h"
#include <fstream>
#include <cstring>
#include <cmath>
#include <vector>
#include <algorithm>
#include <string>
#include <stdexcept>


bool is_hdf5(std::string file_name)
{
    std::ifstream file(file_name, std::ifstream::binary);
    const char hdf5_magic[] = "\x89HDF\x0d\x0a\x1a\x0a";
    char buffer[8];
    file.read(buffer, 8);
    if (!file) throw std::runtime_error("Unable to read file " + file_name);
    bool result = (memcmp(buffer, hdf5_magic, 8)==0);
    file.close();
    return result;
}

Input_data ascii_read(const std::string &file_name)
{
    Input_data input_data;
    char rest[512];
    int result;

    std::ifstream file(file_name);
    if (!file.good()) throw std::runtime_error("File " + file_name + " not found.");
    std::string str;

    std::getline(file, str);
    result = sscanf(str.c_str(), "%d%s", &input_data.step_num, rest);
    if (result!=1) throw std::runtime_error("Error parsing line 1: expected one integer");

    std::getline(file, str);
    result = sscanf(str.c_str(), "%d%s", &input_data.N, rest);
    if (result!=1) throw std::runtime_error("Error parsing line 2: expected one integer");

    std::getline(file, str);
    result = sscanf(str.c_str(), "%lf%s", &input_data.t, rest);
    if (result!=1) throw std::runtime_error("Error parsing line 3: expected one real number");

    input_data.m.resize(input_data.N);
    input_data.x.resize(input_data.N);
    input_data.v.resize(input_data.N);

    int i = -1;
    while (std::getline(file, str)) {
        if (++i > input_data.N) throw std::runtime_error("Error parsing line " + std::to_string(i+4) + ": particle out of range");
        result = sscanf(str.c_str(), "%*s %lf %lf %lf %lf %lf %lf %lf%s", &input_data.m[i], &(input_data.x[i].x), &(input_data.x[i].y), &(input_data.x[i].z), &(input_data.v[i].x), &(input_data.v[i].y), &(input_data.v[i].z), rest);
    }
    file.close();
    return input_data;
}

void ascii_write(const std::string file_name, const int step_num, const int N, const double t, const std::vector<double> &m, const std::vector<double3> &x, const std::vector<double3> &v, int precision)
{
    auto file = std::ofstream(file_name);
    if (!file.is_open()) throw std::runtime_error("Cannot open file for output");
    int id_width = (int)log10(N-1) + 1;
    char string_template[256], output_string[256];
    file << step_num << '\n';
    file << N << '\n';
    file.precision(16);
    file << std::scientific << t << '\n';
    sprintf(string_template, "%%0%dd%%%d.%dE%%%d.%dE%%%d.%dE%%%d.%dE%%%d.%dE%%%d.%dE%%%d.%dE\n", id_width, precision+7, precision, precision+8, precision, precision+8, precision, precision+8, precision, precision+8, precision, precision+8, precision, precision+8, precision);
    for (int i=0; i<N; i++) {
        sprintf(output_string, string_template, i, m[i], x[i][0], x[i][1], x[i][2], v[i][0], v[i][1], v[i][2]);
        file << output_string;
    }
    file.close();
}

Input_data h5_read(const std::string &file_name)
{
#ifdef HAS_HDF5
    Input_data input_data;

    // Open file and root group; count number of top level objects
    hid_t file_id = H5Fopen(file_name.c_str(), H5F_ACC_RDONLY, H5P_DEFAULT);
    hid_t group_id = H5Gopen2(file_id, "/", H5P_DEFAULT);
    H5G_info_t object_info;
    H5Gget_info(group_id, &object_info);
    // Iterate over objects and add the number of each "step" group into a vector
    std::vector<int> step_num_arr;
    for (unsigned int i=0; i<object_info.nlinks; i++) {
        char name_cstr[256];
        H5Gget_objname_by_idx(group_id, i, name_cstr, 256);
        std::string name(name_cstr);
        if (name.substr(0, 5) == "Step#") step_num_arr.push_back(std::stoi(name.substr(5, std::string::npos)));
    }
    // Find the highest step in the file
    input_data.step_num = *max_element(step_num_arr.begin(), step_num_arr.end());

    // Prepare to read the data sets dimensionality data.
    char path[256];
    hid_t attr_id, dataset_id, dataspace_id;
    int ndims;
    hsize_t dims[2];

    sprintf(path, "/Step#%d", input_data.step_num);
    attr_id = H5Aopen_by_name(file_id, path, "Time", H5P_DEFAULT, H5P_DEFAULT);
    H5Aread(attr_id, H5T_NATIVE_DOUBLE, &input_data.t);
    H5Aclose(attr_id);

    sprintf(path, "/Step#%d/MASS", input_data.step_num);
    dataset_id = H5Dopen2(file_id, path, H5P_DEFAULT);
    dataspace_id = H5Dget_space(dataset_id);
    ndims = H5Sget_simple_extent_ndims(dataspace_id);
    if (ndims != 1)
        throw std::runtime_error("Dataset MASS in Step#" + std::to_string(input_data.step_num) + " of file " + file_name + " is " + std::to_string(ndims) + "-dimensional (expected 1-dimensional)");
    H5Sget_simple_extent_dims(dataspace_id, dims, NULL);
    input_data.N = dims[0];

    input_data.m.resize(input_data.N);
    H5Dread(dataset_id, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, input_data.m.data());
    H5Sclose(dataspace_id);
    H5Dclose(dataset_id);

    sprintf(path, "/Step#%d/X", input_data.step_num);
    dataset_id = H5Dopen2(file_id, path, H5P_DEFAULT);
    dataspace_id = H5Dget_space(dataset_id);
    ndims = H5Sget_simple_extent_ndims(dataspace_id);
    if (ndims != 2) throw std::runtime_error("Bad dimensionality in " + std::string(path));
    H5Sget_simple_extent_dims(dataspace_id, dims, NULL);
    if ((dims[0] != input_data.N) || (dims[1] != 3))  throw std::runtime_error("Bad dimensionality");
    input_data.x.resize(input_data.N);
    H5Dread(dataset_id, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, (double*)input_data.x.data());
    H5Sclose(dataspace_id);
    H5Dclose(dataset_id);

    sprintf(path, "/Step#%d/V", input_data.step_num);
    dataset_id = H5Dopen2(file_id, path, H5P_DEFAULT);
    dataspace_id = H5Dget_space(dataset_id);
    ndims = H5Sget_simple_extent_ndims(dataspace_id);
    if (ndims != 2) throw std::runtime_error("Bad dimensionality in" + std::string(path));
    H5Sget_simple_extent_dims(dataspace_id, dims, NULL);
    if ((dims[0] != input_data.N) || (dims[1] != 3)) throw std::runtime_error("Bad dimensionality");
    input_data.v.resize(input_data.N);
    H5Dread(dataset_id, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, (double*)input_data.v.data());
    H5Sclose(dataspace_id);
    H5Dclose(dataset_id);

    H5Gclose(group_id);
    H5Fclose(file_id);
    return input_data;
#else
    throw std::runtime_error("h5_read was called but compiled without HDF5 support");
#endif
}

void h5_write(const std::string file_name, const int step_num, const int N, const double t, const std::vector<double> &m, const std::vector<double3> &x, const std::vector<double3> &v, const std::vector<double> &pot, const std::vector<double3> &acc, const std::vector<double3> &jrk, const int extra_mode, const bool use_double_precision)
{
#ifdef HAS_HDF5
    hid_t file_id, group_id, attribute_id, dataspace_id;
    hsize_t dims[2] = {(hsize_t)N, 3};

    hid_t h5_float_type;
    if (use_double_precision) h5_float_type = H5T_IEEE_F64LE;
    else h5_float_type = H5T_IEEE_F32LE;
    file_id = H5Fcreate(file_name.c_str(), H5F_ACC_TRUNC, H5P_DEFAULT, H5P_DEFAULT);
    char group_name[32];
    sprintf(group_name, "/Step#%d", step_num);
    group_id = H5Gcreate2(file_id, group_name, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
    dataspace_id = H5Screate(H5S_SCALAR);
    attribute_id = H5Acreate2 (group_id, "Time", H5T_IEEE_F64LE, dataspace_id, H5P_DEFAULT, H5P_DEFAULT);
    H5Awrite(attribute_id, H5T_NATIVE_DOUBLE, &t);
    H5Sclose(dataspace_id);

    auto write_dataset = [&](const char dataset_name[], int ndims, double *data) {
        hid_t dataspace_id = H5Screate_simple(ndims, dims, NULL);
        char dataset_path[32];
        sprintf(dataset_path, "%s/%s", group_name, dataset_name);
        hid_t dataset_id = H5Dcreate2(file_id, dataset_path, h5_float_type, dataspace_id, H5P_DEFAULT, H5P_DEFAULT, H5P_DEFAULT);
        H5Dwrite(dataset_id, H5T_NATIVE_DOUBLE, H5S_ALL, H5S_ALL, H5P_DEFAULT, data);
        H5Dclose(dataset_id);
        H5Sclose(dataspace_id);
    };

    write_dataset("MASS", 1, (double*)m.data()); // casting away const...
    write_dataset("X",    2, (double*)x.data());
    write_dataset("V",    2, (double*)v.data());

    bool write_pot = (extra_mode     ) & 1;
    bool write_acc = (extra_mode >> 1) & 1;
    bool write_jrk = (extra_mode >> 2) & 1;
    if (write_pot) write_dataset("POT", 1, (double*)pot.data());
    if (write_acc) write_dataset("ACC", 2, (double*)acc.data());
    if (write_jrk) write_dataset("JRK", 2, (double*)jrk.data());

    H5Gclose(group_id);
    H5Fclose(file_id);
    H5close(); // If we don't do that (HDF5 1.10.5) then the file isn't really closed...
#else
    throw std::runtime_error("h5_write was called but compiled without HDF5 support");
#endif
}