From c4d88037f3b20b195df36bce8b96972a2b93685a Mon Sep 17 00:00:00 2001
From: Sakimori <sakimorihalley@gmail.com>
Date: Tue, 29 Jul 2025 14:33:09 -0400
Subject: [PATCH] added random sampling to test pipeline in various situations

---
 signalgen.py | 48 ++++++++++++++++++++++++++++++++++--------------
 1 file changed, 34 insertions(+), 14 deletions(-)

diff --git a/signalgen.py b/signalgen.py
index 4e4be54..69cd6c4 100755
--- a/signalgen.py
+++ b/signalgen.py
@@ -4,6 +4,7 @@ import logging
 from os import path, chdir, getcwd, makedirs
 from argparse import ArgumentDefaultsHelpFormatter, ArgumentParser
 from typing import Union
+from datetime import datetime
 
 import numpy as np
 import matplotlib.pyplot as plt
@@ -22,6 +23,13 @@ logger = logging.getLogger()
 logFmt = "%(asctime)s - %(name)s - %(levelname)s - %(message)s"
 logging.basicConfig(level=logging.INFO, format=logFmt)
 
+#Set values for random sampling
+MIN_DM = 300
+MAX_DM = 10000
+MIN_WIDTH = 0.001 #1 ms
+MAX_WIDTH = 1 #1 s
+rng = np.random.default_rng()
+
 #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,
@@ -79,6 +87,13 @@ def processList(values):
 			dataName = path.dirname(dir[:-1]) #strip the trailing slash for dirname
 			values.file = path.join(dir,dataname+".fil")
 			addBurst(values)
+
+def randomDMandWidth():
+	"""
+	Helper function that returns (DM, pulseWidth) tuple with bounds set at start of script.
+	"""
+	return (rng.uniform(low=MIN_DM, high=MAX_DM), rng.uniform(low=MIN_WIDTH, high=MAX_WIDTH))
+
             
 def addBurst(values):
 	"""
@@ -97,22 +112,29 @@ def addBurst(values):
 	filterbankObj = Your(path.join("./", f"{basename}_trunc.fil"))
 	spectra = filterbankObj.get_data(0, samples)
 
-	#save pre-injection spectra plot
-	if not values.skipplot:
-		show_dynamic(spectra, f"{basename} Pre-injection Dynamic Spectra", save=True)
-
 	#get bandpass and store in bpWeights
 	bpWeights = create.filter_weights(spectra)
 
 	logger.info(f"{basename} loaded. Sampling pulse {values.nsamp} times.")
 
 	#create pulse
+
+	#check if this is part of a rng run
+	if values.rsamp:
+		#generate the values AND save the values to file for comparison later
+		dm, pWidth = randomDMandWidth()
+		with open(values.thisRunName, "a") as file:
+			file.write(f"{str(round(dm)).ljust(6)} pc/cc {str(round(pWidth,3))} s {basename}\n")
+	else:
+		dm = values.dm
+		pWidth = 0.001
+
 	#first version is very simple, plan on adding more complex injections in future
 	pulseObj = create.SimpleGaussPulse(
-		sigma_time=0.001,
+		sigma_time=pWidth,
 		sigma_freq=350,
 		center_freq = filterbankObj.your_header.center_freq,
-		dm = values.dm,
+		dm = dm,
 		tau = 20,
 		phi=np.pi / 3, #does nothing if nscint = 0
 		spectral_index_alpha=0,
@@ -131,10 +153,6 @@ def addBurst(values):
 		start = samples // 3
 	)
 
-	#and save the new plot
-	if not values.skipplot:
-		logger.info("Saving plot...")
-		show_dynamic(injectedSpectra, f"{basename} Dynamic Spectra and {values.dm} DM Pulse", save=True)
 	#now generate new filterbank file
 	newName = f"{basename}_injected.fil"
 	sigprocObj = make_sigproc_object(
@@ -177,6 +195,9 @@ if __name__ == "__main__":
 	parser.add_argument(
 		"-f", "--file", dest="file", type=str, help="Single filterbank file."
 	)
+	parser.add_argument(
+		"-r", "--random", dest="rsamp", action="store_true", help="Use random DM and pulse widths with bounds set by file. Ignores -d."
+	)
 	parser.add_argument(
 		"-d", "--dm", dest="dm", type=float, help="DM of injected pulse."
 	)
@@ -186,15 +207,12 @@ if __name__ == "__main__":
 	parser.add_argument(
 		"-p", "--plot", action="store_true", help="Just plot file and quit."
 	)
-	parser.add_argument(
-		"-s", "--skipplot", action="store_true", help="Skip plotting for large filterbanks."
-	)
 	parser.set_defaults(dm=250.0)
 	parser.set_defaults(nsamp=int(3e5))
 	parser.set_defaults(listfile=None)
 	parser.set_defaults(file=None)
 	parser.set_defaults(plot=False)
-	parser.set_defaults(skipplot=False)
+	parser.set_defaults(rsamp=False)
 	values = parser.parse_args()
 
 	#set working directory to ignored directory
@@ -203,6 +221,8 @@ if __name__ == "__main__":
 		makedirs(outdir)
 	chdir(outdir)
 
+	values.thisRunName = datetime.now().isoformat(timespec='seconds').replace(":", "-") + ".txt"
+
 	if values.file is not None: #single file takes priority
 		logging.info(f"Running with file {values.file}")
 		if values.plot: