utils.py

import os
import re
import h5py
import numpy as np
import pandas as pd
from tqdm import tqdm
import matplotlib.pyplot as plt
from scipy.integrate import quad
from scipy.constants import G, parsec


# Simulation Parameters

OMEGA_MATTER_0 = 0.3 
OMEGA_LAMBDA_0 = 0.7
OMEGA_0 = OMEGA_MATTER_0 / (OMEGA_MATTER_0 + OMEGA_LAMBDA_0)
h = 0.73   # Dimensionless Hubble Parameter
TimeBetSnapshot = 1.112
TimeOfFirstSnapshot = 0.0078125
INITIAL_REDSHIFT = 127.0
INITIAL_SCALE_FACTOR = 1/(1+INITIAL_REDSHIFT)




# Constants

G_SI = G  # Gravitational constant in SI units (m^3/kg/s^2)
solar_mass = 1.989e30  # Solar mass in kg
megaparsec = 1e6 * parsec  # Megaparsec in meters

H0 = 73.0  # Hubble constant in km/s/Mpc
# H0_SI = H0 * 1000 / (1e6 * parsec)  # Convert H0 to SI units (s^-1)
# rho_crit = (3 * H0_SI**2) / (8 * np.pi * G_SI) # Calculate critical density (rho_crit) in SI units (kg/m^3)

# v_ext_paper = np.array([89, -131, 17])  # External velocity in km/s





# Helper Functions

def critical_density(): 
    """
    Returns critical density in the units: h²kg/m³
    """
    return (3 * 100**2) / (8 * np.pi * G_SI * (10**6) * parsec**2)

def E(a):  # H(a)/H0
    return np.sqrt(OMEGA_MATTER_0 / a**3 + OMEGA_LAMBDA_0)

def integrand(a):
    return 1.0 / (a**3 * E(a)**3)

def D(a):
    integral, _ = quad(integrand, 0, a)
    norm, _ = quad(integrand, 0, 1)
    return E(a) * integral / norm

def f(a):
    return (OMEGA_MATTER_0 / (OMEGA_MATTER_0 + OMEGA_LAMBDA_0 * a**3)) ** 0.55

def H(a):
    """
    Returns H(a) in the units 
    """
    return H0 * E(a)





# Simulation Helper Functions

def mass_each_particle(N=128, grid_spacing=400./256.):
    """
    Returns mass of each particle in the units: 1e10 solar masses/h

    Parameters:
    N: number of particles
    grid_spacing: grid spacing in Mpc/h

    Returns:
    mass: mass of each particle in the units: 1e10 solar masses/h
    """
    L = (N-1)*grid_spacing*1000.  # in kpc/h
    return ( (3*OMEGA_0*(100**2)*(L**3)) / (8*np.pi*G_SI*(N**3)*(10**10)*solar_mass) ) * (10**3 * parsec)

def new_mass_each_particle(N=128, L=200000.0):
    """
    Returns mass of each particle in the units: 1e10 solar masses/h

    Parameters:
    N: number of particles
    L: Box Size in kpc/h

    Returns:
    mass: mass of each particle in the units: 1e10 solar masses/h
    """
    # L = (N-1)*grid_spacing*1000.  # in kpc/h
    return ( (3*OMEGA_0*(100**2)*(L**3)) / (8*np.pi*G_SI*(N**3)*(10**10)*solar_mass) ) * (10**3 * parsec)


def calculate_masstable(N_PARTICLES_LIST, N, grid_spacing=400./256.):
    """
    Returns mass table for each particle type in the units: 1e10 solar masses/h

    Parameters:
    N_PARTICLES_LIST: list of number of particles for each type.
    grid_spacing: grid spacing in Mpc/h

    Returns:
    masstable: mass table for each particle type in the units: 1e10 solar masses/h
    """
    mass = mass_each_particle(N=128, grid_spacing=grid_spacing)
    masstable = np.zeros(6)
    for i in range(len(N_PARTICLES_LIST)):
        if N_PARTICLES_LIST[i] > 0:
            masstable[i] = mass
    return masstable

def calculate_new_masstable(N_PARTICLES_LIST):
    """
    Returns mass table for each particle type in the units: 1e10 solar masses/h

    Parameters:
    N_PARTICLES_LIST: list of number of particles for each type.
    grid_spacing: grid spacing in Mpc/h

    Returns:
    masstable: mass table for each particle type in the units: 1e10 solar masses/h
    """
    mass = new_mass_each_particle()
    masstable = np.zeros(6)
    for i in range(len(N_PARTICLES_LIST)):
        if N_PARTICLES_LIST[i] > 0:
            masstable[i] = mass
    return masstable


# Creating Initial Conditions for Simulations

def create_initial_conditions_from_peculiar_velocity_field(
    df, velocity_columns=['G_PVX', 'G_PVY', 'G_PVZ'], N=128, LGpoint=63, box_size=198437.5, grid_spacing=400./256., remove_external_dipole=False, 
    external_dipole=np.array([89, -131, 17]), velocity_factor=0.12050692015791899, displacement_factor=0.0016137431506938527, 
    plots=False, bins=100, folder_name=None, phi=False
):
    """
    This function creates initial conditions for a cosmological simulation
    based on a peculiar velocity field. It uses the INITIAL_SCALE_FACTOR and other global parameters to create the initial conditions.

    Parameters:
    df : pandas.DataFrame
        DataFrame containing the peculiar velocity data.
        columns should include:
            'G_PVX', 'G_PVY', and 'G_PVZ' for the peculiar velocity components in km/s.
            'GX', 'GY', and 'GZ' for the positions in Mpc/h where (0, 0, 0) is the center of the simulation box or the LG.
    -----------
    velocity_columns : list of str
        List of column names in the DataFrame that contain the peculiar velocity components.
    N : int
        Number of grid points in each dimension.
    LGpoint : int
        Point in the grid where the Local Group is located.
        This point is used to calculate the positions of the particles in the simulation box.
    box_size : float
        Size of the simulation box in kpc/h.
    grid_spacing : float
        Spacing between grid points in Mpc/h.
    remove_external_dipole : bool
        If True, the external dipole is removed from the peculiar velocity field.
    external_dipole : np.ndarray
        External dipole to be removed from the peculiar velocity field.
    velocity_factor : float
        Factor to convert peculiar velocities from the present time to initial conditions.
        This factor is calculated based on the function compute_velocity_factor.
        # It is in units of check the units here
    displacement_factor : float
        Factor to calculate the displacement of the particles from the present time to the initial conditions.
        This factor is calculated based on the function compute_displacement_factor. 
        # It is in units of check the units here
    plots : bool
        If True, histograms of the displacements and peculiar velocities are plotted.
    bins : int
        Number of bins for the histograms.
    folder_name : str
        Name of the folder where the initial conditions will be saved.
        The initial conditions will be saved in a file named 'initial_conditions.hdf5' in this folder.
    phi : bool
        If True, the potential of the particles is included in the initial conditions.
        The potential is calculated based on the function compute_potential.
        # It is in units of check the units here
    
    
    -----------

    Returns:
    df : pandas.DataFrame
        DataFrame with the initial conditions for the simulation.
        The columns 'X', 'Y', and 'Z' are added, representing the positions of the particles today in kpc/h wrt the box.
        The columns 'vx_initial', 'vy_initial', and 'vz_initial' are added, representing the peculiar velocities of the particles in km/s at the initial time.
        The columns 'dx', 'dy', and 'dz' are added, representing the displacements of the particles from the present time to the initial conditions in kpc/h.
        The columns 'Xi', 'Yi', and 'DZ' are added, representing the positions of the particles at the initial time in kpc/h wrt the box.

    """
    assert int(np.round(len(df) ** (1/3))) == N, f"Number of particles {len(df)} is not a cube of an integer. Please check the grid size."

    N_PARTICLES_GAS = 0
    N_PARTICLES_HALO = len(df)
    N_PARTICLES_DISK = 0
    N_PARTICLES_BULGE = 0
    N_PARTICLES_STARS = 0
    N_PARTICLES_BNDRY = 0
    N_PARTICLES_LIST = [N_PARTICLES_GAS, N_PARTICLES_HALO, N_PARTICLES_DISK, N_PARTICLES_BULGE, N_PARTICLES_STARS, N_PARTICLES_BNDRY]
    MASSTABALE = calculate_masstable(N_PARTICLES_LIST, N, grid_spacing=grid_spacing)
    print("Mass table: ", MASSTABALE)
    print("Number of particles: ", N_PARTICLES_LIST)

    df["X"] = (df["GX"] + LGpoint*grid_spacing)*1000.
    df["Y"] = (df["GY"] + LGpoint*grid_spacing)*1000.
    df["Z"] = (df["GZ"] + LGpoint*grid_spacing)*1000.

    if remove_external_dipole:
        df[velocity_columns[0]] = df[velocity_columns[0]] - external_dipole[0]
        df[velocity_columns[1]] = df[velocity_columns[1]] - external_dipole[1]
        df[velocity_columns[2]] = df[velocity_columns[2]] - external_dipole[2]
    
    df["vx_initial"] = df[velocity_columns[0]] * velocity_factor
    df["vy_initial"] = df[velocity_columns[1]] * velocity_factor
    df["vz_initial"] = df[velocity_columns[2]] * velocity_factor

    df["dx"] = displacement_factor * df["vx_initial"] * 1000  # kpc/h
    df["dy"] = displacement_factor * df["vy_initial"] * 1000
    df["dz"] = displacement_factor * df["vz_initial"] * 1000

    df["Xi"] = df["X"] - df["dx"]  # (kpc/h)
    df["Yi"] = df["Y"] - df["dy"]
    df["Zi"] = df["Z"] - df["dz"]
    
    df["vx_gadget"] = df["vx_initial"] / INITIAL_SCALE_FACTOR**0.5 # (km/s)
    df["vy_gadget"] = df["vy_initial"] / INITIAL_SCALE_FACTOR**0.5
    df["vz_gadget"] = df["vz_initial"] / INITIAL_SCALE_FACTOR**0.5

    assert df["X"].values[-1] == box_size, "Last X coordinate is not equal to box size. Please check the box size."

    if plots:
        fig, ax = plt.subplots(3, 3, figsize=(18, 15))
        ax = ax.flatten()
        df["vx_initial"].hist(bins=bins, label="PVX Initial (km/s)", alpha=0.5, ax=ax[0])
        df["vy_initial"].hist(bins=bins, label="PVY Initial (km/s)", alpha=0.5, ax=ax[1])
        df["vz_initial"].hist(bins=bins, label="PVZ Initial (km/s)", alpha=0.5, ax=ax[2])

        df["vx_gadget"].hist(bins=bins, label="PVX Gadget (km/s)", alpha=0.5, ax=ax[3])
        df["vy_gadget"].hist(bins=bins, label="PVY Gadget (km/s)", alpha=0.5, ax=ax[4])
        df["vz_gadget"].hist(bins=bins, label="PVZ Gadget (km/s)", alpha=0.5, ax=ax[5])

        df["dx"].hist(bins=bins, label="dX (kpc/h)", alpha=0.5, ax=ax[6])
        df["dy"].hist(bins=bins, label="dY (kpc/h)", alpha=0.5, ax=ax[7])
        df["dz"].hist(bins=bins, label="dZ (kpc/h)", alpha=0.5, ax=ax[8])

        for a in ax:
            a.legend()
            a.set_ylabel("Count")
        
        plt.tight_layout()
        fig.show()

    file_path = f"Simulations/{folder_name}/initial_conditions.hdf5"

    with h5py.File(file_path, 'w') as f:
        header = f.create_group("Header")
        header.attrs['BoxSize'] = box_size
        header.attrs['Flag_Cooling'] = np.int64(0)
        header.attrs['Flag_Entropy_ICs'] = np.array([0, 0, 0, 0, 0, 0])
        header.attrs['Flag_Feedback'] = np.int64(0)
        header.attrs['Flag_Metals'] = np.int64(0)
        header.attrs['Flag_Sfr'] = np.int64(0)
        header.attrs['Flag_StellarAge'] = np.int64(0)
        header.attrs['HubbleParam'] = h
        header.attrs['MassTable'] = np.array(MASSTABALE)
        header.attrs['NumFilesPerSnapshot'] = np.int64(1)
        header.attrs['NumPart_ThisFile'] = np.array(N_PARTICLES_LIST)
        header.attrs['NumPart_Total'] = np.array(N_PARTICLES_LIST)
        header.attrs['NumPart_Total_HighWord'] = np.array([0, 0, 0, 0, 0, 0])
        header.attrs['Omega0'] = OMEGA_0 #OMEGA_MATTER_0
        header.attrs['OmegaLambda'] = OMEGA_LAMBDA_0
        header.attrs['Redshift'] = INITIAL_REDSHIFT
        header.attrs['Time'] = INITIAL_SCALE_FACTOR

        parttype1 = f.create_group("PartType1")
        
        coordinates_data = np.array(df[["Xi", "Yi", "Zi"]])
        parttype1.create_dataset("Coordinates", dtype=float, shape=(2097152,3), data=coordinates_data)

        velocities_data = np.array(df[["vx_gadget", "vy_gadget", "vz_gadget"]])
        parttype1.create_dataset("Velocities", dtype=float, shape=(2097152,3), data=velocities_data)

        particle_ids = np.array(df.index)
        parttype1.create_dataset("ParticleIDs", dtype=int, data=particle_ids)

        if phi:
            potential_data = np.array(df[["Phi"]])
            parttype1.create_dataset("Potential", dtype=float, data=potential_data)
        
        f.close()
    return df







# View Initial Conditions and Simulation Snapshots

def check_simulation_results(folder_name, snapshotnumber=None, initial_conditions=False, phi=False, N=128, box_size=198437.5, plots=True, bins=200, print_values=True):
    """
    Check the simulation results for a given folder and snapshot number.
    
    Parameters:
    folder_name: str
        The name of the folder containing the simulation results.
    snapshotnumber: str
        The snapshot number to check in string format. If None, the initial conditions are checked.
    initial_conditions:  bool
        If True, check the initial conditions instead of a snapshot. 
    phi: bool
        If True, include the potential in the output DataFrame. Only works when potential is calculated.
    N: int
        The number of particles in the simulation to the power (1/3).
    box_size: float
        The size of the simulation box in kpc/h.
    plots: bool
        If True, plot histograms of the displacements and velocities.
    bins: int
        The number of bins for the histograms.

    ----------

    Returns:
    df: pandas.DataFrame
        DataFrame containing the simulation results with the following columns:
        - X: X coordinates
        - Y: Y coordinates
        - Z: Z coordinates
        - VX: X velocities
        - VY: Y velocities
        - VZ: Z velocities
        - PVX: Proper X velocities
        - PVY: Proper Y velocities
        - PVZ: Proper Z velocities
        - dx: Displacement in X
        - dy: Displacement in Y
        - dz: Displacement in Z
        - phi: Potential (if included)
    """
    if snapshotnumber is not None:
        path = f"./Simulations/{folder_name}/Output/snapshot3halosc_{snapshotnumber}.hdf5"
        time = TimeOfFirstSnapshot * TimeBetSnapshot**int(snapshotnumber)
        z = (1/time) -1
    
    elif initial_conditions:
        path = f"./Simulations/{folder_name}/initial_conditions.hdf5"
        z = INITIAL_REDSHIFT
        time = 1/(1+z)

    if print_values:
        print("Redshift: ", z)
        print("Scale Factor: ", time)
    
    f = h5py.File(path)
    df = pd.DataFrame(np.array(f['PartType1']["Coordinates"]), columns=["X", "Y", "Z"])
    df["VX"], df["VY"], df["VZ"] = np.array(f['PartType1']["Velocities"])[:, 0], np.array(f['PartType1']["Velocities"])[:, 1], np.array(f['PartType1']["Velocities"])[:, 2]
    df["PVX"] = df["VX"] * time**0.5
    df["PVY"] = df["VY"] * time**0.5
    df["PVZ"] = df["VZ"] * time**0.5
    if phi:
        df["phi"] = np.array(f['PartType1']["Potential"])
    df["ids"] = np.array(f['PartType1']["ParticleIDs"])
    df.sort_values(by=['ids'], inplace=True)
    df.reset_index(drop=True, inplace=True)
    df.drop(columns=["ids"], inplace=True)

    # Calculate Displacements
    # if initial_conditions:
    #     X = Y = Z = np.arange(0, N) * (box_size/(N-1))
    #     X, Y, Z = np.meshgrid(X, Y, Z, indexing="ij")
    #     df["Xi"] = X.flatten()
    #     df["Yi"] = Y.flatten()
    #     df["Zi"] = Z.flatten()
    
    # if snapshotnumber is not None:
    #     path_initial = f"./Simulations/{folder_name}/Output/snapshot3halosc_000.hdf5"
    #     ff = h5py.File(path_initial)
    #     idf = pd.DataFrame(np.array(ff['PartType1']["Coordinates"]), columns=["X", "Y", "Z"])
    #     idf["ids"] = np.array(ff['PartType1']["ParticleIDs"])
    #     idf.sort_values(by=['ids'], inplace=True)
    #     idf.reset_index(drop=True, inplace=True)
    #     idf.drop(columns=["ids"], inplace=True)
    #     df["Xi"] = idf["X"]
    #     df["Yi"] = idf["Y"]
    #     df["Zi"] = idf["Z"]


    X = Y = Z = np.arange(0, N) * (box_size/(N-1))
    X, Y, Z = np.meshgrid(X, Y, Z, indexing="ij")
    df["Xi"] = X.flatten()
    df["Yi"] = Y.flatten()
    df["Zi"] = Z.flatten()

    def periodic_displacement(delta, box_size):
        return (delta + box_size / 2) % box_size - box_size / 2

    df["dx"] = periodic_displacement(df["X"] - df["Xi"], box_size)
    df["dy"] = periodic_displacement(df["Y"] - df["Yi"], box_size)
    df["dz"] = periodic_displacement(df["Z"] - df["Zi"], box_size)


    if plots:
        fig, ax = plt.subplots(3, 3, figsize=(18, 15))
        ax = ax.flatten()
        df["dx"].hist(bins=bins, label="dX (kpc/h)", alpha=0.5, ax=ax[0])
        df["dy"].hist(bins=bins, label="dY (kpc/h)", alpha=0.5, ax=ax[1])
        df["dz"].hist(bins=bins, label="dZ (kpc/h)", alpha=0.5, ax=ax[2])

        df["VX"].hist(bins=bins, label="VX (km/s)", alpha=0.5, ax=ax[3])
        df["VY"].hist(bins=bins, label="VY (km/s)", alpha=0.5, ax=ax[4])
        df["VZ"].hist(bins=bins, label="VZ (km/s)", alpha=0.5, ax=ax[5])

        df["PVX"].hist(bins=bins, label="PVX (km/s)", alpha=0.5, ax=ax[6])
        df["PVY"].hist(bins=bins, label="PVY (km/s)", alpha=0.5, ax=ax[7])
        df["PVZ"].hist(bins=bins, label="PVZ (km/s)", alpha=0.5, ax=ax[8])

        for a in ax:
            a.legend()
            a.set_ylabel("Count")
        
        plt.tight_layout()
        fig.show()
    
    return df




def check_new_simulation_results(folder_name, snapshotnumber=None, initial_conditions=False, phi=False, N=128, box_size=200000.0, GS=400./256.*1000.0, plots=True, bins=200, print_values=True):
    """
    Check the simulation results for a given folder and snapshot number.
    
    Parameters:
    folder_name: str
        The name of the folder containing the simulation results.
    snapshotnumber: str
        The snapshot number to check in string format. If None, the initial conditions are checked.
    initial_conditions:  bool
        If True, check the initial conditions instead of a snapshot. 
    phi: bool
        If True, include the potential in the output DataFrame. Only works when potential is calculated.
    N: int
        The number of particles in the simulation to the power (1/3).
    box_size: float
        The size of the simulation box in kpc/h.
    plots: bool
        If True, plot histograms of the displacements and velocities.
    bins: int
        The number of bins for the histograms.

    ----------

    Returns:
    df: pandas.DataFrame
        DataFrame containing the simulation results with the following columns:
        - X: X coordinates
        - Y: Y coordinates
        - Z: Z coordinates
        - VX: X velocities
        - VY: Y velocities
        - VZ: Z velocities
        - PVX: Proper X velocities
        - PVY: Proper Y velocities
        - PVZ: Proper Z velocities
        - dx: Displacement in X
        - dy: Displacement in Y
        - dz: Displacement in Z
        - phi: Potential (if included)
    """
    if snapshotnumber is not None:
        path = f"./Simulations/{folder_name}/Output/snapshot3halosc_{snapshotnumber}.hdf5"
        time = TimeOfFirstSnapshot * TimeBetSnapshot**int(snapshotnumber)
        z = (1/time) -1
    
    elif initial_conditions:
        path = f"./Simulations/{folder_name}/initial_conditions.hdf5"
        z = INITIAL_REDSHIFT
        time = 1/(1+z)
    
    if print_values:
        print("Redshift: ", z)
        print("Scale Factor: ", time)
    
    f = h5py.File(path)
    df = pd.DataFrame(np.array(f['PartType1']["Coordinates"]), columns=["X", "Y", "Z"])
    df["VX"], df["VY"], df["VZ"] = np.array(f['PartType1']["Velocities"])[:, 0], np.array(f['PartType1']["Velocities"])[:, 1], np.array(f['PartType1']["Velocities"])[:, 2]
    df["PVX"] = df["VX"] * time**0.5
    df["PVY"] = df["VY"] * time**0.5
    df["PVZ"] = df["VZ"] * time**0.5
    if phi:
        df["phi"] = np.array(f['PartType1']["Potential"])
    df["ids"] = np.array(f['PartType1']["ParticleIDs"])
    df.sort_values(by=['ids'], inplace=True)
    df.reset_index(drop=True, inplace=True)
    df.drop(columns=["ids"], inplace=True)

    # Calculate Displacements
    # if initial_conditions:
    #     X = Y = Z = np.arange(0, N) * (box_size/(N-1))
    #     X, Y, Z = np.meshgrid(X, Y, Z, indexing="ij")
    #     df["Xi"] = X.flatten()
    #     df["Yi"] = Y.flatten()
    #     df["Zi"] = Z.flatten()
    
    # if snapshotnumber is not None:
    #     path_initial = f"./Simulations/{folder_name}/Output/snapshot3halosc_000.hdf5"
    #     ff = h5py.File(path_initial)
    #     idf = pd.DataFrame(np.array(ff['PartType1']["Coordinates"]), columns=["X", "Y", "Z"])
    #     idf["ids"] = np.array(ff['PartType1']["ParticleIDs"])
    #     idf.sort_values(by=['ids'], inplace=True)
    #     idf.reset_index(drop=True, inplace=True)
    #     idf.drop(columns=["ids"], inplace=True)
    #     df["Xi"] = idf["X"]
    #     df["Yi"] = idf["Y"]
    #     df["Zi"] = idf["Z"]


    # X = Y = Z = np.arange(0, N) * (box_size/(N-1))
    X = Y = Z = np.arange(0, N) * GS + 0.5* GS
    X, Y, Z = np.meshgrid(X, Y, Z, indexing="ij")
    df["Xi"] = X.flatten()
    df["Yi"] = Y.flatten()
    df["Zi"] = Z.flatten()

    def periodic_displacement(delta, box_size):
        return (delta + box_size / 2) % box_size - box_size / 2

    df["dx"] = periodic_displacement(df["X"] - df["Xi"], box_size)
    df["dy"] = periodic_displacement(df["Y"] - df["Yi"], box_size)
    df["dz"] = periodic_displacement(df["Z"] - df["Zi"], box_size)

    # df["dx"] = periodic_displacement(df["Xi"] - df["X"], box_size)
    # df["dy"] = periodic_displacement(df["Yi"] - df["Y"], box_size)
    # df["dz"] = periodic_displacement(df["Zi"] - df["Z"], box_size)


    if plots:
        fig, ax = plt.subplots(3, 3, figsize=(18, 15))
        ax = ax.flatten()
        df["dx"].hist(bins=bins, label="dX (kpc/h)", alpha=0.5, ax=ax[0])
        df["dy"].hist(bins=bins, label="dY (kpc/h)", alpha=0.5, ax=ax[1])
        df["dz"].hist(bins=bins, label="dZ (kpc/h)", alpha=0.5, ax=ax[2])

        df["VX"].hist(bins=bins, label="VX (km/s)", alpha=0.5, ax=ax[3])
        df["VY"].hist(bins=bins, label="VY (km/s)", alpha=0.5, ax=ax[4])
        df["VZ"].hist(bins=bins, label="VZ (km/s)", alpha=0.5, ax=ax[5])

        df["PVX"].hist(bins=bins, label="PVX (km/s)", alpha=0.5, ax=ax[6])
        df["PVY"].hist(bins=bins, label="PVY (km/s)", alpha=0.5, ax=ax[7])
        df["PVZ"].hist(bins=bins, label="PVZ (km/s)", alpha=0.5, ax=ax[8])

        for a in ax:
            a.legend()
            a.set_ylabel("Count")
        
        plt.tight_layout()
        fig.show()
    
    return df

def track_particle_evolution(folder_name, particle_index=1040319, N=128, box_size=200000.0, GS=400./256.*1000.0):
    """
    Tracks and plots the evolution of a particle's properties across snapshots in a simulation folder.
    
    Parameters:
    folder_name: str
        The name of the folder containing the simulation snapshots.
    particle_index: int
        The index (0 to N^3-1) of the particle to track.
    N: int
        Number of particles per dimension (N^3 total).
    box_size: float
        Size of the simulation box in kpc/h.
    GS: float
        Grid spacing in kpc/h.
    
    Returns:
    df_all: pandas.DataFrame
        DataFrame with time evolution of the selected particle.
    """
    snapshots = []
    snapshot_folder = f"./Simulations/{folder_name}/Output"
    regex = re.compile(r"snapshot3halosc_(\d+)\.hdf5")
    
    for fname in sorted(os.listdir(snapshot_folder)):
        match = regex.match(fname)
        if match:
            snapshots.append(match.group(1))
    
    all_data = []

    for snap in tqdm(snapshots, desc="Reading snapshots"):
        df = check_new_simulation_results(folder_name, snapshotnumber=snap, plots=False, N=N, box_size=box_size, GS=GS, print_values=False)
        time = TimeOfFirstSnapshot * TimeBetSnapshot**int(snap)
        a = time  # scale factor
        particle = df.iloc[particle_index].copy()
        particle["a"] = a
        particle["snapshot"] = snap
        all_data.append(particle)
    
    import pandas as pd
    df_all = pd.DataFrame(all_data)
    
    # Plotting evolution
    fig, ax = plt.subplots(2, 3, figsize=(16, 10))
    ax = ax.flatten()
    
    ax[0].plot(df_all["a"], np.abs(df_all["dx"]), label="dx")
    ax[1].plot(df_all["a"], np.abs(df_all["dy"]), label="dy")
    ax[2].plot(df_all["a"], np.abs(df_all["dz"]), label="dz")
    
    ax[3].plot(df_all["a"], np.abs(df_all["PVX"]), label="PVX")
    ax[4].plot(df_all["a"], np.abs(df_all["PVY"]), label="PVY")
    ax[5].plot(df_all["a"], np.abs(df_all["PVZ"]), label="PVZ")

    for a in ax:
        a.set_xlabel("Scale factor a")
        a.set_ylabel("Value")
        a.legend()
    
    plt.suptitle(f"Evolution of Particle {particle_index}")
    plt.tight_layout()
    plt.show()
    
    return df_all



def create_initial_conditions_from_peculiar_velocity_field_for_AdrianScript(
    df, N=128, box_size=198437.5, grid_spacing=400./256., 
    plots=False, bins=100, folder_name=None, phi=False
):
    """
    This function creates initial conditions for a cosmological simulation
    based on a peculiar velocity field. It uses the INITIAL_SCALE_FACTOR and other global parameters to create the initial conditions.

    Parameters:
    df : pandas.DataFrame
        DataFrame containing the peculiar velocity data.
        columns should include:
            'G_PVX', 'G_PVY', and 'G_PVZ' for the peculiar velocity components in km/s.
            'GX', 'GY', and 'GZ' for the positions in Mpc/h where (0, 0, 0) is the center of the simulation box or the LG.
    -----------
    velocity_columns : list of str
        List of column names in the DataFrame that contain the peculiar velocity components.
    N : int
        Number of grid points in each dimension.
    LGpoint : int
        Point in the grid where the Local Group is located.
        This point is used to calculate the positions of the particles in the simulation box.
    box_size : float
        Size of the simulation box in kpc/h.
    grid_spacing : float
        Spacing between grid points in Mpc/h.
    remove_external_dipole : bool
        If True, the external dipole is removed from the peculiar velocity field.
    external_dipole : np.ndarray
        External dipole to be removed from the peculiar velocity field.
    velocity_factor : float
        Factor to convert peculiar velocities from the present time to initial conditions.
        This factor is calculated based on the function compute_velocity_factor.
        # It is in units of check the units here
    displacement_factor : float
        Factor to calculate the displacement of the particles from the present time to the initial conditions.
        This factor is calculated based on the function compute_displacement_factor. 
        # It is in units of check the units here
    plots : bool
        If True, histograms of the displacements and peculiar velocities are plotted.
    bins : int
        Number of bins for the histograms.
    folder_name : str
        Name of the folder where the initial conditions will be saved.
        The initial conditions will be saved in a file named 'initial_conditions.hdf5' in this folder.
    phi : bool
        If True, the potential of the particles is included in the initial conditions.
        The potential is calculated based on the function compute_potential.
        # It is in units of check the units here
    
    
    -----------

    Returns:
    df : pandas.DataFrame
        DataFrame with the initial conditions for the simulation.
        The columns 'X', 'Y', and 'Z' are added, representing the positions of the particles today in kpc/h wrt the box.
        The columns 'vx_initial', 'vy_initial', and 'vz_initial' are added, representing the peculiar velocities of the particles in km/s at the initial time.
        The columns 'dx', 'dy', and 'dz' are added, representing the displacements of the particles from the present time to the initial conditions in kpc/h.
        The columns 'Xi', 'Yi', and 'DZ' are added, representing the positions of the particles at the initial time in kpc/h wrt the box.

    """
    assert int(np.round(len(df) ** (1/3))) == N, f"Number of particles {len(df)} is not a cube of an integer. Please check the grid size."

    N_PARTICLES_GAS = 0
    N_PARTICLES_HALO = len(df)
    N_PARTICLES_DISK = 0
    N_PARTICLES_BULGE = 0
    N_PARTICLES_STARS = 0
    N_PARTICLES_BNDRY = 0
    N_PARTICLES_LIST = [N_PARTICLES_GAS, N_PARTICLES_HALO, N_PARTICLES_DISK, N_PARTICLES_BULGE, N_PARTICLES_STARS, N_PARTICLES_BNDRY]
    MASSTABALE = calculate_masstable(N_PARTICLES_LIST, N, grid_spacing=grid_spacing)
    print("Mass table: ", MASSTABALE)
    print("Number of particles: ", N_PARTICLES_LIST)
    
    df["vx_gadget"] = df["PVX"] / INITIAL_SCALE_FACTOR**0.5 # (km/s)
    df["vy_gadget"] = df["PVY"] / INITIAL_SCALE_FACTOR**0.5
    df["vz_gadget"] = df["PVZ"] / INITIAL_SCALE_FACTOR**0.5

    # assert df["X"].values[-1] == box_size, "Last X coordinate is not equal to box size. Please check the box size."

    if plots:
        fig, ax = plt.subplots(3, 3, figsize=(18, 15))
        ax = ax.flatten()
        df["PVX"].hist(bins=bins, label="PVX Initial (km/s)", alpha=0.5, ax=ax[0])
        df["PVY"].hist(bins=bins, label="PVY Initial (km/s)", alpha=0.5, ax=ax[1])
        df["PVZ"].hist(bins=bins, label="PVZ Initial (km/s)", alpha=0.5, ax=ax[2])

        df["vx_gadget"].hist(bins=bins, label="PVX Gadget (km/s)", alpha=0.5, ax=ax[3])
        df["vy_gadget"].hist(bins=bins, label="PVY Gadget (km/s)", alpha=0.5, ax=ax[4])
        df["vz_gadget"].hist(bins=bins, label="PVZ Gadget (km/s)", alpha=0.5, ax=ax[5])

        df["dx"].hist(bins=bins, label="dX (kpc/h)", alpha=0.5, ax=ax[6])
        df["dy"].hist(bins=bins, label="dY (kpc/h)", alpha=0.5, ax=ax[7])
        df["dz"].hist(bins=bins, label="dZ (kpc/h)", alpha=0.5, ax=ax[8])

        for a in ax:
            a.legend()
            a.set_ylabel("Count")
        
        plt.tight_layout()
        fig.show()

    file_path = f"Simulations/{folder_name}/initial_conditions.hdf5"

    with h5py.File(file_path, 'w') as f:
        header = f.create_group("Header")
        header.attrs['BoxSize'] = box_size
        header.attrs['Flag_Cooling'] = np.int64(0)
        header.attrs['Flag_Entropy_ICs'] = np.array([0, 0, 0, 0, 0, 0])
        header.attrs['Flag_Feedback'] = np.int64(0)
        header.attrs['Flag_Metals'] = np.int64(0)
        header.attrs['Flag_Sfr'] = np.int64(0)
        header.attrs['Flag_StellarAge'] = np.int64(0)
        header.attrs['HubbleParam'] = h
        header.attrs['MassTable'] = np.array(MASSTABALE)
        header.attrs['NumFilesPerSnapshot'] = np.int64(1)
        header.attrs['NumPart_ThisFile'] = np.array(N_PARTICLES_LIST)
        header.attrs['NumPart_Total'] = np.array(N_PARTICLES_LIST)
        header.attrs['NumPart_Total_HighWord'] = np.array([0, 0, 0, 0, 0, 0])
        header.attrs['Omega0'] = OMEGA_0 #OMEGA_MATTER_0
        header.attrs['OmegaLambda'] = OMEGA_LAMBDA_0
        header.attrs['Redshift'] = INITIAL_REDSHIFT
        header.attrs['Time'] = INITIAL_SCALE_FACTOR

        parttype1 = f.create_group("PartType1")
        
        coordinates_data = np.array(df[["X", "Y", "Z"]])
        parttype1.create_dataset("Coordinates", dtype=float, shape=(2097152,3), data=coordinates_data)

        velocities_data = np.array(df[["vx_gadget", "vy_gadget", "vz_gadget"]])
        parttype1.create_dataset("Velocities", dtype=float, shape=(2097152,3), data=velocities_data)

        particle_ids = np.array(df.index)
        parttype1.create_dataset("ParticleIDs", dtype=int, data=particle_ids)

        if phi:
            potential_data = np.array(df[["Phi"]])
            parttype1.create_dataset("Potential", dtype=float, data=potential_data)
        
        f.close()
    return df









### Misc


# def mass_each_particle(N=128, grid_spacing=400000./256.):
#     """
#     Returns mass of each particle in the units: 1e10 solar masses/h

#     Parameters:
#     N: number of particles
#     grid_spacing: grid spacing in kpc/h

#     Returns:
#     mass: mass of each particle in the units: 1e10 solar masses/h
#     """
#     L = (N-1)*grid_spacing
#     return ( (3*OMEGA_0*(100**2)*(L**3)) / (8*np.pi*G_SI*(N**3)*(10**10)*solar_mass) ) * (10**3 * parsec)



# def calculate_masstable(N_PARTICLES_LIST, BOX_SIZE):
#     """
#     Return mass of each halo particle in 1.0e10 solar masses/h
#     """
#     n_halo = N_PARTICLES_LIST[1]
#     volume = (BOX_SIZE * 1000 * parsec / h)**3    # m^3
#     mass_halo = (rho_crit * OMEGA_MATTER_0 * volume) / n_halo  # kg
#     mass_halo = (mass_halo * h) / (solar_mass*1.0e10)  # 1.0e10 solar masses/h
#     return [0., mass_halo, 0., 0., 0., 0.]