"""
Routine to read a CRIPTIC checkpoint
"""
import numpy as np
import os.path as osp
import glob
import warnings
import sys
import h5py
from copy import deepcopy
try:
import astropy.units as u
import astropy.constants as const
from astropy.constants import m_p, m_e, c, h
except:
u = None
from .sr import _Tp, parttypes
[docs]
def readchk(chkname, units=True, ke=True, ndot=True,
sort=False, meta=False, sources_only=False):
"""
Read a CRIPTIC checkpoint
Parameters
chkname : string
The name of the checkpoint to read
units : bool
If True, all quantities are returned as astropy.Quantity
objects with appropriate units; if False data are returned as
flat numpy arrays.
ke : bool
If True, energies will be computed for all momentum
quantities; only allowed if units is also True
ndot : bool
If True, compute number of CRs injected per unit time for
each source
sort : bool
If True, all packets, deleted packets, and sources will be
sorted in order of unique ID
meta : bool
If True, only the metadata are read and returned
sources_only : bool
If True, only sources are returned, not packets
Returns
out : dict
a dict containing the following entries (some optional, may
or may not be included depending on contents of the
checkpoint file being read and the keywords provided)
"t" : float or astropy.Quantity
time of output
"dt" : float or astropy.Quantity
size of next time step
"step" : int
next step number
"uid" : int
value of the unique ID counter
"part_names" : list of string
names of particles that correspond to each particle type
"packets" : dict
a dict containing information on packets, with the following
entries:
"x" : ndarray or astropy.Quantity, size (n,3)
positions of all packets
"ptype" : ndarray, size (n)
particle type numbers of all packets
"uid" : ndarray, size (n)
unique IDs of all packets
"source" : ndarray, size (n)
unique IDs of sources of all packets
"p" : ndarray or astropy.Quantity, size (n)
momenta of all packets
"w" : ndarray, size (n)
statistical weights of all packets
"tinj" : ndarray or astropy.Quantity, size (n)
times at which each packet was injected
"wgtinj" : ndarray or astropy.Quantity, size (n)
statistical weights of each packet at injection
"gr" : ndarray or astropy.Quantity, size (n)
grammage traversed by each packet
"mu " : ndarray, size(n)
pitch angles of CR packets
"sec" : ndarray, size (n)
boolean that is true if packet is a secondary, false if not
"lossmech" : list of string
names of loss mechanisms
"T" : ndarray or astropy.Quantity, size (n)
CR packet kinetic energies
"pres" : ndarray or astropy.Quantity, size(n)
CR pressure divided by k_B at position of each packet;
pressures only include packets with gyroradii >= the
gyroradius of this packet
"presGrad" : ndarray or astropy.Quantity, size (n,3)
gradient of pres
"eden" : ndarray or astropy.Quantity, size(n)
CR energy density at position of each packet, including only
CRs with gyroradii >= the gyroradius of this packet
"edenGrad" : ndarray or astropy.Quantity, size(n,3)
gradient of eden
"nden" : ndarray or astropy.Quantity, size(n)
CR number density at position of each packet, including only
CRs with gyroradii >= the gyroradius of this packet
"ndenGrad" : ndarray or astropy.Quantity, size(n,3)
gradient of nden
delpackets : dict
a dict containing information on packets that were deleted
ce the last checkpoint, with the following entries:
"x" : ndarray or astropy.Quantity, size (m,3)
positions of all packets
"ptype" : ndarray, size (m)
particle types of all packets
"uid" : ndarray, size (m)
unique IDs of all packets
"source" : ndarray, size (m)
unique IDs of sources of all packets
"p" : ndarray or astropy.Quantity, size (m)
momenta of all packets
"w" : ndarray, size (m)
statistical weights of all packets
"tinj" : ndarray or astropy.Quantity, size (m)
times at which each packet was injected
"wgtinj" : ndarray, size (m)
statistical weights of each packet at injection
"gr" : ndarray or astropy.Quantity, size (m)
grammage traversed by each packet
"mu " : ndarray, size(n)
pitch angles of CR packets
"T" : ndarray or astropy.Quantity, size (n)
CR packet kinetic energies
sources : dict
a dict containing information on sources, with the following
entries:
"x" : ndarray or astropy.Quantity, size (i,3)
positions of all sources
"ptype" : ndarray, size(i)
particle types produced by sources
"uid" : ndarray, size (i)
unique IDs of all sources
"p0" : ndarray or astropy.Quantity, size (i)
lower limit of momentum of packets injected by each source
"p1" : ndarray or astropy.Quantity, size (i)
upper limit of momentum of packets injected by each source
"q" : ndarray, size (i)
index of source momentum distribution
"mu0 " : ndarray, size(n)
minimum pitch angle of packets injected by each source
"mu1 " : ndarray, size(n)
maximum pitch angle of packets injected by each source
"coef" : ndarray or astropy.Quantity, size (i)
coefficient of source momentum distribution
"ndot" : ndarray or astropy.Quantity, size (i)
number of CRs injected per unit time by source
losses : dict
a dict containing information on losses, with the
following entries:
"lossrate" : ndarray or astropy.Quantity, size(n,m)
rate of catastrophic loss for all loss processes
"dpdt" : ndarray or astropy.Quantity, size(n,m)
rate of continuous momentum loss for all loss processes
photons : dict
a dict containing information on photon emission, with the
following entries:
"en" : ndarray or astropy.Quantity, size(j)
energies at which photon emission is calculated
"freq" : ndarray or astropy.Quantity, size(j)
frequencies at which photon emission is calculated;
equal to en / h
"dLdE" : ndarray or astropy.Quantity, size(n,m,j)
array giving the specific luminosity dL/dE (units of
erg s^-1 GeV^-1 or J s^-1 GeV^-1 if criptic was run in
MKS mode) of each of the n packets produced via each of
the m loss mechanisms at each of the j energies
"dLdE_sum" : ndarray or astropy.Quantity, size(u,v,m,j)
array giving the specific luminosity dL/dE (units of
erg s^-1 GeV^-1 or J s^-1 GeV^-1 if criptic was run in
MKS mode) produced by different mechanisms and
particle populations; the index u runs from 0 to number
of distinct particle types, and gives emission produced
by that type; the index v is 0 (for primary particles)
or 1 (for secondary particles); the index m runs over
all the emission mechanisms (in the same orders as
lossmech); finally, the index j runs over the j
energies in en
ionization : dict
a dict containing information on ionization rates, with
the following entries:
"targets" : list of string
names of the ionization targets included in the file
"rate" : ndarray or astropy.Quantity, size(n,m)
array of size number of packets x number of ionization
targets, giving the ionization rate (units of 1/time)
for each packet for each target species
dpdt_transport : dict
a dict containing information on rates of momentum change
by transport (as opposed to loss) processes, with the
following entries:
"Fermi2" : ndarray or astropy.Quantity, size(n)
rate of momentum change due to 2nd order Fermi acceleration
"adiabatic" : ndarray or astropy.Quantity, size(n)
rate of momentum change due to adiabatic gain or loss
"streaming" : ndarray or astropy.Quantity, size(n)
rate of momentum change due to streaming gain or loss
"""
# If units have been requested, make sure we have astropy; issue
# warning if not
if units and u is None:
warnings.warn("units requested but astropy not found")
units = False
# If we have been requested to get kinetic energy, make sure we
# have astropy
if ke and not units:
warnings.warn("cannot compute kinetic energies without units")
ke = False
# Open file
f = h5py.File(chkname, mode='r')
# Extract global run metadata
t = f['CRPacket'].attrs['t'][0]
dt = f['CRPacket'].attrs['dt'][0]
step = f['CRPacket'].attrs['step'][0]
cgs = f['CRPacket'].attrs['units'][0] == 0
uid = f['CRPacket'].attrs['uniqueID']
part_names = [ s.decode('ascii') for s in
f['CRPacket'].attrs['ParticleNames'] ]
# Read number of packets, sources, and deleted packets
npacket = f['CRPacket']['pos'].size
if 'CRSource' in f:
nsource = f['CRSource']['pos'].size
else:
nsource = 0
if 'CRPacketDel' in f and not sources_only:
ndel = f['CRPacketDel']['pos'].size
else:
ndel = 0
# Read list of field quantities
if 'FieldQty' in f and not sources_only:
fieldqty = [ s.decode('ascii') for s in
f['FieldQty'].attrs['names']]
nfq = len(fieldqty)
else:
fieldqty = []
nfq = 0
# Pack metadata into output dict
out = { 't' : t,
'dt' : dt,
'step' : step,
'uid' : uid,
'npacket' : npacket,
'ndelpacket' : ndel,
'nsource' : nsource,
'part_names' : part_names
}
# Read list of loss mechanisms
if 'LossMechanisms' in f['CRPacket'].attrs.keys():
out['lossmech'] = [ s.decode('ascii') for s in
f['CRPacket'].attrs['LossMechanisms']]
# Read photon energies at which luminosity is output
if 'PhotonEmission' in f:
out['photons'] = { "en" : f['PhotonEmission'].attrs['photEn'] }
# Read ionization targets
if 'Ionization' in f:
out['ionization'] \
= { "targets" : [ s.decode('ascii') for s in
f['Ionization'].attrs['TargetNames'] ] }
# If requested to just read metadata, apply units and return
if meta:
f.close()
if units:
out['t'] = out['t'] * u.s
out['dt'] = out['dt'] * u.s
out['photons']['en'] = out['photons']['en'] * u.GeV
return out
# Extract packet data
if not sources_only:
x = f['CRPacket']['pos'][:]
ptype = f['CRPacket']['data'][:]['type']
source = f['CRPacket']['data'][:]['src']
uid = f['CRPacket']['data'][:]['uniqueID']
tinj = f['CRPacket']['data'][:]['tInj']
winj = f['CRPacket']['data'][:]['wInj']
p = f['CRPacket']['data'][:]['p']
w = f['CRPacket']['data'][:]['w']
gr = f['CRPacket']['data'][:]['gr']
if 'mu' in f['CRPacket']['data'].dtype.names:
mu = f['CRPacket']['data'][:]['mu']
else:
mu = None
# Extract source data
if nsource > 0:
xsrc = f['CRSource']['pos'][:]
ptypesrc = f['CRSource']['data'][:]['type']
uidsrc = f['CRSource']['data'][:]['uniqueID']
p0src = f['CRSource']['data'][:]['p0']
p1src = f['CRSource']['data'][:]['p1']
qsrc = f['CRSource']['data'][:]['q']
if 'mu0' in f['CRSource']['data'].dtype.names:
mu0src = f['CRSource']['data'][:]['mu0']
mu1src = f['CRSource']['data'][:]['mu1']
else:
mu0src = None
mu1src = None
coefsrc = f['CRSource']['data'][:]['k']
vsrc = f['CRSource']['data'][:]['v']
asrc = f['CRSource']['data'][:]['a']
# Extract deleted packet data
if ndel > 0 and not sources_only:
xdel = f['CRPacketDel']['pos'][:]
ptypedel = f['CRPacketDel']['data'][:]['type']
sourcedel = f['CRPacketDel']['data'][:]['src']
uiddel = f['CRPacketDel']['data'][:]['uniqueID']
tinjdel = f['CRPacketDel']['data'][:]['tInj']
winjdel = f['CRPacketDel']['data'][:]['wInj']
pdel = f['CRPacketDel']['data'][:]['p']
wdel = f['CRPacketDel']['data'][:]['w']
grdel = f['CRPacketDel']['data'][:]['gr']
if 'mu' in f['CRPacketDel']['data'].dtype.names:
mudel = f['CRPacketDel']['data'][:]['mu']
else:
mudel = None
tdel = f['CRPacketDel']['tdel'][:]
# Extract field quantities
if 'FieldQty' in f and not sources_only:
fld = f['FieldQty']['qty'][:]
if 'grad' in f['FieldQty']:
fldgrad = f['FieldQty']['grad'][:]
else:
fldgrad = None
# Extract loss rates
if 'Losses' in f and not sources_only:
dpdt = f['Losses']['dpdt'][:]
lossrate = f['Losses']['lossRate'][:]
else:
dpdt = None
lossrate = None
# Extract photon emission
if 'PhotonEmission' in f and not sources_only:
if 'dLdE' in f['PhotonEmission']:
phot_dLdE = f['PhotonEmission']['dLdE'][:]
else:
phot_dLdE = None
if 'dLdEsummary' in f['PhotonEmission']:
phot_dLdE_sum = f['PhotonEmission']['dLdEsummary'][:]
else:
phot_dLdE_sum = None
# Extract ionization
if 'Ionization' in f and not sources_only:
ionrate = f['Ionization']['ionRate'][:]
else:
ionrate = None
# Extract dp/dt transport terms
if 'PdotTransport' in f and not sources_only:
dpdt_transport = { }
dpdt_transport['Fermi2'] = f['PdotTransport']['dpdt'][:,0]
dpdt_transport['adiabatic'] = f['PdotTransport']['dpdt'][:,1]
dpdt_transport['streaming'] = f['PdotTransport']['dpdt'][:,2]
else:
dpdt_transport = None
# Close file
f.close()
# Decode secondary flag; secondaries are marked by having the
# most significant bit set of the source field set to 1; sec_mask
# is a bit mask that is 1 for this bit and 0 for all other bits
if not sources_only:
sec_mask = np.invert(np.array(np.iinfo(source.dtype).max//2,
dtype=source.dtype))
sec = (source & sec_mask) != 0
source = source & np.invert(sec_mask)
if ndel > 0:
secdel = (sourcedel & sec_mask) != 0
sourcedel = sourcedel & np.invert(sec_mask)
# Sort if requested
if sort:
# Packets
if not sources_only:
idx = np.argsort(uid)
uid = uid[idx]
x = x[idx]
ptype = ptype[idx]
source = source[idx]
tinj = tinj[idx]
winj = winj[idx]
p = p[idx]
w = w[idx]
gr = gr[idx]
if mu is not None:
mu = mu[idx]
if nfq > 0:
fld = fld[idx]
if fldgrad is not None:
fldgrad = fldgrad[idx]
if dpdt is not None:
dpdt = dpdt[idx]
lossrate = lossrate[idx]
if phot_dLdE is not None:
phot_dLdE = phot_dLdE[idx]
if ionrate is not None:
ionrate = ionrate[idx]
if dpdt_transport is not None:
for k in dpdt_transport.keys():
dpdt_transport[k] = dpdt_transport[k][idx]
# Deleted packets
if ndel > 0:
idx = np.argsort(uiddel)
uiddel = uiddel[idx]
xdel = xdel[idx]
ptypedel = ptypedel[idx]
sourcedel = sourcedel[idx]
tinjdel = tinjdel[idx]
winjdel = winjdel[idx]
pdel = pdel[idx]
wdel = wdel[idx]
grdel = grdel[idx]
if mudel is not None:
mudel = mudel[idx]
tdel = tdel[idx]
# Sources
if nsource > 0:
idx = np.argsort(uidsrc)
uidsrc = uidsrc[idx]
xsrc = xsrc[idx]
ptypesrc = ptypesrc[idx]
p0src = p0src[idx]
p1src = p1src[idx]
qsrc = qsrc[idx]
if mu0src is not None:
mu0src = mu0src[idx]
mu1src = mu1src[idx]
coefsrc = coefsrc[idx]
vsrc = vsrc[idx]
asrc = asrc[idx]
# Split up field quantities
pres = None
presGrad = None
eden = None
edenGrad = None
nden = None
ndenGrad = None
for i, nm in enumerate(fieldqty):
if nm == "pres":
pres = fld[:,i]
if fldgrad is not None: presGrad = fldgrad[:,i]
elif nm == "eden":
eden = fld[:,i]
if fldgrad is not None: edenGrad = fldgrad[:,i]
elif nm == "nden":
nden = fld[:,i]
if fldgrad is not None: ndenGrad = fldgrad[:,i]
# Compute source luminosities if requested
if nsource > 0 and ndot:
ndotsrc = np.zeros(coefsrc.shape)
idx = p1src <= p0src
ndotsrc[idx] = coefsrc[idx]
idx = np.logical_and(qsrc == -1.0, p1src > p0src)
ndotsrc[idx] = coefsrc[idx] * np.log(p1src[idx] / p0src[idx])
idx = np.logical_and(qsrc != -1.0, p1src > p0src)
q_ = qsrc[idx]
ndotsrc[idx] = coefsrc[idx] * (p1src[idx]**(q_+1) -
p0src[idx]**(q_+1)) / \
(q_ + 1)
# Apply units
if units:
# Astrophysical quantities are in cgs or SI units
out['t'] = out['t'] * u.s
out['dt'] = out['dt'] * u.s
if not sources_only:
tinj = tinj * u.s
if ndel > 0:
tdel = tdel * u.s
tinjdel = tinjdel*u.s
if nsource > 0:
coefsrc = coefsrc / u.s
if ndot:
ndotsrc = ndotsrc / u.s
if lossrate is not None: lossrate = lossrate / u.s
if 'photons' in out.keys():
out['photons']['en'] *= u.GeV
out['photons']['freq'] = (out['photons']['en'] / h).to(u.GHz)
if ionrate is not None: ionrate = ionrate / u.s
if cgs:
if not sources_only:
x = x * u.cm
gr = gr * u.g
if ndel > 0:
xdel = xdel * u.cm
grdel = grdel * u.g
if nsource > 0:
xsrc = xsrc * u.cm
vsrc = vsrc * u.cm / u.s
asrc = asrc * u.cm / u.s**2
if nden is not None:
nden = nden / u.cm**3
if ndenGrad is not None:
ndenGrad = ndenGrad / u.cm**4
if 'photons' in out.keys():
if phot_dLdE is not None:
phot_dLdE = phot_dLdE * u.erg / u.s / u.GeV
if phot_dLdE_sum is not None:
phot_dLdE_sum = phot_dLdE_sum * u.erg / u.s / u.GeV
else:
if not sources_only:
x = x * u.m
gr = gr * u.kg
if ndel > 0:
xdel = xdel * u.m
grdel = grdel * u.kg
if nsource > 0:
xscr = xsrc * u.m
vsrc = vsrc * u.m / u.s
asrc = asrc * u.m / u.s**2
if nden is not None:
nden = nden / u.m**3
if ndenGrad is not None:
ndenGrad = ndenGrad / u.m**4
if 'photons' in out.keys():
if phot_dLdE is not None:
phot_dLdE = phot_dLdE * u.J / u.s / u.GeV
if phot_dLdE_sum is not None:
phot_dLdE_sum = phot_dLdE_sum * u.J / u.s / u.GeV
# CR quantities are in criptic internal units, where masses
# are in m_p and velocities are in c
if not sources_only:
p = p * m_p * c
if ndel > 0:
pdel = pdel * m_p * c
if nsource > 0:
p0src = p0src * m_p * c
p1src = p1src * m_p * c
if dpdt is not None:
dpdt = dpdt * m_p * c / u.s
if dpdt_transport is not None:
for k in dpdt_transport.keys():
dpdt_transport[k] = dpdt_transport[k] * m_p * c / u.s
if cgs:
if pres is not None:
pres = (pres * m_p * c**2 / u.cm**3 / const.k_B). \
to(u.K / u.cm**3)
if presGrad is not None:
presGrad = (presGrad * m_p * c**2 / u.cm**4 / const.k_B). \
to(u.K / u.cm**4)
if eden is not None:
eden = (eden * m_p * c**2 / u.cm**3).to(u.eV/u.cm**3)
if edenGrad is not None:
edenGrad = (edenGrad * m_p * c**2 / u.cm**4). \
to(u.eV/u.cm**4)
else:
if pres is not None:
pres = (pres * m_p * c**2 / u.m**3 / const.k_B). \
to(u.K / u.m**3)
if presGrad is not None:
presGrad = (presGrad * m_p * c**2 / u.m**4 / const.k_B). \
to(u.K / u.m**3)
if eden is not None:
eden = (eden * m_p * c**2 / u.m**3).to(u.eV / u.m**3)
if edenGrad is not None:
edenGrad = (edenGrad * m_p * c**2 / u.m**4). \
to(u.eV / u.m**4)
# Compute kinetic energies if requested
if ke:
if not sources_only:
T = (_Tp( np.array((p / (m_p*c)).to('')), ptype) *
m_p*c**2).to('GeV')
if ndel > 0:
Tdel = (_Tp( np.array((pdel / (m_p*c)).to('')), ptypedel) *
m_p*c**2).to('GeV')
if nsource > 0:
T0src = (_Tp( np.array((p0src / (m_p*c)).to('')), ptypesrc) *
m_p*c**2).to('GeV')
T1src = (_Tp( np.array((p1src / (m_p*c)).to('')), ptypesrc) *
m_p*c**2).to('GeV')
# Pack data into dicts and return
if sources_only:
packets = { }
else:
packets = { "x" : x,
"ptype" : ptype,
"uid" : uid,
"source" : source,
"p" : p,
"w" : w,
"tinj" : tinj,
"winj" : winj,
"gr" : gr,
"ptype" : ptype,
"sec" : sec
}
if mu is not None:
packets["mu"] = mu
if pres is not None:
packets["pres"] = pres
if presGrad is not None:
packets["presGrad"] = presGrad
if eden is not None:
packets["eden"] = eden
if edenGrad is not None:
packets["edenGrad"] = edenGrad
if nden is not None:
packets["nden"] = nden
if ndenGrad is not None:
packets["ndenGrad"] = ndenGrad
if ke and not sources_only:
packets["T"] = T
if not sources_only:
out['packets'] = packets
if ndel > 0:
delpackets = { "x" : xdel,
"ptype" : ptypedel,
"uid" : uiddel,
"source" : sourcedel,
"p" : pdel,
"w" : wdel,
"tinj" : tinjdel,
"winj" : winjdel,
"gr" : grdel,
"ptype" : ptypedel,
"sec" : secdel
}
if mudel is not None:
delpackets["mu"] = mudel
if ke:
delpackets["T"] = Tdel
out['delpackets'] = delpackets
if nsource > 0:
sources = { "x" : xsrc,
"ptype" : ptypesrc,
"uid" : uidsrc,
"p0" : p0src,
"p1" : p1src,
"q" : qsrc,
"coef" : coefsrc,
"v" : vsrc,
"a" : asrc
}
if mu0src is not None:
sources["mu0"] = mu0src
sources["mu1"] = mu1src
if ke:
sources["T0"] = T0src
sources["T1"] = T1src
if ndot:
sources["ndot"] = ndotsrc
out['sources'] = sources
if lossrate is not None:
out['losses'] = { 'lossrate' : lossrate,
'dpdt' : dpdt }
if 'photons' in out:
if phot_dLdE is not None:
out['photons']['dLdE'] = phot_dLdE
if phot_dLdE_sum is not None:
out['photons']['dLdE_sum'] = phot_dLdE_sum
if ionrate is not None:
out['ionization']['rate'] = ionrate
if dpdt_transport is not None:
out['dpdt_transport'] = dpdt_transport
return out