"""Uniform interface for source extractoin using several methods."""
import logging
import numpy
from astropy.io import fits
from scipy import spatial
from numpy.lib import recfunctions
from astrowisp.utils.file_utilities import (
prepare_file_output,
get_unpacked_fits,
)
from autowisp import source_finder_util
from autowisp.evaluator import Evaluator
from autowisp.exceptions import NoSourcesFoundError
# This still makes sense as a class
# pylint: disable=too-few-public-methods
[docs]
class SourceFinder:
"""Find sources in an image of the night sky and report properties."""
[docs]
@staticmethod
def _create_mock_source_list(fits_fname, configuration):
"""Create a source list with randomly generated properties."""
with fits.open(fits_fname) as fits_file:
# False positive
# pylint: disable=no-member
hdu_index = 0 if fits_file[0].header["NAXIS"] else 1
xresolution = fits_file[hdu_index].header["NAXIS1"]
yresolution = fits_file[hdu_index].header["NAXIS2"]
# pylint: disable=no-member
med_pixel = numpy.median(fits_file[hdu_index].data)
nsources = 1000
result = numpy.empty(
nsources,
dtype=[
(
name,
(numpy.int32 if name in ["id", "npix"] else numpy.float64),
)
for name in source_finder_util.get_srcextract_columns("fistar")
],
)
result["id"] = numpy.arange(nsources)
# False positive
# pylint: disable=no-member
result["x"] = numpy.random.random(nsources) * xresolution
result["y"] = numpy.random.random(nsources) * yresolution
result["bg"] = (1.0 + 0.1 * numpy.random.random(nsources)) * med_pixel
result["flux"] = (
numpy.random.random(nsources)
* configuration["brightness_threshold"]
)
result["amp"] = 0.2 * result["flux"]
result["s/n"] = result["flux"] / result["bg"]
result["s"] = 2.3 + 0.2 * numpy.random.random(nsources)
result["d"] = 0.3 + 0.1 * numpy.random.random(nsources)
result["k"] = 0.3 + 0.1 * numpy.random.random(nsources)
result["npix"] = 20 + (5 * numpy.random.random(nsources)).astype(int)
# pylint: enable=no-member
return result
[docs]
def _add_saturation_flags(self, fits_fname, source_list):
"""Add array of how many pixels near source are saturated."""
with fits.open(fits_fname, "readonly") as fits_file:
saturated_kdtree = spatial.KDTree(
numpy.column_stack(numpy.nonzero(fits_file[2].data)) + 0.5
)
source_coords = numpy.column_stack((source_list["y"], source_list["x"]))
print(
"Match radii: " + repr(numpy.sqrt(source_list["npix"] / numpy.pi))
)
saturated_in_range = saturated_kdtree.query_ball_point(
source_coords, numpy.sqrt(source_list["npix"] / numpy.pi)
)
print(f"Saturated in range shape: {saturated_in_range.shape}")
print(f"Saturated in range: {saturated_in_range!r}")
return recfunctions.append_fields(
source_list,
"nsatpix",
[len(saturated) for saturated in saturated_in_range],
usemask=False,
)
[docs]
def __init__( # pylint: disable=too-many-arguments
self,
*,
tool="hatphot",
brightness_threshold=10,
brightness_quantile=0.999,
brightness_quantile_scale=1.0,
filter_sources="True",
max_sources=0,
allow_overwrite=False,
allow_dir_creation=False,
always_return_sources=False,
):
"""Prepare to use the specified tool and define faint limit."""
self.configuration = {
"tool": tool,
"brightness_threshold": brightness_threshold,
"brightness_quantile": brightness_quantile,
"brightness_quantile_scale": brightness_quantile_scale,
"allow_overwrite": allow_overwrite,
"allow_dir_creation": allow_dir_creation,
"always_return_sources": always_return_sources,
"filter_sources": filter_sources,
"max_sources": max_sources,
}
[docs]
def __call__(self, fits_fname, source_fname=None, **configuration):
"""
Extract the sources from the given frame and save or return them.
Args:
fits_fname(str): The filename of the fits file to extract
sources from. Can be packed or unpacked.
source_fname(str or None): If None, the extract sources are
returned as numpy. field array. Otherwise, this specifies a file
to save the source extractor output.
Returns:
field array or None:
The extracted source from the image, if source_fname was None.
"""
logger = logging.getLogger(__name__)
configuration = {**self.configuration, **configuration}
if configuration["tool"] == "mock":
return self._create_mock_source_list(fits_fname, configuration)
# Calculate brightness_threshold from image quantile if not specified
if not configuration["brightness_threshold"]:
with fits.open(fits_fname) as fits_file:
hdu_index = 0 if fits_file[0].header["NAXIS"] else 1
image_data = fits_file[hdu_index].data
configuration["brightness_threshold"] = configuration[
"brightness_quantile_scale"
] * numpy.quantile(
image_data.flatten(), configuration["brightness_quantile"]
)
logger.info(
"Computed brightness_threshold = %r from %r quantile "
"with scale %r for %r",
configuration["brightness_threshold"],
configuration["brightness_quantile"],
configuration["brightness_quantile_scale"],
fits_fname,
)
start_extraction = getattr(
source_finder_util, "start_" + configuration["tool"]
)
with get_unpacked_fits(fits_fname) as unpacked_fname:
logger.debug(
"Extracting sources from %r to %r",
unpacked_fname,
source_fname,
)
extraction_args = (
unpacked_fname,
configuration["brightness_threshold"],
)
if source_fname:
prepare_file_output(
source_fname,
configuration["allow_overwrite"],
configuration["allow_dir_creation"],
)
with open(source_fname, "wb") as destination:
start_extraction(*extraction_args, destination).wait()
if not configuration["always_return_sources"]:
return None
sources_file = source_fname
else:
logger.debug("Creating extraction process.")
extraction_process = start_extraction(*extraction_args)
logger.debug("Extraction process started.")
sources_file = extraction_process.stdout
logger.debug("Parsing source list from %s", repr(unpacked_fname))
result = numpy.genfromtxt(
sources_file,
names=source_finder_util.get_srcextract_columns(
configuration["tool"]
),
dtype=None,
deletechars="",
missing_values="-",
)
# A single extracted source comes back from genfromtxt as a 0-d
# array, which the finite-value mask below cannot index.
result = numpy.atleast_1d(result)
if result.size == 0:
if not source_fname:
# stdout is already at EOF, so this only reaps the
# extractor rather than leaving it behind on the raise.
extraction_process.communicate()
raise NoSourcesFoundError(
f"Source extraction with {configuration['tool']!r} found "
f"no sources in {fits_fname!r} using brightness threshold "
f"{configuration['brightness_threshold']!r}.",
user_message=(
f"No stars were found in {fits_fname}. The image may "
"be clouded over or badly out of focus, or the "
"brightness threshold may be set too high."
),
details={
"image": fits_fname,
"srcfind_tool": configuration["tool"],
"brightness_threshold": configuration[
"brightness_threshold"
],
"brightness_quantile": configuration[
"brightness_quantile"
],
"brightness_quantile_scale": configuration[
"brightness_quantile_scale"
],
},
)
finite = numpy.ones(result.size, dtype=bool)
for colname in result.dtype.names:
finite = numpy.logical_and(
finite, numpy.isfinite(result[colname])
)
result = result[finite]
logger.debug("Extracted %d stars from %r", result.size, fits_fname)
if configuration["tool"] == "hatphot":
result["x"] -= 0.5
result["y"] -= 0.5
if configuration["filter_sources"] != "True":
result = result[
Evaluator(
result,
{
"brightness_threshold": configuration[
"brightness_threshold"
]
},
)(configuration["filter_sources"])
]
logger.debug(
"Sorting %d filtered sources from %r (%r)",
result.size,
fits_fname,
unpacked_fname,
)
result.sort(order="flux")
result = numpy.flip(result)
if configuration["max_sources"] > 0:
result = result[: configuration["max_sources"]]
if not source_fname:
logger.debug(
"Waiting for source extraction process for %s to finish",
repr(unpacked_fname),
)
extraction_process.communicate()
logger.debug(
"Adding saturation flags for %d sources from %r ",
result.size,
fits_fname,
)
return self._add_saturation_flags(unpacked_fname, result)
# pylint: enable=too-few-public-methods