<commit_before>#!/usr/bin/env python
# DialsIndexer.py
#   Copyright (C) 2014 Diamond Light Source, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# An indexer compatible with XDS but using the DIALS methods.

import os
import sys
import math
import shutil

if not os.environ.has_key('XIA2_ROOT'):
  raise RuntimeError, 'XIA2_ROOT not defined'

if not os.environ['XIA2_ROOT'] in sys.path:
  sys.path.append(os.environ['XIA2_ROOT'])

# wrappers for programs that this needs: XDS

from Wrappers.XDS.XDSXycorr import XDSXycorr as _Xycorr
from Wrappers.XDS.XDSInit import XDSInit as _Init
from Wrappers.XDS.XDSColspot import XDSColspot as _Colspot
from Wrappers.XDS.XDSIdxref import XDSIdxref as _Idxref

# wrappers for programs that this needs: DIALS

from Wrappers.Dials.Import import Import as _Import
from Wrappers.Dials.Spotfinder import Spotfinder as _Spotfinder
from Wrappers.Dials.Index import Index as _Index
from Wrappers.Dials.RefineBravaisSettings import RefineBravaisSettings as \
     _RefineBravaisSettings
from Wrappers.Dials.ExportXDS import ExportXDS as _ExportXDS

from Wrappers.XIA.Diffdump import Diffdump

# interfaces that this must implement to be an indexer

from Schema.Interfaces.Indexer import Indexer
from Schema.Interfaces.FrameProcessor import FrameProcessor

# odds and sods that are needed

from lib.bits import auto_logfiler, nint
from Handlers.Streams import Chatter, Debug, Journal
from Handlers.Flags import Flags
from Handlers.Phil import PhilIndex
from Handlers.Files import FileHandler

class DialsIndexer(FrameProcessor,
                   Indexer):
  def __init__(self):

    # set up the inherited objects

    FrameProcessor.__init__(self)
    Indexer.__init__(self)

    self._working_directory = os.getcwd()

    self._background_images = None

    # place to store working data

    self._data_files = { }
    self._solutions = { }

    # FIXME this is a stupid low resolution limit to use...
    self._indxr_low_resolution = 40.0

    return

  # admin functions

  def set_working_directory(self, working_directory):
    self._working_directory = working_directory
    return

  def get_working_directory(self):
    return self._working_directory

  # factory functions

  def Xycorr(self):
    xycorr = _Xycorr()
    xycorr.set_working_directory(self.get_working_directory())

    xycorr.setup_from_image(self.get_image_name(
        self._indxr_images[0][0]))

    if self.get_distance():
      xycorr.set_distance(self.get_distance())

    if self.get_wavelength():
      xycorr.set_wavelength(self.get_wavelength())

    auto_logfiler(xycorr, 'XYCORR')

    return xycorr

  def Init(self):
    from Handlers.Phil import PhilIndex
    init = _Init(params=PhilIndex.params.xds.init)
    init.set_working_directory(self.get_working_directory())

    init.setup_from_image(self.get_image_name(
        self._indxr_images[0][0]))

    if self.get_distance():
      init.set_distance(self.get_distance())

    if self.get_wavelength():
      init.set_wavelength(self.get_wavelength())

    auto_logfiler(init, 'INIT')

    return init

  def Import(self):
    importer = _Import()
    importer.set_working_directory(self.get_working_directory())
    importer.setup_from_image(self.get_image_name(
      self._indxr_images[0][0]))
    auto_logfiler(importer)
    return importer

  def Spotfinder(self):
    spotfinder = _Spotfinder()
    spotfinder.set_working_directory(self.get_working_directory())
    auto_logfiler(spotfinder)
    return spotfinder

  def Index(self):
    index = _Index()
    index.set_working_directory(self.get_working_directory())
    auto_logfiler(index)
    return index

  def ExportXDS(self):
    export_xds = _ExportXDS()
    export_xds.set_working_directory(self.get_working_directory())
    auto_logfiler(export_xds)
    return export_xds

  def RefineBravaisSettings(self):
    rbs = _RefineBravaisSettings()
    rbs.set_working_directory(self.get_working_directory())
    auto_logfiler(rbs)
    return rbs

  ##########################################

  def _index_prepare(self):

    all_images = self.get_matching_images()
    first = min(all_images)
    last = max(all_images)

    self._indxr_images = [(first, last)]

    dd = Diffdump()
    dd.set_image(self.get_image_name(first))
    header = dd.readheader()
    last_background = int(round(5.0 / header['phi_width'])) - 1 + first

    # next start to process these - first xycorr
    # FIXME run these *afterwards* as then we have a refined detector geometry
    # so the parallax correction etc. should be slightly better.

    xycorr = self.Xycorr()
    xycorr.set_data_range(first, last)
    xycorr.set_background_range(first, last_background)
    xycorr.run()

    for file in ['X-CORRECTIONS.cbf',
                 'Y-CORRECTIONS.cbf']:
      self._data_files[file] = xycorr.get_output_data_file(file)

    # next start to process these - then init

    init = self.Init()

    for file in ['X-CORRECTIONS.cbf',
                 'Y-CORRECTIONS.cbf']:
      init.set_input_data_file(file, self._data_files[file])

    init.set_data_range(first, last)
    init.set_background_range(first, last_background)
    init.run()

    for file in ['BLANK.cbf',
                 'BKGINIT.cbf',
                 'GAIN.cbf']:
      self._data_files[file] = init.get_output_data_file(file)

    # at this stage, break out to run the DIALS code: this sets itself up

    importer = self.Import()
    importer.run()

    # FIXME this should really use the assigned spot finding regions
    spotfinder = self.Spotfinder()
    spotfinder.set_sweep_filename(importer.get_sweep_filename())
    spotfinder.set_scan_ranges([(first, last)])
    if PhilIndex.params.dials.spotfinder.phil_file is not None:
      spotfinder.set_phil_file(PhilIndex.params.dials.spotfinder.phil_file)
    spotfinder.run()

    self._spot_filename = spotfinder.get_spot_filename()
    self._sweep_filename = importer.get_sweep_filename()

    return

  def _index(self):
    # FIXME allow humans to set the indexing method from whatever list...
    # FIXME respect input unit cell / symmetry if set - or if decided from
    # previous indexing cycle
    indexer = self.Index()
    indexer.set_spot_filename(self._spot_filename)
    indexer.set_sweep_filename(self._sweep_filename)
    if PhilIndex.params.dials.index.phil_file is not None:
      indexer.set_phil_file(PhilIndex.params.dials.index.phil_file)

    if self._indxr_input_lattice:
      indexer.set_indexer_input_lattice(self._indxr_input_lattice)
      Debug.write('Set lattice: %s' % self._indxr_input_lattice)

    if self._indxr_input_cell:
      indexer.set_indexer_input_cell(self._indxr_input_cell)
      Debug.write('Set cell: %f %f %f %f %f %f' % \
                  self._indxr_input_cell)
      original_cell = self._indxr_input_cell

    if PhilIndex.params.dials.index.method is None:
      if self._indxr_input_cell is not None:
        method = 'real_space_grid_search'
      else:
        method = 'fft3d'
      Debug.write('Choosing indexing method: %s' % method)
    else:
      method = PhilIndex.params.dials.index.method

    indexer.run(method)

    # FIXME don't keep hold of an indexer: prevents pickling

    self._indexer = indexer

    # FIXME in here should respect the input unit cell and lattice if provided

    # FIXME from this (i) populate the helper table, (ii) export those files
    # which XDS will expect to find, (iii) try to avoid re-running the indexing
    # step if we eliminate a solution as we have all of the refined results
    # already available.

    rbs = self.RefineBravaisSettings()
    rbs.set_experiments_filename(indexer.get_experiments_filename())
    rbs.set_indexed_filename(indexer.get_indexed_filename())
    rbs.run()

    rmsd_p1 = rbs.get_bravais_summary()[1]['rmsd']

    for k in sorted(rbs.get_bravais_summary()):
      summary = rbs.get_bravais_summary()[k]

      # FIXME need to do this better - for the moment only accept lattices
      # where R.M.S. deviation is less than twice P1 R.M.S. deviation.

      if summary['rmsd'] > 2.0 * rmsd_p1:
        continue

      self._solutions[k] = {
        'number':k,
        'mosaic':0.0,
        'metric':summary['max_angular_difference'],
        'rmsd':summary['rmsd'],
        'nspots':summary['nspots'],
        'lattice':summary['bravais'],
        'cell':summary['unit_cell'],
        'experiments_file':summary['experiments_file']
        }

    self._solution = self.get_solution()

    for solution in self._solutions.keys():
      lattice = self._solutions[solution]['lattice']
      if self._indxr_other_lattice_cell.has_key(lattice):
        if self._indxr_other_lattice_cell[lattice]['metric'] < \
          self._solutions[solution]['metric']:
          continue

      self._indxr_other_lattice_cell[lattice] = {
        'metric':self._solutions[solution]['metric'],
        'cell':self._solutions[solution]['cell']}

    self._indxr_lattice = self._solution['lattice']
    self._indxr_cell = tuple(self._solution['cell'])
    self._indxr_mosaic = self._solution['mosaic']

    return

  def _compare_cell(self, c_ref, c_test):
    '''Compare two sets of unit cell constants: if they differ by
    less than 5% / 5 degrees return True, else False.'''

    for j in range(3):
      if math.fabs((c_test[j] - c_ref[j]) / c_ref[j]) > 0.05:
        return False

    for j in range(3, 6):
      if math.fabs(c_test[j] - c_ref[j]) > 5:
        return False

    return True

  def get_solutions():
    return self._solutions

  def get_solution(self):

    import copy

    # FIXME I really need to clean up the code in here...

    if self._indxr_input_lattice is None:
      return copy.deepcopy(
        self._solutions[max(self._solutions.keys())])
    else:
      if self._indxr_input_cell:
        for s in self._solutions.keys():
          if self._solutions[s]['lattice'] == \
            self._indxr_input_lattice:
            if self._compare_cell(
                self._indxr_input_cell,
                self._solutions[s]['cell']):
              return copy.deepcopy(self._solutions[s])
            else:
              del(self._solutions[s])
          else:
            del(self._solutions[s])

        raise RuntimeError, \
          'no solution for lattice %s with given cell' % \
          self._indxr_input_lattice

      else:
        for s in self._solutions.keys():
          if self._solutions[s]['lattice'] == \
            self._indxr_input_lattice:
            return copy.deepcopy(self._solutions[s])
          else:
            del(self._solutions[s])

        raise RuntimeError, 'no solution for lattice %s' % \
          self._indxr_input_lattice

    return

  def _index_finish(self):
    exporter = self.ExportXDS()
    exporter.set_experiments_filename(self.get_solution()['experiments_file'])
    exporter.run()

    for file in ['XPARM.XDS']:
      self._data_files[file] = open(os.path.join(
        self.get_working_directory(), 'xds', file), 'rb').read()

    self._indxr_payload['xds_files'] = self._data_files

    from Wrappers.XDS.XDS import xds_read_xparm
    xparm_dict = xds_read_xparm(
      os.path.join(self.get_working_directory(), 'xds', 'XPARM.XDS'))

    distance = xparm_dict['distance']
    wavelength = xparm_dict['wavelength']
    pixel = xparm_dict['px'], xparm_dict['py']
    beam = xparm_dict['ox'], xparm_dict['oy']

    self._indxr_payload['xds_files'] = self._data_files

    return
<commit_msg>bug fix for narrow wedges<commit_after>#!/usr/bin/env python
# DialsIndexer.py
#   Copyright (C) 2014 Diamond Light Source, Graeme Winter
#
#   This code is distributed under the BSD license, a copy of which is
#   included in the root directory of this package.
#
# An indexer compatible with XDS but using the DIALS methods.

import os
import sys
import math
import shutil

if not os.environ.has_key('XIA2_ROOT'):
  raise RuntimeError, 'XIA2_ROOT not defined'

if not os.environ['XIA2_ROOT'] in sys.path:
  sys.path.append(os.environ['XIA2_ROOT'])

# wrappers for programs that this needs: XDS

from Wrappers.XDS.XDSXycorr import XDSXycorr as _Xycorr
from Wrappers.XDS.XDSInit import XDSInit as _Init
from Wrappers.XDS.XDSColspot import XDSColspot as _Colspot
from Wrappers.XDS.XDSIdxref import XDSIdxref as _Idxref

# wrappers for programs that this needs: DIALS

from Wrappers.Dials.Import import Import as _Import
from Wrappers.Dials.Spotfinder import Spotfinder as _Spotfinder
from Wrappers.Dials.Index import Index as _Index
from Wrappers.Dials.RefineBravaisSettings import RefineBravaisSettings as \
     _RefineBravaisSettings
from Wrappers.Dials.ExportXDS import ExportXDS as _ExportXDS

from Wrappers.XIA.Diffdump import Diffdump

# interfaces that this must implement to be an indexer

from Schema.Interfaces.Indexer import Indexer
from Schema.Interfaces.FrameProcessor import FrameProcessor

# odds and sods that are needed

from lib.bits import auto_logfiler, nint
from Handlers.Streams import Chatter, Debug, Journal
from Handlers.Flags import Flags
from Handlers.Phil import PhilIndex
from Handlers.Files import FileHandler

class DialsIndexer(FrameProcessor,
                   Indexer):
  def __init__(self):

    # set up the inherited objects

    FrameProcessor.__init__(self)
    Indexer.__init__(self)

    self._working_directory = os.getcwd()

    self._background_images = None

    # place to store working data

    self._data_files = { }
    self._solutions = { }

    # FIXME this is a stupid low resolution limit to use...
    self._indxr_low_resolution = 40.0

    return

  # admin functions

  def set_working_directory(self, working_directory):
    self._working_directory = working_directory
    return

  def get_working_directory(self):
    return self._working_directory

  # factory functions

  def Xycorr(self):
    xycorr = _Xycorr()
    xycorr.set_working_directory(self.get_working_directory())

    xycorr.setup_from_image(self.get_image_name(
        self._indxr_images[0][0]))

    if self.get_distance():
      xycorr.set_distance(self.get_distance())

    if self.get_wavelength():
      xycorr.set_wavelength(self.get_wavelength())

    auto_logfiler(xycorr, 'XYCORR')

    return xycorr

  def Init(self):
    from Handlers.Phil import PhilIndex
    init = _Init(params=PhilIndex.params.xds.init)
    init.set_working_directory(self.get_working_directory())

    init.setup_from_image(self.get_image_name(
        self._indxr_images[0][0]))

    if self.get_distance():
      init.set_distance(self.get_distance())

    if self.get_wavelength():
      init.set_wavelength(self.get_wavelength())

    auto_logfiler(init, 'INIT')

    return init

  def Import(self):
    importer = _Import()
    importer.set_working_directory(self.get_working_directory())
    importer.setup_from_image(self.get_image_name(
      self._indxr_images[0][0]))
    auto_logfiler(importer)
    return importer

  def Spotfinder(self):
    spotfinder = _Spotfinder()
    spotfinder.set_working_directory(self.get_working_directory())
    auto_logfiler(spotfinder)
    return spotfinder

  def Index(self):
    index = _Index()
    index.set_working_directory(self.get_working_directory())
    auto_logfiler(index)
    return index

  def ExportXDS(self):
    export_xds = _ExportXDS()
    export_xds.set_working_directory(self.get_working_directory())
    auto_logfiler(export_xds)
    return export_xds

  def RefineBravaisSettings(self):
    rbs = _RefineBravaisSettings()
    rbs.set_working_directory(self.get_working_directory())
    auto_logfiler(rbs)
    return rbs

  ##########################################

  def _index_prepare(self):

    all_images = self.get_matching_images()
    first = min(all_images)
    last = max(all_images)

    self._indxr_images = [(first, last)]

    dd = Diffdump()
    dd.set_image(self.get_image_name(first))
    header = dd.readheader()
    last_background = int(round(5.0 / header['phi_width'])) - 1 + first
    last_background = min(last, last_background)

    # next start to process these - first xycorr
    # FIXME run these *afterwards* as then we have a refined detector geometry
    # so the parallax correction etc. should be slightly better.

    xycorr = self.Xycorr()
    xycorr.set_data_range(first, last)
    xycorr.set_background_range(first, last_background)
    xycorr.run()

    for file in ['X-CORRECTIONS.cbf',
                 'Y-CORRECTIONS.cbf']:
      self._data_files[file] = xycorr.get_output_data_file(file)

    # next start to process these - then init

    init = self.Init()

    for file in ['X-CORRECTIONS.cbf',
                 'Y-CORRECTIONS.cbf']:
      init.set_input_data_file(file, self._data_files[file])

    init.set_data_range(first, last)
    init.set_background_range(first, last_background)
    init.run()

    for file in ['BLANK.cbf',
                 'BKGINIT.cbf',
                 'GAIN.cbf']:
      self._data_files[file] = init.get_output_data_file(file)

    # at this stage, break out to run the DIALS code: this sets itself up

    importer = self.Import()
    importer.run()

    # FIXME this should really use the assigned spot finding regions
    spotfinder = self.Spotfinder()
    spotfinder.set_sweep_filename(importer.get_sweep_filename())
    spotfinder.set_scan_ranges([(first, last)])
    if PhilIndex.params.dials.spotfinder.phil_file is not None:
      spotfinder.set_phil_file(PhilIndex.params.dials.spotfinder.phil_file)
    spotfinder.run()

    self._spot_filename = spotfinder.get_spot_filename()
    self._sweep_filename = importer.get_sweep_filename()

    return

  def _index(self):
    # FIXME allow humans to set the indexing method from whatever list...
    # FIXME respect input unit cell / symmetry if set - or if decided from
    # previous indexing cycle
    indexer = self.Index()
    indexer.set_spot_filename(self._spot_filename)
    indexer.set_sweep_filename(self._sweep_filename)
    if PhilIndex.params.dials.index.phil_file is not None:
      indexer.set_phil_file(PhilIndex.params.dials.index.phil_file)

    if self._indxr_input_lattice:
      indexer.set_indexer_input_lattice(self._indxr_input_lattice)
      Debug.write('Set lattice: %s' % self._indxr_input_lattice)

    if self._indxr_input_cell:
      indexer.set_indexer_input_cell(self._indxr_input_cell)
      Debug.write('Set cell: %f %f %f %f %f %f' % \
                  self._indxr_input_cell)
      original_cell = self._indxr_input_cell

    if PhilIndex.params.dials.index.method is None:
      if self._indxr_input_cell is not None:
        method = 'real_space_grid_search'
      else:
        method = 'fft3d'
      Debug.write('Choosing indexing method: %s' % method)
    else:
      method = PhilIndex.params.dials.index.method

    indexer.run(method)

    # FIXME don't keep hold of an indexer: prevents pickling

    self._indexer = indexer

    # FIXME in here should respect the input unit cell and lattice if provided

    # FIXME from this (i) populate the helper table, (ii) export those files
    # which XDS will expect to find, (iii) try to avoid re-running the indexing
    # step if we eliminate a solution as we have all of the refined results
    # already available.

    rbs = self.RefineBravaisSettings()
    rbs.set_experiments_filename(indexer.get_experiments_filename())
    rbs.set_indexed_filename(indexer.get_indexed_filename())
    rbs.run()

    rmsd_p1 = rbs.get_bravais_summary()[1]['rmsd']

    for k in sorted(rbs.get_bravais_summary()):
      summary = rbs.get_bravais_summary()[k]

      # FIXME need to do this better - for the moment only accept lattices
      # where R.M.S. deviation is less than twice P1 R.M.S. deviation.

      if summary['rmsd'] > 2.0 * rmsd_p1:
        continue

      self._solutions[k] = {
        'number':k,
        'mosaic':0.0,
        'metric':summary['max_angular_difference'],
        'rmsd':summary['rmsd'],
        'nspots':summary['nspots'],
        'lattice':summary['bravais'],
        'cell':summary['unit_cell'],
        'experiments_file':summary['experiments_file']
        }

    self._solution = self.get_solution()

    for solution in self._solutions.keys():
      lattice = self._solutions[solution]['lattice']
      if self._indxr_other_lattice_cell.has_key(lattice):
        if self._indxr_other_lattice_cell[lattice]['metric'] < \
          self._solutions[solution]['metric']:
          continue

      self._indxr_other_lattice_cell[lattice] = {
        'metric':self._solutions[solution]['metric'],
        'cell':self._solutions[solution]['cell']}

    self._indxr_lattice = self._solution['lattice']
    self._indxr_cell = tuple(self._solution['cell'])
    self._indxr_mosaic = self._solution['mosaic']

    return

  def _compare_cell(self, c_ref, c_test):
    '''Compare two sets of unit cell constants: if they differ by
    less than 5% / 5 degrees return True, else False.'''

    for j in range(3):
      if math.fabs((c_test[j] - c_ref[j]) / c_ref[j]) > 0.05:
        return False

    for j in range(3, 6):
      if math.fabs(c_test[j] - c_ref[j]) > 5:
        return False

    return True

  def get_solutions():
    return self._solutions

  def get_solution(self):

    import copy

    # FIXME I really need to clean up the code in here...

    if self._indxr_input_lattice is None:
      return copy.deepcopy(
        self._solutions[max(self._solutions.keys())])
    else:
      if self._indxr_input_cell:
        for s in self._solutions.keys():
          if self._solutions[s]['lattice'] == \
            self._indxr_input_lattice:
            if self._compare_cell(
                self._indxr_input_cell,
                self._solutions[s]['cell']):
              return copy.deepcopy(self._solutions[s])
            else:
              del(self._solutions[s])
          else:
            del(self._solutions[s])

        raise RuntimeError, \
          'no solution for lattice %s with given cell' % \
          self._indxr_input_lattice

      else:
        for s in self._solutions.keys():
          if self._solutions[s]['lattice'] == \
            self._indxr_input_lattice:
            return copy.deepcopy(self._solutions[s])
          else:
            del(self._solutions[s])

        raise RuntimeError, 'no solution for lattice %s' % \
          self._indxr_input_lattice

    return

  def _index_finish(self):
    exporter = self.ExportXDS()
    exporter.set_experiments_filename(self.get_solution()['experiments_file'])
    exporter.run()

    for file in ['XPARM.XDS']:
      self._data_files[file] = open(os.path.join(
        self.get_working_directory(), 'xds', file), 'rb').read()

    self._indxr_payload['xds_files'] = self._data_files

    from Wrappers.XDS.XDS import xds_read_xparm
    xparm_dict = xds_read_xparm(
      os.path.join(self.get_working_directory(), 'xds', 'XPARM.XDS'))

    distance = xparm_dict['distance']
    wavelength = xparm_dict['wavelength']
    pixel = xparm_dict['px'], xparm_dict['py']
    beam = xparm_dict['ox'], xparm_dict['oy']

    self._indxr_payload['xds_files'] = self._data_files

    return
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2010-2011, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# only, as published by the Free Software Foundation.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License version 3 for more details
# (a copy is included in the LICENSE file that accompanied this code).
#
# You should have received a copy of the GNU Lesser General Public License
# version 3 along with OpenQuake.  If not, see
# <http://www.gnu.org/licenses/lgpl-3.0.txt> for a copy of the LGPLv3 License.

"""
Utility functions related to splitting work into tasks.
"""

import itertools
import time

from celery.task.sets import TaskSet

from openquake.job import Job
from openquake import logs
from openquake.utils import config

# Do not create batches of more than DEFAULT_BLOCK_SIZE celery subtasks.
# Celery cannot cope with these and dies.
DEFAULT_BLOCK_SIZE = 4096

def distribute(cardinality, the_task, (name, data), other_args=None,
               flatten_results=False):
    """Runs `the_task` in a task set with the given `cardinality`.

    The given `data` is portioned across the subtasks in the task set.
    The results returned by the subtasks are returned in a list e.g.:
        [result1, result2, ..]
    If each subtask returns a list that will result in list of lists. Please
    set `flatten_results` to `True` if you want the results to be returned in a
    single list.

    :param int cardinality: The size of the task set.
    :param the_task: A `celery` task callable.
    :param str name: The parameter name under which the portioned `data` is to
        be passed to `the_task`.
    :param data: The `data` that is to be portioned and passed to the subtasks
        for processing.
    :param dict other_args: The remaining (keyword) parameters that are to be
        passed to the subtasks.
    :param bool flatten_results: If set, the results will be returned as a
        single list (as opposed to [[results1], [results2], ..]).
    :returns: A list where each element is a result returned by a subtask.
        The result order is the same as the subtask order.
    """
    logs.HAZARD_LOG.info("cardinality: %s" % cardinality)

    block_size = config.get("tasks", "block_size")
    block_size = int(block_size) if block_size else DEFAULT_BLOCK_SIZE

    logs.HAZARD_LOG.info("block_size: %s" % block_size)

    data_length = len(data)
    logs.HAZARD_LOG.info("data_length: %s" % data_length)

    results = []

    for start in xrange(0, data_length, block_size):
        end = start + block_size
        logs.HAZARD_LOG.info("data[%s:%s]" % (start, end))
        iresults = _distribute(cardinality, the_task, name, data[start:end],
                               other_args, flatten_results)
        results.extend(iresults)

    return results

def _distribute(cardinality, a_task, name, data, other_args, flatten_results):
    """Runs `a_task` in a task set with the given `cardinality`.

    The given `data` is portioned across the subtasks in the task set.
    The results returned by the subtasks are returned in a list e.g.:
        [result1, result2, ..]
    If each subtask returns a list that will result in list of lists. Please
    set `flatten_results` to `True` if you want the results to be returned in a
    single list.

    :param int cardinality: The size of the task set.
    :param a_task: A `celery` task callable.
    :param str name: The parameter name under which the portioned `data` is to
        be passed to `a_task`.
    :param data: The `data` that is to be portioned and passed to the subtasks
        for processing.
    :param dict other_args: The remaining (keyword) parameters that are to be
        passed to the subtasks.
    :param bool flatten_results: If set, the results will be returned as a
        single list (as opposed to [[results1], [results2], ..]).
    :returns: A list where each element is a result returned by a subtask.
        The result order is the same as the subtask order.
    """
    # Too many local variables
    # pylint: disable=R0914

    def kwargs(data_portion):
        """
        Construct the full set of keyword parameters for the task to be
        invoked.
        """
        params = {name: data_portion}
        if other_args:
            params.update(other_args)
        return params

    data_length = len(data)
    logs.HAZARD_LOG.info("-data_length: %s" % data_length)

    subtasks = []
    start = 0
    end = chunk_size = int(data_length / float(cardinality))
    if chunk_size == 0:
        # We were given less data items than the number of subtasks
        # specified.  Spawn at least on subtask even if the data to be
        # processed is empty.
        cardinality = data_length if data_length > 0 else 1
        end = chunk_size = 1

    logs.HAZARD_LOG.info("-chunk_size: %s" % chunk_size)

    for _ in xrange(cardinality - 1):
        logs.HAZARD_LOG.info("-data[%s:%s]" % (start, end))
        data_portion = data[start:end]
        subtask = a_task.subtask(**kwargs(data_portion))
        subtasks.append(subtask)
        start = end
        end += chunk_size
    # The last subtask takes the rest of the data.
    data_portion = data[start:]
    subtask = a_task.subtask(**kwargs(data_portion))
    subtasks.append(subtask)

    # At this point we have created all the subtasks and each one got
    # a portion of the data that is to be processed. Now we will create
    # and run the task set.
    logs.HAZARD_LOG.info("-#subtasks: %s" % len(subtasks))
    the_results = _handle_subtasks(subtasks, flatten_results)
    return the_results

def parallelize(
    cardinality, the_task, kwargs, flatten_results=False, index_tasks=True):
    """Runs `the_task` in a task set with the given `cardinality`.

    All subtasks receive the *same* parameters i.e. whatever was passed via
    `kwargs`.

    :param int cardinality: The size of the task set.
    :param the_task: A `celery` task callable.
    :param dict kwargs: The (keyword) parameters that are to be passed
        to *all* subtasks.
    :param bool flatten_results: If set, the results will be returned as a
        single list (as opposed to [[results1], [results2], ..]).
    :param bool index_tasks: If set, each subtask will receive an additional
        `task_index` parameter it can use for the purpose of diversification.
    :returns: A list where each element is a result returned by a subtask.
        The result order is the same as the subtask order.
    :raises WrongTaskParameters: When a task receives a parameter it does not
        know.
    :raises TaskFailed: When at least one subtask fails (raises an exception).
    """
    logs.HAZARD_LOG.info("cardinality: %s" % cardinality)

    assert isinstance(kwargs, dict), "Parameters must be passed in a dict."
    subtasks = []
    for tidx in xrange(cardinality):
        task_args = kwargs
        if index_tasks:
            task_args["task_index"] = tidx
        subtask = the_task.subtask(**task_args)
        subtasks.append(subtask)

    logs.HAZARD_LOG.info("#subtasks: %s" % len(subtasks))

    # At this point we have created all the subtasks.
    the_results = _handle_subtasks(subtasks, flatten_results)
    return the_results

def _handle_subtasks(subtasks, flatten_results):
    """Start a `TaskSet` with the given `subtasks` and wait for it to finish.

    :param subtasks: The subtasks to run
    :type subtasks: [celery_subtask]
    :param bool flatten_results: If set, the results will be returned as a
        single list (as opposed to [[results1], [results2], ..]).
    :returns: A list where each element is a result returned by a subtask.
        The result order is the same as the subtask order.
    :raises WrongTaskParameters: When a task receives a parameter it does not
        know.
    :raises TaskFailed: When at least one subtask fails (raises an exception).
    """
    result = TaskSet(tasks=subtasks).apply_async()

    # Wait for all subtasks to complete.
    while not result.ready():
        time.sleep(0.25)

    the_results = result.join()

    if flatten_results and the_results:
        if isinstance(the_results, list) or isinstance(the_results, tuple):
            the_results = list(itertools.chain(*the_results))

    return the_results

class JobCompletedError(Exception):
    """
    Exception to be thrown by :func:`check_job_status`
    in case of dealing with already completed job.
    """

def check_job_status(job_id):
    """
    Helper function which is intended to be run by celery task functions.

    :raises JobCompletedError:
        If :meth:`~openquake.job.Job.is_job_completed` returns ``True``
        for ``job_id``.
    """
    job = Job.from_kvs(job_id)
    level = job.params.get('debug') if job and job.params else 'warn'
    logs.init_logs_amqp_send(level=level, job_id=job_id)
    if Job.is_job_completed(job_id):
        raise JobCompletedError(job_id)
<commit_msg>switch from info() to debug()<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4

# Copyright (c) 2010-2011, GEM Foundation.
#
# OpenQuake is free software: you can redistribute it and/or modify
# it under the terms of the GNU Lesser General Public License version 3
# only, as published by the Free Software Foundation.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU Lesser General Public License version 3 for more details
# (a copy is included in the LICENSE file that accompanied this code).
#
# You should have received a copy of the GNU Lesser General Public License
# version 3 along with OpenQuake.  If not, see
# <http://www.gnu.org/licenses/lgpl-3.0.txt> for a copy of the LGPLv3 License.

"""
Utility functions related to splitting work into tasks.
"""

import itertools
import time

from celery.task.sets import TaskSet

from openquake.job import Job
from openquake import logs
from openquake.utils import config

# Do not create batches of more than DEFAULT_BLOCK_SIZE celery subtasks.
# Celery cannot cope with these and dies.
DEFAULT_BLOCK_SIZE = 4096

def distribute(cardinality, the_task, (name, data), other_args=None,
               flatten_results=False):
    """Runs `the_task` in a task set with the given `cardinality`.

    The given `data` is portioned across the subtasks in the task set.
    The results returned by the subtasks are returned in a list e.g.:
        [result1, result2, ..]
    If each subtask returns a list that will result in list of lists. Please
    set `flatten_results` to `True` if you want the results to be returned in a
    single list.

    :param int cardinality: The size of the task set.
    :param the_task: A `celery` task callable.
    :param str name: The parameter name under which the portioned `data` is to
        be passed to `the_task`.
    :param data: The `data` that is to be portioned and passed to the subtasks
        for processing.
    :param dict other_args: The remaining (keyword) parameters that are to be
        passed to the subtasks.
    :param bool flatten_results: If set, the results will be returned as a
        single list (as opposed to [[results1], [results2], ..]).
    :returns: A list where each element is a result returned by a subtask.
        The result order is the same as the subtask order.
    """
    logs.HAZARD_LOG.info("cardinality: %s" % cardinality)

    block_size = config.get("tasks", "block_size")
    block_size = int(block_size) if block_size else DEFAULT_BLOCK_SIZE

    logs.HAZARD_LOG.debug("block_size: %s" % block_size)

    data_length = len(data)
    logs.HAZARD_LOG.debug("data_length: %s" % data_length)

    results = []

    for start in xrange(0, data_length, block_size):
        end = start + block_size
        logs.HAZARD_LOG.debug("data[%s:%s]" % (start, end))
        iresults = _distribute(cardinality, the_task, name, data[start:end],
                               other_args, flatten_results)
        results.extend(iresults)

    return results

def _distribute(cardinality, a_task, name, data, other_args, flatten_results):
    """Runs `a_task` in a task set with the given `cardinality`.

    The given `data` is portioned across the subtasks in the task set.
    The results returned by the subtasks are returned in a list e.g.:
        [result1, result2, ..]
    If each subtask returns a list that will result in list of lists. Please
    set `flatten_results` to `True` if you want the results to be returned in a
    single list.

    :param int cardinality: The size of the task set.
    :param a_task: A `celery` task callable.
    :param str name: The parameter name under which the portioned `data` is to
        be passed to `a_task`.
    :param data: The `data` that is to be portioned and passed to the subtasks
        for processing.
    :param dict other_args: The remaining (keyword) parameters that are to be
        passed to the subtasks.
    :param bool flatten_results: If set, the results will be returned as a
        single list (as opposed to [[results1], [results2], ..]).
    :returns: A list where each element is a result returned by a subtask.
        The result order is the same as the subtask order.
    """
    # Too many local variables
    # pylint: disable=R0914

    def kwargs(data_portion):
        """
        Construct the full set of keyword parameters for the task to be
        invoked.
        """
        params = {name: data_portion}
        if other_args:
            params.update(other_args)
        return params

    data_length = len(data)
    logs.HAZARD_LOG.debug("-data_length: %s" % data_length)

    subtasks = []
    start = 0
    end = chunk_size = int(data_length / float(cardinality))
    if chunk_size == 0:
        # We were given less data items than the number of subtasks
        # specified.  Spawn at least on subtask even if the data to be
        # processed is empty.
        cardinality = data_length if data_length > 0 else 1
        end = chunk_size = 1

    logs.HAZARD_LOG.debug("-chunk_size: %s" % chunk_size)

    for _ in xrange(cardinality - 1):
        data_portion = data[start:end]
        subtask = a_task.subtask(**kwargs(data_portion))
        subtasks.append(subtask)
        start = end
        end += chunk_size
    # The last subtask takes the rest of the data.
    data_portion = data[start:]
    subtask = a_task.subtask(**kwargs(data_portion))
    subtasks.append(subtask)

    # At this point we have created all the subtasks and each one got
    # a portion of the data that is to be processed. Now we will create
    # and run the task set.
    logs.HAZARD_LOG.debug("-#subtasks: %s" % len(subtasks))
    the_results = _handle_subtasks(subtasks, flatten_results)
    return the_results

def parallelize(
    cardinality, the_task, kwargs, flatten_results=False, index_tasks=True):
    """Runs `the_task` in a task set with the given `cardinality`.

    All subtasks receive the *same* parameters i.e. whatever was passed via
    `kwargs`.

    :param int cardinality: The size of the task set.
    :param the_task: A `celery` task callable.
    :param dict kwargs: The (keyword) parameters that are to be passed
        to *all* subtasks.
    :param bool flatten_results: If set, the results will be returned as a
        single list (as opposed to [[results1], [results2], ..]).
    :param bool index_tasks: If set, each subtask will receive an additional
        `task_index` parameter it can use for the purpose of diversification.
    :returns: A list where each element is a result returned by a subtask.
        The result order is the same as the subtask order.
    :raises WrongTaskParameters: When a task receives a parameter it does not
        know.
    :raises TaskFailed: When at least one subtask fails (raises an exception).
    """
    logs.HAZARD_LOG.debug("cardinality: %s" % cardinality)

    assert isinstance(kwargs, dict), "Parameters must be passed in a dict."
    subtasks = []
    for tidx in xrange(cardinality):
        task_args = kwargs
        if index_tasks:
            task_args["task_index"] = tidx
        subtask = the_task.subtask(**task_args)
        subtasks.append(subtask)

    logs.HAZARD_LOG.debug("#subtasks: %s" % len(subtasks))

    # At this point we have created all the subtasks.
    the_results = _handle_subtasks(subtasks, flatten_results)
    return the_results

def _handle_subtasks(subtasks, flatten_results):
    """Start a `TaskSet` with the given `subtasks` and wait for it to finish.

    :param subtasks: The subtasks to run
    :type subtasks: [celery_subtask]
    :param bool flatten_results: If set, the results will be returned as a
        single list (as opposed to [[results1], [results2], ..]).
    :returns: A list where each element is a result returned by a subtask.
        The result order is the same as the subtask order.
    :raises WrongTaskParameters: When a task receives a parameter it does not
        know.
    :raises TaskFailed: When at least one subtask fails (raises an exception).
    """
    result = TaskSet(tasks=subtasks).apply_async()

    # Wait for all subtasks to complete.
    while not result.ready():
        time.sleep(0.25)

    the_results = result.join()

    if flatten_results and the_results:
        if isinstance(the_results, list) or isinstance(the_results, tuple):
            the_results = list(itertools.chain(*the_results))

    return the_results

class JobCompletedError(Exception):
    """
    Exception to be thrown by :func:`check_job_status`
    in case of dealing with already completed job.
    """

def check_job_status(job_id):
    """
    Helper function which is intended to be run by celery task functions.

    :raises JobCompletedError:
        If :meth:`~openquake.job.Job.is_job_completed` returns ``True``
        for ``job_id``.
    """
    job = Job.from_kvs(job_id)
    level = job.params.get('debug') if job and job.params else 'warn'
    logs.init_logs_amqp_send(level=level, job_id=job_id)
    if Job.is_job_completed(job_id):
        raise JobCompletedError(job_id)
<|endoftext|>


<commit_before>#!/usr/bin/env python
#########################################################################################
#
# Create mask along z direction.
#
# ---------------------------------------------------------------------------------------
# Copyright (c) 2014 Polytechnique Montreal <www.neuro.polymtl.ca>
# Authors: Julien Cohen-Adad
# Modified: 2014-10-11
#
# About the license: see the file LICENSE.TXT
#########################################################################################

# TODO: scale size in mm.

import sys
import os
import getopt
import commands
import sct_utils as sct
import time
import numpy
import nibabel
from scipy import ndimage
from sct_image import get_orientation
from sct_convert import convert
from msct_image import Image
from sct_image import copy_header, concat_data

# DEFAULT PARAMETERS
class Param:
    def __init__(self):
        self.debug = 0
        self.fname_data = ''
        self.fname_out = ''
        self.method_list = ['coord', 'point', 'centerline', 'center']
        self.method = 'center'  # default method
        self.shape_list = ['cylinder', 'box', 'gaussian']
        self.shape = 'cylinder'  # default shape
        self.size = 41  # in voxel. if gaussian, size corresponds to sigma.
        self.file_prefix = 'mask_'  # output prefix
        self.verbose = 1
        self.remove_tmp_files = 1

# main
#=======================================================================================================================
def main():

    # Parameters for debug mode
    if param.debug:
        print '\n*** WARNING: DEBUG MODE ON ***\n'
        # get path of the testing data
        status, path_sct_data = commands.getstatusoutput('echo $SCT_TESTING_DATA_DIR')
        param.fname_data = path_sct_data+'/mt/mt1.nii.gz'
        param.method = 'point,'+path_sct_data+'/mt/mt1_point.nii.gz' #'centerline,/Users/julien/data/temp/sct_example_data/t2/t2_centerlinerpi.nii.gz'  #coord,68x69'
        param.shape = 'cylinder'
        param.size = 20
        param.remove_tmp_files = 1
        param.verbose = 1
    else:
        # Check input parameters
        try:
            opts, args = getopt.getopt(sys.argv[1:], 'hf:i:m:o:r:s:v:')
        except getopt.GetoptError:
            usage()
        if not opts:
            usage()
        for opt, arg in opts:
            if opt == '-h':
                usage()
            elif opt in '-f':
                param.shape = arg
            elif opt in '-i':
                param.fname_data = arg
            elif opt in '-m':
                param.method = arg
            elif opt in ('-o'):
                param.fname_out = arg
            elif opt in '-r':
                param.remove_tmp_files = int(arg)
            elif opt in '-s':
                param.size = int(arg)
            elif opt in '-v':
                param.verbose = int(arg)

    # run main program
    create_mask()

# create_mask
#=======================================================================================================================
def create_mask():

    fsloutput = 'export FSLOUTPUTTYPE=NIFTI; '  # for faster processing, all outputs are in NIFTI

    # display usage if a mandatory argument is not provided
    if param.fname_data == '' or param.method == '':
        sct.printv('\nERROR: All mandatory arguments are not provided. See usage (add -h).\n', 1, 'error')

    # parse argument for method
    method_list = param.method.replace(' ', '').split(',')  # remove spaces and parse with comma
    # method_list = param.method.split(',')  # parse with comma
    method_type = method_list[0]

    # check existence of method type
    if not method_type in param.method_list:
        sct.printv('\nERROR in '+os.path.basename(__file__)+': Method "'+method_type+'" is not recognized. See usage (add -h).\n', 1, 'error')

    # check method val
    if not method_type == 'center':
        method_val = method_list[1]
    del method_list

    # check existence of shape
    if not param.shape in param.shape_list:
        sct.printv('\nERROR in '+os.path.basename(__file__)+': Shape "'+param.shape+'" is not recognized. See usage (add -h).\n', 1, 'error')

    # check existence of input files
    sct.printv('\ncheck existence of input files...', param.verbose)
    sct.check_file_exist(param.fname_data, param.verbose)
    if method_type == 'centerline':
        sct.check_file_exist(method_val, param.verbose)

    # check if orientation is RPI
    sct.printv('\nCheck if orientation is RPI...', param.verbose)
    ori = get_orientation(param.fname_data, filename=True)
    if not ori == 'RPI':
        sct.printv('\nERROR in '+os.path.basename(__file__)+': Orientation of input image should be RPI. Use sct_image -setorient to put your image in RPI.\n', 1, 'error')

    # display input parameters
    sct.printv('\nInput parameters:', param.verbose)
    sct.printv('  data ..................'+param.fname_data, param.verbose)
    sct.printv('  method ................'+method_type, param.verbose)

    # Extract path/file/extension
    path_data, file_data, ext_data = sct.extract_fname(param.fname_data)

    # Get output folder and file name
    if param.fname_out == '':
        param.fname_out = param.file_prefix+file_data+ext_data
    #fname_out = os.path.abspath(path_out+file_out+ext_out)

    # create temporary folder
    sct.printv('\nCreate temporary folder...', param.verbose)
    path_tmp = sct.slash_at_the_end('tmp.'+time.strftime("%y%m%d%H%M%S"), 1)
    sct.run('mkdir '+path_tmp, param.verbose)

    # Copying input data to tmp folder and convert to nii
    sct.printv('\nCopying input data to tmp folder and convert to nii...', param.verbose)
    convert(param.fname_data, path_tmp+'data.nii')
    # sct.run('cp '+param.fname_data+' '+path_tmp+'data'+ext_data, param.verbose)
    if method_type == 'centerline':
        convert(method_val, path_tmp+'centerline.nii.gz')

    # go to tmp folder
    os.chdir(path_tmp)

    # Get dimensions of data
    sct.printv('\nGet dimensions of data...', param.verbose)
    nx, ny, nz, nt, px, py, pz, pt = Image('data.nii').dim
    sct.printv('  ' + str(nx) + ' x ' + str(ny) + ' x ' + str(nz)+ ' x ' + str(nt), param.verbose)
    # in case user input 4d data
    if nt != 1:
        sct.printv('WARNING in '+os.path.basename(__file__)+': Input image is 4d but output mask will 3D.', param.verbose, 'warning')
        # extract first volume to have 3d reference
        nii = Image('data.nii')
        data3d = nii.data[:,:,:,0]
        nii.data = data3d
        nii.save()

    if method_type == 'coord':
        # parse to get coordinate
        coord = map(int, method_val.split('x'))

    if method_type == 'point':
        # get file name
        fname_point = method_val
        # extract coordinate of point
        sct.printv('\nExtract coordinate of point...', param.verbose)
        status, output = sct.run('sct_label_utils -i '+fname_point+' -t display-voxel', param.verbose)
        # parse to get coordinate
        coord = output[output.find('Position=')+10:-17].split(',')

    if method_type == 'center':
        # set coordinate at center of FOV
        coord = round(float(nx)/2), round(float(ny)/2)

    if method_type == 'centerline':
        # get name of centerline from user argument
        fname_centerline = 'centerline.nii.gz'
    else:
        # generate volume with line along Z at coordinates 'coord'
        sct.printv('\nCreate line...', param.verbose)
        fname_centerline = create_line('data.nii', coord, nz)

    # create mask
    sct.printv('\nCreate mask...', param.verbose)
    centerline = nibabel.load(fname_centerline)  # open centerline
    hdr = centerline.get_header()  # get header
    hdr.set_data_dtype('uint8')  # set imagetype to uint8
    data_centerline = centerline.get_data()  # get centerline
    z_centerline = [iz for iz in range(0, nz, 1) if data_centerline[:, :, iz].any()]
    nz = len(z_centerline)
    # get center of mass of the centerline
    cx = [0] * nz
    cy = [0] * nz
    for iz in range(0, nz, 1):
        cx[iz], cy[iz] = ndimage.measurements.center_of_mass(numpy.array(data_centerline[:, :, z_centerline[iz]]))
    # create 2d masks
    file_mask = 'data_mask'
    for iz in range(nz):
        center = numpy.array([cx[iz], cy[iz]])
        mask2d = create_mask2d(center, param.shape, param.size, nx, ny)
        # Write NIFTI volumes
        img = nibabel.Nifti1Image(mask2d, None, hdr)
        nibabel.save(img, (file_mask+str(iz)+'.nii'))
    # merge along Z
    # cmd = 'fslmerge -z mask '
    im_list = []
    for iz in range(nz):
        im_list.append(Image(file_mask+str(iz)+'.nii'))
    im_out = concat_data(im_list, 2)
    im_out.setFileName('mask.nii.gz')
    im_out.save()

    # copy geometry
    im_dat = Image('data.nii')
    im_mask = Image('mask.nii.gz')
    im_mask = copy_header(im_dat, im_mask)
    im_mask.save()

    # come back to parent folder
    os.chdir('..')

    # Generate output files
    sct.printv('\nGenerate output files...', param.verbose)
    sct.generate_output_file(path_tmp+'mask.nii.gz', param.fname_out)

    # Remove temporary files
    if param.remove_tmp_files == 1:
        sct.printv('\nRemove temporary files...', param.verbose)
        sct.run('rm -rf '+path_tmp, param.verbose, error_exit='warning')

    # to view results
    sct.printv('\nDone! To view results, type:', param.verbose)
    sct.printv('fslview '+param.fname_data+' '+param.fname_out+' -l Red -t 0.5 &', param.verbose, 'info')
    print

# create_line
# ==========================================================================================
def create_line(fname, coord, nz):

    # duplicate volume (assumes input file is nifti)
    sct.run('cp '+fname+' line.nii', param.verbose)

    # set all voxels to zero
    sct.run('sct_maths -i line.nii -mul 0 -o line.nii', param.verbose)

    cmd = 'sct_label_utils -i line.nii -o line.nii -t add -x '
    for iz in range(nz):
        if iz == nz-1:
            cmd += str(int(coord[0]))+','+str(int(coord[1]))+','+str(iz)+',1'
        else:
            cmd += str(int(coord[0]))+','+str(int(coord[1]))+','+str(iz)+',1:'

    sct.run(cmd, param.verbose)

    return 'line.nii'

# create_mask2d
# ==========================================================================================
def create_mask2d(center, shape, size, nx, ny):

    # initialize 2d grid
    xx, yy = numpy.mgrid[:nx, :ny]
    mask2d = numpy.zeros((nx, ny))
    xc = center[0]
    yc = center[1]
    radius = round(float(size+1)/2)  # add 1 because the radius includes the center.

    if shape == 'box':
        mask2d[xc-radius:xc+radius+1, yc-radius:yc+radius+1] = 1

    elif shape == 'cylinder':
        mask2d = ((xx-xc)**2 + (yy-yc)**2 <= radius**2)*1

    elif shape == 'gaussian':
        sigma = float(radius)
        mask2d = numpy.exp(-(((xx-xc)**2)/(2*(sigma**2)) + ((yy-yc)**2)/(2*(sigma**2))))

    # import matplotlib.pyplot as plt
    # plt.imshow(mask2d)
    # plt.show()

    return mask2d

# Print usage
# ==========================================================================================
def usage():
    print """
"""+os.path.basename(__file__)+"""
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Part of the Spinal Cord Toolbox <https://sourceforge.net/projects/spinalcordtoolbox>

DESCRIPTION
  Create mask along z direction.

USAGE
  """+os.path.basename(__file__)+""" -i <data> -m <method,val> -s <size>

MANDATORY ARGUMENTS
  -i <data>        image to create mask on. Only used to get header. Must be 3D.

OPTIONAL ARGUMENTS
  -m <method,val>  method to generate mask and associated value. Default="""+str(param_default.method)+"""
                     coord: X,Y coordinate of center of mask. E.g.: coord,20x15
                     point: volume that contains a single point. E.g.: point,label.nii.gz
                     center: mask is created at center of FOV. In that case, "val" is not required.
                     centerline: volume that contains centerline. E.g.: centerline,my_centerline.nii
  -s <size>        size in voxel. Odd values are better (for mask to be symmetrical). Default="""+str(param_default.size)+"""
                   If shape=gaussian, size corresponds to "sigma".
  -f {box,cylinder,gaussian}  shape of the mask. Default="""+str(param_default.shape)+"""
  -o <output>      name of output mask. Default is "mask_INPUTFILE".
  -r {0,1}         remove temporary files. Default="""+str(param_default.remove_tmp_files)+"""
  -v {0,1}         verbose. Default="""+str(param_default.verbose)+"""
  -h               help. Show this message

EXAMPLE
  """+os.path.basename(__file__)+""" -i dwi_mean.nii -m coord,35x42 -s 20 -f box\n"""

    # exit program
    sys.exit(2)

#=======================================================================================================================
# Start program
#=======================================================================================================================
if __name__ == "__main__":
    param = Param()
    param_default = Param()
    main()
<commit_msg>REF: added feature to sct_create_mask for mask with even size<commit_after>#!/usr/bin/env python
#########################################################################################
#
# Create mask along z direction.
#
# ---------------------------------------------------------------------------------------
# Copyright (c) 2014 Polytechnique Montreal <www.neuro.polymtl.ca>
# Authors: Julien Cohen-Adad
# Modified: 2014-10-11
#
# About the license: see the file LICENSE.TXT
#########################################################################################

# TODO: scale size in mm.

import sys
import os
import getopt
import commands
import sct_utils as sct
import time
import numpy
import nibabel
from scipy import ndimage
from sct_image import get_orientation
from sct_convert import convert
from msct_image import Image
from sct_image import copy_header, concat_data

# DEFAULT PARAMETERS
class Param:
    def __init__(self):
        self.debug = 0
        self.fname_data = ''
        self.fname_out = ''
        self.method_list = ['coord', 'point', 'centerline', 'center']
        self.method = 'center'  # default method
        self.shape_list = ['cylinder', 'box', 'gaussian']
        self.shape = 'cylinder'  # default shape
        self.size = 41  # in voxel. if gaussian, size corresponds to sigma.
        self.even = 0
        self.file_prefix = 'mask_'  # output prefix
        self.verbose = 1
        self.remove_tmp_files = 1

# main
#=======================================================================================================================
def main():

    # Parameters for debug mode
    if param.debug:
        print '\n*** WARNING: DEBUG MODE ON ***\n'
        # get path of the testing data
        status, path_sct_data = commands.getstatusoutput('echo $SCT_TESTING_DATA_DIR')
        param.fname_data = path_sct_data+'/mt/mt1.nii.gz'
        param.method = 'point,'+path_sct_data+'/mt/mt1_point.nii.gz' #'centerline,/Users/julien/data/temp/sct_example_data/t2/t2_centerlinerpi.nii.gz'  #coord,68x69'
        param.shape = 'cylinder'
        param.size = 20
        param.remove_tmp_files = 1
        param.verbose = 1
    else:
        # Check input parameters
        try:
            opts, args = getopt.getopt(sys.argv[1:], 'hf:i:m:o:r:s:e:v:')
        except getopt.GetoptError:
            usage()
        if not opts:
            usage()
        for opt, arg in opts:
            if opt == '-h':
                usage()
            elif opt in '-f':
                param.shape = arg
            elif opt in '-i':
                param.fname_data = arg
            elif opt in '-m':
                param.method = arg
            elif opt in '-o':
                param.fname_out = arg
            elif opt in '-r':
                param.remove_tmp_files = int(arg)
            elif opt in '-s':
                param.size = int(arg)
            elif opt in '-e':
                param.even = int(arg)
            elif opt in '-v':
                param.verbose = int(arg)

    # run main program
    create_mask()

# create_mask
#=======================================================================================================================
def create_mask():

    fsloutput = 'export FSLOUTPUTTYPE=NIFTI; '  # for faster processing, all outputs are in NIFTI

    # display usage if a mandatory argument is not provided
    if param.fname_data == '' or param.method == '':
        sct.printv('\nERROR: All mandatory arguments are not provided. See usage (add -h).\n', 1, 'error')

    # parse argument for method
    method_list = param.method.replace(' ', '').split(',')  # remove spaces and parse with comma
    # method_list = param.method.split(',')  # parse with comma
    method_type = method_list[0]

    # check existence of method type
    if not method_type in param.method_list:
        sct.printv('\nERROR in '+os.path.basename(__file__)+': Method "'+method_type+'" is not recognized. See usage (add -h).\n', 1, 'error')

    # check method val
    if not method_type == 'center':
        method_val = method_list[1]
    del method_list

    # check existence of shape
    if not param.shape in param.shape_list:
        sct.printv('\nERROR in '+os.path.basename(__file__)+': Shape "'+param.shape+'" is not recognized. See usage (add -h).\n', 1, 'error')

    # check existence of input files
    sct.printv('\ncheck existence of input files...', param.verbose)
    sct.check_file_exist(param.fname_data, param.verbose)
    if method_type == 'centerline':
        sct.check_file_exist(method_val, param.verbose)

    # check if orientation is RPI
    sct.printv('\nCheck if orientation is RPI...', param.verbose)
    ori = get_orientation(param.fname_data, filename=True)
    if not ori == 'RPI':
        sct.printv('\nERROR in '+os.path.basename(__file__)+': Orientation of input image should be RPI. Use sct_image -setorient to put your image in RPI.\n', 1, 'error')

    # display input parameters
    sct.printv('\nInput parameters:', param.verbose)
    sct.printv('  data ..................'+param.fname_data, param.verbose)
    sct.printv('  method ................'+method_type, param.verbose)

    # Extract path/file/extension
    path_data, file_data, ext_data = sct.extract_fname(param.fname_data)

    # Get output folder and file name
    if param.fname_out == '':
        param.fname_out = param.file_prefix+file_data+ext_data
    #fname_out = os.path.abspath(path_out+file_out+ext_out)

    # create temporary folder
    sct.printv('\nCreate temporary folder...', param.verbose)
    path_tmp = sct.slash_at_the_end('tmp.'+time.strftime("%y%m%d%H%M%S"), 1)
    sct.run('mkdir '+path_tmp, param.verbose)

    # Copying input data to tmp folder and convert to nii
    sct.printv('\nCopying input data to tmp folder and convert to nii...', param.verbose)
    convert(param.fname_data, path_tmp+'data.nii')
    # sct.run('cp '+param.fname_data+' '+path_tmp+'data'+ext_data, param.verbose)
    if method_type == 'centerline':
        convert(method_val, path_tmp+'centerline.nii.gz')

    # go to tmp folder
    os.chdir(path_tmp)

    # Get dimensions of data
    sct.printv('\nGet dimensions of data...', param.verbose)
    nx, ny, nz, nt, px, py, pz, pt = Image('data.nii').dim
    sct.printv('  ' + str(nx) + ' x ' + str(ny) + ' x ' + str(nz)+ ' x ' + str(nt), param.verbose)
    # in case user input 4d data
    if nt != 1:
        sct.printv('WARNING in '+os.path.basename(__file__)+': Input image is 4d but output mask will 3D.', param.verbose, 'warning')
        # extract first volume to have 3d reference
        nii = Image('data.nii')
        data3d = nii.data[:,:,:,0]
        nii.data = data3d
        nii.save()

    if method_type == 'coord':
        # parse to get coordinate
        coord = map(int, method_val.split('x'))

    if method_type == 'point':
        # get file name
        fname_point = method_val
        # extract coordinate of point
        sct.printv('\nExtract coordinate of point...', param.verbose)
        status, output = sct.run('sct_label_utils -i '+fname_point+' -t display-voxel', param.verbose)
        # parse to get coordinate
        coord = output[output.find('Position=')+10:-17].split(',')

    if method_type == 'center':
        # set coordinate at center of FOV
        coord = round(float(nx)/2), round(float(ny)/2)

    if method_type == 'centerline':
        # get name of centerline from user argument
        fname_centerline = 'centerline.nii.gz'
    else:
        # generate volume with line along Z at coordinates 'coord'
        sct.printv('\nCreate line...', param.verbose)
        fname_centerline = create_line('data.nii', coord, nz)

    # create mask
    sct.printv('\nCreate mask...', param.verbose)
    centerline = nibabel.load(fname_centerline)  # open centerline
    hdr = centerline.get_header()  # get header
    hdr.set_data_dtype('uint8')  # set imagetype to uint8
    data_centerline = centerline.get_data()  # get centerline
    z_centerline = [iz for iz in range(0, nz, 1) if data_centerline[:, :, iz].any()]
    nz = len(z_centerline)
    # get center of mass of the centerline
    cx = [0] * nz
    cy = [0] * nz
    for iz in range(0, nz, 1):
        cx[iz], cy[iz] = ndimage.measurements.center_of_mass(numpy.array(data_centerline[:, :, z_centerline[iz]]))
    # create 2d masks
    file_mask = 'data_mask'
    for iz in range(nz):
        center = numpy.array([cx[iz], cy[iz]])
        mask2d = create_mask2d(center, param.shape, param.size, nx, ny, param.even)
        # Write NIFTI volumes
        img = nibabel.Nifti1Image(mask2d, None, hdr)
        nibabel.save(img, (file_mask+str(iz)+'.nii'))
    # merge along Z
    # cmd = 'fslmerge -z mask '
    im_list = []
    for iz in range(nz):
        im_list.append(Image(file_mask+str(iz)+'.nii'))
    im_out = concat_data(im_list, 2)
    im_out.setFileName('mask.nii.gz')
    im_out.save()

    # copy geometry
    im_dat = Image('data.nii')
    im_mask = Image('mask.nii.gz')
    im_mask = copy_header(im_dat, im_mask)
    im_mask.save()

    # come back to parent folder
    os.chdir('..')

    # Generate output files
    sct.printv('\nGenerate output files...', param.verbose)
    sct.generate_output_file(path_tmp+'mask.nii.gz', param.fname_out)

    # Remove temporary files
    if param.remove_tmp_files == 1:
        sct.printv('\nRemove temporary files...', param.verbose)
        sct.run('rm -rf '+path_tmp, param.verbose, error_exit='warning')

    # to view results
    sct.printv('\nDone! To view results, type:', param.verbose)
    sct.printv('fslview '+param.fname_data+' '+param.fname_out+' -l Red -t 0.5 &', param.verbose, 'info')
    print

# create_line
# ==========================================================================================
def create_line(fname, coord, nz):

    # duplicate volume (assumes input file is nifti)
    sct.run('cp '+fname+' line.nii', param.verbose)

    # set all voxels to zero
    sct.run('sct_maths -i line.nii -mul 0 -o line.nii', param.verbose)

    cmd = 'sct_label_utils -i line.nii -o line.nii -t add -x '
    for iz in range(nz):
        if iz == nz-1:
            cmd += str(int(coord[0]))+','+str(int(coord[1]))+','+str(iz)+',1'
        else:
            cmd += str(int(coord[0]))+','+str(int(coord[1]))+','+str(iz)+',1:'

    sct.run(cmd, param.verbose)

    return 'line.nii'

# create_mask2d
# ==========================================================================================
def create_mask2d(center, shape, size, nx, ny, even=0):

    # initialize 2d grid
    xx, yy = numpy.mgrid[:nx, :ny]
    mask2d = numpy.zeros((nx, ny))
    xc = center[0]
    yc = center[1]
    if even != 0:
        radius = int(size / 2)
    else:
        radius = round(float(size + 1) / 2)  # add 1 because the radius includes the center.

    if shape == 'box':
        if even != 0:
            mask2d[xc - radius:xc + radius, yc - radius:yc + radius] = 1
        else:
            mask2d[xc-radius:xc+radius+1, yc-radius:yc+radius+1] = 1

    elif shape == 'cylinder':
        mask2d = ((xx-xc)**2 + (yy-yc)**2 <= radius**2)*1

    elif shape == 'gaussian':
        sigma = float(radius)
        mask2d = numpy.exp(-(((xx-xc)**2)/(2*(sigma**2)) + ((yy-yc)**2)/(2*(sigma**2))))

    # import matplotlib.pyplot as plt
    # plt.imshow(mask2d)
    # plt.show()

    return mask2d

# Print usage
# ==========================================================================================
def usage():
    print """
"""+os.path.basename(__file__)+"""
~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
Part of the Spinal Cord Toolbox <https://sourceforge.net/projects/spinalcordtoolbox>

DESCRIPTION
  Create mask along z direction.

USAGE
  """+os.path.basename(__file__)+""" -i <data> -m <method,val> -s <size>

MANDATORY ARGUMENTS
  -i <data>        image to create mask on. Only used to get header. Must be 3D.

OPTIONAL ARGUMENTS
  -m <method,val>  method to generate mask and associated value. Default="""+str(param_default.method)+"""
                     coord: X,Y coordinate of center of mask. E.g.: coord,20x15
                     point: volume that contains a single point. E.g.: point,label.nii.gz
                     center: mask is created at center of FOV. In that case, "val" is not required.
                     centerline: volume that contains centerline. E.g.: centerline,my_centerline.nii
  -s <size>        size in voxel. Odd values are better (for mask to be symmetrical). Default="""+str(param_default.size)+"""
                   If shape=gaussian, size corresponds to "sigma".
  -e {0,1}         0: box size is odd. 1: box size is even.
  -f {box,cylinder,gaussian}  shape of the mask. Default="""+str(param_default.shape)+"""
  -o <output>      name of output mask. Default is "mask_INPUTFILE".
  -r {0,1}         remove temporary files. Default="""+str(param_default.remove_tmp_files)+"""
  -v {0,1}         verbose. Default="""+str(param_default.verbose)+"""
  -h               help. Show this message

EXAMPLE
  """+os.path.basename(__file__)+""" -i dwi_mean.nii -m coord,35x42 -s 20 -f box\n"""

    # exit program
    sys.exit(2)

#=======================================================================================================================
# Start program
#=======================================================================================================================
if __name__ == "__main__":
    param = Param()
    param_default = Param()
    main()
<|endoftext|>


<commit_before><commit_msg>OPT: made -o mandatory. fixed #217<commit_after><|endoftext|>


<commit_before>"""all things PeerAssets protocol."""

import warnings
from pypeerassets.kutil import Kutil
from pypeerassets.paproto_pb2 import DeckSpawn as deckspawnproto
from pypeerassets.paproto_pb2 import CardTransfer as cardtransferproto
from pypeerassets.exceptions import RecieverAmountMismatch
from operator import itemgetter
from pypeerassets.card_parsers import *
from enum import Enum
from typing import List, Generator, cast

class IssueMode(Enum):

    NONE = 0x00
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L19
    # No issuance allowed.

    CUSTOM = 0x01
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L20
    # Custom issue mode, verified by client aware of this.

    ONCE = 0x02
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L21
    # A single card_issue transaction allowed.

    MULTI = 0x04
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L22
    # Multiple card_issue transactions allowed.

    MONO = 0x08
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L23
    # All card transaction amounts are equal to 1.

    UNFLUSHABLE = 0x10
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L24
    # The UNFLUSHABLE issue mode invalidates any card transfer transaction except for the card issue transaction.
    # Meaning that only the issuing entity is able to change the balance of a specific address.
    # To correctly calculate the balance of a PeerAssets addres a client should only consider the card transfer
    # transactions originating from the deck owner.

    SUBSCRIPTION = 0x34  # SUBSCRIPTION (34 = 20 | 4 | 10)
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L26
    # The SUBSCRIPTION issue mode marks an address holding tokens as subscribed for a limited timeframe. This timeframe is
    # defined by the balance of the account and the time at which the first cards of this token are received.
    # To check validity of a subscription one should take the timestamp of the first received cards and add the address' balance to it in hours.

    SINGLET = 0x0a  # SINGLET is a combination of ONCE and MONO (2 | 8)
    #  Singlet deck, one MONO card issunce allowed

class Deck:

    def __init__(self, name: str, number_of_decimals: int, issue_mode: int,
                 network: str, production: bool, version: int,
                 asset_specific_data: bytes=None, issuer: str="", time: int=None,
                 id: str=None, tx_confirmations: int=None) -> None:
        '''
        Initialize deck object, load from dictionary Deck(**dict) or initilize
        with kwargs Deck("deck", 3, "ONCE")
        '''

        self.version = version  # protocol version
        self.name = name  # deck name
        self.issue_mode = issue_mode  # deck issue mode
        self.number_of_decimals = number_of_decimals
        self.asset_specific_data = asset_specific_data  # optional metadata for the deck
        self.id = id
        self.issuer = issuer
        self.issue_time = time
        self.confirms = tx_confirmations
        self.network = network
        self.production = production
        if self.network.startswith("t") or 'testnet' in self.network:
            self.testnet = True
        else:
            self.testnet = False

    @property
    def p2th_address(self) -> Optional[str]:
        '''P2TH address of this deck'''

        if self.id:
            return Kutil(network=self.network,
                         privkey=bytearray.fromhex(self.id)).address
        else:
            return None

    @property
    def p2th_wif(self) -> Optional[str]:
        '''P2TH privkey in WIF format'''

        if self.id:
            return Kutil(network=self.network,
                         privkey=bytearray.fromhex(self.id)).wif
        else:
            return None

    @property
    def metainfo_to_protobuf(self) -> bytes:
        '''encode deck into protobuf'''

        deck = deckspawnproto()
        deck.version = self.version
        deck.name = self.name
        deck.number_of_decimals = self.number_of_decimals
        deck.issue_mode = self.issue_mode
        if self.asset_specific_data:
            if not isinstance(self.asset_specific_data, bytes):
                deck.asset_specific_data = self.asset_specific_data.encode()
            else:
                deck.asset_specific_data = self.asset_specific_data

        proto = deck.SerializeToString()

        if len(proto) > 80:
            warnings.warn('\nMetainfo size exceeds maximum of 80bytes that fit into OP_RETURN.')

        return proto

    @property
    def metainfo_to_dict(self) -> dict:
        '''encode deck into dictionary'''

        r = {
            "version": self.version,
            "name": self.name,
            "number_of_decimals": self.number_of_decimals,
            "issue_mode": self.issue_mode
        }

        if self.asset_specific_data:
            r.update({'asset_specific_data': self.asset_specific_data})

        return r

    def __str__(self) -> str:

        r = []
        for key in self.__dict__:
            r.append("{key}='{value}'".format(key=key, value=self.__dict__[key]))

        return ', '.join(r)

class CardTransfer:

    def __init__(self, deck: Deck, receiver: list=[], amount: List[int]=[],
                 version: int=1, blockhash: str=None, txid: str=None,
                 sender: str=None, asset_specific_data: bytes=None,
                 number_of_decimals: int=None, blockseq: int=None,
                 cardseq: int=None, blocknum: int=None,
                 timestamp: int=None, tx_confirmations: int=None,
                 type: str=None) -> None:

        '''CardTransfer object, used when parsing card_transfers from the blockchain
        or when sending out new card_transfer.
        It can be initialized by passing the **kwargs and it will do the parsing,
        or it can be initialized with passed arguments.

        * deck - instance of Deck object
        * receiver - list of receivers
        * amount - list of amounts to be sent, must be integer
        * version - protocol version, default 1
        * txid - transaction ID of CardTransfer
        * sender - transaction sender
        * blockhash - block ID where the tx was first included
        * blockseq - order in which tx was serialized into block
        * timestamp - unix timestamp of the block where it was first included
        * tx_confirmations - number of confirmations of the transaction
        * asset_specific_data - extra metadata
        * number_of_decimals - number of decimals for amount, inherited from Deck object
        : type: card type [CardIssue, CardTransfer, CardBurn]'''

        if not len(amount) == len(receiver):
            raise RecieverAmountMismatch({"error": "carn mmount must match card receiver."})

        self.version = version
        self.deck_id = deck.id
        self.deck_p2th = deck.p2th_address
        self.txid = txid
        self.sender = sender
        self.asset_specific_data = asset_specific_data
        if not number_of_decimals:
            self.number_of_decimals = deck.number_of_decimals
        else:
            self.number_of_decimals = number_of_decimals

        self.receiver = receiver
        self.amount = amount

        if blockhash:
            self.blockhash = blockhash
            self.blockseq = blockseq
            self.timestamp = timestamp
            self.blocknum = blocknum
            self.cardseq = cardseq
            self.confirms = tx_confirmations
        else:
            self.blockhash = ""
            self.blockseq = 0
            self.blocknum = 0
            self.timestamp = 0
            self.cardseq = 0
            self.confirms = 0

        if type:
            self.type = type

        if self.sender == deck.issuer:
            self.type = "CardIssue"
        elif self.receiver[0] == deck.issuer:
            self.type = "CardBurn"
        else:
            self.type = "CardTransfer"

    @property
    def metainfo_to_protobuf(self) -> str:
        '''encode card_transfer info to protobuf'''

        card = cardtransferproto()
        card.version = self.version
        card.amount.extend(self.amount)
        card.number_of_decimals = self.number_of_decimals
        if self.asset_specific_data:
            if not isinstance(self.asset_specific_data, bytes):
                card.asset_specific_data = self.asset_specific_data.encode()
            else:
                card.asset_specific_data = self.asset_specific_data

        proto = card.SerializeToString()

        if len(proto) > 80:
            warnings.warn('\nMetainfo size exceeds maximum of 80bytes that fit into OP_RETURN.')

        return proto

    @property
    def metainfo_to_dict(self) -> dict:
        '''encode card into dictionary'''

        r = {
            "version": self.version,
            "amount": self.amount,
            "number_of_decimals": self.number_of_decimals
        }

        if self.asset_specific_data:
            r.update({'asset_specific_data': self.asset_specific_data})

        return r

    def __str__(self) -> str:

        r = []
        for key in self.__dict__:
            r.append("{key}='{value}'".format(key=key, value=self.__dict__[key]))

        return ', '.join(r)

def validate_card_issue_modes(issue_mode: int, cards: list) -> list:
    """validate cards against deck_issue modes"""

    supported_mask = 63  # sum of all issue_mode values

    if not bool(issue_mode & supported_mask):
        return []  # return empty list

    for i in [1 << x for x in range(len(IssueMode))]:
        if bool(i & issue_mode):

            try:
                parser_fn = cast(
                    Callable[[list], Optional[list]],
                    parsers[IssueMode(i).name]
                )
            except ValueError:
                continue

            parsed_cards = parser_fn(cards)
            if not parsed_cards:
                return []
            cards = parsed_cards

    return cards

class DeckState:

    def __init__(self, cards: Generator) -> None:

        self.cards = cards
        self.total = 0
        self.burned = 0
        self.balances = cast(dict, {})
        self.processed_issues = cast(dict, {})
        self.processed_transfers = cast(dict, {})
        self.processed_burns = cast(dict, {})

        self.calc_state()
        self.checksum = not bool(self.total - sum(self.balances.values()))

    def _process(self, card: dict, ctype: str) -> bool:

        sender = card["sender"]
        receiver = card["receiver"][0]
        amount = card["amount"][0]

        if ctype != 'CardIssue':
            balance_check = sender in self.balances and self.balances[sender] >= amount

            if balance_check:
                self.balances[sender] -= amount

                if 'CardBurn' not in ctype:
                    self._append_balance(amount, receiver)

                return True

            return False

        if 'CardIssue' in ctype:
            self._append_balance(amount, receiver)
            return True

        return False

    def _append_balance(self, amount: int, receiver: str) -> None:

            try:
                self.balances[receiver] += amount
            except KeyError:
                self.balances[receiver] = amount

    def _sort_cards(self, cards: Generator) -> list:
        '''sort cards by blocknum and blockseq'''

        return sorted([card.__dict__ for card in cards],
                            key=itemgetter('blocknum', 'blockseq', 'cardseq'))

    def calc_state(self) -> None:

        for card in self._sort_cards(self.cards):

            # txid + blockseq + cardseq, as unique ID
            cid = card["txid"] + str(card["blockseq"]) + str(card["cardseq"])
            ctype = card["type"]
            amount = card["amount"][0]

            if ctype == 'CardIssue' and cid not in self.processed_issues:
                validate = self._process(card, ctype)
                self.total += amount * validate  # This will set amount to 0 if validate is False
                self.processed_issues[cid] = cid

            if ctype == 'CardTransfer' and cid not in self.processed_transfers:
                self._process(card, ctype)
                self.processed_transfers[cid] = cid

            if ctype == 'CardBurn' and cid not in self.processed_burns:
                validate = self._process(card, ctype)

                self.total -= amount * validate
                self.burned += amount * validate
                self.processed_burns[cid] = cid
<commit_msg>use set over dict for processed cards<commit_after>"""all things PeerAssets protocol."""

import warnings
from pypeerassets.kutil import Kutil
from pypeerassets.paproto_pb2 import DeckSpawn as deckspawnproto
from pypeerassets.paproto_pb2 import CardTransfer as cardtransferproto
from pypeerassets.exceptions import RecieverAmountMismatch
from operator import itemgetter
from pypeerassets.card_parsers import *
from enum import Enum
from typing import List, Generator, cast

class IssueMode(Enum):

    NONE = 0x00
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L19
    # No issuance allowed.

    CUSTOM = 0x01
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L20
    # Custom issue mode, verified by client aware of this.

    ONCE = 0x02
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L21
    # A single card_issue transaction allowed.

    MULTI = 0x04
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L22
    # Multiple card_issue transactions allowed.

    MONO = 0x08
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L23
    # All card transaction amounts are equal to 1.

    UNFLUSHABLE = 0x10
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L24
    # The UNFLUSHABLE issue mode invalidates any card transfer transaction except for the card issue transaction.
    # Meaning that only the issuing entity is able to change the balance of a specific address.
    # To correctly calculate the balance of a PeerAssets addres a client should only consider the card transfer
    # transactions originating from the deck owner.

    SUBSCRIPTION = 0x34  # SUBSCRIPTION (34 = 20 | 4 | 10)
    # https://github.com/PeerAssets/rfcs/blob/master/0001-peerassets-transaction-specification.proto#L26
    # The SUBSCRIPTION issue mode marks an address holding tokens as subscribed for a limited timeframe. This timeframe is
    # defined by the balance of the account and the time at which the first cards of this token are received.
    # To check validity of a subscription one should take the timestamp of the first received cards and add the address' balance to it in hours.

    SINGLET = 0x0a  # SINGLET is a combination of ONCE and MONO (2 | 8)
    #  Singlet deck, one MONO card issunce allowed

class Deck:

    def __init__(self, name: str, number_of_decimals: int, issue_mode: int,
                 network: str, production: bool, version: int,
                 asset_specific_data: bytes=None, issuer: str="", time: int=None,
                 id: str=None, tx_confirmations: int=None) -> None:
        '''
        Initialize deck object, load from dictionary Deck(**dict) or initilize
        with kwargs Deck("deck", 3, "ONCE")
        '''

        self.version = version  # protocol version
        self.name = name  # deck name
        self.issue_mode = issue_mode  # deck issue mode
        self.number_of_decimals = number_of_decimals
        self.asset_specific_data = asset_specific_data  # optional metadata for the deck
        self.id = id
        self.issuer = issuer
        self.issue_time = time
        self.confirms = tx_confirmations
        self.network = network
        self.production = production
        if self.network.startswith("t") or 'testnet' in self.network:
            self.testnet = True
        else:
            self.testnet = False

    @property
    def p2th_address(self) -> Optional[str]:
        '''P2TH address of this deck'''

        if self.id:
            return Kutil(network=self.network,
                         privkey=bytearray.fromhex(self.id)).address
        else:
            return None

    @property
    def p2th_wif(self) -> Optional[str]:
        '''P2TH privkey in WIF format'''

        if self.id:
            return Kutil(network=self.network,
                         privkey=bytearray.fromhex(self.id)).wif
        else:
            return None

    @property
    def metainfo_to_protobuf(self) -> bytes:
        '''encode deck into protobuf'''

        deck = deckspawnproto()
        deck.version = self.version
        deck.name = self.name
        deck.number_of_decimals = self.number_of_decimals
        deck.issue_mode = self.issue_mode
        if self.asset_specific_data:
            if not isinstance(self.asset_specific_data, bytes):
                deck.asset_specific_data = self.asset_specific_data.encode()
            else:
                deck.asset_specific_data = self.asset_specific_data

        proto = deck.SerializeToString()

        if len(proto) > 80:
            warnings.warn('\nMetainfo size exceeds maximum of 80bytes that fit into OP_RETURN.')

        return proto

    @property
    def metainfo_to_dict(self) -> dict:
        '''encode deck into dictionary'''

        r = {
            "version": self.version,
            "name": self.name,
            "number_of_decimals": self.number_of_decimals,
            "issue_mode": self.issue_mode
        }

        if self.asset_specific_data:
            r.update({'asset_specific_data': self.asset_specific_data})

        return r

    def __str__(self) -> str:

        r = []
        for key in self.__dict__:
            r.append("{key}='{value}'".format(key=key, value=self.__dict__[key]))

        return ', '.join(r)

class CardTransfer:

    def __init__(self, deck: Deck, receiver: list=[], amount: List[int]=[],
                 version: int=1, blockhash: str=None, txid: str=None,
                 sender: str=None, asset_specific_data: bytes=None,
                 number_of_decimals: int=None, blockseq: int=None,
                 cardseq: int=None, blocknum: int=None,
                 timestamp: int=None, tx_confirmations: int=None,
                 type: str=None) -> None:

        '''CardTransfer object, used when parsing card_transfers from the blockchain
        or when sending out new card_transfer.
        It can be initialized by passing the **kwargs and it will do the parsing,
        or it can be initialized with passed arguments.

        * deck - instance of Deck object
        * receiver - list of receivers
        * amount - list of amounts to be sent, must be integer
        * version - protocol version, default 1
        * txid - transaction ID of CardTransfer
        * sender - transaction sender
        * blockhash - block ID where the tx was first included
        * blockseq - order in which tx was serialized into block
        * timestamp - unix timestamp of the block where it was first included
        * tx_confirmations - number of confirmations of the transaction
        * asset_specific_data - extra metadata
        * number_of_decimals - number of decimals for amount, inherited from Deck object
        : type: card type [CardIssue, CardTransfer, CardBurn]'''

        if not len(amount) == len(receiver):
            raise RecieverAmountMismatch({"error": "carn mmount must match card receiver."})

        self.version = version
        self.deck_id = deck.id
        self.deck_p2th = deck.p2th_address
        self.txid = txid
        self.sender = sender
        self.asset_specific_data = asset_specific_data
        if not number_of_decimals:
            self.number_of_decimals = deck.number_of_decimals
        else:
            self.number_of_decimals = number_of_decimals

        self.receiver = receiver
        self.amount = amount

        if blockhash:
            self.blockhash = blockhash
            self.blockseq = blockseq
            self.timestamp = timestamp
            self.blocknum = blocknum
            self.cardseq = cardseq
            self.confirms = tx_confirmations
        else:
            self.blockhash = ""
            self.blockseq = 0
            self.blocknum = 0
            self.timestamp = 0
            self.cardseq = 0
            self.confirms = 0

        if type:
            self.type = type

        if self.sender == deck.issuer:
            self.type = "CardIssue"
        elif self.receiver[0] == deck.issuer:
            self.type = "CardBurn"
        else:
            self.type = "CardTransfer"

    @property
    def metainfo_to_protobuf(self) -> str:
        '''encode card_transfer info to protobuf'''

        card = cardtransferproto()
        card.version = self.version
        card.amount.extend(self.amount)
        card.number_of_decimals = self.number_of_decimals
        if self.asset_specific_data:
            if not isinstance(self.asset_specific_data, bytes):
                card.asset_specific_data = self.asset_specific_data.encode()
            else:
                card.asset_specific_data = self.asset_specific_data

        proto = card.SerializeToString()

        if len(proto) > 80:
            warnings.warn('\nMetainfo size exceeds maximum of 80bytes that fit into OP_RETURN.')

        return proto

    @property
    def metainfo_to_dict(self) -> dict:
        '''encode card into dictionary'''

        r = {
            "version": self.version,
            "amount": self.amount,
            "number_of_decimals": self.number_of_decimals
        }

        if self.asset_specific_data:
            r.update({'asset_specific_data': self.asset_specific_data})

        return r

    def __str__(self) -> str:

        r = []
        for key in self.__dict__:
            r.append("{key}='{value}'".format(key=key, value=self.__dict__[key]))

        return ', '.join(r)

def validate_card_issue_modes(issue_mode: int, cards: list) -> list:
    """validate cards against deck_issue modes"""

    supported_mask = 63  # sum of all issue_mode values

    if not bool(issue_mode & supported_mask):
        return []  # return empty list

    for i in [1 << x for x in range(len(IssueMode))]:
        if bool(i & issue_mode):

            try:
                parser_fn = cast(
                    Callable[[list], Optional[list]],
                    parsers[IssueMode(i).name]
                )
            except ValueError:
                continue

            parsed_cards = parser_fn(cards)
            if not parsed_cards:
                return []
            cards = parsed_cards

    return cards

class DeckState:

    def __init__(self, cards: Generator) -> None:

        self.cards = cards
        self.total = 0
        self.burned = 0
        self.balances = cast(dict, {})
        self.processed_issues = set()
        self.processed_transfers = set()
        self.processed_burns = set()

        self.calc_state()
        self.checksum = not bool(self.total - sum(self.balances.values()))

    def _process(self, card: dict, ctype: str) -> bool:

        sender = card["sender"]
        receiver = card["receiver"][0]
        amount = card["amount"][0]

        if ctype != 'CardIssue':
            balance_check = sender in self.balances and self.balances[sender] >= amount

            if balance_check:
                self.balances[sender] -= amount

                if 'CardBurn' not in ctype:
                    self._append_balance(amount, receiver)

                return True

            return False

        if 'CardIssue' in ctype:
            self._append_balance(amount, receiver)
            return True

        return False

    def _append_balance(self, amount: int, receiver: str) -> None:

            try:
                self.balances[receiver] += amount
            except KeyError:
                self.balances[receiver] = amount

    def _sort_cards(self, cards: Generator) -> list:
        '''sort cards by blocknum and blockseq'''

        return sorted([card.__dict__ for card in cards],
                            key=itemgetter('blocknum', 'blockseq', 'cardseq'))

    def calc_state(self) -> None:

        for card in self._sort_cards(self.cards):

            # txid + blockseq + cardseq, as unique ID
            cid = str(card["txid"] + str(card["blockseq"]) + str(card["cardseq"]))
            ctype = card["type"]
            amount = card["amount"][0]

            if ctype == 'CardIssue' and cid not in self.processed_issues:
                validate = self._process(card, ctype)
                self.total += amount * validate  # This will set amount to 0 if validate is False
                self.processed_issues |= {cid}

            if ctype == 'CardTransfer' and cid not in self.processed_transfers:
                self._process(card, ctype)
                self.processed_transfers |= {cid}

            if ctype == 'CardBurn' and cid not in self.processed_burns:
                validate = self._process(card, ctype)

                self.total -= amount * validate
                self.burned += amount * validate
                self.processed_burns |= {cid}
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
__author__ = 'Paweł Sołtysiak'

import numpy as np
import pandas as pd
# import

def readData():
    columnNames = [
        'animalName',
        'hair',
        'feathers',
        'eggs',
        'milk',
        'airborne',
        'aquatic',
        'predator',
        'toothed',
        'backbone',
        'breathes',
        'venomous',
        'fins',
        'legs',
        'tail',
        'domestic',
        'catsize',
        'type'
    ]
    data = pd.read_csv(u'../Datasets/zoo.data', sep=',', names=columnNames, header=None)

    types = {
        1: 'mammal',
        2: 'bird',
        3: 'reptile',
        4: 'fish',
        5: 'amphibian',
        6: 'insect',
        7: 'shellfish',
    }
    data['type'].replace(types, inplace=True)
    return data

if __name__ == "__main__":
    df = readData()
    dataProps = {
        'sum': df.sum(),
        'count': df.count(),
        'mean': df.mean(),
        'std': df.std(),
    }
    print(dataProps['sum'])
    typeGroups = df.groupby('type')
    dataCount = typeGroups.agg(['sum', 'count', 'mean', 'std'])
    # print dataCount.loc['fish']
    #print dataCount.loc['mammal'].loc['legs']
    pass
<commit_msg>created classify method without using legs in consideration<commit_after>from __future__ import division
# -*- coding: utf-8 -*-
__author__ = 'Paweł Sołtysiak'

import numpy as np
import pandas as pd

# import
class BayesZoo:
    numbersToNamesMapping = {
        1: 'mammal',
        2: 'bird',
        3: 'reptile',
        4: 'fish',
        5: 'amphibian',
        6: 'insect',
        7: 'shellfish',
    }
    typeNames = numbersToNamesMapping.values()
    columnNames = [
        'animalName',
        'hair',
        'feathers',
        'eggs',
        'milk',
        'airborne',
        'aquatic',
        'predator',
        'toothed',
        'backbone',
        'breathes',
        'venomous',
        'fins',
        'legs',
        'tail',
        'domestic',
        'catsize',
        'type'
    ]

    def __init__(self):
        self.readData()
        pass

    def calculateBinaryProbality(self, typeName, columnName):
        sumForType = self.typeGroupsAgg.loc[typeName].loc[columnName].loc['sum']
        classCount = self.data[columnName].sum()
        return sumForType / classCount

    def calculateValuesProbality(self, typeName, columnName):
        typeStd = self.typeGroupsAgg.loc[typeName].loc[columnName].loc['std']
        mean = self.typeGroupsAgg.loc[typeName].loc[columnName].loc['mean']
        denominator = np.sqrt(2 * np.pi * typeStd ** 2)
        value = (1 / denominator) * np.exp(-(6 - mean) ** 2 / 2 * typeStd ** 2)

        return value

    def calculateProbality(self, typeName, columnName):
        if columnName == "legs":
            value = self.calculateValuesProbality(typeName, columnName)
        else:
            value = self.calculateBinaryProbality(typeName, columnName)

        return value

    def readData(self):

        self.data = pd.read_csv(u'../Datasets/zoo.data', sep=',', names=self.columnNames, header=None)
        self.data['type'].replace(self.numbersToNamesMapping, inplace=True)

        self.typeGroups = self.data.groupby(['type'])
        self.typeGroupsAgg = self.typeGroups.agg(['sum', 'count', 'mean', 'std'])

    def classify(self, data):
        pForEachType = {key: 0 for (key) in self.typeNames}

        for typeName in self.typeNames:
            pType = b.data['type'].value_counts().loc[typeName] / len(b.data)
            for columnName in self.columnNames[1:-1]:
                if columnName == "legs":
                    continue
                pForClass = self.calculateProbality(typeName, columnName)
                if columnName != "legs" and data[columnName] == 0:
                    pForClass = 1 - pForClass
                pType *= pForClass
            pForEachType[typeName] = pType
        return pForEachType

if __name__ == "__main__":
    b = BayesZoo()
    sample = "0,1,1 0,1,0 1,0,1 1,0,0 2,1,0 0"

    testData = {
        'hair': 1,
        'feathers': 0,
        'eggs': 0,
        'milk': 1,
        'airborne': 0,
        'aquatic': 0,
        'predator': 1,
        'toothed': 1,
        'backbone': 1,
        'breathes': 1,
        'venomous': 0,
        'fins': 0,
        # 'legs':2,
        'tail': 0,
        'domestic': 1,
        'catsize': 0,

    }
    result = b.classify(testData)
    print max(result, key=result.get)

    pass
<|endoftext|>