3769530328

Build Type

push

github

Committed by Robert Stein

Commit Message

lintify flat

Run Details

4 of 4 new or added lines in 1 file covered. (100.0%)

4632 of 6121 relevant lines covered (75.67%)

0.76 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

98.25

/winterdrp/processors/flat.py

"""
Module containing processors for flat calibration
"""
import logging
import sys
from collections.abc import Callable

import numpy as np

from winterdrp.data import Image, ImageBatch
from winterdrp.errors import ImageNotFoundError
from winterdrp.paths import FLAT_FRAME_KEY, LATEST_SAVE_KEY
from winterdrp.processors.base_processor import (
    ProcessorPremadeCache,
    ProcessorWithCache,
)
from winterdrp.processors.utils.image_selector import select_from_images

logger = logging.getLogger(__name__)


def default_select_flat(
    images: ImageBatch,
) -> ImageBatch:
    """
    Select images tagged as flat

    :param images: set of images
    :return: subset of flat images
    """
    return select_from_images(images, target_values="flat")


class FlatCalibrator(ProcessorWithCache):
    """
    Processor to apply flat calibration
    """

    base_key = "flat"

    def __init__(
        self,
        *args,
        x_min: int = 0,
        x_max: int = sys.maxsize,
        y_min: int = 0,
        y_max: int = sys.maxsize,
        flat_nan_threshold: float = 0.0,
        select_flat_images: Callable[[ImageBatch], ImageBatch] = default_select_flat,
        **kwargs,
    ):
        super().__init__(*args, **kwargs)
        self.x_min = x_min
        self.x_max = x_max
        self.y_min = y_min
        self.y_max = y_max
        self.flat_nan_threshold = flat_nan_threshold
        self.select_cache_images = select_flat_images

    def __str__(self) -> str:
        return "Creates a flat image, divides other images by this image."

    def _apply_to_images(
        self,
        batch: ImageBatch,
    ) -> ImageBatch:

        master_flat = self.get_cache_file(batch)
        master_flat_data = master_flat.get_data()

        mask = master_flat_data <= self.flat_nan_threshold

        if np.sum(mask) > 0:
            master_flat_data[mask] = np.nan

        for image in batch:
            data = image.get_data()
            data = data / master_flat_data
            image.set_data(data)
            image[FLAT_FRAME_KEY] = master_flat[LATEST_SAVE_KEY]

        return batch

    def make_image(
        self,
        images: ImageBatch,
    ) -> Image:
        images = self.select_cache_images(images)

        n_frames = len(images)
        if n_frames == 0:
            err = f"Found {n_frames} suitable flats in batch"
            logger.error(err)
            raise ImageNotFoundError(err)

        nx, ny = images[0].get_data().shape

        flats = np.zeros((nx, ny, n_frames))

        for i, img in enumerate(images):
            median = np.nanmedian(
                img.get_data()[self.x_min : self.x_max, self.y_min : self.y_max]
            )
            flats[:, :, i] = img.get_data() / median

        logger.info(f"Median combining {n_frames} flats")
        master_flat = np.nanmedian(flats, axis=2)

        return Image(master_flat, header=images[0].get_header())


class SkyFlatCalibrator(FlatCalibrator):
    """
    Processor to do flat calibration using sky flats
    """

    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs, select_flat_images=self.select_sky_flat)

    @staticmethod
    def select_sky_flat(
        images: ImageBatch,
    ) -> ImageBatch:
        """
        Select science images to use as sky flats

        :param images: set of images
        :return: subset of 'sky' images
        """
        return select_from_images(images, key="obsclass", target_values="science")

    def __str__(self) -> str:
        return (
            "Processor to create a sky flat image, divides other images by this image."
        )


class MasterFlatCalibrator(ProcessorPremadeCache, FlatCalibrator):
    """Processor to do flat calibration with a master flate"""

1	"""
2	Module containing processors for flat calibration
3	"""
4	import logging	1✔
5	import sys	1✔
6	from collections.abc import Callable	1✔
7
8	import numpy as np	1✔
9
10	from winterdrp.data import Image, ImageBatch	1✔
11	from winterdrp.errors import ImageNotFoundError	1✔
12	from winterdrp.paths import FLAT_FRAME_KEY, LATEST_SAVE_KEY	1✔
13	from winterdrp.processors.base_processor import (	1✔
14	ProcessorPremadeCache,
15	ProcessorWithCache,
16	)
17	from winterdrp.processors.utils.image_selector import select_from_images	1✔
18
19	logger = logging.getLogger(__name__)	1✔
20
21
22	def default_select_flat(	1✔
23	images: ImageBatch,
24	) -> ImageBatch:
25	"""
26	Select images tagged as flat
27
28	:param images: set of images
29	:return: subset of flat images
30	"""
31	return select_from_images(images, target_values="flat")	1✔
32
33
34	class FlatCalibrator(ProcessorWithCache):	1✔
35	"""
36	Processor to apply flat calibration
37	"""
38
39	base_key = "flat"	1✔
40
41	def __init__(	1✔
42	self,
43	*args,
44	x_min: int = 0,
45	x_max: int = sys.maxsize,
46	y_min: int = 0,
47	y_max: int = sys.maxsize,
48	flat_nan_threshold: float = 0.0,
49	select_flat_images: Callable[[ImageBatch], ImageBatch] = default_select_flat,
50	**kwargs,
51	):
52	super().__init__(args, *kwargs)	1✔
53	self.x_min = x_min	1✔
54	self.x_max = x_max	1✔
55	self.y_min = y_min	1✔
56	self.y_max = y_max	1✔
57	self.flat_nan_threshold = flat_nan_threshold	1✔
58	self.select_cache_images = select_flat_images	1✔
59
60	def __str__(self) -> str:	1✔
61	return "Creates a flat image, divides other images by this image."	1✔
62
63	def _apply_to_images(	1✔
64	self,
65	batch: ImageBatch,
66	) -> ImageBatch:
67
68	master_flat = self.get_cache_file(batch)	1✔
69	master_flat_data = master_flat.get_data()	1✔
70
71	mask = master_flat_data <= self.flat_nan_threshold	1✔
72
73	if np.sum(mask) > 0:	1✔
74	master_flat_data[mask] = np.nan	1✔
75
76	for image in batch:	1✔
77	data = image.get_data()	1✔
78	data = data / master_flat_data	1✔
79	image.set_data(data)	1✔
80	image[FLAT_FRAME_KEY] = master_flat[LATEST_SAVE_KEY]	1✔
81
82	return batch	1✔
83
84	def make_image(	1✔
85	self,
86	images: ImageBatch,
87	) -> Image:
88	images = self.select_cache_images(images)	1✔
89
90	n_frames = len(images)	1✔
91	if n_frames == 0:	1✔
92	err = f"Found {n_frames} suitable flats in batch"
93	logger.error(err)
94	raise ImageNotFoundError(err)
95
96	nx, ny = images[0].get_data().shape	1✔
97
98	flats = np.zeros((nx, ny, n_frames))	1✔
99
100	for i, img in enumerate(images):	1✔
101	median = np.nanmedian(	1✔
102	img.get_data()[self.x_min : self.x_max, self.y_min : self.y_max]
103	)
104	flats[:, :, i] = img.get_data() / median	1✔
105
106	logger.info(f"Median combining {n_frames} flats")	1✔
107	master_flat = np.nanmedian(flats, axis=2)	1✔
108
109	return Image(master_flat, header=images[0].get_header())	1✔
110
111
112	class SkyFlatCalibrator(FlatCalibrator):	1✔
113	"""
114	Processor to do flat calibration using sky flats
115	"""
116
117	def __init__(self, args, *kwargs):	1✔
118	super().__init__(args, *kwargs, select_flat_images=self.select_sky_flat)	1✔
119
120	@staticmethod	1✔
121	def select_sky_flat(	1✔
122	images: ImageBatch,
123	) -> ImageBatch:
124	"""
125	Select science images to use as sky flats
126
127	:param images: set of images
128	:return: subset of 'sky' images
129	"""
130	return select_from_images(images, key="obsclass", target_values="science")	×
131
132	def __str__(self) -> str:	1✔
133	return (	1✔
134	"Processor to create a sky flat image, divides other images by this image."
135	)
136
137
138	class MasterFlatCalibrator(ProcessorPremadeCache, FlatCalibrator):	1✔
139	"""Processor to do flat calibration with a master flate"""

winter-telescope / winterdrp / 3769530328

Source File Press 'n' to go to next uncovered line, 'b' for previous

Source File
Press 'n' to go to next uncovered line, 'b' for previous