def smooth_vel(wave, spec, outwave, sigma, nsigma=10, inres=0, **extras):
"""Smooth a spectrum in velocity space. This is insanely slow, but general
and correct.
:param wave:
Wavelength vector of the input spectrum.
:param spec:
Flux vector of the input spectrum.
:param outwave:
Desired output wavelength vector.
:param sigma:
Desired velocity resolution (km/s), *not* FWHM.
:param nsigma:
Number of sigma away from the output wavelength to consider in the
integral. If less than zero, all wavelengths are used. Setting this
to some positive number decreses the scaling constant in the O(N_out *
N_in) algorithm used here.
:param inres:
The velocity resolution of the input spectrum (km/s), *not* FWHM.
"""
sigma_eff_sq = sigma**2 - inres**2
# if np.any(sigma_eff_sq) < 0.0:
# raise ValueError("Desired velocity resolution smaller than the value"
# "possible for this input spectrum.".format(inres))
# sigma_eff is in units of sigma_lambda / lambda
sigma_eff = np.sqrt(sigma_eff_sq) / ckms
lnwave = np.log(wave)
flux = np.zeros(len(outwave))
for i, w in enumerate(outwave):
x = (np.log(w) - lnwave) / sigma_eff
# if nsigma > 0:
good = np.abs(x) < nsigma
x = x[good]
_spec = spec[good]
# else:
# _spec = spec
f = np.exp(-0.5 * x**2)
flux[i] = np.trapz(f * _spec, x) / np.trapz(f, x)
return flux
def fitness(z):
# translate decision vector
states, controls, T = unflatten(z)
# time grid
n = states.shape[0]
times = np.linspace(0, T, n)
# objective
L = vmap(lambda state, control: self.lagrangian(
state, control, homotopy, *params))
L = L(states, controls)
J = np.trapz(L, dx=T / (n - 1))
# Lagrangian state dynamics constraints, and boundary constraints
# e0 = self.collocate_lagrangian(states, controls, times, costs, homotopy, *params)
e1 = self.collocate_state(states, controls, times, *params)
e2, e3 = boundaries(states[0, :], states[-1, :])
e = np.hstack((e1.flatten(), e2, e3))**2
# fitness vector
return np.hstack((J, e))
def trapz(y, x=None, dx=1.0, axis: int = -1): if isinstance(y, JaxArray): y = y.value if isinstance(x, JaxArray): x = x.value return jnp.trapz(y, x=x, dx=dx, axis=axis)
def posterior(
self,
inputs: pd.DataFrame,
column: str,
grid: np.ndarray,
batch_size: int = None,
) -> np.ndarray:
"""Calculates posterior distributions for the provided column.
Calculates the conditional posterior distribution, assuming the
data values in the other columns of the DataFrame.
Parameters
----------
inputs : pd.DataFrame
Data on which the posterior distributions are conditioned.
Must have columns matching self.data_columns, *except*
for the column specified for the posterior (see below).
column : str
Name of the column for which the posterior distribution
is calculated. Must be one of the columns in self.data_columns.
However, whether or not this column is one of the columns in
`inputs` is irrelevant.
grid : np.ndarray
Grid on which to calculate the posterior.
batch_size : int, default=None
Size of batches in which to calculate posteriors. If None, all
posteriors are calculated simultaneously. Simultaneous calculation
is faster, but memory intensive for large data sets.
Returns
-------
np.ndarray
Device array of shape (inputs.shape[0], grid.size).
"""
# get the index of the provided column, and remove it from the list
columns = list(self.data_columns)
idx = columns.index(column)
columns.remove(column)
nrows = inputs.shape[0]
batch_size = nrows if batch_size is None else batch_size
# convert data (sans the provided column)
# to an array with columns ordered
X = np.array(inputs[columns].values)
pdfs = np.zeros((nrows, len(grid)))
for batch_idx in range(0, nrows, batch_size):
batch = X[batch_idx:batch_idx + batch_size]
# make a new copy of each row for each value of the column
# for which we are calculating the posterior
batch = np.hstack((
np.repeat(
batch[:, :idx],
len(grid),
axis=0,
),
np.tile(grid, len(batch))[:, None],
np.repeat(
batch[:, idx:],
len(grid),
axis=0,
),
))
# calculate probability densities
log_prob = self._log_prob(batch).reshape((-1, len(grid)))
pdfs = ops.index_update(
pdfs,
ops.index[batch_idx:batch_idx + batch_size, :],
np.exp(log_prob),
indices_are_sorted=True,
unique_indices=True,
)
# reshape so that each row is a posterior
pdfs = pdfs.reshape((nrows, len(grid)))
# normalize so they integrate to one
pdfs = pdfs / np.trapz(y=pdfs, x=grid).reshape(-1, 1)
# set NaN's equal to zero probability
pdfs = np.nan_to_num(pdfs, nan=0.0)
return pdfs
def smooth_wave(wave,
spec,
outwave,
sigma,
nsigma=10,
inres=0,
in_vel=False,
**extras):
"""Smooth a spectrum in wavelength space. This is insanely slow, but
general and correct (except for the treatment of the input resolution if it
is velocity)
:param wave:
Wavelength vector of the input spectrum.
:param spec:
Flux vector of the input spectrum.
:param outwave:
Desired output wavelength vector.
:param sigma:
Desired resolution (*not* FWHM) in wavelength units. This can be a
vector of same length as ``wave``, in which case a wavelength dependent
broadening is calculated
:param nsigma: (optional, default=10)
Number of sigma away from the output wavelength to consider in the
integral. If less than zero, all wavelengths are used. Setting this
to some positive number decreses the scaling constant in the O(N_out *
N_in) algorithm used here.
:param inres: (optional, default: 0.0)
Resolution of the input, in either wavelength units or
lambda/dlambda (c/v). Ignored if <= 0.
:param in_vel: (optional, default: False)
If True, the input spectrum has been smoothed in velocity
space, and ``inres`` is assumed to be in lambda/dlambda.
:returns flux:
The output smoothed flux vector, same length as ``outwave``.
"""
# sigma_eff is in angstroms
if inres <= 0:
sigma_eff_sq = sigma**2
elif in_vel:
# Make an approximate correction for the intrinsic wavelength
# dependent dispersion. This sort of maybe works.
sigma_eff_sq = sigma**2 - (wave / inres)**2
else:
sigma_eff_sq = sigma**2 - inres**2
if np.any(sigma_eff_sq < 0):
raise ValueError("Desired wavelength sigma is lower than the value "
"possible for this input spectrum.")
sigma_eff = np.sqrt(sigma_eff_sq)
flux = np.zeros(len(outwave))
for i, w in enumerate(outwave):
x = (wave - w) / sigma_eff
if nsigma > 0:
good = np.abs(x) < nsigma
x = x[good]
_spec = spec[good]
else:
_spec = spec
f = np.exp(-0.5 * x**2)
flux[i] = np.trapz(f * _spec, x) / np.trapz(f, x)
return flux