25/06/27 進捗報告

並列化を行うことで高速化を行おうとしたが数十万点であれば確かにその効果を感じた。

しかし、今回の数億点ある状況においては高速化されたというより、遅くなったようにさえ感じた。

そもそも、pygrackleの中はFortranで動いており、この高速化はあまり有用ではないことがある程度分かった。

そこで lookup tableをpygrackleで作成し、そのテーブルを線形補間を行うことにした。

修正は必要であるものの一旦のコードが下記。

from pygrackle import chemistry_data, setup_fluid_container
import numpy as np
from pygrackle.utilities.physical_constants import mass_hydrogen_cgs, sec_per_Myr, cm_per_mpc, boltzmann_constant_cgs

def compute_tcool_Gyr(rho_cgs, T, Z, x_e, chem, redshift=0.0):
    chem.a_value = 1.0 / (1.0 + redshift)
    metal_mass_fraction = np.asarray(Z) * 0.0127
    fc = setup_fluid_container(
        chem,
        density=np.asarray(rho_cgs),
        temperature=np.asarray(T),
        metal_mass_fraction=metal_mass_fraction,
        converge=True
    )
    data = fc.finalize_data()
    Lambda = data["cooling_rate"]  # erg cm^3 / s

    n_H = np.asarray(rho_cgs) / (0.6 * mass_hydrogen_cgs)
    n_e = np.asarray(x_e) * n_H
    tcool_s = (1.5 * boltzmann_constant_cgs * np.asarray(T)) / (n_e * -Lambda)  # [s]
    tcool_Gyr = tcool_s / (60 * 60 * 24 * 365 * 1e9) # [Gyr]
    return tcool_Gyr

rho_grid = np.logspace(-28, -22, 24)
T_grid = np.logspace(4, 8, 32)
Z_grid = np.logspace(-3, 0, 8)
xe_grid = np.linspace(0.01, 1.0, 8)

tcool_table = np.zeros((len(rho_grid), len(T_grid), len(Z_grid), len(xe_grid)), dtype=np.float64)

for i, rho in enumerate(rho_grid):
    for j, T in enumerate(T_grid):
        for k, Z in enumerate(Z_grid):
            for l, x_e in enumerate(xe_grid):
                chem = chemistry_data()
                chem.use_grackle = 1
                chem.with_radiative_cooling = 0
                chem.primordial_chemistry = 3
                chem.metal_cooling = 1
                chem.UVbackground = 1
                chem.grackle_data_file = "/home/nishihama/Osaka-gadget4/data/CloudyDatav3/CloudyData_UVB=HM2012.h5"
                chem.comoving_coordinates = 0
                chem.a_units = 1.0
                chem.a_value = 1.0
                chem.density_units = mass_hydrogen_cgs
                chem.length_units = cm_per_mpc
                chem.time_units = sec_per_Myr
                chem.set_velocity_units()
                chem.initialize()
                tcool = compute_tcool_Gyr(rho, T, Z, x_e, chem)
                tcool_table[i, j, k, l] = tcool

np.savez(
    "tcool_lookup_table_4d.v3.npz",
    rho_grid=rho_grid, T_grid=T_grid, Z_grid=Z_grid, xe_grid=xe_grid,
    tcool_table=tcool_table
)
print("4D Lookup table (rho, T, Z, x_e) saved as tcool_lookup_table_4d.npz")

黒色がGrackleで赤色が補完で作成したものである。

綺麗に計算できている。これを採用することにする。

使い方

import numpy as np
from scipy.interpolate import RegularGridInterpolator

# --- テーブル読み込み ---
data = np.load("tcool_lookup_table_4d.v3.npz")
rho_grid = data["rho_grid"]
T_grid = data["T_grid"]
Z_grid = data["Z_grid"]
xe_grid = data["xe_grid"]
tcool_table = data["tcool_table"]

# --- 4D補間関数を作成 ---
interp = RegularGridInterpolator(
    (rho_grid, T_grid, Z_grid, xe_grid),
    tcool_table,
    bounds_error=False, fill_value=np.nan  # グリッド外はNaN
)

補完したい値を代入する時は以下のようにする。

points = np.stack([rho_cgs, T, Z, x_e], axis=-1)
tcool_interp = interp(points)

tcool_interpが結果である。