<commit_before># Copyright 2018 The TensorFlow Authors. 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.
# ==============================================================================
"""A tool to generate api_docs for TensorFlow2.

```
python generate2.py --output_dir=/tmp/out
```

Requires a local installation of:
  https://github.com/tensorflow/docs/tree/master/tools
  tf-nightly-2.0-preview
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from os import path
import textwrap

from absl import app
from absl import flags
from distutils.version import LooseVersion

import tensorflow as tf

from tensorflow_docs.api_generator import doc_controls
from tensorflow_docs.api_generator import doc_generator_visitor
from tensorflow_docs.api_generator import generate_lib
from tensorflow_docs.api_generator import parser

import tensorboard
import tensorflow_estimator
from tensorflow.python.util import tf_export
from tensorflow.python.util import tf_inspect

# Use tensorflow's `tf_inspect`, which is aware of `tf_decorator`.
parser.tf_inspect = tf_inspect

# `tf` has an `__all__` that doesn't list important things like `keras`.
# The doc generator recognizes `__all__` as the list of public symbols.
# So patch `tf.__all__` to list everything.
tf.__all__ = [item_name for item_name, value in tf_inspect.getmembers(tf)]

FLAGS = flags.FLAGS

flags.DEFINE_string(
    "code_url_prefix",
    "/code/stable/tensorflow",
    "A url to prepend to code paths when creating links to defining code")

flags.DEFINE_string(
    "output_dir", "/tmp/out",
    "A directory, where the docs will be output to.")

flags.DEFINE_bool("search_hints", True,
                  "Include meta-data search hints at the top of each file.")

flags.DEFINE_string("site_path", "",
                    "The prefix ({site-path}/api_docs/python/...) used in the "
                    "`_toc.yaml` and `_redirects.yaml` files")

if tf.__version__.startswith('1'):
  PRIVATE_MAP = {
      'tf.contrib.autograph': ['utils', 'operators'],
      'tf.test': ['mock'],
      'tf.contrib.estimator': ['python'],
  }

  DO_NOT_DESCEND_MAP = {
      'tf': ['cli', 'lib', 'wrappers'],
      'tf.contrib': [
          'compiler',
          'grid_rnn',
          # Block contrib.keras to de-clutter the docs
          'keras',
          'labeled_tensor',
          'quantization',
          'session_bundle',
          'slim',
          'solvers',
          'specs',
          'tensor_forest',
          'tensorboard',
          'testing',
          'tfprof',
      ],
      'tf.contrib.bayesflow': [
          'special_math', 'stochastic_gradient_estimators',
          'stochastic_variables'
      ],
      'tf.contrib.ffmpeg': ['ffmpeg_ops'],
      'tf.contrib.graph_editor': [
          'edit', 'match', 'reroute', 'subgraph', 'transform', 'select', 'util'
      ],
      'tf.contrib.keras': ['api', 'python'],
      'tf.contrib.layers': ['feature_column', 'summaries'],
      'tf.contrib.learn': [
          'datasets',
          'head',
          'graph_actions',
          'io',
          'models',
          'monitors',
          'ops',
          'preprocessing',
          'utils',
      ],
      'tf.contrib.util': ['loader'],
  }
else:
  PRIVATE_MAP = {'tf': ['python', 'core', 'compiler', 'examples', 'tools']}
  DO_NOT_DESCEND_MAP = {}
  tf.__doc__ = """
    ## TensorFlow 2.0 RC

    Caution:  This is a developer preview.  You will likely find some bugs,
    performance issues, and more, and we encourage you to tell us about them.
    We value your feedback!

    These docs were generated from the beta build of TensorFlow 2.0.

    You can install the exact version that was used to generate these docs
    with:

    ```
    pip install tensorflow==2.0.0-rc0
    ```
    """

_raw_ops_doc = textwrap.dedent("""\n
  Note: `tf.raw_ops` provides direct/low level access to all TensorFlow ops. See \
  [the RFC](https://github.com/tensorflow/community/blob/master/rfcs/20181225-tf-raw-ops.md)
  for details. Unless you are library writer, you likely do not need to use these
  ops directly.""")

if LooseVersion(tf.__version__) < LooseVersion('2'):
  tf.raw_ops.__doc__ = _raw_ops_doc
  tf.contrib.__doc__ = """
    Contrib module containing volatile or experimental code.

    Warning: The `tf.contrib` module will not be included in TensorFlow 2.0. Many
    of its submodules have been integrated into TensorFlow core, or spun-off into
    other projects like [`tensorflow_io`](https://github.com/tensorflow/io), or
    [`tensorflow_addons`](https://github.com/tensorflow/addons). For instructions
    on how to upgrade see the
    [Migration guide](https://www.tensorflow.org/beta/guide/migration_guide).
    """
else:
  tf.raw_ops.__doc__ += _raw_ops_doc

# The doc generator isn't aware of tf_export.
# So prefix the score tuples with -1 when this is the canonical name, +1
# otherwise. The generator chooses the name with the lowest score.
class TfExportAwareDocGeneratorVisitor(
    doc_generator_visitor.DocGeneratorVisitor):
  """A `tf_export` aware doc_visitor."""

  def _score_name(self, name):
    canonical = tf_export.get_canonical_name_for_symbol(self._index[name])

    canonical_score = 1
    if canonical is not None and name == "tf." + canonical:
      canonical_score = -1

    scores = super(TfExportAwareDocGeneratorVisitor, self)._score_name(name)
    return (canonical_score,) + scores

def _hide_layer_and_module_methods():
  """Hide methods and properties defined in the base classes of keras layers."""
  # __dict__ only sees attributes defined in *this* class, not on parent classes
  module_contents = list(tf.Module.__dict__.items())
  layer_contents = list(tf.keras.layers.Layer.__dict__.items())

  for name, obj in module_contents + layer_contents:
    if name == "__init__":
      continue

    if isinstance(obj, property):
      obj = obj.fget

    if isinstance(obj, (staticmethod, classmethod)):
      obj = obj.__func__

    try:
      doc_controls.do_not_doc_in_subclasses(obj)
    except AttributeError:
      pass

def build_docs(output_dir, code_url_prefix, search_hints=True):
  """Build api docs for tensorflow v2.

  Args:
    output_dir: A string path, where to put the files.
    code_url_prefix: prefix for "Defined in" links.
    search_hints: Bool. Include meta-data search hints at the top of each file.
  """
  _hide_layer_and_module_methods()

  try:
    doc_controls.do_not_generate_docs(tf.tools)
  except AttributeError:
    pass

  try:
    doc_controls.do_not_generate_docs(tf.compat.v1.pywrap_tensorflow)
  except AttributeError:
    pass

  try:
    doc_controls.do_not_generate_docs(tf.pywrap_tensorflow)
  except AttributeError:
    pass

  try:
    doc_controls.do_not_generate_docs(tf.flags)
  except AttributeError:
    pass

  base_dir = path.dirname(tf.__file__)

  base_dirs = (
      base_dir,
      # External packages base directories,
      path.dirname(tensorboard.__file__),
      path.dirname(tensorflow_estimator.__file__),
  )

  code_url_prefixes = (
      code_url_prefix,
      # External packages source repositories,
      "https://github.com/tensorflow/tensorboard/tree/master/tensorboard",
      "https://github.com/tensorflow/estimator/tree/master/tensorflow_estimator",
  )

  if LooseVersion(tf.__version__) < LooseVersion('2'):
    root_title = 'TensorFlow'
  elif LooseVersion(tf.__version__) >= LooseVersion('2'):
    root_title = 'TensorFlow 2.0'

  doc_generator = generate_lib.DocGenerator(
      root_title=root_title,
      py_modules=[("tf", tf)],
      base_dir=base_dirs,
      search_hints=search_hints,
      code_url_prefix=code_url_prefixes,
      site_path=FLAGS.site_path,
      visitor_cls=TfExportAwareDocGeneratorVisitor,
      private_map=PRIVATE_MAP,
      do_not_descend_map=DO_NOT_DESCEND_MAP)

  doc_generator.build(output_dir)

def main(argv):
  del argv
  build_docs(output_dir=FLAGS.output_dir,
             code_url_prefix=FLAGS.code_url_prefix,
             search_hints=FLAGS.search_hints)

if __name__ == "__main__":
  app.run(main)
<commit_msg>Since the packages in tensorflow moved to tensorflow_core, changing the copy.bara.sky rule to reflect that. Also removing compat.v2.compat.v1 aliases from api docs.<commit_after># Copyright 2018 The TensorFlow Authors. 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.
# ==============================================================================
"""A tool to generate api_docs for TensorFlow2.

```
python generate2.py --output_dir=/tmp/out
```

Requires a local installation of:
  https://github.com/tensorflow/docs/tree/master/tools
  tf-nightly-2.0-preview
"""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

from os import path
import textwrap

from absl import app
from absl import flags
from distutils.version import LooseVersion

import tensorflow as tf

from tensorflow_docs.api_generator import doc_controls
from tensorflow_docs.api_generator import doc_generator_visitor
from tensorflow_docs.api_generator import generate_lib
from tensorflow_docs.api_generator import parser

import tensorboard
import tensorflow_estimator
from tensorflow.python.util import tf_export
from tensorflow.python.util import tf_inspect

# Use tensorflow's `tf_inspect`, which is aware of `tf_decorator`.
parser.tf_inspect = tf_inspect

# `tf` has an `__all__` that doesn't list important things like `keras`.
# The doc generator recognizes `__all__` as the list of public symbols.
# So patch `tf.__all__` to list everything.
tf.__all__ = [item_name for item_name, value in tf_inspect.getmembers(tf)]

FLAGS = flags.FLAGS

flags.DEFINE_string(
    "code_url_prefix",
    "/code/stable/tensorflow",
    "A url to prepend to code paths when creating links to defining code")

flags.DEFINE_string(
    "output_dir", "/tmp/out",
    "A directory, where the docs will be output to.")

flags.DEFINE_bool("search_hints", True,
                  "Include meta-data search hints at the top of each file.")

flags.DEFINE_string("site_path", "",
                    "The prefix ({site-path}/api_docs/python/...) used in the "
                    "`_toc.yaml` and `_redirects.yaml` files")

if tf.__version__.startswith('1'):
  PRIVATE_MAP = {
      'tf.contrib.autograph': ['utils', 'operators'],
      'tf.test': ['mock'],
      'tf.contrib.estimator': ['python'],
  }

  DO_NOT_DESCEND_MAP = {
      'tf': ['cli', 'lib', 'wrappers'],
      'tf.contrib': [
          'compiler',
          'grid_rnn',
          # Block contrib.keras to de-clutter the docs
          'keras',
          'labeled_tensor',
          'quantization',
          'session_bundle',
          'slim',
          'solvers',
          'specs',
          'tensor_forest',
          'tensorboard',
          'testing',
          'tfprof',
      ],
      'tf.contrib.bayesflow': [
          'special_math', 'stochastic_gradient_estimators',
          'stochastic_variables'
      ],
      'tf.contrib.ffmpeg': ['ffmpeg_ops'],
      'tf.contrib.graph_editor': [
          'edit', 'match', 'reroute', 'subgraph', 'transform', 'select', 'util'
      ],
      'tf.contrib.keras': ['api', 'python'],
      'tf.contrib.layers': ['feature_column', 'summaries'],
      'tf.contrib.learn': [
          'datasets',
          'head',
          'graph_actions',
          'io',
          'models',
          'monitors',
          'ops',
          'preprocessing',
          'utils',
      ],
      'tf.contrib.util': ['loader'],
  }
else:
  PRIVATE_MAP = {'tf': ['python', 'core', 'compiler', 'examples', 'tools']}
  DO_NOT_DESCEND_MAP = {}
  tf.__doc__ = """
    ## TensorFlow 2.0 RC

    Caution:  This is a developer preview.  You will likely find some bugs,
    performance issues, and more, and we encourage you to tell us about them.
    We value your feedback!

    These docs were generated from the beta build of TensorFlow 2.0.

    You can install the exact version that was used to generate these docs
    with:

    ```
    pip install tensorflow==2.0.0-rc0
    ```
    """

_raw_ops_doc = textwrap.dedent("""\n
  Note: `tf.raw_ops` provides direct/low level access to all TensorFlow ops. See \
  [the RFC](https://github.com/tensorflow/community/blob/master/rfcs/20181225-tf-raw-ops.md)
  for details. Unless you are library writer, you likely do not need to use these
  ops directly.""")

if LooseVersion(tf.__version__) < LooseVersion('2'):
  tf.raw_ops.__doc__ = _raw_ops_doc
  tf.contrib.__doc__ = """
    Contrib module containing volatile or experimental code.

    Warning: The `tf.contrib` module will not be included in TensorFlow 2.0. Many
    of its submodules have been integrated into TensorFlow core, or spun-off into
    other projects like [`tensorflow_io`](https://github.com/tensorflow/io), or
    [`tensorflow_addons`](https://github.com/tensorflow/addons). For instructions
    on how to upgrade see the
    [Migration guide](https://www.tensorflow.org/beta/guide/migration_guide).
    """
else:
  tf.raw_ops.__doc__ += _raw_ops_doc

# The doc generator isn't aware of tf_export.
# So prefix the score tuples with -1 when this is the canonical name, +1
# otherwise. The generator chooses the name with the lowest score.
class TfExportAwareDocGeneratorVisitor(
    doc_generator_visitor.DocGeneratorVisitor):
  """A `tf_export` aware doc_visitor."""

  def _score_name(self, name):
    canonical = tf_export.get_canonical_name_for_symbol(self._index[name])

    canonical_score = 1
    if canonical is not None and name == "tf." + canonical:
      canonical_score = -1

    scores = super(TfExportAwareDocGeneratorVisitor, self)._score_name(name)
    return (canonical_score,) + scores

def _hide_layer_and_module_methods():
  """Hide methods and properties defined in the base classes of keras layers."""
  # __dict__ only sees attributes defined in *this* class, not on parent classes
  module_contents = list(tf.Module.__dict__.items())
  layer_contents = list(tf.keras.layers.Layer.__dict__.items())

  for name, obj in module_contents + layer_contents:
    if name == "__init__":
      continue

    if isinstance(obj, property):
      obj = obj.fget

    if isinstance(obj, (staticmethod, classmethod)):
      obj = obj.__func__

    try:
      doc_controls.do_not_doc_in_subclasses(obj)
    except AttributeError:
      pass

def build_docs(output_dir, code_url_prefix, search_hints=True):
  """Build api docs for tensorflow v2.

  Args:
    output_dir: A string path, where to put the files.
    code_url_prefix: prefix for "Defined in" links.
    search_hints: Bool. Include meta-data search hints at the top of each file.
  """
  _hide_layer_and_module_methods()

  try:
    doc_controls.do_not_generate_docs(tf.tools)
  except AttributeError:
    pass

  try:
    doc_controls.do_not_generate_docs(tf.compat.v1.pywrap_tensorflow)
  except AttributeError:
    pass

  try:
    doc_controls.do_not_generate_docs(tf.pywrap_tensorflow)
  except AttributeError:
    pass

  try:
    doc_controls.do_not_generate_docs(tf.flags)
  except AttributeError:
    pass

  base_dir = path.normpath(path.join(path.dirname(tf.__file__), "../.."))

  base_dirs = (
      base_dir,
      path.normpath(path.join(base_dir, "../../tensorflow")),
      path.dirname(tensorboard.__file__),
      path.dirname(tensorflow_estimator.__file__),
  )

  code_url_prefixes = (
      code_url_prefix,
      # External packages source repositories,
      "https://github.com/tensorflow/tensorboard/tree/master/tensorboard",
      "https://github.com/tensorflow/estimator/tree/master/tensorflow_estimator",
  )

  if LooseVersion(tf.__version__) < LooseVersion('2'):
    root_title = 'TensorFlow'
  elif LooseVersion(tf.__version__) >= LooseVersion('2'):
    root_title = 'TensorFlow 2.0'

  doc_generator = generate_lib.DocGenerator(
      root_title=root_title,
      py_modules=[("tf", tf)],
      base_dir=base_dirs,
      search_hints=search_hints,
      code_url_prefix=code_url_prefixes,
      site_path=FLAGS.site_path,
      visitor_cls=TfExportAwareDocGeneratorVisitor,
      private_map=PRIVATE_MAP,
      do_not_descend_map=DO_NOT_DESCEND_MAP)

  doc_generator.build(output_dir)

def main(argv):
  del argv
  build_docs(output_dir=FLAGS.output_dir,
             code_url_prefix=FLAGS.code_url_prefix,
             search_hints=FLAGS.search_hints)

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


<commit_before>import hashlib
import hmac
import json
from datetime import datetime, timedelta
from base64 import b64encode

import boto
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from urllib.parse import urlparse, unquote, parse_qs, urlencode

from django.conf import settings

class KeyNotFoundException(Exception):
    pass

def get_at(index, t):
    try:
        value = t[index]
    except IndexError:
        value = None
    return value

# NOTE: Don't use constant as it will break ability to change at runtime
# (E.g. tests)
def get_s3upload_destinations():
    """Returns s3upload destinations, converting old format if necessary."""
    destinations = getattr(settings, 'S3UPLOAD_DESTINATIONS', None)
    if destinations is None:
        return None

    # TODO: Remove when older "positional" settings are no longer supported
    converted_destinations = {}
    key_mapping = {
        0: 'key',
        1: 'auth',
        2: 'allowed',
        3: 'acl',
        4: 'bucket',
        5: 'cache_control',
        6: 'content_disposition',
        7: 'server_side_encryption',
    }
    if destinations:
        for dest, dest_value in destinations.items():
            if type(dest_value) is tuple or type(dest_value) is list:
                converted_destinations[dest] = {}
                for index, key_name in key_mapping.items():
                    converted_destinations[dest][key_name] = get_at(index, dest_value)
            else:
                converted_destinations[dest] = dest_value

    return converted_destinations

def create_upload_data(content_type, key, acl, bucket=None, cache_control=None,
                       content_disposition=None, content_length_range=None,
                       server_side_encryption=None, access_key=None,
                       secret_access_key=None, token=None):

    bucket = bucket or settings.AWS_STORAGE_BUCKET_NAME
    region = getattr(settings, 'S3UPLOAD_REGION', None)
    if not region or region == 'us-east-1':
        endpoint = 's3.amazonaws.com'
    else:
        endpoint = 's3-%s.amazonaws.com' % region
    expires_in = datetime.utcnow() + timedelta(seconds=60 * 5)
    expires = expires_in.strftime('%Y-%m-%dT%H:%M:%S.000Z')
    now_date = datetime.utcnow().strftime('%Y%m%dT%H%M%S000Z')
    raw_date = datetime.utcnow().strftime('%Y%m%d')

    policy_dict = {
        "expiration": expires,
        "conditions": [
            {"bucket": bucket},
            {"acl": acl},
            ["starts-with", "$key", ''],
            {"success_action_status": '201'},
            {"x-amz-credential": '%s/%s/%s/s3/aws4_request' % (
                access_key,
                raw_date, region
            )},
            {"x-amz-algorithm": "AWS4-HMAC-SHA256"},
            {"x-amz-date": now_date},
            {"content-type": content_type},
        ]
    }

    if token:
        policy_dict["conditions"].append({"x-amz-security-token": token})

    if cache_control:
        policy_dict['conditions'].append({'Cache-Control': cache_control})

    if content_disposition:
        policy_dict['conditions'].append({
            'Content-Disposition': content_disposition
        })

    if server_side_encryption:
        policy_dict['conditions'].append(
            {'x-amz-server-side-encryption': server_side_encryption}
        )

    if content_length_range:
        policy_dict['conditions'].append(
            [
                'content-length-range',
                content_length_range[0],
                content_length_range[1]
            ]
        )

    policy_object = json.dumps(policy_dict)

    policy = b64encode(
        policy_object.replace('\n', '').replace('\r', '').encode())

    date_key = hmac.new(b'AWS4' + secret_access_key.encode('utf-8'),
                        msg=raw_date.encode('utf-8'),
                        digestmod=hashlib.sha256).digest()

    date_region_key = hmac.new(date_key, msg=region.encode('utf-8'),
                               digestmod=hashlib.sha256).digest()

    date_region_service_key = hmac.new(date_region_key, msg=b's3',
                                       digestmod=hashlib.sha256).digest()

    signing_key = hmac.new(date_region_service_key, msg=b'aws4_request',
                           digestmod=hashlib.sha256).digest()

    signature = hmac.new(signing_key, msg=policy,
                         digestmod=hashlib.sha256).hexdigest()

    structure = getattr(settings, 'S3UPLOAD_URL_STRUCTURE', 'https://{0}/{1}')
    bucket_url = structure.format(endpoint, bucket)

    return_dict = {
        # FIXME: .decode() does nothing, b64decode works but is decoding
        # really intended?
        "policy": policy.decode(),
        "success_action_status": 201,
        "x-amz-credential": "%s/%s/%s/s3/aws4_request" % (
            access_key, raw_date, region
        ),
        "x-amz-date": now_date,
        "x-amz-signature": signature,
        "x-amz-algorithm": "AWS4-HMAC-SHA256",
        "form_action": bucket_url,
        "key": key,
        "acl": acl,
        "content-type": content_type,
    }

    if token:
        return_dict['x-amz-security-token'] = token

    if server_side_encryption:
        return_dict['x-amz-server-side-encryption'] = server_side_encryption

    if cache_control:
        return_dict['Cache-Control'] = cache_control

    if content_disposition:
        return_dict['Content-Disposition'] = content_disposition

    return return_dict

def get_key_from_url(url, bucket_name=settings.AWS_STORAGE_BUCKET_NAME):
    conn = boto.connect_s3()
    bucket = conn.get_bucket(bucket_name)
    path = urlparse(url).path
    # The bucket name might be part of the path,
    # so get the path that comes after the bucket name
    path = path.split(bucket_name)[-1]
    return bucket.get_key(path)

def get_path_from_url(url, bucket_name=settings.AWS_STORAGE_BUCKET_NAME):
    path = urlparse(url).path
    # The bucket name might be part of the path,
    # so get the path that comes after the bucket name
    path = path.split(bucket_name)[-1]
    return path

def get_signed_download_url(
    source_url,
    bucket_name=settings.AWS_STORAGE_BUCKET_NAME,
    ttl=60,  # 60 seconds should be plenty for a redirected URL to load
):
    conn = S3Connection(
        settings.AWS_ACCESS_KEY_ID,
        settings.AWS_SECRET_ACCESS_KEY,
    )

    bucket = conn.get_bucket(bucket_name)
    k = Key(bucket)

    key_obj = get_key_from_url(unquote(source_url), bucket_name)

    try:
        k.key = key_obj.name
    except AttributeError:
        # key has no 'name' because it's not present in the bucket
        raise KeyNotFoundException

    download_url = k.generate_url(
        expires_in=ttl,
        query_auth=True
    )

    return download_url

def remove_signature(url):
    parsed_url = urlparse(url)
    query_vars = parse_qs(parsed_url.query)

    try:
        del query_vars["Signature"]
        del query_vars["Expires"]
        del query_vars["AWSAccessKeyId"]
    except KeyError:
        # Not a signed URL - return original string
        return url

    address_only = url.split("?")[0]

    if len(query_vars) > 0:
        return "{0}?{1}".format(
            address_only,
            urlencode(query_vars, doseq=True)
        )

    return address_only
<commit_msg>Convert signing func to use key instead of URL<commit_after>import hashlib
import hmac
import json
from datetime import datetime, timedelta
from base64 import b64encode

import boto
from boto.s3.connection import S3Connection
from boto.s3.key import Key
from urllib.parse import urlparse, unquote, parse_qs, urlencode

from django.conf import settings

class KeyNotFoundException(Exception):
    pass

def get_at(index, t):
    try:
        value = t[index]
    except IndexError:
        value = None
    return value

# NOTE: Don't use constant as it will break ability to change at runtime
# (E.g. tests)
def get_s3upload_destinations():
    """Returns s3upload destinations, converting old format if necessary."""
    destinations = getattr(settings, 'S3UPLOAD_DESTINATIONS', None)
    if destinations is None:
        return None

    # TODO: Remove when older "positional" settings are no longer supported
    converted_destinations = {}
    key_mapping = {
        0: 'key',
        1: 'auth',
        2: 'allowed',
        3: 'acl',
        4: 'bucket',
        5: 'cache_control',
        6: 'content_disposition',
        7: 'server_side_encryption',
    }
    if destinations:
        for dest, dest_value in destinations.items():
            if type(dest_value) is tuple or type(dest_value) is list:
                converted_destinations[dest] = {}
                for index, key_name in key_mapping.items():
                    converted_destinations[dest][key_name] = get_at(index, dest_value)
            else:
                converted_destinations[dest] = dest_value

    return converted_destinations

def create_upload_data(content_type, key, acl, bucket=None, cache_control=None,
                       content_disposition=None, content_length_range=None,
                       server_side_encryption=None, access_key=None,
                       secret_access_key=None, token=None):

    bucket = bucket or settings.AWS_STORAGE_BUCKET_NAME
    region = getattr(settings, 'S3UPLOAD_REGION', None)
    if not region or region == 'us-east-1':
        endpoint = 's3.amazonaws.com'
    else:
        endpoint = 's3-%s.amazonaws.com' % region
    expires_in = datetime.utcnow() + timedelta(seconds=60 * 5)
    expires = expires_in.strftime('%Y-%m-%dT%H:%M:%S.000Z')
    now_date = datetime.utcnow().strftime('%Y%m%dT%H%M%S000Z')
    raw_date = datetime.utcnow().strftime('%Y%m%d')

    policy_dict = {
        "expiration": expires,
        "conditions": [
            {"bucket": bucket},
            {"acl": acl},
            ["starts-with", "$key", ''],
            {"success_action_status": '201'},
            {"x-amz-credential": '%s/%s/%s/s3/aws4_request' % (
                access_key,
                raw_date, region
            )},
            {"x-amz-algorithm": "AWS4-HMAC-SHA256"},
            {"x-amz-date": now_date},
            {"content-type": content_type},
        ]
    }

    if token:
        policy_dict["conditions"].append({"x-amz-security-token": token})

    if cache_control:
        policy_dict['conditions'].append({'Cache-Control': cache_control})

    if content_disposition:
        policy_dict['conditions'].append({
            'Content-Disposition': content_disposition
        })

    if server_side_encryption:
        policy_dict['conditions'].append(
            {'x-amz-server-side-encryption': server_side_encryption}
        )

    if content_length_range:
        policy_dict['conditions'].append(
            [
                'content-length-range',
                content_length_range[0],
                content_length_range[1]
            ]
        )

    policy_object = json.dumps(policy_dict)

    policy = b64encode(
        policy_object.replace('\n', '').replace('\r', '').encode())

    date_key = hmac.new(b'AWS4' + secret_access_key.encode('utf-8'),
                        msg=raw_date.encode('utf-8'),
                        digestmod=hashlib.sha256).digest()

    date_region_key = hmac.new(date_key, msg=region.encode('utf-8'),
                               digestmod=hashlib.sha256).digest()

    date_region_service_key = hmac.new(date_region_key, msg=b's3',
                                       digestmod=hashlib.sha256).digest()

    signing_key = hmac.new(date_region_service_key, msg=b'aws4_request',
                           digestmod=hashlib.sha256).digest()

    signature = hmac.new(signing_key, msg=policy,
                         digestmod=hashlib.sha256).hexdigest()

    structure = getattr(settings, 'S3UPLOAD_URL_STRUCTURE', 'https://{0}/{1}')
    bucket_url = structure.format(endpoint, bucket)

    return_dict = {
        # FIXME: .decode() does nothing, b64decode works but is decoding
        # really intended?
        "policy": policy.decode(),
        "success_action_status": 201,
        "x-amz-credential": "%s/%s/%s/s3/aws4_request" % (
            access_key, raw_date, region
        ),
        "x-amz-date": now_date,
        "x-amz-signature": signature,
        "x-amz-algorithm": "AWS4-HMAC-SHA256",
        "form_action": bucket_url,
        "key": key,
        "acl": acl,
        "content-type": content_type,
    }

    if token:
        return_dict['x-amz-security-token'] = token

    if server_side_encryption:
        return_dict['x-amz-server-side-encryption'] = server_side_encryption

    if cache_control:
        return_dict['Cache-Control'] = cache_control

    if content_disposition:
        return_dict['Content-Disposition'] = content_disposition

    return return_dict

def get_path_from_url(url, bucket_name=settings.AWS_STORAGE_BUCKET_NAME):
    path = urlparse(url).path
    # The bucket name might be part of the path,
    # so get the path that comes after the bucket name
    path = path.split(bucket_name)[-1]
    return path

def get_key_from_url(url, bucket_name=settings.AWS_STORAGE_BUCKET_NAME):
    conn = boto.connect_s3()
    bucket = conn.get_bucket(bucket_name)
    path = get_path_from_url(url)
    return bucket.get_key(path)

def get_signed_download_url(
    key,
    bucket_name=settings.AWS_STORAGE_BUCKET_NAME,
    ttl=60,  # 60 seconds should be plenty for a redirected URL to load
):
    conn = S3Connection(
        settings.AWS_ACCESS_KEY_ID,
        settings.AWS_SECRET_ACCESS_KEY,
    )

    bucket = conn.get_bucket(bucket_name)
    k = Key(bucket)

    try:
        k.key = key
    except AttributeError:
        # key has no 'name' because it's not present in the bucket
        raise KeyNotFoundException

    download_url = k.generate_url(
        expires_in=ttl,
        query_auth=True
    )

    return download_url

def remove_signature(url):
    parsed_url = urlparse(url)
    query_vars = parse_qs(parsed_url.query)

    try:
        del query_vars["Signature"]
        del query_vars["Expires"]
        del query_vars["AWSAccessKeyId"]
    except KeyError:
        # Not a signed URL - return original string
        return url

    address_only = url.split("?")[0]

    if len(query_vars) > 0:
        return "{0}?{1}".format(
            address_only,
            urlencode(query_vars, doseq=True)
        )

    return address_only
<|endoftext|>


<commit_before>#!/usr/bin/env python2.7
# -*- mode: python; coding: utf-8; -*-

"""Module for generating lexicon using Severyn and Moschitti's method (2014).

"""

##################################################################
# Imports
from __future__ import unicode_literals, print_function

from common import ENCODING, ESC_CHAR, FMAX, FMIN, \
    INFORMATIVE_TAGS, NEGATIVE, POSITIVE, SENT_END_RE, \
    TAB_RE, NONMATCH_RE, MIN_TOK_CNT, check_word, normalize

from collections import Counter
from itertools import chain
from sklearn.feature_extraction import DictVectorizer
from sklearn.pipeline import Pipeline
from sklearn.svm import LinearSVC
import codecs
import sys

##################################################################
# Constants
FASTMODE = False
POS_IDX = 0
NEG_IDX = 1

##################################################################
# Methods
def _toks2feats(a_tweet_toks):
    """Convert set of tweet's tokens to a dictionary of features.

    @param a_tweet_toks - set of tweet's uni- and bigrams

    @return feature dictionary with tweet's tokens as keys and values 1

    """
    return {w: 1. for w in a_tweet_toks}

def _update_ts(a_ts_x, a_ts_y, a_tweet_toks, a_pos, a_neg,
               a_pos_re=NONMATCH_RE, a_neg_re=NONMATCH_RE):
    """Update training set of features and classes.

    @param a_ts_x - training set features
    @param a_ts_y - training set gold classes
    @param a_tweet_toks - set of tweet's uni- and bigrams
    @param a_pos - initial set of positive terms to be expanded
    @param a_neg - initial set of negative terms to be expanded
    @param a_pos_re - regular expression for matching positive terms
    @param a_neg_re - regular expression for matching negative terms

    @return \c void

    @note modifies `a_ts' and `a_tweet_feats' in place

    """
    if not a_tweet_toks:
        return

    tweet = ' '.join(sorted(a_tweet_toks))
    if a_tweet_toks & a_pos or a_pos_re.search(tweet):
        a_ts_x.append(_toks2feats(a_tweet_toks))
        a_ts_y.append(POSITIVE)
    elif a_tweet_toks & a_neg or a_neg_re.search(tweet):
        a_ts_x.append(_toks2feats(a_tweet_toks))
        a_ts_y.append(NEGATIVE)
    a_tweet_toks.clear()

def _prune_ts(a_ts_x, a_ts_y):
    """Remove fetures with too low frequency.

    @param a_ts_x - list of extracted features
    @param a_ts_y - list of gold instance classes

    @return 2-tuple - pruned copies of `a_ts_x' and `a_ts_y'

    """
    tokstat = Counter(x_i
                      for x in a_ts_x
                      for x_i in x.iterkeys()
                      )
    x_ = None
    ts_x = []
    ts_y = []
    for x, y in zip(a_ts_x, a_ts_y):
        x_ = {k: v for k, v in x.iteritems()
              if tokstat[k] >= MIN_TOK_CNT
              }
        if x_:
            ts_x.append(x_)
            ts_y.append(y)
    return (ts_x, ts_y)

def _read_files(a_crp_files, a_pos, a_neg,
                a_pos_re=NONMATCH_RE, a_neg_re=NONMATCH_RE):
    """Read corpus files and populate one-directional co-occurrences.

    @param a_crp_files - files of the original corpus
    @param a_pos - initial set of positive terms to be expanded
    @param a_neg - initial set of negative terms to be expanded
    @param a_pos_re - regular expression for matching positive terms
    @param a_neg_re - regular expression for matching negative terms

    @return 2-tuple - training sets of features and their gold classes

    """
    print("Reading corpus...", end="", file=sys.stderr)
    i = 0
    ts_x = []
    ts_y = []
    tweet_toks = set()
    iform = itag = ilemma = prev_lemma = ""
    for ifname in a_crp_files:
        with codecs.open(ifname, 'r', ENCODING) as ifile:
            prev_lemma = ""
            for iline in ifile:
                iline = iline.strip().lower()
                if iline and iline[0] == ESC_CHAR:
                    if FASTMODE:
                        i += 1
                        if i > 300:
                            break
                    _update_ts(ts_x, ts_y, tweet_toks,
                               a_pos, a_neg, a_pos_re, a_neg_re)
                    prev_lemma = ""
                    continue
                elif not iline or SENT_END_RE.match(iline):
                    prev_lemma = ""
                    continue
                try:
                    iform, itag, ilemma = TAB_RE.split(iline)
                except:
                    print("Invalid line format at line: {:s}".format(
                        repr(iline)), file=sys.stderr
                    )
                    continue
                ilemma = normalize(ilemma)
                if a_pos_re.search(iform) or a_neg_re.search(iform):
                    tweet_toks.add(iform)
                elif a_pos_re.search(ilemma) or a_neg_re.search(ilemma):
                    tweet_toks.add(ilemma)
                elif itag[:2] not in INFORMATIVE_TAGS \
                        or not check_word(ilemma):
                    continue
                else:
                    tweet_toks.add(ilemma)
                if prev_lemma:
                    tweet_toks.add((prev_lemma, ilemma))
            _update_ts(ts_x, ts_y, tweet_toks,
                       a_pos, a_neg, a_pos_re, a_neg_re)
    print(" done", file=sys.stderr)
    return _prune_ts(ts_x, ts_y)

def severyn(a_N, a_crp_files, a_pos, a_neg,
            a_pos_re=NONMATCH_RE, a_neg_re=NONMATCH_RE):
    """Method for generating sentiment lexicons using Severyn's approach.

    @param a_N - number of terms to extract
    @param a_crp_files - files of the original corpus
    @param a_pos - initial set of positive terms to be expanded
    @param a_neg - initial set of negative terms to be expanded
    @param a_pos_re - regular expression for matching positive terms
    @param a_neg_re - regular expression for matching negative terms

    @return list of terms sorted according to their polarity scores

    """
    a_pos = set(normalize(w) for w in a_pos)
    a_neg = set(normalize(w) for w in a_neg)

    vectorizer = DictVectorizer()
    clf = LinearSVC(C=0.3)
    model = Pipeline([("vectorizer", vectorizer),
                      ("LinearSVC", clf)])
    X, Y = _read_files(a_crp_files, a_pos, a_neg,
                       a_pos_re, a_neg_re)
    model.fit(X, Y)

    ret = [(w, POSITIVE, FMAX) for w in a_pos] \
        + [(w, NEGATIVE, FMIN) for w in a_neg]
    coefs = clf.coef_[0]
    for f_name, f_score in zip(vectorizer.get_feature_names(),
                               coefs
                               ):
        if f_name in a_pos or f_name in a_neg:
            continue
        ret.append((f_name,
                    POSITIVE if f_score > 0. else NEGATIVE,
                    f_score))
    ret.sort(key=lambda el: abs(el[-1]), reverse=True)
    return ret
<commit_msg>multi tokens fixed in severyn<commit_after>#!/usr/bin/env python2.7
# -*- mode: python; coding: utf-8; -*-

"""Module for generating lexicon using Severyn and Moschitti's method (2014).

"""

##################################################################
# Imports
from __future__ import unicode_literals, print_function

from common import ENCODING, ESC_CHAR, FMAX, FMIN, \
    INFORMATIVE_TAGS, NEGATIVE, POSITIVE, SENT_END_RE, \
    TAB_RE, NONMATCH_RE, MIN_TOK_CNT, check_word, normalize

from collections import Counter
from itertools import chain
from sklearn.feature_extraction import DictVectorizer
from sklearn.pipeline import Pipeline
from sklearn.svm import LinearSVC
import codecs
import sys

##################################################################
# Constants
FASTMODE = False
POS_IDX = 0
NEG_IDX = 1

##################################################################
# Methods
def _toks2feats(a_tweet_toks):
    """Convert set of tweet's tokens to a dictionary of features.

    @param a_tweet_toks - set of tweet's uni- and bigrams

    @return feature dictionary with tweet's tokens as keys and values 1

    """
    return {w: 1. for w in a_tweet_toks}

def _update_ts(a_ts_x, a_ts_y, a_tweet_toks, a_pos, a_neg,
               a_pos_re=NONMATCH_RE, a_neg_re=NONMATCH_RE):
    """Update training set of features and classes.

    @param a_ts_x - training set features
    @param a_ts_y - training set gold classes
    @param a_tweet_toks - set of tweet's uni- and bigrams
    @param a_pos - initial set of positive terms to be expanded
    @param a_neg - initial set of negative terms to be expanded
    @param a_pos_re - regular expression for matching positive terms
    @param a_neg_re - regular expression for matching negative terms

    @return \c void

    @note modifies `a_ts' and `a_tweet_feats' in place

    """
    if not a_tweet_toks:
        return

    tweet = ' '.join([w
                      if isinstance(w, basestring)
                      else ' '.join(w)
                      for w in sorted(a_tweet_toks)])
    if a_tweet_toks & a_pos or a_pos_re.search(tweet):
        a_ts_x.append(_toks2feats(a_tweet_toks))
        a_ts_y.append(POSITIVE)
    elif a_tweet_toks & a_neg or a_neg_re.search(tweet):
        a_ts_x.append(_toks2feats(a_tweet_toks))
        a_ts_y.append(NEGATIVE)
    a_tweet_toks.clear()

def _prune_ts(a_ts_x, a_ts_y):
    """Remove fetures with too low frequency.

    @param a_ts_x - list of extracted features
    @param a_ts_y - list of gold instance classes

    @return 2-tuple - pruned copies of `a_ts_x' and `a_ts_y'

    """
    tokstat = Counter(x_i
                      for x in a_ts_x
                      for x_i in x.iterkeys()
                      )
    x_ = None
    ts_x = []
    ts_y = []
    for x, y in zip(a_ts_x, a_ts_y):
        x_ = {k: v for k, v in x.iteritems()
              if tokstat[k] >= MIN_TOK_CNT
              }
        if x_:
            ts_x.append(x_)
            ts_y.append(y)
    return (ts_x, ts_y)

def _read_files(a_crp_files, a_pos, a_neg,
                a_pos_re=NONMATCH_RE, a_neg_re=NONMATCH_RE):
    """Read corpus files and populate one-directional co-occurrences.

    @param a_crp_files - files of the original corpus
    @param a_pos - initial set of positive terms to be expanded
    @param a_neg - initial set of negative terms to be expanded
    @param a_pos_re - regular expression for matching positive terms
    @param a_neg_re - regular expression for matching negative terms

    @return 2-tuple - training sets of features and their gold classes

    """
    print("Reading corpus...", end="", file=sys.stderr)
    i = 0
    ts_x = []
    ts_y = []
    tweet_toks = set()
    iform = itag = ilemma = prev_lemma = ""
    for ifname in a_crp_files:
        with codecs.open(ifname, 'r', ENCODING) as ifile:
            prev_lemma = ""
            for iline in ifile:
                iline = iline.strip().lower()
                if iline and iline[0] == ESC_CHAR:
                    if FASTMODE:
                        i += 1
                        if i > 300:
                            break
                    _update_ts(ts_x, ts_y, tweet_toks,
                               a_pos, a_neg, a_pos_re, a_neg_re)
                    prev_lemma = ""
                    continue
                elif not iline or SENT_END_RE.match(iline):
                    prev_lemma = ""
                    continue
                try:
                    iform, itag, ilemma = TAB_RE.split(iline)
                except:
                    print("Invalid line format at line: {:s}".format(
                        repr(iline)), file=sys.stderr
                    )
                    continue
                ilemma = normalize(ilemma)
                if a_pos_re.search(iform) or a_neg_re.search(iform):
                    tweet_toks.add(iform)
                elif a_pos_re.search(ilemma) or a_neg_re.search(ilemma):
                    tweet_toks.add(ilemma)
                elif itag[:2] not in INFORMATIVE_TAGS \
                        or not check_word(ilemma):
                    continue
                else:
                    tweet_toks.add(ilemma)
                if prev_lemma:
                    tweet_toks.add((prev_lemma, ilemma))
                prev_lemma = ilemma
            _update_ts(ts_x, ts_y, tweet_toks,
                       a_pos, a_neg, a_pos_re, a_neg_re)
    print(" done", file=sys.stderr)
    return _prune_ts(ts_x, ts_y)

def severyn(a_N, a_crp_files, a_pos, a_neg,
            a_pos_re=NONMATCH_RE, a_neg_re=NONMATCH_RE):
    """Method for generating sentiment lexicons using Severyn's approach.

    @param a_N - number of terms to extract
    @param a_crp_files - files of the original corpus
    @param a_pos - initial set of positive terms to be expanded
    @param a_neg - initial set of negative terms to be expanded
    @param a_pos_re - regular expression for matching positive terms
    @param a_neg_re - regular expression for matching negative terms

    @return list of terms sorted according to their polarity scores

    """
    a_pos = set(normalize(w) for w in a_pos)
    a_neg = set(normalize(w) for w in a_neg)

    vectorizer = DictVectorizer()
    clf = LinearSVC(C=0.3)
    model = Pipeline([("vectorizer", vectorizer),
                      ("LinearSVC", clf)])
    X, Y = _read_files(a_crp_files, a_pos, a_neg,
                       a_pos_re, a_neg_re)
    model.fit(X, Y)

    ret = [(w, POSITIVE, FMAX) for w in a_pos] \
        + [(w, NEGATIVE, FMIN) for w in a_neg]
    coefs = clf.coef_[0]
    for f_name, f_score in zip(vectorizer.get_feature_names(),
                               coefs
                               ):
        if f_name in a_pos or f_name in a_neg:
            continue
        ret.append((f_name,
                    POSITIVE if f_score > 0. else NEGATIVE,
                    f_score))
    ret.sort(key=lambda el: abs(el[-1]), reverse=True)
    return ret
<|endoftext|>


<commit_before>#!/usr/bin/env python3

""" Smart Automatic Cover Art Downloader : search and download music album covers. """

__version__ = "2.3.0"
__author__ = "desbma"
__license__ = "MPL 2.0"

import argparse
import asyncio
import functools
import logging
import os

from sacad import colored_logging
from sacad import sources
from sacad.cover import CoverSourceResult, HAS_JPEGOPTIM, HAS_OPTIPNG, SUPPORTED_IMG_FORMATS

async def search_and_download(album, artist, format, size, out_filepath, *, size_tolerance_prct, amazon_tlds,
                              no_lq_sources, preserve_format=False):
  """ Search and download a cover, return True if success, False instead. """
  # register sources
  source_args = (size, size_tolerance_prct)
  cover_sources = [sources.LastFmCoverSource(*source_args),
                   sources.DeezerCoverSource(*source_args),
                   sources.AmazonCdCoverSource(*source_args),
                   sources.AmazonDigitalCoverSource(*source_args)]
  for tld in amazon_tlds:
    cover_sources.append(sources.AmazonCdCoverSource(*source_args, tld=tld))
  if not no_lq_sources:
    cover_sources.append(sources.GoogleImagesWebScrapeCoverSource(*source_args))

  # schedule search work
  search_futures = []
  for cover_source in cover_sources:
    coroutine = cover_source.search(album, artist)
    future = asyncio.ensure_future(coroutine)
    search_futures.append(future)

  # wait for it
  await asyncio.wait(search_futures)

  # get results
  results = []
  for future in search_futures:
    source_results = future.result()
    results.extend(source_results)

  # sort results
  results = await CoverSourceResult.preProcessForComparison(results, size, size_tolerance_prct)
  results.sort(reverse=True,
               key=functools.cmp_to_key(functools.partial(CoverSourceResult.compare,
                                                          target_size=size,
                                                          size_tolerance_prct=size_tolerance_prct)))
  if not results:
    logging.getLogger("Main").info("No results")

  # download
  done = False
  for result in results:
    try:
      await result.get(format, size, size_tolerance_prct, out_filepath, preserve_format=preserve_format)
    except Exception as e:
      logging.getLogger("Main").warning("Download of %s failed: %s %s" % (result,
                                                                          e.__class__.__qualname__,
                                                                          e))
      continue
    else:
      done = True
      break

  # cleanup sessions
  close_cr = []
  for cover_source in cover_sources:
    close_cr.append(cover_source.closeSession())
  await asyncio.gather(*close_cr)

  return done

def setup_common_args(arg_parser):
  arg_parser.add_argument("-t",
                          "--size-tolerance",
                          type=int,
                          default=25,
                          dest="size_tolerance_prct",
                          help="""Tolerate this percentage of size difference with the target size.
                                  Note that covers with size above or close to the target size will still be preferred
                                  if available""")
  arg_parser.add_argument("-a",
                          "--amazon-sites",
                          nargs="+",
                          choices=sources.AmazonCdCoverSource.TLDS[1:],
                          default=(),
                          dest="amazon_tlds",
                          help="""Amazon site TLDs to use as search source, in addition to amazon.com""")
  arg_parser.add_argument("-d",
                          "--disable-low-quality-sources",
                          action="store_true",
                          default=False,
                          dest="no_lq_sources",
                          help="""Disable cover sources that may return unreliable results (ie. Google Images).
                                  It will speed up search and improve reliability, but may fail to find results for
                                  some difficult searches.""")
  arg_parser.add_argument("-p",
                          "--preserve-format",
                          action="store_true",
                          default=False,
                          help="Preserve source image format if possible. Target format will still be prefered when sorting results.")

def cl_main():
  # parse args
  arg_parser = argparse.ArgumentParser(description="SACAD v%s. Search and download an album cover." % (__version__),
                                       formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  arg_parser.add_argument("artist",
                          help="Artist to search for")
  arg_parser.add_argument("album",
                          help="Album to search for")
  arg_parser.add_argument("size",
                          type=int,
                          help="Target image size")
  arg_parser.add_argument("out_filepath",
                          help="Output image filepath")
  setup_common_args(arg_parser)
  arg_parser.add_argument("-v",
                          "--verbosity",
                          choices=("quiet", "warning", "normal", "debug"),
                          default="normal",
                          dest="verbosity",
                          help="Level of logging output")
  args = arg_parser.parse_args()
  args.format = os.path.splitext(args.out_filepath)[1][1:].lower()
  try:
    args.format = SUPPORTED_IMG_FORMATS[args.format]
  except KeyError:
    print("Unable to guess image format from extension, or unknown format: %s" % (args.format))
    exit(1)

  # setup logger
  logging_level = {"quiet": logging.CRITICAL + 1,
                   "warning": logging.WARNING,
                   "normal": logging.INFO,
                   "debug": logging.DEBUG}
  logging.getLogger().setLevel(logging_level[args.verbosity])
  if logging_level[args.verbosity] == logging.DEBUG:
    fmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s"
  else:
    fmt = "%(name)s: %(message)s"
  logging_formatter = colored_logging.ColoredFormatter(fmt=fmt)
  logging_handler = logging.StreamHandler()
  logging_handler.setFormatter(logging_formatter)
  logging.getLogger().addHandler(logging_handler)
  if logging_level[args.verbosity] == logging.DEBUG:
    logging.getLogger("asyncio").setLevel(logging.WARNING)
  else:
    logging.getLogger("asyncio").setLevel(logging.CRITICAL + 1)

  # display warning if optipng or jpegoptim are missing
  if not HAS_JPEGOPTIM:
    logging.getLogger("Main").warning("jpegoptim could not be found, JPEG crunching will be disabled")
  if not HAS_OPTIPNG:
    logging.getLogger("Main").warning("optipng could not be found, PNG crunching will be disabled")

  # search and download
  coroutine = search_and_download(args.album,
                                  args.artist,
                                  args.format,
                                  args.size,
                                  args.out_filepath,
                                  size_tolerance_prct=args.size_tolerance_prct,
                                  amazon_tlds=args.amazon_tlds,
                                  no_lq_sources=args.no_lq_sources,
                                  preserve_format=args.preserve_format)
  if hasattr(asyncio, "run"):
    # Python >=3.7.0
    asyncio.run(coroutine)
  else:
    future = asyncio.ensure_future(coroutine)
    asyncio.get_event_loop().run_until_complete(future)

if __name__ == "__main__":
  cl_main()
<commit_msg>Version 2.3.1<commit_after>#!/usr/bin/env python3

""" Smart Automatic Cover Art Downloader : search and download music album covers. """

__version__ = "2.3.1"
__author__ = "desbma"
__license__ = "MPL 2.0"

import argparse
import asyncio
import functools
import logging
import os

from sacad import colored_logging
from sacad import sources
from sacad.cover import CoverSourceResult, HAS_JPEGOPTIM, HAS_OPTIPNG, SUPPORTED_IMG_FORMATS

async def search_and_download(album, artist, format, size, out_filepath, *, size_tolerance_prct, amazon_tlds,
                              no_lq_sources, preserve_format=False):
  """ Search and download a cover, return True if success, False instead. """
  # register sources
  source_args = (size, size_tolerance_prct)
  cover_sources = [sources.LastFmCoverSource(*source_args),
                   sources.DeezerCoverSource(*source_args),
                   sources.AmazonCdCoverSource(*source_args),
                   sources.AmazonDigitalCoverSource(*source_args)]
  for tld in amazon_tlds:
    cover_sources.append(sources.AmazonCdCoverSource(*source_args, tld=tld))
  if not no_lq_sources:
    cover_sources.append(sources.GoogleImagesWebScrapeCoverSource(*source_args))

  # schedule search work
  search_futures = []
  for cover_source in cover_sources:
    coroutine = cover_source.search(album, artist)
    future = asyncio.ensure_future(coroutine)
    search_futures.append(future)

  # wait for it
  await asyncio.wait(search_futures)

  # get results
  results = []
  for future in search_futures:
    source_results = future.result()
    results.extend(source_results)

  # sort results
  results = await CoverSourceResult.preProcessForComparison(results, size, size_tolerance_prct)
  results.sort(reverse=True,
               key=functools.cmp_to_key(functools.partial(CoverSourceResult.compare,
                                                          target_size=size,
                                                          size_tolerance_prct=size_tolerance_prct)))
  if not results:
    logging.getLogger("Main").info("No results")

  # download
  done = False
  for result in results:
    try:
      await result.get(format, size, size_tolerance_prct, out_filepath, preserve_format=preserve_format)
    except Exception as e:
      logging.getLogger("Main").warning("Download of %s failed: %s %s" % (result,
                                                                          e.__class__.__qualname__,
                                                                          e))
      continue
    else:
      done = True
      break

  # cleanup sessions
  close_cr = []
  for cover_source in cover_sources:
    close_cr.append(cover_source.closeSession())
  await asyncio.gather(*close_cr)

  return done

def setup_common_args(arg_parser):
  arg_parser.add_argument("-t",
                          "--size-tolerance",
                          type=int,
                          default=25,
                          dest="size_tolerance_prct",
                          help="""Tolerate this percentage of size difference with the target size.
                                  Note that covers with size above or close to the target size will still be preferred
                                  if available""")
  arg_parser.add_argument("-a",
                          "--amazon-sites",
                          nargs="+",
                          choices=sources.AmazonCdCoverSource.TLDS[1:],
                          default=(),
                          dest="amazon_tlds",
                          help="""Amazon site TLDs to use as search source, in addition to amazon.com""")
  arg_parser.add_argument("-d",
                          "--disable-low-quality-sources",
                          action="store_true",
                          default=False,
                          dest="no_lq_sources",
                          help="""Disable cover sources that may return unreliable results (ie. Google Images).
                                  It will speed up search and improve reliability, but may fail to find results for
                                  some difficult searches.""")
  arg_parser.add_argument("-p",
                          "--preserve-format",
                          action="store_true",
                          default=False,
                          help="Preserve source image format if possible. Target format will still be prefered when sorting results.")

def cl_main():
  # parse args
  arg_parser = argparse.ArgumentParser(description="SACAD v%s. Search and download an album cover." % (__version__),
                                       formatter_class=argparse.ArgumentDefaultsHelpFormatter)
  arg_parser.add_argument("artist",
                          help="Artist to search for")
  arg_parser.add_argument("album",
                          help="Album to search for")
  arg_parser.add_argument("size",
                          type=int,
                          help="Target image size")
  arg_parser.add_argument("out_filepath",
                          help="Output image filepath")
  setup_common_args(arg_parser)
  arg_parser.add_argument("-v",
                          "--verbosity",
                          choices=("quiet", "warning", "normal", "debug"),
                          default="normal",
                          dest="verbosity",
                          help="Level of logging output")
  args = arg_parser.parse_args()
  args.format = os.path.splitext(args.out_filepath)[1][1:].lower()
  try:
    args.format = SUPPORTED_IMG_FORMATS[args.format]
  except KeyError:
    print("Unable to guess image format from extension, or unknown format: %s" % (args.format))
    exit(1)

  # setup logger
  logging_level = {"quiet": logging.CRITICAL + 1,
                   "warning": logging.WARNING,
                   "normal": logging.INFO,
                   "debug": logging.DEBUG}
  logging.getLogger().setLevel(logging_level[args.verbosity])
  if logging_level[args.verbosity] == logging.DEBUG:
    fmt = "%(asctime)s %(levelname)s [%(name)s] %(message)s"
  else:
    fmt = "%(name)s: %(message)s"
  logging_formatter = colored_logging.ColoredFormatter(fmt=fmt)
  logging_handler = logging.StreamHandler()
  logging_handler.setFormatter(logging_formatter)
  logging.getLogger().addHandler(logging_handler)
  if logging_level[args.verbosity] == logging.DEBUG:
    logging.getLogger("asyncio").setLevel(logging.WARNING)
  else:
    logging.getLogger("asyncio").setLevel(logging.CRITICAL + 1)

  # display warning if optipng or jpegoptim are missing
  if not HAS_JPEGOPTIM:
    logging.getLogger("Main").warning("jpegoptim could not be found, JPEG crunching will be disabled")
  if not HAS_OPTIPNG:
    logging.getLogger("Main").warning("optipng could not be found, PNG crunching will be disabled")

  # search and download
  coroutine = search_and_download(args.album,
                                  args.artist,
                                  args.format,
                                  args.size,
                                  args.out_filepath,
                                  size_tolerance_prct=args.size_tolerance_prct,
                                  amazon_tlds=args.amazon_tlds,
                                  no_lq_sources=args.no_lq_sources,
                                  preserve_format=args.preserve_format)
  if hasattr(asyncio, "run"):
    # Python >=3.7.0
    asyncio.run(coroutine)
  else:
    future = asyncio.ensure_future(coroutine)
    asyncio.get_event_loop().run_until_complete(future)

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


<commit_before>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
Script to run numerical simulations.
"""
from __future__ import print_function

import argparse
import os
import shutil
import numpy as np
from tuna.simu.main import SimuParams, DivisionParams
from tuna.simu.ou import OUParams, OUSimulation

# Arguments
parser = argparse.ArgumentParser()
parser.add_argument('-p', '--path', type=str,
                    help='Parent directory in which simulation is stored',
                    default='~/tmptuna')
parser.add_argument('-l', '--label', type=str,
                    help='Label of the experiment/simulation',
                    default='simutest')
parser.add_argument('-s', '--samples', type=int,
                    help='Number of simulated container samples',
                    default=100)
parser.add_argument('--start', type=float,
                    help='Time at which simulation starts',
                    default=0.)
parser.add_argument('--stop', type=float,
                    help='Time at which simulation stops',
                    default=180.)
parser.add_argument('--period', type=float,
                    help=('Time period between two consecutive time-lapse'
                          ' acquisitions'),
                    default=5.)
args = parser.parse_args()

path = os.path.abspath(os.path.expanduser(args.path))
if not os.path.exists(path):
    os.makedirs(path)
print('Path: {}'.format(path))
print('Label: {}'.format(args.label))
# %% SET PARAMETERS FOR NUMERICAL SIMULATION

# general simulation parameters
tstart = args.start  # initial time of simulation
tstop = args.stop  # final time of simulation (data is not recorded beyond this)
dt = args.period  # time interval for data recording (sampling of process)
nsamples = args.samples  # number of containers that will be created by simulation

tree_per_cont = 2  # number of trees per container

# OU parameters
target_value = np.log(2.)/60.  # average value of the random process
spring = 1./30.  # spring constant (brings back random walker to target value)
noise_intensity = 2. * spring * (target_value/10.)**2   # noise intensity for random walk

# DIVISION TIMING
interdivision_mean = 60.  # average of interdivision time
interdivision_std = 6.  # standard deviation for interdivision time
minimum = dt  # minimum value of interdivision time should be >= dt

# %% Instantiating parameters objects

simuParams = SimuParams(nbr_container=nsamples,
                        nbr_colony_per_container=tree_per_cont,
                        start=tstart, stop=tstop, interval=dt)

ouParams = OUParams(target=target_value, spring=spring,
                    noise=noise_intensity)

divParams = DivisionParams(mean=interdivision_mean, std=interdivision_std)

# %% SIMULATION: initiate, define, run/write
np.random.seed()  # seed for the random number generator

exp = OUSimulation(label=args.label,
                   simuParams=simuParams, divisionParams=divParams,
                   ouParams=ouParams)

# check that experiment has been saved before
ans = 'go'
current_name = args.label
exp_path = os.path.join(path, args.label)
while os.path.exists(exp_path) and ans != 'a':
    print('Experiment {} already exists.'.format(current_name))
    ans = raw_input('Override [o], Change experiment name [c], Abort [a]: ')
    if ans == 'c':
        new_name = current_name
        while new_name == current_name:
            new_name = raw_input('NEW NAME: ')
        exp.label = new_name
        exp_path = os.path.join(path, new_name)
    # when overriding, need to erase everything first
    elif ans == 'o':
        shutil.rmtree(exp_path)

# export except if process has been aborted
if ans != 'a':
    exp.raw_text_export(path=path)
<commit_msg>Edit main dosctring<commit_after>#!/usr/bin/env python2
# -*- coding: utf-8 -*-
"""
script: simurun.py

used to run numerical simulations from terminal. Command is:

    python simurun.py [[-p|--path <path>]
                       [-l|--label <exp-name>]
                       [-s|--samples <number-of-samples>]
                       [--start <starting-time-value>]
                       [--stop <stoping-time-value>]
                       [--period <time-interval-between-acquisitions>]
                      ]

All parameters have default values (see code).
"""
from __future__ import print_function

import argparse
import os
import shutil
import numpy as np
from tuna.simu.main import SimuParams, DivisionParams
from tuna.simu.ou import OUParams, OUSimulation

# Arguments
parser = argparse.ArgumentParser()
parser.add_argument('-p', '--path', type=str,
                    help='Parent directory in which simulation is stored',
                    default='~/tmptuna')
parser.add_argument('-l', '--label', type=str,
                    help='Label of the experiment/simulation',
                    default='simutest')
parser.add_argument('-s', '--samples', type=int,
                    help='Number of simulated container samples',
                    default=100)
parser.add_argument('--start', type=float,
                    help='Time at which simulation starts',
                    default=0.)
parser.add_argument('--stop', type=float,
                    help='Time at which simulation stops',
                    default=180.)
parser.add_argument('--period', type=float,
                    help=('Time period between two consecutive time-lapse'
                          ' acquisitions'),
                    default=5.)
args = parser.parse_args()

path = os.path.abspath(os.path.expanduser(args.path))
if not os.path.exists(path):
    os.makedirs(path)
print('Path: {}'.format(path))
print('Label: {}'.format(args.label))
# %% SET PARAMETERS FOR NUMERICAL SIMULATION

# general simulation parameters
tstart = args.start  # initial time of simulation
tstop = args.stop  # final time of simulation (data is not recorded beyond this)
dt = args.period  # time interval for data recording (sampling of process)
nsamples = args.samples  # number of containers that will be created by simulation

tree_per_cont = 2  # number of trees per container

# OU parameters
target_value = np.log(2.)/60.  # average value of the random process
spring = 1./30.  # spring constant (brings back random walker to target value)
noise_intensity = 2. * spring * (target_value/10.)**2   # noise intensity for random walk

# DIVISION TIMING
interdivision_mean = 60.  # average of interdivision time
interdivision_std = 6.  # standard deviation for interdivision time
minimum = dt  # minimum value of interdivision time should be >= dt

# %% Instantiating parameters objects

simuParams = SimuParams(nbr_container=nsamples,
                        nbr_colony_per_container=tree_per_cont,
                        start=tstart, stop=tstop, interval=dt)

ouParams = OUParams(target=target_value, spring=spring,
                    noise=noise_intensity)

divParams = DivisionParams(mean=interdivision_mean, std=interdivision_std)

# %% SIMULATION: initiate, define, run/write
np.random.seed()  # seed for the random number generator

exp = OUSimulation(label=args.label,
                   simuParams=simuParams, divisionParams=divParams,
                   ouParams=ouParams)

# check that experiment has been saved before
ans = 'go'
current_name = args.label
exp_path = os.path.join(path, args.label)
while os.path.exists(exp_path) and ans != 'a':
    print('Experiment {} already exists.'.format(current_name))
    ans = raw_input('Override [o], Change experiment name [c], Abort [a]: ')
    if ans == 'c':
        new_name = current_name
        while new_name == current_name:
            new_name = raw_input('NEW NAME: ')
        exp.label = new_name
        exp_path = os.path.join(path, new_name)
    # when overriding, need to erase everything first
    elif ans == 'o':
        shutil.rmtree(exp_path)

# export except if process has been aborted
if ans != 'a':
    exp.raw_text_export(path=path)
<|endoftext|>


<commit_before>install_requires = [
    'docker-py>=1.2.0',
    'watchdog>=0.8.0',
    'PyYAML>=3.11',
    'PrettyTable>=0.7.0',
    'GitPython>=1.0.1',
    'docopt==0.6.2',
    'schemer==0.2.8',
]

test_requires = [
    'nose>=1.3.0',
    'mock>=1.0.1',
]
<commit_msg>update requirements<commit_after>install_requires = [
    'docker-py>=1.2.0',
    'watchdog>=0.8.0',
    'PyYAML>=3.11',
    'PrettyTable>=0.7.0',
    'GitPython>=1.0.1',
    'docopt==0.6.2',
    'Schemer==0.2.8',
]

test_requires = [
    'nose>=1.3.0',
    'mock>=1.0.1',
]
<|endoftext|>


<commit_before>#!/usr/bin/env python
# Author: Paul Young
# Last Modified: 07/12/2014
# Objectified autorun script for GAMESS->QMCPACK calculation on the molecule specified

template_dir = "QMC_Templates/" # remember the trailing /

"""
Prerequisites:
====
	1. GAMESS is compiled and "rungms" script location is added to PATH
	2. QMCPACK is compiled and binary location is added to PATH
	3. rungms outputs the .dat file to cwd
	4. $(molecule).inp is present in cwd
	5. Directory containing QMC templates is specified

To do:
----
	1. change all os.system to subprocess.call
"""

import os
import subprocess

class MrData:
	def __init__(self,thisMolecule):
		# Tell Mr. Data the name of the molecule you would like to analyze
		# for example, if the GAMESS input is LiH.inp, then "molecule=LiH"
		self.molecule=thisMolecule

		# some file names
		self.inp = self.molecule + ".inp"
		self.out = self.molecule + ".out"
		self.reinp = self.molecule+"_rerun.inp"
		self.reout = self.molecule+"_rerun.out"
		self.temp = "gamess2qmcpack.tmp"
		self.gamess_err = "gamess.err"
		self.gamess_dat = self.molecule+".dat "+self.molecule+"_rerun.dat"

		self.wfs = self.molecule+"_wfs.xml"
		self.opt_wfs=self.molecule+"_opt_wfs.xml"
		self.ptcl = self.molecule+"_ptcl.xml"
		self.convert_out = self.molecule+"_convert.out"

		self.template_main = template_dir+"main.xml"

	# ======================= rungms ======================= #
	def rungms(self,i):
		# REM use ISPHER=1 at $CONTROL in inp
		gamess_err_file = open(self.gamess_err,'w') # log error from GAMESS runs
		if (i==1 or i==0):
			# run GAMESS for the first time
			# ====
			print("============================================")
			print("Executing 1st GAMESS caulation with "+self.inp)
			print("Check progress in "+self.out)
			with open(self.out,'w') as outfile:
				subprocess.call(["rungms",self.inp],stdout=outfile,stderr=gamess_err_file)
		if (i==2 or i==0):
			# creat reinp with MOREAD as GUESS
			# ====
				# find NORB
				# ----
			print("Looking for OPTIMIZED MCSCF orbits")
			gamess_out=open(self.molecule+".dat",'r') # make sure GAMESS outputs to cwd
			os.system("cp "+self.inp +" "+ self.reinp)
			reinpfile=open(self.reinp,'a')
			in_orbit=False
			in_vec=False
			count_NORB=0
			cur=0
			prev=-1
			for line in gamess_out:
				if line.find("OPTIMIZED MCSCF")!=-1:
					in_orbit=True
				if in_orbit and line.find("$VEC")!=-1:
					in_vec=True
				if in_orbit and in_vec:
					cur = line[:2]
					if cur!=prev:
						count_NORB+=1
					prev=cur
					reinpfile.write(line)
					if line.find("$END")==-1:
						NORB=line[:2]
				if line.find("$END")!=-1:
					in_vec=False
					in_orbit=False
			count_NORB-=2 # first transition is cur=$ prev=-1, last transition is cur=$,prev=last
			if (int(NORB)!=int(count_NORB)%100):
                        	print "Mismatch! NORB="+str(NORB)+" <> count_NORB="+str(count_NORB)
				exit(0)
			else:
				NORB=str(count_NORB)
			reinpfile.close()
				# creat reinp
				# ----
			print("Found "+NORB+" orbits, augmenting GUESS to 'GUESS=MOREAD NORB="+NORB+"' in "+self.reinp)
			os.system("sed 's/GUESS=HCORE/GUESS=MOREAD NORB="+
				NORB+"/g' <"+self.reinp+" >"+self.temp)
			os.system("rm "+self.reinp+"; mv "+self.temp+" "+self.reinp)

			# rerun GAMESS with MO as GUESS
			# ====
			print("============================================")
			print("Executing 2nd GAMESS caulation with "+self.reinp)
			print("Check progress in "+self.reout)

			with open(self.reout,'w') as outfile:
				subprocess.call(["rungms",self.reinp],stdout=outfile,stderr=gamess_err_file)

		gamess_err_file.close()

		print("GAMESS calculations completed")
		print("=========================================")

		# blindly reorganize folder for now, should check for completion first.
		os.system("mkdir gamess;mv *.inp *.out *.dat "+self.gamess_err+" gamess")
		os.system("cp gamess/"+self.reout+" .")
	# ======================= convert4qmc ======================= #
	def convert4qmc(self):
		# Convert GAMESS output
		print("Converting GAMESS output to QMCPACK input")
		with open(self.convert_out,"w") as outfile:
			subprocess.call(["convert4qmc","-gamessAscii",self.reout,"-ci",self.reout,"-threshold","0.0001","-add3BodyJ"],stdout=outfile)

		# folder management
		# ====
		os.system("mv sample.Gaussian-G2.xml "+self.wfs+"; mv sample.Gaussian-G2.ptcl.xml "+self.ptcl)
		os.system("mkdir convert; mv *.xml *.combo.dat "+self.reout+" "+self.convert_out+" convert")

	# ======================= cuspCorrection ======================= #

	# --------------- addOption2QMCline --------------- #
	def addOption2QMCline(self,linehead,option):
		outfile=open(self.temp,'w')
		with open(self.wfs,'r') as infile:
			for line in infile:
				if line.find(linehead)!=-1:
					outfile.write(line[:-2]+option)
				else:
					outfile.write(line)
		outfile.close()
		subprocess.call(["rm",self.wfs])
		subprocess.call(["mv",self.temp,self.wfs])

	# --------------- mvQMCblock --------------- #
	def mvQMCblock(self,blockname,filename):
		block_file = open(filename,'w')
		in_block=False
		outfile=open(self.temp,'w')
		with open(self.wfs,'r') as infile:
			for line in infile:
				if in_block or line.find(blockname)!=-1:
					block_file.write(line)
				else:
					outfile.write(line)
				if line.find(blockname)!=-1:
					in_block=not in_block
		block_file.close()
		outfile.close()
		subprocess.call(["rm",self.wfs])
		subprocess.call(["mv",self.temp,self.wfs])

	# --------------- fileline2file --------------- #
	def fileline2file(self,templatefile,linehead,linefile):
		tmp_file=open(self.temp,'w')
		with open(templatefile,'r') as infile:
			for line in infile:
				if line.find(linehead)!=-1:
					with open(linefile,'r') as replacement:
						tmp_file.write(replacement.read())
				else:
					tmp_file.write(line)
		tmp_file.close()
		subprocess.call(["rm",templatefile])
		subprocess.call(["mv",self.temp,templatefile])

	def add_QMCblock_before(self,linehead,block_file):
		tmp_file=open(self.temp,'w')
		with open(self.wfs,'r') as infile:
			for line in infile:
				if line.find(linehead)!=-1:
					with open(block_file,'r') as block:
						tmp_file.write(block.read())
				tmp_file.write(line)
		tmp_file.close()
		subprocess.call(["rm",self.wfs])
		subprocess.call(["mv",self.temp,self.wfs])

	def cuspCorrection(self):
		# add cuspCorrection and remove jastrow in the wavefunction
		os.system("cp convert/"+self.wfs+" convert/"+self.ptcl+" .")
		self.addOption2QMCline("<determinantset"," cuspCorrection=\"yes\">\n")
		self.mvQMCblock("jastrow","jastrow.xml")
		qmcblock = "cuspCorrection"
		os.system("sed 's/molecule/"+self.molecule+"/g' "+self.template_main+" > "+qmcblock+".xml")
		self.fileline2file(qmcblock+".xml","qmcblock",template_dir+qmcblock+".xml")

		print("============================================")
		print("Executing Cusp Correction")
		print("Check progress in "+qmcblock+".out")
		os.system("export OMP_NUM_THREADS=1")
		with open(qmcblock+".out",'w') as outfile:
			subprocess.call(["qmcapp",qmcblock+".xml"],stdout=outfile)

		# folder management
		os.system("mkdir cusp;mv newOrbs* *.xml wftest.000 eloc.dat *.out cusp")

	# ======================= optJastrow ======================= #
	def optJastrow(self):
		# modify the wavefunction
		os.system("cp cusp/"+self.wfs+" cusp/"+self.ptcl+" cusp/*.cuspInfo.xml cusp/jastrow.xml .")
		self.addOption2QMCline("updet"," cuspInfo=\"spo-up.cuspInfo.xml\">\n")
		self.addOption2QMCline("downdet"," cuspInfo=\"spo-dn.cuspInfo.xml\">\n")
		self.add_QMCblock_before("</wavefunction>","jastrow.xml")

		qmcblock="optJastrow"
		os.system("sed 's/molecule/"+self.molecule+"/g' "+self.template_main+" > "+qmcblock+".xml")
		self.fileline2file(qmcblock+".xml","qmcblock",template_dir+qmcblock+".xml")

		print("============================================")
		print("Optimizing Jastrows")
		print("Check progress in "+qmcblock+".out")

		#with open(qmcblock+".out",'w') as outfile:
			#subprocess.call(["qmcapp",qmcblock+".xml"],stdout=outfile)

		# folder management

	# ======================= rundmc ======================= #
	def rundmc(self):
		# write molecule-specified data to qmc template
		qmcblock="dmc"
		os.system("sed 's/molecule/"+self.molecule+"/g' "+self.template_main+" > "+qmcblock+".xml")
		self.fileline2file(qmcblock+".xml","qmcblock",template_dir+qmcblock+".xml")
		os.system("sed -i 's/"+self.wfs+"/"+self.opt_wfs+"/' "+qmcblock+".xml")

		# run dmc
		#print("============================================")
		#print("Running DMC")
		#print("Check progress in "+qmcblock+".out")
		#with open(qmcblock+".out",'w') as outfile:
		#	subprocess.call(["qmcapp",qmcblock+".xml"],stdout=outfile)

		# folder management
		#new_folder="qmc/"+qmcblock
		#os.system("mkdir -p "+new_folder+";mv newOrbs* J1* J2* "+new_folder+";mv "+qmcblock+".xml "+new_folder+"; mv "+qmcblock+".out "+new_folder)

# ======================= main ======================= #
import argparse

def main():
	parser = argparse.ArgumentParser(description='LtCdr. Data at your service. I can analyze the properties of unknown materials if sufficient information regarding its constituents is available.')
	parser.add_argument("molecule", help="molecule to analyze. (name of the rungms input file should be 'molecule.inp')")
	parser.add_argument("-g", "--gms", type=int, choices=[0,1,2], help="run gamess calculation (you may specify if it's the 1st or 2nd calculation, if 0 is given I'll run both)" )
	parser.add_argument("-c", "--convert", action="store_true", help="convert gamess output to qmc input" )
	parser.add_argument("-cusp", "--cuspCorrection", action="store_true", help="perform cusp correction" )
	parser.add_argument("-j", "--optJastrow", action="store_true", help="perform jastrow optimization" )
	parser.add_argument("-d", "--runDMC", action="store_true", help="run Diffusion Monte Carlo" )
	args = parser.parse_args()
	args = parser.parse_args()

	Data=MrData(args.molecule)

	if args.gms==1:
		Data.rungms(1)
	elif args.gms==2:
		Data.rungms(2)
	elif args.gms==0:
		Data.rungms(0)
	else:
		pass

	if args.convert:
		Data.convert4qmc()
	if args.cuspCorrection:
		Data.cuspCorrection()
	if args.optJastrow:
		Data.optJastrow()
	if args.runDMC:
		Data.rundmc()
main()

<commit_msg>v0.43 Mr. Data now runs optimizaion in a folder<commit_after>#!/usr/bin/env python
# Author: Paul Young
# Last Modified: 07/12/2014
# Objectified autorun script for GAMESS->QMCPACK calculation on the molecule specified

template_dir = "QMC_Templates/" # remember the trailing /

"""
Prerequisites:
====
	1. GAMESS is compiled and "rungms" script location is added to PATH
	2. QMCPACK is compiled and binary location is added to PATH
	3. rungms outputs the .dat file to cwd
	4. $(molecule).inp is present in cwd
	5. Directory containing QMC templates is specified

To do:
----
	1. change all os.system to subprocess.call
"""

import os
import subprocess

class MrData:
	def __init__(self,thisMolecule):
		# Tell Mr. Data the name of the molecule you would like to analyze
		# for example, if the GAMESS input is LiH.inp, then "molecule=LiH"
		self.molecule=thisMolecule

		# some file names
		self.inp = self.molecule + ".inp"
		self.out = self.molecule + ".out"
		self.reinp = self.molecule+"_rerun.inp"
		self.reout = self.molecule+"_rerun.out"
		self.temp = "gamess2qmcpack.tmp"
		self.gamess_err = "gamess.err"
		self.gamess_dat = self.molecule+".dat "+self.molecule+"_rerun.dat"

		self.wfs = self.molecule+"_wfs.xml"
		self.opt_wfs=self.molecule+"_opt_wfs.xml"
		self.ptcl = self.molecule+"_ptcl.xml"
		self.convert_out = self.molecule+"_convert.out"

		self.template_main = template_dir+"main.xml"

	# ======================= rungms ======================= #
	def rungms(self,i):
		# REM use ISPHER=1 at $CONTROL in inp
		gamess_err_file = open(self.gamess_err,'w') # log error from GAMESS runs
		if (i==1 or i==0):
			# run GAMESS for the first time
			# ====
			print("============================================")
			print("Executing 1st GAMESS caulation with "+self.inp)
			print("Check progress in "+self.out)
			with open(self.out,'w') as outfile:
				subprocess.call(["rungms",self.inp],stdout=outfile,stderr=gamess_err_file)
		if (i==2 or i==0):
			# creat reinp with MOREAD as GUESS
			# ====
				# find NORB
				# ----
			print("Looking for OPTIMIZED MCSCF orbits")
			gamess_out=open(self.molecule+".dat",'r') # make sure GAMESS outputs to cwd
			os.system("cp "+self.inp +" "+ self.reinp)
			reinpfile=open(self.reinp,'a')
			in_orbit=False
			in_vec=False
			count_NORB=0
			cur=0
			prev=-1
			for line in gamess_out:
				if line.find("OPTIMIZED MCSCF")!=-1:
					in_orbit=True
				if in_orbit and line.find("$VEC")!=-1:
					in_vec=True
				if in_orbit and in_vec:
					cur = line[:2]
					if cur!=prev:
						count_NORB+=1
					prev=cur
					reinpfile.write(line)
					if line.find("$END")==-1:
						NORB=line[:2]
				if line.find("$END")!=-1:
					in_vec=False
					in_orbit=False
			count_NORB-=2 # first transition is cur=$ prev=-1, last transition is cur=$,prev=last
			if (int(NORB)!=int(count_NORB)%100):
                        	print "Mismatch! NORB="+str(NORB)+" <> count_NORB="+str(count_NORB)
				exit(0)
			else:
				NORB=str(count_NORB)
			reinpfile.close()
				# creat reinp
				# ----
			print("Found "+NORB+" orbits, augmenting GUESS to 'GUESS=MOREAD NORB="+NORB+"' in "+self.reinp)
			os.system("sed 's/GUESS=HCORE/GUESS=MOREAD NORB="+
				NORB+"/g' <"+self.reinp+" >"+self.temp)
			os.system("rm "+self.reinp+"; mv "+self.temp+" "+self.reinp)

			# rerun GAMESS with MO as GUESS
			# ====
			print("Executing 2nd GAMESS caulation with "+self.reinp)
			print("Check progress in "+self.reout)

			with open(self.reout,'w') as outfile:
				subprocess.call(["rungms",self.reinp],stdout=outfile,stderr=gamess_err_file)

		gamess_err_file.close()

		print("GAMESS calculations completed")
		print("=========================================")

		# blindly reorganize folder for now, should check for completion first.
		os.system("mkdir gamess;mv *.inp *.out *.dat "+self.gamess_err+" gamess")
		os.system("cp gamess/"+self.reout+" .")
	# ======================= convert4qmc ======================= #
	def convert4qmc(self):
		# Convert GAMESS output
		print("Converting GAMESS output to QMCPACK input")
		with open(self.convert_out,"w") as outfile:
			subprocess.call(["convert4qmc","-gamessAscii",self.reout,"-ci",self.reout,"-threshold","0.0001","-add3BodyJ"],stdout=outfile)

		# folder management
		# ====
		os.system("mv sample.Gaussian-G2.xml "+self.wfs+"; mv sample.Gaussian-G2.ptcl.xml "+self.ptcl)
		os.system("mkdir convert; mv *.xml *.combo.dat "+self.reout+" "+self.convert_out+" convert")

	# ======================= cuspCorrection ======================= #

	# --------------- addOption2QMCline --------------- #
	def addOption2QMCline(self,file2edit,linehead,option):
		outfile=open(self.temp,'w')
		with open(file2edit,'r') as infile:
			for line in infile:
				if line.find(linehead)!=-1:
					outfile.write(line[:-2]+option)
				else:
					outfile.write(line)
		outfile.close()
		subprocess.call(["rm",file2edit])
		subprocess.call(["mv",self.temp,file2edit])

	# --------------- mvQMCblock --------------- #
	def mvQMCblock(self,blockname,filename):
		block_file = open(filename,'w')
		in_block=False
		outfile=open(self.temp,'w')
		with open(self.wfs,'r') as infile:
			for line in infile:
				if in_block or line.find(blockname)!=-1:
					block_file.write(line)
				else:
					outfile.write(line)
				if line.find(blockname)!=-1:
					in_block=not in_block
		block_file.close()
		outfile.close()
		subprocess.call(["rm",self.wfs])
		subprocess.call(["mv",self.temp,self.wfs])

	# --------------- fileline2file --------------- #
	def fileline2file(self,templatefile,linehead,linefile):
		tmp_file=open(self.temp,'w')
		with open(templatefile,'r') as infile:
			for line in infile:
				if line.find(linehead)!=-1:
					with open(linefile,'r') as replacement:
						tmp_file.write(replacement.read())
				else:
					tmp_file.write(line)
		tmp_file.close()
		subprocess.call(["rm",templatefile])
		subprocess.call(["mv",self.temp,templatefile])

	def add_QMCblock_before(self,file2edit,linehead,block_file):
		tmp_file=open(self.temp,'w')
		with open(file2edit,'r') as infile:
			for line in infile:
				if line.find(linehead)!=-1:
					with open(block_file,'r') as block:
						tmp_file.write(block.read())
				tmp_file.write(line)
		tmp_file.close()
		subprocess.call(["rm",file2edit])
		subprocess.call(["mv",self.temp,file2edit])

	def cuspCorrection(self):
		# add cuspCorrection and remove jastrow in the wavefunction
		os.system("cp convert/"+self.wfs+" convert/"+self.ptcl+" .")
		self.addOption2QMCline(self.wfs,"<determinantset"," cuspCorrection=\"yes\">\n")
		self.mvQMCblock("jastrow","jastrow.xml")
		qmcblock = "cuspCorrection"
		os.system("sed 's/molecule/"+self.molecule+"/g' "+self.template_main+" > "+qmcblock+".xml")
		self.fileline2file(qmcblock+".xml","qmcblock",template_dir+qmcblock+".xml")

		print("Executing Cusp Correction")
		print("Check progress in "+qmcblock+".out")
		os.system("export OMP_NUM_THREADS=1")
		with open(qmcblock+".out",'w') as outfile:
			subprocess.call(["qmcapp",qmcblock+".xml"],stdout=outfile)

		# folder management
		os.system("mkdir cusp;mv newOrbs* *.xml wftest.000 eloc.dat *.out cusp")

	# ======================= optJastrow ======================= #
	def optJastrow(self,tag):
		opt_dir="opt"+tag
		os.system("mkdir "+opt_dir)
		print("Setting up Jastrow Optimization in "+opt_dir)

		# copy necessary files
		os.system("cp cusp/"+self.wfs+" cusp/"+self.ptcl+" cusp/*.cuspInfo.xml cusp/jastrow.xml "+opt_dir)

		# modify the wavefunction
		self.addOption2QMCline(opt_dir+"/"+self.wfs,"updet"," cuspInfo=\"spo-up.cuspInfo.xml\">\n")
		self.addOption2QMCline(opt_dir+"/"+self.wfs,"downdet"," cuspInfo=\"spo-dn.cuspInfo.xml\">\n")
		self.add_QMCblock_before(opt_dir+"/"+self.wfs,"</wavefunction>",opt_dir+"/jastrow.xml")

		qmcblock="optJastrow"
		os.system("sed 's/molecule/"+self.molecule+"/g' "+self.template_main+" > "+qmcblock+".xml")
		self.fileline2file(qmcblock+".xml","qmcblock",template_dir+qmcblock+".xml")
		os.system("mv "+qmcblock+".xml "+opt_dir)
		print("============================================")

		#with open(qmcblock+".out",'w') as outfile:
			#subprocess.call(["qmcapp",qmcblock+".xml"],stdout=outfile)

	# ======================= rundmc ======================= #
	def rundmc(self):
		# write molecule-specified data to qmc template
		print("============================================")
		print("Setting up DMC run")
		qmcblock="dmc"
		os.system("sed 's/molecule/"+self.molecule+"/g' "+self.template_main+" > "+qmcblock+".xml")
		self.fileline2file(qmcblock+".xml","qmcblock",template_dir+qmcblock+".xml")
		os.system("sed -i 's/"+self.wfs+"/"+self.opt_wfs+"/' "+qmcblock+".xml")

# ======================= main ======================= #
import argparse

def main():
	parser = argparse.ArgumentParser(description='LtCdr. Data at your service. I can analyze the properties of unknown materials if sufficient information regarding its constituents is available.')
	parser.add_argument("molecule", help="molecule to analyze. (name of the rungms input file should be 'molecule.inp')")
	parser.add_argument("-g", "--gms", type=int, choices=[0,1,2], help="run gamess calculation (you may specify if it's the 1st or 2nd calculation, if 0 is given I'll run both)" )
	parser.add_argument("-c", "--convert", action="store_true", help="convert gamess output to qmc input" )
	parser.add_argument("-cusp", "--cuspCorrection", action="store_true", help="perform cusp correction" )
	parser.add_argument("-j", "--optJastrow", type=str, help="\'-j tag\' perform jastrow optimization in folder opt$tag" )
	parser.add_argument("-d", "--runDMC", action="store_true", help="run Diffusion Monte Carlo" )
	args = parser.parse_args()
	args = parser.parse_args()

	Data=MrData(args.molecule)

	if args.gms==1:
		Data.rungms(1)
	elif args.gms==2:
		Data.rungms(2)
	elif args.gms==0:
		Data.rungms(0)
	else:
		pass

	if args.convert:
		Data.convert4qmc()
	if args.cuspCorrection:
		Data.cuspCorrection()
	if args.optJastrow:
		Data.optJastrow(args.optJastrow)
	if args.runDMC:
		Data.rundmc()
main()

<|endoftext|>