<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2016 GEM Foundation
#
# 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 Affero 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
import cgi
import time
from datetime import date, datetime, timedelta
import itertools
from docutils.examples import html_parts

from openquake.commonlib.datastore import read
from openquake.calculators.views import view_fullreport
from openquake.calculators import base  # needed for registering task_info
from openquake.engine.logs import dbcmd

tablecounter = itertools.count(0)

def html(header_rows):
    """
    Convert a list of tuples describing a table into a HTML string
    """
    name = 'table%d' % next(tablecounter)
    return HtmlTable([map(str, row) for row in header_rows], name).render()

class HtmlTable(object):
    """
    Convert a sequence header+body into a HTML table.
    """
    css = """\
    tr.evenRow { background-color: lightgreen }
    tr.oddRow { }
    th { background-color: lightblue }
    """
    maxrows = 5000
    border = "1"
    summary = ""

    def __init__(self, header_plus_body, name='noname',
                 empty_table='Empty table'):
        header, body = header_plus_body[0], header_plus_body[1:]
        self.name = name
        self.empty_table = empty_table
        rows = []  # rows is a finite sequence of tuples
        for i, row in enumerate(body):
            if i == self.maxrows:
                rows.append(
                    ["Table truncated because too big: more than %s rows" % i])
                break
            rows.append(row)
        self.rows = rows
        self.header = tuple(header)  # horizontal header

    def render(self, dummy_ctxt=None):
        out = "\n%s\n" % "".join(list(self._gen_table()))
        if not self.rows:
            out += '%s' % cgi.escape(self.empty_table, quote=True)
        return out

    def _gen_table(self):
        yield '
\n'\
              % (self.name, self.border, self.summary)
        yield '
\n'
        yield '
%s
\n' % ''.join(
            '%s	\n' % h for h in self.header)
        yield '
\n'
        yield '
'
        for r, row in enumerate(self.rows):
            yield '
\n' % ["even", "odd"][r % 2]
            for col in row:
                yield '%s	\n' % col
            yield '
\n'
        yield '
\n'
        yield '
\n'

JOB_STATS = '''
SELECT id, user_name, start_time, stop_time, status,
strftime('%s', stop_time) - strftime('%s', start_time) AS duration
FROM job WHERE id=?x;
'''

ALL_JOBS = '''
SELECT id, user_name, status, ds_calc_dir FROM job
WHERE start_time >= ?x AND start_time < ?x ORDER BY stop_time
'''

PAGE_TEMPLATE = '''\
<!DOCTYPE html>

<meta charset="UTF-8">

<link rel="stylesheet" href="http://ajax.googleapis.com/ajax/libs/jqueryui/1.11.2/themes/smoothness/jquery-ui.css" />

%s

'''

def make_tabs(tag_ids, tag_status, tag_contents):
    """
    Return a HTML string containing all the tabs we want to display
    """
    templ = '''

%s

%s

'''
    lis = []
    contents = []
    for i, (tag_id, status, tag_content) in enumerate(
            zip(tag_ids, tag_status, tag_contents), 1):
        mark = '.' if status == 'complete' else '!'
        lis.append('
%s%s
' % (i, tag_id, mark))
        contents.append('
%s
' % (
            i, tag_content))
    return templ % ('\n'.join(lis), '\n'.join(contents))

def make_report(isodate='today'):
    """
    Build a HTML report with the computations performed at the given isodate.
    Return the name of the report, which is saved in the current directory.
    """
    if isodate == 'today':
        isodate = date.today()
    else:
        isodate = date(*time.strptime(isodate, '%Y-%m-%d')[:3])
    isodate1 = isodate + timedelta(1)  # +1 day

    tag_ids = []
    tag_status = []
    tag_contents = []

    # the fetcher returns an header which is stripped with [1:]
    jobs = dbcmd(
        'fetch', ALL_JOBS, isodate.isoformat(), isodate1.isoformat())
    page = '
%d job(s) finished before midnight of %s
' % (
        len(jobs), isodate)
    for job_id, user, status, ds_calc in jobs:
        tag_ids.append(job_id)
        tag_status.append(status)
        [stats] = dbcmd('fetch', JOB_STATS, job_id)
        (job_id, user, start_time, stop_time, status, duration) = stats
        try:
            ds = read(job_id, datadir=os.path.dirname(ds_calc))
            txt = view_fullreport('fullreport', ds)
            report = html_parts(txt)
        except Exception as exc:
            report = dict(
                html_title='Could not generate report: %s' % cgi.escape(
                    exc, quote=True),
                fragment='')
        page = report['html_title']
        page += html([stats._fields, stats])
        page += report['fragment']
        tag_contents.append(page)

    page = make_tabs(tag_ids, tag_status, tag_contents) + (
        'Report last updated: %s' % datetime.now())
    fname = 'jobs-%s.html' % isodate
    with open(fname, 'w') as f:
        f.write(PAGE_TEMPLATE % page.encode('utf-8'))
    return fname
<commit_msg>Fixed make_html_report<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2014-2016 GEM Foundation
#
# 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 Affero 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
import cgi
import time
from datetime import date, datetime, timedelta
import itertools
from docutils.examples import html_parts

from openquake.commonlib.datastore import read
from openquake.calculators.views import view_fullreport
from openquake.calculators import base  # needed for registering task_info
from openquake.engine.logs import dbcmd

tablecounter = itertools.count(0)

def html(header_rows):
    """
    Convert a list of tuples describing a table into a HTML string
    """
    name = 'table%d' % next(tablecounter)
    return HtmlTable([map(str, row) for row in header_rows], name).render()

class HtmlTable(object):
    """
    Convert a sequence header+body into a HTML table.
    """
    css = """\
    tr.evenRow { background-color: lightgreen }
    tr.oddRow { }
    th { background-color: lightblue }
    """
    maxrows = 5000
    border = "1"
    summary = ""

    def __init__(self, header_plus_body, name='noname',
                 empty_table='Empty table'):
        header, body = header_plus_body[0], header_plus_body[1:]
        self.name = name
        self.empty_table = empty_table
        rows = []  # rows is a finite sequence of tuples
        for i, row in enumerate(body):
            if i == self.maxrows:
                rows.append(
                    ["Table truncated because too big: more than %s rows" % i])
                break
            rows.append(row)
        self.rows = rows
        self.header = tuple(header)  # horizontal header

    def render(self, dummy_ctxt=None):
        out = "\n%s\n" % "".join(list(self._gen_table()))
        if not self.rows:
            out += '%s' % cgi.escape(self.empty_table, quote=True)
        return out

    def _gen_table(self):
        yield '
\n'\
              % (self.name, self.border, self.summary)
        yield '
\n'
        yield '
%s
\n' % ''.join(
            '%s	\n' % h for h in self.header)
        yield '
\n'
        yield '
'
        for r, row in enumerate(self.rows):
            yield '
\n' % ["even", "odd"][r % 2]
            for col in row:
                yield '%s	\n' % col
            yield '
\n'
        yield '
\n'
        yield '
\n'

JOB_STATS = '''
SELECT id, user_name, start_time, stop_time, status,
strftime('%s', stop_time) - strftime('%s', start_time) AS duration
FROM job WHERE id=?x;
'''

ALL_JOBS = '''
SELECT id, user_name, status, ds_calc_dir FROM job
WHERE start_time >= ?x AND start_time < ?x ORDER BY stop_time
'''

PAGE_TEMPLATE = '''\
<!DOCTYPE html>

<meta charset="UTF-8">

<link rel="stylesheet" href="http://ajax.googleapis.com/ajax/libs/jqueryui/1.11.2/themes/smoothness/jquery-ui.css" />

%s

'''

def make_tabs(tag_ids, tag_status, tag_contents):
    """
    Return a HTML string containing all the tabs we want to display
    """
    templ = '''

%s

%s

'''
    lis = []
    contents = []
    for i, (tag_id, status, tag_content) in enumerate(
            zip(tag_ids, tag_status, tag_contents), 1):
        mark = '.' if status == 'complete' else '!'
        lis.append('
%s%s
' % (i, tag_id, mark))
        contents.append('
%s
' % (
            i, tag_content))
    return templ % ('\n'.join(lis), '\n'.join(contents))

def make_report(isodate='today'):
    """
    Build a HTML report with the computations performed at the given isodate.
    Return the name of the report, which is saved in the current directory.
    """
    if isodate == 'today':
        isodate = date.today()
    else:
        isodate = date(*time.strptime(isodate, '%Y-%m-%d')[:3])
    isodate1 = isodate + timedelta(1)  # +1 day

    tag_ids = []
    tag_status = []
    tag_contents = []

    # the fetcher returns an header which is stripped with [1:]
    jobs = dbcmd(
        'fetch', ALL_JOBS, isodate.isoformat(), isodate1.isoformat())
    page = '
%d job(s) finished before midnight of %s
' % (
        len(jobs), isodate)
    for job_id, user, status, ds_calc in jobs:
        tag_ids.append(job_id)
        tag_status.append(status)
        [stats] = dbcmd('fetch', JOB_STATS, job_id)
        (job_id, user, start_time, stop_time, status, duration) = stats
        try:
            ds = read(job_id, datadir=os.path.dirname(ds_calc))
            txt = view_fullreport('fullreport', ds)
            report = html_parts(txt)
        except Exception as exc:
            report = dict(
                html_title='Could not generate report: %s' % cgi.escape(
                    str(exc), quote=True),
                fragment='')
        page = report['html_title']
        page += html([stats._fields, stats])
        page += report['fragment']
        tag_contents.append(page)

    page = make_tabs(tag_ids, tag_status, tag_contents) + (
        'Report last updated: %s' % datetime.now())
    fname = 'jobs-%s.html' % isodate
    with open(fname, 'w') as f:
        f.write(PAGE_TEMPLATE % page.encode('utf-8'))
    return fname
<|endoftext|>


<commit_before>import cgi
import datetime
import decimal
import itertools
from docutils.examples import html_parts

from openquake.commonlib.datastore import DataStore
from openquake.calculators.views import view_fullreport

tablecounter = itertools.count(0)

def html(header_rows):
    """
    Convert a list of tuples describing a table into a HTML string
    """
    name = 'table%d' % next(tablecounter)
    return HtmlTable([map(fmt, row) for row in header_rows], name).render()

def fmt(number):
    """
    Format numbers in a nice way
    """
    if isinstance(number, datetime.datetime):
        return number.isoformat()[:19]
    elif isinstance(number, datetime.timedelta):
        return str(number)[:7]
    elif isinstance(number, (float, decimal.Decimal)):
        if number > 1000:
            return '{:,d}'.format(int(round(number)))
        elif number < 10:
            return '%6.3f' % number
        else:
            return str(int(number))
    return str(number)

class Fetcher(object):
    """
    Small wrapper over a DB API 2 cursor
    """
    def __init__(self, curs):
        self.curs = curs

    def query(self, templ, *args):
        self.curs.execute(templ, args)
        header = [r[0] for r in self.curs.description]
        return [header] + self.curs.fetchall()

class HtmlTable(object):
    """
    Convert a sequence header+body into a HTML table.
    """
    css = """\
    tr.evenRow { background-color: lightgreen }
    tr.oddRow { }
    th { background-color: lightblue }
    """
    maxrows = 5000
    border = "1"
    summary = ""

    def __init__(self, header_plus_body, name='noname',
                 empty_table='Empty table'):
        header, body = header_plus_body[0], header_plus_body[1:]
        self.name = name
        self.empty_table = empty_table
        rows = []  # rows is a finite sequence of tuples
        for i, row in enumerate(body):
            if i == self.maxrows:
                rows.append(
                    ["Table truncated because too big: more than %s rows" % i])
                break
            rows.append(row)
        self.rows = rows
        self.header = tuple(header)  # horizontal header

    def render(self, dummy_ctxt=None):
        out = "\n%s\n" % "".join(list(self._gen_table()))
        if not self.rows:
            out += '%s' % cgi.escape(self.empty_table, quote=True)
        return out

    def _gen_table(self):
        yield '
\n'\
              % (self.name, self.border, self.summary)
        yield '
\n'
        yield '
%s
\n' % ''.join(
            '%s	\n' % h for h in self.header)
        yield '
\n'
        yield '
'
        for r, row in enumerate(self.rows):
            yield '
\n' % ["even", "odd"][r % 2]
            for col in row:
                yield '%s	\n' % col
            yield '
\n'
        yield '
\n'
        yield '
\n'

JOB_STATS = '''
SELECT oq_job_id, user_name, stop_time, status, duration FROM (
  SELECT oq_job_id, user_name, stop_time, status,
  stop_time - start_time as duration
  FROM uiapi.job_stats AS s
  INNER JOIN uiapi.oq_job AS o
  ON s.oq_job_id=o.id) AS x
WHERE oq_job_id=%s;
'''

JOB_INFO = '''
SELECT * FROM uiapi.job_info WHERE oq_job_id=%s;
'''

ALL_JOBS = '''
SELECT s.oq_job_id, o.user_name, status FROM uiapi.job_stats AS s
INNER JOIN uiapi.oq_job AS o
ON o.id=s.oq_job_id
WHERE stop_time::date = %s OR stop_time IS NULL AND start_time >= %s
ORDER BY stop_time
'''

PAGE_TEMPLATE = '''\

<link rel="stylesheet" href="http://ajax.googleapis.com/ajax/libs/jqueryui/1.11.2/themes/smoothness/jquery-ui.css" />

%s

'''

def make_tabs(tag_ids, tag_status, tag_contents):
    """
    Return a HTML string containing all the tabs we want to display
    """
    templ = '''

%s

%s

'''
    lis = []
    contents = []
    for i, (tag_id, status, tag_content) in enumerate(
            zip(tag_ids, tag_status, tag_contents), 1):
        mark = '.' if status == 'complete' else '!'
        lis.append('
%s%s
' % (i, tag_id, mark))
        contents.append('
%s
' % (
            i, tag_content))
    return templ % ('\n'.join(lis), '\n'.join(contents))

def make_report(conn, isodate='today'):
    """
    Build a HTML report with the computations performed at the given isodate.
    Return the name of the report, which is saved in the current directory.
    """
    if isodate == 'today':
        isodate = datetime.date.today().isoformat()
    curs = conn.cursor()
    fetcher = Fetcher(curs)
    tag_ids = []
    tag_status = []
    tag_contents = []

    jobs = fetcher.query(ALL_JOBS, isodate, isodate)[1:]
    page = '
%d job(s) finished before midnight of %s
' % (
        len(jobs), isodate)
    for job_id, user, status in jobs:
        tag_ids.append(job_id)
        tag_status.append(status)
        stats = fetcher.query(JOB_STATS, job_id)[1:]
        if not stats:
            continue
        (job_id, user, stop_time, status, duration) = stats[0]

        ds = DataStore(job_id, mode='r')
        try:
            report = html_parts(view_fullreport('fullreport', ds))
        except Exception as exc:
            report = dict(
                html_title='Could not generate report: %s' % cgi.escape(
                    str(exc), quote=True),
                fragment='')

        page = report['html_title']

        if 'realizations' in ds:
            tot_rlzs = len(ds['realizations'])
        else:
            tot_rlzs = 1
        data = fetcher.query(JOB_INFO, job_id)
        if data[1:]:
            info_rows = (
                '
Job Info for calculation=%s, %s
'
                '
owner: %s, realizations: %d
' % (
                    job_id, ds.attrs['description'].decode('utf-8'), user, tot_rlzs,
                )) + html(data)
            page += info_rows

        job_stats = html(fetcher.query(JOB_STATS, job_id))
        page += job_stats

        page += report['fragment']

        tag_contents.append(page)

    page = make_tabs(tag_ids, tag_status, tag_contents) + (
        'Report last updated: %s' % datetime.datetime.now())
    fname = 'jobs-%s.html' % isodate
    with open(fname, 'w') as f:
        f.write(PAGE_TEMPLATE % page)
    return fname
<commit_msg>Indentation change<commit_after>import cgi
import datetime
import decimal
import itertools
from docutils.examples import html_parts

from openquake.commonlib.datastore import DataStore
from openquake.calculators.views import view_fullreport

tablecounter = itertools.count(0)

def html(header_rows):
    """
    Convert a list of tuples describing a table into a HTML string
    """
    name = 'table%d' % next(tablecounter)
    return HtmlTable([map(fmt, row) for row in header_rows], name).render()

def fmt(number):
    """
    Format numbers in a nice way
    """
    if isinstance(number, datetime.datetime):
        return number.isoformat()[:19]
    elif isinstance(number, datetime.timedelta):
        return str(number)[:7]
    elif isinstance(number, (float, decimal.Decimal)):
        if number > 1000:
            return '{:,d}'.format(int(round(number)))
        elif number < 10:
            return '%6.3f' % number
        else:
            return str(int(number))
    return str(number)

class Fetcher(object):
    """
    Small wrapper over a DB API 2 cursor
    """
    def __init__(self, curs):
        self.curs = curs

    def query(self, templ, *args):
        self.curs.execute(templ, args)
        header = [r[0] for r in self.curs.description]
        return [header] + self.curs.fetchall()

class HtmlTable(object):
    """
    Convert a sequence header+body into a HTML table.
    """
    css = """\
    tr.evenRow { background-color: lightgreen }
    tr.oddRow { }
    th { background-color: lightblue }
    """
    maxrows = 5000
    border = "1"
    summary = ""

    def __init__(self, header_plus_body, name='noname',
                 empty_table='Empty table'):
        header, body = header_plus_body[0], header_plus_body[1:]
        self.name = name
        self.empty_table = empty_table
        rows = []  # rows is a finite sequence of tuples
        for i, row in enumerate(body):
            if i == self.maxrows:
                rows.append(
                    ["Table truncated because too big: more than %s rows" % i])
                break
            rows.append(row)
        self.rows = rows
        self.header = tuple(header)  # horizontal header

    def render(self, dummy_ctxt=None):
        out = "\n%s\n" % "".join(list(self._gen_table()))
        if not self.rows:
            out += '%s' % cgi.escape(self.empty_table, quote=True)
        return out

    def _gen_table(self):
        yield '
\n'\
              % (self.name, self.border, self.summary)
        yield '
\n'
        yield '
%s
\n' % ''.join(
            '%s	\n' % h for h in self.header)
        yield '
\n'
        yield '
'
        for r, row in enumerate(self.rows):
            yield '
\n' % ["even", "odd"][r % 2]
            for col in row:
                yield '%s	\n' % col
            yield '
\n'
        yield '
\n'
        yield '
\n'

JOB_STATS = '''
SELECT oq_job_id, user_name, stop_time, status, duration FROM (
  SELECT oq_job_id, user_name, stop_time, status,
  stop_time - start_time as duration
  FROM uiapi.job_stats AS s
  INNER JOIN uiapi.oq_job AS o
  ON s.oq_job_id=o.id) AS x
WHERE oq_job_id=%s;
'''

JOB_INFO = '''
SELECT * FROM uiapi.job_info WHERE oq_job_id=%s;
'''

ALL_JOBS = '''
SELECT s.oq_job_id, o.user_name, status FROM uiapi.job_stats AS s
INNER JOIN uiapi.oq_job AS o
ON o.id=s.oq_job_id
WHERE stop_time::date = %s OR stop_time IS NULL AND start_time >= %s
ORDER BY stop_time
'''

PAGE_TEMPLATE = '''\

<link rel="stylesheet" href="http://ajax.googleapis.com/ajax/libs/jqueryui/1.11.2/themes/smoothness/jquery-ui.css" />

%s

'''

def make_tabs(tag_ids, tag_status, tag_contents):
    """
    Return a HTML string containing all the tabs we want to display
    """
    templ = '''

%s

%s

'''
    lis = []
    contents = []
    for i, (tag_id, status, tag_content) in enumerate(
            zip(tag_ids, tag_status, tag_contents), 1):
        mark = '.' if status == 'complete' else '!'
        lis.append('
%s%s
' % (i, tag_id, mark))
        contents.append('
%s
' % (
            i, tag_content))
    return templ % ('\n'.join(lis), '\n'.join(contents))

def make_report(conn, isodate='today'):
    """
    Build a HTML report with the computations performed at the given isodate.
    Return the name of the report, which is saved in the current directory.
    """
    if isodate == 'today':
        isodate = datetime.date.today().isoformat()
    curs = conn.cursor()
    fetcher = Fetcher(curs)
    tag_ids = []
    tag_status = []
    tag_contents = []

    jobs = fetcher.query(ALL_JOBS, isodate, isodate)[1:]
    page = '
%d job(s) finished before midnight of %s
' % (
        len(jobs), isodate)
    for job_id, user, status in jobs:
        tag_ids.append(job_id)
        tag_status.append(status)
        stats = fetcher.query(JOB_STATS, job_id)[1:]
        if not stats:
            continue
        (job_id, user, stop_time, status, duration) = stats[0]

        ds = DataStore(job_id, mode='r')
        try:
            report = html_parts(view_fullreport('fullreport', ds))
        except Exception as exc:
            report = dict(
                html_title='Could not generate report: %s' % cgi.escape(
                    str(exc), quote=True),
                fragment='')

        page = report['html_title']

        if 'realizations' in ds:
            tot_rlzs = len(ds['realizations'])
        else:
            tot_rlzs = 1
        data = fetcher.query(JOB_INFO, job_id)
        if data[1:]:
            info_rows = (
                '
Job Info for calculation=%s, %s
'
                '
owner: %s, realizations: %d
' % (
                    job_id, ds.attrs['description'].decode('utf-8'),
                    user, tot_rlzs)) + html(data)
            page += info_rows

        job_stats = html(fetcher.query(JOB_STATS, job_id))
        page += job_stats

        page += report['fragment']

        tag_contents.append(page)

    page = make_tabs(tag_ids, tag_status, tag_contents) + (
        'Report last updated: %s' % datetime.datetime.now())
    fname = 'jobs-%s.html' % isodate
    with open(fname, 'w') as f:
        f.write(PAGE_TEMPLATE % page)
    return fname
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2022, GEM Foundation
#
# 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 Affero 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
import unittest
import numpy as np
import matplotlib.pyplot as plt
from openquake.baselib.general import DictArray
from openquake.baselib.performance import Monitor
from openquake.hazardlib.calc.mrd import (
    update_mrd, update_mrd_indirect, get_uneven_bins_edges,
    calc_mean_rate_dist)
from openquake.hazardlib.contexts import read_cmakers
from openquake.commonlib import datastore
from openquake.hazardlib.cross_correlation import BakerJayaram2008

PLOT = False
CWD = os.path.dirname(__file__)

class MRD01TestCase(unittest.TestCase):
    """ Computes the mean rate density using a simple PSHA input model """

    def setUp(self):

        # Settings
        fname = os.path.join(CWD, 'data', 'mrd', 'calc_934.hdf5')
        self.imts = ['SA(0.2)', 'SA(1.0)']  # subset of the parent IMTs

        # Load datastore
        self.dstore = datastore.read(fname)
        self.oqp = self.dstore['oqparam']

        # Read the context maker and set the IMTLS
        [self.cmaker] = read_cmakers(self.dstore)
        self.cmaker.imtls = DictArray({k: self.oqp.imtls[k] for k in self.imts})

        # Read contexts
        [self.ctx] = self.cmaker.read_ctxs(self.dstore)

        # Set the cross correlation model
        self.crosscorr = BakerJayaram2008()
        self.rng = np.random.default_rng(42)

    def test_direct(self):

        # Compute the MRD
        imls1 = self.oqp.hazard_imtls[self.imts[0]]
        imls2 = self.oqp.hazard_imtls[self.imts[1]]
        len1 = len(imls1) - 1
        len2 = len(imls2) - 1
        assert len(self.oqp.sites) == 1
        mrd = np.zeros((len1, len2, len(self.cmaker.gsims)))
        update_mrd(self.ctx, self.cmaker, self.crosscorr, mrd, self.rng)

        # Loading Hazard Curves.
        # The poes array is 4D: |sites| x || x |IMTs| x |IMLs|
        poes = self.dstore['hcurves-stats'][:]
        afe = - np.log(1-poes)
        afo = afe[:, :, :, :-1] - afe[:, :, :, 1:]

        imts = list(self.oqp.hazard_imtls)
        idx1 = imts.index(self.imts[0])
        idx2 = imts.index(self.imts[1])

        afo1 = afo[0, 0, idx1, :]
        afo2 = afo[0, 0, idx2, :]

        tmp = self.oqp.hazard_imtls[self.imts[0]]
        c1 = tmp[:-1] + np.diff(tmp) / 2
        tmp = self.oqp.hazard_imtls[self.imts[1]]
        c2 = tmp[:-1] + np.diff(tmp) / 2

        # Compute marginal
        cm1 = imls1[:-1] + np.diff(imls1) / 2
        marg1 = np.squeeze(np.sum(mrd, axis=0))
        cm2 = imls2[:-1] + np.diff(imls2) / 2
        marg2 = np.squeeze(np.sum(mrd, axis=1))

        # Test
        np.testing.assert_almost_equal(marg1, afo1, decimal=5)
        np.testing.assert_almost_equal(marg2, afo2, decimal=5)

        if PLOT:
            plt.title('Direct method test')
            plt.plot(c1, afo1, label=f'HC {self.imts[0]}')
            plt.plot(cm1, marg1, 'o', mfc='None')
            plt.plot(c2, afo2, label=f'HC {self.imts[1]}')
            plt.plot(cm2, marg2, 'o', mfc='None', )
            plt.xlabel('Spectral Acceleration, S$_a$ [g]')
            plt.ylabel('Annual Rate of Occurrence')
            plt.legend()
            plt.grid(which='minor', ls=':', color='lightgrey')
            plt.grid(which='major', ls='--', color='grey')
            plt.show()

    def test_indirect(self):
        # Bin edges
        lefts = [-3, -2, 1, 2]
        numb = [80, 80, 10]
        be_mea = get_uneven_bins_edges(lefts, numb)
        be_sig = np.arange(0.50, 0.70, 0.01)
        imt1, imt2 = self.imts

        # Compute the MRD
        mon = Monitor('multivariate')
        mrd = calc_mean_rate_dist(self.ctx, self.cmaker, self.crosscorr,
                                  imt1, imt2, be_mea, be_sig, self.rng, mon)
        print(mon)

        # Loading Hazard Curves.
        # The poes array is 4D: |sites| x || x |IMTs| x |IMLs|
        poes = self.dstore['hcurves-stats'][:]
        afe = - np.log(1-poes)
        afo = afe[:, :, :, :-1] - afe[:, :, :, 1:]

        imts = list(self.oqp.imtls)
        idx1 = imts.index(self.imts[0])
        idx2 = imts.index(self.imts[1])

        afo1 = afo[0, 0, idx1, :]
        afo2 = afo[0, 0, idx2, :]

        tmp = self.oqp.imtls[self.imts[0]]
        c1 = tmp[:-1] + np.diff(tmp) / 2
        tmp = self.oqp.imtls[self.imts[1]]
        c2 = tmp[:-1] + np.diff(tmp) / 2

        # Compute marginal
        imls1 = self.oqp.imtls[self.imts[0]]
        imls2 = self.oqp.imtls[self.imts[1]]
        cm1 = imls1[:-1] + np.diff(imls1) / 2
        marg1 = np.squeeze(np.sum(mrd, axis=0))
        cm2 = imls2[:-1] + np.diff(imls2) / 2
        marg2 = np.squeeze(np.sum(mrd, axis=1))

        # Test
        np.testing.assert_almost_equal(marg1, afo1, decimal=5)
        np.testing.assert_almost_equal(marg2, afo2, decimal=5)

        if PLOT:
            plt.title('Indirect method test')
            plt.plot(c1, afo1, label=f'HC {self.imts[0]}')
            plt.plot(cm1, marg1, 'o', mfc='None')
            plt.plot(c2, afo2, label=f'HC {self.imts[1]}')
            plt.plot(cm2, marg2, 'o', mfc='None', )
            plt.xlabel('Spectral Acceleration, S$_a$ [g]')
            plt.ylabel('Annual Rate of Occurrence')
            plt.legend()
            plt.grid(which='minor', ls=':', color='lightgrey')
            plt.grid(which='major', ls='--', color='grey')
            plt.show()

    def test_compare(self):
        # Bin edges
        lefts = [-3, -2, 1, 2]
        numb = [80, 80, 10]
        be_mea = get_uneven_bins_edges(lefts, numb)
        be_sig = np.arange(0.50, 0.70, 0.01)

        # Set params
        imls1 = self.oqp.imtls[self.imts[0]]
        imls2 = self.oqp.imtls[self.imts[1]]
        len1 = len(imls1)-1
        len2 = len(imls2)-1
        nsites = len(self.oqp.sites)

        # Compute the MRD: indirect
        imt1, imt2 = self.imts
        mrdi = calc_mean_rate_dist(self.ctx, self.cmaker, self.crosscorr,
                                   imt1, imt2, be_mea, be_sig, self.rng)

        # Compute the MRD: direct
        mrdd = np.zeros((len1, len2, len(self.cmaker.gsims)))
        update_mrd(self.ctx, self.cmaker, self.crosscorr, mrdd, self.rng)

        np.testing.assert_almost_equal(mrdi[:,:, 0], mrdd)

        if PLOT:
            imlc1 = np.diff(imls1)/2 + imls1[:-1]
            imlc2 = np.diff(imls2)/2 + imls2[:-1]

            fig, axs = plt.subplots(1, 1)
            fig.set_size_inches(9, 6)
            plt1 = plt.contourf(np.log(imlc1), np.log(imlc2), mrdd[:, :, 0, 0])
            _ = plt.contour(np.log(imlc1), np.log(imlc2), mrdi[:, :, 0, 0],
                            colors='orange', linestyles='dashed')
            _ = plt.colorbar(mappable=plt1)
            _ = plt.title('MRD')
            axs.set_xlabel(f'{self.imts[0]}')
            axs.set_ylabel(f'{self.imts[1]}')
            plt.show()
<commit_msg>Cleaned direct method too<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2022, GEM Foundation
#
# 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 Affero 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
import unittest
import numpy as np
import matplotlib.pyplot as plt
from openquake.baselib.general import DictArray
from openquake.baselib.performance import Monitor
from openquake.hazardlib.calc.mrd import (
    update_mrd, update_mrd_indirect, get_uneven_bins_edges,
    calc_mean_rate_dist)
from openquake.hazardlib.contexts import read_cmakers
from openquake.commonlib import datastore
from openquake.hazardlib.cross_correlation import BakerJayaram2008

PLOT = False
CWD = os.path.dirname(__file__)

class MRD01TestCase(unittest.TestCase):
    """ Computes the mean rate density using a simple PSHA input model """

    def setUp(self):

        # Settings
        fname = os.path.join(CWD, 'data', 'mrd', 'calc_934.hdf5')
        self.imts = ['SA(0.2)', 'SA(1.0)']  # subset of the parent IMTs

        # Load datastore
        self.dstore = datastore.read(fname)
        self.oqp = self.dstore['oqparam']

        # Read the context maker and set the IMTLS
        [self.cmaker] = read_cmakers(self.dstore)
        self.cmaker.imtls = DictArray({k: self.oqp.imtls[k] for k in self.imts})

        # Read contexts
        [self.ctx] = self.cmaker.read_ctxs(self.dstore)

        # Set the cross correlation model
        self.crosscorr = BakerJayaram2008()
        self.rng = np.random.default_rng(42)

    def test_direct(self):

        # Compute the MRD
        imls1 = self.oqp.hazard_imtls[self.imts[0]]
        imls2 = self.oqp.hazard_imtls[self.imts[1]]
        len1 = len(imls1) - 1
        len2 = len(imls2) - 1
        assert len(self.oqp.sites) == 1
        mrd = np.zeros((len1, len2, len(self.cmaker.gsims)))
        update_mrd(self.ctx, self.cmaker, self.crosscorr, mrd, self.rng)

        # Loading Hazard Curves.
        # The poes array is 4D: |sites| x || x |IMTs| x |IMLs|
        poes = self.dstore['hcurves-stats'][:]
        afe = - np.log(1-poes)
        afo = afe[:, :, :, :-1] - afe[:, :, :, 1:]

        imts = list(self.oqp.hazard_imtls)
        idx1 = imts.index(self.imts[0])
        idx2 = imts.index(self.imts[1])

        afo1 = afo[0, 0, idx1, :]
        afo2 = afo[0, 0, idx2, :]

        tmp = self.oqp.hazard_imtls[self.imts[0]]
        c1 = tmp[:-1] + np.diff(tmp) / 2
        tmp = self.oqp.hazard_imtls[self.imts[1]]
        c2 = tmp[:-1] + np.diff(tmp) / 2

        # Compute marginal
        cm1 = imls1[:-1] + np.diff(imls1) / 2
        marg1 = np.squeeze(np.sum(mrd, axis=0))
        cm2 = imls2[:-1] + np.diff(imls2) / 2
        marg2 = np.squeeze(np.sum(mrd, axis=1))

        # Test
        np.testing.assert_almost_equal(marg1, afo1, decimal=5)
        np.testing.assert_almost_equal(marg2, afo2, decimal=5)

        if PLOT:
            plt.title('Direct method test')
            plt.plot(c1, afo1, label=f'HC {self.imts[0]}')
            plt.plot(cm1, marg1, 'o', mfc='None')
            plt.plot(c2, afo2, label=f'HC {self.imts[1]}')
            plt.plot(cm2, marg2, 'o', mfc='None', )
            plt.xlabel('Spectral Acceleration, S$_a$ [g]')
            plt.ylabel('Annual Rate of Occurrence')
            plt.legend()
            plt.grid(which='minor', ls=':', color='lightgrey')
            plt.grid(which='major', ls='--', color='grey')
            plt.show()

    def test_indirect(self):
        # Bin edges
        lefts = [-3, -2, 1, 2]
        numb = [80, 80, 10]
        be_mea = get_uneven_bins_edges(lefts, numb)
        be_sig = np.arange(0.50, 0.70, 0.01)
        imt1, imt2 = self.imts

        # Compute the MRD
        mon = Monitor('multivariate')
        mrd = calc_mean_rate_dist(self.ctx, self.cmaker, self.crosscorr,
                                  imt1, imt2, be_mea, be_sig, self.rng, mon)
        print(mon)

        # Loading Hazard Curves.
        # The poes array is 4D: |sites| x || x |IMTs| x |IMLs|
        poes = self.dstore['hcurves-stats'][:]
        afe = - np.log(1-poes)
        afo = afe[:, :, :, :-1] - afe[:, :, :, 1:]

        imts = list(self.oqp.imtls)
        idx1 = imts.index(self.imts[0])
        idx2 = imts.index(self.imts[1])

        afo1 = afo[0, 0, idx1, :]
        afo2 = afo[0, 0, idx2, :]

        tmp = self.oqp.imtls[self.imts[0]]
        c1 = tmp[:-1] + np.diff(tmp) / 2
        tmp = self.oqp.imtls[self.imts[1]]
        c2 = tmp[:-1] + np.diff(tmp) / 2

        # Compute marginal
        imls1 = self.oqp.imtls[self.imts[0]]
        imls2 = self.oqp.imtls[self.imts[1]]
        cm1 = imls1[:-1] + np.diff(imls1) / 2
        marg1 = np.squeeze(np.sum(mrd, axis=0))
        cm2 = imls2[:-1] + np.diff(imls2) / 2
        marg2 = np.squeeze(np.sum(mrd, axis=1))

        # Test
        np.testing.assert_almost_equal(marg1, afo1, decimal=5)
        np.testing.assert_almost_equal(marg2, afo2, decimal=5)

        if PLOT:
            plt.title('Indirect method test')
            plt.plot(c1, afo1, label=f'HC {self.imts[0]}')
            plt.plot(cm1, marg1, 'o', mfc='None')
            plt.plot(c2, afo2, label=f'HC {self.imts[1]}')
            plt.plot(cm2, marg2, 'o', mfc='None', )
            plt.xlabel('Spectral Acceleration, S$_a$ [g]')
            plt.ylabel('Annual Rate of Occurrence')
            plt.legend()
            plt.grid(which='minor', ls=':', color='lightgrey')
            plt.grid(which='major', ls='--', color='grey')
            plt.show()

    def test_compare(self):
        # Bin edges
        lefts = [-3, -2, 1, 2]
        numb = [80, 80, 10]
        be_mea = get_uneven_bins_edges(lefts, numb)
        be_sig = np.arange(0.50, 0.70, 0.01)

        # Set params
        imls1 = self.oqp.imtls[self.imts[0]]
        imls2 = self.oqp.imtls[self.imts[1]]
        len1 = len(imls1)-1

        # Compute the MRD: indirect
        imt1, imt2 = self.imts
        mrdi = calc_mean_rate_dist(self.ctx, self.cmaker, self.crosscorr,
                                   imt1, imt2, be_mea, be_sig, self.rng)

        # Compute the MRD: direct
        mrdd = np.zeros((len1, len1, len(self.cmaker.gsims)))
        update_mrd(self.ctx, self.cmaker, self.crosscorr, mrdd, self.rng)

        np.testing.assert_almost_equal(mrdi[:,:, 0], mrdd)

        if PLOT:
            imlc1 = np.diff(imls1)/2 + imls1[:-1]
            imlc2 = np.diff(imls2)/2 + imls2[:-1]

            fig, axs = plt.subplots(1, 1)
            fig.set_size_inches(9, 6)
            plt1 = plt.contourf(np.log(imlc1), np.log(imlc2), mrdd[:, :, 0, 0])
            _ = plt.contour(np.log(imlc1), np.log(imlc2), mrdi[:, :, 0, 0],
                            colors='orange', linestyles='dashed')
            _ = plt.colorbar(mappable=plt1)
            _ = plt.title('MRD')
            axs.set_xlabel(f'{self.imts[0]}')
            axs.set_ylabel(f'{self.imts[1]}')
            plt.show()
<|endoftext|>


<commit_before>from opentrons_sdk.protocol.handlers import ProtocolHandler
from opentrons_sdk.labware import deck, instruments
from opentrons_sdk.containers import placeable

class ContextHandler(ProtocolHandler):

    """
    ContextHandler runs all the stuff on the virtual robot in the background
    and makes relevant data available.
    """

    _deck = None
    _instruments = None  # Axis as keys; Pipette object as vals.

    # Calibration is organized by axis.
    #
    # {
    #   'a': {'_axis': {top, bottom, blowout}, (0,0): {x, y, top, bottom}}
    # }
    _calibration = None

    def setup(self):
        # self._deck = deck.Deck()
        self.setup_deck()
        self._instruments = {}
        self._calibration = {}

    def get_deck_slot_types(self):
        return 'acrylic_slots'

    def get_slot_offsets(self):
        """
        col_offset
        - from bottom left corner of A to bottom corner of B

        row_offset
        - from bottom left corner of 1 to bottom corner of 2
        """
        SLOT_OFFSETS = {
            '3d_printed_slots': {
                'col_offset': 91,
                'row_offset': 134.5
            },
            'acrylic_slots': {
                'col_offset': 96.25,
                'row_offset': 133.3
            }

        }
        slot_settings = SLOT_OFFSETS.get(self.get_deck_slot_types())
        row_offset = slot_settings.get('row_offset')
        col_offset = slot_settings.get('col_offset')
        return (row_offset, col_offset)

    def get_max_robot_rows(self):
        # TODO: dynamically figure out robot rows
        return 3

    def setup_deck(self):
        self.deck = placeable.Deck()
        robot_rows = self.get_max_robot_rows()
        # row_offset, col_offset

        for col_index, col in enumerate('EDCBA'):
            for row_index, row in enumerate(range(robot_rows,1, -1)):
                slot = placeable.Slot()
                self.deck.add(
                    slot, "{}{}".format(col, row),
                    (row_offset * row_index, col_offset * col_index, 0)  # TODO: should z always be zero?
                )

    def add_instrument(self, axis, instrument=None):
        axis = axis.upper()
        # We only have pipettes now so this is pipette-specific.
        self._instruments[axis] = instrument
        self._instruments[axis]._axis = axis

    def get_instrument(self, axis=None, **kwargs):
        if axis is not None:
            axis = self.normalize_axis(axis)
            if axis not in self._instruments:
                raise KeyError(
                    "No instrument assigned to {} axis.".format(axis)
                )
            else:
                inst = self._instruments[axis]
                inst._axis = axis
                if axis in self._calibration:
                    if '_axis' in self._calibration[axis]:
                        inst.calibrate(**self._calibration[axis]['_axis'])
                return self._instruments[axis]
        volume = kwargs.pop('volume', None)
        for k, i in sorted(self._instruments.items()):
            match = True
            if volume and i.supports_volume(volume) is False:
                    continue
            for j, v in kwargs.items():
                if getattr(i, j) != v:
                    match = False
                    break
            if match:
                return self.get_instrument(axis=k)

        return None

    def get_instruments(self):
        """
        :returns: sorted list of (axis, instrument)
        """
        return sorted(self._instruments.items())

    def find_container(self, **filters):
        return self._deck.find_module(**filters)

    def get_containers(self, **filters):
        return self._deck.get_modules(**filters)

    def add_container(self, slot, container_name):
        return self._deck.add_module(slot, container_name)

    def get_only_instrument(self):
        ks = list(self._instruments)
        if len(ks) > 0:
            return self.get_instrument(axis=ks[0])
        if len(ks) is 0:
            raise KeyError("No instruments loaded.")
        else:
            return None

    def normalize_axis(self, axis):
        """
        Returns an axis by axis, after normalizing axis input.

        If axis is none, the first instrument axis if only one instrument is
        attached to the protocol.

        If the axis isn't valid, it raises a KeyError.
        """
        if axis is None:
            raise KeyError("Axis must be specified.")
        axis = axis.upper()
        if axis not in self._instruments:
            raise KeyError("Can't find instrument for axis {}.".format(axis))
        return axis

    def get_axis_calibration(self, axis):
        """
        Initializes and returns calibration for a particular axis.
        """
        if axis is None:
            axis = self.get_only_instrument().axis
        if axis is None:
            raise KeyError(
                "Calibration axis must be specified when multiple " +
                "instruments are loaded."
            )
        axis = self.normalize_axis(axis)
        if axis not in self._calibration:
            self._calibration[axis] = {}
        return self._calibration[axis]

    def calibrate(self, pos, axis=None, x=None, y=None, z=None, top=None,
                  bottom=None):
        if axis is None:
            instrument = self.get_only_instrument()
            if instrument is None:
                raise KeyError(
                    "Calibration axis must be specified when multiple " +
                    "instruments are loaded."
                )
            axis = instrument.axis
        cal = self.get_axis_calibration(axis)
        if pos not in cal:
            cal[pos] = {}
        pos_cal = cal[pos]
        # Roll in all the new calibration changes.
        if x is not None:
            pos_cal['x'] = x
        if y is not None:
            pos_cal['y'] = y
        if z is not None:
            pos_cal['z'] = z
        if top is not None:
            pos_cal['top'] = top
        if bottom is not None:
            pos_cal['bottom'] = bottom

    def calibrate_instrument(self, axis, top=None, blowout=None, droptip=None,
                             bottom=None):
        cal = self.get_axis_calibration(axis)
        if '_instrument' not in cal:
            cal['_instrument'] = {}
        a_cal = cal['_instrument']
        # Roll in all the new calibration changes.
        if top is not None:
            a_cal['top'] = top
        if bottom is not None:
            a_cal['bottom'] = bottom
        if blowout is not None:
            a_cal['blowout'] = blowout
        if droptip is not None:
            a_cal['droptip'] = droptip
        if bottom is not None:
            a_cal['bottom'] = bottom
        self.get_instrument(axis=axis).calibrate(**a_cal)

    def get_coordinates(self, position, axis=None):
        """ Returns the calibrated coordinates for a position. """
        cal = self.get_axis_calibration(axis)
        slot, well = position
        defaults = ({'top': 0, 'bottom': 0, 'x': 0, 'y': 0})
        output = {}
        # Calibration for A1 in this container (x, y, top, bottom).
        slot_cal = {}
        slot_cal.update(defaults)
        slot_cal.update(cal.get((slot), {}))
        # Default offset on x, y calculated from container definition.
        ox, oy = self._deck.slot(slot).get_child(well).coordinates()
        # x, y, top bottom
        well_cal = cal.get((slot, well), {})
        output.update(well_cal)
        # Use calculated offsets if no custom well calibration provided.
        if 'x' not in output:
            output['x'] = slot_cal['x'] + ox
        if 'y' not in output:
            output['y'] = slot_cal['y'] + oy
        # Merge slot and well calibration
        if 'top' not in output:
            output['top'] = slot_cal['top']
        if 'bottom' not in output:
            output['bottom'] = slot_cal['bottom']
        return output

    def get_tip_coordinates(self, pipette):
        name = 'tiprack.{}'.format(pipette.name)
        tiprack = self.find_container(name=name, has_tips=True)
        if tiprack is None:
            raise KeyError("No tiprack found for {}.".format(name))
        tip = tiprack.get_next_tip()
        return self.get_coordinates(tip.address, axis=pipette.axis)

    def get_trash_coordinates(self, axis=None):
        trash = self.find_container(name='point.trash')
        if trash is None:
            raise KeyError("No disposal point (trash) on deck.")
        return self.get_coordinates(trash.address + [(0, 0)], axis)

    def get_volume(self, well):
        slot, well = self._protocol._normalize_address(well)
        return self._deck.slot(slot).well(well).get_volume()

    def transfer(self, start=None, end=None, volume=None, **kwargs):
        start_slot, start_well = start
        end_slot, end_well = end
        start = self._deck.slot(start_slot).well(start_well)
        end = self._deck.slot(end_slot).well(end_well)
        start.transfer(volume, end)

    def transfer_group(self, *args, **kwargs):
        pass

    def distribute(self, *args, **kwargs):
        pass

    def mix(self, *args, **kwargs):
        pass

    def consolidate(self, *args, **kwargs):
        pass
<commit_msg>add slots based on assumed x and y offset, rename for compatability w/old api<commit_after>from opentrons_sdk.protocol.handlers import ProtocolHandler
from opentrons_sdk.labware import deck, instruments
from opentrons_sdk.containers import placeable

class ContextHandler(ProtocolHandler):

    """
    ContextHandler runs all the stuff on the virtual robot in the background
    and makes relevant data available.
    """

    _deck = None
    _instruments = None  # Axis as keys; Pipette object as vals.

    # Calibration is organized by axis.
    #
    # {
    #   'a': {'_axis': {top, bottom, blowout}, (0,0): {x, y, top, bottom}}
    # }
    _calibration = None

    def setup(self):
        self._deck = placeable.Deck()
        self.setup_deck()
        self._instruments = {}
        self._calibration = {}

    def get_deck_slot_types(self):
        return 'acrylic_slots'

    def get_slot_offsets(self):
        """
        col_offset
        - from bottom left corner of A to bottom corner of B

        row_offset
        - from bottom left corner of 1 to bottom corner of 2

        TODO: figure out actual X and Y offsets (from origin)
        """
        SLOT_OFFSETS = {
            '3d_printed_slots': {
                'x_offset': 10,
                'y_offset': 10,
                'col_offset': 91,
                'row_offset': 134.5
            },
            'acrylic_slots': {
                'x_offset': 10,
                'y_offset': 10,
                'col_offset': 96.25,
                'row_offset': 133.3
            }

        }
        slot_settings = SLOT_OFFSETS.get(self.get_deck_slot_types())
        row_offset = slot_settings.get('row_offset')
        col_offset = slot_settings.get('col_offset')
        x_offset = slot_settings.get('x_offset')
        y_offset = slot_settings.get('y_offset')
        return (row_offset, col_offset, x_offset, y_offset)

    def get_max_robot_rows(self):
        # TODO: dynamically figure out robot rows
        return 3

    def setup_deck(self):
        robot_rows = self.get_max_robot_rows()
        row_offset, col_offset, x_offset, y_offset = get_slot_offsets()

        for col_index, col in enumerate('EDCBA'):
            for row_index, row in enumerate(range(robot_rows, 1, -1)):
                slot = placeable.Slot()
                slot_coordinates = (
                    (row_offset * row_index) + x_offset,
                    (col_offset * col_index) + y_offset,
                    0  # TODO: should z always be zero?
                )
                slot_name = "{}{}".format(col, row)
                self.deck.add(slot, slot_name, (slot_coordinates))

    def add_instrument(self, axis, instrument=None):
        axis = axis.upper()
        # We only have pipettes now so this is pipette-specific.
        self._instruments[axis] = instrument
        self._instruments[axis]._axis = axis

    def get_instrument(self, axis=None, **kwargs):
        if axis is not None:
            axis = self.normalize_axis(axis)
            if axis not in self._instruments:
                raise KeyError(
                    "No instrument assigned to {} axis.".format(axis)
                )
            else:
                inst = self._instruments[axis]
                inst._axis = axis
                if axis in self._calibration:
                    if '_axis' in self._calibration[axis]:
                        inst.calibrate(**self._calibration[axis]['_axis'])
                return self._instruments[axis]
        volume = kwargs.pop('volume', None)
        for k, i in sorted(self._instruments.items()):
            match = True
            if volume and i.supports_volume(volume) is False:
                    continue
            for j, v in kwargs.items():
                if getattr(i, j) != v:
                    match = False
                    break
            if match:
                return self.get_instrument(axis=k)

        return None

    def get_instruments(self):
        """
        :returns: sorted list of (axis, instrument)
        """
        return sorted(self._instruments.items())

    def find_container(self, **filters):
        return self._deck.find_module(**filters)

    def get_containers(self, **filters):
        return self._deck.get_modules(**filters)

    def add_container(self, slot, container_name):
        return self._deck.add_module(slot, container_name)

    def get_only_instrument(self):
        ks = list(self._instruments)
        if len(ks) > 0:
            return self.get_instrument(axis=ks[0])
        if len(ks) is 0:
            raise KeyError("No instruments loaded.")
        else:
            return None

    def normalize_axis(self, axis):
        """
        Returns an axis by axis, after normalizing axis input.

        If axis is none, the first instrument axis if only one instrument is
        attached to the protocol.

        If the axis isn't valid, it raises a KeyError.
        """
        if axis is None:
            raise KeyError("Axis must be specified.")
        axis = axis.upper()
        if axis not in self._instruments:
            raise KeyError("Can't find instrument for axis {}.".format(axis))
        return axis

    def get_axis_calibration(self, axis):
        """
        Initializes and returns calibration for a particular axis.
        """
        if axis is None:
            axis = self.get_only_instrument().axis
        if axis is None:
            raise KeyError(
                "Calibration axis must be specified when multiple " +
                "instruments are loaded."
            )
        axis = self.normalize_axis(axis)
        if axis not in self._calibration:
            self._calibration[axis] = {}
        return self._calibration[axis]

    def calibrate(self, pos, axis=None, x=None, y=None, z=None, top=None,
                  bottom=None):
        if axis is None:
            instrument = self.get_only_instrument()
            if instrument is None:
                raise KeyError(
                    "Calibration axis must be specified when multiple " +
                    "instruments are loaded."
                )
            axis = instrument.axis
        cal = self.get_axis_calibration(axis)
        if pos not in cal:
            cal[pos] = {}
        pos_cal = cal[pos]
        # Roll in all the new calibration changes.
        if x is not None:
            pos_cal['x'] = x
        if y is not None:
            pos_cal['y'] = y
        if z is not None:
            pos_cal['z'] = z
        if top is not None:
            pos_cal['top'] = top
        if bottom is not None:
            pos_cal['bottom'] = bottom

    def calibrate_instrument(self, axis, top=None, blowout=None, droptip=None,
                             bottom=None):
        cal = self.get_axis_calibration(axis)
        if '_instrument' not in cal:
            cal['_instrument'] = {}
        a_cal = cal['_instrument']
        # Roll in all the new calibration changes.
        if top is not None:
            a_cal['top'] = top
        if bottom is not None:
            a_cal['bottom'] = bottom
        if blowout is not None:
            a_cal['blowout'] = blowout
        if droptip is not None:
            a_cal['droptip'] = droptip
        if bottom is not None:
            a_cal['bottom'] = bottom
        self.get_instrument(axis=axis).calibrate(**a_cal)

    def get_coordinates(self, position, axis=None):
        """ Returns the calibrated coordinates for a position. """
        cal = self.get_axis_calibration(axis)
        slot, well = position
        defaults = ({'top': 0, 'bottom': 0, 'x': 0, 'y': 0})
        output = {}
        # Calibration for A1 in this container (x, y, top, bottom).
        slot_cal = {}
        slot_cal.update(defaults)
        slot_cal.update(cal.get((slot), {}))
        # Default offset on x, y calculated from container definition.
        ox, oy = self._deck.slot(slot).get_child(well).coordinates()
        # x, y, top bottom
        well_cal = cal.get((slot, well), {})
        output.update(well_cal)
        # Use calculated offsets if no custom well calibration provided.
        if 'x' not in output:
            output['x'] = slot_cal['x'] + ox
        if 'y' not in output:
            output['y'] = slot_cal['y'] + oy
        # Merge slot and well calibration
        if 'top' not in output:
            output['top'] = slot_cal['top']
        if 'bottom' not in output:
            output['bottom'] = slot_cal['bottom']
        return output

    def get_tip_coordinates(self, pipette):
        name = 'tiprack.{}'.format(pipette.name)
        tiprack = self.find_container(name=name, has_tips=True)
        if tiprack is None:
            raise KeyError("No tiprack found for {}.".format(name))
        tip = tiprack.get_next_tip()
        return self.get_coordinates(tip.address, axis=pipette.axis)

    def get_trash_coordinates(self, axis=None):
        trash = self.find_container(name='point.trash')
        if trash is None:
            raise KeyError("No disposal point (trash) on deck.")
        return self.get_coordinates(trash.address + [(0, 0)], axis)

    def get_volume(self, well):
        slot, well = self._protocol._normalize_address(well)
        return self._deck.slot(slot).well(well).get_volume()

    def transfer(self, start=None, end=None, volume=None, **kwargs):
        start_slot, start_well = start
        end_slot, end_well = end
        start = self._deck.slot(start_slot).well(start_well)
        end = self._deck.slot(end_slot).well(end_well)
        start.transfer(volume, end)

    def transfer_group(self, *args, **kwargs):
        pass

    def distribute(self, *args, **kwargs):
        pass

    def mix(self, *args, **kwargs):
        pass

    def consolidate(self, *args, **kwargs):
        pass
<|endoftext|>


<commit_before>#!/usr/bin/env python
#
# Copyright 2011, Google Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Test for end-to-end."""

import logging
import os
import signal
import socket
import subprocess
import sys
import time
import unittest

import set_sys_path  # Update sys.path to locate mod_pywebsocket module.

from test import client_for_testing

# Special message that tells the echo server to start closing handshake
_GOODBYE_MESSAGE = 'Goodbye'

# Test body functions
def _echo_check_procedure(client):
    client.connect()

    client.send_message('test')
    client.assert_receive('test')

    client.send_close()
    client.assert_receive_close()

    client.assert_connection_closed()

def _echo_check_procedure_with_goodbye(client):
    client.connect()

    client.send_message('test')
    client.assert_receive('test')

    client.send_message(_GOODBYE_MESSAGE)
    client.assert_receive(_GOODBYE_MESSAGE)

    client.assert_receive_close()
    client.send_close()

    client.assert_connection_closed()

class EndToEndTest(unittest.TestCase):
    """An end-to-end test that launches pywebsocket standalone server as a
    separate process, connects to it using the client_for_testing module, and
    checks if the server behaves correctly by exchanging opening handshake and
    frames over a TCP connection.
    """

    def setUp(self):
        self.top_dir = os.path.join(os.path.split(__file__)[0], '..')
        os.putenv('PYTHONPATH', os.path.pathsep.join(sys.path))
        self.standalone_command = os.path.join(
            self.top_dir, 'mod_pywebsocket', 'standalone.py')
        self.document_root = os.path.join(self.top_dir, 'example')
        s = socket.socket()
        s.bind(('127.0.0.1', 0))
        (_, self.test_port) = s.getsockname()
        s.close()

        self._options = client_for_testing.ClientOptions()
        self._options.server_host = 'localhost'
        self._options.origin = 'http://localhost'
        self._options.resource = '/echo'
        self._options.server_port = self.test_port

    def _run_python_command(self, commandline):
        return subprocess.Popen([sys.executable] + commandline, close_fds=True,
                                stdout=sys.stdout, stderr=sys.stderr)

    def _run_server(self, allow_draft75=False):
        args = [self.standalone_command,
                '-p', str(self.test_port),
                '-d', self.document_root]

        # Inherit the level set to the root logger by test runner.
        root_logger = logging.getLogger()
        log_level = root_logger.getEffectiveLevel()
        if log_level != logging.NOTSET:
            args.append('--log-level')
            args.append(logging.getLevelName(log_level).lower())

        if allow_draft75:
            args.append('--allow-draft75')

        return self._run_python_command(args)

    def _kill_process(self, pid):
        if sys.platform in ('win32', 'cygwin'):
            subprocess.call(
                ('taskkill.exe', '/f', '/pid', str(pid)), close_fds=True)
        else:
            os.kill(pid, signal.SIGKILL)

    def _run_hybi_test(self, test_function):
        server = self._run_server()
        try:
            # TODO(tyoshino): add some logic to poll the server until it
            # becomes ready
            time.sleep(0.2)

            client = client_for_testing.create_client(self._options)
            try:
                test_function(client)
            finally:
                client.close_socket()
        finally:
            self._kill_process(server.pid)

    def _run_hybi_deflate_test(self, test_function):
        server = self._run_server()
        try:
            time.sleep(0.2)

            self._options.use_deflate = True
            client = client_for_testing.create_client(self._options)
            try:
                test_function(client)
            finally:
                client.close_socket()
        finally:
            self._kill_process(server.pid)

    def test_echo(self):
        self._run_hybi_test(_echo_check_procedure)

    def test_echo_server_close(self):
        self._run_hybi_test(_echo_check_procedure_with_goodbye)

    def test_echo_deflate(self):
        self._run_hybi_deflate_test(_echo_check_procedure)

    def test_echo_deflate_server_close(self):
        self._run_hybi_deflate_test(_echo_check_procedure_with_goodbye)

    def test_close_on_protocol_error(self):
        """Tests that the server sends a close frame with protocol error status
        code when the client sends data with some protocol error.
        """

        def test_function(client):
            client.connect()

            # Intermediate frame without any preceding start of fragmentation
            # frame.
            client.send_frame_of_arbitrary_bytes('\x80\x80', '')
            client.assert_receive_close(
                client_for_testing.STATUS_PROTOCOL_ERROR)

        self._run_hybi_test(test_function)

    def test_close_on_unsupported_frame(self):
        """Tests that the server sends a close frame with unsupported operation
        status code when the client sends data asking some operation that is
        not supported by the server.
        """

        def test_function(client):
            client.connect()

            # Text frame with RSV3 bit raised.
            client.send_frame_of_arbitrary_bytes('\x91\x80', '')
            client.assert_receive_close(
                client_for_testing.STATUS_UNSUPPORTED)

        self._run_hybi_test(test_function)

    def _run_hybi00_test(self, test_function):
        server = self._run_server()
        try:
            time.sleep(0.2)

            client = client_for_testing.create_client_hybi00(self._options)
            try:
                test_function(client)
            finally:
                client.close_socket()
        finally:
            self._kill_process(server.pid)

    def test_echo_hybi00(self):
        self._run_hybi00_test(_echo_check_procedure)

    def test_echo_server_close_hybi00(self):
        self._run_hybi00_test(_echo_check_procedure_with_goodbye)

    def _run_hixie75_test(self, test_function):
        server = self._run_server(allow_draft75=True)
        try:
            time.sleep(0.2)

            client = client_for_testing.create_client_hixie75(self._options)
            try:
                test_function(client)
            finally:
                client.close_socket()
        finally:
            self._kill_process(server.pid)

    def test_echo_hixie75(self):
        """Tests that the server can talk draft-hixie-thewebsocketprotocol-75
        protocol.
        """

        def test_function(client):
            client.connect()

            client.send_message('test')
            client.assert_receive('test')

        self._run_hixie75_test(test_function)

    def test_echo_server_close_hixie75(self):
        """Tests that the server can talk draft-hixie-thewebsocketprotocol-75
        protocol. At the end of message exchanging, the client sends a keyword
        message that requests the server to close the connection, and then
        checks if the connection is really closed.
        """

        def test_function(client):
            client.connect()

            client.send_message('test')
            client.assert_receive('test')

            client.send_message(_GOODBYE_MESSAGE)
            client.assert_receive(_GOODBYE_MESSAGE)

        self._run_hixie75_test(test_function)

if __name__ == '__main__':
    unittest.main()

# vi:sts=4 sw=4 et
<commit_msg>Add example/echo_client.py test into test/test_endtoend.py<commit_after>#!/usr/bin/env python
#
# Copyright 2011, Google Inc.
# All rights reserved.
#
# Redistribution and use in source and binary forms, with or without
# modification, are permitted provided that the following conditions are
# met:
#
#     * Redistributions of source code must retain the above copyright
# notice, this list of conditions and the following disclaimer.
#     * Redistributions in binary form must reproduce the above
# copyright notice, this list of conditions and the following disclaimer
# in the documentation and/or other materials provided with the
# distribution.
#     * Neither the name of Google Inc. nor the names of its
# contributors may be used to endorse or promote products derived from
# this software without specific prior written permission.
#
# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS
# "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT
# LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR
# A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT
# OWNER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL,
# SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT
# LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE,
# DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY
# THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

"""Test for end-to-end."""

import logging
import os
import signal
import socket
import subprocess
import sys
import time
import unittest

import set_sys_path  # Update sys.path to locate mod_pywebsocket module.

from test import client_for_testing

# Special message that tells the echo server to start closing handshake
_GOODBYE_MESSAGE = 'Goodbye'

# Test body functions
def _echo_check_procedure(client):
    client.connect()

    client.send_message('test')
    client.assert_receive('test')

    client.send_close()
    client.assert_receive_close()

    client.assert_connection_closed()

def _echo_check_procedure_with_goodbye(client):
    client.connect()

    client.send_message('test')
    client.assert_receive('test')

    client.send_message(_GOODBYE_MESSAGE)
    client.assert_receive(_GOODBYE_MESSAGE)

    client.assert_receive_close()
    client.send_close()

    client.assert_connection_closed()

class EndToEndTest(unittest.TestCase):
    """An end-to-end test that launches pywebsocket standalone server as a
    separate process, connects to it using the client_for_testing module, and
    checks if the server behaves correctly by exchanging opening handshake and
    frames over a TCP connection.
    """

    def setUp(self):
        self.top_dir = os.path.join(os.path.split(__file__)[0], '..')
        os.putenv('PYTHONPATH', os.path.pathsep.join(sys.path))
        self.standalone_command = os.path.join(
            self.top_dir, 'mod_pywebsocket', 'standalone.py')
        self.document_root = os.path.join(self.top_dir, 'example')
        s = socket.socket()
        s.bind(('127.0.0.1', 0))
        (_, self.test_port) = s.getsockname()
        s.close()

        self._options = client_for_testing.ClientOptions()
        self._options.server_host = 'localhost'
        self._options.origin = 'http://localhost'
        self._options.resource = '/echo'
        self._options.server_port = self.test_port

    def _run_python_command(self, commandline, stdout=None):
        return subprocess.Popen([sys.executable] + commandline, close_fds=True,
                                stdout=stdout)

    def _run_server(self, allow_draft75=False):
        args = [self.standalone_command,
                '-p', str(self.test_port),
                '-d', self.document_root]

        # Inherit the level set to the root logger by test runner.
        root_logger = logging.getLogger()
        log_level = root_logger.getEffectiveLevel()
        if log_level != logging.NOTSET:
            args.append('--log-level')
            args.append(logging.getLevelName(log_level).lower())

        if allow_draft75:
            args.append('--allow-draft75')

        return self._run_python_command(args)

    def _kill_process(self, pid):
        if sys.platform in ('win32', 'cygwin'):
            subprocess.call(
                ('taskkill.exe', '/f', '/pid', str(pid)), close_fds=True)
        else:
            os.kill(pid, signal.SIGKILL)

    def _run_hybi_test(self, test_function):
        server = self._run_server()
        try:
            # TODO(tyoshino): add some logic to poll the server until it
            # becomes ready
            time.sleep(0.2)

            client = client_for_testing.create_client(self._options)
            try:
                test_function(client)
            finally:
                client.close_socket()
        finally:
            self._kill_process(server.pid)

    def _run_hybi_deflate_test(self, test_function):
        server = self._run_server()
        try:
            time.sleep(0.2)

            self._options.use_deflate = True
            client = client_for_testing.create_client(self._options)
            try:
                test_function(client)
            finally:
                client.close_socket()
        finally:
            self._kill_process(server.pid)

    def test_echo(self):
        self._run_hybi_test(_echo_check_procedure)

    def test_echo_server_close(self):
        self._run_hybi_test(_echo_check_procedure_with_goodbye)

    def test_echo_deflate(self):
        self._run_hybi_deflate_test(_echo_check_procedure)

    def test_echo_deflate_server_close(self):
        self._run_hybi_deflate_test(_echo_check_procedure_with_goodbye)

    def test_close_on_protocol_error(self):
        """Tests that the server sends a close frame with protocol error status
        code when the client sends data with some protocol error.
        """

        def test_function(client):
            client.connect()

            # Intermediate frame without any preceding start of fragmentation
            # frame.
            client.send_frame_of_arbitrary_bytes('\x80\x80', '')
            client.assert_receive_close(
                client_for_testing.STATUS_PROTOCOL_ERROR)

        self._run_hybi_test(test_function)

    def test_close_on_unsupported_frame(self):
        """Tests that the server sends a close frame with unsupported operation
        status code when the client sends data asking some operation that is
        not supported by the server.
        """

        def test_function(client):
            client.connect()

            # Text frame with RSV3 bit raised.
            client.send_frame_of_arbitrary_bytes('\x91\x80', '')
            client.assert_receive_close(
                client_for_testing.STATUS_UNSUPPORTED)

        self._run_hybi_test(test_function)

    def _run_hybi00_test(self, test_function):
        server = self._run_server()
        try:
            time.sleep(0.2)

            client = client_for_testing.create_client_hybi00(self._options)
            try:
                test_function(client)
            finally:
                client.close_socket()
        finally:
            self._kill_process(server.pid)

    def test_echo_hybi00(self):
        self._run_hybi00_test(_echo_check_procedure)

    def test_echo_server_close_hybi00(self):
        self._run_hybi00_test(_echo_check_procedure_with_goodbye)

    def _run_hixie75_test(self, test_function):
        server = self._run_server(allow_draft75=True)
        try:
            time.sleep(0.2)

            client = client_for_testing.create_client_hixie75(self._options)
            try:
                test_function(client)
            finally:
                client.close_socket()
        finally:
            self._kill_process(server.pid)

    def test_echo_hixie75(self):
        """Tests that the server can talk draft-hixie-thewebsocketprotocol-75
        protocol.
        """

        def test_function(client):
            client.connect()

            client.send_message('test')
            client.assert_receive('test')

        self._run_hixie75_test(test_function)

    def test_echo_server_close_hixie75(self):
        """Tests that the server can talk draft-hixie-thewebsocketprotocol-75
        protocol. At the end of message exchanging, the client sends a keyword
        message that requests the server to close the connection, and then
        checks if the connection is really closed.
        """

        def test_function(client):
            client.connect()

            client.send_message('test')
            client.assert_receive('test')

            client.send_message(_GOODBYE_MESSAGE)
            client.assert_receive(_GOODBYE_MESSAGE)

        self._run_hixie75_test(test_function)

    def test_example_echo_client(self):
        server = self._run_server()
        try:
            time.sleep(0.2)

            client_command = os.path.join(
                self.top_dir, 'example', 'echo_client.py')
            args = [client_command,
                    '-p', str(self.test_port)]
            pipe_recv, pipe_send = os.pipe()
            self._run_python_command(args, stdout=pipe_send)
            actual = os.read(pipe_recv, 4096).decode("utf-8")
            expected = ('Send: Hello\n' 'Recv: Hello\n'
                u'Send: \u65e5\u672c\n' u'Recv: \u65e5\u672c\n'
                'Send close\n' 'Recv ack\n')
            if actual != expected:
                raise Exception('Unexpected result on example echo client: '
                                '%r (expected) vs %r (actual)' %
                                (expected, actual))
        finally:
            self._kill_process(server.pid)

if __name__ == '__main__':
    unittest.main()

# vi:sts=4 sw=4 et
<|endoftext|>