def test_renormalization_on(host, method, model_shape):
models = setup_renomlization_model(host, model_shape)
result = renormalization(host, models, normalize=True, method=method)
assert result.ndim == models.ndim
assert result.ndim == len(model_shape)
assert result.shape == models.shape
def test_renormalization_off(host, method, model_shape):
expected_shape = tuple(1 for _ in model_shape)
expected_shape = (len(host.flux), *expected_shape[1:])
models = setup_renomlization_model(host, model_shape)
with pytest.warns(UserWarning) as record:
result = renormalization(host, models, normalize=False, method=method)
assert len(record) == 1
assert record[0].message.args[0] == "Not Scalar Re-normalizing to observations!"
assert result.ndim == len(model_shape)
assert result.shape == expected_shape
def test_renormalization_with_invalid_method(host, method):
with pytest.raises(ValueError):
renormalization(host.flux, host.flux, method=method, normalize=True)
def tcm_wrapper(num,
params1,
model2_pars,
alphas,
rvs,
gammas,
obs_spec,
errors=None,
norm=True,
verbose=False,
save_only=True,
chip=None,
prefix=None,
wav_scale=True,
norm_method="scalar"):
"""Wrapper for iteration loop of tcm. params1 fixed, model2_pars are many."""
normalization_limits = [2105, 2185] # small as possible?
if prefix is None:
sf = os.path.join(
simulators.paths["output_dir"], obs_spec.header["OBJECT"].upper(),
"tc_{0}_{1}-{2}_part{6}_host_pars_[{3}_{4}_{5}].csv".format(
obs_spec.header["OBJECT"].upper(),
int(obs_spec.header["MJD-OBS"]), chip, params1[0], params1[1],
params1[2], num))
else:
sf = "{0}_part{4}_host_pars_[{1}_{2}_{3}].csv".format(
prefix, params1[0], params1[1], params1[2], num)
save_filename = sf
if os.path.exists(save_filename) and save_only:
print("''{}' exists, so not repeating calculation.".format(
save_filename))
return None
else:
if not save_only:
broadcast_chisqr_vals = np.empty(len(model2_pars))
for jj, params2 in enumerate(model2_pars):
if verbose:
print("Starting iteration with parameters:\n {0}={1},{2}={3}".
format(num, params1, jj, params2))
mod1_spec = load_starfish_spectrum(params1,
limits=normalization_limits,
hdr=True,
normalize=True,
wav_scale=wav_scale)
mod2_spec = load_starfish_spectrum(params2,
limits=normalization_limits,
hdr=True,
normalize=True,
wav_scale=wav_scale)
# Wavelength selection
rv_limits = observation_rv_limits(obs_spec, rvs, gammas)
mod1_spec.wav_select(*rv_limits)
mod2_spec.wav_select(*rv_limits)
obs_spec = obs_spec.remove_nans()
# One component model with broadcasting over gammas
# two_comp_model(wav, model1, model2, alphas, rvs, gammas)
assert np.allclose(mod1_spec.xaxis, mod2_spec.xaxis)
broadcast_result = two_comp_model(mod1_spec.xaxis,
mod1_spec.flux,
mod2_spec.flux,
alphas=alphas,
rvs=rvs,
gammas=gammas)
broadcast_values = broadcast_result(obs_spec.xaxis)
assert ~np.any(np.isnan(obs_spec.flux)), "Observation is nan"
# RE-NORMALIZATION
if chip == 4:
# Quadratically renormalize anyway
obs_spec = renormalization(obs_spec,
broadcast_values,
normalize=True,
method="quadratic")
obs_flux = renormalization(obs_spec,
broadcast_values,
normalize=norm,
method=norm_method)
# sp_chisquare is much faster but don't think I can add masking.
broadcast_chisquare = chi_squared(obs_flux,
broadcast_values,
error=errors)
# sp_chisquare = stats.chisquare(obs_flux, broadcast_values, axis=0).statistic
# broadcast_chisquare = sp_chisquare
if not save_only:
print(broadcast_chisquare.shape)
print(broadcast_chisquare.ravel()[np.argmin(
broadcast_chisquare)])
broadcast_chisqr_vals[jj] = broadcast_chisquare.ravel()[
np.argmin(broadcast_chisquare)]
npix = obs_flux.shape[0]
save_full_tcm_chisqr(save_filename,
params1,
params2,
alphas,
rvs,
gammas,
broadcast_chisquare,
npix,
verbose=verbose)
if save_only:
return None
else:
return broadcast_chisqr_vals
def bhm_analysis(obs_spec, model_pars, gammas=None, errors=None, prefix=None, verbose=False, chip=None, norm=False,
wav_scale=True, norm_method="scalar"):
"""Run one component model over all parameter combinations in model_pars."""
# Gammas
if gammas is None:
gammas = np.array([0])
elif isinstance(gammas, (float, int)):
gammas = np.asarray(gammas, dtype=np.float32)
if isinstance(model_pars, list):
logging.debug(__("Number of close model_pars returned {0}", len(model_pars)))
# Solution Grids to return
model_chisqr_vals = np.empty(len(model_pars))
model_xcorr_vals = np.empty(len(model_pars))
model_xcorr_rv_vals = np.empty(len(model_pars))
bhm_grid_chisqr_vals = np.empty(len(model_pars))
bhm_grid_gamma = np.empty(len(model_pars))
full_bhm_grid_chisquare = np.empty((len(model_pars), len(gammas)))
normalization_limits = [2105, 2185] # small as possible?
for ii, params in enumerate(tqdm(model_pars)):
if prefix is None:
save_name = os.path.join(
simulators.paths["output_dir"], obs_spec.header["OBJECT"].upper(), "bhm",
"bhm_{0}_{1}_{3}_part{2}.csv".format(
obs_spec.header["OBJECT"].upper(), obs_spec.header["MJD-OBS"], ii, chip))
else:
save_name = os.path.join("{0}_part{1}.csv".format(prefix, ii))
if verbose:
print("Starting iteration with parameter:s\n{}".format(params))
mod_spec = load_starfish_spectrum(params, limits=normalization_limits, hdr=True,
normalize=True, wav_scale=wav_scale)
# Wavelength selection
mod_spec.wav_select(np.min(obs_spec.xaxis) - 5,
np.max(obs_spec.xaxis) + 5) # +- 5nm of obs
obs_spec = obs_spec.remove_nans()
# One component model with broadcasting over gammas
bhm_grid_func = one_comp_model(mod_spec.xaxis, mod_spec.flux, gammas=gammas)
bhm_grid_values = bhm_grid_func(obs_spec.xaxis)
assert ~np.any(np.isnan(obs_spec.flux)), "Observation is nan"
# RENORMALIZATION
if chip == 4:
# Quadratically renormalize anyway
obs_spec = renormalization(obs_spec, bhm_grid_values, normalize=True, method="quadratic")
obs_flux = renormalization(obs_spec, bhm_grid_values, normalize=norm, method=norm_method)
# Simple chi2
bhm_grid_chisquare_old = chi_squared(obs_flux, bhm_grid_values, error=errors)
# Applying arbitrary scalar normalization to continuum
bhm_norm_grid_values, arb_norm = arbitrary_rescale(bhm_grid_values,
*simulators.sim_grid["arb_norm"])
# Calculate Chi-squared
obs_flux = np.expand_dims(obs_flux, -1) # expand on last axis to match rescale
bhm_norm_grid_chisquare = chi_squared(obs_flux, bhm_norm_grid_values, error=errors)
# Take minimum chi-squared value along Arbitrary normalization axis
bhm_grid_chisquare, arbitrary_norms = arbitrary_minimums(bhm_norm_grid_chisquare, arb_norm)
assert np.any(
bhm_grid_chisquare_old >= bhm_grid_chisquare), "All chi2 values are not better or same with arbitrary scaling"
# Interpolate to obs
mod_spec.spline_interpolate_to(obs_spec)
org_model_chi_val = chi_squared(obs_spec.flux, mod_spec.flux)
model_chisqr_vals[ii] = org_model_chi_val # This is gamma = 0 version
# New parameters to explore
bhm_grid_chisqr_vals[ii] = bhm_grid_chisquare[np.argmin(bhm_grid_chisquare)]
bhm_grid_gamma[ii] = gammas[np.argmin(bhm_grid_chisquare)]
full_bhm_grid_chisquare[ii, :] = bhm_grid_chisquare
################
# Find cross correlation RV
# Should run though all models and find best rv to apply uniformly
rvoffset, cc_max = xcorr_peak(obs_spec, mod_spec, plot=False)
if verbose:
print("Cross correlation RV = {}".format(rvoffset))
print("Cross correlation max = {}".format(cc_max))
model_xcorr_vals[ii] = cc_max
model_xcorr_rv_vals[ii] = rvoffset
###################
npix = obs_flux.shape[0]
# print("bhm shape", bhm_grid_chisquare.shape)
save_full_bhm_chisqr(save_name, params, gammas, bhm_grid_chisquare, arbitrary_norms,
npix, rvoffset)
return (model_chisqr_vals, model_xcorr_vals, model_xcorr_rv_vals,
bhm_grid_chisqr_vals, bhm_grid_gamma, full_bhm_grid_chisquare)