HotspotPopulationAnalysis/run_hotspot_search.py at main · chiarabellenghi/HotspotPopulationAnalysis · GitHub

1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
import argparse
import numpy as np
import os
from glob import glob
from data_manager import SkyLLH_scan
from code_utility import print_current_settings, save_pickle

p = argparse.ArgumentParser(
    description="Get and save all hotspots from many background skymaps. "
    "Possible choises: save trials for background parametrization "
    "(list of LocalHotSpotsList objects); save trials for background "
    "trials generation (list of LocalHotSpotsList objects to generate the "
    "backgorund TS distribution + array of all p-values above the minimum "
    "-log10(p-value) threshold from all maps). The two options are splitted "
    "because in principle the background parametrization and trials should be "
    "generated from two independent sets of skymaps.",
    formatter_class=argparse.RawDescriptionHelpFormatter,
)
p.add_argument(
    "--skyscan_dir",
    type=str,
    default="/data/user/cbellenghi/hpa/data/background_skyscans/for_trials_generation",
    help="The path to the folder containing the background healpy skymaps. ",
)
p.add_argument(
    "--output_dir",
    type=str,
    default="./data/background_hotspots",
    help="Path to the folder to store the result.",
)
p.add_argument(
    "--nside",
    type=int,
    default=256,
    help="Healpy `nside` parameter defining the resolution of the skymaps. "
    "Default: 256.",
)
p.add_argument(
    "--log10p_threshold",
    type=float,
    default=2.0,
    help="The significance threshold defining a hotspot in -log10(p-value). "
    "Default: 2.0",
)
p.add_argument(
    "--psi_min",
    type=float,
    default=1.0,
    help="Minimum angular distance between two warm spots (in degrees). "
    "Default: 1.0",
)
p.add_argument(
    "--for_parametrization",
    action="store_true",
    default=False,
    help="Run the search on many background skymaps and store the output in "
    "a list of hotspots per skymap. Default: False",
)

args = vars(p.parse_args())
print_current_settings(args)

filenames = sorted(glob(os.path.join(args["skyscan_dir"], "*")))

hotspots = list()
if args["for_parametrization"]:
    outfile = os.path.join(args["output_dir"], "hs_for_parametrization.pkl")
else:
    bg_pool = None
    outfile = os.path.join(args["output_dir"], "hs_for_bg_ts_distribution.pkl")
for fname in filenames:
    print("Searching hotspots in", os.path.basename(fname))
    skyscan = SkyLLH_scan(path=fname, nside=args["nside"])
    skyscan.prepare_full_scan()
    spots = skyscan.get_hotspots(args["log10p_threshold"], args["psi_min"])
    hotspots.append(spots)
    if not args["for_parametrization"]:
        if bg_pool is None:
            bg_pool = spots["logpVal"]
        else:
            bg_pool = np.concatenate((bg_pool, spots["logpVal"]))

# Make dir if does not exist
os.makedirs(args["output_dir"], exist_ok=True)

# Save the hotspots list
save_pickle(outfile, hotspots)
if not args["for_parametrization"]:
    outfile = os.path.join(args["output_dir"], "bg_hotspots_pool.npy")
    print("\nSaving background hotspot pool in", outfile)
    np.save(outfile, bg_pool)