#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
Run preparation of CBVs for single or several CBV-areas.
.. codeauthor:: Mikkel N. Lund <mikkelnl@phys.au.dk>
.. codeauthor:: Rasmus Handberg <rasmush@phys.au.dk>
"""
import argparse
import os
import logging
from functools import partial
import multiprocessing
import corrections
from corrections.taskmanager import _build_constraints
from corrections.utilities import CadenceType
#--------------------------------------------------------------------------------------------------
[docs]
def main():
# Parse command line arguments:
parser = argparse.ArgumentParser(description='Run preparation of CBVs for single or several CBV-areas.', formatter_class=argparse.ArgumentDefaultsHelpFormatter)
parser.add_argument('-d', '--debug', help='Print debug messages.', action='store_true')
parser.add_argument('-q', '--quiet', help='Only report warnings and errors.', action='store_true')
parser.add_argument('-ip', '--iniplot', help='Make Initial fitting plots.', action='store_true')
parser.add_argument('--threads', type=int, default=None, help="Number of parallel threads to use. If not specified, all available CPUs will be used.")
parser.add_argument('--output', type=str, default=None, help="Directory where output CBVs will be saved.")
group = parser.add_argument_group('Specifying CBVs to calculate')
group.add_argument('--sector', type=int, default=None, help='Sector to create CBVs for.')
group.add_argument('--cadence', type=CadenceType, default='ffi', choices=('ffi', 1800, 600, 120, 20), help='Cadence for the creation of CBVs.')
group.add_argument('--camera', type=int, choices=(1,2,3,4), action='append', default=None, help='TESS Camera. Default is to run all cameras.')
group.add_argument('--ccd', type=int, choices=(1,2,3,4), action='append', default=None, help='TESS CCD. Default is to run all CCDs.')
group.add_argument('-a', '--area', type=int, action='append', default=None, help='Single CBV_area for which to prepare photometry. Default is to run all areas.')
group = parser.add_argument_group('Settings')
group.add_argument('--ncbv', type=int, default=16, help='Number of CBVs to compute')
group.add_argument('--corr', type=float, default=0.5, help='Fraction of most correlated stars to use for CBVs.')
group.add_argument('--snr', type=float, default=5, help='SNR (dB) for selection of CBVs.')
group.add_argument('--el', type=float, default=-0.5, help='Entropy limit for discarting star contribution to CBV.')
parser.add_argument('--version', type=int, required=True, help='Data release number to store in output files.')
parser.add_argument('input_folder', type=str, nargs='?', default=None, help='Directory to create catalog files in.')
args = parser.parse_args()
# Set logging level:
logging_level = logging.INFO
if args.quiet:
logging_level = logging.WARNING
elif args.debug:
logging_level = logging.DEBUG
# Setup logging:
formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
console = logging.StreamHandler()
console.setFormatter(formatter)
logger = logging.getLogger(__name__)
if not logger.hasHandlers():
logger.addHandler(console)
logger.setLevel(logging_level)
logger_parent = logging.getLogger('corrections')
if not logger_parent.hasHandlers():
logger_parent.addHandler(console)
logger_parent.setLevel(logging_level)
# Parse the input folder:
input_folder = args.input_folder
if input_folder is None:
input_folder = os.environ.get('TESSCORR_INPUT')
if input_folder is None or not os.path.exists(input_folder):
parser.error("Invalid input folder")
logger.info("Loading input data from '%s'", input_folder)
# Build list of constraints:
constraints = _build_constraints(
sector=args.sector,
cadence=args.cadence,
camera=args.camera,
ccd=args.ccd,
cbv_area=args.area,
return_list=True)
if not constraints:
constraints = None
# Invoke the TaskManager to ensure that the input TODO-file has the correct columns
# and indicies, which is automatically created by the TaskManager init function.
with corrections.TaskManager(input_folder, cleanup=False, cleanup_constraints=constraints):
pass
# Use the BaseCorrector to search the database for which CBV_AREAS to run:
with corrections.BaseCorrector(input_folder) as bc:
# Search for valid areas:
cbv_areas = [(row['sector'], row['cbv_area']) for row in bc.search_database(select=['sector','cbv_area'], distinct=True, search=constraints, order_by=['sector', 'cbv_area'])]
logger.debug("CBV areas: %s", cbv_areas)
# Stop if there are no CBV-Areas to process:
if not cbv_areas:
logger.info("No CBV-areas found to be processed.")
return
# Number of threads to run in parallel:
threads = args.threads
if not threads:
threads = int(os.environ.get('SLURM_CPUS_PER_TASK', multiprocessing.cpu_count()))
threads = min(threads, len(cbv_areas))
logger.info("Using %d processes.", threads)
# Create wrapper function which only takes a single cbv_area as input:
create_cbv_wrapper = partial(corrections.create_cbv,
input_folder=input_folder,
output_folder=args.output,
cadence=args.cadence,
version=args.version,
threshold_correlation=args.corr,
threshold_snrtest=args.snr,
ncbv=args.ncbv,
threshold_entropy=args.el,
ip=args.iniplot)
# Run the preparation:
if threads > 1:
# Disable printing info messages from the parent function.
# It is going to be all jumbled up anyway.
logger_parent.setLevel(logging.WARNING)
# There is more than one area to process, so let's start
# a process pool and process them in parallel:
with multiprocessing.Pool(threads) as pool:
pool.map(create_cbv_wrapper, cbv_areas)
else:
# Only a single area to process, so let's not bother with
# starting subprocesses, but simply run it directly:
for cbv_area in cbv_areas:
create_cbv_wrapper(cbv_area)
#--------------------------------------------------------------------------------------------------
if __name__ == '__main__':
main()