def test_marginal_ln_likelihood(tmpdir, case): prior, _ = get_prior(case) data, _ = make_data() prior_samples = prior.sample(size=100) joker = TheJoker(prior) # pass JokerSamples instance ll = joker.marginal_ln_likelihood(data, prior_samples) assert len(ll) == len(prior_samples) # save prior samples to a file and pass that instead filename = str(tmpdir / 'samples.hdf5') prior_samples.write(filename, overwrite=True) ll = joker.marginal_ln_likelihood(data, filename) assert len(ll) == len(prior_samples) # make sure batches work: ll = joker.marginal_ln_likelihood(data, filename, n_batches=10) assert len(ll) == len(prior_samples) # NOTE: this makes it so I can't parallelize tests, I think with MultiPool(processes=2) as pool: joker = TheJoker(prior, pool=pool) ll = joker.marginal_ln_likelihood(data, filename) assert len(ll) == len(prior_samples)
def pool(request): multimode = 'None' # multimode = 'Serial' # multimode = 'Multi' # multimode = 'MPI' # setup code pool = None if multimode == 'Serial': from schwimmbad import SerialPool pool = SerialPool() elif multimode == 'Multi': from schwimmbad import MultiPool pool = MultiPool() elif multimode == 'MPI': from schwimmbad import MPIPool pool = MPIPool() if not pool.is_master(): pool.wait() import sys sys.exit(0) # inject class variables request.cls.pool = pool yield # tear down if multimode == 'Multi' or multimode == 'MPI': pool.close()
def sample(self, n_samples): if self.pool is None or _GPU_ENABLED: pool = SerialPool() else: if isinstance(self.pool, int): pool = MultiPool(self.pool) elif isinstance(self.pool, (SerialPool, MultiPool)): pool = self.pool else: raise TypeError( "Does not understand the given multiprocessing pool.") drawn_samples = list( pool.map(self.draw_one_joint_posterior_sample_map, range(n_samples))) pool.close() drawn_zs = [drawn_samples[i][0] for i in range(n_samples)] drawn_inference_posteriors = [ drawn_samples[i][1] for i in range(n_samples) ] drawn_joint_posterior_samples = pd.DataFrame( drawn_inference_posteriors) drawn_joint_posterior_samples["redshift"] = drawn_zs return drawn_joint_posterior_samples
def main(data_path: Path, model_path: Path, mask_path: Path, log_level: int): logging.basicConfig( stream=sys.stdout, level=log_level, datefmt='%Y-%m-%d %H:%M', format='%(asctime)s - %(name)s - %(levelname)s - %(message)s') with open(data_path) as f: data_cfg = yaml.load(f, Loader=yaml.FullLoader) nside = data_cfg['nside'] npix = hp.nside2npix(nside) nmc = data_cfg['monte_carlo'] amplitude_output_shape = (nmc, 2, npix) parameter_output_shape = (nmc, npix) frequencies = np.array(data_cfg['frequencies']) with h5py.File(data_cfg['hdf5_path'], 'r') as f: data = f['maps/monte_carlo/data'][...] cov = f['maps/monte_carlo/cov'][...] masking = lynx.Masking(mask_path) fitting_masks = list(masking.get_fitting_indices()) tasks = get_tasks(data, cov, nmc, frequencies, fitting_masks, model_path) with MultiPool() as pool: results = pool.map(do_fitting, tasks) with h5py.File(data_cfg['hdf5_path'], 'a') as f: for result in results: save_data(f, amplitude_output_shape, parameter_output_shape, result)
def Regrid_PS(PS1, Corners): dim1, dim2 = Corners.shape[1:] dim1 -= 1 dim2 -= 1 global px, py px, py = np.where(PS1) global squares squares = np.array(Make_squares(Corners)) square_num = np.arange(0, len(squares)) points = np.zeros((len(px), 2)) points[:, 0] = px points[:, 1] = py global pspixels pspixels = Footprint_square(Corners, points) global psimage psimage = PS1.copy() pool = MultiPool() values = list(pool.map(Pix_sum, square_num)) pool.close() PS_scene = np.array(values) PS_scene = np.nansum(PS_scene, axis=0) PS_scene = PS_scene.astype('float') PS_scene = PS_scene.reshape(dim1, dim2) return PS_scene
def generate_joint_posterior_samples_from_marginalized_likelihood( joint_result, single_trigger_likelihoods, sep_char="^", ncores=1, ): common_parameters = [ p for p in list(joint_result.posterior.columns) if sep_char not in p ] independent_parameters = [ p for p in list(joint_result.posterior.columns) if sep_char in p ] logger = logging.getLogger(__prog__) logger.info( "Using {} CPU core(s) for generating joint posterior samples from marginalized likelihood" .format(ncores)) import tqdm with MultiPool(ncores) as pool: output_samples = pool.starmap( generate_joint_posterior_sample_from_marginalized_likelihood, tqdm.tqdm([[ row, single_trigger_likelihoods, common_parameters, independent_parameters, sep_char ] for _, row in joint_result.posterior.iterrows()])) # Edit data frame joint_result.posterior = pd.DataFrame(output_samples)
def generate_atlas_in_parallel_chunking(zval, chunksize, nchunks, fname='temp_parallel_atlas', filter_list='filter_list_goodss.dat', filt_dir='internal', priors=[], z_bw=0.05, pg_folder='parallel_atlases/'): N_pregrid = chunksize atlas_vals = [ fname, zval, priors, pg_folder, filter_list, filt_dir, N_pregrid ] time_start = time.time() try: with MultiPool() as pool: values = list(pool.map(partial(gen_pg_parallel, atlas_vals), data)) finally: print('Generated pregrid (%.0f chunks, %.0f sedsperchunk) at zval', zval) print('time taken: %.2f mins.' % ((time.time() - time_start) / 60)) # need to add code here to then concatenate chunks into a single file and delete the individual ones return
def generate_snrs(result, likelihood, ncores=1): logger = logging.getLogger(__prog__) logger.info("Using {} CPU core(s) for computing SNRs".format(ncores)) import tqdm with MultiPool(ncores) as pool: output_samples = pool.starmap( generate_snrs_per_sample, tqdm.tqdm([[row.to_dict(), likelihood] for _, row in result.posterior.iterrows()])) result.posterior = pd.DataFrame(output_samples) return result
def compute_mean_selection_function(selection_function, N_avg, pool=None): if pool is None: pool = SerialPool() elif isinstance(pool, int): pool = MultiPool(pool) elif isinstance(pool, (SerialPool, MultiPool)): pool = pool else: raise TypeError("Does not understand the given multiprocessing pool.") out = list( pool.starmap(selection_function.evaluate, [() for _ in range(N_avg)])) avg = np.average(out) pool.close() return avg
def test_de_dt_integrate(self): n_values = 10 m_1 = np.random.uniform(0, 10, n_values) * u.Msun m_2 = np.random.uniform(0, 10, n_values) * u.Msun f_orb = 10**(np.random.uniform(-5, -1, n_values)) * u.Hz ecc = np.random.uniform(0.0, 0.9, n_values) beta, a_i = evol.check_mass_freq_input(m_1=m_1, m_2=m_2, f_orb_i=f_orb) n_step = 100 c_0 = utils.c_0(a_i=a_i, ecc_i=ecc) timesteps = evol.create_timesteps_array(a_i=a_i, beta=beta, ecc_i=ecc, t_evol=1 * u.yr, n_step=n_step) t_merge = evol.get_t_merge_ecc(ecc_i=ecc, f_orb_i=f_orb, m_1=m_1, m_2=m_2) # remove any bad timesteps that would evolve past the merger bad_timesteps = timesteps >= t_merge[:, np.newaxis] timesteps[bad_timesteps] = -1 * u.Gyr previous = timesteps.max(axis=1).repeat(timesteps.shape[1]) timesteps[bad_timesteps] = previous.reshape( timesteps.shape)[bad_timesteps] # get rid of the units for faster integration c_0 = c_0.to(u.m).value beta = beta.to(u.m**4 / u.s).value timesteps = timesteps.to(u.s).value # integrate by hand: ecc_evol = np.array([ odeint(evol.de_dt, ecc[i], timesteps[i], args=(beta[i], c_0[i])).flatten() for i in range(len(ecc)) ]) # integrate with function: with MultiPool(processes=1) as pool: ecc_pool = np.array( list( pool.map(evol.integrate_de_dt, zip(ecc, timesteps, beta, c_0)))) self.assertTrue(np.allclose(ecc_evol, ecc_pool, equal_nan=True))
def calm2l_mcmc(infile, alfvar, ncpu, outname): """ - for test purpose """ ALFPY_HOME = os.environ['ALFPY_HOME'] samples = np.array( pd.read_csv(infile, delim_whitespace=True, header=None, comment='#'))[:, 1:47] tstart = time.time() with MultiPool() as pool: pwork = partial(worker_m2l, alfvar, key_list) ml_res = pool.map(pwork, samples) ndur = time.time() - tstart pool.close() print('\ncalculating m2l in .mcmc file: {:.2f}minutes'.format(ndur / 60.)) np.savez('{0}results/{1}_mcmcm2l_b.npz'.format(ALFPY_HOME, outname), m2l=ml_res) return np.array(ml_res)
def run_generation(env, parameters): num_agents = parameters.shape[0] rewards = np.zeros(num_agents) ''' for a in range(num_agents): weights = parameters[a,:] rewards[a] = run_episode(env, weights) ''' # replace loop with map operation for parallel processing with MultiPool() as pool: rewards = np.array(pool.map(run_episode, parameters)) best_reward = np.max(rewards) mean_reward = np.mean(rewards) ranking = rewards.argsort() return parameters[ranking[-5:], :], best_reward, mean_reward, parameters[ ranking[-1], :]
def generate_atlas_in_parallel_zgrid(zgrid, atlas_params, dynamic_decouple=True): """ Make a set of atlases given a redshift grid and a list of parameters (including a priors object). Atlas Params: [N_pregrid, priors, fname, store, path, filter_list, filt_dir, z_bw] """ time_start = time.time() try: with MultiPool() as pool: values = list( pool.map( partial(make_atlas_parallel, atlas_params=atlas_params), zgrid)) finally: time_end = time.time() print('time taken [parallel]: %.2f min.' % ((time_end - time_start) / 60))
def generate_atlas_in_parallel_chunking(chunksize, nchunks, fname='temp_parallel_atlas', filter_list='filter_list_goodss.dat', filt_dir='internal', priors=[], pg_folder='pregrids/'): """ Generate chunks of an atlas in parallel and combine them into one big atlas """ chunk_path = pg_folder + 'atlaschunks/' store_path = pg_folder atlas_vals = [fname, priors, chunk_path, filter_list, filt_dir, chunksize] time_start = time.time() data = np.arange(nchunks) try: with MultiPool() as pool: values = list( pool.map(partial(gen_pg_parallel, atlas_vals=atlas_vals), data)) finally: print('Generated pregrid (%.0f chunks, %.0f sedsperchunk)') print('time taken: %.2f mins.' % ((time.time() - time_start) / 60)) combine_pregrid_chunks(fname_base=fname, N_chunks=nchunks, N_pregrid=chunksize, N_param=priors.Nparam, path=chunk_path, store_path=store_path) return
def simulate(self, bandpasses, norm=None, seed=13, Ncpus=None): self.obs = np.array([]) np.random.seed(seed) tstep = 1 tmin = self.model.mintime() tmax = self.model.maxtime() time = np.arange(tmin, tmax + tstep, tstep) wstep = 10 wmin = self.model.minwave() wmax = self.model.maxwave() wavelen = np.arange(wmin, wmax + wstep, wstep) fluxes = self.model.flux(time, wavelen) norm = np.max(fluxes) if norm is None else norm fluxes = fluxes.T * norm / np.max(fluxes) lc = LightCurve(time, wavelen, fluxes) redshifts = list( sncosmo.zdist(self.zmin, self.zmax, time=self.duration, area=self.area)) tasks = list( zip(redshifts, [self] * len(redshifts), [bandpasses] * len(redshifts), [self.bias] * len(redshifts), [lc] * len(redshifts), np.random.randint(2**32 - 1, size=len(redshifts)))) with MultiPool(processes=Ncpus) as pool: observations = np.array(list(pool.map(survey_worker, tasks))) self.obs = observations
def fit_catalog(fit_cat, atlas_path, atlas_fname, output_fname, N_pregrid=10000, N_param=3, z_bw=0.05, f160_cut=100, fit_mask=[], zgrid=[], sfr_uncert_cutoff=2.0): cat_id, cat_zbest, cat_seds, cat_errs, cat_f160, cat_class_star = fit_cat #if not zgrid: if isinstance(zgrid, (np.ndarray)) == False: zgrid = np.arange(np.amin(cat_zbest), np.amax(cat_zbest), z_bw) fit_id = cat_id.copy() fit_logM_50 = np.zeros_like(cat_zbest) fit_logM_MAP = np.zeros_like(cat_zbest) fit_logM_16 = np.zeros_like(cat_zbest) fit_logM_84 = np.zeros_like(cat_zbest) fit_logSFRinst_50 = np.zeros_like(cat_zbest) fit_logSFRinst_MAP = np.zeros_like(cat_zbest) fit_logSFRinst_16 = np.zeros_like(cat_zbest) fit_logSFRinst_84 = np.zeros_like(cat_zbest) fit_logZsol_50 = np.zeros_like(cat_zbest) fit_logZsol_16 = np.zeros_like(cat_zbest) fit_logZsol_84 = np.zeros_like(cat_zbest) fit_Av_50 = np.zeros_like(cat_zbest) fit_Av_16 = np.zeros_like(cat_zbest) fit_Av_84 = np.zeros_like(cat_zbest) fit_zfit_50 = np.zeros_like(cat_zbest) fit_zfit_16 = np.zeros_like(cat_zbest) fit_zfit_84 = np.zeros_like(cat_zbest) fit_logMt_50 = np.zeros_like(cat_zbest) fit_logMt_16 = np.zeros_like(cat_zbest) fit_logMt_84 = np.zeros_like(cat_zbest) fit_logSFR100_50 = np.zeros_like(cat_zbest) fit_logSFR100_16 = np.zeros_like(cat_zbest) fit_logSFR100_84 = np.zeros_like(cat_zbest) fit_nparam = np.zeros_like(cat_zbest) fit_t25_50 = np.zeros_like(cat_zbest) fit_t25_16 = np.zeros_like(cat_zbest) fit_t25_84 = np.zeros_like(cat_zbest) fit_t50_50 = np.zeros_like(cat_zbest) fit_t50_16 = np.zeros_like(cat_zbest) fit_t50_84 = np.zeros_like(cat_zbest) fit_t75_50 = np.zeros_like(cat_zbest) fit_t75_16 = np.zeros_like(cat_zbest) fit_t75_84 = np.zeros_like(cat_zbest) fit_nbands = np.zeros_like(cat_zbest) fit_f160w = np.zeros_like(cat_zbest) fit_stellarity = np.zeros_like(cat_zbest) fit_chi2 = np.zeros_like(cat_zbest) fit_flags = np.zeros_like(cat_zbest) for i in (range(len(zgrid))): print('loading atlas at', zgrid[i]) # for a given redshift slice, zval = zgrid[i] # select the galaxies to be fit z_mask = (cat_zbest < (zval + z_bw / 2)) & (cat_zbest > (zval - z_bw / 2)) & (cat_f160 < f160_cut) fit_ids = np.arange(len(cat_zbest))[z_mask] # for gal_id in fit_ids: # gal_sed = cat_seds[gal_id, 0:] # gal_err = cat_errs[gal_id, 0:] # fit_likelihood, fit_norm_fac = evaluate_sed_likelihood(gal_sed,gal_err,atlas,fit_mask=[], # zbest=None,deltaz=None) # quants = get_quants(fit_likelihood, atlas, fit_norm_fac) print('starting parallel fitting for Ngals = ', len(fit_ids), ' at redshift ', str(zval)) try: # load the atlas fname = atlas_fname + '_zval_%.0f_' % (zgrid[i] * 10000) atlas = load_atlas(fname, N_pregrid, N_param=N_param, path=atlas_path) print('loaded atlas') with MultiPool() as pool: # note: Parallel doesn't work in Python2.6 # if not fit_mask: if isinstance(fit_mask, np.ndarray) == False: all_quants = list( pool.map( partial(fit_gals, catvals=(cat_seds, cat_errs, atlas)), fit_ids)) else: all_quants = list( pool.map( partial(fit_gals, catvals=(cat_seds, cat_errs, fit_mask, atlas)), fit_ids)) print('finished fitting parallel zbest chunk at z=%.3f' % zval) print('starting to put values in arrays') for ii, gal_id in enumerate(fit_ids): gal_sed = cat_seds[gal_id, 0:] gal_err = cat_errs[gal_id, 0:] quants = all_quants[ii][0] fit_likelihood = all_quants[ii][1] # fit_logM_MAP[gal_id] = all_quants[ii][2] # fit_logSFRinst_MAP[gal_id] = all_quants[ii][3] fit_logM_50[gal_id] = quants[0][0] fit_logM_16[gal_id] = quants[0][1] fit_logM_84[gal_id] = quants[0][2] fit_logSFRinst_50[gal_id] = quants[1][0] fit_logSFRinst_16[gal_id] = quants[1][1] fit_logSFRinst_84[gal_id] = quants[1][2] fit_Av_50[gal_id] = quants[2][0] fit_Av_16[gal_id] = quants[2][1] fit_Av_84[gal_id] = quants[2][2] fit_logZsol_50[gal_id] = quants[3][0] fit_logZsol_16[gal_id] = quants[3][1] fit_logZsol_84[gal_id] = quants[3][2] fit_zfit_50[gal_id] = quants[4][0] fit_zfit_16[gal_id] = quants[4][1] fit_zfit_84[gal_id] = quants[4][2] fit_logMt_50[gal_id] = quants[5][0][0] fit_logMt_16[gal_id] = quants[5][1][0] fit_logMt_84[gal_id] = quants[5][2][0] fit_logSFR100_50[gal_id] = quants[5][0][1] fit_logSFR100_16[gal_id] = quants[5][1][1] fit_logSFR100_84[gal_id] = quants[5][2][1] fit_nparam[gal_id] = quants[5][0][2] fit_t25_50[gal_id] = quants[5][0][3] fit_t25_16[gal_id] = quants[5][1][3] fit_t25_84[gal_id] = quants[5][2][3] fit_t50_50[gal_id] = quants[5][0][4] fit_t50_16[gal_id] = quants[5][1][4] fit_t50_84[gal_id] = quants[5][2][4] fit_t75_50[gal_id] = quants[5][0][5] fit_t75_16[gal_id] = quants[5][1][5] fit_t75_84[gal_id] = quants[5][2][5] fit_nbands[gal_id] = np.sum(gal_sed > 0) fit_f160w[gal_id] = cat_f160[gal_id] fit_stellarity[gal_id] = cat_class_star[gal_id] fit_chi2[gal_id] = np.amin(fit_likelihood) # flagging galaxies that either # 1. have nan values for mass # 2. have SFR uncertainties > sfr_uncert_cutoff # 3. are flagged as a star # 4. have extremely large chi2 if np.isnan(quants[0][0]): fit_flags[gal_id] = 1.0 elif (np.abs(fit_logSFRinst_84[gal_id] - fit_logSFRinst_16[gal_id]) > sfr_uncert_cutoff): fit_flags[gal_id] = 2.0 elif (cat_class_star[gal_id] > 0.5): fit_flags[gal_id] = 3.0 elif (fit_chi2[gal_id] > 1000): fit_flags[gal_id] = 4.0 else: fit_flags[gal_id] = 0.0 except: print('couldn\'t fit with pool at z=', zval) print('finishing that') pl.clf() pl.figure(figsize=(12, 6)) pl.hist(cat_zbest[cat_zbest > 0], np.arange(0, 6, z_bw), color='black', alpha=0.3) #pl.hist(fit_zfit_50[fit_zfit_50>0],np.arange(0,6,z_bw),color='royalblue') pl.hist(cat_zbest[fit_zfit_50 > 0], np.arange(0, 6, z_bw), color='royalblue') pl.title('fit %.0f/%.0f galaxies' % (np.sum(fit_zfit_50 > 0), len(cat_zbest))) pl.xlabel('redshift') pl.ylabel('# galaxies') display.clear_output(wait=True) display.display(pl.gcf()) pl.show() #'logSFRinst_MAP':fit_logSFRinst_MAP, #'logM_MAP':fit_logM_MAP, fit_mdict = { 'ID': fit_id, 'logM_50': fit_logM_50, 'logM_16': fit_logM_16, 'logM_84': fit_logM_84, 'logSFRinst_50': fit_logSFRinst_50, 'logSFRinst_16': fit_logSFRinst_16, 'logSFRinst_84': fit_logSFRinst_84, 'logZsol_50': fit_logZsol_50, 'logZsol_16': fit_logZsol_16, 'logZsol_84': fit_logZsol_84, 'Av_50': fit_Av_50, 'Av_16': fit_Av_16, 'Av_84': fit_Av_84, 'zfit_50': fit_zfit_50, 'zfit_16': fit_zfit_16, 'zfit_84': fit_zfit_84, 'logMt_50': fit_logMt_50, 'logMt_16': fit_logMt_16, 'logMt_84': fit_logMt_84, 'logSFR100_50': fit_logSFR100_50, 'logSFR100_16': fit_logSFR100_16, 'logSFR100_84': fit_logSFR100_84, 't25_50': fit_t25_50, 't25_16': fit_t25_16, 't25_84': fit_t25_84, 't50_50': fit_t50_50, 't50_16': fit_t50_16, 't50_84': fit_t50_84, 't75_50': fit_t75_50, 't75_16': fit_t75_16, 't75_84': fit_t75_84, 'nparam': fit_nparam, 'nbands': fit_nbands, 'F160w': fit_f160w, 'stellarity': fit_stellarity, 'chi2': fit_chi2, 'fit_flags': fit_flags } fit_cat = Table(fit_mdict) fit_cat.write(output_fname, format='ascii.commented_header') return
def evol_ecc(ecc_i, t_evol=None, n_step=100, timesteps=None, beta=None, m_1=None, m_2=None, a_i=None, f_orb_i=None, output_vars=['ecc', 'f_orb'], n_proc=1, avoid_merger=True, exact_t_merge=False, t_before=1 * u.Myr, t_merge=None): """Evolve an array of eccentric binaries for ``t_evol`` time This function use Peters & Mathews (1964) Eq. 5.11 and 5.13. Note that all of {``beta``, ``m_1``, ``m_2``, ``ecc_i``, ``a_i``, ``f_orb_i``} must have the same dimensions. Parameters ---------- ecc_i : `float/array` Initial eccentricity t_evol : `float/array` Amount of time for which to evolve each binaries. Required if ``timesteps`` is None. Defaults to merger times. n_steps : `int` Number of timesteps to take between t=0 and t=``t_evol``. Required if ``timesteps`` is None. Defaults to 100. timesteps : `float/array` Array of exact timesteps to take when evolving each binary. Must be monotonically increasing and start with t=0. Either supply a 1D array to use for every binary or a 2D array that has a different array of timesteps for each binary. ``timesteps`` is used in place of ``t_evol`` and ``n_steps`` and takes precedence over them. beta : `float/array` Constant defined in Peters and Mathews (1964) Eq. 5.9. See :meth:`legwork.utils.beta` (if supplied ``m_1`` and `m_2` are ignored) m_1 : `float/array` Primary mass (required if ``beta`` is None or if ``output_vars`` contains a frequency) m_2 : `float/array` Secondary mass (required if ``beta`` is None or if ``output_vars`` contains a frequency) a_i : `float/array` Initial semi-major axis (if supplied ``f_orb_i`` is ignored) f_orb_i : `float/array` Initial orbital frequency (required if ``a_i`` is None) output_vars : `array` List of **ordered** output vars, choose from any of ``timesteps``, ``ecc``, ``a``, ``f_orb`` and ``f_GW`` for which of timesteps, eccentricity, semi-major axis and orbital/GW frequency that you want. Default is [``ecc``, ``f_orb``] n_proc : `int` Number of processors to split eccentricity evolution over, where the default is n_proc=1 avoid_merger : `boolean` Whether to avoid integration around the merger of the binary. Warning: setting this to false will result in many LSODA errors to be outputted since the derivatives get so large. exact_t_merge : `boolean` Whether to calculate the merger time exactly or use a fit (only relevant when ``avoid_merger`` is set to True t_before : `float` How much time before the merger to cutoff the integration (default is 1 Myr - this will prevent all LSODA warnings for e < 0.95, you may need to increase this time if your sample is more eccentric than this) t_merge : `float/array` Merger times for each source to be evolved. Only used when `avoid_merger=True`. If `None` then these will be automatically calculated either approximately or exactly based on the values of `exact_t_merge`. Returns ------- evolution : `array` Array possibly containing eccentricity, semi-major axis, timesteps and frequency evolution. Content determined by ``output_vars`` """ # transform input if only a single source arrayed_args, single_source = utils.ensure_array(m_1, m_2, beta, a_i, f_orb_i, ecc_i) m_1, m_2, beta, a_i, f_orb_i, ecc_i = arrayed_args output_vars = np.array([output_vars]) if isinstance(output_vars, str) else output_vars beta, a_i = check_mass_freq_input(beta=beta, m_1=m_1, m_2=m_2, a_i=a_i, f_orb_i=f_orb_i) if np.isin(output_vars, ["f_orb", "f_GW"]).any() and (m_1 is None or m_2 is None): raise ValueError( "`m_1`` and `m_2` required if `output_vars` contains a frequency") c_0 = utils.c_0(a_i=a_i, ecc_i=ecc_i) timesteps = create_timesteps_array(a_i=a_i, beta=beta, ecc_i=ecc_i, t_evol=t_evol, n_step=n_step, timesteps=timesteps) # if avoiding the merger during integration if avoid_merger: if t_merge is None: # calculate the merger time t_merge = get_t_merge_ecc(ecc_i=ecc_i, a_i=a_i, beta=beta, exact=exact_t_merge).to(u.Gyr) # warn the user if they are evolving past the merger if np.any(timesteps > t_merge[:, np.newaxis]): print( "WARNING: Some timesteps are past the merger of the source and this may produce erroneous", "results in combination with `avoid_merger=True`. Only evolve sources until their merger", "or set `avoid_merger=False`.") # make a mask for any timesteps that are too close to the merger too_close = timesteps >= t_merge[:, np.newaxis] - t_before check = too_close check[:, 0] = True if np.all(check): # pragma: no cover print("WARNING: All timesteps are too close to merger so", "evolution is not possible. Either set `t_before` to a", "smaller time or turn off `avoid_merger`") # ensure that the first timestep is always valid too_close[:, 0] = False if np.any(too_close): # set them all equal to the previous timestep before passing limit timesteps[too_close] = -1 * u.Gyr previous = timesteps.max(axis=1).repeat(timesteps.shape[1]) timesteps[too_close] = previous.reshape(timesteps.shape)[too_close] # get rid of the units for faster integration c_0 = c_0.to(u.m).value beta = beta.to(u.m**4 / u.s).value timesteps = timesteps.to(u.s).value # perform the evolution if n_proc > 1: with MultiPool(processes=n_proc) as pool: ecc_evol = np.array( list( pool.map(integrate_de_dt, zip(ecc_i, timesteps.tolist(), beta, c_0)))) else: ecc_evol = np.array([ odeint(de_dt, ecc_i[i], timesteps[i], args=(beta[i], c_0[i])).flatten() for i in range(len(ecc_i)) ]) c_0 = c_0[:, np.newaxis] * u.m ecc_evol = np.nan_to_num(ecc_evol, nan=0.0) # calculate a_evol if any frequency or separation requested if np.isin(output_vars, ["a", "f_orb", "f_GW"]).any(): a_evol = utils.get_a_from_ecc(ecc_evol, c_0) # calculate f_orb_evol if any frequency requested if np.isin(output_vars, ["f_orb", "f_GW"]).any(): # change merged binaries to extremely small separations a_not0 = np.where(a_evol.value == 0.0, 1e-30 * a_evol.unit, a_evol) f_orb_evol = utils.get_f_orb_from_a(a=a_not0, m_1=m_1[:, np.newaxis], m_2=m_2[:, np.newaxis]) # change frequencies back to 1Hz since LISA can't measure above f_orb_evol = np.where(a_not0.value == 1e-30, 1e2 * u.Hz, f_orb_evol) # construct evolution output evolution = [] for var in output_vars: if var == "timesteps": timesteps = timesteps.flatten() if single_source else timesteps evolution.append((timesteps * u.s).to(u.yr)) elif var == "ecc": ecc_evol = ecc_evol.flatten() if single_source else ecc_evol evolution.append(ecc_evol) elif var == "a": a_evol = a_evol.flatten() if single_source else a_evol evolution.append(a_evol.to(u.AU)) elif var == "f_orb": f_orb_evol = f_orb_evol.flatten() if single_source else f_orb_evol evolution.append(f_orb_evol.to(u.Hz)) elif var == "f_GW": f_orb_evol = f_orb_evol.flatten() if single_source else f_orb_evol evolution.append(2 * f_orb_evol.to(u.Hz)) return evolution if len(evolution) > 1 else evolution[0]
def evolve(cls, initialbinarytable, pool=None, **kwargs): """After setting a number of initial conditions we evolve the system. Parameters ---------- initialbinarytable : DataFrame Initial conditions of the binary **kwargs: There are three ways to tell evolve and thus the fortran what you want all the flags and other BSE specific parameters to be. If you pass both a dictionary of flags and/or a inifile and a table with the BSE parameters in the columns, the column values will be overwritten by what is in the dictionary or ini file. NUMBER 1: PASS A DICTIONARY OF FLAGS BSEDict NUMBER 2: PASS A PANDAS DATA FRAME WITH PARAMS DEFINED AS COLUMNS All you need is the initialbinarytable if the all the BSE parameters are defined as columns NUMBER 3: PASS PATH TO A INI FILE WITH THE FLAGS DEFINED params randomseed : `int`, optional, default let numpy choose for you If you would like the random seed that the underlying fortran code uses to be the same for all of the initial conditions you passed then you can send this keyword argument in. It is recommended to just let numpy choose a random number as the Fortran random seed and then this number will be returned as a column in the initial binary table so that you can reproduce the results. nproc : `int`, optional, default: 1 number of CPUs to use to evolve systems in parallel idx : `int`, optional, default: 0 initial index of the bcm/bpp arrays dtp : `float`, optional: default: tphysf timestep size in Myr for bcm output where tphysf is total evolution time in Myr n_per_block : `int`, optional, default: -1 number of systems to evolve in a block with _evolve_multi_system, to allow larger multiprocessing queues and reduced overhead. If less than 1 use _evolve_single_system Returns ------- output_bpp : DataFrame Evolutionary history of each binary output_bcm : DataFrame Final state of each binary initialbinarytable : DataFrame Initial conditions for each binary """ idx = kwargs.pop('idx', 0) nproc = min(kwargs.pop('nproc', 1), len(initialbinarytable)) n_per_block = kwargs.pop('n_per_block', -1) # There are three ways to tell evolve and thus the fortran # what you want all the flags and other BSE specific # parameters to be # NUMBER 1: PASS A DICTIONARY OF FLAGS BSEDict = kwargs.pop('BSEDict', {}) # NUMBER 2: PASS A PANDAS DATA FRAME WITH PARAMS DEFINED AS COLUMNS # All you need is the initialbinarytable with columns, # If you pass both a dictionary of flags and/or a inifile # and a table with the columns, the column values will be # overwritten by what is in the dictionary or ini file # NUMBER 3: PASS PATH TO A INI FILE WITH THE FLAGS DEFINED params = kwargs.pop('params', None) if BSEDict and params is not None: raise ValueError('Please pass either a dictionary ' 'of BSE flags or a path to an inifle not both.') if params is not None: if not os.path.isfile(params): raise ValueError( "File does not exist, probably supplied incorrect " "path to the inifile.") BSEDict, _, _, _, _ = utils.parse_inifile(params) # error check the parameters you are trying to pass to BSE # if we sent in a table with the parameter names # then we will temporarily create a dictionary # in order to verify that the values in the table # are valid utils.error_check(BSEDict) # check the initial conditions of the system and warn user if # anything is weird about them, such as the star starts # in Roche Lobe overflow utils.check_initial_conditions(initialbinarytable) # assign some columns based on keyword arguments but that # can be overwritten by the params or BSEDict if 'dtp' not in initialbinarytable.keys(): initialbinarytable = initialbinarytable.assign( dtp=kwargs.pop('dtp', initialbinarytable['tphysf'])) if 'randomseed' not in initialbinarytable.keys(): seed = np.random.randint(np.iinfo(np.int32).min, np.iinfo(np.int32).max, size=len(initialbinarytable)) initialbinarytable = initialbinarytable.assign( randomseed=kwargs.pop('randomseed', seed)) if 'bin_num' not in initialbinarytable.keys(): initialbinarytable = initialbinarytable.assign( bin_num=np.arange(idx, idx + len(initialbinarytable))) for k, v in BSEDict.items(): if k in initialbinarytable.keys(): warnings.warn( "The value for {0} in initial binary table is being " "overwritten by the value of {0} from either the params " "file or the BSEDict.".format(k)) # special columns that need to be handled differently if k == 'natal_kick_array': assign_natal_kick_array = [BSEDict['natal_kick_array'] ] * len(initialbinarytable) initialbinarytable = initialbinarytable.assign( natal_kick_array=assign_natal_kick_array) for idx, column_name in enumerate(NATAL_KICK_COLUMNS): for sn_idx in range(2): column_name_sn = column_name + '_{0}'.format(sn_idx + 1) column_values = pd.Series( [BSEDict['natal_kick_array'][sn_idx][idx]] * len(initialbinarytable), index=initialbinarytable.index, name=column_name_sn) kwargs1 = {column_name_sn: column_values} initialbinarytable = initialbinarytable.assign( **kwargs1) elif k == 'qcrit_array': initialbinarytable = initialbinarytable.assign( qcrit_array=[BSEDict['qcrit_array']] * len(initialbinarytable)) for kstar in range(0, 16): columns_values = pd.Series( [BSEDict['qcrit_array'][kstar]] * len(initialbinarytable), index=initialbinarytable.index, name='qcrit_{0}'.format(kstar)) initialbinarytable.loc[:, 'qcrit_{0}'. format(kstar)] = columns_values elif k == 'fprimc_array': columns_values = [BSEDict['fprimc_array'] ] * len(initialbinarytable) initialbinarytable = initialbinarytable.assign( fprimc_array=columns_values) for kstar in range(0, 16): columns_values = pd.Series( [BSEDict['fprimc_array'][kstar]] * len(initialbinarytable), index=initialbinarytable.index, name='fprimc_{0}'.format(kstar)) initialbinarytable.loc[:, 'fprimc_{0}'. format(kstar)] = columns_values else: # assigning values this way work for most of the parameters. kwargs1 = {k: v} initialbinarytable = initialbinarytable.assign(**kwargs1) # Here we perform two checks # First, if the BSE parameters are not in the initial binary table # and either a dictionary or an inifile was not provided # then we need to raise an ValueError and tell the user to provide # either a dictionary or an inifile or add more columns if not BSEDict: if ((not set(INITIAL_BINARY_TABLE_SAVE_COLUMNS).issubset( initialbinarytable.columns)) and (not set(INITIAL_CONDITIONS_PASS_COLUMNS).issubset( initialbinarytable.columns))): raise ValueError( "You are passing BSE parameters as columns in the " "initial binary table but not all BSE parameters are defined. " "Please pass a BSEDict or a params file or make sure " "you have all BSE parameters as columns {0} or {1}.". format(INITIAL_BINARY_TABLE_SAVE_COLUMNS, INITIAL_CONDITIONS_PASS_COLUMNS)) # If you did not supply the natal kick or qcrit_array or fprimc_array in the BSEdict then we construct # it from the initial conditions table if ((pd.Series(FLATTENED_NATAL_KICK_COLUMNS).isin( initialbinarytable.keys()).all()) and ('natal_kick_array' not in BSEDict)): column_values = initialbinarytable[ FLATTENED_NATAL_KICK_COLUMNS].values.reshape( -1, 2, len(NATAL_KICK_COLUMNS)).tolist() initialbinarytable = initialbinarytable.assign( natal_kick_array=column_values) if (pd.Series(QCRIT_COLUMNS).isin( initialbinarytable.keys()).all()) and ('qcrit_array' not in BSEDict): initialbinarytable = initialbinarytable.assign( qcrit_array=initialbinarytable[QCRIT_COLUMNS].values.tolist()) if (pd.Series(FPRIMC_COLUMNS).isin( initialbinarytable.keys()).all()) and ('fprimc_array' not in BSEDict): initialbinarytable = initialbinarytable.assign( fprimc_array=initialbinarytable[FPRIMC_COLUMNS].values.tolist( )) # need to ensure that the order of parameters that we pass to BSE # is correct initial_conditions = initialbinarytable[ INITIAL_CONDITIONS_PASS_COLUMNS].to_dict('records') # we use different columns to save the BSE parameters because some # of the parameters are list/arrays which we instead save as # individual values because it makes saving to HDF5 easier/more efficient. initialbinarytable = initialbinarytable[ INITIAL_BINARY_TABLE_SAVE_COLUMNS] # Allow a user to specify a custom time step sampling for certain parts of the evolution timestep_conditions = kwargs.pop('timestep_conditions', []) set_checkstates(timestep_conditions=timestep_conditions) # check if a pool was passed if pool is None: with MultiPool(processes=nproc) as pool: # evolve systems if n_per_block > 0: initial_conditions = np.asarray(initial_conditions) n_tot = initial_conditions.shape[0] initial_conditions_blocked = [] itr_block = 0 while itr_block < n_tot: itr_next = np.min([n_tot, itr_block + n_per_block]) initial_conditions_blocked.append( initial_conditions[itr_block:itr_next]) itr_block = itr_next output = list( pool.map(_evolve_multi_system, initial_conditions_blocked)) else: output = list( pool.map(_evolve_single_system, initial_conditions)) else: # evolve systems if n_per_block > 0: initial_conditions = np.asarray(initial_conditions) n_tot = initial_conditions.shape[0] initial_conditions_blocked = [] itr_block = 0 while itr_block < n_tot: itr_next = np.min([n_tot, itr_block + n_per_block]) initial_conditions_blocked.append( initial_conditions[itr_block:itr_next]) itr_block = itr_next output = list( pool.map(_evolve_multi_system, initial_conditions_blocked)) else: output = list( pool.map(_evolve_single_system, initial_conditions)) output = np.array(output, dtype=object) bpp_arrays = np.vstack(output[:, 1]) bcm_arrays = np.vstack(output[:, 2]) kick_info_arrays = np.vstack(output[:, 3]) natal_kick_arrays = np.vstack(output[:, 4]) natal_kick_arrays = natal_kick_arrays.reshape( -1, 1, len(FLATTENED_NATAL_KICK_COLUMNS)) for idx, column in enumerate(FLATTENED_NATAL_KICK_COLUMNS): # assigning values this way work for most of the parameters. kwargs1 = {column: natal_kick_arrays[:, :, idx]} initialbinarytable = initialbinarytable.assign(**kwargs1) kick_info = pd.DataFrame(kick_info_arrays, columns=KICK_COLUMNS, index=kick_info_arrays[:, -1].astype(int)) bpp = pd.DataFrame(bpp_arrays, columns=BPP_COLUMNS, index=bpp_arrays[:, -1].astype(int)) bcm = pd.DataFrame(bcm_arrays, columns=BCM_COLUMNS, index=bcm_arrays[:, -1].astype(int)) bcm.merger_type = bcm.merger_type.astype(int).astype(str).apply( lambda x: x.zfill(4)) bcm.bin_state = bcm.bin_state.astype(int) bpp.bin_num = bpp.bin_num.astype(int) bcm.bin_num = bcm.bin_num.astype(int) return bpp, bcm, initialbinarytable, kick_info
def main(): # setup paths for output dump_dir = os.path.join(WORK_DIR, 'hdf5_dump') if not os.path.exists(dump_dir): os.makedirs(dump_dir) # get the options for the code and set the data_release globally (ugly) to # allow MultiPool to work kwargs = plasticc.get_data.parse_getdata_options() global data_release data_release = kwargs.pop('data_release') getter = plasticc.get_data.GetData(data_release) # we can use model as a dummy string to indicate if we are generating # training or test data dummy = kwargs.pop('model') offset = kwargs.pop('offset') limit = kwargs.pop('limit') if dummy == 'training': outfile = os.path.join(dump_dir, 'training_set.hdf5') offset = None else: if limit is None: outfile = os.path.join(dump_dir, 'test_set.hdf5') else: if offset is None: offset = 0 outfile = os.path.join(dump_dir, 'test_n{}_set.hdf5'.format(offset)) # make sure we remove any lingering files if os.path.exists(outfile): os.remove(outfile) _ = kwargs.get('field') # set the header keywords for training and testing # same except for sntype will be removed from test and hostgal_photoz isn't # provided if dummy == 'training': kwargs['columns']=['objid','ptrobs_min','ptrobs_max','ra','decl', 'mwebv', 'mwebv_err',\ 'hostgal_photoz', 'hostgal_photoz', 'hostgal_photoz_err', 'sntype'] else: kwargs['columns']=['objid','ptrobs_min','ptrobs_max','ra','decl', 'mwebv', 'mwebv_err',\ 'hostgal_photoz', 'hostgal_photoz_err', 'sntype'] # set an extrasql query to get just the DDF and WFD objects # sntype for testing = true sntype + 100 if dummy == 'training': extrasql = "AND sntype < 100 AND ((objid LIKE 'WFD%') OR (objid LIKE 'DDF%'))" else: extrasql = "AND sntype > 100 AND ((objid LIKE 'WFD%') OR (objid LIKE 'DDF%'))" # set up options for data retrieval ignoring many of the command-line # options kwargs['extrasql'] = extrasql kwargs['model'] = '%' kwargs['field'] = '%' kwargs['sort'] = True kwargs['shuffle'] = False kwargs['limit'] = None kwargs['get_num_lightcurves'] = True total = getter.get_lcs_headers(**kwargs) total = list(total)[0] kwargs['limit'] = total kwargs['get_num_lightcurves'] = False kwargs['offset'] = offset head = getter.get_lcs_headers(**kwargs) # make a big list of the header - NOTE THAT WE ALWAYS RETRIEVE ALL OBJECTS head = list(head) if dummy == 'training': pass else: # if we're generating test data, if we set a limit, just draw a random # sample else shuffle the full list if limit is not None: head = random.sample(head, limit) else: random.shuffle(head) # convert the selected header entries to a table and remove uncessary columns out = at.Table(rows=head, names=kwargs['columns']) out.remove_columns(['ptrobs_min', 'ptrobs_max']) # galactic objects have -9 as redshift - change to 0 dummy_val = np.repeat(-9, len(out)) ind = np.isclose(out['hostgal_photoz'], dummy_val) out['hostgal_photoz'][ind] = 0. out['hostgal_photoz_err'][ind] = 0. # the object names have the model name in them, so we need to edit them # new name = <FIELD><SNID> orig_name = out['objid'] new_name = [ '{}{}'.format(x.split('_')[0], x.split('_')[-1]) for x in orig_name ] new_name = np.array(new_name) new_name = new_name.astype('bytes') out['objid'] = new_name # if we are generating test data, save a truth table if dummy == 'training': pass else: sntype = out['sntype'] # remove the model type from the output header that goes with the test data out.remove_column('sntype') truth_file = outfile.replace('_set.hdf5', '_truthtable.hdf5') if os.path.exists(truth_file): os.remove(truth_file) # ... saving it in the truth table only orig_name = orig_name.astype(bytes) new_name = new_name.astype(bytes) sntype = sntype.astype(bytes) truth_table = at.Table([orig_name, new_name, sntype], names=['objid', 'shortid', 'sntype']) truth_table.write(truth_file, compression=True, path='truth_table', serialize_meta=False, append=True) # write out the header out.write(outfile, compression=True, path='header', serialize_meta=False, append=True) nmc = len(out) # use a multiprocessing pool to load each light curve and dump to HDF5 with MultiPool() as pool: with tqdm(total=nmc) as pbar: for result in pool.imap(task, head): short_obj, thislc = result thislc.write(outfile, path=short_obj, compression=True, serialize_meta=False, append=True) pbar.update()
def initial_sample( self, M1min=0.08, M2min=0.08, M1max=150.0, M2max=150.0, porb_lo=0.15, porb_hi=8.0, rand_seed=0, size=None, nproc=1, pool=None, mp_seeds=None, ): """Sample initial binary distribution according to Moe & Di Stefano (2017) <http://adsabs.harvard.edu/abs/2017ApJS..230...15M>`_ Parameters ---------- M1min : `float` minimum primary mass to sample [Msun] DEFAULT: 0.08 M2min : `float` minimum secondary mass to sample [Msun] DEFAULT: 0.08 M1max : `float` maximum primary mass to sample [Msun] DEFAULT: 150.0 M2max : `float` maximum primary mass to sample [Msun] DEFAULT: 150.0 porb_lo : `float` minimum orbital period to sample [log10(days)] porb_hi : `float` maximum orbital period to sample [log10(days)] rand_seed : int random seed generator DEFAULT: 0 size : int, optional number of evolution times to sample NOTE: this is set in cosmic-pop call as Nstep Returns ------- primary_mass_list : array array of primary masses with size=size secondary_mass_list : array array of secondary masses with size=size porb_list : array array of orbital periods in days with size=size ecc_list : array array of eccentricities with size=size mass_singles : `float` Total mass in single stars needed to generate population mass_binaries : `float` Total mass in binaries needed to generate population n_singles : `int` Number of single stars needed to generate a population n_binaries : `int` Number of binaries needed to generate a population binfrac_list : array array of binary probabilities based on primary mass and period with size=size """ if pool is None: with MultiPool(processes=nproc) as pool: if mp_seeds is not None: if len(list(mp_seeds)) != nproc: raise ValueError( "Must supply a list of random seeds with length equal to number of processors" ) else: mp_seeds = [ nproc * (task._identity[0] - 1) for task in pool._pool ] inputs = [(M1min, M2min, M1max, M2max, porb_hi, porb_lo, size / nproc, rand_seed + mp_seed) for mp_seed in mp_seeds] worker = Worker() results = list(pool.map(worker, inputs)) else: if mp_seeds is not None: if len(list(mp_seeds)) != nproc: raise ValueError( "Must supply a list of random seeds with length equal to number of processors" ) else: if isinstance(pool, MPIPool): mp_seeds = [nproc * (task - 1) for task in pool.workers] elif isinstance(pool, MultiPool): mp_seeds = [ nproc * (task._identity[0] - 1) for task in pool._pool ] else: mp_seeds = [0 for i in range(nproc)] inputs = [(M1min, M2min, M1max, M2max, porb_hi, porb_lo, size / nproc, rand_seed + mp_seed) for mp_seed in mp_seeds] worker = Worker() results = list(pool.map(worker, inputs)) dat_lists = [[], [], [], [], [], [], [], [], []] for output_list in results: ii = 0 for dat_list in output_list: dat_lists[ii].append(dat_list) ii += 1 primary_mass_list = np.hstack(dat_lists[0]) secondary_mass_list = np.hstack(dat_lists[1]) porb_list = np.hstack(dat_lists[2]) ecc_list = np.hstack(dat_lists[3]) mass_singles = np.sum(dat_lists[4]) mass_binaries = np.sum(dat_lists[5]) n_singles = np.sum(dat_lists[6]) n_binaries = np.sum(dat_lists[7]) binfrac_list = np.hstack(dat_lists[8]) return (primary_mass_list, secondary_mass_list, porb_list, ecc_list, mass_singles, mass_binaries, n_singles, n_binaries, binfrac_list)
def main(): # how many files should we split the test output into nfiles = 10 wfd_thresh = 1000000 # get the options for the code and set the data_release globally (ugly) to # allow MultiPool to work kwargs = plasticc.get_data.parse_getdata_options() global data_release data_release = kwargs.pop('data_release') getter = plasticc.get_data.GetData(data_release) # setup paths for output base_dir = os.path.join(WORK_DIR, 'csv_dump') dump_dir = os.path.join(base_dir, data_release) if not os.path.exists(dump_dir): os.makedirs(dump_dir) # we can use model as a dummy string to indicate if we are generating # training or test data dummy = kwargs.pop('model') offset = kwargs.pop('offset') limit = kwargs.pop('limit') # setup root filenames for output - these get changed by fixpath if dummy == 'training': outfile = os.path.join(dump_dir, 'plasticc_training_set.csv') offset = None else: if limit is None: outfile = os.path.join(dump_dir, 'plasticc_test_set.csv') else: if offset is None: offset = 0 outfile = os.path.join(dump_dir, 'plasticc_test_n{}_set.csv'.format(offset)) # if we're limiting the output, then just dump one file nfiles = 1 # header file is named something sensible and is public header_file = outfile.replace('.csv', '_metadata.csv') header_file = fixpath(header_file) # make sure we remove any lingering files if os.path.exists(outfile): os.remove(outfile) _ = kwargs.get('field') # set the header keywords for training and testing # same except for sntype will be removed from test and hostgal_photoz isn't # provided kwargs['columns']=['objid','ptrobs_min','ptrobs_max','ra','decl', 'mwebv',\ 'hostgal_specz', 'hostgal_photoz', 'hostgal_photoz_err','sntype'] # set an extrasql query to get just the DDF and WFD objects # sntype for testing = true sntype + 100 if dummy == 'training': extrasql = "AND sntype < 100 AND ((objid LIKE 'WFD%') OR (objid LIKE 'DDF%'))" else: extrasql = "AND sntype > 100 AND ((objid LIKE 'WFD%') OR (objid LIKE 'DDF%'))" # set up options for data retrieval ignoring many of the command-line # options - impose cuts later kwargs['extrasql'] = extrasql kwargs['model'] = '%' kwargs['field'] = '%' kwargs['sort'] = True kwargs['shuffle'] = False kwargs['limit'] = None kwargs['get_num_lightcurves'] = True total = getter.get_lcs_headers(**kwargs) total = list(total)[0] kwargs['limit'] = total kwargs['get_num_lightcurves'] = False kwargs['offset'] = offset out = getter.get_lcs_headers(**kwargs) # current as of 20180827 aggregate_types = { 11: 11, 2: 2, 3: 3, 12: 2, 13: 3, 14: 2, 41: 41, 43: 43, 51: 51, 60: 60, 61: 99, 62: 99, 63: 99, 64: 64, 70: 70, 80: 80, 81: 81, 83: 83, 84: 84, 90: 99, 91: 91, 92: 99, 93: 91 } aggregate_names = { 11: 'SNIa-normal', 2: 'SNCC-II', 3: 'SNCC-Ibc', 12: 'SNCC-II', 13: 'SNCC-Ibc', 14: 'SNCC-II', 41: 'SNIa-91bg', 43: 'SNIa-x', 51: 'KN', 60: 'SLSN-I', 61: 'PISN', 62: 'ILOT', 63: 'CART', 64: 'TDE', 70: 'AGN', 80: 'RRlyrae', 81: 'Mdwarf', 83: 'EBE', 84: 'MIRA', 90: 'uLens-Binary', 91: 'uLens-Point', 92: 'uLens-STRING', 93: 'uLens-Point' } if dummy == 'training': pass else: aggregate_types = {x + 100: y for x, y in aggregate_types.items()} print('Aggregating as ', aggregate_types) # make a big list of the header - NOTE THAT WE ALWAYS RETRIEVE ALL OBJECTS out = list(out) if dummy == 'training': # we don't need to shuffle the training set pass else: # if we're generating test data, if we set a limit, just draw a random # sample else shuffle the full list if limit is not None: out = random.sample(out, limit) else: random.shuffle(out) # convert the selected header entries to a table out = at.Table(rows=out, names=kwargs['columns']) # we're not necessariy keeping all the models we simulated - remove any models that are not in keep_types keep_types = aggregate_types.keys() mask = np.array( [True if x in keep_types else False for x in out['sntype']]) out = out[mask] # aggregate types - map to new class numbers (i.e. MODELNUM_PLASTICC) if dummy == 'training': new_type = np.array( [aggregate_types.get(x, None) for x in out['sntype']]) else: new_type = np.array( [aggregate_types.get(x, None) for x in out['sntype']]) out['sntype'] = new_type # type 99 is not included in training train_types = set(aggregate_types.values()) - set([ 99, ]) train_types = list(train_types) # make sure that there are no "other" classes included in the training data if dummy == 'training': mask = np.array( [True if x in train_types else False for x in out['sntype']]) out = out[mask] # randomize the output type ID - keep rare as 99 target_map_file = outfile.replace('.csv', '_targetmap.txt').replace( '_test_set', '').replace('_training_set', '').replace(dump_dir, base_dir) try: target_map_data = at.Table.read(target_map_file, format='ascii') train_types = target_map_data['train_types'] target_types = target_map_data['target_types'] print(f'Restoring Target Map from {target_map_file}') target_map_file = target_map_file.replace(base_dir, dump_dir) target_map_file = fixpath(target_map_file, public=False) target_map_data.write(target_map_file, format='ascii.fixed_width', delimiter=' ', overwrite=True) print(f'Wrote distribution target mapping to file {target_map_file}') except Exception as e: target_types = np.random.choice(99, len(train_types), replace=False).tolist() target_map_data = at.Table([train_types, target_types], names=['train_types', 'target_types']) target_map_data.write(target_map_file, format='ascii.fixed_width', delimiter=' ', overwrite=True) print(f'Wrote target mapping to file {target_map_file}') target_map_file = target_map_file.replace(base_dir, dump_dir) target_map_file = fixpath(target_map_file, public=False) target_map_data.write(target_map_file, format='ascii.fixed_width', delimiter=' ', overwrite=True) print(f'Wrote distribution target mapping to file {target_map_file}') # map the aggregated IDs to random target IDs target_map = dict(zip(train_types, target_types)) target = np.array([target_map.get(x, 99) for x in out['sntype']]) out['target'] = target print('Mapping as {}'.format(target_map)) # orig map file is like target_map (and also private) but includes the rares orig_map_file = outfile.replace('.csv', '_origmap.txt').replace( '_test_set', '').replace('_training_set', '') orig_map_file = fixpath(orig_map_file, public=False) if not os.path.exists(orig_map_file): orig = [] aggregated = [] mapped = [] names = [] for key, val in aggregate_types.items(): name = aggregate_names.get(key, 'Rare') names.append(name) orig.append(key) aggregated.append(val) mapping = target_map.get(val, 99) mapped.append(mapping) orig_map_data = at.Table([orig, aggregated, mapped, names],\ names=['ORIG_NUM', 'MODEL_NUM', 'TARGET', 'MODEL_NAME']) orig_map_data.write(orig_map_file, format='ascii.fixed_width', delimiter=' ', overwrite=True) print(f'Wrote original mapping to file {orig_map_file}') # galactic objects have -9 as redshift - change to NaN # the numpy.isclose should have worked last time.... check this by hand. ind = out['hostgal_photoz'] == -9. out['hostgal_photoz'][ind] = np.nan out['hostgal_photoz_err'][ind] = np.nan ind = out['hostgal_specz'] == -9. out['hostgal_specz'][ind] = np.nan # add galactic coordinates c = SkyCoord(out['ra'], out['decl'], "icrs", unit='deg') gal = c.galactic out['gall'] = gal.l.value out['galb'] = gal.b.value # add distance modulus cosmo = FlatLambdaCDM(70, 0.3) out['distmod'] = cosmo.distmod(out['hostgal_photoz']).value ind = np.isfinite(out['distmod']) out['distmod'][~ind] = np.nan # figure out what fits files the data are in fits_files = [ "LSST_{0}_MODEL{1}/LSST_{0}_{2}_PHOT.FITS".format(*x.split('_')) for x in out['objid'] ] fits_files = np.array(fits_files) # the object names have the model name in them, so we need to edit them # new name = <SNID> orig_name = out['objid'] new_name = np.array([x.split('_')[-1] for x in orig_name], dtype=np.int32) ddf_field = np.zeros(len(new_name), dtype=np.uint8) ind = new_name < wfd_thresh ddf_field[ind] = 1 # preseve the mapping between old name, new name and file name out['object_id'] = new_name out['filename'] = fits_files out['ddf_bool'] = ddf_field # sort things by object id - Rick has already randomized these, so we preserve his order. out.sort('object_id') del new_name del fits_files del ddf_field del target del new_type # if we are generating test data, save a truth table if dummy == 'training': pass else: out_name = out['object_id'] target = out['target'] # remove the model type from the output header that goes with the test data out.remove_column('target') # make sure the truth table actually matches the job presently executing truth_file = outfile.replace('_set.csv', '_truthtable.csv') truth_file = fixpath(truth_file, public=False) if os.path.exists(truth_file): os.remove(truth_file) # write the truth table truth_table = at.Table([out_name, target], names=['object_id', 'target']) truth_table.write(truth_file) print(f'Wrote {truth_file}') nmc = len(out) out_ind = np.arange(nmc) nthreads = max(multiprocessing.cpu_count() // 2 - 1, 1) print(f'Using {nthreads} threads.') if dummy == 'training': batch_inds = np.array_split(out_ind, nthreads) else: # if this is test data, we want to break files up so that DDF and WFD # are in separate files and the number of files is split so we don't # have a giant CSV file batch_inds = [] ind = np.where(out['object_id'] < wfd_thresh)[0] print('DDF objects {}'.format(len(ind))) batch_inds.append(ind) ind = np.where(out['object_id'] >= wfd_thresh)[0] print('WFD objects {}'.format(len(ind))) batch_inds += np.array_split(ind, nfiles) # make batches to load the data batches = [] for ind in batch_inds: # we need the fits file for each object + the object pointers this_batch_lcs = at.Table( [ out['object_id'][ind], out['ptrobs_min'][ind], out['ptrobs_max'][ind], out['filename'][ind] ], names=['object_id', 'ptrobs_min', 'ptrobs_max', 'filename']) batches.append(this_batch_lcs) gc.collect() # create a map from batch number to first objid in each batch # batch number is helpful to name files by batch # this is sequential, but you might imagine more complicated schemes batch_ids = np.arange(len(batches)) + 1 batch_keys = [x['object_id'][0] for x in batches] batch_map = dict(zip(batch_keys, batch_ids)) # do the output if dummy == 'training': # training is simple - dump each batch into one file in sequence outfile = fixpath(outfile, gzip=True) with MultiPool(processes=nthreads) as pool: with tqdm(total=nmc) as pbar: outlines = 'object_id,mjd,passband,flux,flux_err,detected_bool\n' # change to pool.imap so order is preserved in output file # combine all the batches for result in pool.imap(task, batches): _, nbatch, batchlines = result pbar.update(nbatch) outlines += '\n'.join(batchlines) outlines += '\n' outbytes = outlines.encode() gc.collect() # do the output with gzip.open(outfile, 'wb', compresslevel=9) as f: f.write(outbytes) else: # these variables will set the accessed time and modified time to the same numbers for all batches st_atime = None st_mtime = None with tqdm(total=nmc) as pbar: for batch in batches: # for test, the batches each get a separate file batch_key = batch[0]['object_id'] batch_id = batch_map[batch_key] batchfile = outfile.replace('.csv', f'_batch{batch_id}.csv') batchfile = fixpath(batchfile, gzip=True) # for actual mutliprocessing, split up each file's indices into mini batches ind = np.arange(len(batch)) mini_inds = np.array_split(ind, nthreads - 1) mini_batches = [batch[x] for x in mini_inds] nbatch = 0 # combine all the batches outlines = 'object_id,mjd,passband,flux,flux_err,detected_bool\n' with MultiPool(processes=nthreads - 1) as pool: for result in pool.imap(task, mini_batches): _, mini_nbatch, mini_batchlines = result nbatch += mini_nbatch outlines += '\n'.join(mini_batchlines) outlines += '\n' pbar.update(mini_nbatch) outbytes = outlines.encode() gc.collect() # do the output with gzip.open(batchfile, 'wb', compresslevel=9) as f: f.write(outbytes) # get the timestamps of the first batch file if st_atime is None: st = os.stat(batchfile) st_atime = st.st_atime st_mtime = st.st_mtime # change the timestamp of the output os.utime(batchfile, (st_atime, st_mtime)) # remove and rename some columns from the metadata output # this isn't strictly necessary, since we choose exactly what columns to output # but this makes sure astropy also strips any metadata about the columns itself out.remove_columns( ['objid', 'ptrobs_min', 'ptrobs_max', 'filename', 'sntype']) # setup what columns get output into the headers cols = [ 'object_id', 'ra', 'decl', 'gall', 'galb', 'ddf_bool', 'hostgal_specz', 'hostgal_photoz', 'hostgal_photoz_err', 'distmod', 'mwebv' ] if dummy == 'training': cols.append('target') out = out[cols] # fix column precision precision = {'ra':6, 'decl':6, 'gall':6, 'galb':6,\ 'hostgal_specz':4, 'hostgal_photoz':4, 'hostgal_photoz_err':4, 'distmod':4, 'mwebv':3} for col, val in precision.items(): formatstr = f'%.{val}f' out[col].format = formatstr # write out the header out.write(header_file, format='ascii.csv', overwrite=True)
def main(): # setup paths for output dump_dir = os.path.join(WORK_DIR, 'hdf5_dump') if not os.path.exists(dump_dir): os.makedirs(dump_dir) # get the options for the code and set the data_release globally (ugly) to # allow MultiPool to work kwargs = plasticc.get_data.parse_getdata_options() global data_release data_release = kwargs.pop('data_release') getter = plasticc.get_data.GetData(data_release) # we can use model as a dummy string to indicate if we are generating # training or test data dummy = kwargs.pop('model') offset = kwargs.pop('offset') limit = kwargs.pop('limit') if dummy == 'training': outfile = os.path.join(dump_dir, 'training_set.hdf5') offset = None else: if limit is None: outfile = os.path.join(dump_dir, 'test_set.hdf5') else: if offset is None: offset = 0 outfile = os.path.join(dump_dir, 'test_n{}_set.hdf5'.format(offset)) # make sure we remove any lingering files if os.path.exists(outfile): os.remove(outfile) _ = kwargs.get('field') # set the header keywords for training and testing # same except for sntype will be removed from test and hostgal_photoz isn't # provided if dummy == 'training': kwargs['columns']=['objid','ptrobs_min','ptrobs_max','ra','decl', 'mwebv', 'mwebv_err',\ 'hostgal_photoz', 'hostgal_photoz', 'hostgal_photoz_err', 'sntype'] else: kwargs['columns']=['objid','ptrobs_min','ptrobs_max','ra','decl', 'mwebv', 'mwebv_err',\ 'hostgal_photoz', 'hostgal_photoz_err', 'sntype'] # set an extrasql query to get just the DDF and WFD objects # sntype for testing = true sntype + 100 if dummy == 'training': extrasql = "AND sntype < 100 AND ((objid LIKE 'WFD%') OR (objid LIKE 'DDF%'))" else: extrasql = "AND sntype > 100 AND ((objid LIKE 'WFD%') OR (objid LIKE 'DDF%'))" # set up options for data retrieval ignoring many of the command-line # options kwargs['extrasql'] = extrasql kwargs['model'] = '%' kwargs['field'] = '%' kwargs['sort'] = True kwargs['shuffle'] = False kwargs['limit'] = None kwargs['get_num_lightcurves'] = True total = getter.get_lcs_headers(**kwargs) total = list(total)[0] kwargs['limit'] = total kwargs['get_num_lightcurves'] = False kwargs['offset'] = offset out = getter.get_lcs_headers(**kwargs) aggregate_types = {1:1, 2:2, 3:3, 12:2, 13:3, 14:2, 41:41, 43:43, 45:45, 51:51, 60:60, 61:61, 62:62, 63:63, 64:64, 80:80, 81:81, 82:82, 83:83, 84:84, 90:90, 91:91} if dummy == 'training': pass else: aggregate_types = {x+100:y for x,y in aggregate_types.items()} print('Aggregating as ', aggregate_types) # make a big list of the header - NOTE THAT WE ALWAYS RETRIEVE ALL OBJECTS out = list(out) if dummy == 'training': # we don't need to shuffle the training set pass else: # if we're generating test data, if we set a limit, just draw a random # sample else shuffle the full list if limit is not None: out = random.sample(out, limit) else: random.shuffle(out) # convert the selected header entries to a table out = at.Table(rows=out, names=kwargs['columns']) # we're not necessariy keeping all the models we simulated (42 and 50 are going bye bye) keep_types = aggregate_types.keys() mask = np.array([True if x in keep_types else False for x in out['sntype']]) out = out[mask] # aggregate types if dummy=='training': new_type = np.array([aggregate_types.get(x, None) for x in out['sntype']]) else: new_type = np.array([aggregate_types.get(x, None) for x in out['sntype']]) out['sntype'] = new_type # make sure that there are no "other" classes included in the training data if dummy == 'training': # not train types - 45, 60, 61, 62, 63, 64, 90, 91 train_types = (1, 2, 3, 41, 43, 51, 80, 81, 82, 83, 84) mask = np.array([True if x in train_types else False for x in out['sntype']]) out = out[mask] # galactic objects have -9 as redshift - change to 0 dummy_val = np.repeat(-9, len(out)) ind = np.isclose(out['hostgal_photoz'], dummy_val) out['hostgal_photoz'][ind] = 0. out['hostgal_photoz_err'][ind] = 0. # figure out what fits files the data are in fits_files =[ "LSST_{0}_MODEL{1}/LSST_{0}_{2}_PHOT.FITS".format(*x.split('_')) for x in out['objid']] fits_files = np.array(fits_files) uniq_files = np.unique(fits_files) # the object names have the model name in them, so we need to edit them # new name = <FIELD><SNID> orig_name = out['objid'] new_name = [ '{}{}'.format(x.split('_')[0], x.split('_')[-1]) for x in orig_name] new_name = np.array(new_name) out_name = new_name nmc = len(out) # if we are generating test data, save a truth table if dummy == 'training': pass else: sntype = out['sntype'] # remove the model type from the output header that goes with the test data out.remove_column('sntype') truth_file = outfile.replace('_set.hdf5', '_truthtable.hdf5') if os.path.exists(truth_file): os.remove(truth_file) # ... saving it in the truth table only orig_name = orig_name.astype(bytes) new_name = new_name.astype(bytes) truth_table = at.Table([orig_name, new_name, sntype], names=['objid','shortid','sntype']) truth_table.write(truth_file, compression=True, path='truth_table', serialize_meta=False, append=True) # make batches to load the data batches = {} for filename in uniq_files: ind = (fits_files == filename) this_fits_lcs = at.Table([out['objid'][ind], out['ptrobs_min'][ind], out['ptrobs_max'][ind]], names=['objid','ptrobs_min', 'ptrobs_max']) batches[filename] = this_fits_lcs name_lookup = dict(zip(orig_name, out_name)) gc.collect() # do the output failed = [] with MultiPool() as pool: with tqdm(total=nmc) as pbar: for result in pool.imap_unordered(task, batches.items()): this_file_n = len(result.items()) with h5py.File(outfile, 'a') as outf: for true_obj, thislc in result.items(): short_obj = name_lookup[true_obj] retries = 10 notwritten = True overwrite = False while notwritten and retries > 0: try: outf.create_dataset(short_obj, data=thislc, compression='lzf') #thislc.write(outfile, path=short_obj, compression=True, serialize_meta=False, append=True,\ # overwrite=overwrite) notwritten = False except Exception as e: timer.sleep(0.010) overwrite = True retries -= 1 print('{} {}'.format(true_obj, e)) if notwritten: failed.append((true_obj, short_obj)) print("Failed", true_obj) outf.flush() pbar.update(this_file_n) gc.collect() print(failed) # write out the header out['objid'] = out_name.astype(bytes) out.remove_columns(['ptrobs_min', 'ptrobs_max']) out.write(outfile, compression=True, path='header', serialize_meta=False, append=True)
bounds = scipy.optimize.Bounds([ np.arctan(-10), -2, np.arctan(-10), -2, np.arctan(-10), -np.inf, 0, -np.inf, -10. ], [ np.arctan(10), np.inf, np.arctan(10), np.inf, np.arctan(10), np.inf, np.inf, np.inf, leopy.stats.logit(0.30) ]) try: from schwimmbad import MultiPool pool = MultiPool() print('Parallel execution on ' + str(pool.size) + ' processes') except ImportError as error: print('Serial execution as module `schwimmbad` was not found') pool = None my_print_fun = auxiliary.MyPrintFun() my_take_step = auxiliary.MyTakeStep(stepsize=1) minimizer_options = {'disp': True, 'ftol': 1e-8} minimizer_kwargs = { 'method': method, 'bounds': bounds, 'options': minimizer_options, 'args': (pool, ) }
values = [] for x in data: values.append(do_the_processing(x)) #using the map() function which applies the function (passed as the first argument) to each element in the iterable (second argument) #Note: python3 map() returns a generator object, so we call list on this to get out the values values = list(map(do_the_processing, data)) #class-based wrapper for the built-in (serial) map() function from schwimmbad import SerialPool pool = SearlPool() values = list(pool.map(do_the_processing, data)) #utilize multiple cores on the same processor from schwimmbad import MultiPool with MultiPool() as pool: values = list(pool.map(do_the_processing, data)) #when using MPI pool tell all worker processes to wait for tasks from the master process def main(pool, data): values = pool.map(do_the_processing, data) from schwimmbad import MPIPool with MPIPool() as pool: if not pool.is_master(): pool.wait() sys.exit(0) #selecting a pool with command-line arguments