<commit_before>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import datetime
import itertools
import locale
import simplejson as json
import math

locale.setlocale(locale.LC_ALL, '')

BIG_VALUE = 2 ** 60
SMALL_VALUE = - (2 ** 60)

def devisible(a, b):
    if b == 0:
        return False
    return a % b == 0

def hostport(hoststring, default_port=8091):
    """ finds the host and port given a host:port string """
    try:
        host, port = hoststring.split(':')
        port = int(port)
    except ValueError:
        host = hoststring
        port = default_port

    return (host, port)

def time_label(s):
    # -(2**64) -> '-inf'
    # 2**64 -> 'inf'
    # 0 -> '0'
    # 4 -> '4us'
    # 838384 -> '838ms'
    # 8283852 -> '8s'
    if s > BIG_VALUE:
        return 'inf'
    elif s < SMALL_VALUE:
        return '-inf'
    elif s == 0:
        return '0'
    product = 1
    sizes = (('us', 1), ('ms', 1000), ('sec', 1000), ('min', 60))
    sizeMap = []
    for l,sz in sizes:
        product = sz * product
        sizeMap.insert(0, (l, product))
    lbl, factor = itertools.dropwhile(lambda x: x[1] > s, sizeMap).next()
    if devisible(s, factor):
        return '%d %s' % (s / factor, lbl)
    else:
        return '%.*f %s' % (3, s * 1.0/factor, lbl)

def size_label(s):
    if type(s) in (int, long, float, complex) :
        if s == 0:
            return "0"
        sizes=['', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB']
        e = math.floor(math.log(abs(s), 1024))
        suffix = sizes[int(e)]
        if devisible(s, 1024 ** math.floor(e)):
            return '%d %s' % ( s / (1024 ** math.floor(e)), suffix)
        else:
            return "%.*f %s" % (3, s *1.0/(1024 ** math.floor(e)), suffix)
    else:
        return s

def number_label(s):
    if type(s) in (int, long, float, complex) :
        if s < 0:
            s = -s
            flag = "-"
        else:
            flag = ""
        if s < 1:
            return "0"
        sizes=['', 'thousand', 'million', 'billion', 'trillion', 'quadrillion', 'quintillion']
        e = math.floor(math.log(abs(s), 1000))
        if e < 0:
           e = 0
        suffix = sizes[int(e)]
        if devisible(s, 1000 ** math.floor(e)):
            return "%s%d %s" % (flag, s / (1000 ** math.floor(e)), suffix)
        else:
            return "%s%.*f %s" % (flag, 2, s *1.0/(1000 ** math.floor(e)), suffix)
    else:
        return s

def linreg(X, Y):
    """
    Summary
        Linear regression of y = ax + b
    Usage
        real, real, real = linreg(list, list)
    Returns coefficients to the regression line "y=ax+b" from x[] and y[], and R^2 Value
    """
    if len(X) != len(Y):  raise ValueError, 'unequal length'
    N = len(X)
    Sx = Sy = Sxx = Syy = Sxy = 0.0
    for x, y in map(None, X, Y):
        Sx = Sx + x
        Sy = Sy + y
        Sxx = Sxx + x*x
        Syy = Syy + y*y
        Sxy = Sxy + x*y
    det = Sxx * N - Sx * Sx
    if det == 0:
        return 0, 0
    else:
        a, b = (Sxy * N - Sy * Sx)/det, (Sxx * Sy - Sx * Sxy)/det
        return a, b

def two_pass_variance(data):
    n    = 0
    sum1 = 0
    sum2 = 0

    for x in data:
        n    = n + 1
        sum1 = sum1 + x

    mean = sum1/n
    for x in data:
        sum2 = sum2 + (x - mean)*(x - mean)
    if n <= 1:
        return 0
    variance = sum2/(n - 1)
    return variance

def abnormal_extract(vals, threshold, op = '>='):
    abnormal = []
    begin_index = -1
    seg_count = 0

    for index, sample in enumerate(vals):
        ev = evalfunc(sample, threshold, op)
        if ev is None:
            return abnormal
        elif ev:
            if begin_index < 0:
                begin_index = index
            seg_count += 1
        else:
            if begin_index >= 0:
                abnormal.append((begin_index, seg_count))
                begin_index = -1
                seg_count = 0

    if begin_index >= 0:
        abnormal.append((begin_index, seg_count))
    return abnormal

def evalfunc(value, threshold, op):
    rt = None

    func = {'>=' : lambda x, y: x >= y,
            '>' : lambda x, y: x > y,
            '==' : lambda x, y: x == y,
            '<' : lambda x, y: x < y,
            '!=' : lambda x, y: x != y,
           }.get(op, None)
    if func is None:
        return rt
    return func(value, threshold)

def pretty_float(number, precision=2):
    return '%.*f' % (precision, number)

def pretty_print(obj):
    return json.dumps(obj, indent=4, sort_keys=True)

def pretty_datetime(number, timeonly=False):
    if timeonly:
        return str(datetime.datetime.fromtimestamp(number/1000).time())
    else:
        timestamp = datetime.datetime.fromtimestamp(number/1000)
        return timestamp.strftime('%x') + ' ' + str(timestamp.time())
<commit_msg>MB-8366: transform data based specified unit<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

import datetime
import itertools
import locale
import simplejson as json
import math

locale.setlocale(locale.LC_ALL, '')

BIG_VALUE = 2 ** 60
SMALL_VALUE = - (2 ** 60)

def devisible(a, b):
    if b == 0:
        return False
    return a % b == 0

def hostport(hoststring, default_port=8091):
    """ finds the host and port given a host:port string """
    try:
        host, port = hoststring.split(':')
        port = int(port)
    except ValueError:
        host = hoststring
        port = default_port

    return (host, port)

def time_label(s):
    # -(2**64) -> '-inf'
    # 2**64 -> 'inf'
    # 0 -> '0'
    # 4 -> '4us'
    # 838384 -> '838ms'
    # 8283852 -> '8s'
    if s > BIG_VALUE:
        return 'inf'
    elif s < SMALL_VALUE:
        return '-inf'
    elif s == 0:
        return '0'
    product = 1
    sizes = (('us', 1), ('ms', 1000), ('sec', 1000), ('min', 60))
    sizeMap = []
    for l,sz in sizes:
        product = sz * product
        sizeMap.insert(0, (l, product))
    try:
        lbl, factor = itertools.dropwhile(lambda x: x[1] > s, sizeMap).next()
    except StopIteration:
        lbl, factor = sizeMap[-1]
    if devisible(s, factor):
        return '%d %s' % (s / factor, lbl)
    else:
        return '%.*f %s' % (3, s * 1.0/factor, lbl)

def size_label(s):
    if type(s) in (int, long, float, complex) :
        if s == 0:
            return "0"
        sizes=['', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB']
        e = math.floor(math.log(abs(s), 1024))
        suffix = sizes[int(e)]
        if devisible(s, 1024 ** math.floor(e)):
            return '%d %s' % ( s / (1024 ** math.floor(e)), suffix)
        else:
            return "%.*f %s" % (3, s *1.0/(1024 ** math.floor(e)), suffix)
    else:
        return s

def size_convert(s, unit):
    if type(s) in (int, long, float, complex) :
        if s == 0:
            return s
        sizes=['', 'KB', 'MB', 'GB', 'TB', 'PB', 'EB']
        try:
            e = sizes.index(unit.upper())
        except Exception:
            e = 0
        if devisible(s, 1024 ** math.floor(e)):
            return '%d' % (s / (1024 ** math.floor(e)))
        else:
            return "%.*f" % (3, s *1.0/(1024 ** math.floor(e)))
    else:
        return s

def number_label(s):
    if type(s) in (int, long, float, complex) :
        if s < 0:
            s = -s
            flag = "-"
        else:
            flag = ""
        if s < 1:
            return "0"
        sizes=['', 'thousand', 'million', 'billion', 'trillion', 'quadrillion', 'quintillion']
        e = math.floor(math.log(abs(s), 1000))
        if e < 0:
           e = 0
        suffix = sizes[int(e)]
        if devisible(s, 1000 ** math.floor(e)):
            return "%s%d %s" % (flag, s / (1000 ** math.floor(e)), suffix)
        else:
            return "%s%.*f %s" % (flag, 2, s *1.0/(1000 ** math.floor(e)), suffix)
    else:
        return s

def linreg(X, Y):
    """
    Summary
        Linear regression of y = ax + b
    Usage
        real, real, real = linreg(list, list)
    Returns coefficients to the regression line "y=ax+b" from x[] and y[], and R^2 Value
    """
    if len(X) != len(Y):  raise ValueError, 'unequal length'
    N = len(X)
    Sx = Sy = Sxx = Syy = Sxy = 0.0
    for x, y in map(None, X, Y):
        Sx = Sx + x
        Sy = Sy + y
        Sxx = Sxx + x*x
        Syy = Syy + y*y
        Sxy = Sxy + x*y
    det = Sxx * N - Sx * Sx
    if det == 0:
        return 0, 0
    else:
        a, b = (Sxy * N - Sy * Sx)/det, (Sxx * Sy - Sx * Sxy)/det
        return a, b

def two_pass_variance(data):
    n    = 0
    sum1 = 0
    sum2 = 0

    for x in data:
        n    = n + 1
        sum1 = sum1 + x

    mean = sum1/n
    for x in data:
        sum2 = sum2 + (x - mean)*(x - mean)
    if n <= 1:
        return 0
    variance = sum2/(n - 1)
    return variance

def abnormal_extract(vals, threshold, op = '>='):
    abnormal = []
    begin_index = -1
    seg_count = 0

    for index, sample in enumerate(vals):
        ev = evalfunc(sample, threshold, op)
        if ev is None:
            return abnormal
        elif ev:
            if begin_index < 0:
                begin_index = index
            seg_count += 1
        else:
            if begin_index >= 0:
                abnormal.append((begin_index, seg_count))
                begin_index = -1
                seg_count = 0

    if begin_index >= 0:
        abnormal.append((begin_index, seg_count))
    return abnormal

def evalfunc(value, threshold, op):
    rt = None

    func = {'>=' : lambda x, y: x >= y,
            '>' : lambda x, y: x > y,
            '==' : lambda x, y: x == y,
            '<' : lambda x, y: x < y,
            '!=' : lambda x, y: x != y,
           }.get(op, None)
    if func is None:
        return rt
    return func(value, threshold)

def pretty_float(number, precision=2):
    return '%.*f' % (precision, number)

def pretty_print(obj):
    return json.dumps(obj, indent=4, sort_keys=True)

def pretty_datetime(number, timeonly=False):
    if timeonly:
        return str(datetime.datetime.fromtimestamp(number/1000).time())
    else:
        timestamp = datetime.datetime.fromtimestamp(number/1000)
        return timestamp.strftime('%x') + ' ' + str(timestamp.time())
<|endoftext|>


<commit_before># coding: utf-8

import tornado

__all__ = ['sms_send', 'ret_code2desc']

errors = {
        -1: u"没有该用户账户",
        -2: u"密钥不正确（不是用户密码）",
        -3: u"短信数量不足",
        -11: u"该用户被禁用",
        -14: u"短信内容出现非法字符",
        -41: u"手机号码为空",
        -42: u"短信内容为空",
        }

def handler_sms_reqeust(response):
    if response.error:
        print "Oops! error to request the sms service."
    else:
        ret_code = int(response.body)
        print ret_code2desc(ret_code)
        tornado.ioloop.IOLoop.instance().stop()

def sms_send(phone, msg, msg_type):

    assert isinstance(msg_type, int) and isinstance(msg, dict), "apply function with wrong arguments type"

    # sms content
    msg_content = get_sms_template(msg_type) % msg

    # sms api uri
    uri = u"http://utf8.sms.webchinese.cn/?Uid=cloud&Key=q12wer43ui8765tyop09&smsMob=%(phone)s&smsText=%(msg)s"
    uri = uri % {'phone': phone, 'msg': tornado.escape.url_escape(msg_content.encode('utf-8'))}

    http_client = httpclient.AsyncHTTPClient()
    http_client.fetch(uri, handle_request)

    tornado.ioloop.IOLoop.instance().start()

def get_sms_template(msg_type):
    res = u""
    if msg_type == 1:
        res = u"乐帮注册验证码:%(code)s。乐帮感谢你的注册！"
    elif msg_type == 2:
        res = u"%(name)s，你好！重置后的密码:%(password)s。乐帮提醒你尽快修改！"
    return res

def ret_code2desc(ret_code):
    if ret_code < 0:
        return errors[ret_code]
    else:
        return u"成功发送%d条" % ret_code
<commit_msg>fuck this world<commit_after># coding: utf-8

import tornado

__all__ = ['sms_send', 'ret_code2desc']

errors = {
        -1: u"没有该用户账户",
        -2: u"密钥不正确（不是用户密码）",
        -3: u"短信数量不足",
        -11: u"该用户被禁用",
        -14: u"短信内容出现非法字符",
        -41: u"手机号码为空",
        -42: u"短信内容为空",
        }

def handler_sms_reqeust(response):
    if response.error:
        print "Oops! error to request the sms service."
    else:
        ret_code = int(response.body)
        print ret_code2desc(ret_code)
        tornado.ioloop.IOLoop.instance().stop()

def sms_send(phone, msg, msg_type):

    assert isinstance(msg_type, int) and isinstance(msg, dict), "apply function with wrong arguments type"

    # sms content
    msg_content = get_sms_template(msg_type) % msg

    # sms api uri
    uri = u"http://utf8.sms.webchinese.cn/?Uid=cloud&Key=q12wer43ui8765tyop09&smsMob=%(phone)s&smsText=%(msg)s"
    uri = uri % {'phone': phone, 'msg': tornado.escape.url_escape(msg_content.encode('utf-8'))}

    http_client = tornado.httpclient.AsyncHTTPClient()
    http_client.fetch(uri, handler_sms_reqeust)

    tornado.ioloop.IOLoop.instance().start()

def get_sms_template(msg_type):
    res = u""
    if msg_type == 1:
        res = u"乐帮注册验证码:%(code)s。乐帮感谢你的注册！"
    elif msg_type == 2:
        res = u"%(name)s，你好！重置后的密码:%(password)s。乐帮提醒你尽快修改！"
    return res

def ret_code2desc(ret_code):
    if ret_code < 0:
        return errors[ret_code]
    else:
        return u"成功发送%d条" % ret_code
<|endoftext|>


<commit_before>
import json
import codecs
from jsonschema import validate

# load schema and json file
schema_file = './graphson_schema.json'
schema = json.load( codecs.open(schema_file, 'r', encoding='utf-8') )

obj_file = './normal_test_graph.json'
obj = json.load( codecs.open(obj_file, 'r', encoding='utf-8') )

# If no exception is raised by validate(), the instance is valid.
validate(obj, schema)

obj_file = './extended_test_graph.json'
obj = json.load( codecs.open(obj_file, 'r', encoding='utf-8') )

# If no exception is raised by validate(), the instance is valid.
validate(obj, schema)

print 'passed!'<commit_msg>path fix<commit_after>
import json
import codecs
from jsonschema import validate

# load schema and json file
schema_file = './graphson_schema.json'
schema = json.load( codecs.open(schema_file, 'r', encoding='utf-8') )

obj_file = './tests/normal_test_graph.json'
obj = json.load( codecs.open(obj_file, 'r', encoding='utf-8') )

# If no exception is raised by validate(), the instance is valid.
validate(obj, schema)

obj_file = './tests/extended_test_graph.json'
obj = json.load( codecs.open(obj_file, 'r', encoding='utf-8') )

# If no exception is raised by validate(), the instance is valid.
validate(obj, schema)

print 'passed!'<|endoftext|>


<commit_before>#!/usr/bin/env python
""" ImageNet Validation Script

This is intended to be a lean and easily modifiable ImageNet validation script for evaluating pretrained
models or training checkpoints against ImageNet or similarly organized image datasets. It prioritizes
canonical PyTorch, standard Python style, and good performance. Repurpose as you see fit.

Hacked together by Ross Wightman (https://github.com/rwightman)
"""
import argparse
import os
import csv
import glob
import time
import logging
import torch
import torch.nn as nn
import torch.nn.parallel
from collections import OrderedDict
from contextlib import suppress

from timm.models import create_model, apply_test_time_pool, load_checkpoint, is_model, list_models
from timm.data import create_dataset, create_loader, resolve_data_config, RealLabelsImagenet
from timm.utils import accuracy, AverageMeter, natural_key, setup_default_logging, set_jit_legacy

has_apex = False
try:
    from apex import amp
    has_apex = True
except ImportError:
    pass

has_native_amp = False
try:
    if getattr(torch.cuda.amp, 'autocast') is not None:
        has_native_amp = True
except AttributeError:
    pass

torch.backends.cudnn.benchmark = True
_logger = logging.getLogger('validate')

parser = argparse.ArgumentParser(description='PyTorch ImageNet Validation')
parser.add_argument('data', metavar='DIR',
                    help='path to dataset')
parser.add_argument('--dataset', '-d', metavar='NAME', default='',
                    help='dataset type (default: ImageFolder/ImageTar if empty)')
parser.add_argument('--split', metavar='NAME', default='validation',
                    help='dataset split (default: validation)')
parser.add_argument('--model', '-m', metavar='NAME', default='dpn92',
                    help='model architecture (default: dpn92)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 2)')
parser.add_argument('-b', '--batch-size', default=256, type=int,
                    metavar='N', help='mini-batch size (default: 256)')
parser.add_argument('--img-size', default=None, type=int,
                    metavar='N', help='Input image dimension, uses model default if empty')
parser.add_argument('--input-size', default=None, nargs=3, type=int,
                    metavar='N N N', help='Input all image dimensions (d h w, e.g. --input-size 3 224 224), uses model default if empty')
parser.add_argument('--crop-pct', default=None, type=float,
                    metavar='N', help='Input image center crop pct')
parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
                    help='Override mean pixel value of dataset')
parser.add_argument('--std', type=float,  nargs='+', default=None, metavar='STD',
                    help='Override std deviation of of dataset')
parser.add_argument('--interpolation', default='', type=str, metavar='NAME',
                    help='Image resize interpolation type (overrides model)')
parser.add_argument('--num-classes', type=int, default=None,
                    help='Number classes in dataset')
parser.add_argument('--class-map', default='', type=str, metavar='FILENAME',
                    help='path to class to idx mapping file (default: "")')
parser.add_argument('--gp', default=None, type=str, metavar='POOL',
                    help='Global pool type, one of (fast, avg, max, avgmax, avgmaxc). Model default if None.')
parser.add_argument('--log-freq', default=10, type=int,
                    metavar='N', help='batch logging frequency (default: 10)')
parser.add_argument('--checkpoint', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                    help='use pre-trained model')
parser.add_argument('--num-gpu', type=int, default=1,
                    help='Number of GPUS to use')
parser.add_argument('--no-test-pool', dest='no_test_pool', action='store_true',
                    help='disable test time pool')
parser.add_argument('--no-prefetcher', action='store_true', default=False,
                    help='disable fast prefetcher')
parser.add_argument('--pin-mem', action='store_true', default=False,
                    help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
parser.add_argument('--channels-last', action='store_true', default=False,
                    help='Use channels_last memory layout')
parser.add_argument('--amp', action='store_true', default=False,
                    help='Use AMP mixed precision. Defaults to Apex, fallback to native Torch AMP.')
parser.add_argument('--apex-amp', action='store_true', default=False,
                    help='Use NVIDIA Apex AMP mixed precision')
parser.add_argument('--native-amp', action='store_true', default=False,
                    help='Use Native Torch AMP mixed precision')
parser.add_argument('--tf-preprocessing', action='store_true', default=False,
                    help='Use Tensorflow preprocessing pipeline (require CPU TF installed')
parser.add_argument('--use-ema', dest='use_ema', action='store_true',
                    help='use ema version of weights if present')
parser.add_argument('--torchscript', dest='torchscript', action='store_true',
                    help='convert model torchscript for inference')
parser.add_argument('--legacy-jit', dest='legacy_jit', action='store_true',
                    help='use legacy jit mode for pytorch 1.5/1.5.1/1.6 to get back fusion performance')
parser.add_argument('--results-file', default='', type=str, metavar='FILENAME',
                    help='Output csv file for validation results (summary)')
parser.add_argument('--real-labels', default='', type=str, metavar='FILENAME',
                    help='Real labels JSON file for imagenet evaluation')
parser.add_argument('--valid-labels', default='', type=str, metavar='FILENAME',
                    help='Valid label indices txt file for validation of partial label space')

def validate(args):
    # might as well try to validate something
    args.pretrained = args.pretrained or not args.checkpoint
    args.prefetcher = not args.no_prefetcher
    amp_autocast = suppress  # do nothing
    if args.amp:
        if has_apex:
            args.apex_amp = True
        elif has_native_amp:
            args.native_amp = True
        else:
            _logger.warning("Neither APEX or Native Torch AMP is available, using FP32.")
    assert not args.apex_amp or not args.native_amp, "Only one AMP mode should be set."
    if args.native_amp:
        amp_autocast = torch.cuda.amp.autocast

    if args.legacy_jit:
        set_jit_legacy()

    # create model
    model = create_model(
        args.model,
        pretrained=args.pretrained,
        num_classes=args.num_classes,
        in_chans=3,
        global_pool=args.gp,
        scriptable=args.torchscript)
    if args.num_classes is None:
        assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.'
        args.num_classes = model.num_classes

    if args.checkpoint:
        load_checkpoint(model, args.checkpoint, args.use_ema)

    param_count = sum([m.numel() for m in model.parameters()])
    _logger.info('Model %s created, param count: %d' % (args.model, param_count))

    data_config = resolve_data_config(vars(args), model=model)
    model, test_time_pool = (model, False) if args.no_test_pool else apply_test_time_pool(model, data_config)

    if args.torchscript:
        torch.jit.optimized_execution(True)
        model = torch.jit.script(model)

    model = model.cuda()
    if args.apex_amp:
        model = amp.initialize(model, opt_level='O1')

    if args.channels_last:
        model = model.to(memory_format=torch.channels_last)

    if args.num_gpu > 1:
        model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu)))

    criterion = nn.CrossEntropyLoss().cuda()

    dataset = create_dataset(
        root=args.data, name=args.dataset, split=args.split,
        load_bytes=args.tf_preprocessing, class_map=args.class_map)

    if args.valid_labels:
        with open(args.valid_labels, 'r') as f:
            valid_labels = {int(line.rstrip()) for line in f}
            valid_labels = [i in valid_labels for i in range(args.num_classes)]
    else:
        valid_labels = None

    if args.real_labels:
        real_labels = RealLabelsImagenet(dataset.filenames(basename=True), real_json=args.real_labels)
    else:
        real_labels = None

    crop_pct = 1.0 if test_time_pool else data_config['crop_pct']
    loader = create_loader(
        dataset,
        input_size=data_config['input_size'],
        batch_size=args.batch_size,
        use_prefetcher=args.prefetcher,
        interpolation=data_config['interpolation'],
        mean=data_config['mean'],
        std=data_config['std'],
        num_workers=args.workers,
        crop_pct=crop_pct,
        pin_memory=args.pin_mem,
        tf_preprocessing=args.tf_preprocessing)

    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    model.eval()
    with torch.no_grad():
        # warmup, reduce variability of first batch time, especially for comparing torchscript vs non
        input = torch.randn((args.batch_size,) + data_config['input_size']).cuda()
        if args.channels_last:
            input = input.contiguous(memory_format=torch.channels_last)
        model(input)
        end = time.time()
        for batch_idx, (input, target) in enumerate(loader):
            if args.no_prefetcher:
                target = target.cuda()
                input = input.cuda()
            if args.channels_last:
                input = input.contiguous(memory_format=torch.channels_last)

            # compute output
            with amp_autocast():
                output = model(input)

            if valid_labels is not None:
                output = output[:, valid_labels]
            loss = criterion(output, target)

            if real_labels is not None:
                real_labels.add_result(output)

            # measure accuracy and record loss
            acc1, acc5 = accuracy(output.detach(), target, topk=(1, 5))
            losses.update(loss.item(), input.size(0))
            top1.update(acc1.item(), input.size(0))
            top5.update(acc5.item(), input.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if batch_idx % args.log_freq == 0:
                _logger.info(
                    'Test: [{0:>4d}/{1}]  '
                    'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s)  '
                    'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f})  '
                    'Acc@1: {top1.val:>7.3f} ({top1.avg:>7.3f})  '
                    'Acc@5: {top5.val:>7.3f} ({top5.avg:>7.3f})'.format(
                        batch_idx, len(loader), batch_time=batch_time,
                        rate_avg=input.size(0) / batch_time.avg,
                        loss=losses, top1=top1, top5=top5))

    if real_labels is not None:
        # real labels mode replaces topk values at the end
        top1a, top5a = real_labels.get_accuracy(k=1), real_labels.get_accuracy(k=5)
    else:
        top1a, top5a = top1.avg, top5.avg
    results = OrderedDict(
        top1=round(top1a, 4), top1_err=round(100 - top1a, 4),
        top5=round(top5a, 4), top5_err=round(100 - top5a, 4),
        param_count=round(param_count / 1e6, 2),
        img_size=data_config['input_size'][-1],
        cropt_pct=crop_pct,
        interpolation=data_config['interpolation'])

    _logger.info(' * Acc@1 {:.3f} ({:.3f}) Acc@5 {:.3f} ({:.3f})'.format(
       results['top1'], results['top1_err'], results['top5'], results['top5_err']))

    return results

def main():
    setup_default_logging()
    args = parser.parse_args()
    model_cfgs = []
    model_names = []
    if os.path.isdir(args.checkpoint):
        # validate all checkpoints in a path with same model
        checkpoints = glob.glob(args.checkpoint + '/*.pth.tar')
        checkpoints += glob.glob(args.checkpoint + '/*.pth')
        model_names = list_models(args.model)
        model_cfgs = [(args.model, c) for c in sorted(checkpoints, key=natural_key)]
    else:
        if args.model == 'all':
            # validate all models in a list of names with pretrained checkpoints
            args.pretrained = True
            model_names = list_models(pretrained=True)
            model_cfgs = [(n, '') for n in model_names]
        elif not is_model(args.model):
            # model name doesn't exist, try as wildcard filter
            model_names = list_models(args.model)
            model_cfgs = [(n, '') for n in model_names]

    if len(model_cfgs):
        results_file = args.results_file or './results-all.csv'
        _logger.info('Running bulk validation on these pretrained models: {}'.format(', '.join(model_names)))
        results = []
        try:
            start_batch_size = args.batch_size
            for m, c in model_cfgs:
                batch_size = start_batch_size
                args.model = m
                args.checkpoint = c
                result = OrderedDict(model=args.model)
                r = {}
                while not r and batch_size >= args.num_gpu:
                    torch.cuda.empty_cache()
                    try:
                        args.batch_size = batch_size
                        print('Validating with batch size: %d' % args.batch_size)
                        r = validate(args)
                    except RuntimeError as e:
                        if batch_size <= args.num_gpu:
                            print("Validation failed with no ability to reduce batch size. Exiting.")
                            raise e
                        batch_size = max(batch_size // 2, args.num_gpu)
                        print("Validation failed, reducing batch size by 50%")
                result.update(r)
                if args.checkpoint:
                    result['checkpoint'] = args.checkpoint
                results.append(result)
        except KeyboardInterrupt as e:
            pass
        results = sorted(results, key=lambda x: x['top1'], reverse=True)
        if len(results):
            write_results(results_file, results)
    else:
        validate(args)

def write_results(results_file, results):
    with open(results_file, mode='w') as cf:
        dw = csv.DictWriter(cf, fieldnames=results[0].keys())
        dw.writeheader()
        for r in results:
            dw.writerow(r)
        cf.flush()

if __name__ == '__main__':
    main()
<commit_msg>Add validation script update for using test_input_size in model default_cfgs<commit_after>#!/usr/bin/env python
""" ImageNet Validation Script

This is intended to be a lean and easily modifiable ImageNet validation script for evaluating pretrained
models or training checkpoints against ImageNet or similarly organized image datasets. It prioritizes
canonical PyTorch, standard Python style, and good performance. Repurpose as you see fit.

Hacked together by Ross Wightman (https://github.com/rwightman)
"""
import argparse
import os
import csv
import glob
import time
import logging
import torch
import torch.nn as nn
import torch.nn.parallel
from collections import OrderedDict
from contextlib import suppress

from timm.models import create_model, apply_test_time_pool, load_checkpoint, is_model, list_models
from timm.data import create_dataset, create_loader, resolve_data_config, RealLabelsImagenet
from timm.utils import accuracy, AverageMeter, natural_key, setup_default_logging, set_jit_legacy

has_apex = False
try:
    from apex import amp
    has_apex = True
except ImportError:
    pass

has_native_amp = False
try:
    if getattr(torch.cuda.amp, 'autocast') is not None:
        has_native_amp = True
except AttributeError:
    pass

torch.backends.cudnn.benchmark = True
_logger = logging.getLogger('validate')

parser = argparse.ArgumentParser(description='PyTorch ImageNet Validation')
parser.add_argument('data', metavar='DIR',
                    help='path to dataset')
parser.add_argument('--dataset', '-d', metavar='NAME', default='',
                    help='dataset type (default: ImageFolder/ImageTar if empty)')
parser.add_argument('--split', metavar='NAME', default='validation',
                    help='dataset split (default: validation)')
parser.add_argument('--model', '-m', metavar='NAME', default='dpn92',
                    help='model architecture (default: dpn92)')
parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
                    help='number of data loading workers (default: 2)')
parser.add_argument('-b', '--batch-size', default=256, type=int,
                    metavar='N', help='mini-batch size (default: 256)')
parser.add_argument('--img-size', default=None, type=int,
                    metavar='N', help='Input image dimension, uses model default if empty')
parser.add_argument('--input-size', default=None, nargs=3, type=int,
                    metavar='N N N', help='Input all image dimensions (d h w, e.g. --input-size 3 224 224), uses model default if empty')
parser.add_argument('--crop-pct', default=None, type=float,
                    metavar='N', help='Input image center crop pct')
parser.add_argument('--mean', type=float, nargs='+', default=None, metavar='MEAN',
                    help='Override mean pixel value of dataset')
parser.add_argument('--std', type=float,  nargs='+', default=None, metavar='STD',
                    help='Override std deviation of of dataset')
parser.add_argument('--interpolation', default='', type=str, metavar='NAME',
                    help='Image resize interpolation type (overrides model)')
parser.add_argument('--num-classes', type=int, default=None,
                    help='Number classes in dataset')
parser.add_argument('--class-map', default='', type=str, metavar='FILENAME',
                    help='path to class to idx mapping file (default: "")')
parser.add_argument('--gp', default=None, type=str, metavar='POOL',
                    help='Global pool type, one of (fast, avg, max, avgmax, avgmaxc). Model default if None.')
parser.add_argument('--log-freq', default=10, type=int,
                    metavar='N', help='batch logging frequency (default: 10)')
parser.add_argument('--checkpoint', default='', type=str, metavar='PATH',
                    help='path to latest checkpoint (default: none)')
parser.add_argument('--pretrained', dest='pretrained', action='store_true',
                    help='use pre-trained model')
parser.add_argument('--num-gpu', type=int, default=1,
                    help='Number of GPUS to use')
parser.add_argument('--no-test-pool', dest='no_test_pool', action='store_true',
                    help='disable test time pool')
parser.add_argument('--no-prefetcher', action='store_true', default=False,
                    help='disable fast prefetcher')
parser.add_argument('--pin-mem', action='store_true', default=False,
                    help='Pin CPU memory in DataLoader for more efficient (sometimes) transfer to GPU.')
parser.add_argument('--channels-last', action='store_true', default=False,
                    help='Use channels_last memory layout')
parser.add_argument('--amp', action='store_true', default=False,
                    help='Use AMP mixed precision. Defaults to Apex, fallback to native Torch AMP.')
parser.add_argument('--apex-amp', action='store_true', default=False,
                    help='Use NVIDIA Apex AMP mixed precision')
parser.add_argument('--native-amp', action='store_true', default=False,
                    help='Use Native Torch AMP mixed precision')
parser.add_argument('--tf-preprocessing', action='store_true', default=False,
                    help='Use Tensorflow preprocessing pipeline (require CPU TF installed')
parser.add_argument('--use-ema', dest='use_ema', action='store_true',
                    help='use ema version of weights if present')
parser.add_argument('--torchscript', dest='torchscript', action='store_true',
                    help='convert model torchscript for inference')
parser.add_argument('--legacy-jit', dest='legacy_jit', action='store_true',
                    help='use legacy jit mode for pytorch 1.5/1.5.1/1.6 to get back fusion performance')
parser.add_argument('--results-file', default='', type=str, metavar='FILENAME',
                    help='Output csv file for validation results (summary)')
parser.add_argument('--real-labels', default='', type=str, metavar='FILENAME',
                    help='Real labels JSON file for imagenet evaluation')
parser.add_argument('--valid-labels', default='', type=str, metavar='FILENAME',
                    help='Valid label indices txt file for validation of partial label space')

def validate(args):
    # might as well try to validate something
    args.pretrained = args.pretrained or not args.checkpoint
    args.prefetcher = not args.no_prefetcher
    amp_autocast = suppress  # do nothing
    if args.amp:
        if has_apex:
            args.apex_amp = True
        elif has_native_amp:
            args.native_amp = True
        else:
            _logger.warning("Neither APEX or Native Torch AMP is available, using FP32.")
    assert not args.apex_amp or not args.native_amp, "Only one AMP mode should be set."
    if args.native_amp:
        amp_autocast = torch.cuda.amp.autocast

    if args.legacy_jit:
        set_jit_legacy()

    # create model
    model = create_model(
        args.model,
        pretrained=args.pretrained,
        num_classes=args.num_classes,
        in_chans=3,
        global_pool=args.gp,
        scriptable=args.torchscript)
    if args.num_classes is None:
        assert hasattr(model, 'num_classes'), 'Model must have `num_classes` attr if not set on cmd line/config.'
        args.num_classes = model.num_classes

    if args.checkpoint:
        load_checkpoint(model, args.checkpoint, args.use_ema)

    param_count = sum([m.numel() for m in model.parameters()])
    _logger.info('Model %s created, param count: %d' % (args.model, param_count))

    data_config = resolve_data_config(vars(args), model=model, use_test_size=True)
    test_time_pool = False
    if not args.no_test_pool:
        model, test_time_pool = apply_test_time_pool(model, data_config, use_test_size=True)

    if args.torchscript:
        torch.jit.optimized_execution(True)
        model = torch.jit.script(model)

    model = model.cuda()
    if args.apex_amp:
        model = amp.initialize(model, opt_level='O1')

    if args.channels_last:
        model = model.to(memory_format=torch.channels_last)

    if args.num_gpu > 1:
        model = torch.nn.DataParallel(model, device_ids=list(range(args.num_gpu)))

    criterion = nn.CrossEntropyLoss().cuda()

    dataset = create_dataset(
        root=args.data, name=args.dataset, split=args.split,
        load_bytes=args.tf_preprocessing, class_map=args.class_map)

    if args.valid_labels:
        with open(args.valid_labels, 'r') as f:
            valid_labels = {int(line.rstrip()) for line in f}
            valid_labels = [i in valid_labels for i in range(args.num_classes)]
    else:
        valid_labels = None

    if args.real_labels:
        real_labels = RealLabelsImagenet(dataset.filenames(basename=True), real_json=args.real_labels)
    else:
        real_labels = None

    crop_pct = 1.0 if test_time_pool else data_config['crop_pct']
    loader = create_loader(
        dataset,
        input_size=data_config['input_size'],
        batch_size=args.batch_size,
        use_prefetcher=args.prefetcher,
        interpolation=data_config['interpolation'],
        mean=data_config['mean'],
        std=data_config['std'],
        num_workers=args.workers,
        crop_pct=crop_pct,
        pin_memory=args.pin_mem,
        tf_preprocessing=args.tf_preprocessing)

    batch_time = AverageMeter()
    losses = AverageMeter()
    top1 = AverageMeter()
    top5 = AverageMeter()

    model.eval()
    with torch.no_grad():
        # warmup, reduce variability of first batch time, especially for comparing torchscript vs non
        input = torch.randn((args.batch_size,) + data_config['input_size']).cuda()
        if args.channels_last:
            input = input.contiguous(memory_format=torch.channels_last)
        model(input)
        end = time.time()
        for batch_idx, (input, target) in enumerate(loader):
            if args.no_prefetcher:
                target = target.cuda()
                input = input.cuda()
            if args.channels_last:
                input = input.contiguous(memory_format=torch.channels_last)

            # compute output
            with amp_autocast():
                output = model(input)

            if valid_labels is not None:
                output = output[:, valid_labels]
            loss = criterion(output, target)

            if real_labels is not None:
                real_labels.add_result(output)

            # measure accuracy and record loss
            acc1, acc5 = accuracy(output.detach(), target, topk=(1, 5))
            losses.update(loss.item(), input.size(0))
            top1.update(acc1.item(), input.size(0))
            top5.update(acc5.item(), input.size(0))

            # measure elapsed time
            batch_time.update(time.time() - end)
            end = time.time()

            if batch_idx % args.log_freq == 0:
                _logger.info(
                    'Test: [{0:>4d}/{1}]  '
                    'Time: {batch_time.val:.3f}s ({batch_time.avg:.3f}s, {rate_avg:>7.2f}/s)  '
                    'Loss: {loss.val:>7.4f} ({loss.avg:>6.4f})  '
                    'Acc@1: {top1.val:>7.3f} ({top1.avg:>7.3f})  '
                    'Acc@5: {top5.val:>7.3f} ({top5.avg:>7.3f})'.format(
                        batch_idx, len(loader), batch_time=batch_time,
                        rate_avg=input.size(0) / batch_time.avg,
                        loss=losses, top1=top1, top5=top5))

    if real_labels is not None:
        # real labels mode replaces topk values at the end
        top1a, top5a = real_labels.get_accuracy(k=1), real_labels.get_accuracy(k=5)
    else:
        top1a, top5a = top1.avg, top5.avg
    results = OrderedDict(
        top1=round(top1a, 4), top1_err=round(100 - top1a, 4),
        top5=round(top5a, 4), top5_err=round(100 - top5a, 4),
        param_count=round(param_count / 1e6, 2),
        img_size=data_config['input_size'][-1],
        cropt_pct=crop_pct,
        interpolation=data_config['interpolation'])

    _logger.info(' * Acc@1 {:.3f} ({:.3f}) Acc@5 {:.3f} ({:.3f})'.format(
       results['top1'], results['top1_err'], results['top5'], results['top5_err']))

    return results

def main():
    setup_default_logging()
    args = parser.parse_args()
    model_cfgs = []
    model_names = []
    if os.path.isdir(args.checkpoint):
        # validate all checkpoints in a path with same model
        checkpoints = glob.glob(args.checkpoint + '/*.pth.tar')
        checkpoints += glob.glob(args.checkpoint + '/*.pth')
        model_names = list_models(args.model)
        model_cfgs = [(args.model, c) for c in sorted(checkpoints, key=natural_key)]
    else:
        if args.model == 'all':
            # validate all models in a list of names with pretrained checkpoints
            args.pretrained = True
            model_names = list_models(pretrained=True)
            model_cfgs = [(n, '') for n in model_names]
        elif not is_model(args.model):
            # model name doesn't exist, try as wildcard filter
            model_names = list_models(args.model)
            model_cfgs = [(n, '') for n in model_names]

    if len(model_cfgs):
        results_file = args.results_file or './results-all.csv'
        _logger.info('Running bulk validation on these pretrained models: {}'.format(', '.join(model_names)))
        results = []
        try:
            start_batch_size = args.batch_size
            for m, c in model_cfgs:
                batch_size = start_batch_size
                args.model = m
                args.checkpoint = c
                result = OrderedDict(model=args.model)
                r = {}
                while not r and batch_size >= args.num_gpu:
                    torch.cuda.empty_cache()
                    try:
                        args.batch_size = batch_size
                        print('Validating with batch size: %d' % args.batch_size)
                        r = validate(args)
                    except RuntimeError as e:
                        if batch_size <= args.num_gpu:
                            print("Validation failed with no ability to reduce batch size. Exiting.")
                            raise e
                        batch_size = max(batch_size // 2, args.num_gpu)
                        print("Validation failed, reducing batch size by 50%")
                result.update(r)
                if args.checkpoint:
                    result['checkpoint'] = args.checkpoint
                results.append(result)
        except KeyboardInterrupt as e:
            pass
        results = sorted(results, key=lambda x: x['top1'], reverse=True)
        if len(results):
            write_results(results_file, results)
    else:
        validate(args)

def write_results(results_file, results):
    with open(results_file, mode='w') as cf:
        dw = csv.DictWriter(cf, fieldnames=results[0].keys())
        dw.writeheader()
        for r in results:
            dw.writerow(r)
        cf.flush()

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


<commit_before>import numpy, sys, time, pylab, math, re
from astropy.io import fits as pyfits
from numpy import *
from pylab import *
from matplotlib import *
from math import *
import kepio, kepmsg, kepkey, kepfit, kepstat, kepfourier, keplab

def kepdynamic(infile,outfile,fcol,pmin,pmax,nfreq,deltat,nslice,
          plot,plotscale,cmap,clobber,verbose,logfile,status,cmdLine=False):

# startup parameters

    status = 0
    labelsize = 24
    ticksize = 16
    xsize = 12
    ysize = 6
    lcolor = '#0000ff'
    lwidth = 1.0
    fcolor = '#ffff00'
    falpha = 0.2
    numpy.seterr(all="ignore")

# log the call

    hashline = '----------------------------------------------------------------------------'
    kepmsg.log(logfile,hashline,verbose)
    call = 'KEPDYNAMIC -- '
    call += 'infile='+infile+' '
    call += 'outfile='+outfile+' '
    call += 'fcol='+fcol+' '
    call += 'pmin='+str(pmin)+' '
    call += 'pmax='+str(pmax)+' '
    call += 'nfreq='+str(nfreq)+' '
    call += 'deltat='+str(deltat)+' '
    call += 'nslice='+str(nslice)+' '
    plotit = 'n'
    if (plot): plotit = 'y'
    call += 'plot='+plotit+ ' '
    call += 'plotscale='+plotscale+ ' '
    call += 'cmap='+str(cmap)+' '
    overwrite = 'n'
    if (clobber): overwrite = 'y'
    call += 'clobber='+overwrite+ ' '
    chatter = 'n'
    if (verbose): chatter = 'y'
    call += 'verbose='+chatter+' '
    call += 'logfile='+logfile
    kepmsg.log(logfile,call+'\n',verbose)

# start time

    kepmsg.clock('Start time is',logfile,verbose)

# test log file

    logfile = kepmsg.test(logfile)

# error checking

    if status == 0 and pmin >= pmax:
        message = 'ERROR -- KEPDYNAMIC: PMIN must be less than PMAX'
        status = kepmsg.err(logfile,message,verbose)

# clobber output file

    if clobber: status = kepio.clobber(outfile,logfile,verbose)
    if kepio.fileexists(outfile):
        message = 'ERROR -- KEPDYNAMIC: ' + outfile + ' exists. Use clobber'
        status = kepmsg.err(logfile,message,verbose)

# plot color map

    if status == 0 and cmap == 'browse':
        status = keplab.cmap_plot()

# open input file

    if status == 0:
	    instr, status = kepio.openfits(infile,'readonly',logfile,verbose)
    if status == 0:
        tstart, tstop, bjdref, cadence, status = kepio.timekeys(instr,infile,logfile,verbose,status)
    if status == 0:
        try:
            work = instr[0].header['FILEVER']
            cadenom = 1.0
        except:
            cadenom = cadence

# fudge non-compliant FITS keywords with no values

    if status == 0:
	    instr = kepkey.emptykeys(instr,file,logfile,verbose)

# read table columns

    if status == 0:
        barytime, status = kepio.readtimecol(infile,instr[1].data,logfile,verbose)
    if status == 0:
        signal, status = kepio.readfitscol(infile,instr[1].data,fcol,logfile,verbose)
    if status == 0:
        barytime = barytime + bjdref
        signal = signal / cadenom

# remove infinite data from time series

    if status == 0:
	    incols = [barytime, signal]
	    outcols = kepstat.removeinfinlc(signal, incols)
	    barytime = outcols[0]
	    signal = outcols[1]

# period to frequency conversion

    if status == 0:
        fmin = 1.0 / pmax
        fmax = 1.0 / pmin
        deltaf = (fmax - fmin) / nfreq

# determine bounds of time slices

    if status == 0:
        t1 = []; t2 = []
        dt = barytime[-1] - barytime[0]
        dt -= deltat
        if dt < 0:
            message = 'ERROR -- KEPDYNAMIC: time slices are larger than data range'
            status = kepmsg.err(logfile,message,verbose)
        ds = dt / (nslice - 1)
        for i in range(nslice):
            t1.append(barytime[0] + ds * float(i))
            t2.append(barytime[0] + deltat + ds * float(i))

# loop through time slices

    if status == 0:
        dynam = []
        for i in range(nslice):
            x = []; y = []
            for j in range(len(barytime)):
                if (barytime[j] >= t1[i] and barytime[j] <= t2[i]):
                    x.append(barytime[j])
                    y.append(signal[j])
            x = array(x,dtype='float64')
            y = array(y,dtype='float32')
            y = y - median(y)

# determine FT power

	    fr, power = kepfourier.ft(x,y,fmin,fmax,deltaf,False)
            for j in range(len(power)):
                dynam.append(power[j])
            print 'Timeslice: %.4f  Pmax: %.2E' % ((t2[i] + t1[i]) / 2, power.max())

# define shape of results array

        dynam = array(dynam,dtype='float64')
        dynam.shape = len(t1),len(power)

# write output file

    if status == 0:
        instr.append(ImageHDU())
        instr[-1].data = dynam.transpose()
        instr[-1].header.update('EXTNAME','DYNAMIC FT','extension name')
        instr[-1].header.update('WCSAXES',2,'number of WCS axes')
        instr[-1].header.update('CRPIX1',0.5,'reference pixel along axis 1')
        instr[-1].header.update('CRPIX2',0.5,'reference pixel along axis 2')
        instr[-1].header.update('CRVAL1',t1[0],'time at reference pixel (BJD)')
        instr[-1].header.update('CRVAL2',fmin,'frequency at reference pixel (1/day)')
        instr[-1].header.update('CDELT1',(barytime[-1] - barytime[0]) / nslice,
                                'pixel scale in dimension 1 (days)')
        instr[-1].header.update('CDELT2',deltaf,'pixel scale in dimension 2 (1/day)')
        instr[-1].header.update('CTYPE1','BJD','data type of dimension 1')
        instr[-1].header.update('CTYPE2','FREQUENCY','data type of dimension 2')
        instr.writeto(outfile)

# history keyword in output file

    if status == 0:
	    status = kepkey.history(call,instr[0],outfile,logfile,verbose)

# close input file

    if status == 0:
	    status = kepio.closefits(instr,logfile,verbose)

# clean up x-axis unit

    if status == 0:
	time0 = float(int(barytime[0] / 100) * 100.0)
	barytime = barytime - time0
	xlab = 'BJD $-$ %d' % time0

# image intensity min and max

    if status == 0:
        if 'rithmic' in plotscale:
            dynam = numpy.log10(dynam)
        elif 'sq' in plotscale:
            dynam = numpy.sqrt(dynam)
        elif 'logoflog' in plotscale:
            dynam = numpy.log10(numpy.abs(numpy.log10(dynam)))
#        dynam = -dynam
        nstat = 2; pixels = []
        for i in range(dynam.shape[0]):
            for j in range(dynam.shape[1]):
                pixels.append(dynam[i,j])
        pixels = array(sort(pixels),dtype=float32)
        if int(float(len(pixels)) * 0.1 + 0.5) > nstat:
            nstat = int(float(len(pixels)) * 0.1 + 0.5)
        zmin = median(pixels[:nstat])
        zmax = median(pixels[-1:])
        if isnan(zmax):
            zmax = median(pixels[-nstat/2:])
        if isnan(zmax):
            zmax = numpy.nanmax(pixels)

# plot power spectrum

    if status == 0 and plot:
        params = {'backend': 'png',
                  'axes.linewidth': 2.5,
                  'axes.labelsize': labelsize,
                  'axes.font': 'sans-serif',
                  'axes.fontweight' : 'bold',
                  'text.fontsize': 12,
                  'legend.fontsize': 12,
                  'xtick.labelsize': ticksize,
                  'ytick.labelsize': ticksize}
        rcParams.update(params)
        pylab.figure(1,figsize=[xsize,ysize])
        pylab.clf()
        pylab.axes([0.08,0.113,0.91,0.86])
        dynam = dynam.transpose()
        pylab.imshow(dynam,origin='lower',aspect='auto',cmap=cmap,vmin=zmin,vmax=zmax,
                     extent=[barytime[0],barytime[-1],fmin,fmax],interpolation='bilinear')
        xlabel(xlab, {'color' : 'k'})
        ylabel(r'Frequency (d$^{-1}$)', {'color' : 'k'})
        grid()
        pylab.savefig(re.sub('\.\S+','.png',outfile),dpi=100)

# render plot

        if cmdLine:
            pylab.show()
        else:
            pylab.ion()
            pylab.plot([])
            pylab.ioff()

    return status

## end time

    if (status == 0):
	    message = 'KEPDYNAMIC completed at'
    else:
	    message = '\nKEPDYNAMIC aborted at'
    kepmsg.clock(message,logfile,verbose)

# main
if '--shell' in sys.argv:
    import argparse

    parser = argparse.ArgumentParser(description='Construct a dynamic (time-dependent) power spectrum from Kepler time series data')
    parser.add_argument('--shell', action='store_true', help='Are we running from the shell?')

    parser.add_argument('infile', help='Name of input file', type=str)
    parser.add_argument('outfile', help='Name of FITS file to output', type=str)

    parser.add_argument('--datacol', default='SAP_FLUX', help='Name of data column to plot', type=str, dest='fcol')

    parser.add_argument('--pmin', default=0.1, help='Minimum search period [days]', type=float)
    parser.add_argument('--pmax', default=10., help='Maximum search period [days]', type=float)
    parser.add_argument('--nfreq', default=100, help='Number of frequency intervals', type=int)
    parser.add_argument('--deltat', default=10., help='Length of time slice [days]',type=float)
    parser.add_argument('--nslice', default=10., help='Number of time slices', type=int)

    parser.add_argument('--plot', action='store_true', help='Plot result?')
    parser.add_argument('--plotscale', default='logarithmic', help='type of image intensity scale',
        type=str, choices=['linear','logarithmic','squareroot'])
    parser.add_argument('--cmap', default='PuBu', help='image colormap', type=str)

    parser.add_argument('--clobber', action='store_true', help='Overwrite output file?')
    parser.add_argument('--verbose', action='store_true', help='Write to a log file?')
    parser.add_argument('--logfile', '-l', help='Name of ascii log file', default='kepcotrend.log', dest='logfile', type=str)
    parser.add_argument('--status', '-e', help='Exit status (0=good)', default=0, dest='status', type=int)

    args = parser.parse_args()

    cmdLine=True

    kepynamic(args.infile, args.outfile, args.fcol, args.pmin, args.pmax, args.nfreq, args.deltat, args.nslice,
          args.plot,args.plotscale,args.cmap,args.clobber,args.verbose,args.logfile,args.status,cmdLine)

else:
    from pyraf import iraf
    parfile = iraf.osfn("kepler$kepdynamic.par")
    t = iraf.IrafTaskFactory(taskname="kepdynamic", value=parfile, function=kepdynamic)
<commit_msg>fix kepdynamic<commit_after>import sys, time, math, re
from astropy.io import fits as pyfits
from matplotlib import pyplot as plt
import numpy as np
import kepio, kepmsg, kepkey, kepfit, kepstat, kepfourier, keplab

def kepdynamic(infile,outfile,fcol,pmin,pmax,nfreq,deltat,nslice,
               plot,plotscale,cmap,clobber,verbose,logfile,status,cmdLine=False):

# startup parameters

    status = 0
    labelsize = 24
    ticksize = 16
    xsize = 12
    ysize = 6
    lcolor = '#0000ff'
    lwidth = 1.0
    fcolor = '#ffff00'
    falpha = 0.2
    np.seterr(all="ignore")

# log the call

    hashline = '----------------------------------------------------------------------------'
    kepmsg.log(logfile,hashline,verbose)
    call = 'KEPDYNAMIC -- '
    call += 'infile='+infile+' '
    call += 'outfile='+outfile+' '
    call += 'fcol='+fcol+' '
    call += 'pmin='+str(pmin)+' '
    call += 'pmax='+str(pmax)+' '
    call += 'nfreq='+str(nfreq)+' '
    call += 'deltat='+str(deltat)+' '
    call += 'nslice='+str(nslice)+' '
    plotit = 'n'
    if (plot): plotit = 'y'
    call += 'plot='+plotit+ ' '
    call += 'plotscale='+plotscale+ ' '
    call += 'cmap='+str(cmap)+' '
    overwrite = 'n'
    if (clobber): overwrite = 'y'
    call += 'clobber='+overwrite+ ' '
    chatter = 'n'
    if (verbose): chatter = 'y'
    call += 'verbose='+chatter+' '
    call += 'logfile='+logfile
    kepmsg.log(logfile,call+'\n',verbose)

# start time

    kepmsg.clock('Start time is',logfile,verbose)

# test log file

    logfile = kepmsg.test(logfile)

# error checking

    if status == 0 and pmin >= pmax:
        message = 'ERROR -- KEPDYNAMIC: PMIN must be less than PMAX'
        status = kepmsg.err(logfile,message,verbose)

# clobber output file

    if clobber: status = kepio.clobber(outfile,logfile,verbose)
    if kepio.fileexists(outfile):
        message = 'ERROR -- KEPDYNAMIC: ' + outfile + ' exists. Use clobber'
        status = kepmsg.err(logfile,message,verbose)

# plot color map

    if status == 0 and cmap == 'browse':
        status = keplab.cmap_plot()

# open input file

    if status == 0:
	    instr, status = kepio.openfits(infile,'readonly',logfile,verbose)
    if status == 0:
        tstart, tstop, bjdref, cadence, status = kepio.timekeys(instr,infile,logfile,verbose,status)
    if status == 0:
        try:
            work = instr[0].header['FILEVER']
            cadenom = 1.0
        except:
            cadenom = cadence

# fudge non-compliant FITS keywords with no values

    if status == 0:
	    instr = kepkey.emptykeys(instr,file,logfile,verbose)

# read table columns

    if status == 0:
        barytime, status = kepio.readtimecol(infile,instr[1].data,logfile,verbose)
    if status == 0:
        signal, status = kepio.readfitscol(infile,instr[1].data,fcol,logfile,verbose)
    if status == 0:
        barytime = barytime + bjdref
        signal = signal / cadenom

# remove infinite data from time series

    if status == 0:
	    incols = [barytime, signal]
	    outcols = kepstat.removeinfinlc(signal, incols)
	    barytime = outcols[0]
	    signal = outcols[1]

# period to frequency conversion

    if status == 0:
        fmin = 1.0 / pmax
        fmax = 1.0 / pmin
        deltaf = (fmax - fmin) / nfreq

# determine bounds of time slices

    if status == 0:
        t1 = []; t2 = []
        dt = barytime[-1] - barytime[0]
        dt -= deltat
        if dt < 0:
            message = 'ERROR -- KEPDYNAMIC: time slices are larger than data range'
            status = kepmsg.err(logfile,message,verbose)
        ds = dt / (nslice - 1)
        for i in range(nslice):
            t1.append(barytime[0] + ds * float(i))
            t2.append(barytime[0] + deltat + ds * float(i))

# loop through time slices

    if status == 0:
        dynam = []
        for i in range(nslice):
            x = []; y = []
            for j in range(len(barytime)):
                if (barytime[j] >= t1[i] and barytime[j] <= t2[i]):
                    x.append(barytime[j])
                    y.append(signal[j])
            x = np.array(x,dtype='float64')
            y = np.array(y,dtype='float')
            y = y - np.median(y)

# determine FT power

	    fr, power = kepfourier.ft(x,y,fmin,fmax,deltaf,False)
            for j in range(len(power)):
                dynam.append(power[j])
            print 'Timeslice: %.4f  Pmax: %.2E' % ((t2[i] + t1[i]) / 2, power.max())

# define shape of results array

        dynam = np.array(dynam,dtype='float64')
        dynam.shape = len(t1),len(power)

# write output file

    if status == 0:
        instr.append(pyfits.ImageHDU())
        instr[-1].data = dynam.transpose()
        instr[-1].header['EXTNAME'] = ('DYNAMIC FT','extension name')
        instr[-1].header['WCSAXES'] = (2,'number of WCS axes')
        instr[-1].header['CRPIX1' ] = (0.5,'reference pixel along axis 1')
        instr[-1].header['CRPIX2' ] = (0.5,'reference pixel along axis 2')
        instr[-1].header['CRVAL1' ] = (t1[0],'time at reference pixel (BJD)')
        instr[-1].header['CRVAL2' ] = (fmin,'frequency at reference pixel (1/day)')
        instr[-1].header['CDELT1' ] = ((barytime[-1] - barytime[0]) / nslice,
                        'pixel scale in dimension 1 (days)')
        instr[-1].header['CDELT2'] = (deltaf,'pixel scale in dimension 2 (1/day)')
        instr[-1].header['CTYPE1'] = ('BJD','data type of dimension 1')
        instr[-1].header['CTYPE2'] = ('FREQUENCY','data type of dimension 2')
        instr.writeto(outfile)

# history keyword in output file

    if status == 0:
	    status = kepkey.history(call,instr[0],outfile,logfile,verbose)

# close input file

    if status == 0:
	    status = kepio.closefits(instr,logfile,verbose)

# clean up x-axis unit

    if status == 0:
	time0 = float(int(barytime[0] / 100) * 100.0)
	barytime = barytime - time0
	xlab = 'BJD $-$ %d' % time0

# image intensity min and max

    if status == 0:
        if 'rithmic' in plotscale:
            dynam = np.log10(dynam)
        elif 'sq' in plotscale:
            dynam = np.sqrt(dynam)
        elif 'logoflog' in plotscale:
            dynam = np.log10(np.abs(np.log10(dynam)))
#        dynam = -dynam
        nstat = 2; pixels = []
        for i in range(dynam.shape[0]):
            for j in range(dynam.shape[1]):
                pixels.append(dynam[i,j])
        pixels = np.array(np.sort(pixels),dtype='float')
        if int(float(len(pixels)) * 0.1 + 0.5) > nstat:
            nstat = int(float(len(pixels)) * 0.1 + 0.5)
        zmin = np.median(pixels[:nstat])
        zmax = np.median(pixels[-1:])
        if np.isnan(zmax):
            zmax = np.median(pixels[-nstat/2:])
        if np.isnan(zmax):
            zmax = np.nanmax(pixels)

# plot power spectrum

    if status == 0 and plot:
        plt.figure(1,figsize=[xsize,ysize])
        plt.clf()
        plt.axes([0.08,0.113,0.91,0.86])
        dynam = dynam.transpose()
        plt.imshow(dynam,origin='lower',aspect='auto',cmap=cmap,vmin=zmin,vmax=zmax,
                     extent=[barytime[0],barytime[-1],fmin,fmax],interpolation='bilinear')
        plt.xlabel(xlab, {'color' : 'k'})
        plt.ylabel(r'Frequency (d$^{-1}$)', {'color' : 'k'})
        plt.grid()
        plt.savefig(re.sub('\.\S+','.png',outfile),dpi=100)

# render plot

        plt.show()

    return status

## end time

    if (status == 0):
	    message = 'KEPDYNAMIC completed at'
    else:
	    message = '\nKEPDYNAMIC aborted at'
    kepmsg.clock(message,logfile,verbose)

# main
if '--shell' in sys.argv:
    import argparse

    parser = argparse.ArgumentParser(description='Construct a dynamic (time-dependent) power spectrum from Kepler time series data')
    parser.add_argument('--shell', action='store_true', help='Are we running from the shell?')

    parser.add_argument('infile', help='Name of input file', type=str)
    parser.add_argument('outfile', help='Name of FITS file to output', type=str)

    parser.add_argument('--datacol', default='SAP_FLUX', help='Name of data column to plot', type=str, dest='fcol')

    parser.add_argument('--pmin', default=0.1, help='Minimum search period [days]', type=float)
    parser.add_argument('--pmax', default=10., help='Maximum search period [days]', type=float)
    parser.add_argument('--nfreq', default=100, help='Number of frequency intervals', type=int)
    parser.add_argument('--deltat', default=10., help='Length of time slice [days]',type=float)
    parser.add_argument('--nslice', default=10., help='Number of time slices', type=int)

    parser.add_argument('--plot', action='store_true', help='Plot result?')
    parser.add_argument('--plotscale', default='logarithmic', help='type of image intensity scale',
                        type=str, choices=['linear','logarithmic','squareroot'])
    parser.add_argument('--cmap', default='PuBu', help='image colormap', type=str)

    parser.add_argument('--clobber', action='store_true', help='Overwrite output file?')
    parser.add_argument('--verbose', action='store_true', help='Write to a log file?')
    parser.add_argument('--logfile', '-l', help='Name of ascii log file', default='kepcotrend.log',
                        dest='logfile', type=str)
    parser.add_argument('--status', '-e', help='Exit status (0=good)', default=0, dest='status', type=int)

    args = parser.parse_args()

    cmdLine=True

    kepynamic(args.infile, args.outfile, args.fcol, args.pmin, args.pmax, args.nfreq, args.deltat, args.nslice,
          args.plot,args.plotscale,args.cmap,args.clobber,args.verbose,args.logfile,args.status,cmdLine)

else:
    from pyraf import iraf
    parfile = iraf.osfn("kepler$kepdynamic.par")
    t = iraf.IrafTaskFactory(taskname="kepdynamic", value=parfile, function=kepdynamic)
<|endoftext|>


<commit_before>import os
import mimetypes

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

from django.conf import settings
from django.core.files.base import File
from django.core.files.storage import Storage
from django.core.exceptions import ImproperlyConfigured, SuspiciousOperation
from django.utils.encoding import force_unicode, smart_str

try:
    from boto.s3.connection import S3Connection
    from boto.exception import S3ResponseError
    from boto.s3.key import Key
except ImportError:
    raise ImproperlyConfigured, "Could not load Boto's S3 bindings.\
    \nSee http://code.google.com/p/boto/"

ACCESS_KEY_NAME     = getattr(settings, 'AWS_ACCESS_KEY_ID', None)
SECRET_KEY_NAME     = getattr(settings, 'AWS_SECRET_ACCESS_KEY', None)
HEADERS             = getattr(settings, 'AWS_HEADERS', {})
STORAGE_BUCKET_NAME = getattr(settings, 'AWS_STORAGE_BUCKET_NAME', None)
AUTO_CREATE_BUCKET  = getattr(settings, 'AWS_AUTO_CREATE_BUCKET', True)
DEFAULT_ACL         = getattr(settings, 'AWS_DEFAULT_ACL', 'public-read')
BUCKET_ACL          = getattr(settings, 'AWS_BUCKET_ACL', DEFAULT_ACL)
QUERYSTRING_AUTH    = getattr(settings, 'AWS_QUERYSTRING_AUTH', True)
QUERYSTRING_EXPIRE  = getattr(settings, 'AWS_QUERYSTRING_EXPIRE', 3600)
REDUCED_REDUNDANCY  = getattr(settings, 'AWS_REDUCED_REDUNDANCY', False)
LOCATION            = getattr(settings, 'AWS_LOCATION', '')
CUSTOM_DOMAIN       = getattr(settings, 'AWS_S3_CUSTOM_DOMAIN', None)
SECURE_URLS         = getattr(settings, 'AWS_S3_SECURE_URLS', True)
FILE_NAME_CHARSET   = getattr(settings, 'AWS_S3_FILE_NAME_CHARSET', 'utf-8')
FILE_OVERWRITE      = getattr(settings, 'AWS_S3_FILE_OVERWRITE', True)
IS_GZIPPED          = getattr(settings, 'AWS_IS_GZIPPED', False)
GZIP_CONTENT_TYPES  = getattr(settings, 'GZIP_CONTENT_TYPES', (
    'text/css',
    'application/javascript',
    'application/x-javascript'
))

if IS_GZIPPED:
    from gzip import GzipFile

def safe_join(base, *paths):
    """
    A version of django.utils._os.safe_join for S3 paths.

    Joins one or more path components to the base path component intelligently.
    Returns a normalized version of the final path.

    The final path must be located inside of the base path component (otherwise
    a ValueError is raised).

    Paths outside the base path indicate a possible security sensitive operation.
    """
    from urlparse import urljoin
    base_path = force_unicode(base)
    paths = map(lambda p: force_unicode(p), paths)
    final_path = urljoin(base_path + ("/" if not base_path.endswith("/") else ""), *paths)
    # Ensure final_path starts with base_path and that the next character after
    # the final path is '/' (or nothing, in which case final_path must be
    # equal to base_path).
    base_path_len = len(base_path)
    if not final_path.startswith(base_path) \
       or final_path[base_path_len:base_path_len+1] not in ('', '/'):
        raise ValueError('the joined path is located outside of the base path'
                         ' component')
    return final_path

class S3BotoStorage(Storage):
    """Amazon Simple Storage Service using Boto"""

    def __init__(self, bucket=STORAGE_BUCKET_NAME, access_key=None,
                       secret_key=None, bucket_acl=BUCKET_ACL, acl=DEFAULT_ACL, headers=HEADERS,
                       gzip=IS_GZIPPED, gzip_content_types=GZIP_CONTENT_TYPES,
                       querystring_auth=QUERYSTRING_AUTH, querystring_expire=QUERYSTRING_EXPIRE,
                       reduced_redundancy=REDUCED_REDUNDANCY,
                       custom_domain=CUSTOM_DOMAIN, secure_urls=SECURE_URLS,
                       location=LOCATION, file_name_charset=FILE_NAME_CHARSET):
        self.bucket_acl = bucket_acl
        self.bucket_name = bucket
        self.acl = acl
        self.headers = headers
        self.gzip = gzip
        self.gzip_content_types = gzip_content_types
        self.querystring_auth = querystring_auth
        self.querystring_expire = querystring_expire
        self.reduced_redundancy = reduced_redundancy
        self.custom_domain = custom_domain
        self.secure_urls = secure_urls
        self.location = location or ''
        self.location = self.location.lstrip('/')
        self.file_name_charset = file_name_charset

        if not access_key and not secret_key:
             access_key, secret_key = self._get_access_keys()

        self.connection = S3Connection(access_key, secret_key)

    @property
    def bucket(self):
        if not hasattr(self, '_bucket'):
            self._bucket = self._get_or_create_bucket(self.bucket_name)
        return self._bucket

    def _get_access_keys(self):
        access_key = ACCESS_KEY_NAME
        secret_key = SECRET_KEY_NAME
        if (access_key or secret_key) and (not access_key or not secret_key):
            access_key = os.environ.get(ACCESS_KEY_NAME)
            secret_key = os.environ.get(SECRET_KEY_NAME)

        if access_key and secret_key:
            # Both were provided, so use them
            return access_key, secret_key

        return None, None

    def _get_or_create_bucket(self, name):
        """Retrieves a bucket if it exists, otherwise creates it."""
        try:
            return self.connection.get_bucket(name)
        except S3ResponseError, e:
            if AUTO_CREATE_BUCKET:
                bucket = self.connection.create_bucket(name)
                bucket.set_acl(self.bucket_acl)
                return bucket
            raise ImproperlyConfigured, ("Bucket specified by "
            "AWS_STORAGE_BUCKET_NAME does not exist. Buckets can be "
            "automatically created by setting AWS_AUTO_CREATE_BUCKET=True")

    def _clean_name(self, name):
        # Useful for windows' paths
        return os.path.normpath(name).replace('\\', '/')

    def _normalize_name(self, name):
        try:
            return safe_join(self.location, name).lstrip('/')
        except ValueError:
            raise SuspiciousOperation("Attempted access to '%s' denied." % name)

    def _encode_name(self, name):
        return smart_str(name, encoding=self.file_name_charset)

    def _compress_content(self, content):
        """Gzip a given string."""
        zbuf = StringIO()
        zfile = GzipFile(mode='wb', compresslevel=6, fileobj=zbuf)
        zfile.write(content.read())
        zfile.close()
        content.file = zbuf
        return content

    def _open(self, name, mode='rb'):
        name = self._normalize_name(self._clean_name(name))
        f = S3BotoStorageFile(name, mode, self)
        if not f.key:
            raise IOError('File does not exist: %s' % name)
        return f

    def _save(self, name, content):
        cleaned_name = self._clean_name(name)
        name = self._normalize_name(cleaned_name)
        headers = self.headers
        content_type = getattr(content,'content_type', mimetypes.guess_type(name)[0] or Key.DefaultContentType)

        if self.gzip and content_type in self.gzip_content_types:
            content = self._compress_content(content)
            headers.update({'Content-Encoding': 'gzip'})

        content.name = cleaned_name
        k = self.bucket.get_key(self._encode_name(name))
        if not k:
            k = self.bucket.new_key(self._encode_name(name))

        k.set_metadata('Content-Type',content_type)
        k.set_contents_from_file(content, headers=headers, policy=self.acl,
                                 reduced_redundancy=self.reduced_redundancy)
        return cleaned_name

    def delete(self, name):
        name = self._normalize_name(self._clean_name(name))
        self.bucket.delete_key(self._encode_name(name))

    def exists(self, name):
        name = self._normalize_name(self._clean_name(name))
        k = self.bucket.new_key(self._encode_name(name))
        return k.exists()

    def listdir(self, name):
        name = self._normalize_name(self._clean_name(name))
        dirlist = self.bucket.list(self._encode_name(name))
        files = []
        dirs = set()
        base_parts = name.split("/") if name else []
        for item in dirlist:
            parts = item.name.split("/")
            parts = parts[len(base_parts):]
            if len(parts) == 1:
                # File
                files.append(parts[0])
            elif len(parts) > 1:
                # Directory
                dirs.add(parts[0])
        return list(dirs),files

    def size(self, name):
        name = self._normalize_name(self._clean_name(name))
        return self.bucket.get_key(self._encode_name(name)).size

    def url(self, name):
        name = self._normalize_name(self._clean_name(name))
        if self.custom_domain:
            return "%s://%s/%s" % ('https' if self.secure_urls else 'http', self.custom_domain, name)
        else:
            return self.connection.generate_url(self.querystring_expire, method='GET', \
                    bucket=self.bucket.name, key=self._encode_name(name), query_auth=self.querystring_auth, \
                    force_http=not self.secure_urls)

    def get_available_name(self, name):
        """ Overwrite existing file with the same name. """
        if FILE_OVERWRITE:
            name = self._clean_name(name)
            return name
        return super(S3BotoStorage, self).get_available_name(name)

class S3BotoStorageFile(File):
    def __init__(self, name, mode, storage):
        self._storage = storage
        self.name = name[len(self._storage.location):].lstrip('/')
        self._mode = mode
        self.key = storage.bucket.get_key(self._storage._encode_name(name))
        self._is_dirty = False
        self._file = None

    @property
    def size(self):
        return self.key.size

    @property
    def file(self):
        if self._file is None:
            self._file = StringIO()
            if 'r' in self._mode:
                self._is_dirty = False
                self.key.get_contents_to_file(self._file)
                self._file.seek(0)
        return self._file

    def read(self, *args, **kwargs):
        if 'r' not in self._mode:
            raise AttributeError("File was not opened in read mode.")
        return super(S3BotoStorageFile, self).read(*args, **kwargs)

    def write(self, *args, **kwargs):
        if 'w' not in self._mode:
            raise AttributeError("File was opened for read-only access.")
        self._is_dirty = True
        return super(S3BotoStorageFile, self).write(*args, **kwargs)

    def close(self):
        if self._is_dirty:
            self.key.set_contents_from_file(self._file, headers=self._storage.headers, policy=self._storage.acl)
        self.key.close()
<commit_msg>Added AWS_PRELOAD_METADATA options for S3BotoStorage backend to preload bucket metadata.<commit_after>import os
import mimetypes

try:
    from cStringIO import StringIO
except ImportError:
    from StringIO import StringIO

from django.conf import settings
from django.core.files.base import File
from django.core.files.storage import Storage
from django.core.exceptions import ImproperlyConfigured, SuspiciousOperation
from django.utils.encoding import force_unicode, smart_str

try:
    from boto.s3.connection import S3Connection
    from boto.exception import S3ResponseError
    from boto.s3.key import Key
except ImportError:
    raise ImproperlyConfigured("Could not load Boto's S3 bindings.\n"
                               "See http://code.google.com/p/boto/")

ACCESS_KEY_NAME     = getattr(settings, 'AWS_ACCESS_KEY_ID', None)
SECRET_KEY_NAME     = getattr(settings, 'AWS_SECRET_ACCESS_KEY', None)
HEADERS             = getattr(settings, 'AWS_HEADERS', {})
STORAGE_BUCKET_NAME = getattr(settings, 'AWS_STORAGE_BUCKET_NAME', None)
AUTO_CREATE_BUCKET  = getattr(settings, 'AWS_AUTO_CREATE_BUCKET', True)
DEFAULT_ACL         = getattr(settings, 'AWS_DEFAULT_ACL', 'public-read')
BUCKET_ACL          = getattr(settings, 'AWS_BUCKET_ACL', DEFAULT_ACL)
QUERYSTRING_AUTH    = getattr(settings, 'AWS_QUERYSTRING_AUTH', True)
QUERYSTRING_EXPIRE  = getattr(settings, 'AWS_QUERYSTRING_EXPIRE', 3600)
REDUCED_REDUNDANCY  = getattr(settings, 'AWS_REDUCED_REDUNDANCY', False)
LOCATION            = getattr(settings, 'AWS_LOCATION', '')
CUSTOM_DOMAIN       = getattr(settings, 'AWS_S3_CUSTOM_DOMAIN', None)
SECURE_URLS         = getattr(settings, 'AWS_S3_SECURE_URLS', True)
FILE_NAME_CHARSET   = getattr(settings, 'AWS_S3_FILE_NAME_CHARSET', 'utf-8')
FILE_OVERWRITE      = getattr(settings, 'AWS_S3_FILE_OVERWRITE', True)
IS_GZIPPED          = getattr(settings, 'AWS_IS_GZIPPED', False)
PRELOAD_METADATA    = getattr(settings, 'AWS_PRELOAD_METADATA', False)
GZIP_CONTENT_TYPES  = getattr(settings, 'GZIP_CONTENT_TYPES', (
    'text/css',
    'application/javascript',
    'application/x-javascript'
))

if IS_GZIPPED:
    from gzip import GzipFile

def safe_join(base, *paths):
    """
    A version of django.utils._os.safe_join for S3 paths.

    Joins one or more path components to the base path component intelligently.
    Returns a normalized version of the final path.

    The final path must be located inside of the base path component (otherwise
    a ValueError is raised).

    Paths outside the base path indicate a possible security sensitive operation.
    """
    from urlparse import urljoin
    base_path = force_unicode(base)
    paths = map(lambda p: force_unicode(p), paths)
    final_path = urljoin(base_path + ("/" if not base_path.endswith("/") else ""), *paths)
    # Ensure final_path starts with base_path and that the next character after
    # the final path is '/' (or nothing, in which case final_path must be
    # equal to base_path).
    base_path_len = len(base_path)
    if not final_path.startswith(base_path) \
       or final_path[base_path_len:base_path_len+1] not in ('', '/'):
        raise ValueError('the joined path is located outside of the base path'
                         ' component')
    return final_path

class S3BotoStorage(Storage):
    """Amazon Simple Storage Service using Boto"""

    def __init__(self, bucket=STORAGE_BUCKET_NAME, access_key=None,
                       secret_key=None, bucket_acl=BUCKET_ACL, acl=DEFAULT_ACL, headers=HEADERS,
                       gzip=IS_GZIPPED, gzip_content_types=GZIP_CONTENT_TYPES,
                       querystring_auth=QUERYSTRING_AUTH, querystring_expire=QUERYSTRING_EXPIRE,
                       reduced_redundancy=REDUCED_REDUNDANCY,
                       custom_domain=CUSTOM_DOMAIN, secure_urls=SECURE_URLS,
                       location=LOCATION, file_name_charset=FILE_NAME_CHARSET,
                       preload_metadata=PRELOAD_METADATA):
        self.bucket_acl = bucket_acl
        self.bucket_name = bucket
        self.acl = acl
        self.headers = headers
        self.preload_metadata = preload_metadata
        self.gzip = gzip
        self.gzip_content_types = gzip_content_types
        self.querystring_auth = querystring_auth
        self.querystring_expire = querystring_expire
        self.reduced_redundancy = reduced_redundancy
        self.custom_domain = custom_domain
        self.secure_urls = secure_urls
        self.location = location or ''
        self.location = self.location.lstrip('/')
        self.file_name_charset = file_name_charset

        if not access_key and not secret_key:
             access_key, secret_key = self._get_access_keys()

        self.connection = S3Connection(access_key, secret_key)
        self._entries = {}

    @property
    def bucket(self):
        if not hasattr(self, '_bucket'):
            self._bucket = self._get_or_create_bucket(self.bucket_name)
        return self._bucket

    @property
    def entries(self):
        if self.preload_metadata and not self._entries:
            self._entries = dict((self._decode_name(entry.key), entry)
                                for entry in self.bucket.list())
            print self._entries
        return self._entries

    def _get_access_keys(self):
        access_key = ACCESS_KEY_NAME
        secret_key = SECRET_KEY_NAME
        if (access_key or secret_key) and (not access_key or not secret_key):
            access_key = os.environ.get(ACCESS_KEY_NAME)
            secret_key = os.environ.get(SECRET_KEY_NAME)

        if access_key and secret_key:
            # Both were provided, so use them
            return access_key, secret_key

        return None, None

    def _get_or_create_bucket(self, name):
        """Retrieves a bucket if it exists, otherwise creates it."""
        try:
            return self.connection.get_bucket(name)
        except S3ResponseError, e:
            if AUTO_CREATE_BUCKET:
                bucket = self.connection.create_bucket(name)
                bucket.set_acl(self.bucket_acl)
                return bucket
            raise ImproperlyConfigured, ("Bucket specified by "
            "AWS_STORAGE_BUCKET_NAME does not exist. Buckets can be "
            "automatically created by setting AWS_AUTO_CREATE_BUCKET=True")

    def _clean_name(self, name):
        # Useful for windows' paths
        return os.path.normpath(name).replace('\\', '/')

    def _normalize_name(self, name):
        try:
            return safe_join(self.location, name).lstrip('/')
        except ValueError:
            raise SuspiciousOperation("Attempted access to '%s' denied." % name)

    def _encode_name(self, name):
        return smart_str(name, encoding=self.file_name_charset)

    def _decode_name(self, name):
        return force_unicode(name, encoding=self.file_name_charset)

    def _compress_content(self, content):
        """Gzip a given string."""
        zbuf = StringIO()
        zfile = GzipFile(mode='wb', compresslevel=6, fileobj=zbuf)
        zfile.write(content.read())
        zfile.close()
        content.file = zbuf
        return content

    def _open(self, name, mode='rb'):
        name = self._normalize_name(self._clean_name(name))
        f = S3BotoStorageFile(name, mode, self)
        if not f.key:
            raise IOError('File does not exist: %s' % name)
        return f

    def _save(self, name, content):
        cleaned_name = self._clean_name(name)
        name = self._normalize_name(cleaned_name)
        headers = self.headers
        content_type = getattr(content,'content_type', mimetypes.guess_type(name)[0] or Key.DefaultContentType)

        if self.gzip and content_type in self.gzip_content_types:
            content = self._compress_content(content)
            headers.update({'Content-Encoding': 'gzip'})

        content.name = cleaned_name
        k = self.bucket.get_key(self._encode_name(name))
        if not k:
            k = self.bucket.new_key(self._encode_name(name))

        k.set_metadata('Content-Type',content_type)
        k.set_contents_from_file(content, headers=headers, policy=self.acl,
                                 reduced_redundancy=self.reduced_redundancy)
        return cleaned_name

    def delete(self, name):
        name = self._normalize_name(self._clean_name(name))
        self.bucket.delete_key(self._encode_name(name))

    def exists(self, name):
        name = self._normalize_name(self._clean_name(name))
        if self.entries:
            return name in self.entries
        k = self.bucket.new_key(self._encode_name(name))
        return k.exists()

    def listdir(self, name):
        name = self._normalize_name(self._clean_name(name))
        dirlist = self.bucket.list(self._encode_name(name))
        files = []
        dirs = set()
        base_parts = name.split("/") if name else []
        for item in dirlist:
            parts = item.name.split("/")
            parts = parts[len(base_parts):]
            if len(parts) == 1:
                # File
                files.append(parts[0])
            elif len(parts) > 1:
                # Directory
                dirs.add(parts[0])
        return list(dirs),files

    def size(self, name):
        name = self._normalize_name(self._clean_name(name))
        if self.entries:
            entry = self.entries.get(name)
            if entry:
                return entry.size
            return 0
        return self.bucket.get_key(self._encode_name(name)).size

    def modified_time(self, name):
        try:
           from dateutil import parser, tz
        except ImportError:
            raise NotImplementedError()
        name = self._normalize_name(self._clean_name(name))
        entry = self.entries.get(name, self.bucket.get_key(self._encode_name(name)))
        # convert to string to date
        last_modified_date = parser.parse(entry.last_modified)
        # if the date has no timzone, assume UTC
        if last_modified_date.tzinfo == None:
            last_modified_date = last_modified_date.replace(tzinfo=tz.tzutc())
        # convert date to local time w/o timezone
        return last_modified_date.astimezone(tz.tzlocal()).replace(tzinfo=None)

    def url(self, name):
        name = self._normalize_name(self._clean_name(name))
        if self.custom_domain:
            return "%s://%s/%s" % ('https' if self.secure_urls else 'http', self.custom_domain, name)
        else:
            return self.connection.generate_url(self.querystring_expire, method='GET', \
                    bucket=self.bucket.name, key=self._encode_name(name), query_auth=self.querystring_auth, \
                    force_http=not self.secure_urls)

    def get_available_name(self, name):
        """ Overwrite existing file with the same name. """
        if FILE_OVERWRITE:
            name = self._clean_name(name)
            return name
        return super(S3BotoStorage, self).get_available_name(name)

class S3BotoStorageFile(File):
    def __init__(self, name, mode, storage):
        self._storage = storage
        self.name = name[len(self._storage.location):].lstrip('/')
        self._mode = mode
        self.key = storage.bucket.get_key(self._storage._encode_name(name))
        self._is_dirty = False
        self._file = None

    @property
    def size(self):
        return self.key.size

    @property
    def file(self):
        if self._file is None:
            self._file = StringIO()
            if 'r' in self._mode:
                self._is_dirty = False
                self.key.get_contents_to_file(self._file)
                self._file.seek(0)
        return self._file

    def read(self, *args, **kwargs):
        if 'r' not in self._mode:
            raise AttributeError("File was not opened in read mode.")
        return super(S3BotoStorageFile, self).read(*args, **kwargs)

    def write(self, *args, **kwargs):
        if 'w' not in self._mode:
            raise AttributeError("File was opened for read-only access.")
        self._is_dirty = True
        return super(S3BotoStorageFile, self).write(*args, **kwargs)

    def close(self):
        if self._is_dirty:
            self.key.set_contents_from_file(self._file, headers=self._storage.headers, policy=self._storage.acl)
        self.key.close()
<|endoftext|>