|
| 1 | +#!/usr/bin/env python3 |
| 2 | +import os, time, random |
| 3 | +import numpy as np |
| 4 | +from astropy.time import Time |
| 5 | +import dsautils.dsa_store as ds |
| 6 | +import slack_sdk as slack |
| 7 | + |
| 8 | +# ----------------- config ----------------- |
| 9 | +SLACK_CHANNEL = "candidates" |
| 10 | +SLACK_TOKEN_FILE = f"{os.path.expanduser('~')}/.config/slack_api" |
| 11 | +INJ_LIST = "/home/ubuntu/data/injections/injection_list.txt" |
| 12 | +AUDIT_INJECTION_FILENAME = "/operations/T2/injection_audit_results/injections_for_audit.txt" |
| 13 | +PARAMS_TXT = "/home/ubuntu/simulated_frb_params.txt" # columns: DM SNR width_fwhm spec_ind |
| 14 | +TEMPLATES = [ |
| 15 | + "/home/ubuntu/data/burst_0.inject", |
| 16 | + "/home/ubuntu/data/burst_1.inject", |
| 17 | + "/home/ubuntu/data/burst_2.inject", |
| 18 | + "/home/ubuntu/data/burst_3.inject", |
| 19 | + "/home/ubuntu/data/burst_4.inject", |
| 20 | +] |
| 21 | + |
| 22 | +SLEEP_SEC = 600 # 10 minutes |
| 23 | + |
| 24 | +NODES = [17, 18, 19, 20] |
| 25 | +PAIR = {17: 19, 18: 20, 19: 17, 20: 18} |
| 26 | +LOCAL_BEAMS_PER_NODE = 128 |
| 27 | + |
| 28 | +# Legacy output formatting (unchanged) |
| 29 | +FMT_OUT = "%5.9f %d %0.2f %0.1f %0.3f %0.2f %s\n" |
| 30 | + |
| 31 | +# --- SNR→scale mapping (Recovered_SNR ≈ k * scale) --- |
| 32 | +K_DEFAULT = 135.0 # 27/.2 from old script |
| 33 | +SCALE_MIN = 0.075 |
| 34 | +SCALE_MAX = 0.3 |
| 35 | +niterations = 288 |
| 36 | +scale = np.random.uniform(SCALE_MIN, SCALE_MAX, niterations) |
| 37 | +# ------------------------------------------------------ |
| 38 | + |
| 39 | +def ensure_slack(): |
| 40 | + if not os.path.exists(SLACK_TOKEN_FILE): |
| 41 | + raise RuntimeError(f"Could not find file with slack api token at {SLACK_TOKEN_FILE}") |
| 42 | + with open(SLACK_TOKEN_FILE, "r") as sf_handler: |
| 43 | + slack_token = sf_handler.read() |
| 44 | + return slack.WebClient(token=slack_token) |
| 45 | + |
| 46 | +def slack_msg(cli, text): |
| 47 | + try: |
| 48 | + cli.chat_postMessage(channel=SLACK_CHANNEL, text=text) |
| 49 | + except Exception as e: |
| 50 | + print(f"[slack] {e}") |
| 51 | + |
| 52 | +def ensure_injection_list(filename): |
| 53 | + if not os.path.exists(filename): |
| 54 | + with open(filename, "w") as f: |
| 55 | + f.write("# MJD Beam DM SNR Width_fwhm spec_ind FRBno\n") |
| 56 | + |
| 57 | +def global_to_node_local(g): |
| 58 | + if not (0 <= g <= 511): |
| 59 | + raise ValueError("global beam must be in [0, 511]") |
| 60 | + group = g // LOCAL_BEAMS_PER_NODE |
| 61 | + node = NODES[group] |
| 62 | + local = g % LOCAL_BEAMS_PER_NODE |
| 63 | + return node, local |
| 64 | + |
| 65 | +def load_params_and_templates(): |
| 66 | + params = np.genfromtxt(PARAMS_TXT) # (N,4): DM, SNR, width_fwhm, spec_ind |
| 67 | + n = min(len(params), len(TEMPLATES)) |
| 68 | + if n == 0: |
| 69 | + raise RuntimeError("No params/templates found") |
| 70 | + if n < len(params): |
| 71 | + print(f"[warn] Truncating params to {n} to match templates") |
| 72 | + return params[:n], TEMPLATES[:n] |
| 73 | + |
| 74 | +def snr_to_scale(target_snr: float) -> float: |
| 75 | + sc = float(target_snr) / float(K_DEFAULT) |
| 76 | + if sc < SCALE_MIN: sc = SCALE_MIN |
| 77 | + if sc > SCALE_MAX: sc = SCALE_MAX |
| 78 | + return sc |
| 79 | + |
| 80 | +def scale_to_snr(scale: float) -> float: |
| 81 | + return float(scale * K_DEFAULT) |
| 82 | + |
| 83 | +def run(): |
| 84 | + random.seed() |
| 85 | + np.random.seed() |
| 86 | + ensure_injection_list(INJ_LIST) |
| 87 | + ensure_injection_list(AUDIT_INJECTION_FILENAME) |
| 88 | + |
| 89 | + slack_cli = ensure_slack() |
| 90 | + store = ds.DsaStore() |
| 91 | + |
| 92 | + params, templates = load_params_and_templates() |
| 93 | + |
| 94 | + for i in range(niterations): |
| 95 | + current_scale = scale[i] |
| 96 | + print(f"[info] Starting iteration {i+1}/{niterations}") |
| 97 | + for idx, row in enumerate(params): |
| 98 | + DM, SNR, width_fwhm, spec_ind = map(float, row) # SNR is fixed from file (e.g., 15.0) |
| 99 | + template = templates[idx] |
| 100 | + frbno = os.path.splitext(os.path.basename(template))[0].split("_")[-1] |
| 101 | + |
| 102 | + # choose a random global beam, then derive EW/NS pair |
| 103 | + gbeam = random.randint(0, 511) |
| 104 | + print(f"[info] Injecting FRB {frbno} with DM={DM}, SNR={SNR}, into global beam {gbeam} using template {template}") |
| 105 | + |
| 106 | + if gbeam <= 255: |
| 107 | + ew_beam, ns_beam = gbeam, gbeam + 256 |
| 108 | + else: |
| 109 | + ns_beam, ew_beam = gbeam, gbeam - 256 |
| 110 | + |
| 111 | + node, local = global_to_node_local(gbeam) |
| 112 | + partner = PAIR[node] |
| 113 | + |
| 114 | + |
| 115 | + effective_snr = scale_to_snr(current_scale) |
| 116 | + |
| 117 | + # event bookkeeping (legacy append) |
| 118 | + mjd = Time.now().mjd |
| 119 | + with open(INJ_LIST, "a") as f: |
| 120 | + # Beam column uses the global beam you selected |
| 121 | + f.write(FMT_OUT % (mjd, gbeam, DM, effective_snr, width_fwhm, spec_ind, frbno)) |
| 122 | + |
| 123 | + with open(AUDIT_INJECTION_FILENAME, "a") as f: |
| 124 | + f.write(FMT_OUT % (mjd, gbeam, DM, effective_snr, width_fwhm, spec_ind, frbno)) |
| 125 | + |
| 126 | + |
| 127 | + injection_cmd = {"cmd": "inject", "val": f"{local}-{template}-{current_scale:.3f}-"} |
| 128 | + store.put_dict(f"/cmd/corr/{node}", injection_cmd) |
| 129 | + store.put_dict(f"/cmd/corr/{partner}", injection_cmd) |
| 130 | + |
| 131 | + |
| 132 | + slack_msg( |
| 133 | + slack_cli, |
| 134 | + (f"Sending injection to beam {ew_beam} (EW) and {ns_beam} (NS) " |
| 135 | + f"with DM={DM:.1f} and SNR={effective_snr:.1f}" |
| 136 | + ) |
| 137 | + ) |
| 138 | + |
| 139 | + time.sleep(SLEEP_SEC) |
| 140 | + |
| 141 | +if __name__ == "__main__": |
| 142 | + run() |
0 commit comments