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

"""
This module provides classes for plotting PhaseDiagram objects.
"""

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__status__ = "Production"
__date__ = "Sep 23, 2011"

import math
import numpy as np
import re
import itertools

from pymatgen.phasediagram.pdanalyzer import PDAnalyzer
from pymatgen.util.string_utils import latexify

class PDPlotter(object):
    '''
    A plotter class for phase diagrams.
    '''

    def __init__(self, phasediagram, show_unstable=True):
        """
        Args:
            phasediagram:
                A PhaseDiagram object.
        """
        self._pd = phasediagram
        self._dim = len(self._pd.elements)
        self.lines = uniquelines(self._pd.facets)
        self.show_unstable = show_unstable
        if self._dim < 2 or self._dim > 4:
            raise ValueError("Only 2-4 components supported!")

    @property
    def pd_plot_data(self):
        '''
        Plot data for phase diagram.
        2-comp - Full hull with energies
        3/4-comp - Projection into 2D or 3D gibbs triangle.

        Returns:
            (lines, stable_entries, unstable_entries):
                - lines is a list of list of coordinates for lines in the PD.
                - stable_entries is a {coordinate : entry} for each stable node
                  in the phase diagram. (Each coordinate can only have one
                  stable phase)
                - unstable_entries is a {entry: coordinates} for all unstable
                  nodes in the phase diagram.
        '''
        pd = self._pd
        entries = pd.qhull_entries
        data = np.array(pd.qhull_data)
        facetlines = self.lines
        lines = list()
        stable_entries = dict()
        for line in facetlines:
            entry1 = entries[line[0]]
            entry2 = entries[line[1]]
            if self._dim == 2:
                x = [data[line[0]][0], data[line[1]][0]]
                y = [pd.get_form_energy_per_atom(entry1),
                     pd.get_form_energy_per_atom(entry2)]
                coord = [x, y]
            elif self._dim == 3:
                coord = triangular_coord(data[line, 0:2])
            else:
                coord = tet_coord(data[line, 0:3])
            lines.append(coord)
            labelcoord = list(zip(*coord))
            stable_entries[labelcoord[0]] = entry1
            stable_entries[labelcoord[1]] = entry2

        allentries = pd.all_entries
        alldata = np.array(pd.all_entries_hulldata)
        unstable_entries = dict()
        stable = pd.stable_entries
        for i in xrange(0, len(allentries)):
            entry = allentries[i]
            if entry not in stable:
                if self._dim == 2:
                    x = [alldata[i][0], alldata[i][0]]
                    y = [pd.get_form_energy_per_atom(entry),
                         pd.get_form_energy_per_atom(entry)]
                    coord = [x, y]
                elif self._dim == 3:
                    coord = triangular_coord([alldata[i, 0:2], alldata[i, 0:2]])
                else:
                    coord = tet_coord([alldata[i, 0:3], alldata[i, 0:3],
                                       alldata[i, 0:3]])
                labelcoord = list(zip(*coord))
                unstable_entries[entry] = labelcoord[0]

        return (lines, stable_entries, unstable_entries)

    def show(self):
        """
        Draws the phase diagram using Matplotlib and show it.
        """
        if self._dim < 4:
            plt = self._get_2d_plot()
        elif self._dim == 4:
            plt = self._get_3d_plot()
        plt.show()

    def _get_2d_plot(self):
        '''
        Shows the plot using pylab.  Usually I won't do imports in methods,
        but since plotting is a fairly expensive library to load and not all
        machines have matplotlib installed, I have done it this way.
        '''
        from pymatgen.util.plotting_utils import get_publication_quality_plot
        plt = get_publication_quality_plot(8, 6)
        from matplotlib.font_manager import FontProperties
        (lines, labels, unstable) = self.pd_plot_data
        for x, y in lines:
            plt.plot(x, y, 'bo-', linewidth=3, markeredgecolor='b',
                     markerfacecolor='r', markersize=10)
        font = FontProperties()
        font.set_weight('bold')
        font.set_size(24)

        if len(self._pd.elements) == 3:
            plt.axis('equal')
            plt.xlim((-0.1, 1.2))
            plt.ylim((-0.1, 1.0))
            plt.axis('off')
            center = (0.5, math.sqrt(3) / 6)
        else:
            allcoords = labels.keys()
            miny = min([c[1] for c in allcoords])
            ybuffer = max(abs(miny) * 0.1, 0.1)
            plt.xlim((-0.1, 1.1))
            plt.ylim((miny - ybuffer, ybuffer))
            center = (0.5, miny / 2)

        for coords in sorted(labels.keys(), key=lambda x:-x[1]):
            entry = labels[coords]
            label = entry.name
            vec = (np.array(coords) - center)
            vec = vec / np.linalg.norm(vec) * 10
            valign = "bottom" if vec[1] > 0 else "top"
            if vec[0] < -0.01:
                halign = "right"
            elif vec[0] > 0.01:
                halign = "left"
            else:
                halign = "center"

            plt.annotate(latexify(label), coords, xytext=vec,
                             textcoords='offset points',
                             horizontalalignment=halign,
                             verticalalignment=valign,
                             fontproperties=font)

        if self.show_unstable:
            font = FontProperties()
            font.set_size(16)
            for entry, coords in unstable.items():
                vec = (np.array(coords) - center)
                vec = vec / np.linalg.norm(vec) * 10
                plt.plot(coords[0], coords[1], 'bx', linewidth=3,
                         markeredgecolor='b', markerfacecolor='b',
                         markersize=10)
                plt.annotate(latexify(entry.name), coords, xytext=vec,
                             textcoords='offset points',
                             horizontalalignment=halign, color="b",
                             verticalalignment=valign,
                             fontproperties=font)
        F = plt.gcf()
        F.set_size_inches((8, 6.4))
        return plt

    def _get_3d_plot(self):
        '''
        Shows the plot using pylab.  Usually I won't do imports in methods,
        but since plotting is a fairly expensive library to load and not all
        machines have matplotlib installed, I have done it this way.
        '''
        import matplotlib.pyplot as plt
        import mpl_toolkits.mplot3d.axes3d as p3
        from matplotlib.font_manager import FontProperties
        fig = plt.figure()
        ax = p3.Axes3D(fig)
        font = FontProperties()
        font.set_weight('bold')
        font.set_size(20)
        (lines, labels, unstable) = self.pd_plot_data
        count = 1
        newlabels = list()
        for x, y, z in lines:
            ax.plot(x, y, z, 'bo-', linewidth=3, markeredgecolor='b',
                    markerfacecolor='r', markersize=10)
        for coords in sorted(labels.keys()):
            entry = labels[coords]
            label = entry.name
            if len(entry.composition.elements) == 1:
                # commented out options are only for matplotlib 1.0.  Removed them so that most ppl can use this class.
                ax.text(coords[0], coords[1], coords[2], label)#, horizontalalignment=halign, verticalalignment=valign, fontproperties=font)
            else:
                ax.text(coords[0], coords[1], coords[2], str(count))#, horizontalalignment=halign, verticalalignment=valign, fontproperties=font)
                newlabels.append("{} : {}".format(count, latexify(label)))
                count += 1
        plt.figtext(0.01, 0.01, '\n'.join(newlabels))
        ax.axis('off')
        return plt

    def write_image(self, stream, image_format="svg"):
        '''
        Writes the phase diagram to an image in a stream.

        Args:
            stream:
                stream to write to. Can be a file stream or a StringIO stream.
            image_format
                format for image. Can be any of matplotlib supported formats.
                Defaults to svg for best results for vector graphics.
        '''
        if self._dim < 4:
            plt = self._get_2d_plot()
        elif self._dim == 4:
            plt = self._get_3d_plot()

        f = plt.gcf()
        f.set_size_inches((12, 10))

        plt.savefig(stream, format=image_format)

    def plot_chempot_range_map(self, elements):
        """
        Plot chemical potential range map. Currently works only for 3-component
        PDs.

        Args:
            elements:
                Sequence of elements to be considered as independent variables.
                E.g., if you want to show the stability ranges of all Li-Co-O
                phases wrt to uLi and uO, you will supply
                [Element("Li"), Element("O")]
        """
        from pymatgen.util.plotting_utils import get_publication_quality_plot
        from pymatgen.util.coord_utils import in_coord_list
        plt = get_publication_quality_plot(12, 8)
        analyzer = PDAnalyzer(self._pd)
        chempot_ranges = analyzer.get_chempot_range_map(elements)
        missing_lines = {}

        for entry, lines in chempot_ranges.items():
            center_x = 0
            center_y = 0
            coords = []
            for line in lines:
                (x, y) = line.coords.transpose()
                plt.plot(x, y, 'k')
                center_x += sum(x)
                center_y += sum(y)
                for coord in line.coords:
                    if not in_coord_list(coords, coord):
                        coords.append(coord.tolist())
                    else:
                        coords.remove(coord.tolist())
            comp = entry.composition
            frac_sum = sum([comp.get_atomic_fraction(el) for el in elements])
            if coords and frac_sum < 0.99:
                missing_lines[entry] = coords
            else:
                xy = (center_x / 2 / len(lines), center_y / 2 / len(lines))
                plt.annotate(latexify(entry.name), xy, fontsize=22)

        ax = plt.gca()
        xlim = ax.get_xlim()
        ylim = ax.get_ylim()

        for entry, coords in missing_lines.items():
            center_x = 0
            center_y = 0
            comp = entry.composition
            if not comp.is_element:
                for coord in coords:
                    x = None
                    y = None
                    if entry.composition.get_atomic_fraction(elements[0]) < 0.01:
                        x = [coord[0], min(xlim)]
                        y = [coord[1], coord[1]]
                    elif entry.composition.get_atomic_fraction(elements[1]) < 0.01:
                        x = [coord[0], coord[0]]
                        y = [coord[1], min(ylim)]
                    if x and y:
                        plt.plot(x, y, 'k')
                        center_x += sum(x)
                        center_y += sum(y)
            else:
                center_x = sum(coord[0] for coord in coords) * 2 + xlim[0]
                center_y = sum(coord[1] for coord in coords) * 2 + ylim[0]
            xy = (center_x / 2 / len(coords), center_y / 2 / len(coords))
            plt.annotate(latexify(entry.name), xy, horizontalalignment="center",
                         verticalalignment="center", fontsize=22)

        plt.xlabel("$\mu_{{{0}}} - \mu_{{{0}}}^0$ (eV)".format(elements[0].symbol))
        plt.ylabel("$\mu_{{{0}}} - \mu_{{{0}}}^0$ (eV)".format(elements[1].symbol))
        plt.tight_layout()
        plt.show()

def uniquelines(q):
    '''
    Given all the facets, convert it into a set of unique lines.  Specifically
    used for converting convex hull facets into line pairs of coordinates.

    Args:
        q:
            A 2-dim sequence, where each row represents a facet. E.g.,
            [[1,2,3],[3,6,7],...]

    Returns:
        setoflines:
            A set of tuple of lines.  E.g., ((1,2), (1,3), (2,3), ....)
    '''
    setoflines = set()
    for facets in q:
        for line in itertools.combinations(facets, 2):
            setoflines.add(tuple(sorted(line)))
    return setoflines

def triangular_coord(coord):
    '''
    Convert a 2D coordinate into a triangle-based coordinate system for a
    prettier phase diagram.

    Args:
        coordinate:
            coordinate used in the convex hull computation.

    Returns:
        coordinates in a triangular-based coordinate system.
    '''
    unitvec = np.array([[1, 0], [0.5, math.sqrt(3) / 2]])
    result = np.dot(np.array(coord), unitvec)
    return result.transpose()

def tet_coord(coord):
    '''
    Convert a 3D coordinate into a tetrahedron based coordinate system for a
    prettier phase diagram.

    Args:
        coordinate:
            coordinate used in the convex hull computation.

    Returns:
        coordinates in a tetrahedron-based coordinate system.
    '''
    unitvec = np.array([[1, 0, 0], [0.5, math.sqrt(3) / 2, 0],
                        [0.5, 1.0 / 3.0 * math.sqrt(3) / 2, math.sqrt(6) / 3]])
    result = np.dot(np.array(coord), unitvec)
    return result.transpose()
<commit_msg>Add comments and further improve color scheme.<commit_after>#!/usr/bin/env python

"""
This module provides classes for plotting PhaseDiagram objects.
"""

__author__ = "Shyue Ping Ong"
__copyright__ = "Copyright 2011, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Shyue Ping Ong"
__email__ = "shyue@mit.edu"
__status__ = "Production"
__date__ = "Sep 23, 2011"

import math
import numpy as np
import re
import itertools

from pymatgen.phasediagram.pdanalyzer import PDAnalyzer
from pymatgen.util.string_utils import latexify

class PDPlotter(object):
    '''
    A plotter class for phase diagrams.
    '''

    def __init__(self, phasediagram, show_unstable=True):
        """
        Args:
            phasediagram:
                A PhaseDiagram object.
        """
        self._pd = phasediagram
        self._dim = len(self._pd.elements)
        self.lines = uniquelines(self._pd.facets)
        self.show_unstable = show_unstable
        if self._dim < 2 or self._dim > 4:
            raise ValueError("Only 2-4 components supported!")

    @property
    def pd_plot_data(self):
        '''
        Plot data for phase diagram.
        2-comp - Full hull with energies
        3/4-comp - Projection into 2D or 3D gibbs triangle.

        Returns:
            (lines, stable_entries, unstable_entries):
                - lines is a list of list of coordinates for lines in the PD.
                - stable_entries is a {coordinate : entry} for each stable node
                  in the phase diagram. (Each coordinate can only have one
                  stable phase)
                - unstable_entries is a {entry: coordinates} for all unstable
                  nodes in the phase diagram.
        '''
        pd = self._pd
        entries = pd.qhull_entries
        data = np.array(pd.qhull_data)
        facetlines = self.lines
        lines = list()
        stable_entries = dict()
        for line in facetlines:
            entry1 = entries[line[0]]
            entry2 = entries[line[1]]
            if self._dim == 2:
                x = [data[line[0]][0], data[line[1]][0]]
                y = [pd.get_form_energy_per_atom(entry1),
                     pd.get_form_energy_per_atom(entry2)]
                coord = [x, y]
            elif self._dim == 3:
                coord = triangular_coord(data[line, 0:2])
            else:
                coord = tet_coord(data[line, 0:3])
            lines.append(coord)
            labelcoord = list(zip(*coord))
            stable_entries[labelcoord[0]] = entry1
            stable_entries[labelcoord[1]] = entry2

        allentries = pd.all_entries
        alldata = np.array(pd.all_entries_hulldata)
        unstable_entries = dict()
        stable = pd.stable_entries
        for i in xrange(0, len(allentries)):
            entry = allentries[i]
            if entry not in stable:
                if self._dim == 2:
                    x = [alldata[i][0], alldata[i][0]]
                    y = [pd.get_form_energy_per_atom(entry),
                         pd.get_form_energy_per_atom(entry)]
                    coord = [x, y]
                elif self._dim == 3:
                    coord = triangular_coord([alldata[i, 0:2], alldata[i, 0:2]])
                else:
                    coord = tet_coord([alldata[i, 0:3], alldata[i, 0:3],
                                       alldata[i, 0:3]])
                labelcoord = list(zip(*coord))
                unstable_entries[entry] = labelcoord[0]

        return (lines, stable_entries, unstable_entries)

    def show(self):
        """
        Draws the phase diagram using Matplotlib and show it.
        """
        if self._dim < 4:
            plt = self._get_2d_plot()
        elif self._dim == 4:
            plt = self._get_3d_plot()
        plt.show()

    def _get_2d_plot(self):
        '''
        Shows the plot using pylab.  Usually I won't do imports in methods,
        but since plotting is a fairly expensive library to load and not all
        machines have matplotlib installed, I have done it this way.
        '''
        from pymatgen.util.plotting_utils import get_publication_quality_plot
        plt = get_publication_quality_plot(8, 6)
        from matplotlib.font_manager import FontProperties
        (lines, labels, unstable) = self.pd_plot_data
        for x, y in lines:
            plt.plot(x, y, 'ko-', linewidth=3, markeredgecolor='k',
                     markerfacecolor='b', markersize=10)
        font = FontProperties()
        font.set_weight('bold')
        font.set_size(24)

        # Sets a nice layout depending on the type of PD. Also defines a
        # "center" for the PD, which then allows the annotations to be spread
        # out in a nice manner.
        if len(self._pd.elements) == 3:
            plt.axis('equal')
            plt.xlim((-0.1, 1.2))
            plt.ylim((-0.1, 1.0))
            plt.axis('off')
            center = (0.5, math.sqrt(3) / 6)
        else:
            allcoords = labels.keys()
            miny = min([c[1] for c in allcoords])
            ybuffer = max(abs(miny) * 0.1, 0.1)
            plt.xlim((-0.1, 1.1))
            plt.ylim((miny - ybuffer, ybuffer))
            center = (0.5, miny / 2)

        for coords in sorted(labels.keys(), key=lambda x:-x[1]):
            entry = labels[coords]
            label = entry.name

            # The follow defines an offset for the annotatino text emanating
            # from the center of the PD. Results in fairly nice layouts for the
            # most part.
            vec = (np.array(coords) - center)
            vec = vec / np.linalg.norm(vec) * 10
            valign = "bottom" if vec[1] > 0 else "top"
            if vec[0] < -0.01:
                halign = "right"
            elif vec[0] > 0.01:
                halign = "left"
            else:
                halign = "center"

            plt.annotate(latexify(label), coords, xytext=vec,
                             textcoords='offset points',
                             horizontalalignment=halign,
                             verticalalignment=valign,
                             fontproperties=font)

        if self.show_unstable:
            font = FontProperties()
            font.set_size(16)
            for entry, coords in unstable.items():
                vec = (np.array(coords) - center)
                vec = vec / np.linalg.norm(vec) * 10
                plt.plot(coords[0], coords[1], 'rx', linewidth=3,
                         markeredgecolor='r', markerfacecolor='r',
                         markersize=10)
                plt.annotate(latexify(entry.name), coords, xytext=vec,
                             textcoords='offset points',
                             horizontalalignment=halign, color="b",
                             verticalalignment=valign,
                             fontproperties=font)
        F = plt.gcf()
        F.set_size_inches((8, 6.4))
        return plt

    def _get_3d_plot(self):
        '''
        Shows the plot using pylab.  Usually I won't do imports in methods,
        but since plotting is a fairly expensive library to load and not all
        machines have matplotlib installed, I have done it this way.
        '''
        import matplotlib.pyplot as plt
        import mpl_toolkits.mplot3d.axes3d as p3
        from matplotlib.font_manager import FontProperties
        fig = plt.figure()
        ax = p3.Axes3D(fig)
        font = FontProperties()
        font.set_weight('bold')
        font.set_size(20)
        (lines, labels, unstable) = self.pd_plot_data
        count = 1
        newlabels = list()
        for x, y, z in lines:
            ax.plot(x, y, z, 'bo-', linewidth=3, markeredgecolor='b',
                    markerfacecolor='r', markersize=10)
        for coords in sorted(labels.keys()):
            entry = labels[coords]
            label = entry.name
            if len(entry.composition.elements) == 1:
                # commented out options are only for matplotlib 1.0.  Removed them so that most ppl can use this class.
                ax.text(coords[0], coords[1], coords[2], label)#, horizontalalignment=halign, verticalalignment=valign, fontproperties=font)
            else:
                ax.text(coords[0], coords[1], coords[2], str(count))#, horizontalalignment=halign, verticalalignment=valign, fontproperties=font)
                newlabels.append("{} : {}".format(count, latexify(label)))
                count += 1
        plt.figtext(0.01, 0.01, '\n'.join(newlabels))
        ax.axis('off')
        return plt

    def write_image(self, stream, image_format="svg"):
        '''
        Writes the phase diagram to an image in a stream.

        Args:
            stream:
                stream to write to. Can be a file stream or a StringIO stream.
            image_format
                format for image. Can be any of matplotlib supported formats.
                Defaults to svg for best results for vector graphics.
        '''
        if self._dim < 4:
            plt = self._get_2d_plot()
        elif self._dim == 4:
            plt = self._get_3d_plot()

        f = plt.gcf()
        f.set_size_inches((12, 10))

        plt.savefig(stream, format=image_format)

    def plot_chempot_range_map(self, elements):
        """
        Plot chemical potential range map. Currently works only for 3-component
        PDs.

        Args:
            elements:
                Sequence of elements to be considered as independent variables.
                E.g., if you want to show the stability ranges of all Li-Co-O
                phases wrt to uLi and uO, you will supply
                [Element("Li"), Element("O")]
        """
        from pymatgen.util.plotting_utils import get_publication_quality_plot
        from pymatgen.util.coord_utils import in_coord_list
        plt = get_publication_quality_plot(12, 8)
        analyzer = PDAnalyzer(self._pd)
        chempot_ranges = analyzer.get_chempot_range_map(elements)
        missing_lines = {}

        for entry, lines in chempot_ranges.items():
            center_x = 0
            center_y = 0
            coords = []
            for line in lines:
                (x, y) = line.coords.transpose()
                plt.plot(x, y, 'k')
                center_x += sum(x)
                center_y += sum(y)
                for coord in line.coords:
                    if not in_coord_list(coords, coord):
                        coords.append(coord.tolist())
                    else:
                        coords.remove(coord.tolist())
            comp = entry.composition
            frac_sum = sum([comp.get_atomic_fraction(el) for el in elements])
            if coords and frac_sum < 0.99:
                missing_lines[entry] = coords
            else:
                xy = (center_x / 2 / len(lines), center_y / 2 / len(lines))
                plt.annotate(latexify(entry.name), xy, fontsize=22)

        ax = plt.gca()
        xlim = ax.get_xlim()
        ylim = ax.get_ylim()

        for entry, coords in missing_lines.items():
            center_x = 0
            center_y = 0
            comp = entry.composition
            if not comp.is_element:
                for coord in coords:
                    x = None
                    y = None
                    if entry.composition.get_atomic_fraction(elements[0]) < 0.01:
                        x = [coord[0], min(xlim)]
                        y = [coord[1], coord[1]]
                    elif entry.composition.get_atomic_fraction(elements[1]) < 0.01:
                        x = [coord[0], coord[0]]
                        y = [coord[1], min(ylim)]
                    if x and y:
                        plt.plot(x, y, 'k')
                        center_x += sum(x)
                        center_y += sum(y)
            else:
                center_x = sum(coord[0] for coord in coords) * 2 + xlim[0]
                center_y = sum(coord[1] for coord in coords) * 2 + ylim[0]
            xy = (center_x / 2 / len(coords), center_y / 2 / len(coords))
            plt.annotate(latexify(entry.name), xy, horizontalalignment="center",
                         verticalalignment="center", fontsize=22)

        plt.xlabel("$\mu_{{{0}}} - \mu_{{{0}}}^0$ (eV)".format(elements[0].symbol))
        plt.ylabel("$\mu_{{{0}}} - \mu_{{{0}}}^0$ (eV)".format(elements[1].symbol))
        plt.tight_layout()
        plt.show()

def uniquelines(q):
    '''
    Given all the facets, convert it into a set of unique lines.  Specifically
    used for converting convex hull facets into line pairs of coordinates.

    Args:
        q:
            A 2-dim sequence, where each row represents a facet. E.g.,
            [[1,2,3],[3,6,7],...]

    Returns:
        setoflines:
            A set of tuple of lines.  E.g., ((1,2), (1,3), (2,3), ....)
    '''
    setoflines = set()
    for facets in q:
        for line in itertools.combinations(facets, 2):
            setoflines.add(tuple(sorted(line)))
    return setoflines

def triangular_coord(coord):
    '''
    Convert a 2D coordinate into a triangle-based coordinate system for a
    prettier phase diagram.

    Args:
        coordinate:
            coordinate used in the convex hull computation.

    Returns:
        coordinates in a triangular-based coordinate system.
    '''
    unitvec = np.array([[1, 0], [0.5, math.sqrt(3) / 2]])
    result = np.dot(np.array(coord), unitvec)
    return result.transpose()

def tet_coord(coord):
    '''
    Convert a 3D coordinate into a tetrahedron based coordinate system for a
    prettier phase diagram.

    Args:
        coordinate:
            coordinate used in the convex hull computation.

    Returns:
        coordinates in a tetrahedron-based coordinate system.
    '''
    unitvec = np.array([[1, 0, 0], [0.5, math.sqrt(3) / 2, 0],
                        [0.5, 1.0 / 3.0 * math.sqrt(3) / 2, math.sqrt(6) / 3]])
    result = np.dot(np.array(coord), unitvec)
    return result.transpose()
<|endoftext|>


<commit_before>"""
Backend for NCS VPN module.

Author: Henrik Thostrup Jensen <htj at nordu.net>
Copyright: NORDUnet(2011-2013)
"""

import base64
import random

from twisted.python import log
from twisted.web.error import Error as WebError

from opennsa import constants as cnt, config
from opennsa.backends.common import genericbackend
from opennsa.protocols.shared import httpclient

# basic payload
#
#<service xmlns="http://tail-f.com/ns/ncs" >
#  <object-id>nsi-vpn</object-id>
#  <type>
#    <vpn xmlns="http://nordu.net/ns/ncs/vpn">
#      <side-a>
#        <router>routerA</router>
#        <interface>interface1</interface>
#      </side-a>
#      <side-b>
#        <router></router>
#        <interface>ge-1/0/1</interface>
#      </side-b>
#      <vlan>1720</vlan>
#    </vpn>
#  </type>
#</service>
#
# encapsulation type can be ethernet or ethernet-vlan
# vlan must be specified if encapsulation-type is ethernet-vlan, otherwise not
#
# the payload must be posted to the services url, e.g.,:
# http://localhost:8080/api/running/services
#
# To tear down the VPN, do a DELETE against
# "http://localhost:8080/api/running/services/service/nsi-vpn"
#
# The connection id -> object-id mapping is hence rather important to remember, but it can be the

NCS_TIMEOUT = 45 # ncs typically spends 28-31 seconds creating/deleting a vpn

ETHERNET_VPN_PAYLOAD_BASE = """
<bod xmlns="http://nordu.net/ns/ncs/vpn">
    <service-name>%(service_name)s</service-name>
    <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
    </side-a>
    <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
  </side-b>
  <vlan>%(vlan)i</vlan>
</bod>
"""

ETHERNET_VLAN_VPN_PAYLOAD_BASE = """
<bod xmlns="http://nordu.net/ns/ncs/vpn">
    <service-name>%(service_name)s</service-name>
    <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
    </side-a>
    <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
  </side-b>
  <vlan>%(vlan)i</vlan>
</bod>
"""

ETHERNET_VLAN_REWRITE_VPN_PAYLOAD_BASE = """
<bod xmlns="http://nordu.net/ns/ncs/vpn">
    <service-name>%(service_name)s</service-name>
    <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
    </side-a>
    <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
    </side-b>
    <vlan-side-a>%(vlan_a)i</vlan-side-a>
    <vlan-side-b>%(vlan_b)i</vlan-side-b>
</bod>
"""

LOG_SYSTEM = 'opennsa.ncsvpn'

class NCSVPNTarget(object):

    def __init__(self, router, interface, vlan=None):
        self.router = router
        self.interface = interface
        self.vlan = vlan

    def __str__(self):
        if self.vlan:
            return '<NCSVPNTarget %s/%s#%i>' % (self.router, self.interface, self.vlan)
        else:
            return '<NCSVPNTarget %s/%s>' % (self.router, self.interface)

def createVPNPayload(service_name, source_target, dest_target):

    intps = {
        'service_name'  : service_name,
        'router_a'      : source_target.router,
        'interface_a'   : source_target.interface,
        'router_b'      : dest_target.router,
        'interface_b'   : dest_target.interface
    }

    if source_target.vlan and dest_target.vlan:
        #assert source_target.vlan == dest_target.vlan, 'VLANs must match (until we get rewrite in place)'
        if source_target.vlan == dest_target.vlan:
            intps['vlan'] = source_target.vlan
            payload = ETHERNET_VLAN_VPN_PAYLOAD_BASE % intps
        else:
            intps['vlan_a'] = source_target.vlan
            intps['vlan_b'] = dest_target.vlan
            payload = ETHERNET_VLAN_REWRITE_VPN_PAYLOAD_BASE % intps
    else:
        payload = ETHERNET_VPN_PAYLOAD_BASE % intps

    return payload

def _extractErrorMessage(failure):
    # used to extract error messages from http requests
    if isinstance(failure.value, WebError):
        return failure.value.response
    else:
        return failure.getErrorMessage()

class NCSVPNConnectionManager:

    def __init__(self, ncs_services_url, user, password, port_map, log_system):
        self.ncs_services_url = ncs_services_url
        self.user             = user
        self.password         = password
        self.port_map         = port_map
        self.log_system       = log_system

    def getResource(self, port, label_type, label_value):
        assert label_type in (None, cnt.ETHERNET_VLAN), 'Label must be None or VLAN'
        return port + ':' + str(label_value) # port contains router and port

    def getTarget(self, port, label_type, label_value):
        assert label_type in (None, cnt.ETHERNET_VLAN), 'Label must be None or VLAN'
        if label_type == cnt.ETHERNET_VLAN:
            vlan = int(label_value)
            assert 1 <= vlan <= 4095, 'Invalid label value for vlan: %s' % label_value

        ri = self.port_map[port]
        router, interface = ri.split(':')
        return NCSVPNTarget(router, interface, vlan)

    def createConnectionId(self, source_target, dest_target):
        return 'ON-' + str(random.randint(100000,999999))

    def canSwapLabel(self, label_type):
        return label_type == cnt.ETHERNET_VLAN

    def _createAuthzHeader(self):
        return 'Basic ' + base64.b64encode( self.user + ':' + self.password)

    def _createHeaders(self):
        headers = {}
        headers['Content-Type'] = 'application/vnd.yang.data+xml'
        headers['Authorization'] = self._createAuthzHeader()
        return headers

    def setupLink(self, connection_id, source_target, dest_target, bandwidth):
        payload = createVPNPayload(connection_id, source_target, dest_target)
        headers = self._createHeaders()

        def linkUp(_):
            log.msg('Link %s -> %s up' % (source_target, dest_target), system=self.log_system)

        def error(failure):
            log.msg('Error bringing up link %s -> %s' % (source_target, dest_target), system=self.log_system)
            log.msg('Message: %s' % _extractErrorMessage(failure), system=self.log_system)
            return failure

        d = httpclient.httpRequest(self.ncs_services_url, payload, headers, method='POST', timeout=NCS_TIMEOUT)
        d.addCallbacks(linkUp, error)
        return d

    def teardownLink(self, connection_id, source_target, dest_target, bandwidth):
        service_url = self.ncs_services_url + '/bod/' + connection_id
        headers = self._createHeaders()

        def linkDown(_):
            log.msg('Link %s -> %s down' % (source_target, dest_target), system=self.log_system)

        def error(failure):
            log.msg('Error bringing down link %s -> %s' % (source_target, dest_target), system=self.log_system)
            log.msg('Message: %s' % _extractErrorMessage(failure), system=self.log_system)
            return failure

        d = httpclient.httpRequest(service_url, None, headers, method='DELETE', timeout=NCS_TIMEOUT)
        d.addCallbacks(linkDown, error)
        return d

def NCSVPNBackend(network_name, nrm_ports, parent_requester, cfg):

    name = 'NCS VPN %s' % network_name
    nrm_map  = dict( [ (p.name, p) for p in nrm_ports ] ) # for the generic backend
    port_map = dict( [ (p.name, p.interface) for p in nrm_ports ] ) # for the nrm backend

    # extract config items
    ncs_services_url = str(cfg[config.NCS_SERVICES_URL]) # convert from unicode
    user             = cfg[config.NCS_USER]
    password         = cfg[config.NCS_PASSWORD]

    cm = NCSVPNConnectionManager(ncs_services_url, user, password, port_map, name)
    return genericbackend.GenericBackend(network_name, nrm_map, cm, parent_requester, name)

<commit_msg>increase ncs timeout to 60 seconds, as 45 is not always enough<commit_after>"""
Backend for NCS VPN module.

Author: Henrik Thostrup Jensen <htj at nordu.net>
Copyright: NORDUnet(2011-2013)
"""

import base64
import random

from twisted.python import log
from twisted.web.error import Error as WebError

from opennsa import constants as cnt, config
from opennsa.backends.common import genericbackend
from opennsa.protocols.shared import httpclient

# basic payload
#
#<service xmlns="http://tail-f.com/ns/ncs" >
#  <object-id>nsi-vpn</object-id>
#  <type>
#    <vpn xmlns="http://nordu.net/ns/ncs/vpn">
#      <side-a>
#        <router>routerA</router>
#        <interface>interface1</interface>
#      </side-a>
#      <side-b>
#        <router></router>
#        <interface>ge-1/0/1</interface>
#      </side-b>
#      <vlan>1720</vlan>
#    </vpn>
#  </type>
#</service>
#
# encapsulation type can be ethernet or ethernet-vlan
# vlan must be specified if encapsulation-type is ethernet-vlan, otherwise not
#
# the payload must be posted to the services url, e.g.,:
# http://localhost:8080/api/running/services
#
# To tear down the VPN, do a DELETE against
# "http://localhost:8080/api/running/services/service/nsi-vpn"
#
# The connection id -> object-id mapping is hence rather important to remember, but it can be the

NCS_TIMEOUT = 60 # ncs typically spends 25-32 seconds creating/deleting a vpn, sometimes a bit more

ETHERNET_VPN_PAYLOAD_BASE = """
<bod xmlns="http://nordu.net/ns/ncs/vpn">
    <service-name>%(service_name)s</service-name>
    <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
    </side-a>
    <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
  </side-b>
  <vlan>%(vlan)i</vlan>
</bod>
"""

ETHERNET_VLAN_VPN_PAYLOAD_BASE = """
<bod xmlns="http://nordu.net/ns/ncs/vpn">
    <service-name>%(service_name)s</service-name>
    <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
    </side-a>
    <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
  </side-b>
  <vlan>%(vlan)i</vlan>
</bod>
"""

ETHERNET_VLAN_REWRITE_VPN_PAYLOAD_BASE = """
<bod xmlns="http://nordu.net/ns/ncs/vpn">
    <service-name>%(service_name)s</service-name>
    <side-a>
        <router>%(router_a)s</router>
        <interface>%(interface_a)s</interface>
    </side-a>
    <side-b>
        <router>%(router_b)s</router>
        <interface>%(interface_b)s</interface>
    </side-b>
    <vlan-side-a>%(vlan_a)i</vlan-side-a>
    <vlan-side-b>%(vlan_b)i</vlan-side-b>
</bod>
"""

LOG_SYSTEM = 'opennsa.ncsvpn'

class NCSVPNTarget(object):

    def __init__(self, router, interface, vlan=None):
        self.router = router
        self.interface = interface
        self.vlan = vlan

    def __str__(self):
        if self.vlan:
            return '<NCSVPNTarget %s/%s#%i>' % (self.router, self.interface, self.vlan)
        else:
            return '<NCSVPNTarget %s/%s>' % (self.router, self.interface)

def createVPNPayload(service_name, source_target, dest_target):

    intps = {
        'service_name'  : service_name,
        'router_a'      : source_target.router,
        'interface_a'   : source_target.interface,
        'router_b'      : dest_target.router,
        'interface_b'   : dest_target.interface
    }

    if source_target.vlan and dest_target.vlan:
        #assert source_target.vlan == dest_target.vlan, 'VLANs must match (until we get rewrite in place)'
        if source_target.vlan == dest_target.vlan:
            intps['vlan'] = source_target.vlan
            payload = ETHERNET_VLAN_VPN_PAYLOAD_BASE % intps
        else:
            intps['vlan_a'] = source_target.vlan
            intps['vlan_b'] = dest_target.vlan
            payload = ETHERNET_VLAN_REWRITE_VPN_PAYLOAD_BASE % intps
    else:
        payload = ETHERNET_VPN_PAYLOAD_BASE % intps

    return payload

def _extractErrorMessage(failure):
    # used to extract error messages from http requests
    if isinstance(failure.value, WebError):
        return failure.value.response
    else:
        return failure.getErrorMessage()

class NCSVPNConnectionManager:

    def __init__(self, ncs_services_url, user, password, port_map, log_system):
        self.ncs_services_url = ncs_services_url
        self.user             = user
        self.password         = password
        self.port_map         = port_map
        self.log_system       = log_system

    def getResource(self, port, label_type, label_value):
        assert label_type in (None, cnt.ETHERNET_VLAN), 'Label must be None or VLAN'
        return port + ':' + str(label_value) # port contains router and port

    def getTarget(self, port, label_type, label_value):
        assert label_type in (None, cnt.ETHERNET_VLAN), 'Label must be None or VLAN'
        if label_type == cnt.ETHERNET_VLAN:
            vlan = int(label_value)
            assert 1 <= vlan <= 4095, 'Invalid label value for vlan: %s' % label_value

        ri = self.port_map[port]
        router, interface = ri.split(':')
        return NCSVPNTarget(router, interface, vlan)

    def createConnectionId(self, source_target, dest_target):
        return 'ON-' + str(random.randint(100000,999999))

    def canSwapLabel(self, label_type):
        return label_type == cnt.ETHERNET_VLAN

    def _createAuthzHeader(self):
        return 'Basic ' + base64.b64encode( self.user + ':' + self.password)

    def _createHeaders(self):
        headers = {}
        headers['Content-Type'] = 'application/vnd.yang.data+xml'
        headers['Authorization'] = self._createAuthzHeader()
        return headers

    def setupLink(self, connection_id, source_target, dest_target, bandwidth):
        payload = createVPNPayload(connection_id, source_target, dest_target)
        headers = self._createHeaders()

        def linkUp(_):
            log.msg('Link %s -> %s up' % (source_target, dest_target), system=self.log_system)

        def error(failure):
            log.msg('Error bringing up link %s -> %s' % (source_target, dest_target), system=self.log_system)
            log.msg('Message: %s' % _extractErrorMessage(failure), system=self.log_system)
            return failure

        d = httpclient.httpRequest(self.ncs_services_url, payload, headers, method='POST', timeout=NCS_TIMEOUT)
        d.addCallbacks(linkUp, error)
        return d

    def teardownLink(self, connection_id, source_target, dest_target, bandwidth):
        service_url = self.ncs_services_url + '/bod/' + connection_id
        headers = self._createHeaders()

        def linkDown(_):
            log.msg('Link %s -> %s down' % (source_target, dest_target), system=self.log_system)

        def error(failure):
            log.msg('Error bringing down link %s -> %s' % (source_target, dest_target), system=self.log_system)
            log.msg('Message: %s' % _extractErrorMessage(failure), system=self.log_system)
            return failure

        d = httpclient.httpRequest(service_url, None, headers, method='DELETE', timeout=NCS_TIMEOUT)
        d.addCallbacks(linkDown, error)
        return d

def NCSVPNBackend(network_name, nrm_ports, parent_requester, cfg):

    name = 'NCS VPN %s' % network_name
    nrm_map  = dict( [ (p.name, p) for p in nrm_ports ] ) # for the generic backend
    port_map = dict( [ (p.name, p.interface) for p in nrm_ports ] ) # for the nrm backend

    # extract config items
    ncs_services_url = str(cfg[config.NCS_SERVICES_URL]) # convert from unicode
    user             = cfg[config.NCS_USER]
    password         = cfg[config.NCS_PASSWORD]

    cm = NCSVPNConnectionManager(ncs_services_url, user, password, port_map, name)
    return genericbackend.GenericBackend(network_name, nrm_map, cm, parent_requester, name)

<|endoftext|>


<commit_before>"""Basic Network Types."""

from pyof.foundation.base import GenericStruct
from pyof.foundation.basic_types import (BinaryData, FixedTypeList, HWAddress,
                                         UBInt8, UBInt16, UBInt64)

from pyof.foundation.exceptions import PackException

__all__ = ('Ethernet', 'LLDP', 'LLDP_TLV', 'ListOfTLVs')

class Ethernet(GenericStruct):
    destination = HWAddress()
    source = HWAddress()
    type = UBInt16()
    data = BinaryData()

    def __init__(self, destination=None, source=None, type=None, data=b''):
        super().__init__()
        self.destination = destination
        self.source = source
        self.type = type
        self.data = data

    def __hash__(self):
        return hash(self.pack())

    def get_hash(self):
        return self.__hash__()

class LLDP_TLV(GenericStruct):
    """TLV structure of LLDP packets."""

    #: This field actually have 7 bits, so on pack/unpack we will need to do
    #: some magic.
    type = UBInt8()
    #: Defines the length of subtype + value This field actually have 9 bits.
    length = UBInt8()
    #: Binary data to be passed along. No matter the type and length, here you
    #: will always need to add content as binary data. So, if you have
    #: something that is not binary (bytes), then pack it befor adding.
    value = BinaryData()

    def __init__(self, type, value=None):
        self.type = type
        self.value = value if value is not None else b''

    def _type_is_valid(self):
        """Check if the type value can be represented in 7 bits."""
        # 7 bits can represent 128 values, from 0 to 127.
        if self.type >= 0 and self.type < 128:
            return True
        else:
            return False

    def _length_is_valid(self):
        """Check if the length value can be represented in 9 bits."""
        # 9 bits can represent 512 values, from 0 to 511
        if self.length >= 0 and self.length < 511:
            return True
        else:
            return False

    def update_length(self):
        """Updates the length attribute of the instance."""
        self.length = len(BinaryData().pack(self.value))

    def pack(self, value=None):
        """Pack the LLDP_TLV in a binary representation.

        Returns:
            bytes: Binary representation of the struct object.

        Raises:
            :exc:`~.exceptions.ValidationError`: If validation fails.
        """
        if value is None:
            output = b''

            # Packing the 7bits type field
            if not self._type_is_valid():
                error_msg = '%s is not a valid value for TLV_type field'
                raise PackException(error_msg, self.type)

            self.update_length()

            if not self._length_is_valid():
                error_msg = '%s is not a valid value for TLV_length field'
                raise PackException(error_msg, self.length)

            # Here we are gathering the 7 bits from self.type with the higher
            # bit from self.length to pack it in one byte.
            output += UBInt8().pack(self.type << 1 | self.length >> 8)
            # Here we are removing the higher bit from length to pack it as
            # 1 byte also.
            output += UBInt8().pack(self.length & ((1 << 8) - 1))
            output += BinaryData().pack(self.value)

            return output
        elif isinstance(value, type(self)):
            return value.pack()
        else:
            msg = "{} is not an instance of {}".format(value,
                                                       type(self).__name__)
            raise PackException(msg)

    def unpack(self, buffer, offset=0):
        buffer = buffer[offset:]
        type_and_length = UBInt16().unpack(buffer[0:16]).value
        self.type = type_and_length >> 9
        self.length = type_and_length & ((1 << 10) - 1)
        self.value = BinaryData().unpack(buffer[16:16+self.length])

    def get_size(self, value=None):
        if isinstance(value, type(self)):
            return value.get_size()
        else:
            return 2 + len(BinaryData().pack(self.value))

class ListOfTLVs(FixedTypeList):
    """List of instances of LLDP_TLV class.

    Represented by instances of LLDP_TLV and used on :class:`LLDP` objects.
    """

    def __init__(self, items=None):
        """The constructor takes the option paramenter below.

        Args: items(:class:`list`, :class:`LLDP_TLV`): One :class:`LLDP_TLV`
            instance or list.
        """
        super().__init__(pyof_class=LLDP_TLV,
                         items=items)

class LLDP(GenericStruct):
    ethernet_header = Ethernet(type=35020)
    lldp_tlv_chassis_id = LLDP_TLV(type=1)
    lldp_tlv_port = LLDP_TLV(type=2)
    lldp_tlv_ttl = LLDP_TLV(type=3)
    # list_of_extra_tlvs = ListOfTLVs()
    lldp_tlv_end = LLDP_TLV(type=0, length=0, value=0)

    def __init__(self, source=None, tlvs=None):
        super().__init__()
        # 01-80-C2-00-00-00 or 01-80-C2-00-00-03 are also valid values for LLDP
        # multicast
        self.ethernet_header.destination = '01-80-C2-00-00-0E'
        self.ethernet_header.source = source
<commit_msg>Improvements on LLDP class<commit_after>"""Basic Network Types."""

from pyof.foundation.base import GenericStruct
from pyof.foundation.basic_types import (BinaryData, FixedTypeList, HWAddress,
                                         UBInt8, UBInt16, UBInt64)

from pyof.foundation.exceptions import PackException

__all__ = ('Ethernet', 'LLDP', 'LLDP_TLV', 'ListOfTLVs')

class Ethernet(GenericStruct):
    destination = HWAddress()
    source = HWAddress()
    type = UBInt16()
    data = BinaryData()

    def __init__(self, destination=None, source=None, type=None, data=b''):
        super().__init__()
        self.destination = destination
        self.source = source
        self.type = type
        self.data = data

    def __hash__(self):
        return hash(self.pack())

    def get_hash(self):
        return self.__hash__()

class LLDP_TLV(GenericStruct):
    """TLV structure of LLDP packets."""

    #: This field actually have 7 bits, so on pack/unpack we will need to do
    #: some magic.
    type = UBInt8()
    #: Defines the length of subtype + value This field actually have 9 bits.
    length = UBInt8()
    #: Binary data to be passed along. No matter the type and length, here you
    #: will always need to add content as binary data. So, if you have
    #: something that is not binary (bytes), then pack it befor adding.
    value = BinaryData()

    def __init__(self, type, value=None):
        self.type = type
        self.value = value if value is not None else b''

    def _type_is_valid(self):
        """Check if the type value can be represented in 7 bits."""
        # 7 bits can represent 128 values, from 0 to 127.
        if self.type >= 0 and self.type < 128:
            return True
        else:
            return False

    def _length_is_valid(self):
        """Check if the length value can be represented in 9 bits."""
        # 9 bits can represent 512 values, from 0 to 511
        if self.length >= 0 and self.length < 511:
            return True
        else:
            return False

    def update_length(self):
        """Updates the length attribute of the instance."""
        self.length = len(BinaryData().pack(self.value))

    def pack(self, value=None):
        """Pack the LLDP_TLV in a binary representation.

        Returns:
            bytes: Binary representation of the struct object.

        Raises:
            :exc:`~.exceptions.ValidationError`: If validation fails.
        """
        if value is None:
            output = b''

            # Packing the 7bits type field
            if not self._type_is_valid():
                error_msg = '%s is not a valid value for TLV_type field'
                raise PackException(error_msg, self.type)

            self.update_length()

            if not self._length_is_valid():
                error_msg = '%s is not a valid value for TLV_length field'
                raise PackException(error_msg, self.length)

            # Here we are gathering the 7 bits from self.type with the higher
            # bit from self.length to pack it in one byte.
            output += UBInt8().pack(self.type << 1 | self.length >> 8)
            # Here we are removing the higher bit from length to pack it as
            # 1 byte also.
            output += UBInt8().pack(self.length & ((1 << 8) - 1))
            output += BinaryData().pack(self.value)

            return output
        elif isinstance(value, type(self)):
            return value.pack()
        else:
            msg = "{} is not an instance of {}".format(value,
                                                       type(self).__name__)
            raise PackException(msg)

    def unpack(self, buffer, offset=0):
        buffer = buffer[offset:]
        type_and_length = UBInt16().unpack(buffer[0:16]).value
        self.type = type_and_length >> 9
        self.length = type_and_length & ((1 << 10) - 1)
        self.value = BinaryData().unpack(buffer[16:16+self.length])

    def get_size(self, value=None):
        if isinstance(value, type(self)):
            return value.get_size()
        else:
            return 2 + len(BinaryData().pack(self.value))

class ListOfTLVs(FixedTypeList):
    """List of instances of LLDP_TLV class.

    Represented by instances of LLDP_TLV and used on :class:`LLDP` objects.
    """

    def __init__(self, items=None):
        """The constructor takes the option paramenter below.

        Args: items(:class:`list`, :class:`LLDP_TLV`): One :class:`LLDP_TLV`
            instance or list.
        """
        super().__init__(pyof_class=LLDP_TLV,
                         items=items)

class LLDP(GenericStruct):
    """LLDP class."""

    ethernet_header = Ethernet(type=35020)
    tlv_chassis_id = LLDP_TLV(type=1)
    tlv_port_id = LLDP_TLV(type=2)
    #: TTL time is given in seconds, between 0 and 65535
    tlv_ttl = LLDP_TLV(type=3)
    extra_tlvs = ListOfTLVs()
    tlv_end = LLDP_TLV(type=0, value=0)

    def __init__(self, source=None, chassis_id=None, port_id=None, ttl=30,
                 extra_tlvs=None):
        super().__init__()
        # 01-80-C2-00-00-00 or 01-80-C2-00-00-03 are also valid values for LLDP
        # multicast
        self.ethernet_header.destination = '01:80:C2:00:00:0E'
        self.ethernet_header.source = source
        #: For TLV_chassis_id we are using subtype 6 (InterfaceName) and
        #: passing, for now, the switch dpid. And the datapath_id of the
        #: switches are defined as UBInt64 on the OpenFlow spec (v0x01).
        self.tlv_chassis_id.value = b'\x06' + UBInt64().pack(chassis_id)
        #: For TLV_port_id we are using subtype 5 (InterfaceName) and
        #: passing, for now, the port_no.
        self.tlv_port_id.value = b'\x05' + UBInt16().pack(port_id)
        self.tlv_ttl.value = UBInt16().pack(ttl)
        self.extra_tlvs = [] if extra_tlvs is None else extra_tlvs
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-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/>.

"""Engine: A collection of fundamental functions for initializing and running
calculations."""

import os
import sys
import time
import getpass
import itertools
import operator
import traceback
from datetime import datetime

import celery.task.control

import openquake.engine

from django.core import exceptions
from django import db as django_db

from openquake.engine import logs
from openquake.server.db import models
from openquake.engine.utils import config, tasks
from openquake.engine.celery_node_monitor import CeleryNodeMonitor
from openquake.server.db.schema.upgrades import upgrader

from openquake.commonlib import readinput, valid, datastore, export
from openquake.calculators import base

def get_calc_id(job_id=None):
    """
    Return the latest calc_id by looking both at the datastore
    and the database.
    """
    calcs = datastore.get_calc_ids(datastore.DATADIR)
    calc_id = 1 if not calcs else calcs[-1]
    if job_id is None:
        try:
            job_id = models.OqJob.objects.latest('id').id
        except exceptions.ObjectDoesNotExist:
            job_id = 1
    return max(calc_id, job_id)

INPUT_TYPES = set(dict(models.INPUT_TYPE_CHOICES))

UNABLE_TO_DEL_HC_FMT = 'Unable to delete hazard calculation: %s'
UNABLE_TO_DEL_RC_FMT = 'Unable to delete risk calculation: %s'

TERMINATE = valid.boolean(
    config.get('celery', 'terminate_workers_on_revoke') or 'false')

USE_CELERY = valid.boolean(config.get('celery', 'use_celery') or 'false')

class InvalidCalculationID(Exception):
    pass

RISK_HAZARD_MAP = dict(
    scenario_risk=['scenario', 'scenario_risk'],
    scenario_damage=['scenario', 'scenario_damage'],
    classical_risk=['classical', 'classical_risk'],
    classical_bcr=['classical', 'classical_bcr'],
    classical_damage=['classical', 'classical_damage'],
    event_based_risk=['event_based', 'event_based_risk'])

# this is called only if USE_CELERY is True
def cleanup_after_job(job, terminate, task_ids=()):
    """
    Release the resources used by an openquake job.
    In particular revoke the running tasks (if any).

    :param int job_id: the job id
    :param bool terminate: the celery revoke command terminate flag
    :param task_ids: celery task IDs
    """
    # Using the celery API, terminate and revoke and terminate any running
    # tasks associated with the current job.
    if task_ids:
        logs.LOG.warn('Revoking %d tasks', len(task_ids))
    else:  # this is normal when OQ_NO_DISTRIBUTE=1
        logs.LOG.debug('No task to revoke')
    for tid in task_ids:
        celery.task.control.revoke(tid, terminate=terminate)
        logs.LOG.debug('Revoked task %s', tid)

def create_job(user_name="openquake", hc_id=None):
    """
    Create job for the given user, return it.

    :param str username:
        Username of the user who owns/started this job. If the username doesn't
        exist, a user record for this name will be created.
    :param hc_id:
        If not None, then the created job is a risk job
    :returns:
        :class:`openquake.server.db.models.OqJob` instance.
    """
    calc_id = get_calc_id() + 1
    job = models.OqJob.objects.create(
        id=calc_id,
        description='A job',
        user_name=user_name,
        ds_calc_dir=os.path.join(datastore.DATADIR, 'calc_%s' % calc_id))
    if hc_id:
        job.hazard_calculation = models.OqJob.objects.get(pk=hc_id)
    job.save()
    return job

# used by bin/openquake and openquake.server.views
def run_calc(job, log_level, log_file, exports, hazard_calculation_id=None):
    """
    Run a calculation.

    :param job:
        :class:`openquake.server.db.model.OqJob` instance
    :param str log_level:
        The desired logging level. Valid choices are 'debug', 'info',
        'progress', 'warn', 'error', and 'critical'.
    :param str log_file:
        Complete path (including file name) to file where logs will be written.
        If `None`, logging will just be printed to standard output.
    :param exports:
        A comma-separated string of export types.
    """
    # first of all check the database version and exit if the db is outdated
    upgrader.check_versions(django_db.connections['admin'])
    with logs.handle(job, log_level, log_file):  # run the job
        tb = 'None\n'
        try:
            _do_run_calc(job, exports, hazard_calculation_id)
            job.status = 'complete'
        except:
            tb = traceback.format_exc()
            try:
                logs.LOG.critical(tb)
                job.status = 'failed'
            except:  # an OperationalError may always happen
                sys.stderr.write(tb)
            raise
        finally:
            # try to save the job stats on the database and then clean up;
            # if there was an error in the calculation, this part may fail;
            # in such a situation, we simply log the cleanup error without
            # taking further action, so that the real error can propagate
            try:
                job.is_running = False
                job.stop_time = datetime.utcnow()
                job.save()
                if USE_CELERY:
                    cleanup_after_job(
                        job, TERMINATE, tasks.OqTaskManager.task_ids)
            except:
                # log the finalization error only if there is no real error
                if tb == 'None\n':
                    logs.LOG.error('finalizing', exc_info=True)
        expose_outputs(job.calc.datastore, job)
    return job.calc

# keep this as a private function, since it is mocked by engine_test.py
def _do_run_calc(job, exports, hazard_calculation_id):
    job.calc.run(exports=exports, hazard_calculation_id=hazard_calculation_id)

def del_calc(job_id):
    """
    Delete a calculation and all associated outputs.

    :param job_id:
        ID of a :class:`~openquake.server.db.models.OqJob`.
    """
    try:
        job = models.OqJob.objects.get(id=job_id)
    except exceptions.ObjectDoesNotExist:
        raise RuntimeError('Unable to delete hazard calculation: '
                           'ID=%s does not exist' % job_id)

    user = getpass.getuser()
    if job.user_name == user:
        # we are allowed to delete this

        # but first, check if any risk calculations are referencing any of our
        # outputs, or the hazard calculation itself
        msg = UNABLE_TO_DEL_HC_FMT % (
            'The following risk calculations are referencing this hazard'
            ' calculation: %s')

        assoc_outputs = models.OqJob.objects.filter(hazard_calculation=job)
        if assoc_outputs.count() > 0:
            raise RuntimeError(
                msg % ', '.join(str(x.id) for x in assoc_outputs))

        # No risk calculation are referencing what we want to delete.
        # Carry on with the deletion. Notice that we cannot use job.delete()
        # directly because Django is so stupid that it reads from the database
        # all the records to delete before deleting them: thus, it runs out
        # of memory for large calculations
        curs = models.getcursor('admin')
        curs.execute('DELETE FROM uiapi.oq_job WHERE id=%s', (job_id,))
    else:
        # this doesn't belong to the current user
        raise RuntimeError(UNABLE_TO_DEL_HC_FMT % 'Access denied')
    try:
        os.remove(job.ds_calc_dir + '.hdf5')
    except:
        pass
    else:
        print('Removed %s' % job.ds_calc_dir + '.hdf5')

def list_outputs(job_id, full=True):
    """
    List the outputs for a given
    :class:`~openquake.server.db.models.OqJob`.

    :param job_id:
        ID of a calculation.
    :param bool full:
        If True produce a full listing, otherwise a short version
    """
    outputs = get_outputs(job_id)
    if models.oqparam(job_id).calculation_mode == 'scenario':
        # ignore SES output
        outputs = [o for o in outputs if o.output_type != 'ses']
    print_outputs_summary(outputs, full)

def print_results(job_id, duration, list_outputs):
    print 'Calculation %d completed in %d seconds. Results:' % (
        job_id, duration)
    list_outputs(job_id, full=False)

def print_outputs_summary(outputs, full=True):
    """
    List of :class:`openquake.server.db.models.Output` objects.
    """
    if len(outputs) > 0:
        truncated = False
        print '  id | output_type | name'
        for output_type, group in itertools.groupby(
                sorted(outputs, key=operator.attrgetter('output_type')),
                key=operator.attrgetter('output_type')):
            outs = sorted(group, key=operator.attrgetter('display_name'))
            for i, o in enumerate(outs):
                if not full and i >= 10:
                    print ' ... | %s | %d additional output(s)' % (
                        o.get_output_type_display(), len(outs) - 10)
                    truncated = True
                    break
                print '%4d | %s | %s' % (
                    o.id, o.get_output_type_display(), o.display_name)
        if truncated:
            print ('Some outputs where not shown. You can see the full list '
                   'with the command\n`oq-engine --list-outputs`')

# this function is called only by openquake_cli.py, not by the engine server
def run_job(cfg_file, log_level, log_file, exports='',
            hazard_output_id=None, hazard_calculation_id=None):
    """
    Run a job using the specified config file and other options.

    :param str cfg_file:
        Path to calculation config (INI-style) files.
    :param str log_level:
        'debug', 'info', 'warn', 'error', or 'critical'
    :param str log_file:
        Path to log file.
    :param exports:
        A comma-separated string of export types requested by the user.
        Currently only 'xml' is supported.
    """
    # first of all check the database version and exit if the db is outdated
    upgrader.check_versions(django_db.connections['admin'])
    with CeleryNodeMonitor(openquake.engine.no_distribute(), interval=3):
        job = job_from_file(
            cfg_file, getpass.getuser(), log_level, exports,
            hazard_output_id=hazard_output_id,
            hazard_calculation_id=hazard_calculation_id)
        t0 = time.time()
        run_calc(job, log_level, log_file, exports,
                 hazard_calculation_id=hazard_calculation_id)
        duration = time.time() - t0
        if job.status == 'complete':
            print_results(job.id, duration, list_outputs)
        else:
            sys.exit('Calculation %s failed' % job.id)
    return job

DISPLAY_NAME = dict(dmg_by_asset='dmg_by_asset_and_collapse_map')

def expose_outputs(dstore, job):
    """
    Build a correspondence between the outputs in the datastore and the
    ones in the database.

    :param dstore: a datastore instance
    :param job: an OqJob instance
    """
    exportable = set(ekey[0] for ekey in export.export)
    oq = job.get_oqparam()

    # small hack: remove the sescollection outputs from scenario
    # calculators, as requested by Vitor
    calcmode = oq.calculation_mode
    if 'scenario' in calcmode and 'sescollection' in exportable:
        exportable.remove('sescollection')
    uhs = oq.uniform_hazard_spectra
    if uhs and 'hmaps' in dstore:
        models.Output.objects.create_output(
            job, 'uhs', output_type='datastore', ds_key='uhs')

    for key in dstore:
        if key in exportable:
            if key == 'realizations' and len(dstore['realizations']) == 1:
                continue  # there is no point in exporting a single realization
                ds_key=key)

def check_hazard_risk_consistency(haz_job, risk_mode):
    """
    Make sure that the provided hazard job is the right one for the
    current risk calculator.

    :param job:
        an OqJob instance referring to the previous hazard calculation
    :param risk_mode:
        the `calculation_mode` string of the current risk calculation
    """
    # check for obsolete calculation_mode
    if risk_mode in ('classical', 'event_based', 'scenario'):
        raise ValueError('Please change calculation_mode=%s into %s_risk '
                         'in the .ini file' % (risk_mode, risk_mode))

    # check calculation_mode consistency
    prev_mode = haz_job.get_oqparam().calculation_mode
    ok_mode = RISK_HAZARD_MAP[risk_mode]
    if prev_mode not in ok_mode:
        raise InvalidCalculationID(
            'In order to run a risk calculation of kind %r, '
            'you need to provide a calculation of kind %r, '
            'but you provided a %r instead' %
            (risk_mode, ok_mode, prev_mode))

@django_db.transaction.atomic
def job_from_file(cfg_file, username, log_level='info', exports='',
                  hazard_output_id=None, hazard_calculation_id=None, **extras):
    """
    Create a full job profile from a job config file.

    :param str cfg_file:
        Path to the job.ini files.
    :param str username:
        The user who will own this job profile and all results.
    :param str log_level:
        Desired log level.
    :param exports:
        Comma-separated sting of desired export types
    :param hazard_output_id:
        Hazard output ID
    :param hazard_calculation_id:
        Hazard calculation ID
    :params extras:
        Extra parameters (used only in the tests to override the params)

    :returns:
        :class:`openquake.server.db.models.OqJob` object
    :raises:
        `RuntimeError` if the input job configuration is not valid
    """
    # create the current job
    job = create_job(username, hazard_calculation_id)
    with logs.handle(job, log_level):
        # read calculation params and create the calculation profile
        params = readinput.get_params([cfg_file])
        params.update(extras)
        # build and validate an OqParam object
        oq = readinput.get_oqparam(params)
        calc = base.calculators(oq, calc_id=job.id)
        calc.save_params()
        job.description = oq.description
        job.calc = calc
    return job

# this is patched in the tests
def get_outputs(job_id):
    """
    :param job_id:
        ID of a calculation.
    :returns:
        A sequence of :class:`openquake.server.db.models.Output` objects
    """
    return models.Output.objects.filter(oq_job=job_id)
<commit_msg>Fixed a merge error<commit_after># -*- coding: utf-8 -*-
# vim: tabstop=4 shiftwidth=4 softtabstop=4
#
# Copyright (C) 2010-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/>.

"""Engine: A collection of fundamental functions for initializing and running
calculations."""

import os
import sys
import time
import getpass
import itertools
import operator
import traceback
from datetime import datetime

import celery.task.control

import openquake.engine

from django.core import exceptions
from django import db as django_db

from openquake.engine import logs
from openquake.server.db import models
from openquake.engine.utils import config, tasks
from openquake.engine.celery_node_monitor import CeleryNodeMonitor
from openquake.server.db.schema.upgrades import upgrader

from openquake.commonlib import readinput, valid, datastore, export
from openquake.calculators import base

def get_calc_id(job_id=None):
    """
    Return the latest calc_id by looking both at the datastore
    and the database.
    """
    calcs = datastore.get_calc_ids(datastore.DATADIR)
    calc_id = 1 if not calcs else calcs[-1]
    if job_id is None:
        try:
            job_id = models.OqJob.objects.latest('id').id
        except exceptions.ObjectDoesNotExist:
            job_id = 1
    return max(calc_id, job_id)

INPUT_TYPES = set(dict(models.INPUT_TYPE_CHOICES))

UNABLE_TO_DEL_HC_FMT = 'Unable to delete hazard calculation: %s'
UNABLE_TO_DEL_RC_FMT = 'Unable to delete risk calculation: %s'

TERMINATE = valid.boolean(
    config.get('celery', 'terminate_workers_on_revoke') or 'false')

USE_CELERY = valid.boolean(config.get('celery', 'use_celery') or 'false')

class InvalidCalculationID(Exception):
    pass

RISK_HAZARD_MAP = dict(
    scenario_risk=['scenario', 'scenario_risk'],
    scenario_damage=['scenario', 'scenario_damage'],
    classical_risk=['classical', 'classical_risk'],
    classical_bcr=['classical', 'classical_bcr'],
    classical_damage=['classical', 'classical_damage'],
    event_based_risk=['event_based', 'event_based_risk'])

# this is called only if USE_CELERY is True
def cleanup_after_job(job, terminate, task_ids=()):
    """
    Release the resources used by an openquake job.
    In particular revoke the running tasks (if any).

    :param int job_id: the job id
    :param bool terminate: the celery revoke command terminate flag
    :param task_ids: celery task IDs
    """
    # Using the celery API, terminate and revoke and terminate any running
    # tasks associated with the current job.
    if task_ids:
        logs.LOG.warn('Revoking %d tasks', len(task_ids))
    else:  # this is normal when OQ_NO_DISTRIBUTE=1
        logs.LOG.debug('No task to revoke')
    for tid in task_ids:
        celery.task.control.revoke(tid, terminate=terminate)
        logs.LOG.debug('Revoked task %s', tid)

def create_job(user_name="openquake", hc_id=None):
    """
    Create job for the given user, return it.

    :param str username:
        Username of the user who owns/started this job. If the username doesn't
        exist, a user record for this name will be created.
    :param hc_id:
        If not None, then the created job is a risk job
    :returns:
        :class:`openquake.server.db.models.OqJob` instance.
    """
    calc_id = get_calc_id() + 1
    job = models.OqJob.objects.create(
        id=calc_id,
        description='A job',
        user_name=user_name,
        ds_calc_dir=os.path.join(datastore.DATADIR, 'calc_%s' % calc_id))
    if hc_id:
        job.hazard_calculation = models.OqJob.objects.get(pk=hc_id)
    job.save()
    return job

# used by bin/openquake and openquake.server.views
def run_calc(job, log_level, log_file, exports, hazard_calculation_id=None):
    """
    Run a calculation.

    :param job:
        :class:`openquake.server.db.model.OqJob` instance
    :param str log_level:
        The desired logging level. Valid choices are 'debug', 'info',
        'progress', 'warn', 'error', and 'critical'.
    :param str log_file:
        Complete path (including file name) to file where logs will be written.
        If `None`, logging will just be printed to standard output.
    :param exports:
        A comma-separated string of export types.
    """
    # first of all check the database version and exit if the db is outdated
    upgrader.check_versions(django_db.connections['admin'])
    with logs.handle(job, log_level, log_file):  # run the job
        tb = 'None\n'
        try:
            _do_run_calc(job, exports, hazard_calculation_id)
            job.status = 'complete'
        except:
            tb = traceback.format_exc()
            try:
                logs.LOG.critical(tb)
                job.status = 'failed'
            except:  # an OperationalError may always happen
                sys.stderr.write(tb)
            raise
        finally:
            # try to save the job stats on the database and then clean up;
            # if there was an error in the calculation, this part may fail;
            # in such a situation, we simply log the cleanup error without
            # taking further action, so that the real error can propagate
            try:
                job.is_running = False
                job.stop_time = datetime.utcnow()
                job.save()
                if USE_CELERY:
                    cleanup_after_job(
                        job, TERMINATE, tasks.OqTaskManager.task_ids)
            except:
                # log the finalization error only if there is no real error
                if tb == 'None\n':
                    logs.LOG.error('finalizing', exc_info=True)
        expose_outputs(job.calc.datastore, job)
    return job.calc

# keep this as a private function, since it is mocked by engine_test.py
def _do_run_calc(job, exports, hazard_calculation_id):
    job.calc.run(exports=exports, hazard_calculation_id=hazard_calculation_id)

def del_calc(job_id):
    """
    Delete a calculation and all associated outputs.

    :param job_id:
        ID of a :class:`~openquake.server.db.models.OqJob`.
    """
    try:
        job = models.OqJob.objects.get(id=job_id)
    except exceptions.ObjectDoesNotExist:
        raise RuntimeError('Unable to delete hazard calculation: '
                           'ID=%s does not exist' % job_id)

    user = getpass.getuser()
    if job.user_name == user:
        # we are allowed to delete this

        # but first, check if any risk calculations are referencing any of our
        # outputs, or the hazard calculation itself
        msg = UNABLE_TO_DEL_HC_FMT % (
            'The following risk calculations are referencing this hazard'
            ' calculation: %s')

        assoc_outputs = models.OqJob.objects.filter(hazard_calculation=job)
        if assoc_outputs.count() > 0:
            raise RuntimeError(
                msg % ', '.join(str(x.id) for x in assoc_outputs))

        # No risk calculation are referencing what we want to delete.
        # Carry on with the deletion. Notice that we cannot use job.delete()
        # directly because Django is so stupid that it reads from the database
        # all the records to delete before deleting them: thus, it runs out
        # of memory for large calculations
        curs = models.getcursor('admin')
        curs.execute('DELETE FROM uiapi.oq_job WHERE id=%s', (job_id,))
    else:
        # this doesn't belong to the current user
        raise RuntimeError(UNABLE_TO_DEL_HC_FMT % 'Access denied')
    try:
        os.remove(job.ds_calc_dir + '.hdf5')
    except:
        pass
    else:
        print('Removed %s' % job.ds_calc_dir + '.hdf5')

def list_outputs(job_id, full=True):
    """
    List the outputs for a given
    :class:`~openquake.server.db.models.OqJob`.

    :param job_id:
        ID of a calculation.
    :param bool full:
        If True produce a full listing, otherwise a short version
    """
    outputs = get_outputs(job_id)
    if models.oqparam(job_id).calculation_mode == 'scenario':
        # ignore SES output
        outputs = [o for o in outputs if o.output_type != 'ses']
    print_outputs_summary(outputs, full)

def print_results(job_id, duration, list_outputs):
    print 'Calculation %d completed in %d seconds. Results:' % (
        job_id, duration)
    list_outputs(job_id, full=False)

def print_outputs_summary(outputs, full=True):
    """
    List of :class:`openquake.server.db.models.Output` objects.
    """
    if len(outputs) > 0:
        truncated = False
        print '  id | output_type | name'
        for output_type, group in itertools.groupby(
                sorted(outputs, key=operator.attrgetter('output_type')),
                key=operator.attrgetter('output_type')):
            outs = sorted(group, key=operator.attrgetter('display_name'))
            for i, o in enumerate(outs):
                if not full and i >= 10:
                    print ' ... | %s | %d additional output(s)' % (
                        o.get_output_type_display(), len(outs) - 10)
                    truncated = True
                    break
                print '%4d | %s | %s' % (
                    o.id, o.get_output_type_display(), o.display_name)
        if truncated:
            print ('Some outputs where not shown. You can see the full list '
                   'with the command\n`oq-engine --list-outputs`')

# this function is called only by openquake_cli.py, not by the engine server
def run_job(cfg_file, log_level, log_file, exports='',
            hazard_output_id=None, hazard_calculation_id=None):
    """
    Run a job using the specified config file and other options.

    :param str cfg_file:
        Path to calculation config (INI-style) files.
    :param str log_level:
        'debug', 'info', 'warn', 'error', or 'critical'
    :param str log_file:
        Path to log file.
    :param exports:
        A comma-separated string of export types requested by the user.
        Currently only 'xml' is supported.
    """
    # first of all check the database version and exit if the db is outdated
    upgrader.check_versions(django_db.connections['admin'])
    with CeleryNodeMonitor(openquake.engine.no_distribute(), interval=3):
        job = job_from_file(
            cfg_file, getpass.getuser(), log_level, exports,
            hazard_output_id=hazard_output_id,
            hazard_calculation_id=hazard_calculation_id)
        t0 = time.time()
        run_calc(job, log_level, log_file, exports,
                 hazard_calculation_id=hazard_calculation_id)
        duration = time.time() - t0
        if job.status == 'complete':
            print_results(job.id, duration, list_outputs)
        else:
            sys.exit('Calculation %s failed' % job.id)
    return job

DISPLAY_NAME = dict(dmg_by_asset='dmg_by_asset_and_collapse_map')

def expose_outputs(dstore, job):
    """
    Build a correspondence between the outputs in the datastore and the
    ones in the database.

    :param dstore: a datastore instance
    :param job: an OqJob instance
    """
    exportable = set(ekey[0] for ekey in export.export)
    oq = job.get_oqparam()

    # small hack: remove the sescollection outputs from scenario
    # calculators, as requested by Vitor
    calcmode = oq.calculation_mode
    if 'scenario' in calcmode and 'sescollection' in exportable:
        exportable.remove('sescollection')
    uhs = oq.uniform_hazard_spectra
    if uhs and 'hmaps' in dstore:
        models.Output.objects.create_output(
            job, 'uhs', output_type='datastore', ds_key='uhs')

    for key in dstore:
        if key in exportable:
            if key == 'realizations' and len(dstore['realizations']) == 1:
                continue  # there is no point in exporting a single realization
            models.Output.objects.create_output(
                job, DISPLAY_NAME.get(key, key), output_type='datastore',
                ds_key=key)

def check_hazard_risk_consistency(haz_job, risk_mode):
    """
    Make sure that the provided hazard job is the right one for the
    current risk calculator.

    :param job:
        an OqJob instance referring to the previous hazard calculation
    :param risk_mode:
        the `calculation_mode` string of the current risk calculation
    """
    # check for obsolete calculation_mode
    if risk_mode in ('classical', 'event_based', 'scenario'):
        raise ValueError('Please change calculation_mode=%s into %s_risk '
                         'in the .ini file' % (risk_mode, risk_mode))

    # check calculation_mode consistency
    prev_mode = haz_job.get_oqparam().calculation_mode
    ok_mode = RISK_HAZARD_MAP[risk_mode]
    if prev_mode not in ok_mode:
        raise InvalidCalculationID(
            'In order to run a risk calculation of kind %r, '
            'you need to provide a calculation of kind %r, '
            'but you provided a %r instead' %
            (risk_mode, ok_mode, prev_mode))

@django_db.transaction.atomic
def job_from_file(cfg_file, username, log_level='info', exports='',
                  hazard_output_id=None, hazard_calculation_id=None, **extras):
    """
    Create a full job profile from a job config file.

    :param str cfg_file:
        Path to the job.ini files.
    :param str username:
        The user who will own this job profile and all results.
    :param str log_level:
        Desired log level.
    :param exports:
        Comma-separated sting of desired export types
    :param hazard_output_id:
        Hazard output ID
    :param hazard_calculation_id:
        Hazard calculation ID
    :params extras:
        Extra parameters (used only in the tests to override the params)

    :returns:
        :class:`openquake.server.db.models.OqJob` object
    :raises:
        `RuntimeError` if the input job configuration is not valid
    """
    # create the current job
    job = create_job(username, hazard_calculation_id)
    with logs.handle(job, log_level):
        # read calculation params and create the calculation profile
        params = readinput.get_params([cfg_file])
        params.update(extras)
        # build and validate an OqParam object
        oq = readinput.get_oqparam(params)
        calc = base.calculators(oq, calc_id=job.id)
        calc.save_params()
        job.description = oq.description
        job.calc = calc
    return job

# this is patched in the tests
def get_outputs(job_id):
    """
    :param job_id:
        ID of a calculation.
    :returns:
        A sequence of :class:`openquake.server.db.models.Output` objects
    """
    return models.Output.objects.filter(oq_job=job_id)
<|endoftext|>