diff --git a/pta_replicator/red_noise.py b/pta_replicator/red_noise.py index 100ebdd..ea5cab9 100644 --- a/pta_replicator/red_noise.py +++ b/pta_replicator/red_noise.py @@ -103,6 +103,41 @@ def create_fourier_design_matrix_red(toas: np.ndarray, nmodes: int = 30, return F, Ffreqs +def create_fourier_design_matrix_red_chromatic(toas: np.ndarray, nmodes: int = 30, + Tspan: float = None, logf: bool = False, + fmin: float = None, fmax: float = None, + pshift: bool = False, libstempo_convention: bool = False, modes: np.ndarray = None, + fref=1400, alpha=2.0) -> tuple: + """ + Construct chromatic fourier design matrix + + Parameters + ---------- + toas [array]: Vector of time series in seconds. + nmodes [int]: Number of fourier coefficients to use. + Tspan [float]: Option to us some other Tspan [s] + logf [bool]: Use log frequency spacing. + fmin [float]: Lower sampling frequency. + fmax [float]: Upper sampling frequency. + pshift [bool]: Option to add random phase shift. + libstempo_convention [bool]: Use libstempo's convention where toas-toas[0] is used in F matrix instead of just toas, and cosine comes before sine + modes [array]: Option to provide explicit list or array of sampling frequencies. + fref [float]: Reference frequency for chromaticity. + alpha [float]: Chromatic index. + + Returns + ------- + F [array]: fourier design matrix, [NTOAs x 2 nfreqs]. + f [array]: Sampling frequencies, [2 nfreqs]. + """ + F, Ffreqs = create_fourier_design_matrix_red(toas=toas, nmodes=nmodes, Tspan=Tspan, + logf=logf, fmin=fmin, fmax=fmax, pshift=pshift, + libstempo_convention=libstempo_convention, modes=modes) + chromaticity = (Ffreqs / fref)**-alpha + + return F * chromaticity[:, None], Ffreqs + + def add_red_noise(psr: SimulatedPulsar, log10_amplitude: float, spectral_index: float, components: int = 30, seed: int = None, modes: np.ndarray = None, Tspan: float = None, libstempo_convention: bool = True): @@ -133,6 +168,44 @@ def add_red_noise(psr: SimulatedPulsar, log10_amplitude: float, spectral_index: psr.toas.adjust_TOAs(TimeDelta(dt.to('day'))) psr.update_residuals() + + +def add_chromatic_noise(psr: SimulatedPulsar, log10_amplitude: float, spectral_index: float, alpha: float, + components: int = 30, seed: int = None, signal_name: str = 'chrom', + modes: np.ndarray = None, Tspan: float = None, libstempo_convention: bool = False): + """Add chromatic red noise with P(f) = A^2 / (12 pi^2) (f * year)^-gamma, + using `components` Fourier bases. + Optionally take a pseudorandom-number-generator seed.""" + psr.update_added_signals('{}_{}_noise'.format(psr.name, signal_name), + {'amplitude': log10_amplitude, 'spectral_index': spectral_index}) + A = 10**(log10_amplitude) + gamma = spectral_index + # nobs = len(psr.toas.table) + + fyr = 1 / YEAR_IN_SEC + + if seed is not None: + np.random.seed(seed) + if modes is not None: + print('Must use linear spacing.') + + toas = np.array(psr.toas.table['tdbld'], dtype='float64') * DAY_IN_SEC #to sec + Tspan = toas.max() - toas.min() + F, freqs = create_fourier_design_matrix_red_chromatic(toas, alpha=alpha, Tspan=Tspan, nmodes=components, modes=modes, libstempo_convention=libstempo_convention) + prior = A**2 * (freqs/fyr)**(-gamma) / (12 * np.pi**2 * Tspan) * YEAR_IN_SEC**3 + y = np.sqrt(prior) * np.random.randn(freqs.size) + dt = np.dot(F,y) * u.s + psr.toas.adjust_TOAs(TimeDelta(dt.to('day'))) + psr.update_residuals() + + +def add_dm_noise(psr: SimulatedPulsar, log10_amplitude: float, spectral_index: float, + components: int = 30, seed: int = None, + modes: np.ndarray = None, Tspan: float = None, libstempo_convention: bool = False): + 'Convenience function for adding DM noise.' + add_chromatic_noise(psr=psr, log10_amplitude=log10_amplitude, spectral_index=spectral_index, + alpha=2., components=components, seed=seed, signal_name='dm', + modes=modes, Tspan=Tspan, libstempo_convention=libstempo_convention) def add_gwb(