<commit_before>from scipy import pi as pi
import scipy as _sp

def generic_conductance(target, transport_type, pore_diffusivity,
                        throat_diffusivity, pore_area, throat_area,
                        conduit_lengths, conduit_shape_factors, **kwargs):
    r"""
    Calculate the generic conductance (could be mass, thermal, electrical,
    or hydraylic) of conduits in the network, where a conduit is
    ( 1/2 pore - full throat - 1/2 pore ).

    Parameters
    ----------
    target : OpenPNM Object
        The object which this model is associated with. This controls the
        length of the calculated array, and also provides access to other
        necessary properties.

    pore_diffusivity : string
        Dictionary key of the pore diffusivity values.

    throat_diffusivity : string
        Dictionary key of the throat diffusivity values.

    pore_area : string
        Dictionary key of the pore area values.

    throat_area : string
        Dictionary key of the throat area values.

    shape_factor : string
        Dictionary key of the conduit shape factor values.

    Returns
    -------
    g : ndarray
        Array containing conductance values for conduits in the geometry
        attached to the given physics object.

    Notes
    -----
    (1) This function requires that all the necessary phase properties already
    be calculated.

    (2) This function calculates the specified property for the *entire*
    network then extracts the values for the appropriate throats at the end.

    (3) This function assumes cylindrical throats with constant cross-section
    area. Corrections for different shapes and variable cross-section area can
    be imposed by passing the proper shape factor.

    (4) shape_factor depends on the physics of the problem, i.e. diffusion-like
    processes and fluid flow need different shape factors.

    """
    network = target.project.network
    throats = network.map_throats(throats=target.Ts, origin=target)
    phase = target.project.find_phase(target)
    cn = network['throat.conns'][throats]
    # Getting equivalent areas
    A1 = network[pore_area][cn[:, 0]]
    At = network[throat_area][throats]
    A2 = network[pore_area][cn[:, 1]]
    # Getting conduit lengths
    L1 = network[conduit_lengths + '.pore1'][throats]
    Lt = network[conduit_lengths + '.throat'][throats]
    L2 = network[conduit_lengths + '.pore2'][throats]
    # Preallocating g
    g1, g2, gt = _sp.zeros((3, len(Lt)))
    # Setting g to inf when Li = 0 (ex. boundary pores)
    # INFO: This is needed since area could also be zero, which confuses NumPy
    m1, m2, mt = [Li != 0 for Li in [L1, L2, Lt]]
    g1[~m1] = g2[~m2] = gt[~mt] = _sp.inf
    # If ionic conductance, compute conductance and return it
    if transport_type == 'ionic':
        g1[m1] = (A1)[m1] * L1[m1]
        g2[m2] = (A2)[m2] * L2[m2]
        gt[mt] = (At)[mt] * Lt[mt]
        return 1/((g1+g2+gt)/(L1+L2+Lt))
    # Getting shape factors
    try:
        SF1 = phase[conduit_shape_factors+'.pore1'][throats]
        SFt = phase[conduit_shape_factors+'.throat'][throats]
        SF2 = phase[conduit_shape_factors+'.pore2'][throats]
    except KeyError:
        SF1 = SF2 = SFt = 1.0
    # Interpolate pore phase property values to throats
    try:
        Dt = phase[throat_diffusivity][throats]
    except KeyError:
        Dt = phase.interpolate_data(propname=pore_diffusivity)[throats]
    try:
        D1 = phase[pore_diffusivity][cn[:, 0]]
        D2 = phase[pore_diffusivity][cn[:, 1]]
    except KeyError:
        D1 = phase.interpolate_data(propname=throat_diffusivity)[cn[:, 0]]
        D2 = phase.interpolate_data(propname=throat_diffusivity)[cn[:, 1]]
    # Find g for half of pore 1, throat, and half of pore 2
    if transport_type == 'flow':
        g1[m1] = A1[m1]**2 / (8*pi*D1*L1)[m1]
        g2[m2] = A2[m2]**2 / (8*pi*D2*L2)[m2]
        gt[mt] = At[mt]**2 / (8*pi*Dt*Lt)[mt]
    elif transport_type == 'flow_power_law':
        for k, v in kwargs.items():
            if k == 'pore_consistency':
                pore_consistency = v
            elif k == 'throat_consistency':
                throat_consistency = v
            elif k == 'pore_flow_index':
                pore_flow_index = v
            elif k == 'throat_flow_index':
                throat_flow_index = v
            elif k == 'pore_pressure':
                pore_pressure = v

        # Check if pressure field exists
        try:
            phase[pore_pressure]
        except KeyError:
            phase[pore_pressure] = 0
        P = phase[pore_pressure]

        # Interpolate pore phase property values to throats
        try:
            Ct = phase[throat_consistency][throats]
        except KeyError:
            Ct = phase.interpolate_data(propname=pore_consistency)[throats]
        try:
            nt = phase[throat_flow_index][throats]
        except KeyError:
            nt = phase.interpolate_data(propname=pore_flow_index)[throats]
        # Interpolate throat phase property values to pores
        try:
            C1 = phase[pore_consistency][cn[:, 0]]
            C2 = phase[pore_consistency][cn[:, 1]]
        except KeyError:
            C1 = phase.interpolate_data(propname=throat_consistency)[cn[:, 0]]
            C2 = phase.interpolate_data(propname=throat_consistency)[cn[:, 1]]
        try:
            n1 = phase[pore_flow_index][cn[:, 0]]
            n2 = phase[pore_flow_index][cn[:, 1]]
        except KeyError:
            n1 = phase.interpolate_data(propname=throat_flow_index)[cn[:, 0]]
            n2 = phase.interpolate_data(propname=throat_flow_index)[cn[:, 1]]
        # Interpolate pore pressure values to throats
        Pt = phase.interpolate_data(propname=pore_pressure)[throats]

        # Pressure differences dP
        dP1 = _sp.absolute(P[cn[:, 0]]-Pt)
        dP2 = _sp.absolute(P[cn[:, 1]]-Pt)
        dPt = _sp.absolute(_sp.diff(P[cn], axis=1).squeeze())

        # Apparent viscosities
        mu1 = (dP1**(1-1/n1)[m1] * C1**(1/n1)[m1] / ((4*n1/(3*n1+1)) *
               (2*L1/((A1/pi)**0.5))**(1-1/n1))[m1])

        mu2 = (dP2**(1-1/n2)[m2] * C2**(1/n2)[m2] / ((4*n2/(3*n2+1)) *
               (2*L2/((A2/pi)**0.5))**(1-1/n2))[m2])

        mut = (dPt**(1-1/nt)[mt] * Ct**(1/nt)[mt] / ((4*nt/(3*nt+1)) *
               (2*Lt/((At/pi)**0.5))**(1-1/nt))[mt])

        # Bound the apparent viscosity
        vis_min = 1e-08
        vis_max = 1e+04
        mu1[mu1 < vis_min] = vis_min
        mu1[mu1 > vis_max] = vis_max
        mu2[mu2 < vis_min] = vis_min
        mu2[mu2 > vis_max] = vis_max
        mut[mut < vis_min] = vis_min
        mut[mut > vis_max] = vis_max

        phase['throat.viscosity_eff'] = mut

        g1[m1] = A1[m1]**2 / ((8*pi*L1)[m1]*mu1)
        g2[m2] = A2[m2]**2 / ((8*pi*L2)[m2]*mu2)
        gt[mt] = At[mt]**2 / ((8*pi*Lt)[mt]*mut)

    elif transport_type == 'diffusion':
        g1[m1] = (D1*A1)[m1] / L1[m1]
        g2[m2] = (D2*A2)[m2] / L2[m2]
        gt[mt] = (Dt*At)[mt] / Lt[mt]
    elif transport_type == 'taylor_aris_diffusion':
        for k, v in kwargs.items():
            if k == 'pore_pressure':
                pore_pressure = v
            elif k == 'throat_hydraulic_conductance':
                throat_hydraulic_conductance = v
        P = phase[pore_pressure]
        gh = phase[throat_hydraulic_conductance]
        Qt = -gh*_sp.diff(P[cn], axis=1).squeeze()

        u1 = Qt[m1]/A1[m1]
        u2 = Qt[m2]/A2[m2]
        ut = Qt[mt]/At[mt]

        Pe1 = u1 * ((4*A1[m1]/_sp.pi)**0.5) / D1[m1]
        Pe2 = u2 * ((4*A2[m2]/_sp.pi)**0.5) / D2[m2]
        Pet = ut * ((4*At[mt]/_sp.pi)**0.5) / Dt[mt]

        g1[m1] = D1[m1]*(1+(Pe1**2)/192)*A1[m1] / L1[m1]
        g2[m2] = D2[m2]*(1+(Pe2**2)/192)*A2[m2] / L2[m2]
        gt[mt] = Dt[mt]*(1+(Pet**2)/192)*At[mt] / Lt[mt]
    elif transport_type == 'dispersion':
        for k, v in kwargs.items():
            if k == 'pore_pressure':
                pore_pressure = v
            elif k == 'throat_hydraulic_conductance':
                throat_hydraulic_conductance = v
            elif k == 'throat_diffusive_conductance':
                throat_diffusive_conductance = v
            elif k == 's_scheme':
                s_scheme = v

        P = phase[pore_pressure]
        gh = phase[throat_hydraulic_conductance]
        gd = phase[throat_diffusive_conductance]
        gd = _sp.tile(gd, 2)

        Qij = -gh*_sp.diff(P[cn], axis=1).squeeze()
        Qij = _sp.append(Qij, -Qij)

        Peij = Qij/gd
        Peij[(Peij < 1e-10) & (Peij >= 0)] = 1e-10
        Peij[(Peij > -1e-10) & (Peij <= 0)] = -1e-10
        Qij = Peij*gd

        if s_scheme == 'upwind':
            w = gd + _sp.maximum(0, -Qij)
        elif s_scheme == 'hybrid':
            w = _sp.maximum(0, _sp.maximum(-Qij, gd-Qij/2))
        elif s_scheme == 'powerlaw':
            w = gd * _sp.maximum(0, (1 - 0.1*_sp.absolute(Peij))**5) + \
                _sp.maximum(0, -Qij)
        elif s_scheme == 'exponential':
            w = -Qij / (1 - _sp.exp(Peij))
        else:
            raise Exception('Unrecognized discretization scheme: ' + s_scheme)
        w = _sp.reshape(w, (network.Nt, 2), order='F')
        return w

    elif transport_type == 'ad_dif_mig':
        for k, v in kwargs.items():
            if k == 'pore_pressure':
                pore_pressure = v
            if k == 'pore_potential':
                pore_potential = v
            elif k == 'throat_hydraulic_conductance':
                throat_hydraulic_conductance = v
            elif k == 'throat_diffusive_conductance':
                throat_diffusive_conductance = v
            elif k == 'throat_valence':
                throat_valence = v
            elif k == 'throat_temperature':
                throat_temperature = v
            elif k == 's_scheme':
                s_scheme = v

        P = phase[pore_pressure]
        V = phase[pore_potential]
        gh = phase[throat_hydraulic_conductance]
        gd = phase[throat_diffusive_conductance]
        gd = _sp.tile(gd, 2)
        z = phase[throat_valence]
        T = phase[throat_temperature]
        D = Dt
        F = 96485.3329
        R = 8.3145

        S = (A1*L1+A2*L2+At*Lt)/(L1+L2+Lt)
        L = L1 + Lt + L2

        # Advection
        Qij = -gh*_sp.diff(P[cn], axis=1).squeeze()
        Qij = _sp.append(Qij, -Qij)

        # Migration
        grad_V = _sp.diff(V[cn], axis=1).squeeze() / L
        mig = ((z*F*D*S)/(R*T)) * grad_V
        mig = _sp.append(mig, -mig)

        # Advection-migration
        adv_mig = Qij-mig

        # Peclet number (includes advection and migration)
        Peij_adv_mig = adv_mig/gd
        Peij_adv_mig[(Peij_adv_mig < 1e-10) & (Peij_adv_mig >= 0)] = 1e-10
        Peij_adv_mig[(Peij_adv_mig > -1e-10) & (Peij_adv_mig <= 0)] = -1e-10

        # Corrected advection-migration
        adv_mig = Peij_adv_mig*gd

        if s_scheme == 'upwind':
            w = gd + _sp.maximum(0, -adv_mig)
        elif s_scheme == 'hybrid':
            w = _sp.maximum(0, _sp.maximum(-adv_mig, gd-adv_mig/2))
        elif s_scheme == 'powerlaw':
            w = (gd * _sp.maximum(0, (1 - 0.1*_sp.absolute(Peij_adv_mig))**5) +
                 _sp.maximum(0, -adv_mig))
        elif s_scheme == 'exponential':
            w = -adv_mig / (1 - _sp.exp(Peij_adv_mig))
        else:
            raise Exception('Unrecognized discretization scheme: ' + s_scheme)
        w = _sp.reshape(w, (network.Nt, 2), order='F')
        return w

    else:
        raise Exception('Unknown keyword for "transport_type", can only be' +
                        ' "flow", "diffusion", "taylor_aris_diffusion",' +
                        ' "ad_dif_mig" or "ionic"')
    # Apply shape factors and calculate the final conductance
    return (1/gt/SFt + 1/g1/SF1 + 1/g2/SF2)**(-1)
<commit_msg>fixing merge error<commit_after>from scipy import pi as pi
import scipy as _sp

def generic_conductance(target, transport_type, pore_diffusivity,
                        throat_diffusivity, pore_area, throat_area,
                        conduit_lengths, conduit_shape_factors, **kwargs):
    r"""
    Calculate the generic conductance (could be mass, thermal, electrical,
    or hydraylic) of conduits in the network, where a conduit is
    ( 1/2 pore - full throat - 1/2 pore ).

    Parameters
    ----------
    target : OpenPNM Object
        The object which this model is associated with. This controls the
        length of the calculated array, and also provides access to other
        necessary properties.

    pore_diffusivity : string
        Dictionary key of the pore diffusivity values.

    throat_diffusivity : string
        Dictionary key of the throat diffusivity values.

    pore_area : string
        Dictionary key of the pore area values.

    throat_area : string
        Dictionary key of the throat area values.

    shape_factor : string
        Dictionary key of the conduit shape factor values.

    Returns
    -------
    g : ndarray
        Array containing conductance values for conduits in the geometry
        attached to the given physics object.

    Notes
    -----
    (1) This function requires that all the necessary phase properties already
    be calculated.

    (2) This function calculates the specified property for the *entire*
    network then extracts the values for the appropriate throats at the end.

    (3) This function assumes cylindrical throats with constant cross-section
    area. Corrections for different shapes and variable cross-section area can
    be imposed by passing the proper shape factor.

    (4) shape_factor depends on the physics of the problem, i.e. diffusion-like
    processes and fluid flow need different shape factors.

    """
    network = target.project.network
    throats = network.map_throats(throats=target.Ts, origin=target)
    phase = target.project.find_phase(target)
    cn = network['throat.conns'][throats]
    # Getting equivalent areas
    A1 = network[pore_area][cn[:, 0]]
    At = network[throat_area][throats]
    A2 = network[pore_area][cn[:, 1]]
    # Getting conduit lengths
    L1 = network[conduit_lengths + '.pore1'][throats]
    Lt = network[conduit_lengths + '.throat'][throats]
    L2 = network[conduit_lengths + '.pore2'][throats]
    # Preallocating g
    g1, g2, gt = _sp.zeros((3, len(Lt)))
    # Setting g to inf when Li = 0 (ex. boundary pores)
    # INFO: This is needed since area could also be zero, which confuses NumPy
    m1, m2, mt = [Li != 0 for Li in [L1, L2, Lt]]
    g1[~m1] = g2[~m2] = gt[~mt] = _sp.inf
    # If ionic conductance, compute conductance and return it
    if transport_type == 'ionic':
        g1[m1] = (A1)[m1] * L1[m1]
        g2[m2] = (A2)[m2] * L2[m2]
        gt[mt] = (At)[mt] * Lt[mt]
        return 1/((g1+g2+gt)/(L1+L2+Lt))
    # Getting shape factors
    try:
        SF1 = phase[conduit_shape_factors+'.pore1'][throats]
        SFt = phase[conduit_shape_factors+'.throat'][throats]
        SF2 = phase[conduit_shape_factors+'.pore2'][throats]
    except KeyError:
        SF1 = SF2 = SFt = 1.0
    # Interpolate pore phase property values to throats
    try:
        Dt = phase[throat_diffusivity][throats]
    except KeyError:
        Dt = phase.interpolate_data(propname=pore_diffusivity)[throats]
    try:
        D1 = phase[pore_diffusivity][cn[:, 0]]
        D2 = phase[pore_diffusivity][cn[:, 1]]
    except KeyError:
        D1 = phase.interpolate_data(propname=throat_diffusivity)[cn[:, 0]]
        D2 = phase.interpolate_data(propname=throat_diffusivity)[cn[:, 1]]
    # Find g for half of pore 1, throat, and half of pore 2
    if transport_type == 'flow':
        g1[m1] = A1[m1]**2 / (8*pi*D1*L1)[m1]
        g2[m2] = A2[m2]**2 / (8*pi*D2*L2)[m2]
        gt[mt] = At[mt]**2 / (8*pi*Dt*Lt)[mt]
    elif transport_type == 'flow_power_law':
        for k, v in kwargs.items():
            if k == 'pore_consistency':
                pore_consistency = v
            elif k == 'throat_consistency':
                throat_consistency = v
            elif k == 'pore_flow_index':
                pore_flow_index = v
            elif k == 'throat_flow_index':
                throat_flow_index = v
            elif k == 'pore_pressure':
                pore_pressure = v

        # Check if pressure field exists
        try:
            phase[pore_pressure]
        except KeyError:
            phase[pore_pressure] = 0
        P = phase[pore_pressure]

        # Interpolate pore phase property values to throats
        try:
            Ct = phase[throat_consistency][throats]
        except KeyError:
            Ct = phase.interpolate_data(propname=pore_consistency)[throats]
        try:
            nt = phase[throat_flow_index][throats]
        except KeyError:
            nt = phase.interpolate_data(propname=pore_flow_index)[throats]
        # Interpolate throat phase property values to pores
        try:
            C1 = phase[pore_consistency][cn[:, 0]]
            C2 = phase[pore_consistency][cn[:, 1]]
        except KeyError:
            C1 = phase.interpolate_data(propname=throat_consistency)[cn[:, 0]]
            C2 = phase.interpolate_data(propname=throat_consistency)[cn[:, 1]]
        try:
            n1 = phase[pore_flow_index][cn[:, 0]]
            n2 = phase[pore_flow_index][cn[:, 1]]
        except KeyError:
            n1 = phase.interpolate_data(propname=throat_flow_index)[cn[:, 0]]
            n2 = phase.interpolate_data(propname=throat_flow_index)[cn[:, 1]]
        # Interpolate pore pressure values to throats
        Pt = phase.interpolate_data(propname=pore_pressure)[throats]

        # Pressure differences dP
        dP1 = _sp.absolute(P[cn[:, 0]]-Pt)
        dP2 = _sp.absolute(P[cn[:, 1]]-Pt)
        dPt = _sp.absolute(_sp.diff(P[cn], axis=1).squeeze())

        # Apparent viscosities
        mu1 = (dP1**(1-1/n1)[m1] * C1**(1/n1)[m1] / ((4*n1/(3*n1+1)) *
               (2*L1/((A1/pi)**0.5))**(1-1/n1))[m1])

        mu2 = (dP2**(1-1/n2)[m2] * C2**(1/n2)[m2] / ((4*n2/(3*n2+1)) *
               (2*L2/((A2/pi)**0.5))**(1-1/n2))[m2])

        mut = (dPt**(1-1/nt)[mt] * Ct**(1/nt)[mt] / ((4*nt/(3*nt+1)) *
               (2*Lt/((At/pi)**0.5))**(1-1/nt))[mt])

        # Bound the apparent viscosity
        vis_min = 1e-08
        vis_max = 1e+04
        mu1[mu1 < vis_min] = vis_min
        mu1[mu1 > vis_max] = vis_max
        mu2[mu2 < vis_min] = vis_min
        mu2[mu2 > vis_max] = vis_max
        mut[mut < vis_min] = vis_min
        mut[mut > vis_max] = vis_max

        phase['throat.viscosity_eff'] = mut

        g1[m1] = A1[m1]**2 / ((8*pi*L1)[m1]*mu1)
        g2[m2] = A2[m2]**2 / ((8*pi*L2)[m2]*mu2)
        gt[mt] = At[mt]**2 / ((8*pi*Lt)[mt]*mut)

    elif transport_type == 'diffusion':
        g1[m1] = (D1*A1)[m1] / L1[m1]
        g2[m2] = (D2*A2)[m2] / L2[m2]
        gt[mt] = (Dt*At)[mt] / Lt[mt]
    elif transport_type == 'taylor_aris_diffusion':
        for k, v in kwargs.items():
            if k == 'pore_pressure':
                pore_pressure = v
            elif k == 'throat_hydraulic_conductance':
                throat_hydraulic_conductance = v
        P = phase[pore_pressure]
        gh = phase[throat_hydraulic_conductance]
        Qt = -gh*_sp.diff(P[cn], axis=1).squeeze()

        u1 = Qt[m1]/A1[m1]
        u2 = Qt[m2]/A2[m2]
        ut = Qt[mt]/At[mt]

        Pe1 = u1 * ((4*A1[m1]/_sp.pi)**0.5) / D1[m1]
        Pe2 = u2 * ((4*A2[m2]/_sp.pi)**0.5) / D2[m2]
        Pet = ut * ((4*At[mt]/_sp.pi)**0.5) / Dt[mt]

        g1[m1] = D1[m1]*(1+(Pe1**2)/192)*A1[m1] / L1[m1]
        g2[m2] = D2[m2]*(1+(Pe2**2)/192)*A2[m2] / L2[m2]
        gt[mt] = Dt[mt]*(1+(Pet**2)/192)*At[mt] / Lt[mt]
    elif transport_type == 'dispersion':
        for k, v in kwargs.items():
            if k == 'pore_pressure':
                pore_pressure = v
            elif k == 'throat_hydraulic_conductance':
                throat_hydraulic_conductance = v
            elif k == 'throat_diffusive_conductance':
                throat_diffusive_conductance = v
            elif k == 's_scheme':
                s_scheme = v

        P = phase[pore_pressure]
        gh = phase[throat_hydraulic_conductance]
        gd = phase[throat_diffusive_conductance]
        gd = _sp.tile(gd, 2)

        Qij = -gh*_sp.diff(P[cn], axis=1).squeeze()
        Qij = _sp.append(Qij, -Qij)

        Peij = Qij/gd
        Peij[(Peij < 1e-10) & (Peij >= 0)] = 1e-10
        Peij[(Peij > -1e-10) & (Peij <= 0)] = -1e-10
        Qij = Peij*gd

        if s_scheme == 'upwind':
            w = gd + _sp.maximum(0, -Qij)
        elif s_scheme == 'hybrid':
            w = _sp.maximum(0, _sp.maximum(-Qij, gd-Qij/2))
        elif s_scheme == 'powerlaw':
            w = gd * _sp.maximum(0, (1 - 0.1*_sp.absolute(Peij))**5) + \
                _sp.maximum(0, -Qij)
        elif s_scheme == 'exponential':
            w = -Qij / (1 - _sp.exp(Peij))
        else:
            raise Exception('Unrecognized discretization scheme: ' + s_scheme)
        w = _sp.reshape(w, (network.Nt, 2), order='F')
        return w
    elif transport_type == 'ad_dif_mig':
        for k, v in kwargs.items():
            if k == 'pore_pressure':
                pore_pressure = v
            if k == 'pore_potential':
                pore_potential = v
            elif k == 'throat_hydraulic_conductance':
                throat_hydraulic_conductance = v
            elif k == 'throat_diffusive_conductance':
                throat_diffusive_conductance = v
            elif k == 'throat_valence':
                throat_valence = v
            elif k == 'throat_temperature':
                throat_temperature = v
            elif k == 's_scheme':
                s_scheme = v

        P = phase[pore_pressure]
        V = phase[pore_potential]
        gh = phase[throat_hydraulic_conductance]
        gd = phase[throat_diffusive_conductance]
        gd = _sp.tile(gd, 2)
        z = phase[throat_valence]
        T = phase[throat_temperature]
        D = Dt
        F = 96485.3329
        R = 8.3145

        S = (A1*L1+A2*L2+At*Lt)/(L1+L2+Lt)
        L = L1 + Lt + L2

        # Advection
        Qij = -gh*_sp.diff(P[cn], axis=1).squeeze()
        Qij = _sp.append(Qij, -Qij)

        # Migration
        grad_V = _sp.diff(V[cn], axis=1).squeeze() / L
        mig = ((z*F*D*S)/(R*T)) * grad_V
        mig = _sp.append(mig, -mig)

        # Advection-migration
        adv_mig = Qij-mig

        # Peclet number (includes advection and migration)
        Peij_adv_mig = adv_mig/gd
        Peij_adv_mig[(Peij_adv_mig < 1e-10) & (Peij_adv_mig >= 0)] = 1e-10
        Peij_adv_mig[(Peij_adv_mig > -1e-10) & (Peij_adv_mig <= 0)] = -1e-10

        # Corrected advection-migration
        adv_mig = Peij_adv_mig*gd

        if s_scheme == 'upwind':
            w = gd + _sp.maximum(0, -adv_mig)
        elif s_scheme == 'hybrid':
            w = _sp.maximum(0, _sp.maximum(-adv_mig, gd-adv_mig/2))
        elif s_scheme == 'powerlaw':
            w = (gd * _sp.maximum(0, (1 - 0.1*_sp.absolute(Peij_adv_mig))**5) +
                 _sp.maximum(0, -adv_mig))
        elif s_scheme == 'exponential':
            w = -adv_mig / (1 - _sp.exp(Peij_adv_mig))
        else:
            raise Exception('Unrecognized discretization scheme: ' + s_scheme)
        w = _sp.reshape(w, (network.Nt, 2), order='F')
        return w
    else:
        raise Exception('Unknown keyword for "transport_type", can only be' +
                        ' "flow", "diffusion", "taylor_aris_diffusion",' +
                        ' "dispersion", "ad_dif_mig" or "ionic"')
    # Apply shape factors and calculate the final conductance
    return (1/gt/SFt + 1/g1/SF1 + 1/g2/SF2)**(-1)
<|endoftext|>


<commit_before>import os
import datetime
import importlib

import astropy.units as u
import astropy.coordinates as coords
from astropy.time import Time
from astropy.coordinates import SkyCoord, EarthLocation, AltAz

from . import mount as mount
from . import camera as camera
from . import scheduler as scheduler

from .utils.config import load_config
from .utils.logger import has_logger
from .utils import error as error

@has_logger
class Observatory(object):

    """
    Main Observatory class
    """

    def __init__(self, config=None):
        """
        Starts up the observatory. Reads config file (TODO), sets up location,
        dates, mount, cameras, and weather station
        """
        assert config is not None, self.logger.warning("Config not set for observatory")

        self.logger.info('Initializing observatory')

        self.config = config

        self.location = None
        self.horizon = 30 * u.degree
        self.utc_offset = 0 * u.minutes
        self.elevation = 0 * u.meters
        self.air_pressure = 0 * u.millibarye

        self.mount = None
        self.scheduler = None
        self.cameras = list()

       # Setup information about site location
        self.logger.info('\t Setting up observatory details')
        self._setup_observatory()

        # Create default mount and cameras. Should be read in by config file
        self.logger.info('\t Setting up mount')
        self._create_mount()

        self.logger.info('\t Setting up cameras')
        self._create_cameras()

        self.logger.info('\t Setting up scheduler')
        self._create_scheduler()

    def get_target(self):
        """ Gets the next target from the scheduler """

        target = self.scheduler.get_target(self)

        return target

    def heartbeat(self):
        """
        Touch a file each time signaling life
        """
        self.logger.debug('Touching heartbeat file')
        with open(self.heartbeat_filename, 'w') as fileobject:
            fileobject.write(str(datetime.datetime.now()) + "\n")

    def horizon(self, alt, az):
        '''Function to evaluate whether a particular alt, az is
        above the horizon

        NOTE: This could be done better with pyephem
        '''
        assert isinstance(alt, u.Quantity)
        assert isinstance(az, u.Quantity)

        horizon = float(self.config.get('site.horizon', 0)) * u.deg

        if alt > horizon:
            return True
        else:
            return False

    def is_dark(self):
        """
        Need to calculate day/night for site.

        NOTE: This could be done better with pyephem
        """
        self.logger.debug('Calculating is_dark.')

        self.site.date = ephem.now()
        self.sun.compute(self.site)

        dark_horizon = float(self.config.get('twilight_horizon', -12))

        self.is_dark = self.sun.alt < dark_horizon

        return self.is_dark

    def _setup_observatory(self, start_date=Time.now()):
        """
        Sets up the site and location details, for the observatory.

        These items are read from the 'site' config directive and include:
        * lat (latitude)
        * lon (longitude)
        * elevation
        # * horizon

        """
        self.logger.info('Setting up site details of observatory')
        earth_location = None

        if 'site' in self.config:
            config_site = self.config.get('site')

            lat = config_site.get('lat') * u.degree
            lon = config_site.get('lon') * u.degree

            self.elevation = config_site.get('elevation', 0) * u.meter
            self.horizon = config_site.get('horizon', 0) * u.degree
            self.air_pressure = config_site.get('pressure', 680) * u.millibarye

            # Create an astropy EarthLocation
            earth_location = coords.EarthLocation(
                lat=lat,
                lon=lon,
                height=self.elevation,
            )
        else:
            raise error.Error(msg='Bad site information')

        self.location = earth_location

    def _create_mount(self, mount_info=None):
        """Creates a mount object.

        Details for the creation of the mount object are held in the
        configuration file or can be passed to the method.

        This method ensures that the proper mount type is loaded.

        Note:
            This does not actually make a serial connection to the mount. To do so,
            call the 'mount.connect()' explicitly.

        Args:
            mount_info (dict):  Configuration items for the mount.

        Returns:
            panoptes.mount:     Returns a sub-class of the mount type
        """
        if mount_info is None:
            mount_info = self.config.get('mount')

        model = mount_info.get('model')

        self.logger.info('Creating mount: {}'.format(model))

        mount = None

        # Actually import the model of mount
        try:
            module = importlib.import_module('.{}'.format(model), package='panoptes.mount')
        except ImportError as err:
            self.logger.warning('ImportError. Check that the mount module exists and that all dependencies are installed')
            raise error.NotFound(model)

        # Make the mount include site information
        mount = module.Mount(config=mount_info, site=self.site)

        return mount

    def _create_cameras(self, camera_info=None):
        """Creates a camera object(s)

        Creates a camera for each camera item listed in the config. Ensures the
        appropriate camera module is loaded.

        Note:
            This does not actually make a usb connection to the camera. To do so,
            call the 'camear.connect()' explicitly.

        Args:
            camera_info (dict): Configuration items for the cameras.

        Returns:
            list: A list of created camera objects.
        """
        camera_info = self.config.get('cameras')

        for camera in camera_info:
            # Actually import the model of camera
            try:
                module = importlib.import_module('.{}'.format(camera.get('model')), 'panoptes.camera')
                cameras.append(module.Camera(config=camera))

            except ImportError as err:
                raise error.NotFound(msg=model)

        self.cameras = cameras

    def _create_scheduler(self):
        """ Sets up the scheduler that will be used by the observatory """

        # Read the targets from the file
        targets_path = os.path.join(
            self.config.get('base_dir'),
            self.config.get('targets', 'default_targets.yaml')
        )

        self.logger.info('\t Scheduler file: {}'.format(targets_path))
        self.scheduler = scheduler.Scheduler(target_list_file=targets_path)
<commit_msg>Add the gmt_offset correctly<commit_after>import os
import datetime
import importlib

import astropy.units as u
import astropy.coordinates as coords
from astropy.time import Time
from astropy.coordinates import SkyCoord, EarthLocation, AltAz

from . import mount as mount
from . import camera as camera
from . import scheduler as scheduler

from .utils.config import load_config
from .utils.logger import has_logger
from .utils import error as error

@has_logger
class Observatory(object):

    """
    Main Observatory class
    """

    def __init__(self, config=None):
        """
        Starts up the observatory. Reads config file (TODO), sets up location,
        dates, mount, cameras, and weather station
        """
        assert config is not None, self.logger.warning("Config not set for observatory")

        self.logger.info('Initializing observatory')

        self.config = config

        self.location = None
        self.horizon = 30 * u.degree
        self.gmt_offset = 0 * u.minutes
        self.elevation = 0 * u.meters
        self.air_pressure = 0 * u.millibarye

        self.mount = None
        self.scheduler = None
        self.cameras = list()

       # Setup information about site location
        self.logger.info('\t Setting up observatory details')
        self._setup_observatory()

        # Create default mount and cameras. Should be read in by config file
        self.logger.info('\t Setting up mount')
        self._create_mount()

        self.logger.info('\t Setting up cameras')
        self._create_cameras()

        self.logger.info('\t Setting up scheduler')
        self._create_scheduler()

    def get_target(self):
        """ Gets the next target from the scheduler """

        target = self.scheduler.get_target(self)

        return target

    def heartbeat(self):
        """
        Touch a file each time signaling life
        """
        self.logger.debug('Touching heartbeat file')
        with open(self.heartbeat_filename, 'w') as fileobject:
            fileobject.write(str(datetime.datetime.now()) + "\n")

    def horizon(self, alt, az):
        '''Function to evaluate whether a particular alt, az is
        above the horizon

        NOTE: This could be done better with pyephem
        '''
        assert isinstance(alt, u.Quantity)
        assert isinstance(az, u.Quantity)

        horizon = float(self.config.get('site.horizon', 0)) * u.deg

        if alt > horizon:
            return True
        else:
            return False

    def is_dark(self):
        """
        Need to calculate day/night for site.

        NOTE: This could be done better with pyephem
        """
        self.logger.debug('Calculating is_dark.')

        self.site.date = ephem.now()
        self.sun.compute(self.site)

        dark_horizon = float(self.config.get('twilight_horizon', -12))

        self.is_dark = self.sun.alt < dark_horizon

        return self.is_dark

    def _setup_observatory(self, start_date=Time.now()):
        """
        Sets up the site and location details, for the observatory.

        These items are read from the 'site' config directive and include:
        * lat (latitude)
        * lon (longitude)
        * elevation
        # * horizon

        """
        self.logger.info('Setting up site details of observatory')
        earth_location = None

        if 'site' in self.config:
            config_site = self.config.get('site')

            self.gmt_offset = config_site.get('gmt_offset') * u.minutes

            lat = config_site.get('lat') * u.degree
            lon = config_site.get('lon') * u.degree

            self.elevation = config_site.get('elevation', 0) * u.meter
            self.horizon = config_site.get('horizon', 0) * u.degree
            self.air_pressure = config_site.get('pressure', 680) * u.millibarye

            # Create an astropy EarthLocation
            earth_location = coords.EarthLocation(
                lat=lat,
                lon=lon,
                height=self.elevation,
            )
        else:
            raise error.Error(msg='Bad site information')

        self.location = earth_location

    def _create_mount(self, mount_info=None):
        """Creates a mount object.

        Details for the creation of the mount object are held in the
        configuration file or can be passed to the method.

        This method ensures that the proper mount type is loaded.

        Note:
            This does not actually make a serial connection to the mount. To do so,
            call the 'mount.connect()' explicitly.

        Args:
            mount_info (dict):  Configuration items for the mount.

        Returns:
            panoptes.mount:     Returns a sub-class of the mount type
        """
        if mount_info is None:
            mount_info = self.config.get('mount')

        model = mount_info.get('model')

        self.logger.info('Creating mount: {}'.format(model))

        mount = None

        # Actually import the model of mount
        try:
            module = importlib.import_module('.{}'.format(model), package='panoptes.mount')
        except ImportError as err:
            self.logger.warning('ImportError. Check that the mount module exists and that all dependencies are installed')
            raise error.NotFound(model)

        # Make the mount include site information
        mount = module.Mount(config=mount_info, site=self.site)

        return mount

    def _create_cameras(self, camera_info=None):
        """Creates a camera object(s)

        Creates a camera for each camera item listed in the config. Ensures the
        appropriate camera module is loaded.

        Note:
            This does not actually make a usb connection to the camera. To do so,
            call the 'camear.connect()' explicitly.

        Args:
            camera_info (dict): Configuration items for the cameras.

        Returns:
            list: A list of created camera objects.
        """
        camera_info = self.config.get('cameras')

        for camera in camera_info:
            # Actually import the model of camera
            try:
                module = importlib.import_module('.{}'.format(camera.get('model')), 'panoptes.camera')
                cameras.append(module.Camera(config=camera))

            except ImportError as err:
                raise error.NotFound(msg=model)

        self.cameras = cameras

    def _create_scheduler(self):
        """ Sets up the scheduler that will be used by the observatory """

        # Read the targets from the file
        targets_path = os.path.join(
            self.config.get('base_dir'),
            self.config.get('targets', 'default_targets.yaml')
        )

        self.logger.info('\t Scheduler file: {}'.format(targets_path))
        self.scheduler = scheduler.Scheduler(target_list_file=targets_path)
<|endoftext|>


<commit_before>from collections import deque

import numpy as np
import pickle
from mujoco_py import MujocoException

from baselines.her.util import convert_episode_to_batch_major, store_args

class RolloutWorker:

    @store_args
    def __init__(self, make_env, policy, dims, logger, T, rollout_batch_size=1,
                 exploit=False, use_target_net=False, compute_Q=False, noise_eps=0,
                 random_eps=0, history_len=100, render=False, **kwargs):
        """Rollout worker generates experience by interacting with one or many environments.

        Args:
            make_env (function): a factory function that creates a new instance of the environment
                when called
            policy (object): the policy that is used to act
            dims (dict of ints): the dimensions for observations (o), goals (g), and actions (u)
            logger (object): the logger that is used by the rollout worker
            rollout_batch_size (int): the number of parallel rollouts that should be used
            exploit (boolean): whether or not to exploit, i.e. to act optimally according to the
                current policy without any exploration
            use_target_net (boolean): whether or not to use the target net for rollouts
            compute_Q (boolean): whether or not to compute the Q values alongside the actions
            noise_eps (float): scale of the additive Gaussian noise
            random_eps (float): probability of selecting a completely random action
            history_len (int): length of history for statistics smoothing
            render (boolean): whether or not to render the rollouts
        """
        self.envs = [make_env() for _ in range(rollout_batch_size)]
        assert self.T > 0

        self.info_keys = [key.replace('info_', '') for key in dims.keys() if key.startswith('info_')]

        self.success_history = deque(maxlen=history_len)
        self.Q_history = deque(maxlen=history_len)

        self.n_episodes = 0
        self.g = np.empty((self.rollout_batch_size, self.dims['g']), np.float32)  # goals
        self.initial_o = np.empty((self.rollout_batch_size, self.dims['o']), np.float32)  # observations
        self.initial_ag = np.empty((self.rollout_batch_size, self.dims['g']), np.float32)  # achieved goals
        self.reset_all_rollouts()
        self.clear_history()

    def reset_rollout(self, i):
        """Resets the `i`-th rollout environment, re-samples a new goal, and updates the `initial_o`
        and `g` arrays accordingly.
        """
        obs = self.envs[i].reset()
        self.initial_o[i] = obs['observation']
        self.initial_ag[i] = obs['achieved_goal']
        self.g[i] = obs['desired_goal']

    def reset_all_rollouts(self):
        """Resets all `rollout_batch_size` rollout workers.
        """
        for i in range(self.rollout_batch_size):
            self.reset_rollout(i)

    def generate_rollouts(self):
        """Performs `rollout_batch_size` rollouts in parallel for time horizon `T` with the current
        policy acting on it accordingly.
        """
        self.reset_all_rollouts()

        # compute observations
        o = np.empty((self.rollout_batch_size, self.dims['o']), np.float32)  # observations
        ag = np.empty((self.rollout_batch_size, self.dims['g']), np.float32)  # achieved goals
        o[:] = self.initial_o
        ag[:] = self.initial_ag

        # generate episodes
        obs, achieved_goals, acts, goals, successes = [], [], [], [], []
        info_values = [np.empty((self.T, self.rollout_batch_size, self.dims['info_' + key]), np.float32) for key in self.info_keys]
        Qs = []
        for t in range(self.T):
            policy_output = self.policy.get_actions(
                o, ag, self.g,
                compute_Q=self.compute_Q,
                noise_eps=self.noise_eps if not self.exploit else 0.,
                random_eps=self.random_eps if not self.exploit else 0.,
                use_target_net=self.use_target_net)

            if self.compute_Q:
                u, Q = policy_output
                Qs.append(Q)
            else:
                u = policy_output

            if u.ndim == 1:
                # The non-batched case should still have a reasonable shape.
                u = u.reshape(1, -1)

            o_new = np.empty((self.rollout_batch_size, self.dims['o']))
            ag_new = np.empty((self.rollout_batch_size, self.dims['g']))
            success = np.zeros(self.rollout_batch_size)
            # compute new states and observations
            for i in range(self.rollout_batch_size):
                try:
                    # We fully ignore the reward here because it will have to be re-computed
                    # for HER.
                    curr_o_new, _, _, info = self.envs[i].step(u[i])
                    if 'is_success' in info:
                        success[i] = info['is_success']
                    o_new[i] = curr_o_new['observation']
                    ag_new[i] = curr_o_new['achieved_goal']
                    for idx, key in enumerate(self.info_keys):
                        info_values[idx][t, i] = info[key]
                    if self.render:
                        self.envs[i].render()
                except MujocoException as e:
                    return self.generate_rollouts()

            if np.isnan(o_new).any():
                self.logger.warning('NaN caught during rollout generation. Trying again...')
                self.reset_all_rollouts()
                return self.generate_rollouts()

            obs.append(o.copy())
            achieved_goals.append(ag.copy())
            successes.append(success.copy())
            acts.append(u.copy())
            goals.append(self.g.copy())
            o[...] = o_new
            ag[...] = ag_new
        obs.append(o.copy())
        achieved_goals.append(ag.copy())
        self.initial_o[:] = o

        episode = dict(o=obs,
                       u=acts,
                       g=goals,
                       ag=achieved_goals)
        for key, value in zip(self.info_keys, info_values):
            episode['info_{}'.format(key)] = value

        # stats
        successful = np.array(successes)[-1, :]
        assert successful.shape == (self.rollout_batch_size,)
        success_rate = np.mean(successful)
        self.success_history.append(success_rate)
        if self.compute_Q:
            self.Q_history.append(np.mean(Qs))
        self.n_episodes += self.rollout_batch_size

        return convert_episode_to_batch_major(episode)

    def clear_history(self):
        """Clears all histories that are used for statistics
        """
        self.success_history.clear()
        self.Q_history.clear()

    def current_success_rate(self):
        return np.mean(self.success_history)

    def current_mean_Q(self):
        return np.mean(self.Q_history)

    def save_policy(self, path):
        """Pickles the current policy for later inspection.
        """
        with open(path, 'wb') as f:
            pickle.dump(self.policy, f)

    def logs(self, prefix='worker'):
        """Generates a dictionary that contains all collected statistics.
        """
        logs = []
        logs += [('success_rate', np.mean(self.success_history))]
        if self.compute_Q:
            logs += [('mean_Q', np.mean(self.Q_history))]
        logs += [('episode', self.n_episodes)]

        if prefix is not '' and not prefix.endswith('/'):
            return [(prefix + '/' + key, val) for key, val in logs]
        else:
            return logs

    def seed(self, seed):
        """Seeds each environment with a distinct seed derived from the passed in global seed.
        """
        for idx, env in enumerate(self.envs):
            env.seed(seed + 1000 * idx)
<commit_msg>fixed warning (#464)<commit_after>from collections import deque

import numpy as np
import pickle
from mujoco_py import MujocoException

from baselines.her.util import convert_episode_to_batch_major, store_args

class RolloutWorker:

    @store_args
    def __init__(self, make_env, policy, dims, logger, T, rollout_batch_size=1,
                 exploit=False, use_target_net=False, compute_Q=False, noise_eps=0,
                 random_eps=0, history_len=100, render=False, **kwargs):
        """Rollout worker generates experience by interacting with one or many environments.

        Args:
            make_env (function): a factory function that creates a new instance of the environment
                when called
            policy (object): the policy that is used to act
            dims (dict of ints): the dimensions for observations (o), goals (g), and actions (u)
            logger (object): the logger that is used by the rollout worker
            rollout_batch_size (int): the number of parallel rollouts that should be used
            exploit (boolean): whether or not to exploit, i.e. to act optimally according to the
                current policy without any exploration
            use_target_net (boolean): whether or not to use the target net for rollouts
            compute_Q (boolean): whether or not to compute the Q values alongside the actions
            noise_eps (float): scale of the additive Gaussian noise
            random_eps (float): probability of selecting a completely random action
            history_len (int): length of history for statistics smoothing
            render (boolean): whether or not to render the rollouts
        """
        self.envs = [make_env() for _ in range(rollout_batch_size)]
        assert self.T > 0

        self.info_keys = [key.replace('info_', '') for key in dims.keys() if key.startswith('info_')]

        self.success_history = deque(maxlen=history_len)
        self.Q_history = deque(maxlen=history_len)

        self.n_episodes = 0
        self.g = np.empty((self.rollout_batch_size, self.dims['g']), np.float32)  # goals
        self.initial_o = np.empty((self.rollout_batch_size, self.dims['o']), np.float32)  # observations
        self.initial_ag = np.empty((self.rollout_batch_size, self.dims['g']), np.float32)  # achieved goals
        self.reset_all_rollouts()
        self.clear_history()

    def reset_rollout(self, i):
        """Resets the `i`-th rollout environment, re-samples a new goal, and updates the `initial_o`
        and `g` arrays accordingly.
        """
        obs = self.envs[i].reset()
        self.initial_o[i] = obs['observation']
        self.initial_ag[i] = obs['achieved_goal']
        self.g[i] = obs['desired_goal']

    def reset_all_rollouts(self):
        """Resets all `rollout_batch_size` rollout workers.
        """
        for i in range(self.rollout_batch_size):
            self.reset_rollout(i)

    def generate_rollouts(self):
        """Performs `rollout_batch_size` rollouts in parallel for time horizon `T` with the current
        policy acting on it accordingly.
        """
        self.reset_all_rollouts()

        # compute observations
        o = np.empty((self.rollout_batch_size, self.dims['o']), np.float32)  # observations
        ag = np.empty((self.rollout_batch_size, self.dims['g']), np.float32)  # achieved goals
        o[:] = self.initial_o
        ag[:] = self.initial_ag

        # generate episodes
        obs, achieved_goals, acts, goals, successes = [], [], [], [], []
        info_values = [np.empty((self.T, self.rollout_batch_size, self.dims['info_' + key]), np.float32) for key in self.info_keys]
        Qs = []
        for t in range(self.T):
            policy_output = self.policy.get_actions(
                o, ag, self.g,
                compute_Q=self.compute_Q,
                noise_eps=self.noise_eps if not self.exploit else 0.,
                random_eps=self.random_eps if not self.exploit else 0.,
                use_target_net=self.use_target_net)

            if self.compute_Q:
                u, Q = policy_output
                Qs.append(Q)
            else:
                u = policy_output

            if u.ndim == 1:
                # The non-batched case should still have a reasonable shape.
                u = u.reshape(1, -1)

            o_new = np.empty((self.rollout_batch_size, self.dims['o']))
            ag_new = np.empty((self.rollout_batch_size, self.dims['g']))
            success = np.zeros(self.rollout_batch_size)
            # compute new states and observations
            for i in range(self.rollout_batch_size):
                try:
                    # We fully ignore the reward here because it will have to be re-computed
                    # for HER.
                    curr_o_new, _, _, info = self.envs[i].step(u[i])
                    if 'is_success' in info:
                        success[i] = info['is_success']
                    o_new[i] = curr_o_new['observation']
                    ag_new[i] = curr_o_new['achieved_goal']
                    for idx, key in enumerate(self.info_keys):
                        info_values[idx][t, i] = info[key]
                    if self.render:
                        self.envs[i].render()
                except MujocoException as e:
                    return self.generate_rollouts()

            if np.isnan(o_new).any():
                self.logger.warn('NaN caught during rollout generation. Trying again...')
                self.reset_all_rollouts()
                return self.generate_rollouts()

            obs.append(o.copy())
            achieved_goals.append(ag.copy())
            successes.append(success.copy())
            acts.append(u.copy())
            goals.append(self.g.copy())
            o[...] = o_new
            ag[...] = ag_new
        obs.append(o.copy())
        achieved_goals.append(ag.copy())
        self.initial_o[:] = o

        episode = dict(o=obs,
                       u=acts,
                       g=goals,
                       ag=achieved_goals)
        for key, value in zip(self.info_keys, info_values):
            episode['info_{}'.format(key)] = value

        # stats
        successful = np.array(successes)[-1, :]
        assert successful.shape == (self.rollout_batch_size,)
        success_rate = np.mean(successful)
        self.success_history.append(success_rate)
        if self.compute_Q:
            self.Q_history.append(np.mean(Qs))
        self.n_episodes += self.rollout_batch_size

        return convert_episode_to_batch_major(episode)

    def clear_history(self):
        """Clears all histories that are used for statistics
        """
        self.success_history.clear()
        self.Q_history.clear()

    def current_success_rate(self):
        return np.mean(self.success_history)

    def current_mean_Q(self):
        return np.mean(self.Q_history)

    def save_policy(self, path):
        """Pickles the current policy for later inspection.
        """
        with open(path, 'wb') as f:
            pickle.dump(self.policy, f)

    def logs(self, prefix='worker'):
        """Generates a dictionary that contains all collected statistics.
        """
        logs = []
        logs += [('success_rate', np.mean(self.success_history))]
        if self.compute_Q:
            logs += [('mean_Q', np.mean(self.Q_history))]
        logs += [('episode', self.n_episodes)]

        if prefix is not '' and not prefix.endswith('/'):
            return [(prefix + '/' + key, val) for key, val in logs]
        else:
            return logs

    def seed(self, seed):
        """Seeds each environment with a distinct seed derived from the passed in global seed.
        """
        for idx, env in enumerate(self.envs):
            env.seed(seed + 1000 * idx)
<|endoftext|>


<commit_before># -*- coding: utf8 -*-
"""Functions to uncompact details inside field values.

.. note::

    :func:`re.compile` operations are inside of expand functions:

        1. readability
        2. module-level function bytecode is cached in python
        3. the last used compiled regexes are cached
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals, with_statement)

import re
import zhon.pinyin
import zhon.hanzi

from unihan_etl.constants import SPACE_DELIMITED_FIELDS

def expand_kDefinition(value):
    return [c.strip() for c in value.split(';')]

def expand_kMandarin(value):
    cn = value[0]
    if len(value) == 1:
        tw = value[0]
    else:
        tw = value[1]
    return {
        "zh-Hans": cn,
        "zh-Hant": tw
    }

def expand_kTotalStrokes(value):
    cn = value[0]
    if len(value) == 1:
        tw = value[0]
    else:
        tw = value[1]
    return {
        "zh-Hans": int(cn),
        "zh-Hant": int(tw)
    }

def expand_kHanYu(value):
    for i, v in enumerate(value):
        value[i] = {
            "volume": int(v[0]),
            "page": int(v[1:5]),
            "character": int(v[6:8]),
            "virtual": int(v[8])
        }
    return value

def expand_kHanyuPinyin(value):
    location_pattern = re.compile(r"""
        (?P<volume>[1-8])
        (?P<page>[0-9]{4})\.
        (?P<character>[0-3][0-9])
        (?P<virtual>[0-3])
    """, re.X)

    for i, v in enumerate(value):
        v = [s.split(',') for s in v.split(':')]
        value[i] = {
            "locations": v[0],
            "readings": v[1]
        }

        for n, loc in enumerate(value[i]['locations']):
            m = location_pattern.match(loc).groupdict()
            value[i]['locations'][n] = {
                "volume": int(m['volume']),
                "page": int(m['page']),
                "character": int(m['character']),
                "virtual": int(m['virtual'])
            }
    return value

def expand_kXHC1983(value):
    pattern = re.compile(r"""
        (?P<page>[0-9]{4})\.
        (?P<character>[0-9]{2})
        (?P<entry>[0-9]{1})
        (?P<substituted>\*?)
    """, re.X)

    for i, v in enumerate(value):
        vals = v.split(':')
        value[i] = {
            "locations": vals[0].split(','),
            "reading": vals[1],
        }

        for n, loc in enumerate(value[i]['locations']):
            m = pattern.match(loc).groupdict()
            value[i]['locations'][n] = {
                "page": int(m['page']),
                "character": int(m['character']),
                "entry": int(m['entry']),
                "substituted": m['substituted'] == "*"
            }
    return value

def expand_kCheungBauer(value):
    pattern = re.compile(r"""
        (?P<radical>[0-9]{3})\/(?P<strokes>[0-9]{2});
        (?P<cangjie>[A-Z]*);
        (?P<readings>[a-z1-6\[\]\/,]+)
    """, re.X)
    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "cangjie": m['cangjie'] or None,
            "readings": m['readings'].split(',')
        }
    return value

def expand_kRSAdobe_Japan1_6(value):
    pattern = re.compile(r"""
        (?P<type>[CV])\+
        (?P<cid>[0-9]{1,5})\+
        (?P<radical>[1-9][0-9]{0,2})\.
        (?P<strokes>[1-9][0-9]?)\.
        (?P<strokes_residue>[0-9]{1,2})
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()

        value[i] = {
            "type": m['type'],
            "cid": int(m['cid']),
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "strokes-residue": int(m['strokes_residue'])
        }
    return value

def expand_kCihaiT(value):
    for i, v in enumerate(value):
        vals = v.split('.')
        value[i] = {
            "page": int(vals[0]),
            "row": int(vals[1][0]),
            "character": int(vals[1][1:3]),
        }
    return value

def expand_kDaeJaweon(value):
    pattern = re.compile(r"""
        (?P<page>[0-9]{4})\.
        (?P<character>[0-9]{2})
        (?P<virtual>[01])
    """, re.X)
    m = pattern.match(value).groupdict()

    value = {
        "page": int(m['page']),
        "character": int(m['character']),
        "virtual": int(m['virtual']),
    }
    return value

def expand_kFenn(value):
    pattern = re.compile(r"""
        (?P<phonetic>[0-9]+a?)
        (?P<frequency>[A-KP*])
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict(v)
        value[i] = {
            "phonetic": int(m['phonetic']),
            "frequency": m['frequency']
        }
    return value

def expand_kHanyuPinlu(value):
    pattern = re.compile(r"""
        (?P<phonetic>[a-z{}]+)
        \((?P<frequency>[0-9]+)\)
    """.format(zhon.pinyin.lowercase), re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "phonetic": m['phonetic'],
            "frequency": int(m['frequency'])
        }
    return value

def expand_kHDZRadBreak(value):
    rad, loc = value.split(':')

    pattern = re.compile(r"""
        (?P<radical>[{}]+)
        \[(?P<ucn>U\+2F[0-9A-D][0-9A-F])\]
    """.format(zhon.hanzi.radicals), re.X)
    m = pattern.match(rad).groupdict()

    return {
        "radical": m['radical'],
        "ucn": m['ucn'],
        "location": loc
    }

def expand_kSBGY(value):
    for i, v in enumerate(value):
        vals = v.split('.')
        value[i] = {
            "page": int(vals[0]),
            "character": int(vals[1])
        }
    return value

def _expand_kRSGeneric(value):
    pattern = re.compile(r"""
        (?P<radical>[1-9][0-9]{0,2})
        (?P<simplified>\'?)\.
        (?P<strokes>-?[0-9]{1,2})
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "simplified": m['simplified'] == "'"
        }
    return value

expand_kRSUnicode = _expand_kRSGeneric
expand_kRSJapanese = _expand_kRSGeneric
expand_kRSKangXi = _expand_kRSGeneric
expand_kRSKanWa = _expand_kRSGeneric
expand_kRSKorean = _expand_kRSGeneric

def _expand_kIRG_GenericSource(value):
    v = value.split('-')
    return {
        "source": v[0],
        "location": v[1]
    }

expand_kIRG_GSource = _expand_kIRG_GenericSource
expand_kIRG_HSource = _expand_kIRG_GenericSource
expand_kIRG_JSource = _expand_kIRG_GenericSource
expand_kIRG_KPSource = _expand_kIRG_GenericSource
expand_kIRG_KSource = _expand_kIRG_GenericSource
expand_kIRG_MSource = _expand_kIRG_GenericSource
expand_kIRG_TSource = _expand_kIRG_GenericSource
expand_kIRG_USource = _expand_kIRG_GenericSource
expand_kIRG_VSource = _expand_kIRG_GenericSource

def expand_kGSR(value):
    pattern = re.compile(r"""
        (?P<set>[0-9]{4})
        (?P<letter>[a-vx-z])
        (?P<apostrophe>\')?
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "set": int(m['set']),
            "letter": m['letter'],
            "apostrophe": m['apostrophe'] == "'"
        }
    return value

def expand_field(field, fvalue):
    """Return structured value of information in UNIHAN field.

    :param field: field name
    :type field: str
    :param fvalue: value of field
    :type favalue: str
    :returns: list or dict of expanded field information per UNIHAN's
        documentation
    :rtype: list or dict
    """
    if field in SPACE_DELIMITED_FIELDS and fvalue:
        fvalue = fvalue.split(' ')

    try:
        expansion_func = eval('expand_%s' % field)
        return expansion_func(fvalue)
    except NameError:
        pass

    return fvalue
<commit_msg>use regex for kCihaiT<commit_after># -*- coding: utf8 -*-
"""Functions to uncompact details inside field values.

.. note::

    :func:`re.compile` operations are inside of expand functions:

        1. readability
        2. module-level function bytecode is cached in python
        3. the last used compiled regexes are cached
"""

from __future__ import (absolute_import, division, print_function,
                        unicode_literals, with_statement)

import re
import zhon.pinyin
import zhon.hanzi

from unihan_etl.constants import SPACE_DELIMITED_FIELDS

def expand_kDefinition(value):
    return [c.strip() for c in value.split(';')]

def expand_kMandarin(value):
    cn = value[0]
    if len(value) == 1:
        tw = value[0]
    else:
        tw = value[1]
    return {
        "zh-Hans": cn,
        "zh-Hant": tw
    }

def expand_kTotalStrokes(value):
    cn = value[0]
    if len(value) == 1:
        tw = value[0]
    else:
        tw = value[1]
    return {
        "zh-Hans": int(cn),
        "zh-Hant": int(tw)
    }

def expand_kHanYu(value):
    for i, v in enumerate(value):
        value[i] = {
            "volume": int(v[0]),
            "page": int(v[1:5]),
            "character": int(v[6:8]),
            "virtual": int(v[8])
        }
    return value

def expand_kHanyuPinyin(value):
    location_pattern = re.compile(r"""
        (?P<volume>[1-8])
        (?P<page>[0-9]{4})\.
        (?P<character>[0-3][0-9])
        (?P<virtual>[0-3])
    """, re.X)

    for i, v in enumerate(value):
        v = [s.split(',') for s in v.split(':')]
        value[i] = {
            "locations": v[0],
            "readings": v[1]
        }

        for n, loc in enumerate(value[i]['locations']):
            m = location_pattern.match(loc).groupdict()
            value[i]['locations'][n] = {
                "volume": int(m['volume']),
                "page": int(m['page']),
                "character": int(m['character']),
                "virtual": int(m['virtual'])
            }
    return value

def expand_kXHC1983(value):
    pattern = re.compile(r"""
        (?P<page>[0-9]{4})\.
        (?P<character>[0-9]{2})
        (?P<entry>[0-9]{1})
        (?P<substituted>\*?)
    """, re.X)

    for i, v in enumerate(value):
        vals = v.split(':')
        value[i] = {
            "locations": vals[0].split(','),
            "reading": vals[1],
        }

        for n, loc in enumerate(value[i]['locations']):
            m = pattern.match(loc).groupdict()
            value[i]['locations'][n] = {
                "page": int(m['page']),
                "character": int(m['character']),
                "entry": int(m['entry']),
                "substituted": m['substituted'] == "*"
            }
    return value

def expand_kCheungBauer(value):
    pattern = re.compile(r"""
        (?P<radical>[0-9]{3})\/(?P<strokes>[0-9]{2});
        (?P<cangjie>[A-Z]*);
        (?P<readings>[a-z1-6\[\]\/,]+)
    """, re.X)
    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "cangjie": m['cangjie'] or None,
            "readings": m['readings'].split(',')
        }
    return value

def expand_kRSAdobe_Japan1_6(value):
    pattern = re.compile(r"""
        (?P<type>[CV])\+
        (?P<cid>[0-9]{1,5})\+
        (?P<radical>[1-9][0-9]{0,2})\.
        (?P<strokes>[1-9][0-9]?)\.
        (?P<strokes_residue>[0-9]{1,2})
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()

        value[i] = {
            "type": m['type'],
            "cid": int(m['cid']),
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "strokes-residue": int(m['strokes_residue'])
        }
    return value

def expand_kCihaiT(value):
    pattern = re.compile(r"""
        (?P<page>[1-9][0-9]{0,3})\.
        (?P<row>[0-9]{1})
        (?P<character>[0-9]{2})
    """, re.X)
    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "page": int(m['page']),
            "row": int(m['row']),
            "character": int(m['character']),
        }
    return value

def expand_kDaeJaweon(value):
    pattern = re.compile(r"""
        (?P<page>[0-9]{4})\.
        (?P<character>[0-9]{2})
        (?P<virtual>[01])
    """, re.X)
    m = pattern.match(value).groupdict()

    value = {
        "page": int(m['page']),
        "character": int(m['character']),
        "virtual": int(m['virtual']),
    }
    return value

def expand_kFenn(value):
    pattern = re.compile(r"""
        (?P<phonetic>[0-9]+a?)
        (?P<frequency>[A-KP*])
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict(v)
        value[i] = {
            "phonetic": int(m['phonetic']),
            "frequency": m['frequency']
        }
    return value

def expand_kHanyuPinlu(value):
    pattern = re.compile(r"""
        (?P<phonetic>[a-z{}]+)
        \((?P<frequency>[0-9]+)\)
    """.format(zhon.pinyin.lowercase), re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "phonetic": m['phonetic'],
            "frequency": int(m['frequency'])
        }
    return value

def expand_kHDZRadBreak(value):
    rad, loc = value.split(':')

    pattern = re.compile(r"""
        (?P<radical>[{}]+)
        \[(?P<ucn>U\+2F[0-9A-D][0-9A-F])\]
    """.format(zhon.hanzi.radicals), re.X)
    m = pattern.match(rad).groupdict()

    return {
        "radical": m['radical'],
        "ucn": m['ucn'],
        "location": loc
    }

def expand_kSBGY(value):
    for i, v in enumerate(value):
        vals = v.split('.')
        value[i] = {
            "page": int(vals[0]),
            "character": int(vals[1])
        }
    return value

def _expand_kRSGeneric(value):
    pattern = re.compile(r"""
        (?P<radical>[1-9][0-9]{0,2})
        (?P<simplified>\'?)\.
        (?P<strokes>-?[0-9]{1,2})
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "radical": int(m['radical']),
            "strokes": int(m['strokes']),
            "simplified": m['simplified'] == "'"
        }
    return value

expand_kRSUnicode = _expand_kRSGeneric
expand_kRSJapanese = _expand_kRSGeneric
expand_kRSKangXi = _expand_kRSGeneric
expand_kRSKanWa = _expand_kRSGeneric
expand_kRSKorean = _expand_kRSGeneric

def _expand_kIRG_GenericSource(value):
    v = value.split('-')
    return {
        "source": v[0],
        "location": v[1]
    }

expand_kIRG_GSource = _expand_kIRG_GenericSource
expand_kIRG_HSource = _expand_kIRG_GenericSource
expand_kIRG_JSource = _expand_kIRG_GenericSource
expand_kIRG_KPSource = _expand_kIRG_GenericSource
expand_kIRG_KSource = _expand_kIRG_GenericSource
expand_kIRG_MSource = _expand_kIRG_GenericSource
expand_kIRG_TSource = _expand_kIRG_GenericSource
expand_kIRG_USource = _expand_kIRG_GenericSource
expand_kIRG_VSource = _expand_kIRG_GenericSource

def expand_kGSR(value):
    pattern = re.compile(r"""
        (?P<set>[0-9]{4})
        (?P<letter>[a-vx-z])
        (?P<apostrophe>\')?
    """, re.X)

    for i, v in enumerate(value):
        m = pattern.match(v).groupdict()
        value[i] = {
            "set": int(m['set']),
            "letter": m['letter'],
            "apostrophe": m['apostrophe'] == "'"
        }
    return value

def expand_field(field, fvalue):
    """Return structured value of information in UNIHAN field.

    :param field: field name
    :type field: str
    :param fvalue: value of field
    :type favalue: str
    :returns: list or dict of expanded field information per UNIHAN's
        documentation
    :rtype: list or dict
    """
    if field in SPACE_DELIMITED_FIELDS and fvalue:
        fvalue = fvalue.split(' ')

    try:
        expansion_func = eval('expand_%s' % field)
        return expansion_func(fvalue)
    except NameError:
        pass

    return fvalue
<|endoftext|>


<commit_before><commit_msg>Added _SUCCESS to the process<commit_after><|endoftext|>


<commit_before>'''
A module for shelling out

Keep in mind that this module is insecure, in that it can give whomever has
access to the master root execution access to all salt minions
'''

import logging
import os
import subprocess
import tempfile

# Set up logging
log = logging.getLogger(__name__)

# Set the default working directory to the home directory
# of the user salt-minion is running as.  Default:  /root
DEFAULT_CWD = os.path.expanduser('~')

# Set up the default outputters
__outputter__ = {
                 'run': 'txt'
                 }

def _is_exec(path):
    '''
    Return true if the passed path exists and is execuatable
    '''
    return os.path.exists(path) and os.access(path, os.X_OK)

def run(cmd, cwd=DEFAULT_CWD):
    '''
    Execute the passed command and return the output as a string

    CLI Example::

        salt '*' cmd.run "ls -l | grep foo | awk '{print $2}'"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    out = subprocess.Popen(cmd,
            shell=True,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT).communicate()[0]
    log.debug(out)
    return out

def run_stdout(cmd, cwd=DEFAULT_CWD):
    '''
    Execute a command, and only return the standard out

    CLI Example::

        salt '*' cmd.run_stdout "ls -l | grep foo | awk '{print $2}'"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    stdout = subprocess.Popen(cmd,
            shell=True,
            cwd=cwd,
            stdout=subprocess.PIPE).communicate()[0]
    log.debug(stdout)
    return stdout

def run_stderr(cmd, cwd=DEFAULT_CWD):
    '''
    Execute a command and only return the standard error

    CLI Example::

        salt '*' cmd.run_stderr "ls -l | grep foo | awk '{print $2}'"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    stderr = subprocess.Popen(cmd,
            shell=True,
            cwd=cwd,
            stderr=subprocess.PIPE).communicate()[0]
    log.debug(stderr)
    return stderr

def run_all(cmd, cwd=DEFAULT_CWD):
    '''
    Execute the passed command and return a dict of return data

    CLI Example::

        salt '*' cmd.run_all "ls -l | grep foo | awk '{print $2}'"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    ret = {}
    proc = subprocess.Popen(cmd,
            shell=True,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
    out = proc.communicate()
    ret['stdout'] = out[0]
    ret['stderr'] = out[1]
    ret['retcode'] = proc.returncode
    ret['pid'] = proc.pid
    if not ret['retcode']:
        log.error('Command {0} failed'.format(cmd))
        log.error('stdout: {0}'.format(ret['stdout']))
        log.error('stderr: {0}'.format(ret['stderr']))
    else:
        log.info('stdout: {0}'.format(ret['stdout']))
        log.info('stderr: {0}'.format(ret['stderr']))
    return ret

def retcode(cmd, cwd=DEFAULT_CWD):
    '''
    Execute a shell command and return the command's return code.

    CLI Example::

        salt '*' cmd.retcode "file /bin/bash"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    return subprocess.call(cmd, shell=True, cwd=cwd)

def has_exec(cmd):
    '''
    Returns true if the executable is available on the minion, false otherwise

    CLI Example::

        salt '*' cmd.has_exec cat
    '''
    if cmd.startswith('/'):
        return _is_exec(cmd)
    for path in os.environ['PATH'].split(os.pathsep):
        fn_ = os.path.join(path, cmd)
        if _is_exec(fn_):
            return True
    return False

def exec_code(lang, code, cwd=DEFAULT_CWD):
    '''
    Pass in two strings, the first naming the executable language, aka -
    python2, python3, ruby, perl, lua, etc. the second string containing
    the code you wish to execute. The stdout and stderr will be returned

    CLI Example::

        salt '*' cmd.exec_code ruby 'puts "cheese"'
    '''
    fd, cfn = tempfile.mkstemp()
    open(cfn, 'w+').write(code)
    return subprocess.Popen(lang + ' ' + cfn,
            shell=True,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT).communicate()[0]
<commit_msg>Cleanups of the cmd module<commit_after>'''
A module for shelling out

Keep in mind that this module is insecure, in that it can give whomever has
access to the master root execution access to all salt minions
'''

import logging
import os
import subprocess
import tempfile
import salt.utils

# Set up logging
log = logging.getLogger(__name__)

# Set the default working directory to the home directory
# of the user salt-minion is running as.  Default:  /root
DEFAULT_CWD = os.path.expanduser('~')

# Set up the default outputters
__outputter__ = {
    'run': 'txt',
}

def _is_exec(path):
    '''
    Return true if the passed path exists and is execuatable
    '''
    return os.path.exists(path) and os.access(path, os.X_OK)

def run(cmd, cwd=DEFAULT_CWD):
    '''
    Execute the passed command and return the output as a string

    CLI Example::

        salt '*' cmd.run "ls -l | awk '/foo/{print $2}'"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    out = subprocess.Popen(cmd,
            shell=True,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT).communicate()[0]
    log.debug(out)
    return out

def run_stdout(cmd, cwd=DEFAULT_CWD):
    '''
    Execute a command, and only return the standard out

    CLI Example::

        salt '*' cmd.run_stdout "ls -l | awk '/foo/{print $2}'"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    stdout = subprocess.Popen(cmd,
            shell=True,
            cwd=cwd,
            stdout=subprocess.PIPE).communicate()[0]
    log.debug(stdout)
    return stdout

def run_stderr(cmd, cwd=DEFAULT_CWD):
    '''
    Execute a command and only return the standard error

    CLI Example::

        salt '*' cmd.run_stderr "ls -l | awk '/foo/{print $2}'"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    stderr = subprocess.Popen(cmd,
            shell=True,
            cwd=cwd,
            stderr=subprocess.PIPE).communicate()[0]
    log.debug(stderr)
    return stderr

def run_all(cmd, cwd=DEFAULT_CWD):
    '''
    Execute the passed command and return a dict of return data

    CLI Example::

        salt '*' cmd.run_all "ls -l | awk '/foo/{print $2}'"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    ret = {}
    proc = subprocess.Popen(cmd,
            shell=True,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.PIPE)
    out = proc.communicate()
    ret['stdout'] = out[0]
    ret['stderr'] = out[1]
    ret['retcode'] = proc.returncode
    ret['pid'] = proc.pid
    if not ret['retcode']:
        log.error('Command {0} failed'.format(cmd))
        log.error('stdout: {0}'.format(ret['stdout']))
        log.error('stderr: {0}'.format(ret['stderr']))
    else:
        log.info('stdout: {0}'.format(ret['stdout']))
        log.info('stderr: {0}'.format(ret['stderr']))
    return ret

def retcode(cmd, cwd=DEFAULT_CWD):
    '''
    Execute a shell command and return the command's return code.

    CLI Example::

        salt '*' cmd.retcode "file /bin/bash"
    '''
    log.info('Executing command {0} in directory {1}'.format(cmd, cwd))
    return subprocess.call(cmd, shell=True, cwd=cwd)

def has_exec(cmd):
    '''
    Returns true if the executable is available on the minion, false otherwise

    CLI Example::

        salt '*' cmd.has_exec cat
    '''
    return bool(salt.utils.which(cmd))

def which(cmd):
    '''
    Returns the path of an executable available on the minion, None otherwise

    CLI Example::

        salt '*' cmd.which cat
    '''
    return salt.utils.which(cmd)

def exec_code(lang, code, cwd=DEFAULT_CWD):
    '''
    Pass in two strings, the first naming the executable language, aka -
    python2, python3, ruby, perl, lua, etc. the second string containing
    the code you wish to execute. The stdout and stderr will be returned

    CLI Example::

        salt '*' cmd.exec_code ruby 'puts "cheese"'
    '''
    fd, cfn = tempfile.mkstemp()
    open(cfn, 'w+').write(code)
    return subprocess.Popen(lang + ' ' + cfn,
            shell=True,
            cwd=cwd,
            stdout=subprocess.PIPE,
            stderr=subprocess.STDOUT).communicate()[0]
<|endoftext|>


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

# Copyright (c) 2011 Stefaan Lippens

'''
Git update hook Nitty Committy
==============================
Git update hook that manages a commit message histogram
and complains about or blocks pushes with messages that are too common.

Usage:
-----
Copy/symlink update.py to hooks/update in the git repo to protect.

'''

import os
import sys
import subprocess
import sqlite3
import logging

# TODO: support for pushing new branches (00000)
# TODO: support for non-fast forward pushes
# TODO: normalize messages (capitalization, non alphanum characters, ...)

def git_log(begin, end):
    proc = subprocess.Popen(['git', 'log', '--no-merges', '--format=%an:%cn:%s', begin + '..' + end], stdout=subprocess.PIPE)
    stdout = proc.communicate()[0]
    assert proc.returncode == 0
    log = []
    for line in stdout.split('\n'):
        if len(line) > 1:
            log.append(tuple(line.split(':', 2)))
    return log

class MessageHistogram(object):
    '''
    Wrapper class that manages a message histrogram
    '''

    def __init__(self, db_file_name):
        logging.debug('MessageHistogram from file ' + db_file_name)

        if not os.path.exists(db_file_name):
            MessageHistogram.create_database(db_file_name)

        self._conn = sqlite3.connect(db_file_name)

    @classmethod
    def create_database(cls, db_file_name):
        logging.debug('Creating database in file "{0}".'.format(db_file_name))
        conn = sqlite3.connect(db_file_name)
        c = conn.cursor()
        c.execute('''
            CREATE TABLE message_histogram (
                message varchar(160),
                count int unsigned,
                UNIQUE (message)
            )
        ''')
        conn.commit()

    def increase(self, message):
        '''Increase the count for a message'''
        c = self._conn.cursor()
        try:
            c.execute('INSERT INTO message_histogram (message, count) VALUES (:msg, 1)', {'msg': message})
        except sqlite3.IntegrityError:
            c.execute('UPDATE message_histogram SET count = count + 1 WHERE message = :msg', {'msg': message})
        self._conn.commit()
        logging.debug('Increased count for message "{0}"'.format(message))

    def get_top_n_messages(self, n=10):
        '''Get the top N messages.'''
        c = self._conn.cursor()
        c.execute('SELECT message, count FROM message_histogram ORDER BY count DESC LIMIT ?', (n,))
        return c.fetchall()

    def in_top_n(self, message, n=10):
        '''Check whether the given message is in the top N (and return the given count if so).'''
        top_n = self.get_top_n_messages(n)
        for msg, count in top_n:
            if msg == message:
                return count

def main():
    logging.basicConfig(level=logging.DEBUG)

    logging.debug('sys.argv = {0!r}'.format(sys.argv))
    (ref, current, new) = sys.argv[1:4]

    if current == '0000000000000000000000000000000000000000':
        # Pushing a new branch/tag: not easy to do a git log here.
        # TODO: support checking of new branches
        sys.exit(0)

    # Get log messages
    log = git_log(current, new)

    # Load commit message database
    filename = os.path.splitext(__file__)[0] + '.db.sqlite'
    histogram = MessageHistogram(filename)

    for (author, committer, msg) in log:
        if histogram.in_top_n(msg):
            print 'Warning: I don\'t like this commit message (by {author}): "{msg}"'.format(msg=msg, author=author)
        histogram.increase(msg)

if __name__ == '__main__':
    try:
        main()
    except Exception, e:
        print sys.argv[0], 'failed with exception:', repr(e)

<commit_msg>added normalization of messages<commit_after>#!/usr/bin/env python

# Copyright (c) 2011 Stefaan Lippens

'''
Git update hook Nitty Committy
==============================
Git update hook that manages a commit message histogram
and complains about or blocks pushes with messages that are too common.

Usage:
-----
Copy/symlink update.py to hooks/update in the git repo to protect.

'''

import os
import sys
import subprocess
import sqlite3
import logging
import re

# TODO: support for pushing new branches (00000)
# TODO: support for non-fast forward pushes

def git_log(begin, end):
    proc = subprocess.Popen(['git', 'log', '--no-merges', '--format=%an:%cn:%s', begin + '..' + end], stdout=subprocess.PIPE)
    stdout = proc.communicate()[0]
    assert proc.returncode == 0
    log = []
    for line in stdout.split('\n'):
        if len(line) > 1:
            log.append(tuple(line.split(':', 2)))
    return log

def normalize_message(message):
    '''
    Normalize a message by collapsing whitespace, removing capitalization, unimportant characters, ...
    '''
    message = message.lower()
    message = re.sub('[^a-z0-9 ]', '', message)
    message = re.sub('\s+', ' ', message)
    return message

class MessageHistogram(object):
    '''
    Wrapper class that manages a message histrogram
    '''

    def __init__(self, db_file_name):
        logging.debug('MessageHistogram from file ' + db_file_name)

        if not os.path.exists(db_file_name):
            MessageHistogram.create_database(db_file_name)

        self._conn = sqlite3.connect(db_file_name)

    @classmethod
    def create_database(cls, db_file_name):
        logging.debug('Creating database in file "{0}".'.format(db_file_name))
        conn = sqlite3.connect(db_file_name)
        c = conn.cursor()
        c.execute('''
            CREATE TABLE message_histogram (
                message varchar(160),
                count int unsigned,
                UNIQUE (message)
            )
        ''')
        conn.commit()

    def increase(self, message):
        '''Increase the count for a message'''
        c = self._conn.cursor()
        try:
            c.execute('INSERT INTO message_histogram (message, count) VALUES (:msg, 1)', {'msg': message})
        except sqlite3.IntegrityError:
            c.execute('UPDATE message_histogram SET count = count + 1 WHERE message = :msg', {'msg': message})
        self._conn.commit()
        logging.debug('Increased count for message "{0}"'.format(message))

    def get_top_n_messages(self, n=10):
        '''Get the top N messages.'''
        c = self._conn.cursor()
        c.execute('SELECT message, count FROM message_histogram ORDER BY count DESC LIMIT ?', (n,))
        return c.fetchall()

    def in_top_n(self, message, n=10):
        '''Check whether the given message is in the top N (and return the given count if so).'''
        top_n = self.get_top_n_messages(n)
        for msg, count in top_n:
            if msg == message:
                return count

def main():
    logging.basicConfig(level=logging.DEBUG)

    logging.debug('sys.argv = {0!r}'.format(sys.argv))
    (ref, current, new) = sys.argv[1:4]

    if current == '0000000000000000000000000000000000000000':
        # Pushing a new branch/tag: not easy to do a git log here.
        # TODO: support checking of new branches
        sys.exit(0)

    # Get log messages
    log = git_log(current, new)

    # Load commit message database
    filename = os.path.splitext(__file__)[0] + '.db.sqlite'
    histogram = MessageHistogram(filename)

    for (author, committer, msg) in log:
        msg = normalize_message(msg)
        if histogram.in_top_n(msg):
            print 'Warning: I don\'t like this commit message (by {author}): "{msg}"'.format(msg=msg, author=author)
        histogram.increase(msg)

if __name__ == '__main__':
    try:
        main()
    except Exception, e:
        print sys.argv[0], 'failed with exception:', repr(e)

<|endoftext|>


<commit_before><commit_msg>Name.appendSegment() now works on Python3<commit_after><|endoftext|>


<commit_before>#-*- coding: utf-8 -*-

import os
import re
import logging
import datetime
from google.appengine.ext.webapp.util import run_wsgi_app as run
from google.appengine.ext import webapp
from google.appengine.api import memcache, urlfetch
from google.appengine.ext.webapp import template
from django.utils import simplejson

from libs import twitter
from utils import be_god
from utils.handler import WebHandler
from utils.render import render
import config
import dbs

week = 604800

register = template.create_template_register()
template.register_template_library('god.twitter')
@register.filter
def at(value):
    value = re.sub(r'@([a-zA-Z0-9\_]+)',r'@\1', value)
    return value

class Twitter(object):
    def __init__(self, token=None):
        self.consumer = twitter.oauth.Consumer(config.twitter_key, config.twitter_secret)
        self.token = token
        self.method = twitter.oauth.SignatureMethod_HMAC_SHA1()

    def set_qs_token(self, qs):
        self.token = twitter.oauth.Token.from_string(qs)

    def get_url(self, url, method='GET', param={}):
        req = twitter.oauth.Request.from_consumer_and_token(
            self.consumer, token=self.token, http_method=method,
            http_url = url, parameters=param)
        req.sign_request(self.method, self.consumer, self.token)
        return req.to_url()
    def fetch(self, url):
        try:
            res = urlfetch.fetch(url)
        except urlfetch.DownloadError, e:
            logging.error('Twitter: ' + str(e))
            raise Exception('Twitter: ' + str(e))
        return res.content
    def set_qs_api(self, qs, input_encoding=None):
        self.token = twitter.oauth.Token.from_string(qs)
        api = twitter.Api(self.consumer.key, self.consumer.secret,
                          self.token.key, self.token.secret, input_encoding)
        return api

def get_path(ua, name):
    path = os.path.join(config.ROOT, 'god', 'tpl', name)
    return path

class Dashboard(WebHandler):
    @be_god
    def get(self):
        ua = self.request.headers.get('User-Agent', 'bot')
        path = get_path(ua, 'twitter_dashboard.html')
        rdic = {}
        source = self.request.get('from', None)
        message = ''
        if source:
            message = self.session.get('message','')
            self.session.delete('message')
        rdic['message'] = message
        qs = dbs.Vigo.get('oauth_twitter')
        if not qs:
            return self.redirect('/god/twitter/login')
        api = Twitter().set_qs_api(qs)
        statuses = memcache.get('twitter$home')
        if statuses is None:
            statuses = api.GetFriendsTimeline(count=30, retweets=True)
            for status in statuses:
                status.datetime = datetime.datetime.strptime(status.created_at, '%a %b %d %H:%M:%S +0000 %Y')
            memcache.set('twitter$home', statuses, 240)
        rdic['statuses'] = statuses
        return self.response.out.write(render(path, rdic))

class UserStatus(WebHandler):
    @be_god
    def get(self, username):
        ua = self.request.headers.get('User-Agent', 'bot')
        path = get_path(ua, 'twitter_user.html')
        rdic = {}
        rdic['username'] = username
        qs = dbs.Vigo.get('oauth_twitter')
        if not qs:
            return self.redirect('/god/twitter/login')
        api = Twitter().set_qs_api(qs)
        statuses = memcache.get('twitter$status$' + username)
        if statuses is None:
            try:
                statuses = api.GetUserTimeline(count=30, screen_name=username)
            except twitter.TwitterError, e:
                logging.error(str(e))
                return self.redirect('/god/twitter')
            for status in statuses:
                status.datetime = datetime.datetime.strptime(status.created_at, '%a %b %d %H:%M:%S +0000 %Y')
            memcache.set('twitter$status$' + username, statuses, 240)
        rdic['statuses'] = statuses
        #profile = memcache.get('twitter$profile$' + username)
        #if profile is None:
        #    profile = api.GetUser(username)
        #    memcache.set('twitter$profile$'+username, profile, 86400)
        #rdic['profile'] = profile
        return self.response.out.write(render(path, rdic))

class Login(WebHandler):
    @be_god
    def get(self):
        auth = Twitter()
        url = auth.get_url(twitter.REQUEST_TOKEN_URL)
        qs = auth.fetch(url)
        self.session['twitter_qs'] = qs
        auth.set_qs_token(qs)
        to = auth.get_url(twitter.AUTHORIZATION_URL)
        return self.redirect(to)

class Auth(WebHandler):
    @be_god
    def get(self):
        qs = dbs.Vigo.get('oauth_twitter')
        if qs:
            return self.response.out.write('authed')
        qs = self.session.get('twitter_qs','')
        if not qs:
            self.session['message'] = 'Request Twitter access token failed'
            return self.redirect('/god?from=twitter')
        auth = Twitter()
        auth.set_qs_token(qs)
        url = auth.get_url(twitter.ACCESS_TOKEN_URL)
        dbs.Vigo.set('oauth_twitter', qs)
        self.session['message'] = 'Twitter Auth Success'
        return self.redirect('/god/twitter?from=auth')

class Status(WebHandler):
    @be_god
    def post(self):
        content = self.request.get('text', '')
        if len(content) > 140:
            content = content[:133] + '...'
        qs = dbs.Vigo.get('oauth_twitter')
        api = Twitter().set_qs_api(qs, 'utf-8')
        try:
            api.PostUpdate(content)
            data = {'text':'Post To Twitter Success'}
        except twitter.TwitterError, e:
            data = {'text':str(e)}
        self.response.headers['Content-Type'] = 'application/json'
        self.response.out.write(simplejson.dumps(data))

apps = webapp.WSGIApplication(
    [
        ('/god/twitter', Dashboard),
        ('/god/twitter/login', Login),
        ('/god/twitter/auth', Auth),
        ('/god/twitter/status', Status),
        ('/god/twitter/user/(.*)', UserStatus),
    ],
    debug = config.DEBUG,
)

if '__main__' == __name__:
    run(apps)
<commit_msg>big bug..  how can i forget to fetch qs  for oauth<commit_after>#-*- coding: utf-8 -*-

import os
import re
import logging
import datetime
from google.appengine.ext.webapp.util import run_wsgi_app as run
from google.appengine.ext import webapp
from google.appengine.api import memcache, urlfetch
from google.appengine.ext.webapp import template
from django.utils import simplejson

from libs import twitter
from utils import be_god
from utils.handler import WebHandler
from utils.render import render
import config
import dbs

week = 604800

register = template.create_template_register()
template.register_template_library('god.twitter')
@register.filter
def at(value):
    value = re.sub(r'@([a-zA-Z0-9\_]+)',r'@\1', value)
    return value

class Twitter(object):
    def __init__(self, token=None):
        self.consumer = twitter.oauth.Consumer(config.twitter_key, config.twitter_secret)
        self.token = token
        self.method = twitter.oauth.SignatureMethod_HMAC_SHA1()

    def set_qs_token(self, qs):
        self.token = twitter.oauth.Token.from_string(qs)

    def get_url(self, url, method='GET', param={}):
        req = twitter.oauth.Request.from_consumer_and_token(
            self.consumer, token=self.token, http_method=method,
            http_url = url, parameters=param)
        req.sign_request(self.method, self.consumer, self.token)
        return req.to_url()
    def fetch(self, url):
        try:
            res = urlfetch.fetch(url)
        except urlfetch.DownloadError, e:
            logging.error('Twitter: ' + str(e))
            raise Exception('Twitter: ' + str(e))
        return res.content
    def set_qs_api(self, qs, input_encoding=None):
        self.token = twitter.oauth.Token.from_string(qs)
        api = twitter.Api(self.consumer.key, self.consumer.secret,
                          self.token.key, self.token.secret, input_encoding)
        return api

def get_path(ua, name):
    path = os.path.join(config.ROOT, 'god', 'tpl', name)
    return path

class Dashboard(WebHandler):
    @be_god
    def get(self):
        ua = self.request.headers.get('User-Agent', 'bot')
        path = get_path(ua, 'twitter_dashboard.html')
        rdic = {}
        source = self.request.get('from', None)
        message = ''
        if source:
            message = self.session.get('message','')
            self.session.delete('message')
        rdic['message'] = message
        qs = dbs.Vigo.get('oauth_twitter')
        if not qs:
            return self.redirect('/god/twitter/login')
        api = Twitter().set_qs_api(qs)
        statuses = memcache.get('twitter$home')
        if statuses is None:
            statuses = api.GetFriendsTimeline(count=30, retweets=True)
            for status in statuses:
                status.datetime = datetime.datetime.strptime(status.created_at, '%a %b %d %H:%M:%S +0000 %Y')
            memcache.set('twitter$home', statuses, 240)
        rdic['statuses'] = statuses
        return self.response.out.write(render(path, rdic))

class UserStatus(WebHandler):
    @be_god
    def get(self, username):
        ua = self.request.headers.get('User-Agent', 'bot')
        path = get_path(ua, 'twitter_user.html')
        rdic = {}
        rdic['username'] = username
        qs = dbs.Vigo.get('oauth_twitter')
        if not qs:
            return self.redirect('/god/twitter/login')
        api = Twitter().set_qs_api(qs)
        statuses = memcache.get('twitter$status$' + username)
        if statuses is None:
            try:
                statuses = api.GetUserTimeline(count=30, screen_name=username)
            except twitter.TwitterError, e:
                logging.error(str(e))
                return self.redirect('/god/twitter')
            for status in statuses:
                status.datetime = datetime.datetime.strptime(status.created_at, '%a %b %d %H:%M:%S +0000 %Y')
            memcache.set('twitter$status$' + username, statuses, 240)
        rdic['statuses'] = statuses
        #profile = memcache.get('twitter$profile$' + username)
        #if profile is None:
        #    profile = api.GetUser(username)
        #    memcache.set('twitter$profile$'+username, profile, 86400)
        #rdic['profile'] = profile
        return self.response.out.write(render(path, rdic))

class Login(WebHandler):
    @be_god
    def get(self):
        auth = Twitter()
        url = auth.get_url(twitter.REQUEST_TOKEN_URL)
        qs = auth.fetch(url)
        self.session['twitter_qs'] = qs
        auth.set_qs_token(qs)
        to = auth.get_url(twitter.AUTHORIZATION_URL)
        return self.redirect(to)

class Auth(WebHandler):
    @be_god
    def get(self):
        qs = dbs.Vigo.get('oauth_twitter')
        if qs:
            return self.response.out.write('authed')
        qs = self.session.get('twitter_qs','')
        if not qs:
            self.session['message'] = 'Request Twitter access token failed'
            return self.redirect('/god?from=twitter')
        auth = Twitter()
        auth.set_qs_token(qs)
        url = auth.get_url(twitter.ACCESS_TOKEN_URL)
        qs = auth.fetch(url)
        dbs.Vigo.set('oauth_twitter', qs)
        self.session['message'] = 'Twitter Auth Success'
        return self.redirect('/god/twitter?from=auth')

class Status(WebHandler):
    @be_god
    def post(self):
        content = self.request.get('text', '')
        if len(content) > 140:
            content = content[:133] + '...'
        qs = dbs.Vigo.get('oauth_twitter')
        api = Twitter().set_qs_api(qs, 'utf-8')
        try:
            api.PostUpdate(content)
            data = {'text':'Post To Twitter Success'}
        except twitter.TwitterError, e:
            data = {'text':str(e)}
        self.response.headers['Content-Type'] = 'application/json'
        self.response.out.write(simplejson.dumps(data))

apps = webapp.WSGIApplication(
    [
        ('/god/twitter', Dashboard),
        ('/god/twitter/login', Login),
        ('/god/twitter/auth', Auth),
        ('/god/twitter/status', Status),
        ('/god/twitter/user/(.*)', UserStatus),
    ],
    debug = config.DEBUG,
)

if '__main__' == __name__:
    run(apps)
<|endoftext|>