#!/usr/bin/env python3
"""Utilities for plotting individual lightcurves."""
from itertools import product, count
import sys
import numpy
import pandas
from pytransit import RoadRunnerModel
from configargparse import ArgumentParser, DefaultsFormatter
from autowisp.light_curves.light_curve_file import LightCurveFile
from autowisp.evaluator import LightCurveEvaluator
from autowisp.database.interface import set_project_home
# TODO:Document all expected entries in configuration for `get_plot_data()`
[docs]
def evaluate_model(model, lc_eval, expression_params, shift_to=None):
"""Return the model evaluated for the given lightcurve."""
args = [
lc_eval(expression.format_map(expression_params))
for expression in model.get("args", [])
]
kwargs = {
arg_name: lc_eval(expression.format_map(expression_params))
for arg_name, expression in model.get("kwargs", {}).items()
}
model_values = globals()[model["type"] + "_model"](*args, **kwargs)
if shift_to is not None:
shift_to = lc_eval(shift_to.format_map(expression_params))
assert len(shift_to) == len(model_values)
model_values += numpy.nanmedian(shift_to - model_values)
return model_values
# pylint: disable=too-many-arguments
[docs]
def optimize_substitutions(
lc_eval, *, find_best, minimize, y_expression, model, expression_params
):
"""
Find the values of LC substitution params that minimize an expression.
Updates ``lc_evals.lc_substitutions`` with the best values found.
Args:
lc_eval(LightCurveEvaluator): Allows evaluating the expression to
minimize.
find_best(iterable): Iterable of 2-tuples with the first entry
in each tuple being a substitution parameters that need to be
optimized and the second entry containing an iterable of all
possible values for that parameter. All possible combinations
are tried.
minimize(str): Expression that is evaluated for each combination of
values from ``find_best`` to select the combination for which
``minimize`` evaluates to the smallest value.
Returns:
the smallest value of the ``minimize`` expression found.
"""
key_order = [key for key, _ in find_best]
best_combination = None
best_found = None
best_model = None
for combination in product(*(values for _, values in find_best)):
lc_eval.update_substitutions(zip(key_order, combination))
if model is not None:
model_values = evaluate_model(model, lc_eval, expression_params)
lc_eval.symtable["model_diff"] = (
lc_eval(y_expression.format_map(expression_params))
- model_values
)
minimize_val = lc_eval(minimize.format_map(expression_params))
if best_found is None or minimize_val < best_found:
best_found = minimize_val
best_combination = combination
if model is not None:
best_model = model_values
print(f"Best substitutions: {dict(zip(key_order, best_combination))!r}")
print(f"Best value: {best_found!r}")
lc_eval.update_substitutions(zip(key_order, best_combination))
return best_found, best_model
# pylint: enable=too-many-arguments
[docs]
def set_substitutions(
lc_eval, configuration, lightcurve, photometry_mode, **optimize_kwargs
):
"""Set the LC path substitutions ``lc_eval`` should use."""
if configuration["lc_substitutions"].get("magfit_iteration", 0) < 0:
lc_eval.update_substitutions(
{
"magfit_iteration": configuration["lc_substitutions"][
"magfit_iteration"
]
+ lightcurve.get_num_magfit_iterations(
photometry_mode,
lc_eval.lc_points_selection,
aperture_index=0,
**lc_eval.lc_substitutions,
)
}
)
if configuration["find_best"]:
(minimize_value, best_model) = optimize_substitutions(
lc_eval,
find_best=configuration["find_best"],
minimize=configuration["minimize"],
**optimize_kwargs,
)
lc_eval.symtable["best_model"] = best_model
else:
minimize_value = None
return minimize_value
[docs]
def evaluate_expressions(expressions, lc_eval, photometry_mode):
"""Evaluate the required expressions for the current lightcurve."""
result = {}
aperture_indices = lc_eval.lc_substitutions.get("aperture_index", None)
print("Aperture indices: ", aperture_indices)
if aperture_indices is None:
aperture_indices = count()
else:
aperture_indices = [aperture_indices]
for ap_ind in aperture_indices:
lc_eval.update_substitutions({"aperture_index": ap_ind})
try:
for var_name, var_expr in expressions.items():
print(f"Evaluating {var_expr!r} for aperture {ap_ind}")
result[var_name.format(aperture_index=ap_ind)] = lc_eval(
var_expr.format(mode=photometry_mode)
)
except OSError:
if isinstance(aperture_indices, list):
raise
break
return result
# pylint: disable=too-many-arguments
[docs]
def get_sphotref_result(
*,
single_photref_fname,
lightcurve,
expressions,
configuration,
model,
lc_eval,
):
"""Get the plot data for a single photometric reference."""
print(f"Single photref: {single_photref_fname!r}")
best_minimize = None
sphotref_result = {}
for photometry_mode in configuration["photometry_modes"]:
sphotref_dset_key = photometry_mode + ".magfit.cfg.single_photref"
lc_eval.lc_points_selection = None
lc_eval.update_substitutions({"aperture_index": 0})
lc_points_selection = lc_eval(
sphotref_dset_key + " == " + repr(single_photref_fname)
)
del lc_eval.lc_substitutions["aperture_index"]
if configuration["selection"] is not None:
lc_points_selection = numpy.logical_and(
lc_eval(configuration["selection"] or "True"),
lc_points_selection,
)
lc_eval.lc_points_selection = lc_points_selection
minimize_value = set_substitutions(
lc_eval,
configuration,
lightcurve,
photometry_mode,
y_expression=(
None if model is None else expressions[model["quantity"]]
),
model=model,
expression_params={"mode": photometry_mode},
)
if minimize_value is not None:
sphotref_result["best_model"] = lc_eval.symtable["best_model"]
if best_minimize is None or minimize_value < best_minimize:
best_minimize = minimize_value
sphotref_result.update(
evaluate_expressions(expressions, lc_eval, photometry_mode)
)
return sphotref_result
# pylint: enable=too-many-arguments
[docs]
def get_plot_data(lc_fname, expressions, configuration, model=None):
"""Read relevant data from the lightcurve."""
result = {}
if model and model.get("shift_to") is True:
model["shift_to"] = expressions[model["quantity"]]
with LightCurveFile(lc_fname, "r") as lightcurve:
lc_eval = LightCurveEvaluator(
lightcurve, **configuration["lc_substitutions"]
)
all_sphotref_fnames = set()
for photometry_mode in configuration["photometry_modes"]:
all_sphotref_fnames |= set(
lightcurve.get_dataset(
photometry_mode + ".magfit.cfg.single_photref",
aperture_index=0,
)
)
for single_photref_fname in all_sphotref_fnames:
result[single_photref_fname.decode()] = get_sphotref_result(
single_photref_fname=single_photref_fname,
lightcurve=lightcurve,
expressions=expressions,
configuration=configuration,
model=model,
lc_eval=lc_eval,
)
return result
[docs]
def calculate_combined(plot_data, match_id_key, aggregation_function):
"""Create a specified plot of the given lightcurve."""
fixed_order_data = list(plot_data.values())
combined_data = {}
match_ids = pandas.concat(
(pandas.Series(data[match_id_key]) for data in fixed_order_data),
ignore_index=True,
)
for var_name in fixed_order_data[0].keys():
try:
combined_data[var_name] = (
pandas.concat(
(
pandas.Series(data[var_name])
for data in fixed_order_data
),
ignore_index=True,
)
.groupby(match_ids)
.agg(aggregation_function)
.to_numpy()
)
except TypeError:
pass
return combined_data
[docs]
def transit_model(times, **params):
"""Calculate the magnitude change of exoplanet with given parameters."""
model = RoadRunnerModel("quadratic")
model.set_data(times)
print(
f"Evaluating transit model for parameters: {params!r} "
f"for times: {times!r}."
)
mag_change = -2.5 * numpy.log10(model.evaluate(**params))
return mag_change
[docs]
def parse_command_line():
"""Return the command line configuration."""
def parse_range(range_str):
"""Parse optimize option value."""
key, range_str = range_str.split(":")
start, stop = map(int, range_str.split(".."))
return key, range(start, stop + 1)
def parse_format(format_str):
"""Parse the format string for output."""
formatter = f"{{{format_str}}}"
print("Formatter: ", formatter)
return formatter.format
parser = ArgumentParser(
description="Extract information from a lightcurve for plotting or "
"other analysis.",
default_config_files=[],
formatter_class=DefaultsFormatter,
ignore_unknown_config_file_keys=True,
)
parser.add_argument(
"lc_fname",
type=str,
nargs="?",
help="Path to the lightcurve file to plot.",
)
parser.add_argument(
"--config-file",
"--config",
"-c",
is_config_file=True,
help="Path to configuration file.",
)
parser.add_argument(
"--list-lc-quantities",
action="store_true",
help="List all available lightcurve quantities and exit.",
)
parser.add_argument(
"--expression",
"-e",
type=str,
action="append",
default=[],
help="Add another expression of lightcurve quantities to save. Use "
"``--list-lc-quantities`` to see available quantities and brief "
"descriptions. Should be formatted as <expression_name>=<expression>. "
"If the quantity is aperture dependent, and ``--find-best`` is "
"not specified, the expression name should contain {aperture_index} "
"substitution to allow the value for all apertures to be saved. "
"Otherwise, only one of the apertures will be saved.",
)
parser.add_argument(
"--photometry-modes",
nargs="+",
default=["apphot"],
help="The photometry modes to search for best lightcurve.",
)
parser.add_argument(
"--find-best",
nargs="*",
default=[("aperture_index", range(8))],
type=parse_range,
help="Lightcurve substitutions to optimize to find the smallest value "
"of the expression specified in ``--minimize-expression``. The format "
"is ``substitution_key:<min>..<max>``. If multiple keys are given all "
"possible combinations are tested.",
)
parser.add_argument(
"--minimize",
default=(
"nanmedian(abs({mode}.magfit.magnitude - "
"nanmedian({mode}.magfit.magnitude)))"
),
help="The expression to minimize by iterating over all possibilities "
"specified by ``--optimize``.",
)
parser.add_argument(
"--bin",
default=None,
nargs=2,
help="Binning to apply. Should be an expression of lightcurve "
"quantities followed by a binning function. Some examples include:\n"
"\t* ``fitsheader.rawfname`` ``nanmedian`` to plot the color "
"median of the measurements in different channels of an each image.\n"
"\t* ``(skypos.BJD * 24 * 12).astype(int)`` ``nanmedian`` to bin in 5 "
"min increments. By default no binning is applied.",
)
parser.add_argument(
"--output",
"-o",
default=None,
help="The filename under which to save the extracted light curve data.",
)
parser.add_argument(
"--na-rep",
default="",
help="The string to use for representing NaN values in the output.",
)
parser.add_argument(
"--sep",
default=",",
help="The separator to use in the output file.",
)
parser.add_argument(
"--float-format",
type=parse_format,
default=None,
help="The format to use for floating point numbers in the output.",
)
result = parser.parse_args()
if result.list_lc_quantities:
print(format_lc_quantities())
sys.exit(0)
if not getattr(result, "lc_fname", False):
print("No lightcurve file specified.")
parser.print_usage()
sys.exit(1)
return result
[docs]
def main():
"""Avoid polluting global scope."""
configuration = vars(parse_command_line())
set_project_home(configuration.get("project_home"))
out_fname = configuration.pop("output")
if out_fname:
configuration["lc_substitutions"] = {}
configuration["selection"] = None
expressions = dict(
expression.split("=", 1)
for expression in configuration.pop("expression")
)
data_by_sphotref = get_plot_data(
configuration.pop("lc_fname"),
expressions=expressions,
configuration=configuration,
model=None,
)
data_by_sphotref = next(iter(data_by_sphotref.values()))
data_by_sphotref.pop("best_model", None)
data_by_sphotref = pandas.DataFrame.from_dict(data_by_sphotref)
data_by_sphotref.to_csv(
out_fname,
na_rep=configuration["na_rep"],
sep=configuration["sep"],
float_format=configuration["float_format"],
index=False,
)
if __name__ == "__main__":
main()