<commit_before>#!/usr/bin/env python
"""
TODO:   Check behavior when cubes have missing channels.
        Implement parallelization in velocity.
"""

import numpy as np
import glob
import sys
import argparse
import logging
from astropy.io import fits
from crrlpy import crrls
from crrlpy import utils
from scipy.interpolate import RegularGridInterpolator
from datetime import datetime
startTime = datetime.now()

def parse_cube_list(cubes):
    """
    """

    logger = logging.getLogger(__name__)
    
    if glob.glob(cubes):
        return glob.glob(cubes)

    else:
        if len(cubes.split(',')) > 1:
            return cubes.split(',')
        else:
            logger.info('Input cube list not understood.')

def stack_cubes(cubes, outfits, vmax, vmin, dv, weight_list=None, v_axis=3, overwrite=False, algo='channel'):
    """
    """
    
    logger = logging.getLogger(__name__)
    
    cubel = cubes #glob.glob(cubes)
    logger.debug(cubel)
    
    if dv == 0:
        logger.info('Velocity width not specified.'), 
        logger.info('Will try to determine from input data.')
        for i,cube in enumerate(cubel):
            hdu = fits.open(cube)
            head = hdu[0].header
            x = crrls.get_axis(head, v_axis)
            if i == 0:
                dv = utils.get_min_sep(x)
                vmax_min = max(x)
                vmin_max = min(x)
            else:
                dv = max(dv, utils.get_min_sep(x))
                vmax_min = min(vmax_min, max(x))
                vmin_max = max(vmin_max, min(x))
            logger.debug('Cube: {0}'.format(cube))
            logger.debug('Cube velocity limits: {0} {1} {2}'.format(min(x), max(x), utils.get_min_sep(x)))
        logger.info('Will use a velocity width of {0}'.format(dv))
        
        # Check velocity ranges to avoid latter crashes.
        if vmax_min < vmax:
            logger.info('Requested maximum velocity is larger '\
                        'than one of the cubes velocity axis.')
            #logger.info('Conflicting cube: {0}'.format(cube))
            logger.info('v_max={0}, v_max_min={1}'.format(vmax, vmax_min))
            logger.info('Will now exit')
            sys.exit(1)
        if vmin_max > vmin:
            logger.info('Requested minimum velocity is smaller '\
                        'than one of the cubes velocity axis.')
            #logger.info('Conflicting cube: {0}'.format(cube))
            logger.info('v_min={0}, v_min_max={1}'.format(vmin, vmin_max))
            logger.info('Will now exit')
            sys.exit(1)

    shape = hdu[0].shape
    if len(shape) > 3:
        logger.info('Will drop first axis.')
        s = 1
    else:
        s = 0
    
    nvaxis = np.arange(vmin, vmax, dv)
    stack = np.ma.zeros((len(nvaxis),) + shape[s+1:])
    covrg = np.zeros((len(nvaxis),) + shape[s+1:])

    # Check if there is a weight list
    if weight_list:
        try:
            wl = np.loadtxt(weight_list, dtype=[('fits', np.str_, 256), ('w', np.float64)])
        except ValueError:
            wl = np.loadtxt(weight_list, dtype=[('fits', np.str_, 256), ('w', np.str_, 256)])
    weight = np.ma.zeros((len(nvaxis),) + shape[s+1:])
        
    logger.info('Output stack will have dimensions: {0}'.format(stack.shape))
    
    for i,cube in enumerate(cubel):
        
        logger.info('Adding cube {0}/{1}'.format(i, len(cubel)-1))

        # Load the data
        hdu = fits.open(cube)
        head = hdu[0].header
        data = np.ma.masked_invalid(hdu[0].data)

        # Check the weight
        if weight_list:
            w = wl['w'][np.where(wl['fits'] == cube)]
        else:
            w = 1
        logger.info('Will use a weight of {0}.'.format(w))
        try:
            aux_weight = np.ones(stack.shape)*w
        except TypeError:
            w = np.ma.masked_invalid(fits.open(w[0])[0].data)
        #aux_weight = np.ones(shape[s+1:])*w
        
        if len(data.shape) > 3:
            logger.info('Will drop first axis.')
            data = data[0]
        
        # Get the cube axes
        ra = crrls.get_axis(head, 1)
        de = crrls.get_axis(head, 2)
        ve = crrls.get_axis(head, v_axis)

        logger.info('RA axis limits: {0} {1}'.format(min(ra), max(ra)))
        logger.info('DEC axis limits: {0} {1}'.format(min(de), max(de)))
        logger.info('VEL axis limits: {0} {1}'.format(min(ve), max(ve)))
       
        vmin_idx = utils.best_match_indx(vmin, ve)
        vmax_idx = utils.best_match_indx(vmax, ve)
        [vmin_idx,vmax_idx] = sorted([vmin_idx,vmax_idx])
        vmin_idx -= 1
        vmax_idx += 1
        data_ = data[vmin_idx:vmax_idx]
        ve_ = ve[vmin_idx:vmax_idx]<commit_msg>Fixed shape 0 velocity axis while stacking cubes<commit_after>#!/usr/bin/env python
"""
TODO:   Check behavior when cubes have missing channels.
        Implement parallelization in velocity.
"""

import numpy as np
import glob
import sys
import argparse
import logging
from astropy.io import fits
from crrlpy import crrls
from crrlpy import utils
from scipy.interpolate import RegularGridInterpolator
from datetime import datetime
startTime = datetime.now()

def parse_cube_list(cubes):
    """
    """

    logger = logging.getLogger(__name__)
    
    if glob.glob(cubes):
        return glob.glob(cubes)

    else:
        if len(cubes.split(',')) > 1:
            return cubes.split(',')
        else:
            logger.info('Input cube list not understood.')

def stack_cubes(cubes, outfits, vmax, vmin, dv, weight_list=None, v_axis=3, overwrite=False, algo='channel'):
    """
    """
    
    logger = logging.getLogger(__name__)
    
    cubel = cubes #glob.glob(cubes)
    logger.debug(cubel)
    
    if dv == 0:
        logger.info('Velocity width not specified.'), 
        logger.info('Will try to determine from input data.')
        for i,cube in enumerate(cubel):
            hdu = fits.open(cube)
            head = hdu[0].header
            x = crrls.get_axis(head, v_axis)
            if i == 0:
                dv = utils.get_min_sep(x)
                vmax_min = max(x)
                vmin_max = min(x)
            else:
                dv = max(dv, utils.get_min_sep(x))
                vmax_min = min(vmax_min, max(x))
                vmin_max = max(vmin_max, min(x))
            logger.debug('Cube: {0}'.format(cube))
            logger.debug('Cube velocity limits: {0} {1} {2}'.format(min(x), max(x), utils.get_min_sep(x)))
        logger.info('Will use a velocity width of {0}'.format(dv))
        
        # Check velocity ranges to avoid latter crashes.
        if vmax_min < vmax:
            logger.info('Requested maximum velocity is larger '\
                        'than one of the cubes velocity axis.')
            #logger.info('Conflicting cube: {0}'.format(cube))
            logger.info('v_max={0}, v_max_min={1}'.format(vmax, vmax_min))
            logger.info('Will now exit')
            sys.exit(1)
        if vmin_max > vmin:
            logger.info('Requested minimum velocity is smaller '\
                        'than one of the cubes velocity axis.')
            #logger.info('Conflicting cube: {0}'.format(cube))
            logger.info('v_min={0}, v_min_max={1}'.format(vmin, vmin_max))
            logger.info('Will now exit')
            sys.exit(1)

    shape = hdu[0].shape
    if len(shape) > 3:
        logger.info('Will drop first axis.')
        s = 1
    else:
        s = 0
    
    nvaxis = np.arange(vmin, vmax, dv)
    stack = np.ma.zeros((len(nvaxis),) + shape[s+1:])
    covrg = np.zeros((len(nvaxis),) + shape[s+1:])

    # Check if there is a weight list
    if weight_list:
        try:
            wl = np.loadtxt(weight_list, dtype=[('fits', np.str_, 256), ('w', np.float64)])
        except ValueError:
            wl = np.loadtxt(weight_list, dtype=[('fits', np.str_, 256), ('w', np.str_, 256)])
    weight = np.ma.zeros((len(nvaxis),) + shape[s+1:])
        
    logger.info('Output stack will have dimensions: {0}'.format(stack.shape))
    
    for i,cube in enumerate(cubel):
        
        logger.info('Adding cube {0}/{1}'.format(i, len(cubel)-1))

        # Load the data
        hdu = fits.open(cube)
        head = hdu[0].header
        data = np.ma.masked_invalid(hdu[0].data)

        # Check the weight
        if weight_list:
            w = wl['w'][np.where(wl['fits'] == cube)]
        else:
            w = 1
        logger.info('Will use a weight of {0}.'.format(w))
        try:
            aux_weight = np.ones(stack.shape)*w
        except TypeError:
            w = np.ma.masked_invalid(fits.open(w[0])[0].data)
        #aux_weight = np.ones(shape[s+1:])*w
        
        if len(data.shape) > 3:
            logger.info('Will drop first axis.')
            data = data[0]
        
        # Get the cube axes
        ra = crrls.get_axis(head, 1)
        de = crrls.get_axis(head, 2)
        ve = crrls.get_axis(head, v_axis)

        logger.info('RA axis limits: {0} {1}'.format(min(ra), max(ra)))
        logger.info('DEC axis limits: {0} {1}'.format(min(de), max(de)))
        logger.info('VEL axis limits: {0} {1}'.format(min(ve), max(ve)))
       
        vmin_idx = utils.best_match_indx(vmin, ve)
        vmax_idx = utils.best_match_indx(vmax, ve)
        [vmin_idx,vmax_idx] = sorted([vmin_idx,vmax_idx])
        if vmin_idx > 0: vmin_idx -= 1
        if vmax_idx < len(ve): vmax_idx += 1
        data_ = data[vmin_idx:vmax_idx]
        ve_ = ve[vmin_idx:vmax_idx]
{'fraction_non_alphanumeric': 0.13837540931657336, 'fraction_numerical': 0.01267561001373191, 'mean_word_length': 2.084039087947883, 'pattern_counts': {'":': 0, '<': 6, '<?xml version=': 0, '>': 12, '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': '476939', 'n_tokens_mistral': 3311, 'n_tokens_neox': 3180, 'n_words': 734}