<commit_before>import numpy as np
import matplotlib.pylab as plt
import astropy.units as u
from astropy import log
import pyspeckit

# TODO: make informative log.info messages in
#       the methods of the SubCube class

# TODO: redefine SpectralCube.py#L795-L823 to
#       include the option of choosing between
#       the guesses in the middle of fiteach run

class SubCube(pyspeckit.Cube):
    """
    An extension of Cube, tinkered to be an instance of MultiCube, from which 
    it receives references to instances of pyspeckit.Cube that do not depend 
    on a spectral model chosen (so that parent MultiCube doesn't weigh so much)

    Is designed to have methods that operate within a single spectral model.
    """
    def __init__(self, *args, **kwargs):
        super(SubCube, self).__init__(*args, **kwargs)

        # because that UnitConversionError pops up way too often
        if self.xarr.velocity_convention is None:
            self.xarr.velocity_convention = 'radio'
        
        # so I either define some things as `None` 
        # or I'll have to call hasattr or them...
        # TODO: which is a more Pythonic approach?
        # A: probably the hasattr method, see here:  
        # http://programmers.stackexchange.com/questions/
        # 254576/is-it-a-good-practice-to-declare-instance
        # -variables-as-none-in-a-class-in-python
        self.guess_grid = None
        self.model_grid = None

    def info(self):
        def getinfo(target):
            try:
                return getattr(getattr(self,target), 'shape')
            except AttributeError:
                return 'N/A'
                pass
        log.info("Shapes of the arrays:\n"
                 "--> Data cube:\t{}\n".format(getinfo('cube'))+\
                 "--> Guess grid:\t{}\n".format(getinfo('guess_grid'))+\
                 "--> Model grid:\t{}\n".format(getinfo('model_grid'))+\
                 "--> Result rms:\t{}\n".format(getinfo('_residual_rms'))+\
                 "--> SNR map:\t{}\n".format(getinfo('snr_map')))

    def update_model(self, fit_type='gaussian'):
        """
        Tie a model to a SubCube. Didn't test it
        on anything but gaussian fitter so far.
        Yeap, I don't understand how models work.
        """
        try:
            allowed_fitters = self.specfit.Registry.multifitters
            self.specfit.fitter = allowed_fitters[fit_type]
        except KeyError:
            # TODO: get other models through add_fitter Registry method!
            raise ValueError('Unsupported fit type: %s\n'
                             'Choose one from %s' 
                             % (fit_type, allowed_fitters.keys()))
        self.specfit.fittype = fit_type
        self.fittype = fit_type

    def make_guess_grid(self, minpars, maxpars, finesse, 
            fixed=None, limitedmin=None, limitedmax=None):
        """
        Given parameter ranges and a finesse parameter, generate a grid of 
        guesses in a parameter space to be iterated upon in self.best_guess
        Maybe if parlimits arg is None we can look into parinfo?

        Parameters
        ----------
        minpars : an iterable containing minimal parameter values

        maxpars : an iterable containing maximal parameter values

        finesse : an integer or 1xNpars list/array setting the size 
                  of cells between minimal and maximal values in 
                  the resulting guess grid

        fixed : an iterable of booleans setting whether or not to fix the
                fitting parameters. Will be passed to Cube.fiteach, defaults
                to an array of False-s.

        limitedmin : an iterable of booleans controlling if the fit fixed
                     the minimal boundary of from minpars.

        limitedmax : an iterable of booleans controlling if the fit fixed
                     the maximal boundary of from maxpars.

        Returns
        -------
        guess_grid : a grid of guesses to use for SubCube.generate_model

        In addition, it saves a number of variables under self as a dictionary
        passed later to Cube.fiteach as additional arguments, with keywords:
        ['fixed', 'limitedmin', 'limitedmax', 'minpars', 'maxpars']
        """
        minpars, maxpars = np.asarray([minpars, maxpars])
        truths, falses = np.ones(minpars.shape, dtype=bool), \
                         np.zeros(minpars.shape, dtype=bool)
        
        fixed = falses if fixed is None else fixed
        limitedmin = truths if limitedmin is None else limitedmin
        limitedmax = truths if limitedmax is None else limitedmax
        # TODO: why does 'fixed' break the gaussian fitter?
        self.fiteach_args = {#'fixed'     : fixed,
                             'limitedmin': limitedmin,
                             'limitedmax': limitedmax,
                             'minpars'   : minpars,
                             'maxpars'   : maxpars    }
 
        # TODO: make sure you return the same shape!
        input_shape = minpars.shape
        npars = minpars.size

        # conformity for finesse: int or np.array goes in and np.array goes out
        finesse = np.atleast_1d(finesse) * np.ones(npars)
            
        par_space = []
        for i_len, i_min, i_max in zip(finesse, minpars, maxpars):
            par_space.append(np.linspace(i_min, i_max, i_len))

        nguesses = np.prod(map(len,par_space))
        guess_grid = np.array(np.meshgrid(*par_space)).reshape(npars, nguesses).T
        
        # this flips all the dimensions, might have some troubles should 
        # I want to scale this up later for a multidimensional case
        self.guess_grid = guess_grid

        return guess_grid

    def generate_model(self, guess_grid=None):
        """
        Generates a grid of spectral models matching the
        shape of the input guess_grid array. Can take the
        following numpy arrays as an input:

        Parameters
        ----------
        guess_grid : numpy.array
                     A grid of input parameters. 
                     Can be one of the following:
                     1) An (M,)-shaped 1d array of model parameters
                        (see pyspeckit docs for details)
                     2) An (N, M) array to compute N models 
                        for M sets of parameters
                     3) A guess cube of (Y, X, M) size
                     4) An (N, Y, X, M)-shaped array, to 
                        iterate over cubes of guesses.
                     If not set, SubCube.guess_grid is used.

        Returns
        -------
        model_grid : a grid of spectral models, following the 
                     shape of guess_grid. Also saved as an 
                     instance under SubCube.model_grid
        """
        if guess_grid is None:
            try:
                guess_grid = self.guess_grid
            except AttributeError:
                raise RuntimeError("Can't find the guess grid to use,")

        # safeguards preventing wrong output shapes
        npars = self.specfit.fitter.npars
        guess_grid = np.atleast_2d(guess_grid)
        grid_shape = guess_grid.shape[:-1]
        if ((len(grid_shape)>1 and grid_shape[-2:]!=self.cube.shape[1:]) or
            (len(grid_shape)>3) or (guess_grid.shape[-1]!=npars)):
            raise ValueError("Invalid shape for the guess_grid, "
                             "check the docsting for details.")

        model_grid = np.empty(shape=grid_shape+tuple([self.xarr.size]))
        # NOTE: this for loop is the performance bottleneck!
        # would be nice if I could broadcast grid_grid to n_modelfunc...
        for idx in np.ndindex(grid_shape):
            model_grid[idx] = \
                self.specfit.get_full_model(pars=guess_grid[idx])

        self.model_grid = model_grid
        return model_grid<commit_msg>allow for multiple velocity components in generate_model<commit_after>import numpy as np
import matplotlib.pylab as plt
import astropy.units as u
from astropy import log
import pyspeckit

# TODO: make informative log.info messages in
#       the methods of the SubCube class

# TODO: redefine SpectralCube.py#L795-L823 to
#       include the option of choosing between
#       the guesses in the middle of fiteach run

class SubCube(pyspeckit.Cube):
    """
    An extension of Cube, tinkered to be an instance of MultiCube, from which 
    it receives references to instances of pyspeckit.Cube that do not depend 
    on a spectral model chosen (so that parent MultiCube doesn't weigh so much)

    Is designed to have methods that operate within a single spectral model.
    """
    def __init__(self, *args, **kwargs):
        super(SubCube, self).__init__(*args, **kwargs)

        # because that UnitConversionError pops up way too often
        if self.xarr.velocity_convention is None:
            self.xarr.velocity_convention = 'radio'
        
        # so I either define some things as `None` 
        # or I'll have to call hasattr or them...
        # TODO: which is a more Pythonic approach?
        # A: probably the hasattr method, see here:  
        # http://programmers.stackexchange.com/questions/
        # 254576/is-it-a-good-practice-to-declare-instance
        # -variables-as-none-in-a-class-in-python
        self.guess_grid = None
        self.model_grid = None

    def info(self):
        def getinfo(target):
            try:
                return getattr(getattr(self,target), 'shape')
            except AttributeError:
                return 'N/A'
                pass
        log.info("Shapes of the arrays:\n"
                 "--> Data cube:\t{}\n".format(getinfo('cube'))+\
                 "--> Guess grid:\t{}\n".format(getinfo('guess_grid'))+\
                 "--> Model grid:\t{}\n".format(getinfo('model_grid'))+\
                 "--> Result rms:\t{}\n".format(getinfo('_residual_rms'))+\
                 "--> SNR map:\t{}\n".format(getinfo('snr_map')))

    def update_model(self, fit_type='gaussian'):
        """
        Tie a model to a SubCube. Didn't test it
        on anything but gaussian fitter so far.
        Yeap, I don't understand how models work.
        """
        try:
            allowed_fitters = self.specfit.Registry.multifitters
            self.specfit.fitter = allowed_fitters[fit_type]
        except KeyError:
            # TODO: get other models through add_fitter Registry method!
            raise ValueError('Unsupported fit type: %s\n'
                             'Choose one from %s' 
                             % (fit_type, allowed_fitters.keys()))
        self.specfit.fittype = fit_type
        self.fittype = fit_type

    def make_guess_grid(self, minpars, maxpars, finesse, 
            fixed=None, limitedmin=None, limitedmax=None):
        """
        Given parameter ranges and a finesse parameter, generate a grid of 
        guesses in a parameter space to be iterated upon in self.best_guess
        Maybe if parlimits arg is None we can look into parinfo?

        Parameters
        ----------
        minpars : an iterable containing minimal parameter values

        maxpars : an iterable containing maximal parameter values

        finesse : an integer or 1xNpars list/array setting the size 
                  of cells between minimal and maximal values in 
                  the resulting guess grid

        fixed : an iterable of booleans setting whether or not to fix the
                fitting parameters. Will be passed to Cube.fiteach, defaults
                to an array of False-s.

        limitedmin : an iterable of booleans controlling if the fit fixed
                     the minimal boundary of from minpars.

        limitedmax : an iterable of booleans controlling if the fit fixed
                     the maximal boundary of from maxpars.

        Returns
        -------
        guess_grid : a grid of guesses to use for SubCube.generate_model

        In addition, it saves a number of variables under self as a dictionary
        passed later to Cube.fiteach as additional arguments, with keywords:
        ['fixed', 'limitedmin', 'limitedmax', 'minpars', 'maxpars']
        """
        minpars, maxpars = np.asarray([minpars, maxpars])
        truths, falses = np.ones(minpars.shape, dtype=bool), \
                         np.zeros(minpars.shape, dtype=bool)
        
        fixed = falses if fixed is None else fixed
        limitedmin = truths if limitedmin is None else limitedmin
        limitedmax = truths if limitedmax is None else limitedmax
        # TODO: why does 'fixed' break the gaussian fitter?
        self.fiteach_args = {#'fixed'     : fixed,
                             'limitedmin': limitedmin,
                             'limitedmax': limitedmax,
                             'minpars'   : minpars,
                             'maxpars'   : maxpars    }
 
        # TODO: make sure you return the same shape!
        input_shape = minpars.shape
        npars = minpars.size

        # conformity for finesse: int or np.array goes in and np.array goes out
        finesse = np.atleast_1d(finesse) * np.ones(npars)
            
        par_space = []
        for i_len, i_min, i_max in zip(finesse, minpars, maxpars):
            par_space.append(np.linspace(i_min, i_max, i_len))

        nguesses = np.prod(map(len,par_space))
        guess_grid = np.array(np.meshgrid(*par_space)).reshape(npars, nguesses).T
        
        # this flips all the dimensions, might have some troubles should 
        # I want to scale this up later for a multidimensional case
        self.guess_grid = guess_grid

        return guess_grid

    def generate_model(self, guess_grid=None):
        """
        Generates a grid of spectral models matching the
        shape of the input guess_grid array. Can take the
        following numpy arrays as an input:

        Parameters
        ----------
        guess_grid : numpy.array
                     A grid of input parameters. 
                     Can be one of the following:
                     1) An (M,)-shaped 1d array of model parameters
                        (see pyspeckit docs for details)
                     2) An (N, M) array to compute N models 
                        for M sets of parameters
                     3) A guess cube of (Y, X, M) size
                     4) An (N, Y, X, M)-shaped array, to 
                        iterate over cubes of guesses.
                     If not set, SubCube.guess_grid is used.

        Returns
        -------
        model_grid : a grid of spectral models, following the 
                     shape of guess_grid. Also saved as an 
                     instance under SubCube.model_grid
        """
        if guess_grid is None:
            try:
                guess_grid = self.guess_grid
            except AttributeError:
                raise RuntimeError("Can't find the guess grid to use,")

        # safeguards preventing wrong output shapes
        npars = self.specfit.fitter.npars
        guess_grid = np.atleast_2d(guess_grid)
        grid_shape = guess_grid.shape[:-1]
        if ((len(grid_shape)>1 and grid_shape[-2:]!=self.cube.shape[1:]) or
            (len(grid_shape)>3) or (guess_grid.shape[-1]%npars)):
            raise ValueError("Invalid shape for the guess_grid, "
                             "check the docsting for details.")

        model_grid = np.empty(shape=grid_shape+tuple([self.xarr.size]))
        # NOTE: this for loop is the performance bottleneck!
        # would be nice if I could broadcast grid_grid to n_modelfunc...
        for idx in np.ndindex(grid_shape):
            model_grid[idx] = \
                self.specfit.get_full_model(pars=guess_grid[idx])

        self.model_grid = model_grid
        return model_grid
{'fraction_non_alphanumeric': 0.08077166350411132, 'fraction_numerical': 0.003289057558507274, 'mean_word_length': 1.8442165857168555, 'pattern_counts': {'":': 0, '<': 3, '<?xml version=': 0, '>': 17, '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': '1713757', 'n_tokens_mistral': 4423, 'n_tokens_neox': 4222, 'n_words': 1453}