<commit_before>'''
Jam arrange: Arrangement algorithm
Author: Numstack (Pty) Ltd
Copyright 2016
Objective: Identify mp3, wav, etc. files and arange them accordingly in desired local storage folder.
'''

'''
Importing all necessary libraries
'''
import os
import os.path
from os import listdir
from os.path import isfile, join
import shutil
from tinytag import TinyTag

'''
Repalcing special characters in file name
'''

def replace_special_chars_files(temp):
    reserved_chars = ['"',"'"]
    for i in range(len(reserved_chars)):
        if reserved_chars[i] in temp:
            temp = temp.replace(reserved_chars[i],' ')
    return temp

'''
Repalcing special characters in a dir name
'''

def replace_special_chars(temp):
    reserved_chars = ['<', '>', ':', '"', '/', '\\', '|', '?', '*']  # see tutorial for explanation
    if temp is not None:
        for i in range(len(reserved_chars)):
            if reserved_chars[i] in temp:
                temp = temp.replace(reserved_chars[i], ' ')
    else:
        pass
    return temp

'''
Collect audio files' details in chosen directory
'''

def collect_audio(target):
    '''
    finding all files in dir and subdir
    '''
    root = target
    path = os.path.join(root, "targetdirectory")
    all_files_in_dir = []
    all_files_dir = []
    for path, subdirs, files in os.walk(root):
        for name in files:
            all_files_in_dir.append(replace_special_chars(str(name)))
            all_files_dir.append(path)
    audio_files = []
    '''
                            Extracted Audio file details
    1. Artist
    2. Title (song name)
    3. Album title
    4. Path on local storage
    5. Track number
    6. Total number of tracks on specific album
    7. Tracking id generation
    '''
    audio_file_deets = [[], [], [], [], [], [], []]
    files_not_parsed = []
    count = 0
    for i in range(len(all_files_in_dir)):
        if all_files_in_dir[i].endswith(('.mp3', '.wav', '.MP3', '.wma', '.WMA', '.WAV', '.mp4', '.MP4')) == True:
            temp = os.path.join(all_files_dir[i],str(all_files_in_dir[i]))
            current_audiofile = TinyTag.get(temp)
            if current_audiofile is not None:
                curr_artist = replace_special_chars(current_audiofile.artist)
                if curr_artist == '':
                    curr_artist = None
                    audio_file_deets[0].append(curr_artist)
                else:
                    audio_file_deets[0].append(curr_artist)
                audio_file_deets[1].append(current_audiofile.title)
                audio_file_deets[2].append(current_audiofile.album)
                audio_file_deets[3].append(temp)
                audio_file_deets[4].append(current_audiofile.track)
                audio_file_deets[5].append(current_audiofile.track_total)
                audio_file_deets[6].append(count)
                count = count + 1
                audio_files.append(all_files_in_dir[i])
            if current_audiofile == None:
                files_not_parsed.append(temp)
    return audio_file_deets, files_not_parsed, all_files_in_dir, all_files_dir

'''
Making a list of all known artist names to use later when creating song
storage folders.
'''

def music_handling(audio_file_deets):
    artist_names = []
    for i in range(len(audio_file_deets[0])):
        '''
                            NB: Attention!!!!!!!!!!!
        Unicode testing function must enter here once completed:
        isinstance(a, unicode)
        '''
        if audio_file_deets[0][i] is not None:
            artist_names.append(str(audio_file_deets[0][i].strip()
                                    .encode('utf-8')))
        if audio_file_deets[0][i] is None:
            artist_names.append(i)

    folder_titles = []
    unknown_songs = [[],[]]

    for i in range(len(artist_names)):
        temp = [x for x, val in enumerate(artist_names)
                if val == artist_names[i]]
        if len(temp) > 1:
            if artist_names[i] in folder_titles:
                pass
            else:
                folder_titles.append(artist_names[i])
        elif len(temp) == 1:
            temp2 = isinstance(artist_names[i], str)
            if temp2 == True:
                folder_titles.append(artist_names[i])
            else:
                unknown_songs[0].append(artist_names[i])
                unknown_songs[1].append(audio_file_deets[3][i])

    return artist_names, folder_titles, unknown_songs

'''
Make folders for all songs with known artist meta data.
'''

def making_arranged_dir(target, folder_titles):
    moveto = os.path.join(target, 'Arranged files')
    os.mkdir(os.path.join(target, 'Arranged files'))
    new_dirs = []
    for i in range(len(folder_titles)):
        temp = os.path.join(moveto, folder_titles[i])
        if temp.strip().lower() not in new_dirs:
            os.mkdir(temp)
            new_dirs.append(temp.strip())
    unknown_dir = os.path.join(moveto, 'Unknown artists').strip()
    os.mkdir(os.path.join(moveto, 'Unknown artists'))
    return moveto, new_dirs, unknown_dir

'''
Move files to respective folders
'''

def movefiles():
    pass

'''
Delete files while testing
'''

def delete_files(moveto):
    user_input = raw_input('''would you like to delete newly created
                           files: [y/n]\n''')
    if user_input.strip() == 'y':
        shutil.rmtree(moveto)

'''
copy all files to right places
'''

def copy_to_arranged(new_dirs, song_locations, folder_titles, artist_names, unknown_dir):
    for i in range(len(artist_names)):
        for j in range(len(folder_titles)):
            str_test = isinstance(artist_names[i], str)
            if (artist_names[i] == folder_titles[j]) and (str_test == True):
                shutil.copy(song_locations[i].strip(), new_dirs[j].strip())

'''
Copy all files that need meta data identification, to a new repository.
Allowing the audio fingerprinting process to be simplified and centrallised,
'''

def copy_all_unknowns(unknown_songs, unknown_dir):
    for i in range(len(unknown_songs[1])):
        shutil.copy(unknown_songs[1][i], unknown_dir)

'''
Rename files: If tags i.e. meta data doesn't match visible file name on pc
it can be difficult for the user to know what song they are looking at.
'''

def rename_files():
    pass

'''
MAIN CODE AND MAIN FUNCTION
'''

if __name__ == "__main__":
    # target = "C:\\Users\\User\\Music\\test folder2"
    # target = "C:\\Users\\User\\Music\\Hip Hop\\2013 Hit singles\\Release 1.1"
    target = "/home/victor/Music"
    audio_details = collect_audio(target)
    handle_files = music_handling(audio_details[0])
    create_folders = making_arranged_dir(target, handle_files[1])
    copy_to_arranged(create_folders[1], audio_details[0][3], handle_files[1], handle_files[0], create_folders[2])
    copy_all_unknowns(handle_files[2], create_folders[2])

##    delete_arrangement = delete_files(create_folders[0]) #Delete demo files
<commit_msg>added delete unknown folder for all songs known.<commit_after>'''
Jam arrange: Arrangement algorithm
Author: Numstack (Pty) Ltd
Copyright 2016
Objective: Identify mp3, wav, etc. files and arange them accordingly in desired local storage folder.
'''

'''
Importing all necessary libraries
'''
import os
import os.path
from os import listdir
from os.path import isfile, join
import shutil
from tinytag import TinyTag

'''
Repalcing special characters in file name
'''

def replace_special_chars_files(temp):
    reserved_chars = ['"',"'"]
    for i in range(len(reserved_chars)):
        if reserved_chars[i] in temp:
            temp = temp.replace(reserved_chars[i],' ')
    return temp

'''
Repalcing special characters in a dir name
'''

def replace_special_chars(temp):
    reserved_chars = ['<', '>', ':', '"', '/', '\\', '|', '?', '*']  # see tutorial for explanation
    if temp is not None:
        for i in range(len(reserved_chars)):
            if reserved_chars[i] in temp:
                temp = temp.replace(reserved_chars[i], ' ')
    else:
        pass
    return temp

'''
Collect audio files' details in chosen directory
'''

def collect_audio(target):
    '''
    finding all files in dir and subdir
    '''
    root = target
    path = os.path.join(root, "targetdirectory")
    all_files_in_dir = []
    all_files_dir = []
    for path, subdirs, files in os.walk(root):
        for name in files:
            all_files_in_dir.append(replace_special_chars(str(name)))
            all_files_dir.append(path)
    audio_files = []
    '''
                            Extracted Audio file details
    1. Artist
    2. Title (song name)
    3. Album title
    4. Path on local storage
    5. Track number
    6. Total number of tracks on specific album
    7. Tracking id generation
    '''
    audio_file_deets = [[], [], [], [], [], [], []]
    files_not_parsed = []
    count = 0
    for i in range(len(all_files_in_dir)):
        if all_files_in_dir[i].endswith(('.mp3', '.wav', '.MP3', '.wma', '.WMA', '.WAV', '.mp4', '.MP4')) == True:
            temp = os.path.join(all_files_dir[i],str(all_files_in_dir[i]))
            current_audiofile = TinyTag.get(temp)
            if current_audiofile is not None:
                curr_artist = replace_special_chars(current_audiofile.artist)
                if curr_artist == '':
                    curr_artist = None
                    audio_file_deets[0].append(curr_artist)
                else:
                    audio_file_deets[0].append(curr_artist)
                audio_file_deets[1].append(current_audiofile.title)
                audio_file_deets[2].append(current_audiofile.album)
                audio_file_deets[3].append(temp)
                audio_file_deets[4].append(current_audiofile.track)
                audio_file_deets[5].append(current_audiofile.track_total)
                audio_file_deets[6].append(count)
                count = count + 1
                audio_files.append(all_files_in_dir[i])
            if current_audiofile == None:
                files_not_parsed.append(temp)
    return audio_file_deets, files_not_parsed, all_files_in_dir, all_files_dir

'''
Making a list of all known artist names to use later when creating song
storage folders.
'''

def music_handling(audio_file_deets):
    artist_names = []
    for i in range(len(audio_file_deets[0])):
        '''
                            NB: Attention!!!!!!!!!!!
        Unicode testing function must enter here once completed:
        isinstance(a, unicode)
        '''
        if audio_file_deets[0][i] is not None:
            artist_names.append(str(audio_file_deets[0][i].strip()
                                    .encode('utf-8')))
        if audio_file_deets[0][i] is None:
            artist_names.append(i)

    folder_titles = []
    unknown_songs = [[],[]]

    for i in range(len(artist_names)):
        temp = [x for x, val in enumerate(artist_names)
                if val == artist_names[i]]
        if len(temp) > 1:
            if artist_names[i] in folder_titles:
                pass
            else:
                folder_titles.append(artist_names[i])
        elif len(temp) == 1:
            temp2 = isinstance(artist_names[i], str)
            if temp2 == True:
                folder_titles.append(artist_names[i])
            else:
                unknown_songs[0].append(artist_names[i])
                unknown_songs[1].append(audio_file_deets[3][i])

    return artist_names, folder_titles, unknown_songs

'''
Make folders for all songs with known artist meta data.
'''

def making_arranged_dir(target, folder_titles):
    moveto = os.path.join(target, 'Arranged files')
    os.mkdir(os.path.join(target, 'Arranged files'))
    new_dirs = []
    for i in range(len(folder_titles)):
        temp = os.path.join(moveto, folder_titles[i])
        if temp.strip().lower() not in new_dirs:
            os.mkdir(temp)
            new_dirs.append(temp.strip())
    unknown_dir = os.path.join(moveto, 'Unknown artists').strip()
    os.mkdir(os.path.join(moveto, 'Unknown artists'))
    return moveto, new_dirs, unknown_dir

'''
Move files to respective folders
'''

def movefiles():
    pass

'''
Delete unknown folder if no unknown songs found
'''
def clean_unknown_folder(unknown_songs, unknown_dir):
    if unknown_songs == [[],[]]:
        shutil.rmtree(unknown_dir)
    else:
        pass

'''
Delete files while testing
'''

def delete_files(moveto):
    user_input = raw_input('''would you like to delete newly created
                           files: [y/n]\n''')
    if user_input.strip() == 'y':
        shutil.rmtree(moveto)

'''
copy all files to right places
'''

def copy_to_arranged(new_dirs, song_locations, folder_titles, artist_names, unknown_dir):
    for i in range(len(artist_names)):
        for j in range(len(folder_titles)):
            str_test = isinstance(artist_names[i], str)
            if (artist_names[i] == folder_titles[j]) and (str_test == True):
                shutil.copy(song_locations[i].strip(), new_dirs[j].strip())

'''
Copy all files that need meta data identification, to a new repository.
Allowing the audio fingerprinting process to be simplified and centrallised,
'''

def copy_all_unknowns(unknown_songs, unknown_dir):
    for i in range(len(unknown_songs[1])):
        shutil.copy(unknown_songs[1][i], unknown_dir)

'''
Rename files: If tags i.e. meta data doesn't match visible file name on pc
it can be difficult for the user to know what song they are looking at.
'''

def rename_files():
    pass

'''
MAIN CODE AND MAIN FUNCTION
'''

if __name__ == "__main__":
    target = "/Users/victor/Music/test"
    audio_details = collect_audio(target)
    handle_files = music_handling(audio_details[0])
    create_folders = making_arranged_dir(target, handle_files[1])
    copy_to_arranged(create_folders[1], audio_details[0][3], handle_files[1], handle_files[0], create_folders[2])
    copy_all_unknowns(handle_files[2], create_folders[2])
    clean_arranged_folder = clean_unknown_folder(handle_files[2], create_folders[2])

##    delete_arrangement = delete_files(create_folders[0]) #Delete demo files
<|endoftext|>


<commit_before>
'''
Routines for feathering large cubes together using uvcombine.

Some of these *should* be moved into uvcombine in good time.
'''

from spectral_cube import SpectralCube
from uvcombine.uvcombine import feather_simple, feather_compare
from astropy import log
import astropy.units as u

import numpy as np

from .io_utils import create_huge_fits
from .progressbar import _map_context

def feather_cube(cube_hi, cube_lo, verbose=True, save_feather=True,
                 save_name=None, num_cores=1, chunk=100,
                 restfreq=1.42040575177 * u.GHz,
                 weights=1.,
                 feather_kwargs={}):
    '''
    Feather two cubes together. The spectral axis of the cubes *must match*.
    This function is specifically for cube too large to fit in memory. The
    feathered planes are saved directly to a new FITS file.
    '''

    # Ensure the spectral axes are the same.
    if not np.allclose(cube_hi.spectral_axis.to(cube_lo._spectral_unit).value,
                       cube_lo.spectral_axis.value):
        raise Warning("The spectral axes do not match. Spectrally regrid the "
                      "cubes to be the same before feathering.")

    if save_feather:
        if save_name is None:
            raise TypeError("save_name must be given the save the shifted "
                            "cube.")

        output_hdu = create_huge_fits(save_name, cube_hi.header,
                                      return_hdu=True)
    else:
        newcube = np.empty(cube_hi.shape, dtype=cube_hi[:1, 0, 0].dtype)

    num_chans = cube_hi.shape[0]
    chunked_channels = get_channel_chunks(num_chans, chunk)

    freq_axis = cube_hi.with_spectral_unit(u.Hz, velocity_convention='radio',
                                           rest_value=restfreq).spectral_axis

    if weights is not None:
        feather_kwargs['weights'] = weights

    for i, chunk_chans in enumerate(chunked_channels):

        log.info("On chunk {0} of {1}".format(i + 1, len(chunked_channels)))

        changen = ((chan, cube_hi[chan:chan + 1], cube_lo[chan:chan + 1],
                    freq_axis[chan], feather_kwargs) for chan in chunk_chans)

        with _map_context(num_cores, verbose=verbose,
                          num_jobs=len(chunk_chans)) as map:
            log.info("Feathering")
            output = map(_feather, changen)

        for chan, shifted_arr in output:
            if save_feather:
                output_hdu[0].data[chan] = shifted_arr.real
            else:
                newcube[chan] = shifted_arr.real

        if save_feather:
            output_hdu.flush()

    if save_feather:

        # Append the beam table, if needed
        if hasattr(cube_hi, 'beams'):
            from spectral_cube.cube_utils import beams_to_bintable
            output_hdu.append(beams_to_bintable(cube_hi.beams))

        output_hdu.close()
    else:
        return SpectralCube(newcube, wcs=cube_hi.wcs,
                            header=cube_hi.header, meta=cube_hi.meta)

def _feather(args):
    '''
    Feather 2D images together.
    '''
    chan, plane_hi, plane_lo, freq, kwargs = args

    # TODO: When Jy/beam is supported in SpectralCube, just let uvcombine
    # match the units. BUT we also need Slice objects to retain a 1 element
    # spectral axis!
    if plane_hi.unit != plane_lo.unit:
        if plane_hi.unit == u.K:
            if hasattr(plane_hi, 'beams'):
                beam = plane_hi.beams[0]
            else:
                beam = plane_hi.beam

            plane_hi = plane_hi.to(u.Jy / u.beam, plane_hi.beam.jtok_equiv(freq))

        if plane_lo.unit == u.K:
            if hasattr(plane_lo, 'beams'):
                beam = plane_lo.beams[0]
            else:
                beam = plane_lo.beam

            plane_lo = plane_lo.to(u.Jy / u.beam, beam.jtok_equiv(freq))

    if hasattr(plane_hi, 'hdu'):
        plane_hi_hdu = plane_hi.hdu
    else:
        plane_hi_hdu = plane_hi.hdulist[0]
        plane_hi_hdu.header.update(plane_hi.beams[0].to_header_keywords())

    if hasattr(plane_lo, 'hdu'):
        plane_lo_hdu = plane_lo.hdu
    else:
        plane_lo_hdu = plane_lo.hdulist[0]
        plane_lo_hdu.header.update(plane_lo.beams[0].to_header_keywords())

    feathered = feather_simple(plane_hi_hdu, plane_lo_hdu, **kwargs)

    # Expect that the interferometer image will cover a smaller region.
    # Cut to that region.
    feathered[np.isnan(plane_hi[0])] = np.NaN

    return chan, feathered

def feather_compare_cube(cube_hi, cube_lo, LAS, lowresfwhm=None,
                         restfreq=1.42040575177 * u.GHz,
                         verbose=True, num_cores=1, chunk=100,
                         weights=1.):
    '''
    Record the ratios of the flux in the overlap region between the cubes.

    **Assumes velocity axes use the radio convention!**

    '''

    # Ensure the spectral axes are the same.
    if not np.allclose(cube_hi.spectral_axis.to(cube_lo._spectral_unit).value,
                       cube_lo.spectral_axis.value):
        raise Warning("The spectral axes do not match. Spectrally regrid the "
                      "cubes to be the same before feathering.")

    num_chans = cube_hi.shape[0]
    chunked_channels = get_channel_chunks(num_chans, chunk)

    radii = []
    ratios = []
    highres_pts = []
    lowres_pts = []

    if lowresfwhm is None:
        lowresfwhm = cube_lo.beam.major.to(u.arcsec)
    else:
        lowresfwhm = lowresfwhm.to(u.arcsec)

    freq_axis = cube_hi.with_spectral_unit(u.Hz, velocity_convention='radio',
                                           rest_value=restfreq).spectral_axis

    for i, chunk_chans in enumerate(chunked_channels):

        log.info("On chunk {0} of {1}".format(i + 1, len(chunked_channels)))

        changen = ((chan, cube_hi[chan], cube_lo[chan],
                    LAS, lowresfwhm, freq_axis[chan],
                    weights) for chan in chunk_chans)

        with _map_context(num_cores, verbose=verbose,
                          num_jobs=len(chunk_chans)) as map:
            output = map(_compare, changen)

        chan_out = np.array([out[0] for out in output])

        for jj in chan_out.argsort():

            radii.append(output[jj][1][0])
            ratios.append(output[jj][1][1])
            highres_pts.append(output[jj][1][2])
            lowres_pts.append(output[jj][1][3])

    return radii, ratios, highres_pts, lowres_pts

def _compare(args):

    chan, plane_hi, plane_lo, LAS, lowresfwhm, freq, weights = args

    out = feather_compare(plane_hi.to(u.K, freq=freq).hdu,
                          plane_lo.hdu,
                          return_samples=True, doplot=False,
                          LAS=LAS, SAS=lowresfwhm,
                          lowresfwhm=lowresfwhm,
                          weights=weights)

    return chan, out

def flux_recovery(cube_hi, cube_lo, frequency=1.42040575177 * u.GHz,
                  verbose=True, doplot=False,
                  num_cores=1, chunk=100, mask=None):
    '''
    Calculate the fraction of flux recovered between the high-resolution
    (interferometer) cube and the low-resolution (single-dish) cube.
    '''

    # Ensure the spectral axes are the same.
    if not np.allclose(cube_hi.spectral_axis.to(cube_lo._spectral_unit).value,
                       cube_lo.spectral_axis.value):
        raise Warning("The spectral axes do not match. Spectrally regrid the "
                      "cubes to be the same before feathering.")

    if mask is not None:
        if mask.shape != cube_hi[0].shape:
            raise ValueError("mask must have the same shape as the "
                             "high-resolution cube.")

    num_chans = cube_hi.shape[0]
    chunked_channels = get_channel_chunks(num_chans, chunk)

    total_hires = np.empty(num_chans) * u.Jy
    total_lores = np.empty(num_chans) * u.Jy

    for i, chunk_chans in enumerate(chunked_channels):

        log.info("On chunk {0} of {1}".format(i + 1, len(chunked_channels)))

        changen = ((chan, cube_hi[chan], cube_lo[chan],
                    frequency, mask) for chan in chunk_chans)

        with _map_context(num_cores, verbose=verbose)as map:
            output = map(_totalplanes, changen)

        chan_out = np.array([out[0] for out in output])

        for jj in chan_out.argsort():

            total_hires[chan_out[jj]] = output[jj][1]
            total_lores[chan_out[jj]] = output[jj][2]

    if doplot:
        import matplotlib.pyplot as plt
        plt.subplot(121)
        plt.plot(total_hires.value, label='High-res')
        plt.plot(total_lores.value, label='Low-res')
        plt.ylabel("Total Intensity (Jy)")
        plt.legend()
        plt.subplot(122)
        plt.plot(total_hires.value / total_lores.value)
        plt.ylabel("Fraction of recovered")
        plt.xlabel("Channel")

    return total_hires, total_lores

def _totalplanes(args):

    chan, plane_hi, plane_lo, freq, mask = args

    if mask is not None:
        if not plane_hi.wcs.wcs.compare(plane_lo.wcs.wcs):
            # Reproject the low resolution onto the high

            plane_lo = plane_lo.reproject(plane_hi.header)

        # !! Update to Jy/bm (or equiv) when spectral-cube can handle it
        total_hi = \
            np.nansum(plane_hi[mask].to(u.Jy, plane_hi.beam.jtok_equiv(freq)))
        total_lo = \
            np.nansum(plane_lo[mask].to(u.Jy, plane_lo.beam.jtok_equiv(freq)))
    else:
        total_hi = \
            np.nansum(plane_hi.to(u.Jy, plane_hi.beam.jtok_equiv(freq)))
        total_lo = \
            np.nansum(plane_lo.to(u.Jy, plane_lo.beam.jtok_equiv(freq)))

    # Convert from Jy/beam to actual Jy
    total_hi *= (1 / plane_hi.beam.sr.to(u.deg**2)) * \
        (plane_hi.header["CDELT2"] * u.deg)**2
    total_lo *= (1 / plane_lo.beam.sr.to(u.deg**2)) * \
        (plane_lo.header["CDELT2"] * u.deg)**2

    return chan, total_hi, total_lo

def get_channel_chunks(num_chans, chunk):
    '''
    Parameters
    ----------
    num_chans : int
        Number of channels
    chunk : int
        Size of chunks

    Returns
    -------
    chunked_channels : list of np.ndarray
        List of channels in chunks of the given size.
    '''
    channels = np.arange(num_chans)
    chunked_channels = \
        np.array_split(channels,
                       [chunk * i for i in xrange(num_chans / chunk)])
    if chunked_channels[-1].size == 0:
        chunked_channels = chunked_channels[:-1]
    if chunked_channels[0].size == 0:
        chunked_channels = chunked_channels[1:]

    return chunked_channels
<commit_msg>Default weights to None<commit_after>
'''
Routines for feathering large cubes together using uvcombine.

Some of these *should* be moved into uvcombine in good time.
'''

from spectral_cube import SpectralCube
from uvcombine.uvcombine import feather_simple, feather_compare
from astropy import log
import astropy.units as u

import numpy as np

from .io_utils import create_huge_fits
from .progressbar import _map_context

def feather_cube(cube_hi, cube_lo, verbose=True, save_feather=True,
                 save_name=None, num_cores=1, chunk=100,
                 restfreq=1.42040575177 * u.GHz,
                 weights=None,
                 feather_kwargs={}):
    '''
    Feather two cubes together. The spectral axis of the cubes *must match*.
    This function is specifically for cube too large to fit in memory. The
    feathered planes are saved directly to a new FITS file.
    '''

    # Ensure the spectral axes are the same.
    if not np.allclose(cube_hi.spectral_axis.to(cube_lo._spectral_unit).value,
                       cube_lo.spectral_axis.value):
        raise Warning("The spectral axes do not match. Spectrally regrid the "
                      "cubes to be the same before feathering.")

    if save_feather:
        if save_name is None:
            raise TypeError("save_name must be given the save the shifted "
                            "cube.")

        output_hdu = create_huge_fits(save_name, cube_hi.header,
                                      return_hdu=True)
    else:
        newcube = np.empty(cube_hi.shape, dtype=cube_hi[:1, 0, 0].dtype)

    num_chans = cube_hi.shape[0]
    chunked_channels = get_channel_chunks(num_chans, chunk)

    freq_axis = cube_hi.with_spectral_unit(u.Hz, velocity_convention='radio',
                                           rest_value=restfreq).spectral_axis

    if weights is not None:
        feather_kwargs['weights'] = weights

    for i, chunk_chans in enumerate(chunked_channels):

        log.info("On chunk {0} of {1}".format(i + 1, len(chunked_channels)))

        changen = ((chan, cube_hi[chan:chan + 1], cube_lo[chan:chan + 1],
                    freq_axis[chan], feather_kwargs) for chan in chunk_chans)

        with _map_context(num_cores, verbose=verbose,
                          num_jobs=len(chunk_chans)) as map:
            log.info("Feathering")
            output = map(_feather, changen)

        for chan, shifted_arr in output:
            if save_feather:
                output_hdu[0].data[chan] = shifted_arr.real
            else:
                newcube[chan] = shifted_arr.real

        if save_feather:
            output_hdu.flush()

    if save_feather:

        # Append the beam table, if needed
        if hasattr(cube_hi, 'beams'):
            from spectral_cube.cube_utils import beams_to_bintable
            output_hdu.append(beams_to_bintable(cube_hi.beams))

        output_hdu.close()
    else:
        return SpectralCube(newcube, wcs=cube_hi.wcs,
                            header=cube_hi.header, meta=cube_hi.meta)

def _feather(args):
    '''
    Feather 2D images together.
    '''
    chan, plane_hi, plane_lo, freq, kwargs = args

    # TODO: When Jy/beam is supported in SpectralCube, just let uvcombine
    # match the units. BUT we also need Slice objects to retain a 1 element
    # spectral axis!
    if plane_hi.unit != plane_lo.unit:
        if plane_hi.unit == u.K:
            if hasattr(plane_hi, 'beams'):
                beam = plane_hi.beams[0]
            else:
                beam = plane_hi.beam

            plane_hi = plane_hi.to(u.Jy / u.beam, plane_hi.beam.jtok_equiv(freq))

        if plane_lo.unit == u.K:
            if hasattr(plane_lo, 'beams'):
                beam = plane_lo.beams[0]
            else:
                beam = plane_lo.beam

            plane_lo = plane_lo.to(u.Jy / u.beam, beam.jtok_equiv(freq))

    if hasattr(plane_hi, 'hdu'):
        plane_hi_hdu = plane_hi.hdu
    else:
        plane_hi_hdu = plane_hi.hdulist[0]
        plane_hi_hdu.header.update(plane_hi.beams[0].to_header_keywords())

    if hasattr(plane_lo, 'hdu'):
        plane_lo_hdu = plane_lo.hdu
    else:
        plane_lo_hdu = plane_lo.hdulist[0]
        plane_lo_hdu.header.update(plane_lo.beams[0].to_header_keywords())

    feathered = feather_simple(plane_hi_hdu, plane_lo_hdu, **kwargs)

    # Expect that the interferometer image will cover a smaller region.
    # Cut to that region.
    feathered[np.isnan(plane_hi[0])] = np.NaN

    return chan, feathered

def feather_compare_cube(cube_hi, cube_lo, LAS, lowresfwhm=None,
                         restfreq=1.42040575177 * u.GHz,
                         verbose=True, num_cores=1, chunk=100,
                         weights=1.):
    '''
    Record the ratios of the flux in the overlap region between the cubes.

    **Assumes velocity axes use the radio convention!**

    '''

    # Ensure the spectral axes are the same.
    if not np.allclose(cube_hi.spectral_axis.to(cube_lo._spectral_unit).value,
                       cube_lo.spectral_axis.value):
        raise Warning("The spectral axes do not match. Spectrally regrid the "
                      "cubes to be the same before feathering.")

    num_chans = cube_hi.shape[0]
    chunked_channels = get_channel_chunks(num_chans, chunk)

    radii = []
    ratios = []
    highres_pts = []
    lowres_pts = []

    if lowresfwhm is None:
        lowresfwhm = cube_lo.beam.major.to(u.arcsec)
    else:
        lowresfwhm = lowresfwhm.to(u.arcsec)

    freq_axis = cube_hi.with_spectral_unit(u.Hz, velocity_convention='radio',
                                           rest_value=restfreq).spectral_axis

    for i, chunk_chans in enumerate(chunked_channels):

        log.info("On chunk {0} of {1}".format(i + 1, len(chunked_channels)))

        changen = ((chan, cube_hi[chan], cube_lo[chan],
                    LAS, lowresfwhm, freq_axis[chan],
                    weights) for chan in chunk_chans)

        with _map_context(num_cores, verbose=verbose,
                          num_jobs=len(chunk_chans)) as map:
            output = map(_compare, changen)

        chan_out = np.array([out[0] for out in output])

        for jj in chan_out.argsort():

            radii.append(output[jj][1][0])
            ratios.append(output[jj][1][1])
            highres_pts.append(output[jj][1][2])
            lowres_pts.append(output[jj][1][3])

    return radii, ratios, highres_pts, lowres_pts

def _compare(args):

    chan, plane_hi, plane_lo, LAS, lowresfwhm, freq, weights = args

    out = feather_compare(plane_hi.to(u.K, freq=freq).hdu,
                          plane_lo.hdu,
                          return_samples=True, doplot=False,
                          LAS=LAS, SAS=lowresfwhm,
                          lowresfwhm=lowresfwhm,
                          weights=weights)

    return chan, out

def flux_recovery(cube_hi, cube_lo, frequency=1.42040575177 * u.GHz,
                  verbose=True, doplot=False,
                  num_cores=1, chunk=100, mask=None):
    '''
    Calculate the fraction of flux recovered between the high-resolution
    (interferometer) cube and the low-resolution (single-dish) cube.
    '''

    # Ensure the spectral axes are the same.
    if not np.allclose(cube_hi.spectral_axis.to(cube_lo._spectral_unit).value,
                       cube_lo.spectral_axis.value):
        raise Warning("The spectral axes do not match. Spectrally regrid the "
                      "cubes to be the same before feathering.")

    if mask is not None:
        if mask.shape != cube_hi[0].shape:
            raise ValueError("mask must have the same shape as the "
                             "high-resolution cube.")

    num_chans = cube_hi.shape[0]
    chunked_channels = get_channel_chunks(num_chans, chunk)

    total_hires = np.empty(num_chans) * u.Jy
    total_lores = np.empty(num_chans) * u.Jy

    for i, chunk_chans in enumerate(chunked_channels):

        log.info("On chunk {0} of {1}".format(i + 1, len(chunked_channels)))

        changen = ((chan, cube_hi[chan], cube_lo[chan],
                    frequency, mask) for chan in chunk_chans)

        with _map_context(num_cores, verbose=verbose)as map:
            output = map(_totalplanes, changen)

        chan_out = np.array([out[0] for out in output])

        for jj in chan_out.argsort():

            total_hires[chan_out[jj]] = output[jj][1]
            total_lores[chan_out[jj]] = output[jj][2]

    if doplot:
        import matplotlib.pyplot as plt
        plt.subplot(121)
        plt.plot(total_hires.value, label='High-res')
        plt.plot(total_lores.value, label='Low-res')
        plt.ylabel("Total Intensity (Jy)")
        plt.legend()
        plt.subplot(122)
        plt.plot(total_hires.value / total_lores.value)
        plt.ylabel("Fraction of recovered")
        plt.xlabel("Channel")

    return total_hires, total_lores

def _totalplanes(args):

    chan, plane_hi, plane_lo, freq, mask = args

    if mask is not None:
        if not plane_hi.wcs.wcs.compare(plane_lo.wcs.wcs):
            # Reproject the low resolution onto the high

            plane_lo = plane_lo.reproject(plane_hi.header)

        # !! Update to Jy/bm (or equiv) when spectral-cube can handle it
        total_hi = \
            np.nansum(plane_hi[mask].to(u.Jy, plane_hi.beam.jtok_equiv(freq)))
        total_lo = \
            np.nansum(plane_lo[mask].to(u.Jy, plane_lo.beam.jtok_equiv(freq)))
    else:
        total_hi = \
            np.nansum(plane_hi.to(u.Jy, plane_hi.beam.jtok_equiv(freq)))
        total_lo = \
            np.nansum(plane_lo.to(u.Jy, plane_lo.beam.jtok_equiv(freq)))

    # Convert from Jy/beam to actual Jy
    total_hi *= (1 / plane_hi.beam.sr.to(u.deg**2)) * \
        (plane_hi.header["CDELT2"] * u.deg)**2
    total_lo *= (1 / plane_lo.beam.sr.to(u.deg**2)) * \
        (plane_lo.header["CDELT2"] * u.deg)**2

    return chan, total_hi, total_lo

def get_channel_chunks(num_chans, chunk):
    '''
    Parameters
    ----------
    num_chans : int
        Number of channels
    chunk : int
        Size of chunks

    Returns
    -------
    chunked_channels : list of np.ndarray
        List of channels in chunks of the given size.
    '''
    channels = np.arange(num_chans)
    chunked_channels = \
        np.array_split(channels,
                       [chunk * i for i in xrange(num_chans / chunk)])
    if chunked_channels[-1].size == 0:
        chunked_channels = chunked_channels[:-1]
    if chunked_channels[0].size == 0:
        chunked_channels = chunked_channels[1:]

    return chunked_channels
<|endoftext|>


<commit_before>
from __future__ import absolute_import, division

from desispec.io import read_frame
from desispec.io import read_fiberflat
from desispec.io import read_sky
from desispec.io import shorten_filename
from desispec.io.fluxcalibration import read_stdstar_models
from desispec.io.fluxcalibration import write_flux_calibration
from desispec.fiberflat import apply_fiberflat
from desispec.sky import subtract_sky
from desispec.fluxcalibration import compute_flux_calibration, isStdStar
from desiutil.log import get_logger
from desitarget.targets import main_cmx_or_sv
from desispec.fiberbitmasking import get_fiberbitmasked_frame

import argparse
import os
import os.path
import numpy as np
import sys

def parse(options=None):
    parser = argparse.ArgumentParser(description="Compute the flux calibration for a DESI frame using precomputed spectro-photometrically calibrated stellar models.")

    parser.add_argument('--infile', type = str, default = None, required=True,
                        help = 'path of DESI exposure frame fits file')
    parser.add_argument('--fiberflat', type = str, default = None, required=True,
                        help = 'path of DESI fiberflat fits file')
    parser.add_argument('--sky', type = str, default = None, required=True,
                        help = 'path of DESI sky fits file')
    parser.add_argument('--models', type = str, default = None, required=True,
                        help = 'path of spetro-photometric stellar spectra fits file')
    parser.add_argument('--chi2cut', type = float, default = 0., required=False,
                        help = 'apply a reduced chi2 cut for the selection of stars')
    parser.add_argument('--chi2cut-nsig', type = float, default = 0., required=False,
                        help = 'discard n-sigma outliers from the reduced chi2 of the standard star fit')
    parser.add_argument('--color', type = str, default = None, required=False,
                        help = 'color used for filtering. Can be G-R R-Z or GAIA-BP-RP or GAIA-G-RP')
    parser.add_argument('--min-color', type = float, default = None, required=False,
                        help = 'only consider stars with color greater than this')
    parser.add_argument('--delta-color-cut', type = float, default = 0.2, required=False,
                        help = 'discard model stars with different broad-band color from imaging')
    parser.add_argument('--nostdcheck', dest='nostdcheck',
                        help='Do not check the standards against flags in the FIBERMAP; just use objects in the model file', action='store_true')
    parser.add_argument('--outfile', type = str, default = None, required=True,
                        help = 'path of DESI flux calbration fits file')
    parser.add_argument('--qafile', type=str, default=None, required=False,
                        help='path of QA file.')
    parser.add_argument('--qafig', type = str, default = None, required=False,
                        help = 'path of QA figure file')
    parser.add_argument('--highest-throughput', type = int, default = 0, required=False,
                        help = 'use this number of stars ranked by highest throughput to normalize transmission (for DESI commissioning)')
    parser.add_argument('--seeing-fwhm', type = float, default = 1.1, required=False,
                        help = 'seeing FWHM in arcsec, used for fiberloss correction')
    parser.set_defaults(nostdcheck=False)
    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)
    return args

def main(args) :

    log=get_logger()

    cmd = ['desi_compute_fluxcalibration',]
    for key, value in args.__dict__.items():
        if value is not None:
            cmd += ['--'+key, str(value)]
    cmd = ' '.join(cmd)
    log.info(cmd)

    log.info("read frame")
    # read frame
    frame = read_frame(args.infile)

    # Set fibermask flagged spectra to have 0 flux and variance
    frame = get_fiberbitmasked_frame(frame, bitmask='flux',ivar_framemask=True)

    log.info("apply fiberflat")
    # read fiberflat
    fiberflat = read_fiberflat(args.fiberflat)

    # apply fiberflat
    apply_fiberflat(frame, fiberflat)

    log.info("subtract sky")
    # read sky
    skymodel=read_sky(args.sky)

    # subtract sky
    subtract_sky(frame, skymodel)

    log.info("compute flux calibration")

    # read models
    model_flux, model_wave, model_fibers, model_metadata=read_stdstar_models(args.models)

    ok=np.ones(len(model_metadata),dtype=bool)

    if args.chi2cut > 0 :
        log.info("apply cut CHI2DOF<{}".format(args.chi2cut))
        good = (model_metadata["CHI2DOF"]<args.chi2cut)
        bad  = ~good
        ok  &= good
        if np.any(bad) :
            log.info(" discard {} stars with CHI2DOF= {}".format(np.sum(bad),list(model_metadata["CHI2DOF"][bad])))

    legacy_filters = ('G-R', 'R-Z')
    gaia_filters = ('GAIA-BP-RP', 'GAIA-G-RP')
    model_column_list = model_metadata.columns.names
    if args.color is None:
        if 'MODEL_G-R' in model_column_list:
            color = 'G-R'
        elif 'MODEL_GAIA-BP-RP' in model_column_list:
            log.info('Using Gaia filters')
            color ='GAIA-BP-RP'
        else:
            log.error("Can't find either G-R or BP-RP color in the model file.")
            sys.exit(15)
    else:
        if args.color not in legacy_filters and args.color not in gaia_filters:
            log.error('Color name {} is not allowed, must be one of {} {}'.format(args.color, legacy_filters,gaia_filters))
            sys.exit(14)
        color = args.color
        if color not in model_column_list:
            # This should't happen
            log.error('The color {} was not computed in the models'.format(color))
            sys.exit(16)

    if args.delta_color_cut > 0 :
        log.info("apply cut |delta color|<{}".format(args.delta_color_cut))
        good = (np.abs(model_metadata["MODEL_"+color]-model_metadata["DATA_"+color])<args.delta_color_cut)
        bad  = ok&(~good)
        ok  &= good
        if np.any(bad) :
            vals=model_metadata["MODEL_"+color][bad]-model_metadata["DATA_"+color][bad]
            log.info(" discard {} stars with dcolor= {}".format(np.sum(bad),list(vals)))

    if args.min_color is not None :
        log.info("apply cut DATA_{}>{}".format(color, args.min_color))
        good = (model_metadata["DATA_{}".format(color)]>args.min_color)
        bad  = ok&(~good)
        ok  &= good
        if np.any(bad) :
            vals=model_metadata["DATA_{}".format(color)][bad]
            log.info(" discard {} stars with {}= {}".format(np.sum(bad),color,list(vals)))

    if args.chi2cut_nsig > 0 :
        # automatically reject stars that ar chi2 outliers
        mchi2=np.median(model_metadata["CHI2DOF"])
        rmschi2=np.std(model_metadata["CHI2DOF"])
        maxchi2=mchi2+args.chi2cut_nsig*rmschi2
        log.info("apply cut CHI2DOF<{} based on chi2cut_nsig={}".format(maxchi2,args.chi2cut_nsig))
        good = (model_metadata["CHI2DOF"]<=maxchi2)
        bad  = ok&(~good)
        ok  &= good
        if np.any(bad) :
            log.info(" discard {} stars with CHI2DOF={}".format(np.sum(bad),list(model_metadata["CHI2DOF"][bad])))

    ok=np.where(ok)[0]
    if ok.size == 0 :
        log.error("selection cuts discarded all stars")
        sys.exit(12)
    nstars=model_flux.shape[0]
    nbad=nstars-ok.size
    if nbad>0 :
        log.warning("discarding %d star(s) out of %d because of cuts"%(nbad,nstars))
        model_flux=model_flux[ok]
        model_fibers=model_fibers[ok]
        model_metadata=model_metadata[:][ok]

    stdcheck = not args.nostdcheck

    # check that the model_fibers are actually standard stars
    fibermap = frame.fibermap

    ## check whether star fibers from args.models are consistent with fibers from fibermap
    ## if not print the OBJTYPE from fibermap for the fibers numbers in args.models and exit
    if stdcheck:
        fibermap_std_indices = np.where(isStdStar(fibermap))[0]
        if np.any(~np.in1d(model_fibers%500, fibermap_std_indices)):
            target_colnames, target_masks, survey = main_cmx_or_sv(fibermap)
            colname =  target_colnames[0]
            for i in model_fibers%500:
                log.error("inconsistency with spectrum {}, OBJTYPE={}, {}={} in fibermap".format(
                i, fibermap["OBJTYPE"][i], colname, fibermap[colname][i]))
            sys.exit(12)
    else:
        fibermap_std_indices = model_fibers % 500
    # Make sure the fibers of interest aren't entirely masked.
    if np.sum(np.sum(frame.ivar[model_fibers%500, :] == 0, axis=1) == frame.nwave) == len(model_fibers):
        log.warning('All standard-star spectra are masked!')
        return

    fluxcalib = compute_flux_calibration(frame, model_wave, model_flux, model_fibers%500, highest_throughput_nstars = args.highest_throughput, exposure_seeing_fwhm = args.seeing_fwhm, stdcheck=stdcheck)

    # QA
    if (args.qafile is not None):

        from desispec.io import write_qa_frame
        from desispec.io.qa import load_qa_frame
        from desispec.qa import qa_plots

        log.info("performing fluxcalib QA")
        # Load
        qaframe = load_qa_frame(args.qafile, frame_meta=frame.meta, flavor=frame.meta['FLAVOR'])
        # Run
        #import pdb; pdb.set_trace()
        qaframe.run_qa('FLUXCALIB', (frame, fluxcalib))
        # Write
        if args.qafile is not None:
            write_qa_frame(args.qafile, qaframe)
            log.info("successfully wrote {:s}".format(args.qafile))
        # Figure(s)
        if args.qafig is not None:
            qa_plots.frame_fluxcalib(args.qafig, qaframe, frame, fluxcalib)

    # record inputs
    frame.meta['IN_FRAME'] = shorten_filename(args.infile)
    frame.meta['IN_SKY']   = shorten_filename(args.sky)
    frame.meta['FIBERFLT'] = shorten_filename(args.fiberflat)
    frame.meta['STDMODEL'] = shorten_filename(args.models)

    # write result
    write_flux_calibration(args.outfile, fluxcalib, header=frame.meta)

    log.info("successfully wrote %s"%args.outfile)
<commit_msg>consider an error on the g-r reddening of 0.2*E(B-V)<commit_after>
from __future__ import absolute_import, division

from desispec.io import read_frame
from desispec.io import read_fiberflat
from desispec.io import read_sky
from desispec.io import shorten_filename
from desispec.io.fluxcalibration import read_stdstar_models
from desispec.io.fluxcalibration import write_flux_calibration
from desispec.fiberflat import apply_fiberflat
from desispec.sky import subtract_sky
from desispec.fluxcalibration import compute_flux_calibration, isStdStar
from desiutil.log import get_logger
from desitarget.targets import main_cmx_or_sv
from desispec.fiberbitmasking import get_fiberbitmasked_frame

import argparse
import os
import os.path
import numpy as np
import sys

def parse(options=None):
    parser = argparse.ArgumentParser(description="Compute the flux calibration for a DESI frame using precomputed spectro-photometrically calibrated stellar models.")

    parser.add_argument('--infile', type = str, default = None, required=True,
                        help = 'path of DESI exposure frame fits file')
    parser.add_argument('--fiberflat', type = str, default = None, required=True,
                        help = 'path of DESI fiberflat fits file')
    parser.add_argument('--sky', type = str, default = None, required=True,
                        help = 'path of DESI sky fits file')
    parser.add_argument('--models', type = str, default = None, required=True,
                        help = 'path of spetro-photometric stellar spectra fits file')
    parser.add_argument('--chi2cut', type = float, default = 0., required=False,
                        help = 'apply a reduced chi2 cut for the selection of stars')
    parser.add_argument('--chi2cut-nsig', type = float, default = 0., required=False,
                        help = 'discard n-sigma outliers from the reduced chi2 of the standard star fit')
    parser.add_argument('--color', type = str, default = None, required=False,
                        help = 'color used for filtering. Can be G-R R-Z or GAIA-BP-RP or GAIA-G-RP')
    parser.add_argument('--min-color', type = float, default = None, required=False,
                        help = 'only consider stars with color greater than this')
    parser.add_argument('--delta-color-cut', type = float, default = 0.2, required=False,
                        help = 'discard model stars with different broad-band color from imaging')
    parser.add_argument('--nostdcheck', dest='nostdcheck',
                        help='Do not check the standards against flags in the FIBERMAP; just use objects in the model file', action='store_true')
    parser.add_argument('--outfile', type = str, default = None, required=True,
                        help = 'path of DESI flux calbration fits file')
    parser.add_argument('--qafile', type=str, default=None, required=False,
                        help='path of QA file.')
    parser.add_argument('--qafig', type = str, default = None, required=False,
                        help = 'path of QA figure file')
    parser.add_argument('--highest-throughput', type = int, default = 0, required=False,
                        help = 'use this number of stars ranked by highest throughput to normalize transmission (for DESI commissioning)')
    parser.add_argument('--seeing-fwhm', type = float, default = 1.1, required=False,
                        help = 'seeing FWHM in arcsec, used for fiberloss correction')
    parser.set_defaults(nostdcheck=False)
    args = None
    if options is None:
        args = parser.parse_args()
    else:
        args = parser.parse_args(options)
    return args

def main(args) :

    log=get_logger()

    cmd = ['desi_compute_fluxcalibration',]
    for key, value in args.__dict__.items():
        if value is not None:
            cmd += ['--'+key, str(value)]
    cmd = ' '.join(cmd)
    log.info(cmd)

    log.info("read frame")
    # read frame
    frame = read_frame(args.infile)

    # Set fibermask flagged spectra to have 0 flux and variance
    frame = get_fiberbitmasked_frame(frame, bitmask='flux',ivar_framemask=True)

    log.info("apply fiberflat")
    # read fiberflat
    fiberflat = read_fiberflat(args.fiberflat)

    # apply fiberflat
    apply_fiberflat(frame, fiberflat)

    log.info("subtract sky")
    # read sky
    skymodel=read_sky(args.sky)

    # subtract sky
    subtract_sky(frame, skymodel)

    log.info("compute flux calibration")

    # read models
    model_flux, model_wave, model_fibers, model_metadata=read_stdstar_models(args.models)

    ok=np.ones(len(model_metadata),dtype=bool)

    if args.chi2cut > 0 :
        log.info("apply cut CHI2DOF<{}".format(args.chi2cut))
        good = (model_metadata["CHI2DOF"]<args.chi2cut)
        bad  = ~good
        ok  &= good
        if np.any(bad) :
            log.info(" discard {} stars with CHI2DOF= {}".format(np.sum(bad),list(model_metadata["CHI2DOF"][bad])))

    legacy_filters = ('G-R', 'R-Z')
    gaia_filters = ('GAIA-BP-RP', 'GAIA-G-RP')
    model_column_list = model_metadata.columns.names
    if args.color is None:
        if 'MODEL_G-R' in model_column_list:
            color = 'G-R'
        elif 'MODEL_GAIA-BP-RP' in model_column_list:
            log.info('Using Gaia filters')
            color ='GAIA-BP-RP'
        else:
            log.error("Can't find either G-R or BP-RP color in the model file.")
            sys.exit(15)
    else:
        if args.color not in legacy_filters and args.color not in gaia_filters:
            log.error('Color name {} is not allowed, must be one of {} {}'.format(args.color, legacy_filters,gaia_filters))
            sys.exit(14)
        color = args.color
        if color not in model_column_list:
            # This should't happen
            log.error('The color {} was not computed in the models'.format(color))
            sys.exit(16)

    if args.delta_color_cut > 0 :
        # check dust extinction values for those stars
        stars_ebv = np.array(frame.fibermap[model_fibers % 500]["EBV"])
        log.info("min max E(B-V) for std stars = {:4.3f} {:4.3f}".format(np.min(stars_ebv),np.max(stars_ebv)))
        star_gr_reddening_relative_error = 0.2 * stars_ebv
        log.info("Consider a g-r reddening sys. error in the range {:4.3f} {:4.3f}".format(np.min(star_gr_reddening_relative_error),np.max(star_gr_reddening_relative_error)))

        log.info("apply cut |delta color|<{}+reddening error".format(args.delta_color_cut))
        good = (np.abs(model_metadata["MODEL_"+color]-model_metadata["DATA_"+color])<args.delta_color_cut+star_gr_reddening_relative_error)
        bad  = ok&(~good)
        ok  &= good
        if np.any(bad) :
            vals=model_metadata["MODEL_"+color][bad]-model_metadata["DATA_"+color][bad]
            log.info(" discard {} stars with dcolor= {}".format(np.sum(bad),list(vals)))

    if args.min_color is not None :
        log.info("apply cut DATA_{}>{}".format(color, args.min_color))
        good = (model_metadata["DATA_{}".format(color)]>args.min_color)
        bad  = ok&(~good)
        ok  &= good
        if np.any(bad) :
            vals=model_metadata["DATA_{}".format(color)][bad]
            log.info(" discard {} stars with {}= {}".format(np.sum(bad),color,list(vals)))

    if args.chi2cut_nsig > 0 :
        # automatically reject stars that ar chi2 outliers
        mchi2=np.median(model_metadata["CHI2DOF"])
        rmschi2=np.std(model_metadata["CHI2DOF"])
        maxchi2=mchi2+args.chi2cut_nsig*rmschi2
        log.info("apply cut CHI2DOF<{} based on chi2cut_nsig={}".format(maxchi2,args.chi2cut_nsig))
        good = (model_metadata["CHI2DOF"]<=maxchi2)
        bad  = ok&(~good)
        ok  &= good
        if np.any(bad) :
            log.info(" discard {} stars with CHI2DOF={}".format(np.sum(bad),list(model_metadata["CHI2DOF"][bad])))

    ok=np.where(ok)[0]
    if ok.size == 0 :
        log.error("selection cuts discarded all stars")
        sys.exit(12)
    nstars=model_flux.shape[0]
    nbad=nstars-ok.size
    if nbad>0 :
        log.warning("discarding %d star(s) out of %d because of cuts"%(nbad,nstars))
        model_flux=model_flux[ok]
        model_fibers=model_fibers[ok]
        model_metadata=model_metadata[:][ok]

    stdcheck = not args.nostdcheck

    # check that the model_fibers are actually standard stars
    fibermap = frame.fibermap

    ## check whether star fibers from args.models are consistent with fibers from fibermap
    ## if not print the OBJTYPE from fibermap for the fibers numbers in args.models and exit
    if stdcheck:
        fibermap_std_indices = np.where(isStdStar(fibermap))[0]
        if np.any(~np.in1d(model_fibers%500, fibermap_std_indices)):
            target_colnames, target_masks, survey = main_cmx_or_sv(fibermap)
            colname =  target_colnames[0]
            for i in model_fibers%500:
                log.error("inconsistency with spectrum {}, OBJTYPE={}, {}={} in fibermap".format(
                i, fibermap["OBJTYPE"][i], colname, fibermap[colname][i]))
            sys.exit(12)
    else:
        fibermap_std_indices = model_fibers % 500
    # Make sure the fibers of interest aren't entirely masked.
    if np.sum(np.sum(frame.ivar[model_fibers%500, :] == 0, axis=1) == frame.nwave) == len(model_fibers):
        log.warning('All standard-star spectra are masked!')
        return

    fluxcalib = compute_flux_calibration(frame, model_wave, model_flux, model_fibers%500, highest_throughput_nstars = args.highest_throughput, exposure_seeing_fwhm = args.seeing_fwhm, stdcheck=stdcheck)

    # QA
    if (args.qafile is not None):

        from desispec.io import write_qa_frame
        from desispec.io.qa import load_qa_frame
        from desispec.qa import qa_plots

        log.info("performing fluxcalib QA")
        # Load
        qaframe = load_qa_frame(args.qafile, frame_meta=frame.meta, flavor=frame.meta['FLAVOR'])
        # Run
        #import pdb; pdb.set_trace()
        qaframe.run_qa('FLUXCALIB', (frame, fluxcalib))
        # Write
        if args.qafile is not None:
            write_qa_frame(args.qafile, qaframe)
            log.info("successfully wrote {:s}".format(args.qafile))
        # Figure(s)
        if args.qafig is not None:
            qa_plots.frame_fluxcalib(args.qafig, qaframe, frame, fluxcalib)

    # record inputs
    frame.meta['IN_FRAME'] = shorten_filename(args.infile)
    frame.meta['IN_SKY']   = shorten_filename(args.sky)
    frame.meta['FIBERFLT'] = shorten_filename(args.fiberflat)
    frame.meta['STDMODEL'] = shorten_filename(args.models)

    # write result
    write_flux_calibration(args.outfile, fluxcalib, header=frame.meta)

    log.info("successfully wrote %s"%args.outfile)
<|endoftext|>


<commit_before>#!/usr/bin/env python

"""
CUDA Library Installer

Installs the latest CUDA library supported by CuPy.
"""

# This script will also be used as a standalone script when building wheels.
# Keep the script runnable without CuPy dependency.

import argparse
import json
import os
import platform
import shutil
import sys
import tempfile
import urllib.request

_cudnn_records = []
_cutensor_records = []
_nccl_records = []
library_records = {}

def _make_cudnn_url(public_version, filename):
    # https://developer.download.nvidia.com/compute/redist/cudnn/v8.0.2/cudnn-11.0-linux-x64-v8.0.2.39.tgz
    return (
        'https://developer.download.nvidia.com/compute/redist/cudnn' +
        '/v{}/{}'.format(public_version, filename))

def _make_cudnn_record(
        cuda_version, public_version, filename_linux, filename_windows):
    major_version = public_version.split('.')[0]

    if major_version == '7':
        suffix_list = ['']
    elif major_version == '8':
        # Dependency order is documented at:
        # https://docs.nvidia.com/deeplearning/cudnn/api/index.html
        suffix_list = ['', '_ops_infer', '_ops_train',
                       '_cnn_infer', '_cnn_train',
                       '_adv_infer', '_adv_train']
    else:
        raise AssertionError

    return {
        'cuda': cuda_version,
        'cudnn': public_version,
        'assets': {
            'Linux': {
                'url': _make_cudnn_url(public_version, filename_linux),
                'filenames': [f'libcudnn{suffix}.so.{public_version}'
                              for suffix in suffix_list]
            },
            'Windows': {
                'url': _make_cudnn_url(public_version, filename_windows),
                'filenames': [f'cudnn{suffix}64_{major_version}.dll'
                              for suffix in suffix_list]
            },
        }
    }

# Latest cuDNN versions: https://developer.nvidia.com/rdp/cudnn-download
_cudnn_records.append(_make_cudnn_record(
    '11.4', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '11.3', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '11.2', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '11.1', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '11.0', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '10.2', '8.2.4',
    'cudnn-10.2-linux-x64-v8.2.4.15.tgz',
    'cudnn-10.2-windows10-x64-v8.2.4.15.zip'))
library_records['cudnn'] = _cudnn_records

def _make_cutensor_url(public_version, filename):
    # https://developer.download.nvidia.com/compute/cutensor/1.2.2/local_installers/libcutensor-linux-x86_64-1.2.2.5.tar.gz
    return (
        'https://developer.download.nvidia.com/compute/cutensor/' +
        '{}/local_installers/{}'.format(public_version, filename))

def _make_cutensor_record(
        cuda_version, public_version, filename_linux, filename_windows):
    return {
        'cuda': cuda_version,
        'cutensor': public_version,
        'assets': {
            'Linux': {
                'url': _make_cutensor_url(public_version, filename_linux),
                'filenames': ['libcutensor.so.{}'.format(public_version)],
            },
            'Windows': {
                'url': _make_cutensor_url(public_version, filename_windows),
                'filenames': ['cutensor.dll'],
            },
        }
    }

_cutensor_records.append(_make_cutensor_record(
    '11.4', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '11.3', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '11.2', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '11.1', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '11.0', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '10.2', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
library_records['cutensor'] = _cutensor_records

def _make_nccl_url(public_version, filename):
    # https://developer.download.nvidia.com/compute/redist/nccl/v2.8/nccl_2.8.4-1+cuda11.2_x86_64.txz
    return (
        'https://developer.download.nvidia.com/compute/redist/nccl/' +
        'v{}/{}'.format(public_version, filename))

def _make_nccl_record(
        cuda_version, full_version, public_version, filename_linux):
    return {
        'cuda': cuda_version,
        'nccl': full_version,
        'assets': {
            'Linux': {
                'url': _make_nccl_url(public_version, filename_linux),
                'filenames': ['libnccl.so.{}'.format(full_version)],
            },
        },
    }

_nccl_records.append(_make_nccl_record(
    '11.4', '2.11.4', '2.11',
    'nccl_2.11.4-1+cuda11.4_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '11.3', '2.9.9', '2.9',
    'nccl_2.9.9-1+cuda11.3_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '11.2', '2.8.4', '2.8',
    'nccl_2.8.4-1+cuda11.2_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '11.1', '2.8.4', '2.8',
    'nccl_2.8.4-1+cuda11.1_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '11.0', '2.11.4', '2.11',
    'nccl_2.11.4-1+cuda11.0_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '10.2', '2.11.4', '2.11',
    'nccl_2.11.4-1+cuda10.2_x86_64.txz'))
library_records['nccl'] = _nccl_records

def install_lib(cuda, prefix, library):
    record = None
    lib_records = library_records
    for record in lib_records[library]:
        if record['cuda'] == cuda:
            break
    else:
        raise RuntimeError('''
The CUDA version specified is not supported.
Should be one of {}.'''.format(str([x['cuda'] for x in lib_records[library]])))
    if prefix is None:
        prefix = os.path.expanduser('~/.cupy/cuda_lib')
    destination = calculate_destination(prefix, cuda, library, record[library])

    if os.path.exists(destination):
        raise RuntimeError('''
The destination directory {} already exists.
Remove the directory first if you want to reinstall.'''.format(destination))

    target_platform = platform.system()
    asset = record['assets'].get(target_platform, None)
    if asset is None:
        raise RuntimeError('''
The current platform ({}) is not supported.'''.format(target_platform))

    print('Installing {} {} for CUDA {} to: {}'.format(
        library, record[library], record['cuda'], destination))

    url = asset['url']
    print('Downloading {}...'.format(url))
    with tempfile.TemporaryDirectory() as tmpdir:
        with open(os.path.join(tmpdir, os.path.basename(url)), 'wb') as f:
            with urllib.request.urlopen(url) as response:
                f.write(response.read())
        print('Extracting...')
        outdir = os.path.join(tmpdir, 'extract')
        shutil.unpack_archive(f.name, outdir)
        print('Installing...')
        if library == 'cudnn':
            shutil.move(os.path.join(outdir, 'cuda'), destination)
        elif library == 'cutensor':
            if cuda.startswith('11.') and cuda != '11.0':
                cuda = '11'
            shutil.move(
                os.path.join(outdir, 'libcutensor', 'include'),
                os.path.join(destination, 'include'))
            shutil.move(
                os.path.join(outdir, 'libcutensor', 'lib', cuda),
                os.path.join(destination, 'lib'))
            if cuda == '10.1':
                license = 'license.pdf'  # v1.2.2
            else:
                license = 'license.txt'  # v1.3.0
            shutil.move(
                os.path.join(outdir, 'libcutensor', license), destination)
        elif library == 'nccl':
            subdir = os.listdir(outdir)  # ['nccl_2.8.4-1+cuda11.2_x86_64']
            assert len(subdir) == 1
            shutil.move(os.path.join(outdir, subdir[0]), destination)
        else:
            assert False
        print('Cleaning up...')
    print('Done!')

def calculate_destination(prefix, cuda, lib, lib_ver):
    """Calculates the installation directory.

    ~/.cupy/cuda_lib/{cuda_version}/{library_name}/{library_version}
    """
    return os.path.join(prefix, cuda, lib, lib_ver)

def main(args):
    parser = argparse.ArgumentParser()

    # TODO(kmaehashi): support NCCL
    parser.add_argument('--library',
                        choices=['cudnn', 'cutensor', 'nccl'],
                        required=True,
                        help='Library to install')
    parser.add_argument('--cuda', type=str, required=True,
                        help='CUDA version')
    parser.add_argument('--prefix', type=str, default=None,
                        help='Install destination')
    parser.add_argument('--action', choices=['install', 'dump'],
                        default='install',
                        help='Action to perform')
    params = parser.parse_args(args)

    if params.prefix is not None:
        params.prefix = os.path.abspath(params.prefix)

    if params.action == 'install':
        install_lib(params.cuda, params.prefix, params.library)
    elif params.action == 'dump':
        print(json.dumps(library_records[params.library], indent=4))
    else:
        assert False

if __name__ == '__main__':
    main(sys.argv[1:])
<commit_msg>limit architecture<commit_after>#!/usr/bin/env python

"""
CUDA Library Installer

Installs the latest CUDA library supported by CuPy.
"""

# This script will also be used as a standalone script when building wheels.
# Keep the script runnable without CuPy dependency.

import argparse
import json
import os
import platform
import shutil
import sys
import tempfile
import urllib.request

_cudnn_records = []
_cutensor_records = []
_nccl_records = []
library_records = {}

def _make_cudnn_url(public_version, filename):
    # https://developer.download.nvidia.com/compute/redist/cudnn/v8.0.2/cudnn-11.0-linux-x64-v8.0.2.39.tgz
    return (
        'https://developer.download.nvidia.com/compute/redist/cudnn' +
        '/v{}/{}'.format(public_version, filename))

def _make_cudnn_record(
        cuda_version, public_version, filename_linux, filename_windows):
    major_version = public_version.split('.')[0]

    if major_version == '7':
        suffix_list = ['']
    elif major_version == '8':
        # Dependency order is documented at:
        # https://docs.nvidia.com/deeplearning/cudnn/api/index.html
        suffix_list = ['', '_ops_infer', '_ops_train',
                       '_cnn_infer', '_cnn_train',
                       '_adv_infer', '_adv_train']
    else:
        raise AssertionError

    return {
        'cuda': cuda_version,
        'cudnn': public_version,
        'assets': {
            'Linux': {
                'url': _make_cudnn_url(public_version, filename_linux),
                'filenames': [f'libcudnn{suffix}.so.{public_version}'
                              for suffix in suffix_list]
            },
            'Windows': {
                'url': _make_cudnn_url(public_version, filename_windows),
                'filenames': [f'cudnn{suffix}64_{major_version}.dll'
                              for suffix in suffix_list]
            },
        }
    }

# Latest cuDNN versions: https://developer.nvidia.com/rdp/cudnn-download
_cudnn_records.append(_make_cudnn_record(
    '11.4', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '11.3', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '11.2', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '11.1', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '11.0', '8.2.4',
    'cudnn-11.4-linux-x64-v8.2.4.15.tgz',
    'cudnn-11.4-windows-x64-v8.2.4.15.zip'))
_cudnn_records.append(_make_cudnn_record(
    '10.2', '8.2.4',
    'cudnn-10.2-linux-x64-v8.2.4.15.tgz',
    'cudnn-10.2-windows10-x64-v8.2.4.15.zip'))
library_records['cudnn'] = _cudnn_records

def _make_cutensor_url(public_version, filename):
    # https://developer.download.nvidia.com/compute/cutensor/1.2.2/local_installers/libcutensor-linux-x86_64-1.2.2.5.tar.gz
    return (
        'https://developer.download.nvidia.com/compute/cutensor/' +
        '{}/local_installers/{}'.format(public_version, filename))

def _make_cutensor_record(
        cuda_version, public_version, filename_linux, filename_windows):
    return {
        'cuda': cuda_version,
        'cutensor': public_version,
        'assets': {
            'Linux': {
                'url': _make_cutensor_url(public_version, filename_linux),
                'filenames': ['libcutensor.so.{}'.format(public_version)],
            },
            'Windows': {
                'url': _make_cutensor_url(public_version, filename_windows),
                'filenames': ['cutensor.dll'],
            },
        }
    }

_cutensor_records.append(_make_cutensor_record(
    '11.4', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '11.3', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '11.2', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '11.1', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '11.0', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
_cutensor_records.append(_make_cutensor_record(
    '10.2', '1.3.1',
    'libcutensor-linux-x86_64-1.3.1.3.tar.gz',
    'libcutensor-windows-x86_64-1.3.1.3.zip'))
library_records['cutensor'] = _cutensor_records

def _make_nccl_url(public_version, filename):
    # https://developer.download.nvidia.com/compute/redist/nccl/v2.8/nccl_2.8.4-1+cuda11.2_x86_64.txz
    return (
        'https://developer.download.nvidia.com/compute/redist/nccl/' +
        'v{}/{}'.format(public_version, filename))

def _make_nccl_record(
        cuda_version, full_version, public_version, filename_linux):
    return {
        'cuda': cuda_version,
        'nccl': full_version,
        'assets': {
            'Linux': {
                'url': _make_nccl_url(public_version, filename_linux),
                'filenames': ['libnccl.so.{}'.format(full_version)],
            },
        },
    }

_nccl_records.append(_make_nccl_record(
    '11.4', '2.11.4', '2.11',
    'nccl_2.11.4-1+cuda11.4_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '11.3', '2.9.9', '2.9',
    'nccl_2.9.9-1+cuda11.3_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '11.2', '2.8.4', '2.8',
    'nccl_2.8.4-1+cuda11.2_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '11.1', '2.8.4', '2.8',
    'nccl_2.8.4-1+cuda11.1_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '11.0', '2.11.4', '2.11',
    'nccl_2.11.4-1+cuda11.0_x86_64.txz'))
_nccl_records.append(_make_nccl_record(
    '10.2', '2.11.4', '2.11',
    'nccl_2.11.4-1+cuda10.2_x86_64.txz'))
library_records['nccl'] = _nccl_records

def install_lib(cuda, prefix, library):
    if os.uname().machine != 'x86_64':
        raise RuntimeError('''
Currently this tool only supports x86_64 architecture.''')
    record = None
    lib_records = library_records
    for record in lib_records[library]:
        if record['cuda'] == cuda:
            break
    else:
        raise RuntimeError('''
The CUDA version specified is not supported.
Should be one of {}.'''.format(str([x['cuda'] for x in lib_records[library]])))
    if prefix is None:
        prefix = os.path.expanduser('~/.cupy/cuda_lib')
    destination = calculate_destination(prefix, cuda, library, record[library])

    if os.path.exists(destination):
        raise RuntimeError('''
The destination directory {} already exists.
Remove the directory first if you want to reinstall.'''.format(destination))

    target_platform = platform.system()
    asset = record['assets'].get(target_platform, None)
    if asset is None:
        raise RuntimeError('''
The current platform ({}) is not supported.'''.format(target_platform))

    print('Installing {} {} for CUDA {} to: {}'.format(
        library, record[library], record['cuda'], destination))

    url = asset['url']
    print('Downloading {}...'.format(url))
    with tempfile.TemporaryDirectory() as tmpdir:
        with open(os.path.join(tmpdir, os.path.basename(url)), 'wb') as f:
            with urllib.request.urlopen(url) as response:
                f.write(response.read())
        print('Extracting...')
        outdir = os.path.join(tmpdir, 'extract')
        shutil.unpack_archive(f.name, outdir)
        print('Installing...')
        if library == 'cudnn':
            shutil.move(os.path.join(outdir, 'cuda'), destination)
        elif library == 'cutensor':
            if cuda.startswith('11.') and cuda != '11.0':
                cuda = '11'
            shutil.move(
                os.path.join(outdir, 'libcutensor', 'include'),
                os.path.join(destination, 'include'))
            shutil.move(
                os.path.join(outdir, 'libcutensor', 'lib', cuda),
                os.path.join(destination, 'lib'))
            if cuda == '10.1':
                license = 'license.pdf'  # v1.2.2
            else:
                license = 'license.txt'  # v1.3.0
            shutil.move(
                os.path.join(outdir, 'libcutensor', license), destination)
        elif library == 'nccl':
            subdir = os.listdir(outdir)  # ['nccl_2.8.4-1+cuda11.2_x86_64']
            assert len(subdir) == 1
            shutil.move(os.path.join(outdir, subdir[0]), destination)
        else:
            assert False
        print('Cleaning up...')
    print('Done!')

def calculate_destination(prefix, cuda, lib, lib_ver):
    """Calculates the installation directory.

    ~/.cupy/cuda_lib/{cuda_version}/{library_name}/{library_version}
    """
    return os.path.join(prefix, cuda, lib, lib_ver)

def main(args):
    parser = argparse.ArgumentParser()

    # TODO(kmaehashi): support NCCL
    parser.add_argument('--library',
                        choices=['cudnn', 'cutensor', 'nccl'],
                        required=True,
                        help='Library to install')
    parser.add_argument('--cuda', type=str, required=True,
                        help='CUDA version')
    parser.add_argument('--prefix', type=str, default=None,
                        help='Install destination')
    parser.add_argument('--action', choices=['install', 'dump'],
                        default='install',
                        help='Action to perform')
    params = parser.parse_args(args)

    if params.prefix is not None:
        params.prefix = os.path.abspath(params.prefix)

    if params.action == 'install':
        install_lib(params.cuda, params.prefix, params.library)
    elif params.action == 'dump':
        print(json.dumps(library_records[params.library], indent=4))
    else:
        assert False

if __name__ == '__main__':
    main(sys.argv[1:])
<|endoftext|>


<commit_before>import py

here = py.magic.autopath().dirpath()
update_website = here.join('../../bin/_update_website.py').pyimport()

def test_rsync():
    temp = py.test.ensuretemp('update_website_rsync')
    pkgpath = temp.join('pkg')
    apipath = temp.join('apigen')
    pkgpath.ensure('foo/bar.txt', file=True).write('baz')
    pkgpath.ensure('spam/eggs.txt', file=True).write('spam')
    apipath.ensure('api/foo.html', file=True).write('<html />')
    apipath.ensure('source/spam.html', file=True).write('<html />')

    rsyncpath = temp.join('rsync')
    assert not rsyncpath.check()
    gateway = py.execnet.PopenGateway()
    update_website.rsync(pkgpath, apipath, gateway, rsyncpath.strpath)
    assert rsyncpath.check(dir=True)
    assert rsyncpath.join('pkg').check(dir=True)
    assert rsyncpath.join('pkg/spam/eggs.txt').read() == 'spam'
    assert rsyncpath.join('apigen').check(dir=True)
    assert rsyncpath.join('apigen/api/foo.html').read() == '<html />'

def setup_pkg(testname):
    temp = py.test.ensuretemp(testname)
    pkgpath = temp.ensure('pkg', dir=True)
    pyfile = pkgpath.ensure('mod.py').write(py.code.Source("""
        def foo(x):
            return x + 1
    """))
    testfile = pkgpath.ensure('test/test_mod.py').write(py.code.Source("""
        from pkg.sub import foo
        def test_foo():
            assert foo(1) == 2
    """))
    initfile = pkgpath.ensure('__init__.py').write(py.code.Source("""\
        import py
        from py.__.initpkg import initpkg
        initpkg(__name__, exportdefs={
            'sub.foo': ('./mod.py', 'foo'),
        })
    """))
    return pkgpath

def test_run_tests():
    if py.std.sys.platform == "win32":
        py.test.skip("update_website is not supposed to be run from win32")
    pkgpath = setup_pkg('update_website_run_tests')
    errors = update_website.run_tests(pkgpath)
    assert not errors
    assert pkgpath.join('../apigen').check(dir=True)
    assert pkgpath.join('../apigen/api/sub.foo.html').check(file=True)

def test_run_tests_failure():
    pkgpath = setup_pkg('update_website_run_tests_failure')
    assert not pkgpath.join('../apigen').check(dir=True)
    pkgpath.ensure('../apigen', file=True)
    errors = update_website.run_tests(pkgpath)
    assert errors # some error message

<commit_msg>[svn r38143] Removed some confusing output to stdout/stderr (output of failing tests, which are part of a test that tests whether a set of tests ran by a nested py.test process fails... so the failing output is correct).<commit_after>import py
import sys

here = py.magic.autopath().dirpath()
update_website = here.join('../../bin/_update_website.py').pyimport()

def test_rsync():
    temp = py.test.ensuretemp('update_website_rsync')
    pkgpath = temp.join('pkg')
    apipath = temp.join('apigen')
    pkgpath.ensure('foo/bar.txt', file=True).write('baz')
    pkgpath.ensure('spam/eggs.txt', file=True).write('spam')
    apipath.ensure('api/foo.html', file=True).write('<html />')
    apipath.ensure('source/spam.html', file=True).write('<html />')

    rsyncpath = temp.join('rsync')
    assert not rsyncpath.check()
    gateway = py.execnet.PopenGateway()
    update_website.rsync(pkgpath, apipath, gateway, rsyncpath.strpath)
    assert rsyncpath.check(dir=True)
    assert rsyncpath.join('pkg').check(dir=True)
    assert rsyncpath.join('pkg/spam/eggs.txt').read() == 'spam'
    assert rsyncpath.join('apigen').check(dir=True)
    assert rsyncpath.join('apigen/api/foo.html').read() == '<html />'

def setup_pkg(testname):
    temp = py.test.ensuretemp(testname)
    pkgpath = temp.ensure('pkg', dir=True)
    pyfile = pkgpath.ensure('mod.py').write(py.code.Source("""
        def foo(x):
            return x + 1
    """))
    testfile = pkgpath.ensure('test/test_mod.py').write(py.code.Source("""
        from pkg.sub import foo
        def test_foo():
            assert foo(1) == 2
    """))
    initfile = pkgpath.ensure('__init__.py').write(py.code.Source("""\
        import py
        from py.__.initpkg import initpkg
        initpkg(__name__, exportdefs={
            'sub.foo': ('./mod.py', 'foo'),
        })
    """))
    return pkgpath

def test_run_tests():
    if py.std.sys.platform == "win32":
        py.test.skip("update_website is not supposed to be run from win32")
    pkgpath = setup_pkg('update_website_run_tests')
    errors = update_website.run_tests(pkgpath)
    assert not errors
    assert pkgpath.join('../apigen').check(dir=True)
    assert pkgpath.join('../apigen/api/sub.foo.html').check(file=True)

def test_run_tests_failure():
    if py.std.sys.platform == "win32":
        py.test.skip("update_website is not supposed to be run from win32")
    pkgpath = setup_pkg('update_website_run_tests_failure')
    assert not pkgpath.join('../apigen').check(dir=True)
    pkgpath.ensure('../apigen', file=True)
    errors = update_website.run_tests(pkgpath, '> /dev/null 2>&1')
    assert errors # some error message

<|endoftext|>


<commit_before>from __future__ import unicode_literals

from django.apps import apps
from django.conf import settings
from django.urls import reverse, NoReverseMatch
from django.utils.translation import ugettext_lazy as _
from fluent_dashboard.dashboard import modules, FluentIndexDashboard, FluentAppIndexDashboard
import six

from waldur_core import __version__
from waldur_core.core import models as core_models, WaldurExtension
from waldur_core.logging import models as logging_models
from waldur_core.structure import models as structure_models, SupportedServices

class CustomIndexDashboard(FluentIndexDashboard):
    """
    Custom index dashboard for admin site.
    """
    title = _('Waldur administration')

    def _get_installed_plugin_info(self):
        links = []

        for ext in WaldurExtension.get_extensions():
            app_config = self._get_app_config(ext.django_app())
            if not app_config:
                # App is not found
                continue

            name = self._get_app_name(app_config)
            version = self._get_app_version(app_config)

            links.append(
                {
                    'title': '%s %s' % (name, version),
                    'url': 'http://docs.waldur.com/',
                    'external': True,
                    'attrs': {'target': '_blank'},
                }
            )

        # Order plugins by title
        links = sorted(links, key=lambda item: item['title'].lower())

        # Core is the most important component, therefore
        # it should be the the first item in the list
        links.insert(0, {
            'title': _('Waldur Core %s') % __version__,
            'url': 'http://docs.waldur.com/',
            'external': True,
            'attrs': {'target': '_blank'},
        })

        return links

    def _get_app_config(self, app_name):
        """
        Returns an app config for the given name, not by label.
        """

        matches = [app_config for app_config in apps.get_app_configs()
                   if app_config.name == app_name]
        if not matches:
            return
        return matches[0]

    def _get_app_name(self, app_config):
        """
        Strip redundant prefixes, because some apps
        don't specify prefix, while others use deprecated prefix.
        """

        return app_config.verbose_name\
            .replace('Waldur', '')\
            .strip()

    def _get_app_version(self, app_config):
        """
        Some plugins ship multiple applications and extensions.
        However all of them have the same version, because they are released together.
        That's why only-top level module is used to fetch version information.
        """

        base_name = app_config.__module__.split('.')[0]
        module = __import__(base_name)
        return getattr(module, '__version__', 'N/A')

    def _get_quick_access_info(self):
        """
        Returns a list of ListLink items to be added to Quick Access tab.
        Contains:
        - links to Organizations, Projects and Users;
        - a link to shared service settings;
        - custom configured links in admin/settings FLUENT_DASHBOARD_QUICK_ACCESS_LINKS attribute;
        """
        quick_access_links = []

        # add custom links
        quick_access_links.extend(settings.FLUENT_DASHBOARD_QUICK_ACCESS_LINKS)

        for model in (structure_models.Project,
                      structure_models.Customer,
                      core_models.User,
                      structure_models.SharedServiceSettings,
                      logging_models.Report,):
            quick_access_links.append(self._get_link_to_model(model))

        return quick_access_links

    def _get_erred_resource_link(self, model, erred_amount, erred_state):
        result = self._get_link_to_model(model)
        result['title'] = _('%(num)s %(resources)s in ERRED state') % {
            'num': erred_amount,
            'resources': result['title']
        }
        if 'url' in result:
            result['url'] = '%s?shared__exact=1&state__exact=%s' % (result['url'], erred_state)
        return result

    def _get_link_to_model(self, model):
        result = {
            'title': six.text_type(model._meta.verbose_name_plural).capitalize(),
            'external': True,
            'attrs': {'target': '_blank'},
        }
        try:
            result['url'] = reverse('admin:%s_%s_changelist' % (model._meta.app_label, model._meta.model_name))
        except NoReverseMatch:
            pass
        return result

    def _get_link_to_instance(self, instance):
        result = {
            'title': six.text_type(instance),
            'external': True,
            'attrs': {'target': '_blank'},
        }
        try:
            result['url'] = reverse('admin:%s_%s_change' % (instance._meta.app_label, instance._meta.model_name),
                                    args=(instance.pk,))
        except NoReverseMatch:
            pass
        return result

    def _get_erred_shared_settings_module(self):
        """
        Returns a LinkList based module which contains link to shared service setting instances in ERRED state.
        """
        result_module = modules.LinkList(title=_('Shared provider settings in erred state'))
        result_module.template = 'admin/dashboard/erred_link_list.html'
        erred_state = structure_models.SharedServiceSettings.States.ERRED

        queryset = structure_models.SharedServiceSettings.objects
        settings_in_erred_state = queryset.filter(state=erred_state).count()

        if settings_in_erred_state:
            result_module.title = '%s (%s)' % (result_module.title, settings_in_erred_state)
            for service_settings in queryset.filter(state=erred_state).iterator():
                module_child = self._get_link_to_instance(service_settings)
                module_child['error'] = service_settings.error_message
                result_module.children.append(module_child)
        else:
            result_module.pre_content = _('Nothing found.')

        return result_module

    def _get_erred_resources_module(self):
        """
        Returns a list of links to resources which are in ERRED state and linked to a shared service settings.
        """
        result_module = modules.LinkList(title=_('Resources in erred state'))
        erred_state = structure_models.NewResource.States.ERRED
        children = []

        resource_models = SupportedServices.get_resource_models()
        resources_in_erred_state_overall = 0
        for resource_type, resource_model in resource_models.items():
            queryset = resource_model.objects.filter(service_project_link__service__settings__shared=True)
            erred_amount = queryset.filter(state=erred_state).count()
            if erred_amount:
                resources_in_erred_state_overall = resources_in_erred_state_overall + erred_amount
                link = self._get_erred_resource_link(resource_model, erred_amount, erred_state)
                children.append(link)

        if resources_in_erred_state_overall:
            result_module.title = '%s (%s)' % (result_module.title, resources_in_erred_state_overall)
            result_module.children = children
        else:
            result_module.pre_content = _('Nothing found.')

        return result_module

    def __init__(self, **kwargs):
        FluentIndexDashboard.__init__(self, **kwargs)

        self.children.append(modules.LinkList(
            _('Installed components'),
            layout='stacked',
            enabled=False,
            draggable=True,
            deletable=True,
            collapsible=True,
            children=self._get_installed_plugin_info()
        ))

        self.children.append(modules.LinkList(
            _('Quick access'),
            children=self._get_quick_access_info())
        )

        self.children.append(self._get_erred_shared_settings_module())
        self.children.append(self._get_erred_resources_module())

class CustomAppIndexDashboard(FluentAppIndexDashboard):
    def __init__(self, app_title, models, **kwargs):
        super(CustomAppIndexDashboard, self).__init__(app_title, models, **kwargs)
        path = self._get_app_models_path()
        self.children = [modules.ModelList(title=app_title, models=[path])]

    def _get_app_models_path(self):
        return '%s.models.*' % self.app_title.replace(' ', '.', 1).replace(' ', '_').lower()
<commit_msg>Skip links without URLs in Django admin [WAL-2580]<commit_after>from __future__ import unicode_literals

from django.apps import apps
from django.conf import settings
from django.urls import reverse, NoReverseMatch
from django.utils.translation import ugettext_lazy as _
from fluent_dashboard.dashboard import modules, FluentIndexDashboard, FluentAppIndexDashboard
import six

from waldur_core import __version__
from waldur_core.core import models as core_models, WaldurExtension
from waldur_core.logging import models as logging_models
from waldur_core.structure import models as structure_models, SupportedServices

class CustomIndexDashboard(FluentIndexDashboard):
    """
    Custom index dashboard for admin site.
    """
    title = _('Waldur administration')

    def _get_installed_plugin_info(self):
        links = []

        for ext in WaldurExtension.get_extensions():
            app_config = self._get_app_config(ext.django_app())
            if not app_config:
                # App is not found
                continue

            name = self._get_app_name(app_config)
            version = self._get_app_version(app_config)

            links.append(
                {
                    'title': '%s %s' % (name, version),
                    'url': 'http://docs.waldur.com/',
                    'external': True,
                    'attrs': {'target': '_blank'},
                }
            )

        # Order plugins by title
        links = sorted(links, key=lambda item: item['title'].lower())

        # Core is the most important component, therefore
        # it should be the the first item in the list
        links.insert(0, {
            'title': _('Waldur Core %s') % __version__,
            'url': 'http://docs.waldur.com/',
            'external': True,
            'attrs': {'target': '_blank'},
        })

        return links

    def _get_app_config(self, app_name):
        """
        Returns an app config for the given name, not by label.
        """

        matches = [app_config for app_config in apps.get_app_configs()
                   if app_config.name == app_name]
        if not matches:
            return
        return matches[0]

    def _get_app_name(self, app_config):
        """
        Strip redundant prefixes, because some apps
        don't specify prefix, while others use deprecated prefix.
        """

        return app_config.verbose_name\
            .replace('Waldur', '')\
            .strip()

    def _get_app_version(self, app_config):
        """
        Some plugins ship multiple applications and extensions.
        However all of them have the same version, because they are released together.
        That's why only-top level module is used to fetch version information.
        """

        base_name = app_config.__module__.split('.')[0]
        module = __import__(base_name)
        return getattr(module, '__version__', 'N/A')

    def _get_quick_access_info(self):
        """
        Returns a list of ListLink items to be added to Quick Access tab.
        Contains:
        - links to Organizations, Projects and Users;
        - a link to shared service settings;
        - custom configured links in admin/settings FLUENT_DASHBOARD_QUICK_ACCESS_LINKS attribute;
        """
        quick_access_links = []

        # add custom links
        quick_access_links.extend(settings.FLUENT_DASHBOARD_QUICK_ACCESS_LINKS)

        for model in (structure_models.Project,
                      structure_models.Customer,
                      core_models.User,
                      structure_models.SharedServiceSettings,
                      logging_models.Report,):
            link = self._get_link_to_model(model)
            if 'url' in link:
                quick_access_links.append(link)

        return quick_access_links

    def _get_erred_resource_link(self, model, erred_amount, erred_state):
        result = self._get_link_to_model(model)
        result['title'] = _('%(num)s %(resources)s in ERRED state') % {
            'num': erred_amount,
            'resources': result['title']
        }
        if 'url' in result:
            result['url'] = '%s?shared__exact=1&state__exact=%s' % (result['url'], erred_state)
        return result

    def _get_link_to_model(self, model):
        result = {
            'title': six.text_type(model._meta.verbose_name_plural).capitalize(),
            'external': True,
            'attrs': {'target': '_blank'},
        }
        try:
            result['url'] = reverse('admin:%s_%s_changelist' % (model._meta.app_label, model._meta.model_name))
        except NoReverseMatch:
            pass
        return result

    def _get_link_to_instance(self, instance):
        result = {
            'title': six.text_type(instance),
            'external': True,
            'attrs': {'target': '_blank'},
        }
        try:
            result['url'] = reverse('admin:%s_%s_change' % (instance._meta.app_label, instance._meta.model_name),
                                    args=(instance.pk,))
        except NoReverseMatch:
            pass
        return result

    def _get_erred_shared_settings_module(self):
        """
        Returns a LinkList based module which contains link to shared service setting instances in ERRED state.
        """
        result_module = modules.LinkList(title=_('Shared provider settings in erred state'))
        result_module.template = 'admin/dashboard/erred_link_list.html'
        erred_state = structure_models.SharedServiceSettings.States.ERRED

        queryset = structure_models.SharedServiceSettings.objects
        settings_in_erred_state = queryset.filter(state=erred_state).count()

        if settings_in_erred_state:
            result_module.title = '%s (%s)' % (result_module.title, settings_in_erred_state)
            for service_settings in queryset.filter(state=erred_state).iterator():
                module_child = self._get_link_to_instance(service_settings)
                module_child['error'] = service_settings.error_message
                result_module.children.append(module_child)
        else:
            result_module.pre_content = _('Nothing found.')

        return result_module

    def _get_erred_resources_module(self):
        """
        Returns a list of links to resources which are in ERRED state and linked to a shared service settings.
        """
        result_module = modules.LinkList(title=_('Resources in erred state'))
        erred_state = structure_models.NewResource.States.ERRED
        children = []

        resource_models = SupportedServices.get_resource_models()
        resources_in_erred_state_overall = 0
        for resource_type, resource_model in resource_models.items():
            queryset = resource_model.objects.filter(service_project_link__service__settings__shared=True)
            erred_amount = queryset.filter(state=erred_state).count()
            if erred_amount:
                resources_in_erred_state_overall = resources_in_erred_state_overall + erred_amount
                link = self._get_erred_resource_link(resource_model, erred_amount, erred_state)
                if 'url' in link:
                    children.append(link)

        if resources_in_erred_state_overall:
            result_module.title = '%s (%s)' % (result_module.title, resources_in_erred_state_overall)
            result_module.children = children
        else:
            result_module.pre_content = _('Nothing found.')

        return result_module

    def __init__(self, **kwargs):
        FluentIndexDashboard.__init__(self, **kwargs)

        self.children.append(modules.LinkList(
            _('Installed components'),
            layout='stacked',
            enabled=False,
            draggable=True,
            deletable=True,
            collapsible=True,
            children=self._get_installed_plugin_info()
        ))

        self.children.append(modules.LinkList(
            _('Quick access'),
            children=self._get_quick_access_info())
        )

        self.children.append(self._get_erred_shared_settings_module())
        self.children.append(self._get_erred_resources_module())

class CustomAppIndexDashboard(FluentAppIndexDashboard):
    def __init__(self, app_title, models, **kwargs):
        super(CustomAppIndexDashboard, self).__init__(app_title, models, **kwargs)
        path = self._get_app_models_path()
        self.children = [modules.ModelList(title=app_title, models=[path])]

    def _get_app_models_path(self):
        return '%s.models.*' % self.app_title.replace(' ', '.', 1).replace(' ', '_').lower()
<|endoftext|>