<commit_before>#!/usr/bin/python
# unpacklzexe.py by Sam Russell <sam.h.russell@gmail.com>
# Written May 2015

import struct
import json

def parseheader(header):
  fields = struct.unpack('<HHHHHHHHHHHHHH', header)
  headerdata = {}
  headerdata['signature'] = ''.join([chr(fields[0] & 0xFF), chr(fields[0] >> 8)])
  headerdata['partpage'] = fields[1]
  headerdata['pagecnt'] = fields[2]
  headerdata['relocnt'] = fields[3]
  headerdata['hdrsize'] = fields[4]
  headerdata['minalloc'] = fields[5]
  headerdata['maxalloc'] = fields[6]
  headerdata['initss'] = fields[7]
  headerdata['initsp'] = fields[8]
  headerdata['chksum'] = fields[9]
  headerdata['initip'] = fields[10]
  headerdata['initcs'] = fields[11]
  headerdata['tabloff'] = fields[12]
  headerdata['overlayno'] = fields[13]

  return headerdata

def generateheader(headerdata):
  headerbase = struct.pack('<BBHHHHHHHHHHHHH',
                            ord(headerdata['signature'][0]),
                            ord(headerdata['signature'][1]),
                            headerdata['partpage'],
                            headerdata['pagecnt'],
                            headerdata['relocnt'],
                            headerdata['hdrsize'],
                            headerdata['minalloc'],
                            headerdata['maxalloc'],
                            headerdata['initss'],
                            headerdata['initsp'],
                            headerdata['chksum'],
                            headerdata['initip'],
                            headerdata['initcs'],
                            headerdata['tabloff'],
                            headerdata['overlayno'])
  return headerbase

def checksignaturelz91(checkstring):
  if checkstring == "LZ91":
    return True
  return False

def unpacklz91data(indata):
  outdata = ""
  si = 0
  dx = 0x10
  bp = struct.unpack('<H', indata[si:si+2])[0]
  si = si + 2
  bit = 0
  while True:
    bit = bp & 1
    bp = bp >> 1
    dx = dx - 1
    if dx == 0:
      bp = struct.unpack('<H', indata[si:si+2])[0]
      si = si + 2
      dx = 0x10
    if bit == 1:
      outdata = outdata + indata[si]
      si = si + 1
      continue
    cx = 0
    bit = bp & 1
    bp = bp >> 1
    dx = dx - 1
    if dx == 0:
      bp = struct.unpack('<H', indata[si:si+2])[0]
      si = si + 2
      dx = 0x10
    if bit == 0:
      bit = bp & 1
      bp = bp >> 1
      dx = dx - 1
      if dx == 0:
        bp = struct.unpack('<H', indata[si:si+2])[0]
        si = si + 2
        dx = 0x10
      cx = ((cx << 1) + bit) & 0xFFFF
      bit = bp & 1
      bp = bp >> 1
      dx = dx - 1
      if dx == 0:
        bp = struct.unpack('<H', indata[si:si+2])[0]
        si = si + 2
        dx = 0x10
      cx = ((cx << 1) + bit) & 0xFFFF
      cx = cx + 2
      tempbyte = struct.unpack('<B', indata[si:si+1])[0]
      si = si + 1
      bx = tempbyte - 0x100
      for i in range(cx):
        tempbyte = outdata[bx]
        outdata = outdata + tempbyte
      continue
    ax = struct.unpack('<H', indata[si:si+2])[0]
    si = si + 2
    bx = ax
    bh = bx >> 8
    bh = bh >> 3
    bx = (bh << 8) + (bx & 0xff)
    bx = bx - 0x2000
    ah = ax >> 8
    ah = ah & 0x7
    if ah != 0:
      cx = ah + 2
      for i in range(cx):
        tempbyte = outdata[bx]
        outdata = outdata + tempbyte
      continue
    al = struct.unpack('<B', indata[si:si+1])[0]
    si = si + 1
    if al == 0:
      break
    if al != 1:
      cx = al + 1
      for i in range(cx):
        tempbyte = outdata[bx]
        outdata = outdata + tempbyte
      continue
    # byte 0x01 is a signal to rephrase the current mem pointer
    # this makes sense under DOS, but is unnecessary here with no 16-bit segments
    #if al != 1:
    #  print "si: %04X di:%04X al should be 1 but is %02X" % (si, len(outdata), al)
    #print "pretending to realign segments"
  return outdata

def unpacklz91reloc(relocdata):
  relocoutstr = bytearray()
  si = 0
  dx = 0
  di = 0
  while True:
    al = struct.unpack("<B", relocdata[si:si+1])[0]
    si = si + 1
    if al == 0:
      ax = struct.unpack("<H", relocdata[si:si+2])[0]
      si = si + 2
      if ax == 0:
        dx = dx + 0xfff
        dx = dx % 0x10000
        es = dx
        continue
      if ax == 1:
        break
    else:
      ax = al
    di = di + ax
    ax = di
    di = di & 0xf
    ax = ax >>4
    dx = dx + ax
    es = dx
    relocoutstr = relocoutstr + struct.pack('<HH', di, es)
  return relocoutstr

def unpacklz91(data, outfile):
  header = data[:0x1C]
  headerdata = parseheader(header)
  # size of header is hdrsize x 0x10
  headersize = headerdata['hdrsize'] * 0x10
  # loader stub is at CS:0000 + hdrsize
  loadercs = headerdata['initcs']
  loaderoffset = loadercs * 0x10 + headersize
  #print "Loader data is at %02X" % loaderoffset
  # ss is at 0x06 from start
  # sp is at 0x04 from start
  # cs is at 0x02 from start
  # ip is at 0x00 from start
  exeip = struct.unpack('<H', data[loaderoffset + 0x00:loaderoffset + 0x02])[0]
  execs = struct.unpack('<H', data[loaderoffset + 0x02:loaderoffset + 0x04])[0]
  exesp = struct.unpack('<H', data[loaderoffset + 0x04:loaderoffset + 0x06])[0]
  exess = struct.unpack('<H', data[loaderoffset + 0x06:loaderoffset + 0x08])[0]
  #print "EXE loads at %04X:%04X with SS:SP %04X:%04X" % (execs, exeip, exess, exesp)
  # now just to unpack :)
  # packed data is data[0x20:loaderoffset]
  # packed reloc is data[loaderoffset+0x158]
  packeddata = data[0x20:loaderoffset]
  packedreloc = data[loaderoffset+0x158:]
  unpackeddata = unpacklz91data(packeddata)
  unpackedreloc = unpacklz91reloc(packedreloc)
  # now just to generate new header and build exe
  # deal with partpage and pagecnt after we do relocs
  # number of relocs = len(unpackedreloc)/4
  headerdata['relocnt'] = len(unpackedreloc)/4
  # header is now 0x1C + relocnt*4 and round up to nearest 0x10
  headersize = headerdata['relocnt']*4 + 0x1C
  extra = headersize % 0x10
  if extra > 0:
    headersize = headersize + 0x10 - extra
  # remember we need padding on end of header
  padding = ''.join(['\x00' for x in range(0x10 - extra)])
  headerdata['hdrsize'] = headersize/0x10
  # can now calculate total file size
  filesize = headersize + len(unpackeddata)
  headerdata['partpage'] = filesize % 0x200
  headerdata['pagecnt'] = filesize / 0x200
  if headerdata['partpage'] > 0:
    headerdata['pagecnt'] = headerdata['pagecnt'] + 1
  lessmemory = len(unpackeddata) - len(packeddata)
  headerdata['minalloc'] = headerdata['minalloc'] - lessmemory/0x10
  headerdata['maxalloc'] = 0xffff
  headerdata['chksum'] = 0
  # file data for loading
  headerdata['initss'] = exess
  headerdata['initsp'] = exesp
  headerdata['initip'] = exeip
  headerdata['initcs'] = execs
  headerdata['tabloff'] = 0x1c
  headerdata['overlayno'] = 0
  # we are now good to go
  newheader = generateheader(headerdata)
  # now write newheader + unpackedreloc + padding + unpackeddata
  fout = open(outfile, 'w')
  fout.write(newheader + unpackedreloc + padding + unpackeddata)
  return 1

def unpacklzexe(data, outfile):
  if checksignaturelz91(data[0x1c:0x20]):
    print "Detected LZ91 file"
    return unpacklz91(data, outfile)
  return None

if __name__ == '__main__':
  # run test suite
  import argparse
  parser = argparse.ArgumentParser('Provide an LZEXE packed file')
  parser.add_argument('infile')
  parser.add_argument('outfile')
  args = parser.parse_args()
  filename = args.infile
  f = open(filename, 'r')
  data = f.read()
  unpacked = unpacklzexe(data, args.outfile)
  if unpacked:
    print "Done"
  else:
    print "Not a valid LZ91 file, could not unpack"

<commit_msg>bug when header size is a multiple of 0x10<commit_after>#!/usr/bin/python
# unpacklzexe.py by Sam Russell <sam.h.russell@gmail.com>
# Written May 2015

import struct
import json

def parseheader(header):
  fields = struct.unpack('<HHHHHHHHHHHHHH', header)
  headerdata = {}
  headerdata['signature'] = ''.join([chr(fields[0] & 0xFF), chr(fields[0] >> 8)])
  headerdata['partpage'] = fields[1]
  headerdata['pagecnt'] = fields[2]
  headerdata['relocnt'] = fields[3]
  headerdata['hdrsize'] = fields[4]
  headerdata['minalloc'] = fields[5]
  headerdata['maxalloc'] = fields[6]
  headerdata['initss'] = fields[7]
  headerdata['initsp'] = fields[8]
  headerdata['chksum'] = fields[9]
  headerdata['initip'] = fields[10]
  headerdata['initcs'] = fields[11]
  headerdata['tabloff'] = fields[12]
  headerdata['overlayno'] = fields[13]

  return headerdata

def generateheader(headerdata):
  headerbase = struct.pack('<BBHHHHHHHHHHHHH',
                            ord(headerdata['signature'][0]),
                            ord(headerdata['signature'][1]),
                            headerdata['partpage'],
                            headerdata['pagecnt'],
                            headerdata['relocnt'],
                            headerdata['hdrsize'],
                            headerdata['minalloc'],
                            headerdata['maxalloc'],
                            headerdata['initss'],
                            headerdata['initsp'],
                            headerdata['chksum'],
                            headerdata['initip'],
                            headerdata['initcs'],
                            headerdata['tabloff'],
                            headerdata['overlayno'])
  return headerbase

def checksignaturelz91(checkstring):
  if checkstring == "LZ91":
    return True
  return False

def unpacklz91data(indata):
  outdata = ""
  si = 0
  dx = 0x10
  bp = struct.unpack('<H', indata[si:si+2])[0]
  si = si + 2
  bit = 0
  while True:
    bit = bp & 1
    bp = bp >> 1
    dx = dx - 1
    if dx == 0:
      bp = struct.unpack('<H', indata[si:si+2])[0]
      si = si + 2
      dx = 0x10
    if bit == 1:
      outdata = outdata + indata[si]
      si = si + 1
      continue
    cx = 0
    bit = bp & 1
    bp = bp >> 1
    dx = dx - 1
    if dx == 0:
      bp = struct.unpack('<H', indata[si:si+2])[0]
      si = si + 2
      dx = 0x10
    if bit == 0:
      bit = bp & 1
      bp = bp >> 1
      dx = dx - 1
      if dx == 0:
        bp = struct.unpack('<H', indata[si:si+2])[0]
        si = si + 2
        dx = 0x10
      cx = ((cx << 1) + bit) & 0xFFFF
      bit = bp & 1
      bp = bp >> 1
      dx = dx - 1
      if dx == 0:
        bp = struct.unpack('<H', indata[si:si+2])[0]
        si = si + 2
        dx = 0x10
      cx = ((cx << 1) + bit) & 0xFFFF
      cx = cx + 2
      tempbyte = struct.unpack('<B', indata[si:si+1])[0]
      si = si + 1
      bx = tempbyte - 0x100
      for i in range(cx):
        tempbyte = outdata[bx]
        outdata = outdata + tempbyte
      continue
    ax = struct.unpack('<H', indata[si:si+2])[0]
    si = si + 2
    bx = ax
    bh = bx >> 8
    bh = bh >> 3
    bx = (bh << 8) + (bx & 0xff)
    bx = bx - 0x2000
    ah = ax >> 8
    ah = ah & 0x7
    if ah != 0:
      cx = ah + 2
      for i in range(cx):
        tempbyte = outdata[bx]
        outdata = outdata + tempbyte
      continue
    al = struct.unpack('<B', indata[si:si+1])[0]
    si = si + 1
    if al == 0:
      break
    if al != 1:
      cx = al + 1
      for i in range(cx):
        tempbyte = outdata[bx]
        outdata = outdata + tempbyte
      continue
    # byte 0x01 is a signal to rephrase the current mem pointer
    # this makes sense under DOS, but is unnecessary here with no 16-bit segments
    #if al != 1:
    #  print "si: %04X di:%04X al should be 1 but is %02X" % (si, len(outdata), al)
    #print "pretending to realign segments"
  return outdata

def unpacklz91reloc(relocdata):
  relocoutstr = bytearray()
  si = 0
  dx = 0
  di = 0
  while True:
    al = struct.unpack("<B", relocdata[si:si+1])[0]
    si = si + 1
    if al == 0:
      ax = struct.unpack("<H", relocdata[si:si+2])[0]
      si = si + 2
      if ax == 0:
        dx = dx + 0xfff
        dx = dx % 0x10000
        es = dx
        continue
      if ax == 1:
        break
    else:
      ax = al
    di = di + ax
    ax = di
    di = di & 0xf
    ax = ax >>4
    dx = dx + ax
    es = dx
    relocoutstr = relocoutstr + struct.pack('<HH', di, es)
  return relocoutstr

def unpacklz91(data, outfile):
  header = data[:0x1C]
  headerdata = parseheader(header)
  # size of header is hdrsize x 0x10
  headersize = headerdata['hdrsize'] * 0x10
  # loader stub is at CS:0000 + hdrsize
  loadercs = headerdata['initcs']
  loaderoffset = loadercs * 0x10 + headersize
  #print "Loader data is at %02X" % loaderoffset
  # ss is at 0x06 from start
  # sp is at 0x04 from start
  # cs is at 0x02 from start
  # ip is at 0x00 from start
  exeip = struct.unpack('<H', data[loaderoffset + 0x00:loaderoffset + 0x02])[0]
  execs = struct.unpack('<H', data[loaderoffset + 0x02:loaderoffset + 0x04])[0]
  exesp = struct.unpack('<H', data[loaderoffset + 0x04:loaderoffset + 0x06])[0]
  exess = struct.unpack('<H', data[loaderoffset + 0x06:loaderoffset + 0x08])[0]
  #print "EXE loads at %04X:%04X with SS:SP %04X:%04X" % (execs, exeip, exess, exesp)
  # now just to unpack :)
  # packed data is data[0x20:loaderoffset]
  # packed reloc is data[loaderoffset+0x158]
  packeddata = data[0x20:loaderoffset]
  packedreloc = data[loaderoffset+0x158:]
  unpackeddata = unpacklz91data(packeddata)
  unpackedreloc = unpacklz91reloc(packedreloc)
  # now just to generate new header and build exe
  # deal with partpage and pagecnt after we do relocs
  # number of relocs = len(unpackedreloc)/4
  headerdata['relocnt'] = len(unpackedreloc)/4
  # header is now 0x1C + relocnt*4 and round up to nearest 0x10
  headersize = headerdata['relocnt']*4 + 0x1C
  extra = headersize % 0x10
  padbytes = 0
  if extra > 0:
    padbytes = 0x10 - extra
  headersize = headersize + padbytes
  # remember we need padding on end of header
  padding = ''.join(['\x00' for x in range(padbytes)])
  headerdata['hdrsize'] = headersize/0x10
  # can now calculate total file size
  filesize = headersize + len(unpackeddata)
  headerdata['partpage'] = filesize % 0x200
  headerdata['pagecnt'] = filesize / 0x200
  if headerdata['partpage'] > 0:
    headerdata['pagecnt'] = headerdata['pagecnt'] + 1
  lessmemory = len(unpackeddata) - len(packeddata)
  headerdata['minalloc'] = headerdata['minalloc'] - lessmemory/0x10
  headerdata['maxalloc'] = 0xffff
  headerdata['chksum'] = 0
  # file data for loading
  headerdata['initss'] = exess
  headerdata['initsp'] = exesp
  headerdata['initip'] = exeip
  headerdata['initcs'] = execs
  headerdata['tabloff'] = 0x1c
  headerdata['overlayno'] = 0
  # we are now good to go
  newheader = generateheader(headerdata)
  # now write newheader + unpackedreloc + padding + unpackeddata
  fout = open(outfile, 'w')
  fout.write(newheader + unpackedreloc + padding + unpackeddata)
  return 1

def unpacklzexe(data, outfile):
  if checksignaturelz91(data[0x1c:0x20]):
    print "Detected LZ91 file"
    return unpacklz91(data, outfile)
  return None

if __name__ == '__main__':
  # run test suite
  import argparse
  parser = argparse.ArgumentParser('Provide an LZEXE packed file')
  parser.add_argument('infile')
  parser.add_argument('outfile')
  args = parser.parse_args()
  filename = args.infile
  f = open(filename, 'r')
  data = f.read()
  unpacked = unpacklzexe(data, args.outfile)
  if unpacked:
    print "Done"
  else:
    print "Not a valid LZ91 file, could not unpack"

<|endoftext|>


<commit_before>import itertools
import logging
import time
import random

def retry(f, *allowed_exceptions, allowed_exception_fn=None, times=6, sleep=1, exponent=1, silent=False):
    def fn(*a, **kw):
        for i in itertools.count():
            try:
                return f(*a, **kw)
            except allowed_exceptions:
                raise
            except Exception as e:
                if allowed_exception_fn and allowed_exception_fn(e):
                    raise
                if i == times:
                    raise
                if not silent:
                    logging.info(f'retrying: {f.__module__}.{f.__name__}, because of: {e}')
                if exponent:
                    amount = (sleep * i) ** exponent
                else:
                    amount = sleep
                time.sleep(amount + random.random())
    return fn
<commit_msg>allow retry max seconds<commit_after>import itertools
import logging
import time

def retry(f, *allowed_exceptions, allowed_exception_fn=None, times=6, sleep=1, exponent=1, silent=False, max_seconds=0):
    def fn(*a, **kw):
        start = time.time()
        for i in itertools.count():
            try:
                return f(*a, **kw)
            except allowed_exceptions:
                raise
            except Exception as e:
                if allowed_exception_fn and allowed_exception_fn(e):
                    raise
                duration = time.time() - start
                if i == times or max_seconds and duration >= max_seconds:
                    raise
                if not silent:
                    logging.info(f'retrying: {f.__module__}.{f.__name__}, because of: {e}')
                if exponent:
                    amount = (sleep * i) ** exponent
                else:
                    amount = sleep
                if max_seconds and duration + amount > max_seconds:
                    amount = max_seconds - duration
                time.sleep(amount)
    return fn
<|endoftext|>


<commit_before>from __future__ import absolute_import, division
from dials.array_family import flex # import dependency

class Test(object):

  def __init__(self):
    from os.path import join
    import libtbx.load_env
    try:
      dials_regression = libtbx.env.dist_path('dials_regression')
    except KeyError, e:
      print 'SKIP: dials_regression not configured'
      exit(0)

    self.lcls_path = join(dials_regression, "image_examples/LCLS_cspad_nexus")
    self.sacla_path = join(dials_regression, "image_examples/SACLA_MPCCD_Cheetah")

  def run(self):
    self.test_cspad_cbf_in_memory()
    self.test_sacla_h5()

  def test_cspad_cbf_in_memory(self):
    from os.path import join, exists
    import os, dxtbx
    from uuid import uuid4
    from dials.command_line.stills_process import phil_scope, Processor
    from libtbx.phil import parse
    from dxtbx.imageset import MemImageSet
    from dxtbx.datablock import DataBlockFactory
    from dxtbx.format.FormatCBFCspad import FormatCBFCspadInMemory
    import cPickle as pickle

    dirname ='tmp_%s' % uuid4().hex
    os.mkdir(dirname)
    os.chdir(dirname)

    assert exists(join(self.lcls_path, 'idx-20130301060858801.cbf'))

    f = open("process_lcls.phil", 'w')
    f.write("""
      dispatch.squash_errors = False
      spotfinder {
        filter.min_spot_size=2
        threshold.xds.gain=25
        threshold.xds.global_threshold=100
      }
      indexing {
        known_symmetry {
          space_group = P6122
          unit_cell = 92.9 92.9 130.4 90 90 120
        }
        method=fft1d
        refinement_protocol.d_min_start=1.7
      }
      integration {
        integrator=stills
        profile.fitting=False
        background {
          algorithm = simple
          simple {
            model.algorithm = linear2d
            outlier.algorithm = plane
          }
        }
      }
      refinement {
        parameterisation {
          beam.fix=all
          detector.fix=all
        }
        reflections {
          outlier.algorithm=null
          weighting_strategy.override=stills
        }
      }
      profile {
        gaussian_rs {
          min_spots.overall = 0
        }
      }
      """)
    f.close()
    params = phil_scope.fetch(parse(file_name="process_lcls.phil")).extract()
    params.output.datablock_filename = None
    processor = Processor(params)
    mem_img = dxtbx.load(join(self.lcls_path, 'idx-20130301060858801.cbf'))
    raw_data = mem_img.get_raw_data() # cache the raw data to prevent swig errors
    mem_img = FormatCBFCspadInMemory(mem_img._cbf_handle)
    mem_img._raw_data = raw_data
    mem_img._cbf_handle = None # drop the file handle to prevent swig errors
    imgset = MemImageSet([mem_img])
    datablock = DataBlockFactory.from_imageset(imgset)[0]
    processor.process_datablock("20130301060858801", datablock) # index/integrate the image
    result = "idx-20130301060858801_integrated.pickle"
    n_refls = range(140,152) # large ranges to handle platform-specific differences
    table = pickle.load(open(result, 'rb'))
    assert len(table) in n_refls, len(table)
    assert 'id' in table
    assert (table['id'] == 0).count(False) == 0

    print 'OK'

  def test_sacla_h5(self, in_memory=False):
    from os.path import join, exists
    from libtbx import easy_run
    import os
    from uuid import uuid4

    dirname ='tmp_%s' % uuid4().hex
    os.mkdir(dirname)
    os.chdir(dirname)

    assert exists(join(self.sacla_path, 'run266702-0-subset.h5'))

    geometry_path = join(self.sacla_path, 'refined_experiments_level1.json')
    assert exists(geometry_path)

    f = open("process_sacla.phil", 'w')
    f.write("""
      dispatch.squash_errors = False
      input.reference_geometry=%s
      indexing {
        known_symmetry {
          space_group = P43212
          unit_cell = 78.9 78.9 38.1 90 90 90
        }
        method = fft1d
        refinement_protocol.d_min_start = 2.2
      }

      spotfinder {
        filter.min_spot_size = 2
        threshold {
          xds {
            gain = 5.46 # from dials.estimate_gain run266702-0-subset.h5 max_images=4
            global_threshold = 50
          }
        }
      }

      refinement {
        parameterisation {
          beam.fix = all
          detector.fix_list = Dist,Tau1
          auto_reduction {
            action = fix
            min_nref_per_parameter = 1
          }
          crystal {
            unit_cell {
              restraints {
                tie_to_target {
                  values = 78.9,78.9,38.1,90,90,90
                  sigmas = 1,1,1,0,0,0
                }
              }
            }
          }
        }
      }
      integration {
        integrator = stills
        profile.fitting = False
        background {
          algorithm = simple
          simple {
            model.algorithm = linear2d
            outlier.algorithm = tukey
          }
        }
      }
      profile {
        gaussian_rs {
          min_spots.overall = 0
        }
      }
      """%geometry_path)
    f.close()

    # Call dials.stills_process
    result = easy_run.fully_buffered([
      'dials.stills_process',
      join(self.sacla_path, 'run266702-0-subset.h5'),
      'process_sacla.phil',
    ]).raise_if_errors()
    result.show_stdout()

    import cPickle as pickle
    # Frame 1 no longer indexing after cctbx r25607 which made wavelengths be on a per-image basis
    #for result, n_refls in zip(["idx-run266702-0-subset_00000_integrated.pickle",
    #                            "idx-run266702-0-subset_00001_integrated.pickle"],
    #                            [range(109,114), range(80,85)]): # large ranges to handle platform-specific differences
    #for result, n_refls in zip(["idx-run266702-0-subset_00000_integrated.pickle"],
    #                            [range(109,114)]): # large ranges to handle platform-specific differences
    # dxtbx r25668 and 25669 flip X axis in the SACLA format class and changed indexing results.
    #for result, n_refls in zip(["idx-run266702-0-subset_00001_integrated.pickle"],
    #                            [range(90,96)]): # large ranges to handle platform-specific differences
    # 02/12/17 Handle change to stills_process refining after indexing plus new spotfinding params
    #for result, n_refls in zip(["idx-run266702-0-subset_00000_integrated.pickle",
    #                            "idx-run266702-0-subset_00001_integrated.pickle",
    #                            "idx-run266702-0-subset_00003_integrated.pickle"],
    #                            [range(75,90),range(220,230),range(285,295)]): # large ranges to handle platform-specific differences
    # 02/14/17 Further changes to stills_process: resetting rejected reflections before re-refinement
    #for result, n_refls in zip(["idx-run266702-0-subset_00000_integrated.pickle",
    #                            "idx-run266702-0-subset_00001_integrated.pickle",
    #                            "idx-run266702-0-subset_00003_integrated.pickle"],
    #                            [range(80,95),range(225,235),range(235,245)]): # large ranges to handle platform-specific differences
    # 02/21/17 Changes to stills_process: refine during indexing instead of after. Also used refined metrology from Rahel
    #for result, n_refls in zip(["idx-run266702-0-subset_00001_integrated.pickle",
    #                            "idx-run266702-0-subset_00003_integrated.pickle"],
    #                            [range(600,610),range(505,520)]): # large ranges to handle platform-specific differences
    # 04/25/17 Changes after reverting sign_error_27Feb2014_through_15Feb2017 in xfel/mono_simulation/max_like.py
    for result, n_refls in zip(["idx-run266702-0-subset_00001_integrated.pickle",
                                "idx-run266702-0-subset_00003_integrated.pickle"],
                                [range(565,580),range(495,510)]): # large ranges to handle platform-specific differences
      table = pickle.load(open(result, 'rb'))
      assert len(table) in n_refls, len(table)
      assert 'id' in table
      assert (table['id'] == 0).count(False) == 0
    print 'OK'

if __name__ == '__main__':
  from dials.test import cd_auto
  with cd_auto(__file__):
    test = Test()
    test.run()
<commit_msg>Sync test with dxtbx-refactor changes<commit_after>from __future__ import absolute_import, division
from dials.array_family import flex # import dependency

class Test(object):

  def __init__(self):
    from os.path import join
    import libtbx.load_env
    try:
      dials_regression = libtbx.env.dist_path('dials_regression')
    except KeyError, e:
      print 'SKIP: dials_regression not configured'
      exit(0)

    self.lcls_path = join(dials_regression, "image_examples/LCLS_cspad_nexus")
    self.sacla_path = join(dials_regression, "image_examples/SACLA_MPCCD_Cheetah")

  def run(self):
    self.test_cspad_cbf_in_memory()
    self.test_sacla_h5()

  def test_cspad_cbf_in_memory(self):
    from os.path import join, exists
    import os, dxtbx
    from uuid import uuid4
    from dials.command_line.stills_process import phil_scope, Processor
    from libtbx.phil import parse
    from dxtbx.imageset import ImageSet, ImageSetData, MemReader
    from dxtbx.datablock import DataBlockFactory
    from dxtbx.format.FormatCBFCspad import FormatCBFCspadInMemory
    import cPickle as pickle

    dirname ='tmp_%s' % uuid4().hex
    os.mkdir(dirname)
    os.chdir(dirname)

    assert exists(join(self.lcls_path, 'idx-20130301060858801.cbf'))

    f = open("process_lcls.phil", 'w')
    f.write("""
      dispatch.squash_errors = False
      spotfinder {
        filter.min_spot_size=2
        threshold.xds.gain=25
        threshold.xds.global_threshold=100
      }
      indexing {
        known_symmetry {
          space_group = P6122
          unit_cell = 92.9 92.9 130.4 90 90 120
        }
        method=fft1d
        refinement_protocol.d_min_start=1.7
      }
      integration {
        integrator=stills
        profile.fitting=False
        background {
          algorithm = simple
          simple {
            model.algorithm = linear2d
            outlier.algorithm = plane
          }
        }
      }
      refinement {
        parameterisation {
          beam.fix=all
          detector.fix=all
        }
        reflections {
          outlier.algorithm=null
          weighting_strategy.override=stills
        }
      }
      profile {
        gaussian_rs {
          min_spots.overall = 0
        }
      }
      """)
    f.close()
    params = phil_scope.fetch(parse(file_name="process_lcls.phil")).extract()
    params.output.datablock_filename = None
    processor = Processor(params)
    mem_img = dxtbx.load(join(self.lcls_path, 'idx-20130301060858801.cbf'))
    raw_data = mem_img.get_raw_data() # cache the raw data to prevent swig errors
    mem_img = FormatCBFCspadInMemory(mem_img._cbf_handle)
    mem_img._raw_data = raw_data
    mem_img._cbf_handle = None # drop the file handle to prevent swig errors
    imgset = ImageSet(ImageSetData(MemReader([mem_img])))
    imgset.set_beam(mem_img.get_beam())
    imgset.set_detector(mem_img.get_detector())
    datablock = DataBlockFactory.from_imageset(imgset)[0]
    processor.process_datablock("20130301060858801", datablock) # index/integrate the image
    result = "idx-20130301060858801_integrated.pickle"
    n_refls = range(140,152) # large ranges to handle platform-specific differences
    table = pickle.load(open(result, 'rb'))
    assert len(table) in n_refls, len(table)
    assert 'id' in table
    assert (table['id'] == 0).count(False) == 0

    print 'OK'

  def test_sacla_h5(self, in_memory=False):
    from os.path import join, exists
    from libtbx import easy_run
    import os
    from uuid import uuid4

    dirname ='tmp_%s' % uuid4().hex
    os.mkdir(dirname)
    os.chdir(dirname)

    assert exists(join(self.sacla_path, 'run266702-0-subset.h5'))

    geometry_path = join(self.sacla_path, 'refined_experiments_level1.json')
    assert exists(geometry_path)

    f = open("process_sacla.phil", 'w')
    f.write("""
      dispatch.squash_errors = False
      input.reference_geometry=%s
      indexing {
        known_symmetry {
          space_group = P43212
          unit_cell = 78.9 78.9 38.1 90 90 90
        }
        method = fft1d
        refinement_protocol.d_min_start = 2.2
      }

      spotfinder {
        filter.min_spot_size = 2
        threshold {
          xds {
            gain = 5.46 # from dials.estimate_gain run266702-0-subset.h5 max_images=4
            global_threshold = 50
          }
        }
      }

      refinement {
        parameterisation {
          beam.fix = all
          detector.fix_list = Dist,Tau1
          auto_reduction {
            action = fix
            min_nref_per_parameter = 1
          }
          crystal {
            unit_cell {
              restraints {
                tie_to_target {
                  values = 78.9,78.9,38.1,90,90,90
                  sigmas = 1,1,1,0,0,0
                }
              }
            }
          }
        }
      }
      integration {
        integrator = stills
        profile.fitting = False
        background {
          algorithm = simple
          simple {
            model.algorithm = linear2d
            outlier.algorithm = tukey
          }
        }
      }
      profile {
        gaussian_rs {
          min_spots.overall = 0
        }
      }
      """%geometry_path)
    f.close()

    # Call dials.stills_process
    result = easy_run.fully_buffered([
      'dials.stills_process',
      join(self.sacla_path, 'run266702-0-subset.h5'),
      'process_sacla.phil',
    ]).raise_if_errors()
    result.show_stdout()

    import cPickle as pickle
    # Frame 1 no longer indexing after cctbx r25607 which made wavelengths be on a per-image basis
    #for result, n_refls in zip(["idx-run266702-0-subset_00000_integrated.pickle",
    #                            "idx-run266702-0-subset_00001_integrated.pickle"],
    #                            [range(109,114), range(80,85)]): # large ranges to handle platform-specific differences
    #for result, n_refls in zip(["idx-run266702-0-subset_00000_integrated.pickle"],
    #                            [range(109,114)]): # large ranges to handle platform-specific differences
    # dxtbx r25668 and 25669 flip X axis in the SACLA format class and changed indexing results.
    #for result, n_refls in zip(["idx-run266702-0-subset_00001_integrated.pickle"],
    #                            [range(90,96)]): # large ranges to handle platform-specific differences
    # 02/12/17 Handle change to stills_process refining after indexing plus new spotfinding params
    #for result, n_refls in zip(["idx-run266702-0-subset_00000_integrated.pickle",
    #                            "idx-run266702-0-subset_00001_integrated.pickle",
    #                            "idx-run266702-0-subset_00003_integrated.pickle"],
    #                            [range(75,90),range(220,230),range(285,295)]): # large ranges to handle platform-specific differences
    # 02/14/17 Further changes to stills_process: resetting rejected reflections before re-refinement
    #for result, n_refls in zip(["idx-run266702-0-subset_00000_integrated.pickle",
    #                            "idx-run266702-0-subset_00001_integrated.pickle",
    #                            "idx-run266702-0-subset_00003_integrated.pickle"],
    #                            [range(80,95),range(225,235),range(235,245)]): # large ranges to handle platform-specific differences
    # 02/21/17 Changes to stills_process: refine during indexing instead of after. Also used refined metrology from Rahel
    #for result, n_refls in zip(["idx-run266702-0-subset_00001_integrated.pickle",
    #                            "idx-run266702-0-subset_00003_integrated.pickle"],
    #                            [range(600,610),range(505,520)]): # large ranges to handle platform-specific differences
    # 04/25/17 Changes after reverting sign_error_27Feb2014_through_15Feb2017 in xfel/mono_simulation/max_like.py
    for result, n_refls in zip(["idx-run266702-0-subset_00001_integrated.pickle",
                                "idx-run266702-0-subset_00003_integrated.pickle"],
                                [range(565,580),range(495,510)]): # large ranges to handle platform-specific differences
      table = pickle.load(open(result, 'rb'))
      assert len(table) in n_refls, len(table)
      assert 'id' in table
      assert (table['id'] == 0).count(False) == 0
    print 'OK'

if __name__ == '__main__':
  from dials.test import cd_auto
  with cd_auto(__file__):
    test = Test()
    test.run()
<|endoftext|>


<commit_before>"""
This file contains rows that provide row objects commonly employed by FatStax
"""
import fatstax

class CloudStaxExportRow(fatstax.BaseRow):

    def _subdir(self, max_category):
        t = "Category {}"
        for x in range(max_category):
            try:
                resp = self.getcolumn(t.format(x+1))
                if resp:
                    yield resp
                else:
                    yield None
            except KeyError:
                yield None

    def category_dir(self, max_category):
        """Returns a list of the catagories in order"""
        return list(self._subdir(max_category))

        <commit_msg>Removed Irrelevant and old directory.<commit_after><|endoftext|>


<commit_before># ！/usr/bin/python3
# -*- coding: utf-8 -*-
# Copyright (c) 2017-2020 Rhilip <rhilipruan@gmail.com>

import re
import time
from flask import Blueprint, request, jsonify
from app import mysql

ptboard_blueprint = Blueprint('ptboard', __name__)

@ptboard_blueprint.route("/ptboard")
def ptboard():
    search_raw = request.args.get("search") or ""
    order_raw = request.args.get("order") or "desc"
    limit_raw = request.args.get("limit") or 50
    offset_raw = request.args.get("offset") or 0

    t0 = time.clock()

    search = re.sub(r"[ _\-,.]", " ", search_raw)
    search = re.sub(r"\'", r"''", search)
    search = search.split()
    search = search[:10]
    if search:
        key = ["`title` LIKE '%{key}%'".format(key=i) for i in search]
        opt = " AND ".join(key)
        total_data = mysql.exec("SELECT COUNT(*) FORM `rss_pt_site` WHERE {opt}".format(opt=opt))
    else:
        opt = "1=1"
        total_data = mysql.exec("SELECT `TABLE_ROWS` FROM `information_schema`.`TABLES` "
                                "WHERE `TABLE_NAME`='rss_pt_site'")[0]

    order = "desc" if order_raw not in ["desc", "asc"] else order_raw
    try:
        limit = int(limit_raw)
        if limit > 200:
            limit = 200
    except ValueError or TypeError:
        limit = 50
    try:
        offset = int(offset_raw)
    except ValueError or TypeError:
        offset = 0

    sql = ("SELECT `site`,`sid`,`title`,`link`,`pubDate` FROM `rss_pt_site` "
           "WHERE {opt} ORDER BY `pubDate` {_da} LIMIT {_offset}, {_limit}".format(opt=opt, _da=order.upper(),
                                                                                _offset=offset, _limit=limit)
           )
    rows_data = mysql.exec(sql=sql, r_dict=True, fetch_all=True)

    return jsonify({
        "cost": time.clock() - t0,
        "rows": rows_data,
        "total": total_data
    })
<commit_msg>Fix typeerror<commit_after># ！/usr/bin/python3
# -*- coding: utf-8 -*-
# Copyright (c) 2017-2020 Rhilip <rhilipruan@gmail.com>

import re
import time
from flask import Blueprint, request, jsonify
from app import mysql

ptboard_blueprint = Blueprint('ptboard', __name__)

@ptboard_blueprint.route("/ptboard")
def ptboard():
    search_raw = request.args.get("search") or ""
    order_raw = request.args.get("order") or "desc"
    limit_raw = request.args.get("limit") or 50
    offset_raw = request.args.get("offset") or 0

    t0 = time.clock()

    search = re.sub(r"[ _\-,.]", " ", search_raw)
    search = re.sub(r"\'", r"''", search)
    search = search.split()
    search = search[:10]
    if search:
        key = ["`title` LIKE '%{key}%'".format(key=i) for i in search]
        opt = " AND ".join(key)
        total_data = mysql.exec("SELECT COUNT(*) FROM `rss_pt_site` WHERE {opt}".format(opt=opt))
    else:
        opt = "1=1"
        total_data = mysql.exec("SELECT `TABLE_ROWS` FROM `information_schema`.`TABLES` "
                                "WHERE `TABLE_NAME`='rss_pt_site'")[0]

    order = "desc" if order_raw not in ["desc", "asc"] else order_raw
    try:
        limit = int(limit_raw)
        if limit > 200:
            limit = 200
    except ValueError or TypeError:
        limit = 50
    try:
        offset = int(offset_raw)
    except ValueError or TypeError:
        offset = 0

    sql = ("SELECT `site`,`sid`,`title`,`link`,`pubDate` FROM `rss_pt_site` "
           "WHERE {opt} ORDER BY `pubDate` {_da} LIMIT {_offset}, {_limit}".format(opt=opt, _da=order.upper(),
                                                                                _offset=offset, _limit=limit)
           )
    rows_data = mysql.exec(sql=sql, r_dict=True, fetch_all=True)

    return jsonify({
        "cost": time.clock() - t0,
        "rows": rows_data,
        "total": total_data
    })
<|endoftext|>


<commit_before>"""

Many things copied from Johannes Berg's `MoinMoin Latex support`_

.. `MoinMoin Latex support`: http://johannes.sipsolutions.net/Projects/new-moinmoin-latex

"""

import subprocess, os, sys, shutil, tempfile, resource, md5

os.system('touch /tmp/foo')

OUT_PATH = "/var/www-pub/root/pTSc0V/testwiki/math"
OUT_URI_BASE = "http://192.168.0.100/pTSc0V/testwiki/math/"

DVIPNG = "dvipng"
LIMITED_LATEX = os.path.join(os.path.abspath(os.path.dirname(__file__)),
                             "limited-latex.py")
LATEX_TEMPLATE = r"""
\documentclass[10pt]{article}
\pagestyle{empty}
\usepackage[utf8]{inputenc}
\usepackage{amsmath,amssymb}
%(prologue)s
\begin{document}
%(raw)s
\end{document}
"""
LATEX_ARGS = ["--interaction=nonstopmode"]
DVIPNG_ARGS = ["-bgTransparent", "-Ttight", "--noghostscript", "-l1"]

# -----------------------------------------------------------------------------
# Running LaTeX safely
# -----------------------------------------------------------------------------

def latex_to_png(prologue, in_text, out_png):
    out_png = os.path.abspath(out_png)
    cwd = os.getcwd()
    tmp_dir = tempfile.mkdtemp()
    try:
        os.chdir(tmp_dir)
        f = open('foo.tex', 'w')
        f.write(LATEX_TEMPLATE % dict(raw=in_text.encode('utf-8'),
                                      prologue=prologue))
        f.close()
        ret = subprocess.call([LIMITED_LATEX] + LATEX_ARGS + ['foo.tex'],
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if ret != 0:
            raise RuntimeError("Executing latex failed")
        ret = subprocess.call([DVIPNG] + DVIPNG_ARGS + ['-o', 'foo.png',
                                                        'foo.dvi'],
                              stdout=subprocess.PIPE, stderr=subprocess.PIPE)
        if ret != 0:
            raise RuntimeError("Executing dvipng failed")
        shutil.copy('foo.png', out_png)
    finally:
        os.chdir(cwd)
        shutil.rmtree(tmp_dir)

def latex_to_uri(in_text):
    file_basename = md5.new(in_text.encode('utf-8')).hexdigest() + ".png"
    file_name = os.path.join(OUT_PATH, file_basename)
    if not os.path.isfile(file_name):
        latex_to_png("", in_text, file_name)
    uri = OUT_URI_BASE + file_basename
    return uri

# -----------------------------------------------------------------------------
# Roles and directives
# -----------------------------------------------------------------------------
import MoinMoin.parser.rst
import docutils
import docutils.core, docutils.nodes, docutils.parsers.rst

def math_role(name, rawtext, text, lineno, inliner, options={}, content=[]):
    uri = latex_to_uri(ur'$%s$' % text)
    img = docutils.nodes.image("", uri=uri)
    return [img], []

def math_directive(name, arguments, options, content, lineno,
                   content_offset, block_text, state, state_machine):
    uri = latex_to_uri(ur'\begin{align*}%s\end{align*}' % u''.join(content))
    img = docutils.nodes.image("", uri=uri, align='center')
    return [img]

# -----------------------------------------------------------------------------
# Register to reStructuredText engine
# -----------------------------------------------------------------------------

math_directive.arguments = (
    1, # number of required arguments
    1, # number of optional arguments
    False # whether final argument can contain whitespace
)
math_directive.options = {
}
math_directive.arguments = (
    0, # number of required arguments
    0, # number of optional arguments
    False # whether final argument can contain whitespace
)
math_directive.options = {}
math_directive.content = True # whether content is allowed

docutils.parsers.rst.directives.register_directive('math', math_directive)
docutils.parsers.rst.roles.register_local_role('math', math_role)
<commit_msg>Make it actually work<commit_after>"""

Many things copied from Johannes Berg's `MoinMoin Latex support`_

.. `MoinMoin Latex support`: http://johannes.sipsolutions.net/Projects/new-moinmoin-latex

"""

import subprocess, os, sys, shutil, tempfile, resource, md5

os.system('touch /tmp/foo')

OUT_PATH = "/var/www-pub/root/pTSc0V/testwiki/math"
OUT_URI_BASE = "http://192.168.0.100/pTSc0V/testwiki/math/"

DVIPNG = "dvipng"
LIMITED_LATEX = os.path.join(os.path.abspath(os.path.dirname(__file__)),
                             "limited-latex.py")
LATEX_TEMPLATE = r"""
\documentclass[10pt]{article}
\pagestyle{empty}
\usepackage[utf8]{inputenc}
\usepackage{amsmath,amssymb}
%(prologue)s
\begin{document}
%(raw)s
\end{document}
"""
LATEX_ARGS = ["--interaction=nonstopmode"]
DVIPNG_ARGS = ["-bgTransparent", "-Ttight", "--noghostscript", "-l1"]

# -----------------------------------------------------------------------------
# Running LaTeX safely
# -----------------------------------------------------------------------------

def exec_cmd(cmd, ok_return_value=0, show_cmd=True, echo=False, **kw):
    """
    Run given command and check return value.
    Return concatenated input and output.
    """
    try:
        if show_cmd:
            print '%s$' % os.getcwd(),
            if isinstance(cmd, str):
                print cmd
            else:
                print ' '.join(cmd)
            print '[running ...]'
        p = subprocess.Popen(cmd,
                             stdout=subprocess.PIPE,
                             stderr=subprocess.PIPE,
                             stdin=subprocess.PIPE, **kw)
        out, err = p.communicate()
    except OSError, e:
        raise RuntimeError("Command %s failed: %s" % (' '.join(cmd), e))

    if ok_return_value is not None and p.returncode != ok_return_value:
        raise RuntimeError("Command %s failed (code %d): %s"
                           % (' '.join(cmd), p.returncode, out + err))
    if echo: print out + err
    return out + err

def latex_to_png(prologue, in_text, out_png):
    out_png = os.path.abspath(out_png)
    cwd = os.getcwd()
    tmp_dir = tempfile.mkdtemp()
    try:
        os.chdir(tmp_dir)
        f = open('foo.tex', 'w')
        f.write(LATEX_TEMPLATE % dict(raw=in_text.encode('utf-8'),
                                      prologue=prologue))
        f.close()
        exec_cmd([LIMITED_LATEX] + LATEX_ARGS + ['foo.tex'], show_cmd=False)
        exec_cmd([DVIPNG] + DVIPNG_ARGS + ['-o', 'foo.png', 'foo.dvi'],
                 show_cmd=False)
        shutil.copy('foo.png', out_png)
    except OSError, e:
        raise RuntimeError(str(e))
    finally:
        os.chdir(cwd)
        shutil.rmtree(tmp_dir)

def latex_to_uri(in_text):
    file_basename = md5.new(in_text.encode('utf-8')).hexdigest() + ".png"
    file_name = os.path.join(OUT_PATH, file_basename)
    if not os.path.isfile(file_name):
        latex_to_png("", in_text, file_name)
    uri = OUT_URI_BASE + file_basename
    return uri

# -----------------------------------------------------------------------------
# Roles and directives
# -----------------------------------------------------------------------------
import MoinMoin.parser.rst
import docutils
import docutils.utils
import docutils.core, docutils.nodes, docutils.parsers.rst

def math_role(name, rawtext, text, lineno, inliner, options={}, content=[]):
    try:
        uri = latex_to_uri(ur'$%s$' % text)
        img = docutils.nodes.image("", uri=uri)
        return [img], []
    except RuntimeError, e:
        item = docutils.nodes.literal(text=str(text))
        return [item], []

def math_directive(name, arguments, options, content, lineno,
                   content_offset, block_text, state, state_machine):
    try:
        uri = latex_to_uri(ur'\begin{align*}%s\end{align*}' % u'\n'.join(content))
        img = docutils.nodes.image("", uri=uri, align='center')
        return [img]
    except RuntimeError, e:
        item = docutils.nodes.literal_block(text=u"\n".join(content))
        return [item]

# -----------------------------------------------------------------------------
# Register to reStructuredText engine
# -----------------------------------------------------------------------------

math_directive.arguments = (
    1, # number of required arguments
    1, # number of optional arguments
    False # whether final argument can contain whitespace
)
math_directive.options = {
}
math_directive.arguments = (
    0, # number of required arguments
    0, # number of optional arguments
    False # whether final argument can contain whitespace
)
math_directive.options = {}
math_directive.content = True # whether content is allowed

docutils.parsers.rst.directives.register_directive('math', math_directive)
docutils.parsers.rst.roles.register_local_role('math', math_role)
<|endoftext|>


<commit_before>#comment notes
#1. Currently supporting direction independent gains. Future versions may
#   support direction dependent gains with list of such objects.
#2.
class UVCal(UVBase):
    """ A class defining a calibration """
    def __init__(self):
        self._Nfreqs = uvp.UVParameter('Nfreqs', description='Number of frequency channels',
                                       expected_type=int)
        self._Npols = uvp.UVParameter('Npols', description='Number of polarizations',
                                      expected_type=int)
        self._Ntimes = uvp.UVParameter('Ntimes', description='Number of times',
                                       expected_type=int)
        self._history = uvp.UVParameter('history', description='String of history, units English',
                                        form='str', expected_type=str)

        desc = ('Number of antennas with data present (i.e. number of unique '
                'entries in ant_1_array and ant_2_array). May be smaller ' +
                'than the number of antennas in the array')
        self._Nants_data = uvp.UVParameter('Nants_data', description=desc,
                                           expected_type=int)

        desc = ('List of antenna names, shape (Nants_telescope), '
                'with numbers given by antenna_numbers (which can be matched '
                'to ant_1_array and ant_2_array). There must be one entry '
                'here for each unique entry in ant_1_array and '
                'ant_2_array, but there may be extras as well.')
        self._antenna_names = uvp.UVParameter('antenna_names', description=desc,
                                              form=('Nants_telescope',),
                                              expected_type=str)

        desc = ('List of integer antenna numbers corresponding to antenna_names, '
                'shape (Nants_telescope). There must be one '
                'entry here for each unique entry in ant_1_array and '
                'ant_2_array, but there may be extras as well.')
        self._antenna_numbers = uvp.UVParameter('antenna_numbers', description=desc,
                                                form=('Nants_telescope',),
                                                expected_type=int)

        desc = ('Number of antennas in the array. May be larger ' +
                'than the number of antennas with data')
        self._Nants_telescope = uvp.UVParameter('Nants_telescope',
                                                description=desc, expected_type=int)

        desc = 'Array of frequencies, shape (Nspws, Nfreqs), units Hz'
        self._freq_array = uvp.UVParameter('freq_array', description=desc,
                                           form=('Nspws', 'Nfreqs'),
                                           expected_type=np.float,
                                           tols=1e-3)  # mHz

        desc = ('Array of polarization integers, shape (Npols). '
                'AIPS Memo 117 says: stokes 1:4 (I,Q,U,V);  '
                'circular -1:-4 (RR,LL,RL,LR); linear -5:-8 (XX,YY,XY,YX)')
        self._polarization_array = uvp.UVParameter('polarization_array',
                                                   description=desc,
                                                   expected_type=int,
                                                   sane_vals=list(np.arange(-8, 0)) + list(np.arange(1, 5)),
                                                   form=('Npols',))

        desc = ('Array of times, center of integration, shape (Nblts), ' +
                'units Julian Date')
        self._time_array = uvp.UVParameter('time_array', description=desc,
                                           form=('Nblts',),
                                           expected_type=np.float,
                                           tols=1e-3 / (60.0 * 60.0 * 24.0))

        desc = ('The convention for applying he calibration solutions to data.'
                'Divide/multiply uncalibrated data by gains.')
        self._gain_convention = uvp.UVParameter('gain_convention', form='str',
                                           expected_type=str,
                                           description=desc,
                                           sane_vals=['divide', 'multiply'])

        desc = ('Array of flags, True is flagged. shape: (Nants_data, Nfreqs, Ntimes, '
                'Npols), type = bool.')
        self._flag_array = uvp.UVParameter('flag_array', description=desc,
                                           form=('Nants_data', 'Nfreqs',
                                                 'Ntimes', 'Npols'),
                                           expected_type=np.bool)

        desc = ('Array of qualities, shape: (Nants_data, Nfreqs, Ntimes, '
                'Npols), type = float.')
        self._quality_array = uvp.UVParameter('quality_array', description=desc,
                                           form=('Nants_data', 'Nfreqs',
                                                 'Ntimes', 'Npols'),
                                           expected_type=np.float)
        desc = ('Delay or gain switch parameter. Values are delay/gain.')
        self._delay_gain_switch = uvp.UVParameter('delay_gain_switch', form='str',
                                           expected_type=str,
                                           description=desc,
                                           sane_vals=['delay', 'gain'])

        desc = ('Polarization orientation. Values are E/N/Unknown.')
        self._x_orientation = uvp.UVParameter('x_orientation', description=desc,
                                            expected_type=str,
                                            sane_vals=['E', 'N', 'U'])
#--- optional parameters ---
        desc = ('Array of gains, shape: (Nants_data, Nfreqs, Ntimes, '
                'Npols), type = complex float.')
        self._gain_array = uvp.UVParameter('gain_array', description=desc,
                                           required=False,
                                           form=('Nants_data', 'Nfreqs',
                                                 'Ntimes', 'Npols'),
                                           expected_type=np.complex)

        desc = ('Array of delays. shape= (Nants_data, Ntimes, Npols), type = float')
        self._delay_array = uvp.UVParameter('delay_array', description=desc,
                                            required=False,
                                            form=('Nants_data', 'Ntimes',
                                                   'Npols')
                                            expected_type=np.float)

#--- truly optional parameters ---
        desc = ('Array of input flags, True is flagged. shape: (Nants_data, '
        'Nfreqs, Ntimes, Npols), type = bool.')
        self._input_flag_array = uvp.UVParameter('input_flag_array',
                                           description=desc,
                                           required=False,
                                           form=('Nants_data', 'Nfreqs',
                                                 'Ntimes', 'Npols'),
                                           expected_type=np.bool)
        super(UVData, self).__init__()
        #need to have either gain_array or delay_array
<commit_msg>update imports in cal.py and super call.<commit_after>#comment notes
#1. Currently supporting direction independent gains. Future versions may
#   support direction dependent gains with list of such objects.
#2.
import numpy as np
from uvbase import UVBase
import parameter as uvp

class UVCal(UVBase):
    """ A class defining a calibration """
    def __init__(self):
        self._Nfreqs = uvp.UVParameter('Nfreqs', description='Number of frequency channels',
                                       expected_type=int)
        self._Npols = uvp.UVParameter('Npols', description='Number of polarizations',
                                      expected_type=int)
        self._Ntimes = uvp.UVParameter('Ntimes', description='Number of times',
                                       expected_type=int)
        self._history = uvp.UVParameter('history', description='String of history, units English',
                                        form='str', expected_type=str)

        desc = ('Number of antennas with data present (i.e. number of unique '
                'entries in ant_1_array and ant_2_array). May be smaller ' +
                'than the number of antennas in the array')
        self._Nants_data = uvp.UVParameter('Nants_data', description=desc,
                                           expected_type=int)

        desc = ('List of antenna names, shape (Nants_telescope), '
                'with numbers given by antenna_numbers (which can be matched '
                'to ant_1_array and ant_2_array). There must be one entry '
                'here for each unique entry in ant_1_array and '
                'ant_2_array, but there may be extras as well.')
        self._antenna_names = uvp.UVParameter('antenna_names', description=desc,
                                              form=('Nants_telescope',),
                                              expected_type=str)

        desc = ('List of integer antenna numbers corresponding to antenna_names, '
                'shape (Nants_telescope). There must be one '
                'entry here for each unique entry in ant_1_array and '
                'ant_2_array, but there may be extras as well.')
        self._antenna_numbers = uvp.UVParameter('antenna_numbers', description=desc,
                                                form=('Nants_telescope',),
                                                expected_type=int)

        desc = ('Number of antennas in the array. May be larger ' +
                'than the number of antennas with data')
        self._Nants_telescope = uvp.UVParameter('Nants_telescope',
                                                description=desc, expected_type=int)

        desc = 'Array of frequencies, shape (Nspws, Nfreqs), units Hz'
        self._freq_array = uvp.UVParameter('freq_array', description=desc,
                                           form=('Nspws', 'Nfreqs'),
                                           expected_type=np.float,
                                           tols=1e-3)  # mHz

        desc = ('Array of polarization integers, shape (Npols). '
                'AIPS Memo 117 says: stokes 1:4 (I,Q,U,V);  '
                'circular -1:-4 (RR,LL,RL,LR); linear -5:-8 (XX,YY,XY,YX)')
        self._polarization_array = uvp.UVParameter('polarization_array',
                                                   description=desc,
                                                   expected_type=int,
                                                   sane_vals=list(np.arange(-8, 0)) + list(np.arange(1, 5)),
                                                   form=('Npols',))

        desc = ('Array of times, center of integration, shape (Nblts), ' +
                'units Julian Date')
        self._time_array = uvp.UVParameter('time_array', description=desc,
                                           form=('Nblts',),
                                           expected_type=np.float,
                                           tols=1e-3 / (60.0 * 60.0 * 24.0))

        desc = ('The convention for applying he calibration solutions to data.'
                'Divide/multiply uncalibrated data by gains.')
        self._gain_convention = uvp.UVParameter('gain_convention', form='str',
                                           expected_type=str,
                                           description=desc,
                                           sane_vals=['divide', 'multiply'])

        desc = ('Array of flags, True is flagged. shape: (Nants_data, Nfreqs, Ntimes, '
                'Npols), type = bool.')
        self._flag_array = uvp.UVParameter('flag_array', description=desc,
                                           form=('Nants_data', 'Nfreqs',
                                                 'Ntimes', 'Npols'),
                                           expected_type=np.bool)

        desc = ('Array of qualities, shape: (Nants_data, Nfreqs, Ntimes, '
                'Npols), type = float.')
        self._quality_array = uvp.UVParameter('quality_array', description=desc,
                                           form=('Nants_data', 'Nfreqs',
                                                 'Ntimes', 'Npols'),
                                           expected_type=np.float)
        desc = ('Delay or gain switch parameter. Values are delay/gain.')
        self._delay_gain_switch = uvp.UVParameter('delay_gain_switch', form='str',
                                           expected_type=str,
                                           description=desc,
                                           sane_vals=['delay', 'gain'])

        desc = ('Polarization orientation. Values are E/N/Unknown.')
        self._x_orientation = uvp.UVParameter('x_orientation', description=desc,
                                            expected_type=str,
                                            sane_vals=['E', 'N', 'U'])
#--- optional parameters ---
        desc = ('Array of gains, shape: (Nants_data, Nfreqs, Ntimes, '
                'Npols), type = complex float.')
        self._gain_array = uvp.UVParameter('gain_array', description=desc,
                                           required=False,
                                           form=('Nants_data', 'Nfreqs',
                                                 'Ntimes', 'Npols'),
                                           expected_type=np.complex)

        desc = ('Array of delays. shape= (Nants_data, Ntimes, Npols), type = float')
        self._delay_array = uvp.UVParameter('delay_array', description=desc,
                                            required=False,
                                            form=('Nants_data', 'Ntimes',
                                                   'Npols'),
                                            expected_type=np.float)

#--- truly optional parameters ---
        desc = ('Array of input flags, True is flagged. shape: (Nants_data, '
        'Nfreqs, Ntimes, Npols), type = bool.')
        self._input_flag_array = uvp.UVParameter('input_flag_array',
                                           description=desc,
                                           required=False,
                                           form=('Nants_data', 'Nfreqs',
                                                 'Ntimes', 'Npols'),
                                           expected_type=np.bool)
        super(UVCal, self).__init__()
        #need to have either gain_array or delay_array
<|endoftext|>


<commit_before>from scipy.interpolate import RegularGridInterpolator
from scipy import interpolate, optimize, integrate

import multiprocessing as mp
import matplotlib.pyplot as plt
import numpy as np

import iofile
import rotations

#from iofile import *

def z_slides(x, y, z, function):
    return function([x, y, z])

class Calculation:

    def __init__(self,
                 electronic_density,
                 ranges,
                 order=1,
                 center=(0, 0, 0),
                 align = (0, 0, 1),
                 radial=False,
                 r_grid=(20, 20, 20),
                 ):

        self._electronic_density = None
        self._order = order
        self._center = center
        self._align = align
        self._ranges_cart = ranges
        self._radial = radial

        self._density = None
        self._total_overlap = None
        self._measure = None

        self._radial_grid = r_grid

        self._fn_electronic_density = None
        self._fn_overlap = None
        self._fn_density2 = None
        self._fn_density = None

        print('Boundary cartesian')
        print(' x: {0} {1}'.format(ranges[0][0], ranges[0][-1]))
        print(' y: {0} {1}'.format(ranges[1][0], ranges[1][-1]))
        print(' z: {0} {1}'.format(ranges[2][0], ranges[2][-1]))

        self._fn_electronic_density = RegularGridInterpolator(self._ranges_cart, electronic_density, bounds_error=False, fill_value=0)

        if radial:
            print('Using cylindrical')

            r1 = np.max([ranges[0][-1]-center[0],center[0]-ranges[0][1]])
            r2 = np.max([ranges[1][-1]-center[1],center[1]-ranges[1][1]])
            r3 = np.max([ranges[2][-1]-center[2],center[2]-ranges[2][1]])

            max_rad = np.max([r1, r2, r3])

            r = np.linspace(0, float(max_rad), self._radial_grid[0])
            long_c = np.linspace(np.linalg.norm(center)-max_rad,np.linalg.norm(center)+max_rad, self._radial_grid[2])

            angle = np.linspace(0, 2 * np.pi, self._radial_grid[1])
     #       r = linspace(0, 15, self._radial_grid[2])

            x = long_c
            y = angle
            z = r
            self._ranges = [x, y, z]

            print(' z:     {0} {1}'.format(self._ranges[0][0], self._ranges[0][-1]))
            print(' angle: {0} {1}'.format(self._ranges[1][0], self._ranges[1][-1]))
            print(' r:     {0} {1}'.format(self._ranges[2][0], self._ranges[2][-1]))

        else:
            self._ranges = self._ranges_cart

    def plot_full(self, step, rotation=0):
        x, y, z = self._ranges_cart

        Y, X = np.meshgrid(y, x)

        density_slide = self.fn_electronic_density(np.array([[rotations.rotate_align_z([i, j, 0],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align
                                                                                ) for j in y] for i in x]))

        for val in np.arange(self._ranges[0][0]+step, self._ranges[0][-1], step):
            density_slide += self.fn_electronic_density(np.array([[rotations.rotate_align_z([i, j, val],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align
                                                                                ) for j in y] for i in x]))
        density_slide *=step
    #    levels = np.linspace(0, 0.5, 40)
        plt.xlabel("Bohr")
        plt.ylabel("Bohr")
        plt.contourf(X, Y, density_slide)
        plt.colorbar().set_label("Electronic density (e-/Bohr^2)")
        plt.show()

    def plot_full_rad(self, step, rotation=0):
        x, y, z = self._ranges
        Z, Y = np.meshgrid(z, y)
        density_slide = self.fn_electronic_density(np.array([[rotations.rotate_align_z([0, j, k],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial
                                                                                ) for k in z] for j in y]))

        for val in np.arange(self._ranges[0][0]+step, self._ranges[0][-1], step):
            density_slide += self.fn_electronic_density(np.array([[rotations.rotate_align_z([i, j, val],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial
                                                                                ) for j in y] for i in x]))
        density_slide *=step

        plt.contourf(Y, Z, density_slide)
        if self._radial:
            plt.xlabel("Angle (Radian)")
            plt.ylabel("Radius (Bohr)")
        else:
            plt.xlabel("Bohr")
            plt.ylabel("Bohr")
        plt.colorbar().set_label("Electronic density (e-/Bohr^3)")
        plt.show()

    def plot_slide(self, val, rotation=0):
        x, y, z = self._ranges_cart

        Y, X = np.meshgrid(y, x)
        density_slide = self.fn_electronic_density(np.array([[rotations.rotate_align_z([i, j, val],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align
                                                                                ) for j in y] for i in x]))
    #    levels = np.linspace(0, 0.5, 40)
        plt.xlabel("Bohr")
        plt.ylabel("Bohr")
        plt.contourf(X, Y, density_slide)
        plt.colorbar().set_label("Electronic density (e-/Bohr^3)")
        plt.show()

    def plot_slide_rad(self, val, rotation=0):
        x, y, z = self._ranges
        Z, Y = np.meshgrid(z, y)
        density_slide = self.fn_electronic_density(np.array([[rotations.rotate_align_z([val, j, k],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial
                                                                                ) for k in z] for j in y]))

        plt.contourf(Y, Z, density_slide)
        if self._radial:
            plt.xlabel("Angle (Radian)")
            plt.ylabel("Radius (Bohr)")
        else:
            plt.xlabel("Bohr")
            plt.ylabel("Bohr")
        plt.colorbar().set_label("Electronic density (e-/Bohr^3)")
        plt.show()

    def get_density(self):
        if self._density is None:
            x, y, z = self._ranges
            self._density = self.fn_electronic_density(np.array([[[rotations.rotate_align_z([i, j, k],
                                                                                0,
                                                                                center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial,
                                                                                ) for k in z] for j in y] for i in x]))
        return self._density

    def get_total_overlap(self):
        if self._total_overlap is None:
       #     print('getting_overlap')
            x, y, z = self._ranges
            overlappings = []
            for i in range(self._order):
  #              print('overlap: {0}'.format(i))
                angle = 2*np.pi/self._order * i
                dens_rot = self.fn_electronic_density(np.array([[[rotations.rotate_align_z([i, j, k],
                                                                                angle, center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial
                                                                                ) for k in z] for j in y] for i in x]))

                overlappings.append(np.multiply(self.get_density(), dens_rot))

            self._total_overlap = sum(overlappings)/self._order
        return self._total_overlap

    def get_measure(self, n_points=None, epsabs=1e2, epsrel=1e2, measure_error=1E-5):

        if self._measure is None:

            if n_points:
                measure_points = np.linspace(self._ranges[2][0], self._ranges[2][-1], n_points)
            else:
                measure_points = np.linspace(self._ranges[2][0], self._ranges[2][-1], self._radial_grid[2])

            minx = self._ranges[0][0]
            maxx = self._ranges[0][-1]
            miny= self._ranges[1][0]
            maxy = self._ranges[1][-1]
       #     print(minx, maxx)
       #     print(miny, maxy)

            print('#     coordinate     measure(C{0})       overlap        density2        density'.format(self._order))
            measure = {'symmetry' : [], 'coordinate' : [], 'overlap' : [], 'density2' : [], 'density' : []}
            for z_slide in measure_points:
                [integral_overlap, error_overlap] = integrate.dblquad(z_slides,
                                                                      miny ,maxy, lambda x: minx, lambda x: maxx,
                                                                      args=(z_slide, self.fn_overlap),
                                                                      epsabs=epsabs, epsrel=epsrel)

                [integral_density2, error_density2] = integrate.dblquad(z_slides,
                                                                       miny ,maxy, lambda x: minx, lambda x: maxx,
                                                                       args=(z_slide, self.fn_density2),
                                                                       epsabs=epsabs, epsrel=epsrel)

                [integral_density, error_density] = integrate.dblquad(z_slides,
                                                                       miny ,maxy, lambda x: minx, lambda x: maxx,
                                                                       args=(z_slide, self.fn_density),
                                                                       epsabs=epsabs, epsrel=epsrel)

                if abs(integral_density2-integral_overlap) < measure_error:
                    normalized_measure =0
                else:
                    normalized_measure = (1-integral_overlap/integral_density2)*100

                if self._radial:
                    integral_density *= z_slide
                    integral_density2 *= z_slide
                    integral_overlap *=z_slide

                print('{0:15.8f} {1:15.8f} {2:15.8f} {3:15.8f} {4:15.8f}'.format(z_slide,
                                                                                normalized_measure,
                                                                                integral_overlap,
                                                                                integral_density2,
                                                                                integral_density))

                measure['symmetry'].append(normalized_measure)
                measure['coordinate'].append(z_slide)
                measure['density2'].append(integral_density2)
                measure['density'].append(integral_density)
                measure['overlap'].append(integral_overlap)

            self._measure = measure

        return self._measure

    def plot_measure(self):

        fig, ax1 = plt.subplots()

        ax2 = ax1.twinx()
        ax3 = ax1.twiny()

        ax1.set_xlabel('Bohr')
        ax1.set_ylabel('Asymmetry %')
        ax2.set_ylabel('Density (e-/Bohr)')
        ax3.set_xlabel('Electrons')

        ax3.xaxis.grid()

        data = self.get_measure()

        step = data['coordinate'][1] - data['coordinate'][0]

        fn = interpolate.interp1d(data['coordinate'], np.cumsum(data['density'])*step ,bounds_error=False)
        def fun(x, n):
            return fn(x) - n

        total_electron = int(fun(data['coordinate'][-1], 0))

        n_step = abs(total_electron/10)+1
        positions = [optimize.bisect(fun, data['coordinate'][0], data['coordinate'][-1], args=(i, )) for i in range(1, total_electron+1, n_step)]

        positions = (np.array(positions) - data['coordinate'][0])/(data['coordinate'][-1] - data['coordinate'][0])
        plt.xticks(positions, range(1, total_electron+1, n_step))

        asim_p, = ax1.plot(data['coordinate'], data['symmetry'],
                           label='Asymmetry C{0}'.format(self._order), color='red', linewidth=4)
        #ax2.plot(data['coordinate'], data['density2'], label='Density2')
        dens_p, =ax2.plot(data['coordinate'], data['density'], label='Density', linestyle='--')
        #ax2.plot(data['coordinate'], data['overlap'], label='Overlap')
        plt.ylim(0,100)
        plt.legend(handles=[asim_p, dens_p],loc=1)

        plt.show()

    #Properties
    @property
    def fn_electronic_density(self):
        return self._fn_electronic_density

    @property
    def fn_overlap(self):
        if self._fn_overlap is None:
            x, y, z = self._ranges
            total_overlap = self.get_total_overlap()
            self._fn_overlap = RegularGridInterpolator((x, y, z), total_overlap, bounds_error=False)
        return self._fn_overlap

    @property
    def fn_density2(self):
        if self._fn_density2 is None:
            x, y, z = self._ranges

            density = self.get_density()
            self._fn_density2 = RegularGridInterpolator((x, y, z), np.square(density), bounds_error=False)
        return self._fn_density2

    @property
    def fn_density(self):
        if self._fn_density is None:
            x, y, z = self._ranges

            density = self.get_density()
            self._fn_density = RegularGridInterpolator((x, y, z), density, bounds_error=False)
        return self._fn_density

if __name__ == "__main__":
    #ranges, cube_density = get_density_cube('C3B9.cube')

    ranges, cube_density = iofile.get_density_gaussian('potential')
    #ranges, cube_density = get_density_gaussian()

    calculation = Calculation(cube_density, ranges, order=3, align=[0, 0, 1], center=[0, 0, 0], radial=True)

    calculation.plot_slide(2, rotation=0)

    calculation.plot_slide_rad(0, rotation=0)

    calculation.plot_measure(30)<commit_msg>Added full integration measure<commit_after>from scipy.interpolate import RegularGridInterpolator
from scipy import interpolate, optimize, integrate

import multiprocessing as mp
import matplotlib.pyplot as plt
import numpy as np

import iofile
import rotations

#from iofile import *

def z_slides(x, y, z, function):
    return function([x, y, z])

class Calculation:

    def __init__(self,
                 electronic_density,
                 ranges,
                 order=1,
                 center=(0, 0, 0),
                 align = (0, 0, 1),
                 radial=False,
                 r_grid=(20, 20, 20),
                 ):

        self._electronic_density = None
        self._order = order
        self._center = center
        self._align = align
        self._ranges_cart = ranges
        self._radial = radial

        self._density = None
        self._total_overlap = None
        self._measure = None

        self._radial_grid = r_grid

        self._fn_electronic_density = None
        self._fn_overlap = None
        self._fn_density2 = None
        self._fn_density = None

        print('Boundary cartesian')
        print(' x: {0} {1}'.format(ranges[0][0], ranges[0][-1]))
        print(' y: {0} {1}'.format(ranges[1][0], ranges[1][-1]))
        print(' z: {0} {1}'.format(ranges[2][0], ranges[2][-1]))

        self._fn_electronic_density = RegularGridInterpolator(self._ranges_cart, electronic_density, bounds_error=False, fill_value=0)

        if radial:
            print('Using cylindrical')

            r1 = np.max([ranges[0][-1]-center[0],center[0]-ranges[0][1]])
            r2 = np.max([ranges[1][-1]-center[1],center[1]-ranges[1][1]])
            r3 = np.max([ranges[2][-1]-center[2],center[2]-ranges[2][1]])

            max_rad = np.max([r1, r2, r3])

            r = np.linspace(0, float(max_rad), self._radial_grid[0])
            long_c = np.linspace(np.linalg.norm(center)-max_rad,np.linalg.norm(center)+max_rad, self._radial_grid[2])

            angle = np.linspace(0, 2 * np.pi, self._radial_grid[1])
     #       r = linspace(0, 15, self._radial_grid[2])

            x = long_c
            y = angle
            z = r
            self._ranges = [x, y, z]

            print(' z:     {0} {1}'.format(self._ranges[0][0], self._ranges[0][-1]))
            print(' angle: {0} {1}'.format(self._ranges[1][0], self._ranges[1][-1]))
            print(' r:     {0} {1}'.format(self._ranges[2][0], self._ranges[2][-1]))

        else:
            self._ranges = self._ranges_cart

    def plot_full(self, step, rotation=0):
        x, y, z = self._ranges_cart

        Y, X = np.meshgrid(y, x)

        density_slide = self.fn_electronic_density(np.array([[rotations.rotate_align_z([i, j, 0],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align
                                                                                ) for j in y] for i in x]))

        for val in np.arange(self._ranges[0][0]+step, self._ranges[0][-1], step):
            density_slide += self.fn_electronic_density(np.array([[rotations.rotate_align_z([i, j, val],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align
                                                                                ) for j in y] for i in x]))
        density_slide *=step
    #    levels = np.linspace(0, 0.5, 40)
        plt.xlabel("Bohr")
        plt.ylabel("Bohr")
        plt.contourf(X, Y, density_slide)
        plt.colorbar().set_label("Electronic density (e-/Bohr^2)")
        plt.show()

    def plot_full_rad(self, step, rotation=0):
        x, y, z = self._ranges
        Z, Y = np.meshgrid(z, y)
        density_slide = self.fn_electronic_density(np.array([[rotations.rotate_align_z([0, j, k],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial
                                                                                ) for k in z] for j in y]))

        for val in np.arange(self._ranges[0][0]+step, self._ranges[0][-1], step):
            density_slide += self.fn_electronic_density(np.array([[rotations.rotate_align_z([i, j, val],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial
                                                                                ) for j in y] for i in x]))
        density_slide *=step

        plt.contourf(Y, Z, density_slide)
        if self._radial:
            plt.xlabel("Angle (Radian)")
            plt.ylabel("Radius (Bohr)")
        else:
            plt.xlabel("Bohr")
            plt.ylabel("Bohr")
        plt.colorbar().set_label("Electronic density (e-/Bohr^3)")
        plt.show()

    def plot_slide(self, val, rotation=0):
        x, y, z = self._ranges_cart

        Y, X = np.meshgrid(y, x)
        density_slide = self.fn_electronic_density(np.array([[rotations.rotate_align_z([i, j, val],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align
                                                                                ) for j in y] for i in x]))
    #    levels = np.linspace(0, 0.5, 40)
        plt.xlabel("Bohr")
        plt.ylabel("Bohr")
        plt.contourf(X, Y, density_slide)
        plt.colorbar().set_label("Electronic density (e-/Bohr^3)")
        plt.show()

    def plot_slide_rad(self, val, rotation=0):
        x, y, z = self._ranges
        Z, Y = np.meshgrid(z, y)
        density_slide = self.fn_electronic_density(np.array([[rotations.rotate_align_z([val, j, k],
                                                                                rotation,
                                                                                center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial
                                                                                ) for k in z] for j in y]))

        plt.contourf(Y, Z, density_slide)
        if self._radial:
            plt.xlabel("Angle (Radian)")
            plt.ylabel("Radius (Bohr)")
        else:
            plt.xlabel("Bohr")
            plt.ylabel("Bohr")
        plt.colorbar().set_label("Electronic density (e-/Bohr^3)")
        plt.show()

    def get_density(self):
        if self._density is None:
            x, y, z = self._ranges
            self._density = self.fn_electronic_density(np.array([[[rotations.rotate_align_z([i, j, k],
                                                                                0,
                                                                                center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial,
                                                                                ) for k in z] for j in y] for i in x]))
        return self._density

    def get_total_overlap(self):
        if self._total_overlap is None:
       #     print('getting_overlap')
            x, y, z = self._ranges
            overlappings = []
            for i in range(self._order):
  #              print('overlap: {0}'.format(i))
                angle = 2*np.pi/self._order * i
                dens_rot = self.fn_electronic_density(np.array([[[rotations.rotate_align_z([i, j, k],
                                                                                angle, center=self._center,
                                                                                align=self._align,
                                                                                radial=self._radial
                                                                                ) for k in z] for j in y] for i in x]))

                overlappings.append(np.multiply(self.get_density(), dens_rot))

            self._total_overlap = sum(overlappings)/self._order
        return self._total_overlap

    def get_measure(self, n_points=None, epsabs=1e2, epsrel=1e2, measure_error=1E-5):

        if self._measure is None:

            if n_points:
                measure_points = np.linspace(self._ranges[2][0], self._ranges[2][-1], n_points)
            else:
                measure_points = np.linspace(self._ranges[2][0], self._ranges[2][-1], self._radial_grid[2])

            minx = self._ranges[0][0]
            maxx = self._ranges[0][-1]
            miny= self._ranges[1][0]
            maxy = self._ranges[1][-1]
       #     print(minx, maxx)
       #     print(miny, maxy)

            print('#     coordinate     measure(C{0})       overlap        density2        density'.format(self._order))
            measure = {'symmetry' : [], 'coordinate' : [], 'overlap' : [], 'density2' : [], 'density' : []}
            for z_slide in measure_points:
                [integral_overlap, error_overlap] = integrate.dblquad(z_slides,
                                                                      miny ,maxy, lambda x: minx, lambda x: maxx,
                                                                      args=(z_slide, self.fn_overlap),
                                                                      epsabs=epsabs, epsrel=epsrel)

                [integral_density2, error_density2] = integrate.dblquad(z_slides,
                                                                       miny ,maxy, lambda x: minx, lambda x: maxx,
                                                                       args=(z_slide, self.fn_density2),
                                                                       epsabs=epsabs, epsrel=epsrel)

                [integral_density, error_density] = integrate.dblquad(z_slides,
                                                                       miny ,maxy, lambda x: minx, lambda x: maxx,
                                                                       args=(z_slide, self.fn_density),
                                                                       epsabs=epsabs, epsrel=epsrel)

                if abs(integral_density2-integral_overlap) < measure_error:
                    normalized_measure =0
                else:
                    normalized_measure = (1-integral_overlap/integral_density2)*100

                if self._radial:
                    integral_density *= z_slide
                    integral_density2 *= z_slide
                    integral_overlap *=z_slide

                print('{0:15.8f} {1:15.8f} {2:15.8f} {3:15.8f} {4:15.8f}'.format(z_slide,
                                                                                normalized_measure,
                                                                                integral_overlap,
                                                                                integral_density2,
                                                                                integral_density))

                measure['symmetry'].append(normalized_measure)
                measure['coordinate'].append(z_slide)
                measure['density2'].append(integral_density2)
                measure['density'].append(integral_density)
                measure['overlap'].append(integral_overlap)

            # Total measure
            total_overlap = integrate.simps(measure['overlap'], measure['coordinate'],even='avg')
            total_density2 = integrate.simps(measure['density2'], measure['coordinate'],even='avg')
            total_symmetry = (1-total_overlap/total_density2)*100
            measure.update({'total_symmetry' : total_symmetry})
            print ('\nTotal measure: {0}'.format(total_symmetry))

            self._measure = measure

        return self._measure

    def plot_measure(self):

        fig, ax1 = plt.subplots()

        ax2 = ax1.twinx()
        ax3 = ax1.twiny()

        ax1.set_xlabel('Bohr')
        ax1.set_ylabel('Asymmetry %')
        ax2.set_ylabel('Density (e-/Bohr)')
        ax3.set_xlabel('Electrons')

        ax3.xaxis.grid()

        data = self.get_measure()

        step = data['coordinate'][1] - data['coordinate'][0]

        fn = interpolate.interp1d(data['coordinate'], np.cumsum(data['density'])*step ,bounds_error=False)
        def fun(x, n):
            return fn(x) - n

        total_electron = int(fun(data['coordinate'][-1], 0))

        n_step = abs(total_electron/10)+1
        positions = [optimize.bisect(fun, data['coordinate'][0], data['coordinate'][-1], args=(i, )) for i in range(1, total_electron+1, n_step)]

        positions = (np.array(positions) - data['coordinate'][0])/(data['coordinate'][-1] - data['coordinate'][0])
        plt.xticks(positions, range(1, total_electron+1, n_step))

        asim_p, = ax1.plot(data['coordinate'], data['symmetry'],
                           label='Asymmetry C{0}'.format(self._order), color='red', linewidth=4)
        #ax2.plot(data['coordinate'], data['density2'], label='Density2')
        dens_p, =ax2.plot(data['coordinate'], data['density'], label='Density', linestyle='--')
        #ax2.plot(data['coordinate'], data['overlap'], label='Overlap')
        plt.ylim(0,100)
        plt.legend(handles=[asim_p, dens_p],loc=1)

        plt.show()

    #Properties
    @property
    def fn_electronic_density(self):
        return self._fn_electronic_density

    @property
    def fn_overlap(self):
        if self._fn_overlap is None:
            x, y, z = self._ranges
            total_overlap = self.get_total_overlap()
            self._fn_overlap = RegularGridInterpolator((x, y, z), total_overlap, bounds_error=False)
        return self._fn_overlap

    @property
    def fn_density2(self):
        if self._fn_density2 is None:
            x, y, z = self._ranges

            density = self.get_density()
            self._fn_density2 = RegularGridInterpolator((x, y, z), np.square(density), bounds_error=False)
        return self._fn_density2

    @property
    def fn_density(self):
        if self._fn_density is None:
            x, y, z = self._ranges

            density = self.get_density()
            self._fn_density = RegularGridInterpolator((x, y, z), density, bounds_error=False)
        return self._fn_density

if __name__ == "__main__":
    #ranges, cube_density = get_density_cube('C3B9.cube')

    ranges, cube_density = iofile.get_density_gaussian('potential')
    #ranges, cube_density = get_density_gaussian()

    calculation = Calculation(cube_density, ranges, order=3, align=[0, 0, 1], center=[0, 0, 0], radial=True)

    calculation.plot_slide(2, rotation=0)

    calculation.plot_slide_rad(0, rotation=0)

    calculation.plot_measure(30)<|endoftext|>