def test_convolution_can_accept_worker_pool(num_proc): n = 20 x = np.linspace(2.0, 2.3, n) y = np.random.randn(n) with Parallel(num_proc) as parallel: convolution(x, y, vsini=1, R=100_000, band="K", num_procs=parallel, verbose=False)
def convolve_and_resample( wav: ndarray, flux: ndarray, vsini: float, R: float, band: str, sampling: float, **conv_kwargs, ) -> Tuple[ndarray, ndarray]: """Convolve and resample functions together. Returns ------- wav_grid: ndarray Resampled wavelength array sampled_flux: ndarray Convolved and resampled flux array """ wav_band, __, convolved_flux = convolution(wav, flux, vsini, R, band, **conv_kwargs) # Re-sample to sampling per resolution element. wav_grid = log_resample(wav_band, sampling, R) sampled_flux = np.interp(wav_grid, wav_band, convolved_flux) return wav_grid, sampled_flux
def augment_spectra_parallel(spectra, wav, new_wav, vrot_list, noise_list, vrad_list, instrument_res, trailing_zeros_l=0, trailing_zeros_r=0): """ Augments (in parallel) a list of spectra with rotational velocity, radial velocity, noise, and resolution degradation. :param spectra: list of spectra :param wav: list of synthetic wavelength grids :param new_wav: wavelength grid to rebin the synthetic wavelength grid to :param vrot_list: a list, same length as spectra list, of rotational velocities (km/s) to apply :param noise_list: a list, same length as spectra list, of maximum noise (fraction of flux) to apply :param vrad_list: a list, same length as spectra list, of radial velocities (km/s) to apply :param instrument_res: instrumental resolution to degrade the synthetic spectra to :param trailing_zeros_l: maximum # of trailing zeros to add on left end of spectrum :param trailing_zeros_r: maximum # of trailing zeros to add on right end of spectrum :return: a list of modified input spectra """ @contextmanager def poolcontext(*args, **kwargs): pool = multiprocessing.Pool(*args, **kwargs) yield pool pool.terminate() num_spectra = np.shape(spectra)[0] num_cpu = multiprocessing.cpu_count() pool_size = num_cpu if num_spectra >= num_cpu else num_spectra print('[INFO] Pool size: {}'.format(pool_size)) # Degrade resolution and apply rotational broadening for i in range(len(spectra)): epsilon = random.uniform(0, 1.) _, _, flux = convolution(wav=wav[i], flux=spectra[i], vsini=vrot_list[i], R=instrument_res, epsilon=epsilon, normalize=True, num_procs=10) spectra[i] = flux pool_arg_list = [(spectra[i], wav[i], new_wav, vrot_list[i], noise_list[i], vrad_list[i], instrument_res, trailing_zeros_l, trailing_zeros_r) for i in range(num_spectra)] with poolcontext(processes=pool_size) as pool: results = pool.starmap(augment_spectrum, pool_arg_list) return results
def convolve_spectra( spectrum, band, vsini, R, epsilon: float = 0.6, fwhm_lim: float = 5.0, num_procs: Optional[int] = None, results_dir: str = results_dir, normalize: bool = True, output_name: Optional[str] = None, ) -> int: """Load Spectrum, apply convolution and then save results.""" print("Reading the data...") wav, flux = read_spectrum(spectrum) # In microns and photon flux. print("Done.") wav_band, flux_band, convolved_flux = convolution( wav, flux, vsini, R, band, epsilon=epsilon, fwhm_lim=fwhm_lim, num_procs=num_procs, normalize=normalize, ) if not normalize: norm_ = "_unnormalized" else: norm_ = "" if output_name is None: name_model = name_assignment(spectrum) filename = "{0}Spectrum_{1}_{2}band_vsini{3:3.1f}_R{4:d}k{5}.dat".format( results_dir, name_model, band, vsini, R / 1000, norm_) else: filename = os.path.join(results_dir, output_name) save_convolution_results(filename, wav_band, flux_band, convolved_flux) return 0
def test_convolution_can_accept_int(num_proc): n = 5 x = np.linspace(2.0, 2.3, n) y = np.random.randn(n) convolution(x, y, vsini=1, R=1000, band="K", num_procs=num_proc)
def test_convolution_can_accept_None(): n = 20 x = np.linspace(2.0, 2.3, n) y = np.random.randn(n) convolution(x, y, vsini=1, R=100_000, band="K", num_procs=None)