<commit_before>#!/usr/bin/env python
# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2016, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# 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 Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import copy
import random

import matplotlib.pyplot as plt
import numpy as np

from nupic.bindings.math import GetNTAReal
realDType = GetNTAReal()
uintType = "uint32"

def percentOverlap(x1, x2):
  """
  Computes the percentage of overlap between vectors x1 and x2.

  @param x1   (array) binary vector
  @param x2   (array) binary vector
  @param size (int)   length of binary vectors

  @return percentOverlap (float) percentage overlap between x1 and x2
  """
  nonZeroX1 = np.count_nonzero(x1)
  nonZeroX2 = np.count_nonzero(x2)

  percentOverlap = 0
  minX1X2 = min(nonZeroX1, nonZeroX2)
  if minX1X2 > 0:
    overlap = float(np.dot(x1.T, x2))
    percentOverlap = overlap / minX1X2

  return percentOverlap

def addNoiseToVector(inputVector, noiseLevel, vectorType):
  """
  Add noise to SDRs
  @param inputVector (array) binary vector to be corrupted
  @param noiseLevel  (float) amount of noise to be applied on the vector.
  @param vectorType  (string) "sparse" or "dense"
  """
  if vectorType == 'sparse':
    corruptSparseVector(inputVector, noiseLevel)
  elif vectorType == 'dense':
    corruptDenseVector(inputVector, noiseLevel)
  else:
    raise ValueError("vectorType must be 'sparse' or 'dense' ")

def corruptDenseVector(vector, noiseLevel):
  """
  Corrupts a binary vector by inverting noiseLevel percent of its bits.

  @param vector     (array) binary vector to be corrupted
  @param noiseLevel (float) amount of noise to be applied on the vector.
  """
  size = len(vector)
  for i in range(size):
    rnd = random.random()
    if rnd < noiseLevel:
      if vector[i] == 1:
        vector[i] = 0
      else:
        vector[i] = 1

def corruptSparseVector(sdr, noiseLevel):
  """
  Add noise to sdr by turning off numNoiseBits active bits and turning on
  numNoiseBits in active bits
  @param sdr        (array) Numpy array of the  SDR
  @param noiseLevel (float) amount of noise to be applied on the vector.
  """

  numNoiseBits = int(noiseLevel * np.sum(sdr))
  if numNoiseBits <= 0:
    return sdr
  activeBits = np.where(sdr > 0)[0]
  inActiveBits = np.where(sdr == 0)[0]

  turnOffBits = np.random.permutation(activeBits)
  turnOnBits = np.random.permutation(inActiveBits)
  turnOffBits = turnOffBits[:numNoiseBits]
  turnOnBits = turnOnBits[:numNoiseBits]

  sdr[turnOffBits] = 0
  sdr[turnOnBits] = 1

def calculateOverlapCurve(sp, inputVectors):
  """
  Evalulate noise robustness of SP for a given set of SDRs
  @param sp a spatial pooler instance
  @param inputVectors list of arrays.
  :return:
  """
  columnNumber = np.prod(sp.getColumnDimensions())
  numInputVector, inputSize = inputVectors.shape

  outputColumns = np.zeros((numInputVector, columnNumber), dtype=uintType)
  outputColumnsCorrupted = np.zeros((numInputVector, columnNumber),
                                    dtype=uintType)

  noiseLevelList = np.linspace(0, 1.0, 21)
  inputOverlapScore = np.zeros((numInputVector, len(noiseLevelList)))
  outputOverlapScore = np.zeros((numInputVector, len(noiseLevelList)))
  for i in range(numInputVector):
    for j in range(len(noiseLevelList)):
      inputVectorCorrupted = copy.deepcopy(inputVectors[i][:])
      corruptSparseVector(inputVectorCorrupted, noiseLevelList[j])

      sp.compute(inputVectors[i][:], False, outputColumns[i][:])
      sp.compute(inputVectorCorrupted, False,
                 outputColumnsCorrupted[i][:])

      inputOverlapScore[i][j] = percentOverlap(inputVectors[i][:],
                                               inputVectorCorrupted)
      outputOverlapScore[i][j] = percentOverlap(outputColumns[i][:],
                                                outputColumnsCorrupted[i][:])

  return inputOverlapScore, outputOverlapScore

def classifySPoutput(targetOutputColumns, outputColumns):
  """
  Classify the SP output
  @param targetOutputColumns (list) The target outputs, corresponding to
                                    different classes
  @param outputColumns (array) The current output
  @return classLabel (int) classification outcome
  """
  numTargets, numDims = targetOutputColumns.shape
  overlap = np.zeros((numTargets,))
  for i in range(numTargets):
    overlap[i] = percentOverlap(outputColumns, targetOutputColumns[i, :])
  classLabel = np.argmax(overlap)
  return classLabel

def classificationAccuracyVsNoise(sp, inputVectors, noiseLevelList):
  """
  Evaluate whether the SP output is classifiable, with varying amount of noise
  @param sp a spatial pooler instance
  @param inputVectors (list) list of input SDRs
  @param noiseLevelList (list) list of noise levels
  :return:
  """
  numInputVector, inputSize = inputVectors.shape

  if sp is None:
    targetOutputColumns = copy.deepcopy(inputVectors)
  else:
    columnNumber = np.prod(sp.getColumnDimensions())
    # calculate target output given the uncorrupted input vectors
    targetOutputColumns = np.zeros((numInputVector, columnNumber),
                                   dtype=uintType)
    for i in range(numInputVector):
      sp.compute(inputVectors[i][:], False, targetOutputColumns[i][:])

  outcomes = np.zeros((len(noiseLevelList), numInputVector))
  for i in range(len(noiseLevelList)):
    for j in range(numInputVector):
      corruptedInputVector = copy.deepcopy(inputVectors[j][:])
      corruptSparseVector(corruptedInputVector, noiseLevelList[i])
      # addNoiseToVector(corruptedInputVector, noiseLevelList[i], inputVectorType)

      if sp is None:
        outputColumns = copy.deepcopy(corruptedInputVector)
      else:
        outputColumns = np.zeros(sp.getColumnDimensions(), dtype=uintType)
        sp.compute(corruptedInputVector, False, outputColumns)

      predictedClassLabel = classifySPoutput(targetOutputColumns, outputColumns)
      outcomes[i][j] = predictedClassLabel == j

  predictionAccuracy = np.mean(outcomes, 1)
  return predictionAccuracy

def inspectSpatialPoolerStats(sp, inputVectors, saveFigPrefix=None):
  """
  Inspect the statistics of a spatial pooler given a set of input vectors
  @param sp: an spatial pooler instance
  @param inputVectors: a set of input vectors
  """
  numInputVector, inputSize = inputVectors.shape
  numColumns = np.prod(sp.getColumnDimensions())

  outputColumns = np.zeros((numInputVector, numColumns), dtype=uintType)
  inputOverlap = np.zeros((numInputVector, numColumns), dtype=uintType)

  connectedCounts = np.zeros((numColumns, ), dtype=uintType)
  sp.getConnectedCounts(connectedCounts)

  winnerInputOverlap = np.zeros(numInputVector)
  for i in range(numInputVector):
    sp.compute(inputVectors[i][:], False, outputColumns[i][:])
    inputOverlap[i][:] = sp.getOverlaps()
    winnerInputOverlap[i] = np.mean(inputOverlap[i][np.where(outputColumns[i][:] > 0)[0]])
  avgInputOverlap = np.mean(inputOverlap, 0)

  activationProb = np.mean(outputColumns.astype(realDType), 0)

  # fig, axs = plt.subplots(2, 1)
  # axs[0].imshow(inputVectors[:, :200], cmap='gray')
  # axs[0].set_ylabel('sample #')
  # axs[0].set_title('input vectors')
  # axs[1].imshow(outputColumns[:, :200], cmap='gray')
  # axs[1].set_ylabel('sample #')
  # axs[1].set_title('output vectors')
  # if saveFigPrefix is not None:
  #   plt.savefig('figures/{}_example_input_output.pdf'.format(saveFigPrefix))

  fig, axs = plt.subplots(2, 2)
  axs[0, 0].hist(connectedCounts)
  axs[0, 0].set_xlabel('# Connected Synapse')

  axs[0, 1].hist(winnerInputOverlap)
  axs[0, 1].set_xlabel('# winner input overlap')

  axs[1, 0].hist(activationProb)
  axs[1, 0].set_xlabel('activation prob')

  axs[1, 1].plot(connectedCounts, activationProb, '.')
  axs[1, 1].set_xlabel('connection #')
  axs[1, 1].set_ylabel('activation freq')
  plt.tight_layout()

  if saveFigPrefix is not None:
    plt.savefig('figures/{}_network_stats.pdf'.format(saveFigPrefix))
  return fig

def calculateEntropy(activeColumns):
  """
  calculate entropy given activation history
  @param activeColumns (array) 2D numpy array of activation history
  @return entropy (flaot) entropy
  """
  MIN_ACTIVATION_PROB = 0.000001
  activationProb = np.mean(activeColumns, 0)
  activationProb[activationProb < MIN_ACTIVATION_PROB] = MIN_ACTIVATION_PROB
  activationProb = activationProb / np.sum(activationProb)

  entropy = -np.dot(activationProb, np.log2(activationProb))
  return entropy

def calculateInputOverlapMat(inputVectors, sp):
  numColumns = np.product(sp.getColumnDimensions())
  numInputVector, inputSize = inputVectors.shape
  overlapMat = np.zeros((numColumns, numInputVector))
  for c in range(numColumns):
    connectedSynapses = np.zeros((inputSize, ), dtype=uintType)
    sp.getConnectedSynapses(c, connectedSynapses)
    for i in range(numInputVector):
      overlapMat[c, i] = percentOverlap(connectedSynapses, inputVectors[i, :])
  return overlapMat

def calculateStability(activeColumnsCurrentEpoch, activeColumnsPreviousEpoch):
  activeColumnsStable = np.logical_and(activeColumnsCurrentEpoch,
                                       activeColumnsPreviousEpoch)
  stability = np.mean(np.sum(activeColumnsStable, 1))/\
              np.mean(np.sum(activeColumnsCurrentEpoch))
  return stability<commit_msg>fix stability metric<commit_after>#!/usr/bin/env python
# ----------------------------------------------------------------------
# Numenta Platform for Intelligent Computing (NuPIC)
# Copyright (C) 2016, Numenta, Inc.  Unless you have an agreement
# with Numenta, Inc., for a separate license for this software code, the
# following terms and conditions apply:
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero Public License version 3 as
# published by the Free Software Foundation.
#
# 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 Affero Public License for more details.
#
# You should have received a copy of the GNU Affero Public License
# along with this program.  If not, see http://www.gnu.org/licenses.
#
# http://numenta.org/licenses/
# ----------------------------------------------------------------------

import copy
import random

import matplotlib.pyplot as plt
import numpy as np

from nupic.bindings.math import GetNTAReal
realDType = GetNTAReal()
uintType = "uint32"

def percentOverlap(x1, x2):
  """
  Computes the percentage of overlap between vectors x1 and x2.

  @param x1   (array) binary vector
  @param x2   (array) binary vector
  @param size (int)   length of binary vectors

  @return percentOverlap (float) percentage overlap between x1 and x2
  """
  nonZeroX1 = np.count_nonzero(x1)
  nonZeroX2 = np.count_nonzero(x2)

  percentOverlap = 0
  minX1X2 = min(nonZeroX1, nonZeroX2)
  if minX1X2 > 0:
    overlap = float(np.dot(x1.T, x2))
    percentOverlap = overlap / minX1X2

  return percentOverlap

def addNoiseToVector(inputVector, noiseLevel, vectorType):
  """
  Add noise to SDRs
  @param inputVector (array) binary vector to be corrupted
  @param noiseLevel  (float) amount of noise to be applied on the vector.
  @param vectorType  (string) "sparse" or "dense"
  """
  if vectorType == 'sparse':
    corruptSparseVector(inputVector, noiseLevel)
  elif vectorType == 'dense':
    corruptDenseVector(inputVector, noiseLevel)
  else:
    raise ValueError("vectorType must be 'sparse' or 'dense' ")

def corruptDenseVector(vector, noiseLevel):
  """
  Corrupts a binary vector by inverting noiseLevel percent of its bits.

  @param vector     (array) binary vector to be corrupted
  @param noiseLevel (float) amount of noise to be applied on the vector.
  """
  size = len(vector)
  for i in range(size):
    rnd = random.random()
    if rnd < noiseLevel:
      if vector[i] == 1:
        vector[i] = 0
      else:
        vector[i] = 1

def corruptSparseVector(sdr, noiseLevel):
  """
  Add noise to sdr by turning off numNoiseBits active bits and turning on
  numNoiseBits in active bits
  @param sdr        (array) Numpy array of the  SDR
  @param noiseLevel (float) amount of noise to be applied on the vector.
  """

  numNoiseBits = int(noiseLevel * np.sum(sdr))
  if numNoiseBits <= 0:
    return sdr
  activeBits = np.where(sdr > 0)[0]
  inActiveBits = np.where(sdr == 0)[0]

  turnOffBits = np.random.permutation(activeBits)
  turnOnBits = np.random.permutation(inActiveBits)
  turnOffBits = turnOffBits[:numNoiseBits]
  turnOnBits = turnOnBits[:numNoiseBits]

  sdr[turnOffBits] = 0
  sdr[turnOnBits] = 1

def calculateOverlapCurve(sp, inputVectors):
  """
  Evalulate noise robustness of SP for a given set of SDRs
  @param sp a spatial pooler instance
  @param inputVectors list of arrays.
  :return:
  """
  columnNumber = np.prod(sp.getColumnDimensions())
  numInputVector, inputSize = inputVectors.shape

  outputColumns = np.zeros((numInputVector, columnNumber), dtype=uintType)
  outputColumnsCorrupted = np.zeros((numInputVector, columnNumber),
                                    dtype=uintType)

  noiseLevelList = np.linspace(0, 1.0, 21)
  inputOverlapScore = np.zeros((numInputVector, len(noiseLevelList)))
  outputOverlapScore = np.zeros((numInputVector, len(noiseLevelList)))
  for i in range(numInputVector):
    for j in range(len(noiseLevelList)):
      inputVectorCorrupted = copy.deepcopy(inputVectors[i][:])
      corruptSparseVector(inputVectorCorrupted, noiseLevelList[j])

      sp.compute(inputVectors[i][:], False, outputColumns[i][:])
      sp.compute(inputVectorCorrupted, False,
                 outputColumnsCorrupted[i][:])

      inputOverlapScore[i][j] = percentOverlap(inputVectors[i][:],
                                               inputVectorCorrupted)
      outputOverlapScore[i][j] = percentOverlap(outputColumns[i][:],
                                                outputColumnsCorrupted[i][:])

  return inputOverlapScore, outputOverlapScore

def classifySPoutput(targetOutputColumns, outputColumns):
  """
  Classify the SP output
  @param targetOutputColumns (list) The target outputs, corresponding to
                                    different classes
  @param outputColumns (array) The current output
  @return classLabel (int) classification outcome
  """
  numTargets, numDims = targetOutputColumns.shape
  overlap = np.zeros((numTargets,))
  for i in range(numTargets):
    overlap[i] = percentOverlap(outputColumns, targetOutputColumns[i, :])
  classLabel = np.argmax(overlap)
  return classLabel

def classificationAccuracyVsNoise(sp, inputVectors, noiseLevelList):
  """
  Evaluate whether the SP output is classifiable, with varying amount of noise
  @param sp a spatial pooler instance
  @param inputVectors (list) list of input SDRs
  @param noiseLevelList (list) list of noise levels
  :return:
  """
  numInputVector, inputSize = inputVectors.shape

  if sp is None:
    targetOutputColumns = copy.deepcopy(inputVectors)
  else:
    columnNumber = np.prod(sp.getColumnDimensions())
    # calculate target output given the uncorrupted input vectors
    targetOutputColumns = np.zeros((numInputVector, columnNumber),
                                   dtype=uintType)
    for i in range(numInputVector):
      sp.compute(inputVectors[i][:], False, targetOutputColumns[i][:])

  outcomes = np.zeros((len(noiseLevelList), numInputVector))
  for i in range(len(noiseLevelList)):
    for j in range(numInputVector):
      corruptedInputVector = copy.deepcopy(inputVectors[j][:])
      corruptSparseVector(corruptedInputVector, noiseLevelList[i])

      if sp is None:
        outputColumns = copy.deepcopy(corruptedInputVector)
      else:
        outputColumns = np.zeros(sp.getColumnDimensions(), dtype=uintType)
        sp.compute(corruptedInputVector, False, outputColumns)

      predictedClassLabel = classifySPoutput(targetOutputColumns, outputColumns)
      outcomes[i][j] = predictedClassLabel == j

  predictionAccuracy = np.mean(outcomes, 1)
  return predictionAccuracy

def inspectSpatialPoolerStats(sp, inputVectors, saveFigPrefix=None):
  """
  Inspect the statistics of a spatial pooler given a set of input vectors
  @param sp: an spatial pooler instance
  @param inputVectors: a set of input vectors
  """
  numInputVector, inputSize = inputVectors.shape
  numColumns = np.prod(sp.getColumnDimensions())

  outputColumns = np.zeros((numInputVector, numColumns), dtype=uintType)
  inputOverlap = np.zeros((numInputVector, numColumns), dtype=uintType)

  connectedCounts = np.zeros((numColumns, ), dtype=uintType)
  sp.getConnectedCounts(connectedCounts)

  winnerInputOverlap = np.zeros(numInputVector)
  for i in range(numInputVector):
    sp.compute(inputVectors[i][:], False, outputColumns[i][:])
    inputOverlap[i][:] = sp.getOverlaps()
    winnerInputOverlap[i] = np.mean(inputOverlap[i][np.where(outputColumns[i][:] > 0)[0]])
  avgInputOverlap = np.mean(inputOverlap, 0)

  activationProb = np.mean(outputColumns.astype(realDType), 0)

  # fig, axs = plt.subplots(2, 1)
  # axs[0].imshow(inputVectors[:, :200], cmap='gray')
  # axs[0].set_ylabel('sample #')
  # axs[0].set_title('input vectors')
  # axs[1].imshow(outputColumns[:, :200], cmap='gray')
  # axs[1].set_ylabel('sample #')
  # axs[1].set_title('output vectors')
  # if saveFigPrefix is not None:
  #   plt.savefig('figures/{}_example_input_output.pdf'.format(saveFigPrefix))

  fig, axs = plt.subplots(2, 2)
  axs[0, 0].hist(connectedCounts)
  axs[0, 0].set_xlabel('# Connected Synapse')

  axs[0, 1].hist(winnerInputOverlap)
  axs[0, 1].set_xlabel('# winner input overlap')

  axs[1, 0].hist(activationProb)
  axs[1, 0].set_xlabel('activation prob')

  axs[1, 1].plot(connectedCounts, activationProb, '.')
  axs[1, 1].set_xlabel('connection #')
  axs[1, 1].set_ylabel('activation freq')
  plt.tight_layout()

  if saveFigPrefix is not None:
    plt.savefig('figures/{}_network_stats.pdf'.format(saveFigPrefix))
  return fig

def calculateEntropy(activeColumns):
  """
  calculate entropy given activation history
  @param activeColumns (array) 2D numpy array of activation history
  @return entropy (flaot) entropy
  """
  MIN_ACTIVATION_PROB = 0.000001
  activationProb = np.mean(activeColumns, 0)
  activationProb[activationProb < MIN_ACTIVATION_PROB] = MIN_ACTIVATION_PROB
  activationProb = activationProb / np.sum(activationProb)

  entropy = -np.dot(activationProb, np.log2(activationProb))
  return entropy

def calculateInputOverlapMat(inputVectors, sp):
  numColumns = np.product(sp.getColumnDimensions())
  numInputVector, inputSize = inputVectors.shape
  overlapMat = np.zeros((numColumns, numInputVector))
  for c in range(numColumns):
    connectedSynapses = np.zeros((inputSize, ), dtype=uintType)
    sp.getConnectedSynapses(c, connectedSynapses)
    for i in range(numInputVector):
      overlapMat[c, i] = percentOverlap(connectedSynapses, inputVectors[i, :])
  return overlapMat

def calculateStability(activeColumnsCurrentEpoch, activeColumnsPreviousEpoch):
  activeColumnsStable = np.logical_and(activeColumnsCurrentEpoch,
                                       activeColumnsPreviousEpoch)
  stability = np.mean(np.sum(activeColumnsStable, 1))/\
              np.mean(np.sum(activeColumnsCurrentEpoch, 1))
  return stability<|endoftext|>


<commit_before>

class Store(object):

    #Properties
    #context_aware
    #transaction_aware

    def __init__(self, identifier, configuration=None):
        """
        identifier: URIRef of the Store. Defaults to CWD
        configuration: string containing infomation open can use to
        connect to datastore.
        """
        self.identifier = identifier
        if configuration:
            self.open(configuration)

    def open(self, configuration, create=True):
        """
        Opens the store specified by the configuration string. If
	create is True a store will be created if it does not already
	exist. If create is False and a store does not already exist
	an exception is raised. An exception is also raised if a store
	exists, but there is insufficient permissions to open the
	store.
        """

    def close(self, commit_pending_transaction=False):
        """ """

    def destroy(self, configuration):
        """ """

    def add(self, (subject, predicate, object), context=None, quoted=False):
        """ Add a triple to the store of triples. """

    def remove(self, (subject, predicate, object), context):
        """ Remove a triple from the store """

    def triples(self, (subject, predicate, object), context=None):
        """
        A generator over all the triples matching pattern. Pattern can
        be any objects for comparing against nodes in the store, for
        example, RegExLiteral, Date? DateRange?
        """

    # varients of triples will be done if / when optimization is needed

    def __len__(self, context=None):
        """ Number of statements in the store. """

    def contexts(self, triple=None):
        """ """

    def remove_context(self, identifier):
        """ """

    # Optional Namespace methods

    def bind(self, prefix, namespace):
        """ """

    def prefix(self, namespace):
        """ """

    def namespace(self, prefix):
        """ """

    def namespaces(self):
        """ """

    # Optional Transactional methods

    def commit(self):
        """ """

    def rollback(self):
        """ """

<commit_msg>added formula_aware mention<commit_after>

class Store(object):

    #Properties
    #context_aware
    #formula_aware
    #transaction_aware

    def __init__(self, identifier, configuration=None):
        """
        identifier: URIRef of the Store. Defaults to CWD
        configuration: string containing infomation open can use to
        connect to datastore.
        """
        self.identifier = identifier
        if configuration:
            self.open(configuration)

    def open(self, configuration, create=True):
        """
        Opens the store specified by the configuration string. If
	create is True a store will be created if it does not already
	exist. If create is False and a store does not already exist
	an exception is raised. An exception is also raised if a store
	exists, but there is insufficient permissions to open the
	store.
        """

    def close(self, commit_pending_transaction=False):
        """ """

    def destroy(self, configuration):
        """ """

    def add(self, (subject, predicate, object), context=None, quoted=False):
        """ Add a triple to the store of triples. """

    def remove(self, (subject, predicate, object), context):
        """ Remove a triple from the store """

    def triples(self, (subject, predicate, object), context=None):
        """
        A generator over all the triples matching pattern. Pattern can
        be any objects for comparing against nodes in the store, for
        example, RegExLiteral, Date? DateRange?
        """

    # varients of triples will be done if / when optimization is needed

    def __len__(self, context=None):
        """ Number of statements in the store. """

    def contexts(self, triple=None):
        """ """

    def remove_context(self, identifier):
        """ """

    # Optional Namespace methods

    def bind(self, prefix, namespace):
        """ """

    def prefix(self, namespace):
        """ """

    def namespace(self, prefix):
        """ """

    def namespaces(self):
        """ """

    # Optional Transactional methods

    def commit(self):
        """ """

    def rollback(self):
        """ """

<|endoftext|>


<commit_before>#
#  ENVISIoN
#
#  Copyright (c) 2019 Anton Hjert
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#  1. Redistributions of source code must retain the above copyright notice, this
#  list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright notice,
#  this list of conditions and the following disclaimer in the documentation
#  and/or other materials provided with the distribution.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
#  ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##############################################################################################
#
#  Alterations to this file by
#
#  To the extent possible under law, the person who associated CC0
#  with the alterations to this file has waived all copyright and related
#  or neighboring rights to the alterations made to this file.
#
#  You should have received a copy of the CC0 legalcode along with
#  this work.  If not, see
#  <http://creativecommons.org/publicdomain/zero/1.0/>.

import wx, sys, os, h5py
from parameter_utils import *
from generalCollapsible import GeneralCollapsible
from volumeControlCollapsible import VolumeControlCollapsible
from backgroundCollapsible import BackgroundCollapsible
class ChargeFrame(GeneralCollapsible):
    def __init__(self, parent):
        super().__init__(parent, label = "Charge")
        self.slice=False

        # Setup volume rendering controls
        self.volumeCollapsible = VolumeControlCollapsible(self.GetPane(), "Volume Rendering")
        self.add_sub_collapsible(self.volumeCollapsible)

        self.sliceBackgroundCollapsibe = BackgroundCollapsible(self.GetPane(), "Slice background")
        self.sliceBackgroundCollapsibe.type = 'SliceBackground'

        # Setup background controls
        self.backgroundCollapsibe = BackgroundCollapsible(self.GetPane(), "Background")
        self.add_sub_collapsible(self.backgroundCollapsibe)

        self.sliceBox = wx.CheckBox(self.GetPane(), label="Slice")
        self.add_item(self.sliceBox)

        # Override default binding
        # Note that function should be called "on_collapse" to work
        self.Bind(wx.EVT_COLLAPSIBLEPANE_CHANGED, self.on_collapse)
        self.sliceBox.Bind(wx.EVT_CHECKBOX,self.on_check)

    def on_check(self,event):
        if self.slice:
            self.slice = False
            self.backgroundCollapsibe.type = 'Background'
            self.backgroundCollapsibe.SetLabel('Background')
            self.remove_sub_collapsible(3)
        else:
            self.slice = True
            self.add_sub_collapsible(self.sliceBackgroundCollapsibe)
            self.backgroundCollapsibe.SetLabel('Volume background')
            self.backgroundCollapsibe.type = 'VolumeBackground'
        self.update_collapse()
        clear_processor_network()
        self.start_vis()

    def on_collapse(self, event = None):
        self.update_collapse()
        # Needs to be called to update the layout properly
        if self.IsCollapsed():
            # Disable charge vis
            clear_processor_network()
        else:
            #Start Charge vis
            self.start_vis()

    def start_vis(self):
        if '/{}'.format('CHG') in h5py.File(self.parent_collapsible.path, 'r'):
            envision.inviwo.charge(self.parent_collapsible.path,
                                iso = None, slice = self.slice,
                                xpos = 0, ypos = 0)
            if self.slice:
                self.set_canvas_pos('SliceCanvas')
                self.set_canvas_pos('VolumeCanvas')
            else:
                self.set_canvas_pos()
        else:
            self.open_message('The file of choice does not contain Charge-data',
                                'Visualization failed!')
        <commit_msg>bug fixes in frameCharge<commit_after>#
#  ENVISIoN
#
#  Copyright (c) 2019 Anton Hjert
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#  1. Redistributions of source code must retain the above copyright notice, this
#  list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright notice,
#  this list of conditions and the following disclaimer in the documentation
#  and/or other materials provided with the distribution.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
#  ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##############################################################################################
#
#  Alterations to this file by
#
#  To the extent possible under law, the person who associated CC0
#  with the alterations to this file has waived all copyright and related
#  or neighboring rights to the alterations made to this file.
#
#  You should have received a copy of the CC0 legalcode along with
#  this work.  If not, see
#  <http://creativecommons.org/publicdomain/zero/1.0/>.

import wx, sys, os, h5py
from parameter_utils import *
from generalCollapsible import GeneralCollapsible
from volumeControlCollapsible import VolumeControlCollapsible
from backgroundCollapsible import BackgroundCollapsible
class ChargeFrame(GeneralCollapsible):
    def __init__(self, parent):
        super().__init__(parent, label = "Charge")
        self.slice=False

        # Setup volume rendering controls
        self.volumeCollapsible = VolumeControlCollapsible(self.GetPane(), "Volume Rendering")
        self.add_sub_collapsible(self.volumeCollapsible)

        self.sliceBackgroundCollapsibe = BackgroundCollapsible(self.GetPane(), "Slice background")
        self.sliceBackgroundCollapsibe.type = 'SliceBackground'

        # Setup background controls
        self.backgroundCollapsibe = BackgroundCollapsible(self.GetPane(), "Background")
        self.add_sub_collapsible(self.backgroundCollapsibe)

        self.sliceBox = wx.CheckBox(self.GetPane(), label="Slice")
        self.add_item(self.sliceBox)

        # Override default binding
        # Note that function should be called "on_collapse" to work
        self.Bind(wx.EVT_COLLAPSIBLEPANE_CHANGED, self.on_collapse)
        self.sliceBox.Bind(wx.EVT_CHECKBOX,self.on_check)

    def on_check(self,event):
        if self.slice:
            self.slice = False
            self.backgroundCollapsibe.type = 'Background'
            self.backgroundCollapsibe.SetLabel('Background')
            self.remove_sub_collapsible(3)
        else:
            self.slice = True
            self.add_sub_collapsible(self.sliceBackgroundCollapsibe)
            self.backgroundCollapsibe.SetLabel('Volume background')
            self.backgroundCollapsibe.type = 'VolumeBackground'
        self.update_collapse()
        clear_processor_network()
        self.start_vis()

    def on_collapse(self, event = None):
        self.update_collapse()
        # Needs to be called to update the layout properly
        if self.IsCollapsed():
            # Disable charge vis
            clear_processor_network()
        else:
            #Start Charge vis
            self.start_vis()

    def start_vis(self):
        if '/{}'.format('CHG') in h5py.File(self.parent_collapsible.path, 'r'):
            envision.inviwo.charge(self.parent_collapsible.path,
                                iso = None, slice = self.slice,
                                xpos = 0, ypos = 0)
            if self.slice:
                self.set_canvas_pos('SliceCanvas')
                self.set_canvas_pos('VolumeCanvas')
            else:
                self.set_canvas_pos()
        else:
            self.open_message('The file of choice does not contain Charge-data',
                                'Visualization failed!')
        <|endoftext|>


<commit_before>#
#  ENVISIoN
#
#  Copyright (c) 2017-2018 Josef Adamsson, Robert Cranston, David Hartman, Denise Härnström, Fredrik Segerhammar, Anders Rehult, Viktor Bernholtz, Elvis Jakobsson
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#  1. Redistributions of source code must retain the above copyright notice, this
#  list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright notice,
#  this list of conditions and the following disclaimer in the documentation
#  and/or other materials provided with the distribution.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
#  ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##############################################################################################
#
#  Alterations to this file by Anders Rehult and Viktor Bernholtz
#
#  To the extent possible under law, the person who associated CC0 with
#  the alterations to this file has waived all copyright and related
#  or neighboring rights to the alterations made to this file.
#
#  You should have received a copy of the CC0 legalcode along with
#  this work.  If not, see
#  <http://creativecommons.org/publicdomain/zero/1.0/>.

import numpy as np
import h5py
import array as arr

def _write_coordinates(h5file, atom_count, coordinates_list, elements, path):
    with h5py.File(h5file, "a") as h5:
        p=0
        for n in range(0,len(atom_count)):
            dataset_name = path+'/Atoms/'+format(n,'04d')
            h5.create_dataset(
                dataset_name,
                data=np.asarray(coordinates_list[p:atom_count[n]+p]),
                dtype=np.float32
            )
            h5[dataset_name].attrs["element"] = elements[n]
            p=p+atom_count[n]
    return

def _write_basis(h5file, basis):
    with h5py.File(h5file, "a") as h5:
        if not "/basis" in h5:
            h5.create_dataset('/basis', data=basis, dtype=np.float32)
    return

def _write_scaling_factor(h5file, scaling_factor):
    with h5py.File(h5file, "a") as h5:
        if not "/scaling_factor" in h5:
            h5.create_dataset('/scaling_factor', data=scaling_factor, dtype=np.float32)
    return

def _write_md(h5file, atom_count, coordinates_list, elements, step):
    with h5py.File(h5file, "a") as h5:
        p=0
        for n in range(0,len(atom_count)):
            dataset_name = '/MD/Atoms/'+format(n,'04d')
            if step == 0:
                h5.create_dataset(
                    dataset_name,
                    data=np.asarray(coordinates_list[p:atom_count[n]+p]),
                    dtype=np.float32,
                    maxshape=(None, 3)
                    )
                h5[dataset_name].attrs["element"] = elements[n]
                h5[dataset_name].attrs["atoms"] = atom_count[n]
                p=p+atom_count[n]
            else:
                dataset = h5[dataset_name]
                dataset.resize((step+1)*atom_count[n],axis=0)
                start = step*atom_count[n]
                dataset[start:] = np.asarray(coordinates_list[p:atom_count[n]+p])
                p=p+atom_count[n]
    return

def _write_steps(h5file, steps):
    with h5py.File(h5file, "a") as h5:
        h5['/MD'].attrs["steps"] = steps
    return

def _write_bandstruct(h5file, band_data, kval_list):
    with h5py.File(h5file, "a") as h5:
        h5.create_dataset('BandStructure/KPoints', data=np.array(kval_list), dtype = np.float32)
        for i, band in enumerate(band_data):
            dataset = h5.create_dataset('Bandstructure/Bands/{}/{}'.format(i, 'Energy'), data=np.array(band), dtype = np.float32)
            dataset.attrs['Unit'] = 'eV'
            dataset.attrs['QuantitySymbol'] = '$E$'
            dataset.attrs['QuantityName'] = 'Energy'
            dataset.attrs['VariableName'] = 'Band {}'.format(i)
            dataset.attrs['VariableSymbol'] = '$B_{{{}}}$'.format(i)

def _write_fermisurface(h5file, kval_list, fermi_energy, reciprocal_lattice_vectors):
    with h5py.File(h5file,"a") as h5:
        for i in range(0, len(kval_list)):
            dataset = h5.create_dataset('/FermiSurface/KPoints/{}/Energy'.format(i),
                                        data = np.array(kval_list[i].energies),
                                        dtype = np.float32)
            dataset.attrs['Unit'] = 'eV'
            dataset.attrs['QuantitySymbol'] = '$E$'
            dataset.attrs['QuantityName'] = 'Energy'
            dataset.attrs['VariableName'] = 'Band {}'.format(i)
            dataset.attrs['VariableSymbol'] = '$B_{{{}}}$'.format(i)

            h5.create_dataset('/FermiSurface/KPoints/{}/Coordinates'.format(i),
                              data = np.array(kval_list[i].coordinates),
                              dtype = np.float32)

        h5.create_dataset('/FermiSurface/FermiEnergy',
                          data = np.float(fermi_energy),
                          dtype = np.float32)

        for i in range(0, len(reciprocal_lattice_vectors)):
            h5.create_dataset('/FermiSurface/ReciprocalLatticeVectors/{}'.format(i),
                              data = np.array([float(x) for x in reciprocal_lattice_vectors[i]]),
                              dtype = np.float32)

def _write_dos(h5file, total, partial, total_data, partial_list, fermi_energy):

    def set_attrs(dataset, VariableName = '', VariableSymbol = '', QuantityName = '', QuantitySymbol = '', Unit = ''):
        dataset.attrs.update({
            'VariableName' : VariableName,
            'VariableSymbol' : VariableSymbol,
            'QuantityName' : QuantityName,
            'QuantitySymbol' : QuantitySymbol,
            'Unit' : Unit
        })

    with h5py.File(h5file, "a") as h5:
        dataset = h5.create_dataset('FermiEnergy', data = np.array(fermi_energy), dtype = np.float32)
        set_attrs(dataset, 'Fermi Energy', '$E_f$', 'Energy', '$E$', 'eV')
        for i, (name, data) in enumerate(zip(total, total_data)):
            dataset = h5.create_dataset('DOS/Total/{}'.format(name), data=np.array(data), dtype = np.float32)
            if name == 'Energy':
                set_attrs(dataset, 'Energy', '$E$', 'Energy', '$E$', 'eV')
            else:
                if name.startswith('Integrated'):
                    set_attrs(dataset, name, '', 'Integrated Density of States', '$\\int D$', 'states')
                else:
                    set_attrs(dataset, name, '', 'Density of States', '$D$', 'states/eV')
        for i, partial_data in enumerate(partial_list):
            # Write all available DOS. Stop when partial or partial_data runs out (if directories for for example f-DOS
            # exist but no corresponding data exists, stop.)
            for (name, data) in zip(partial, partial_data):
                dataset = h5.create_dataset('DOS/Partial/{}/{}'.format(i, name), data=np.array(data), dtype = np.float32)
                if name == 'Energy':
                    set_attrs(dataset, 'Energy', '$E$', 'Energy', '$E$', 'eV')
                else:
                    if name.startswith('Integrated'):
                        set_attrs(dataset, name, '', 'Integrated Density of States', '$\\int D$', 'states')
                    else:
                        set_attrs(dataset, name, '', 'Density of States', '$D$', 'states/eV')

def _write_volume(h5file, i, array, data_dim, hdfgroup):
    with h5py.File(h5file, "a") as h5:
        if array:
            h5.create_dataset('{}/{}'.format(hdfgroup, i), data = np.reshape(array, (data_dim[2],data_dim[1],data_dim[0])), dtype=np.float32)
        else:
            h5['{}/final'.format(hdfgroup)] = h5['{}/{}'.format(hdfgroup, i-1)]
    return

def _write_incar(h5file, incar_data):
    with h5py.File(h5file, 'a') as h5:
        for key, value in incar_data.items():
            h5.create_dataset("/incar/{}".format(key), data=value)

def _write_parcharges(h5file, array_tot, data_dim_tot, array_mag, data_dim_mag, band_nr):
    with h5py.File(h5file, "a") as h5:
        h5.create_dataset('PARCHG/Bands/{}/total'.format(band_nr), data = np.reshape(array_tot, (data_dim_tot[2],data_dim_tot[1],data_dim_tot[0])), dtype=np.float32)
        if not np.size(array_mag) == 0:
            h5.create_dataset('PARCHG/Bands/{}/magnetic'.format(band_nr), data = np.reshape(array_mag, (data_dim_mag[2],data_dim_mag[1],data_dim_mag[0])), dtype=np.float32)
    return

def _write_pcdat(h5file, pcdat_data, APACO_val, NPACO_val):
    #   The function is called to write data from PCDAT to HDF5-file. A dataset is created for each element in the system.

    #    Parameters
    #    __________
    #    h5file: str
    #        String containing path to HDF5-file.

    #    pcdat_data:
    #        Is a dictionary with the structure {'element_type':[PKF_values]}. PKF_values are the average number of atoms, of specific element_type. If the system has elements 'Si', 'Au'       and 'K', the dictionary will look like {'Si':float[x], 'Au':float[x], 'K':float[x]} where the float[x] is a list with the PKF values.
    #
    #    APACO_val:
    #        The value of APACO in INCAR if set, default value 16 (Å), otherwise. It sets the maximum distance in the evaluation of the pair-correlation function
    #
    #    NPACO_val:
    #        The value of NPACO in INCAR if set, default value 256 slots otherwise. It sets the number of slots in the pair-correlation function written to PCDAT.

    #    Return
    #    ______
    #    Bool: True if parsed, False otherwise.

    #WARNING!!!!!Just for BP4- CLEAR THIS WHEN DONEEE!!!!
    for i in range(257):
        pcdat_data['Cu'].pop(512-i)

    #x-values for PKF
    with h5py.File(h5file, "a") as h5:
        dset_name = "PairCorrelationFunc/Iterations"
        normal_arr = arr.array('f', [])

        for i in range(0, NPACO_val):
            x = (APACO_val / NPACO_val) * i
            normal_arr.append(x)

        value = np.asarray(normal_arr)
        h5.create_dataset(dset_name, data=value, dtype=np.float32)

        # create dataset for every element
        for n in range(len(pcdat_data)):
            # By making a dict a list, it's keys are accessable
            element_symbol = list(pcdat_data)[n]
            dset_name = "PairCorrelationFunc/Elements/{}".format(element_symbol)
            value = np.asarray(pcdat_data[element_symbol])
            dset = h5.create_dataset(dset_name, data=value, dtype=np.float32)
            h5[dset_name].attrs["element"] = element_symbol

    return None

<commit_msg>Made changes to parsing system PKF. Now it parses paircorrelation function with accordance to different time frames.<commit_after>#
#  ENVISIoN
#
#  Copyright (c) 2017-2018 Josef Adamsson, Robert Cranston, David Hartman, Denise Härnström, Fredrik Segerhammar, Anders Rehult, Viktor Bernholtz, Elvis Jakobsson
#  All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are met:
#
#  1. Redistributions of source code must retain the above copyright notice, this
#  list of conditions and the following disclaimer.
#  2. Redistributions in binary form must reproduce the above copyright notice,
#  this list of conditions and the following disclaimer in the documentation
#  and/or other materials provided with the distribution.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
#  DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE LIABLE FOR
#  ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##############################################################################################
#
#  Alterations to this file by Anders Rehult and Viktor Bernholtz
#
#  To the extent possible under law, the person who associated CC0 with
#  the alterations to this file has waived all copyright and related
#  or neighboring rights to the alterations made to this file.
#
#  You should have received a copy of the CC0 legalcode along with
#  this work.  If not, see
#  <http://creativecommons.org/publicdomain/zero/1.0/>.

import numpy as np
import h5py
import array as arr

def _write_coordinates(h5file, atom_count, coordinates_list, elements, path):
    with h5py.File(h5file, "a") as h5:
        p=0
        for n in range(0,len(atom_count)):
            dataset_name = path+'/Atoms/'+format(n,'04d')
            h5.create_dataset(
                dataset_name,
                data=np.asarray(coordinates_list[p:atom_count[n]+p]),
                dtype=np.float32
            )
            h5[dataset_name].attrs["element"] = elements[n]
            p=p+atom_count[n]
    return

def _write_basis(h5file, basis):
    with h5py.File(h5file, "a") as h5:
        if not "/basis" in h5:
            h5.create_dataset('/basis', data=basis, dtype=np.float32)
    return

def _write_scaling_factor(h5file, scaling_factor):
    with h5py.File(h5file, "a") as h5:
        if not "/scaling_factor" in h5:
            h5.create_dataset('/scaling_factor', data=scaling_factor, dtype=np.float32)
    return

def _write_md(h5file, atom_count, coordinates_list, elements, step):
    with h5py.File(h5file, "a") as h5:
        p=0
        for n in range(0,len(atom_count)):
            dataset_name = '/MD/Atoms/'+format(n,'04d')
            if step == 0:
                h5.create_dataset(
                    dataset_name,
                    data=np.asarray(coordinates_list[p:atom_count[n]+p]),
                    dtype=np.float32,
                    maxshape=(None, 3)
                    )
                h5[dataset_name].attrs["element"] = elements[n]
                h5[dataset_name].attrs["atoms"] = atom_count[n]
                p=p+atom_count[n]
            else:
                dataset = h5[dataset_name]
                dataset.resize((step+1)*atom_count[n],axis=0)
                start = step*atom_count[n]
                dataset[start:] = np.asarray(coordinates_list[p:atom_count[n]+p])
                p=p+atom_count[n]
    return

def _write_steps(h5file, steps):
    with h5py.File(h5file, "a") as h5:
        h5['/MD'].attrs["steps"] = steps
    return

def _write_bandstruct(h5file, band_data, kval_list):
    with h5py.File(h5file, "a") as h5:
        h5.create_dataset('BandStructure/KPoints', data=np.array(kval_list), dtype = np.float32)
        for i, band in enumerate(band_data):
            dataset = h5.create_dataset('Bandstructure/Bands/{}/{}'.format(i, 'Energy'), data=np.array(band), dtype = np.float32)
            dataset.attrs['Unit'] = 'eV'
            dataset.attrs['QuantitySymbol'] = '$E$'
            dataset.attrs['QuantityName'] = 'Energy'
            dataset.attrs['VariableName'] = 'Band {}'.format(i)
            dataset.attrs['VariableSymbol'] = '$B_{{{}}}$'.format(i)

def _write_fermisurface(h5file, kval_list, fermi_energy, reciprocal_lattice_vectors):
    with h5py.File(h5file,"a") as h5:
        for i in range(0, len(kval_list)):
            dataset = h5.create_dataset('/FermiSurface/KPoints/{}/Energy'.format(i),
                                        data = np.array(kval_list[i].energies),
                                        dtype = np.float32)
            dataset.attrs['Unit'] = 'eV'
            dataset.attrs['QuantitySymbol'] = '$E$'
            dataset.attrs['QuantityName'] = 'Energy'
            dataset.attrs['VariableName'] = 'Band {}'.format(i)
            dataset.attrs['VariableSymbol'] = '$B_{{{}}}$'.format(i)

            h5.create_dataset('/FermiSurface/KPoints/{}/Coordinates'.format(i),
                              data = np.array(kval_list[i].coordinates),
                              dtype = np.float32)

        h5.create_dataset('/FermiSurface/FermiEnergy',
                          data = np.float(fermi_energy),
                          dtype = np.float32)

        for i in range(0, len(reciprocal_lattice_vectors)):
            h5.create_dataset('/FermiSurface/ReciprocalLatticeVectors/{}'.format(i),
                              data = np.array([float(x) for x in reciprocal_lattice_vectors[i]]),
                              dtype = np.float32)

def _write_dos(h5file, total, partial, total_data, partial_list, fermi_energy):

    def set_attrs(dataset, VariableName = '', VariableSymbol = '', QuantityName = '', QuantitySymbol = '', Unit = ''):
        dataset.attrs.update({
            'VariableName' : VariableName,
            'VariableSymbol' : VariableSymbol,
            'QuantityName' : QuantityName,
            'QuantitySymbol' : QuantitySymbol,
            'Unit' : Unit
        })

    with h5py.File(h5file, "a") as h5:
        dataset = h5.create_dataset('FermiEnergy', data = np.array(fermi_energy), dtype = np.float32)
        set_attrs(dataset, 'Fermi Energy', '$E_f$', 'Energy', '$E$', 'eV')
        for i, (name, data) in enumerate(zip(total, total_data)):
            dataset = h5.create_dataset('DOS/Total/{}'.format(name), data=np.array(data), dtype = np.float32)
            if name == 'Energy':
                set_attrs(dataset, 'Energy', '$E$', 'Energy', '$E$', 'eV')
            else:
                if name.startswith('Integrated'):
                    set_attrs(dataset, name, '', 'Integrated Density of States', '$\\int D$', 'states')
                else:
                    set_attrs(dataset, name, '', 'Density of States', '$D$', 'states/eV')
        for i, partial_data in enumerate(partial_list):
            # Write all available DOS. Stop when partial or partial_data runs out (if directories for for example f-DOS
            # exist but no corresponding data exists, stop.)
            for (name, data) in zip(partial, partial_data):
                dataset = h5.create_dataset('DOS/Partial/{}/{}'.format(i, name), data=np.array(data), dtype = np.float32)
                if name == 'Energy':
                    set_attrs(dataset, 'Energy', '$E$', 'Energy', '$E$', 'eV')
                else:
                    if name.startswith('Integrated'):
                        set_attrs(dataset, name, '', 'Integrated Density of States', '$\\int D$', 'states')
                    else:
                        set_attrs(dataset, name, '', 'Density of States', '$D$', 'states/eV')

def _write_volume(h5file, i, array, data_dim, hdfgroup):
    with h5py.File(h5file, "a") as h5:
        if array:
            h5.create_dataset('{}/{}'.format(hdfgroup, i), data = np.reshape(array, (data_dim[2],data_dim[1],data_dim[0])), dtype=np.float32)
        else:
            h5['{}/final'.format(hdfgroup)] = h5['{}/{}'.format(hdfgroup, i-1)]
    return

def _write_incar(h5file, incar_data):
    with h5py.File(h5file, 'a') as h5:
        for key, value in incar_data.items():
            h5.create_dataset("/incar/{}".format(key), data=value)

def _write_parcharges(h5file, array_tot, data_dim_tot, array_mag, data_dim_mag, band_nr):
    with h5py.File(h5file, "a") as h5:
        h5.create_dataset('PARCHG/Bands/{}/total'.format(band_nr), data = np.reshape(array_tot, (data_dim_tot[2],data_dim_tot[1],data_dim_tot[0])), dtype=np.float32)
        if not np.size(array_mag) == 0:
            h5.create_dataset('PARCHG/Bands/{}/magnetic'.format(band_nr), data = np.reshape(array_mag, (data_dim_mag[2],data_dim_mag[1],data_dim_mag[0])), dtype=np.float32)
    return

def _write_pcdat(h5file, pcdat_data, APACO_val, NPACO_val):
    #   The function is called to write data from PCDAT to HDF5-file. A dataset is created for each element in the system.

    #    Parameters
    #    __________
    #    h5file: str
    #        String containing path to HDF5-file.

    #    pcdat_data:
    #        Is a dictionary with the structure {'element_type':[PKF_values]}. PKF_values are the average number of atoms, of specific element_type. If the system has elements 'Si', 'Au'       and 'K', the dictionary will look like {'Si':float[x], 'Au':float[x], 'K':float[x]} where the float[x] is a list with the PKF values.
    #
    #    APACO_val:
    #        The value of APACO in INCAR if set, default value 16 (Å), otherwise. It sets the maximum distance in the evaluation of the pair-correlation function
    #
    #    NPACO_val:
    #        The value of NPACO in INCAR if set, default value 256 slots otherwise. It sets the number of slots in the pair-correlation function written to PCDAT.

    #    Return
    #    ______
    #    Bool: True if parsed, False otherwise.

    with h5py.File(h5file, "a") as h5:
        dset_name = "PairCorrelationFunc/Iterations"
        normal_arr = arr.array('f', [])

        #creating x-values, equally spaced (space APACO_val/NPACO_val)
        for i in range(0, NPACO_val):
            x = (APACO_val / NPACO_val) * i
            normal_arr.append(x)
        value = np.asarray(normal_arr)
        h5.create_dataset(dset_name, data=value, dtype=np.float32)

        # create dataset for paircorrelation function for every element.
        for n in range(len(pcdat_data)):
            # By making a dict a list, its keys are accessable
            element_symbol = list(pcdat_data)[n]
            dset_name = "PairCorrelationFunc/Elements/{}".format(element_symbol)

            #Iterate for every time frame.
            original_list = pcdat_data[element_symbol]
            value_list = []

            for i in range(len(original_list)):
                fill_list = []
                for i in range(NPACO_val):
                    fill_list.append(original_list[0])
                    del original_list[0]

                value_list.append(fill_list)

                if len(original_list) == 0:
                    break
                else:
                    del original_list[0]

            # if NPACO= 256, there are 3 time frames (t1,t2 and t3) then value.shape = (3,256)
            # values from row one correspond to pair correlation of t1, row 2 for t2 e.t.c.
            value = np.asarray(value_list)
            h5.create_dataset(dset_name, data=value, dtype=np.float32)
            h5[dset_name].attrs["element"] = element_symbol

    return None

<|endoftext|>


<commit_before>from __future__ import unicode_literals

import frappe
from frappe import _
from frappe.utils import date_diff, add_months, today, getdate, add_days, flt
from erpnext.accounts.utils import get_account_currency
from erpnext.accounts.general_ledger import make_gl_entries

def validate_service_stop_date(doc):
	''' Validates service_stop_date for Purchase Invoice and Sales Invoice '''

	enable_check = "enable_deferred_revenue" \
		if doc.doctype=="Sales Invoice" else "enable_deferred_expense"

	old_stop_dates = {}
	old_doc = frappe.db.get_all("{0} Item".format(doc.doctype),
		{"parent": doc.name}, ["name", "service_stop_date"])

	for d in old_doc:
		old_stop_dates[d.name] = d.service_stop_date or ""

	for item in doc.items:
		if not item.get(enable_check): continue

		if item.service_stop_date:
			if date_diff(item.service_stop_date, item.service_start_date) < 0:
				frappe.throw(_("Service Stop Date cannot be before Service Start Date"))

			if date_diff(item.service_stop_date, item.service_end_date) > 0:
				frappe.throw(_("Service Stop Date cannot be after Service End Date"))

		if old_stop_dates and old_stop_dates[item.name] and item.service_stop_date!=old_stop_dates[item.name]:
			frappe.throw(_("Cannot change Service Stop Date for item in row {0}".format(item.idx)))

def convert_deferred_expense_to_expense(start_date=None, end_date=None):
	# check for the purchase invoice for which GL entries has to be done
	invoices = frappe.db.sql_list('''
		select distinct parent from `tabPurchase Invoice Item` where service_start_date<=%s and service_end_date>=%s
		and enable_deferred_expense = 1 and docstatus = 1 and ifnull(amount, 0) > 0
	''', (end_date or today(), start_date or add_months(today(), -1)))

	# For each invoice, book deferred expense
	for invoice in invoices:
		doc = frappe.get_doc("Purchase Invoice", invoice)
		book_deferred_income_or_expense(doc, start_date, end_date)

def convert_deferred_revenue_to_income(start_date=None, end_date=None):
	# check for the sales invoice for which GL entries has to be done
	invoices = frappe.db.sql_list('''
		select distinct parent from `tabSales Invoice Item` where service_start_date<=%s and service_end_date>=%s
		and enable_deferred_revenue = 1 and docstatus = 1 and ifnull(amount, 0) > 0
	''', (end_date or today(), start_date or add_months(today(), -1)))

	# For each invoice, book deferred revenue
	for invoice in invoices:
		doc = frappe.get_doc("Sales Invoice", invoice)
		book_deferred_income_or_expense(doc, start_date, end_date)

def get_booking_dates(doc, item, start_date=None, end_date=None):
	deferred_account = "deferred_revenue_account" if doc.doctype=="Sales Invoice" else "deferred_expense_account"
	last_gl_entry, skip = False, False

	booking_end_date = getdate(add_days(today(), -1)) if not end_date else end_date
	if booking_end_date < item.service_start_date or \
		(item.service_stop_date and booking_end_date.month > item.service_stop_date.month):
		return None, None, None, True
	elif booking_end_date >= item.service_end_date:
		last_gl_entry = True
		booking_end_date = item.service_end_date
	elif item.service_stop_date and item.service_stop_date <= booking_end_date:
		last_gl_entry = True
		booking_end_date = item.service_stop_date

	booking_start_date = getdate(add_months(today(), -1)) if not start_date else start_date
	booking_start_date = booking_start_date \
		if booking_start_date > item.service_start_date else item.service_start_date

	prev_gl_entry = frappe.db.sql('''
		select name, posting_date from `tabGL Entry` where company=%s and account=%s and
		voucher_type=%s and voucher_no=%s and voucher_detail_no=%s
		order by posting_date desc limit 1
	''', (doc.company, item.get(deferred_account), doc.doctype, doc.name, item.name), as_dict=True)

	if not prev_gl_entry and item.service_start_date < booking_start_date:
		booking_start_date = item.service_start_date
	elif prev_gl_entry:
		booking_start_date = getdate(add_days(prev_gl_entry[0].posting_date, 1))
		skip = True if booking_start_date > booking_end_date else False

	return last_gl_entry, booking_start_date, booking_end_date, skip

def calculate_amount_and_base_amount(doc, item, last_gl_entry, total_days, total_booking_days):
	account_currency = get_account_currency(item.expense_account)

	if doc.doctype == "Sales Invoice":
		total_credit_debit, total_credit_debit_currency = "debit", "debit_in_account_currency"
		deferred_account = "deferred_revenue_account"
	else:
		total_credit_debit, total_credit_debit_currency = "credit", "credit_in_account_currency"
		deferred_account = "deferred_expense_account"

	amount, base_amount = 0, 0
	if not last_gl_entry:
		base_amount = flt(item.base_net_amount*total_booking_days/flt(total_days), item.precision("base_net_amount"))
		if account_currency==doc.company_currency:
			amount = base_amount
		else:
			amount = flt(item.net_amount*total_booking_days/flt(total_days), item.precision("net_amount"))
	else:
		gl_entries_details = frappe.db.sql('''
			select sum({0}) as total_credit, sum({1}) as total_credit_in_account_currency, voucher_detail_no
			from `tabGL Entry` where company=%s and account=%s and voucher_type=%s and voucher_no=%s and voucher_detail_no=%s
			group by voucher_detail_no
		'''.format(total_credit_debit, total_credit_debit_currency),
			(doc.company, item.get(deferred_account), doc.doctype, doc.name, item.name), as_dict=True)

		already_booked_amount = gl_entries_details[0].total_credit if gl_entries_details else 0
		base_amount = flt(item.base_net_amount - already_booked_amount, item.precision("base_net_amount"))
		if account_currency==doc.company_currency:
			amount = base_amount
		else:
			already_booked_amount_in_account_currency = gl_entries_details[0].total_credit_in_account_currency if gl_entries_details else 0
			amount = flt(item.net_amount - already_booked_amount_in_account_currency, item.precision("net_amount"))

	return amount, base_amount

def book_deferred_income_or_expense(doc, start_date=None, end_date=None):
	# book the expense/income on the last day, but it will be trigger on the 1st of month at 12:00 AM
	# start_date: 1st of the last month or the start date
	# end_date: end_date or today-1

	gl_entries = []
	for item in doc.get('items'):
		skip = False
		last_gl_entry, booking_start_date, booking_end_date, skip = \
			get_booking_dates(doc, item, start_date, end_date)

		if skip: continue
		total_days = date_diff(item.service_end_date, item.service_start_date)
		total_booking_days = date_diff(booking_end_date, booking_start_date) + 1

		account_currency = get_account_currency(item.expense_account)
		amount, base_amount = calculate_amount_and_base_amount(doc, item, last_gl_entry, total_days, total_booking_days)

		if doc.doctype == "Sales Invoice":
			against, project = doc.customer, doc.project
			credit_account, debit_account = item.income_account, item.deferred_revenue_account
		else:
			against, project = doc.supplier, item.project
			credit_account, debit_account = item.deferred_expense_account, item.expense_account

		# GL Entry for crediting the amount in the deferred expense
		gl_entries.append(
			doc.get_gl_dict({
				"account": credit_account,
				"against": against,
				"credit": base_amount,
				"credit_in_account_currency": amount,
				"cost_center": item.cost_center,
				'posting_date': booking_end_date,
				'project': project
			}, account_currency)
		)
		# GL Entry to debit the amount from the expense
		gl_entries.append(
			doc.get_gl_dict({
				"account": debit_account,
				"against": against,
				"debit": base_amount,
				"debit_in_account_currency": amount,
				"cost_center": item.cost_center,
				"voucher_detail_no": item.name,
				'posting_date': booking_end_date,
				'project': project
			}, account_currency)
		)

	if gl_entries:
		make_gl_entries(gl_entries, cancel=(doc.docstatus == 2), merge_entries=True)
<commit_msg>add voucher detail no while creating gl entry<commit_after>from __future__ import unicode_literals

import frappe
from frappe import _
from frappe.utils import date_diff, add_months, today, getdate, add_days, flt
from erpnext.accounts.utils import get_account_currency
from erpnext.accounts.general_ledger import make_gl_entries

def validate_service_stop_date(doc):
	''' Validates service_stop_date for Purchase Invoice and Sales Invoice '''

	enable_check = "enable_deferred_revenue" \
		if doc.doctype=="Sales Invoice" else "enable_deferred_expense"

	old_stop_dates = {}
	old_doc = frappe.db.get_all("{0} Item".format(doc.doctype),
		{"parent": doc.name}, ["name", "service_stop_date"])

	for d in old_doc:
		old_stop_dates[d.name] = d.service_stop_date or ""

	for item in doc.items:
		if not item.get(enable_check): continue

		if item.service_stop_date:
			if date_diff(item.service_stop_date, item.service_start_date) < 0:
				frappe.throw(_("Service Stop Date cannot be before Service Start Date"))

			if date_diff(item.service_stop_date, item.service_end_date) > 0:
				frappe.throw(_("Service Stop Date cannot be after Service End Date"))

		if old_stop_dates and old_stop_dates[item.name] and item.service_stop_date!=old_stop_dates[item.name]:
			frappe.throw(_("Cannot change Service Stop Date for item in row {0}".format(item.idx)))

def convert_deferred_expense_to_expense(start_date=None, end_date=None):
	# check for the purchase invoice for which GL entries has to be done
	invoices = frappe.db.sql_list('''
		select distinct parent from `tabPurchase Invoice Item` where service_start_date<=%s and service_end_date>=%s
		and enable_deferred_expense = 1 and docstatus = 1 and ifnull(amount, 0) > 0
	''', (end_date or today(), start_date or add_months(today(), -1)))

	# For each invoice, book deferred expense
	for invoice in invoices:
		doc = frappe.get_doc("Purchase Invoice", invoice)
		book_deferred_income_or_expense(doc, start_date, end_date)

def convert_deferred_revenue_to_income(start_date=None, end_date=None):
	# check for the sales invoice for which GL entries has to be done
	invoices = frappe.db.sql_list('''
		select distinct parent from `tabSales Invoice Item` where service_start_date<=%s and service_end_date>=%s
		and enable_deferred_revenue = 1 and docstatus = 1 and ifnull(amount, 0) > 0
	''', (end_date or today(), start_date or add_months(today(), -1)))

	# For each invoice, book deferred revenue
	for invoice in invoices:
		doc = frappe.get_doc("Sales Invoice", invoice)
		book_deferred_income_or_expense(doc, start_date, end_date)

def get_booking_dates(doc, item, start_date=None, end_date=None):
	deferred_account = "deferred_revenue_account" if doc.doctype=="Sales Invoice" else "deferred_expense_account"
	last_gl_entry, skip = False, False

	booking_end_date = getdate(add_days(today(), -1)) if not end_date else end_date
	if booking_end_date < item.service_start_date or \
		(item.service_stop_date and booking_end_date.month > item.service_stop_date.month):
		return None, None, None, True
	elif booking_end_date >= item.service_end_date:
		last_gl_entry = True
		booking_end_date = item.service_end_date
	elif item.service_stop_date and item.service_stop_date <= booking_end_date:
		last_gl_entry = True
		booking_end_date = item.service_stop_date

	booking_start_date = getdate(add_months(today(), -1)) if not start_date else start_date
	booking_start_date = booking_start_date \
		if booking_start_date > item.service_start_date else item.service_start_date

	prev_gl_entry = frappe.db.sql('''
		select name, posting_date from `tabGL Entry` where company=%s and account=%s and
		voucher_type=%s and voucher_no=%s and voucher_detail_no=%s
		order by posting_date desc limit 1
	''', (doc.company, item.get(deferred_account), doc.doctype, doc.name, item.name), as_dict=True)

	if not prev_gl_entry and item.service_start_date < booking_start_date:
		booking_start_date = item.service_start_date
	elif prev_gl_entry:
		booking_start_date = getdate(add_days(prev_gl_entry[0].posting_date, 1))
		skip = True if booking_start_date > booking_end_date else False

	return last_gl_entry, booking_start_date, booking_end_date, skip

def calculate_amount_and_base_amount(doc, item, last_gl_entry, total_days, total_booking_days):
	account_currency = get_account_currency(item.expense_account)

	if doc.doctype == "Sales Invoice":
		total_credit_debit, total_credit_debit_currency = "debit", "debit_in_account_currency"
		deferred_account = "deferred_revenue_account"
	else:
		total_credit_debit, total_credit_debit_currency = "credit", "credit_in_account_currency"
		deferred_account = "deferred_expense_account"

	amount, base_amount = 0, 0
	if not last_gl_entry:
		base_amount = flt(item.base_net_amount*total_booking_days/flt(total_days), item.precision("base_net_amount"))
		if account_currency==doc.company_currency:
			amount = base_amount
		else:
			amount = flt(item.net_amount*total_booking_days/flt(total_days), item.precision("net_amount"))
	else:
		gl_entries_details = frappe.db.sql('''
			select sum({0}) as total_credit, sum({1}) as total_credit_in_account_currency, voucher_detail_no
			from `tabGL Entry` where company=%s and account=%s and voucher_type=%s and voucher_no=%s and voucher_detail_no=%s
			group by voucher_detail_no
		'''.format(total_credit_debit, total_credit_debit_currency),
			(doc.company, item.get(deferred_account), doc.doctype, doc.name, item.name), as_dict=True)

		already_booked_amount = gl_entries_details[0].total_credit if gl_entries_details else 0
		base_amount = flt(item.base_net_amount - already_booked_amount, item.precision("base_net_amount"))
		if account_currency==doc.company_currency:
			amount = base_amount
		else:
			already_booked_amount_in_account_currency = gl_entries_details[0].total_credit_in_account_currency if gl_entries_details else 0
			amount = flt(item.net_amount - already_booked_amount_in_account_currency, item.precision("net_amount"))

	return amount, base_amount

def book_deferred_income_or_expense(doc, start_date=None, end_date=None):
	# book the expense/income on the last day, but it will be trigger on the 1st of month at 12:00 AM
	# start_date: 1st of the last month or the start date
	# end_date: end_date or today-1

	gl_entries = []
	for item in doc.get('items'):
		skip = False
		last_gl_entry, booking_start_date, booking_end_date, skip = \
			get_booking_dates(doc, item, start_date, end_date)

		if skip: continue
		total_days = date_diff(item.service_end_date, item.service_start_date)
		total_booking_days = date_diff(booking_end_date, booking_start_date) + 1

		account_currency = get_account_currency(item.expense_account)
		amount, base_amount = calculate_amount_and_base_amount(doc, item, last_gl_entry, total_days, total_booking_days)

		if doc.doctype == "Sales Invoice":
			against, project = doc.customer, doc.project
			credit_account, debit_account = item.income_account, item.deferred_revenue_account
		else:
			against, project = doc.supplier, item.project
			credit_account, debit_account = item.deferred_expense_account, item.expense_account

		# GL Entry for crediting the amount in the deferred expense
		gl_entries.append(
			doc.get_gl_dict({
				"account": credit_account,
				"against": against,
				"credit": base_amount,
				"credit_in_account_currency": amount,
				"cost_center": item.cost_center,
				"voucher_detail_no": item.name,
				'posting_date': booking_end_date,
				'project': project
			}, account_currency)
		)
		# GL Entry to debit the amount from the expense
		gl_entries.append(
			doc.get_gl_dict({
				"account": debit_account,
				"against": against,
				"debit": base_amount,
				"debit_in_account_currency": amount,
				"cost_center": item.cost_center,
				"voucher_detail_no": item.name,
				'posting_date': booking_end_date,
				'project': project
			}, account_currency)
		)

	if gl_entries:
		make_gl_entries(gl_entries, cancel=(doc.docstatus == 2), merge_entries=True)
<|endoftext|>


<commit_before>"""The pyexperiment module - quick and clean experiments with Python.
"""
<commit_msg>Add version to __init__.py<commit_after>"""The pyexperiment module - quick and clean experiments with Python.
"""
__version__ = '0.1.19'
<|endoftext|>


<commit_before><commit_msg>Fix 2.3 compatibility.<commit_after><|endoftext|>


<commit_before><commit_msg>nif_neuron fixed import and namespace bugs<commit_after><|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""ECG (waveform) Dicom module

Read and plot images from DICOM ECG waveforms.
"""

from datetime import datetime
from matplotlib import pylab as plt
from scipy.signal import butter, lfilter
import numpy as np
import dicom
import struct
import cStringIO
import requests

try:
    from ecgconfig import WADOSERVER, LAYOUT, PRODUCER
except ImportError:
    WADOSERVER = "http://example.com"
    LAYOUT = {'3X4_1': [[0, 3, 6, 9],
                        [1, 4, 7, 10],
                        [2, 5, 8, 11],
                        [1]],
              '3X4':   [[0, 3, 6, 9],
                        [1, 4, 7, 10],
                        [2, 5, 8, 11]],
              '12X1':  [[0],
                        [1],
                        [2],
                        [3],
                        [4],
                        [5],
                        [6],
                        [7],
                        [8],
                        [9],
                        [10],
                        [11]]}

    PRODUCER = "Example Producer"

__author__ = "Marco De Benedetto"
__email__ = "debe@galliera.it"

def butter_lowpass(highcut, sampfreq, order=5):
    nyquist_freq = .5 * sampfreq
    high = highcut / nyquist_freq
    num, denom = butter(order, high, btype='lowpass')
    return num, denom

def butter_lowpass_filter(data, highcut, sampfreq, order):
    num, denom = butter_lowpass(highcut, sampfreq, order=order)
    return lfilter(num, denom, data)

class ECG(object):

    paper_w, paper_h = 297.0, 210.0

    # Dimensions in mm of plot area
    width = 250.0
    height = 170.0
    margin_left = margin_right = .5 * (paper_w - width)
    margin_bottom = 10

    # Normalized in [0, 1]
    left = margin_left/paper_w
    right = left+width/paper_w
    bottom = margin_bottom/paper_h
    top = bottom+height/paper_h

    def __init__(self, source):

        def err(msg):
            raise(Exception(msg))

        def wadoget(stu, ser, obj):

            payload = {'requestType': 'WADO', 'studyUID': stu,
                       'seriesUID': ser, 'objectUID': obj}
            headers = {'content-type': 'application/json'}

            resp = requests.get(WADOSERVER, params=payload, headers=headers)
            return cStringIO.StringIO(resp.content)

        if isinstance(source, dict):
            # dictionary of stu, ser, obj
            if set(source.keys()) == set(('stu', 'ser', 'obj')):
                inputdata = wadoget(**source)
            else:
                err("source must be a dictionary of stu, ser and obj")
        elif isinstance(source, basestring) or getattr(source, 'getvalue'):
            # it is a filename or a (StringIO or cStringIO buffer)
            inputdata = source
        else:
            # What is it?
            err("`source´ must be a path/to/file.ext string\n" +
                "or a dictionary of stu, ser and obj")

        self.dicom = dicom.read_file(inputdata)
        sequence_item = self.dicom.WaveformSequence[0]

        assert(sequence_item.WaveformSampleInterpretation == 'SS')
        assert(sequence_item.WaveformBitsAllocated == 16)

        self.channel_definitions = sequence_item.ChannelDefinitionSequence
        self.wavewform_data = sequence_item.WaveformData
        self.channels_no = sequence_item.NumberOfWaveformChannels
        self.samples = sequence_item.NumberOfWaveformSamples
        self.sampling_frequency = sequence_item.SamplingFrequency
        self.duration = self.samples / self.sampling_frequency
        self.mm_s = self.width / self.duration
        self.signals = self._signals()
        self.fig, self.axis = self.create_figure()

    def create_figure(self):
        """
        Prepare figure and axes
        """

        # Init figure and axes
        fig = plt.figure(tight_layout=False)
        axes = fig.add_subplot(1, 1, 1)
        # axes.set_frame_on(False)
        fig.subplots_adjust(left=self.left, right=self.right,
                            top=self.top, bottom=self.bottom)
        axes.set_ylim([0, self.height])
        # we want to plot N points, where N=number of samples
        axes.set_xlim([0, self.samples-1])
        return fig, axes

    def _signals(self):
        """
        sequence_item := dicom.dataset.FileDataset.WaveformData[n]
        Return a list of signals.
        """

        factor = np.zeros(self.channels_no) + 1
        baseln = np.zeros(self.channels_no)
        units = []
        for idx in range(self.channels_no):
            definition = self.channel_definitions[idx]

            assert(definition.WaveformBitsStored == 16)

            if definition.get('ChannelSensitivity'):
                factor[idx] = (
                    float(definition.ChannelSensitivity) *
                    float(definition.ChannelSensitivityCorrectionFactor))
            if definition.get('ChannelBaseline'):
                baseln[idx] = float(definition.get('ChannelBaseline'))

            units.append(
                definition.ChannelSensitivityUnitsSequence[0].CodeValue)

        signals = np.asarray(
            struct.unpack('<' + str(len(self.wavewform_data)/2) + 'h',
                          self.wavewform_data), dtype=np.float32).reshape(
                              self.samples, self.channels_no).transpose()

        for channel in range(self.channels_no):
            signals[channel] = (
                (signals[channel] + baseln[channel]) * factor[channel]
            )

        high = 40.0

        # conversion factor to obtain millivolts values
        millivolts = {'uV': 1000.0, 'mV': 1.0}

        for i, signal in enumerate(signals):
            signals[i] = butter_lowpass_filter(
                np.asarray(signal),
                high,
                self.sampling_frequency,
                order=1) / millivolts[units[i]]

        return signals

    def draw_grid(self):

        #self.axis.xaxis.set_minor_locator(plt.LinearLocator(self.width+1))
        self.axis.xaxis.set_major_locator(plt.LinearLocator(self.width/5+1))
        #self.axis.yaxis.set_minor_locator(plt.LinearLocator(self.height+1))
        self.axis.yaxis.set_major_locator(plt.LinearLocator(self.height/5+1))

        color = {'minor': '#ff5333', 'major': '#d43d1a'}
        linewidth = {'minor': .2, 'major': .6}
        alpha = 1

        for axe in 'x', 'y':
            for which in 'major', 'minor':
                self.axis.grid(which=which, axis=axe,
                               linewidth=linewidth[which],
                               linestyle='-', color=color[which], alpha=alpha)
                self.axis.tick_params(which=which, axis=axe,
                                      color=color[which],
                                      bottom='off', top='off',
                                      left='off', right='off')

        self.axis.set_xticklabels([])
        self.axis.set_yticklabels([])

    def legend(self):

        ecgdata = {}
        for was in self.dicom.WaveformAnnotationSequence:
            if was.get('ConceptNameCodeSequence'):
                cncs = was.ConceptNameCodeSequence[0]
                if cncs.CodeMeaning in ('QT Interval',
                                        'QTc Interval',
                                        'RR Interval',
                                        'QRS Duration',
                                        'QRS Axis',
                                        'T Axis',
                                        'P Axis',
                                        'PR Interval'):
                    ecgdata[cncs.CodeMeaning] = str(was.NumericValue)

        bpm = ''
        if ecgdata.get('RR Interval'):
            bpm = str(int(round(60.0 / self.duration *
                      self.samples /
                      int(ecgdata['RR Interval'])))) + ' BPM\n'

        return 'Ventr. Freq.: ' + \
               bpm + \
               'PR Interval: ' + ecgdata.get('PR Interval', '') + ' ms\n' + \
               'QRS Duration: ' + ecgdata.get('QRS Duration', '') + ' ms\n' + \
               'QT/QTc: ' + \
               ecgdata.get('QT Interval', '') + '/' + \
               ecgdata.get('QTc Interval', '') + ' ms\n' + \
               'P-R-T Axis: ' + \
               ecgdata.get('P Axis', '') + ' ' + \
               ecgdata.get('QRS Axis', '') + ' ' + \
               ecgdata.get('T Axis', '')

    def print_info(self):

        try:
            pat_surname, pat_firstname = self.dicom.PatientName.split('^')
        except ValueError:
            pat_surname = self.dicom.PatientName
            pat_firstname = ''

        pat_name = ' '.join((pat_surname, pat_firstname.title()))

        pat_id = self.dicom.PatientID
        pat_sex = self.dicom.PatientSex
        pat_bdate = datetime.strptime(
            self.dicom.PatientBirthDate, '%Y%m%d').strftime("%e %b %Y")

        info = "%s (%s) sex: %s\nbirthdate: %s" % (pat_name, pat_id,
                                                   pat_sex, pat_bdate)
        plt.figtext(0.08, 0.865, info, fontsize=12)

        info = "total time: %s s | sample freq.: %s Hz" % (
            self.duration, self.sampling_frequency)
        plt.figtext(0.08, 0.02, info, fontsize=10)

        # TODO: the bandpass filter 0.05-40 Hz will have to became a parameter
        info = "%s mm/s %s mm/mV 0.05-40 Hz" % (self.mm_s, self.mm_mv)
        plt.figtext(0.78, 0.02, info, fontsize=10)

        plt.figtext(0.5, 0.865, self.legend(), fontsize=10)

        info = PRODUCER + '\nStudy date: '
        date_str = (self.dicom.InstanceCreationDate +
                    self.dicom.InstanceCreationTime)
        info += datetime.strftime(datetime.strptime(date_str,
                                                    '%Y%m%d%H%M%S'),
                                  '%d %b %Y %H:%M')
        plt.figtext(0.03, 0.94, info, fontsize=12)

    def save(self, outputfile=None, outformat=None):

        def _save(output):
            plt.savefig(output, dpi=300, format=outformat,
                        papertype='a4', orientation='landscape')

        if outputfile:
            _save(outputfile)
        else:
            output = cStringIO.StringIO()
            _save(output)
            return output.getvalue()

    def plot(self, layoutid, mm_mv):

        self.mm_mv = mm_mv

        layout = LAYOUT[layoutid]
        rows = len(layout)

        for numrow, row in enumerate(layout):
            columns = len(row)
            h_delta = self.samples / columns
            signal = np.ndarray(0)
            row_height = self.height / rows
            v_delta = round(self.height * (1 - 1.0/(rows*2)) -
                            numrow*(self.height/rows))

            v_delta = (v_delta + 2.5) - (v_delta + 2.5) % 5
            chunk_size = int(self.samples/len(row))

            for numcol, signum in enumerate(row):
                left = numcol*chunk_size
                right = (1+numcol)*chunk_size
                h = h_delta * numcol
                plt.plot([h, h], [v_delta-row_height/2.6,
                                  v_delta+row_height/2.6],
                         'k-', lw=1, color='blue', zorder=50)
                signal = np.concatenate((
                    signal,
                    mm_mv*self.signals[signum][left:right])
                )
                cseq = self.channel_definitions[signum].ChannelSourceSequence
                meaning = cseq[0].CodeMeaning.replace(
                    'Lead', '').replace('(Einthoven)', '')
                self.axis.text(h+40, v_delta+row_height/3,
                               meaning, color='b', zorder=50)

            self.axis.plot(signal+v_delta, linewidth=.6, color='black',
                           antialiased=True, zorder=10)

        # A4 size in inches
        self.fig.set_size_inches(11.69, 8.27)

    def draw(self, layoutid, mm_mv=10.0):
        self.draw_grid()
        self.plot(layoutid, mm_mv)
        self.print_info()
<commit_msg>Draw signals chunks directly.<commit_after># -*- coding: utf-8 -*-
"""ECG (waveform) Dicom module

Read and plot images from DICOM ECG waveforms.
"""

from datetime import datetime
from matplotlib import pylab as plt
from scipy.signal import butter, lfilter
import numpy as np
import dicom
import struct
import cStringIO
import requests

try:
    from ecgconfig import WADOSERVER, LAYOUT, PRODUCER
except ImportError:
    WADOSERVER = "http://example.com"
    LAYOUT = {'3X4_1': [[0, 3, 6, 9],
                        [1, 4, 7, 10],
                        [2, 5, 8, 11],
                        [1]],
              '3X4':   [[0, 3, 6, 9],
                        [1, 4, 7, 10],
                        [2, 5, 8, 11]],
              '12X1':  [[0],
                        [1],
                        [2],
                        [3],
                        [4],
                        [5],
                        [6],
                        [7],
                        [8],
                        [9],
                        [10],
                        [11]]}

    PRODUCER = "Example Producer"

__author__ = "Marco De Benedetto"
__email__ = "debe@galliera.it"

def butter_lowpass(highcut, sampfreq, order=5):
    nyquist_freq = .5 * sampfreq
    high = highcut / nyquist_freq
    num, denom = butter(order, high, btype='lowpass')
    return num, denom

def butter_lowpass_filter(data, highcut, sampfreq, order):
    num, denom = butter_lowpass(highcut, sampfreq, order=order)
    return lfilter(num, denom, data)

class ECG(object):

    paper_w, paper_h = 297.0, 210.0

    # Dimensions in mm of plot area
    width = 250.0
    height = 170.0
    margin_left = margin_right = .5 * (paper_w - width)
    margin_bottom = 10

    # Normalized in [0, 1]
    left = margin_left/paper_w
    right = left+width/paper_w
    bottom = margin_bottom/paper_h
    top = bottom+height/paper_h

    def __init__(self, source):

        def err(msg):
            raise(Exception(msg))

        def wadoget(stu, ser, obj):

            payload = {'requestType': 'WADO', 'studyUID': stu,
                       'seriesUID': ser, 'objectUID': obj}
            headers = {'content-type': 'application/json'}

            resp = requests.get(WADOSERVER, params=payload, headers=headers)
            return cStringIO.StringIO(resp.content)

        if isinstance(source, dict):
            # dictionary of stu, ser, obj
            if set(source.keys()) == set(('stu', 'ser', 'obj')):
                inputdata = wadoget(**source)
            else:
                err("source must be a dictionary of stu, ser and obj")
        elif isinstance(source, basestring) or getattr(source, 'getvalue'):
            # it is a filename or a (StringIO or cStringIO buffer)
            inputdata = source
        else:
            # What is it?
            err("`source´ must be a path/to/file.ext string\n" +
                "or a dictionary of stu, ser and obj")

        self.dicom = dicom.read_file(inputdata)
        sequence_item = self.dicom.WaveformSequence[0]

        assert(sequence_item.WaveformSampleInterpretation == 'SS')
        assert(sequence_item.WaveformBitsAllocated == 16)

        self.channel_definitions = sequence_item.ChannelDefinitionSequence
        self.wavewform_data = sequence_item.WaveformData
        self.channels_no = sequence_item.NumberOfWaveformChannels
        self.samples = sequence_item.NumberOfWaveformSamples
        self.sampling_frequency = sequence_item.SamplingFrequency
        self.duration = self.samples / self.sampling_frequency
        self.mm_s = self.width / self.duration
        self.signals = self._signals()
        self.fig, self.axis = self.create_figure()

    def create_figure(self):
        """
        Prepare figure and axes
        """

        # Init figure and axes
        fig = plt.figure(tight_layout=False)
        axes = fig.add_subplot(1, 1, 1)
        # axes.set_frame_on(False)
        fig.subplots_adjust(left=self.left, right=self.right,
                            top=self.top, bottom=self.bottom)
        axes.set_ylim([0, self.height])
        # we want to plot N points, where N=number of samples
        axes.set_xlim([0, self.samples-1])
        return fig, axes

    def _signals(self):
        """
        sequence_item := dicom.dataset.FileDataset.WaveformData[n]
        Return a list of signals.
        """

        factor = np.zeros(self.channels_no) + 1
        baseln = np.zeros(self.channels_no)
        units = []
        for idx in range(self.channels_no):
            definition = self.channel_definitions[idx]

            assert(definition.WaveformBitsStored == 16)

            if definition.get('ChannelSensitivity'):
                factor[idx] = (
                    float(definition.ChannelSensitivity) *
                    float(definition.ChannelSensitivityCorrectionFactor))
            if definition.get('ChannelBaseline'):
                baseln[idx] = float(definition.get('ChannelBaseline'))

            units.append(
                definition.ChannelSensitivityUnitsSequence[0].CodeValue)

        signals = np.asarray(
            struct.unpack('<' + str(len(self.wavewform_data)/2) + 'h',
                          self.wavewform_data), dtype=np.float32).reshape(
                              self.samples, self.channels_no).transpose()

        for channel in range(self.channels_no):
            signals[channel] = (
                (signals[channel] + baseln[channel]) * factor[channel]
            )

        high = 40.0

        # conversion factor to obtain millivolts values
        millivolts = {'uV': 1000.0, 'mV': 1.0}

        for i, signal in enumerate(signals):
            signals[i] = butter_lowpass_filter(
                np.asarray(signal),
                high,
                self.sampling_frequency,
                order=1) / millivolts[units[i]]

        return signals

    def draw_grid(self):

        #self.axis.xaxis.set_minor_locator(plt.LinearLocator(self.width+1))
        self.axis.xaxis.set_major_locator(plt.LinearLocator(self.width/5+1))
        #self.axis.yaxis.set_minor_locator(plt.LinearLocator(self.height+1))
        self.axis.yaxis.set_major_locator(plt.LinearLocator(self.height/5+1))

        color = {'minor': '#ff5333', 'major': '#d43d1a'}
        linewidth = {'minor': .2, 'major': .6}
        alpha = 1

        for axe in 'x', 'y':
            for which in 'major', 'minor':
                self.axis.grid(which=which, axis=axe,
                               linewidth=linewidth[which],
                               linestyle='-', color=color[which], alpha=alpha)
                self.axis.tick_params(which=which, axis=axe,
                                      color=color[which],
                                      bottom='off', top='off',
                                      left='off', right='off')

        self.axis.set_xticklabels([])
        self.axis.set_yticklabels([])

    def legend(self):

        ecgdata = {}
        for was in self.dicom.WaveformAnnotationSequence:
            if was.get('ConceptNameCodeSequence'):
                cncs = was.ConceptNameCodeSequence[0]
                if cncs.CodeMeaning in ('QT Interval',
                                        'QTc Interval',
                                        'RR Interval',
                                        'QRS Duration',
                                        'QRS Axis',
                                        'T Axis',
                                        'P Axis',
                                        'PR Interval'):
                    ecgdata[cncs.CodeMeaning] = str(was.NumericValue)

        bpm = ''
        if ecgdata.get('RR Interval'):
            bpm = str(int(round(60.0 / self.duration *
                      self.samples /
                      int(ecgdata['RR Interval'])))) + ' BPM\n'

        return 'Ventr. Freq.: ' + \
               bpm + \
               'PR Interval: ' + ecgdata.get('PR Interval', '') + ' ms\n' + \
               'QRS Duration: ' + ecgdata.get('QRS Duration', '') + ' ms\n' + \
               'QT/QTc: ' + \
               ecgdata.get('QT Interval', '') + '/' + \
               ecgdata.get('QTc Interval', '') + ' ms\n' + \
               'P-R-T Axis: ' + \
               ecgdata.get('P Axis', '') + ' ' + \
               ecgdata.get('QRS Axis', '') + ' ' + \
               ecgdata.get('T Axis', '')

    def print_info(self):

        try:
            pat_surname, pat_firstname = self.dicom.PatientName.split('^')
        except ValueError:
            pat_surname = self.dicom.PatientName
            pat_firstname = ''

        pat_name = ' '.join((pat_surname, pat_firstname.title()))

        pat_id = self.dicom.PatientID
        pat_sex = self.dicom.PatientSex
        pat_bdate = datetime.strptime(
            self.dicom.PatientBirthDate, '%Y%m%d').strftime("%e %b %Y")

        info = "%s (%s) sex: %s\nbirthdate: %s" % (pat_name, pat_id,
                                                   pat_sex, pat_bdate)
        plt.figtext(0.08, 0.865, info, fontsize=12)

        info = "total time: %s s | sample freq.: %s Hz" % (
            self.duration, self.sampling_frequency)
        plt.figtext(0.08, 0.02, info, fontsize=10)

        # TODO: the bandpass filter 0.05-40 Hz will have to became a parameter
        info = "%s mm/s %s mm/mV 0.05-40 Hz" % (self.mm_s, self.mm_mv)
        plt.figtext(0.78, 0.02, info, fontsize=10)

        plt.figtext(0.5, 0.865, self.legend(), fontsize=10)

        info = PRODUCER + '\nStudy date: '
        date_str = (self.dicom.InstanceCreationDate +
                    self.dicom.InstanceCreationTime)
        info += datetime.strftime(datetime.strptime(date_str,
                                                    '%Y%m%d%H%M%S'),
                                  '%d %b %Y %H:%M')
        plt.figtext(0.03, 0.94, info, fontsize=12)

    def save(self, outputfile=None, outformat=None):

        def _save(output):
            plt.savefig(output, dpi=300, format=outformat,
                        papertype='a4', orientation='landscape')

        if outputfile:
            _save(outputfile)
        else:
            output = cStringIO.StringIO()
            _save(output)
            return output.getvalue()

    def plot(self, layoutid, mm_mv):

        self.mm_mv = mm_mv

        layout = LAYOUT[layoutid]
        rows = len(layout)

        for numrow, row in enumerate(layout):

            columns = len(row)
            h_delta = self.samples / columns
            row_height = self.height / rows
            v_delta = round(self.height * (1 - 1.0/(rows*2)) -
                            numrow*(self.height/rows))

            v_delta = (v_delta + 2.5) - (v_delta + 2.5) % 5

            # Lenght of a signal chunk
            chunk_size = int(self.samples/len(row))
            for numcol, signum in enumerate(row):
                left = numcol*chunk_size
                right = (1+numcol)*chunk_size

                # The signal chunk, vertical shifted and
                # scaled by mm/mV factor
                signal = v_delta + mm_mv * self.signals[signum][left:right]
                self.axis.plot(range(left, right), signal,
                               linewidth=0.6, color='black', zorder=10)

                cseq = self.channel_definitions[signum].ChannelSourceSequence
                meaning = cseq[0].CodeMeaning.replace(
                    'Lead', '').replace('(Einthoven)', '')
                self.axis.text(h+40, v_delta+row_height/3,
                               meaning, color='b', zorder=50)

            self.axis.plot(signal+v_delta, linewidth=.6, color='black',
                           antialiased=True, zorder=10)

        # A4 size in inches
        self.fig.set_size_inches(11.69, 8.27)

    def draw(self, layoutid, mm_mv=10.0):
        self.draw_grid()
        self.plot(layoutid, mm_mv)
        self.print_info()
<|endoftext|>