<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
A pedestrian version of The Cannon.
"""

from __future__ import (division, print_function, absolute_import,
                        unicode_literals)

__all__ = ["CannonModel"]

import logging
import numpy as np
import scipy.optimize as op
import sharedmem

from . import (model, utils)

logger = logging.getLogger(__name__)


class CannonModel(model.BaseCannonModel):
    """
    A generalised Cannon model for the estimation of arbitrary stellar labels.

    :param labelled_set:
        A set of labelled objects. The most common input form is a table with
        columns as labels, and stars/objects as rows.

    :type labelled_set:
        :class:`~astropy.table.Table`, numpy structured array

    :param normalized_flux:
        An array of normalized fluxes for stars in the labelled set, given as
        shape `(num_stars, num_pixels)`. The `num_stars` should match the number
        of rows in `labelled_set`.

    :type normalized_flux:
        :class:`np.ndarray`

    :param normalized_ivar:
        An array of inverse variances on the normalized fluxes for stars in the
        labelled set. The shape of the `normalized_ivar` array should match that
        of `normalized_flux`.

    :type normalized_ivar:
        :class:`np.ndarray`

    :param dispersion: [optional]
        The dispersion values corresponding to the given pixels. If provided, 
        this should have length `num_pixels`.

    :param threads: [optional]
        Specify the number of parallel threads to use. If `threads > 1`, the
        training and prediction phases will be automagically parallelised.

    :param pool: [optional]
        Specify an optional multiprocessing pool to map jobs onto.
        This argument is only used if specified and if `threads > 1`.
    """
    def __init__(self, *args, **kwargs):
        super(CannonModel, self).__init__(*args, **kwargs)


    @model.requires_model_description
    def train(self, fixed_scatter=False, progressbar=True, **kwargs):
        """
        Train the model based on the labelled set using the given vectorizer.

        :param fixed_scatter: [optional]
            Fix the scatter terms and do not solve for them during the training
            phase. If set to `True`, the `s2` attribute must be already set.

        :param progressbar: [optional]
            Show a progress bar.
        """
        
        if fixed_scatter and self.s2 is None:
            raise ValueError("intrinsic pixel variance (s2) must be set "
                             "before training if fixed_scatter is set to True")

        # Initialize the scatter.
        p0_scatter = np.sqrt(self.s2) if fixed_scatter \
            else 0.01 * np.ones_like(self.dispersion)

        # Prepare details about any progressbar to show.
        M, N = self.normalized_flux.shape
        message = None if not progressbar else \
            "Training {0}-label {1} with {2} stars and {3} pixels/star".format(
                len(self.vectorizer.label_names), type(self).__name__, M, N)

        # Prepare the method and arguments.
        fitter = kwargs.pop("function", _fit_pixel)
        args = [self.normalized_flux.T, self.normalized_ivar.T, p0_scatter]
        args.extend(kwargs.pop("additional_args", []))
        kwds = {
            "fixed_scatter": fixed_scatter,
            "design_matrix": self.design_matrix
        }
        #kwds.update(kwargs)

        # Only use shared memory if necessary because it is a performance hit.
        if self.pool is None:
            mapper = map
        else:
            mapper = self.pool.map
            if kwargs.get("shared_memory", True):
                shared_args = []
                for arg in args:
                    if isinstance(arg, np.ndarray):
                        shared_arg = sharedmem.empty_like(arg)
                        shared_arg[:] = arg
                    else:
                        shared_arg = arg
                    shared_args.append(shared_arg)
                args = shared_args

                # Aaaand don't forget the big one!
                dm = sharedmem.empty_like(self.design_matrix)
                dm[:] = self.design_matrix
                kwds["design_matrix"] = dm
        
        # Wrap the function so we can parallelize it out.
        f = utils.wrapper(fitter, None, kwds, N, message=message)

        # Time for work.
        results = np.array(mapper(f, [row for row in zip(*args)]))
        
        # Unpack the results.
        self.theta, self.s2 = (results[:, :-1], results[:, -1]**2)
        return None


    @model.requires_training_wheels
    def predict(self, labels, **kwargs):
        """
        Predict spectra from the trained model, given the labels.

        :param labels:
            The label values to predict model spectra of. The length and order
            should match what is required of the vectorizer
            (`CannonModel.vectorizer.label_names`).
        """
        return np.dot(self.theta, self.vectorizer(labels).T).T


    @model.requires_training_wheels
    def fit(self, normalized_flux, normalized_ivar, **kwargs):
        """
        Solve the labels for the given normalized fluxes and inverse variances.

        :param normalized_flux:
            The normalized fluxes. These should be on the same dispersion scale
            as the trained data.

        :param normalized_ivar:
            The inverse variances of the normalized flux values. This should
            have the same shape as `normalized_flux`.
<commit_msg>Chunk the data in parallel<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

"""
A pedestrian version of The Cannon.
"""

from __future__ import (division, print_function, absolute_import,
                        unicode_literals)

__all__ = ["CannonModel"]

import cPickle as pickle
import logging
import numpy as np
import scipy.optimize as op
import sharedmem

from . import (model, utils)

logger = logging.getLogger(__name__)


class CannonModel(model.BaseCannonModel):
    """
    A generalised Cannon model for the estimation of arbitrary stellar labels.

    :param labelled_set:
        A set of labelled objects. The most common input form is a table with
        columns as labels, and stars/objects as rows.

    :type labelled_set:
        :class:`~astropy.table.Table`, numpy structured array

    :param normalized_flux:
        An array of normalized fluxes for stars in the labelled set, given as
        shape `(num_stars, num_pixels)`. The `num_stars` should match the number
        of rows in `labelled_set`.

    :type normalized_flux:
        :class:`np.ndarray`

    :param normalized_ivar:
        An array of inverse variances on the normalized fluxes for stars in the
        labelled set. The shape of the `normalized_ivar` array should match that
        of `normalized_flux`.

    :type normalized_ivar:
        :class:`np.ndarray`

    :param dispersion: [optional]
        The dispersion values corresponding to the given pixels. If provided, 
        this should have length `num_pixels`.

    :param threads: [optional]
        Specify the number of parallel threads to use. If `threads > 1`, the
        training and prediction phases will be automagically parallelised.

    :param pool: [optional]
        Specify an optional multiprocessing pool to map jobs onto.
        This argument is only used if specified and if `threads > 1`.
    """
    def __init__(self, *args, **kwargs):
        super(CannonModel, self).__init__(*args, **kwargs)


    @model.requires_model_description
    def train(self, fixed_scatter=False, progressbar=True, **kwargs):
        """
        Train the model based on the labelled set using the given vectorizer.

        :param fixed_scatter: [optional]
            Fix the scatter terms and do not solve for them during the training
            phase. If set to `True`, the `s2` attribute must be already set.

        :param progressbar: [optional]
            Show a progress bar.
        """
        
        if fixed_scatter and self.s2 is None:
            raise ValueError("intrinsic pixel variance (s2) must be set "
                             "before training if fixed_scatter is set to True")

        # Initialize the scatter.
        p0_scatter = np.sqrt(self.s2) if fixed_scatter \
            else 0.01 * np.ones_like(self.dispersion)

        # Prepare details about any progressbar to show.
        M, N = self.normalized_flux.shape
        message = None if not progressbar else \
            "Training {0}-label {1} with {2} stars and {3} pixels/star".format(
                len(self.vectorizer.label_names), type(self).__name__, M, N)

        # Prepare the method and arguments.
        fitter = kwargs.pop("function", _fit_pixel)
        kwds = {
            "fixed_scatter": fixed_scatter,
            "design_matrix": self.design_matrix
        }
        #kwds.update(kwargs)

        # Only use shared memory if necessary because it is a performance hit.
        if self.pool is None:
            mapper = map
        else:
            mapper = self.pool.map
            if kwargs.get("shared_memory", True):
                """
                shared_args = []
                args = [self.normalized_flux.T, self.normalized_ivar.T, p0_scatter]
                args.extend(kwargs.get("additional_args", []))
                for arg in args:
                    if isinstance(arg, np.ndarray):
                        shared_arg = sharedmem.empty_like(arg)
                        shared_arg[:] = arg
                    else:
                        shared_arg = arg
                    shared_args.append(shared_arg)
                args = shared_args
                """
                # Aaaand don't forget the big one!
                dm = sharedmem.empty_like(self.design_matrix)
                dm[:] = self.design_matrix.copy()
                kwds["design_matrix"] = dm
        
        # Wrap the function so we can parallelize it out.
        f = utils.wrapper(fitter, None, kwds, N, message=message)

        # Calculate chunk size, etc.
        chunks = kwargs.pop("chunks", 10)
        chunk_size = int(np.ceil(len(self.dispersion)/float(chunks)))
        theta = np.zeros((len(self.dispersion), 1 + len(self.vectorizer.terms)))
        scatter = np.zeros((len(self.dispersion)))
        for i in range(chunks):
            # Time for work.
            logger.info("Chunks: {0} w/ chunk size: {1}".format(chunks,
                chunk_size))
            a, b = (i * chunk_size, (i + 1) * chunk_size)

            sm_nf = sharedmem.empty_like(self.normalized_flux.T[a:b])
            sm_ni = sharedmem.empty_like(self.normalized_ivar.T[a:b])
            sm_p0 = sharedmem.empty_like(p0_scatter[a:b])

            sm_nf[:] = self.normalized_flux.T[a:b]
            sm_ni[:] = self.normalized_ivar.T[a:b]
            sm_p0[:] = p0_scatter[a:b]
            args = [sm_nf, sm_ni, sm_p0]
            args.extend(kwargs.get("additional_args", []))
            
            #args = [self.normalized_flux.T[a:b], self.normalized_ivar.T[a:b], p0_scatter[a:b]]
            #args.extend(kwargs.get("additional_args", []))
        
            #_ = args[a:b]
            results = mapper(f, [row for row in zip(*args)])
            results = np.array(results)
            
            # Save these:
            logger.info("Saving to temp.pkl and re-chunking..")
            theta[a:b] = results[a:b, :-1]
            scatter[a:b] = results[a:b, -1]

            with open("temp.pkl", "wb") as fp:
                pickle.dump((theta, scatter), fp, -1)

        self.theta = theta
        self.s2 = scatter**2

        # Unpack the results.
        #self.theta, self.s2 = (results[:, :-1], results[:, -1]**2)
        return None


    @model.requires_training_wheels
    def predict(self, labels, **kwargs):
        """
        Predict spectra from the trained model, given the labels.

        :param labels:
            The label values to predict model spectra of. The length and order
            should match what is required of the vectorizer
            (`CannonModel.vectorizer.label_names`).
        """
        return np.dot(self.theta, self.vectorizer(labels).T).T


    @model.requires_training_wheels
    def fit(self, normalized_flux, normalized_ivar, **kwargs):
        """
        Solve the labels for the given normalized fluxes and inverse variances.

        :param normalized_flux:
            The normalized fluxes. These should be on the same dispersion scale
            as the trained data.

        :param normalized_ivar:
            The inverse variances of the normalized flux values. This should
            have the same shape as `normalized_flux`.

{'fraction_non_alphanumeric': 0.08988412151581585, 'fraction_numerical': 0.004462887566551832, 'mean_word_length': 2.074127557160048, 'pattern_counts': {'":': 4, '<': 3, '<?xml version=': 0, '>': 7, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '1738083', 'n_tokens_mistral': 3666, 'n_tokens_neox': 3471, 'n_words': 1146}