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)