<commit_before>from os.path import join, abspath, dirname, getsize, curdir
from netCDF4 import Dataset
from flyingpigeon import config
import logging
LOGGER = logging.getLogger("PYWPS")

def has_Lambert_Conformal(resource):
    """
    Check if grid is organised as Lambert_Conformal

    :param resource: file to be checked

    :return Boolean: True/False
    """
    if type(resource) != list:
        resource = [resource]
    for nc in resource:
        ds = Dataset(nc)
        if 'Lambert_Conformal' not in ds.variables.keys():
            return False
    return True

def call(resource=[], variable=None, dimension_map=None, calc=None,
         calc_grouping=None, conform_units_to=None, crs=None, memory_limit=None,  prefix=None,
         regrid_destination=None, regrid_options='bil', level_range=None,
         geom=None, output_format_options=None, search_radius_mult=2.,
         select_nearest=False, select_ugid=None, spatial_wrapping=None,
         t_calendar=None, time_region=None,
         time_range=None, dir_output=None, output_format='nc'):
    '''
    ocgis operation call

    :param resource:
    :param variable: variable in the input file to be picked
    :param dimension_map: dimension map in case of unconventional storage of data
    :param calc: ocgis calc syntax for calculation partion
    :param calc_grouping: time aggregate grouping
    :param conform_units_to:
    :param crs: coordinate reference system
    :param memory_limit: limit the amount of data to be loaded into the memory at once \
        if None (default) free memory is detected by birdhouse
    :param level_range: subset of given levels
    :param prefix: string for the file base name
    :param regrid_destination: file path with netCDF file with grid for output file
    :param geom: name of shapefile stored in birdhouse shape cabinet
    :param output_format_options: output options for netCDF e.g compression level()
    :param regrid_destination: file containing the targed grid (griddes.txt or netCDF file)
    :param regrid_options: methods for regridding:
                          'bil' = Bilinear interpolation
                          'bic' = Bicubic interpolation
                          'dis' = Distance-weighted average remapping
                          'nn' = nearest neighbour
                          'con' = First-order conservative remapping
                          'laf' = largest area fraction reamapping
    :param search_radius_mult: search radius for point geometries. All included gridboxes will be returned
    :param select_nearest: nearest neighbour selection for point geometries
    :param select_ugid: ugid for appropriate polygons
    :param spatial_wrapping: how to handle coordinates in case of subsets, options: None (default), 'wrap', 'unwrap'
    :param time_region: select single month
    :param time_range: sequence of two datetime.datetime objects to mark start and end point
    :param dir_output (default= curdir):
    :param output_format:
    :return: output file path
    '''
    LOGGER.info('Start ocgis module call function')
    from ocgis import OcgOperations, RequestDataset, env, DimensionMap, crs
    from ocgis.util.large_array import compute
    from datetime import datetime as dt
    import uuid

    # prepare the environment
    env.DIR_SHPCABINET = config.shapefiles_path()
    env.OVERWRITE = True
    # env.DIR_OUTPUT = dir_output
    # LOGGER.debug(' **** env.DIR_OUTPUT  = %s ' % env.DIR_OUTPUT)

    if dir_output is None:
        dir_output = abspath(curdir)

    # check time_range fromat:

    if time_range is not None:
        try:
            LOGGER.debug('time_range type= %s , %s ' % (type(time_range[0]), type(time_range[1])))
            if type(time_range[0] is 'datetime.date'):
                time_range = [dt.combine(time_range[0], dt.min.time()),
                              dt.combine(time_range[1], dt.min.time())]
            LOGGER.debug('time_range changed to type= %s , %s ' % (type(time_range[0]), type(time_range[1])))
        except:
            LOGGER.exception('failed to confert data to datetime')

    if spatial_wrapping == 'wrap':
        spatial_reorder = True
    else:
        spatial_reorder = False
    LOGGER.debug('spatial_reorder: %s and spatial_wrapping: %s ' % (spatial_reorder, spatial_wrapping))

    if prefix is None:
        prefix = str(uuid.uuid1())
        env.PREFIX = prefix
    #
    # if output_format_options is False:
    #     output_format_options = None
    # elif output_format_options is True:
    #     output_format_options = {'data_model': 'NETCDF4',  # NETCDF4_CLASSIC
    #                              'variable_kwargs': {'zlib': True, 'complevel': 9}}
    # else:
    if output_format_options is not None:
        LOGGER.info('output_format_options are set to %s ' % (output_format_options))

    if type(resource) != list:
        resource = list([resource])
    # execute ocgis
    LOGGER.info('Execute ocgis module call function')

    # # needed for some AFR-44 data
    # dimension_map = DimensionMap()
    # dimension_map.set_variable('x', 'lon', dimension='rlon')
    # dimension_map.set_variable('y', 'lat', dimension='rlat')
    # dimension_map.set_variable('time', 'time', dimension='time')
    # crs=crs.Spherical()

    try:
        LOGGER.debug('call module curdir = %s ' % abspath(curdir))
        rd = RequestDataset(resource, variable=variable, level_range=level_range,
                            dimension_map=dimension_map, conform_units_to=conform_units_to,
                            time_region=time_region, t_calendar=t_calendar, time_range=time_range, crs=crs)

        # from ocgis.constants import DimensionMapKey
        # rd.dimension_map.set_bounds(DimensionMapKey.TIME, None)

        ops = OcgOperations(dataset=rd,
                            output_format_options=output_format_options,
                            dir_output=dir_output,
                            spatial_wrapping=spatial_wrapping,
                            spatial_reorder=spatial_reorder,
                            # regrid_destination=rd_regrid,
                            # options=options,
                            calc=calc,
                            calc_grouping=calc_grouping,
                            geom=geom,
                            output_format=output_format,
                            prefix=prefix,
                            search_radius_mult=search_radius_mult,
                            select_nearest=select_nearest,
                            select_ugid=select_ugid,
                            add_auxiliary_files=False)
        LOGGER.info('OcgOperations set')
    except:
        LOGGER.exception('failed to setup OcgOperations')
        return None

    try:
        LOGGER.info('ocgis module call as ops.execute()')
        geom_file = ops.execute()
    except:
        LOGGER.exception('failed to execute ocgis operation')
        return None
    #
    # try:
    #     from numpy import sqrt
    #     from flyingpigeon.utils import FreeMemory
    #
    #     if memory_limit is None:
    #         f = FreeMemory()
    #         mem_kb = f.user_free
    #         mem_mb = mem_kb / 1024.
    #         mem_limit = mem_mb / 2.  # set limit to half of the free memory
    #     else:
    #         mem_limit = memory_limit
    #
    #     if mem_limit >= 1024. * 4:
    #         mem_limit = 1024. * 4
    #         # 475.0 MB for openDAP
    #
    #     LOGGER.info('memory_limit = %s Mb' % (mem_limit))
    #
    #     data_kb = ops.get_base_request_size()['total']
    #     data_mb = data_kb / 1024.
    #
    #     # data_kb = size['total']/reduce(lambda x,y: x*y,size['variables'][variable]['value']['shape'])
    #     LOGGER.info('data_mb  = %s Mb' % (data_mb))
    #
    #     if data_mb <= mem_limit:  # input is smaler than the half of free memory size
    #         try:
    #             LOGGER.info('ocgis module call as ops.execute()')
    #             geom_file = ops.execute()
    #         except Exception as e:
    #             LOGGER.debug('failed to execute ocgis operation')
    #             raise
    #             return None
    #
    #     else:
    #         ##########################
    #         # calcultion of chunk size
    #         ##########################
    #         try:
    #             size = ops.get_base_request_size()
    #             nb_time_coordinates_rd = size['variables'][variable]['temporal']['shape'][0]
    #             element_in_kb = size['total']/reduce(lambda x, y: x*y, size['variables'][variable]['value']['shape'])
    #             element_in_mb = element_in_kb / 1024.
    #             tile_dim = sqrt(mem_limit/(element_in_mb*nb_time_coordinates_rd))  # maximum chunk size
    #
    #             LOGGER.info('ocgis module call compute with chunks')
    #             if calc is None:
    #                 calc = '%s=%s*1' % (variable, variable)
    #                 LOGGER.info('calc set to = %s ' % calc)
    #             ops = OcgOperations(dataset=rd,
    #                                 output_format_options=output_format_options,
    #                                 dir_output=dir_output,
    #                                 spatial_wrapping=spatial_wrapping,
    #                                 spatial_reorder=spatial_reorder,
    #                                 # regrid_destination=rd_regrid,
    #                                 # options=options,
    #                                 calc=calc,
    #                                 calc_grouping=calc_grouping,
    #                                 geom=geom,
    #                                 output_format=output_format,
    #                                 prefix=prefix,
    #                                 search_radius_mult=search_radius_mult,
    #                                 select_nearest=select_nearest,
    #                                 select_ugid=select_ugid,
    #                                 add_auxiliary_files=False)
    #             geom_file = compute(ops, tile_dimension=int(tile_dim), verbose=True)
    #             print 'ocgis calculated'
    #         except Exception as e:
    #             LOGGER.debug('failed to compute ocgis with chunks')
    #             raise
    #             return None
    #     LOGGER.info('Succeeded with ocgis module call function')
    # except:
    #     LOGGER.exception('failed to compare dataload with free memory, calling as execute instead')

    ############################################
    # remapping according to regrid informations
    ############################################
    if regrid_destination is not None:
        try:
            from tempfile import mkstemp
            from cdo import Cdo
            cdo = Cdo()
            output = '%s.nc' % uuid.uuid1()
            remap = 'remap%s' % regrid_options
            call = [op for op in dir(cdo) if remap in op]
            cmd = "output = cdo.%s('%s',input='%s', output='%s')" \
                  % (str(call[0]), regrid_destination, geom_file, output)
            exec cmd
        except Exception as e:
            LOGGER.debug('failed to remap')
            raise
            return None
    else:
        output = geom_file

    # try:
    #     from flyingpigeon.utils import unrotate_pole
    #     lat, lon = unrotate_pole(output)
    # except:
    #     LOGGER.exception('failed to unrotate pole')
    return output

def eval_timerange(resource, time_range):
    """
    quality checker if given time_range is covered by timesteps in resource files

    :param resource: input netCDF files
    :param time_range: start and end date of time range [datetime,datetime]

    :returns [datetime,datetime]: time_range
    """
    from flyingpigeon.utils import get_time

    LOGGER.info('time_range: %s' % time_range)

    if type(resource) != str:
        resource.sort()
    time = get_time(resource)
    start = time[0]
    end = time[-1]

    if (time_range[0] > start or time_range[0] < end):
        LOGGER.debug('time range start %s not in input dataset covering: %s to %s' % (time_range[0], start, end))
        time_range[0] = start
    LOGGER.debug('time_range start changed to first timestep of dataset')
    if (time_range[1] > end or time_range[1] < start):
        LOGGER.debug('time range end %s not in input dataset covering: %s to %s' % (time_range[0], start, end))
        time_range[1] = end
    LOGGER.debug('time_range end changed to last timestep of dataset')
    if (time_range[0] > time_range[1]):
        time_range = reversed(time_range)
        LOGGER.debug('time range reversed! start was later than end ')
    LOGGER.info('time range start and end set')
    return time_range

# # check memory load
# from os import stat
#   if memory_limit == None:
#     f = FreeMemory()
#     mem_kb = f.user_free
#     mem_mb = mem_kb / 1024.
#     mem_limit = mem_mb / 2. # set limit to half of the free memory
#   else:
#     mem_limit = memory_limit
#
#   if mem_limit >= 1024. * 4:
#     mem_limit = 1024. * 4
#     # 475.0 MB for openDAP
#
#   #if type(resource) == list :
#     #data_kb =  stat(resource[0]).st_size * len(resource)
#   #else:
#     #data_kb =  stat(resource).st_size
#   size = ops.get_base_request_size()['total']
#   data_kb = size['total']/reduce(lambda x,y: x*y,size['variables'][variable]['value']['shape'])
#   data_mb = data_kb / 1024.
#
#   if variable == None:
#     variable = rd.variable
#     LOGGER.info('%s as variable dedected' % (variable))
<commit_msg>crs in ocgis_modul<commit_after>from os.path import join, abspath, dirname, getsize, curdir
from netCDF4 import Dataset
from flyingpigeon import config
import logging
LOGGER = logging.getLogger("PYWPS")

def has_Lambert_Conformal(resource):
    """
    Check if grid is organised as Lambert_Conformal

    :param resource: file to be checked

    :return Boolean: True/False
    """
    if type(resource) != list:
        resource = [resource]
    for nc in resource:
        ds = Dataset(nc)
        if 'Lambert_Conformal' not in ds.variables.keys():
            return False
    return True

def call(resource=[], variable=None, dimension_map=None, calc=None,
         calc_grouping=None, conform_units_to=None, crs=None, memory_limit=None,  prefix=None,
         regrid_destination=None, regrid_options='bil', level_range=None,
         geom=None, output_format_options=None, search_radius_mult=2.,
         select_nearest=False, select_ugid=None, spatial_wrapping=None,
         t_calendar=None, time_region=None,
         time_range=None, dir_output=None, output_format='nc'):
    '''
    ocgis operation call

    :param resource:
    :param variable: variable in the input file to be picked
    :param dimension_map: dimension map in case of unconventional storage of data
    :param calc: ocgis calc syntax for calculation partion
    :param calc_grouping: time aggregate grouping
    :param conform_units_to:
    :param crs: coordinate reference system
    :param memory_limit: limit the amount of data to be loaded into the memory at once \
        if None (default) free memory is detected by birdhouse
    :param level_range: subset of given levels
    :param prefix: string for the file base name
    :param regrid_destination: file path with netCDF file with grid for output file
    :param geom: name of shapefile stored in birdhouse shape cabinet
    :param output_format_options: output options for netCDF e.g compression level()
    :param regrid_destination: file containing the targed grid (griddes.txt or netCDF file)
    :param regrid_options: methods for regridding:
                          'bil' = Bilinear interpolation
                          'bic' = Bicubic interpolation
                          'dis' = Distance-weighted average remapping
                          'nn' = nearest neighbour
                          'con' = First-order conservative remapping
                          'laf' = largest area fraction reamapping
    :param search_radius_mult: search radius for point geometries. All included gridboxes will be returned
    :param select_nearest: nearest neighbour selection for point geometries
    :param select_ugid: ugid for appropriate polygons
    :param spatial_wrapping: how to handle coordinates in case of subsets, options: None (default), 'wrap', 'unwrap'
    :param time_region: select single month
    :param time_range: sequence of two datetime.datetime objects to mark start and end point
    :param dir_output (default= curdir):
    :param output_format:
    :return: output file path
    '''
    LOGGER.info('Start ocgis module call function')
    from ocgis import OcgOperations, RequestDataset, env, DimensionMap, crs
    from ocgis.util.large_array import compute
    from datetime import datetime as dt
    import uuid

    # prepare the environment
    env.DIR_SHPCABINET = config.shapefiles_path()
    env.OVERWRITE = True
    # env.DIR_OUTPUT = dir_output
    # LOGGER.debug(' **** env.DIR_OUTPUT  = %s ' % env.DIR_OUTPUT)

    if dir_output is None:
        dir_output = abspath(curdir)

    # check time_range fromat:

    if time_range is not None:
        try:
            LOGGER.debug('time_range type= %s , %s ' % (type(time_range[0]), type(time_range[1])))
            if type(time_range[0] is 'datetime.date'):
                time_range = [dt.combine(time_range[0], dt.min.time()),
                              dt.combine(time_range[1], dt.min.time())]
            LOGGER.debug('time_range changed to type= %s , %s ' % (type(time_range[0]), type(time_range[1])))
        except:
            LOGGER.exception('failed to confert data to datetime')

    if spatial_wrapping == 'wrap':
        spatial_reorder = True
    else:
        spatial_reorder = False
    LOGGER.debug('spatial_reorder: %s and spatial_wrapping: %s ' % (spatial_reorder, spatial_wrapping))

    if prefix is None:
        prefix = str(uuid.uuid1())
        env.PREFIX = prefix
    #
    # if output_format_options is False:
    #     output_format_options = None
    # elif output_format_options is True:
    #     output_format_options = {'data_model': 'NETCDF4',  # NETCDF4_CLASSIC
    #                              'variable_kwargs': {'zlib': True, 'complevel': 9}}
    # else:
    if output_format_options is not None:
        LOGGER.info('output_format_options are set to %s ' % (output_format_options))

    if type(resource) != list:
        resource = list([resource])
    # execute ocgis
    LOGGER.info('Execute ocgis module call function')

    # # needed for some AFR-44 data
    # dimension_map = DimensionMap()
    # dimension_map.set_variable('x', 'lon', dimension='rlon')
    # dimension_map.set_variable('y', 'lat', dimension='rlat')
    # dimension_map.set_variable('time', 'time', dimension='time')
    # crs=crs.Spherical()

    try:
        LOGGER.debug('call module curdir = %s ' % abspath(curdir))
        rd = RequestDataset(resource, variable=variable, level_range=level_range,
                            dimension_map=dimension_map, conform_units_to=conform_units_to,
                            time_region=time_region, t_calendar=t_calendar, time_range=time_range, crs=crs)

        # from ocgis.constants import DimensionMapKey
        # rd.dimension_map.set_bounds(DimensionMapKey.TIME, None)

        ops = OcgOperations(dataset=rd,
                            output_format_options=output_format_options,
                            dir_output=dir_output,
                            spatial_wrapping=spatial_wrapping,
                            spatial_reorder=spatial_reorder,
                            # regrid_destination=rd_regrid,
                            # options=options,
                            calc=calc,
                            calc_grouping=calc_grouping,
                            geom=geom,
                            output_format=output_format,
                            prefix=prefix,
                            search_radius_mult=search_radius_mult,
                            select_nearest=select_nearest,
                            select_ugid=select_ugid,
                            add_auxiliary_files=False)
        LOGGER.info('OcgOperations set')
    except:
        LOGGER.exception('failed to setup OcgOperations')
        return None

    try:
        LOGGER.info('ocgis module call as ops.execute()')
        geom_file = ops.execute()
    except:
        LOGGER.exception('failed to execute ocgis operation')
        return None

    #
    # try:
    #     from numpy import sqrt
    #     from flyingpigeon.utils import FreeMemory
    #
    #     if memory_limit is None:
    #         f = FreeMemory()
    #         mem_kb = f.user_free
    #         mem_mb = mem_kb / 1024.
    #         mem_limit = mem_mb / 2.  # set limit to half of the free memory
    #     else:
    #         mem_limit = memory_limit
    #
    #     if mem_limit >= 1024. * 4:
    #         mem_limit = 1024. * 4
    #         # 475.0 MB for openDAP
    #
    #     LOGGER.info('memory_limit = %s Mb' % (mem_limit))
    #
    #     data_kb = ops.get_base_request_size()['total']
    #     data_mb = data_kb / 1024.
    #
    #     # data_kb = size['total']/reduce(lambda x,y: x*y,size['variables'][variable]['value']['shape'])
    #     LOGGER.info('data_mb  = %s Mb' % (data_mb))
    #
    #     if data_mb <= mem_limit:  # input is smaler than the half of free memory size
    #         try:
    #             LOGGER.info('ocgis module call as ops.execute()')
    #             geom_file = ops.execute()
    #         except Exception as e:
    #             LOGGER.debug('failed to execute ocgis operation')
    #             raise
    #             return None
    #
    #     else:
    #         ##########################
    #         # calcultion of chunk size
    #         ##########################
    #         try:
    #             size = ops.get_base_request_size()
    #             nb_time_coordinates_rd = size['variables'][variable]['temporal']['shape'][0]
    #             element_in_kb = size['total']/reduce(lambda x, y: x*y, size['variables'][variable]['value']['shape'])
    #             element_in_mb = element_in_kb / 1024.
    #             tile_dim = sqrt(mem_limit/(element_in_mb*nb_time_coordinates_rd))  # maximum chunk size
    #
    #             LOGGER.info('ocgis module call compute with chunks')
    #             if calc is None:
    #                 calc = '%s=%s*1' % (variable, variable)
    #                 LOGGER.info('calc set to = %s ' % calc)
    #             ops = OcgOperations(dataset=rd,
    #                                 output_format_options=output_format_options,
    #                                 dir_output=dir_output,
    #                                 spatial_wrapping=spatial_wrapping,
    #                                 spatial_reorder=spatial_reorder,
    #                                 # regrid_destination=rd_regrid,
    #                                 # options=options,
    #                                 calc=calc,
    #                                 calc_grouping=calc_grouping,
    #                                 geom=geom,
    #                                 output_format=output_format,
    #                                 prefix=prefix,
    #                                 search_radius_mult=search_radius_mult,
    #                                 select_nearest=select_nearest,
    #                                 select_ugid=select_ugid,
    #                                 add_auxiliary_files=False)
    #             geom_file = compute(ops, tile_dimension=int(tile_dim), verbose=True)
    #             print 'ocgis calculated'
    #         except Exception as e:
    #             LOGGER.debug('failed to compute ocgis with chunks')
    #             raise
    #             return None
    #     LOGGER.info('Succeeded with ocgis module call function')
    # except:
    #     LOGGER.exception('failed to compare dataload with free memory, calling as execute instead')

    ############################################
    # remapping according to regrid informations
    ############################################
    if regrid_destination is not None:
        try:
            from tempfile import mkstemp
            from cdo import Cdo
            cdo = Cdo()
            output = '%s.nc' % uuid.uuid1()
            remap = 'remap%s' % regrid_options
            call = [op for op in dir(cdo) if remap in op]
            cmd = "output = cdo.%s('%s',input='%s', output='%s')" \
                  % (str(call[0]), regrid_destination, geom_file, output)
            exec cmd
        except Exception as e:
            LOGGER.debug('failed to remap')
            raise
            return None
    else:
        output = geom_file

    # try:
    #     from flyingpigeon.utils import unrotate_pole
    #     lat, lon = unrotate_pole(output)
    # except:
    #     LOGGER.exception('failed to unrotate pole')
    return output

def eval_timerange(resource, time_range):
    """
    quality checker if given time_range is covered by timesteps in resource files

    :param resource: input netCDF files
    :param time_range: start and end date of time range [datetime,datetime]

    :returns [datetime,datetime]: time_range
    """
    from flyingpigeon.utils import get_time

    LOGGER.info('time_range: %s' % time_range)

    if type(resource) != str:
        resource.sort()
    time = get_time(resource)
    start = time[0]
    end = time[-1]

    if (time_range[0] > start or time_range[0] < end):
        LOGGER.debug('time range start %s not in input dataset covering: %s to %s' % (time_range[0], start, end))
        time_range[0] = start
    LOGGER.debug('time_range start changed to first timestep of dataset')
    if (time_range[1] > end or time_range[1] < start):
        LOGGER.debug('time range end %s not in input dataset covering: %s to %s' % (time_range[0], start, end))
        time_range[1] = end
    LOGGER.debug('time_range end changed to last timestep of dataset')
    if (time_range[0] > time_range[1]):
        time_range = reversed(time_range)
        LOGGER.debug('time range reversed! start was later than end ')
    LOGGER.info('time range start and end set')
    return time_range

# # check memory load
# from os import stat
#   if memory_limit == None:
#     f = FreeMemory()
#     mem_kb = f.user_free
#     mem_mb = mem_kb / 1024.
#     mem_limit = mem_mb / 2. # set limit to half of the free memory
#   else:
#     mem_limit = memory_limit
#
#   if mem_limit >= 1024. * 4:
#     mem_limit = 1024. * 4
#     # 475.0 MB for openDAP
#
#   #if type(resource) == list :
#     #data_kb =  stat(resource[0]).st_size * len(resource)
#   #else:
#     #data_kb =  stat(resource).st_size
#   size = ops.get_base_request_size()['total']
#   data_kb = size['total']/reduce(lambda x,y: x*y,size['variables'][variable]['value']['shape'])
#   data_mb = data_kb / 1024.
#
#   if variable == None:
#     variable = rd.variable
#     LOGGER.info('%s as variable dedected' % (variable))
<|endoftext|>


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

# The MIT License (MIT)

# Copyright (c) 2016-2017 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

import matplotlib
import warnings
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    matplotlib.use('Agg')
import matplotlib.pyplot as plt

import os
import os.path
import datetime
import numpy as np
from pyannote.audio.util import mkdir_p
from keras.callbacks import Callback

class LoggingCallback(Callback):
    """Logging callback

    Parameters
    ----------
    log_dir : str
    log : list of tuples, optional
        Defaults to [('train', 'loss')]
    extract_embedding : func
        Function that takes Keras model as input and returns the actual model.
        This is useful for embedding that are not directly optimized by Keras.
        Defaults to identity function.
    restart : boolean, optional
        Indicates that this training is a restart, not a cold start (default).
    """

    ARCHITECTURE_YML = '{log_dir}/architecture.yml'
    WEIGHTS_DIR = '{log_dir}/weights'
    WEIGHTS_H5 = '{log_dir}/weights/{epoch:04d}.h5'
    WEIGHTS_PT = '{log_dir}/weights/{epoch:04d}.pt'
    OPTIMIZER_PT = '{log_dir}/weights/{epoch:04d}.optimizer.pt'
    LOG_TXT = '{log_dir}/{name}.{subset}.txt'
    LOG_PNG = '{log_dir}/{name}.{subset}.png'
    LOG_EPS = '{log_dir}/{name}.{subset}.eps'

    def __init__(self, log_dir, log=None, extract_embedding=None,
                 restart=False, backend='keras'):
        super(LoggingCallback, self).__init__()

        # make sure path is absolute
        self.log_dir = os.path.realpath(log_dir)

        if extract_embedding is None:
            extract_embedding = lambda model: model
        self.extract_embedding = extract_embedding

        # create log_dir directory
        mkdir_p(self.log_dir)

        # this will fail if the directory already exists
        # and this is OK  because 'weights' directory
        # usually contains the output of very long computations
        # and you do not want to erase them by mistake :/
        self.restart = restart
        if not self.restart:
            weights_dir = self.WEIGHTS_DIR.format(log_dir=self.log_dir)
            os.makedirs(weights_dir)

        if log is None:
            log = [('train', 'loss')]
        self.log = log

        self.values = {}
        for subset, name in self.log:
            if subset not in self.values:
                self.values[subset] = {}
            if name not in self.values[subset]:
                self.values[subset][name] = []

        self.backend = backend

    def get_loss(self, epoch, subset, logs={}):
        if subset != 'train':
            raise ValueError('"loss" is only available for "train"')
        value = logs['loss']
        minimize = True
        return value, minimize

    def get_accuracy(self, epoch, subset, logs={}):
        if subset != 'train':
            raise ValueError('"accuracy" is only available for "train"')
        value = logs['acc']
        minimize = False
        return value, minimize

    def get_value(self, epoch, name, subset, logs={}):
        get_value = getattr(self, 'get_' + name)
        value, minimize = get_value(epoch, subset, logs=logs)
        return value, minimize

    def on_epoch_end(self, epoch, logs={}):
        """Save weights (and various curves) after each epoch"""

        # keep track of when the epoch ended
        now = datetime.datetime.now().isoformat()

        # save model after this epoch

        if self.backend == 'keras':
            import keras.models

            weights_h5 = self.WEIGHTS_H5.format(log_dir=self.log_dir,
                                                epoch=epoch)
            # overwrite only in case of a restart
            # include optimizer only every 10 epochs
            keras.models.save_model(self.model, weights_h5,
                                    overwrite=self.restart,
                                    include_optimizer=(epoch % 10 == 0))

        elif self.backend == 'pytorch':
            import torch

            weights_pt = self.WEIGHTS_PT.format(
                log_dir=self.log_dir, epoch=epoch)
            torch.save(self.model.state_dict(), weights_pt)

            optimizer_pt = self.OPTIMIZER_PT.format(
                log_dir=self.log_dir, epoch=epoch)
            torch.save(self.optimizer.state_dict(), optimizer_pt)

        for subset, name in self.log:

            value, minimize = self.get_value(epoch, name, subset, logs=logs)

            # keep track of value after last epoch
            self.values[subset][name].append(value)
            values = self.values[subset][name]

            # write value to file
            log_txt = self.LOG_TXT.format(
                log_dir=self.log_dir, name=name, subset=subset)
            TXT_TEMPLATE = '{epoch:d} ' + now + ' {value:.3f}\n'

            mode = 'a' if epoch > 0 else 'w'
            try:
                with open(log_txt, mode) as fp:
                    fp.write(TXT_TEMPLATE.format(epoch=epoch, value=value))
                    fp.flush()
            except Exception as e:
                pass

            # keep track of 'best value' model
            if minimize:
                best_epoch = np.argmin(values)
                best_value = np.min(values)
            else:
                best_epoch = np.argmax(values)
                best_value = np.max(values)

            # plot values to file and mark best value so far
            plt.plot(values, 'b')
            plt.plot([best_epoch], [best_value], 'bo')
            plt.plot([0, epoch], [best_value, best_value], 'k--')
            plt.grid(True)

            plt.xlabel('epoch')

            YLABEL = '{name} on {subset}'
            ylabel = YLABEL.format(name=name, subset=subset)
            plt.ylabel(ylabel)

            TITLE = '{name} = {best_value:.3f} on {subset} @ epoch #{best_epoch:d}'
            title = TITLE.format(name=name, subset=subset, best_value=best_value, best_epoch=best_epoch)
            plt.title(title)

            plt.tight_layout()

            # save plot as PNG
            log_png = self.LOG_PNG.format(
                log_dir=self.log_dir, name=name, subset=subset)
            try:
                plt.savefig(log_png, dpi=75)
            except Exception as e:
                pass

            # save plot as EPS
            log_eps = self.LOG_EPS.format(
                log_dir=self.log_dir, name=name, subset=subset)
            try:
                plt.savefig(log_eps)
            except Exception as e:
                pass

            plt.close()

class BaseLogger(Callback):
    """Callback that accumulates epoch averages of metrics.
    This callback is automatically applied to every Keras model.
    """

    def on_epoch_begin(self, epoch, logs=None):
        self.seen = 0
        self.totals = {}

    def on_batch_end(self, batch, logs=None):
        logs = logs or {}
        batch_size = logs.get('size', 0)
        self.seen += batch_size

        for k, v in logs.items():
            if k in self.totals:
                try:
                    self.totals[k] += v * batch_size
                except ValueError as e:
                    pass
            else:
                self.totals[k] = v * batch_size

    def on_epoch_end(self, epoch, logs=None):
        if logs is not None:
            for k in self.params['metrics']:
                if k in self.totals:
                    # Make value available to next callbacks.
                    logs[k] = self.totals[k] / self.seen

class Debugging(Callback):

    GRADIENT_PNG = '{log_dir}/gradient.png'

    def __init__(self):
        super(Debugging, self).__init__()
        self.min_gradient = []
        self.max_gradient = []
        self.avg_gradient = []
        self.pct_gradient = []

    def on_batch_end(self, batch, logs=None):

        min_gradient = np.min(np.abs(logs['gradient']))
        max_gradient = np.max(np.abs(logs['gradient']))
        avg_gradient = np.average(np.abs(logs['gradient']))

        self.min_gradient.append(min_gradient)
        self.max_gradient.append(max_gradient)
        self.avg_gradient.append(avg_gradient)

        # plot values to file and mark best value so far
        plt.plot(self.min_gradient, 'b', label='min')
        plt.plot(self.max_gradient, 'r', label='max')
        plt.plot(self.avg_gradient, 'g', label='avg')
        plt.legend()
        plt.title('Gradients')
        plt.tight_layout()

        # save plot as PNG
        gradient_png = self.GRADIENT_PNG.format(log_dir=logs['log_dir'])
        try:
            plt.savefig(gradient_png, dpi=75)
        except Exception as e:
            pass

        plt.close()
<commit_msg>improve: increase precision in LoggingCallback output file<commit_after>#!/usr/bin/env python
# encoding: utf-8

# The MIT License (MIT)

# Copyright (c) 2016-2017 CNRS

# Permission is hereby granted, free of charge, to any person obtaining a copy
# of this software and associated documentation files (the "Software"), to deal
# in the Software without restriction, including without limitation the rights
# to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the Software is
# furnished to do so, subject to the following conditions:

# The above copyright notice and this permission notice shall be included in
# all copies or substantial portions of the Software.

# THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
# IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
# FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
# AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
# LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
# OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE
# SOFTWARE.

# AUTHORS
# Hervé BREDIN - http://herve.niderb.fr

import matplotlib
import warnings
with warnings.catch_warnings():
    warnings.simplefilter("ignore")
    matplotlib.use('Agg')
import matplotlib.pyplot as plt

import os
import os.path
import datetime
import numpy as np
from pyannote.audio.util import mkdir_p
from keras.callbacks import Callback

class LoggingCallback(Callback):
    """Logging callback

    Parameters
    ----------
    log_dir : str
    log : list of tuples, optional
        Defaults to [('train', 'loss')]
    extract_embedding : func
        Function that takes Keras model as input and returns the actual model.
        This is useful for embedding that are not directly optimized by Keras.
        Defaults to identity function.
    restart : boolean, optional
        Indicates that this training is a restart, not a cold start (default).
    """

    ARCHITECTURE_YML = '{log_dir}/architecture.yml'
    WEIGHTS_DIR = '{log_dir}/weights'
    WEIGHTS_H5 = '{log_dir}/weights/{epoch:04d}.h5'
    WEIGHTS_PT = '{log_dir}/weights/{epoch:04d}.pt'
    OPTIMIZER_PT = '{log_dir}/weights/{epoch:04d}.optimizer.pt'
    LOG_TXT = '{log_dir}/{name}.{subset}.txt'
    LOG_PNG = '{log_dir}/{name}.{subset}.png'
    LOG_EPS = '{log_dir}/{name}.{subset}.eps'

    def __init__(self, log_dir, log=None, extract_embedding=None,
                 restart=False, backend='keras'):
        super(LoggingCallback, self).__init__()

        # make sure path is absolute
        self.log_dir = os.path.realpath(log_dir)

        if extract_embedding is None:
            extract_embedding = lambda model: model
        self.extract_embedding = extract_embedding

        # create log_dir directory
        mkdir_p(self.log_dir)

        # this will fail if the directory already exists
        # and this is OK  because 'weights' directory
        # usually contains the output of very long computations
        # and you do not want to erase them by mistake :/
        self.restart = restart
        if not self.restart:
            weights_dir = self.WEIGHTS_DIR.format(log_dir=self.log_dir)
            os.makedirs(weights_dir)

        if log is None:
            log = [('train', 'loss')]
        self.log = log

        self.values = {}
        for subset, name in self.log:
            if subset not in self.values:
                self.values[subset] = {}
            if name not in self.values[subset]:
                self.values[subset][name] = []

        self.backend = backend

    def get_loss(self, epoch, subset, logs={}):
        if subset != 'train':
            raise ValueError('"loss" is only available for "train"')
        value = logs['loss']
        minimize = True
        return value, minimize

    def get_accuracy(self, epoch, subset, logs={}):
        if subset != 'train':
            raise ValueError('"accuracy" is only available for "train"')
        value = logs['acc']
        minimize = False
        return value, minimize

    def get_value(self, epoch, name, subset, logs={}):
        get_value = getattr(self, 'get_' + name)
        value, minimize = get_value(epoch, subset, logs=logs)
        return value, minimize

    def on_epoch_end(self, epoch, logs={}):
        """Save weights (and various curves) after each epoch"""

        # keep track of when the epoch ended
        now = datetime.datetime.now().isoformat()

        # save model after this epoch

        if self.backend == 'keras':
            import keras.models

            weights_h5 = self.WEIGHTS_H5.format(log_dir=self.log_dir,
                                                epoch=epoch)
            # overwrite only in case of a restart
            # include optimizer only every 10 epochs
            keras.models.save_model(self.model, weights_h5,
                                    overwrite=self.restart,
                                    include_optimizer=(epoch % 10 == 0))

        elif self.backend == 'pytorch':
            import torch

            weights_pt = self.WEIGHTS_PT.format(
                log_dir=self.log_dir, epoch=epoch)
            torch.save(self.model.state_dict(), weights_pt)

            optimizer_pt = self.OPTIMIZER_PT.format(
                log_dir=self.log_dir, epoch=epoch)
            torch.save(self.optimizer.state_dict(), optimizer_pt)

        for subset, name in self.log:

            value, minimize = self.get_value(epoch, name, subset, logs=logs)

            # keep track of value after last epoch
            self.values[subset][name].append(value)
            values = self.values[subset][name]

            # write value to file
            log_txt = self.LOG_TXT.format(
                log_dir=self.log_dir, name=name, subset=subset)
            TXT_TEMPLATE = '{epoch:d} ' + now + ' {value:.8f}\n'

            mode = 'a' if epoch > 0 else 'w'
            try:
                with open(log_txt, mode) as fp:
                    fp.write(TXT_TEMPLATE.format(epoch=epoch, value=value))
                    fp.flush()
            except Exception as e:
                pass

            # keep track of 'best value' model
            if minimize:
                best_epoch = np.argmin(values)
                best_value = np.min(values)
            else:
                best_epoch = np.argmax(values)
                best_value = np.max(values)

            # plot values to file and mark best value so far
            plt.plot(values, 'b')
            plt.plot([best_epoch], [best_value], 'bo')
            plt.plot([0, epoch], [best_value, best_value], 'k--')
            plt.grid(True)

            plt.xlabel('epoch')

            YLABEL = '{name} on {subset}'
            ylabel = YLABEL.format(name=name, subset=subset)
            plt.ylabel(ylabel)

            TITLE = '{name} = {best_value:.8f} on {subset} @ epoch #{best_epoch:d}'
            title = TITLE.format(name=name, subset=subset, best_value=best_value, best_epoch=best_epoch)
            plt.title(title)

            plt.tight_layout()

            # save plot as PNG
            log_png = self.LOG_PNG.format(
                log_dir=self.log_dir, name=name, subset=subset)
            try:
                plt.savefig(log_png, dpi=75)
            except Exception as e:
                pass

            # save plot as EPS
            log_eps = self.LOG_EPS.format(
                log_dir=self.log_dir, name=name, subset=subset)
            try:
                plt.savefig(log_eps)
            except Exception as e:
                pass

            plt.close()

class BaseLogger(Callback):
    """Callback that accumulates epoch averages of metrics.
    This callback is automatically applied to every Keras model.
    """

    def on_epoch_begin(self, epoch, logs=None):
        self.seen = 0
        self.totals = {}

    def on_batch_end(self, batch, logs=None):
        logs = logs or {}
        batch_size = logs.get('size', 0)
        self.seen += batch_size

        for k, v in logs.items():
            if k in self.totals:
                try:
                    self.totals[k] += v * batch_size
                except ValueError as e:
                    pass
            else:
                self.totals[k] = v * batch_size

    def on_epoch_end(self, epoch, logs=None):
        if logs is not None:
            for k in self.params['metrics']:
                if k in self.totals:
                    # Make value available to next callbacks.
                    logs[k] = self.totals[k] / self.seen

class Debugging(Callback):

    GRADIENT_PNG = '{log_dir}/gradient.png'

    def __init__(self):
        super(Debugging, self).__init__()
        self.min_gradient = []
        self.max_gradient = []
        self.avg_gradient = []
        self.pct_gradient = []

    def on_batch_end(self, batch, logs=None):

        min_gradient = np.min(np.abs(logs['gradient']))
        max_gradient = np.max(np.abs(logs['gradient']))
        avg_gradient = np.average(np.abs(logs['gradient']))

        self.min_gradient.append(min_gradient)
        self.max_gradient.append(max_gradient)
        self.avg_gradient.append(avg_gradient)

        # plot values to file and mark best value so far
        plt.plot(self.min_gradient, 'b', label='min')
        plt.plot(self.max_gradient, 'r', label='max')
        plt.plot(self.avg_gradient, 'g', label='avg')
        plt.legend()
        plt.title('Gradients')
        plt.tight_layout()

        # save plot as PNG
        gradient_png = self.GRADIENT_PNG.format(log_dir=logs['log_dir'])
        try:
            plt.savefig(gradient_png, dpi=75)
        except Exception as e:
            pass

        plt.close()
<|endoftext|>


<commit_before>import os
import nipype.pipeline.engine as pe
from nipype.interfaces.utility import IdentityInterface, Function
from nipype.interfaces.ants.registration import Registration
from nipype.interfaces.ants import AverageImages, ApplyTransforms
from nipype.interfaces.dcmstack import CopyMeta
from nipype.interfaces import fsl
from nipype.interfaces.dcmstack import DcmStack, MergeNifti, CopyMeta
from nipype.interfaces.utility import Select, IdentityInterface, Function
from ..helpers.xnat_helper import PROJECT
from ..interfaces import XNATDownload, XNATUpload, MriVolCluster
from ..helpers import xnat_helper, file_helper
from ..helpers.ast_config import read_config

import logging
logger = logging.getLogger(__name__)

DEF_CONFIG_FILE = os.path.join(os.path.dirname(__file__), '..', '..', 'conf', 'registration.cfg')

REG_PREFIX = 'reg'
"""The XNAT registration reconstruction name prefix."""

def run(*session_specs, **opts):
    """
    Registers the scan NiFTI images for the given sessions as follows:

    - Download the NiFTI scans from XNAT

    - Make a mask to subtract extraneous tissue

    - Mask each input scan

    - Make a template by averaging the masked images

    - Create an affine and non-rigid transform for each image

    - Reslice the masked image with the transforms

    - Upload the mask and the resliced images

    The NiFTI scan images for each session are downloaded from XNAT into the
    ``scans`` subdirectory of the ``base_dir`` specified in the options
    (default is the current directory).

    The average is taken on the middle half of the NiFTI scan images.
    These images are averaged into a fixed reference template image.

    The options include the Pyxnat Workflow initialization options, as well as
    the following options:

    - ``config``: an optional configuration file

    The configuration file can contain the following sections:

    - ``FSLMriVolCluster``: the FSL ``MriVolCluster`` interface options

    - ``ANTSAverage``: the ANTS ``Average`` interface options

    - ``ANTSRegistration``: the ANTS ``Registration`` interface options

    - ``ANTSApplyTransforms``: the ANTS ``ApplyTransforms`` interface options

    The default registration applies an affine followed by a symmetric normalization
    transform.

    :param session_specs: the XNAT (subject, session) name tuples to register
    :param opts: the workflow options
    :return: the resliced XNAT (subject, session, reconstruction) designator tuples
    """

    # The work directory.
    work = opts.get('base_dir') or os.getcwd()
    # The scan image downloaad location.
    dest = os.path.join(work, 'scans')

    # Run the workflow on each session.
    recon_specs = [_register(sbj, sess, dest, **opts) for sbj, sess in session_specs]

    return recon_specs

def _register(subject, session, dest, **opts):
    """
    Builds and runs the registration workflow.

    :param subject: the XNAT subject label
    :param session: the XNAT session label
    :param dest: the scan download directory
    :param opts: the workflow options
    :return: the warpd XNAT (subject, session, reconstruction) designator tuple
    """
    # Download the scan images.
    tgt = os.path.join(dest, subject, session)
    images = _download_scans(subject, session, tgt)

    # Make a unique registration reconstruction name. This permits more than one
    # registration to be stored for each input image without a name conflict.
    recon = _generate_name(REG_PREFIX)

    # Make the workflow
    workflow = create_workflow(subject, session, recon, images, **opts)
    # Execute the workflow
    workflow.run()

    # Return the recon specification.
    return (subject, session, recon)

def _download_scans(subject, session, dest):
    """
    Download the NIFTI scan files for the given session.

    :param subject: the XNAT subject label
    :param session: the XNAT session label
    :param dest: the destination directory path
    :return: the download file paths
    """

    with xnat_helper.connection() as xnat:
        return xnat.download(PROJECT, subject, session, dest=dest, container_type='scan', format='NIFTI')

def create_workflow(subject, session, recon, images, **opts):
    """
    Creates the nipype workflow for the given session images.

    :param subject: the XNAT subject label
    :param session: the XNAT session label
    :param recon: the XNAT registration reconstruction label
    :param images: the input session scan NiFTI stacks
    :param opts: the workflow options
    :return: the registration Workflow object
    """
    msg = 'Creating the %s %s registration workflow' % (subject, session)
    if opts:
        msg = msg + ' with options %s' % opts
    logger.debug("%s...", msg)

    # The configuration.
    cfg_file = opts.pop('config', None)
    if cfg_file:
        cfg = read_config(DEF_CONFIG_FILE, cfg_file)
    else:
        cfg = read_config(DEF_CONFIG_FILE)
    cfg_opts = dict(cfg)

    # The mask step options.
    mask_opts = cfg_opts.get('FSLMriVolCluster', {})
    # The average step options.
    avg_opts = cfg_opts.get('ANTSAverage', {})
    # The registration step options.
    reg_opts = cfg_opts.get('ANTSRegistration', {})
    # The reslice step options.
    reslice_opts = cfg_opts.get('ANTSApplyTransforms', {})

    workflow = pe.Workflow(name='register', **opts)

    # The workflow input image iterator.
    input_spec = pe.Node(IdentityInterface(fields=['image']), name='input_spec')
    input_spec.inputs.subject = subject
    input_spec.inputs.session = session
    input_spec.iterables = dict(image=images).items()

    # Merge the DCE data to 4D.
    dce_merge = pe.Node(MergeNifti(), name='dce_merge')
    dce_merge.inputs.out_format = 'dce_series'
    dce_merge.inputs.in_files = images

    # Upload the 4D image to XNAT.
    upload_4d = pe.Node(XNATUpload(project=PROJECT, reconstruction='4d', format='NIFTI'),
        name='upload_4d')
    upload_4d.inputs.subject = subject
    upload_4d.inputs.session = session
    workflow.connect(dce_merge, 'out_file', upload_4d, 'in_files')

    # Get a mean image from the DCE data.
    dce_mean = pe.Node(fsl.MeanImage(), name='dce_mean')
    workflow.connect(dce_merge, 'out_file', dce_mean, 'in_file')

    # Find the center of gravity from the mean image.
    find_cog = pe.Node(fsl.ImageStats(), name='find_cog')
    find_cog.inputs.op_string = '-C'
    workflow.connect(dce_mean, 'out_file', find_cog, 'in_file')

    # Zero everything posterior to the center of gravity on mean image.
    crop_back = pe.Node(fsl.ImageMaths(), name='crop_back')
    workflow.connect(dce_mean, 'out_file', crop_back, 'in_file')
    workflow.connect(find_cog, ('out_stat', _gen_crop_op_string), crop_back, 'op_string')

    # Find large clusters of empty space on the cropped image.
    cluster_mask = pe.Node(MriVolCluster(**mask_opts), name='cluster_mask')
    workflow.connect(crop_back, 'out_file', cluster_mask, 'in_file')

    # Convert the cluster labels to a binary mask.
    binarize = pe.Node(fsl.BinaryMaths(), name='binarize')
    binarize.inputs.operation = 'min'
    binarize.inputs.operand_value = 1
    workflow.connect(cluster_mask, 'out_cluster_file', binarize, 'in_file')

    # Invert the binary mask.
    inv_mask = pe.Node(fsl.maths.MathsCommand(), name='inv_mask')
    inv_mask.inputs.args = '-sub 1 -mul -1'
    inv_mask.inputs.out_file = "%s_%s.nii.gz" % (session.lower(), 'mask')
    workflow.connect(binarize, 'out_file', inv_mask, 'in_file')

    # Upload the mask to XNAT.
    upload_mask = pe.Node(XNATUpload(project=PROJECT, reconstruction='mask', format='NIFTI'),
        name='upload_mask')
    upload_mask.inputs.subject = subject
    upload_mask.inputs.session = session
    workflow.connect(inv_mask, 'out_file', upload_mask, 'in_files')

    # Make the ANTS template.
    average = pe.Node(AverageImages(**avg_opts), name='average')
    # Use the middle half of the images.
    offset = len(images) / 4
    average.inputs.images = sorted(images)[offset:len(images)-offset]

    # Register the images to create the warp and affine transformations.
    register = pe.Node(Registration(**reg_opts), name='register')
    workflow.connect(input_spec, 'image', register, 'moving_image')
    workflow.connect(average, 'output_average_image', register, 'fixed_image')
    workflow.connect(inv_mask, 'out_file', register, 'fixed_image_mask')
    workflow.connect(inv_mask, 'out_file', register, 'moving_image_mask')

    # Apply the transforms to the input image.
    reslice = pe.Node(ApplyTransforms(**reslice_opts), name='reslice')
    workflow.connect(input_spec, 'image', reslice, 'input_image')
    workflow.connect(input_spec, ('image', _gen_reslice_filename), reslice, 'output_image')
    workflow.connect(average, 'output_average_image', reslice, 'reference_image')
    workflow.connect(register, 'forward_transforms', reslice, 'transforms')

    # Copy the DICOM meta-data.
    copy_meta = pe.Node(CopyMeta(), name='copy_meta')
    workflow.connect(input_spec, 'image', copy_meta, 'src_file')
    workflow.connect(reslice, 'output_image', copy_meta, 'dest_file')

    # Upload the resliced image to XNAT.
    upload_reg = pe.Node(XNATUpload(project=PROJECT, format='NIFTI'),
        name='upload_reg')
    upload_reg.inputs.subject = subject
    upload_reg.inputs.session = session
    upload_reg.inputs.reconstruction = recon
    workflow.connect(copy_meta, 'dest_file', upload_reg, 'in_files')

    return workflow

def _run_workflow(workflow):
    """
    Executes the given workflow.

    :param workflow: the workflow to run
    """
    # If debug is set, then diagram the workflow graph.
    if logger.level <= logging.DEBUG:
        fname = "%s.dot" % workflow.name
        if workflow.base_dir:
            grf = os.path.join(workflow.base_dir, fname)
        else:
            grf = fname
        workflow.write_graph(dotfilename=grf)
        logger.debug("The %s workflow graph is depicted at %s.png." % (workflow.name, grf))

    # Run the workflow.
    workflow.run()

def _generate_name(prefix):
    """
    :param: the name prefix
    :return: a unique name which starts with the given prefix
    """
    # The name suffix.
    suffix = file_helper.generate_file_name()

    return "%s_%s" % (prefix, suffix)

def _gen_reslice_filename(in_file):
    """
    :param in_file: the input scan image filename
    :return: the registered image filename
    """
    import re

    base, ext = re.match('(.*?)(\.[.\w]+)*$', in_file).groups()
    fname = "%s_%s" % (base, 'reg')
    if ext:
        return fname + ext
    else:
        return fname

def _gen_crop_op_string(cog):
    """
    :param cog: the center of gravity
    :return: the crop -roi option
    """
    return "-roi 0 -1 %d -1 0 -1 0 -1" % cog[1]

def _crop_posterior(image, cog):
    from nipype.interfaces import fsl

    crop_back = fsl.ImageMaths()
    crop_back.inputs.op_string = '-roi 0 -1 %d -1 0 -1 0 -1' % cog[1]
    crop_back.inputs.in_file = image
    return crop_back.run().outputs.out_file
<commit_msg>Inline _generate_name.<commit_after>import os
import nipype.pipeline.engine as pe
from nipype.interfaces.utility import IdentityInterface, Function
from nipype.interfaces.ants.registration import Registration
from nipype.interfaces.ants import AverageImages, ApplyTransforms
from nipype.interfaces.dcmstack import CopyMeta
from nipype.interfaces import fsl
from nipype.interfaces.dcmstack import DcmStack, MergeNifti, CopyMeta
from nipype.interfaces.utility import Select, IdentityInterface, Function
from ..helpers.xnat_helper import PROJECT
from ..interfaces import XNATDownload, XNATUpload, MriVolCluster
from ..helpers import xnat_helper, file_helper
from ..helpers.ast_config import read_config

import logging
logger = logging.getLogger(__name__)

DEF_CONFIG_FILE = os.path.join(os.path.dirname(__file__), '..', '..', 'conf', 'registration.cfg')

REG_PREFIX = 'reg'
"""The XNAT registration reconstruction name prefix."""

def run(*session_specs, **opts):
    """
    Registers the scan NiFTI images for the given sessions as follows:

    - Download the NiFTI scans from XNAT

    - Make a mask to subtract extraneous tissue

    - Mask each input scan

    - Make a template by averaging the masked images

    - Create an affine and non-rigid transform for each image

    - Reslice the masked image with the transforms

    - Upload the mask and the resliced images

    The NiFTI scan images for each session are downloaded from XNAT into the
    ``scans`` subdirectory of the ``base_dir`` specified in the options
    (default is the current directory).

    The average is taken on the middle half of the NiFTI scan images.
    These images are averaged into a fixed reference template image.

    The options include the Pyxnat Workflow initialization options, as well as
    the following options:

    - ``config``: an optional configuration file

    The configuration file can contain the following sections:

    - ``FSLMriVolCluster``: the FSL ``MriVolCluster`` interface options

    - ``ANTSAverage``: the ANTS ``Average`` interface options

    - ``ANTSRegistration``: the ANTS ``Registration`` interface options

    - ``ANTSApplyTransforms``: the ANTS ``ApplyTransforms`` interface options

    The default registration applies an affine followed by a symmetric normalization
    transform.

    :param session_specs: the XNAT (subject, session) name tuples to register
    :param opts: the workflow options
    :return: the resliced XNAT (subject, session, reconstruction) designator tuples
    """

    # The work directory.
    work = opts.get('base_dir') or os.getcwd()
    # The scan image downloaad location.
    dest = os.path.join(work, 'scans')

    # Run the workflow on each session.
    recon_specs = [_register(sbj, sess, dest, **opts) for sbj, sess in session_specs]

    return recon_specs

def _register(subject, session, dest, **opts):
    """
    Builds and runs the registration workflow.

    :param subject: the XNAT subject label
    :param session: the XNAT session label
    :param dest: the scan download directory
    :param opts: the workflow options
    :return: the warpd XNAT (subject, session, reconstruction) designator tuple
    """
    # Download the scan images.
    tgt = os.path.join(dest, subject, session)
    images = _download_scans(subject, session, tgt)

    # Make a unique registration reconstruction name. This permits more than one
    # registration to be stored for each input image without a name conflict.
    recon = "%s_%s" % (REG_PREFIX, file_helper.generate_file_name())

    # Make the workflow
    workflow = create_workflow(subject, session, recon, images, **opts)
    # Execute the workflow
    workflow.run()

    # Return the recon specification.
    return (subject, session, recon)

def _download_scans(subject, session, dest):
    """
    Download the NIFTI scan files for the given session.

    :param subject: the XNAT subject label
    :param session: the XNAT session label
    :param dest: the destination directory path
    :return: the download file paths
    """

    with xnat_helper.connection() as xnat:
        return xnat.download(PROJECT, subject, session, dest=dest, container_type='scan', format='NIFTI')

def create_workflow(subject, session, recon, images, **opts):
    """
    Creates the nipype workflow for the given session images.

    :param subject: the XNAT subject label
    :param session: the XNAT session label
    :param recon: the XNAT registration reconstruction label
    :param images: the input session scan NiFTI stacks
    :param opts: the workflow options
    :return: the registration Workflow object
    """
    msg = 'Creating the %s %s registration workflow' % (subject, session)
    if opts:
        msg = msg + ' with options %s' % opts
    logger.debug("%s...", msg)

    # The configuration.
    cfg_file = opts.pop('config', None)
    if cfg_file:
        cfg = read_config(DEF_CONFIG_FILE, cfg_file)
    else:
        cfg = read_config(DEF_CONFIG_FILE)
    cfg_opts = dict(cfg)

    # The mask step options.
    mask_opts = cfg_opts.get('FSLMriVolCluster', {})
    # The average step options.
    avg_opts = cfg_opts.get('ANTSAverage', {})
    # The registration step options.
    reg_opts = cfg_opts.get('ANTSRegistration', {})
    # The reslice step options.
    reslice_opts = cfg_opts.get('ANTSApplyTransforms', {})

    workflow = pe.Workflow(name='register', **opts)

    # The workflow input image iterator.
    input_spec = pe.Node(IdentityInterface(fields=['image']), name='input_spec')
    input_spec.inputs.subject = subject
    input_spec.inputs.session = session
    input_spec.iterables = dict(image=images).items()

    # Merge the DCE data to 4D.
    dce_merge = pe.Node(MergeNifti(), name='dce_merge')
    dce_merge.inputs.out_format = 'dce_series'
    dce_merge.inputs.in_files = images

    # Upload the 4D image to XNAT.
    upload_4d = pe.Node(XNATUpload(project=PROJECT, reconstruction='4d', format='NIFTI'),
        name='upload_4d')
    upload_4d.inputs.subject = subject
    upload_4d.inputs.session = session
    workflow.connect(dce_merge, 'out_file', upload_4d, 'in_files')

    # Get a mean image from the DCE data.
    dce_mean = pe.Node(fsl.MeanImage(), name='dce_mean')
    workflow.connect(dce_merge, 'out_file', dce_mean, 'in_file')

    # Find the center of gravity from the mean image.
    find_cog = pe.Node(fsl.ImageStats(), name='find_cog')
    find_cog.inputs.op_string = '-C'
    workflow.connect(dce_mean, 'out_file', find_cog, 'in_file')

    # Zero everything posterior to the center of gravity on mean image.
    crop_back = pe.Node(fsl.ImageMaths(), name='crop_back')
    workflow.connect(dce_mean, 'out_file', crop_back, 'in_file')
    workflow.connect(find_cog, ('out_stat', _gen_crop_op_string), crop_back, 'op_string')

    # Find large clusters of empty space on the cropped image.
    cluster_mask = pe.Node(MriVolCluster(**mask_opts), name='cluster_mask')
    workflow.connect(crop_back, 'out_file', cluster_mask, 'in_file')

    # Convert the cluster labels to a binary mask.
    binarize = pe.Node(fsl.BinaryMaths(), name='binarize')
    binarize.inputs.operation = 'min'
    binarize.inputs.operand_value = 1
    workflow.connect(cluster_mask, 'out_cluster_file', binarize, 'in_file')

    # Invert the binary mask.
    inv_mask = pe.Node(fsl.maths.MathsCommand(), name='inv_mask')
    inv_mask.inputs.args = '-sub 1 -mul -1'
    inv_mask.inputs.out_file = "%s_%s.nii.gz" % (session.lower(), 'mask')
    workflow.connect(binarize, 'out_file', inv_mask, 'in_file')

    # Upload the mask to XNAT.
    upload_mask = pe.Node(XNATUpload(project=PROJECT, reconstruction='mask', format='NIFTI'),
        name='upload_mask')
    upload_mask.inputs.subject = subject
    upload_mask.inputs.session = session
    workflow.connect(inv_mask, 'out_file', upload_mask, 'in_files')

    # Make the ANTS template.
    average = pe.Node(AverageImages(**avg_opts), name='average')
    # Use the middle half of the images.
    offset = len(images) / 4
    average.inputs.images = sorted(images)[offset:len(images)-offset]

    # Register the images to create the warp and affine transformations.
    register = pe.Node(Registration(**reg_opts), name='register')
    workflow.connect(input_spec, 'image', register, 'moving_image')
    workflow.connect(average, 'output_average_image', register, 'fixed_image')
    workflow.connect(inv_mask, 'out_file', register, 'fixed_image_mask')
    workflow.connect(inv_mask, 'out_file', register, 'moving_image_mask')

    # Apply the transforms to the input image.
    reslice = pe.Node(ApplyTransforms(**reslice_opts), name='reslice')
    workflow.connect(input_spec, 'image', reslice, 'input_image')
    workflow.connect(input_spec, ('image', _gen_reslice_filename), reslice, 'output_image')
    workflow.connect(average, 'output_average_image', reslice, 'reference_image')
    workflow.connect(register, 'forward_transforms', reslice, 'transforms')

    # Copy the DICOM meta-data.
    copy_meta = pe.Node(CopyMeta(), name='copy_meta')
    workflow.connect(input_spec, 'image', copy_meta, 'src_file')
    workflow.connect(reslice, 'output_image', copy_meta, 'dest_file')

    # Upload the resliced image to XNAT.
    upload_reg = pe.Node(XNATUpload(project=PROJECT, format='NIFTI'),
        name='upload_reg')
    upload_reg.inputs.subject = subject
    upload_reg.inputs.session = session
    upload_reg.inputs.reconstruction = recon
    workflow.connect(copy_meta, 'dest_file', upload_reg, 'in_files')

    return workflow

def _run_workflow(workflow):
    """
    Executes the given workflow.

    :param workflow: the workflow to run
    """
    # If debug is set, then diagram the workflow graph.
    if logger.level <= logging.DEBUG:
        fname = "%s.dot" % workflow.name
        if workflow.base_dir:
            grf = os.path.join(workflow.base_dir, fname)
        else:
            grf = fname
        workflow.write_graph(dotfilename=grf)
        logger.debug("The %s workflow graph is depicted at %s.png." % (workflow.name, grf))

    # Run the workflow.
    workflow.run()

def _gen_reslice_filename(in_file):
    """
    :param in_file: the input scan image filename
    :return: the registered image filename
    """
    import re

    base, ext = re.match('(.*?)(\.[.\w]+)*$', in_file).groups()
    fname = "%s_%s" % (base, 'reg')
    if ext:
        return fname + ext
    else:
        return fname

def _gen_crop_op_string(cog):
    """
    :param cog: the center of gravity
    :return: the crop -roi option
    """
    return "-roi 0 -1 %d -1 0 -1 0 -1" % cog[1]

def _crop_posterior(image, cog):
    from nipype.interfaces import fsl

    crop_back = fsl.ImageMaths()
    crop_back.inputs.op_string = '-roi 0 -1 %d -1 0 -1 0 -1' % cog[1]
    crop_back.inputs.in_file = image
    return crop_back.run().outputs.out_file
<|endoftext|>


<commit_before>from nose.tools import eq_, assert_not_equals
import os
import fmn.lib.models
import fmn.lib.tests

class TestBasics(fmn.lib.tests.Base):
    def test_setup_and_teardown(self):
        pass

    def test_user_get_or_create(self):
        user1 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user2 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user3 = fmn.lib.models.User.get_or_create(
            self.sess, openid="toshio.id.fedoraproject",
            openid_url="http://toshio.id.fedoraproject.org/",
        )
        eq_(user1.openid, user2.openid)
        eq_(user1, user2)
        assert_not_equals(user1, user3)

    def test_context_create(self):
        context1 = fmn.lib.models.Context.create(
            self.sess, name="irc", description="Internet Relay Chat",
            detail_name="irc nick", icon="user")
        context2 = fmn.lib.models.Context.get(self.sess, name="irc")
        context3 = fmn.lib.models.Context.create(
            self.sess, name="android", description="Google Cloud Messaging",
            detail_name="registration id", icon="phone")
        eq_(context1, context2)
        assert_not_equals(context1, context3)

    def test_user_all(self):
        user1 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user2 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user3 = fmn.lib.models.User.get_or_create(
            self.sess, openid="toshio.id.fedoraproject",
            openid_url="http://toshio.id.fedoraproject.org/",
        )
        eq_(len(fmn.lib.models.User.all(self.sess)), 2)

    def test_context_all(self):
        context1 = fmn.lib.models.Context.create(
            self.sess, name="irc", description="Internet Relay Chat",
            detail_name="irc nick", icon="user")
        context2 = fmn.lib.models.Context.create(
            self.sess, name="android", description="Google Cloud Messaging",
            detail_name="registration id", icon="phone")
        eq_(len(fmn.lib.models.Context.all(self.sess)), 2)

    def test_queued_message_enqueue(self):
        user1 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user2 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        context1 = fmn.lib.models.Context.create(
            self.sess, name="irc", description="Internet Relay Chat",
            detail_name="irc nick", icon="user")
        context2 = fmn.lib.models.Context.create(
            self.sess, name="android", description="Google Cloud Messaging",
            detail_name="registration id", icon="phone")

        raw_msg = {'testing': 'foobar'}

        obj1 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj2 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj3 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)

        assert obj1 != obj2 and obj2 != obj3 and obj1 != obj3
        eq_(fmn.lib.models.QueuedMessage.count_for(
            self.sess, user1, context1), 3)

    def test_queued_message_earliest(self):
        user1 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user2 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        context1 = fmn.lib.models.Context.create(
            self.sess, name="irc", description="Internet Relay Chat",
            detail_name="irc nick", icon="user")
        context2 = fmn.lib.models.Context.create(
            self.sess, name="android", description="Google Cloud Messaging",
            detail_name="registration id", icon="phone")

        raw_msg = {'testing': 'foobar'}

        obj1 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj2 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj3 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)

        assert obj1 != obj2 and obj2 != obj3 and obj1 != obj3
        eq_(fmn.lib.models.QueuedMessage.earliest_for(
            self.sess, user1, context1), obj1)
<commit_msg>Beef up those tests.<commit_after>from nose.tools import eq_, assert_not_equals
import os
import fmn.lib.models
import fmn.lib.tests

class TestBasics(fmn.lib.tests.Base):
    def test_setup_and_teardown(self):
        pass

    def test_user_get_or_create(self):
        user1 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user2 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user3 = fmn.lib.models.User.get_or_create(
            self.sess, openid="toshio.id.fedoraproject",
            openid_url="http://toshio.id.fedoraproject.org/",
        )
        eq_(user1.openid, user2.openid)
        eq_(user1, user2)
        assert_not_equals(user1, user3)

    def test_context_create(self):
        context1 = fmn.lib.models.Context.create(
            self.sess, name="irc", description="Internet Relay Chat",
            detail_name="irc nick", icon="user")
        context2 = fmn.lib.models.Context.get(self.sess, name="irc")
        context3 = fmn.lib.models.Context.create(
            self.sess, name="android", description="Google Cloud Messaging",
            detail_name="registration id", icon="phone")
        eq_(context1, context2)
        assert_not_equals(context1, context3)

    def test_user_all(self):
        user1 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user2 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user3 = fmn.lib.models.User.get_or_create(
            self.sess, openid="toshio.id.fedoraproject",
            openid_url="http://toshio.id.fedoraproject.org/",
        )
        eq_(len(fmn.lib.models.User.all(self.sess)), 2)

    def test_context_all(self):
        context1 = fmn.lib.models.Context.create(
            self.sess, name="irc", description="Internet Relay Chat",
            detail_name="irc nick", icon="user")
        context2 = fmn.lib.models.Context.create(
            self.sess, name="android", description="Google Cloud Messaging",
            detail_name="registration id", icon="phone")
        eq_(len(fmn.lib.models.Context.all(self.sess)), 2)

    def test_queued_message_enqueue(self):
        user1 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user2 = fmn.lib.models.User.get_or_create(
            self.sess, openid="toshio.id.fedoraproject",
            openid_url="http://toshio.id.fedoraproject.org/",
        )
        context1 = fmn.lib.models.Context.create(
            self.sess, name="irc", description="Internet Relay Chat",
            detail_name="irc nick", icon="user")
        context2 = fmn.lib.models.Context.create(
            self.sess, name="android", description="Google Cloud Messaging",
            detail_name="registration id", icon="phone")

        raw_msg = {'testing': 'foobar'}

        obj1 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj2 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj3 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj4 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user2, context1, raw_msg)
        obj5 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context2, raw_msg)

        assert obj1 != obj2 and obj2 != obj3 and obj1 != obj3
        eq_(fmn.lib.models.QueuedMessage.count_for(
            self.sess, user1, context1), 3)

    def test_queued_message_earliest(self):
        user1 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        user2 = fmn.lib.models.User.get_or_create(
            self.sess, openid="ralph.id.fedoraproject",
            openid_url="http://ralph.id.fedoraproject.org/",
        )
        context1 = fmn.lib.models.Context.create(
            self.sess, name="irc", description="Internet Relay Chat",
            detail_name="irc nick", icon="user")
        context2 = fmn.lib.models.Context.create(
            self.sess, name="android", description="Google Cloud Messaging",
            detail_name="registration id", icon="phone")

        raw_msg = {'testing': 'foobar'}

        obj1 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj2 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj3 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context1, raw_msg)
        obj4 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user2, context1, raw_msg)
        obj5 = fmn.lib.models.QueuedMessage.enqueue(
            self.sess, user1, context2, raw_msg)

        assert obj1 != obj2 and obj2 != obj3 and obj1 != obj3
        eq_(fmn.lib.models.QueuedMessage.earliest_for(
            self.sess, user1, context1), obj1)
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import os
from os.path import join
import re
import tempfile

import pybamview
from flask import (Blueprint, current_app, redirect, render_template,
                   request, Response, url_for)

from ..utils import message

NUC_TO_COLOR = {
    "A": "red",
    'a': "red",
    "C": "blue",
    "c": "blue",
    "G": "green",
    "g": "green",
    "T": "orange",
    "t": "orange",
    "N": "gray",
    "n": "gray",
    "-": "white",
    ".": "gray",
}

blueprint = Blueprint("browser", __name__, template_folder="templates",
                      static_folder="static")

@blueprint.route('/')
def listsamples(methods=['POST','GET']):
    # If a single file, jump right to the view for that file
    BAM = current_app.config['BAM']
    BAMDIR = current_app.config['BAMDIR']
    if BAM:
        if not os.path.exists(BAM) or not pybamview.CheckBam(BAM):
            return render_template("error.html", message="%s not a valid BAM file"%BAM)
        try:
            samplesToBam = pybamview.GetSamplesFromBamFiles([os.path.abspath(BAM)])
        except ValueError, e:
            return render_template("error.html", message="Problem parsing BAM file: %s"%e, title="PyBamView - %s"%BAM)

        if not os.path.exists(BAM+".bai"):
            return render_template("error.html", message="No index found for %s"%BAM)
        if len(samplesToBam.keys()) == 0:
            return render_template("error.html", message="No samples found in BAM file", title="PyBamView - %s"%BAM)
        argstring = "&".join(["samplebams=%s:%s"%(sample, os.path.abspath(BAM)) for sample in samplesToBam])
        return redirect(url_for(".display_bam")+"?"+argstring)
    # If given a directory to look for bams, determine which samples are present
    else:
        bamfiles = request.args.getlist("bamfiles")
        samples = request.args.getlist("samples")
        if len(bamfiles) > 0 and len(samples) > 0:
            return display_bam(samples)
        try:
            files = os.listdir(BAMDIR)
        except OSError: files = []
        bamfiles = [f for f in files if re.match(".*.bam$", f) is not None]
        bamfiles = [f for f in bamfiles if f+".bai" in files]
        try:
            samplesToBam = pybamview.GetSamplesFromBamFiles([os.path.join(os.path.abspath(BAMDIR), b) for b in bamfiles])
        except ValueError, e:
            return render_template("error.html", message="Problem parsing BAM file: %s"%e, title="PyBamView - %s"%BAMDIR)
        return render_template("index.html", samplesToBam=samplesToBam, title="PyBamView - %s"%BAMDIR)

@blueprint.route('/bamview', methods=['POST', 'GET'])
def display_bam():
    print('Display bam!!!!!!!!')
    samplebams = request.args.getlist("samplebams")
    zoomlevel = request.args.get("zoomlevel")
    print(samplebams)
    print(zoomlevel)
    if len(samplebams) == 0:
        samples_toinclude = list(set(request.args.getlist("samples")))
        bamfiles_toinclude = list(set(request.args.getlist("bamfiles")))
    else:
        samples_toinclude = []
        bamfiles_toinclude = []
        for s in samplebams:
            samples_toinclude.append(s.split(":")[0])
            for item in s.split(":")[1:]:
                bamfiles_toinclude.extend(item.split(","))
    region = request.args.get("region", pybamview.GetDefaultLocation(bamfiles_toinclude))
    return display_bam_region(list(set(bamfiles_toinclude)), samples_toinclude, region, zoomlevel)

def display_bam_region(bamfiles, samples, region, zoomlevel):
    bamdir = current_app.config["BAMDIR"]
    reffile = current_app.config["REFFILE"]
    settings = current_app.config["SETTINGS"]
    bamfile_to_bamview = current_app.config["BAMFILE_TO_BAMVIEW"]
    if region == "error":
        return render_template("error.html", message="No aligned reads found in the selected BAM files")
    for bam in bamfiles:
        if not os.path.exists(join(bamdir,bam)):
            message("bam file %s does not exist"%join(bamdir, bam), "warning")
    if ";".join(bamfiles) not in bamfile_to_bamview:
        bv = pybamview.BamView([join(bamdir, bam) for bam in bamfiles], reffile)
        bamfile_to_bamview[";".join(bamfiles)] = bv
    else:
        bv = bamfile_to_bamview[";".join(bamfiles)]
    try:
        chrom, pos = region.split(":")
        pos = int(pos)
    except:
        return render_template("error.html", message="Invalid region specified %s. Region must be of the form chrom:position."%region)
    bv.LoadAlignmentGrid(chrom, pos, _samples=samples, _settings=settings)
    positions = bv.GetPositions(pos)
    region = "%s:%s" % (chrom, pos)
    return render_template("bamview.html",
        title="PyBamView - %s" % region,
        BAM_FILES=bamfiles,
        REF_FILE=reffile,
        REGION=region,
        SAMPLES=bv.GetSamples(),
        SAMPLE_HASHES=bv.GetSampleHashes(),
        BUFFER=settings["NUMCHAR"],
        MAXZOOM=settings["MAXZOOM"],
        ZOOMLEVEL=zoomlevel,
        REFERENCE=bv.GetReferenceTrack(pos),
        ALIGNMENTS=bv.GetAlignmentTrack(pos),
        TARGETPOS=pos,
        POSITIONS=positions,
        NUC_TO_COLOR=NUC_TO_COLOR,
        CHROM=chrom,
        TARGET_LIST=current_app.config.get("TARGET_LIST", [])
    )

@blueprint.route('/snapshot', methods=['POST', 'GET'])
def snapshot():
    settings = current_app.config['SETTINGS']
    samples = request.form.getlist("samples")
    alignments_by_sample = {}
    for s in samples:
        alignments_by_sample[s] = request.form["alignment_%s"%s]
    region = request.form["region"]
    region.replace("%3A",":")
    startpos = request.form["startpos"]
    reference = request.form["reference"]
    zoom = request.form["zoomlevel"]
    zoomlevel = float(zoom)
    if zoomlevel < 0:
        zoomlevel = -1/zoomlevel
    chrom, start = region.split(":")
    region = "%s-%s"%(max(int(int(start)-25/zoomlevel),0), int(int(start)+50+25/zoomlevel))
    return render_template("snapshot.html", title="Pybamview - take snapshot", SAMPLES=samples, ZOOMLEVEL=zoom, \
                               CHROM=chrom, REGION=region, MINSTART=int(start)-settings["LOADCHAR"]/2, MAXEND=int(start)+settings["LOADCHAR"]/2,\
                               REFERENCE_TRACK=reference, ALIGN_BY_SAMPLE=alignments_by_sample, STARTPOS=startpos)

@blueprint.route('/export', methods=["POST"])
def export():
    svg_xml = request.form.get("data", "Invalid data")
    input_file = tempfile.NamedTemporaryFile(delete=False)
    input_file.write(svg_xml)
    input_file.close()
    filename = request.form.get("filename", "pybamview_snapshot.pdf")
    output_file_name = os.path.join("/tmp", filename)
    cmd = "rsvg-convert -o %s -f pdf %s"%(output_file_name, input_file.name)
    retcode = os.system(cmd)
    if retcode != 0:
        message("Error exporting to PDF. Is rsvg-convert installed?", "warning")
    if os.path.exists(output_file_name):
        pdf_data = open(output_file_name, "r").read()
    else: pdf_data = ""
    response = Response(pdf_data, mimetype="application/x-pdf")
    response.headers["Content-Disposition"] = "attachment; filename=%s"%filename
    return response
<commit_msg>Removed print statements<commit_after># -*- coding: utf-8 -*-
import os
from os.path import join
import re
import tempfile

import pybamview
from flask import (Blueprint, current_app, redirect, render_template,
                   request, Response, url_for)

from ..utils import message

NUC_TO_COLOR = {
    "A": "red",
    'a': "red",
    "C": "blue",
    "c": "blue",
    "G": "green",
    "g": "green",
    "T": "orange",
    "t": "orange",
    "N": "gray",
    "n": "gray",
    "-": "white",
    ".": "gray",
}

blueprint = Blueprint("browser", __name__, template_folder="templates",
                      static_folder="static")

@blueprint.route('/')
def listsamples(methods=['POST','GET']):
    # If a single file, jump right to the view for that file
    BAM = current_app.config['BAM']
    BAMDIR = current_app.config['BAMDIR']
    if BAM:
        if not os.path.exists(BAM) or not pybamview.CheckBam(BAM):
            return render_template("error.html", message="%s not a valid BAM file"%BAM)
        try:
            samplesToBam = pybamview.GetSamplesFromBamFiles([os.path.abspath(BAM)])
        except ValueError, e:
            return render_template("error.html", message="Problem parsing BAM file: %s"%e, title="PyBamView - %s"%BAM)

        if not os.path.exists(BAM+".bai"):
            return render_template("error.html", message="No index found for %s"%BAM)
        if len(samplesToBam.keys()) == 0:
            return render_template("error.html", message="No samples found in BAM file", title="PyBamView - %s"%BAM)
        argstring = "&".join(["samplebams=%s:%s"%(sample, os.path.abspath(BAM)) for sample in samplesToBam])
        return redirect(url_for(".display_bam")+"?"+argstring)
    # If given a directory to look for bams, determine which samples are present
    else:
        bamfiles = request.args.getlist("bamfiles")
        samples = request.args.getlist("samples")
        if len(bamfiles) > 0 and len(samples) > 0:
            return display_bam(samples)
        try:
            files = os.listdir(BAMDIR)
        except OSError: files = []
        bamfiles = [f for f in files if re.match(".*.bam$", f) is not None]
        bamfiles = [f for f in bamfiles if f+".bai" in files]
        try:
            samplesToBam = pybamview.GetSamplesFromBamFiles([os.path.join(os.path.abspath(BAMDIR), b) for b in bamfiles])
        except ValueError, e:
            return render_template("error.html", message="Problem parsing BAM file: %s"%e, title="PyBamView - %s"%BAMDIR)
        return render_template("index.html", samplesToBam=samplesToBam, title="PyBamView - %s"%BAMDIR)

@blueprint.route('/bamview', methods=['POST', 'GET'])
def display_bam():
    samplebams = request.args.getlist("samplebams")
    zoomlevel = request.args.get("zoomlevel")
    if len(samplebams) == 0:
        samples_toinclude = list(set(request.args.getlist("samples")))
        bamfiles_toinclude = list(set(request.args.getlist("bamfiles")))
    else:
        samples_toinclude = []
        bamfiles_toinclude = []
        for s in samplebams:
            samples_toinclude.append(s.split(":")[0])
            for item in s.split(":")[1:]:
                bamfiles_toinclude.extend(item.split(","))
    region = request.args.get("region", pybamview.GetDefaultLocation(bamfiles_toinclude))
    return display_bam_region(list(set(bamfiles_toinclude)), samples_toinclude, region, zoomlevel)

def display_bam_region(bamfiles, samples, region, zoomlevel):
    bamdir = current_app.config["BAMDIR"]
    reffile = current_app.config["REFFILE"]
    settings = current_app.config["SETTINGS"]
    bamfile_to_bamview = current_app.config["BAMFILE_TO_BAMVIEW"]
    if region == "error":
        return render_template("error.html", message="No aligned reads found in the selected BAM files")
    for bam in bamfiles:
        if not os.path.exists(join(bamdir,bam)):
            message("bam file %s does not exist"%join(bamdir, bam), "warning")
    if ";".join(bamfiles) not in bamfile_to_bamview:
        bv = pybamview.BamView([join(bamdir, bam) for bam in bamfiles], reffile)
        bamfile_to_bamview[";".join(bamfiles)] = bv
    else:
        bv = bamfile_to_bamview[";".join(bamfiles)]
    try:
        chrom, pos = region.split(":")
        pos = int(pos)
    except:
        return render_template("error.html", message="Invalid region specified %s. Region must be of the form chrom:position."%region)
    bv.LoadAlignmentGrid(chrom, pos, _samples=samples, _settings=settings)
    positions = bv.GetPositions(pos)
    region = "%s:%s" % (chrom, pos)
    return render_template("bamview.html",
        title="PyBamView - %s" % region,
        BAM_FILES=bamfiles,
        REF_FILE=reffile,
        REGION=region,
        SAMPLES=bv.GetSamples(),
        SAMPLE_HASHES=bv.GetSampleHashes(),
        BUFFER=settings["NUMCHAR"],
        MAXZOOM=settings["MAXZOOM"],
        ZOOMLEVEL=zoomlevel,
        REFERENCE=bv.GetReferenceTrack(pos),
        ALIGNMENTS=bv.GetAlignmentTrack(pos),
        TARGETPOS=pos,
        POSITIONS=positions,
        NUC_TO_COLOR=NUC_TO_COLOR,
        CHROM=chrom,
        TARGET_LIST=current_app.config.get("TARGET_LIST", [])
    )

@blueprint.route('/snapshot', methods=['POST', 'GET'])
def snapshot():
    settings = current_app.config['SETTINGS']
    samples = request.form.getlist("samples")
    alignments_by_sample = {}
    for s in samples:
        alignments_by_sample[s] = request.form["alignment_%s"%s]
    region = request.form["region"]
    region.replace("%3A",":")
    startpos = request.form["startpos"]
    reference = request.form["reference"]
    zoom = request.form["zoomlevel"]
    zoomlevel = float(zoom)
    if zoomlevel < 0:
        zoomlevel = -1/zoomlevel
    chrom, start = region.split(":")
    region = "%s-%s"%(max(int(int(start)-25/zoomlevel),0), int(int(start)+50+25/zoomlevel))
    return render_template("snapshot.html", title="Pybamview - take snapshot", SAMPLES=samples, ZOOMLEVEL=zoom, \
                               CHROM=chrom, REGION=region, MINSTART=int(start)-settings["LOADCHAR"]/2, MAXEND=int(start)+settings["LOADCHAR"]/2,\
                               REFERENCE_TRACK=reference, ALIGN_BY_SAMPLE=alignments_by_sample, STARTPOS=startpos)

@blueprint.route('/export', methods=["POST"])
def export():
    svg_xml = request.form.get("data", "Invalid data")
    input_file = tempfile.NamedTemporaryFile(delete=False)
    input_file.write(svg_xml)
    input_file.close()
    filename = request.form.get("filename", "pybamview_snapshot.pdf")
    output_file_name = os.path.join("/tmp", filename)
    cmd = "rsvg-convert -o %s -f pdf %s"%(output_file_name, input_file.name)
    retcode = os.system(cmd)
    if retcode != 0:
        message("Error exporting to PDF. Is rsvg-convert installed?", "warning")
    if os.path.exists(output_file_name):
        pdf_data = open(output_file_name, "r").read()
    else: pdf_data = ""
    response = Response(pdf_data, mimetype="application/x-pdf")
    response.headers["Content-Disposition"] = "attachment; filename=%s"%filename
    return response
<|endoftext|>