From 6582c6498c804fb0ce514dcc7868311d1d4f7f59 Mon Sep 17 00:00:00 2001 From: Kaze Wong Date: Tue, 26 Dec 2023 14:41:44 -0500 Subject: [PATCH] update injecting detector noise --- src/jimgw/detector.py | 31 +++++++++++++++++++++++++------ 1 file changed, 25 insertions(+), 6 deletions(-) diff --git a/src/jimgw/detector.py b/src/jimgw/detector.py index bc859b99..eb8176ad 100644 --- a/src/jimgw/detector.py +++ b/src/jimgw/detector.py @@ -371,18 +371,37 @@ def antenna_pattern(self, ra: Float, dec: Float, psi: Float, gmst: Float) -> dic def inject_signal( self, key: PRNGKeyArray, - freqs: Array, - h_sky: dict, - params: dict, + freqs: Float[Array, " n_sample"], + h_sky: dict[str, Float[Array, " n_sample"]], + params: dict[str, Float], psd_file: str = "", ) -> None: - """ """ + """ + Inject a signal into the detector data. + + Parameters + ---------- + key : PRNGKeyArray + JAX PRNG key. + freqs : Float[Array, " n_sample"] + Array of frequencies. + h_sky : dict[str, Float[Array, " n_sample"]] + Array of waveforms in the sky frame. The key is the polarization mode. + params : dict[str, Float] + Dictionary of parameters. + psd_file : str + Path to the PSD file. + + Returns + ------- + None + """ self.frequencies = freqs self.psd = self.load_psd(freqs, psd_file) key, subkey = jax.random.split(key, 2) var = self.psd / (4 * (freqs[1] - freqs[0])) - noise_real = jax.random.normal(key, shape=freqs.shape) * jnp.sqrt(var) - noise_imag = jax.random.normal(subkey, shape=freqs.shape) * jnp.sqrt(var) + noise_real = jax.random.normal(key, shape=freqs.shape) * jnp.sqrt(var / 2.0) + noise_imag = jax.random.normal(subkey, shape=freqs.shape) * jnp.sqrt(var / 2.0) align_time = jnp.exp( -1j * 2 * jnp.pi * freqs * (params["epoch"] + params["t_c"]) )