convert_cellh5_inidat.py

'''
Created by Erick Urquilla, Department of Physics and Astronomy, University of Tennessee, Knoxville.
This script convert the cell hdf5 files into a particle_input.dat file that can be used as an initial condition
for the Emu code.
The cell hdf5 files are generated by the script write_particles_per_cell.py
The particle data is stored in hdf5 files with the following labels: cell_i_j_k.h5
where i, j, k are the cell indexes in the x, y, z directions respectively.
The ddf5 files has to be in the directory this script is run.
'''

import glob
import h5py
import numpy as np
import sys
import os
importpath = os.path.dirname(os.path.realpath(__file__))
sys.path.append(importpath)
sys.path.append(importpath+"/../initial_conditions")
from initial_condition_tools import uniform_sphere, write_particles
import amrex_plot_tools as amrex

def read_hdf5_file(file_name):
    """
    Reads an HDF5 file and returns its datasets as a dictionary.

    Parameters:
        file_name (str): Path to the HDF5 file.

    Returns:
        dict: A dictionary where keys are dataset names and values are NumPy arrays.
    """
    with h5py.File(file_name, 'r') as file:
        return {dataset: np.array(file[dataset]) for dataset in file.keys()}

# Find all files matching the pattern
file_pattern = 'cell_*_*_*.h5'
files = glob.glob(file_pattern)

data = read_hdf5_file(files[0])

labels_nf2=['pos_x','pos_y','pos_z', 'time', 'x', 'y', 'z', 'pupx', 'pupy', 'pupz', 'pupt', 'N00_Re', 'N01_Re', 'N01_Im', 'N11_Re', 'N00_Rebar', 'N01_Rebar', 'N01_Imbar', 'N11_Rebar', 'TrHN', 'Vphase']
labels_nf3=['pos_x','pos_y','pos_z','time','x', 'y', 'z', 'pupx', 'pupy', 'pupz', 'pupt', 'N00_Re', 'N01_Re', 'N01_Im', 'N02_Re', 'N02_Im', 'N11_Re', 'N12_Re', 'N12_Im' ,'N22_Re', 'N00_Rebar', 'N01_Rebar', 'N01_Imbar', 'N02_Rebar', 'N02_Imbar', 'N11_Rebar', 'N12_Rebar' ,'N12_Imbar', 'N22_Rebar', 'TrHN', 'Vphase']

NF = 0
if len(data.keys()) == len(labels_nf2):
    labels = labels_nf2
    NF = 2
elif len(data.keys()) == len(labels_nf3):
    labels = labels_nf3
    NF = 3
else:
    print(f'len(data.keys()) = {len(data.keys())} does not match any of the known labels')
    exit(1)

# Get variable keys
rkey, ikey = amrex.get_particle_keys(NF, ignore_pos=True)

# Generate the number of variables that describe each particle
n_variables = len(rkey)

# Generate the number of particles
n_particles = len(data[labels[0]])

# Initialize a NumPy array to store all particles
particles = np.zeros((n_particles, n_variables))

# Volume of the simulation box
volume = 1e5**3 # ccm

for label in rkey:

    print(f'Writing {label} to particles')
    particles[:, rkey[label]] = data[label]
    
    if (label.startswith('N') | label.startswith('Vphase')):
        particles[:, rkey[label]] /= volume

# Write particles initial condition file
write_particles(np.array(particles), NF, "particle_input.dat")