<commit_before># -*- coding: utf-8 -*-
"""Supports the Michelson Interferometer for Global High-resolution
Thermospheric Imaging (MIGHTI) instrument onboard the Ionospheric
CONnection Explorer (ICON) satellite.  Accesses local data in
netCDF format.

Parameters
----------
platform : string
    'icon'
name : string
    'mighti'
tag : string
    Supports 'los_wind_green', 'los_wind_red', 'vector_wind_green',
    'vector_wind_red', 'temperature'.  Note that not every data product
    available for every sat_id
sat_id : string
    '', 'a', or 'b'

Warnings
--------
- The cleaning parameters for the instrument are still under development.
- Only supports level-2 data.

Example
-------
::

    import pysat
    mighti = pysat.Instrument('icon', 'mighti', 'vector_wind_green',
                              clean_level='clean')
    mighti.download(dt.datetime(2020, 1, 30), dt.datetime(2020, 1, 31))
    mighti.load(2020, 2)

By default, pysat removes the ICON level tags from variable names, ie,
ICON_L27_Ion_Density becomes Ion_Density.  To retain the original names, use
::
    mighti = pysat.Instrument(platform='icon', name='mighti',
                              tag='vector_wind_green', clean_level='clean',
                              keep_original_names=True)

Authors
---------
Originated from EUV support.
Jeff Klenzing, Mar 17, 2018, Goddard Space Flight Center
Russell Stoneback, Mar 23, 2018, University of Texas at Dallas
Conversion to MIGHTI, Oct 8th, 2028, University of Texas at Dallas

"""

from __future__ import print_function
from __future__ import absolute_import

import datetime as dt
import functools
import logging
import numpy as np

import pysat
from pysat.instruments.methods import general as mm_gen
from pysat.instruments.methods import icon as mm_icon

logger = logging.getLogger(__name__)

platform = 'icon'
name = 'mighti'
tags = {'los_wind_green': 'Line of sight wind data -- Green Line',
        'los_wind_red': 'Line of sight wind data -- Red Line',
        'vector_wind_green': 'Vector wind data -- Green Line',
        'vector_wind_red': 'Vector wind data -- Red Line',
        'temperature': 'Neutral temperature data'}
sat_ids = {'': ['vector_wind_green', 'vector_wind_red'],
           'a': ['los_wind_green', 'los_wind_red', 'temperature'],
           'b': ['los_wind_green', 'los_wind_red', 'temperature']}
_test_dates = {jj: {kk: dt.datetime(2020, 1, 2) for kk in sat_ids[jj]}
               for jj in sat_ids.keys()}
_test_download_travis = {jj: {kk: False for kk in sat_ids[jj]}
                         for jj in sat_ids.keys()}
pandas_format = False

datestr = '{year:04d}-{month:02d}-{day:02d}_v{version:02d}r{revision:03d}'
fname1 = 'ICON_L2-1_MIGHTI-{id:s}_LOS-Wind-{color:s}_{date:s}.NC'
fname2 = 'ICON_L2-2_MIGHTI_Vector-Wind-{color:s}_{date:s}.NC'
fname3 = 'ICON_L2-3_MIGHTI-{id:s}_Temperature_{date:s}.NC'
supported_tags = {'': {'vector_wind_green': fname2.format(color='Green',
                                                          date=datestr),
                       'vector_wind_red': fname2.format(color='Red',
                                                        date=datestr)},
                  'a': {'los_wind_green': fname1.format(id='A', color='Green',
                                                        date=datestr),
                        'los_wind_red': fname1.format(id='A', color='Red',
                                                      date=datestr),
                        'temperature': fname3.format(id='A', date=datestr)},
                  'b': {'los_wind_green': fname1.format(id='B', color='Green',
                                                        date=datestr),
                        'los_wind_red': fname1.format(id='B', color='Red',
                                                      date=datestr),
                        'temperature': fname3.format(id='B', date=datestr)}}

# use the CDAWeb methods list files routine
list_files = functools.partial(mm_gen.list_files,
                               supported_tags=supported_tags)

# support download routine
dirstr = '/pub/LEVEL.2/MIGHTI{id:s}'
dirdatestr = '{year:4d}/{doy:03d}/'
ids = {'': '',
       'a': '-A',
       'b': '-B'}
products = {'vector_wind_green': 'Vector-Winds/',
            'vector_wind_red': 'Vector-Winds/',
            'los_wind_green': 'LOS-Winds/',
            'los_wind_red': 'LOS-Winds/',
            'temperature': 'Temperature/'}
datestr = '{year:04d}-{month:02d}-{day:02d}'

download_tags = {}
for skey in supported_tags.keys():
    download_tags[skey] = {}
    for tkey in supported_tags[skey].keys():
        fname = supported_tags[skey][tkey]

        download_tags[skey][tkey] = {'dir': dirstr.format(id=ids[skey]),
                                     'remote_fname': ''.join((products[tkey],
                                                              fname))}

download = functools.partial(mm_icon.ssl_download, supported_tags=download_tags)

# support listing files on SSL
list_remote_files = functools.partial(mm_icon.list_remote_files,
                                      supported_tags=download_tags)

def init(self):
    """Initializes the Instrument object with instrument specific values.

    Runs once upon instantiation.

    Parameters
    -----------
    inst : (pysat.Instrument)
        Instrument class object

    """

    logger.info(mm_icon.ackn_str)
    self.meta.acknowledgements = mm_icon.ackn_str
    self.meta.references = ''.join((mm_icon.refs['mission'],
                                    mm_icon.refs['mighti']))

    pass

def default(inst):
    """Default routine to be applied when loading data.  Adjusts epoch timestamps
    to datetimes and removes variable preambles.

    """

    mm_gen.convert_timestamp_to_datetime(inst, sec_mult=1.0e-3)
    if not inst.kwargs['keep_original_names']:
        remove_preamble(inst)

def remove_preamble(inst):
    """Removes preambles in variable names"""

    target = {'los_wind_green': 'ICON_L21_',
              'los_wind_red': 'ICON_L21_',
              'vector_wind_green': 'ICON_L22_',
              'vector_wind_red': 'ICON_L22_',
              'temperature': ['ICON_L1_MIGHTI_{}_'.format(inst.sat_id.upper()),
                              'ICON_L23_MIGHTI_{}_'.format(inst.sat_id.upper()),
                              'ICON_L23_']}
    mm_gen.remove_leading_text(inst, target=target[inst.tag])

    return

def load(fnames, tag=None, sat_id=None, keep_original_names=False):
    """Loads ICON FUV data using pysat into pandas.

    This routine is called as needed by pysat. It is not intended
    for direct user interaction.

    Parameters
    ----------
    fnames : array-like
        iterable of filename strings, full path, to data files to be loaded.
        This input is nominally provided by pysat itself.
    tag : string
        tag name used to identify particular data set to be loaded.
        This input is nominally provided by pysat itself.
    sat_id : string
        Satellite ID used to identify particular data set to be loaded.
        This input is nominally provided by pysat itself.
    keep_original_names : boolean
        if True then the names as given in the netCDF ICON file
        will be used as is. If False, a preamble is removed.

    Returns
    -------
    data, metadata
        Data and Metadata are formatted for pysat. Data is a pandas
        DataFrame while metadata is a pysat.Meta instance.

    Note
    ----
    Any additional keyword arguments passed to pysat.Instrument
    upon instantiation are passed along to this routine.

    Examples
    --------
    ::
        inst = pysat.Instrument('icon', 'fuv')
        inst.load(2020, 1)

    """

    return pysat.utils.load_netcdf4(fnames, epoch_name='Epoch',
                                    units_label='Units',
                                    name_label='Long_Name',
                                    notes_label='Var_Notes',
                                    desc_label='CatDesc',
                                    plot_label='FieldNam',
                                    axis_label='LablAxis',
                                    scale_label='ScaleTyp',
                                    min_label='ValidMin',
                                    max_label='ValidMax',
                                    fill_label='FillVal',
                                    pandas_format=pandas_format)

def clean(inst):
    """Provides data cleaning based upon clean_level.

    clean_level is set upon Instrument instantiation to
    one of the following:

    'Clean'
    'Dusty'
    'Dirty'
    'None'

    Routine is called by pysat, and not by the end user directly.

    Parameters
    -----------
    inst : (pysat.Instrument)
        Instrument class object, whose attribute clean_level is used to return
        the desired level of data selectivity.

    Note
    ----
        Supports 'clean', 'dusty', 'dirty', 'none'

    """

    if inst.tag.find('los') >= 0:
        # dealing with LOS winds
        var = 'Line_of_Sight_Wind'
        if 'Wind_Quality' in inst.variables:
            qual_flag = 'Quality_Flags'
        else:
            qual_flag = 'ICON_L21_Quality_Flags'
            var = 'ICON_L21_' + var
        if inst.clean_level in ['clean', 'dusty']:
            # find location with any of the flags set
            idx, idy, = np.where(inst[qual_flag].any(axis=2))
            inst[idx, idy, var] = np.nan
        else:
            # dirty and worse lets everything through
            pass

    elif inst.tag.find('vector') >= 0:
        # vector winds area
        vars = ['Zonal_Wind', 'Meridional_Wind']
        if 'Wind_Quality' in inst.variables:
            qual_flag = 'Wind_Quality'
        else:
            qual_flag = 'ICON_L22_Wind_Quality'
            vars = ['ICON_L22_' + x for x in vars]
        if inst.clean_level == 'clean':
            for var in vars:
                inst[var] = inst[var].where(inst[qual_flag] == 1.0)
        elif inst.clean_level == 'dusty':
            for var in vars:
                inst[var] = inst[var].where(inst[qual_flag] >= 0.5)
        else:
            # dirty lets everything through
            pass

    elif inst.tag.find('temp') >= 0:
        # neutral temperatures
        var = 'Temperature'
        saa_flag = 'Quality_Flag_South_Atlantic_Anomaly'
        cal_flag = 'Quality_Flag_Bad_Calibration'
        if saa_flag not in inst.variables:
            saa_flag = '_'.join(('ICON_L1_MIGHTI', inst.sat_id.upper(),
                                 saa_flag))
            cal_flag = '_'.join(('ICON_L1_MIGHTI', inst.sat_id.upper(),
                                 cal_flag))
            var = '_'.JOIN(('ICON_L23_MIGHTI', inst.sat_id.upper(), var))
        if inst.clean_level in ['clean', 'dusty']:
            # SAA
            idx, = np.where(inst[saa_flag.format(s=inst.sat_id.upper())] > 0)
            inst[:, idx, var] = np.nan
            # Calibration file
            idx, = np.where(inst[cal_flag.format(s=inst.sat_id.upper())] > 0)
            inst[:, idx, var] = np.nan
        else:
            # dirty and worse lets everything through
            pass

    return
<commit_msg>STY: remove unused string format<commit_after># -*- coding: utf-8 -*-
"""Supports the Michelson Interferometer for Global High-resolution
Thermospheric Imaging (MIGHTI) instrument onboard the Ionospheric
CONnection Explorer (ICON) satellite.  Accesses local data in
netCDF format.

Parameters
----------
platform : string
    'icon'
name : string
    'mighti'
tag : string
    Supports 'los_wind_green', 'los_wind_red', 'vector_wind_green',
    'vector_wind_red', 'temperature'.  Note that not every data product
    available for every sat_id
sat_id : string
    '', 'a', or 'b'

Warnings
--------
- The cleaning parameters for the instrument are still under development.
- Only supports level-2 data.

Example
-------
::

    import pysat
    mighti = pysat.Instrument('icon', 'mighti', 'vector_wind_green',
                              clean_level='clean')
    mighti.download(dt.datetime(2020, 1, 30), dt.datetime(2020, 1, 31))
    mighti.load(2020, 2)

By default, pysat removes the ICON level tags from variable names, ie,
ICON_L27_Ion_Density becomes Ion_Density.  To retain the original names, use
::
    mighti = pysat.Instrument(platform='icon', name='mighti',
                              tag='vector_wind_green', clean_level='clean',
                              keep_original_names=True)

Authors
---------
Originated from EUV support.
Jeff Klenzing, Mar 17, 2018, Goddard Space Flight Center
Russell Stoneback, Mar 23, 2018, University of Texas at Dallas
Conversion to MIGHTI, Oct 8th, 2028, University of Texas at Dallas

"""

from __future__ import print_function
from __future__ import absolute_import

import datetime as dt
import functools
import logging
import numpy as np

import pysat
from pysat.instruments.methods import general as mm_gen
from pysat.instruments.methods import icon as mm_icon

logger = logging.getLogger(__name__)

platform = 'icon'
name = 'mighti'
tags = {'los_wind_green': 'Line of sight wind data -- Green Line',
        'los_wind_red': 'Line of sight wind data -- Red Line',
        'vector_wind_green': 'Vector wind data -- Green Line',
        'vector_wind_red': 'Vector wind data -- Red Line',
        'temperature': 'Neutral temperature data'}
sat_ids = {'': ['vector_wind_green', 'vector_wind_red'],
           'a': ['los_wind_green', 'los_wind_red', 'temperature'],
           'b': ['los_wind_green', 'los_wind_red', 'temperature']}
_test_dates = {jj: {kk: dt.datetime(2020, 1, 2) for kk in sat_ids[jj]}
               for jj in sat_ids.keys()}
_test_download_travis = {jj: {kk: False for kk in sat_ids[jj]}
                         for jj in sat_ids.keys()}
pandas_format = False

datestr = '{year:04d}-{month:02d}-{day:02d}_v{version:02d}r{revision:03d}'
fname1 = 'ICON_L2-1_MIGHTI-{id:s}_LOS-Wind-{color:s}_{date:s}.NC'
fname2 = 'ICON_L2-2_MIGHTI_Vector-Wind-{color:s}_{date:s}.NC'
fname3 = 'ICON_L2-3_MIGHTI-{id:s}_Temperature_{date:s}.NC'
supported_tags = {'': {'vector_wind_green': fname2.format(color='Green',
                                                          date=datestr),
                       'vector_wind_red': fname2.format(color='Red',
                                                        date=datestr)},
                  'a': {'los_wind_green': fname1.format(id='A', color='Green',
                                                        date=datestr),
                        'los_wind_red': fname1.format(id='A', color='Red',
                                                      date=datestr),
                        'temperature': fname3.format(id='A', date=datestr)},
                  'b': {'los_wind_green': fname1.format(id='B', color='Green',
                                                        date=datestr),
                        'los_wind_red': fname1.format(id='B', color='Red',
                                                      date=datestr),
                        'temperature': fname3.format(id='B', date=datestr)}}

# use the CDAWeb methods list files routine
list_files = functools.partial(mm_gen.list_files,
                               supported_tags=supported_tags)

# support download routine
dirstr = '/pub/LEVEL.2/MIGHTI{id:s}'
dirdatestr = '{year:4d}/{doy:03d}/'
ids = {'': '',
       'a': '-A',
       'b': '-B'}
products = {'vector_wind_green': 'Vector-Winds/',
            'vector_wind_red': 'Vector-Winds/',
            'los_wind_green': 'LOS-Winds/',
            'los_wind_red': 'LOS-Winds/',
            'temperature': 'Temperature/'}
datestr = '{year:04d}-{month:02d}-{day:02d}'

download_tags = {}
for skey in supported_tags.keys():
    download_tags[skey] = {}
    for tkey in supported_tags[skey].keys():
        fname = supported_tags[skey][tkey]

        download_tags[skey][tkey] = {'dir': dirstr.format(id=ids[skey]),
                                     'remote_fname': ''.join((products[tkey],
                                                              fname))}

download = functools.partial(mm_icon.ssl_download, supported_tags=download_tags)

# support listing files on SSL
list_remote_files = functools.partial(mm_icon.list_remote_files,
                                      supported_tags=download_tags)

def init(self):
    """Initializes the Instrument object with instrument specific values.

    Runs once upon instantiation.

    Parameters
    -----------
    inst : (pysat.Instrument)
        Instrument class object

    """

    logger.info(mm_icon.ackn_str)
    self.meta.acknowledgements = mm_icon.ackn_str
    self.meta.references = ''.join((mm_icon.refs['mission'],
                                    mm_icon.refs['mighti']))

    pass

def default(inst):
    """Default routine to be applied when loading data.  Adjusts epoch timestamps
    to datetimes and removes variable preambles.

    """

    mm_gen.convert_timestamp_to_datetime(inst, sec_mult=1.0e-3)
    if not inst.kwargs['keep_original_names']:
        remove_preamble(inst)

def remove_preamble(inst):
    """Removes preambles in variable names"""

    target = {'los_wind_green': 'ICON_L21_',
              'los_wind_red': 'ICON_L21_',
              'vector_wind_green': 'ICON_L22_',
              'vector_wind_red': 'ICON_L22_',
              'temperature': ['ICON_L1_MIGHTI_{}_'.format(inst.sat_id.upper()),
                              'ICON_L23_MIGHTI_{}_'.format(inst.sat_id.upper()),
                              'ICON_L23_']}
    mm_gen.remove_leading_text(inst, target=target[inst.tag])

    return

def load(fnames, tag=None, sat_id=None, keep_original_names=False):
    """Loads ICON FUV data using pysat into pandas.

    This routine is called as needed by pysat. It is not intended
    for direct user interaction.

    Parameters
    ----------
    fnames : array-like
        iterable of filename strings, full path, to data files to be loaded.
        This input is nominally provided by pysat itself.
    tag : string
        tag name used to identify particular data set to be loaded.
        This input is nominally provided by pysat itself.
    sat_id : string
        Satellite ID used to identify particular data set to be loaded.
        This input is nominally provided by pysat itself.
    keep_original_names : boolean
        if True then the names as given in the netCDF ICON file
        will be used as is. If False, a preamble is removed.

    Returns
    -------
    data, metadata
        Data and Metadata are formatted for pysat. Data is a pandas
        DataFrame while metadata is a pysat.Meta instance.

    Note
    ----
    Any additional keyword arguments passed to pysat.Instrument
    upon instantiation are passed along to this routine.

    Examples
    --------
    ::
        inst = pysat.Instrument('icon', 'fuv')
        inst.load(2020, 1)

    """

    return pysat.utils.load_netcdf4(fnames, epoch_name='Epoch',
                                    units_label='Units',
                                    name_label='Long_Name',
                                    notes_label='Var_Notes',
                                    desc_label='CatDesc',
                                    plot_label='FieldNam',
                                    axis_label='LablAxis',
                                    scale_label='ScaleTyp',
                                    min_label='ValidMin',
                                    max_label='ValidMax',
                                    fill_label='FillVal',
                                    pandas_format=pandas_format)

def clean(inst):
    """Provides data cleaning based upon clean_level.

    clean_level is set upon Instrument instantiation to
    one of the following:

    'Clean'
    'Dusty'
    'Dirty'
    'None'

    Routine is called by pysat, and not by the end user directly.

    Parameters
    -----------
    inst : (pysat.Instrument)
        Instrument class object, whose attribute clean_level is used to return
        the desired level of data selectivity.

    Note
    ----
        Supports 'clean', 'dusty', 'dirty', 'none'

    """

    if inst.tag.find('los') >= 0:
        # dealing with LOS winds
        var = 'Line_of_Sight_Wind'
        if 'Wind_Quality' in inst.variables:
            qual_flag = 'Quality_Flags'
        else:
            qual_flag = 'ICON_L21_Quality_Flags'
            var = 'ICON_L21_' + var
        if inst.clean_level in ['clean', 'dusty']:
            # find location with any of the flags set
            idx, idy, = np.where(inst[qual_flag].any(axis=2))
            inst[idx, idy, var] = np.nan
        else:
            # dirty and worse lets everything through
            pass

    elif inst.tag.find('vector') >= 0:
        # vector winds area
        vars = ['Zonal_Wind', 'Meridional_Wind']
        if 'Wind_Quality' in inst.variables:
            qual_flag = 'Wind_Quality'
        else:
            qual_flag = 'ICON_L22_Wind_Quality'
            vars = ['ICON_L22_' + x for x in vars]
        if inst.clean_level == 'clean':
            for var in vars:
                inst[var] = inst[var].where(inst[qual_flag] == 1.0)
        elif inst.clean_level == 'dusty':
            for var in vars:
                inst[var] = inst[var].where(inst[qual_flag] >= 0.5)
        else:
            # dirty lets everything through
            pass

    elif inst.tag.find('temp') >= 0:
        # neutral temperatures
        var = 'Temperature'
        saa_flag = 'Quality_Flag_South_Atlantic_Anomaly'
        cal_flag = 'Quality_Flag_Bad_Calibration'
        if saa_flag not in inst.variables:
            saa_flag = '_'.join(('ICON_L1_MIGHTI', inst.sat_id.upper(),
                                 saa_flag))
            cal_flag = '_'.join(('ICON_L1_MIGHTI', inst.sat_id.upper(),
                                 cal_flag))
            var = '_'.JOIN(('ICON_L23_MIGHTI', inst.sat_id.upper(), var))
        if inst.clean_level in ['clean', 'dusty']:
            # SAA
            idx, = np.where(inst[saa_flag] > 0)
            inst[:, idx, var] = np.nan
            # Calibration file
            idx, = np.where(inst[cal_flag] > 0)
            inst[:, idx, var] = np.nan
        else:
            # dirty and worse lets everything through
            pass

    return
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""Test parsing of full VCF record lines
"""

import io
import sys

from vcfpy import parser

__author__ = "Manuel Holtgrewe <manuel.holtgrewe@bihealth.de>"

MEDIUM_HEADER = """
##fileformat=VCFv4.3
##fileDate=20090805
##source=myImputationProgramV3.1
##reference=file:///seq/references/1000GenomesPilot-NCBI36.fasta
##contig=<ID=20,length=62435964,assembly=B36,md5=f126cdf8a6e0c7f379d618ff66beb2da,species="Homo sapiens",taxonomy=x>
##phasing=partial
##INFO=<ID=NS,Number=1,Type=Integer,Description="Number of Samples With Data">
##INFO=<ID=DP,Number=1,Type=Integer,Description="Total Depth">
##INFO=<ID=AF,Number=A,Type=Float,Description="Allele Frequency">
##INFO=<ID=AA,Number=1,Type=String,Description="Ancestral Allele">
##INFO=<ID=DB,Number=0,Type=Flag,Description="dbSNP membership, build 129">
##INFO=<ID=H2,Number=0,Type=Flag,Description="HapMap2 membership">
##INFO=<ID=ANNO,Number=.,Type=String,Description="Additional annotation">
##FILTER=<ID=q10,Description="Quality below 10">
##FILTER=<ID=s50,Description="Less than 50% of samples have data">
##FILTER=<ID=FOO,Description="A test filter">
##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">
##FORMAT=<ID=GQ,Number=1,Type=Integer,Description="Genotype Quality">
##FORMAT=<ID=DP,Number=1,Type=Integer,Description="Read Depth">
##FORMAT=<ID=HQ,Number=2,Type=Integer,Description="Haplotype Quality">
##FORMAT=<ID=FT,Number=1,Type=String,Description="Call-wise filters">
#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tNA00001\tNA00002\tNA00003
""".lstrip()

def vcf_parser(lines):
    return parser.Parser(io.StringIO(MEDIUM_HEADER + lines), "<builtin>")

def test_parse_minimal_record():
    # Setup parser with stock header and lines to parse
    LINES = "20\t1\t.\tC\tG\t.\t.\t.\tGT\t0/1\t0/2\t.\n"
    p = vcf_parser(LINES)
    p.parse_header()
    # Perform the actual test
    if sys.version_info < (3, 6):
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, [], OrderedDict(), ['GT'], "
            "[Call('NA00001', OrderedDict([('GT', '0/1')])), Call('NA00002', OrderedDict([('GT', '0/2')])), "
            "Call('NA00003', OrderedDict([('GT', None)]))])"
        )
    else:
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, [], {}, ['GT'], "
            "[Call('NA00001', {'GT': '0/1'}), Call('NA00002', {'GT': '0/2'}), Call('NA00003', {'GT': None})])"
        )
    RESULT = p.parse_next_record()
    assert str(RESULT) == EXPECTED

def test_parse_record_with_info():
    # Setup parser with stock header and lines to parse
    LINES = "20\t1\t.\tC\tG\t.\t.\tAA=G\tGT\t0/1\t0/1\t.\n"
    p = vcf_parser(LINES)
    p.parse_header()
    # Perform the actual test
    if sys.version_info < (3, 6):
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, [], OrderedDict([('AA', 'G')]), "
            "['GT'], [Call('NA00001', OrderedDict([('GT', '0/1')])), Call('NA00002', OrderedDict([('GT', '0/1')])),"
            " Call('NA00003', OrderedDict([('GT', None)]))])"
        )
    else:
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, [], {'AA': 'G'}, "
            "['GT'], [Call('NA00001', {'GT': '0/1'}), Call('NA00002', {'GT': '0/1'}), Call('NA00003', {'GT': None})])"
        )
    RESULT = p.parse_next_record()
    assert str(RESULT) == EXPECTED

def test_parse_record_with_escaping():
    # Setup parser with stock header and lines to parse
    LINES = (
        "20\t100\t.\tC\tG\t.\t.\tANNO=Here%2Care%25some chars,"
        "%2525\tGT:FT\t0/1:FOO\t0/0:.\t1/1:.\n"
    )
    p = vcf_parser(LINES)
    p.parse_header()
    # Perform the actual test
    if sys.version_info < (3, 6):
        EXPECTED = (
            "Record('20', 100, [], 'C', [Substitution(type_='SNV', value='G')], None, [], "
            "OrderedDict([('ANNO', ['Here,are%some chars', '%25'])]), ['GT', 'FT'], "
            "[Call('NA00001', OrderedDict([('GT', '0/1'), ('FT', ['FOO'])])),"
            " Call('NA00002', OrderedDict([('GT', '0/0'), ('FT', [])])),"
            " Call('NA00003', OrderedDict([('GT', '1/1'), ('FT', [])]))])"
        )
    else:
        EXPECTED = (
            "Record('20', 100, [], 'C', [Substitution(type_='SNV', value='G')], None, [], "
            "{'ANNO': ['Here,are%some chars', '%25']}, ['GT', 'FT'], "
            "[Call('NA00001', {'GT': '0/1', 'FT': ['FOO']}), Call('NA00002', {'GT': '0/0', 'FT': []}), "
            "Call('NA00003', {'GT': '1/1', 'FT': []})])"
        )
    RESULT = p.parse_next_record()
    assert str(RESULT) == EXPECTED

def test_parse_record_with_filter_warning():
    # Setup parser with stock header and lines to parse
    LINES = "20\t1\t.\tC\tG\t.\tREX\t.\tGT:FT\t0/1:.\t0/2:BAR\t.:.\n"
    p = vcf_parser(LINES)
    p.parse_header()
    # Perform the actual test
    if sys.version_info < (3, 6):
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, ['REX'], OrderedDict(), "
            "['GT', 'FT'], [Call('NA00001', OrderedDict([('GT', '0/1'), ('FT', [])])),"
            " Call('NA00002', OrderedDict([('GT', '0/2'), ('FT', ['BAR'])])),"
            " Call('NA00003', OrderedDict([('GT', None), ('FT', [])]))])"
        )
    else:
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, ['REX'], {}, "
            "['GT', 'FT'], [Call('NA00001', {'GT': '0/1', 'FT': []}),"
            " Call('NA00002', {'GT': '0/2', 'FT': ['BAR']}), Call('NA00003', {'GT': None, 'FT': []})])"
        )
    RESULT = p.parse_next_record()
    assert str(RESULT) == EXPECTED
<commit_msg>Another test for genotype filters.<commit_after># -*- coding: utf-8 -*-
"""Test parsing of full VCF record lines
"""

import io
import sys

from vcfpy import parser

__author__ = "Manuel Holtgrewe <manuel.holtgrewe@bihealth.de>"

MEDIUM_HEADER = """
##fileformat=VCFv4.3
##fileDate=20090805
##source=myImputationProgramV3.1
##reference=file:///seq/references/1000GenomesPilot-NCBI36.fasta
##contig=<ID=20,length=62435964,assembly=B36,md5=f126cdf8a6e0c7f379d618ff66beb2da,species="Homo sapiens",taxonomy=x>
##phasing=partial
##INFO=<ID=NS,Number=1,Type=Integer,Description="Number of Samples With Data">
##INFO=<ID=DP,Number=1,Type=Integer,Description="Total Depth">
##INFO=<ID=AF,Number=A,Type=Float,Description="Allele Frequency">
##INFO=<ID=AA,Number=1,Type=String,Description="Ancestral Allele">
##INFO=<ID=DB,Number=0,Type=Flag,Description="dbSNP membership, build 129">
##INFO=<ID=H2,Number=0,Type=Flag,Description="HapMap2 membership">
##INFO=<ID=ANNO,Number=.,Type=String,Description="Additional annotation">
##FILTER=<ID=q10,Description="Quality below 10">
##FILTER=<ID=s50,Description="Less than 50% of samples have data">
##FILTER=<ID=FOO,Description="A test filter">
##FILTER=<ID=BAZ,Description="A second test filter">
##FORMAT=<ID=GT,Number=1,Type=String,Description="Genotype">
##FORMAT=<ID=GQ,Number=1,Type=Integer,Description="Genotype Quality">
##FORMAT=<ID=DP,Number=1,Type=Integer,Description="Read Depth">
##FORMAT=<ID=HQ,Number=2,Type=Integer,Description="Haplotype Quality">
##FORMAT=<ID=FT,Number=1,Type=String,Description="Call-wise filters">
#CHROM\tPOS\tID\tREF\tALT\tQUAL\tFILTER\tINFO\tFORMAT\tNA00001\tNA00002\tNA00003
""".lstrip()

def vcf_parser(lines):
    return parser.Parser(io.StringIO(MEDIUM_HEADER + lines), "<builtin>")

def test_parse_minimal_record():
    # Setup parser with stock header and lines to parse
    LINES = "20\t1\t.\tC\tG\t.\t.\t.\tGT\t0/1\t0/2\t.\n"
    p = vcf_parser(LINES)
    p.parse_header()
    # Perform the actual test
    if sys.version_info < (3, 6):
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, [], OrderedDict(), ['GT'], "
            "[Call('NA00001', OrderedDict([('GT', '0/1')])), Call('NA00002', OrderedDict([('GT', '0/2')])), "
            "Call('NA00003', OrderedDict([('GT', None)]))])"
        )
    else:
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, [], {}, ['GT'], "
            "[Call('NA00001', {'GT': '0/1'}), Call('NA00002', {'GT': '0/2'}), Call('NA00003', {'GT': None})])"
        )
    RESULT = p.parse_next_record()
    assert str(RESULT) == EXPECTED

def test_parse_record_with_info():
    # Setup parser with stock header and lines to parse
    LINES = "20\t1\t.\tC\tG\t.\t.\tAA=G\tGT\t0/1\t0/1\t.\n"
    p = vcf_parser(LINES)
    p.parse_header()
    # Perform the actual test
    if sys.version_info < (3, 6):
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, [], OrderedDict([('AA', 'G')]), "
            "['GT'], [Call('NA00001', OrderedDict([('GT', '0/1')])), Call('NA00002', OrderedDict([('GT', '0/1')])),"
            " Call('NA00003', OrderedDict([('GT', None)]))])"
        )
    else:
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, [], {'AA': 'G'}, "
            "['GT'], [Call('NA00001', {'GT': '0/1'}), Call('NA00002', {'GT': '0/1'}), Call('NA00003', {'GT': None})])"
        )
    RESULT = p.parse_next_record()
    assert str(RESULT) == EXPECTED

def test_parse_record_with_escaping():
    # Setup parser with stock header and lines to parse
    LINES = (
        "20\t100\t.\tC\tG\t.\t.\tANNO=Here%2Care%25some chars,"
        "%2525\tGT:FT\t0/1:FOO\t0/0:.\t1/1:.\n"
    )
    p = vcf_parser(LINES)
    p.parse_header()
    # Perform the actual test
    if sys.version_info < (3, 6):
        EXPECTED = (
            "Record('20', 100, [], 'C', [Substitution(type_='SNV', value='G')], None, [], "
            "OrderedDict([('ANNO', ['Here,are%some chars', '%25'])]), ['GT', 'FT'], "
            "[Call('NA00001', OrderedDict([('GT', '0/1'), ('FT', ['FOO'])])),"
            " Call('NA00002', OrderedDict([('GT', '0/0'), ('FT', [])])),"
            " Call('NA00003', OrderedDict([('GT', '1/1'), ('FT', [])]))])"
        )
    else:
        EXPECTED = (
            "Record('20', 100, [], 'C', [Substitution(type_='SNV', value='G')], None, [], "
            "{'ANNO': ['Here,are%some chars', '%25']}, ['GT', 'FT'], "
            "[Call('NA00001', {'GT': '0/1', 'FT': ['FOO']}), Call('NA00002', {'GT': '0/0', 'FT': []}), "
            "Call('NA00003', {'GT': '1/1', 'FT': []})])"
        )
    RESULT = p.parse_next_record()
    assert str(RESULT) == EXPECTED

def test_parse_record_with_filter_warning():
    # Setup parser with stock header and lines to parse
    LINES = "20\t1\t.\tC\tG\t.\tREX\t.\tGT:FT\t0/1:.\t0/2:BAR\t.:.\n"
    p = vcf_parser(LINES)
    p.parse_header()
    # Perform the actual test
    if sys.version_info < (3, 6):
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, ['REX'], OrderedDict(), "
            "['GT', 'FT'], [Call('NA00001', OrderedDict([('GT', '0/1'), ('FT', [])])),"
            " Call('NA00002', OrderedDict([('GT', '0/2'), ('FT', ['BAR'])])),"
            " Call('NA00003', OrderedDict([('GT', None), ('FT', [])]))])"
        )
    else:
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, ['REX'], {}, "
            "['GT', 'FT'], [Call('NA00001', {'GT': '0/1', 'FT': []}),"
            " Call('NA00002', {'GT': '0/2', 'FT': ['BAR']}), Call('NA00003', {'GT': None, 'FT': []})])"
        )
    RESULT = p.parse_next_record()
    assert str(RESULT) == EXPECTED

def test_parse_record_with_filter_no_warning():
    # Setup parser with stock header and lines to parse
    LINES = "20\t1\t.\tC\tG\t.\tREX\t.\tGT:FT\t0/1:.\t0/2:FOO;BAZ\t.:.\n"
    p = vcf_parser(LINES)
    p.parse_header()
    # Perform the actual test
    if sys.version_info < (3, 6):
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, ['REX'], OrderedDict(), "
            "['GT', 'FT'], [Call('NA00001', OrderedDict([('GT', '0/1'), ('FT', [])])),"
            " Call('NA00002', OrderedDict([('GT', '0/2'), ('FT', ['FOO', 'BAZ'])])),"
            " Call('NA00003', OrderedDict([('GT', None), ('FT', [])]))])"
        )
    else:
        EXPECTED = (
            "Record('20', 1, [], 'C', [Substitution(type_='SNV', value='G')], None, ['REX'], {}, "
            "['GT', 'FT'], [Call('NA00001', {'GT': '0/1', 'FT': []}),"
            " Call('NA00002', {'GT': '0/2', 'FT': ['FOO', 'BAZ']}), Call('NA00003', {'GT': None, 'FT': []})])"
        )
    RESULT = p.parse_next_record()
    assert str(RESULT) == EXPECTED
<|endoftext|>


<commit_before>#! /usr/bin/env python
# -*- coding: utf-8 -*-
import subprocess
from time import sleep
from random import randint, choice
from test import Test, TestCase

SLEEP_INTERVAL = (1, 25)  # Interval for random sleep periods between tests.
OONI_BINARY = "/home/analyser/ooni-probe/bin/ooniprobe"  # Path to ooniprobe binary.
PROBE_TEST = "siteprobe.py"
DNS_TEST = "dnscompare.py"
TCP_TEST = "tcpconnect.py"
PING_TEST = "ping.py"
TRACEROUTE_TEST = "traceroute.py"
TOR_MIRRORS = ["https://www.unicorncloud.org/public/torproject.org/", "https://mirror.ml/tor/", "https://www.oignon.net/", "https://tor.hackthissite.org/"]
TOR_DIRECTORY_AUTHORITIES = [
                             { "86.59.21.38": 80 },
                             { "86.59.21.38": 443 },
                             { "76.73.17.194": 9030 },
                             { "76.73.17.194": 9090 },
                             { "171.25.193.9": 443 },
                             { "171.25.193.9": 80 },
                             { "dannenberg.ccc.de": 80 },
                             { "193.23.244.244": 443 },
                             { "208.83.223.34": 443 },
                             { "208.83.223.34": 80 },
                             { "128.31.0.34": 9131 },
                             { "128.31.0.34": 9101 },
                             { "194.109.206.212": 80 },
                             { "194.109.206.212": 443 },
                             { "212.112.245.170": 80 },
                             { "212.112.245.170": 443 },
                             { "154.35.32.5": 80 },
                             { "154.35.32.5": 443 }
                            ]

class Analyser():

    def runTest(self, test, argv):
        if len(argv) > 1:
            args = (OONI_BINARY, "-n", test, argv[0], argv[1])
        else:
            args = (OONI_BINARY, "-n", test)
        print "Running test %s" % test
        popen = subprocess.Popen(args, stdout=subprocess.PIPE)
        popen.wait()
        output = popen.stdout.read()
        return output

    def probeSite(self, args):
        return self.runTest(PROBE_TEST, args)

    def compareDNS(self, args):
        return self.runTest(DNS_TEST, args)

    def testTCPConnection(self, args):
        return self.runTest(TCP_TEST, [])

    def ping(self, argv):
        return self.runTest(PING_TEST, argv)

    def traceroute(self, args):
        return self.runTest(TRACEROUTE_TEST, args)

    def siteReachable(self, url):
        return "Site reachable" in self.probeSite(["-u", url])

def probeTorSite():
    if not analyser.siteReachable("https://www.torproject.org"):
        print "Tor site not reachable - running tests"
        tests = [
                    {  analyser.testTCPConnection: ["-t", "bridges.torproject.org"] },
                    {  analyser.compareDNS: [] },
                    {  analyser.probeSite: ["-u", choice(TOR_MIRRORS)] },
                    {  analyser.testTCPConnection: [] }
                ]
        runTests(tests)
    else:
        print "Tor site reachable"

def probeDirectoryAuthorities():
    for server in TOR_DIRECTORY_AUTHORITIES:
        host = server.keys()[0]
        port = server[host]
        url = makeURL(host, port)

        if analyser.siteReachable(url):
            print "Consensus successfully fetched"
            return

    print "Failed to fetch consensus - running tests"
    tests = [
                { analyser.ping: ["-f", "directory_authorities.txt"] },
                { analyser.traceroute: ["-f", "directory_authorities.txt"] }
            ]

    runTests(tests)

def makeURL(host, port):
    if port == 80:
         url = "http://%s/tor/status-vote/current/consensus.z" % (host)
    elif port == 443:
         url = "https://%s/tor/status-vote/current/consensus.z" % (host)
    else:
         url = "http://%s:%s/tor/status-vote/current/consensus.z" % (host, str(port))
    return url

def runTests(tests):
    for x in range(len(tests)):
        test, args = randPop(tests)
        test(args)
        sleep(randint(SLEEP_INTERVAL[0], SLEEP_INTERVAL[1]))

def randPop(tests):
    t = choice(tests)
    test = t.keys()[0]
    args = t[test]
    tests.remove(t)
    return test, args

if __name__ == "__main__":
    analyser = Analyser()
    tests = TestCase()
    tests.addTests( [Test(TCP_TEST, [], OONI_BINARY)])
    tests.addTests( [Test(PROBE_TEST, [], OONI_BINARY)])
    tests.run()
    tests.verifyResults()
    #probeTorSite()
    #probeDirectoryAuthorities()
<commit_msg>Rewrite module to use the new Test-classes<commit_after>#! /usr/bin/env python
# -*- coding: utf-8 -*-
import subprocess
from time import sleep
from random import randint, choice
from test import *

SLEEP_INTERVAL = (1, 25)  # Interval for random sleep periods between tests.
OONI_BINARY = "ooniprobe"
PROBE_TEST = "siteprobe.py"
DNS_TEST = "dnscompare.py"
TCP_TEST = "tcpconnect.py"
PING_TEST = "ping.py"
TRACEROUTE_TEST = "traceroute.py"
TOR_SITE_URL = "https://www.torproject.org"
TOR_BRIDGES_URL =  "bridges.torproject.org"
TOR_MIRRORS = ["https://www.unicorncloud.org/public/torproject.org/", "https://mirror.ml/tor/", "https://www.oignon.net/", "https://tor.hackthissite.org/"]
TOR_DIRECTORY_AUTHORITIES = [
                             { "86.59.21.38": 80 },
                             { "86.59.21.38": 443 },
                             { "76.73.17.194": 9030 },
                             { "76.73.17.194": 9090 },
                             { "171.25.193.9": 443 },
                             { "171.25.193.9": 80 },
                             { "dannenberg.ccc.de": 80 },
                             { "193.23.244.244": 443 },
                             { "208.83.223.34": 443 },
                             { "208.83.223.34": 80 },
                             { "128.31.0.34": 9131 },
                             { "128.31.0.34": 9101 },
                             { "194.109.206.212": 80 },
                             { "194.109.206.212": 443 },
                             { "212.112.245.170": 80 },
                             { "212.112.245.170": 443 },
                             { "154.35.32.5": 80 },
                             { "154.35.32.5": 443 }
                            ]

def probeTorSite():
    if not siteReachable(TOR_SITE_URL):
        print "Tor site not reachable - running tests"
        tests = TestCase(tests = [ test for test in [
                                                    TCPTest(target = TOR_BRIDGES_URL),
                                                    DNSTest(),
                                                    SiteProbe(target=choice(TOR_MIRRORS)),
                                                    TCPTest()
                                                    ] ] )
        tests.run()
        tests.printResults()
    else:
        print "Tor site reachable"

def probeDirectoryAuthorities():
    for server in TOR_DIRECTORY_AUTHORITIES:
        host = server.keys()[0]
        port = server[host]
        url = makeURL(host, port)

        if siteReachable(url):
            print "Consensus successfully fetched"
            return

    print "Failed to fetch consensus - running tests"
    tests = TestCase(tests = [ test for test in [
                                                PingTest(),
                                                Traceroute()
                                                ] ] )
    tests.run()
    tests.printResults()

def makeURL(host, port):
    if port == 80:
         url = "http://%s/tor/status-vote/current/consensus.z" % (host)
    elif port == 443:
         url = "https://%s/tor/status-vote/current/consensus.z" % (host)
    else:
         url = "http://%s:%s/tor/status-vote/current/consensus.z" % (host, str(port))
    return url

def siteReachable(url):
    probe = SiteProbe(target=url)
    probe.run()
    return probe.status == "OK"

if __name__ == "__main__":
    probeTorSite()
    probeDirectoryAuthorities()

<|endoftext|>


<commit_before>from django.conf.urls import patterns, include, url
from django.contrib import admin
from django.conf import settings
from images.views import ImageUploaderView
from images.views import ImageView
from django.views.decorators.csrf import csrf_exempt

urlpatterns = patterns('',
    #url(r'^$', csrf_exempt(ImageUploaderView.as_view())),
    url(r'^$', ImageUploaderView.as_view()),
    url(r'^admin/', include(admin.site.urls)),
    url(r'^image/(?P<image_identifier>[^/]+)/$', ImageView.as_view(), name='image'),
)

# apply authentication middleware on everything but the image urls
from urlmiddleware.conf import middleware, mpatterns
middlewarepatterns = mpatterns('',
    middleware(r'^(?!image)/', 'django.contrib.sessions.middleware.SessionMiddleware'),
    middleware(r'^(?!image)/', 'django.contrib.auth.middleware.AuthenticationMiddleware'),
    middleware(r'^(?!image)/', 'django.contrib.auth.middleware.SessionAuthenticationMiddleware'),
)

# Serving Locally Stored Media Files
# ----------------------------------
# media files are served from the media directory via dj_static. The configuration
# can be found in frame/wsgi.py
<commit_msg>Fixed routing regular expression issue<commit_after>from django.conf.urls import patterns, include, url
from django.contrib import admin
from django.conf import settings
from images.views import ImageUploaderView
from images.views import ImageView
from django.views.decorators.csrf import csrf_exempt

urlpatterns = patterns('',
    #url(r'^$', csrf_exempt(ImageUploaderView.as_view())),
    url(r'^$', ImageUploaderView.as_view()),
    url(r'^admin/', include(admin.site.urls)),
    url(r'^image/(?P<image_identifier>[^/]+)/$', ImageView.as_view(), name='image'),
)

# apply authentication middleware on everything but the image urls
from urlmiddleware.conf import middleware, mpatterns
middlewarepatterns = mpatterns('',
    middleware(r'^(?!image)', 'django.contrib.sessions.middleware.SessionMiddleware'),
    middleware(r'^(?!image)', 'django.contrib.auth.middleware.AuthenticationMiddleware'),
    middleware(r'^(?!image)', 'django.contrib.auth.middleware.SessionAuthenticationMiddleware'),
)

# Serving Locally Stored Media Files
# ----------------------------------
# media files are served from the media directory via dj_static. The configuration
# can be found in frame/wsgi.py
<|endoftext|>


<commit_before>from tempfile import NamedTemporaryFile
import openpyxl

class JSONReaderError(Exception):
    pass

class HeaderValueError(Exception):
    pass

class IteratorJSONReader(object):
    """
    >>> def normalize(it):
    ...     r = []
    ...     for row in IteratorJSONReader(it):
    ...         r.append(sorted(row.items()))
    ...     return r
    >>> normalize([])
    []
    >>> normalize([['A', 'B', 'C'], ['1', '2', '3']])
    [[('A', '1'), ('B', '2'), ('C', '3')]]
    >>> normalize([['A', 'data: key', 'user 1', 'user 2', 'is-ok?'],
    ...     ['1', '2', '3', '4', 'yes']])
    [[('A', '1'), ('data', {'key': '2'}), ('is-ok', True), ('user', ['3', '4'])]]
    """

    def __init__(self, rows):
        # you can only call __iter__ once
        self._rows = iter(rows)
        try:
            self.headers = list(self._rows.next())
        except StopIteration:
            self.headers = []
        self.fieldnames = self.get_fieldnames()

    def row_to_json(self, row):
        obj = {}
        for value, header in zip(row, self.headers):
            self.set_field_value(obj, header, value)
        return obj

    def __iter__(self):
        try:
            for row in self._rows:
                yield self.row_to_json(row)
        finally:
            del self._rows

    def get_fieldnames(self):
        obj = {}
        for field, value in zip(self.headers, [''] * len(self.headers)):
            if not isinstance(field, basestring):
                raise HeaderValueError(u'Field %s is not a string.' % field)
            self.set_field_value(obj, field, value)
        return obj.keys()

    @classmethod
    def set_field_value(cls, obj, field, value):
        if isinstance(value, basestring):
            value = value.strip()
        # try dict
        try:
            field, subfield = field.split(':')
        except Exception:
            pass
        else:
            field = field.strip()

            if field not in obj:
                obj[field] = {}

            cls.set_field_value(obj[field], subfield, value)
            return

        # try list
        try:
            field, _ = field.split()
        except Exception:
            pass
        else:
            dud = {}
            cls.set_field_value(dud, field, value)
            (field, value), = dud.items()

            if field not in obj:
                obj[field] = []
            elif not isinstance(obj[field], list):
                obj[field] = [obj[field]]
            if value not in (None, ''):
                obj[field].append(value)
            return

        # else flat

        # try boolean
        try:
            field, nothing = field.split('?')
            assert(nothing.strip() == '')
        except Exception:
            pass
        else:
            try:
                value = {
                    'yes': True,
                    'true': True,
                    'no': False,
                    'false': False,
                    '': False,
                    None: False,
                }[value.lower() if hasattr(value, 'lower') else value]
            except KeyError:
                raise JSONReaderError(
                    'Values for field %s must be "yes" or "no", not "%s"' % (
                        field, value)
                )

        # set for any flat type
        field = field.strip()
        if field in obj:
            raise JSONReaderError(
                'You have a repeat field: %s' % field
            )

        obj[field] = value

class WorksheetNotFound(Exception):

    def __init__(self, title):
        self.title = title
        super(WorksheetNotFound, self).__init__()

class WorksheetJSONReader(IteratorJSONReader):

    def __init__(self, worksheet, title=None):
        width = 0
        self.title = title
        self.worksheet = worksheet
        try:
            header_row = self.worksheet.iter_rows().next()
        except StopIteration:
            header_row = []
        for cell in header_row:
            if cell.value is None:
                break
            else:
                width += 1
        self.worksheet.calculate_dimension(force=True)

        def iterator():
            def _convert_float(value):
                """
                excel doesn't distinguish between 1 and 1.0
                if it can be an integer assume it is
                """
                if isinstance(value, float) and int(value) == value:
                    return int(value)
                else:
                    return value or ''
            for row in self.worksheet.iter_rows():
                cell_values = [
                    _convert_float(cell.value)
                    for cell in row[:width]
                ]
                if not any(cell_values):
                    break
                yield cell_values
        super(WorksheetJSONReader, self).__init__(iterator())

class WorkbookJSONReader(object):

    def __init__(self, f):
        if isinstance(f, basestring):
            filename = f
        elif not isinstance(f, file):
            tmp = NamedTemporaryFile(mode='wb', suffix='.xlsx', delete=False)
            filename = tmp.name
            tmp.write(f.read())
            tmp.close()
        else:
            filename = f

        self.wb = openpyxl.load_workbook(filename, use_iterators=True)
        self.worksheets_by_title = {}
        self.worksheets = []

        for worksheet in self.wb.worksheets:
            ws = WorksheetJSONReader(worksheet, title=worksheet.title)
            self.worksheets_by_title[worksheet.title] = ws
            self.worksheets.append(ws)

    def get_worksheet(self, title=None, index=None):
        if title is not None and index is not None:
            raise TypeError("Can only get worksheet by title *or* index")
        if title:
            try:
                return self.worksheets_by_title[title]
            except KeyError:
                raise WorksheetNotFound(title=title)
        elif index:
            try:
                return self.worksheets[index]
            except IndexError:
                raise WorksheetNotFound(title=index)
        else:
            try:
                return self.worksheets[0]
            except IndexError:
                raise WorksheetNotFound(title=0)

    def work_book_headers_as_tuples(self):
        """
        Returns raw sheet headers in following format

          (("employee", ("id", "name", "gender")),
           ("building", ("id", "name", "address")))
        """
        all_sheet_headers = []
        for sheet in self.worksheets:
            all_sheet_headers.append(
                (sheet.title, tuple(sheet.headers))
            )
        return tuple(all_sheet_headers)

    def work_book_data_as_tuples(self):
        """
        Note: This is useful to get xlsx file's data into easy readible format
        Exists only to migrate tests using xlsx files to use following format

        Returns raw sheet data in following format

        (("employee", (("1", "cory", "m"),
                       ("2", "christian", "m"),
                       ("3", "amelia", "f"))),
         ("building", (("1", "dimagi", "585 mass ave."),
                       ("2", "old dimagi", "529 main st."))))
        """
        all_sheet_data = []
        for sheet in self.worksheets:
            current_sheet_data = []
            for row in sheet:
                values = [row.get(header) for header in sheet.headers]
                current_sheet_data.append(tuple(values))
            all_sheet_data.append(
                (sheet.title, tuple(current_sheet_data))
            )
        return tuple(all_sheet_data)

def flatten_json_to_path(obj, path=()):
    if isinstance(obj, dict):
        for key, value in obj.items():
            for item in flatten_json_to_path(value, path + (key,)):
                yield item
    elif isinstance(obj, list):
        for key, value in enumerate(obj):
            for item in flatten_json_to_path(value, path + (key,)):
                yield item
    else:
        yield (path, obj)

def format_header(path, value):
    # pretty sure making a string-builder would be slower than concatenation
    s = path[0]
    for p in path[1:]:
        if isinstance(p, basestring):
            s += ': %s' % p
        elif isinstance(p, int):
            s += ' %s' % (p + 1)
    if isinstance(value, bool):
        s += '?'
        value = 'yes' if value else 'no'
    return s, value

def flatten_json(obj):
    for key, value in flatten_json_to_path(obj):
        yield format_header(key, value)

def json_to_headers(obj):
    return [key for key, value in sorted(flatten_json(obj), key=lambda t: alphanumeric_sort_key(t[0]))]

def alphanumeric_sort_key(key):
    """
    Sort the given iterable in the way that humans expect.
    Thanks to http://stackoverflow.com/a/2669120/240553
    """
    import re
    convert = lambda text: int(text) if text.isdigit() else text
    return [convert(c) for c in re.split('([0-9]+)', key)]
<commit_msg>extract data from excel sheets, not formulas<commit_after>from tempfile import NamedTemporaryFile
import openpyxl

class JSONReaderError(Exception):
    pass

class HeaderValueError(Exception):
    pass

class IteratorJSONReader(object):
    """
    >>> def normalize(it):
    ...     r = []
    ...     for row in IteratorJSONReader(it):
    ...         r.append(sorted(row.items()))
    ...     return r
    >>> normalize([])
    []
    >>> normalize([['A', 'B', 'C'], ['1', '2', '3']])
    [[('A', '1'), ('B', '2'), ('C', '3')]]
    >>> normalize([['A', 'data: key', 'user 1', 'user 2', 'is-ok?'],
    ...     ['1', '2', '3', '4', 'yes']])
    [[('A', '1'), ('data', {'key': '2'}), ('is-ok', True), ('user', ['3', '4'])]]
    """

    def __init__(self, rows):
        # you can only call __iter__ once
        self._rows = iter(rows)
        try:
            self.headers = list(self._rows.next())
        except StopIteration:
            self.headers = []
        self.fieldnames = self.get_fieldnames()

    def row_to_json(self, row):
        obj = {}
        for value, header in zip(row, self.headers):
            self.set_field_value(obj, header, value)
        return obj

    def __iter__(self):
        try:
            for row in self._rows:
                yield self.row_to_json(row)
        finally:
            del self._rows

    def get_fieldnames(self):
        obj = {}
        for field, value in zip(self.headers, [''] * len(self.headers)):
            if not isinstance(field, basestring):
                raise HeaderValueError(u'Field %s is not a string.' % field)
            self.set_field_value(obj, field, value)
        return obj.keys()

    @classmethod
    def set_field_value(cls, obj, field, value):
        if isinstance(value, basestring):
            value = value.strip()
        # try dict
        try:
            field, subfield = field.split(':')
        except Exception:
            pass
        else:
            field = field.strip()

            if field not in obj:
                obj[field] = {}

            cls.set_field_value(obj[field], subfield, value)
            return

        # try list
        try:
            field, _ = field.split()
        except Exception:
            pass
        else:
            dud = {}
            cls.set_field_value(dud, field, value)
            (field, value), = dud.items()

            if field not in obj:
                obj[field] = []
            elif not isinstance(obj[field], list):
                obj[field] = [obj[field]]
            if value not in (None, ''):
                obj[field].append(value)
            return

        # else flat

        # try boolean
        try:
            field, nothing = field.split('?')
            assert(nothing.strip() == '')
        except Exception:
            pass
        else:
            try:
                value = {
                    'yes': True,
                    'true': True,
                    'no': False,
                    'false': False,
                    '': False,
                    None: False,
                }[value.lower() if hasattr(value, 'lower') else value]
            except KeyError:
                raise JSONReaderError(
                    'Values for field %s must be "yes" or "no", not "%s"' % (
                        field, value)
                )

        # set for any flat type
        field = field.strip()
        if field in obj:
            raise JSONReaderError(
                'You have a repeat field: %s' % field
            )

        obj[field] = value

class WorksheetNotFound(Exception):

    def __init__(self, title):
        self.title = title
        super(WorksheetNotFound, self).__init__()

class WorksheetJSONReader(IteratorJSONReader):

    def __init__(self, worksheet, title=None):
        width = 0
        self.title = title
        self.worksheet = worksheet
        try:
            header_row = self.worksheet.iter_rows().next()
        except StopIteration:
            header_row = []
        for cell in header_row:
            if cell.value is None:
                break
            else:
                width += 1
        self.worksheet.calculate_dimension(force=True)

        def iterator():
            def _convert_float(value):
                """
                excel doesn't distinguish between 1 and 1.0
                if it can be an integer assume it is
                """
                if isinstance(value, float) and int(value) == value:
                    return int(value)
                else:
                    return value or ''
            for row in self.worksheet.iter_rows():
                cell_values = [
                    _convert_float(cell.value)
                    for cell in row[:width]
                ]
                if not any(cell_values):
                    break
                yield cell_values
        super(WorksheetJSONReader, self).__init__(iterator())

class WorkbookJSONReader(object):

    def __init__(self, f):
        if isinstance(f, basestring):
            filename = f
        elif not isinstance(f, file):
            tmp = NamedTemporaryFile(mode='wb', suffix='.xlsx', delete=False)
            filename = tmp.name
            tmp.write(f.read())
            tmp.close()
        else:
            filename = f

        self.wb = openpyxl.load_workbook(filename, use_iterators=True, data_only=True)
        self.worksheets_by_title = {}
        self.worksheets = []

        for worksheet in self.wb.worksheets:
            ws = WorksheetJSONReader(worksheet, title=worksheet.title)
            self.worksheets_by_title[worksheet.title] = ws
            self.worksheets.append(ws)

    def get_worksheet(self, title=None, index=None):
        if title is not None and index is not None:
            raise TypeError("Can only get worksheet by title *or* index")
        if title:
            try:
                return self.worksheets_by_title[title]
            except KeyError:
                raise WorksheetNotFound(title=title)
        elif index:
            try:
                return self.worksheets[index]
            except IndexError:
                raise WorksheetNotFound(title=index)
        else:
            try:
                return self.worksheets[0]
            except IndexError:
                raise WorksheetNotFound(title=0)

    def work_book_headers_as_tuples(self):
        """
        Returns raw sheet headers in following format

          (("employee", ("id", "name", "gender")),
           ("building", ("id", "name", "address")))
        """
        all_sheet_headers = []
        for sheet in self.worksheets:
            all_sheet_headers.append(
                (sheet.title, tuple(sheet.headers))
            )
        return tuple(all_sheet_headers)

    def work_book_data_as_tuples(self):
        """
        Note: This is useful to get xlsx file's data into easy readible format
        Exists only to migrate tests using xlsx files to use following format

        Returns raw sheet data in following format

        (("employee", (("1", "cory", "m"),
                       ("2", "christian", "m"),
                       ("3", "amelia", "f"))),
         ("building", (("1", "dimagi", "585 mass ave."),
                       ("2", "old dimagi", "529 main st."))))
        """
        all_sheet_data = []
        for sheet in self.worksheets:
            current_sheet_data = []
            for row in sheet:
                values = [row.get(header) for header in sheet.headers]
                current_sheet_data.append(tuple(values))
            all_sheet_data.append(
                (sheet.title, tuple(current_sheet_data))
            )
        return tuple(all_sheet_data)

def flatten_json_to_path(obj, path=()):
    if isinstance(obj, dict):
        for key, value in obj.items():
            for item in flatten_json_to_path(value, path + (key,)):
                yield item
    elif isinstance(obj, list):
        for key, value in enumerate(obj):
            for item in flatten_json_to_path(value, path + (key,)):
                yield item
    else:
        yield (path, obj)

def format_header(path, value):
    # pretty sure making a string-builder would be slower than concatenation
    s = path[0]
    for p in path[1:]:
        if isinstance(p, basestring):
            s += ': %s' % p
        elif isinstance(p, int):
            s += ' %s' % (p + 1)
    if isinstance(value, bool):
        s += '?'
        value = 'yes' if value else 'no'
    return s, value

def flatten_json(obj):
    for key, value in flatten_json_to_path(obj):
        yield format_header(key, value)

def json_to_headers(obj):
    return [key for key, value in sorted(flatten_json(obj), key=lambda t: alphanumeric_sort_key(t[0]))]

def alphanumeric_sort_key(key):
    """
    Sort the given iterable in the way that humans expect.
    Thanks to http://stackoverflow.com/a/2669120/240553
    """
    import re
    convert = lambda text: int(text) if text.isdigit() else text
    return [convert(c) for c in re.split('([0-9]+)', key)]
<|endoftext|>


<commit_before>from rest_framework import serializers

from api import fields
from share import models

class BaseShareSerializer(serializers.ModelSerializer):

    def __init__(self, *args, **kwargs):
        # super hates my additional kwarg
        sparse = kwargs.pop('sparse', False)
        super(BaseShareSerializer, self).__init__(*args, **kwargs)
        # remove version fields
        # easier than specifying excludes for every model serializer

        for k, v in tuple(self.fields.items()):
            if sparse or 'version' in k:
                # if they asked for sparse remove all fields but
                # the @id and @type

                # if the field has version, kill it too.
                self.fields.pop(k)

        # remove typedmodel nonsense
        if 'type' in self.fields.keys():
            self.fields.pop('type')

        # our id field is better
        if 'id' in self.fields.keys():
            self.fields.pop('id')

        # changes, be gone!
        if 'change' in self.fields.keys():
            self.fields.pop('change')

        # add fields with improper names
        self.fields.update({
            '@id': serializers.HyperlinkedIdentityField(
                # view-name: person-detail
                'api:{}-detail'.format(self.Meta.model._meta.model_name),
                lookup_field='pk'
            ),
            '@type': fields.TypeField(),
        })

class VenueSerializer(BaseShareSerializer):
    class Meta:
        model = models.Venue

class OrganizationSerializer(BaseShareSerializer):
    class Meta:
        model = models.Organization

class InstitutionSerializer(BaseShareSerializer):
    class Meta:
        model = models.Institution

class PersonEmailSerializer(BaseShareSerializer):
    class Meta:
        model = models.PersonEmail

class IdentifierSerializer(BaseShareSerializer):
    class Meta:
        model = models.Identifier

class PersonSerializer(BaseShareSerializer):
    # no emails on purpose
    identifiers = IdentifierSerializer(many=True)
    affiliations = OrganizationSerializer(sparse=True, many=True)
    class Meta:
        model = models.Person

class AffiliationSerializer(BaseShareSerializer):
    person = PersonSerializer(sparse=True)
    organization = OrganizationSerializer(sparse=True)

    class Meta:
        models = models.Affiliation

class ContributorSerializer(BaseShareSerializer):
    person = PersonSerializer()
    cited_name = serializers.ReadOnlyField(source='contributor.cited_name')
    order_cited = serializers.ReadOnlyField(source='contributor.order_cited')
    url = serializers.ReadOnlyField(source='contributor.url')
    class Meta:
        model = models.Contributor

class FunderSerializer(BaseShareSerializer):
    class Meta:
        model = models.Funder

class AwardSerializer(BaseShareSerializer):
    class Meta:
        model = models.Award

class TaxonomySerializer(BaseShareSerializer):
    class Meta:
        model = models.Taxonomy

class TagSerializer(BaseShareSerializer):
    taxonomy = TaxonomySerializer()
    class Meta:
        model = models.Tag

class AbstractCreativeWorkSerializer(BaseShareSerializer):
    tags = TagSerializer(many=True)
    contributors = ContributorSerializer(source='contributor_set', many=True)
    institutions = InstitutionSerializer(sparse=True, many=True)
    venues = VenueSerializer(sparse=True, many=True)
    awards = AwardSerializer(sparse=True, many=True)
    subject = TagSerializer(sparse=True)

class CreativeWorkSerializer(AbstractCreativeWorkSerializer):
    class Meta:
        model = models.CreativeWork

class PreprintSerializer(AbstractCreativeWorkSerializer):
    class Meta:
        model = models.Preprint

class ManuscriptSerializer(AbstractCreativeWorkSerializer):
    class Meta:
        model = models.Manuscript

class Link(BaseShareSerializer):
    class Meta:
        model = models.Link
<commit_msg>Add slightly more DRF way of excluding fields we do not want<commit_after>from rest_framework import serializers

from api import fields
from share import models

class BaseShareSerializer(serializers.ModelSerializer):

    def __init__(self, *args, **kwargs):
        # super hates my additional kwarg
        sparse = kwargs.pop('sparse', False)
        super(BaseShareSerializer, self).__init__(*args, **kwargs)
        # remove version fields
        # easier than specifying excludes for every model serializer

        for k, v in tuple(self.fields.items()):
            if sparse or 'version' in k:
                # if they asked for sparse remove all fields but
                # the @id and @type

                # if the field has version, kill it too.
                self.fields.pop(k)

        # remove typedmodel nonsense
        if 'type' in self.fields.keys():
            self.fields.pop('type')

        # add fields with improper names
        self.fields.update({
            '@id': serializers.HyperlinkedIdentityField(
                # view-name: person-detail
                'api:{}-detail'.format(self.Meta.model._meta.model_name),
                lookup_field='pk'
            ),
            '@type': fields.TypeField(),
        })
    class Meta:
        exclude = ('change', 'id')

class VenueSerializer(BaseShareSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Venue

class OrganizationSerializer(BaseShareSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Organization

class InstitutionSerializer(BaseShareSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Institution

class PersonEmailSerializer(BaseShareSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.PersonEmail

class IdentifierSerializer(BaseShareSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Identifier

class PersonSerializer(BaseShareSerializer):
    # no emails on purpose
    identifiers = IdentifierSerializer(many=True)
    affiliations = OrganizationSerializer(sparse=True, many=True)
    class Meta(BaseShareSerializer.Meta):
        model = models.Person

class AffiliationSerializer(BaseShareSerializer):
    person = PersonSerializer(sparse=True)
    organization = OrganizationSerializer(sparse=True)

    class Meta(BaseShareSerializer.Meta):
        models = models.Affiliation

class ContributorSerializer(BaseShareSerializer):
    person = PersonSerializer()
    cited_name = serializers.ReadOnlyField(source='contributor.cited_name')
    order_cited = serializers.ReadOnlyField(source='contributor.order_cited')
    url = serializers.ReadOnlyField(source='contributor.url')
    # TODO find a way to do this, or don't
    # creative_work = CreativeWorkSerializer(sparse=True)
    class Meta(BaseShareSerializer.Meta):
        model = models.Contributor

class FunderSerializer(BaseShareSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Funder

class AwardSerializer(BaseShareSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Award

class TaxonomySerializer(BaseShareSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Taxonomy

class TagSerializer(BaseShareSerializer):
    taxonomy = TaxonomySerializer()
    class Meta(BaseShareSerializer.Meta):
        model = models.Tag

class AbstractCreativeWorkSerializer(BaseShareSerializer):
    tags = TagSerializer(many=True)
    contributors = ContributorSerializer(source='contributor_set', many=True)
    institutions = InstitutionSerializer(sparse=True, many=True)
    venues = VenueSerializer(sparse=True, many=True)
    awards = AwardSerializer(sparse=True, many=True)
    subject = TagSerializer(sparse=True)

class CreativeWorkSerializer(AbstractCreativeWorkSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.CreativeWork

class PreprintSerializer(AbstractCreativeWorkSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Preprint

class ManuscriptSerializer(AbstractCreativeWorkSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Manuscript

class Link(BaseShareSerializer):
    class Meta(BaseShareSerializer.Meta):
        model = models.Link
<|endoftext|>


<commit_before># vim: tabstop=4 shiftwidth=4 softtabstop=4

from lxml import etree

from opengem import producer
from opengem import shapes

GML = 'http://www.opengis.net/gml'

class ShamlFile(producer.FileProducer):
    def _parse(self):
        for event, element in etree.iterparse(self.file,
                                              events=('start', 'end')):
            # We currently only care about the tag name, not the namespace
            tag = element.tag.split('}')[-1]

            if event == 'end' and tag == 'BranchLevelList':
                self._build_branching_levels(element)
            elif event == 'start' and tag == 'EndBranchMapping':
                self._set_branch(element.get('label'))
            elif event == 'end' and tag == 'Source':
                yield (self._to_region(element), self._to_source(element))

            # To save space we could clear elements we no longer care about
            # by doing element.clear(), for example on the Source elements

    def _build_branching_levels(self, branching_levels):
        pass

    def _set_branch(self, branch_label):
        self._current_branch = branch_label

    def _to_region(self, element):
        pos_list = element.find('.//{%s}posList' % GML)

        # build up the region for this element
        # for now we assume we are getting lat-long, not lat-long-depth.
        parts = pos_list.text.strip().split()
        coords = []
        while parts:
            coords.append((float(parts.pop(0)), float(parts.pop(0))))

        linestring = shapes.LineString(coords)
        return linestring

    def _to_source(self, element):
        # TODO(termie): map the element into a source, using
        #               self._current_branch as context
        return element
<commit_msg>Deleted unused sketch producer<commit_after><|endoftext|>


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

import unittest
import numpy as np

from pyscf.pbc import gto as pbcgto
from pyscf.pbc import scf as pbcscf
from pyscf.pbc import df as pdf
import pyscf.pbc.mp
import pyscf.pbc.mp.kmp2

cell = pbcgto.Cell()
cell.unit = 'B'
L = 7
cell.atom.extend([['Be', (L/2.,  L/2., L/2.)]])
cell.a = 7 * np.identity(3)
cell.a[1,0] = 5.0

cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade-q2'
cell.mesh = [12]*3
cell.verbose = 5
cell.output = '/dev/null'
cell.build()

def build_h_cell():
    # Returns FCC H cell.
    cell = pbcgto.Cell()
    cell.atom = [['H', (0.000000000, 0.000000000, 0.000000000)],
                 ['H', (0.000000000, 0.500000000, 0.250000000)],
                 ['H', (0.500000000, 0.500000000, 0.500000000)],
                 ['H', (0.500000000, 0.000000000, 0.750000000)]]
    cell.unit = 'Bohr'
    cell.a = [[1.,0.,0.],[0.,1.,0],[0,0,2.2]]
    cell.verbose = 7
    cell.spin = 0
    cell.charge = 0
    cell.basis = [[0, [1.0, 1]],]
    cell.pseudo = 'gth-pade'
    cell.output = '/dev/null'
    cell.max_memory = 1000
    for i in range(len(cell.atom)):
        cell.atom[i][1] = tuple(np.dot(np.array(cell.atom[i][1]),np.array(cell.a)))
    cell.build()
    return cell

def run_kcell(cell, nk):
    abs_kpts = cell.make_kpts(nk, wrap_around=True)
    kmf = pbcscf.KRHF(cell, abs_kpts)
    kmf.conv_tol = 1e-12
    ekpt = kmf.scf()
    mp = pyscf.pbc.mp.kmp2.KMP2(kmf).run()
    return ekpt, mp.e_corr

def run_kcell_complex(cell, nk):
    abs_kpts = cell.make_kpts(nk, wrap_around=True)
    kmf = pbcscf.KRHF(cell, abs_kpts)
    kmf.conv_tol = 1e-12
    ekpt = kmf.scf()
    kmf.mo_coeff = [kmf.mo_coeff[i].astype(np.complex128) for i in range(np.prod(nk))]
    mp = pyscf.pbc.mp.kmp2.KMP2(kmf).run()
    return ekpt, mp.e_corr

class KnownValues(unittest.TestCase):
    def test_111(self):
        nk = (1, 1, 1)
        escf, emp = run_kcell(cell,nk)
        self.assertAlmostEqual(escf, -1.2061049658473704, 9)
        self.assertAlmostEqual(emp, -5.44597932944397e-06, 9)
        escf, emp = run_kcell_complex(cell,nk)
        self.assertAlmostEqual(emp, -5.44597932944397e-06, 9)

    def test_311_high_cost(self):
        nk = (3, 1, 1)
        escf, emp = run_kcell(cell,nk)
        self.assertAlmostEqual(escf, -1.0585001200928885, 9)
        self.assertAlmostEqual(emp, -7.9832274354253814e-06, 9)

    def test_h4_fcc_k2_frozen(self):
        '''Metallic hydrogen fcc lattice with frozen lowest lying occupied
        and highest lying virtual orbitals.  Checks versus a corresponding
        supercell calculation.

        NOTE: different versions of the davidson may converge to a different
        solution for the k-point IP/EA eom.  If you're getting the wrong
        root, check to see if it's contained in the supercell set of
        eigenvalues.'''
        cell = build_h_cell()

        nmp = [2, 1, 1]

        kmf = pbcscf.KRHF(cell)
        kmf.kpts = cell.make_kpts(nmp, scaled_center=[0.0,0.0,0.0])
        e = kmf.kernel()

        frozen = [[0, 3], []]
        mymp = pyscf.pbc.mp.kmp2.KMP2(kmf, frozen=frozen)
        ekmp2, _ = mymp.kernel()
        self.assertAlmostEqual(ekmp2, -0.022416773725207319, 6)
        self.assertAlmostEqual(mymp.e_tot, 2.155470531550681, 6)

        # Start of supercell calculations
        from pyscf.pbc.tools.pbc import super_cell
        supcell = super_cell(cell, nmp)
        supcell.build()
        mf = pbcscf.KRHF(supcell)
        e = mf.kernel()

        mysmp = pyscf.pbc.mp.kmp2.KMP2(mf, frozen=[0, 7])
        emp2, _ = mysmp.kernel()
        emp2 /= np.prod(nmp)
        self.assertAlmostEqual(emp2, -0.022416773725207319, 6)

    def test_h4_fcc_k2_frozen_df_nocc(self):
        '''Metallic hydrogen fcc lattice.  Test diff. nocc at k points.

        NOTE: different versions of the davidson may converge to a different
        solution for the k-point IP/EA eom.  If you're getting the wrong
        root, check to see if it's contained in the supercell set of
        eigenvalues.'''
        cell = build_h_cell()

        nmp = [2, 1, 1]

        kmf = pbcscf.KRHF(cell)
        kmf.with_df = pdf.GDF(cell)
        kmf.kpts = cell.make_kpts(nmp, scaled_center=[0.0,0.0,0.0])
        e = kmf.kernel()

        frozen = [[0, 3], []]
        mymp = pyscf.pbc.mp.kmp2.KMP2(kmf, frozen=frozen)
        ekmp2, _ = mymp.kernel()
        self.assertAlmostEqual(ekmp2, -0.016333989667540873, 6)
        self.assertAlmostEqual(mymp.e_tot, 2.329841282521279, 6)

    def test_rdm1(self):
        cell = pbcgto.Cell()
        cell.atom = '''Al 0 0 0'''
        cell.basis = 'gth-szv'
        cell.pseudo = 'gth-pade'
        cell.a = '''
        2.47332919 0 1.42797728
        0.82444306 2.33187713 1.42797728
        0, 0, 2.85595455
        '''
        cell.unit = 'angstrom'
        cell.build()
        cell.verbose = 4
        cell.incore_anyway = True

        abs_kpts = cell.make_kpts((2, 1, 1), scaled_center=(.1, .2, .3))
        kmf = pbcscf.KRHF(cell, abs_kpts)
        kmf.conv_tol = 1e-12
        kmf.kernel()
        mp = pyscf.pbc.mp.kmp2.KMP2(kmf)
        mp.kernel(with_t2=True)
        self.assertAlmostEqual(mp.e_corr, -0.00162057921874043)
        dm = mp.make_rdm1()
        np.testing.assert_allclose(np.trace(dm[0]) + np.trace(dm[1]), 6)
        for kdm in dm:
            np.testing.assert_allclose(kdm, kdm.conj().T)

if __name__ == '__main__':
    print("Full kpoint test")
    unittest.main()

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

import unittest
import numpy as np

from pyscf.pbc import gto as pbcgto
from pyscf.pbc import scf as pbcscf
import pyscf.pbc.mp
import pyscf.pbc.mp.kmp2

cell = pbcgto.Cell()
cell.unit = 'B'
L = 7
cell.atom.extend([['Be', (L/2.,  L/2., L/2.)]])
cell.a = 7 * np.identity(3)
cell.a[1,0] = 5.0

cell.basis = 'gth-szv'
cell.pseudo = 'gth-pade-q2'
cell.mesh = [12]*3
cell.verbose = 5
cell.output = '/dev/null'
cell.build()

def build_h_cell():
    # Returns FCC H cell.
    cell = pbcgto.Cell()
    cell.atom = [['H', (0.000000000, 0.000000000, 0.000000000)],
                 ['H', (0.000000000, 0.500000000, 0.250000000)],
                 ['H', (0.500000000, 0.500000000, 0.500000000)],
                 ['H', (0.500000000, 0.000000000, 0.750000000)]]
    cell.unit = 'Bohr'
    cell.a = [[1.,0.,0.],[0.,1.,0],[0,0,2.2]]
    cell.verbose = 7
    cell.spin = 0
    cell.charge = 0
    cell.basis = [[0, [1.0, 1]],]
    cell.pseudo = 'gth-pade'
    cell.output = '/dev/null'
    cell.max_memory = 1000
    for i in range(len(cell.atom)):
        cell.atom[i][1] = tuple(np.dot(np.array(cell.atom[i][1]),np.array(cell.a)))
    cell.build()
    return cell

def run_kcell(cell, nk):
    abs_kpts = cell.make_kpts(nk, wrap_around=True)
    kmf = pbcscf.KRHF(cell, abs_kpts)
    kmf.conv_tol = 1e-12
    ekpt = kmf.scf()
    mp = pyscf.pbc.mp.kmp2.KMP2(kmf).run()
    return ekpt, mp.e_corr

def run_kcell_complex(cell, nk):
    abs_kpts = cell.make_kpts(nk, wrap_around=True)
    kmf = pbcscf.KRHF(cell, abs_kpts)
    kmf.conv_tol = 1e-12
    ekpt = kmf.scf()
    kmf.mo_coeff = [kmf.mo_coeff[i].astype(np.complex128) for i in range(np.prod(nk))]
    mp = pyscf.pbc.mp.kmp2.KMP2(kmf).run()
    return ekpt, mp.e_corr

class KnownValues(unittest.TestCase):
    def test_111(self):
        nk = (1, 1, 1)
        escf, emp = run_kcell(cell,nk)
        self.assertAlmostEqual(escf, -1.2061049658473704, 9)
        self.assertAlmostEqual(emp, -5.44597932944397e-06, 9)
        escf, emp = run_kcell_complex(cell,nk)
        self.assertAlmostEqual(emp, -5.44597932944397e-06, 9)

    def test_311_high_cost(self):
        nk = (3, 1, 1)
        escf, emp = run_kcell(cell,nk)
        self.assertAlmostEqual(escf, -1.0585001200928885, 9)
        self.assertAlmostEqual(emp, -7.9832274354253814e-06, 9)

    def test_h4_fcc_k2_frozen(self):
        '''Metallic hydrogen fcc lattice with frozen lowest lying occupied
        and highest lying virtual orbitals.  Checks versus a corresponding
        supercell calculation.

        NOTE: different versions of the davidson may converge to a different
        solution for the k-point IP/EA eom.  If you're getting the wrong
        root, check to see if it's contained in the supercell set of
        eigenvalues.'''
        cell = build_h_cell()

        nmp = [2, 1, 1]

        kmf = pbcscf.KRHF(cell)
        kmf.kpts = cell.make_kpts(nmp, scaled_center=[0.0,0.0,0.0])
        e = kmf.kernel()

        frozen = [[0, 3], []]
        mymp = pyscf.pbc.mp.kmp2.KMP2(kmf, frozen=frozen)
        ekmp2, _ = mymp.kernel()
        self.assertAlmostEqual(ekmp2, -0.022416773725207319, 6)
        self.assertAlmostEqual(mymp.e_tot, 2.155470531550681, 6)

        # Start of supercell calculations
        from pyscf.pbc.tools.pbc import super_cell
        supcell = super_cell(cell, nmp)
        supcell.build()
        mf = pbcscf.KRHF(supcell)
        e = mf.kernel()

        mysmp = pyscf.pbc.mp.kmp2.KMP2(mf, frozen=[0, 7])
        emp2, _ = mysmp.kernel()
        emp2 /= np.prod(nmp)
        self.assertAlmostEqual(emp2, -0.022416773725207319, 6)

    def test_h4_fcc_k2_frozen_df_nocc(self):
        '''Metallic hydrogen fcc, test different nocc at k points.'''
        cell = build_h_cell()

        nmp = [2, 1, 1]

        kmf = pbcscf.KRHF(cell).density_fit()
        kmf.kpts = cell.make_kpts(nmp, scaled_center=[0.0,0.0,0.0])
        e = kmf.kernel()

        frozen = [[0, 3], []]
        mymp = pyscf.pbc.mp.kmp2.KMP2(kmf, frozen=frozen)
        ekmp2, _ = mymp.kernel()
        self.assertAlmostEqual(ekmp2, -0.016333989667540873, 6)
        self.assertAlmostEqual(mymp.e_tot, 2.329841282521279, 6)

    def test_rdm1(self):
        cell = pbcgto.Cell()
        cell.atom = '''Al 0 0 0'''
        cell.basis = 'gth-szv'
        cell.pseudo = 'gth-pade'
        cell.a = '''
        2.47332919 0 1.42797728
        0.82444306 2.33187713 1.42797728
        0, 0, 2.85595455
        '''
        cell.unit = 'angstrom'
        cell.build()
        cell.verbose = 4
        cell.incore_anyway = True

        abs_kpts = cell.make_kpts((2, 1, 1), scaled_center=(.1, .2, .3))
        kmf = pbcscf.KRHF(cell, abs_kpts)
        kmf.conv_tol = 1e-12
        kmf.kernel()
        mp = pyscf.pbc.mp.kmp2.KMP2(kmf)
        mp.kernel(with_t2=True)
        self.assertAlmostEqual(mp.e_corr, -0.00162057921874043)
        dm = mp.make_rdm1()
        np.testing.assert_allclose(np.trace(dm[0]) + np.trace(dm[1]), 6)
        for kdm in dm:
            np.testing.assert_allclose(kdm, kdm.conj().T)

if __name__ == '__main__':
    print("Full kpoint test")
    unittest.main()

<|endoftext|>


<commit_before><commit_msg>exclude current doc while validating (#13450)<commit_after><|endoftext|>