diff --git a/zdm/grid.py b/zdm/grid.py index 5f87ced..fc0432f 100644 --- a/zdm/grid.py +++ b/zdm/grid.py @@ -28,6 +28,14 @@ def __init__(self, survey, state, zDMgrid, zvals, dmvals, smear_mask, wdist, pre Defines the parameters of the analysis Note, each grid holds the *same* copy so modifying it in one place affects them all. + zvals (np.1darray, float): + redshift values of the grid. These are "bin centres", + representing ranges from +- dz/2. + dmvals (np.1darray, float: + DM values of the grid. These are These are "bin centres", + representing ranges from +- dDM/2. + smear_mask (np.1darray, float): + 1D array wdist (bool): If True, allow for a distribution of widths prev_grid (grid.Grid): @@ -43,7 +51,7 @@ def __init__(self, survey, state, zDMgrid, zvals, dmvals, smear_mask, wdist, pre self.b_fractions = None # State self.state = state - + self.MCinit = False self.source_function = cos.choose_source_evolution_function( state.FRBdemo.source_evolution ) @@ -137,12 +145,64 @@ def load_grid(self, gridfile, zfile, dmfile): self.dmvals = io.load_data(dmfile) self.check_grid() self.volume_grid() - - def check_grid(self): - + + def get_dm_coeffs(self, DMlist): + """ + Returns indices and coefficients for interpolating between DM values + + dmlist: np.ndarray of dispersion measures (extragalactic!) + """ + # get indices in dm space + kdms=DMlist/self.ddm - 0.5 # idea: if DMobs = ddm, we are half-way between bin 0 and bin 1 + Bin0 = np.where(kdms < 0.)[0] # if DMs are in the lower half of the lowest bin, use lowest bin only + kdms[Bin0] = 0. + idms1=kdms.astype('int') # rounds down + idms2=idms1+1 + dkdms2=kdms-idms1 # applies to idms2, i.e. the upper bin. If DM = ddm, then this should be 0.5 + dkdms1 = 1.-dkdms2 # applies to idms1 + return idms1,idms2,dkdms1,dkdms2 + + def get_z_coeffs(self,zlist): + """ + Returns indices and coefficients for interpolating between z values + + zlist: np.ndarray of dispersion measures (extragalactic!) + """ + + kzs=zlist/self.dz - 0.5 + Bin0 = np.where(kzs < 0.)[0] + kzs[Bin0] = 0. + izs1=kzs.astype('int') + izs2=izs1+1 + dkzs2=kzs-izs1 # applies to izs2 + dkzs1 = 1. - dkzs2 + + # checks for values which are too large + toobigz = np.where(zlist > self.zvals[-1] + self.dz/2.)[0] + if len(toobigz) > 0: + raise ValueError("Redshift values ",zlist[toobigz], + " too large for grid max of ",self.zvals[-1] + self.dz/2.) + + # checks for zs in top half of top bin - only works because of above bin + topbin = np.where(zlist > self.zvals[-1])[0] + if len(topbin) > 0: + izs2[topbin] = self.zvals.size-1 + izs1[topbin] = self.zvals.size-2 + dkzs2[topbin] = 1. + dkzs1[topbin] = 0. + + return izs1, izs2, dkzs1, dkzs2 + + def check_grid(self,TOLERANCE = 1e-6): + """ + Check that the grid values are behaving as expected + + TOLERANCE: defines the max relative difference in expected + and found values that will be tolerated + """ self.nz = self.zvals.size self.ndm = self.dmvals.size - + # check to see if these are log-spaced if (self.zvals[-1] - self.zvals[-2]) / (self.zvals[1] - self.zvals[0]) > 1.01: if ( @@ -155,7 +215,7 @@ def check_grid(self): else: self.zlog = False self.dz = self.zvals[1] - self.zvals[0] - + self.ddm = self.dmvals[1] - self.dmvals[0] shape = self.grid.shape if shape[0] != self.nz: @@ -178,17 +238,26 @@ def check_grid(self): expectation = self.dz * np.arange(0, self.nz) + self.zvals[0] diff = self.zvals - expectation maxoff = np.max(diff ** 2) - if maxoff > 1e-6 * self.dz: + if maxoff > TOLERANCE * self.dz: raise ValueError( "Maximum non-linearity in z-grid of ", maxoff ** 0.5, "detected, aborting", ) - + + # Ensures that log-spaced bins are truly bin centres + if not self.zlog and np.abs(self.zvals[0] - self.dz/2.) > TOLERANCE*self.dz: + raise ValueError( + "Linear z-grids *must* begin at dz/2. e.g. 0.05,0.15,0.25 etc, ", + " first value ",self.zvals[0]," expected to be half of spacing ", + self.dz,", aborting..." + ) + + expectation = self.ddm * np.arange(0, self.ndm) + self.dmvals[0] diff = self.dmvals - expectation maxoff = np.max(diff ** 2) - if maxoff > 1e-6 * self.ddm: + if maxoff > TOLERANCE * self.ddm: raise ValueError( "Maximum non-linearity in dm-grid of ", maxoff ** 0.5, @@ -460,14 +529,14 @@ def GenMCSample(self, N, Poisson=False): If Poisson=True, then interpret N as a Poisson expectation value Otherwise, generate precisely N FRBs - Generated values are DM, z, B, w, and SNR + Generated values are [MCz, MCDM, MCb, MCs, MCw] NOTE: the routine GenMCFRB does not know 'w', merely which w bin it generates. """ # Boost? if self.state.energy.luminosity_function in [1, 2]: - Emax_boost = 2.0 + Emax_boost = 3.0 else: Emax_boost = 0.0 @@ -477,22 +546,92 @@ def GenMCSample(self, N, Poisson=False): else: NFRB = int(N) # just to be sure... sample = [] - pwb = None # feeds this back to save time. Lots of time. + for i in np.arange(NFRB): if (i % 100) == 0: print(i) # Regen if the survey would not find this FRB - frb, pwb = self.GenMCFRB(pwb, Emax_boost=Emax_boost) + frb = self.GenMCFRB(Emax_boost) + # This is a pretty naive method of generation. while frb[1] > self.survey.max_dm: - frb, pwb = self.GenMCFRB(pwb, Emax_boost=Emax_boost) + print("Regenerating MC FRB with too high DM ",frb[1],self.survey.max_dm) + frb = self.GenMCFRB(Emax_boost) sample.append(frb) sample = np.array(sample) return sample - - def GenMCFRB(self, pwb=None, Emax_boost=0.0): + + + def initMC(self): + """ + Initialises the MC sample, if it has not been doen already + This uses a great deal of RAM - hence, do not do this lightly! + """ + + # shorthand + lEmin = self.state.energy.lEmin + lEmax = self.state.energy.lEmax + gamma = self.state.energy.gamma + Emin = 10 ** lEmin + Emax = 10 ** lEmax + + # grid of beam values, weights + nw = self.eff_weights.size + nb = self.beam_b.size + + # holds array of probabilities in w,b space + pwb = np.zeros([nw * nb]) + rates = [] + pzcs = [] + + # gets list of DM probabilities to set to zero due to + # the survey missing these FRBs + if self.survey.max_dm is not None: + setDMzero = np.where(self.dmvals +self.ddm/2. > self.survey.max_dm)[0] + + # Generates a joint distribution in B,w + for i, b in enumerate(self.beam_b): + for j, w in enumerate(self.eff_weights): + # each of the following is a 2D array over DM, z which we sum to generate B,w values + pzDM = self.array_cum_lf( + self.thresholds[j, :, :] / b, + Emin, Emax, gamma) + + # sets to zero if we have a max survey DM + if self.survey.max_dm is not None: + pzDM [:,setDMzero] = 0. + + # weighted pzDM + wb_fraction = (self.beam_o[i] * w * pzDM) + pdv = np.multiply(wb_fraction.T, self.dV).T + rate = pdv * self.sfr_smear + rates.append(rate) + pwb[i * nw + j] = np.sum(rate) + + pz = np.sum(rate, axis=1) + pzc = np.cumsum(pz) + pzc /= pzc[-1] + + pzcs.append(pzc) + + # generates cumulative distribution for sampling w,b + pwbc = np.cumsum(pwb) + pwbc /= pwbc[-1] + + # saves cumulative distributions for sampling + self.MCpwbc = pwbc + + # saves individal wxb zDM rates for sampling these distributions + self.MCrates = rates + + # saves projections onto z-axis + self.MCpzcs = pzcs + + self.MCinit = True + + def GenMCFRB(self, Emax_boost): """ Generates a single FRB according to the grid distributions @@ -514,125 +653,192 @@ def GenMCFRB(self, pwb=None, Emax_boost=0.0): Returns: tuple: FRBparams=[MCz,MCDM,MCb,j,MCs], pwb values + These are: + MCz: redshift + MCDM: dispersion measure (extragalactic) + MCb: beam value + j: + MCs: SNR/SNRth value of FRB + MCw: width value of FRB + [MCz, MCDM, MCb, j, MCs, MCw] """ - + # shorthand lEmin = self.state.energy.lEmin lEmax = self.state.energy.lEmax gamma = self.state.energy.gamma Emin = 10 ** lEmin Emax = 10 ** lEmax - + # grid of beam values, weights nw = self.eff_weights.size nb = self.beam_b.size - - # we do this to allow efficient recalculation of this when generating many FRBs - if pwb is not None: - pwbc = np.cumsum(pwb) - pwbc /= pwbc[-1] - else: - pwb = np.zeros([nw * nb]) - - # Generates a joint distribution in B,w - for i, b in enumerate(self.beam_b): - for j, w in enumerate(self.eff_weights): - # each of the following is a 2D array over DM, z which we sum to generate B,w values - wb_fraction = ( - self.beam_o[i] - * w - * self.array_cum_lf( - self.thresholds[j, :, :] / b, Emin, Emax, gamma - ) - ) - pdv = np.multiply(wb_fraction.T, self.dV).T - rate = pdv * self.sfr_smear - pwb[i * nw + j] = np.sum(rate) - pwbc = np.cumsum(pwb) - pwbc /= pwbc[-1] - + + if not self.MCinit: + self.initMC() + # sample distribution in w,b # we do NOT interpolate here - we treat these as qualitative values # i.e. as if there was an irregular grid of them r = np.random.rand(1)[0] - which = np.where(pwbc > r)[0][0] + which = np.where(self.MCpwbc > r)[0][0] i = int(which / nw) j = which - i * nw MCb = self.beam_b[i] MCw = self.eff_weights[j] - - # calculate zdm distribution for sampled w,b only - pzDM = self.array_cum_lf(self.thresholds[j, :, :] / MCb, Emin, Emax, gamma) - wb_fraction = self.array_cum_lf( - self.thresholds[j, :, :] / MCb, Emin, Emax, gamma - ) - pdv = np.multiply(wb_fraction.T, self.dV).T - pzDM = pdv * self.sfr_smear - - # sample distribution in z,DM - pz = np.sum(pzDM, axis=1) - pzc = np.cumsum(pz) - pzc /= pzc[-1] + + # get p(z,DM) distribution for this b,w + pzDM = self.MCrates[which] + + pzc = self.MCpzcs[which] + r = np.random.rand(1)[0] + + # sampling in DM and z space + # First choose z: pzc is the cumulative distribution in z + # for all dm + # each probability represents the p(bin), i.e. z-dz/2. to z+dz/2 + # first, find the bin where the cumulative probability is higher + # than the sampled amount. + # Alternative method: just use the distribution from the bin, + # then multiply the resulting DM linearly with deltaz/z. + # Would be better at low z, worse at high z iz2 = np.where(pzc > r)[0][0] if iz2 > 0: iz1 = iz2 - 1 + iz3 = iz2 + 1 dr = r - pzc[iz1] - kz2 = dr / (pzc[iz2] - pzc[iz1]) # fraction of way to second value - kz1 = 1.0 - kz2 - MCz = self.zvals[iz1] * kz1 + self.zvals[iz2] * kz2 - pDM = pzDM[iz1, :] * kz1 + pzDM[iz2, :] * kz2 + fz = dr / (pzc[iz2] - pzc[iz1]) # fraction of way to upper z value + + # move a fraction of kz2 between z-dz/0.5 and z + dz/0.5 + #MCz = self.zvals[iz2] + (kz2-0.5)*dz + + # weigts between iz1 and iz2 + if fz < 0.5: + kz1 = 0.5 - fz + kz2 = 0.5 + fz + kz3 = 0. + iz3 = 0 # dummy + elif iz2 == self.zvals.size-1: + # we are in the last bin - don't extrapolate, just use it + kz1 = 0. + kz2 = 1. + kz3 = 0. + iz1 = 0 # dummy + iz3 = 0 # dummy + else: + kz1 = 0. + kz2 = (1.5-fz) + kz3 = fz-0.5 + iz1 = 0 #dummy + pDM = pzDM[iz1, :] * kz1 + pzDM[iz2, :] * kz2 + pzDM[iz3, :] * kz3 + MCz = self.zvals[iz1] * kz1 + self.zvals[iz2] * kz2 + self.zvals[iz3]*kz3 else: # we perform a simple linear interpolation in z from 0 to minimum bin - kz2 = r / pzc[iz2] - kz1 = 1.0 - kz2 - MCz = self.zvals[iz2] * kz2 - pDM = pzDM[iz2, :] # just use the value of lowest bin - + fz = r / pzc[iz2] + kz1 = 0. + kz2 = 1. + kz3 = 0. + iz1 = 0 # dummy + iz3 = 0 # dummy + MCz = (self.zvals[iz2] + self.dz/0.5) * fz + # Just use the value of lowest bin. + # This is a gross and repugnant approximation + pDM = pzDM[iz2, :] + # NOW DO dm - # pDM=pzDM[k,:] + # DM represents the distribution for the centre of z-bin pDMc = np.cumsum(pDM) pDMc /= pDMc[-1] r = np.random.rand(1)[0] iDM2 = np.where(pDMc > r)[0][0] if iDM2 > 0: iDM1 = iDM2 - 1 + iDM3 = iDM2 + 1 dDM = r - pDMc[iDM1] - kDM2 = dDM / (pDMc[iDM2] - pDMc[iDM1]) - kDM1 = 1.0 - kDM2 - MCDM = self.dmvals[iDM1] * kDM1 + self.dmvals[iDM2] * kDM2 - if iz2 > 0: - Eth = ( - self.thresholds[j, iz1, iDM1] * kz1 * kDM1 - + self.thresholds[j, iz1, iDM2] * kz1 * kDM2 - + self.thresholds[j, iz2, iDM1] * kz2 * kDM1 - + self.thresholds[j, iz2, iDM2] * kz2 * kDM2 - ) + # fraction of value through DM bin + fDM = dDM / (pDMc[iDM2] - pDMc[iDM1]) + + # get the MC DM through interpolation + if fDM < 0.5: + kDM1 = 0.5 - fDM + kDM2 = 0.5 + fDM + kDM3 = 0. + iDM3 = 0 # dummy + # sets iDM3 to be safe at 0 + MCDM = self.dmvals[iDM1] * kDM1 + self.dmvals[iDM2] * kDM2 + elif iDM2 == self.dmvals.size-1: + kDM1 = 0. + kDM2 = 1. # for future use, not here + kDM3 = 0. + iDM1 = 0 # dummy + iDM3 = 0 # dummy + MCDM = self.dmvals[iDM2] + (fDM - 0.5)*self.dDM # upper DM bins else: - Eth = ( - self.thresholds[j, iz2, iDM1] * kDM1 - + self.thresholds[j, iz2, iDM2] * kDM2 - ) - Eth *= kz2 ** 2 # assume threshold goes as Eth~z^2 in the near Universe + kDM1 = 0. + kDM2 = 1.5-fDM + kDM3 = fDM - 0.5 + iDM1 = 0 # dummy + + MCDM = self.dmvals[iDM3] * kDM3 + self.dmvals[iDM2] * kDM2 + + #MCDM = self.dmvals[iDM1] * kDM1 + self.dmvals[iDM2] * kDM2 + #if iz2 > 0: + # Eth = ( + # self.thresholds[j, iz1, iDM1] * kz1 * kDM1 + # + self.thresholds[j, iz1, iDM2] * kz1 * kDM2 + # + self.thresholds[j, iz1, iDM3] * kz1 * kDM3 + # + self.thresholds[j, iz2, iDM1] * kz2 * kDM1 + # + self.thresholds[j, iz2, iDM2] * kz2 * kDM2 + # + self.thresholds[j, iz2, iDM3] * kz2 * kDM3 + # + self.thresholds[j, iz3, iDM1] * kz3 * kDM1 + # + self.thresholds[j, iz3, iDM2] * kz3 * kDM2 + # + self.thresholds[j, iz3, iDM3] * kz3 * kDM3 + # ) + #else: + # Eth = ( + # self.thresholds[j, iz2, iDM1] * kDM1 + # + self.thresholds[j, iz2, iDM2] * kDM2 + # ) + # Eth *= kz2 ** 2 # assume threshold goes as Eth~z^2 in the near Universe else: # interpolate linearly from 0 to the minimum value - kDM2 = r / pDMc[iDM2] - MCDM = self.dmvals[iDM2] * kDM2 - if iz2 > 0: # ignore effect of lowest DM bin on threshold - Eth = ( - self.thresholds[j, iz1, iDM2] * kz1 - + self.thresholds[j, iz2, iDM2] * kz2 - ) - else: - Eth = self.thresholds[j, iz2, iDM2] * kDM2 - Eth *= kz2 ** 2 # assume threshold goes as Eth~z^2 in the near Universe - + fDM = r / pDMc[iDM2] + MCDM = (self.dmvals[iDM2] + self.ddm/2.) * fDM + kDM1 = 0. + kDM2 = 1. + kDM3 = 0. + iDM1 = 0 #dummy + iDM3 = 0 #dummy + + #if iz2 > 0: # ignore effect of lowest DM bin on threshold + # Eth = ( + # self.thresholds[j, iz1, iDM2] * kz1 + # + self.thresholds[j, iz2, iDM2] * kz2 + # ) + #else: + # Eth = self.thresholds[j, iz2, iDM2] * kDM2 + # Eth *= kz2 ** 2 # assume threshold goes as Eth~z^2 in the near Universe + + # This is constructed such that weights and iz, iDM will work out + # for all cases of the above. Note that only four of these terms at + # most will ever be non-zero. + Eth = self.thresholds[j, iz1, iDM1] * kz1 * kDM1 \ + + self.thresholds[j, iz1, iDM2] * kz1 * kDM2 \ + + self.thresholds[j, iz1, iDM3] * kz1 * kDM3 \ + + self.thresholds[j, iz2, iDM1] * kz2 * kDM1 \ + + self.thresholds[j, iz2, iDM2] * kz2 * kDM2 \ + + self.thresholds[j, iz2, iDM3] * kz2 * kDM3 \ + + self.thresholds[j, iz3, iDM1] * kz3 * kDM1 \ + + self.thresholds[j, iz3, iDM2] * kz3 * kDM2 \ + + self.thresholds[j, iz3, iDM3] * kz3 * kDM3 \ + # now account for beamshape Eth /= MCb # NOW GET snr # Eth=self.thresholds[j,k,l]/MCb - Es = np.logspace(np.log10(Eth), np.log10(Emax) + Emax_boost, 1000) + Es = np.logspace(np.log10(Eth), lEmax + Emax_boost, 1000) PEs = self.vector_cum_lf(Es, Emin, Emax, gamma) PEs /= PEs[0] # normalises: this is now cumulative distribution from 1 to 0 r = np.random.rand(1)[0] @@ -644,8 +850,8 @@ def GenMCFRB(self, pwb=None, Emax_boost=0.0): MCE = 10 ** (np.log10(Es[iE1]) * kE1 + np.log10(Es[iE2]) * kE2) MCs = MCE / Eth - FRBparams = [MCz, MCDM, MCb, j, MCs, MCw] - return FRBparams, pwb + FRBparams = [MCz, MCDM, MCb, MCs, MCw] + return FRBparams def build_sz(self): pass @@ -706,7 +912,11 @@ def update(self, vparams: dict, ALL=False, prev_grid=None): reset_cos, get_zdm, calc_dV = False, False, False smear_mask, smear_dm, calc_pdv, set_evol = False, False, False, False new_sfr_smear, new_pdv_smear, calc_thresh = False, False, False - + + # if we are updating a grid, the MC will in-general need to be + # re-initialised + self.MCinit = False + # Cosmology -- Only H0 so far if self.chk_upd_param("H0", vparams, update=True): reset_cos = True diff --git a/zdm/iteration.py b/zdm/iteration.py index 890deb5..731092c 100644 --- a/zdm/iteration.py +++ b/zdm/iteration.py @@ -295,18 +295,20 @@ def calc_likelihoods_1D(grid,survey,doplot=False,norm=True,psnr=True,Pn=True,dol # TODO: this is slow - should collapse only used columns pdm=np.sum(rates,axis=0) - ddm=dmvals[1]-dmvals[0] - kdms=DMobs/ddm - idms1=kdms.astype('int') - idms2=idms1+1 - dkdms=kdms-idms1 + #ddm=dmvals[1]-dmvals[0] + #kdms=DMobs/ddm + #idms1=kdms.astype('int') + #idms2=idms1+1 + #dkdms=kdms-idms1 + + idms1,idms2,dkdms1,dkdms2 = grid.get_dm_coeffs(DMobs) if grid.state.MW.sigmaDMG == 0.0: if np.any(DMobs < 0): raise ValueError("Negative DMobs with no uncertainty") # Linear interpolation - pvals=pdm[idms1]*(1.-dkdms) + pdm[idms2]*dkdms + pvals=pdm[idms1]*dkdms1 + pdm[idms2]*dkdms2 else: dm_weights, iweights = calc_DMG_weights(DMobs, survey.DMhalos[nozlist], survey.DMGs[nozlist], dmvals, grid.state.MW.sigmaDMG, grid.state.MW.sigmaHalo, grid.state.MW.percentDMG, grid.state.MW.logu) @@ -373,15 +375,17 @@ def calc_likelihoods_1D(grid,survey,doplot=False,norm=True,psnr=True,Pn=True,dol psnr=np.zeros([DMobs.size]) # has already been cut to non-localised number # Evaluate thresholds at the exact DMobs - kdmobs=(survey.DMs - np.median(survey.DMGs + survey.DMhalos))/ddm - kdmobs=kdmobs[nozlist] - kdmobs[kdmobs<0] = 0 - idmobs1=kdmobs.astype('int') - idmobs2=idmobs1+1 - dkdmobs=kdmobs-idmobs1 # applies to idms2 - + #kdmobs=(survey.DMs - np.median(survey.DMGs + survey.DMhalos))/ddm + #kdmobs=kdmobs[nozlist] + #kdmobs[kdmobs<0] = 0 + #idmobs1=kdmobs.astype('int') + #idmobs2=idmobs1+1 + #dkdmobs=kdmobs-idmobs1 # applies to idms2 + DMEGmeans = survey.DMs[nozlist] - np.median(survey.DMGs + survey.DMhalos) + idmobs1,idmobs2,dkdmobs1,dkdmobs2 = grid.get_dm_coeffs(DMEGmeans) + # Linear interpolation - Eths = grid.thresholds[:,:,idmobs1]*(1.-dkdmobs) + grid.thresholds[:,:,idmobs2]*dkdmobs + Eths = grid.thresholds[:,:,idmobs1]*dkdmobs1 + grid.thresholds[:,:,idmobs2]*dkdmobs2 ##### IGNORE THIS, PVALS NOW CONTAINS CORRECT NORMALISATION ###### # we have previously calculated p(DM), normalised by the global sum over all DM (i.e. given 1 FRB detection) @@ -390,7 +394,7 @@ def calc_likelihoods_1D(grid,survey,doplot=False,norm=True,psnr=True,Pn=True,dol # p(snr,DM)/p(DM) * p(DM)/b(burst) # get a vector of rates as a function of z #rs = rates[:,idms1[j]]*(1.-dkdms[j])+ rates[:,idms2[j]]*dkdms[j] - rs = rates[:,idms1]*(1.-dkdms)+ rates[:,idms2]*dkdms + rs = rates[:,idms1]*dkdms1+ rates[:,idms2]*dkdms2 #norms=np.sum(rs,axis=0)/global_norm norms=pvals @@ -414,7 +418,7 @@ def calc_likelihoods_1D(grid,survey,doplot=False,norm=True,psnr=True,Pn=True,dol # this are the grid smearing factors incorporating pcosmic and the host contributions if grid.state.MW.sigmaDMG == 0.0: # Linear interpolation - sg = grid.sfr_smear[:,idms1]*(1.-dkdms)+ grid.sfr_smear[:,idms2]*dkdms + sg = grid.sfr_smear[:,idms1]*dkdms1+ grid.sfr_smear[:,idms2]*dkdms2 else: sg = np.zeros([grid.sfr_smear.shape[0], len(idms1)]) # For each FRB @@ -641,39 +645,50 @@ def calc_likelihoods_2D(grid,survey, else: norm=1. + # in the grid, each z and dm value represents the centre of a bin, with p(z,DM) + # giving the probability of finding the FRB in the range z +- dz/2, dm +- dm/2. + # threshold for when we shift from lower to upper is if z < zcentral, + # weight slowly shifts from lower to upper bin + # get indices in dm space - ddm=dmvals[1]-dmvals[0] - kdms=DMobs/ddm - idms1=kdms.astype('int') - idms2=idms1+1 - dkdms=kdms-idms1 # applies to idms2 - - # get indices in z space - dz=zvals[1]-zvals[0] - kzs=Zobs/dz - izs1=kzs.astype('int') - izs2=izs1+1 - dkzs=kzs-izs1 # applies to izs2 - # print("izs and idms:", izs1, idms1, izs2, idms2, flush=True) - + #ddm=dmvals[1]-dmvals[0] + #kdms=DMobs/ddm - 0.5 # idea: if DMobs = ddm, we are half-way between bin 0 and bin 1 + #Bin0 = np.where(kdms < 0.)[0] # if DMs are in the lower half of the lowest bin, use lowest bin only + #kdms[Bin0] = 0. + #idms1=kdms.astype('int') # rounds down + #idms2=idms1+1 + #dkdms=kdms-idms1 # applies to idms2, i.e. the upper bin. If DM = ddm, then this should be 0.5 + + idms1,idms2,dkdms1,dkdms2 = grid.get_dm_coeffs(DMobs) + + # get indices in z space. See dm comments above. This will not work for log-spaced z + #dz=zvals[1]-zvals[0] + #kzs=Zobs/dz - 0.5 + #Bin0 = np.where(kzs < 0.)[0] + #kzs[Bin0] = 0. + #izs1=kzs.astype('int') + #izs2=izs1+1 + #dkzs=kzs-izs1 # applies to izs2 + + izs1,izs2,dkzs1,dkzs2 = grid.get_z_coeffs(Zobs) + # Calculate probability - if grid.state.MW.sigmaDMG == 0.0: if np.any(DMobs < 0): raise ValueError("Negative DMobs with no uncertainty") # Linear interpolation - pvals = rates[izs1,idms1]*(1.-dkdms)*(1-dkzs) - pvals += rates[izs2,idms1]*(1.-dkdms)*dkzs - pvals += rates[izs1,idms2]*dkdms*(1-dkzs) - pvals += rates[izs2,idms2]*dkdms*dkzs + pvals = rates[izs1,idms1]*dkdms1*dkzs1 + pvals += rates[izs2,idms1]*dkdms1*dkzs2 + pvals += rates[izs1,idms2]*dkdms2*dkzs1 + pvals += rates[izs2,idms2]*dkdms2*dkzs2 else: dm_weights, iweights = calc_DMG_weights(DMobs, survey.DMhalos[zlist], survey.DMGs[zlist], dmvals, grid.state.MW.sigmaDMG, grid.state.MW.sigmaHalo, grid.state.MW.percentDMG, grid.state.MW.logu) pvals = np.zeros(len(izs1)) for i in range(len(izs1)): - pvals[i] = np.sum(rates[izs1[i],iweights[i]] * dm_weights[i] * (1.-dkzs[i]) - + rates[izs2[i],iweights[i]] * dm_weights[i] * dkzs[i]) + pvals[i] = np.sum(rates[izs1[i],iweights[i]] * dm_weights[i] * dkzs1[i] + + rates[izs2[i],iweights[i]] * dm_weights[i] * dkzs2[i]) bad= pvals <= 0. flg_bad = False @@ -698,21 +713,20 @@ def calc_likelihoods_2D(grid,survey, # calculating p(DM) and p(z) if dolist==5: # calculates p(dm) - pdmvals = np.sum(rates[:,idms1],axis=0)*(1.-dkdms) - pdmvals += np.sum(rates[:,idms2],axis=0)*dkdms + pdmvals = np.sum(rates[:,idms1],axis=0)*dkdms1 + pdmvals += np.sum(rates[:,idms2],axis=0)*dkdms2 # implicit calculation of p(z|DM) from p(z,DM)/p(DM) #neither on the RHS is normalised so this is OK! pzgdmvals = pvals/pdmvals #calculates p(z) - pzvals = np.sum(rates[izs1,:],axis=1)*(1.-dkzs) - pzvals += np.sum(rates[izs2,:],axis=1)*dkzs + pzvals = np.sum(rates[izs1,:],axis=1)*dkzs1 + pzvals += np.sum(rates[izs2,:],axis=1)*dkzs2 # implicit calculation of p(z|DM) from p(z,DM)/p(DM) pdmgzvals = pvals/pzvals - for array in pdmvals,pzgdmvals,pzvals,pdmgzvals: bad=np.array(np.where(array <= 0.)) if bad.size > 0: @@ -783,7 +797,7 @@ def calc_likelihoods_2D(grid,survey, # - we want to make FRB width analogous to beam, THEN # - we need an analogous 'beam' (i.e. width) distribution to integrate over, # which gives the normalisation - + if psnr: # NOTE: to break this into a p(SNR|b) p(b) term, we first take # the relative likelihood of the threshold b value compare @@ -798,21 +812,26 @@ def calc_likelihoods_2D(grid,survey, gamma=grid.state.energy.gamma # Evaluate thresholds at the exact DMobs - kdmobs=(survey.DMs - np.median(survey.DMGs + survey.DMhalos))/ddm - kdmobs=kdmobs[zlist] - kdmobs[kdmobs<0] = 0 - idmobs1=kdmobs.astype('int') - idmobs2=idmobs1+1 - dkdmobs=kdmobs-idmobs1 # applies to idms2 - + # The thresholds have already been calculated at mean values + # of the below quantities. Hence, we use the DM relative to + # those means, not the actual DMEG for that FRB + #kdmobs=(survey.DMs - np.median(survey.DMGs + survey.DMhalos))/ddm + #kdmobs=kdmobs[zlist] + #kdmobs[kdmobs<0] = 0 + #idmobs1=kdmobs.astype('int') + #idmobs2=idmobs1+1 + #dkdmobs=kdmobs-idmobs1 # applies to idms2 + DMEGmeans = survey.DMs[zlist] - np.median(survey.DMGs + survey.DMhalos) + idmobs1,idmobs2,dkdmobs1,dkdmobs2 = grid.get_dm_coeffs(DMEGmeans) + # Linear interpolation - Eths = grid.thresholds[:,izs1,idmobs1]*(1.-dkdmobs)*(1-dkzs) - Eths += grid.thresholds[:,izs2,idmobs1]*(1.-dkdmobs)*dkzs - Eths += grid.thresholds[:,izs1,idmobs2]*dkdmobs*(1-dkzs) - Eths += grid.thresholds[:,izs2,idmobs2]*dkdmobs*dkzs + Eths = grid.thresholds[:,izs1,idmobs1]*dkdmobs1*dkzs1 + Eths += grid.thresholds[:,izs2,idmobs1]*dkdmobs1*dkzs2 + Eths += grid.thresholds[:,izs1,idmobs2]*dkdmobs2*dkzs1 + Eths += grid.thresholds[:,izs2,idmobs2]*dkdmobs2*dkzs2 - FtoE = grid.FtoE[izs1]*(1.-dkzs) - FtoE += grid.FtoE[izs2]*dkzs + FtoE = grid.FtoE[izs1]*dkzs1 + FtoE += grid.FtoE[izs2]*dkzs2 beam_norm=np.sum(survey.beam_o) @@ -820,12 +839,15 @@ def calc_likelihoods_2D(grid,survey, # We integrate p(snr|b,w) p(b,w) db dw. # I have no idea how this could be multidimensional psnr=np.zeros(Eths.shape[1]) + for i,b in enumerate(survey.beam_b): bEths=Eths/b # array of shape NFRB, 1/b bEobs=bEths*survey.Ss[zlist] + for j,w in enumerate(grid.eff_weights): temp=grid.array_diff_lf(bEobs[j,:],Emin,Emax,gamma) # * FtoE #one dim in beamshape, one dim in FRB - psnr += temp.T*survey.beam_o[i]*w*bEths[j,:] #multiplies by beam factors and weight + this_psnr = temp.T*survey.beam_o[i]*w*bEths[j,:] #multiplies by beam factors and weight + psnr += this_psnr # at this stage, we have the amplitude from diff power law # summed over beam and weight @@ -834,24 +856,23 @@ def calc_likelihoods_2D(grid,survey, # comparable to the 1D case - otherwise it is superfluous if grid.state.MW.sigmaDMG == 0.0: # Linear interpolation - sg = grid.sfr_smear[izs1,idms1]*(1.-dkdms)*(1-dkzs) - sg += grid.sfr_smear[izs2,idms1]*(1.-dkdms)*dkzs - sg += grid.sfr_smear[izs1,idms2]*dkdms*(1-dkzs) - sg += grid.sfr_smear[izs2,idms2]*dkdms*dkzs + sg = grid.sfr_smear[izs1,idms1]*dkdms1*dkzs1 + sg += grid.sfr_smear[izs2,idms1]*dkdms1*dkzs2 + sg += grid.sfr_smear[izs1,idms2]*dkdms2*dkzs1 + sg += grid.sfr_smear[izs2,idms2]*dkdms2*dkzs2 else: sg = np.zeros(len(izs1)) # For each FRB for i in range(len(izs1)): - sg[i] = np.sum(grid.sfr_smear[izs1[i],iweights[i]] * dm_weights[i] * (1.-dkzs[i]) - + grid.sfr_smear[izs2[i],iweights[i]] * dm_weights[i] * dkzs[i]) / np.sum(dm_weights[i]) + sg[i] = np.sum(grid.sfr_smear[izs1[i],iweights[i]] * dm_weights[i] * dkzs1[i] + + grid.sfr_smear[izs2[i],iweights[i]] * dm_weights[i] * dkzs2[i]) / np.sum(dm_weights[i]) - dV = grid.dV[izs1]*(1-dkzs) + grid.dV[izs2]*dkzs + dV = grid.dV[izs1]*dkzs1 + grid.dV[izs2]*dkzs2 # at this stage, sg and dV account for the DM distribution and SFR; # dV is the volume elements # we just multiply these together sgV = sg*dV wzpsnr = psnr.T*sgV - # this step weights psnr by the volumetric values @@ -2012,7 +2033,11 @@ def CalculateConstant(grid,survey): return constant def CalculateIntegral(rates,survey): - """ Calculates the total expected number of FRBs for that rate array and survey """ + """ + Calculates the total expected number of FRBs for that rate array and survey + + This does NOT include the aboslute number of FRBs (through the log-constant) + """ # check that the survey has a defined observation time if survey.TOBS is not None: diff --git a/zdm/misc_functions.py b/zdm/misc_functions.py index 3f44e54..1952a7b 100644 --- a/zdm/misc_functions.py +++ b/zdm/misc_functions.py @@ -2154,9 +2154,11 @@ def get_zdm_grid( MC: generate via Monte Carlo using dlas.monte_dm nz (int, optional): Size of grid in redshift. Defaults to 500. zmin (float,optional): Minimum z. Used only for log-spaced grids. - zmax (float, optional): Maximum z. Defaults to 5. + zmax (float, optional): Maximum z. Defaults to 5. Represents the + upper edge of the maximum zbin. ndm (int, optional): Size of grid in DM. Defaults to 1400. - dmmax ([type], optional): Maximum DM of grid. Defaults to 7000.. + dmmax ([type], optional): Maximum DM of grid. Defaults to 7000. + Represents the upper edge of the max bin in the DM grid. datdir (str, optional): Directory to load/save grid data. Defaults to 'GridData'. tag (str, optional): Label for grids (unique identifier). Defaults to "". orig (bool, optional): Use original calculations for @@ -2230,9 +2232,13 @@ def get_zdm_grid( ) # labelled pickled files with H0 if new: - + ddm = dmmax / ndm - + # the DMvals and the zvals generated below + # represent bin centres. i.e. characteristic values. + # Probabilities then derived will correspond + # to p(zbin-0.5*dz < z < zbin+0.5*dz) etc. + if zlog: # generates a pseudo-log spacing # grid values increase with \sqrt(log) @@ -2241,14 +2247,19 @@ def get_zdm_grid( zvals = np.logspace(lzmin, lzmax, nz) else: dz = zmax / nz - zvals = (np.arange(nz) + 1) * dz - dmvals = (np.arange(ndm) + 1) * ddm - - dmmeans = dmvals[1:] - (dmvals[1] - dmvals[0]) / 2.0 + zvals = (np.arange(nz) + 0.5) * dz + dmvals = (np.arange(ndm) + 0.5) * ddm + + # Deprecated. dmvals now mean bin centre values + # dmmeans used to be those bin centres + #dmmeans = dmvals[1:] - (dmvals[1] - dmvals[0]) / 2.0 + + # initialises zDM grid zdmgrid = np.zeros([nz, ndm]) if method == "MC": # generate DM grid from the models + # NOT CHECKED if verbose: print("Generating the zdm Monte Carlo grid") nfrb = 10000 @@ -2270,6 +2281,7 @@ def get_zdm_grid( if orig: C0s = pcosmic.make_C0_grid(zvals, state.IGM.logF) else: + # interpolate C0 as a function of log10F f_C0_3 = cosmic.grab_C0_spline() actual_F = 10 ** (state.IGM.logF) sigma = actual_F / np.sqrt(zvals) diff --git a/zdm/pcosmic.py b/zdm/pcosmic.py index f304501..61da886 100644 --- a/zdm/pcosmic.py +++ b/zdm/pcosmic.py @@ -47,11 +47,16 @@ def pcosmic(delta, z, logF, C0): delta: = DM_cosmic/ i.e. fractional cosmic DM - alpha, beta: fitted parameters - constraints: std dev must be f*z^0.5 + z: redshift (sigma depends on this) - A: arbitrary amplitude of relative probability, ignore... + logF: log10 of the fluctuation constant, F + + C0: constant to be optimised + + alpha, beta: these are fitted parameters to be optimised + + constraints: std dev must be F*z^0.5 """ ### logF compensation @@ -77,14 +82,23 @@ def p_delta_DM(z, F, C0, deltas=None, dmin=1e-3, dmax=10, ndelta=10000): def iterate_C0(z, F, C0=1, Niter=10): - """ Iteratively solves for C_0 as a function of z and F """ + """ + Iteratively solves for C_0 as a function of z and F + + C0 goes through 10 iterations, where each iteration + uses the prior value of C0 to calculate C0. + + """ dmin = 1e-3 dmax = 10 ndelta = 10000 + # these represent central bin values of delta = DM/ deltas = np.linspace(dmin, dmax, ndelta) + bin_w = deltas[1] - deltas[0] for i in np.arange(Niter): + # pcosmic is a probability density + # hence, we should calculate this at bin centres pdeltas = pcosmic(deltas, z, F, C0) - bin_w = deltas[1] - deltas[0] norm = bin_w * np.sum(pdeltas) mean = bin_w * np.sum(pdeltas * deltas) / norm C0 += mean - 1.0 @@ -102,12 +116,15 @@ def make_C0_grid(zeds, F): return C0s -def get_mean_DM(zeds: np.ndarray, state: parameters.State): +def get_mean_DM(zeds: np.ndarray, state: parameters.State,Plot=False): """ Gets mean average z to which can be applied deltas Args: zeds (np.ndarray): redshifts (must be linearly spaced) - state (parameters.State): + These zeds are assumed to represent mid-points of + bins, i.e. from 0.5, 1.5, 2.5 etc dz + state (parameters.State): + Plot (bool): create a test plot of DM vs z Returns: np.ndarray: DM_cosmic @@ -117,14 +134,58 @@ def get_mean_DM(zeds: np.ndarray, state: parameters.State): H0=state.cosmo.H0, Ob0=state.cosmo.Omega_b, Om0=state.cosmo.Omega_m ) # - zmax = zeds[-1] + dz = zeds[1]-zeds[0] + zmax = zeds[-1] + dz/2. # top of uppermost zbin nz = zeds.size - DMbar, zeval = igm.average_DM(zmax, cosmo=cosmo, cumul=True, neval=nz + 1) - - # Check - assert np.allclose(zeds, zeval[1:]) - # wrong dimension - return DMbar[1:].value + + # this routine offsets zeval by 1 unit. That is, DM[i] + # is the mean cosmic DM at zeval[i+1] + # we want this for every 0.5, 1.5 etc + # hence, we evaluate at 2*nz+1, + tempDMbar, zeval = igm.average_DM(zmax, cosmo=cosmo, cumul=True, neval=2*nz + 1) + + # we now exract the DMbar that we actually want! + # the zeroeth DMbar corresponds to zeval[1] which + # since we calculate too many, is zeds[1] + DMbar = tempDMbar[:-1:2] + + # performs a test to check if igm.average_DM has been fixed yet or not + if np.abs(DMbar[0]/DMbar[1] - 1./3.) > 1e-2: + print("DMbar is not scaling as expected! Central bins ", + zeds[0]," and ",zeds[1]," have respective DM of ", + DMbar[0]," and ",DMbar[1]," . Expected the second ", + "value to be ",DMbar[0]*3.," . Perhaps ", + igm.average_DM," has been fixed?",DMbar[0]/DMbar[1] - 1./3.) + exit() + + if Plot: + plt.figure() + plt.xlabel('z') + plt.ylabel('DM') + dz = zeval[1]-zeval[0] + plt.plot(zeds,DMbar,marker='+',label="wanted") + plt.plot(zeval+dz,tempDMbar,linestyle=":",marker='x',label="eval") + plt.xlim(0,0.1) + plt.ylim(0,100) + plt.legend() + plt.tight_layout() + plt.show() + + plt.xlim(4.9,5.) + plt.ylim(4900,5500) + plt.tight_layout() + plt.show() + + plt.close() + + + # Remove this check - now replaced as above + # assert np.allclose(zeds, zeval[1:]) + + # now returns the actual values, since + # we have modified DMbar to exclude the + # zero value already + return DMbar.value def get_log_mean_DM(zeds: np.ndarray, state: parameters.State): @@ -194,9 +255,9 @@ def get_pDM(z, F, DMgrid, zgrid, Fgrid, C0grid, zlog=False): """ C0 = get_C0(z, F, zgrid, Fgrid, C0grid) if zlog: - DMbar = get_mean_DM(z) - else: DMbar = get_log_mean_DM(z) + else: + DMbar = get_mean_DM(z) deltas = DMgrid / DMbar # in units of fractional DM pDM = pcosmic(deltas, z, F, C0) return pDM @@ -209,8 +270,9 @@ def get_pDM_grid( state C0grid: C0 values obtained by convergence DMgrid: range of DMs for which we are generating a histogram + This represent bin centres zgrid: redshifts. These do not have to be in any particular - order or spacing. + order or spacing. We just iterature through these zlog (bool): True if zs are log-spaced False if linearly spaced @@ -277,18 +339,31 @@ def plot_mean(zvals, saveas, title="Mean DM"): def get_dm_mask(dmvals, params, zvals=None, plot=False): - """ Generates a mask over which to integrate the lognormal - Apply this mask as DM[i] = DM[set[i]]*mask[i] - DMvals: these give local probabilities of p(DM). + """ + Generates a mask over which to integrate the lognormal + distribution of FRM host galaxy DM contributions. It's + essentially just a probability distribution of p(DMhost), + such that p(DMeg = DMhost + DMcosmic) can be quickly + calculated, as DM[i] = DM[set[i]]*mask[i] + + DMvals (np.ndarray): DMs over which to calculate the mask. + These represent local probabilities of p(DM), i.e. + the probability of getting a DM between + DMval - dDM/2. and DMval + dDM/2. + + params [vector, 2]: mean and sigma of the lognormal (log10) + host galaxy DM distribution + + zvals [np.ndarray]: redshift values at which to calculate this. + If None: return a single, redshift-independent vector. + If not None: return a mask for each value of z, with + DMhost reduced by the (1+z) value. + In future: add a parameter to scale this as (1+z)^xi. + We simply assign lognormal values at the midpoints The renormalisation constants then give some idea of the error in this procedure - This requires parameters to be passed as a vector - - dmvals: numpy array storing the dms over which to calculate the mask - - params [vector, 2]: mean and sigma of the lognormal (log10) distribution """ if len(params) != 2: @@ -297,7 +372,7 @@ def get_dm_mask(dmvals, params, zvals=None, plot=False): params, " (expected log10mean, log10sigma)", ) - exit() + # expect the params to be log10 of actual values for simplicity # this converts to natural log logmean = params[0] / 0.4342944619 @@ -307,9 +382,11 @@ def get_dm_mask(dmvals, params, zvals=None, plot=False): ##### first generates a mask from the lognormal distribution ##### # in theory allows a mask up to length of the DN values, but will - # get truncated - # the first value has half weight (0 to 0.5) - # the rest have width of 1 + # get truncated. + # The first value has half weight (0dDM to 0.5dDM) and represents + # adding no new DM. The rest have width of dDM and represent + # adding an integer number of dDM intervals. + mask = np.zeros([dmvals.size]) if zvals is not None: ndm = dmvals.size @@ -318,7 +395,8 @@ def get_dm_mask(dmvals, params, zvals=None, plot=False): for j, z in enumerate(zvals): # with each redshift, we reduce the effects of a 'host' contribution by (1+z) # this means that we divide the value of logmean by 1/(1+z) - # or equivalently, we multiply the ddm by this factor + # or equivalently, we multiply the ddm by this factor, since a + # measurable increase of dDM means an intrinsic dDM*(1+z) # here we choose the latter, but it is the same mask[j, :] = integrate_pdm(ddm * (1.0 + z), ndm, logmean, logsigma) mask[j, :] /= np.sum(mask[j, :]) # the mask must integrate to unity @@ -366,6 +444,8 @@ def integrate_pdm(ddm, ndm, logmean, logsigma, quick=True, plot=False): Assigns probabilities of DM smearing (e.g. due to the host galaxy contribution) to a histogram in dm space. + Here, the resulting mask assumes DM values of 0 (0 to 0.5), 1( 0.5 to 1.5) etc. + Two methods: quick (use central values of DM bins), and slow (integrate bins) Arguments: @@ -386,8 +466,12 @@ def integrate_pdm(ddm, ndm, logmean, logsigma, quick=True, plot=False): Returns: mask (np.ndarray) """ - # do this for the z=0 case - + # do this for the z=0 case (handling of z>0 can be performed at + # when calling the routine by multiplying dDM values) + + # normalisation constant of a normal distribution. + # Normalisation should probably be redone afterwards + # anyway. norm = (2.0 * np.pi) ** -0.5 / logsigma # csum=pdm @@ -405,6 +489,9 @@ def integrate_pdm(ddm, ndm, logmean, logsigma, quick=True, plot=False): if plot or not quick: m2 = np.zeros([ndm]) args = (logmean, logsigma, norm) + # performs integration of first bin in log space + # Does this for the first bin: probability from + # "0" (-logsigma*10) to ddm*0.5 pdm, err = sp.integrate.quad( loglognormal_dlog, np.log(ddm * 0.5) - logsigma * 10, @@ -412,7 +499,9 @@ def integrate_pdm(ddm, ndm, logmean, logsigma, quick=True, plot=False): args=args, ) m2[0] = pdm - + + # performs the integration for all other bins; + # goes from lower to upper bin bounds for i in np.arange(1, ndm): # if csum > CSUMCUT: # imax=i @@ -439,5 +528,6 @@ def integrate_pdm(ddm, ndm, logmean, logsigma, quick=True, plot=False): plt.savefig("dm_mask_comparison_plot.pdf") plt.close() print("Generated plot of dm masks, exiting...") - exit() # quit to avoid infinite plots + # Quit to avoid infinite plots. This is just a saftey measure. + exit() return mask diff --git a/zdm/scripts/mc_generation_test.py b/zdm/scripts/mc_generation_test.py new file mode 100644 index 0000000..747ea53 --- /dev/null +++ b/zdm/scripts/mc_generation_test.py @@ -0,0 +1,244 @@ +""" +This script is designed to test that MC event generation is working. + +It initiatlises a grid, and then generates MC events from it. + +It saves these to two 1D histograms in z and DM, and a 2D +histogram in zDM space (other MC info is currently discarded). + +It then generates plots to compare the MC generated events to +the grid they were generated from. + +Because generating sufficient MC takes quite some time, +it saves all data, and adds to it with each iteration. +Thus it can be run with 10,100,1000 etc MC events, +and then statistics can be accumulated. + +The MC events are saved in the following npy files: +- dmhist.npy A 1D histogram of dispersion measure +- totalcount.npy Single integer: number of MC events generated +- zdmhist.npy A 2D histogram in zDM space +- zhist.npy A 1D histogram of redshift +Note that for this to make sense, the dimensions of the original +grid should stay constant. + +Plots get generated at each iteration. These are: +- pz.pdf Compares expected and generated p(z) distributions +- pzerr.pdf Statistical error in the above in units of sigma +- pDM.pdf Compares expected and generated p(DM) distributions +- pDMerr.pdf Statistical error in the above in units of sigma +- MCzDM.pdf Generated 2D histogram of MC generated events +- rel_zDM_err.pdf Relative error in the above (units of sigma) +- grid_expectation.pdf Grid expected zDM distribution + + + +""" +import os + +from zdm import cosmology as cos +from zdm import misc_functions +from zdm import parameters +from zdm import survey +from zdm import pcosmic +from zdm import iteration as it +from zdm import loading +from zdm import io +from astropy.cosmology import Planck18 +import numpy as np +from zdm import survey +from matplotlib import pyplot as plt + +def main(): + + # Initialise surveys and grids + sdir='../data/Surveys/' + name='parkes_mb_class_I_and_II' + + state = parameters.State() + state.set_astropy_cosmo(Planck18) + + # approximate best-fit values from recent analysis + param_dict={'sfr_n': 0.8808527057055584, 'alpha': 0.7895161131856694, 'lmean': 2.1198711983468064, 'lsigma': 0.44944780033763343, 'lEmax': 41.18671139482926, + 'lEmin': 39.81049090314043, 'gamma': -1.1558450520609953, 'H0': 54.6887137195215, 'halo_method': 0, 'sigmaDMG': 0.0, 'sigmaHalo': 0.0} + state.update_params(param_dict) + names=["CRAFT_ICS_1300"] + surveys, grids = loading.surveys_and_grids(survey_names = names, repeaters=False, + init_state=state) + + s=surveys[0] + g=grids[0] + + ############ sets up the histograms ######### + + # OK, the zvals and DM vals represent FRBs found in bins *centred* at those values + dz = g.zvals[1]-g.zvals[0] + dDM = g.dmvals[1] - g.dmvals[0] + NZ = g.zvals.size + NDM = g.dmvals.size + zbins = np.linspace(g.zvals[0] - dz/2., g.zvals[-1] + dz/2.,NZ+1) + DMbins = np.linspace(g.dmvals[0] - dDM/2., g.dmvals[-1] + dDM/2.,NDM+1) + + zDMhistfile = "zdmhist.npy" + zhistfile = "dmhist.npy" + DMhistfile = "zhist.npy" + countfile = "totalcount.npy" + + if os.path.exists(zDMhistfile): + zDMhist = np.load(zDMhistfile) + DMhist = np.load(DMhistfile) + zhist = np.load(zhistfile) + NFRB = np.load(countfile) + else: + zDMhist = np.zeros([NZ,NDM]) + zhist = np.zeros([NZ]) + DMhist = np.zeros([NDM]) + NFRB=0 + + ############ generates the MCMC ######### + # array values are DM, z, B, w, and SNR + NMC = 100 + frbs = g.GenMCSample(NMC) + frbs = np.array(frbs) + + zs = frbs[:,0] + DMs = frbs[:,1] + snrs = frbs[:,3] + bs = frbs[:,2] + ws = frbs[:,4] + + # check it lives in the right bins + lowz = np.where(zs < zbins[0])[0] + highz = np.where(zs > zbins[-1])[0] + lowDM = np.where(DMs < DMbins[0])[0] + highDM = np.where(DMs > DMbins[-1])[0] + print("FRBs out of range: ",len(lowz),len(highz),len(lowDM),len(highDM)) + if len(lowz)>0: + print(zs[lowz]) + if len(highz)>0: + print(zs[highz]) + if len(lowDM)>0: + print(DMs[lowDM]) + if len(highDM)>0: + print(DMs[highDM]) + ############ histogram the data ######### + + tempzDMhist,xb,yb = np.histogram2d(zs,DMs,bins=[zbins,DMbins]) + tempzhist,xb = np.histogram(zs,bins=zbins) + tempDMhist,yb = np.histogram(DMs,bins=DMbins) + + zDMhist += tempzDMhist + zhist += tempzhist + DMhist += tempDMhist + + np.save(zDMhistfile,zDMhist) + np.save(zhistfile,zhist) + np.save(DMhistfile,DMhist) + + NFRB += NMC + np.save(countfile,NFRB) + + print("MC done, total number of simulated FRBs: ",NFRB) + ############Normalisation and Errors ######## + # z-only plot + + + # gets expectation + pz = np.sum(g.rates,axis=1) + pdm = np.sum(g.rates,axis=0) + pz /= np.sum(pz) + pdm /= np.sum(pdm) + pzdm = g.rates / np.sum(g.rates) + + zDMhisterr = zDMhist **0.5 + norm = np.sum(zDMhist) + zDMhisterr /= norm + zDMhist /= norm + + zhisterr = zhist **0.5 + norm = np.sum(zhist) + zhisterr /= norm + zhist /= norm + + DMhisterr = DMhist **0.5 + norm = np.sum(DMhist) + DMhisterr /= norm + DMhist /= norm + + # z-only plot + + plt.figure() + plt.plot(g.zvals,pz,linewidth=3,linestyle="-",color="black",label="expectation") + plt.plot(g.zvals,zhist,linestyle = "-",linewidth=2,label="MC") + plt.plot(g.zvals,zhist-zhisterr,linestyle = ":",linewidth=1,label="err",color=plt.gca().lines[-1].get_color()) + plt.plot(g.zvals,zhist+zhisterr,linestyle = ":",linewidth=1,color=plt.gca().lines[-1].get_color()) + plt.xlabel("z") + plt.ylabel("p(z) dz") + plt.xlim(0,2) + plt.legend() + plt.tight_layout() + plt.savefig("pz.pdf") + plt.close() + + plt.figure() + altzhisterr = (pz*NFRB)**0.5/NFRB + altrel_error = (zhist - pz)/altzhisterr + rel_error = (zhist - pz)/zhisterr + plt.plot(g.zvals,rel_error,linewidth=1,linestyle="-") + plt.xlabel("z") + plt.xlim(0,2) + plt.ylabel("$\\Delta p(z) / \\sigma p(z)$") + plt.tight_layout() + plt.savefig("pzerr.pdf") + plt.close() + + + # DM-only plot + plt.figure() + plt.plot(g.dmvals,pdm,linewidth=3,linestyle="-",color="black",label="expectation") + plt.plot(g.dmvals,DMhist,linestyle = "-",linewidth=2,label="MC") + plt.plot(g.dmvals,DMhist-DMhisterr,linestyle = ":",linewidth=1,label="err",color=plt.gca().lines[-1].get_color()) + plt.plot(g.dmvals,DMhist+DMhisterr,linestyle = ":",linewidth=1,color=plt.gca().lines[-1].get_color()) + plt.xlabel("DM") + plt.ylabel("p(DM) dDM") + plt.legend() + plt.tight_layout() + plt.savefig("pDM.pdf") + plt.close() + + + plt.figure() + altDMhisterr = (pdm*NFRB)**0.5/NFRB + altrel_error = (DMhist - pdm)/altDMhisterr + rel_error = (DMhist - pdm)/DMhisterr + plt.plot(g.dmvals,rel_error,linewidth=1,linestyle="-") + plt.xlabel("DM") + plt.ylabel("$\\Delta p({\\rm DM}) / \\sigma p({\\rm DM})$") + plt.tight_layout() + plt.savefig("pDMerr.pdf") + plt.close() + + + + # zDM 2D plot + # plots the p(DMEG (host + cosmic)|z) grid + misc_functions.plot_grid_2(zDMhist,g.zvals,g.dmvals, + name='MCzDM.pdf',norm=3,log=False, + label='$\\log_{10} p({\\rm DM}_{\\rm IGM} + {\\rm DM}_{\\rm host}|z)$ [a.u.]', + project=False) + + misc_functions.plot_grid_2(g.rates,g.zvals,g.dmvals, + name='grid_expectation.pdf',norm=3,log=False, + label='$\\log_{10} p({\\rm DM}_{\\rm IGM} + {\\rm DM}_{\\rm host}|z)$ [a.u.]', + project=False) + + expectation = g.rates / np.sum(g.rates) + zDMerr = (expectation * NFRB)**0.5 / NFRB + rel_err = (zDMhist - expectation)/zDMhisterr + + misc_functions.plot_grid_2(rel_err,g.zvals,g.dmvals, + name='rel_zDM_err.pdf',norm=3,log=False, + label='$\\sigma$ deviation', + project=False) + +main() diff --git a/zdm/survey.py b/zdm/survey.py index 7e732b3..772838d 100644 --- a/zdm/survey.py +++ b/zdm/survey.py @@ -1026,21 +1026,25 @@ def init_beam(self,plot=False, else: print("No beam found to initialise...") - def calc_max_dm(self,Nchan=336): + def calc_max_dm(self): ''' - Calculates the maximum searched DM + Calculates the maximum searched DM. + + Calculates bandwidth using ''' fbar=self.meta['FBAR'] t_res=self.meta['TRES'] nu_res=self.meta['FRES'] max_idt=self.meta['MAX_IDT'] max_dm=self.meta['MAX_DM'] - + print(max_idt, max_dm) if max_dm is None and max_idt is not None: k_DM=4.149 #ms GHz^2 pc^-1 cm^3 - f_low = fbar - (Nchan/2. - 1)*nu_res - f_high = fbar + (Nchan/2. - 1)*nu_res + #f_low = fbar - (Nchan/2. - 1)*nu_res + #f_high = fbar + (Nchan/2. - 1)*nu_res + f_low = fbar - self.meta['BW']/2. # bottom of lowest band + f_high = fbar + self.meta['BW']/2. # top of highest band max_dt = t_res * max_idt max_dm = max_dt / (k_DM * ((f_low/1e3)**(-2) - (f_high/1e3)**(-2)))