bbrzycki / setigen

Python library for generating and injecting artificial narrow-band signals into radio frequency data
https://setigen.readthedocs.io
MIT License
25 stars 11 forks source link

Gist: using setigen to generate FRBs #25

Open telegraphic opened 3 years ago

telegraphic commented 3 years ago

I've been using setigen as a base to generate mock FRBs. Here's some example code:

def gauss2(x, a, x0, fwhm):
    """ Generate gaussian^2 pulse 

    Args:
        x: time series (np.arange())
        a: amplitude of pulse. (total integrated power across pulse)
        x0: index of peak value
        fwhm: width of signal in samples
    """
    sigma = (fwhm / 2) / np.sqrt(2*np.log(2))
    return a / np.sqrt(2*np.pi*sigma**2) * np.exp(-(x-x0)**2 / (2*sigma**2))

def generate_frb(frame, frb_params):
    """ Simple FRB generator 

    Args:
        frame (setigen.Frame): A filterbank frame of data, from setigen
        frb_params (dict): A dictionary with FRB parameters in it

    Notes:
        frb_params should contain width (in s), desired SNR, start time (in s), and DM
    """

    frb_width, frb_snr, frb_t0, frb_dm = frb_params['width'], frb_params['snr'], frb_params['t0'], frb_params['dm'], 

    frb_rms  = frame.get_intensity(snr=frb_snr)
    fch1 = frame.get_frequency(0)

    ## Generate pulse delays - array has same length as freqs
    width_in_chans = frb_width / frame.dt
    t0_in_samps = (frb_t0 / frame.dt) - frame.ts[0]
    tdel_in_samps = 4.15e-3 * frb_dm * ((fch1/1e9)**(-2) - (frame.fs/1e9)**(-2)) / frame.dt
    t0_in_samps = t0_in_samps + tdel_in_samps 

    # Time axis in samples
    t = np.arange(frame.tchans)

    # Use meshgrid to create (N_time, N_freq) array
    t2d, tdel2d = np.meshgrid(t, t0_in_samps)

    # Create pulse profile via 2D time (N_time,1) and time_delay (N_freq, 1) arrays
    profile = gauss2(t2d, frb_rms, tdel2d, width_in_chans)

    return profile.T

## Example usage
frame = stg.Frame(fchans=1024*u.pixel,
                  tchans=1024*2*u.pixel,
                  df=100.0*u.kHz,
                  dt=1*u.ms,
                  fch1=300*u.MHz)

noise = frame.add_noise(x_mean=10, noise_type='chi2')

frb_params = {'snr': 10000.0, 
              't0': 500e-3, 
              'dm': 3, 
              'width': 10e-3}
p = generate_frb(frame, frb_params)
frame.data += p

frb_params = {'snr': 10000.0, 
              't0': 1000e-3, 
              'dm': 10, 
              'width': 10e-3}
p = generate_frb(frame, frb_params)
frame.data += p

frame.plot()

Which gives the following output:

image

Just a proof of concept, but might be worth adding some functionality? This could be useful for FRB studies, and also for SETI seraches for artificial pulses (e.g. different exponents, negative DMs).

bbrzycki commented 2 years ago

This looks great! Moving towards broadband signals is an important direction for setigen. The main challenge will be balancing which methods are general enough to apply to both narrow and broadband injection, since many currently implicitly assume narrowband (for example, technically the get_intensity function converts to and from SNR assuming integration along the time direction). But having something analogous to the current spectrogram injection method for quick broadband injection could be really useful; I'm planning on spending some time brainstorming about this and perhaps we can ultimately use a version of what you came up with already!