calculate_cp.py 4.79 KB
Newer Older
1 2 3
#!/usr/bin/env python3
# -*- coding: utf-8 -*-

4
import dill as pickle
5 6 7 8 9 10 11
import matplotlib.pyplot as plt
import numpy as np
import astropy.units as u
from nikamap import NikaMap
#from completness import *
import myfunctions as mfct
from time import clock
LUSTIG Peter's avatar
LUSTIG Peter committed
12 13 14
from multiprocess import Pool
from functools import partial
import sys
15
from astropy.table import vstack
LUSTIG Peter's avatar
LUSTIG Peter committed
16

17 18 19
#import sys
import completness as cp
import os
LUSTIG Peter's avatar
LUSTIG Peter committed
20
import matplotlib.pyplot as plt
21 22 23 24 25 26

plt.close('all')

flux_bins = 2
flux_range = [0.1, 10]
threshold_bins = 5
LUSTIG Peter's avatar
LUSTIG Peter committed
27 28
threshold_range = [2, 10]
#threshold_range = [3, 5]
29

30
nsim = 10
LUSTIG Peter's avatar
LUSTIG Peter committed
31 32 33 34 35
flux = 10 * u.mJy
noisefile = "../../N2CLS/ForMarseille/HLS091828_common_mode_kids_out/map_JK.fits"
method="unknown"
outfile = None
outdir = "results/"
36

37
ncore = 4
38

LUSTIG Peter's avatar
LUSTIG Peter committed
39 40
'''
import argparse
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
desc = "Place false sources with a defined flux on Jackknife maps, detect sources and calculate the purity and completness of the detection."
parser = argparse.ArgumentParser(description=desc)
#parser.add_argument("-n", "--noisefile", action="store",
#                    dest="noisefile",
#                    help="Name of the Noise Containing Input File")
parser.add_argument(action="store",
                    dest="noisefile",
                    help="Name of the Noise Containing Input File")
parser.add_argument("-s", "--simulations", action="store",
                    dest="nsim", type=int, default=10,
                    help="Number of Simulations")
parser.add_argument("-f", "--flux", action="store",
                    dest="flux", type=float, default=10.,
                    help="Flux of the Created Sources in mJy")
parser.add_argument("-m", "--method", action="store",
                    dest="method", default="unknown",
                    help="Name of the Reduction Method for Outfile Header")
parser.add_argument("-o", "--outfile", action="store",
                    dest="outfile", default=None,
                    help="Output File Name. If not specified: completness-method-flux_[flux]-nsim_[nsim].fits")
parser.add_argument("-d", "--outdir", action="store",
                    dest="outdir", default="",
                    help="Output directory. Default: This folder.")


args=parser.parse_args()
nsim = args.nsim
flux = args.flux * u.mJy
noisefile = args.noisefile
method=args.method
outfile = args.outfile
outdir = args.outdir
#flux bins
#flux range
#threshold bins
#threshold range
LUSTIG Peter's avatar
LUSTIG Peter committed
77 78
'''

79 80 81 82


############### Begin Calculation #############

LUSTIG Peter's avatar
LUSTIG Peter committed
83
sim_per_core = mfct.Distribute_Njobs(ncore, nsim)
84 85 86 87 88 89 90 91 92 93 94 95 96 97 98 99 100 101 102 103

data = NikaMap.read(noisefile)

fluxes = np.logspace(np.log10(flux_range[0]), np.log10(flux_range[1]), flux_bins)*u.mJy
fluxes_edges = np.logspace(np.log10(flux_range[0]), np.log10(flux_range[1]), flux_bins + 1)*u.mJy

threshold = np.linspace(threshold_range[0], threshold_range[1], threshold_bins)
threshold_edges = np.linspace(threshold_range[0], threshold_range[1], threshold_bins+1)

# Create 4D wcs
shape_4D, wcs_4D = cp.completness_purity_wcs(data.shape, data.wcs, bins=20,
                                          flux_range=flux_range, flux_bins=flux_bins, flux_log=True,
                                          threshold_range=threshold_range, threshold_bins=threshold_bins)

# Testing the lower edges
wcs_threshold = wcs_4D.sub([3])
assert np.all(np.abs(wcs_threshold.all_pix2world(np.arange(threshold_bins+1)-0.5, 0) - threshold_edges) < 1e-15)
wcs_flux = wcs_4D.sub([4])
assert np.all(np.abs(wcs_flux.all_pix2world(np.arange(flux_bins+1)-0.5, 0) - fluxes_edges.value) < 1e-13)

LUSTIG Peter's avatar
LUSTIG Peter committed
104
T0 = clock()
105

LUSTIG Peter's avatar
LUSTIG Peter committed
106 107 108 109 110 111 112 113 114 115 116 117 118 119 120 121 122
print('------------')
#data.plot()
#plt.show(block=True)
helpfunc = partial(cp.completness_purity_2, flux, **{"nsources":8**2, "within":(0, 1),
                              "shape":shape_4D[0:3],
                              "wcs":wcs_4D.sub([1, 2, 3]), "jk_map":data})


pool = Pool(ncore)
print(sim_per_core)
tmpres = pool.map(helpfunc, sim_per_core)
tmpres = list(zip(*tmpres))

all_comp = np.sum(tmpres[0], axis=0)
all_comp_n = np.sum(tmpres[1], axis=0)
all_pur = np.sum(tmpres[2], axis=0)
all_pur_n = np.sum(tmpres[3], axis=0)
123
sourceresult = vstack(tmpres[4])
LUSTIG Peter's avatar
LUSTIG Peter committed
124 125 126 127 128 129 130

print(all_comp.shape)
print(all_comp_n.shape)

with np.errstate(divide='ignore', invalid='ignore'):
    completness = all_comp / all_comp_n[..., np.newaxis]
    purity = all_pur / all_pur_n
131

LUSTIG Peter's avatar
LUSTIG Peter committed
132 133 134

if (method == "unknown" and "/" in noisefile):
    method = noisefile.split("/")[-2]
135
    
LUSTIG Peter's avatar
LUSTIG Peter committed
136 137 138 139 140 141
if outfile is None:
    outfile = "completness-" + method + "-flux_{:04.0f}-nsim_{:03}.fits".format(flux.value, nsim)

if outdir != "":
    if outdir[-1] != "/":
        outdir = outdir + "/"
142
    
LUSTIG Peter's avatar
LUSTIG Peter committed
143 144
    if not os.path.exists(outdir):
        os.makedirs(outdir)
145
    
LUSTIG Peter's avatar
LUSTIG Peter committed
146 147 148
    outfile = outdir + outfile


149
mfct.WriteResult(outfile, threshold, completness, purity, nsim, wcs_4D, method, flux, sourceresult)
LUSTIG Peter's avatar
LUSTIG Peter committed
150
cp.Plot_CompPur(completness, purity, threshold, nsim=nsim, savename=outfile[0:-5]+"-plot.png")
151

LUSTIG Peter's avatar
LUSTIG Peter committed
152 153
T = clock()-T0
print('Total calculation time: {:.0f}s'.format(T))