""" the following is the final plan; this is not implemented as of 20100425. ----------------------------------------------------------------------- Treat as Keep in Add to compatible stable output EL/A sum with solabu (no decays) ionlist DEPENDS YES DEPENDS CHECK addions NO YES NO NO keepions YES YES NO NO stabions YES YES YES IONS: NO; EL/A:CHECK orgions NO YES NO CHECK ----------------------------------------------------------------------- """ from __future__ import division from __future__ import print_function from __future__ import with_statement import os import os.path import datetime from collections import defaultdict from copy import deepcopy import numpy as np from byte2human import byte2human from time2human import time2human from version2human import version2human from isotope import Ion, SolAbu, DecIon, VOID, KepIon, PROTON, NEUTRON, ALPHA from utils import contract, project # from physconst import * class Decay(object): """ Object to compute decays for a given set of ions. Initialize the mapping with input ion list and desired output format. Then use the *map* function to do the mapping. *decions* contains the list of output ions. Currently np.array of isotope.Ions. What is not implememented at this time are to pass isobar, element, or isotones as output in ionlist etc. """ def __init__(self, ions = None, mass_in = True, mass_out = True, isobars = False, elements = False, isotones = False, solprod = False, stable = False, solions = False, sort = True, ionlist = None, addions = None, keepions = None, stabions = None, orgions = None, solabu = None, keepall = True, decayfile = '/u/alex/kepler/local_data/decay.dat' ): """ Initialize decay map. ions - list of input ions. Currently np.array of isotope.Ion. Probably should allow (or requre?) composition instead? Allow any string, array of strings, ... Same for all ion lists. mass_in - whether input is mass or abundance. mass_out - whether output is mass or abundance. isobars - return result mapped to isobars [only stable isotopes] elements - return result mapped to elements [only stable isotopes] isotones - return result mapped to isotones [only stable isotopes] solabu - solar abundance set to use: filename, token, or isotope.SolAbu solprod - return result in terms of solar production factor [only stable isotopes] solions - return exactly ions in solar abundace set [plus addions, stabions, keepions] ionlist - output *exactly* these isotops stable - output only stable isotopes, discard chains addions - output these extra ions, but not add to EL/A sums keepions - ions to keep as stable but not add ion EL/A sums stabions - ions to fully treat as stable - may conflict with solprod orgions - ions for which to return the initial value w/o decays sort - sort output ions (Z>A>E), mostly relevant when customized ions are provided. keepall - if set to false, do no keep all ions even if stabke set is used. decayfile - decay file to use. [should add allowing to pass object] """ start_time = datetime.datetime.now() self.ftag = ' ['+self.__class__.__name__+'] ' # ions = np.array([Ion('pb220')]) # ions = np.array([Ion('n13'),Ion('o14'), Ion('pb209')]) # ions = np.array([Ion('n13'),Ion('o14')]) # We probbaly need to add some more convenince, allowing # string arrays as well. and not just nd.arrays # And add abundance sets, etc. assert type(ions) == np.ndarray, self.ftag + "ions need to be np.array type" assert ions.ndim == 1, self.ftag + "ions need to be 1D array" assert issubclass(type(ions[0]),Ion), self.ftag + "ions must be of type Ion" # self.sortmap = ions.argsort() self.ions = ions self.amax = max((i.A for i in ions)) self.zmax = max((i.Z for i in ions)) # TODO: add "keepions" to maxima? self.decdata = DecayData(decayfile, self.amax, self.zmax) # self.decdata = DecayData(decayfile) # now let's add keepions # this modfies decdata. # Should we make deep copy first? if keepions != None: assert type(keepions) == np.ndarray, "keepions need to be np.array type" assert keepions.ndim == 1, "keepions need to be 1D array" assert issubclass(type(keepions[0]),Ion), self.ftag + "keepions must be of type Ion" # the strategy is modify the decdata for ion in keepions: self.decdata.add_stable(ion) # ...and stabions if stabions != None: assert type(keepions) == np.ndarray, "stabions need to be np.array type" assert keepions.ndim == 1, "stabions need to be 1D array" assert issubclass(type(stabions[0]),Ion), self.ftag + "stabions must be of type Ion" # the strategy is modify the decdata for ion in keepions: self.decdata.add_stable(ion) self.dectable = {} # construct table from ions d0 = [self.iter_add_dectable(ion) for ion in self.ions] # let us just use indices into array to speed things up self.d = np.ndarray(len(self.dectable), dtype = object) for decay in self.dectable.values(): self.d[decay[0][1]] = decay # construct decay matrix nions = len(ions) ndec = len(self.dectable) self.decmatrix = np.zeros([nions,ndec], dtype=np.float64) for i,ion in enumerate(self.ions): self.iter_add_decmatrix(i, np.float64(1), d0[i]) ionlist = np.array([ion for ion in self.dectable.keys()], dtype = object) indices = np.array([decay[0][1] for decay in self.dectable.values()], dtype = np.int64) m = ionlist.argsort() self.decions = ionlist[m] self.decmatrix = self.decmatrix[:,indices[m]] if mass_out: f_out = np.array([i.A for i in self.decions]) self.decmatrix *= f_out[np.newaxis,:] if mass_in: f_in = 1/np.array([i.A for i in self.ions]) self.decmatrix *= f_in[:,np.newaxis] # we need to check whether we need solar abundances. need_solabu = solions or solprod # since we can get the list of stable ions from the decay table, # we do not really need this for definition of stable isotopes. if need_solabu: if issubclass(type(solabu),str): solabu = SolAbu(solabu) elif solabu == None: solabu = SolAbu() assert issubclass(type(solabu),SolAbu), self.ftag + "Need solar abundace data." # let us assure solar abundance pattern is sorted assert not np.any(solabu.iso.argsort() - np.arange(len(solabu))), self.ftag + "Solar pattern not sorted." # now we construct the map to solar k = 0 smap = [] rmap = [] for i,iso in enumerate(solabu.iso): while self.decions[k] < iso: k += 1 if k == len(self.decions): break if k == len(self.decions): break if self.decions[k] == iso: smap.append(k) rmap.append(i) # find out whether we need list of stable ions need_stable = stable need_isobars = isobars need_isotopes = isotopes need_elements = elements # add things from ionlist if ionlist: need_isobar |= np.any(Ion.ufunc_is_isobar(ionlist)) need_isotone |= np.any(Ion.ufunc_is_isotone(ionlist)) need_element |= np.any(Ion.ufunc_is_elemnet(ionlist)) if addions: need_isobar |= np.any(Ion.ufunc_is_isobar(addions)) need_isotone |= np.any(Ion.ufunc_is_isotone(addions)) need_element |= np.any(Ion.ufunc_is_elemnet(addions)) if need_elements or need_isobars or need_isotones: need_stable = True if need_stable and not need_solabu: smap = [i for i,ion in enumerate(self.decions) if len(self.dectable[ion]) == 1] ### [WORK HERE] # before we do any of the following, we still need to make # sure to only include stable isotopes. # proably best to keep old array and construct new one piece # by piece. if need_stable or need_solabu: self.decions = self.decions[smap] self.decmatrix = self.decmatrix[:,smap] # add missing solar ions if solions and len(solabu) > len(self.decions): ndec = len(solabu) d = np.zeros([nions,ndec], dtype=np.float64) d[:,rmap] = self.decmatrix[:,:] self.decions = solabu.iso self.decmatrix = d ### add missing ions from ion list # .... # can ion list be elements, isobars, isotones? # yes! if elements: decionz = [i.Z for i in self.decions] z, zinv = np.unique(decionz, return_inverse = True) # m = np.zeros([len(self.ions),len(z)]) # for i,j in enumerate(zinv): # m[:,j] += self.decmatrix[:,i] # self.decmatrix = m self.decmatrix = contract(self.decmatrix, zinv, axis = -1) self.decions = np.array([Ion(Z = Z) for Z in z]) if isobars: deciona = Ion.ufunc_A(self.decions) # a, ainv = np.unique(deciona, return_inverse = True) # m = np.zeros([len(self.ions),len(a)]) # for i,j in enumerate(ainv): # m[:,j] += self.decmatrix[:,i] # self.decmatrix = m self.decmatrix, a = project(self.decmatrix, deciona, return_values = True, axis = -1) self.decions = np.array([Ion(A = A) for A in a]) if isotones: decionn = [i.N for i in self.decions] n, ninv = np.unique(decionn, return_inverse = True) m = np.zeros([len(self.ions),len(n)]) for i,j in enumerate(ninv): m[:,j] += self.decmatrix[:,i] self.decmatrix = m self.decions = np.array([Ion(A = N, N = N) for N in n]) if solprod: pass end_time = datetime.datetime.now() load_time = end_time - start_time print(self.ftag + 'matrix constructed in ' + time2human(load_time.total_seconds())) def iter_add_decmatrix(self, i, weight, j): self.decmatrix[i,j] += weight for b, jj in self.d[j][1:]: w = b * weight self.iter_add_decmatrix(i, w, jj) def iter_add_dectable(self, ion): if not self.dectable.has_key(ion): idx0 = len(self.dectable) declist = [(np.float64(1), idx0)] self.dectable[ion] = declist for b,d in self.decdata(ion): # how to deal with de-excitation? # [this should work] D = ion + d if D != ion: idx = self.iter_add_dectable(D) declist.append((b, idx)) # decay of decay product particles: # this may be particular relevant for sf, but also n for i,n in d.hadrons().items(): idx = self.iter_add_dectable(i) declist.append((b*n, idx)) return idx0 return self.dectable[ion][0][1] def __call__(self, abu): """ Map a data set and return result. Add checks? """ return np.dot(abu,self.decmatrix) # one change we may want to do is to make a copy of the raw data # before we modify the decay with "add_stable" class DecayData(object): """ Class to interface with decay.dat file. """ def __init__(self, filename, amax = 999, zmax = 999, verbose = False): """ Load decay file. TODO: Add isomeres. 'g' and 'm' additions. Ion class needs to recognize these. 'g' decay means E -= 1 """ self.ftag = ' ['+self.__class__.__name__+'] ' self.decayfile = os.path.expandvars(os.path.expanduser(filename)) start_time = datetime.datetime.now() decdata = {} with open(self.decayfile,'rn') as f: print(self.ftag + 'Loading {} ({})' .format(f.name, byte2human(os.fstat(f.fileno()).st_size))) for s in f: if s.startswith(';'): continue reac = s.split() # we shall accept "empty" as stable.... assert 1 <= len(reac) <= 3, ftag + "Decay file entries need to be 2 or 3 units." ion = Ion(reac[0]) assert ion != Ion.VOID, ftag + "Ion name not valid." if len(reac) <= 2: br = np.float64(1.) else: br = np.float64(reac[1]) if len(reac) > 1: decay = reac[-1] else: decay = '' if (ion.A > amax) and (ion.Z > zmax): continue d = DecIon(decay) assert d != VOID if not decdata.has_key(ion): decdata[ion] = [] decdata[ion].append((br,d)) self.decdata = decdata end_time = datetime.datetime.now() load_time = end_time - start_time print(self.ftag + 'data loaded in ' + time2human(load_time.total_seconds())) def get_decay(self, ion): """ Return decay table entry. If the decay is not in stored data, extrapolate from last available decay at same element (Z value). Use gs decays only for extrapolation. """ try: return self.decdata[ion] except: Z = ion.Z a = np.array([i.A for i in self.decdata.iterkeys() if i.Z == Z and i.E == 0]) assert len(a) > 0, self.ftag + "Problem finiding isotopes for decay." A = min(max(a), max(min(a), ion.A)) return self.decdata[Ion(Z = Z, A = A)] def __call__(self, ion): return self.get_decay(ion) def __getitem__(self, ion): return self.get_decay(ion) def add_stable(self, ion): """ Flag an ion as stable, extending decays from edge. This may require updates to handel isomeric states. """ Z = ion.Z A = ion.A a = np.array([i.A for i in self.decdata.iterkeys() if i.Z == Z and i.E == 0]) assert len(a) > 0, self.ftag + "Error: Stable ion outside range." amin = min(a) amax = max(a) for a in xrange(A-1, amin): self.decdata[Ion(Z = Z, A = a)] = deepcopy(self.decdata[Ion(A = amin, Z = Z)]) for a in xrange(amax+1, A+2): self.decdata[Ion(Z = Z, A = a)] = deepcopy(self.decdata[Ion(A = amax, Z = Z)]) # finally overwrite with "stable" self.decdata[ion] = [(np.float64(1),DecIon('s'))] class MapBurn(object): """ Map isotope distribution to APPROX-19 network ideally this should become similar to Decay, maybe inherit from common abstract base class? """ approx19 = ('nt1' , 'h1' , 'pn1' , 'he3' , 'he4' , 'c12' , 'n14' , 'o16' , 'ne20', 'mg24', 'si28', 's32' , 'ar36', 'ca40', 'ti44', 'cr48', 'fe52', 'fe54', 'ni56') approx19_ions = [Ion(ion) for ion in approx19] def __init__(self, ions): """ Create mapping matrix. MAY NEED TO ADD mass/mol mapping. Currently mass only. """ self.decmatrix = np.zeros((len(ions),19), dtype=np.float64) self.decions = np.array([KepIon(i) for i in self.approx19]) self.ions = ions for i,ion in enumerate(ions): if ion == NEUTRON: self.decmatrx[i,0] = 1. elif ion == PROTON: self.decmatrix[i,1] = 1. elif ion == ALPHA: self.decmatrix[i,4] = 1. elif ion == self.approx19_ions[18]: self.decmatrix[i,18] = 1. elif ion.Z < 6: self.decmatrix[i,3] = 1. elif ion.Z < 7: self.decmatrix[i,5] = 1. elif ion.Z < 8: self.decmatrix[i,6] = 1. elif ion.Z < 10: self.decmatrix[i,7] = 1. elif ion.Z < 12: self.decmatrix[i,8] = 1. elif ion.Z < 14: self.decmatrix[i,9] = 1. elif ion.Z < 16: self.decmatrix[i,10] = 1. elif ion.Z < 18: self.decmatrix[i,11] = 1. elif ion.Z < 20: self.decmatrix[i,12] = 1. elif ion.Z < 22: self.decmatrix[i,13] = 1. elif ion.Z < 24: self.decmatrix[i,14] = 1. elif ion.Z < 26: self.decmatrix[i,15] = 1. elif (ion.Z < 26) and (ion.A < 54): self.decmatrix[i,16] = 1. else: self.decmatrix[i,17] = 1. def __call__(self, abu): """ Map a data set and return result. Add checks? """ return np.dot(abu,self.decmatrix)