added random sampling to test pipeline in various situations
This commit is contained in:
parent
2182108909
commit
c4d88037f3
48
signalgen.py
48
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,
|
||||
|
@ -80,6 +88,13 @@ def processList(values):
|
|||
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):
|
||||
"""
|
||||
Entry point if values.file is set.
|
||||
|
@ -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:
|
||||
|
|
Loading…
Reference in a new issue