def test_run_tessphot_mpi(PRIVATE_INPUT_DIR):

    stars_to_test = ([260795451, 'tpf'], [260795451, 'ffi'])

    with TaskManager(PRIVATE_INPUT_DIR) as tm:
        tm.cursor.execute("UPDATE todolist SET status=1;")
        for starid, datasource in stars_to_test:
            tm.cursor.execute(
                "UPDATE todolist SET status=NULL WHERE starid=? AND datasource=?;",
                [starid, datasource])
        tm.conn.commit()
        tm.cursor.execute(
            "SELECT COUNT(*) FROM todolist WHERE status IS NULL;")
        num = tm.cursor.fetchone()[0]

    print(num)
    assert num == len(stars_to_test)

    out, err, exitcode = capture_cli(
        'run_tessphot_mpi.py',
        mpiexec=True,
        params=['--debug', '--version=0', PRIVATE_INPUT_DIR])

    assert ' - INFO - %d tasks to be run' % num in err
    assert ' - INFO - Master starting with 1 workers' in err
    assert ' - DEBUG - Got data from worker 1: {' in err
    assert ' - INFO - Worker 1 exited.' in err
    assert ' - INFO - Master finishing' in err
    assert exitcode == 0
Beispiel #2
0
def test_get_tasks(PRIVATE_TODO_FILE):
    with TaskManager(PRIVATE_TODO_FILE, overwrite=True) as tm:
        task = tm.get_task(starid=267211065)
        assert task['priority'] == 1

        # Call with non-existing starid:
        task = tm.get_task(starid=-1234567890)
        assert task is None
Beispiel #3
0
def test_taskmanager_no_more_tasks(PRIVATE_TODO_FILE):
    with TaskManager(PRIVATE_TODO_FILE) as tm:
        # Set all the tasks as completed:
        tm.cursor.execute("UPDATE todolist SET status=1;")
        tm.conn.commit()

        # When we now ask for a new task, there shouldn't be any:
        assert tm.get_task() is None
        assert tm.get_random_task() is None
        assert tm.get_number_tasks() == 0
Beispiel #4
0
def test_taskmanager_skip_targets2(PRIVATE_TODO_FILE):
    with TaskManager(PRIVATE_TODO_FILE) as tm:

        # Start task with a random task:
        task = tm.get_task(starid=261522674, datasource='ffi')  # Tmag = 14.574
        print(task)

        # Make a fake result we can save;
        tm.start_task(task['priority'])
        result = task.copy()
        result['status'] = STATUS.OK
        result['time'] = 6.14
        result['worker_wait_time'] = 2.0
        result['method_used'] = 'aperture'
        result['details'] = {
            'skip_targets': [267211065]  # Tmag = 2.216
        }

        tm.save_result(result)

        # This time the processed target (the faint one) should end up marked as SKIPPED:
        tm.cursor.execute(
            "SELECT * FROM todolist LEFT JOIN diagnostics ON todolist.priority=diagnostics.priority WHERE todolist.priority=?;",
            [task['priority']])
        row = tm.cursor.fetchall()
        assert len(row) == 1
        row = row[0]
        print(dict(row))
        assert row['priority'] == task['priority']
        assert row['starid'] == task['starid']
        assert row['sector'] == task['sector']
        assert row['camera'] == task['camera']
        assert row['ccd'] == task['ccd']
        assert row['elaptime'] == 6.14
        assert row['method_used'] == 'aperture'
        assert row['status'] == STATUS.SKIPPED.value

        # And the bright target should not have STATUS set, so it can be processed later on:
        tm.cursor.execute(
            "SELECT * FROM todolist LEFT JOIN diagnostics ON todolist.priority=diagnostics.priority WHERE starid=267211065 AND datasource='ffi';"
        )
        row3 = tm.cursor.fetchall()
        assert len(row3) == 1
        row3 = row3[0]
        print(dict(row3))
        assert row3['status'] is None

        # There should now be exactly one entry in the photometry_skipped table:
        tm.cursor.execute("SELECT * FROM photometry_skipped;")
        row = tm.cursor.fetchall()
        assert len(row) == 1
        row = row[0]
        print(dict(row))
        assert row['priority'] == task['priority']
        assert row['skipped_by'] == row3['priority']
Beispiel #5
0
def test_taskmanager_skip_targets1(PRIVATE_TODO_FILE):
    with TaskManager(PRIVATE_TODO_FILE, overwrite=True, cleanup=True) as tm:

        # Start task with a random task:
        task = tm.get_task(starid=267211065, datasource='ffi')  # Tmag = 2.216
        #task = tm.get_task(starid=261522674) # Tmag = 14.574
        print(task)

        # Make a fake result we can save:
        tm.start_task(task['priority'])
        result = task.copy()
        result['status'] = STATUS.OK
        result['time'] = 6.14
        result['worker_wait_time'] = 2.0
        result['method_used'] = 'aperture'
        result['details'] = {
            'skip_targets': [261522674]  # Tmag = 14.574
        }

        tm.save_result(result)

        tm.cursor.execute(
            "SELECT * FROM todolist LEFT JOIN diagnostics ON todolist.priority=diagnostics.priority WHERE todolist.priority=?;",
            [task['priority']])
        row = tm.cursor.fetchall()
        assert len(row) == 1
        row = row[0]
        print(dict(row))
        assert row['priority'] == task['priority']
        assert row['starid'] == task['starid']
        assert row['sector'] == task['sector']
        assert row['camera'] == task['camera']
        assert row['ccd'] == task['ccd']
        assert row['elaptime'] == 6.14
        assert row['method_used'] == 'aperture'
        assert row['status'] == STATUS.OK.value

        tm.cursor.execute(
            "SELECT * FROM todolist LEFT JOIN diagnostics ON todolist.priority=diagnostics.priority WHERE starid=261522674 AND datasource='ffi';"
        )
        row2 = tm.cursor.fetchall()
        assert len(row2) == 1
        row2 = row2[0]
        print(dict(row2))
        assert row2['status'] == STATUS.SKIPPED.value

        # There should now be exactly one entry in the photometry_skipped table:
        tm.cursor.execute("SELECT * FROM photometry_skipped;")
        row = tm.cursor.fetchall()
        assert len(row) == 1
        row = row[0]
        print(dict(row))
        assert row['priority'] == row2['priority']
        assert row['skipped_by'] == task['priority']
Beispiel #6
0
def test_taskmanager_constraints(PRIVATE_TODO_FILE):

    constraints = {'datasource': 'tpf', 'priority': 1}
    with TaskManager(PRIVATE_TODO_FILE,
                     overwrite=True,
                     cleanup_constraints=constraints) as tm:
        task = tm.get_task(**constraints)
        numtasks = tm.get_number_tasks(**constraints)
        print(task)
        assert task['starid'] == 267211065, "Task1 should be None"
        assert numtasks == 1, "Task1 search should give no results"

    constraints = {'datasource': 'tpf', 'priority': 1, 'camera': None}
    with TaskManager(PRIVATE_TODO_FILE,
                     overwrite=True,
                     cleanup_constraints=constraints) as tm:
        task2 = tm.get_task(**constraints)
        numtasks2 = tm.get_number_tasks(**constraints)
        print(task2)
        assert task2 == task, "Tasks should be identical"
        assert numtasks2 == 1, "Task2 search should give no results"

    constraints = {'datasource': 'ffi', 'priority': 2}
    with TaskManager(PRIVATE_TODO_FILE,
                     overwrite=True,
                     cleanup_constraints=constraints) as tm:
        task = tm.get_task(**constraints)
        numtasks = tm.get_number_tasks(**constraints)
        print(task)
        assert task['priority'] == 2, "Task2 should be #2"
        assert task['datasource'] == 'ffi'
        assert task['camera'] == 3
        assert task['ccd'] == 2
        assert numtasks == 1, "Priority search should give one results"

    constraints = {'datasource': 'ffi', 'priority': 2, 'camera': 3, 'ccd': 2}
    with TaskManager(PRIVATE_TODO_FILE,
                     overwrite=True,
                     cleanup_constraints=constraints) as tm:
        task3 = tm.get_task(**constraints)
        numtasks3 = tm.get_number_tasks(**constraints)
        print(task3)
        assert task3 == task, "Tasks should be identical"
        assert numtasks3 == 1, "Task3 search should give one results"

    constraints = ['priority=17']
    with TaskManager(PRIVATE_TODO_FILE, cleanup_constraints=constraints) as tm:
        task4 = tm.get_task(priority=17)
        numtasks4 = tm.get_number_tasks(priority=17)
        print(task4)
        assert task4['priority'] == 17, "Task4 should be #17"
        assert numtasks4 == 1, "Priority search should give one results"

    constraints = {'starid': 267211065}
    with TaskManager(PRIVATE_TODO_FILE, cleanup_constraints=constraints) as tm:
        numtasks5 = tm.get_number_tasks(**constraints)
        assert numtasks5 == 2
        task5 = tm.get_task(**constraints)
        assert task5['priority'] == 1
Beispiel #7
0
def test_taskmanager_skip_targets_secondary1(PRIVATE_TODO_FILE):
    with TaskManager(PRIVATE_TODO_FILE) as tm:

        # Start task with a random task:
        task = tm.get_task(starid=261522674, datasource='ffi')  # Tmag = 14.574
        print(task)

        # Make a fake result we can save;
        tm.start_task(task['priority'])
        result = task.copy()
        result['datasource'] = 'tpf:267211065'
        result['status'] = STATUS.OK
        result['time'] = 6.14
        result['worker_wait_time'] = 2.0
        result['method_used'] = 'aperture'
        result['details'] = {
            'skip_targets': [267211065]  # Tmag = 2.216
        }

        tm.save_result(result)

        # For a secondary target with the primary in skipped_targets, the secondary should
        # end up marked as skipped with a warning in the errors column:
        tm.cursor.execute(
            "SELECT * FROM todolist LEFT JOIN diagnostics ON todolist.priority=diagnostics.priority WHERE todolist.priority=?;",
            [task['priority']])
        row = tm.cursor.fetchall()
        assert len(row) == 1
        row = row[0]
        print(dict(row))
        assert row['priority'] == task['priority']
        assert row['starid'] == task['starid']
        assert row['sector'] == task['sector']
        assert row['camera'] == task['camera']
        assert row['ccd'] == task['ccd']
        assert row['elaptime'] == 6.14
        assert row['method_used'] == 'aperture'
        assert row['status'] == STATUS.SKIPPED.value

        # There should now be exactly one entry in the photometry_skipped table:
        tm.cursor.execute("SELECT * FROM photometry_skipped;")
        row = tm.cursor.fetchall()
        assert len(row) == 1
        row = row[0]
        print(dict(row))
        assert row['priority'] == task['priority']
        assert row['skipped_by'] == 1
Beispiel #8
0
def test_taskmanager():
	"""Test of background estimator"""

	# Load the first image in the input directory:
	INPUT_DIR = os.path.join(os.path.dirname(__file__), 'input')

	# Find the shape of the original image:
	with TaskManager(INPUT_DIR, overwrite=True) as tm:
		# Get the number of tasks:
		numtasks = tm.get_number_tasks()
		print(numtasks)
		assert(numtasks == 168642)

		# Get the first task in the TODO file:
		task1 = tm.get_task()
		print(task1)

		# Check that it contains what we know it should:
		# The first priority in the TODO file is the following:
		assert(task1['priority'] == 1)
		assert(task1['starid'] == 267211065)
		assert(task1['camera'] == 3)
		assert(task1['ccd'] == 2)
		assert(task1['datasource'] == 'tpf')
		assert(task1['sector'] == 1)

		# Start task with priority=1:
		tm.start_task(1)

		# Get the next task, which should be the one with priority=2:
		task2 = tm.get_task()
		print(task2)

		assert(task2['priority'] == 2)
		assert(task2['starid'] == 267211065)
		assert(task2['camera'] == 3)
		assert(task2['ccd'] == 2)
		assert(task2['datasource'] == 'ffi')
		assert(task2['sector'] == 1)

		# Check that the status did actually change in the todolist:
		tm.cursor.execute("SELECT status FROM todolist WHERE priority=1;")
		task1_status = tm.cursor.fetchone()['status']
		print(task1_status)

		assert(task1_status == STATUS.STARTED.value)
def test_taskmanager():
    """Test of background estimator"""

    # Load the first image in the input directory:
    INPUT_DIR = os.path.join(os.path.dirname(__file__), 'input')

    # Find the shape of the original image:
    with TaskManager(INPUT_DIR) as tm:
        # Get the first task in the TODO file:
        task = tm.get_task()
        print(task)

        # Check that it contains what we know it should:
        # The first priority in the TODO file is the following:
        assert (task['priority'] == 1)
        assert (task['starid'] == 284853659)
        assert (task['camera'] == 2)
        assert (task['ccd'] == 2)
Beispiel #10
0
def test_taskmanager(PRIVATE_TODO_FILE):
    """Test of TaskManager"""
    with TaskManager(PRIVATE_TODO_FILE, overwrite=True) as tm:
        # Get the number of tasks:
        numtasks = tm.get_number_tasks()
        print(numtasks)
        assert (numtasks == 168642)

        # Get the first task in the TODO file:
        task1 = tm.get_task()
        print(task1)

        # Check that it contains what we know it should:
        # The first priority in the TODO file is the following:
        assert (task1['priority'] == 1)
        assert (task1['starid'] == 267211065)
        assert (task1['camera'] == 3)
        assert (task1['ccd'] == 2)
        assert (task1['datasource'] == 'tpf')
        assert (task1['sector'] == 1)

        # Start task with priority=1:
        tm.start_task(1)

        # Get the next task, which should be the one with priority=2:
        task2 = tm.get_task()
        print(task2)

        assert (task2['priority'] == 2)
        assert (task2['starid'] == 267211065)
        assert (task2['camera'] == 3)
        assert (task2['ccd'] == 2)
        assert (task2['datasource'] == 'ffi')
        assert (task2['sector'] == 1)

        # Check that the status did actually change in the todolist:
        tm.cursor.execute("SELECT status FROM todolist WHERE priority=1;")
        task1_status = tm.cursor.fetchone()['status']
        print(task1_status)

        assert (task1_status == STATUS.STARTED.value)
Beispiel #11
0
	logger = logging.getLogger(__name__)
	logger.addHandler(console)
	logger.setLevel(logging_level)

	inp = args.input_folder
	input_folder = os.path.abspath(os.path.dirname(inp))
	logger.info("Renaming files in %s", input_folder)
	if not os.path.isdir(input_folder):
		parser.error("Input folder does not exist")

	tqdm_settings = {'disable': not logger.isEnabledFor(logging.INFO)}

	regex_old = re.compile(r'^tess(\d+)-s(\d+)-c(\d+)-dr(\d+)-v(\d+)-tasoc_lc\.fits\.gz$')
	regex_new = re.compile(r'^tess(\d+)-s(\d+)-[1-4]-[1-4]-c(\d+)-dr(\d+)-v(\d+)-tasoc_lc\.fits\.gz$')

	with TaskManager(input_folder) as tm:

		tm.cursor.execute("SELECT todolist.priority,lightcurve,camera,ccd FROM todolist INNER JOIN diagnostics ON todolist.priority=diagnostics.priority WHERE lightcurve IS NOT NULL;")
		results = tm.cursor.fetchall()

		try:
			for k, row in enumerate(tqdm(results, **tqdm_settings)):

				fpath = os.path.join(input_folder, row['lightcurve'])
				if not os.path.isfile(fpath):
					logger.error("File not found: '%s'", row['lightcurve'])
					continue

				bname = os.path.basename(fpath)
				m = regex_old.match(bname)
				if m:
Beispiel #12
0
def main():
	# Parse command line arguments:
	parser = argparse.ArgumentParser(description='Run TESS Photometry in parallel using MPI.')
	#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('-p', '--plot', help='Save plots when running.', action='store_true')
	args = parser.parse_args()

	# Get paths to input and output files from environment variables:
	input_folder = os.environ.get('TESSPHOT_INPUT', os.path.join(os.path.dirname(__file__), 'tests', 'input'))
	output_folder = os.environ.get('TESSPHOT_OUTPUT', os.path.abspath('.'))
	todo_file = os.path.join(input_folder, 'todo.sqlite')

	# Define MPI message tags
	tags = enum.IntEnum('tags', ('READY', 'DONE', 'EXIT', 'START'))

	# Initializations and preliminaries
	comm = MPI.COMM_WORLD   # get MPI communicator object
	size = comm.size        # total number of processes
	rank = comm.rank        # rank of this process
	status = MPI.Status()   # get MPI status object

	if rank == 0:
		# Master process executes code below
		from photometry import TaskManager

		try:
			with TaskManager(todo_file, cleanup=True, summary=os.path.join(output_folder, 'summary.json')) as tm:
				# Get list of tasks:
				numtasks = tm.get_number_tasks()
				tm.logger.info("%d tasks to be run", numtasks)

				# Start the master loop that will assign tasks
				# to the workers:
				num_workers = size - 1
				closed_workers = 0
				tm.logger.info("Master starting with %d workers", num_workers)
				while closed_workers < num_workers:
					# Ask workers for information:
					data = comm.recv(source=MPI.ANY_SOURCE, tag=MPI.ANY_TAG, status=status)
					source = status.Get_source()
					tag = status.Get_tag()

					if tag == tags.DONE:
						# The worker is done with a task
						tm.logger.info("Got data from worker %d: %s", source, data)
						tm.save_result(data)

					if tag in (tags.DONE, tags.READY):
						# Worker is ready, so send it a task
						task = tm.get_task()
						if task:
							task_index = task['priority']
							tm.start_task(task_index)
							comm.send(task, dest=source, tag=tags.START)
							tm.logger.info("Sending task %d to worker %d", task_index, source)
						else:
							comm.send(None, dest=source, tag=tags.EXIT)

					elif tag == tags.EXIT:
						# The worker has exited
						tm.logger.info("Worker %d exited.", source)
						closed_workers += 1

					else:
						# This should never happen, but just to
						# make sure we don't run into an infinite loop:
						raise Exception("Master received an unknown tag: '{0}'".format(tag))

				tm.logger.info("Master finishing")

		except:
			# If something fails in the master
			print(traceback.format_exc().strip())
			comm.Abort(1)

	else:
		# Worker processes execute code below
		from photometry import tessphot
		from timeit import default_timer

		# Configure logging within photometry:
		formatter = logging.Formatter('%(asctime)s - %(levelname)s - %(message)s')
		console = logging.StreamHandler()
		console.setFormatter(formatter)
		logger = logging.getLogger('photometry')
		logger.addHandler(console)
		logger.setLevel(logging.WARNING)

		try:
			# Send signal that we are ready for task:
			comm.send(None, dest=0, tag=tags.READY)

			while True:
				# Receive a task from the master:
				task = comm.recv(source=0, tag=MPI.ANY_TAG, status=status)
				tag = status.Get_tag()

				if tag == tags.START:
					# Do the work here
					result = task.copy()
					del task['priority'], task['tmag']

					t1 = default_timer()
					pho = tessphot(input_folder=input_folder, output_folder=output_folder, plot=args.plot, **task)
					t2 = default_timer()

					# Construct result message:
					result.update({
						'status': pho.status,
						'time': t2 - t1,
						'details': pho._details
					})

					# Send the result back to the master:
					comm.send(result, dest=0, tag=tags.DONE)

					# Attempt some cleanup:
					# TODO: Is this even needed?
					del pho, task, result

				elif tag == tags.EXIT:
					# We were told to EXIT, so lets do that
					break

				else:
					# This should never happen, but just to
					# make sure we don't run into an infinite loop:
					raise Exception("Worker received an unknown tag: '{0}'".format(tag))

		except:
			logger.exception("Something failed in worker")

		finally:
			comm.send(None, dest=0, tag=tags.EXIT)
Beispiel #13
0
def test_taskmanager_invalid():
    with pytest.raises(FileNotFoundError):
        TaskManager(os.path.join(INPUT_DIR, 'does-not-exists'))
Beispiel #14
0
def test_taskmanager_summary(PRIVATE_TODO_FILE):
    with tempfile.TemporaryDirectory() as tmpdir:
        summary_file = os.path.join(tmpdir, 'summary.json')
        with TaskManager(PRIVATE_TODO_FILE,
                         overwrite=True,
                         summary=summary_file,
                         summary_interval=2) as tm:
            # Load the summary file:
            with open(summary_file, 'r') as fid:
                j = json.load(fid)

            assert tm.summary_counter == 0  # Counter should start at zero

            # Everytning should be really empty:
            print(j)
            assert j['numtasks'] == 168642
            assert j['UNKNOWN'] == 0
            assert j['OK'] == 0
            assert j['ERROR'] == 0
            assert j['WARNING'] == 0
            assert j['ABORT'] == 0
            assert j['STARTED'] == 0
            assert j['SKIPPED'] == 0
            assert j['tasks_run'] == 0
            assert j['slurm_jobid'] is None
            assert j['last_error'] is None
            assert j['mean_elaptime'] is None

            # Start task with priority=1:
            task = tm.get_random_task()
            print(task)
            tm.start_task(task['priority'])
            tm.write_summary()

            with open(summary_file, 'r') as fid:
                j = json.load(fid)

            print(j)
            assert j['numtasks'] == 168642
            assert j['UNKNOWN'] == 0
            assert j['OK'] == 0
            assert j['ERROR'] == 0
            assert j['WARNING'] == 0
            assert j['ABORT'] == 0
            assert j['STARTED'] == 1
            assert j['SKIPPED'] == 0
            assert j['tasks_run'] == 0
            assert j['slurm_jobid'] is None
            assert j['last_error'] is None
            assert j['mean_elaptime'] is None

            # Make a fake result we can save;
            result = task.copy()
            result['status'] = STATUS.OK
            result['time'] = 3.14
            result['worker_wait_time'] = 1.0
            result['method_used'] = 'aperture'

            # Save the result:
            tm.save_result(result)
            assert tm.summary_counter == 1  # We saved once, so counter should have gone up one
            tm.write_summary()

            # Check that the diagnostics were updated:
            tm.cursor.execute("SELECT * FROM diagnostics WHERE priority=?;",
                              [task['priority']])
            row = tm.cursor.fetchone()
            print(dict(row))
            assert row['priority'] == task['priority']
            assert 'starid' not in row  # It should no longer be there (after version 5.0)
            assert row['elaptime'] == 3.14
            assert row['method_used'] == 'aperture'

            # Load the summary file after "running the task":
            with open(summary_file, 'r') as fid:
                j = json.load(fid)

            print(j)
            assert j['numtasks'] == 168642
            assert j['UNKNOWN'] == 0
            assert j['OK'] == 1
            assert j['ERROR'] == 0
            assert j['WARNING'] == 0
            assert j['ABORT'] == 0
            assert j['STARTED'] == 0
            assert j['SKIPPED'] == 0
            assert j['tasks_run'] == 1
            assert j['slurm_jobid'] is None
            assert j['last_error'] is None
            assert j['mean_elaptime'] == 3.14
            assert j['mean_worker_waittime'] == 1.0

            task = tm.get_random_task()
            tm.start_task(task['priority'])

            # Make a fake result we can save;
            result = task.copy()
            result['status'] = STATUS.ERROR
            result['time'] = 6.14
            result['worker_wait_time'] = 2.0
            result['method_used'] = 'halo'
            result['details'] = {'errors': ['dummy error 1', 'dummy error 2']}

            # Save the result:
            tm.save_result(result)
            assert tm.summary_counter == 0  # We saved again, so summary_counter should be zero
            tm.write_summary()

            # Check that the diagnostics were updated:
            tm.cursor.execute("SELECT * FROM diagnostics WHERE priority=?;",
                              [task['priority']])
            row = tm.cursor.fetchone()
            print(dict(row))
            assert row['priority'] == task['priority']
            assert 'starid' not in row  # It should no longer be there (after version 5.0)
            assert row['elaptime'] == 6.14
            assert row['method_used'] == 'halo'
            assert row['errors'] == "dummy error 1\ndummy error 2"

            # Load the summary file after "running the task":
            with open(summary_file, 'r') as fid:
                j = json.load(fid)

            print(j)
            assert j['numtasks'] == 168642
            assert j['UNKNOWN'] == 0
            assert j['OK'] == 1
            assert j['ERROR'] == 1
            assert j['WARNING'] == 0
            assert j['ABORT'] == 0
            assert j['STARTED'] == 0
            assert j['SKIPPED'] == 0
            assert j['tasks_run'] == 2
            assert j['slurm_jobid'] is None
            assert j['last_error'] == "dummy error 1\ndummy error 2"
            assert j['mean_elaptime'] == 3.44
            assert j['mean_worker_waittime'] == 1.1
Beispiel #15
0
def test_taskmanager_constraints_invalid(PRIVATE_TODO_FILE):
    with pytest.raises(ValueError) as e:
        with TaskManager(PRIVATE_TODO_FILE, cleanup_constraints='invalid'):
            pass
    assert str(e.value) == 'cleanup_constraints should be dict or list'
Beispiel #16
0
def main():
    # Parse command line arguments:
    parser = argparse.ArgumentParser(
        description='Run TESS Photometry pipeline on single star.')
    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('-o',
                        '--overwrite',
                        help='Overwrite existing results.',
                        action='store_true')
    parser.add_argument('-p',
                        '--plot',
                        help='Save plots when running.',
                        action='store_true')
    parser.add_argument(
        '-t',
        '--test',
        help='Use test data and ignore TESSPHOT_INPUT environment variable.',
        action='store_true')
    parser.add_argument('-m',
                        '--method',
                        choices=('aperture', 'psf', 'linpsf', 'halo'),
                        default=None,
                        help='Photometric method to use.')

    group = parser.add_argument_group('Filter which targets to run')
    group.add_argument(
        '--all',
        help=
        'Run all stars, one by one. Please consider using the MPI program instead.',
        action='store_true')
    group.add_argument('-r',
                       '--random',
                       help='Run on random target from TODO-list.',
                       action='store_true')
    group.add_argument('--priority',
                       type=int,
                       help='Priority of target.',
                       nargs='?',
                       default=None)
    group.add_argument('--starid',
                       type=int,
                       help='TIC identifier of target.',
                       nargs='?',
                       default=None)
    group.add_argument(
        '--datasource',
        type=str,
        choices=('ffi', 'tpf'),
        default=None,
        help='Data-source to load.',
    )
    group.add_argument('--camera',
                       type=int,
                       choices=(1, 2, 3, 4),
                       default=None,
                       help='TESS Camera. Default is to run all cameras.')
    group.add_argument('--ccd',
                       type=int,
                       choices=(1, 2, 3, 4),
                       default=None,
                       help='TESS CCD. Default is to run all CCDs.')

    parser.add_argument('--version',
                        type=int,
                        help='Data release number to store in output files.',
                        nargs='?',
                        default=None)
    parser.add_argument('--output',
                        type=str,
                        help='Directory to put lightcurves into.',
                        nargs='?',
                        default=None)
    parser.add_argument('input_folder',
                        type=str,
                        help='Directory to create catalog files in.',
                        nargs='?',
                        default=None)
    args = parser.parse_args()

    # Make sure at least one setting is given:
    if not args.all and args.starid is None and args.priority is None and not args.random:
        parser.error("Please select either a specific STARID or RANDOM.")

    # 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__)
    logger.addHandler(console)
    logger.setLevel(logging_level)
    logger_parent = logging.getLogger('photometry')
    logger_parent.addHandler(console)
    logger_parent.setLevel(logging_level)

    # Get input and output folder from environment variables:
    test_folder = os.path.abspath(
        os.path.join(os.path.dirname(__file__), 'tests', 'input'))
    if args.test:
        input_folder = test_folder
    elif args.input_folder:
        input_folder = args.input_folder
    else:
        input_folder = os.environ.get('TESSPHOT_INPUT', test_folder)

    if os.path.isfile(input_folder):
        input_folder = os.path.dirname(input_folder)

    if args.output:
        output_folder = args.output
    else:
        output_folder = os.environ.get(
            'TESSPHOT_OUTPUT', os.path.join(input_folder, 'lightcurves'))

    logger.info("Loading input data from '%s'", input_folder)
    logger.info("Putting output data in '%s'", output_folder)

    # Constraints on which targets to process:
    constraints = {
        'camera': args.camera,
        'ccd': args.ccd,
        'datasource': args.datasource,
        'starid': args.starid,
        'priority': args.priority
    }

    # Create partial function of tessphot, setting the common keywords:
    f = functools.partial(tessphot,
                          input_folder=input_folder,
                          output_folder=output_folder,
                          plot=args.plot,
                          version=args.version)

    # Run the program:
    with TaskManager(input_folder,
                     overwrite=args.overwrite,
                     cleanup_constraints=constraints) as tm:
        while True:
            if args.random:
                task = tm.get_random_task()
            else:
                task = tm.get_task(**constraints)

            if task is None:
                parser.error("No task found matching constraints.")
                break

            # If specific method provided, overwrite whatever
            # was defined in the TODO-file:
            if args.method:
                task['method'] = args.method

            result = task.copy()
            del task['priority'], task['tmag']

            t1 = default_timer()
            pho = f(**task)
            t2 = default_timer()

            # Construct result message:
            result.update({
                'status': pho.status,
                'method_used': pho.method,
                'time': t2 - t1,
                'details': pho._details
            })
            tm.save_result(result)

            if not args.all:
                break