def toa_score(template,
pcs,
coeffs,
profile,
ts=None,
tol=sqrt(np.finfo(np.float64).eps)):
'''
Calculate a maximum-likelihood TOA given a template and a PCA model of pulse shape variations.
Uses the dot-product-based method of Osłowski (2011).
`pcs`: The principal components (unit vectors), as rows of an array.
`coeffs`: Coefficients of principal component dot products to use in correcter.
`ts`: Evenly-spaced array of phase values corresponding to the profile.
Sets the units of the TOA. If this is `None`, the TOA is reported in bins.
`tol`: Relative tolerance for optimization (in bins).
'''
n = len(profile)
if ts is None:
ts = np.arange(n)
dt = float(ts[1] - ts[0])
k = len(pcs)
result = toa_fourier(template, profile, ts=ts, tol=tol)
initial_toa = result.toa
ampl = result.ampl
template_shifted = fft_roll(template, initial_toa / dt)
pcs_shifted = fft_roll(pcs, initial_toa / dt)
scores = np.dot(pcs_shifted, profile)
correcter = np.dot(coeffs, scores)
toa = initial_toa - correcter
return ToaScoreResult(toa=toa, ampl=ampl, scores=scores)
def test_toa_recovery(func,
template,
n,
rms_toa,
SNR=np.inf,
ts=None,
tol=sqrt(eps)):
'''
Test function for `toa_ws()` and `toa_fourier()`.
Attempts to recover `n` TOAs at a given SNR and returns the RMS error.
`func`: Function to test (`toa_ws` or `toa_fourier`).
`template`: Template to use. Test profiles will be generated by
shifting it.
`n`: Number of test profiles to generate.
`rms_toa`: RMS TOA for test profiles.
`tol`: Relative tolerance for optimization.
'''
if ts is None:
ts = np.arange(len(template))
dt = ts[1] - ts[0]
dtoas = []
toa_errs = []
for i in range(n):
true_toa = rms_toa * randn()
profile = fft_roll(template, true_toa / dt)
if np.isfinite(SNR):
profile += randn(len(profile)) / SNR
result = func(template, profile, ts=ts, tol=tol)
toa_estimate = result.toa
dtoas.append(toa_estimate - true_toa)
toa_errs.append(result.error)
dtoas = np.array(dtoas)
toa_errs = np.array(toa_errs)
return np.sqrt(np.mean(dtoas**2)), np.sqrt(np.mean(toa_errs**2))
def chisq_fn(shift):
shifted_profile = fft_roll(profile, -shift)
pc_ampls = pcs @ shifted_profile
chisq = shifted_profile @ shifted_profile
chisq -= (template @ shifted_profile)**2 / (template @ template)
chisq -= np.sum(pc_ampls**2)
chisq /= noise_level**2
chisq += np.sum(pc_ampls**2 / (sgvals**2 + noise_level**2))
return chisq
def toa_fourier(template, profile, ts=None, noise_level=None, tol=sqrt(eps)):
'''
Calculate a TOA by maximizing the CCF of the template and the profile
in the frequency domain. Searches within the interval between the sample
below and the sample above the argmax of the circular CCF.
`ts`: Evenly-spaced array of phase values corresponding to the profile.
Sets the units of the TOA. If this is `None`, the TOA is reported
in bins.
`tol`: Relative tolerance for optimization (in bins).
`noise_level`: Off-pulse noise, in the same units as the profile.
Used in calculating error. If not supplied, noise level will be
estimated as the standard deviation of the profile residual.
'''
n = len(profile)
if ts is None:
ts = np.arange(n)
dt = float(ts[1] - ts[0])
template_fft = fft(template)
profile_fft = fft(profile)
phase_per_bin = -2j * pi * fftfreq(n)
circular_ccf = irfft(rfft(profile) * np.conj(rfft(template)), n)
ccf_argmax = np.argmax(circular_ccf)
if ccf_argmax > n / 2:
ccf_argmax -= n
ccf_max = ccf_argmax * dt
def ccf_fourier(tau):
phase = phase_per_bin * tau / dt
ccf = np.inner(profile_fft, exp(-phase) * np.conj(template_fft)) / n
return ccf.real
brack = (ccf_max - dt, ccf_max, ccf_max + dt)
toa = minimize_scalar(lambda tau: -ccf_fourier(tau),
method='Brent',
bracket=brack,
tol=tol * dt).x
assert brack[0] < toa < brack[-1]
template_shifted = fft_roll(template, toa / dt)
b = np.dot(template_shifted, profile) / np.dot(template, template)
residual = profile - b * template_shifted
ampl = b * np.max(template_shifted)
if noise_level is None:
noise_level = offpulse_rms(profile, profile.size // 4)
snr = ampl / noise_level
w_eff = np.sqrt(n * dt / np.trapz(np.gradient(template, ts)**2, ts))
error = w_eff / (snr * sqrt(n))
return ToaResult(toa=toa, error=error, ampl=ampl)
def chisq_fn(shift):
shifted_profile = fft_roll(profile, -shift)
t_ampl = template @ shifted_profile / (template @ template)
td_ampl = template_deriv @ shifted_profile / (
template_deriv @ template_deriv)
pc_ampls = pcs @ shifted_profile
pred_td_ampl = coeffs @ pc_ampls
chisq = (td_ampl - pred_td_ampl)**2 / (t_ampl * scatter)**2
chisq += np.sum(pc_ampls**2 / (t_ampl * sgvals)**2)
chisq += (t_ampl - template_ampl)**2 / template_scatter**2
return chisq
def chisq_fn(shift):
shifted_profile = fft_roll(profile, -shift)
t_ampl = template @ shifted_profile / np.sqrt(template @ template)
td_ampl = template_deriv @ shifted_profile / np.sqrt(
template_deriv @ template_deriv)
pc_ampls = pcs @ shifted_profile
pred_td_ampl = coeffs @ pc_ampls * np.sqrt(
template_deriv @ template_deriv)
chisq = (td_ampl -
pred_td_ampl)**2 / (noise_level**2 + scatter**2 *
(template_deriv @ template_deriv))
chisq += np.sum(pc_ampls**2 / (sgvals**2 + noise_level**2))
return chisq
def gen_data(spec, n_profiles, npprof, n_bins, SNR, drift_bins):
'''
Generated simulated data based on a pulse specification.
'''
phase = np.linspace(-1 / 2, 1 / 2, n_bins, endpoint=False)
profiles = gen_profiles(phase,
spec=spec,
n_profiles=n_profiles,
npprof=npprof,
SNR=SNR)
shifts = drift_bins / n_profiles * np.arange(n_profiles)
shifts -= np.mean(shifts)
for i, profile in enumerate(profiles):
profiles[i] = fft_roll(profile, shifts[i])
return phase, profiles
def toa_ws(template, profile, ts=None, noise_level=None, tol=sqrt(eps)):
'''
Calculate a TOA by maximizing the Whittaker-Shannon interpolant of the
CCF between `template` and `profile`. Searches within the interval
between the sample below and the sample above the argmax of the CCF.
`ts`: Evenly-spaced array of phase values corresponding to the profile.
Sets the units of the TOA. If this is `None`, the TOA is reported
in bins.
`tol`: Relative tolerance for optimization.
`noise_level`: Off-pulse noise, in the same units as the profile.
Used in calculating error. If not supplied, noise level will be
estimated as the standard deviation of the profile residual.
'''
n = len(profile)
if ts is None:
ts = np.arange(n)
dt = ts[1] - ts[0]
lags = np.arange(-len(ts) + 1, len(ts)) * dt
ccf = np.correlate(profile, template, mode='full')
ccf_max = lags[np.argmax(ccf)]
interpolant = interp_ws(ccf, lags)
brack = (ccf_max - dt, ccf_max, ccf_max + dt)
toa = minimize_scalar(lambda t: -interpolant(t),
method='Brent',
bracket=brack,
tol=tol).x
assert brack[0] < toa < brack[-1]
template_shifted = fft_roll(template, toa / dt)
b = np.dot(template_shifted, profile) / np.dot(template, template)
residual = profile - b * template_shifted
ampl = b * np.max(template_shifted)
if noise_level is None:
noise_level = offpulse_rms(profile, profile.size // 4)
snr = ampl / noise_level
w_eff = np.sqrt(n * dt / np.trapz(np.gradient(template, ts)**2, ts))
error = w_eff / (snr * sqrt(n))
return ToaResult(toa=toa, error=error, ampl=ampl)
def get_template(profiles, n_iter=1):
'''
Create a template by iteratively aligning and averaging profiles.
Starts by averaging the middle 10% of profiles and iterates `n_iter` times,
each time aligning the pulses using the previous template and averaging
to create a new template.
'''
n = profiles.shape[0]
n_5_percent = n // 20
sl = slice(n // 2 - n_5_percent, n // 2 + n_5_percent)
template = np.mean(profiles[sl], axis=0)
toas = np.empty(profiles.shape[0])
for i in range(n_iter):
for i, profile in enumerate(profiles):
result = toa_fourier(template, profile)
toas[i] = result.toa
profiles_aligned = np.empty_like(profiles)
for j, profile in enumerate(profiles):
profiles_aligned[j] = fft_roll(profile, -toas[j])
template = np.mean(profiles_aligned, axis=0)
return template
def toa_pca_prior(template,
pcs,
weights,
profile,
ts=None,
tol=sqrt(np.finfo(np.float64).eps),
plot=False):
'''
Calculate a maximum-likelihood TOA given a template and a PCA model of pulse shape variations.
`pcs`: The principal components (unit vectors), as rows of an array.
`ts`: Evenly-spaced array of phase values corresponding to the profile.
Sets the units of the TOA. If this is `None`, the TOA is reported in bins.
`tol`: Relative tolerance for optimization (in bins).
'''
n = len(profile)
if ts is None:
ts = np.arange(n)
dt = float(ts[1] - ts[0])
k = len(pcs)
template_fft = fft(template)
profile_fft = fft(profile)
pcs_fft = fft(pcs)
phase_per_bin = -2j * pi * fftfreq(n)
circular_ccf = irfft(rfft(profile) * np.conj(rfft(template)), n) * dt
sq_ccf = circular_ccf**2 / (np.sum(template**2) * dt)
for i in range(k):
pcfft = irfft(rfft(profile) * np.conj(rfft(pcs[i]))) * sqrt(dt)
sq_ccf += pcfft**2 / (1 + weights[i])
ccf_argmax = np.argmax(sq_ccf)
ccf_max_val = sq_ccf[ccf_argmax]
ccf_max = ccf_argmax * dt
if ccf_argmax > n / 2:
ccf_max -= n * dt
def modified_squared_ccf(tau):
phase = phase_per_bin * tau / dt
ccf = np.inner(profile_fft,
exp(-phase) * np.conj(template_fft)) * dt / n
sq_ccf = ccf.real**2 / (np.sum(template**2) * dt)
for i in range(k):
pc_fft = pcs_fft[i]
weight = weights[i]
pccf = np.inner(profile_fft,
exp(-phase) * np.conj(pc_fft)) * sqrt(dt) / n
sq_ccf += pccf.real**2 / (1 + weight)
return sq_ccf
brack = (ccf_max - dt, ccf_max, ccf_max + dt)
toa = minimize_scalar(lambda tau: -modified_squared_ccf(tau),
method='Brent',
bracket=brack,
tol=tol * dt).x
assert brack[0] < toa < brack[-1]
template_shifted = fft_roll(template, toa / dt)
b = np.dot(template_shifted, profile) / np.dot(template, template)
ampl = b * np.max(template_shifted)
pcs_shifted = fft_roll(pcs, toa / dt)
scores = np.dot(pcs_shifted, profile)
return ToaPcaResult(toa=toa, ampl=ampl, scores=scores)
def extract_pcs(profiles,
n_pcs,
initial_template=None,
return_all=True,
use_trend=True):
'''
Extract a template and principal components from a set of profiles.
An initial template can be supplied; if not, the default strategy is to average
the middle 10 percent of profiles to get an initial template.
Inputs
------
profiles: The profiles, as rows of a 2-D array.
n_pcs: The number of principal components to use in the model.
n_iter: The number of iterations to perform.
initial_template: The initial template (see above).
return_all: Return all principal components (instead of the first `n_pcs`).
use_trend: If `False`, align profiles using their individual TOAs,
ignoring n_iter. If `True`, align using a linear trend (default).
Outputs
-------
template: The final template
pcs: The final principal components
sgvals: The singular values (characteristic amplitudes) corresponding to the
principal components.
scores: The PC scores of the training data.
dtoas: The differences between the TOAs and the linear trend.
'''
n_profiles = profiles.shape[0]
profile_number = np.arange(n_profiles)
if initial_template is None:
initial_template = get_template(profiles, n_iter=0)
toas = np.zeros(n_profiles)
for i, profile in enumerate(profiles):
result = toa_fourier(initial_template, profile)
toas[i] = result.toa
resids = np.empty_like(profiles)
if use_trend:
trend_coeffs = np.polyfit(profile_number, toas, 1)
trend = np.polyval(trend_coeffs, profile_number)
profiles_aligned = np.empty_like(profiles)
for j, profile in enumerate(profiles):
profiles_aligned[j] = fft_roll(profile, -trend[j])
template = np.mean(profiles_aligned, axis=0)
for j, profile in enumerate(profiles_aligned):
ampl = np.dot(profile, template) / np.dot(template, template)
resids[j] = profile - ampl * template
u, s, pcs = svd(resids, full_matrices=return_all)
sgvals = s / np.sqrt(resids.shape[0])
scores = np.dot(pcs, profiles_aligned.T)
dtoas = toas - trend
else:
profiles_aligned = np.empty_like(profiles)
for j, profile in enumerate(profiles):
profiles_aligned[j] = fft_roll(profile, -toas[j])
template = np.mean(profiles_aligned, axis=0)
for j, profile in enumerate(profiles_aligned):
ampl = np.dot(profile, template) / np.dot(template, template)
resids[j] = profile - ampl * template
u, s, pcs = svd(resids, full_matrices=return_all)
sgvals = s / np.sqrt(resids.shape[0])
# Trend used only for computing ΔTOAs
trend_coeffs = np.polyfit(profile_number, toas, 1)
trend = np.polyval(trend_coeffs, profile_number)
scores = np.dot(pcs, profiles_aligned.T)
dtoas = toas - trend
return template, pcs, sgvals, scores, dtoas