def test_many_fluxes(self, mock_data):
dv = calculate_dv(mock_data[0])
new_wave = create_log_lam_grid(dv, mock_data[0].min(),
mock_data[0].max())["wl"]
flux_stack = np.tile(mock_data[1], (4, 1))
fluxes = resample(mock_data[0], flux_stack, new_wave)
assert fluxes.shape == (4, len(new_wave))
def __init__(
self,
emulator: Union[str, Emulator],
data: Union[str, Spectrum],
grid_params: Sequence[float],
max_deque_len: int = 100,
name: str = "SpectrumModel",
**params,
):
if isinstance(emulator, str):
emulator = Emulator.load(emulator)
if isinstance(data, str):
data = Spectrum.load(data)
if len(data) > 1:
raise ValueError(
"Multiple orders detected in data, please use EchelleModel")
self.emulator: Emulator = emulator
self.data_name = data.name
self.data = data[0]
dv = calculate_dv(self.data.wave)
self.min_dv_wave = create_log_lam_grid(dv, self.emulator.wl.min(),
self.emulator.wl.max())["wl"]
self.bulk_fluxes = resample(self.emulator.wl,
self.emulator.bulk_fluxes,
self.min_dv_wave)
self.residuals = deque(maxlen=max_deque_len)
# manually handle cheb coeffs to offset index by 1
chebs = params.pop("cheb", [])
cheb_idxs = [str(i) for i in range(1, len(chebs) + 1)]
params["cheb"] = dict(zip(cheb_idxs, chebs))
# load rest of params into FlatterDict
self.params = FlatterDict(params)
self.frozen = []
self.name = name
# Unpack the grid parameters
self.n_grid_params = len(grid_params)
self.grid_params = grid_params
# None means "yet to be calculated", do not use NaN
self._lnprob = None
self._glob_cov = None
self._loc_cov = None
self.log = logging.getLogger(self.__class__.__name__)
self.flux_scalar_func = flux_scalar = LinearNDInterpolator(
self.emulator.grid_points, self.emulator.flux_scalar)
def test_resample(self, mock_data):
dv = calculate_dv(mock_data[0])
new_wave = create_log_lam_grid(dv, mock_data[0].min(),
mock_data[0].max())["wl"]
flux = resample(*mock_data, new_wave)
assert flux.shape == new_wave.shape
def test_bad_waves(self, mock_data, wave):
with pytest.raises(ValueError):
resample(*mock_data, wave)
def resample_final(w, f):
return (self.wl_final, resample(w, f, self.wl_final))
def resample_loglam(w, f):
return (wl_loglam, resample(w, f, wl_loglam)) # - fl
def mock_data(mock_hdf5_interface):
wave = mock_hdf5_interface.wl
new_wave = create_log_lam_grid(1e3, wave.min(), wave.max())["wl"]
flux = resample(wave, next(mock_hdf5_interface.fluxes), new_wave)
yield new_wave, flux
def __call__(self):
"""
Performs the transformations according to the parameters available in
``self.params``
Returns
-------
flux, cov : tuple
The transformed flux and covariance matrix from the model
"""
wave = self.min_dv_wave
fluxes = self.bulk_fluxes
if "vsini" in self.params:
fluxes = rotational_broaden(wave, fluxes, self.params["vsini"])
if "vz" in self.params:
wave = doppler_shift(wave, self.params["vz"])
fluxes = resample(wave, fluxes, self.data.wave)
if "Av" in self.params:
fluxes = extinct(self.data.wave, fluxes, self.params["Av"])
if "cheb" in self.params:
# force constant term to be 1 to avoid degeneracy with log_scale
coeffs = [1, *self.cheb]
fluxes = chebyshev_correct(self.data.wave, fluxes, coeffs)
# Scale factor from emulator normalization
flux_scalar = self.flux_scalar_func(self.grid_params)[0]
# Only rescale flux_mean and flux_std
if "log_scale" in self.params:
scale = np.exp(self.params["log_scale"]) * flux_scalar
fluxes[-2:] = rescale(fluxes[-2:], scale)
weights, weights_cov = self.emulator(self.grid_params)
L, flag = cho_factor(weights_cov, overwrite_a=True)
# Decompose the bulk_fluxes (see emulator/emulator.py for the ordering)
*eigenspectra, flux_mean, flux_std = fluxes
# Complete the reconstruction
X = eigenspectra * flux_std
flux = weights @ X + flux_mean
# Renorm to data flux if no "log_scale" provided
if "log_scale" not in self.params:
factor = _get_renorm_factor(self.data.wave, flux * flux_scalar,
self.data.flux) * flux_scalar
flux = rescale(flux, factor)
X = rescale(X, factor)
cov = X.T @ cho_solve((L, flag), X)
# Poisson Noise Scaling
if "global_cov:sigma_amp" in self.params:
poisson_scale = self.params["global_cov:sigma_amp"]
else:
poisson_scale = 1
# Trivial covariance
np.fill_diagonal(cov,
cov.diagonal() + (poisson_scale * self.data.sigma**2))
# Trival and global covariance
if "global_cov" in self.params:
if "global_cov" not in self.frozen or self._glob_cov is None:
if "global_cov:log_amp" in self.params.keys():
ag = np.exp(self.params["global_cov:log_amp"])
else:
ag = self.params["global_cov:amp"]
if "global_cov:log_ls" in self.params.keys():
lg = np.exp(self.params["global_cov:log_ls"])
else:
lg = self.params["global_cov:ls"]
T = self.params["T"]
self._glob_cov = global_covariance_matrix(
self.data.wave, T, ag, lg)
if self._glob_cov is not None:
cov += self._glob_cov
# Local covariance
if "local_cov" in self.params:
if "local_cov" not in self.frozen or self._loc_cov is None:
self._loc_cov = 0
for kernel in self.params.as_dict()["local_cov"]:
mu = kernel["mu"]
amplitude = np.exp(kernel["log_amp"])
sigma = np.exp(kernel["log_sigma"])
self._loc_cov += local_covariance_matrix(
self.data.wave, amplitude, mu, sigma)
if self._loc_cov is not None:
cov += self._loc_cov
return flux, cov