Advanced Computing Platform for Theoretical Physics

Commit bee41c3e authored by Chen Xia's avatar Chen Xia
Browse files

Add HelMod and fits_to_hdf5

parent 6f9c30fc
from ._bind_galpropcc import *
from . import tools
from . import populators
from .helmod import HelMod
""" HelMod solar modulation """
import tempfile
import subprocess
import numpy as np
from ._bind_galpropcc import SolarModulation
class HelMod(SolarModulation):
def __init__(self, particle, time_window):
self.particle = particle
self.time_window = time_window
SolarModulation.__init__(self)
def __call__(self, ekin_toa, ekin_is, flux_is, ZA):
with tempfile.NamedTemporaryFile(delete=True) as tmpf:
tmpname = tmpf.name
np.savetxt(tmpname, np.vstack([ekin_is, flux_is]).T)
subprocess.run(['helmod', '-i', tmpname, '-o', tmpname,
'-t', self.time_window, '-p', self.particle],
check=True)
_, flux_toa = np.loadtxt(tmpname, unpack=True)
return flux_toa
......@@ -2,7 +2,7 @@
import os
__all__ = ['dump_gcr']
__all__ = ['dump_gcr', 'fits_to_hdf5']
def dump_gcr(g, outfile, particles="all", mode='w-', **kwargs):
"""Dump full cosmic ray distribution.
......@@ -52,3 +52,53 @@ def dump_gcr(g, outfile, particles="all", mode='w-', **kwargs):
group.attrs['Z'] = gcr.Z()
group.attrs['A'] = gcr.A()
group.attrs['unit'] = r'cm^-2 s^-1 sr^-1 GeV^-1'
def fits_to_hdf5(fitsfilename, hdf5filename, mode='w-', dtype=None, **kwargs):
try:
import numpy as np
import h5py
from astropy.io import fits
except ImportError:
print("To convert GCR fits file to hdf5 format please install h5py "
"and astropy.")
return
outfile = os.path.abspath(hdf5filename)
os.makedirs(os.path.dirname(outfile), exist_ok=True)
dtype = np.float64 if dtype is None else dtype
with h5py.File(outfile, mode, **kwargs) as hdf5:
with fits.open(fitsfilename) as hdul:
data = hdul["PRIMARY"].data
if "NUCLEI" in hdul:
nuclei = hdul["NUCLEI"].data
hdf5.attrs["dim"] = data.ndim - 1
galz = np.asarray(hdul["GAL-Z"].data, dtype)
energy = np.asarray(hdul["Energy"].data, dtype)
dset = hdf5.create_dataset('z', data=galz)
dset.attrs['unit'] = 'kpc'
dset = hdf5.create_dataset('Ekin', data=energy * 1e-3)
dset.attrs['unit'] = 'GeV/n'
if data.ndim == 5:
galx = np.asarray(hdul["GAL-X"].data, dtype)
galy = np.asarray(hdul["GAL-Y"].data, dtype)
dset = hdf5.create_dataset('x', data=galx)
dset.attrs['unit'] = 'kpc'
dset = hdf5.create_dataset('y', data=galy)
dset.attrs['unit'] = 'kpc'
transpo = (3, 2, 1, 0)
dimension = "(x, y, z, Ekin/n)"
if data.ndim == 4:
galr = np.asarray(hdul["GAL-R"].data, dtype)
dset = hdf5.create_dataset('r', data=galr)
dset.attrs['unit'] = 'kpc'
transpo = (2, 1, 0)
dimension = "(r, z, Ekin/n)"
for i in range(len(nuclei)):
nucname, Z, A, _, _ = nuclei[i]
group = hdf5.create_group(nucname)
store = np.transpose(data[i, ...], transpo) / energy**2 * 1e3
group.create_dataset('flux', data=store)
group.attrs['Z'] = Z
group.attrs['A'] = A
group.attrs['unit'] = r'cm^-2 s^-1 sr^-1 GeV^-1'
group.attrs['dimention'] = dimension
Markdown is supported
0% or .
You are about to add 0 people to the discussion. Proceed with caution.
Finish editing this message first!
Please register or to comment