<commit_before>from . import DEFAULT_PRFDIR
from .utils import channel_to_module_output
from abc import abstractmethod
import os
import glob
import math
import scipy
import numpy as np
import tqdm
from astropy.io import fits as pyfits
from oktopus.posterior import PoissonPosterior


__all__ = ['KeplerPRFPhotometry', 'KeplerSceneModel', 'KeplerPRF', 'get_initial_guesses']


class KeplerPRFPhotometry(object):
    """
    This class performs PRF Photometry on a target pixel file from
    NASA's Kepler/K2 missions.

    Attributes
    ----------
    scene_model : instance of KeplerSceneModel
    priors : instance of oktopus.JointPrior
        Priors on the parameters that will be estimated
    loss_function : subclass of oktopus.LossFunction
        Noise distribution associated with each random measurement
    """

    def __init__(self, scene_model, prior, loss_function=PoissonPosterior):
        self.scene_model = scene_model
        self.prior = prior
        self.loss_function = loss_function
        self.opt_params = np.array([])
        self.residuals = np.array([])
        self.uncertainties = np.array([])

    def fit(self, tpf_flux, x0=None, cadences='all', **kwargs):
        """
        Fits the scene model to the given data in ``tpf_flux``.

        Parameters
        ----------
        tpf_flux : array-like
            A pixel flux time-series, e.g, KeplerTargetPixelFile.flux.
            Such that (time, row, column) represents the shape of ``tpf_flux``.
        x0 : array-like or None
            Initial guesses on the parameters.
        cadences : array-like of ints or str
            A list or array that contains the cadences which will be fitted.
            Default is to fit all cadences.
        kwargs : dict
            Dictionary of additional parameters to be passed to
            `scipy.optimize.minimize`.
        """
        if x0 is None:
            x0 = self.prior.mean

        if cadences == 'all':
            cadences = range(tpf_flux.shape[0])

        for t in tqdm.tqdm(cadences):
            loss = self.loss_function(tpf_flux[t], self.scene_model, prior=self.prior)
            opt_params = loss.fit(x0=x0, **kwargs).x
            residuals = tpf_flux[t] - self.scene_model(*opt_params)
            self.opt_params = np.append(self.opt_params, opt_params)
            self.residuals = np.append(self.residuals, residuals)
        self.opt_params = self.opt_params.reshape((tpf_flux.shape[0], len(x0)))
        self.residuals = self.residuals.reshape(tpf_flux.shape)

    def generate_residuals_movie(self):
        pass


class KeplerSceneModel(object):
    """
    This class builds a generic model for a Kepler scene.

    Attributes
    ----------
    prf_model : instance of KeplerPRF
        An instance of the KeplerPRF class
    n_sources : int
        Number of sources to be modeled with ``prf_model``
    bkg_model : callable
        A function that models the background variation.
        Default is a constant background
    """

    def __init__(self, prf_model, n_sources, bkg_model=lambda bkg: np.array([bkg])):
        self.prf_model = prf_model
        self.n_sources = n_sources
        self.bkg_model = bkg_model

    def __call__(self, *params):
        return self.evaluate(*params)

    def evaluate(self, *params):
        """
        Parameters
        ----------
        flux : scalar or array-like
            Total integrated flux of the PRF model
        center_col, center_row : scalar or array-like
            Column and row coordinates of the center
        scale_col, scale_row : scalar or array-like
            Pixel scale in the column and row directions
        bkg_params : scalar or array-like
            Parameters for the background model
        """
        self.mixture_model = []
        for i in range(self.n_sources):
            self.mixture_model.append(self.prf_model(params[i],
                                                     params[i + self.n_sources],
                                                     params[i + 2 * self.n_sources]))
        self.scene_model = np.sum(self.mixture_model, axis=0) + self.bkg_model(params[-1])

        return self.scene_model


class KeplerPRF(object):
    """
    Kepler's Pixel Response Function

    This class provides the necessary interface to load Kepler PRF
    calibration files and to create a model that can be fit as a function
    of flux, center positions, width, and rotation angle.

    Attributes
    ----------
    channel : int
        KeplerTargetPixelFile.channel
    shape : (int, int)
        KeplerTargetPixelFile.shape[1:]
    column : int
        KeplerTargetPixelFile.column
    row : int
        KeplerTargetPixelFile.row
    prf_files_dir : str
        Relative or aboslute path to a directory containing the Pixel Response
        Function calibration files produced during Kepler data comissioning.

    Examples
    --------
    Objects from the KeplerPRF class are defined by a channel number, a pair of
    dimensions (the size of the image), and a reference coordinate (bottom left
    corner). In this example, we create a KeplerPRF object located at channel
    #44 with dimension equals 10 x 10, reference row and column coordinate
    equals (5, 5). After the object has been created, we may translate it to a
    given center coordinate. Additionally, we can specify total flux, pixel
    scales, and rotation around the object's center.

    >>> import math
    >>> import matplotlib.pyplot as plt
    >>> from pyke import KeplerPRF
    >>> kepprf = KeplerPRF(channel=44, shape=(10, 10), column=5, row=5)
    >>> kepprf(flux=1000, center_col=10, center_row=10, scale_col=0.7)

    Notes
    -----
    The PRF calibration files can be downloaded here: http://archive.stsci.edu/missions/kepler/fpc/prf/
    """

    def __init__(self, channel, shape, column, row, prf_files_dir=DEFAULT_PRFDIR):
        self.prf_files_dir = prf_files_dir
        self.channel = channel
        self.shape = shape
        self.column = column
        self.row = row
        self.col_coord, self.row_coord, self.interpolate = self._prepare_prf()

    def __call__(self, flux, center_col, center_row, scale_col, scale_row,
                 rotation_angle):<commit_msg>update docstring example<commit_after>from . import DEFAULT_PRFDIR
from .utils import channel_to_module_output
from abc import abstractmethod
import os
import glob
import math
import scipy
import numpy as np
import tqdm
from astropy.io import fits as pyfits
from oktopus.posterior import PoissonPosterior


__all__ = ['KeplerPRFPhotometry', 'KeplerSceneModel', 'KeplerPRF', 'get_initial_guesses']


class KeplerPRFPhotometry(object):
    """
    This class performs PRF Photometry on a target pixel file from
    NASA's Kepler/K2 missions.

    Attributes
    ----------
    scene_model : instance of KeplerSceneModel
    priors : instance of oktopus.JointPrior
        Priors on the parameters that will be estimated
    loss_function : subclass of oktopus.LossFunction
        Noise distribution associated with each random measurement
    """

    def __init__(self, scene_model, prior, loss_function=PoissonPosterior):
        self.scene_model = scene_model
        self.prior = prior
        self.loss_function = loss_function
        self.opt_params = np.array([])
        self.residuals = np.array([])
        self.uncertainties = np.array([])

    def fit(self, tpf_flux, x0=None, cadences='all', **kwargs):
        """
        Fits the scene model to the given data in ``tpf_flux``.

        Parameters
        ----------
        tpf_flux : array-like
            A pixel flux time-series, e.g, KeplerTargetPixelFile.flux.
            Such that (time, row, column) represents the shape of ``tpf_flux``.
        x0 : array-like or None
            Initial guesses on the parameters.
        cadences : array-like of ints or str
            A list or array that contains the cadences which will be fitted.
            Default is to fit all cadences.
        kwargs : dict
            Dictionary of additional parameters to be passed to
            `scipy.optimize.minimize`.
        """
        if x0 is None:
            x0 = self.prior.mean

        if cadences == 'all':
            cadences = range(tpf_flux.shape[0])

        for t in tqdm.tqdm(cadences):
            loss = self.loss_function(tpf_flux[t], self.scene_model, prior=self.prior)
            opt_params = loss.fit(x0=x0, **kwargs).x
            residuals = tpf_flux[t] - self.scene_model(*opt_params)
            self.opt_params = np.append(self.opt_params, opt_params)
            self.residuals = np.append(self.residuals, residuals)
        self.opt_params = self.opt_params.reshape((tpf_flux.shape[0], len(x0)))
        self.residuals = self.residuals.reshape(tpf_flux.shape)

    def generate_residuals_movie(self):
        pass


class KeplerSceneModel(object):
    """
    This class builds a generic model for a Kepler scene.

    Attributes
    ----------
    prf_model : instance of KeplerPRF
        An instance of the KeplerPRF class
    n_sources : int
        Number of sources to be modeled with ``prf_model``
    bkg_model : callable
        A function that models the background variation.
        Default is a constant background
    """

    def __init__(self, prf_model, n_sources, bkg_model=lambda bkg: np.array([bkg])):
        self.prf_model = prf_model
        self.n_sources = n_sources
        self.bkg_model = bkg_model

    def __call__(self, *params):
        return self.evaluate(*params)

    def evaluate(self, *params):
        """
        Parameters
        ----------
        flux : scalar or array-like
            Total integrated flux of the PRF model
        center_col, center_row : scalar or array-like
            Column and row coordinates of the center
        scale_col, scale_row : scalar or array-like
            Pixel scale in the column and row directions
        rotation_angle : float
            Rotation angle in radians
        bkg_params : scalar or array-like
            Parameters for the background model
        """
        self.mixture_model = []
        for i in range(self.n_sources):
            self.mixture_model.append(self.prf_model(params[i],
                                                     params[i + self.n_sources],
                                                     params[i + 2 * self.n_sources],
                                                     params[i + 3 * self.n_sources],
                                                     params[i + 4 * self.n_sources],
                                                     params[i + 5 * self.n_sources]))
        self.scene_model = np.sum(self.mixture_model, axis=0) + self.bkg_model(params[-1])

        return self.scene_model


class KeplerPRF(object):
    """
    Kepler's Pixel Response Function

    This class provides the necessary interface to load Kepler PRF
    calibration files and to create a model that can be fit as a function
    of flux, center positions, width, and rotation angle.

    Attributes
    ----------
    channel : int
        KeplerTargetPixelFile.channel
    shape : (int, int)
        KeplerTargetPixelFile.shape[1:]
    column : int
        KeplerTargetPixelFile.column
    row : int
        KeplerTargetPixelFile.row
    prf_files_dir : str
        Relative or aboslute path to a directory containing the Pixel Response
        Function calibration files produced during Kepler data comissioning.

    Examples
    --------
    Objects from the KeplerPRF class are defined by a channel number, a pair of
    dimensions (the size of the image), and a reference coordinate (bottom left
    corner). In this example, we create a KeplerPRF object located at channel
    #44 with dimension equals 10 x 10, reference row and column coordinate
    equals (5, 5). After the object has been created, we may translate it to a
    given center coordinate. Additionally, we can specify total flux, pixel
    scales, and rotation around the object's center.

    >>> import math
    >>> import matplotlib.pyplot as plt
    >>> from pyke import KeplerPRF
    >>> kepprf = KeplerPRF(channel=44, shape=(10, 10), column=5, row=5)
    >>> prf = kepprf(flux=1000, center_col=10, center_row=10, scale_col=0.7, rotation_angle=math.pi/2)
    >>> plt.imshow(prf, origin='lower')

    Notes
    -----
    The PRF calibration files can be downloaded here: http://archive.stsci.edu/missions/kepler/fpc/prf/
    """

    def __init__(self, channel, shape, column, row, prf_files_dir=DEFAULT_PRFDIR):
        self.prf_files_dir = prf_files_dir
        self.channel = channel
        self.shape = shape
        self.column = column
        self.row = row
        self.col_coord, self.row_coord, self.interpolate = self._prepare_prf()

    def __call__(self, flux, center_col, center_row, scale_col, scale_row,
                 rotation_angle):
{'fraction_non_alphanumeric': 0.08370146477563357, 'fraction_numerical': 0.006510113926993723, 'mean_word_length': 2.2868059093224655, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 36, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 16, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '7086009', 'n_tokens_mistral': 3745, 'n_tokens_neox': 3547, 'n_words': 1150}