<commit_before>
# -*- coding: utf-8 -*-
#/***************************************************************************
# Irmt
#                                 A QGIS plugin
# OpenQuake Integrated Risk Modelling Toolkit
#                              -------------------
#        begin                : 2013-10-24
#        copyright            : (C) 2014 by GEM Foundation
#        email                : devops@openquake.org
# ***************************************************************************/
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

import os
from qgis.core import (QgsVectorLayer,
                       QgsFeature,
                       QgsPoint,
                       QgsGeometry,
                       QgsMapLayerRegistry,
                       QgsSymbolV2,
                       QgsVectorGradientColorRampV2,
                       QgsGraduatedSymbolRendererV2,
                       QgsRendererRangeV2,
                       )
from PyQt4.QtCore import pyqtSlot, QDir

from PyQt4.QtGui import (QDialogButtonBox,
                         QDialog,
                         QFileDialog,
                         QColor,
                         )
from openquake.baselib import hdf5
from svir.ui.ui_visualize_oq_output import Ui_VisualizeOqOutputDialog
from svir.utilities.shared import DEBUG
from svir.utilities.utils import LayerEditingManager, WaitCursorManager
from svir.calculations.calculate_utils import add_numeric_attribute

class VisualizeOqOutputDialog(QDialog):
    """
    FIXME
    """
    def __init__(self, iface):
        self.iface = iface
        QDialog.__init__(self)
        # Set up the user interface from Designer.
        self.ui = Ui_VisualizeOqOutputDialog()
        self.ui.setupUi(self)
        # Disable ok_button until all comboboxes are filled
        self.ok_button = self.ui.buttonBox.button(QDialogButtonBox.Ok)
        self.ok_button.setDisabled(True)
        self.ui.open_hdfview_btn.setDisabled(True)

    @pyqtSlot(str)
    def on_hdf5_path_le_textChanged(self):
        self.ui.open_hdfview_btn.setDisabled(
            self.ui.hdf5_path_le.text() == '')

    @pyqtSlot()
    def on_open_hdfview_btn_clicked(self):
        file_path = self.ui.hdf5_path_le.text()
        if file_path:
            to_run = "hdfview " + file_path
            # FIXME make system independent
            os.system(to_run)

    @pyqtSlot()
    def on_file_browser_tbn_clicked(self):
        self.hdf5_path = self.open_file_dialog()

    @pyqtSlot(str)
    def on_rlz_cbx_currentIndexChanged(self):
        self.dataset = self.hmaps.get(self.ui.rlz_cbx.currentText())
        self.imts = {}
        for name in self.dataset.dtype.names[2:]:
            imt, poe = name.split('~')
            if imt not in self.imts:
                self.imts[imt] = [poe]
            else:
                self.imts[imt].append(poe)
        self.ui.imt_cbx.clear()
        self.ui.imt_cbx.setEnabled(True)
        self.ui.imt_cbx.addItems(self.imts.keys())

    @pyqtSlot(str)
    def on_imt_cbx_currentIndexChanged(self):
        imt = self.ui.imt_cbx.currentText()
        self.ui.poe_cbx.clear()
        self.ui.poe_cbx.setEnabled(True)
        self.ui.poe_cbx.addItems(self.imts[imt])

    @pyqtSlot(str)
    def on_poe_cbx_currentIndexChanged(self):
        self.set_ok_button()

    def open_file_dialog(self):
        """
        Open a file dialog to select the data file to be loaded
        """
        text = self.tr('Select oq-engine output to import')
        filters = self.tr('HDF5 files (*.hdf5)')
        hdf5_path = QFileDialog.getOpenFileName(
            self, text, QDir.homePath(), filters)
        if hdf5_path:
            self.hdf5_path = hdf5_path
            self.ui.hdf5_path_le.setText(self.hdf5_path)
            self.populate_rlz_cbx()

    def populate_rlz_cbx(self):
        # FIXME: will the file be closed correctly?
        # with hdf5.File(self.hdf5_path, 'r') as hf:
        self.hfile = hdf5.File(self.hdf5_path, 'r')
        self.hmaps = self.hfile.get('hmaps')
        self.rlzs = self.hmaps.keys()
        self.ui.rlz_cbx.clear()
        self.ui.rlz_cbx.setEnabled(True)
        self.ui.rlz_cbx.addItems(self.rlzs)

    def set_ok_button(self):
        self.ok_button.setEnabled(self.ui.poe_cbx.currentIndex != -1)

    def build_layer(self):
        rlz = self.ui.rlz_cbx.currentText()
        imt = self.ui.imt_cbx.currentText()
        poe = self.ui.poe_cbx.currentText()
        self.field_name = '%s~%s' % (imt, poe)
        array = self.dataset.value[['lon', 'lat', self.field_name]]

        layer_name = "%s_%s_%s" % (rlz, imt, poe)
        # create layer
        self.layer = QgsVectorLayer(
            "Point?crs=epsg:4326", layer_name, "memory")
        # NOTE: if we use shapefiles, we need to make sure ~ is fine,
        #       otherwise we have to replace it with something like _
        # NOTE: add_numeric_attribute uses LayerEditingManager
        add_numeric_attribute(self.field_name, self.layer)
        pr = self.layer.dataProvider()
        with LayerEditingManager(self.layer, 'Reading hdf5', DEBUG):
            feats = []
            for row in array:
                # add a feature
                feat = QgsFeature(self.layer.pendingFields())
                lon, lat, value = row
                # NOTE: without casting to float, it produces a null
                #       because it does not recognize the numpy type
                feat.setAttribute(self.field_name, float(value))
                feat.setGeometry(QgsGeometry.fromPoint(QgsPoint(lon, lat)))
                feats.append(feat)
            (res, outFeats) = pr.addFeatures(feats)
        # add self.layer to the legend
        QgsMapLayerRegistry.instance().addMapLayer(self.layer)
        self.iface.setActiveLayer(self.layer)
        self.iface.zoomToActiveLayer()

    def style_layer(self):
        color1 = QColor("#FFEBEB")
        color2 = QColor("red")
        classes_count = 10
        ramp = QgsVectorGradientColorRampV2(color1, color2)
        symbol = QgsSymbolV2.defaultSymbol(self.layer.geometryType())
        symbol.setAlpha(1)  # opacity
        graduated_renderer = QgsGraduatedSymbolRendererV2.createRenderer(
            self.layer,
            self.field_name,
            classes_count,
            # QgsGraduatedSymbolRendererV2.Quantile,
            QgsGraduatedSymbolRendererV2.EqualInterval,
            symbol,
            ramp)
        graduated_renderer.updateRangeLowerValue(0, 0.0)
        symbol_zeros = QgsSymbolV2.defaultSymbol(self.layer.geometryType())
        symbol_zeros.setColor(QColor("green"))
        symbol_zeros.setAlpha(0.3)  # opacity
        zeros_min = 0.0
        zeros_max = 0.0
        range_zeros = QgsRendererRangeV2(
            zeros_min, zeros_max, symbol_zeros,
            " %.4f - %.4f" % (zeros_min, zeros_max), True)
        graduated_renderer.addClassRange(range_zeros)
        graduated_renderer.moveClass(classes_count, 0)
        self.layer.setRendererV2(graduated_renderer)
        self.layer.setLayerTransparency(30)  # percent
        self.layer.triggerRepaint()
        self.iface.legendInterface().refreshLayerSymbology(
            self.layer)
        self.iface.mapCanvas().refresh()

    def accept(self):
        with WaitCursorManager('Creating layer...', self.iface):
            self.build_layer()
        self.hfile.close()
        self.style_layer()
        self.close()

    # FIXME: also cancel should close the hdf5 file
<commit_msg>set thinner outline and set symbol_zeros color to lighter green<commit_after>
# -*- coding: utf-8 -*-
#/***************************************************************************
# Irmt
#                                 A QGIS plugin
# OpenQuake Integrated Risk Modelling Toolkit
#                              -------------------
#        begin                : 2013-10-24
#        copyright            : (C) 2014 by GEM Foundation
#        email                : devops@openquake.org
# ***************************************************************************/
#
# OpenQuake is free software: you can redistribute it and/or modify it
# under the terms of the GNU Affero General Public License as published
# by the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# OpenQuake is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU Affero General Public License
# along with OpenQuake.  If not, see <http://www.gnu.org/licenses/>.

import os
from qgis.core import (QgsVectorLayer,
                       QgsFeature,
                       QgsPoint,
                       QgsGeometry,
                       QgsMapLayerRegistry,
                       QgsSymbolV2,
                       QgsVectorGradientColorRampV2,
                       QgsGraduatedSymbolRendererV2,
                       QgsRendererRangeV2,
                       )
from PyQt4.QtCore import pyqtSlot, QDir

from PyQt4.QtGui import (QDialogButtonBox,
                         QDialog,
                         QFileDialog,
                         QColor,
                         )
from openquake.baselib import hdf5
from svir.ui.ui_visualize_oq_output import Ui_VisualizeOqOutputDialog
from svir.utilities.shared import DEBUG
from svir.utilities.utils import LayerEditingManager, WaitCursorManager
from svir.calculations.calculate_utils import add_numeric_attribute

class VisualizeOqOutputDialog(QDialog):
    """
    FIXME
    """
    def __init__(self, iface):
        self.iface = iface
        QDialog.__init__(self)
        # Set up the user interface from Designer.
        self.ui = Ui_VisualizeOqOutputDialog()
        self.ui.setupUi(self)
        # Disable ok_button until all comboboxes are filled
        self.ok_button = self.ui.buttonBox.button(QDialogButtonBox.Ok)
        self.ok_button.setDisabled(True)
        self.ui.open_hdfview_btn.setDisabled(True)

    @pyqtSlot(str)
    def on_hdf5_path_le_textChanged(self):
        self.ui.open_hdfview_btn.setDisabled(
            self.ui.hdf5_path_le.text() == '')

    @pyqtSlot()
    def on_open_hdfview_btn_clicked(self):
        file_path = self.ui.hdf5_path_le.text()
        if file_path:
            to_run = "hdfview " + file_path
            # FIXME make system independent
            os.system(to_run)

    @pyqtSlot()
    def on_file_browser_tbn_clicked(self):
        self.hdf5_path = self.open_file_dialog()

    @pyqtSlot(str)
    def on_rlz_cbx_currentIndexChanged(self):
        self.dataset = self.hmaps.get(self.ui.rlz_cbx.currentText())
        self.imts = {}
        for name in self.dataset.dtype.names[2:]:
            imt, poe = name.split('~')
            if imt not in self.imts:
                self.imts[imt] = [poe]
            else:
                self.imts[imt].append(poe)
        self.ui.imt_cbx.clear()
        self.ui.imt_cbx.setEnabled(True)
        self.ui.imt_cbx.addItems(self.imts.keys())

    @pyqtSlot(str)
    def on_imt_cbx_currentIndexChanged(self):
        imt = self.ui.imt_cbx.currentText()
        self.ui.poe_cbx.clear()
        self.ui.poe_cbx.setEnabled(True)
        self.ui.poe_cbx.addItems(self.imts[imt])

    @pyqtSlot(str)
    def on_poe_cbx_currentIndexChanged(self):
        self.set_ok_button()

    def open_file_dialog(self):
        """
        Open a file dialog to select the data file to be loaded
        """
        text = self.tr('Select oq-engine output to import')
        filters = self.tr('HDF5 files (*.hdf5)')
        hdf5_path = QFileDialog.getOpenFileName(
            self, text, QDir.homePath(), filters)
        if hdf5_path:
            self.hdf5_path = hdf5_path
            self.ui.hdf5_path_le.setText(self.hdf5_path)
            self.populate_rlz_cbx()

    def populate_rlz_cbx(self):
        # FIXME: will the file be closed correctly?
        # with hdf5.File(self.hdf5_path, 'r') as hf:
        self.hfile = hdf5.File(self.hdf5_path, 'r')
        self.hmaps = self.hfile.get('hmaps')
        self.rlzs = self.hmaps.keys()
        self.ui.rlz_cbx.clear()
        self.ui.rlz_cbx.setEnabled(True)
        self.ui.rlz_cbx.addItems(self.rlzs)

    def set_ok_button(self):
        self.ok_button.setEnabled(self.ui.poe_cbx.currentIndex != -1)

    def build_layer(self):
        rlz = self.ui.rlz_cbx.currentText()
        imt = self.ui.imt_cbx.currentText()
        poe = self.ui.poe_cbx.currentText()
        self.field_name = '%s~%s' % (imt, poe)
        array = self.dataset.value[['lon', 'lat', self.field_name]]

        layer_name = "%s_%s_%s" % (rlz, imt, poe)
        # create layer
        self.layer = QgsVectorLayer(
            "Point?crs=epsg:4326", layer_name, "memory")
        # NOTE: if we use shapefiles, we need to make sure ~ is fine,
        #       otherwise we have to replace it with something like _
        # NOTE: add_numeric_attribute uses LayerEditingManager
        add_numeric_attribute(self.field_name, self.layer)
        pr = self.layer.dataProvider()
        with LayerEditingManager(self.layer, 'Reading hdf5', DEBUG):
            feats = []
            for row in array:
                # add a feature
                feat = QgsFeature(self.layer.pendingFields())
                lon, lat, value = row
                # NOTE: without casting to float, it produces a null
                #       because it does not recognize the numpy type
                feat.setAttribute(self.field_name, float(value))
                feat.setGeometry(QgsGeometry.fromPoint(QgsPoint(lon, lat)))
                feats.append(feat)
            (res, outFeats) = pr.addFeatures(feats)
        # add self.layer to the legend
        QgsMapLayerRegistry.instance().addMapLayer(self.layer)
        self.iface.setActiveLayer(self.layer)
        self.iface.zoomToActiveLayer()

    def style_layer(self):
        color1 = QColor("#FFEBEB")
        color2 = QColor("red")
        classes_count = 10
        ramp = QgsVectorGradientColorRampV2(color1, color2)
        symbol = QgsSymbolV2.defaultSymbol(self.layer.geometryType())
        # see properties at:
        # https://qgis.org/api/qgsmarkersymbollayerv2_8cpp_source.html#l01073
        symbol = symbol.createSimple({'outline_width': '0.000001'})
        symbol.setAlpha(1)  # opacity
        graduated_renderer = QgsGraduatedSymbolRendererV2.createRenderer(
            self.layer,
            self.field_name,
            classes_count,
            # QgsGraduatedSymbolRendererV2.Quantile,
            QgsGraduatedSymbolRendererV2.EqualInterval,
            symbol,
            ramp)
        graduated_renderer.updateRangeLowerValue(0, 0.0)
        symbol_zeros = QgsSymbolV2.defaultSymbol(self.layer.geometryType())
        symbol_zeros = symbol.createSimple({'outline_width': '0.000001'})
        symbol_zeros.setColor(QColor(222, 255, 222))
        zeros_min = 0.0
        zeros_max = 0.0
        range_zeros = QgsRendererRangeV2(
            zeros_min, zeros_max, symbol_zeros,
            " %.4f - %.4f" % (zeros_min, zeros_max), True)
        graduated_renderer.addClassRange(range_zeros)
        graduated_renderer.moveClass(classes_count, 0)
        self.layer.setRendererV2(graduated_renderer)
        self.layer.setLayerTransparency(30)  # percent
        self.layer.triggerRepaint()
        self.iface.legendInterface().refreshLayerSymbology(
            self.layer)
        self.iface.mapCanvas().refresh()

    def accept(self):
        with WaitCursorManager('Creating layer...', self.iface):
            self.build_layer()
        self.hfile.close()
        self.style_layer()
        self.close()

    # FIXME: also cancel should close the hdf5 file
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# Copyright: 2011, Grigoriy Petukhov
# Author: Grigoriy Petukhov (http://lorien.name)
# License: BSD
import logging
import os
import urllib
from random import randint, choice
from copy import copy
import threading
from urlparse import urljoin
from copy import deepcopy
import time
import re

from html import make_unicode, find_refresh_url
import user_agent

__all__ = ['Grab', 'GrabError', 'DataNotFound']

RE_META_CHARSET = re.compile(r'<meta[^>]+content\s*=\s*[^>]+charset=([-\w]+)', re.I)

GLOBAL_STATE = {'request_counter': 0}
DEFAULT_EXTENSIONS = ['grab.ext.pycurl', 'grab.ext.lxml', 'grab.ext.lxml_form',
                      'grab.ext.django']
logger = logging.getLogger('grab')

class GrabError(IOError):
    pass

class DataNotFound(Exception):
    pass

class ImproperlyConfigured(Exception):
    pass

def default_config():
    return dict(
        timeout = 15,
        connect_timeout = 10,
        proxy = None,
        proxy_type = None,
        proxy_userpwd = None,
        post = None,
        payload = None,
        method = None,
        headers = {},
        charset = 'utf-8',
        user_agent = choice(user_agent.variants),
        reuse_cookies = True,
        reuse_referer = True,
        cookies = {},
        cookiefile = None,
        referer = None,
        guess_encodings = ['windows-1251', 'koi8-r', 'utf-8'],
        log_file = None,
        log_dir = False,
        follow_refresh = False,
        nohead = False,
        nobody = False,
        debug = False,
        debug_post = False,
        encoding = 'gzip',
    )

class Response(object):
    """
    HTTP Response.
    """

    def __init__(self):
        self.status = None
        self.code = None
        self.head = None
        self.body = None
        self.headers =None
        self.time = None
        self.url = None
        self.cookies = None
        self.charset = 'utf-8'

    def parse(self):
        """
        This method is called after Grab instance performes network request.
        """
        self.headers = {}
        for line in self.head.split('\n'):
            line = line.rstrip('\r')
            if line:
                if line.startswith('HTTP'):
                    self.status = line
                else:
                    try:
                        name, value = [x.strip() for x in line.split(':', 1)]
                        self.headers[name] = value
                    except ValueError, ex:
                        logging.error('Invalid header line: %s' % line, exc_info=ex)

        self.detect_charset()

    def detect_charset(self):
        charset = None

        # Try to extract charset from http-equiv meta tag
        if self.body:
            pos = self.body.lower().find('')
            if pos > -1:
                html_head = self.body.lower()[:pos]
                if html_head.find('http-equiv') > -1:
                    try:
                        charset = RE_META_CHARSET.search(html_head).group(1)
                    except AttributeError:
                        pass

        if not charset:
            if 'Content-Type' in self.headers:
                pos = self.headers['Content-Type'].find('charset=')
                if pos > -1:
                    charset = self.headers['Content-Type'][(pos + 8):]

        if charset:
            self.charset = charset

    #@property
    #def unicode_body(self):
        ## TODO: probably should be in Response class
        ## or in new extension which also should include search and assert_pattern methods
        #if not self._unicode_body:
            #self._unicode_body = make_unicode(self.response.body, self.config['guess_encodings'])
        #return self._unicode_body

class Grab(object):
    # Shortcut to grab.GrabError
    Error = GrabError

    # Attributes which should be processed when clone
    # of Grab instance is creating
    clonable_attributes = ['request_headers']

    # Info about loaded extensions
    extensions = []
    transport_extension = None

    """
    Public methods
    """

    def __init__(self, extensions=DEFAULT_EXTENSIONS, extra_extensions=None, **kwargs):
        """
        By default grab instance is initiated with some extensions. You can
        override list of extension in ``extensions`` attribute or specify additional
        extensions in ``extra_extensions`` attribute. All other named arguments will
        be treated as settings.
        """

        if extra_extensions:
            extensions = extensions + extra_extensions
        for mod_path in extensions:
            # Can't win in the fight with relative imports...
            # Doing simple hack
            if mod_path.startswith('grab.'):
                mod_path = mod_path[5:]
            mod = __import__(mod_path, globals(), locals(), ['foo'])
            self.load_extension(mod.Extension)
        self.config = default_config()
        for ext in self.extensions:
            try:
                self.config.update(ext.extra_default_config())
            except AttributeError:
                pass
        self.default_headers = self.common_headers()
        self.trigger_extensions('init')
        self.reset()
        if kwargs:
            self.setup(**kwargs)

    def reset(self):
        """
        Reset all attributes which could be modified during previous request
        or which is not initialized yet if this is the new Grab instance.

        This methods is automatically called before each network request.
        """

        self.response = Response()
        self.request_headers = None
        self.trigger_extensions('reset')

    def clone(self):
        """
        Create clone of Grab instance.

        Cloned instance will have the same state: cookies, referer, response data
        """

        g = Grab()
        g.config = deepcopy(self.config)
        # Important for pycurl transport
        # By default pycurl use own implementation
        # of cookies. So when we create new Grab instance
        # and hence new pycurl instance and configure it
        # with ``config`` of old Grab instance then new
        # pycurl instance does not know
        # anything about cookies (because info about them
        # does not contains in config - it contains in some
        # internal structure of old pycurl instance).
        # What is why we explicitly configure cookies
        # of new pycurl instance - we know where to get them from:
        # cookies are always processed in Response instance
        g.setup(cookies=self.response.cookies)
        g.response = deepcopy(self.response)
        for key in self.clonable_attributes:
            setattr(g, key, getattr(self, key))
        return g

    def setup(self, **kwargs):
        """
        Setting up Grab instance configuration.
        """

        if 'url' in kwargs:
            if self.config.get('url'):
                url = urljoin(self.config['url'], kwargs['url'])
                kwargs['url'] = url
        self.config.update(kwargs)

    def go(self, url):
        """
        Go to ``url``

        Args:
            :url: could be absolute or relative. If relative then t will be appended to the
                absolute URL of previous request.
        """

        return self.request(url=url)

    def request(self, **kwargs):
        """
        Perform network request.

        You can specify grab settings in ``**kwargs``.

        Returns: ``grab.Response`` objects.
        """

        # Reset the state setted by prevous request
        self.reset()
        self.increase_request_counter()
        if kwargs:
            self.setup(**kwargs)
        self.process_config()

        self.transport_extension.request(self)
        if self.config['debug_post']:
            if self.config['post']:
                items = sorted(self.config['post'].items(), key=lambda x: x[0])
                rows = '\n'.join('%-25s: %s' % x for x in items)
                logging.debug('POST request:\n%s\n' % rows)

        # It's vital to delete old POST data after request is performed.
        # If POST data remains when next request will try to use them again!
        # This is not what typical user waits.
        self.old_config = deepcopy(self.config)
        self.config['post'] = None
        self.config['payload'] = None
        self.config['method'] = None

        self.prepare_response()

        # TODO: raise GrabWarning if self.config['http_warnings']
        #if 400 <= self.response_code:
            #raise IOError('Response code is %s: ' % self.response_code)

        if self.config['log_file']:
            open(self.config['log_file'], 'w').write(self.response.body)

        self.save_dumps()

        if self.config['cookiefile']:
            self.dump_cookies(self.config['cookiefile'])

        if self.config['reuse_referer']:
            self.config['referer'] = self.response.url

        # TODO: check max redirect count
        if self.config['follow_refresh']:
            url = find_refresh_url(self.response.body)
            if url:
                return self.request(url=url)

        return self.response

    def search(self, anchor):
        """
        Search for string or regular expression.

        Args:
            :anchor: could be byte string, unicode string or regular expression.

        Returns: first found fragment or ``None``
        """
        if hasattr(anchor, 'finditer'):
            return anchor.search(self.response.body) or None
        else:
            if isinstance(anchor, unicode):
                anchor = anchor.encode(self.config['charset'])
            return anchor if self.response.body.find(anchor) > -1 else None

    def assert_pattern(self, anchor):
        """
        Test that document contains ``anchor`` pattern.

        Args:
            :anchor: could be byte string, unicode string or regular expression.

        If nothing found raise ``DataNotFound`` exception.
        """

        if not self.search(anchor):
            raise DataNotFound(u'Could not found pattern: %s' % anchor)

    def reload(self):
        """
        Make the same network request again.

        All cookies, POST data, headers will be sent again.
        """

        self.go('')

    def repeat_request(self):
        """
        WTF: What the difference with ``reload`` method?
        """

        self.config = self.old_config
        self.request()

    def sleep(self, limit1=None, limit2=None):
        if limit1 is None and limit2 is None:
            limit1 = 0
            limit2 = 1
        elif limit1 is None or limit2 is None:
            raise Exception('sleep function accepts 0 or 2 arguments')
        limit1 = int(limit1 * 1000)
        limit2 = int(limit2 * 1000)
        sleep_time = randint(limit1, limit2) / 1000.0
        logger.debug('Sleeping for %f seconds' % sleep_time)
        time.sleep(sleep_time)

    """
    Private methods
    """

    @property
    def request_counter(self):
        # This counter will used in enumerating network queries.
        # Its value will be displayed in logging messages and also used
        # in names of dumps
        # I use mutable module variable to allow different
        # instances of Grab maintain single counter
        # This could be helpful in debuggin when your script
        # creates multiple Grab instances - in case of shared counter
        # grab instances do not overwrite dump logs
        return GLOBAL_STATE['request_counter']

    def trigger_extensions(self, event):
        for ext in self.extensions:
            try:
                getattr(ext, 'extra_%s' % event)(self)
            except AttributeError:
                pass

    def process_config(self):
        raise ImproperlyConfigured('process_config method should be redefined by transport extension')

    def extract_cookies(self):
           raise ImproperlyConfigured('extract_cookies method should be redefined by transport extension')

    def prepare_response(self):
           raise ImproperlyConfigured('prepare_response method should be redefined by transport extension')

    def load_extension(self, ext_class):
        for attr in ext_class.export_attributes:
            self.add_to_class(attr, ext_class.__dict__[attr])
        ext = ext_class()
        self.extensions.append(ext)
        if getattr(ext, 'transport', None):
            self.transport_extension = ext

    @classmethod
    def add_to_class(cls, name, obj):
        setattr(cls, name, obj)

    def common_headers(self):
        """
        Build headers which sends typical browser.
        """

        return {
            'Accept': 'text/xml,application/xml,application/xhtml+xml,text/html;q=0.9,text/plain;q=0.8,image/png,*/*;q=0.%d' % randint(2, 5),
            'Accept-Language': 'en-us;q=0.%d,en,ru;q=0.%d' % (randint(5, 9), randint(1, 4)),
            'Accept-Charset': 'utf-8,windows-1251;q=0.7,*;q=0.%d' % randint(5, 7),
            'Keep-Alive': '300',
        }

    def increase_request_counter(self):
        GLOBAL_STATE['request_counter'] += 1

    def save_dumps(self):
        if self.config['log_dir']:
            tname = threading.currentThread().getName().lower()
            if tname == 'mainthread':
                tname = ''
            else:
                tname = '-%s' % tname
            fname = os.path.join(self.config['log_dir'], '%02d%s.log' % (self.request_counter, tname))
            open(fname, 'w').write(self.response.head)

            fext = 'html'
            #dirs = self.response_url().split('//')[1].strip().split('/')
            #if len(dirs) > 1:
                #fext = dirs[-1].split('.')[-1]

            fname = os.path.join(self.config['log_dir'], '%02d%s.%s' % (self.request_counter, tname, fext))
            open(fname, 'w').write(self.response.body)

    def urlencode(self, items):
        """
        Smart urlencode which know how to process unicode strings and None values.
        """

        if isinstance(items, dict):
            items = items.items()
        def process(item):
            key, value = item
            if isinstance(value, unicode):
                value = value.encode(self.config['charset'])
            elif value is None:
                value = ''
            return key, value
        return urllib.urlencode(map(process, items))

<commit_msg>Implement hammer mode<commit_after># -*- coding: utf-8 -*-
# Copyright: 2011, Grigoriy Petukhov
# Author: Grigoriy Petukhov (http://lorien.name)
# License: BSD
import logging
import os
import urllib
from random import randint, choice
from copy import copy
import threading
from urlparse import urljoin
from copy import deepcopy
import time
import re

from html import make_unicode, find_refresh_url
import user_agent

__all__ = ['Grab', 'GrabError', 'DataNotFound']

RE_META_CHARSET = re.compile(r'<meta[^>]+content\s*=\s*[^>]+charset=([-\w]+)', re.I)

GLOBAL_STATE = {'request_counter': 0}
DEFAULT_EXTENSIONS = ['grab.ext.pycurl', 'grab.ext.lxml', 'grab.ext.lxml_form',
                      'grab.ext.django']
logger = logging.getLogger('grab')

class GrabError(IOError):
    pass

class DataNotFound(Exception):
    pass

class ImproperlyConfigured(Exception):
    pass

def default_config():
    return dict(
        proxy = None,
        proxy_type = None,
        proxy_userpwd = None,
        post = None,
        payload = None,
        method = None,
        headers = {},
        charset = 'utf-8',
        user_agent = choice(user_agent.variants),
        reuse_cookies = True,
        reuse_referer = True,
        cookies = {},
        cookiefile = None,
        referer = None,
        guess_encodings = ['windows-1251', 'koi8-r', 'utf-8'],
        log_file = None,
        log_dir = False,
        follow_refresh = False,
        nohead = False,
        nobody = False,
        debug = False,
        debug_post = False,
        encoding = 'gzip',
        # Timeouts
        timeout = 15,
        connect_timeout = 10,
        hammer_mode = False,
        hammer_timeouts = ((2, 5), (5, 10), (10, 20), (15, 30))
    )

class Response(object):
    """
    HTTP Response.
    """

    def __init__(self):
        self.status = None
        self.code = None
        self.head = None
        self.body = None
        self.headers =None
        self.time = None
        self.url = None
        self.cookies = None
        self.charset = 'utf-8'

    def parse(self):
        """
        This method is called after Grab instance performes network request.
        """
        self.headers = {}
        for line in self.head.split('\n'):
            line = line.rstrip('\r')
            if line:
                if line.startswith('HTTP'):
                    self.status = line
                else:
                    try:
                        name, value = [x.strip() for x in line.split(':', 1)]
                        self.headers[name] = value
                    except ValueError, ex:
                        logging.error('Invalid header line: %s' % line, exc_info=ex)

        self.detect_charset()

    def detect_charset(self):
        charset = None

        # Try to extract charset from http-equiv meta tag
        if self.body:
            pos = self.body.lower().find('')
            if pos > -1:
                html_head = self.body.lower()[:pos]
                if html_head.find('http-equiv') > -1:
                    try:
                        charset = RE_META_CHARSET.search(html_head).group(1)
                    except AttributeError:
                        pass

        if not charset:
            if 'Content-Type' in self.headers:
                pos = self.headers['Content-Type'].find('charset=')
                if pos > -1:
                    charset = self.headers['Content-Type'][(pos + 8):]

        if charset:
            self.charset = charset

    #@property
    #def unicode_body(self):
        ## TODO: probably should be in Response class
        ## or in new extension which also should include search and assert_pattern methods
        #if not self._unicode_body:
            #self._unicode_body = make_unicode(self.response.body, self.config['guess_encodings'])
        #return self._unicode_body

class Grab(object):
    # Shortcut to grab.GrabError
    Error = GrabError

    # Attributes which should be processed when clone
    # of Grab instance is creating
    clonable_attributes = ['request_headers']

    # Info about loaded extensions
    extensions = []
    transport_extension = None

    """
    Public methods
    """

    def __init__(self, extensions=DEFAULT_EXTENSIONS, extra_extensions=None, **kwargs):
        """
        By default grab instance is initiated with some extensions. You can
        override list of extension in ``extensions`` attribute or specify additional
        extensions in ``extra_extensions`` attribute. All other named arguments will
        be treated as settings.
        """

        if extra_extensions:
            extensions = extensions + extra_extensions
        for mod_path in extensions:
            # Can't win in the fight with relative imports...
            # Doing simple hack
            if mod_path.startswith('grab.'):
                mod_path = mod_path[5:]
            mod = __import__(mod_path, globals(), locals(), ['foo'])
            self.load_extension(mod.Extension)
        self.config = default_config()
        for ext in self.extensions:
            try:
                self.config.update(ext.extra_default_config())
            except AttributeError:
                pass
        self.default_headers = self.common_headers()
        self.trigger_extensions('init')
        self.reset()
        if kwargs:
            self.setup(**kwargs)

    def reset(self):
        """
        Reset all attributes which could be modified during previous request
        or which is not initialized yet if this is the new Grab instance.

        This methods is automatically called before each network request.
        """

        self.response = Response()
        self.request_headers = None
        self.trigger_extensions('reset')

    def clone(self):
        """
        Create clone of Grab instance.

        Cloned instance will have the same state: cookies, referer, response data
        """

        g = Grab()
        g.config = deepcopy(self.config)
        # Important for pycurl transport
        # By default pycurl use own implementation
        # of cookies. So when we create new Grab instance
        # and hence new pycurl instance and configure it
        # with ``config`` of old Grab instance then new
        # pycurl instance does not know
        # anything about cookies (because info about them
        # does not contains in config - it contains in some
        # internal structure of old pycurl instance).
        # What is why we explicitly configure cookies
        # of new pycurl instance - we know where to get them from:
        # cookies are always processed in Response instance
        g.setup(cookies=self.response.cookies)
        g.response = deepcopy(self.response)
        for key in self.clonable_attributes:
            setattr(g, key, getattr(self, key))
        return g

    def setup(self, **kwargs):
        """
        Setting up Grab instance configuration.
        """

        if 'url' in kwargs:
            if self.config.get('url'):
                url = urljoin(self.config['url'], kwargs['url'])
                kwargs['url'] = url
        self.config.update(kwargs)

    def go(self, url):
        """
        Go to ``url``

        Args:
            :url: could be absolute or relative. If relative then t will be appended to the
                absolute URL of previous request.
        """

        return self.request(url=url)

    def request(self, **kwargs):
        """
        Perform network request.

        You can specify grab settings in ``**kwargs``.

        Returns: ``grab.Response`` objects.
        """

        if self.config['hammer_mode']:
            hammer_timeouts = list(self.config['hammer_timeouts'])
            connect_timeout, total_timeout = hammer_timeouts.pop(0)
            self.setup(connect_timeout=connect_timeout, timeout=total_timeout)

        while True:
            try:
                # Reset the state setted by prevous request
                self.reset()
                self.increase_request_counter()
                if kwargs:
                    self.setup(**kwargs)
                self.process_config()

                self.transport_extension.request(self)
            except GrabError, ex:
                # In hammer mode try to use next timeouts
                if self.config['hammer_mode'] and ex[0] == 28:
                    # If not more timeouts
                    # then raise an error
                    if not hammer_timeouts:
                        raise
                    else:
                        connect_timeout, total_timeout = hammer_timeouts.pop(0)
                        self.setup(connect_timeout=connect_timeout, timeout=total_timeout)
                        logging.debug('Trying another timeouts. Connect: %d sec., total: %d sec.' % (connect_timeout, total_timeout))
                # If we are not in hammer mode
                # Then just raise an error
                else:
                    raise
            else:
                # Break the infinite loop in case of success reponse
                break

        if self.config['debug_post']:
            if self.config['post']:
                items = sorted(self.config['post'].items(), key=lambda x: x[0])
                rows = '\n'.join('%-25s: %s' % x for x in items)
                logging.debug('POST request:\n%s\n' % rows)

        # It's vital to delete old POST data after request is performed.
        # If POST data remains when next request will try to use them again!
        # This is not what typical user waits.
        self.old_config = deepcopy(self.config)
        self.config['post'] = None
        self.config['payload'] = None
        self.config['method'] = None

        self.prepare_response()

        # TODO: raise GrabWarning if self.config['http_warnings']
        #if 400 <= self.response_code:
            #raise IOError('Response code is %s: ' % self.response_code)

        if self.config['log_file']:
            open(self.config['log_file'], 'w').write(self.response.body)

        self.save_dumps()

        if self.config['cookiefile']:
            self.dump_cookies(self.config['cookiefile'])

        if self.config['reuse_referer']:
            self.config['referer'] = self.response.url

        # TODO: check max redirect count
        if self.config['follow_refresh']:
            url = find_refresh_url(self.response.body)
            if url:
                return self.request(url=url)

        return self.response

    def search(self, anchor):
        """
        Search for string or regular expression.

        Args:
            :anchor: could be byte string, unicode string or regular expression.

        Returns: first found fragment or ``None``
        """
        if hasattr(anchor, 'finditer'):
            return anchor.search(self.response.body) or None
        else:
            if isinstance(anchor, unicode):
                anchor = anchor.encode(self.config['charset'])
            return anchor if self.response.body.find(anchor) > -1 else None

    def assert_pattern(self, anchor):
        """
        Test that document contains ``anchor`` pattern.

        Args:
            :anchor: could be byte string, unicode string or regular expression.

        If nothing found raise ``DataNotFound`` exception.
        """

        if not self.search(anchor):
            raise DataNotFound(u'Could not found pattern: %s' % anchor)

    def reload(self):
        """
        Make the same network request again.

        All cookies, POST data, headers will be sent again.
        """

        self.go('')

    def repeat_request(self):
        """
        WTF: What the difference with ``reload`` method?
        """

        self.config = self.old_config
        self.request()

    def sleep(self, limit1=None, limit2=None):
        if limit1 is None and limit2 is None:
            limit1 = 0
            limit2 = 1
        elif limit1 is None or limit2 is None:
            raise Exception('sleep function accepts 0 or 2 arguments')
        limit1 = int(limit1 * 1000)
        limit2 = int(limit2 * 1000)
        sleep_time = randint(limit1, limit2) / 1000.0
        logger.debug('Sleeping for %f seconds' % sleep_time)
        time.sleep(sleep_time)

    """
    Private methods
    """

    @property
    def request_counter(self):
        # This counter will used in enumerating network queries.
        # Its value will be displayed in logging messages and also used
        # in names of dumps
        # I use mutable module variable to allow different
        # instances of Grab maintain single counter
        # This could be helpful in debuggin when your script
        # creates multiple Grab instances - in case of shared counter
        # grab instances do not overwrite dump logs
        return GLOBAL_STATE['request_counter']

    def trigger_extensions(self, event):
        for ext in self.extensions:
            try:
                getattr(ext, 'extra_%s' % event)(self)
            except AttributeError:
                pass

    def process_config(self):
        raise ImproperlyConfigured('process_config method should be redefined by transport extension')

    def extract_cookies(self):
           raise ImproperlyConfigured('extract_cookies method should be redefined by transport extension')

    def prepare_response(self):
           raise ImproperlyConfigured('prepare_response method should be redefined by transport extension')

    def load_extension(self, ext_class):
        for attr in ext_class.export_attributes:
            self.add_to_class(attr, ext_class.__dict__[attr])
        ext = ext_class()
        self.extensions.append(ext)
        if getattr(ext, 'transport', None):
            self.transport_extension = ext

    @classmethod
    def add_to_class(cls, name, obj):
        setattr(cls, name, obj)

    def common_headers(self):
        """
        Build headers which sends typical browser.
        """

        return {
            'Accept': 'text/xml,application/xml,application/xhtml+xml,text/html;q=0.9,text/plain;q=0.8,image/png,*/*;q=0.%d' % randint(2, 5),
            'Accept-Language': 'en-us;q=0.%d,en,ru;q=0.%d' % (randint(5, 9), randint(1, 4)),
            'Accept-Charset': 'utf-8,windows-1251;q=0.7,*;q=0.%d' % randint(5, 7),
            'Keep-Alive': '300',
        }

    def increase_request_counter(self):
        GLOBAL_STATE['request_counter'] += 1

    def save_dumps(self):
        if self.config['log_dir']:
            tname = threading.currentThread().getName().lower()
            if tname == 'mainthread':
                tname = ''
            else:
                tname = '-%s' % tname
            fname = os.path.join(self.config['log_dir'], '%02d%s.log' % (self.request_counter, tname))
            open(fname, 'w').write(self.response.head)

            fext = 'html'
            #dirs = self.response_url().split('//')[1].strip().split('/')
            #if len(dirs) > 1:
                #fext = dirs[-1].split('.')[-1]

            fname = os.path.join(self.config['log_dir'], '%02d%s.%s' % (self.request_counter, tname, fext))
            open(fname, 'w').write(self.response.body)

    def urlencode(self, items):
        """
        Smart urlencode which know how to process unicode strings and None values.
        """

        if isinstance(items, dict):
            items = items.items()
        def process(item):
            key, value = item
            if isinstance(value, unicode):
                value = value.encode(self.config['charset'])
            elif value is None:
                value = ''
            return key, value
        return urllib.urlencode(map(process, items))

<|endoftext|>


<commit_before>import subprocess
from subprocess import call
import os
import re
from mem import mem

def announce(msg, ignore_quiet=False):
    if mem.Mode['noisy'] == 'noisy' or ignore_quiet:
        if not mem.quiet:
            print "\n>>> " + msg + "\n"

def mention(msg, ignore_quiet=False):
    if mem.Mode['noisy'] == 'noisy' or ignore_quiet:
        if not mem.quiet:
            print ">> " + msg

def note(msg, ignore_quiet=False):
    if mem.Mode['noisy'] == 'noisy' or ignore_quiet:
        if not mem.quiet:
            print "> " + msg

def quietcall(command):
    if mem.quiet:
        nullhandle = open(os.devnull, 'w')
        call(command, stdout=nullhandle)
        nullhandle.close()
    else:
        call(command)

def find_bridge_ip():

    done = False
    try:
        command = ["ip", "addr", "show", "dev", "docker0"]
        sp = subprocess.Popen(command, stdout=subprocess.PIPE)

        output = subprocess.check_output(('grep', 'inet'), stdin=sp.stdout).strip().split()[1].split('/')[0]

        done = True

    except OSError as e:
        # if e.errno == os.errno.ENOENT:
        #     # handle file not found error.
        #     print "File not found"
        #     done = False
        # else:
        #     # Something else went wrong
        #     raise
        print ("WARN: Could not use 'ip addr | grep inet' to find bridge ip")

    if not done:
        try:
            command = ["ifconfig", "docker0"]

            sp = subprocess.Popen(command, stdout=subprocess.PIPE)

            output = subprocess.check_output(('grep', 'inet'), stdin=sp.stdout).strip().split(':')[1].split()[0]

        except OSError as e:
            print("WARN: Could not use 'ifconfig to find bridge ip")

    if not done:
        print("WARN: Continuing without support for BRIDGE_IP expansion")

    sp.wait()

    # ensure we have a valid ip
    p = re.compile(
            '^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])$')
    if not p.match(output):
        raise Exception(output + " is not a valid IP address for BridgeIP")

    return output

def touch(fname, times=None):
    with open(fname, 'a'):
        os.utime(fname, times)

<commit_msg>make bridge ip optional<commit_after>import subprocess
from subprocess import call
import os
import re
from mem import mem

def announce(msg, ignore_quiet=False):
    if mem.Mode['noisy'] == 'noisy' or ignore_quiet:
        if not mem.quiet:
            print "\n>>> " + msg + "\n"

def mention(msg, ignore_quiet=False):
    if mem.Mode['noisy'] == 'noisy' or ignore_quiet:
        if not mem.quiet:
            print ">> " + msg

def note(msg, ignore_quiet=False):
    if mem.Mode['noisy'] == 'noisy' or ignore_quiet:
        if not mem.quiet:
            print "> " + msg

def quietcall(command):
    if mem.quiet:
        nullhandle = open(os.devnull, 'w')
        call(command, stdout=nullhandle)
        nullhandle.close()
    else:
        call(command)

def find_bridge_ip():

    done = False
    try:
        command = ["ip", "addr", "show", "dev", "docker0"]
        sp = subprocess.Popen(command, stdout=subprocess.PIPE)

        output = subprocess.check_output(('grep', 'inet'), stdin=sp.stdout).strip().split()[1].split('/')[0]

        done = True

    except OSError as e:
        print ("WARN: OSError Could not use 'ip addr | grep inet' to find bridge ip " + e.message)
    except subprocess.CalledProcessError as e:
        print ("WARN: CalledProcessError Could not use 'ip addr | grep inet' to find bridge ip " + e.message)

    if not done:
        try:
            command = ["ifconfig", "docker0"]

            sp = subprocess.Popen(command, stdout=subprocess.PIPE)

            output = subprocess.check_output(('grep', 'inet'), stdin=sp.stdout).strip().split(':')[1].split()[0]

        except OSError as e:
            print("WARN: OSError Could not use 'ifconfig to find bridge ip " + e.message)
        except subprocess.CalledProcessError as e:
            print("WARN: CalledProcessError Could not use 'ifconfig to find bridge ip " + e.message)

    if not done:
        print("WARN: Continuing without support for BRIDGE_IP expansion")

    sp.wait()

    # ensure we have a valid ip
    p = re.compile(
            '^(([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])\.){3}([0-9]|[1-9][0-9]|1[0-9]{2}|2[0-4][0-9]|25[0-5])$')
    if not p.match(output):
        raise Exception(output + " is not a valid IP address for BridgeIP")

    return output

def touch(fname, times=None):
    with open(fname, 'a'):
        os.utime(fname, times)

<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
from __future__ import unicode_literals
from django.db.models import get_model
from django.core.exceptions import ImproperlyConfigured, ObjectDoesNotExist
from django.forms import widgets
from django.template.loader import select_template
from django.utils.module_loading import import_by_path
from django.utils.translation import ugettext_lazy as _
from django.utils.safestring import mark_safe
from cms.plugin_pool import plugin_pool
from cms.utils.compat.dj import python_2_unicode_compatible
from cmsplugin_cascade.fields import PartialFormField
from cmsplugin_cascade.plugin_base import CascadePluginBase
from cmsplugin_cascade.link.forms import LinkForm
from cmsplugin_cascade.link.fields import LinkSearchField
from cmsplugin_cascade.link.plugin_base import LinkPluginBase, LinkElementMixin
from cmsplugin_cascade.utils import resolve_dependencies
from shop import settings as shop_settings
from shop.models.product import ProductModel

class ShopPluginBase(CascadePluginBase):
    module = "Shop"
    require_parent = False
    allow_children = False

@python_2_unicode_compatible
class ShopLinkElementMixin(LinkElementMixin):
    def __str__(self):
        return self.plugin_class.get_identifier(self)

class ShopLinkPluginBase(ShopPluginBase):
    """
    Base plugin for arbitrary buttons used during various checkout pages.
    """
    allow_children = False
    fields = (('link_type', 'cms_page',), 'glossary',)
    glossary_field_map = {'link': ('link_type', 'cms_page',)}
    allow_children = False
    parent_classes = []
    require_parent = False

    class Media:
        js = resolve_dependencies('shop/js/admin/shoplinkplugin.js')

    @classmethod
    def get_link(cls, obj):
        link = obj.glossary.get('link', {})
        if link.get('type') == 'cmspage':
            if 'model' in link and 'pk' in link:
                if not hasattr(obj, '_link_model'):
                    Model = get_model(*link['model'].split('.'))
                    try:
                        obj._link_model = Model.objects.get(pk=link['pk'])
                    except Model.DoesNotExist:
                        obj._link_model = None
                if obj._link_model:
                    return obj._link_model.get_absolute_url()
        else:
            # use the link type as special action keyword
            return link.get('type')

    def get_ring_bases(self):
        bases = super(ShopLinkPluginBase, self).get_ring_bases()
        bases.append('LinkPluginBase')
        return bases

class ShopButtonPluginBase(ShopLinkPluginBase):
    """
    Base plugin for arbitrary buttons used during various checkout pages.
    """
    fields = ('link_content', ('link_type', 'cms_page',), 'glossary',)

    class Media:
        css = {'all': ('cascade/css/admin/bootstrap.min.css', 'cascade/css/admin/bootstrap-theme.min.css',)}
        js = resolve_dependencies('shop/js/admin/shoplinkplugin.js')

    @classmethod
    def get_identifier(cls, obj):
        return mark_safe(obj.glossary.get('link_content', ''))

class CatalogLinkForm(LinkForm):
    """
    Alternative implementation of `cmsplugin_cascade.TextLinkForm`, which allows to link onto
    the Product model, using its method ``get_absolute_url``.

    Note: In this form class the field ``product`` is missing. It is added later, when the shop's
    Product knows about its materialized model.
    """
    LINK_TYPE_CHOICES = (('cmspage', _("CMS Page")), ('product', _("Product")),
                         ('exturl', _("External URL")), ('email', _("Mail To")),)

    def clean_product(self):
        if self.cleaned_data.get('link_type') == 'product':
            app_label = self.ProductModel._meta.app_label
            self.cleaned_data['link_data'] = {
                'type': 'product',
                'model': '{0}.{1}'.format(app_label, self.ProductModel.__name__),
                'pk': self.cleaned_data['product'] and self.cleaned_data['product'].pk or None,
            }

    def set_initial_product(self, initial):
        try:
            Model = get_model(*initial['link']['model'].split('.'))
            initial['product'] = Model.objects.get(pk=initial['link']['pk'])
        except (KeyError, ObjectDoesNotExist):
            pass

    @classmethod
    def get_form_class(cls):
        # must add field `product` on the fly, because during the declaration this form class
        # the MaterializedModel of the product is not known yet.
        product = LinkSearchField(required=False, label='',
            queryset=ProductModel.objects.all(),
            search_fields=getattr(ProductModel, 'search_fields'),
            help_text=_("An internal link onto a product from the shop"))
        return type(str('LinkForm'), (cls,), {'ProductModel': ProductModel, 'product': product})

class CatalogLinkPluginBase(LinkPluginBase):
    """
    Modified implementation of ``cmsplugin_cascade.link.LinkPluginBase`` which adds link type
    "Product", to set links onto arbitrary products of this shop.
    """
#     glossary_fields = (
#         PartialFormField('title',
#             widgets.TextInput(),
#             label=_("Title"),
#             help_text=_("Link's Title")
#         ),
#     ) + LinkPluginBase.glossary_fields
    glossary_field_map = {'link': ('link_type', 'cms_page', 'product', 'ext_url', 'mail_to',)}

    class Media:
        js = resolve_dependencies('shop/js/admin/shoplinkplugin.js')

class DialogFormPluginBase(ShopPluginBase):
    """
    Base class for all plugins adding a dialog form to a placeholder field.
    """
    require_parent = True
    parent_classes = ('BootstrapColumnPlugin',)
    CHOICES = (('form', _("Form dialog")), ('summary', _("Summary")),)
    glossary_fields = (
        PartialFormField('render_type',
            widgets.RadioSelect(choices=CHOICES),
            label=_("Render as"),
            initial='form',
            help_text=_("A dialog can also be rendered as a box containing a read-only summary."),
        ),
    )

    @classmethod
    def register_plugin(cls, plugin):
        """
        Register plugins derived from this class with this function instead of
        `plugin_pool.register_plugin`, so that dialog plugins without a corresponding
        form class are not registered.
        """
        if not issubclass(plugin, cls):
            msg = "Can not register plugin class `{}`, since is does not inherit from `{}`."
            raise ImproperlyConfigured(msg.format(plugin.__name__, cls.__name__))
        if plugin.get_form_class() is None:
            msg = "Can not register plugin class `{}`, since is does not define a `form_class`."
            raise ImproperlyConfigured(msg.format(plugin.__name__))
        plugin_pool.register_plugin(plugin)

    @classmethod
    def get_form_class(cls):
        return getattr(cls, 'form_class', None)

    def __init__(self, *args, **kwargs):
        super(DialogFormPluginBase, self).__init__(*args, **kwargs)
        self.FormClass = import_by_path(self.get_form_class())

    def get_form_data(self, request):
        """
        Returns data to initialize the corresponding dialog form.
        This method must return a dictionary containing either `instance` - a Python object to
        initialize the form class for this plugin, or `initial` - a dictionary containing initial
        form data, or if both are set, values from `initial` override those of `instance`.
        """
        return {}

    def get_render_template(self, context, instance, placeholder):
        template_names = [
            '{0}/checkout/{1}'.format(shop_settings.APP_LABEL, self.template_leaf_name),
            'shop/checkout/{}'.format(self.template_leaf_name),
        ]
        return select_template(template_names)

    def render(self, context, instance, placeholder):
        """
        Return the context to render a DialogFormPlugin
        """
        request = context['request']
        form_data = self.get_form_data(request)
        request._plugin_order = getattr(request, '_plugin_order', 0) + 1
        if not isinstance(form_data.get('initial'), dict):
            form_data['initial'] = {}
        form_data['initial'].update(plugin_id=instance.id, plugin_order=request._plugin_order)
        bound_form = self.FormClass(**form_data)
        context[bound_form.form_name] = bound_form
        return super(DialogFormPluginBase, self).render(context, instance, placeholder)
<commit_msg>Currently ProcessStepPlugin does not work with TansparentMixin<commit_after># -*- coding: utf-8 -*-
from __future__ import unicode_literals
from django.db.models import get_model
from django.core.exceptions import ImproperlyConfigured, ObjectDoesNotExist
from django.forms import widgets
from django.template.loader import select_template
from django.utils.module_loading import import_by_path
from django.utils.translation import ugettext_lazy as _
from django.utils.safestring import mark_safe
from cms.plugin_pool import plugin_pool
from cms.utils.compat.dj import python_2_unicode_compatible
from cmsplugin_cascade.fields import PartialFormField
from cmsplugin_cascade.plugin_base import CascadePluginBase
from cmsplugin_cascade.link.forms import LinkForm
from cmsplugin_cascade.link.fields import LinkSearchField
from cmsplugin_cascade.link.plugin_base import LinkPluginBase, LinkElementMixin
from cmsplugin_cascade.utils import resolve_dependencies
from shop import settings as shop_settings
from shop.models.product import ProductModel

class ShopPluginBase(CascadePluginBase):
    module = "Shop"
    require_parent = False
    allow_children = False

@python_2_unicode_compatible
class ShopLinkElementMixin(LinkElementMixin):
    def __str__(self):
        return self.plugin_class.get_identifier(self)

class ShopLinkPluginBase(ShopPluginBase):
    """
    Base plugin for arbitrary buttons used during various checkout pages.
    """
    allow_children = False
    fields = (('link_type', 'cms_page',), 'glossary',)
    glossary_field_map = {'link': ('link_type', 'cms_page',)}
    allow_children = False
    parent_classes = []
    require_parent = False

    class Media:
        js = resolve_dependencies('shop/js/admin/shoplinkplugin.js')

    @classmethod
    def get_link(cls, obj):
        link = obj.glossary.get('link', {})
        if link.get('type') == 'cmspage':
            if 'model' in link and 'pk' in link:
                if not hasattr(obj, '_link_model'):
                    Model = get_model(*link['model'].split('.'))
                    try:
                        obj._link_model = Model.objects.get(pk=link['pk'])
                    except Model.DoesNotExist:
                        obj._link_model = None
                if obj._link_model:
                    return obj._link_model.get_absolute_url()
        else:
            # use the link type as special action keyword
            return link.get('type')

    def get_ring_bases(self):
        bases = super(ShopLinkPluginBase, self).get_ring_bases()
        bases.append('LinkPluginBase')
        return bases

class ShopButtonPluginBase(ShopLinkPluginBase):
    """
    Base plugin for arbitrary buttons used during various checkout pages.
    """
    fields = ('link_content', ('link_type', 'cms_page',), 'glossary',)

    class Media:
        css = {'all': ('cascade/css/admin/bootstrap.min.css', 'cascade/css/admin/bootstrap-theme.min.css',)}
        js = resolve_dependencies('shop/js/admin/shoplinkplugin.js')

    @classmethod
    def get_identifier(cls, obj):
        return mark_safe(obj.glossary.get('link_content', ''))

class CatalogLinkForm(LinkForm):
    """
    Alternative implementation of `cmsplugin_cascade.TextLinkForm`, which allows to link onto
    the Product model, using its method ``get_absolute_url``.

    Note: In this form class the field ``product`` is missing. It is added later, when the shop's
    Product knows about its materialized model.
    """
    LINK_TYPE_CHOICES = (('cmspage', _("CMS Page")), ('product', _("Product")),
                         ('exturl', _("External URL")), ('email', _("Mail To")),)

    def clean_product(self):
        if self.cleaned_data.get('link_type') == 'product':
            app_label = self.ProductModel._meta.app_label
            self.cleaned_data['link_data'] = {
                'type': 'product',
                'model': '{0}.{1}'.format(app_label, self.ProductModel.__name__),
                'pk': self.cleaned_data['product'] and self.cleaned_data['product'].pk or None,
            }

    def set_initial_product(self, initial):
        try:
            Model = get_model(*initial['link']['model'].split('.'))
            initial['product'] = Model.objects.get(pk=initial['link']['pk'])
        except (KeyError, ObjectDoesNotExist):
            pass

    @classmethod
    def get_form_class(cls):
        # must add field `product` on the fly, because during the declaration this form class
        # the MaterializedModel of the product is not known yet.
        product = LinkSearchField(required=False, label='',
            queryset=ProductModel.objects.all(),
            search_fields=getattr(ProductModel, 'search_fields'),
            help_text=_("An internal link onto a product from the shop"))
        return type(str('LinkForm'), (cls,), {'ProductModel': ProductModel, 'product': product})

class CatalogLinkPluginBase(LinkPluginBase):
    """
    Modified implementation of ``cmsplugin_cascade.link.LinkPluginBase`` which adds link type
    "Product", to set links onto arbitrary products of this shop.
    """
#     glossary_fields = (
#         PartialFormField('title',
#             widgets.TextInput(),
#             label=_("Title"),
#             help_text=_("Link's Title")
#         ),
#     ) + LinkPluginBase.glossary_fields
    glossary_field_map = {'link': ('link_type', 'cms_page', 'product', 'ext_url', 'mail_to',)}

    class Media:
        js = resolve_dependencies('shop/js/admin/shoplinkplugin.js')

class DialogFormPluginBase(ShopPluginBase):
    """
    Base class for all plugins adding a dialog form to a placeholder field.
    """
    require_parent = True
    parent_classes = ('BootstrapColumnPlugin', 'ProcessStepPlugin',)
    CHOICES = (('form', _("Form dialog")), ('summary', _("Summary")),)
    glossary_fields = (
        PartialFormField('render_type',
            widgets.RadioSelect(choices=CHOICES),
            label=_("Render as"),
            initial='form',
            help_text=_("A dialog can also be rendered as a box containing a read-only summary."),
        ),
    )

    @classmethod
    def register_plugin(cls, plugin):
        """
        Register plugins derived from this class with this function instead of
        `plugin_pool.register_plugin`, so that dialog plugins without a corresponding
        form class are not registered.
        """
        if not issubclass(plugin, cls):
            msg = "Can not register plugin class `{}`, since is does not inherit from `{}`."
            raise ImproperlyConfigured(msg.format(plugin.__name__, cls.__name__))
        if plugin.get_form_class() is None:
            msg = "Can not register plugin class `{}`, since is does not define a `form_class`."
            raise ImproperlyConfigured(msg.format(plugin.__name__))
        plugin_pool.register_plugin(plugin)

    @classmethod
    def get_form_class(cls):
        return getattr(cls, 'form_class', None)

    def __init__(self, *args, **kwargs):
        super(DialogFormPluginBase, self).__init__(*args, **kwargs)
        self.FormClass = import_by_path(self.get_form_class())

    def get_form_data(self, request):
        """
        Returns data to initialize the corresponding dialog form.
        This method must return a dictionary containing either `instance` - a Python object to
        initialize the form class for this plugin, or `initial` - a dictionary containing initial
        form data, or if both are set, values from `initial` override those of `instance`.
        """
        return {}

    def get_render_template(self, context, instance, placeholder):
        template_names = [
            '{0}/checkout/{1}'.format(shop_settings.APP_LABEL, self.template_leaf_name),
            'shop/checkout/{}'.format(self.template_leaf_name),
        ]
        return select_template(template_names)

    def render(self, context, instance, placeholder):
        """
        Return the context to render a DialogFormPlugin
        """
        request = context['request']
        form_data = self.get_form_data(request)
        request._plugin_order = getattr(request, '_plugin_order', 0) + 1
        if not isinstance(form_data.get('initial'), dict):
            form_data['initial'] = {}
        form_data['initial'].update(plugin_id=instance.id, plugin_order=request._plugin_order)
        bound_form = self.FormClass(**form_data)
        context[bound_form.form_name] = bound_form
        return super(DialogFormPluginBase, self).render(context, instance, placeholder)
<|endoftext|>


<commit_before>
## Agent architecture using keras rl

### Model
# Agents representation of the environment. ( How the agent thinks the environment works)

#### 1. 256 , 127, 256 are the channels- depth of the first layer, one can be colour, edges)
#### 2. Kernel size is the size of the matrix it will be use to make the convolution ( impair size is better)
#### 3. strides are the jumps the model make ( Ask A)

def create_model(input, actions):
  model = Sequential()
  model.add(Convolutional2D(256, kernel_size=(5,5), input_shape=input))
  model.add(Activation('relu'))

  model.add(Convolution2D(127, kernel_size=(3,3))
  model.add(Activation('relu'))

  model.add(Convolution2D(32, kernel_size=(5,5))
  model.add(Activation('relu'))

  model.add(Flatten())
  model.add(Dense(actions))
  model.add(Activation('softmax')

  model.compile(loss="categorical_crossentropy",
  optimizer="adam",
  metrics=["accuracy"])

  return model

### Policy
# Agent´s behaviour function. How the agent pick actions

policy = LinearAnnealedPolicy(EpsGreedyQPolicy(), attr="eps", value_max=1, value_min=0.7, value_test=.0, nb_steps=1e6)

### Agent
#

dqn = DQNAgent(model=model, nb_actions=action, memory=memory, nb_steps_warmup=50, target_model_update=1e-2, policy=policy)
<commit_msg>Update DQN_Agent.py<commit_after>

from keras.models import Sequential
from keras.layers import Dense, Flatten, Convolution2D
from keras.optimizers import Adam, Adamax, Nadam
import numpy as np

from pysc2.env import sc2_env, environment
from pysc2.lib import actions
from pysc2.lib import features

from rl.memory import SequentialMemory
from rl.policy import LinearAnnealedPolicy, EpsGreedyQPolicy
from rl.core import Processor
from rl.agents.dqn import DQNAgent

## Agent architecture using keras rl

### Model
# Agents representation of the environment. ( How the agent thinks the environment works)

#### 1. 256 , 127, 256 are the channels- depth of the first layer, one can be colour, edges)
#### 2. Kernel size is the size of the matrix it will be use to make the convolution ( impair size is better)
#### 3. strides are the jumps the model make ( Ask A)

def create_model(input, actions):
  model = Sequential()
  model.add(Convolutional2D(256, kernel_size=(5,5), input_shape=input))
  model.add(Activation('relu'))

  model.add(Convolution2D(127, kernel_size=(3,3))
  model.add(Activation('relu'))

  model.add(Convolution2D(32, kernel_size=(5,5))
  model.add(Activation('relu'))

  model.add(Flatten())
  model.add(Dense(actions))
  model.add(Activation('softmax')

  model.compile(loss="categorical_crossentropy",
  optimizer="adam",
  metrics=["accuracy"])

  return model

### Policy
# Agent´s behaviour function. How the agent pick actions

policy = LinearAnnealedPolicy(EpsGreedyQPolicy(), attr="eps", value_max=1, value_min=0.7, value_test=.0, nb_steps=1e6)

### Agent
#

dqn = DQNAgent(model=model, nb_actions=action, memory=memory, nb_steps_warmup=50, target_model_update=1e-2, policy=policy)
<|endoftext|>


<commit_before>from SimPEG import Survey, Problem, Utils, np, sp
from simpegEM.Utils import SrcUtils
from simpegEM.Utils.EMUtils import omega

####################################################
# Receivers
####################################################

class RxFDEM(Survey.BaseRx):

    knownRxTypes = {
                    'exr':['e', 'Ex', 'real'],
                    'eyr':['e', 'Ey', 'real'],
                    'ezr':['e', 'Ez', 'real'],
                    'exi':['e', 'Ex', 'imag'],
                    'eyi':['e', 'Ey', 'imag'],
                    'ezi':['e', 'Ez', 'imag'],

                    'bxr':['b', 'Fx', 'real'],
                    'byr':['b', 'Fy', 'real'],
                    'bzr':['b', 'Fz', 'real'],
                    'bxi':['b', 'Fx', 'imag'],
                    'byi':['b', 'Fy', 'imag'],
                    'bzi':['b', 'Fz', 'imag'],

                    'jxr':['j', 'Fx', 'real'],
                    'jyr':['j', 'Fy', 'real'],
                    'jzr':['j', 'Fz', 'real'],
                    'jxi':['j', 'Fx', 'imag'],
                    'jyi':['j', 'Fy', 'imag'],
                    'jzi':['j', 'Fz', 'imag'],

                    'hxr':['h', 'Ex', 'real'],
                    'hyr':['h', 'Ey', 'real'],
                    'hzr':['h', 'Ez', 'real'],
                    'hxi':['h', 'Ex', 'imag'],
                    'hyi':['h', 'Ey', 'imag'],
                    'hzi':['h', 'Ez', 'imag'],
                   }
    radius = None

    def __init__(self, locs, rxType):
        Survey.BaseRx.__init__(self, locs, rxType)

    @property
    def projField(self):
        """Field Type projection (e.g. e b ...)"""
        return self.knownRxTypes[self.rxType][0]

    @property
    def projGLoc(self):
        """Grid Location projection (e.g. Ex Fy ...)"""
        return self.knownRxTypes[self.rxType][1]

    @property
    def projComp(self):
        """Component projection (real/imag)"""
        return self.knownRxTypes[self.rxType][2]

    def projectFields(self, src, mesh, u):
        P = self.getP(mesh)
        u_part_complex = u[src, self.projField]
        # get the real or imag component
        real_or_imag = self.projComp
        u_part = getattr(u_part_complex, real_or_imag)
        return P*u_part

    def projectFieldsDeriv(self, src, mesh, u, v, adjoint=False):
        P = self.getP(mesh)

        if not adjoint:
            Pv_complex = P * v
            real_or_imag = self.projComp
            Pv = getattr(Pv_complex, real_or_imag)
        elif adjoint:
            Pv_real = P.T * v

            real_or_imag = self.projComp
            if real_or_imag == 'imag':
                Pv = 1j*Pv_real
            elif real_or_imag == 'real':
                Pv = Pv_real.astype(complex)
            else:
                raise NotImplementedError('must be real or imag')

        return Pv

####################################################
# Sources
####################################################

class SrcFDEM(Survey.BaseSrc):
    freq = None
    rxPair = RxFDEM
    knownSrcTypes = ['Simple', 'MagDipole'] #TODO: Do we want to just classify them by Magnetic, Electric, Both?

class SrcFDEM_Simple_e(SrcFDEM):
    """
        Simple electric source. It is defined by the user provided vector S_e

        :param numpy.array S_e: electric source term
        :param float freq: frequency
        :param rxList: receiver list
    """

    def __init__(self, S_e, freq, rxList):
        self.S_e = S_e
        self.freq = float(freq)
        SrcFDEM.__init__(self, None, 'Simple', rxList)

    def getSource(self, prob):
        return None, self.S_e

    def getSourceDeriv(self, prob, v, adjoint = False):
        return None, None

class SrcFDEM_Simple_m(SrcFDEM):
    """
        Simple magnetic source. It is defined by the user provided vector S_m

        :param numpy.array S_m: magnetic source term
        :param float freq: frequency
        :param rxList: receiver list
    """

    def __init__(self, S_m, freq, rxList):
        self.S_m = S_m
        self.freq = float(freq)
        SrcFDEM.__init__(self, None, 'Simple', rxList)

    def getSource(self, prob):
        return self.S_m, None

    def getSourceDeriv(self, prob, v, adjoint = False):
        return None, None

class SrcFDEM_Simple(SrcFDEM):
    """
        Simple source. It is defined by the user provided vectors S_m, S_e

        :param numpy.array S_m: magnetic source term
        :param numpy.array S_e: electric source term
        :param float freq: frequency
        :param rxList: receiver list
    """
    def __init__(self, S_m, S_e, freq, rxList):
        self.S_m = S_m
        self.S_e = S_e
        SrcFDEM.__init__(self, None, 'Simple', rxList)

    def getSource(self, prob):
        return self.S_m, self.S_e

    def getSourceDeriv(self, prob, v, adjoint=None):
        return None, None

class SrcFDEM_MagDipole(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, loc, freq, rxList, orientation='Z'):
        self.freq = float(freq)
        self.orientation = orientation
        SrcFDEM.__init__(self, loc, 'MagDipole', rxList)

    def getSource(self, prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl.T

        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            a = SrcUtils.MagneticDipoleVectorPotential(src.loc, gridY, 'y')

        else:
            srcfct = SrcUtils.MagneticDipoleVectorPotential
            ax = srcfct(self.loc, gridX, 'x')
            ay = srcfct(self.loc, gridY, 'y')
            az = srcfct(self.loc, gridZ, 'z')
            a = np.concatenate((ax, ay, az))

        S_m = -1j*omega(self.freq)*C*a

        return S_m, None

    def getSourceDeriv(self, prob, v, adjoint=None):
        return None, None

class SrcFDEM_MagDipole_Bfield(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, loc, freq, rxList, orientation='Z'):
        self.freq = float(freq)
        self.orientation = orientation
        SrcFDEM.__init__(self, loc, 'MagDipole', rxList)

    def getSource(self, prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl.T

        srcfct = SrcUtils.MagneticDipoleFields
        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            bx = srcfct(self.loc, gridX, 'x')
            bz = srcfct(self.loc, gridZ, 'z')
            b = np.concatenate((bx,bz))
        else:
            bx = srcfct(self.loc, gridX, 'x')
            by = srcfct(self.loc, gridY, 'y')
            bz = srcfct(self.loc, gridZ, 'z')
            b = np.concatenate((bx,by,bz))

        return -1j*omega(self.freq)*b, None

    def getSourceDeriv(self, prob, v, adjoint=None):
        return None, None

class SrcFDEM_CircularLoop(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, loc, freq, rxList, orientation='Z', radius = 1.):
        self.freq = float(freq)
        self.orientation = orientation
        self.radius = radius
        SrcFDEM.__init__(self, loc, 'MagDipole', rxList)

    def getSource(self, prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl.T

        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            a = SrcUtils.MagneticDipoleVectorPotential(src.loc, gridY, 'y', self.radius)

        else:
            srcfct = SrcUtils.MagneticDipoleVectorPotential
            ax = srcfct(self.loc, gridX, 'x', self.radius)
            ay = srcfct(self.loc, gridY, 'y', self.radius)
            az = srcfct(self.loc, gridZ, 'z', self.radius)
            a = np.concatenate((ax, ay, az))

        return -1j*omega(self.freq)*C*a

####################################################
# Survey
####################################################

class SurveyFDEM(Survey.BaseSurvey):
    """
        docstring for SurveyFDEM
    """

    srcPair = SrcFDEM

    def __init__(self, srcList, **kwargs):
        # Sort these by frequency
        self.srcList = srcList
        Survey.BaseSurvey.__init__(self, **kwargs)

        _freqDict = {}
        for src in srcList:
            if src.freq not in _freqDict:
                _freqDict[src.freq] = []
            _freqDict[src.freq] += [src]

        self._freqDict = _freqDict
        self._freqs = sorted([f for f in self._freqDict])

    @property
    def freqs(self):
        """Frequencies"""
        return self._freqs

    @property
    def nFreq(self):
        """Number of frequencies"""
        return len(self._freqDict)

    @property
    def nSrcByFreq(self):
        if getattr(self, '_nSrcByFreq', None) is None:
            self._nSrcByFreq = {}
            for freq in self.freqs:
                self._nSrcByFreq[freq] = len(self.getSources(freq))
        return self._nSrcByFreq

    def getSources(self, freq):
        """Returns the sources associated with a specific frequency."""
        assert freq in self._freqDict, "The requested frequency is not in this survey."
        return self._freqDict[freq]

    def projectFields(self, u):
        data = Survey.Data(self)
        for src in self.srcList:
            for rx in src.rxList:
                data[src, rx] = rx.projectFields(src, self.mesh, u)
        return data

    def projectFieldsDeriv(self, u):
        raise Exception('Use Sources to project fields deriv.')
<commit_msg>Simple --> RawVec<commit_after>from SimPEG import Survey, Problem, Utils, np, sp
from simpegEM.Utils import SrcUtils
from simpegEM.Utils.EMUtils import omega

####################################################
# Receivers
####################################################

class RxFDEM(Survey.BaseRx):

    knownRxTypes = {
                    'exr':['e', 'Ex', 'real'],
                    'eyr':['e', 'Ey', 'real'],
                    'ezr':['e', 'Ez', 'real'],
                    'exi':['e', 'Ex', 'imag'],
                    'eyi':['e', 'Ey', 'imag'],
                    'ezi':['e', 'Ez', 'imag'],

                    'bxr':['b', 'Fx', 'real'],
                    'byr':['b', 'Fy', 'real'],
                    'bzr':['b', 'Fz', 'real'],
                    'bxi':['b', 'Fx', 'imag'],
                    'byi':['b', 'Fy', 'imag'],
                    'bzi':['b', 'Fz', 'imag'],

                    'jxr':['j', 'Fx', 'real'],
                    'jyr':['j', 'Fy', 'real'],
                    'jzr':['j', 'Fz', 'real'],
                    'jxi':['j', 'Fx', 'imag'],
                    'jyi':['j', 'Fy', 'imag'],
                    'jzi':['j', 'Fz', 'imag'],

                    'hxr':['h', 'Ex', 'real'],
                    'hyr':['h', 'Ey', 'real'],
                    'hzr':['h', 'Ez', 'real'],
                    'hxi':['h', 'Ex', 'imag'],
                    'hyi':['h', 'Ey', 'imag'],
                    'hzi':['h', 'Ez', 'imag'],
                   }
    radius = None

    def __init__(self, locs, rxType):
        Survey.BaseRx.__init__(self, locs, rxType)

    @property
    def projField(self):
        """Field Type projection (e.g. e b ...)"""
        return self.knownRxTypes[self.rxType][0]

    @property
    def projGLoc(self):
        """Grid Location projection (e.g. Ex Fy ...)"""
        return self.knownRxTypes[self.rxType][1]

    @property
    def projComp(self):
        """Component projection (real/imag)"""
        return self.knownRxTypes[self.rxType][2]

    def projectFields(self, src, mesh, u):
        P = self.getP(mesh)
        u_part_complex = u[src, self.projField]
        # get the real or imag component
        real_or_imag = self.projComp
        u_part = getattr(u_part_complex, real_or_imag)
        return P*u_part

    def projectFieldsDeriv(self, src, mesh, u, v, adjoint=False):
        P = self.getP(mesh)

        if not adjoint:
            Pv_complex = P * v
            real_or_imag = self.projComp
            Pv = getattr(Pv_complex, real_or_imag)
        elif adjoint:
            Pv_real = P.T * v

            real_or_imag = self.projComp
            if real_or_imag == 'imag':
                Pv = 1j*Pv_real
            elif real_or_imag == 'real':
                Pv = Pv_real.astype(complex)
            else:
                raise NotImplementedError('must be real or imag')

        return Pv

####################################################
# Sources
####################################################

class SrcFDEM(Survey.BaseSrc):
    freq = None
    rxPair = RxFDEM
    knownSrcTypes = ['RawVec', 'MagDipole'] #TODO: remove known source types from base simepeg

class SrcFDEM_RawVec_e(SrcFDEM):
    """
        RawVec electric source. It is defined by the user provided vector S_e

        :param numpy.array S_e: electric source term
        :param float freq: frequency
        :param rxList: receiver list
    """

    def __init__(self, S_e, freq, rxList):
        self.S_e = S_e
        self.freq = float(freq)
        SrcFDEM.__init__(self, None, 'RawVec', rxList)

    def getSource(self, prob):
        return None, self.S_e

    def getSourceDeriv(self, prob, v, adjoint = False):
        return None, None

class SrcFDEM_RawVec_m(SrcFDEM):
    """
        RawVec magnetic source. It is defined by the user provided vector S_m

        :param numpy.array S_m: magnetic source term
        :param float freq: frequency
        :param rxList: receiver list
    """

    def __init__(self, S_m, freq, rxList):
        self.S_m = S_m
        self.freq = float(freq)
        SrcFDEM.__init__(self, None, 'RawVec', rxList)

    def getSource(self, prob):
        return self.S_m, None

    def getSourceDeriv(self, prob, v, adjoint = False):
        return None, None

class SrcFDEM_RawVec(SrcFDEM):
    """
        RawVec source. It is defined by the user provided vectors S_m, S_e

        :param numpy.array S_m: magnetic source term
        :param numpy.array S_e: electric source term
        :param float freq: frequency
        :param rxList: receiver list
    """
    def __init__(self, S_m, S_e, freq, rxList):
        self.S_m = S_m
        self.S_e = S_e
        SrcFDEM.__init__(self, None, 'RawVec', rxList)

    def getSource(self, prob):
        return self.S_m, self.S_e

    def getSourceDeriv(self, prob, v, adjoint=None):
        return None, None

class SrcFDEM_MagDipole(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, loc, freq, rxList, orientation='Z', moment=1.):
        self.freq = float(freq)
        self.orientation = orientation
        SrcFDEM.__init__(self, loc, 'MagDipole', rxList)

    def getSource(self, prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl.T

        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            a = SrcUtils.MagneticDipoleVectorPotential(src.loc, gridY, 'y')

        else:
            srcfct = SrcUtils.MagneticDipoleVectorPotential
            ax = srcfct(self.loc, gridX, 'x')
            ay = srcfct(self.loc, gridY, 'y')
            az = srcfct(self.loc, gridZ, 'z')
            a = np.concatenate((ax, ay, az))

        S_m = -1j*omega(self.freq)*C*a

        return S_m, None

    def getSourceDeriv(self, prob, v, adjoint=None):
        return None, None

class SrcFDEM_MagDipole_Bfield(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, loc, freq, rxList, orientation='Z'):
        self.freq = float(freq)
        self.orientation = orientation
        SrcFDEM.__init__(self, loc, 'MagDipole', rxList)

    def getSource(self, prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl.T

        srcfct = SrcUtils.MagneticDipoleFields
        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            bx = srcfct(self.loc, gridX, 'x')
            bz = srcfct(self.loc, gridZ, 'z')
            b = np.concatenate((bx,bz))
        else:
            bx = srcfct(self.loc, gridX, 'x')
            by = srcfct(self.loc, gridY, 'y')
            bz = srcfct(self.loc, gridZ, 'z')
            b = np.concatenate((bx,by,bz))

        return -1j*omega(self.freq)*b, None

    def getSourceDeriv(self, prob, v, adjoint=None):
        return None, None

class SrcFDEM_CircularLoop(SrcFDEM):

    #TODO: right now, orientation doesn't actually do anything! The methods in SrcUtils should take care of that
    def __init__(self, loc, freq, rxList, orientation='Z', radius = 1.):
        self.freq = float(freq)
        self.orientation = orientation
        self.radius = radius
        SrcFDEM.__init__(self, loc, 'MagDipole', rxList)

    def getSource(self, prob):
        eqLocs = prob._eqLocs

        if eqLocs is 'FE':
            gridX = prob.mesh.gridEx
            gridY = prob.mesh.gridEy
            gridZ = prob.mesh.gridEz
            C = prob.mesh.edgeCurl

        elif eqLocs is 'EF':
            gridX = prob.mesh.gridFx
            gridY = prob.mesh.gridFy
            gridZ = prob.mesh.gridFz
            C = prob.mesh.edgeCurl.T

        if prob.mesh._meshType is 'CYL':
            if not prob.mesh.isSymmetric:
                # TODO ?
                raise NotImplementedError('Non-symmetric cyl mesh not implemented yet!')
            a = SrcUtils.MagneticDipoleVectorPotential(src.loc, gridY, 'y', self.radius)

        else:
            srcfct = SrcUtils.MagneticDipoleVectorPotential
            ax = srcfct(self.loc, gridX, 'x', self.radius)
            ay = srcfct(self.loc, gridY, 'y', self.radius)
            az = srcfct(self.loc, gridZ, 'z', self.radius)
            a = np.concatenate((ax, ay, az))

        return -1j*omega(self.freq)*C*a

####################################################
# Survey
####################################################

class SurveyFDEM(Survey.BaseSurvey):
    """
        docstring for SurveyFDEM
    """

    srcPair = SrcFDEM

    def __init__(self, srcList, **kwargs):
        # Sort these by frequency
        self.srcList = srcList
        Survey.BaseSurvey.__init__(self, **kwargs)

        _freqDict = {}
        for src in srcList:
            if src.freq not in _freqDict:
                _freqDict[src.freq] = []
            _freqDict[src.freq] += [src]

        self._freqDict = _freqDict
        self._freqs = sorted([f for f in self._freqDict])

    @property
    def freqs(self):
        """Frequencies"""
        return self._freqs

    @property
    def nFreq(self):
        """Number of frequencies"""
        return len(self._freqDict)

    @property
    def nSrcByFreq(self):
        if getattr(self, '_nSrcByFreq', None) is None:
            self._nSrcByFreq = {}
            for freq in self.freqs:
                self._nSrcByFreq[freq] = len(self.getSources(freq))
        return self._nSrcByFreq

    def getSources(self, freq):
        """Returns the sources associated with a specific frequency."""
        assert freq in self._freqDict, "The requested frequency is not in this survey."
        return self._freqDict[freq]

    def projectFields(self, u):
        data = Survey.Data(self)
        for src in self.srcList:
            for rx in src.rxList:
                data[src, rx] = rx.projectFields(src, self.mesh, u)
        return data

    def projectFieldsDeriv(self, u):
        raise Exception('Use Sources to project fields deriv.')
<|endoftext|>


<commit_before># Copyright 2021 Google Research.
#
# 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.

# Lint as: python3
"""Validates all utterance templates added to this directory.

Verifies that all rows have 2 columns separated by tabs and the number of slot
placeholders match in the template key and the template string.
"""

import os

from absl.testing import absltest
from absl.testing import parameterized

def get_template_paths():
  template_dir = os.path.dirname(os.path.abspath(__file__))
  template_paths = []
  for root, _, files in os.walk(template_dir):
    template_paths.extend(
        os.path.join(root, f) for f in files if f.endswith(".tsv"))
  return template_paths

@parameterized.parameters(get_template_paths())
class TemplateSanityTest(parameterized.TestCase):

  def test_valid_rows(self, template_path):
    with open(template_path) as f:
      for line in f:
        vals = line.strip().split("\t")
        self.assertLen(vals, 2)
        self.assertEqual(vals[0].count("@"), vals[1].count("@"))

if __name__ == "__main__":
  absltest.main()
<commit_msg>Remove unused comments related to Python 2 compatibility.<commit_after># Copyright 2021 Google Research.
#
# 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.

"""Validates all utterance templates added to this directory.

Verifies that all rows have 2 columns separated by tabs and the number of slot
placeholders match in the template key and the template string.
"""

import os

from absl.testing import absltest
from absl.testing import parameterized

def get_template_paths():
  template_dir = os.path.dirname(os.path.abspath(__file__))
  template_paths = []
  for root, _, files in os.walk(template_dir):
    template_paths.extend(
        os.path.join(root, f) for f in files if f.endswith(".tsv"))
  return template_paths

@parameterized.parameters(get_template_paths())
class TemplateSanityTest(parameterized.TestCase):

  def test_valid_rows(self, template_path):
    with open(template_path) as f:
      for line in f:
        vals = line.strip().split("\t")
        self.assertLen(vals, 2)
        self.assertEqual(vals[0].count("@"), vals[1].count("@"))

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


<commit_before># encoding: utf-8

from datetime import datetime

__author__ = "Tsuyoshi Hombashi"
__copyright__ = "Copyright 2016-{}, {}".format(datetime.now().year, __author__)
__license__ = "MIT License"
__version__ = "0.45.5"
__maintainer__ = __author__
__email__ = "tsuyoshi.hombashi@gmail.com"
<commit_msg>Update __copyright__<commit_after># encoding: utf-8

__author__ = "Tsuyoshi Hombashi"
__copyright__ = "Copyright 2016, {}".format(__author__)
__license__ = "MIT License"
__version__ = "0.45.5"
__maintainer__ = __author__
__email__ = "tsuyoshi.hombashi@gmail.com"
<|endoftext|>