first pass

signalgen.py

@@ -3,6 +3,14 @@
 import logging
 from os import path
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
+import numpy as np
+
+import matplotlib.pyplot as plt
+from jess.dispersion import dedisperse
+from jess.fitters import median_fitter
+from scipy.stats import median_abs_deviation
+from your import Your, Writer
+from will import create, inject
 
 
 
@@ -11,6 +19,43 @@ logger = logging.getLogger()
 logFmt = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 logging.basicConfig(level=logging.INFO, format=logFmt)
 
+#This function from https://josephwkania.github.io/will/examples/inject_pulse.html#Show-how-we-can-inject-a-pulse-into-a-GREENBURST-filterbank.
+def show_dynamic(
+    dynamic_spectra: np.ndarray,
+    title: Union[str, None] = None,
+    save: Union[bool, None] = False,
+) -> None:
+    """
+    Show a dynamic spectra by first flattening it
+    in frequency. Do this by getting the medians of
+    each channel and then run a median filter along the
+    bandpass.
+
+    Then set the limits of the imshow so we get good detail
+    for the majority of the data.
+
+    Args:
+        dynmaic_spectra - the dynamic spectra to plot
+
+        title - Title of plot
+
+        save - Save the plot as `title` + `.png`
+    """
+    spectra_mads = median_fitter(np.median(dynamic_spectra, axis=0))
+    flat = dynamic_spectra - spectra_mads
+    std = median_abs_deviation(flat, axis=None)
+    med = np.median(flat)
+    plt.figure(figsize=(20, 10))
+    plt.imshow(flat.T, vmin=med - 3 * std, vmax=med + 6 * std, aspect="auto")
+    plt.xlabel("Time Sample #", size=20)
+    plt.ylabel("Channel #", size=20)
+    plt.colorbar()
+    plt.tight_layout()
+    if title is not None:
+        plt.title(title, size=28)
+    if save:
+        plt.savefig(title.replace(" ", "_") + ".png", dpi=75, bbox_inches="tight")
+
 def processList(values):
 	"""
 	Entry point if values.file is None and values.listfile is set. 
@@ -26,7 +71,48 @@ def processList(values):
 def addBurst(values):
 	"""
 	Entry point if values.file is set.
+	--listfile will enter into this function multiple times for each line.
 	"""
+	filterbankObj = Your(values.file)
+	samples = filterbankObj.your_header.nspectra
+	spectra = filterbankObj.get_data(0, samples)
+	#save pre-injection spectra plot
+	show_dynamic(spectra, f"{path.basename(values.file)} Pre-injection Dynamic Spectra", save=True)
+
+	#get bandpass and store in bpWeights
+	bpWeights = create.filter_weights(spectra)
+
+	#create pulse
+	#first version is very simple, plan on adding more complex injections in future
+	pulseObj = create.SimpleGaussPulse(
+		sigma_time=0.001,
+		sigma_freq=350,
+		center_freq = filterbankObj.your_header.center_freq,
+		dm = values.dm,
+		tau = 20,
+		phi=np.pi / 3, #does nothing if nscint = 0
+		spectral_index_alpha=0,
+		chan_freqs = filterbankObj.chan_freqs,
+		tsamp = filterbankObj.your_header.tsamp,
+		nscint=0,
+		bandpass = bpWeights
+	)
+	pulse = pulseObj.sample_pulse(nsamp=int(3e5)) #sample it 300000 times
+
+	#inject pulse
+	injectedSpectra = inject.inject_constant_into_file(
+		yr_input = filterbankObj,
+		pulse = pulse,
+		start = samples // 4,
+    	gulp = samples
+	)
+
+	#and save the new plot and filterbank file
+	show_dynamic(spectra, f"{path.basename(values.file)} Dynamic Spectra and {values.dm} DM Pulse", save=True)
+	filWriter = Writer(filterbankObj, outdir="./", outname = f"{path.basename(values.file)}_injected")
+	filWriter.to_fil(data=injectedSpectra)
+	logger.info(f"{path.basename(values.file)}_injected.fil successfully written.")
+
 
 
 
@@ -41,6 +127,10 @@ if __name__ == "__main__":
 	parser.add_argument(
 		"-f", "--file", dest="file", type=str, help="Single filterbank file."
 	)
+	parser.add_argument(
+		"-d", "--dm", dest="dm", type=float, help="DM of injected pulse."
+	)
+	parser.set_defaults(dm=250.0)
 	parser.set_defaults(listfile=None)
 	parser.set_defaults(file=None)
 	values = parser.parse_args()