<commit_before>from __future__ import absolute_import
from os import path
import itertools
from collections import defaultdict
from functools import partial
import json
from gusto.diagnostics import Diagnostics
from sys import exit
from abc import ABCMeta, abstractmethod
from firedrake import FiniteElement, TensorProductElement, HDiv, \
    FunctionSpace, MixedFunctionSpace, VectorFunctionSpace, \
    interval, triangle, Function, Mesh, functionspaceimpl,\
    Expression, File, TestFunction, TrialFunction, inner, div, FacetNormal, \
    ds_tb, dx, solve, op2, par_loop, READ, WRITE, DumbCheckpoint, \
    FILE_CREATE, FILE_UPDATE
import numpy as np

class State(object):
    """
    Build a model state to keep the variables in, and specify parameters.

    :arg mesh: The :class:`Mesh` to use.
    :arg vertical_degree: integer, the degree for spaces in the vertical
    (specifies the degree for the pressure space, other spaces are inferred)
    defaults to 1.
    :arg horizontal_degree: integer, the degree for spaces in the horizontal
    (specifies the degree for the pressure space, other spaces are inferred)
    defaults to 1.
    :arg family: string, specifies the velocity space family to use.
    Options:
    "RT": The Raviart-Thomas family (default, recommended for quads)
    "BDM": The BDM family
    "BDFM": The BDFM family
    :arg timestepping: class containing timestepping parameters
    :arg output: class containing output parameters
    :arg parameters: class containing physical parameters
    :arg diagnostics: class containing diagnostic methods
    :arg fieldlist: list of prognostic field names

    """
    __metaclass__ = ABCMeta

    def __init__(self, mesh, vertical_degree=1, horizontal_degree=1,
                 family="RT", z=None, k=None, Omega=None, mu=None,
                 timestepping=None,
                 output=None,
                 parameters=None,
                 diagnostics=None,
                 fieldlist=None,
                 diagnostic_fields=[]):

        self.z = z
        self.k = k
        self.Omega = Omega
        self.mu = mu
        self.timestepping = timestepping
        self.output = output
        self.parameters = parameters
        if fieldlist is None:
            raise RuntimeError("You must provide a fieldlist containing the names of the prognostic fields")
        else:
            self.fieldlist = fieldlist
        if diagnostics is not None:
            self.diagnostics = diagnostics
        else:
            self.diagnostics = Diagnostics(*fieldlist)
        self.diagnostic_fields = diagnostic_fields

        # The mesh
        self.mesh = mesh

        # Build the spaces
        self._build_spaces(mesh, vertical_degree,
                           horizontal_degree, family)

        # Allocate state
        self._allocate_state()
        self.field_dict = {name: func for (name, func) in
                           zip(self.fieldlist, self.xn.split())}

        self.dumpfile = None

    def dump(self, t=0, pickup=False):
        """
        Dump output
        :arg t: the current model time (default is zero).
        :arg pickup: recover state from the checkpointing file if true,
        otherwise dump and checkpoint to disk. (default is False).
        """

        # default behaviour is to dump all prognostic fields
        if self.output.dumplist is None:
            self.output.dumplist = self.fieldlist

        # if there are fields to be dumped in latlon coordinates,
        # setup the latlon coordinate mesh
        if len(self.output.dumplist_latlon) > 0:
            field_dict_ll = {}
            mesh_ll = get_latlon_mesh(self.mesh)

        funcs = self.xn.split()
        field_dict = {name: func for (name, func) in zip(self.fieldlist, funcs)}
        to_dump = [] #fields to output to dump and checkpoint
        to_pickup = [] #fields to pick up from checkpoint
        for name, f in field_dict.iteritems():
            if name in self.output.dumplist:
                to_dump.append(f)
                to_pickup.append(f)
            f.rename(name=name)

        for diagnostic in self.diagnostic_fields:
            to_dump.append(diagnostic(self))

        steady_state_dump_err = defaultdict(bool)
        steady_state_dump_err.update(self.output.steady_state_dump_err)
        for name, f, f_init in zip(self.fieldlist, funcs, self.x_init.split()):
            if steady_state_dump_err[name]:
                err = Function(f.function_space(), name=name+'err').assign(f-f_init)
                field_dict[name+"err"] = err
                self.diagnostics.register(name+"err")
                to_dump.append(err)
                f_init.rename(f.name()+"_init")
                to_dump.append(f_init)
                to_pickup.append(f_init)

        meanfields = defaultdict(lambda: None)
        meanfields.update(self.output.meanfields)
        for name, meanfield in meanfields.iteritems():
            if meanfield is not None:
                field = field_dict[name]
                diff = Function(
                    field.function_space(),
                    name=field.name()+"_perturbation").assign(field - meanfield)
                self.diagnostics.register(name+"perturbation")
                field_dict[name+"perturbation"] = diff
                to_dump.append(diff)
            mean_name=field.name()+"_bar"
            meanfield.rename(name=mean_name)
            to_dump.append(meanfield)
            to_pickup.append(meanfield)

        # make functions on latlon mesh, as specified by dumplist_latlon
        to_dump_latlon = []
        for name in self.output.dumplist_latlon:
            f = field_dict[name]
            f_ll = Function(functionspaceimpl.WithGeometry(f.function_space(), mesh_ll), val=f.topological, name=name+'_ll')
            field_dict_ll[name] = f_ll
            to_dump_latlon.append(f_ll)

        self.dumpdir = path.join("results", self.output.dirname)
        outfile = path.join(self.dumpdir, "field_output.pvd")
        if self.dumpfile is None:
            if self.mesh.comm.rank == 0 and path.exists(self.dumpdir):
                exit("results directory '%s' already exists" % self.dumpdir)
            self.dumpcount = itertools.count()
            self.dumpfile = File(outfile, project_output=self.output.project_fields, comm=self.mesh.comm)
            self.diagnostic_data = defaultdict(partial(defaultdict, list))

            # make output file for fields on latlon mesh if required
            if len(self.output.dumplist_latlon) > 0:
                outfile_latlon = path.join(self.dumpdir, "field_output_latlon.pvd")
                self.dumpfile_latlon = File(outfile_latlon, project_output=self.output.project_fields,
                                            comm=self.mesh.comm)

        if(pickup):
            #Open the checkpointing file for writing
            chkfile = path.join(self.dumpdir, "chkpt")
            with DumbCheckpoint(chkfile, mode=FILE_READ) as chk:
                #Recover all the fields from the checkpoint
                for field in to_pickup:
                    chk.load(field)
                chk.read_attribute("/","time",t)

        elif (next(self.dumpcount) % self.output.dumpfreq) == 0:

            self.dumpfile.write(*to_dump)

            if len(self.output.dumplist_latlon) > 0:
                self.dumpfile_latlon.write(*to_dump_latlon)

            for name in self.diagnostics.fields:
                data = self.diagnostics.l2(field_dict[name])
                self.diagnostic_data[name]["l2"].append(data)

            #Open the checkpointing file
            chkfile = path.join(self.dumpdir, "chkpt")
            with DumbCheckpoint(chkfile, mode=FILE_CREATE) as chk:
                #Dump all the fields to a checkpoint
                for field in to_dump:
                    chk.store(field)
                chk.write_attribute("/","time",t)

        return t

    def diagnostic_dump(self):

        with open(path.join(self.dumpdir, "diagnostics.json"), "w") as f:
            f.write(json.dumps(self.diagnostic_data, indent=4))

    def initialise(self, initial_conditions):
        """
        Initialise state variables
        """

        for x, ic in zip(self.x_init.split(), initial_conditions):
            x.assign(ic)

    @abstractmethod
    def _build_spaces(self, mesh, vertical_degree, horizontal_degree, family):

        """
        Build function spaces:
        """
        pass

    def _allocate_state(self):
        """
        Construct Functions to store the state variables.
        """

        W = self.W
        self.xn = Function(W)
        self.x_init = Function(W)
        self.xstar = Function(W)
        self.xp = Function(W)
        self.xnp1 = Function(W)
        self.xrhs = Function(W)
        self.dy = Function(W)

class CompressibleState(State):

    def __init__(self, mesh, vertical_degree=1, horizontal_degree=1,
                 family="RT", z=None, k=None, Omega=None, mu=None,
                 timestepping=None,
                 output=None,
                 parameters=None,
                 diagnostics=None,
                 fieldlist=None,
                 diagnostic_fields=[],
                 on_sphere=False):

        super(CompressibleState, self).__init__(mesh=mesh,
                                                vertical_degree=vertical_degree,
                                                horizontal_degree=horizontal_degree,
                                                family=family,
                                                z=z, k=k, Omega=Omega, mu=mu,
                                                timestepping=timestepping,
                                                output=output,
                                                parameters=parameters,
                                                diagnostics=diagnostics,
                                                fieldlist=fieldlist,
                                                diagnostic_fields=diagnostic_fields)

        # build the geopotential
        if parameters.geopotential:
            V = FunctionSpace(mesh, "CG", 1)
            if on_sphere:
                self.Phi = Function(V).interpolate(Expression("pow(x[0]*x[0]+x[1]*x[1]+x[2]*x[2],0.5)"))
            else:
                self.Phi = Function(V).interpolate(Expression("x[1]"))
            self.Phi *= parameters.g

        if self.k is None:
            # build the vertical normal
            w = TestFunction(self.Vv)
            u = TrialFunction(self.Vv)
            self.k = Function(self.Vv)
            n = FacetNormal(self.mesh)
            krhs = -div(w)*self.z*dx + inner(w,n)*self.z*ds_tb
            klhs = inner(w,u)*dx
            solve(klhs == krhs, self.k)

    def set_reference_profiles(self, rho_ref, theta_ref):
        """
        Initialise reference profiles
        :arg rho_ref: :class:`.Function` object, initial rho
        :arg theta_ref: :class:`.Function` object, initial theta
        """

        self.rhobar = Function(self.V[1])
        self.thetabar = Function(self.V[2])

        self.rhobar.project(rho_ref)
        self.thetabar.project(theta_ref)

    def _build_spaces(self, mesh, vertical_degree, horizontal_degree, family):
        """
        Build:
        velocity space self.V2,
        pressure space self.V3,
        temperature space self.Vt,
        mixed function space self.W = (V2,V3,Vt)
        """

        # horizontal base spaces
        cell = mesh._base_mesh.ufl_cell()
        if(cell.cellname() == 'triangle'):
            cell = triangle
        if(cell.cellname() == 'interval'):
            cell = interval
        S1 = FiniteElement(family, cell, horizontal_degree+1)
        S2 = FiniteElement("DG", cell, horizontal_degree)

        # vertical base spaces
        T0 = FiniteElement("CG", interval, vertical_degree+1)
        T1 = FiniteElement("DG", interval, vertical_degree)

        # build spaces V2, V3, Vt
        V2h_elt = HDiv(TensorProductElement(S1, T1))
        V2t_elt = TensorProductElement(S2, T0)
        V3_elt = TensorProductElement(S2, T1)
        V2v_elt = HDiv(V2t_elt)
        V2_elt = V2h_elt + V2v_elt

        self.V_elt = [0,0,0]
        self.V_elt[0] = V2_elt
        self.V_elt[1] = V3_elt
        self.V_elt[2] = V2t_elt

        self.V = [0,0,0]
        self.V[0] = FunctionSpace(mesh, V2_elt)
        self.V[1] = FunctionSpace(mesh, V3_elt)
        self.V[2] = FunctionSpace(mesh, V2t_elt)

        self.Vv = FunctionSpace(mesh, V2v_elt)

        self.W = MixedFunctionSpace((self.V[0], self.V[1], self.V[2]))

class ShallowWaterState(State):

    def _build_spaces(self, mesh, vertical_degree, horizontal_degree, family):

        if vertical_degree is not None:
            raise ValueError('Mesh is not extruded in the vertical for shallow water')

        cell = mesh.ufl_cell().cellname()

        V1_elt = FiniteElement(family, cell, horizontal_degree+1)

        self.V = [0,0]
        self.V[0] = FunctionSpace(mesh,V1_elt)
        self.V[1] = FunctionSpace(mesh,"DG",horizontal_degree)

        self.W = MixedFunctionSpace((self.V[0], self.V[1]))

def get_latlon_mesh(mesh):
    coords_orig = mesh.coordinates
    mesh_dg_fs = VectorFunctionSpace(mesh, "DG", 1)
    coords_dg = Function(mesh_dg_fs)
    coords_latlon = Function(mesh_dg_fs)
    par_loop("""
    for (int i=0; i<3; i++) {
    for (int j=0; j<3; j++) {
    dg[i][j] = cg[i][j];
    }
    }
    """, dx, {'dg': (coords_dg, WRITE),
              'cg': (coords_orig, READ)})

    # lat-lon 'x' = atan2(y, x)
    coords_latlon.dat.data[:,0] = np.arctan2(coords_dg.dat.data[:,1], coords_dg.dat.data[:,0])
    # lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2))
    coords_latlon.dat.data[:,1] = np.arcsin(coords_dg.dat.data[:,2]/np.sqrt(coords_dg.dat.data[:,0]**2 + coords_dg.dat.data[:,1]**2 + coords_dg.dat.data[:,2]**2))
    coords_latlon.dat.data[:,2] = 0.0

    kernel = op2.Kernel("""
    #define PI 3.141592653589793
    #define TWO_PI 6.283185307179586
    void splat_coords(double **coords) {
        double diff0 = (coords[0][0] - coords[1][0]);
        double diff1 = (coords[0][0] - coords[2][0]);
        double diff2 = (coords[1][0] - coords[2][0]);

        if (fabs(diff0) > PI || fabs(diff1) > PI || fabs(diff2) > PI) {
            const int sign0 = coords[0][0] < 0 ? -1 : 1;
            const int sign1 = coords[1][0] < 0 ? -1 : 1;
            const int sign2 = coords[2][0] < 0 ? -1 : 1;
            if (sign0 < 0) {
                coords[0][0] += TWO_PI;
            }
            if (sign1 < 0) {
                coords[1][0] += TWO_PI;
            }
            if (sign2 < 0) {
                coords[2][0] += TWO_PI;
            }
        }
    }""", "splat_coords")

    op2.par_loop(kernel, coords_latlon.cell_set,
                 coords_latlon.dat(op2.RW, coords_latlon.cell_node_map()))
    return Mesh(coords_latlon)
<commit_msg>make a backup chkpt in case chkpting fails!<commit_after>from __future__ import absolute_import
from os import path
import itertools
from collections import defaultdict
from functools import partial
import json
from gusto.diagnostics import Diagnostics
from sys import exit
from abc import ABCMeta, abstractmethod
from firedrake import FiniteElement, TensorProductElement, HDiv, \
    FunctionSpace, MixedFunctionSpace, VectorFunctionSpace, \
    interval, triangle, Function, Mesh, functionspaceimpl,\
    Expression, File, TestFunction, TrialFunction, inner, div, FacetNormal, \
    ds_tb, dx, solve, op2, par_loop, READ, WRITE, DumbCheckpoint, \
    FILE_CREATE, FILE_UPDATE
import numpy as np

class State(object):
    """
    Build a model state to keep the variables in, and specify parameters.

    :arg mesh: The :class:`Mesh` to use.
    :arg vertical_degree: integer, the degree for spaces in the vertical
    (specifies the degree for the pressure space, other spaces are inferred)
    defaults to 1.
    :arg horizontal_degree: integer, the degree for spaces in the horizontal
    (specifies the degree for the pressure space, other spaces are inferred)
    defaults to 1.
    :arg family: string, specifies the velocity space family to use.
    Options:
    "RT": The Raviart-Thomas family (default, recommended for quads)
    "BDM": The BDM family
    "BDFM": The BDFM family
    :arg timestepping: class containing timestepping parameters
    :arg output: class containing output parameters
    :arg parameters: class containing physical parameters
    :arg diagnostics: class containing diagnostic methods
    :arg fieldlist: list of prognostic field names

    """
    __metaclass__ = ABCMeta

    def __init__(self, mesh, vertical_degree=1, horizontal_degree=1,
                 family="RT", z=None, k=None, Omega=None, mu=None,
                 timestepping=None,
                 output=None,
                 parameters=None,
                 diagnostics=None,
                 fieldlist=None,
                 diagnostic_fields=[]):

        self.z = z
        self.k = k
        self.Omega = Omega
        self.mu = mu
        self.timestepping = timestepping
        self.output = output
        self.parameters = parameters
        if fieldlist is None:
            raise RuntimeError("You must provide a fieldlist containing the names of the prognostic fields")
        else:
            self.fieldlist = fieldlist
        if diagnostics is not None:
            self.diagnostics = diagnostics
        else:
            self.diagnostics = Diagnostics(*fieldlist)
        self.diagnostic_fields = diagnostic_fields

        # The mesh
        self.mesh = mesh

        # Build the spaces
        self._build_spaces(mesh, vertical_degree,
                           horizontal_degree, family)

        # Allocate state
        self._allocate_state()
        self.field_dict = {name: func for (name, func) in
                           zip(self.fieldlist, self.xn.split())}

        self.dumpfile = None

    def dump(self, t=0, pickup=False):
        """
        Dump output
        :arg t: the current model time (default is zero).
        :arg pickup: recover state from the checkpointing file if true,
        otherwise dump and checkpoint to disk. (default is False).
        """

        # default behaviour is to dump all prognostic fields
        if self.output.dumplist is None:
            self.output.dumplist = self.fieldlist

        # if there are fields to be dumped in latlon coordinates,
        # setup the latlon coordinate mesh
        if len(self.output.dumplist_latlon) > 0:
            field_dict_ll = {}
            mesh_ll = get_latlon_mesh(self.mesh)

        funcs = self.xn.split()
        field_dict = {name: func for (name, func) in zip(self.fieldlist, funcs)}
        to_dump = [] #fields to output to dump and checkpoint
        to_pickup = [] #fields to pick up from checkpoint
        for name, f in field_dict.iteritems():
            if name in self.output.dumplist:
                to_dump.append(f)
                to_pickup.append(f)
            f.rename(name=name)

        for diagnostic in self.diagnostic_fields:
            to_dump.append(diagnostic(self))

        steady_state_dump_err = defaultdict(bool)
        steady_state_dump_err.update(self.output.steady_state_dump_err)
        for name, f, f_init in zip(self.fieldlist, funcs, self.x_init.split()):
            if steady_state_dump_err[name]:
                err = Function(f.function_space(), name=name+'err').assign(f-f_init)
                field_dict[name+"err"] = err
                self.diagnostics.register(name+"err")
                to_dump.append(err)
                f_init.rename(f.name()+"_init")
                to_dump.append(f_init)
                to_pickup.append(f_init)

        meanfields = defaultdict(lambda: None)
        meanfields.update(self.output.meanfields)
        for name, meanfield in meanfields.iteritems():
            if meanfield is not None:
                field = field_dict[name]
                diff = Function(
                    field.function_space(),
                    name=field.name()+"_perturbation").assign(field - meanfield)
                self.diagnostics.register(name+"perturbation")
                field_dict[name+"perturbation"] = diff
                to_dump.append(diff)
            mean_name=field.name()+"_bar"
            meanfield.rename(name=mean_name)
            to_dump.append(meanfield)
            to_pickup.append(meanfield)

        # make functions on latlon mesh, as specified by dumplist_latlon
        to_dump_latlon = []
        for name in self.output.dumplist_latlon:
            f = field_dict[name]
            f_ll = Function(functionspaceimpl.WithGeometry(f.function_space(), mesh_ll), val=f.topological, name=name+'_ll')
            field_dict_ll[name] = f_ll
            to_dump_latlon.append(f_ll)

        self.dumpdir = path.join("results", self.output.dirname)
        outfile = path.join(self.dumpdir, "field_output.pvd")
        if self.dumpfile is None:
            if self.mesh.comm.rank == 0 and path.exists(self.dumpdir):
                exit("results directory '%s' already exists" % self.dumpdir)
            self.dumpcount = itertools.count()
            self.dumpfile = File(outfile, project_output=self.output.project_fields, comm=self.mesh.comm)
            self.diagnostic_data = defaultdict(partial(defaultdict, list))

            # make output file for fields on latlon mesh if required
            if len(self.output.dumplist_latlon) > 0:
                outfile_latlon = path.join(self.dumpdir, "field_output_latlon.pvd")
                self.dumpfile_latlon = File(outfile_latlon, project_output=self.output.project_fields,
                                            comm=self.mesh.comm)

        if(pickup):
            #Open the checkpointing file for writing
            chkfile = path.join(self.dumpdir, "chkpt")
            with DumbCheckpoint(chkfile, mode=FILE_READ) as chk:
                #Recover all the fields from the checkpoint
                for field in to_pickup:
                    chk.load(field)
                chk.read_attribute("/","time",t)

        elif (next(self.dumpcount) % self.output.dumpfreq) == 0:

            self.dumpfile.write(*to_dump)

            if len(self.output.dumplist_latlon) > 0:
                self.dumpfile_latlon.write(*to_dump_latlon)

            for name in self.diagnostics.fields:
                data = self.diagnostics.l2(field_dict[name])
                self.diagnostic_data[name]["l2"].append(data)

            #Open the checkpointing file (backup version)
            files = ["chkptbk", "chkpt"]

            chkfile = path.join(self.dumpdir, "chkptbk")
            with DumbCheckpoint(chkfile, mode=FILE_CREATE) as chk:
                #Dump all the fields to a checkpoint
                for field in to_dump:
                    chk.store(field)
                chk.write_attribute("/","time",t)

            #Open the checkpointing file (main version)
            chkfile = path.join(self.dumpdir, "chkpt")
            with DumbCheckpoint(chkfile, mode=FILE_CREATE) as chk:
                #Dump all the fields to a checkpoint
                for field in to_dump:
                    chk.store(field)
                chk.write_attribute("/","time",t)

        return t

    def diagnostic_dump(self):

        with open(path.join(self.dumpdir, "diagnostics.json"), "w") as f:
            f.write(json.dumps(self.diagnostic_data, indent=4))

    def initialise(self, initial_conditions):
        """
        Initialise state variables
        """

        for x, ic in zip(self.x_init.split(), initial_conditions):
            x.assign(ic)

    @abstractmethod
    def _build_spaces(self, mesh, vertical_degree, horizontal_degree, family):

        """
        Build function spaces:
        """
        pass

    def _allocate_state(self):
        """
        Construct Functions to store the state variables.
        """

        W = self.W
        self.xn = Function(W)
        self.x_init = Function(W)
        self.xstar = Function(W)
        self.xp = Function(W)
        self.xnp1 = Function(W)
        self.xrhs = Function(W)
        self.dy = Function(W)

class CompressibleState(State):

    def __init__(self, mesh, vertical_degree=1, horizontal_degree=1,
                 family="RT", z=None, k=None, Omega=None, mu=None,
                 timestepping=None,
                 output=None,
                 parameters=None,
                 diagnostics=None,
                 fieldlist=None,
                 diagnostic_fields=[],
                 on_sphere=False):

        super(CompressibleState, self).__init__(mesh=mesh,
                                                vertical_degree=vertical_degree,
                                                horizontal_degree=horizontal_degree,
                                                family=family,
                                                z=z, k=k, Omega=Omega, mu=mu,
                                                timestepping=timestepping,
                                                output=output,
                                                parameters=parameters,
                                                diagnostics=diagnostics,
                                                fieldlist=fieldlist,
                                                diagnostic_fields=diagnostic_fields)

        # build the geopotential
        if parameters.geopotential:
            V = FunctionSpace(mesh, "CG", 1)
            if on_sphere:
                self.Phi = Function(V).interpolate(Expression("pow(x[0]*x[0]+x[1]*x[1]+x[2]*x[2],0.5)"))
            else:
                self.Phi = Function(V).interpolate(Expression("x[1]"))
            self.Phi *= parameters.g

        if self.k is None:
            # build the vertical normal
            w = TestFunction(self.Vv)
            u = TrialFunction(self.Vv)
            self.k = Function(self.Vv)
            n = FacetNormal(self.mesh)
            krhs = -div(w)*self.z*dx + inner(w,n)*self.z*ds_tb
            klhs = inner(w,u)*dx
            solve(klhs == krhs, self.k)

    def set_reference_profiles(self, rho_ref, theta_ref):
        """
        Initialise reference profiles
        :arg rho_ref: :class:`.Function` object, initial rho
        :arg theta_ref: :class:`.Function` object, initial theta
        """

        self.rhobar = Function(self.V[1])
        self.thetabar = Function(self.V[2])

        self.rhobar.project(rho_ref)
        self.thetabar.project(theta_ref)

    def _build_spaces(self, mesh, vertical_degree, horizontal_degree, family):
        """
        Build:
        velocity space self.V2,
        pressure space self.V3,
        temperature space self.Vt,
        mixed function space self.W = (V2,V3,Vt)
        """

        # horizontal base spaces
        cell = mesh._base_mesh.ufl_cell()
        if(cell.cellname() == 'triangle'):
            cell = triangle
        if(cell.cellname() == 'interval'):
            cell = interval
        S1 = FiniteElement(family, cell, horizontal_degree+1)
        S2 = FiniteElement("DG", cell, horizontal_degree)

        # vertical base spaces
        T0 = FiniteElement("CG", interval, vertical_degree+1)
        T1 = FiniteElement("DG", interval, vertical_degree)

        # build spaces V2, V3, Vt
        V2h_elt = HDiv(TensorProductElement(S1, T1))
        V2t_elt = TensorProductElement(S2, T0)
        V3_elt = TensorProductElement(S2, T1)
        V2v_elt = HDiv(V2t_elt)
        V2_elt = V2h_elt + V2v_elt

        self.V_elt = [0,0,0]
        self.V_elt[0] = V2_elt
        self.V_elt[1] = V3_elt
        self.V_elt[2] = V2t_elt

        self.V = [0,0,0]
        self.V[0] = FunctionSpace(mesh, V2_elt)
        self.V[1] = FunctionSpace(mesh, V3_elt)
        self.V[2] = FunctionSpace(mesh, V2t_elt)

        self.Vv = FunctionSpace(mesh, V2v_elt)

        self.W = MixedFunctionSpace((self.V[0], self.V[1], self.V[2]))

class ShallowWaterState(State):

    def _build_spaces(self, mesh, vertical_degree, horizontal_degree, family):

        if vertical_degree is not None:
            raise ValueError('Mesh is not extruded in the vertical for shallow water')

        cell = mesh.ufl_cell().cellname()

        V1_elt = FiniteElement(family, cell, horizontal_degree+1)

        self.V = [0,0]
        self.V[0] = FunctionSpace(mesh,V1_elt)
        self.V[1] = FunctionSpace(mesh,"DG",horizontal_degree)

        self.W = MixedFunctionSpace((self.V[0], self.V[1]))

def get_latlon_mesh(mesh):
    coords_orig = mesh.coordinates
    mesh_dg_fs = VectorFunctionSpace(mesh, "DG", 1)
    coords_dg = Function(mesh_dg_fs)
    coords_latlon = Function(mesh_dg_fs)
    par_loop("""
    for (int i=0; i<3; i++) {
    for (int j=0; j<3; j++) {
    dg[i][j] = cg[i][j];
    }
    }
    """, dx, {'dg': (coords_dg, WRITE),
              'cg': (coords_orig, READ)})

    # lat-lon 'x' = atan2(y, x)
    coords_latlon.dat.data[:,0] = np.arctan2(coords_dg.dat.data[:,1], coords_dg.dat.data[:,0])
    # lat-lon 'y' = asin(z/sqrt(x^2 + y^2 + z^2))
    coords_latlon.dat.data[:,1] = np.arcsin(coords_dg.dat.data[:,2]/np.sqrt(coords_dg.dat.data[:,0]**2 + coords_dg.dat.data[:,1]**2 + coords_dg.dat.data[:,2]**2))
    coords_latlon.dat.data[:,2] = 0.0

    kernel = op2.Kernel("""
    #define PI 3.141592653589793
    #define TWO_PI 6.283185307179586
    void splat_coords(double **coords) {
        double diff0 = (coords[0][0] - coords[1][0]);
        double diff1 = (coords[0][0] - coords[2][0]);
        double diff2 = (coords[1][0] - coords[2][0]);

        if (fabs(diff0) > PI || fabs(diff1) > PI || fabs(diff2) > PI) {
            const int sign0 = coords[0][0] < 0 ? -1 : 1;
            const int sign1 = coords[1][0] < 0 ? -1 : 1;
            const int sign2 = coords[2][0] < 0 ? -1 : 1;
            if (sign0 < 0) {
                coords[0][0] += TWO_PI;
            }
            if (sign1 < 0) {
                coords[1][0] += TWO_PI;
            }
            if (sign2 < 0) {
                coords[2][0] += TWO_PI;
            }
        }
    }""", "splat_coords")

    op2.par_loop(kernel, coords_latlon.cell_set,
                 coords_latlon.dat(op2.RW, coords_latlon.cell_node_map()))
    return Mesh(coords_latlon)
<|endoftext|>


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

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Sample conversion script for 1,000,000 VA data.

The 1,000,000 VA data is comprised of VCFs with every position
called.  This script converts those files to gVCF format.

Assumptions:
- one sample per input VCF

TODO(Cuiping and Denis):
 (1) Add logic for VCF block records such as depth and genotype quality.  For
     more detail, see https://sites.google.com/site/gvcftools/home/about-gvcf
     and http://www.broadinstitute.org/gatk/guide/article?id=4017
 (2) Use your favority python unit test framework to add some tests for the
     logic to collapse contiguous non-variant calls into a VCF block record.

This script can be run standalone:
   zcat chr21.vcf.gz | ./gvcf-mapper.py

Or via the debugger:
   python -mpdb ./gvcf-mapper.py chr1.vcf

It can also be run as a Hadoop Streaming job:
  hadoop jar /path/to/your/hadoop-streaming-*.jar \
  -libjars /home/deflaux/custom.jar \
  -outputformat com.custom.CustomMultiOutputFormat \
  -mapper /home/deflaux/gvcf-mapper.py -file /home/deflaux/gvcf-mapper.py \
  -reducer org.apache.hadoop.mapred.lib.IdentityReducer \
  -input inputpath \
  -output outputpath

See also https://developers.google.com/hadoop/

This is a modified version of the simple gvcf conversion script.  It uses the reduce step and a custom output format to consolidate data for each sample to a smaller number of files.  For more detail, see
    http://stackoverflow.com/questions/7449756/get-input-file-name-in-streaming-hadoop-program
    http://stackoverflow.com/questions/18541503/multiple-output-files-for-hadoop-streaming-with-python-mapper
"""

import os
import re
import sys
import pdb

# Constants
INPUT_FILE_KEY = "map_input_file"
SAMPLE_ID_PATTERN = r"/(LP\d{7}-DNA_\w\d{2})/"

# VCF Fields
# http://www.1000genomes.org/wiki/Analysis/Variant%20Call%20Format/vcf-variant-call-format-version-41
CHROM = 0
POS = 1
ID = 2
REF = 3
ALT = 4
QUAL = 5
FILTER = 6
INFO = 7
FORMAT = 8
GENOTYPE = 9

vcf_count = 0
filtered_count = 0
total_count = 0

# Using global variables in an effort to keep this script simple
g_start_block = None
g_end_block = None
min_qual = None
min_depth = None
min_mq = None
max_mq0 = None
ref_block = None

def main():
  """Entry point to the script."""

  global sample_id
  sample_id = None
  sample_id_re = re.compile(SAMPLE_ID_PATTERN)

  # Basic parsing of command line arguments to allow a filename
  # to be passed when running this code in the debugger.
  file_handle = sys.stdin
  if 2 <= len(sys.argv):
    path = sys.argv[1]
    file_handle = open(path, "r")
  else:
    path = os.environ[INPUT_FILE_KEY]
    print >> sys.stderr, path

  match = sample_id_re.search(path)
  sample_id = match.group(1)

  # Loop over each line of the VCF
  line = file_handle.readline()
  while line:
    total_count =+ 1
    line = line.strip()
    if not line:
      continue

    if "#" == line[0]:
      # This is a header line, skip it
      line = file_handle.readline()
      continue

    fields = line.split("\t")
    if is_variant(fields):
        if meets_filter_criteria(fields) == True:
          # This is a variant, emit the preceeding non-variant region VCF block, if
          # applicable, followed by this VCF line
          emit_block(sample_id)
          emit(sample_id, line)
        else:
          accumulate_block(fields, no_call=True)
    else:
      # Gather information about this VCF line in our non-variant region
      if meets_filter_criteria(fields) == True:
        accumulate_block(fields)
      else:
        accumulate_block(fields, no_call=True)

    line = file_handle.readline()

  # Emit the final block, if applicable
  emit_block(sample_id)

def accumulate_block(fields, no_call=False):
  """Accumulates data from each record in a non-variant region.

  Every consecutive VCF record in a non-variant region is passed to
  this function.

  You might modify it to observe the minimum or average quality score
  for records in this region.

  You might also decide within this function that these fields should
  be part of a new VCF block record instead of the current one.  If
  so, call emit_block() and then begin accumulating state for the new
  VCF block record.

  Args:
    fields: (string array) fields from one line of the VCF

  Returns: n/a

  Side Effects:
    global variables are modified
  """
  global g_start_block
  global g_end_block
  global min_qual
  global min_depth
  global min_mq
  global max_mq0
  global ref_block

  # Check to see if the current block matches what we want to add to it
  if no_call is True and ref_block is True:
    emit_block(sample_id)
  elif no_call is False and ref_block is False:
    emit_block(sample_id)

  # If this is the first line in a no call block set the genotype to ./.
  if no_call is True and ref_block is None:
    fields[GENOTYPE] = './.'
    ref_block = False
  elif no_call is False and ref_block is None:
    ref_block = True

  # Set start and end points of block as necessary
  if g_start_block is None:
    g_start_block = fields
    g_end_block = fields
  elif g_end_block is not None and int(fields[POS]) > int(g_end_block[POS]) + 1:
    emit_block(sample_id)
    g_start_block = fields
    # Emit resets ref_block, need to set again
    if no_call is True:
      ref_block = False
    else:
      ref_block = True
  else:
    g_end_block = fields

  # Check max/min values of certain metrics
  variant_info_dict = info_to_dict(fields)
  min_qual = check_values(min_qual, fields[QUAL], "QUAL")
  min_depth = check_values(min_depth, variant_info_dict['DP'], "DP")
  min_mq = check_values(min_mq, variant_info_dict['MQ'], "MQ")
  max_mq0 = check_values(max_mq0, variant_info_dict['MQ0'], "MQ0", min=False)

def check_values(existing, new, metric, min=True):
  if existing is None:
    existing = new
  else:
    type = metric_type(metric)
    existing = type(existing)
    new = type(new)
    if min is True:
      if existing > new:
        existing = new
    else:
      if existing < new:
        existing = new
  return existing

def metric_type(metric):
  metric_key = {
    "QUAL": float,
    "DP": int,
    "MQ": float,
    "MQ0": int
  }
  if not metric in metric_key:
    raise Exception("No type exists for %s" % metric)
  return metric_key[metric]

def emit_block(key):
  """Emits the current non-variant block, if applicable.

  The current implementation naively uses all values from the first
  VCF record in the block overwriting the INFO value with the end of
  the block.

  You might modify this function to utilize more accumulated state about
  the block such as minimum or average quality score.

  You might also merge the END INFO key/value pairs with any other
  existing INFO key/values.

  Args:
    key: (string) the key upon which to sort prior to the reduce step

  Returns
  : n/a

  Side Effects:
    global variables are modified
    a VCF line is written to stdout
  """
  global g_start_block
  global g_end_block
  global min_qual
  global min_depth
  global min_mq
  global max_mq0
  global ref_block

  if g_start_block is None:
    return

  block_fields = g_start_block

  if g_end_block is None:
    block_fields[INFO] = "END=" + g_start_block[POS]
  else:
    block_fields[INFO] = "END=" + g_end_block[POS]

  block_fields[INFO] += ";DP=" + str(min_depth)
  block_fields[INFO] += ";MQ=" + str(min_mq)
  block_fields[INFO] += ";MQ0=" + str(max_mq0)
  block_fields[QUAL] = str(min_qual)

  if ref_block is False:
    block_fields[FORMAT] = "GT"
    block_fields[GENOTYPE] = "./."

  line = "\t".join(block_fields)
  emit(key, line)

  # Reset our block state
  g_start_block = None
  g_end_block = None
  min_depth = None
  min_qual = None
  min_mq = None
  max_mq0 = None
  ref_block = None

def emit(key, line):
  """Emits a VCF line to stdout.

  Args:
    key: (string) the key upon which to sort prior to the reduce step
    line: (string)

  Returns: n/a

  Side Effects:
    a VCF line is written to stdout
  """
  print "%s\t%s" % (key, line)

def is_variant(fields):
  """Determines whether or not the VCF fields constitute variant.

  Args:
    fields: (string array) fields from one line of the VCF

  Returns:
    (boolean) whether or not the VCF field constitute a variant

  Side Effects:
    a VCF line is written to stdout
  """
  if "0|0" in fields[GENOTYPE] or "0/0" in fields[GENOTYPE]:
    return False
  return True

def meets_filter_criteria(fields):
  """
   filter criteria 1: (returns true if)
   if ALT field is . (Mean a reference call) then
   check that the values if the INFO field meet the following requirements:
   1) MQ0 less than 4
   2) MQ greater than or = to 30
   3) QUAL greater than or = to 30
  """
  if fields[ALT] != ".":
    if fields[FILTER] == "PASS":
      return True
    else:
      return False
  if fields[ALT] == ".":    # Means it is a reference call Now to check the values meet our threshold

    variant_info_dict = info_to_dict(fields)

    ## Processing the above dict to meet criteria.
    if float(variant_info_dict['MQ0']) < 4 and float(variant_info_dict['MQ']) >= 30 and float(fields[QUAL]) >= 30:
      vcf_count =+ 1
      return True
    else:
      filtered_count =+ 1
      return False
  else:
    vcf_count =+ 1
    return True
  return True

def info_to_dict(fields):
  ##Preprocessing: Takes the string value of INFO column (field[INFO]) and converts into dict as Key, Value based on the =
  info_string_items = [s for s in fields[INFO].split(';') if s]
  variant_info_dict = {}
  for item in info_string_items:
    try:
      key,value = item.split('=')
      variant_info_dict[key] = value
    except:
      pass  #Exception handling for keys without values, Specifically 'DB' in the INFO string
  return variant_info_dict

if __name__ == "__main__":
  main()
<commit_msg>Bug fix for passing the input file as a parameter.<commit_after>#!/usr/bin/env python

# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Sample conversion script for 1,000,000 VA data.

The 1,000,000 VA data is comprised of VCFs with every position
called.  This script converts those files to gVCF format.

Assumptions:
- one sample per input VCF

TODO(Cuiping and Denis):
 (1) Add logic for VCF block records such as depth and genotype quality.  For
     more detail, see https://sites.google.com/site/gvcftools/home/about-gvcf
     and http://www.broadinstitute.org/gatk/guide/article?id=4017
 (2) Use your favority python unit test framework to add some tests for the
     logic to collapse contiguous non-variant calls into a VCF block record.

This script can be run standalone:
   zcat chr21.vcf.gz | ./gvcf-mapper.py

Or via the debugger:
   python -mpdb ./gvcf-mapper.py chr1.vcf

It can also be run as a Hadoop Streaming job:
  hadoop jar /path/to/your/hadoop-streaming-*.jar \
  -libjars /home/deflaux/custom.jar \
  -outputformat com.custom.CustomMultiOutputFormat \
  -mapper /home/deflaux/gvcf-mapper.py -file /home/deflaux/gvcf-mapper.py \
  -reducer org.apache.hadoop.mapred.lib.IdentityReducer \
  -input inputpath \
  -output outputpath

See also https://developers.google.com/hadoop/

This is a modified version of the simple gvcf conversion script.  It uses the reduce step and a custom output format to consolidate data for each sample to a smaller number of files.  For more detail, see
    http://stackoverflow.com/questions/7449756/get-input-file-name-in-streaming-hadoop-program
    http://stackoverflow.com/questions/18541503/multiple-output-files-for-hadoop-streaming-with-python-mapper
"""

import os
import re
import sys
import pdb

# Constants
INPUT_FILE_KEY = "map_input_file"
SAMPLE_ID_PATTERN = r"/(LP\d{7}-DNA_\w\d{2})/"

# VCF Fields
# http://www.1000genomes.org/wiki/Analysis/Variant%20Call%20Format/vcf-variant-call-format-version-41
CHROM = 0
POS = 1
ID = 2
REF = 3
ALT = 4
QUAL = 5
FILTER = 6
INFO = 7
FORMAT = 8
GENOTYPE = 9

vcf_count = 0
filtered_count = 0
total_count = 0

# Using global variables in an effort to keep this script simple
g_start_block = None
g_end_block = None
min_qual = None
min_depth = None
min_mq = None
max_mq0 = None
ref_block = None

def main():
  """Entry point to the script."""

  global sample_id
  sample_id = None
  sample_id_re = re.compile(SAMPLE_ID_PATTERN)

  # Basic parsing of command line arguments to allow a filename
  # to be passed when running this code in the debugger.
  path = None
  file_handle = sys.stdin
  if 2 <= len(sys.argv):
    path = sys.argv[1]
    file_handle = open(path, "r")
  else:
    path = os.environ[INPUT_FILE_KEY]
    print >> sys.stderr, path

  if path is not None:
    match = sample_id_re.search(path)
    if match:
        sample_id = match.group(1)

  # Loop over each line of the VCF
  line = file_handle.readline()
  while line:
    total_count =+ 1
    line = line.strip()
    if not line:
      continue

    if "#" == line[0]:
      # This is a header line, skip it
      line = file_handle.readline()
      continue

    fields = line.split("\t")
    if is_variant(fields):
        if meets_filter_criteria(fields) == True:
          # This is a variant, emit the preceeding non-variant region VCF block, if
          # applicable, followed by this VCF line
          emit_block(sample_id)
          emit(sample_id, line)
        else:
          accumulate_block(fields, no_call=True)
    else:
      # Gather information about this VCF line in our non-variant region
      if meets_filter_criteria(fields) == True:
        accumulate_block(fields)
      else:
        accumulate_block(fields, no_call=True)

    line = file_handle.readline()

  # Emit the final block, if applicable
  emit_block(sample_id)

def accumulate_block(fields, no_call=False):
  """Accumulates data from each record in a non-variant region.

  Every consecutive VCF record in a non-variant region is passed to
  this function.

  You might modify it to observe the minimum or average quality score
  for records in this region.

  You might also decide within this function that these fields should
  be part of a new VCF block record instead of the current one.  If
  so, call emit_block() and then begin accumulating state for the new
  VCF block record.

  Args:
    fields: (string array) fields from one line of the VCF

  Returns: n/a

  Side Effects:
    global variables are modified
  """
  global g_start_block
  global g_end_block
  global min_qual
  global min_depth
  global min_mq
  global max_mq0
  global ref_block

  # Check to see if the current block matches what we want to add to it
  if no_call is True and ref_block is True:
    emit_block(sample_id)
  elif no_call is False and ref_block is False:
    emit_block(sample_id)

  # If this is the first line in a no call block set the genotype to ./.
  if no_call is True and ref_block is None:
    fields[GENOTYPE] = './.'
    ref_block = False
  elif no_call is False and ref_block is None:
    ref_block = True

  # Set start and end points of block as necessary
  if g_start_block is None:
    g_start_block = fields
    g_end_block = fields
  elif g_end_block is not None and int(fields[POS]) > int(g_end_block[POS]) + 1:
    emit_block(sample_id)
    g_start_block = fields
    # Emit resets ref_block, need to set again
    if no_call is True:
      ref_block = False
    else:
      ref_block = True
  else:
    g_end_block = fields

  # Check max/min values of certain metrics
  variant_info_dict = info_to_dict(fields)
  min_qual = check_values(min_qual, fields[QUAL], "QUAL")
  min_depth = check_values(min_depth, variant_info_dict['DP'], "DP")
  min_mq = check_values(min_mq, variant_info_dict['MQ'], "MQ")
  max_mq0 = check_values(max_mq0, variant_info_dict['MQ0'], "MQ0", min=False)

def check_values(existing, new, metric, min=True):
  if existing is None:
    existing = new
  else:
    type = metric_type(metric)
    existing = type(existing)
    new = type(new)
    if min is True:
      if existing > new:
        existing = new
    else:
      if existing < new:
        existing = new
  return existing

def metric_type(metric):
  metric_key = {
    "QUAL": float,
    "DP": int,
    "MQ": float,
    "MQ0": int
  }
  if not metric in metric_key:
    raise Exception("No type exists for %s" % metric)
  return metric_key[metric]

def emit_block(key):
  """Emits the current non-variant block, if applicable.

  The current implementation naively uses all values from the first
  VCF record in the block overwriting the INFO value with the end of
  the block.

  You might modify this function to utilize more accumulated state about
  the block such as minimum or average quality score.

  You might also merge the END INFO key/value pairs with any other
  existing INFO key/values.

  Args:
    key: (string) the key upon which to sort prior to the reduce step

  Returns
  : n/a

  Side Effects:
    global variables are modified
    a VCF line is written to stdout
  """
  global g_start_block
  global g_end_block
  global min_qual
  global min_depth
  global min_mq
  global max_mq0
  global ref_block

  if g_start_block is None:
    return

  block_fields = g_start_block

  if g_end_block is None:
    block_fields[INFO] = "END=" + g_start_block[POS]
  else:
    block_fields[INFO] = "END=" + g_end_block[POS]

  block_fields[INFO] += ";DP=" + str(min_depth)
  block_fields[INFO] += ";MQ=" + str(min_mq)
  block_fields[INFO] += ";MQ0=" + str(max_mq0)
  block_fields[QUAL] = str(min_qual)

  if ref_block is False:
    block_fields[FORMAT] = "GT"
    block_fields[GENOTYPE] = "./."

  line = "\t".join(block_fields)
  emit(key, line)

  # Reset our block state
  g_start_block = None
  g_end_block = None
  min_depth = None
  min_qual = None
  min_mq = None
  max_mq0 = None
  ref_block = None

def emit(key, line):
  """Emits a VCF line to stdout.

  Args:
    key: (string) the key upon which to sort prior to the reduce step
    line: (string)

  Returns: n/a

  Side Effects:
    a VCF line is written to stdout
  """
  print "%s\t%s" % (key, line)

def is_variant(fields):
  """Determines whether or not the VCF fields constitute variant.

  Args:
    fields: (string array) fields from one line of the VCF

  Returns:
    (boolean) whether or not the VCF field constitute a variant

  Side Effects:
    a VCF line is written to stdout
  """
  if "0|0" in fields[GENOTYPE] or "0/0" in fields[GENOTYPE]:
    return False
  return True

def meets_filter_criteria(fields):
  """
   filter criteria 1: (returns true if)
   if ALT field is . (Mean a reference call) then
   check that the values if the INFO field meet the following requirements:
   1) MQ0 less than 4
   2) MQ greater than or = to 30
   3) QUAL greater than or = to 30
  """
  if fields[ALT] != ".":
    if fields[FILTER] == "PASS":
      return True
    else:
      return False
  if fields[ALT] == ".":    # Means it is a reference call Now to check the values meet our threshold

    variant_info_dict = info_to_dict(fields)

    ## Processing the above dict to meet criteria.
    if float(variant_info_dict['MQ0']) < 4 and float(variant_info_dict['MQ']) >= 30 and float(fields[QUAL]) >= 30:
      vcf_count =+ 1
      return True
    else:
      filtered_count =+ 1
      return False
  else:
    vcf_count =+ 1
    return True
  return True

def info_to_dict(fields):
  ##Preprocessing: Takes the string value of INFO column (field[INFO]) and converts into dict as Key, Value based on the =
  info_string_items = [s for s in fields[INFO].split(';') if s]
  variant_info_dict = {}
  for item in info_string_items:
    try:
      key,value = item.split('=')
      variant_info_dict[key] = value
    except:
      pass  #Exception handling for keys without values, Specifically 'DB' in the INFO string
  return variant_info_dict

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


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

# Copyright 2014, Rackspace US, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Example usages
# python swift-recon.py replication --ring-type account
# python swift-recon.py replication --ring-type container
# python swift-recon.py replication --ring-type object
# python swift-recon.py async-pendings
# python swift-recon.py md5
# python swift-recon.py quarantine

import argparse
import maas_common
import re
import subprocess

class ParseError(maas_common.MaaSException):
    pass

class CommandNotRecognized(maas_common.MaaSException):
    pass

def recon_output(for_ring, options=None):
    """Run swift-recon and filter out extraneous printed lines.

    ::

        >>> recon_output('account', '-r')
        ['[2014-11-21 00:25:16] Checking on replication',
         '[replication_failure] low: 0, high: 0, avg: 0.0, total: 0, Failed: 0.0%, no_result: 0, reported: 2',
         '[replication_success] low: 2, high: 4, avg: 3.0, total: 6, Failed: 0.0%, no_result: 0, reported: 2',
         '[replication_time] low: 0, high: 0, avg: 0.0, total: 0, Failed: 0.0%, no_result: 0, reported: 2',
         '[replication_attempted] low: 1, high: 2, avg: 1.5, total: 3, Failed: 0.0%, no_result: 0, reported: 2',
         'Oldest completion was 2014-11-21 00:24:51 (25 seconds ago) by 192.168.31.1:6002.',
         'Most recent completion was 2014-11-21 00:24:56 (20 seconds ago) by 192.168.31.2:6002.']

    :param str for_ring: Which ring to run swift-recon on
    :param list options: Command line options with which to run swift-recon
    :returns: Strings from output that are most important
    :rtype: list
    """
    command = ['swift-recon', for_ring]
    command.extend(options or [])
    out = subprocess.check_output(command)
    return filter(lambda s: s and not s.startswith(('==', '-')),
                  out.split('\n'))

def stat_regexp_generator(data):
    """Generate a regeular expression that will swift-recon stats.

    Lines printed by swift-recon look like::

        [data] low: 0, high: 0, avg: 0.0, total: 0, Failed: 0.0%, no_result: 0, reported: 0

    Where data above is the value of the ``data`` parameter passed to the
    function.
    """
    expression = """\s+low:\s+(?P<low>\d+),     # parse out the low result
                    \s+high:\s+(?P<high>\d+),   # parse out the high result
                    \s+avg:\s+(?P<avg>\d+.\d+), # you get the idea now
                    \s+total:\s+(?P<total>\d+),
                    \s+Failed:\s+(?P<failed>\d+.\d+%),
                    \s+no_result:\s+(?P<no_result>\d+),
                    \s+reported:\s+(?P<reported>\d+)"""
    return re.compile('\[' + data + '\]' + expression, re.VERBOSE)

def _parse_into_dict(line, parsed_by):
    match = parsed_by.match(line)
    if match:
        return match.groupdict()
    else:
        raise ParseError("Cannot parse '{0}' for statistics.".format(line))

def recon_stats_dicts(for_ring, options, starting_with, parsed_by):
    """Return a list of dictionaries of parsed statistics.

    Swift-recon has a standard format for it's statistics:

    [name] low: 0, high: 0, avg: 0.0, total: 0, Failed: 0.0%, no_result: 0, reported: 0

    Using the regular expression passed by the user in ``parsed_by``, we parse
    this out and return dictionaries of lines that start with
    ``starting_with``. This is after we get the recon output for the ring and
    options.

    :param str for_ring: Which ring to run swift-recon on
    :param list options: Command line options with which to run swift-recon
    :param str starting_with: String with which to filter lines
    :param parsed_by: Compiled regular expression to match and parse with
    :type parsed_by: _sre.SRE_Pattern (the result of calling re.compile)
    :returns: list of dictionaries of parse data
    """
    return map(lambda l: _parse_into_dict(l, parsed_by),
               filter(lambda s: s.startswith(starting_with),
                      recon_output(for_ring, options)))

def swift_replication(for_ring):
    """Parse swift-recon's replication statistics and return them.

    ::

        >>> swift_replication('account')
        {'attempted': {'avg': '1.5',
                       'failed': '0.0%',
                       'high': '2',
                       'low': '1',
                       'no_result': '0',
                       'reported": '5',
                       'total': '3'},
         'failure': {'avg': '0.0',
                     'failed': '0.0%',
                     'high': '0',
                     'low': '0',
                     'no_result': '0',
                     'reported": '5',
                     'total': '0'},
         'success': {'avg': '3.0',
                     'failed': '0.0%',
                     'high': '4',
                     'low': '2',
                     'no_result': '0',
                     'reported": '5',
                     'total': '6'},
         'time': {'avg': '0.0',
                  'failed': '0.0%',
                  'high': '0',
                  'low': '0',
                  'no_result': '0',
                  'reported": '5',
                  'total': '0'}}

    :param str for_ring: Which ring to run swift-recon on
    :returns: Dictionary of attempted, failed, success, and time statistics
    :rtype: dict
    """
    regexp = stat_regexp_generator(r'replication_(?P<replication_type>\w+)')
    replication_dicts = recon_stats_dicts(for_ring, ['-r'], '[replication_',
                                          regexp)

    # reduce could work here but would require an enclosed function which is
    # less readable than this loop
    replication_statistics = {}
    for rep_dict in replication_dicts:
        replication_statistics[rep_dict.pop('replication_type')] = rep_dict

    return replication_statistics

def swift_async():
    """Parse swift-recon's async pendings statistics and return them.

    ::

        >>> swift_async()
        {'avg': '0.0',
         'failed': '0.0%',
         'high': '0',
         'low': '0',
         'no_result': '0',
         'reported': '2',
         'total': '0'}

    :returns: Dictionary of average, failed, high, low, no_result, reported,
        and total statistics.
    """
    regexp = stat_regexp_generator('async_pending')
    async_dicts = recon_stats_dicts('object', ['-a'], '[async_pending]',
                                    regexp)
    stats = {}
    for async_dict in async_dicts:
        if async_dict:
            stats = async_dict
            # Break will skip the for-loop's else block
            break
    else:
        # If we didn't find a non-empty dict, error out
        maas_common.status_err(
            'No data could be collected about pending async operations'
            )
    return {'async': stats}

def swift_quarantine():
    """Parse swift-recon's quarantined objects and return them.

    ::

        >>> swift_quarantine()
        {'accounts': {'avg': '0.0',
                      'failed': '0.0%',
                      'high': '0',
                      'low': '0',
                      'no_result': '0',
                      'reported': '2',
                      'total': '0'},
         'containers': {'avg': '0.0',
                        'failed': '0.0%',
                        'high': '0',
                        'low': '0',
                        'no_result': '0',
                        'reported': '2',
                        'total': '0'},
         'objects': {'avg': '0.0',
                     'failed': '0.0%',
                     'high': '0',
                     'low': '0',
                     'no_result': '0',
                     'reported': '2',
                     'total': '0'}}

    :returns: Dictionary of objects, accounts, and containers.
    """
    regexp = stat_regexp_generator('quarantined_(?P<ring>\w+)')
    quarantined_dicts = recon_stats_dicts('-q', [], '[quarantined_',
                                          regexp)

    quarantined_statistics = {}
    for quar_dict in quarantined_dicts:
        quarantined_statistics[quar_dict.pop('ring')] = quar_dict

    return quarantined_statistics

def swift_md5():
    """Parse swift-recon's md5 check output and return it.

    ::

        >>> swift_md5()
        {'ring': {'errors': '0', 'success': '2', 'total': '2'},
         'swiftconf': {'errors': '0', 'success': '2', 'total': '2'}}

    :returns: Dictioanry
    """
    check_re = re.compile('Checking\s+(?P<check>[^\s]+)\s+md5sums?')
    error_re = re.compile('https?://(?P<address>[^:]+):\d+')
    result_re = re.compile(
        '(?P<success>\d+)/(?P<total>\d+)[^\d]+(?P<errors>\d+).*'
        )
    output = recon_output('--md5')  # We need to pass --md5 as a string here
    md5_statistics = {}
    checking_dict = {}
    for line in output:
        check_match = check_re.search(line)
        if check_match and not checking_dict:
            checking_dict = check_match.groupdict()
            # First line of a grouping we care about, might as well skip all
            # other checks in the loop
            continue
        # If there was an error checking the md5sum, error out immediately
        if line.startswith('!!'):
            error_dict = error_re.search(line).groupdict()
            maas_common.status_err('md5 mismatch for {0} on host {1}'.format(
                checking_dict.get('check'), error_dict['address']
                ))
        results_match = result_re.match(line)
        if results_match:
            check_name = checking_dict['check'].replace('.', '_')
            md5_statistics[check_name] = results_match.groupdict()
            checking_dict = {}

    return md5_statistics

def print_nested_stats(statistics):
    """Print out nested statistics.

    Nested statistics would be retrieved from ``swift_quarantine`` or
    ``swift_replication``.

    The following will show what you'll see

    ::

        >>> a = {'accounts': {'avg': '0.0',
                              'failed': '0.0%',
                              'high': '0',
                              'low': '0',
                              'no_result': '0',
                              'reported': '2',
                              'total': '0'}}
        >>> print_nested_stats(a)
        metric accounts_avg double 0.0
        metric accounts_failed double 0.0
        metric accounts_high uint64 0
        metric accounts_low uint64 0
        metric accounts_no_result uint64 0
        metric accounts_total uint64 0

    """
    for ring, stats in statistics.items():
        print_stats(ring, stats)

metrics_per_stat = {
    'avg': lambda name, val: maas_common.metric(name, 'double', val),
    'failed': lambda name, val: maas_common.metric(name, 'double', val[:-1])
    }

DEFAULT_METRIC = lambda name, val: maas_common.metric(name, 'uint64', val)

def print_stats(prefix, statistics):
    """Print out statistics.

    """
    for name, value in statistics.items():
        metric = metrics_per_stat.get(name, DEFAULT_METRIC)
        metric('{0}_{1}'.format(prefix, name), value)

def make_parser():
    parser = argparse.ArgumentParser(
        description='Process and print swift-recon statistics'
        )
    parser.add_argument('recon',
                        help='Which statistics to collect. Acceptable recon: '
                             '"async-pendings", "md5", "quarantine", '
                             '"replication"')
    parser.add_argument('--ring-type', dest='ring',
                        help='Which ring to run statistics for. Only used by '
                             'replication recon.')
    return parser

def get_stats_from(args):
    stats = {}
    if args.recon == 'async-pendings':
        stats = swift_async()
    elif args.recon == 'md5':
        stats = swift_md5()
    elif args.recon == 'quarantine':
        stats = swift_quarantine()
    elif args.recon == 'replication':
        if args.ring not in {"account", "container", "object"}:
            maas_common.status_err('no ring provided to check')
        stats = swift_replication(args.ring)
    else:
        raise CommandNotRecognized('unrecognized command "{0}"'.format(
            args.recon))
    return stats

def main():
    parser = make_parser()
    args = parser.parse_args()

    try:
        stats = get_stats_from(args)
    except (ParseError, CommandNotRecognized) as e:
        maas_common.status_error(str(e))

    if stats:
        maas_common.status_ok()
        print_nested_stats(stats)

if __name__ == '__main__':
    with maas_common.print_output():
        main()
<commit_msg>Fix swift-recon.py<commit_after>#!/usr/bin/env python

# Copyright 2014, Rackspace US, Inc.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

# Example usages
# python swift-recon.py replication --ring-type account
# python swift-recon.py replication --ring-type container
# python swift-recon.py replication --ring-type object
# python swift-recon.py async-pendings
# python swift-recon.py md5
# python swift-recon.py quarantine

import argparse
import maas_common
import re
import subprocess

class ParseError(maas_common.MaaSException):
    pass

class CommandNotRecognized(maas_common.MaaSException):
    pass

def recon_output(for_ring, options=None):
    """Run swift-recon and filter out extraneous printed lines.

    ::

        >>> recon_output('account', '-r')
        ['[2014-11-21 00:25:16] Checking on replication',
         '[replication_failure] low: 0, high: 0, avg: 0.0, total: 0, Failed: 0.0%, no_result: 0, reported: 2',
         '[replication_success] low: 2, high: 4, avg: 3.0, total: 6, Failed: 0.0%, no_result: 0, reported: 2',
         '[replication_time] low: 0, high: 0, avg: 0.0, total: 0, Failed: 0.0%, no_result: 0, reported: 2',
         '[replication_attempted] low: 1, high: 2, avg: 1.5, total: 3, Failed: 0.0%, no_result: 0, reported: 2',
         'Oldest completion was 2014-11-21 00:24:51 (25 seconds ago) by 192.168.31.1:6002.',
         'Most recent completion was 2014-11-21 00:24:56 (20 seconds ago) by 192.168.31.2:6002.']

    :param str for_ring: Which ring to run swift-recon on
    :param list options: Command line options with which to run swift-recon
    :returns: Strings from output that are most important
    :rtype: list
    """
    command = ['swift-recon', for_ring]
    command.extend(options or [])
    out = subprocess.check_output(command)
    return filter(lambda s: s and not s.startswith(('==', '-')),
                  out.split('\n'))

def stat_regexp_generator(data):
    """Generate a regeular expression that will swift-recon stats.

    Lines printed by swift-recon look like::

        [data] low: 0, high: 0, avg: 0.0, total: 0, Failed: 0.0%, no_result: 0, reported: 0

    Where data above is the value of the ``data`` parameter passed to the
    function.
    """
    expression = """\s+low:\s+(?P<low>\d+),     # parse out the low result
                    \s+high:\s+(?P<high>\d+),   # parse out the high result
                    \s+avg:\s+(?P<avg>\d+.\d+), # you get the idea now
                    \s+total:\s+(?P<total>\d+),
                    \s+Failed:\s+(?P<failed>\d+.\d+%),
                    \s+no_result:\s+(?P<no_result>\d+),
                    \s+reported:\s+(?P<reported>\d+)"""
    return re.compile('\[' + data + '\]' + expression, re.VERBOSE)

def _parse_into_dict(line, parsed_by):
    match = parsed_by.match(line)
    if match:
        return match.groupdict()
    else:
        raise ParseError("Cannot parse '{0}' for statistics.".format(line))

def recon_stats_dicts(for_ring, options, starting_with, parsed_by):
    """Return a list of dictionaries of parsed statistics.

    Swift-recon has a standard format for it's statistics:

    [name] low: 0, high: 0, avg: 0.0, total: 0, Failed: 0.0%, no_result: 0, reported: 0

    Using the regular expression passed by the user in ``parsed_by``, we parse
    this out and return dictionaries of lines that start with
    ``starting_with``. This is after we get the recon output for the ring and
    options.

    :param str for_ring: Which ring to run swift-recon on
    :param list options: Command line options with which to run swift-recon
    :param str starting_with: String with which to filter lines
    :param parsed_by: Compiled regular expression to match and parse with
    :type parsed_by: _sre.SRE_Pattern (the result of calling re.compile)
    :returns: list of dictionaries of parse data
    """
    return map(lambda l: _parse_into_dict(l, parsed_by),
               filter(lambda s: s.startswith(starting_with),
                      recon_output(for_ring, options)))

def swift_replication(for_ring):
    """Parse swift-recon's replication statistics and return them.

    ::

        >>> swift_replication('account')
        {'attempted': {'avg': '1.5',
                       'failed': '0.0%',
                       'high': '2',
                       'low': '1',
                       'no_result': '0',
                       'reported": '5',
                       'total': '3'},
         'failure': {'avg': '0.0',
                     'failed': '0.0%',
                     'high': '0',
                     'low': '0',
                     'no_result': '0',
                     'reported": '5',
                     'total': '0'},
         'success': {'avg': '3.0',
                     'failed': '0.0%',
                     'high': '4',
                     'low': '2',
                     'no_result': '0',
                     'reported": '5',
                     'total': '6'},
         'time': {'avg': '0.0',
                  'failed': '0.0%',
                  'high': '0',
                  'low': '0',
                  'no_result': '0',
                  'reported": '5',
                  'total': '0'}}

    :param str for_ring: Which ring to run swift-recon on
    :returns: Dictionary of attempted, failed, success, and time statistics
    :rtype: dict
    """
    regexp = stat_regexp_generator(r'replication_(?P<replication_type>\w+)')
    replication_dicts = recon_stats_dicts(for_ring, ['-r'], '[replication_',
                                          regexp)

    # reduce could work here but would require an enclosed function which is
    # less readable than this loop
    replication_statistics = {}
    for rep_dict in replication_dicts:
        replication_statistics[rep_dict.pop('replication_type')] = rep_dict

    return replication_statistics

def swift_async():
    """Parse swift-recon's async pendings statistics and return them.

    ::

        >>> swift_async()
        {'avg': '0.0',
         'failed': '0.0%',
         'high': '0',
         'low': '0',
         'no_result': '0',
         'reported': '2',
         'total': '0'}

    :returns: Dictionary of average, failed, high, low, no_result, reported,
        and total statistics.
    """
    regexp = stat_regexp_generator('async_pending')
    async_dicts = recon_stats_dicts('object', ['-a'], '[async_pending]',
                                    regexp)
    stats = {}
    for async_dict in async_dicts:
        if async_dict:
            stats = async_dict
            # Break will skip the for-loop's else block
            break
    else:
        # If we didn't find a non-empty dict, error out
        maas_common.status_err(
            'No data could be collected about pending async operations'
            )
    return {'async': stats}

def swift_quarantine():
    """Parse swift-recon's quarantined objects and return them.

    ::

        >>> swift_quarantine()
        {'accounts': {'avg': '0.0',
                      'failed': '0.0%',
                      'high': '0',
                      'low': '0',
                      'no_result': '0',
                      'reported': '2',
                      'total': '0'},
         'containers': {'avg': '0.0',
                        'failed': '0.0%',
                        'high': '0',
                        'low': '0',
                        'no_result': '0',
                        'reported': '2',
                        'total': '0'},
         'objects': {'avg': '0.0',
                     'failed': '0.0%',
                     'high': '0',
                     'low': '0',
                     'no_result': '0',
                     'reported': '2',
                     'total': '0'}}

    :returns: Dictionary of objects, accounts, and containers.
    """
    regexp = stat_regexp_generator('quarantined_(?P<ring>\w+)')
    quarantined_dicts = recon_stats_dicts('-q', [], '[quarantined_',
                                          regexp)

    quarantined_statistics = {}
    for quar_dict in quarantined_dicts:
        quarantined_statistics[quar_dict.pop('ring')] = quar_dict

    return quarantined_statistics

def swift_md5():
    """Parse swift-recon's md5 check output and return it.

    ::

        >>> swift_md5()
        {'ring': {'errors': '0', 'success': '2', 'total': '2'},
         'swiftconf': {'errors': '0', 'success': '2', 'total': '2'}}

    :returns: Dictioanry
    """
    check_re = re.compile('Checking\s+(?P<check>[^\s]+)\s+md5sums?')
    error_re = re.compile('https?://(?P<address>[^:]+):\d+')
    result_re = re.compile(
        '(?P<success>\d+)/(?P<total>\d+)[^\d]+(?P<errors>\d+).*'
        )
    output = recon_output('--md5')  # We need to pass --md5 as a string here
    md5_statistics = {}
    checking_dict = {}
    for line in output:
        check_match = check_re.search(line)
        if check_match and not checking_dict:
            checking_dict = check_match.groupdict()
            # First line of a grouping we care about, might as well skip all
            # other checks in the loop
            continue
        # If there was an error checking the md5sum, error out immediately
        if line.startswith('!!'):
            error_dict = error_re.search(line).groupdict()
            maas_common.status_err('md5 mismatch for {0} on host {1}'.format(
                checking_dict.get('check'), error_dict['address']
                ))
        results_match = result_re.match(line)
        if results_match:
            check_name = checking_dict['check'].replace('.', '_')
            md5_statistics[check_name] = results_match.groupdict()
            checking_dict = {}

    return md5_statistics

def print_nested_stats(statistics):
    """Print out nested statistics.

    Nested statistics would be retrieved from ``swift_quarantine`` or
    ``swift_replication``.

    The following will show what you'll see

    ::

        >>> a = {'accounts': {'avg': '0.0',
                              'failed': '0.0%',
                              'high': '0',
                              'low': '0',
                              'no_result': '0',
                              'reported': '2',
                              'total': '0'}}
        >>> print_nested_stats(a)
        metric accounts_avg double 0.0
        metric accounts_failed double 0.0
        metric accounts_high uint64 0
        metric accounts_low uint64 0
        metric accounts_no_result uint64 0
        metric accounts_total uint64 0

    """
    for ring, stats in statistics.items():
        print_stats(ring, stats)

metrics_per_stat = {
    'avg': lambda name, val: maas_common.metric(name, 'double', val),
    'failed': lambda name, val: maas_common.metric(name, 'double', val[:-1])
    }

DEFAULT_METRIC = lambda name, val: maas_common.metric(name, 'uint64', val)

def print_stats(prefix, statistics):
    """Print out statistics.

    """
    for name, value in statistics.items():
        metric = metrics_per_stat.get(name, DEFAULT_METRIC)
        metric('{0}_{1}'.format(prefix, name), value)

def make_parser():
    parser = argparse.ArgumentParser(
        description='Process and print swift-recon statistics'
        )
    parser.add_argument('recon',
                        help='Which statistics to collect. Acceptable recon: '
                             '"async-pendings", "md5", "quarantine", '
                             '"replication"')
    parser.add_argument('--ring-type', dest='ring',
                        help='Which ring to run statistics for. Only used by '
                             'replication recon.')
    return parser

def get_stats_from(args):
    stats = {}
    if args.recon == 'async-pendings':
        stats = swift_async()
    elif args.recon == 'md5':
        stats = swift_md5()
    elif args.recon == 'quarantine':
        stats = swift_quarantine()
    elif args.recon == 'replication':
        if args.ring not in {"account", "container", "object"}:
            maas_common.status_err('no ring provided to check')
        stats = swift_replication(args.ring)
    else:
        raise CommandNotRecognized('unrecognized command "{0}"'.format(
            args.recon))
    return stats

def main():
    parser = make_parser()
    args = parser.parse_args()

    try:
        stats = get_stats_from(args)
    except (ParseError, CommandNotRecognized) as e:
        maas_common.status_err(str(e))

    if stats:
        maas_common.status_ok()
        print_nested_stats(stats)

if __name__ == '__main__':
    with maas_common.print_output():
        main()
<|endoftext|>