<commit_before># -*- coding: utf-8 -*-

# Copyright (C) 2006-2016  Music Technology Group - Universitat Pompeu Fabra
#
# This file is part of Essentia
#
# Essentia is free software: you can redistribute it and/or modify it under
# the terms of the GNU Affero General Public License as published by the Free
# Software Foundation (FSF), either version 3 of the License, or (at your
# option) any later version.
#
# This program 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 General Public License for more
# details.
#
# You should have received a copy of the Affero GNU General Public License
# version 3 along with this program. If not, see http://www.gnu.org/licenses/

#
# Essentia documentation build configuration file, created by
# sphinx-quickstart on Thu Nov 11 16:09:32 2010.
#
# This file is execfile()d with the current directory set to its containing dir.
#
# Note that not all possible configuration values are present in this
# autogenerated file.
#
# All configuration values have a default; values that are commented out
# serve to show the default.

import sys, os

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#sys.path.insert(0, os.path.abspath('.'))

# -- General configuration -----------------------------------------------------

# If your documentation needs a minimal Sphinx version, state it here.
#needs_sphinx = '1.0'

# Add any Sphinx extension module names here, as strings. They can be extensions
# coming with Sphinx (named 'sphinx.ext.*') or your custom ones.
extensions = [ 'sphinx.ext.viewcode',
               'sphinxcontrib.doxylink' ]

# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']

# The suffix of source filenames.
source_suffix = '.rst'

# The encoding of source files.
#source_encoding = 'utf-8-sig'

# The master toctree document.
master_doc = 'contents'

# General information about the project.
project = u'Essentia'
copyright = u'2006-2017, Universitat Pompeu Fabra'

# The version info for the project you're documenting, acts as replacement for
# |version| and |release|, also used in various other places throughout the
# built documents.
#
# The short X.Y version.
version = '2.1-dev'
# The full version, including alpha/beta/rc tags.
release = '2.1-dev'

# The language for content autogenerated by Sphinx. Refer to documentation
# for a list of supported languages.
#language = None

# There are two options for replacing |today|: either, you set today to some
# non-false value, then it is used:
#today = ''
# Else, today_fmt is used as the format for a strftime call.
#today_fmt = '%B %d, %Y'

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
exclude_patterns = ['_build']

# The reST default role (used for this markup: `text`) to use for all documents.
#default_role = None

# If true, '()' will be appended to :func: etc. cross-reference text.
#add_function_parentheses = True

# If true, the current module name will be prepended to all description
# unit titles (such as .. function::).
#add_module_names = True

# If true, sectionauthor and moduleauthor directives will be shown in the
# output. They are ignored by default.
#show_authors = False

# The name of the Pygments (syntax highlighting) style to use.
pygments_style = 'friendly'

# A list of ignored prefixes for module index sorting.
#modindex_common_prefix = []

# -- Options for HTML output ---------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
html_theme = 'sphinxdoc_mtg'
html_theme_path = ['_templates']
html_theme_options = {
    "nosidebar": "true"
}

# Theme options are theme-specific and customize the look and feel of a theme
# further.  For a list of options available for each theme, see the
# documentation.
#html_theme_options = {}

# Add any paths that contain custom themes here, relative to this directory.
#html_theme_path = []

# The name for this set of Sphinx documents.  If None, it defaults to
# "<project> v<release> documentation".
#html_title = None

# A shorter title for the navigation bar.  Default is the same as html_title.
#html_short_title = None

# The name of an image file (relative to this directory) to place at the top
# of the sidebar.
#html_logo = None
#html_logo = 'mtg_logo.png'

# The name of an image file (within the static path) to use as favicon of the
# docs.  This file should be a Windows icon file (.ico) being 16x16 or 32x32
# pixels large.
html_favicon = 'favicon.ico'

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['_static']

# If not '', a 'Last updated on:' timestamp is inserted at every page bottom,
# using the given strftime format.
#html_last_updated_fmt = '%b %d, %Y'

# If true, SmartyPants will be used to convert quotes and dashes to
# typographically correct entities.
#html_use_smartypants = True

# Custom sidebar templates, maps document names to template names.
#html_sidebars = { 'index': [ 'indexsidebar.html', 'searchbox.html' ] }
#html_sidebars = { '**': [] } # disable sidebar on all pages

# Additional templates that should be rendered to pages, maps page names to
# template names.
html_additional_pages = { 'index': 'index.html',
                          'algorithms_reference': 'algorithms_reference.html',
                          'applications': 'applications.html',
                          'documentation': 'documentation.html' }

execfile('essentia_reference.py')
html_additional_pages.update(essentia_algorithms)

# If false, no module index is generated.
#html_domain_indices = True

# If false, no index is generated.
html_use_index = False

# If true, the index is split into individual pages for each letter.
#html_split_index = False

# If true, links to the reST sources are added to the pages.
#html_show_sourcelink = True

# If true, "Created using Sphinx" is shown in the HTML footer. Default is True.
#html_show_sphinx = True

# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True.
#html_show_copyright = True

# If true, an OpenSearch description file will be output, and all pages will
# contain a <link> tag referring to it.  The value of this option must be the
# base URL from which the finished HTML is served.
#html_use_opensearch = ''

# This is the file name suffix for HTML files (e.g. ".xhtml").
#html_file_suffix = None

# Output file base name for HTML help builder.
htmlhelp_basename = 'Essentiadoc'

# -- Options for LaTeX output --------------------------------------------------

# The paper size ('letter' or 'a4').
#latex_paper_size = 'letter'

# The font size ('10pt', '11pt' or '12pt').
#latex_font_size = '10pt'

# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title, author, documentclass [howto/manual]).
latex_documents = [
  ('index', 'Essentia.tex', u'Essentia Documentation',
   u'Nicolas Wack', 'manual'),
]

# The name of an image file (relative to this directory) to place at the top of
# the title page.
#latex_logo = None

# For "manual" documents, if this is true, then toplevel headings are parts,
# not chapters.
#latex_use_parts = False

# If true, show page references after internal links.
#latex_show_pagerefs = False

# If true, show URL addresses after external links.
#latex_show_urls = False

# Additional stuff for the LaTeX preamble.
#latex_preamble = ''

# Documents to append as an appendix to all manuals.
#latex_appendices = []

# If false, no module index is generated.
#latex_domain_indices = True

# -- Options for manual page output --------------------------------------------

# One entry per manual page. List of tuples
# (source start file, name, description, authors, manual section).
man_pages = [
    ('index', 'essentia', u'Essentia Documentation',
     [u'Nicolas Wack'], 1)
]

doxylink = {
    'essentia': ('EssentiaDoxygen.tag', 'doxygen')
}
<commit_msg>Small improvements in config code formatting<commit_after># -*- coding: utf-8 -*-

# Copyright (C) 2006-2016  Music Technology Group - Universitat Pompeu Fabra
#
# This file is part of Essentia
#
# Essentia is free software: you can redistribute it and/or modify it under
# the terms of the GNU Affero General Public License as published by the Free
# Software Foundation (FSF), either version 3 of the License, or (at your
# option) any later version.
#
# This program 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 General Public License for more
# details.
#
# You should have received a copy of the Affero GNU General Public License
# version 3 along with this program. If not, see http://www.gnu.org/licenses/

#
# Essentia documentation build configuration file, created by
# sphinx-quickstart on Thu Nov 11 16:09:32 2010.
#
# This file is execfile()d with the current directory set to its containing dir.
#
# Note that not all possible configuration values are present in this
# autogenerated file.
#
# All configuration values have a default; values that are commented out
# serve to show the default.

import sys, os

# If extensions (or modules to document with autodoc) are in another directory,
# add these directories to sys.path here. If the directory is relative to the
# documentation root, use os.path.abspath to make it absolute, like shown here.
#sys.path.insert(0, os.path.abspath('.'))

# -- General configuration -----------------------------------------------------

# If your documentation needs a minimal Sphinx version, state it here.
#needs_sphinx = '1.0'

# Add any Sphinx extension module names here, as strings. They can be extensions
# coming with Sphinx (named 'sphinx.ext.*') or your custom ones.
extensions = ['sphinx.ext.viewcode',
               'sphinxcontrib.doxylink']

# Add any paths that contain templates here, relative to this directory.
templates_path = ['_templates']

# The suffix of source filenames.
source_suffix = '.rst'

# The encoding of source files.
#source_encoding = 'utf-8-sig'

# The master toctree document.
master_doc = 'contents'

# General information about the project.
project = u'Essentia'
copyright = u'2006-2017, Universitat Pompeu Fabra'

# The version info for the project you're documenting, acts as replacement for
# |version| and |release|, also used in various other places throughout the
# built documents.
#
# The short X.Y version.
version = '2.1-dev'
# The full version, including alpha/beta/rc tags.
release = '2.1-dev'

# The language for content autogenerated by Sphinx. Refer to documentation
# for a list of supported languages.
#language = None

# There are two options for replacing |today|: either, you set today to some
# non-false value, then it is used:
#today = ''
# Else, today_fmt is used as the format for a strftime call.
#today_fmt = '%B %d, %Y'

# List of patterns, relative to source directory, that match files and
# directories to ignore when looking for source files.
exclude_patterns = ['_build']

# The reST default role (used for this markup: `text`) to use for all documents.
#default_role = None

# If true, '()' will be appended to :func: etc. cross-reference text.
#add_function_parentheses = True

# If true, the current module name will be prepended to all description
# unit titles (such as .. function::).
#add_module_names = True

# If true, sectionauthor and moduleauthor directives will be shown in the
# output. They are ignored by default.
#show_authors = False

# The name of the Pygments (syntax highlighting) style to use.
pygments_style = 'friendly'

# A list of ignored prefixes for module index sorting.
#modindex_common_prefix = []

# -- Options for HTML output ---------------------------------------------------

# The theme to use for HTML and HTML Help pages.  See the documentation for
# a list of builtin themes.
html_theme = 'sphinxdoc_mtg'

# Add any paths that contain custom themes here, relative to this directory.
html_theme_path = ['_templates']

# Theme options are theme-specific and customize the look and feel of a theme
# further.  For a list of options available for each theme, see the
# documentation.
html_theme_options = {
    "nosidebar": "true"
}

# The name for this set of Sphinx documents.  If None, it defaults to
# "<project> v<release> documentation".
#html_title = None

# A shorter title for the navigation bar.  Default is the same as html_title.
#html_short_title = None

# The name of an image file (relative to this directory) to place at the top
# of the sidebar.
#html_logo = None
#html_logo = 'mtg_logo.png'

# The name of an image file (within the static path) to use as favicon of the
# docs.  This file should be a Windows icon file (.ico) being 16x16 or 32x32
# pixels large.
html_favicon = 'favicon.ico'

# Add any paths that contain custom static files (such as style sheets) here,
# relative to this directory. They are copied after the builtin static files,
# so a file named "default.css" will overwrite the builtin "default.css".
html_static_path = ['_static']

# If not '', a 'Last updated on:' timestamp is inserted at every page bottom,
# using the given strftime format.
#html_last_updated_fmt = '%b %d, %Y'

# If true, SmartyPants will be used to convert quotes and dashes to
# typographically correct entities.
#html_use_smartypants = True

# Custom sidebar templates, maps document names to template names.
#html_sidebars = { 'index': [ 'indexsidebar.html', 'searchbox.html' ] }
#html_sidebars = { '**': [] } # disable sidebar on all pages

# Additional templates that should be rendered to pages, maps page names to
# template names.
html_additional_pages = {'index': 'index.html',
                         'algorithms_reference': 'algorithms_reference.html',
                         'applications': 'applications.html',
                         'documentation': 'documentation.html'}

execfile('essentia_reference.py')
html_additional_pages.update(essentia_algorithms)

# If false, no module index is generated.
#html_domain_indices = True

# If false, no index is generated.
html_use_index = False

# If true, the index is split into individual pages for each letter.
#html_split_index = False

# If true, links to the reST sources are added to the pages.
#html_show_sourcelink = True

# If true, "Created using Sphinx" is shown in the HTML footer. Default is True.
#html_show_sphinx = True

# If true, "(C) Copyright ..." is shown in the HTML footer. Default is True.
#html_show_copyright = True

# If true, an OpenSearch description file will be output, and all pages will
# contain a <link> tag referring to it.  The value of this option must be the
# base URL from which the finished HTML is served.
#html_use_opensearch = ''

# This is the file name suffix for HTML files (e.g. ".xhtml").
#html_file_suffix = None

# Output file base name for HTML help builder.
htmlhelp_basename = 'Essentiadoc'

# -- Options for LaTeX output --------------------------------------------------

# The paper size ('letter' or 'a4').
#latex_paper_size = 'letter'

# The font size ('10pt', '11pt' or '12pt').
#latex_font_size = '10pt'

# Grouping the document tree into LaTeX files. List of tuples
# (source start file, target name, title, author, documentclass [howto/manual]).
latex_documents = [
  ('index', 'Essentia.tex', u'Essentia Documentation',
   u'Nicolas Wack', 'manual'),
]

# The name of an image file (relative to this directory) to place at the top of
# the title page.
#latex_logo = None

# For "manual" documents, if this is true, then toplevel headings are parts,
# not chapters.
#latex_use_parts = False

# If true, show page references after internal links.
#latex_show_pagerefs = False

# If true, show URL addresses after external links.
#latex_show_urls = False

# Additional stuff for the LaTeX preamble.
#latex_preamble = ''

# Documents to append as an appendix to all manuals.
#latex_appendices = []

# If false, no module index is generated.
#latex_domain_indices = True

# -- Options for manual page output --------------------------------------------

# One entry per manual page. List of tuples
# (source start file, name, description, authors, manual section).
man_pages = [
    ('index', 'essentia', u'Essentia Documentation',
     [u'Nicolas Wack'], 1)
]

doxylink = {
    'essentia': ('EssentiaDoxygen.tag', 'doxygen')
}
<|endoftext|>


<commit_before># -*- coding: utf8 -*-

import copy
import hashlib
import logging

from ethereum.utils import ecrecover_to_pub
from py_ecc.secp256k1 import N as secp256k1n
from rlp.utils import ALL_BYTES

from mythril.laser.ethereum.util import bytearray_to_int, sha3

class NativeContractException(Exception):
    pass

def int_to_32bytes(i):   # used because int can't fit as bytes function's input
    o = [0] * 32
    for x in range(32):
        o[31 - x] = i & 0xff
        i >>= 8
    return bytes(o)

def extract32(data, i):
    if i >= len(data):
        return 0
    o = data[i: min(i + 32, len(data))]
    o.extend(bytearray(32 - len(o)))
    return bytearray_to_int(o)

def ecrecover(data):
    try:
        data = bytearray(data)
    except TypeError:
        raise NativeContractException
    message = b''.join([ALL_BYTES[x] for x in data[0:32]])
    v = extract32(data, 32)
    r = extract32(data, 64)
    s = extract32(data, 96)
    if r >= secp256k1n or s >= secp256k1n or v < 27 or v > 28:
        return []
    try:
        pub = ecrecover_to_pub(message, v, r, s)
    except Exception as e:
        logging.info("An error has occured while extracting public key: "+e)
        return []
    o = [0] * 12 + [x for x in sha3(pub)[-20:]]
    return o

def sha256(data):
    try:
        data = bytes(data)
    except TypeError:
        raise NativeContractException
    return hashlib.sha256(data).digest()

def ripemd160(data):
    try:
        data = bytes(data)
    except TypeError:
        raise NativeContractException
    return 12*[0]+[i for i in hashlib.new('ripemd160', data).digest()]

def identity(data):
    return copy.copy(data)

def native_contracts(address, data):
    """
    takes integer address 1, 2, 3, 4
    """

    functions = (ecrecover, sha256, ripemd160, identity)
    return functions[address-1](data)
<commit_msg>Move the v, r, s for catching exception<commit_after># -*- coding: utf8 -*-

import copy
import hashlib
import logging

from ethereum.utils import ecrecover_to_pub
from py_ecc.secp256k1 import N as secp256k1n
from rlp.utils import ALL_BYTES

from mythril.laser.ethereum.util import bytearray_to_int, sha3

class NativeContractException(Exception):
    pass

def int_to_32bytes(i):   # used because int can't fit as bytes function's input
    o = [0] * 32
    for x in range(32):
        o[31 - x] = i & 0xff
        i >>= 8
    return bytes(o)

def extract32(data, i):
    if i >= len(data):
        return 0
    o = data[i: min(i + 32, len(data))]
    o.extend(bytearray(32 - len(o)))
    return bytearray_to_int(o)

def ecrecover(data):

    try:
        data = bytearray(data)
        v = extract32(data, 32)
        r = extract32(data, 64)
        s = extract32(data, 96)
    except TypeError:
        raise NativeContractException

    message = b''.join([ALL_BYTES[x] for x in data[0:32]])
    if r >= secp256k1n or s >= secp256k1n or v < 27 or v > 28:
        return []
    try:
        pub = ecrecover_to_pub(message, v, r, s)
    except Exception as e:
        logging.info("An error has occured while extracting public key: "+e)
        return []
    o = [0] * 12 + [x for x in sha3(pub)[-20:]]
    return o

def sha256(data):
    try:
        data = bytes(data)
    except TypeError:
        raise NativeContractException
    return hashlib.sha256(data).digest()

def ripemd160(data):
    try:
        data = bytes(data)
    except TypeError:
        raise NativeContractException
    return 12*[0]+[i for i in hashlib.new('ripemd160', data).digest()]

def identity(data):
    return copy.copy(data)

def native_contracts(address, data):
    """
    takes integer address 1, 2, 3, 4
    """

    functions = (ecrecover, sha256, ripemd160, identity)
    return functions[address-1](data)
<|endoftext|>


<commit_before>import abc
import glob
import json
import csv
import os
from itertools import chain
from functools import reduce
from operator import lt, gt

from nalaf.structures.data import Entity, Relation
from nalaf.utils import MUT_CLASS_ID

class AnnotationReader:
    """
    Abstract class for annotating the dataset.
    Subclasses that inherit this class should:
    * Be named [Name]AnnotationReader
    * Implement the abstract method annotate
    * Append new items to the list field "annotations" of each Part in the dataset
    """

    @abc.abstractmethod
    def annotate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        return

class AnnJsonAnnotationReader(AnnotationReader):
    """
    Reads the annotations from .ann.json format.

    Implements the abstract class Annotator.
    """

    def __init__(self, directory, read_just_mutations=True, delete_incomplete_docs=True, is_predicted=False, read_relations=False):
        self.directory = directory
        """the directory containing *.ann.json files"""
        self.read_just_mutations = read_just_mutations
        """whether to read in only mutation entities"""
        self.delete_incomplete_docs = delete_incomplete_docs
        """whether to delete documents from the dataset that are not marked as 'anncomplete'"""
        self.is_predicted = is_predicted
        """whether the annotation is predicted or real, which determines where it will be saved"""
        self.read_relations = read_relations
        """whether relations should be read as well"""

    def annotate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        if not os.path.isdir(self.directory):
            filenames = [self.directory]
        else:
            filenames = glob.glob(str(self.directory + "/*.ann.json"))

        for filename in filenames:
            with open(filename, 'r', encoding="utf-8") as file:
                try:
                    basename = os.path.basename(filename)
                    if '-' in basename:
                        doc_id = filename.split('-')[-1].replace('.ann.json', '')
                    else:
                        doc_id = basename.replace('.ann.json', '')

                    ann_json = json.load(file)
                    document = dataset.documents[doc_id]
                    if ann_json['anncomplete']:
                        for entity in ann_json['entities']:
                            # if read_just_mutations is False
                            if not self.read_just_mutations or entity['classId'] == MUT_CLASS_ID:
                                # if it is predicted put it in predicted_annotations else in annotations
                                if self.is_predicted:
                                    ann = Entity(entity['classId'], entity['offsets'][0]['start'],
                                                     entity['offsets'][0]['text'], entity['confidence']['prob'])
                                    document.parts[entity['part']].predicted_annotations.append(ann)
                                else:
                                    ann = Entity(entity['classId'], entity['offsets'][0]['start'],
                                                     entity['offsets'][0]['text'])
                                    document.parts[entity['part']].annotations.append(ann)
                        if self.read_relations:
                            for relation in ann_json['relations']:
                                part = document.parts[relation['entities'][0].split('|')[0]]
                                e1_start = int(relation['entities'][0].split('|')[1].split(',')[0])
                                e2_start = int(relation['entities'][1].split('|')[1].split(',')[0])
                                e1_end = int(relation['entities'][0].split('|')[1].split(',')[1])
                                e2_end = int(relation['entities'][1].split('|')[1].split(',')[1])
                                e1_text = part.text[e1_start:e1_end]
                                e2_text = part.text[e2_start:e2_end]
                                part.relations.append(Relation(e1_start, e2_start, e1_text, e2_text, relation['classId']))
                    elif self.delete_incomplete_docs:
                        del dataset.documents[doc_id]
                    elif self.is_predicted:  # anncomplete=False, delete_incomplete_docs=False, is_predicted=True
                        for entity in ann_json['entities']:
                            if not self.read_just_mutations or entity['classId'] == MUT_CLASS_ID:
                                ann = Entity(entity['classId'], entity['offsets'][0]['start'],
                                                 entity['offsets'][0]['text'], entity['confidence']['prob'])
                                document.parts[entity['part']].predicted_annotations.append(ann)

                    # delete parts that are not annotatable
                    for part_id, part in document.parts.items():
                        if part_id not in ann_json['annotatable']['parts']:
                            del document.parts[part_id]

                except KeyError:
                    # TODO to be removed when external tagtog part_id is fixed, see issue #113
                    pass
        # some docs may not even have a corresponding .ann.json file
        # delete those as well
        if self.delete_incomplete_docs:
            dataset.documents = {doc_id: doc for doc_id, doc in dataset.documents.items()
                                 if sum(len(part.annotations) for part in doc.parts.values()) > 0}

class AnnJsonMergerAnnotationReader(AnnotationReader):
    """
    Merges annotations from several annotators.

    The scheme for merging is:
    1. Start with the annotation from the first annotator
    2. Merge the annotations of the next one, one by one

    There are several available strategies for merging:
    1. Union or intersection
    2. Shortest entity, longest entity or priority
    """
    def __init__(self, directory, strategy='union', entity_strategy='shortest', read_just_mutations=True,
                 filter_below_iaa_threshold=False, iaa_threshold=0.8):
        self.directory = directory
        """
        the directory containing several sub-directories with .ann.json files
        corresponding to the annotations of each annotator
        """
        self.strategy = strategy
        """
        the general merging strategy, can be:
        * intersection: only merge overlapping entities
        * union: merge not existing entities as well
        """
        self.entity_strategy = entity_strategy
        if self.entity_strategy == 'shortest':
            self.operator = lt
        elif self.entity_strategy == 'longest':
            self.operator = gt
        elif self.entity_strategy != 'priority':
            raise ValueError('entity_strategy must be "shortest" or "longest" or "priority"')
        """
        the merging strategy for entities when they are overlapping, can be:
        * longest: takes the longest overlapping entity
        * shortest: takes the shortest overlapping entity
        * priority: take the entity from the annotator with higher priority in the following order:
            (abojchevski, Ectelion, sanjeevkrn, Shpendi, jmcejuela, cuhlig, ANKIT)
        """
        self.priority = ['Ectelion', 'abojchevski', 'sanjeevkrn', 'Shpendi']
        self.read_just_mutations = read_just_mutations
        """whether to read in only mutation entities"""
        self.filter_below_iaa_threshold = filter_below_iaa_threshold
        self.iaa_threshold = iaa_threshold

    def annotate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        if self.entity_strategy == 'priority':
            annotators = self.priority
        else:
            annotators = os.listdir(self.directory)
        self.__merge(dataset, annotators)

    def __append_union(self, merged, entities_x, entities_y):
        # if the strategy is union
        # append the ones that are not overlapping with the already merged ones
        if self.strategy == 'union':
            existing = [Entity(entity['classId'], entity['offsets'][0]['start'], entity['offsets'][0]['text'])
                        for entity in merged]
            for entity in chain(entities_x, entities_y):
                ann = Entity(entity['classId'], entity['offsets'][0]['start'], entity['offsets'][0]['text'])
                if ann not in existing:
                    merged.append(entity)

    def __merge_priority(self, entities_x, entities_y):
        merged = []
        merged_indices_x = []
        merged_indices_y = []

        for index_x, entity_x in enumerate(entities_x):
            for index_y, entity_y in enumerate(entities_y):
                if entity_x['part'] == entity_y['part']:
                    ann_x = Entity(entity_x['classId'], entity_x['offsets'][0]['start'], entity_x['offsets'][0]['text'])
                    ann_y = Entity(entity_y['classId'], entity_y['offsets'][0]['start'], entity_y['offsets'][0]['text'])

                    # if they are the same or overlap
                    # use the first once since that one has higher priority
                    if ann_x == ann_y:
                        if index_x not in merged_indices_x and index_y not in merged_indices_y:
                            merged_indices_x.append(index_x)
                            merged_indices_y.append(index_y)
                            merged.append(entity_x)

        self.__append_union(merged, entities_x, entities_y)
        return merged

    def __merge_pair(self, entities_x, entities_y):
        merged = []
        merged_indices_x = {}
        merged_indices_y = {}

        for index_x, entity_x in enumerate(entities_x):
            for index_y, entity_y in enumerate(entities_y):
                # if they have the same part_id
                if entity_x['part'] == entity_y['part']:
                    ann_x = Entity(entity_x['classId'], entity_x['offsets'][0]['start'], entity_x['offsets'][0]['text'])
                    ann_y = Entity(entity_y['classId'], entity_y['offsets'][0]['start'], entity_y['offsets'][0]['text'])

                    # if they are the same or overlap
                    if ann_x == ann_y:
                        # if neither of them haven't been matched before
                        if index_x not in merged_indices_x and index_y not in merged_indices_y:
                            if self.operator(len(ann_x.text), len(ann_y.text)):
                                merged.append(entity_x)
                                merged_indices_x[index_x] = len(merged), ann_x
                                merged_indices_y[index_y] = len(merged), ann_x
                            else:
                                merged.append(entity_y)
                                merged_indices_x[index_x] = len(merged), ann_y
                                merged_indices_y[index_y] = len(merged), ann_y
                        # if we already matched them before
                        else:
                            # try to see if we have a more suitable match now
                            if index_x in merged_indices_x:
                                index, ann_existing = merged_indices_x[index_x]
                            else:
                                index, ann_existing = merged_indices_y[index_y]
                            if self.operator(len(ann_x.text), len(ann_y.text)):
                                ann_new, entity_new = ann_x, entity_x
                            else:
                                ann_new, entity_new = ann_y, entity_y
                            if self.operator(len(ann_new.text), len(ann_existing.text)):
                                merged[index-1] = entity_new

        self.__append_union(merged, entities_x, entities_y)
        return merged

    def __is_acceptable(self, doc_id, doc, annotators):
        if len(annotators) == 1:
            return True

        from itertools import combinations
        from nalaf.structures.data import Dataset
        from nalaf.learning.evaluators import MentionLevelEvaluator
        import math

        agreement = []
        for first, second in combinations(annotators, 2):
            data = Dataset()
            data.documents[doc_id] = doc

            AnnJsonAnnotationReader(first).annotate(data)
            AnnJsonAnnotationReader(second, is_predicted=True).annotate(data)
            results = MentionLevelEvaluator().evaluate(data)
            if not math.isnan(results[-1]):
                agreement.append(results[-1])

        # clean the doc from any annotations we added to calculate agreement
        for part in doc.parts.values():
            part.annotations = []
            part.predicted_annotations = []

        return agreement and sum(agreement)/len(agreement) >= self.iaa_threshold

    def __merge(self, dataset, annotators):
        for doc_id, doc in dataset.documents.items():
            annotator_entities = {}
            # find the annotations that are marked complete by any annotator
            filenames = []

            annotatable_parts = set()
            for annotator in annotators:
                # either once or zero times
                for filename in glob.glob(os.path.join(os.path.join(self.directory, annotator), '*{}*.ann.json'.format(doc_id))):
                    with open(filename, 'r', encoding='utf-8') as file:
                        filenames.append(filename)
                        ann_json = json.load(file)
                        if ann_json['anncomplete']:
                            annotatable_parts |= set(ann_json['annotatable']['parts'])
                            annotator_entities[annotator] = ann_json['entities']

            if self.filter_below_iaa_threshold and not self.__is_acceptable(doc_id, doc, filenames):
                del dataset.documents[doc_id]
                continue

            # if there is at least once set of annotations
            if len(annotator_entities) > 0:
                Entity.equality_operator = 'exact_or_overlapping'
                if self.entity_strategy == 'priority':
                    merged = reduce(self.__merge_priority, [annotator_entities[x] for x in self.priority
                                                            if x in annotator_entities])
                else:
                    merged = reduce(self.__merge_pair, annotator_entities.values())

                for entity in merged:
                    try:
                        part = doc.parts[entity['part']]
                    except KeyError:
                        # TODO: Remove once the tagtog bug is fixed
                        break
                    if not self.read_just_mutations or entity['classId'] == MUT_CLASS_ID:
                        part.annotations.append(
                            Entity(entity['classId'], entity['offsets'][0]['start'], entity['offsets'][0]['text']))

                # delete parts that are not annotatable
                for part_id, part in doc.parts.items():
                    if part_id not in annotatable_parts:
                        del doc.parts[part_id]

            # delete documents with no annotations
            else:
                del dataset.documents[doc_id]

class SETHAnnotationReader(AnnotationReader):
    """
    Reads the annotations from the SETH-corpus (http://rockt.github.io/SETH/)
    Format: filename = PMID
    Tab separated:
        annotation_type = T# or R# or E# (entity or relationship or ?)
        mention = space separated:
            entity_type start end
        text

        entity_type = one of the following: 'SNP', 'Gene', 'RS'

    We map:
        SNP to e_2 (mutation entity)
        Gene to e_1 (protein entity)
        RS to e_2 (mutation entity)

    Implements the abstract class Annotator.
    """

    def __init__(self, directory):
        self.directory = directory
        """the directory containing *.ann files"""

    def annotate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        for filename in glob.glob(str(self.directory + "/*.ann")):
            with open(filename, 'r', encoding='utf-8') as file:
                reader = csv.reader(file, delimiter='\t')

                pmid = os.path.basename(filename).replace('.ann', '')
                document = dataset.documents[pmid]

                for row in reader:
                    if row[0].startswith('T'):
                        entity_type, start, end = row[1].split()

                        if entity_type == 'SNP' or entity_type == 'RS':
                            ann = Entity(MUT_CLASS_ID, start, row[2])
                            document.parts['abstract'].annotations.append(ann)
                        elif entity_type == 'Gene':
                            ann = Entity('e_1', start, row[2])
                            document.parts['abstract'].annotations.append(ann)
<commit_msg>fix bug in iaa filter<commit_after>import abc
import glob
import json
import csv
import os
from itertools import chain
from functools import reduce
from operator import lt, gt

from nalaf.structures.data import Entity, Relation
from nalaf.utils import MUT_CLASS_ID

class AnnotationReader:
    """
    Abstract class for annotating the dataset.
    Subclasses that inherit this class should:
    * Be named [Name]AnnotationReader
    * Implement the abstract method annotate
    * Append new items to the list field "annotations" of each Part in the dataset
    """

    @abc.abstractmethod
    def annotate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        return

class AnnJsonAnnotationReader(AnnotationReader):
    """
    Reads the annotations from .ann.json format.

    Implements the abstract class Annotator.
    """

    def __init__(self, directory, read_just_mutations=True, delete_incomplete_docs=True, is_predicted=False, read_relations=False):
        self.directory = directory
        """the directory containing *.ann.json files"""
        self.read_just_mutations = read_just_mutations
        """whether to read in only mutation entities"""
        self.delete_incomplete_docs = delete_incomplete_docs
        """whether to delete documents from the dataset that are not marked as 'anncomplete'"""
        self.is_predicted = is_predicted
        """whether the annotation is predicted or real, which determines where it will be saved"""
        self.read_relations = read_relations
        """whether relations should be read as well"""

    def annotate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        if not os.path.isdir(self.directory):
            filenames = [self.directory]
        else:
            filenames = glob.glob(str(self.directory + "/*.ann.json"))

        for filename in filenames:
            with open(filename, 'r', encoding="utf-8") as file:
                try:
                    basename = os.path.basename(filename)
                    if '-' in basename:
                        doc_id = filename.split('-')[-1].replace('.ann.json', '')
                    else:
                        doc_id = basename.replace('.ann.json', '')

                    ann_json = json.load(file)
                    document = dataset.documents[doc_id]
                    if ann_json['anncomplete']:
                        for entity in ann_json['entities']:
                            # if read_just_mutations is False
                            if not self.read_just_mutations or entity['classId'] == MUT_CLASS_ID:
                                # if it is predicted put it in predicted_annotations else in annotations
                                if self.is_predicted:
                                    ann = Entity(entity['classId'], entity['offsets'][0]['start'],
                                                     entity['offsets'][0]['text'], entity['confidence']['prob'])
                                    document.parts[entity['part']].predicted_annotations.append(ann)
                                else:
                                    ann = Entity(entity['classId'], entity['offsets'][0]['start'],
                                                     entity['offsets'][0]['text'])
                                    document.parts[entity['part']].annotations.append(ann)
                        if self.read_relations:
                            for relation in ann_json['relations']:
                                part = document.parts[relation['entities'][0].split('|')[0]]
                                e1_start = int(relation['entities'][0].split('|')[1].split(',')[0])
                                e2_start = int(relation['entities'][1].split('|')[1].split(',')[0])
                                e1_end = int(relation['entities'][0].split('|')[1].split(',')[1])
                                e2_end = int(relation['entities'][1].split('|')[1].split(',')[1])
                                e1_text = part.text[e1_start:e1_end]
                                e2_text = part.text[e2_start:e2_end]
                                part.relations.append(Relation(e1_start, e2_start, e1_text, e2_text, relation['classId']))
                    elif self.delete_incomplete_docs:
                        del dataset.documents[doc_id]
                    elif self.is_predicted:  # anncomplete=False, delete_incomplete_docs=False, is_predicted=True
                        for entity in ann_json['entities']:
                            if not self.read_just_mutations or entity['classId'] == MUT_CLASS_ID:
                                ann = Entity(entity['classId'], entity['offsets'][0]['start'],
                                                 entity['offsets'][0]['text'], entity['confidence']['prob'])
                                document.parts[entity['part']].predicted_annotations.append(ann)

                    # delete parts that are not annotatable
                    for part_id, part in document.parts.items():
                        if part_id not in ann_json['annotatable']['parts']:
                            del document.parts[part_id]

                except KeyError:
                    # TODO to be removed when external tagtog part_id is fixed, see issue #113
                    pass
        # some docs may not even have a corresponding .ann.json file
        # delete those as well
        if self.delete_incomplete_docs:
            dataset.documents = {doc_id: doc for doc_id, doc in dataset.documents.items()
                                 if sum(len(part.annotations) for part in doc.parts.values()) > 0}

class AnnJsonMergerAnnotationReader(AnnotationReader):
    """
    Merges annotations from several annotators.

    The scheme for merging is:
    1. Start with the annotation from the first annotator
    2. Merge the annotations of the next one, one by one

    There are several available strategies for merging:
    1. Union or intersection
    2. Shortest entity, longest entity or priority
    """
    def __init__(self, directory, strategy='union', entity_strategy='shortest', read_just_mutations=True,
                 filter_below_iaa_threshold=False, iaa_threshold=0.8):
        self.directory = directory
        """
        the directory containing several sub-directories with .ann.json files
        corresponding to the annotations of each annotator
        """
        self.strategy = strategy
        """
        the general merging strategy, can be:
        * intersection: only merge overlapping entities
        * union: merge not existing entities as well
        """
        self.entity_strategy = entity_strategy
        if self.entity_strategy == 'shortest':
            self.operator = lt
        elif self.entity_strategy == 'longest':
            self.operator = gt
        elif self.entity_strategy != 'priority':
            raise ValueError('entity_strategy must be "shortest" or "longest" or "priority"')
        """
        the merging strategy for entities when they are overlapping, can be:
        * longest: takes the longest overlapping entity
        * shortest: takes the shortest overlapping entity
        * priority: take the entity from the annotator with higher priority in the following order:
            (abojchevski, Ectelion, sanjeevkrn, Shpendi, jmcejuela, cuhlig, ANKIT)
        """
        self.priority = ['Ectelion', 'abojchevski', 'sanjeevkrn', 'Shpendi']
        self.read_just_mutations = read_just_mutations
        """whether to read in only mutation entities"""
        self.filter_below_iaa_threshold = filter_below_iaa_threshold
        self.iaa_threshold = iaa_threshold

    def annotate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        if self.entity_strategy == 'priority':
            annotators = self.priority
        else:
            annotators = os.listdir(self.directory)
        self.__merge(dataset, annotators)

    def __append_union(self, merged, entities_x, entities_y):
        # if the strategy is union
        # append the ones that are not overlapping with the already merged ones
        if self.strategy == 'union':
            existing = [Entity(entity['classId'], entity['offsets'][0]['start'], entity['offsets'][0]['text'])
                        for entity in merged]
            for entity in chain(entities_x, entities_y):
                ann = Entity(entity['classId'], entity['offsets'][0]['start'], entity['offsets'][0]['text'])
                if ann not in existing:
                    merged.append(entity)

    def __merge_priority(self, entities_x, entities_y):
        merged = []
        merged_indices_x = []
        merged_indices_y = []

        for index_x, entity_x in enumerate(entities_x):
            for index_y, entity_y in enumerate(entities_y):
                if entity_x['part'] == entity_y['part']:
                    ann_x = Entity(entity_x['classId'], entity_x['offsets'][0]['start'], entity_x['offsets'][0]['text'])
                    ann_y = Entity(entity_y['classId'], entity_y['offsets'][0]['start'], entity_y['offsets'][0]['text'])

                    # if they are the same or overlap
                    # use the first once since that one has higher priority
                    if ann_x == ann_y:
                        if index_x not in merged_indices_x and index_y not in merged_indices_y:
                            merged_indices_x.append(index_x)
                            merged_indices_y.append(index_y)
                            merged.append(entity_x)

        self.__append_union(merged, entities_x, entities_y)
        return merged

    def __merge_pair(self, entities_x, entities_y):
        merged = []
        merged_indices_x = {}
        merged_indices_y = {}

        for index_x, entity_x in enumerate(entities_x):
            for index_y, entity_y in enumerate(entities_y):
                # if they have the same part_id
                if entity_x['part'] == entity_y['part']:
                    ann_x = Entity(entity_x['classId'], entity_x['offsets'][0]['start'], entity_x['offsets'][0]['text'])
                    ann_y = Entity(entity_y['classId'], entity_y['offsets'][0]['start'], entity_y['offsets'][0]['text'])

                    # if they are the same or overlap
                    if ann_x == ann_y:
                        # if neither of them haven't been matched before
                        if index_x not in merged_indices_x and index_y not in merged_indices_y:
                            if self.operator(len(ann_x.text), len(ann_y.text)):
                                merged.append(entity_x)
                                merged_indices_x[index_x] = len(merged), ann_x
                                merged_indices_y[index_y] = len(merged), ann_x
                            else:
                                merged.append(entity_y)
                                merged_indices_x[index_x] = len(merged), ann_y
                                merged_indices_y[index_y] = len(merged), ann_y
                        # if we already matched them before
                        else:
                            # try to see if we have a more suitable match now
                            if index_x in merged_indices_x:
                                index, ann_existing = merged_indices_x[index_x]
                            else:
                                index, ann_existing = merged_indices_y[index_y]
                            if self.operator(len(ann_x.text), len(ann_y.text)):
                                ann_new, entity_new = ann_x, entity_x
                            else:
                                ann_new, entity_new = ann_y, entity_y
                            if self.operator(len(ann_new.text), len(ann_existing.text)):
                                merged[index-1] = entity_new

        self.__append_union(merged, entities_x, entities_y)
        return merged

    def __is_acceptable(self, doc_id, doc, annotators):
        if len(annotators) == 1:
            return True

        from itertools import combinations
        from nalaf.structures.data import Dataset
        from nalaf.learning.evaluators import MentionLevelEvaluator
        import math

        agreement = []
        for first, second in combinations(annotators, 2):
            data = Dataset()
            data.documents[doc_id] = doc

            AnnJsonAnnotationReader(first).annotate(data)
            AnnJsonAnnotationReader(second, is_predicted=True).annotate(data)
            results = MentionLevelEvaluator().evaluate(data)
            if not math.isnan(results[-1]):
                agreement.append(results[-1])

        # clean the doc from any annotations we added to calculate agreement
        for part in doc.parts.values():
            part.annotations = []
            part.predicted_annotations = []

        return agreement and sum(agreement)/len(agreement) >= self.iaa_threshold

    def __merge(self, dataset, annotators):
        for doc_id, doc in dataset.documents.items():
            annotator_entities = {}
            # find the annotations that are marked complete by any annotator
            filenames = []

            annotatable_parts = set()
            for annotator in annotators:
                # either once or zero times
                for filename in glob.glob(os.path.join(os.path.join(self.directory, annotator), '*{}*.ann.json'.format(doc_id))):
                    with open(filename, 'r', encoding='utf-8') as file:
                        ann_json = json.load(file)
                        if ann_json['anncomplete']:
                            filenames.append(filename)
                            annotatable_parts |= set(ann_json['annotatable']['parts'])
                            annotator_entities[annotator] = ann_json['entities']

            if self.filter_below_iaa_threshold and not self.__is_acceptable(doc_id, doc, filenames):
                del dataset.documents[doc_id]
                continue

            # if there is at least once set of annotations
            if len(annotator_entities) > 0:
                Entity.equality_operator = 'exact_or_overlapping'
                if self.entity_strategy == 'priority':
                    merged = reduce(self.__merge_priority, [annotator_entities[x] for x in self.priority
                                                            if x in annotator_entities])
                else:
                    merged = reduce(self.__merge_pair, annotator_entities.values())

                for entity in merged:
                    try:
                        part = doc.parts[entity['part']]
                    except KeyError:
                        # TODO: Remove once the tagtog bug is fixed
                        break
                    if not self.read_just_mutations or entity['classId'] == MUT_CLASS_ID:
                        part.annotations.append(
                            Entity(entity['classId'], entity['offsets'][0]['start'], entity['offsets'][0]['text']))

                # delete parts that are not annotatable
                for part_id, part in doc.parts.items():
                    if part_id not in annotatable_parts:
                        del doc.parts[part_id]

            # delete documents with no annotations
            else:
                del dataset.documents[doc_id]

class SETHAnnotationReader(AnnotationReader):
    """
    Reads the annotations from the SETH-corpus (http://rockt.github.io/SETH/)
    Format: filename = PMID
    Tab separated:
        annotation_type = T# or R# or E# (entity or relationship or ?)
        mention = space separated:
            entity_type start end
        text

        entity_type = one of the following: 'SNP', 'Gene', 'RS'

    We map:
        SNP to e_2 (mutation entity)
        Gene to e_1 (protein entity)
        RS to e_2 (mutation entity)

    Implements the abstract class Annotator.
    """

    def __init__(self, directory):
        self.directory = directory
        """the directory containing *.ann files"""

    def annotate(self, dataset):
        """
        :type dataset: nalaf.structures.data.Dataset
        """
        for filename in glob.glob(str(self.directory + "/*.ann")):
            with open(filename, 'r', encoding='utf-8') as file:
                reader = csv.reader(file, delimiter='\t')

                pmid = os.path.basename(filename).replace('.ann', '')
                document = dataset.documents[pmid]

                for row in reader:
                    if row[0].startswith('T'):
                        entity_type, start, end = row[1].split()

                        if entity_type == 'SNP' or entity_type == 'RS':
                            ann = Entity(MUT_CLASS_ID, start, row[2])
                            document.parts['abstract'].annotations.append(ann)
                        elif entity_type == 'Gene':
                            ann = Entity('e_1', start, row[2])
                            document.parts['abstract'].annotations.append(ann)
<|endoftext|>


<commit_before># local imports
from nautilus.network import CRUDHandler, combineActionHandlers
from nautilus.api import create_model_schema
from nautilus.conventions.services import model_service_name
from nautilus.network.actionHandlers import noop_handler
from .service import Service

class ModelService(Service):
    """
        This service acts as the primary data store in your cloud. It manages
        the records of a single model by listening for actions that indicate
        a record mutation as well as emitting actions when the mutations are
        successful. The external API is automatically generated to match the
        given model.

        Args:
            model (nautilus.BaseModel): The nautilus model to manage.
            additonal_action_handler (optional, function): An action handler
                to be called alongside the internal ones.
    """

    def __init__(self, model, additonal_action_handler = noop_handler, **kwargs):
        # save a reference to the model this service is managing
        self.model = model

        # the schema to add to the service
        schema = create_model_schema(model)

        # the action handler is a combination
        action_handler = combineActionHandlers(
            # of the given one
            additonal_action_handler,
            # and a crud handler
            CRUDHandler(model)
        )

        name = kwargs.pop('name', None) or model_service_name(model)

        # create the service
        super().__init__(
            schema = schema,
            actionHandler = action_handler,
            name = name,
            **kwargs
        )
<commit_msg>updated model service doc string<commit_after># local imports
from nautilus.network import CRUDHandler, combineActionHandlers
from nautilus.api import create_model_schema
from nautilus.conventions.services import model_service_name
from nautilus.network.actionHandlers import noop_handler
from .service import Service

class ModelService(Service):
    """
        This service acts as the primary data store in your cloud. It manages
        the records of a single model by listening for actions that indicate
        a record mutation as well as emitting actions when the mutations have
        finished (whether successfully or not). The external API is
        automatically generated to match the given model.

        Args:
            model (nautilus.BaseModel): The nautilus model to manage.
            additonal_action_handler (optional, function): An action handler
                to be called alongside the internal ones.
    """

    def __init__(self, model, additonal_action_handler = noop_handler, **kwargs):
        # save a reference to the model this service is managing
        self.model = model

        # the schema to add to the service
        schema = create_model_schema(model)

        # the action handler is a combination
        action_handler = combineActionHandlers(
            # of the given one
            additonal_action_handler,
            # and a crud handler
            CRUDHandler(model)
        )

        name = kwargs.pop('name', None) or model_service_name(model)

        # create the service
        super().__init__(
            schema = schema,
            actionHandler = action_handler,
            name = name,
            **kwargs
        )
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Copyright 2019 NeuroData (http://neurodata.io)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# ndmg_dwi_pipeline.py
# Repackaged for native space tractography by Derek Pisner
# Email: dpisner@utexas.edu
# Originally created by Greg Kiar and Will Gray Roncal on 2016-01-27.

from __future__ import print_function
import shutil
import warnings
warnings.filterwarnings("ignore", message="numpy.dtype size changed")
from argparse import ArgumentParser
from datetime import datetime
#from ndmg.stats.qa_regdti import *
import time
from ndmg.stats.qa_tensor import *
from ndmg.stats.qa_fibers import *
from ndmg.stats.qa_mri import qa_mri
from ndmg.utils import gen_utils as mgu
from ndmg.register import gen_reg as mgr
from ndmg.track import gen_track as mgt
from ndmg.graph import gen_graph as mgg
import ndmg.preproc as mgp
import numpy as np
import nibabel as nib
import os
from ndmg.graph import biggraph as ndbg
import traceback
from ndmg.utils.bids_utils import name_resource
import sys
from dipy.io.dpy import Dpy

os.environ["MPLCONFIGDIR"] = "/tmp/"

def ndmg_dwi_worker(dwi, bvals, bvecs, t1w, atlas, mask, labels, outdir,
                    clean=False, big=False):
    """
    Creates a brain graph from MRI data
    """
    startTime = datetime.now()
    fmt = '_adj.csv'
    # Create derivative output directories
    namer = name_resource(dwi, t1w, atlas, outdir)

    paths = {'prep_m': "dwi/preproc",
             'prep_a': "anat/preproc",
             'reg_m': "dwi/registered",
             'reg_a': "anat/registered",
             'tensor': "dwi/tensor",
             'fiber': "dwi/fiber",
             'voxelg': "dwi/voxel-connectomes",
             'conn': "dwi/roi-connectomes"}

    opt_dirs = ['prep_m', 'prep_a', 'reg_m', 'reg_a']
    clean_dirs = ['tensor', 'fiber']
    label_dirs = ['conn']  # create label level granularity

    namer.add_dirs(paths, labels, label_dirs)
    qc_stats = "{}/{}_stats.csv".format(namer.dirs['qa']['base'],
        namer.get_mod_source())

    # Create derivative output file names
    reg_dname = "{}_{}".format(namer.get_mod_source(),
        namer.get_template_info())
    reg_aname = "{}_{}".format(namer.get_anat_source(),
        namer.get_template_info())
    streams = namer.name_derivative(namer.dirs['output']['fiber'],
        "{}_streamlines.dpy".format(reg_dname))

    if big:
        voxel = namer.name_derivative(namer.dirs['output']['voxel'],
            "{}_voxel-connectome.npz".format(reg_dname))
        print("Voxelwise Fiber Graph: {}".format(voxel))

    # Again, connectomes are different
    if not isinstance(labels, list):
        labels = [labels]
    connectomes = [namer.name_derivative(
        namer.dirs['output']['conn'][namer.get_label(lab)],
        "{}_{}_measure-spatial-ds.{}".format(namer.get_mod_source(),
            namer.get_label(lab), fmt)) for lab in labels]

    print("Connectomes downsampled to given labels: " +
          ", ".join(connectomes))

    qc_dwi = qa_mri(namer, 'dwi')  # for quality control
    # -------- Preprocessing Steps --------------------------------- #
    # Perform eddy correction
    start_time = time.time()
    dwi_prep = "{}/eddy_corrected_data.nii.gz".format(namer.dirs['output']['prep_m'])
    cmd='eddy_correct ' + dwi + ' ' + dwi_prep + ' 0'
#    os.system(cmd)

    print("Rotating b-vectors and generating gradient table...")
    eddy_rot_param = "{}/eddy_corrected_data.ecclog".format(namer.dirs['output']['prep_m'])
    bvec_scaled = "{}/bvec_scaled.bvec".format(namer.dirs['output']['prep_m'])
    bvec_rotated = "{}/bvec_rotated.bvec".format(namer.dirs['output']['prep_m'])
    bval = "{}/bval.bval".format(namer.dirs['output']['prep_m'])
#    shutil.copyfile(bvals, bval)

    # Rotate bvecs
    cmd='bash fdt_rotate_bvecs ' + bvecs + ' ' + bvec_rotated + ' ' + eddy_rot_param
#    os.system(cmd)

    # Rescale bvecs
#    mgp.rescale_bvec(bvec_rotated, bvec_scaled)

    [gtab, nodif_B0, nodif_B0_mask] = mgu.make_gtab_and_bmask(bvals, bvec_scaled, dwi_prep, namer.dirs['output']['prep_m'])
    print("%s%s%s" % ('Preprocessing runtime: ', str(np.round(time.time() - start_time, 1)), 's'))
    # -------- Registration Steps ----------------------------------- #
    vox_size = '2mm'
    reg = mgr.dmri_reg(outdir, nodif_B0, nodif_B0_mask, t1w, vox_size, simple=False)
    # Perform anatomical segmentation
    start_time = time.time()
#    reg.gen_tissue()
    print("%s%s%s" % ('gen_tissue runtime: ', str(np.round(time.time() - start_time, 1)), 's'))
    # align t1w to dwi
    start_time = time.time()
#    reg.t1w2dwi_align()
    print("%s%s%s" % ('t1w2dwi_align runtime: ', str(np.round(time.time() - start_time, 1)), 's'))
    # align tissue classifiers
    start_time = time.time()
#    reg.tissue2dwi_align()
    print("%s%s%s" % ('tissue2dwi_align runtime: ', str(np.round(time.time() - start_time, 1)), 's'))

    # -------- Tensor Fitting and Fiber Tractography ---------------- #
    # Compute tensors and track fiber streamlines
    print("Beginning tractography...")
    trct = mgt.run_track(dwi_prep, nodif_B0_mask, reg.gm_in_dwi, reg.vent_csf_in_dwi, reg.wm_in_dwi, reg.wm_in_dwi_bin, gtab)
    [streamlines, tens] = trct.run()

    tracks = [sl for sl in streamlines if len(sl) > 1]

    # Save streamlines to disk
    print('Streamlines: ' + streams)
    dpw = Dpy(streams, 'w')
    dpw.write_tracks(streamlines)
    dpw.close()

    # As we've only tested VTK plotting on MNI152 aligned data...
    if nib.load(mask).get_data().shape == (182, 218, 182):
        try:
            visualize_fibs(streamlines, streams, aligned_atlas, namer.dirs['qa']['fiber'], 0.02, 2000)
        except:
            print("Fiber QA failed - VTK for Python not configured properly.")

    # -------- Big Graph Generation --------------------------------- #
    # Generate big graphs from streamlines
    if big:
        print("Making Voxelwise Graph...")
        bg1 = ndbg()
        bg1.make_graph(tracks)
        bg1.save_graph(voxel)

    # ------- Connectome Estimation --------------------------------- #
    # Generate graphs from streamlines for each parcellation
    for idx, label in enumerate(labels):
        print("Generating graph for {} parcellation...".format(label))
	print("%s%s" % ('Applying native-space alignment to ', labels[idx]))
	# align atlas to t1w to dwi
        labels_im_file = reg.atlas2t1w2dwi_align(labels[idx])
	print('Aligned Atlas: ' + labels_im_file)
	labels_im = nib.load(labels_im_file)
	g1 = mgg.graph_tools(attr=len(np.unique(labels_im.get_data()))-1, rois=labels[idx], streamlines=streamlines)
        g1.make_graph()
        g1.summary()
        g1.save_graph(connectomes[idx])

    exe_time = datetime.now() - startTime
    qc_dwi.save(qc_stats, exe_time)

    # Clean temp files
    if clean:
        print("Cleaning up intermediate files... ")
        del_dirs = [namer.dirs['tmp']['base']] + \
            [namer.dirs['output'][k] for k in opt_dirs]
        cmd = "rm -rf {}".format(" ".format(del_dirs))
        mgu.execute_cmd(cmd)

    print("Execution took: {}".format(exe_time))
    print("Complete!")
    sys.exit(0)

def ndmg_dwi_pipeline(dwi, bvals, bvecs, t1w, atlas, mask, labels, outdir,
                      clean=False, big=False):
    """
    A wrapper for the worker to make our pipeline more robust to errors.
    """
    try:
        ndmg_dwi_worker(dwi, bvals, bvecs, t1w, atlas, mask, labels, outdir,
                        clean, big)
    except Exception, e:
        print(traceback.format_exc())
        os.exit()
    finally:
        try:
            os.exit()
        except Exception, e:
            os.exit()
    return

def main():
    parser = ArgumentParser(description="This is an end-to-end connectome \
                            estimation pipeline from sMRI and DTI images")
    parser.add_argument("dwi", action="store", help="Nifti DTI image stack")
    parser.add_argument("bval", action="store", help="DTI scanner b-values")
    parser.add_argument("bvec", action="store", help="DTI scanner b-vectors")
    parser.add_argument("t1w", action="store", help="Nifti T1w MRI image")
    parser.add_argument("atlas", action="store", help="Nifti T1 MRI atlas")
    parser.add_argument("mask", action="store", help="Nifti binary mask of \
                        brain space in the atlas")
    parser.add_argument("outdir", action="store", help="Path to which \
                        derivatives will be stored")
    parser.add_argument("labels", action="store", nargs="*", help="Nifti \
                        labels of regions of interest in atlas space")
    parser.add_argument("-c", "--clean", action="store_true", default=False,
                        help="Whether or not to delete intemediates")
    parser.add_argument("-b", "--big", action="store_true", default=False,
                        help="whether or not to produce voxelwise big graph")
    result = parser.parse_args()

    # Create output directory
    print("Creating output directory: {}".format(result.outdir))
    print("Creating output temp directory: {}/tmp".format(result.outdir))
    mgu.utils.execute_cmd("mkdir -p {} {}/tmp".format(result.outdir, result.outdir))

    ndmg_dwi_pipeline(result.dwi, result.bval, result.bvec, result.t1w,
                      result.atlas, result.mask, result.labels, result.outdir,
                      result.clean, result.big)

if __name__ == "__main__":
    main()
<commit_msg>pipeline debugging, Dcoker corr, hierarchy cleanup<commit_after>#!/usr/bin/env python
# Copyright 2019 NeuroData (http://neurodata.io)
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#

# ndmg_dwi_pipeline.py
# Repackaged for native space tractography by Derek Pisner
# Email: dpisner@utexas.edu
# Originally created by Greg Kiar and Will Gray Roncal on 2016-01-27.

from __future__ import print_function
import shutil
import warnings
warnings.filterwarnings("ignore", message="numpy.dtype size changed")
from argparse import ArgumentParser
from datetime import datetime
#from ndmg.stats.qa_regdti import *
import time
from ndmg.stats.qa_tensor import *
from ndmg.stats.qa_fibers import *
from ndmg.stats.qa_mri import qa_mri
from ndmg.utils import gen_utils as mgu
from ndmg.register import gen_reg as mgr
from ndmg.track import gen_track as mgt
from ndmg.graph import gen_graph as mgg
import ndmg.preproc as mgp
import numpy as np
import nibabel as nib
import os
from ndmg.graph import biggraph as ndbg
import traceback
from ndmg.utils.bids_utils import name_resource
import sys
from dipy.io.dpy import Dpy

os.environ["MPLCONFIGDIR"] = "/tmp/"

def ndmg_dwi_worker(dwi, bvals, bvecs, t1w, atlas, mask, labels, outdir,
                    clean=False, big=False):
    """
    Creates a brain graph from MRI data
    """
    startTime = datetime.now()
    fmt = '_adj.csv'
    # Create derivative output directories
    namer = name_resource(dwi, t1w, atlas, outdir)

    paths = {'prep_m': "dwi/preproc",
             'prep_a': "anat/preproc",
             'reg_m': "dwi/registered",
             'reg_a': "anat/registered",
             'tensor': "dwi/tensor",
             'fiber': "dwi/fiber",
             'voxelg': "dwi/voxel-connectomes",
             'conn': "dwi/roi-connectomes"}

    opt_dirs = ['prep_m', 'prep_a', 'reg_m', 'reg_a']
    clean_dirs = ['tensor', 'fiber']
    label_dirs = ['conn']  # create label level granularity

    namer.add_dirs(paths, labels, label_dirs)
    qc_stats = "{}/{}_stats.csv".format(namer.dirs['qa']['base'],
        namer.get_mod_source())

    # Create derivative output file names
    reg_dname = "{}_{}".format(namer.get_mod_source(),
        namer.get_template_info())
    reg_aname = "{}_{}".format(namer.get_anat_source(),
        namer.get_template_info())
    streams = namer.name_derivative(namer.dirs['output']['fiber'],
        "{}_streamlines.dpy".format(reg_dname))

    if big:
        voxel = namer.name_derivative(namer.dirs['output']['voxel'],
            "{}_voxel-connectome.npz".format(reg_dname))
        print("Voxelwise Fiber Graph: {}".format(voxel))

    # Again, connectomes are different
    if not isinstance(labels, list):
        labels = [labels]
    connectomes = [namer.name_derivative(
        namer.dirs['output']['conn'][namer.get_label(lab)],
        "{}_{}_measure-spatial-ds.{}".format(namer.get_mod_source(),
            namer.get_label(lab), fmt)) for lab in labels]

    print("Connectomes downsampled to given labels: " +
          ", ".join(connectomes))

    qc_dwi = qa_mri(namer, 'dwi')  # for quality control
    # -------- Preprocessing Steps --------------------------------- #
    # Perform eddy correction
    start_time = time.time()
    dwi_prep = "{}/eddy_corrected_data.nii.gz".format(namer.dirs['output']['prep_m'])
    cmd='eddy_correct ' + dwi + ' ' + dwi_prep + ' 0'
    os.system(cmd)

    print("Rotating b-vectors and generating gradient table...")
    eddy_rot_param = "{}/eddy_corrected_data.ecclog".format(namer.dirs['output']['prep_m'])
    bvec_scaled = "{}/bvec_scaled.bvec".format(namer.dirs['output']['prep_m'])
    bvec_rotated = "{}/bvec_rotated.bvec".format(namer.dirs['output']['prep_m'])
    bval = "{}/bval.bval".format(namer.dirs['output']['prep_m'])
    shutil.copyfile(bvals, bval)

    # Rotate bvecs
    cmd='bash fdt_rotate_bvecs ' + bvecs + ' ' + bvec_rotated + ' ' + eddy_rot_param
    os.system(cmd)

    # Rescale bvecs
    mgp.rescale_bvec(bvec_rotated, bvec_scaled)

    [gtab, nodif_B0, nodif_B0_mask] = mgu.make_gtab_and_bmask(bvals, bvec_scaled, dwi_prep, namer.dirs['output']['prep_m'])
    print("%s%s%s" % ('Preprocessing runtime: ', str(np.round(time.time() - start_time, 1)), 's'))
    # -------- Registration Steps ----------------------------------- #
    vox_size = '2mm'
    reg = mgr.dmri_reg(outdir, nodif_B0, nodif_B0_mask, t1w, vox_size, simple=False)
    # Perform anatomical segmentation
    start_time = time.time()
    reg.gen_tissue()
    print("%s%s%s" % ('gen_tissue runtime: ', str(np.round(time.time() - start_time, 1)), 's'))
    # align t1w to dwi
    start_time = time.time()
    reg.t1w2dwi_align()
    print("%s%s%s" % ('t1w2dwi_align runtime: ', str(np.round(time.time() - start_time, 1)), 's'))
    # align tissue classifiers
    start_time = time.time()
    reg.tissue2dwi_align()
    print("%s%s%s" % ('tissue2dwi_align runtime: ', str(np.round(time.time() - start_time, 1)), 's'))

    # -------- Tensor Fitting and Fiber Tractography ---------------- #
    # Compute tensors and track fiber streamlines
    print("Beginning tractography...")
    trct = mgt.run_track(dwi_prep, nodif_B0_mask, reg.gm_in_dwi, reg.vent_csf_in_dwi, reg.wm_in_dwi, reg.wm_in_dwi_bin, gtab)
    [streamlines, tens] = trct.run()

    tracks = [sl for sl in streamlines if len(sl) > 1]

    # Save streamlines to disk
    print('Streamlines: ' + streams)
    dpw = Dpy(streams, 'w')
    dpw.write_tracks(streamlines)
    dpw.close()

    # As we've only tested VTK plotting on MNI152 aligned data...
    if nib.load(mask).get_data().shape == (182, 218, 182):
        try:
            visualize_fibs(streamlines, streams, aligned_atlas, namer.dirs['qa']['fiber'], 0.02, 2000)
        except:
            print("Fiber QA failed - VTK for Python not configured properly.")

    # -------- Big Graph Generation --------------------------------- #
    # Generate big graphs from streamlines
    if big:
        print("Making Voxelwise Graph...")
        bg1 = ndbg()
        bg1.make_graph(tracks)
        bg1.save_graph(voxel)

    # ------- Connectome Estimation --------------------------------- #
    # Generate graphs from streamlines for each parcellation
    for idx, label in enumerate(labels):
        print("Generating graph for {} parcellation...".format(label))
	print("%s%s" % ('Applying native-space alignment to ', labels[idx]))
	# align atlas to t1w to dwi
        labels_im_file = reg.atlas2t1w2dwi_align(labels[idx])
	print('Aligned Atlas: ' + labels_im_file)
	labels_im = nib.load(labels_im_file)
	g1 = mgg.graph_tools(attr=len(np.unique(labels_im.get_data()))-1, rois=labels[idx], streamlines=streamlines)
        g1.make_graph()
        g1.summary()
        g1.save_graph(connectomes[idx])

    exe_time = datetime.now() - startTime
    qc_dwi.save(qc_stats, exe_time)

    # Clean temp files
    if clean:
        print("Cleaning up intermediate files... ")
        del_dirs = [namer.dirs['tmp']['base']] + \
            [namer.dirs['output'][k] for k in opt_dirs]
        cmd = "rm -rf {}".format(" ".format(del_dirs))
        mgu.execute_cmd(cmd)

    print("Execution took: {}".format(exe_time))
    print("Complete!")
    sys.exit(0)

def ndmg_dwi_pipeline(dwi, bvals, bvecs, t1w, atlas, mask, labels, outdir,
                      clean=False, big=False):
    """
    A wrapper for the worker to make our pipeline more robust to errors.
    """
    try:
        ndmg_dwi_worker(dwi, bvals, bvecs, t1w, atlas, mask, labels, outdir,
                        clean, big)
    except Exception, e:
        print(traceback.format_exc())
        os.exit()
    finally:
        try:
            os.exit()
        except Exception, e:
            os.exit()
    return

def main():
    parser = ArgumentParser(description="This is an end-to-end connectome \
                            estimation pipeline from sMRI and DTI images")
    parser.add_argument("dwi", action="store", help="Nifti DTI image stack")
    parser.add_argument("bval", action="store", help="DTI scanner b-values")
    parser.add_argument("bvec", action="store", help="DTI scanner b-vectors")
    parser.add_argument("t1w", action="store", help="Nifti T1w MRI image")
    parser.add_argument("atlas", action="store", help="Nifti T1 MRI atlas")
    parser.add_argument("mask", action="store", help="Nifti binary mask of \
                        brain space in the atlas")
    parser.add_argument("outdir", action="store", help="Path to which \
                        derivatives will be stored")
    parser.add_argument("labels", action="store", nargs="*", help="Nifti \
                        labels of regions of interest in atlas space")
    parser.add_argument("-c", "--clean", action="store_true", default=False,
                        help="Whether or not to delete intemediates")
    parser.add_argument("-b", "--big", action="store_true", default=False,
                        help="whether or not to produce voxelwise big graph")
    result = parser.parse_args()

    # Create output directory
    print("Creating output directory: {}".format(result.outdir))
    print("Creating output temp directory: {}/tmp".format(result.outdir))
    mgu.utils.execute_cmd("mkdir -p {} {}/tmp".format(result.outdir, result.outdir))

    ndmg_dwi_pipeline(result.dwi, result.bval, result.bvec, result.t1w,
                      result.atlas, result.mask, result.labels, result.outdir,
                      result.clean, result.big)

if __name__ == "__main__":
    main()
<|endoftext|>


<commit_before># This script is used to import Northern Ireland areas, from the data in the
# National Statistics Postcode Database.
# http://www.ons.gov.uk/about-statistics/geography/products/geog-products-postcode/nspd/
#
# XXX This is incomplete, it needs to know which things have had boundary changes
# like import_boundary_line does. Hopefully just using new GSS codes by the time
# NI has any boundary changes.
#
# The fields of NSPD Open CSV file are:
#   Postcode (7), Postcode (8), Postcode (sp), Start date, End date, County,
#   council, ward, usertype, Easting, Northing, quality, SHA, IT cluster,
#   country, GOR, Stats region, Parliamentary constituency, Euro region,
#   TEC/LEC, Travel to Work area, Primary Care Org, NUTS, 1991 census ED,
#   1991 census ED, ED indicator, Pre-July 2006 SHA, LEA, Pre 2002 Health
#   Authority, 1991 ward code, 1991 ward code, 1998 ward code, 2005 stats ward,
#   OA code, OA indicator, CAS ward, National Park, SOA (Lower), Datazone, SOA
#   (Middle), Urban/rural, Urban/rural, Urban/rural, Intermediate, SOA (NI), OA
#   classification, Pre October 2006 PCO

import csv
from django.contrib.gis.geos import Point
from django.core.management.base import LabelCommand
from mapit.postcodes.models import Postcode
from mapit.areas.models import Area, Generation

class Command(LabelCommand):
    help = 'Creates/updates Northern Ireland areas, using the NSPD'
    args = '<NSPD CSV file>'

    def handle_label(self, file, **options):
        current_generation = Generation.objects.current()
        new_generation = Generation.objects.new()
        if not new_generation:
            raise Exception, "No new generation to be used for import!"

        euro_area, created = Area.objects.get_or_create(country='N', type='EUR',
            generation_low__lte=current_generation, generation_high__gte=current_generation,
            defaults = { 'generation_low': new_generation, 'generation_high': new_generation }
        )
        euro_area.generation_high = new_generation
        euro_area.save()
        euro_area.names.get_or_create(type='S', name='Northern Ireland')

        # Read in ward name -> electoral area name/area
        ni_eas = csv.reader(open('../../data/ni-electoral-areas.csv'))
        ni_eas.next()
        ward_to_electoral_area = {}
        e = {}
        for district, electoral_area, ward, dummy in ni_eas:
            if not district:
                district = last_district
            if not electoral_area:
                electoral_area = last_electoral_area
            last_district = district
            last_electoral_area = electoral_area
            if electoral_area not in e:
                ea = Area.objects.get_or_create_with_name(
                    country='N', type='LGE', name_type='M', name=electoral_area,
                )
                e[electoral_area] = ea
            ward_to_electoral_area.setdefault(district, {})[ward] = e[electoral_area]

        # Read in new ONS code to names
        snac = csv.reader(open('../../data/snac-2009-ni-cons2ward.csv'))
        snac.next()
        code_to_area = {}
        name_to_area = {}
        ward_to_parl = {}
        for parl_code, parl_name, ward_code, ward_name, district_code, district_name in snac:
            if district_name not in ward_to_electoral_area:
                raise Exception, "District %s is missing" % district_name
            if ward_name not in ward_to_electoral_area[district_name]:
                raise Exception, "Ward %s, district %s is missing" % (ward_name, district_name)

            ward_code = ward_code.replace(' ', '')

            if district_code not in code_to_area:
                district_area = Area.objects.get_or_create_with_code(
                    country='N', type='LGD', code_type='ons', code=district_code,
                )
                district_area.names.get_or_create(type='S', name=district_name)
                code_to_area[district_code] = district_area

            if ward_code not in code_to_area:
                ward_area = Area.objects.get_or_create_with_code(
                    country='N', type='LGW', code_type='ons', code=ward_code
                )
                ward_area.names.get_or_create(type='S', name=ward_name)
                ward_area.parent_area = ward_to_electoral_area[district_name][ward_name]
                ward_area.save()
                ward_area.parent_area.parent_area = code_to_area[district_code]
                ward_area.parent_area.save()
                code_to_area[ward_code] = ward_area

            if ward_code == '95S24': continue # Derryaghy

            if parl_code not in code_to_area:
                parl_area = Area.objects.get_or_create_with_code(
                    country='N', type='WMC', code_type='ons', code=parl_code,
                )
                parl_area.names.get_or_create(type='S', name=parl_name)
                code_to_area[parl_code] = parl_area
                name_to_area[parl_name] = parl_area

            ward_to_parl[ward_code] = code_to_area[parl_code]

        # Read in old SNAC for NI Assembly boundaries, still the same until 2011
        snac = csv.reader(open('../../data/snac-2003-ni-cons2ward.csv'))
        snac.next()
        ward_to_assembly = {}
        for parl_code, parl_name, ward_code, ward_name, district_code, district_name in snac:
            ward_code = ward_code.replace(' ', '')
            if 'NIE' + parl_code not in code_to_area:
                nia_area = Area.objects.get_or_create_with_name(
                    country='N', type='NIE', name_type='S', name=parl_name,
                )
                code_to_area['NIE' + parl_code] = nia_area
            ward_to_assembly[ward_code] = code_to_area['NIE' + parl_code]

        # The manual fix for the split ward in the 2010 boundaries
        derryaghy = csv.reader(open('../../data/Derryaghy-postcodes.csv'))
        derryaghy_fix = {}
        for postcode, parl_name in derryaghy:
            derryaghy_fix[postcode] = name_to_area[parl_name]

<commit_msg>When adding NI WMCs, put in new ONS code as well as old.<commit_after># This script is used to import Northern Ireland areas, from the data in the
# National Statistics Postcode Database.
# http://www.ons.gov.uk/about-statistics/geography/products/geog-products-postcode/nspd/
#
# XXX This is incomplete, it needs to know which things have had boundary changes
# like import_boundary_line does. Hopefully just using new GSS codes by the time
# NI has any boundary changes.
#
# The fields of NSPD Open CSV file are:
#   Postcode (7), Postcode (8), Postcode (sp), Start date, End date, County,
#   council, ward, usertype, Easting, Northing, quality, SHA, IT cluster,
#   country, GOR, Stats region, Parliamentary constituency, Euro region,
#   TEC/LEC, Travel to Work area, Primary Care Org, NUTS, 1991 census ED,
#   1991 census ED, ED indicator, Pre-July 2006 SHA, LEA, Pre 2002 Health
#   Authority, 1991 ward code, 1991 ward code, 1998 ward code, 2005 stats ward,
#   OA code, OA indicator, CAS ward, National Park, SOA (Lower), Datazone, SOA
#   (Middle), Urban/rural, Urban/rural, Urban/rural, Intermediate, SOA (NI), OA
#   classification, Pre October 2006 PCO

import csv, re
from django.contrib.gis.geos import Point
from django.core.management.base import LabelCommand
from mapit.postcodes.models import Postcode
from mapit.areas.models import Area, Generation

class Command(LabelCommand):
    help = 'Creates/updates Northern Ireland areas, using the NSPD'
    args = '<NSPD CSV file>'

    def handle_label(self, file, **options):
        current_generation = Generation.objects.current()
        new_generation = Generation.objects.new()
        if not new_generation:
            raise Exception, "No new generation to be used for import!"

        euro_area, created = Area.objects.get_or_create(country='N', type='EUR',
            generation_low__lte=current_generation, generation_high__gte=current_generation,
            defaults = { 'generation_low': new_generation, 'generation_high': new_generation }
        )
        euro_area.generation_high = new_generation
        euro_area.save()
        euro_area.names.get_or_create(type='S', name='Northern Ireland')

        # Read in ward name -> electoral area name/area
        ni_eas = csv.reader(open('../../data/ni-electoral-areas.csv'))
        ni_eas.next()
        ward_to_electoral_area = {}
        e = {}
        for district, electoral_area, ward, dummy in ni_eas:
            if not district:
                district = last_district
            if not electoral_area:
                electoral_area = last_electoral_area
            last_district = district
            last_electoral_area = electoral_area
            if electoral_area not in e:
                ea = Area.objects.get_or_create_with_name(
                    country='N', type='LGE', name_type='M', name=electoral_area,
                )
                e[electoral_area] = ea
            ward_to_electoral_area.setdefault(district, {})[ward] = e[electoral_area]

        # Read in new ONS code to names
        snac = csv.reader(open('../../data/snac-2009-ni-cons2ward.csv'))
        snac.next()
        code_to_area = {}
        name_to_area = {}
        ward_to_parl = {}
        for parl_code, parl_name, ward_code, ward_name, district_code, district_name in snac:
            if district_name not in ward_to_electoral_area:
                raise Exception, "District %s is missing" % district_name
            if ward_name not in ward_to_electoral_area[district_name]:
                raise Exception, "Ward %s, district %s is missing" % (ward_name, district_name)

            ward_code = ward_code.replace(' ', '')

            if district_code not in code_to_area:
                district_area = Area.objects.get_or_create_with_code(
                    country='N', type='LGD', code_type='ons', code=district_code,
                )
                district_area.names.get_or_create(type='S', name=district_name)
                code_to_area[district_code] = district_area

            if ward_code not in code_to_area:
                ward_area = Area.objects.get_or_create_with_code(
                    country='N', type='LGW', code_type='ons', code=ward_code
                )
                ward_area.names.get_or_create(type='S', name=ward_name)
                ward_area.parent_area = ward_to_electoral_area[district_name][ward_name]
                ward_area.save()
                ward_area.parent_area.parent_area = code_to_area[district_code]
                ward_area.parent_area.save()
                code_to_area[ward_code] = ward_area

            if ward_code == '95S24': continue # Derryaghy

            if parl_code not in code_to_area:
                parl_area = Area.objects.get_or_create_with_code(
                    country='N', type='WMC', code_type='ons', code=parl_code,
                )
                parl_area.names.get_or_create(type='S', name=parl_name)
                new_code = re.sub('^7', 'N060000', parl_code)
                parl_area.codes.get_or_create(type='gss', code=new_code)
                code_to_area[parl_code] = parl_area
                name_to_area[parl_name] = parl_area

            ward_to_parl[ward_code] = code_to_area[parl_code]

        # Read in old SNAC for NI Assembly boundaries, still the same until 2011
        snac = csv.reader(open('../../data/snac-2003-ni-cons2ward.csv'))
        snac.next()
        ward_to_assembly = {}
        for parl_code, parl_name, ward_code, ward_name, district_code, district_name in snac:
            ward_code = ward_code.replace(' ', '')
            if 'NIE' + parl_code not in code_to_area:
                nia_area = Area.objects.get_or_create_with_name(
                    country='N', type='NIE', name_type='S', name=parl_name,
                )
                code_to_area['NIE' + parl_code] = nia_area
            ward_to_assembly[ward_code] = code_to_area['NIE' + parl_code]

        # The manual fix for the split ward in the 2010 boundaries
        derryaghy = csv.reader(open('../../data/Derryaghy-postcodes.csv'))
        derryaghy_fix = {}
        for postcode, parl_name in derryaghy:
            derryaghy_fix[postcode] = name_to_area[parl_name]

<|endoftext|>