<commit_before>""" Module for basic plots related to FRB host and foreground galaxies"""
import os
import numpy as np
from IPython import embed

from pkg_resources import resource_filename

from matplotlib import pyplot as plt

from astropy.io import fits
from astropy.table import Table

from frb.figures import utils

primus_path = os.path.join(resource_filename('frb', 'data'), 'Public')

def sub_bpt(ax_BPT, galaxies, clrs, markers, show_kewley=True, SDSS_clr='BuGn',
            show_legend=True, bptdat=None):
    """
    Generate a BPT diagram

    To use this code, you must download the SDSS_BPT_stellar_mass.fits file from
    https://drive.google.com/open?id=1yHlfsvcRPXK73F6hboT1nM4bRF59ESab
    and put it in data/Public/SDSS

    Args:
        ax_BPT (matplotlib.Axis):
        galaxies (list):
          List of FRBGalaxy objects
        clrs (list):
            List of colors
        markers (list):
            List of markers
        show_kewley (bool, optional):
            Show the BPT lines?
        SDSS_clr (str, optional):
          Set the color map for SDSS
        show_legend (bool, optional):
          Show a legend
        bptdat (Table like):
            SDSS BPT data

    Returns:
        ax_BPT is modified in place

    """

    # Read in data
    if bptdat is None:
        sdss_file = os.path.join(resource_filename('frb', 'data'), 'Public', 'SDSS', 'SDSS_BPT_stellar_mass.fits')
        if not os.path.isfile(sdss_file):
            print("See the method notes to download the SDSS data!")
            return
        hdulist = fits.open(sdss_file)
        bptdat = hdulist[1].data

    # Select only non zero entries and SNR over 5
    lines = np.array(bptdat.names)[[("flux" in name) & ("err" not in name) for name in bptdat.names]]
    line_err = np.array(bptdat.names)[["flux_err" in name for name in bptdat.names]]
    select = {}
    for line, err in zip(lines, line_err):
        select[line] = bptdat[line] / bptdat[err] >= 5

    # SDSS
    bpt1 = select['oiii_5007_flux'] & select['h_beta_flux'] & select['nii_6584_flux'] & select['h_alpha_flux']
    y = bptdat['oiii_5007_flux'][bpt1] / bptdat['h_beta_flux'][bpt1]
    x = bptdat['nii_6584_flux'][bpt1] / bptdat['h_alpha_flux'][bpt1]

    xbins = 100
    ybins = 100
    # Plot
    counts, xedges, yedges = np.histogram2d(np.log10(x), np.log10(y), bins=(xbins, ybins))
    cm = plt.get_cmap(SDSS_clr)
    mplt = ax_BPT.pcolormesh(xedges, yedges, np.log10(counts.transpose()), cmap=cm)

    # Loop on the Galaxies
    for kk,galaxy in enumerate(galaxies):
        # Parse the emission lines
        NII, NII_err = galaxy.calc_nebular_lum('[NII] 6584')
        Ha, Ha_err = galaxy.calc_nebular_lum('Halpha')
        Hb, Hb_err = galaxy.calc_nebular_lum('Hbeta')
        try:
            OIII, OIII_err = galaxy.calc_nebular_lum('[OIII] 5007')
        except:
            import pdb; pdb.set_trace()
        #
        x0 = (NII/Ha).decompose().value
        y0 = (OIII/Hb).decompose().value
        x0_err = x0 * np.sqrt((NII_err / NII).decompose().value**2 +
                (Ha_err/Ha).decompose().value**2)
        y0_err = y0 * np.sqrt((OIII_err / OIII).decompose().value**2 +
                              (Hb_err/Hb).decompose().value**2)
        # Require at least 20% error
        x0_err = max(x0_err, 0.2*x0)
        y0_err = max(y0_err, 0.2*y0)

        logx, xerr = utils.log_me(x0, x0_err)
        logy, yerr = utils.log_me(y0, y0_err)
        # Upper limit on [NII]/Ha?
        if NII_err.value < 0.:
            xerr = None
            # Left arrow
            plt.arrow(logx, logy, -0.05, 0., fc=clrs[kk], ec=clrs[kk],
                      head_width=0.02, head_length=0.05)
        # Plot
        ax_BPT.errorbar([logx], [logy], xerr=xerr, yerr=yerr,
                        color=clrs[kk], marker=markers[kk], markersize="8",
                        capsize=3, label=galaxy.name)

    # Standard curves
    demarc = lambda x: 0.61 / (x - 0.05) + 1.3  # Kauffman et al 2003, MNRAS, 346, 4, pp. 1055-1077. Eq 1
    demarc_kewley = lambda x: 0.61 / (
                x - 0.47) + 1.19  # Kewley F., Dopita M., Sutherland R., Heisler C., Trevena J., 2001, ApJ, 556,121
    demarc_liner = lambda x: 1.01 * x + 0.48  # Cid Fernandes et al 2010, MNRAS, 403,1036 Eq 10
    ax_BPT.plot(np.linspace(-2, 0), demarc(np.linspace(-2, 0)), "k-", lw=2)#, label="Kauffman et al 2003")
    if show_kewley:
        ax_BPT.plot(np.linspace(-2, 0.25), demarc_kewley(np.linspace(-2, 0.25)), "k--", lw=2)#, label="Kewley et al 2001")
    ax_BPT.plot(np.linspace(-0.43, 0.5), demarc_liner(np.linspace(-0.43, 0.5)), "k--", lw=2)#, label="Cid Fernandes et al 2010")

    # Labels
    lsz = 13.
    ax_BPT.annotate(r"\textbf{Star-forming}", (-1.30, 0), fontsize=lsz)
    ax_BPT.annotate(r"\textbf{LINER}", (0.23, 0), fontsize=lsz)
    ax_BPT.annotate(r"\textbf{Seyfert}", (-0.5, 1), fontsize=lsz)

    # Legend
    if show_legend:
        ax_BPT.legend(loc="lower left")
    # Axes
    ax_BPT.set_xlabel(r"$\log \, ({\rm [N\textsc{ii}]/H\,\alpha)}$")
    ax_BPT.set_ylabel(r"$\log \, ({\rm [O\textsc{iii}]/H\,\beta)}$")
    ax_BPT.set_xlim(-1.5, 0.5)
    ax_BPT.set_ylim(-1, 1.2)
    utils.set_fontsize(ax_BPT, 13.)

def sub_sfms(ax_M, galaxies, clrs, markers):
    """
    Generate a SF vs. M* plot on top of PRIMUS galaxies

    Args:
        ax_M (matplotlib.axis):
        galaxies (list):
            List of FRB.galaxies.frbgalaxy.FRBGalaxy objects
        clrs (list):
            List of matplotlib colors
        markers (list):
            List of matplotlib marker types

    """
    utils.set_mplrc()

    # Load up
    primus_zcat = Table.read(os.path.join(primus_path, 'PRIMUS_2013_zcat_v1.fits.gz'))
    primus_mass = Table.read(os.path.join(primus_path, 'PRIMUS_2014_mass_v1.fits.gz'))

    gdz = (primus_zcat['Z'] > 0.2) & (primus_zcat['Z'] < 0.4)
    gd_mag = primus_zcat['SDSS_ABSMAG'][:,0] != 0.

    good_mass = primus_mass['ISGOOD'] == 1

    # PRIMUS
    # Photometry
    gd_color = gdz & gd_mag
    u_r = primus_zcat['SDSS_ABSMAG'][gd_color,0] - primus_zcat['SDSS_ABSMAG'][gd_color,2]
    rmag = primus_zcat['SDSS_ABSMAG'][gd_color,2]

    # Mass/SFR
    gd_msfr = good_mass & gdz
    mass = primus_mass['MASS'][gd_msfr]
    sfr = primus_mass['SFR'][gd_msfr]

    # Plot
    ms = 22.

    # Histogram
    xbins = 50
    ybins = 50
    counts, xedges, yedges = np.histogram2d(mass, sfr, bins=(xbins, ybins))
    #cm = plt.get_cmap('Reds')
    cm = plt.get_cmap('Greys')

    # SF
    mplt = ax_M.pcolormesh(xedges, yedges, counts.transpose(), cmap=cm)

    # Relation
    logm_star = np.linspace(8, 12)
    logsfr = lambda logm_star: -0.49 + 0.65 * (logm_star - 10) + 1.07 * (0.35 - 0.1)
    ax_M.plot(logm_star, logsfr(logm_star), "k--", lw=3)#, label="Moustakas et al 2013")

    # Galaxies
    for kk,galaxy in enumerate(galaxies):
        # M*
        if 'Mstar' in galaxy.derived.keys():
            logM, sig_logM = utils.log_me(galaxy.derived['Mstar'], galaxy.derived['Mstar_err'])
        elif 'Mstar_spec' in galaxy.derived.keys():
            logM, sig_logM = np.log10(galaxy.derived['Mstar_spec']), 0.3
        else:
            continue
        # SFR
        if 'SFR_nebular_err' in galaxy.derived.keys():
            logS, sig_logS = utils.log_me(galaxy.derived['SFR_nebular'], galaxy.derived['SFR_nebular_err'])
        else:
            logS, sig_logS = utils.log_me(galaxy.derived['SFR_nebular'], 0.3*galaxy.derived['SFR_nebular'])
        # Plot
        ax_M.errorbar([logM], [logS], xerr=sig_logM, yerr=sig_logS,
                      color=clrs[kk], marker=markers[kk],
                     markersize="12", capsize=3, label=galaxy.name)
        if sig_logS is None:
            # Down arrow
            plt.arrow(logM, logS, 0., -0.2, fc=clrs[kk], ec=clrs[kk],
                  head_width=0.02*4, head_length=0.05*2)

    ax_M.annotate(r"\textbf{Star forming}", (8.5, 0.8), fontsize=13.)
    ax_M.annotate(r"\textbf{Quiescent}", (11, -0.9), fontsize=13.)
    ax_M.set_xlabel("$\log \, (M_*/M_\odot)$")
    ax_M.set_ylabel("$\log \, SFR (M_\odot$/yr)")
    ax_M.legend(loc='lower left')
    ax_M.set_xlim(7.5, 11.8)
    ax_M.set_ylim(-2.5, 1.2)

def sub_color_mag(ax, galaxies, clrs, markers):
    """
    Generate a color-magnitude diagram using PRIMUS data
    and FRB galaxies

    Args:
        ax (matplotlib.Axis):
        galaxies (list):
            List of FRB.galaxies.frbgalaxy.FRBGalaxy objects
        clrs (list):
            List of matplotlib colors
        markers (list):
            List of matplotlib marker types

    Returns:

    """

    # Load up
    primus_zcat = Table.read(os.path.join(primus_path, 'PRIMUS_2013_zcat_v1.fits.gz'))
    #primus_mass = Table.read(os.path.join(primus_path, 'PRIMUS_2014_mass_v1.fits'))

    gdz = (primus_zcat['Z'] > 0.2) & (primus_zcat['Z'] < 0.4)
    gd_mag = primus_zcat['SDSS_ABSMAG'][:,0] != 0.

    # PRIMUS
    # Photometry
    gd_color = gdz & gd_mag
    u_r = primus_zcat['SDSS_ABSMAG'][gd_color,0] - primus_zcat['SDSS_ABSMAG'][gd_color,2]
    rmag = primus_zcat['SDSS_ABSMAG'][gd_color,2]

    xbins = 100
    ybins = 100
    counts, xedges, yedges = np.histogram2d(rmag, u_r, bins=(xbins, ybins))
    cm = plt.get_cmap('Greys')
    mplt = ax.pcolormesh(xedges, yedges, counts.transpose(), cmap=cm)
    '''
    cb = plt.colorbar(mplt, fraction=0.030, pad=0.04)
    cb.set_label('PRIMUS survey')
    '''
    for kk,galaxy in enumerate(galaxies):
        ax.errorbar([galaxy.derived['M_r']], [galaxy.derived['u-r']],
                    xerr=[galaxy.derived['M_r_err']],
                    yerr=[galaxy.derived['u-r_err']],
                    color=clrs[kk], marker=markers[kk],
                    markersize="5", capsize=3,
                    label=galaxy.name)
    # Label
    plt.ylabel(r"$u-r \textbf{(Rest-frame)}$")
    plt.xlabel(r"$r \textbf{(Rest-frame)}$")
    ax.legend(loc='lower right')

    ax.set_ylim(0.0, 3.3)
    ax.set_xlim(-15.5, -23)

    utils.set_fontsize(ax,11.)

<commit_msg>galaxy fig<commit_after>""" Module for basic plots related to FRB host and foreground galaxies"""
import os
import numpy as np
from IPython import embed

from pkg_resources import resource_filename

from matplotlib import pyplot as plt

from astropy.io import fits
from astropy.table import Table

from frb.figures import utils

primus_path = os.path.join(resource_filename('frb', 'data'), 'Public')

def sub_bpt(ax_BPT, galaxies, clrs, markers, show_kewley=True, SDSS_clr='BuGn',
            show_legend=True, bptdat=None):
    """
    Generate a BPT diagram

    To use this code, you must download the SDSS_BPT_stellar_mass.fits file from
    https://drive.google.com/open?id=1yHlfsvcRPXK73F6hboT1nM4bRF59ESab
    and put it in data/Public/SDSS

    Args:
        ax_BPT (matplotlib.Axis):
        galaxies (list):
          List of FRBGalaxy objects
        clrs (list):
            List of colors
        markers (list):
            List of markers
        show_kewley (bool, optional):
            Show the BPT lines?
        SDSS_clr (str, optional):
          Set the color map for SDSS
        show_legend (bool, optional):
          Show a legend
        bptdat (Table like):
            SDSS BPT data

    Returns:
        ax_BPT is modified in place

    """

    # Read in data
    if bptdat is None:
        sdss_file = os.path.join(resource_filename('frb', 'data'), 'Public', 'SDSS', 'SDSS_BPT_stellar_mass.fits')
        if not os.path.isfile(sdss_file):
            print("See the method notes to download the SDSS data!")
            return
        hdulist = fits.open(sdss_file)
        bptdat = hdulist[1].data

    # Select only non zero entries and SNR over 5
    lines = np.array(bptdat.names)[[("flux" in name) & ("err" not in name) for name in bptdat.names]]
    line_err = np.array(bptdat.names)[["flux_err" in name for name in bptdat.names]]
    select = {}
    for line, err in zip(lines, line_err):
        select[line] = bptdat[line] / bptdat[err] >= 5

    # SDSS
    bpt1 = select['oiii_5007_flux'] & select['h_beta_flux'] & select['nii_6584_flux'] & select['h_alpha_flux']
    y = bptdat['oiii_5007_flux'][bpt1] / bptdat['h_beta_flux'][bpt1]
    x = bptdat['nii_6584_flux'][bpt1] / bptdat['h_alpha_flux'][bpt1]

    xbins = 100
    ybins = 100
    # Plot
    counts, xedges, yedges = np.histogram2d(np.log10(x), np.log10(y), bins=(xbins, ybins))
    cm = plt.get_cmap(SDSS_clr)
    mplt = ax_BPT.pcolormesh(xedges, yedges, np.log10(counts.transpose()), cmap=cm)

    # Loop on the Galaxies
    for kk,galaxy in enumerate(galaxies):
        # Parse the emission lines
        NII, NII_err = galaxy.calc_nebular_lum('[NII] 6584')
        Ha, Ha_err = galaxy.calc_nebular_lum('Halpha')
        Hb, Hb_err = galaxy.calc_nebular_lum('Hbeta')
        try:
            OIII, OIII_err = galaxy.calc_nebular_lum('[OIII] 5007')
        except:
            import pdb; pdb.set_trace()
        #
        x0 = (NII/Ha).decompose().value
        y0 = (OIII/Hb).decompose().value
        x0_err = x0 * np.sqrt((NII_err / NII).decompose().value**2 +
                (Ha_err/Ha).decompose().value**2)
        y0_err = y0 * np.sqrt((OIII_err / OIII).decompose().value**2 +
                              (Hb_err/Hb).decompose().value**2)
        # Require at least 20% error
        x0_err = max(x0_err, 0.2*x0)
        y0_err = max(y0_err, 0.2*y0)

        logx, xerr = utils.log_me(x0, x0_err)
        logy, yerr = utils.log_me(y0, y0_err)
        # Upper limit on [NII]/Ha?
        if NII_err.value < 0.:
            xerr = None
            # Left arrow
            plt.arrow(logx, logy, -0.05, 0., fc=clrs[kk], ec=clrs[kk],
                      head_width=0.02, head_length=0.05)
        # Plot
        ax_BPT.errorbar([logx], [logy], xerr=xerr, yerr=yerr,
                        color=clrs[kk], marker=markers[kk], markersize="8",
                        capsize=3, label=galaxy.name)

    # Standard curves
    demarc = lambda x: 0.61 / (x - 0.05) + 1.3  # Kauffman et al 2003, MNRAS, 346, 4, pp. 1055-1077. Eq 1
    demarc_kewley = lambda x: 0.61 / (
                x - 0.47) + 1.19  # Kewley F., Dopita M., Sutherland R., Heisler C., Trevena J., 2001, ApJ, 556,121
    demarc_liner = lambda x: 1.01 * x + 0.48  # Cid Fernandes et al 2010, MNRAS, 403,1036 Eq 10
    ax_BPT.plot(np.linspace(-2, 0), demarc(np.linspace(-2, 0)), "k-", lw=2)#, label="Kauffman et al 2003")
    if show_kewley:
        ax_BPT.plot(np.linspace(-2, 0.25), demarc_kewley(np.linspace(-2, 0.25)), "k--", lw=2)#, label="Kewley et al 2001")
    ax_BPT.plot(np.linspace(-0.43, 0.5), demarc_liner(np.linspace(-0.43, 0.5)), "k--", lw=2)#, label="Cid Fernandes et al 2010")

    # Labels
    lsz = 13.
    ax_BPT.annotate(r"\textbf{Star-forming}", (-1.30, 0), fontsize=lsz)
    ax_BPT.annotate(r"\textbf{LINER}", (0.23, 0), fontsize=lsz)
    ax_BPT.annotate(r"\textbf{Seyfert}", (-0.5, 1), fontsize=lsz)

    # Legend
    if show_legend:
        ax_BPT.legend(loc="lower left")
    # Axes
    ax_BPT.set_xlabel(r"$\log \, ({\rm [N\textsc{ii}]/H\,\alpha)}$")
    ax_BPT.set_ylabel(r"$\log \, ({\rm [O\textsc{iii}]/H\,\beta)}$")
    ax_BPT.set_xlim(-1.5, 0.5)
    ax_BPT.set_ylim(-1, 1.2)
    utils.set_fontsize(ax_BPT, 13.)

def sub_sfms(ax_M, galaxies, clrs, markers, show_legend=True):
    """
    Generate a SF vs. M* plot on top of PRIMUS galaxies

    Args:
        ax_M (matplotlib.axis):
        galaxies (list):
            List of FRB.galaxies.frbgalaxy.FRBGalaxy objects
        clrs (list or str):
            List of matplotlib colors
            if 'auto', sources are colored according to their SFR
        markers (list):
            List of matplotlib marker types
        show_legend (bool, optional):

    """
    utils.set_mplrc()

    # Load up
    primus_zcat = Table.read(os.path.join(primus_path, 'PRIMUS_2013_zcat_v1.fits.gz'))
    primus_mass = Table.read(os.path.join(primus_path, 'PRIMUS_2014_mass_v1.fits.gz'))

    gdz = (primus_zcat['Z'] > 0.2) & (primus_zcat['Z'] < 0.4)
    gd_mag = primus_zcat['SDSS_ABSMAG'][:,0] != 0.

    good_mass = primus_mass['ISGOOD'] == 1

    # PRIMUS
    # Photometry
    gd_color = gdz & gd_mag
    u_r = primus_zcat['SDSS_ABSMAG'][gd_color,0] - primus_zcat['SDSS_ABSMAG'][gd_color,2]
    rmag = primus_zcat['SDSS_ABSMAG'][gd_color,2]

    # Mass/SFR
    gd_msfr = good_mass & gdz
    mass = primus_mass['MASS'][gd_msfr]
    sfr = primus_mass['SFR'][gd_msfr]

    # Plot
    ms = 22.

    # Histogram
    xbins = 50
    ybins = 50
    counts, xedges, yedges = np.histogram2d(mass, sfr, bins=(xbins, ybins))
    #cm = plt.get_cmap('Reds')
    cm = plt.get_cmap('Greys')

    # SF
    mplt = ax_M.pcolormesh(xedges, yedges, counts.transpose(), cmap=cm)

    # Relation
    logm_star = np.linspace(8, 12)
    logsfr = lambda logm_star: -0.49 + 0.65 * (logm_star - 10) + 1.07 * (0.35 - 0.1)
    ax_M.plot(logm_star, logsfr(logm_star), "k--", lw=3)#, label="Moustakas et al 2013")

    # Galaxies
    for kk,galaxy in enumerate(galaxies):
        # M*
        if 'Mstar' in galaxy.derived.keys():
            logM, sig_logM = utils.log_me(galaxy.derived['Mstar'], galaxy.derived['Mstar_err'])
        elif 'Mstar_spec' in galaxy.derived.keys():
            logM, sig_logM = np.log10(galaxy.derived['Mstar_spec']), 0.3
        else:
            continue
        # SFR
        # Try Nebular first
        if 'SFR_nebular' in galaxy.derived.keys():
            if 'SFR_nebular_err' in galaxy.derived.keys():
                logS, sig_logS = utils.log_me(galaxy.derived['SFR_nebular'], galaxy.derived['SFR_nebular_err'])
            else:
                logS, sig_logS = utils.log_me(galaxy.derived['SFR_nebular'], 0.3*galaxy.derived['SFR_nebular'])
        elif 'SFR_photom' in galaxy.derived.keys():
            # Demand 3 sigma
            logS, sig_logS = utils.log_me(galaxy.derived['SFR_photom'], galaxy.derived['SFR_photom_err'])
            if galaxy.derived['SFR_photom'] < 3*galaxy.derived['SFR_photom_err']:
                sig_logS = None
        else:
            continue
        # Color
        if isinstance(clrs, str):
            assert clrs == 'auto'
            if sig_logS is None:
                clr = 'red'
            else:
                clr = 'blue'
        else:
            clr = clrs[kk]
        # Plot
        ax_M.errorbar([logM], [logS], xerr=sig_logM, yerr=sig_logS,
                      color=clr, marker=markers[kk],
                     markersize="12", capsize=3, label=galaxy.name)
        if sig_logS is None:
            # Down arrow
            plt.arrow(logM, logS, 0., -0.2, fc=clr, ec=clr,
                  head_width=0.02*4, head_length=0.05*2)

    ax_M.annotate(r"\textbf{Star forming}", (8.5, 0.8), fontsize=13.)
    ax_M.annotate(r"\textbf{Quiescent}", (11, -0.9), fontsize=13.)
    ax_M.set_xlabel("$\log \, (M_*/M_\odot)$")
    ax_M.set_ylabel("$\log \, SFR (M_\odot$/yr)")
    if show_legend:
        ax_M.legend(loc='lower left')
    ax_M.set_xlim(7.5, 11.8)
    ax_M.set_ylim(-2.4, 1.6)

def sub_color_mag(ax, galaxies, clrs, markers):
    """
    Generate a color-magnitude diagram using PRIMUS data
    and FRB galaxies

    Args:
        ax (matplotlib.Axis):
        galaxies (list):
            List of FRB.galaxies.frbgalaxy.FRBGalaxy objects
        clrs (list):
            List of matplotlib colors
        markers (list):
            List of matplotlib marker types

    Returns:

    """

    # Load up
    primus_zcat = Table.read(os.path.join(primus_path, 'PRIMUS_2013_zcat_v1.fits.gz'))
    #primus_mass = Table.read(os.path.join(primus_path, 'PRIMUS_2014_mass_v1.fits'))

    gdz = (primus_zcat['Z'] > 0.2) & (primus_zcat['Z'] < 0.4)
    gd_mag = primus_zcat['SDSS_ABSMAG'][:,0] != 0.

    # PRIMUS
    # Photometry
    gd_color = gdz & gd_mag
    u_r = primus_zcat['SDSS_ABSMAG'][gd_color,0] - primus_zcat['SDSS_ABSMAG'][gd_color,2]
    rmag = primus_zcat['SDSS_ABSMAG'][gd_color,2]

    xbins = 100
    ybins = 100
    counts, xedges, yedges = np.histogram2d(rmag, u_r, bins=(xbins, ybins))
    cm = plt.get_cmap('Greys')
    mplt = ax.pcolormesh(xedges, yedges, counts.transpose(), cmap=cm)
    '''
    cb = plt.colorbar(mplt, fraction=0.030, pad=0.04)
    cb.set_label('PRIMUS survey')
    '''
    for kk,galaxy in enumerate(galaxies):
        ax.errorbar([galaxy.derived['M_r']], [galaxy.derived['u-r']],
                    xerr=[galaxy.derived['M_r_err']],
                    yerr=[galaxy.derived['u-r_err']],
                    color=clrs[kk], marker=markers[kk],
                    markersize="5", capsize=3,
                    label=galaxy.name)
    # Label
    plt.ylabel(r"$u-r \textbf{(Rest-frame)}$")
    plt.xlabel(r"$r \textbf{(Rest-frame)}$")
    ax.legend(loc='lower right')

    ax.set_ylim(0.0, 3.3)
    ax.set_xlim(-15.5, -23)

    utils.set_fontsize(ax,11.)

<|endoftext|>


<commit_before># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import json
from ... tenant.tenantcontext import *

class CompleteTenantEvent:

    def __init__(self):
        self.tenants = []
        """ :type : list[Tenant]  """
        self.tenant_list_json = None
        """ :type : str  """

    @staticmethod
    def create_from_json(json_str):
        json_obj = json.loads(json_str)
        instance = CompleteTenantEvent()
        instance.tenants = []

        tenants_str = json_obj["tenants"] if "tenants" in json_obj else None
        instance.tenant_list_json = tenants_str
        if tenants_str is not None:
            for tenant_str in tenants_str:
                tenant_obj = Tenant(int(tenant_str["tenantId"]), tenant_str["tenantDomain"])
                for service_name in tenant_str["serviceNameSubscriptionMap"]:
                    sub_str = tenant_str["serviceNameSubscriptionMap"][service_name]
                    sub = Subscription(sub_str["serviceName"], sub_str["clusterIds"])
                    for domain_name in sub_str["subscriptionDomainMap"]:
                        subdomain_str = sub_str["subscriptionDomainMap"][domain_name]
                        sub.add_subscription_domain(domain_name, subdomain_str["applicationContext"])
                    tenant_obj.add_subscription(sub)
                instance.tenants.append(tenant_obj)

        return instance

class TenantSubscribedEvent:

    def __init__(self):
        self.tenant_id = None
        """ :type : int  """
        self.service_name = None
        """ :type : str  """
        self.cluster_ids = None
        """ :type : list[str]  """

    @staticmethod
    def create_from_json(json_str):
        json_obj = json.loads(json_str)
        instance = TenantSubscribedEvent()

        instance.tenant_id = json_obj["tenantId"] if "tenantId" in json_obj else None
        instance.service_name = json_obj["serviceName"] if "serviceName" in json_obj else None
        instance.cluster_ids = json_obj["clusterIds"] if "clusterIds" in json_obj else None

        return instance

class TenantUnsubscribedEvent:

    def __init__(self):
        self.tenant_id = None
        """ :type : int  """
        self.service_name = None
        """ :type : str  """
        self.cluster_ids = None
        """ :type : list[str]  """

    @staticmethod
    def create_from_json(json_str):
        json_obj = json.loads(json_str)
        instance = TenantUnsubscribedEvent()

        instance.tenant_id = json_obj["tenantId"] if "tenantId" in json_obj else None
        instance.service_name = json_obj["serviceName"] if "serviceName" in json_obj else None
        instance.cluster_ids = json_obj["clusterIds"] if "clusterIds" in json_obj else None

        return instance<commit_msg>Fixing Complete Tenant Event for Python CA<commit_after># Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you under the Apache License, Version 2.0 (the
# "License"); you may not use this file except in compliance
# with the License.  You may obtain a copy of the License at
#
#   http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing,
# software distributed under the License is distributed on an
# "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY
# KIND, either express or implied.  See the License for the
# specific language governing permissions and limitations
# under the License.

import json
from ... tenant.tenantcontext import *

class CompleteTenantEvent:

    def __init__(self):
        self.tenants = []
        """ :type : list[Tenant]  """
        self.tenant_list_json = None
        """ :type : str  """

    @staticmethod
    def create_from_json(json_str):
        json_obj = json.loads(json_str)
        instance = CompleteTenantEvent()
        instance.tenants = []

        tenants_str = json_obj["tenants"] if "tenants" in json_obj else None
        instance.tenant_list_json = tenants_str
        if tenants_str is not None:
            for tenant_str in tenants_str:
                tenant_obj = Tenant(int(tenant_str["tenantId"]), tenant_str["tenantDomain"])
                instance.tenants.append(tenant_obj)

        return instance

class TenantSubscribedEvent:

    def __init__(self):
        self.tenant_id = None
        """ :type : int  """
        self.service_name = None
        """ :type : str  """
        self.cluster_ids = None
        """ :type : list[str]  """

    @staticmethod
    def create_from_json(json_str):
        json_obj = json.loads(json_str)
        instance = TenantSubscribedEvent()

        instance.tenant_id = json_obj["tenantId"] if "tenantId" in json_obj else None
        instance.service_name = json_obj["serviceName"] if "serviceName" in json_obj else None
        instance.cluster_ids = json_obj["clusterIds"] if "clusterIds" in json_obj else None

        return instance

class TenantUnsubscribedEvent:

    def __init__(self):
        self.tenant_id = None
        """ :type : int  """
        self.service_name = None
        """ :type : str  """
        self.cluster_ids = None
        """ :type : list[str]  """

    @staticmethod
    def create_from_json(json_str):
        json_obj = json.loads(json_str)
        instance = TenantUnsubscribedEvent()

        instance.tenant_id = json_obj["tenantId"] if "tenantId" in json_obj else None
        instance.service_name = json_obj["serviceName"] if "serviceName" in json_obj else None
        instance.cluster_ids = json_obj["clusterIds"] if "clusterIds" in json_obj else None

        return instance<|endoftext|>


<commit_before>#!/bin/python
##
## generate GCI tasks for updating citations
##

import getopt
import os
import sys

def usage():
  print "\
Usage: citations.py -[hy:p:]\n\
  -h --help           print this help\n\
  -y --year           year to generate\n\
  -p --pages          number of pages in search results for the given year\n\
"

def get_mentor(n):
  mentors = []
  mentors.append("joel_sherrill")
  mentors.append("jenniferaverettoarcorp")
  mentors.append("chrisjohns")
  mentors.append("gedare")
  return mentors[n % len(mentors)]

def generate_tasks(year, pages):
  for p in xrange(pages):
    title = "Update Citations " + str(year) + "-" + str(p)
    description = "This task consists of updating the RTEMS Wiki page located\
 at http://wiki.rtems.org/wiki/index.php/RTEMSReferences according to the\
 instructions given at\
 http://wiki.rtems.org/wiki/index.php/GoogleCodeInProjects#Update_List_of_Citations\
 for the year " + str(year) + " and page " + str(p) + ". You will need to\
 request a user account on the wiki and be careful to avoid making updates\
 that conflict or delete other relevant citations. You will upload a text\
 file containing the entries that you added to the wiki page."
    time_to_complete = "72"
    mentors = get_mentor(year+p)
    type = "Outreach/Research"
    tags = "wiki"

    print(title + ", " + description + ", " + time_to_complete + ", " +
        mentors + ", " + type + ", \"" + tags + "\"\n")

def main():
  # default args
  year = 0
  pages = 0

  # Process args
  try:
    opts, args = getopt.getopt(sys.argv[1:], "hy:p:",
        ["help", "year", "pages"])
  except getopt.GetoptError, err:
    print str(err)
    usage()
    sys.exit(2)
  for opt, arg in opts:
    if opt in ("-h", "--help"):
      usage()
      sys.exit()
    elif opt in ("-y", "--year"):
      year = int(arg)
    elif opt in ("-p", "--pages"):
      pages = int(arg)
    else:
      usage()
      assert False, "unhandled option"

  if year == 0 or pages == 0:
    usage()
    sys.exit(0)

  generate_tasks(year, pages)

if __name__ == "__main__":
  main()
<commit_msg>Remove extra newline between printed task output<commit_after>#!/bin/python
##
## generate GCI tasks for updating citations
##

import getopt
import os
import sys

def usage():
  print "\
Usage: citations.py -[hy:p:]\n\
  -h --help           print this help\n\
  -y --year           year to generate\n\
  -p --pages          number of pages in search results for the given year\n\
"

def get_mentor(n):
  mentors = []
  mentors.append("joel_sherrill")
  mentors.append("jenniferaverettoarcorp")
  mentors.append("chrisjohns")
  mentors.append("gedare")
  return mentors[n % len(mentors)]

def generate_tasks(year, pages):
  for p in xrange(pages):
    title = "Update Citations " + str(year) + "-" + str(p)
    description = "This task consists of updating the RTEMS Wiki page located\
 at http://wiki.rtems.org/wiki/index.php/RTEMSReferences according to the\
 instructions given at\
 http://wiki.rtems.org/wiki/index.php/GoogleCodeInProjects#Update_List_of_Citations\
 for the year " + str(year) + " and page " + str(p) + ". You will need to\
 request a user account on the wiki and be careful to avoid making updates\
 that conflict or delete other relevant citations. You will upload a text\
 file containing the entries that you added to the wiki page."
    time_to_complete = "72"
    mentors = get_mentor(year+p)
    type = "Outreach/Research"
    tags = "wiki"

    print(title + ", " + description + ", " + time_to_complete + ", " +
        mentors + ", " + type + ", \"" + tags + "\"")

def main():
  # default args
  year = 0
  pages = 0

  # Process args
  try:
    opts, args = getopt.getopt(sys.argv[1:], "hy:p:",
        ["help", "year", "pages"])
  except getopt.GetoptError, err:
    print str(err)
    usage()
    sys.exit(2)
  for opt, arg in opts:
    if opt in ("-h", "--help"):
      usage()
      sys.exit()
    elif opt in ("-y", "--year"):
      year = int(arg)
    elif opt in ("-p", "--pages"):
      pages = int(arg)
    else:
      usage()
      assert False, "unhandled option"

  if year == 0 or pages == 0:
    usage()
    sys.exit(0)

  generate_tasks(year, pages)

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


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

import argparse
import cPickle
import traceback
import logging
import time
import sys

import numpy

import experiments.rnnencdec
from experiments.rnnencdec import RNNEncoderDecoder, parse_input

logger = logging.getLogger(__name__)

class Timer(object):

    def __init__(self):
        self.total = 0

    def start(self):
        self.start_time = time.time()

    def finish(self):
        self.total += time.time() - self.start_time

class BeamSearch(object):

    def __init__(self, enc_dec):
        self.enc_dec = enc_dec
        state = self.enc_dec.state
        self.eos_id = state['null_sym_target']
        self.unk_id = state['unk_sym_target']

    def compile(self):
        self.comp_repr = self.enc_dec.create_representation_computer()
        self.comp_init_states = self.enc_dec.create_initializers()
        self.comp_next_probs = self.enc_dec.create_next_probs_computer()
        self.comp_next_states = self.enc_dec.create_next_states_computer()

    def search(self, seq, n_samples, ignore_unk=False, minlen=1):
        c = self.comp_repr(seq)[0]
        states = map(lambda x : x[None, :], self.comp_init_states(c))
        dim = states[0].shape[1]

        num_levels = len(states)

        fin_trans = []
        fin_costs = []

        trans = [[]]
        costs = [0.0]

        for k in range(3 * len(seq)):
            if n_samples == 0:
                break
            beam_size = len(trans)
            last_words = (numpy.array(map(lambda t : t[-1], trans))
                    if k > 0
                    else numpy.zeros(beam_size, dtype="int64"))
            log_probs = numpy.log(self.comp_next_probs(c, k, last_words, *states)[0])

            # Adjust log probs according to search restrictions
            if ignore_unk:
                log_probs[:,self.unk_id] = -numpy.inf
            if k < minlen:
                log_probs[:,self.eos_id] = -numpy.inf

            # Find the best options by calling argpartition of flatten array
            next_costs = numpy.array(costs)[:, None] - log_probs
            flat_next_costs = next_costs.flatten()
            best_costs_indices = numpy.argpartition(
                    flat_next_costs.flatten(),
                    n_samples)[:n_samples]

            # Decypher flatten indices
            voc_size = log_probs.shape[1]
            trans_indices = best_costs_indices / voc_size
            word_indices = best_costs_indices % voc_size
            costs = flat_next_costs[best_costs_indices]

            new_trans = [[]] * n_samples
            new_costs = numpy.zeros(n_samples)
            new_states = [numpy.zeros((n_samples, dim), dtype="float32") for level
                    in range(num_levels)]
            inputs = numpy.zeros(n_samples, dtype="int64")
            for i, (orig_idx, next_word, next_cost) in enumerate(
                    zip(trans_indices, word_indices, costs)):
                new_trans[i] = trans[orig_idx] + [next_word]
                new_costs[i] = next_cost
                for level in range(num_levels):
                    new_states[level][i] = states[level][orig_idx]
                inputs[i] = next_word

            new_states = self.comp_next_states(c, k, inputs, *new_states)

            trans = []
            costs = []
            indices = []
            for i in range(n_samples):
                if new_trans[i][-1] != self.enc_dec.state['null_sym_target']:
                    trans.append(new_trans[i])
                    costs.append(new_costs[i])
                    indices.append(i)
                else:
                    n_samples -= 1
                    fin_trans.append(new_trans[i])
                    fin_costs.append(new_costs[i])
            states = map(lambda x : x[indices], new_states)

        # Dirty tricks to obtain any translation
        if not len(fin_trans):
            if ignore_unk:
                logger.warning("Did not manage without UNK")
                return self.search(seq, n_samples, False, minlen)
            elif n_samples < 500:
                logger.warning("Still no translations: try beam size {}".format(n_samples * 2))
                return self.search(seq, n_samples * 2, False, minlen)
            else:
                logger.error("Translation failed")

        fin_trans = numpy.array(fin_trans)[numpy.argsort(fin_costs)]
        fin_costs = numpy.array(sorted(fin_costs))
        return fin_trans, fin_costs

def indices_to_words(i2w, seq):
    sen = []
    for k in xrange(len(seq)):
        if i2w[seq[k]] == '<eol>':
            break
        sen.append(i2w[seq[k]])
    return sen

def sample(lm_model, seq, n_samples,
        sampler=None, beam_search=None,
        ignore_unk=False, normalize=False,
        alpha=1, verbose=False):
    if beam_search:
        sentences = []
        trans, costs = beam_search.search(seq, n_samples,
                ignore_unk=ignore_unk, minlen=len(seq) / 2)
        if normalize:
            counts = [len(s) for s in trans]
            costs = [co / cn for co, cn in zip(costs, counts)]
        for i in range(len(trans)):
            sen = indices_to_words(lm_model.word_indxs, trans[i])
            sentences.append(" ".join(sen))
        for i in range(len(costs)):
            if verbose:
                print "{}: {}".format(costs[i], sentences[i])
        return sentences, costs, trans
    elif sampler:
        sentences = []
        all_probs = []
        costs = []

        values, cond_probs = sampler(n_samples, 3 * (len(seq) - 1), alpha, seq)
        for sidx in xrange(n_samples):
            sen = []
            for k in xrange(values.shape[0]):
                if lm_model.word_indxs[values[k, sidx]] == '<eol>':
                    break
                sen.append(lm_model.word_indxs[values[k, sidx]])
            sentences.append(" ".join(sen))
            probs = cond_probs[:, sidx]
            probs = numpy.array(cond_probs[:len(sen) + 1, sidx])
            all_probs.append(numpy.exp(-probs))
            costs.append(-numpy.sum(probs))
        if normalize:
            counts = [len(s.strip().split(" ")) for s in sentences]
            costs = [co / cn for co, cn in zip(costs, counts)]
        sprobs = numpy.argsort(costs)
        if verbose:
            for pidx in sprobs:
                print "{}: {} {} {}".format(pidx, -costs[pidx], all_probs[pidx], sentences[pidx])
            print
        return sentences, costs
    else:
        raise Exception("I don't know what to do")

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--state", help="State to use")
    parser.add_argument("--state-fn", help="Initialization function for state", default="prototype_state")
    parser.add_argument("--beam-search", help="Beam size, turns on beam-search", type=int)
    parser.add_argument("--ignore-unk", default=False, action="store_true",
            help="Ignore unknown words")
    parser.add_argument("--source", help="File of source sentences", default="")
    parser.add_argument("--trans", help="File to save translations in", default="")
    parser.add_argument("--normalize", help="Normalize log-prob with the word count",
        action="store_true", default=False)
    parser.add_argument("--verbose", action="store_true",
        default=False, help="Be verbose")
    parser.add_argument("model_path", help="Path to the model")
    parser.add_argument("changes",  nargs="?", help="Changes to state", default="")
    return parser.parse_args()

def main():
    args = parse_args()

    state = getattr(experiments.rnnencdec, args.state_fn)()
    if hasattr(args, 'state') and args.state:
        with open(args.state) as src:
            state.update(cPickle.load(src))
    state.update(eval("dict({})".format(args.changes)))

    logging.basicConfig(level=getattr(logging, state['level']), format="%(asctime)s: %(name)s: %(levelname)s: %(message)s")

    rng = numpy.random.RandomState(state['seed'])
    enc_dec = RNNEncoderDecoder(state, rng, skip_init=True)
    enc_dec.build()
    lm_model = enc_dec.create_lm_model()
    lm_model.load(args.model_path)
    indx_word = cPickle.load(open(state['word_indx'],'rb'))

    sampler = None
    beam_search = None
    if args.beam_search:
        beam_search = BeamSearch(enc_dec)
        beam_search.compile()
    else:
        sampler = enc_dec.create_sampler(many_samples=True)

    idict_src = cPickle.load(open(state['indx_word'],'r'))

    if args.source != "" and args.trans != "":
        # Actually only beam search is currently supported here
        assert beam_search

        fsrc = open(args.source, 'r')
        ftrans = open(args.trans, 'w')

        start_time = time.time()

        n_samples = args.beam_search
        total_cost = 0.0
        logging.debug("Beam size: {}".format(n_samples))
        for i, line in enumerate(fsrc):
            seqin = line.strip()
            seq,parsed_in = parse_input(state, indx_word, seqin, idx2word=idict_src)
            if args.verbose:
                print "Parsed Input:", parsed_in
            trans, costs, _ = sample(lm_model, seq, n_samples, sampler=sampler,
                    beam_search=beam_search, ignore_unk=args.ignore_unk, normalize=args.normalize)
            best = numpy.argmin(costs)
            print >>ftrans, trans[best]
            if args.verbose:
                print "Translation:", trans[best]
            total_cost += costs[best]
            if (i + 1)  % 100 == 0:
                ftrans.flush()
                logger.debug("Current speed is {} per sentence".
                        format((time.time() - start_time) / (i + 1)))
        print "Total cost of the translations: {}".format(total_cost)

        fsrc.close()
        ftrans.close()
    else:
        while True:
            try:
                seqin = raw_input('Input Sequence: ')
                n_samples = int(raw_input('How many samples? '))
                alpha = None
                if not args.beam_search:
                    alpha = float(raw_input('Inverse Temperature? '))
                seq,parsed_in = parse_input(state, indx_word, seqin, idx2word=idict_src)
                print "Parsed Input:", parsed_in
            except Exception:
                print "Exception while parsing your input:"
                traceback.print_exc()
                continue

            sample(lm_model, seq, n_samples, sampler=sampler,
                    beam_search=beam_search,
                    ignore_unk=args.ignore_unk, normalize=args.normalize,
                    alpha=alpha, verbose=True)

if __name__ == "__main__":
    main()
<commit_msg>sample return three arguments even when beamsearch is off<commit_after>#!/usr/bin/env python

import argparse
import cPickle
import traceback
import logging
import time
import sys

import numpy

import experiments.rnnencdec
from experiments.rnnencdec import RNNEncoderDecoder, parse_input

logger = logging.getLogger(__name__)

class Timer(object):

    def __init__(self):
        self.total = 0

    def start(self):
        self.start_time = time.time()

    def finish(self):
        self.total += time.time() - self.start_time

class BeamSearch(object):

    def __init__(self, enc_dec):
        self.enc_dec = enc_dec
        state = self.enc_dec.state
        self.eos_id = state['null_sym_target']
        self.unk_id = state['unk_sym_target']

    def compile(self):
        self.comp_repr = self.enc_dec.create_representation_computer()
        self.comp_init_states = self.enc_dec.create_initializers()
        self.comp_next_probs = self.enc_dec.create_next_probs_computer()
        self.comp_next_states = self.enc_dec.create_next_states_computer()

    def search(self, seq, n_samples, ignore_unk=False, minlen=1):
        c = self.comp_repr(seq)[0]
        states = map(lambda x : x[None, :], self.comp_init_states(c))
        dim = states[0].shape[1]

        num_levels = len(states)

        fin_trans = []
        fin_costs = []

        trans = [[]]
        costs = [0.0]

        for k in range(3 * len(seq)):
            if n_samples == 0:
                break
            beam_size = len(trans)
            last_words = (numpy.array(map(lambda t : t[-1], trans))
                    if k > 0
                    else numpy.zeros(beam_size, dtype="int64"))
            log_probs = numpy.log(self.comp_next_probs(c, k, last_words, *states)[0])

            # Adjust log probs according to search restrictions
            if ignore_unk:
                log_probs[:,self.unk_id] = -numpy.inf
            if k < minlen:
                log_probs[:,self.eos_id] = -numpy.inf

            # Find the best options by calling argpartition of flatten array
            next_costs = numpy.array(costs)[:, None] - log_probs
            flat_next_costs = next_costs.flatten()
            best_costs_indices = numpy.argpartition(
                    flat_next_costs.flatten(),
                    n_samples)[:n_samples]

            # Decypher flatten indices
            voc_size = log_probs.shape[1]
            trans_indices = best_costs_indices / voc_size
            word_indices = best_costs_indices % voc_size
            costs = flat_next_costs[best_costs_indices]

            new_trans = [[]] * n_samples
            new_costs = numpy.zeros(n_samples)
            new_states = [numpy.zeros((n_samples, dim), dtype="float32") for level
                    in range(num_levels)]
            inputs = numpy.zeros(n_samples, dtype="int64")
            for i, (orig_idx, next_word, next_cost) in enumerate(
                    zip(trans_indices, word_indices, costs)):
                new_trans[i] = trans[orig_idx] + [next_word]
                new_costs[i] = next_cost
                for level in range(num_levels):
                    new_states[level][i] = states[level][orig_idx]
                inputs[i] = next_word

            new_states = self.comp_next_states(c, k, inputs, *new_states)

            trans = []
            costs = []
            indices = []
            for i in range(n_samples):
                if new_trans[i][-1] != self.enc_dec.state['null_sym_target']:
                    trans.append(new_trans[i])
                    costs.append(new_costs[i])
                    indices.append(i)
                else:
                    n_samples -= 1
                    fin_trans.append(new_trans[i])
                    fin_costs.append(new_costs[i])
            states = map(lambda x : x[indices], new_states)

        # Dirty tricks to obtain any translation
        if not len(fin_trans):
            if ignore_unk:
                logger.warning("Did not manage without UNK")
                return self.search(seq, n_samples, False, minlen)
            elif n_samples < 500:
                logger.warning("Still no translations: try beam size {}".format(n_samples * 2))
                return self.search(seq, n_samples * 2, False, minlen)
            else:
                logger.error("Translation failed")

        fin_trans = numpy.array(fin_trans)[numpy.argsort(fin_costs)]
        fin_costs = numpy.array(sorted(fin_costs))
        return fin_trans, fin_costs

def indices_to_words(i2w, seq):
    sen = []
    for k in xrange(len(seq)):
        if i2w[seq[k]] == '<eol>':
            break
        sen.append(i2w[seq[k]])
    return sen

def sample(lm_model, seq, n_samples,
        sampler=None, beam_search=None,
        ignore_unk=False, normalize=False,
        alpha=1, verbose=False):
    if beam_search:
        sentences = []
        trans, costs = beam_search.search(seq, n_samples,
                ignore_unk=ignore_unk, minlen=len(seq) / 2)
        if normalize:
            counts = [len(s) for s in trans]
            costs = [co / cn for co, cn in zip(costs, counts)]
        for i in range(len(trans)):
            sen = indices_to_words(lm_model.word_indxs, trans[i])
            sentences.append(" ".join(sen))
        for i in range(len(costs)):
            if verbose:
                print "{}: {}".format(costs[i], sentences[i])
        return sentences, costs, trans
    elif sampler:
        sentences = []
        all_probs = []
        costs = []

        values, cond_probs = sampler(n_samples, 3 * (len(seq) - 1), alpha, seq)
        for sidx in xrange(n_samples):
            sen = []
            for k in xrange(values.shape[0]):
                if lm_model.word_indxs[values[k, sidx]] == '<eol>':
                    break
                sen.append(lm_model.word_indxs[values[k, sidx]])
            sentences.append(" ".join(sen))
            probs = cond_probs[:, sidx]
            probs = numpy.array(cond_probs[:len(sen) + 1, sidx])
            all_probs.append(numpy.exp(-probs))
            costs.append(-numpy.sum(probs))
        if normalize:
            counts = [len(s.strip().split(" ")) for s in sentences]
            costs = [co / cn for co, cn in zip(costs, counts)]
        sprobs = numpy.argsort(costs)
        if verbose:
            for pidx in sprobs:
                print "{}: {} {} {}".format(pidx, -costs[pidx], all_probs[pidx], sentences[pidx])
            print
        return sentences, costs, None
    else:
        raise Exception("I don't know what to do")

def parse_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--state", help="State to use")
    parser.add_argument("--state-fn", help="Initialization function for state", default="prototype_state")
    parser.add_argument("--beam-search", help="Beam size, turns on beam-search", type=int)
    parser.add_argument("--ignore-unk", default=False, action="store_true",
            help="Ignore unknown words")
    parser.add_argument("--source", help="File of source sentences", default="")
    parser.add_argument("--trans", help="File to save translations in", default="")
    parser.add_argument("--normalize", help="Normalize log-prob with the word count",
        action="store_true", default=False)
    parser.add_argument("--verbose", action="store_true",
        default=False, help="Be verbose")
    parser.add_argument("model_path", help="Path to the model")
    parser.add_argument("changes",  nargs="?", help="Changes to state", default="")
    return parser.parse_args()

def main():
    args = parse_args()

    state = getattr(experiments.rnnencdec, args.state_fn)()
    if hasattr(args, 'state') and args.state:
        with open(args.state) as src:
            state.update(cPickle.load(src))
    state.update(eval("dict({})".format(args.changes)))

    logging.basicConfig(level=getattr(logging, state['level']), format="%(asctime)s: %(name)s: %(levelname)s: %(message)s")

    rng = numpy.random.RandomState(state['seed'])
    enc_dec = RNNEncoderDecoder(state, rng, skip_init=True)
    enc_dec.build()
    lm_model = enc_dec.create_lm_model()
    lm_model.load(args.model_path)
    indx_word = cPickle.load(open(state['word_indx'],'rb'))

    sampler = None
    beam_search = None
    if args.beam_search:
        beam_search = BeamSearch(enc_dec)
        beam_search.compile()
    else:
        sampler = enc_dec.create_sampler(many_samples=True)

    idict_src = cPickle.load(open(state['indx_word'],'r'))

    if args.source != "" and args.trans != "":
        # Actually only beam search is currently supported here
        assert beam_search

        fsrc = open(args.source, 'r')
        ftrans = open(args.trans, 'w')

        start_time = time.time()

        n_samples = args.beam_search
        total_cost = 0.0
        logging.debug("Beam size: {}".format(n_samples))
        for i, line in enumerate(fsrc):
            seqin = line.strip()
            seq,parsed_in = parse_input(state, indx_word, seqin, idx2word=idict_src)
            if args.verbose:
                print "Parsed Input:", parsed_in
            trans, costs, _ = sample(lm_model, seq, n_samples, sampler=sampler,
                    beam_search=beam_search, ignore_unk=args.ignore_unk, normalize=args.normalize)
            best = numpy.argmin(costs)
            print >>ftrans, trans[best]
            if args.verbose:
                print "Translation:", trans[best]
            total_cost += costs[best]
            if (i + 1)  % 100 == 0:
                ftrans.flush()
                logger.debug("Current speed is {} per sentence".
                        format((time.time() - start_time) / (i + 1)))
        print "Total cost of the translations: {}".format(total_cost)

        fsrc.close()
        ftrans.close()
    else:
        while True:
            try:
                seqin = raw_input('Input Sequence: ')
                n_samples = int(raw_input('How many samples? '))
                alpha = None
                if not args.beam_search:
                    alpha = float(raw_input('Inverse Temperature? '))
                seq,parsed_in = parse_input(state, indx_word, seqin, idx2word=idict_src)
                print "Parsed Input:", parsed_in
            except Exception:
                print "Exception while parsing your input:"
                traceback.print_exc()
                continue

            sample(lm_model, seq, n_samples, sampler=sampler,
                    beam_search=beam_search,
                    ignore_unk=args.ignore_unk, normalize=args.normalize,
                    alpha=alpha, verbose=True)

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


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

# This file is part of Archivematica.
#
# Copyright 2010-2013 Artefactual Systems Inc. <http://artefactual.com>
#
# Archivematica is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Archivematica is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Archivematica.  If not, see <http://www.gnu.org/licenses/>.

# @package Archivematica
# @subpackage MCPServer
# @author Joseph Perry <joseph@artefactual.com>

# ~DOC~
#
# --- This is the MCP (master control program) ---
# The intention of this program is to provide a centralized automated distributed system for performing an arbitrary set of tasks on a directory.
# Distributed in that the work can be performed on more than one physical computer simultaneously.
# Centralized in that there is one centre point for configuring flow through the system.
# Automated in that the tasks performed will be based on the config files and instantiated for each of the targets.
#
# It loads configurations from the database.
#
# stdlib, alphabetical by import source
import logging
import logging.config
import getpass
import os
import signal
import sys
import threading
import time
import uuid

import django
django.setup()

from django.conf import settings as django_settings
from django.db.models import Q

# This project, alphabetical by import source
import watchDirectory
from utils import log_exceptions

from jobChain import jobChain
from unitSIP import unitSIP
from unitDIP import unitDIP
from unitFile import unitFile
from unitTransfer import unitTransfer
from utils import isUUID
import RPCServer

from django_mysqlpool import auto_close_db
import databaseFunctions
from archivematicaFunctions import unicodeToStr
import dicts

from main.models import Job, SIP, Task, WatchedDirectory

countOfCreateUnitAndJobChainThreaded = 0

# time to sleep to allow db to be updated with the new location of a SIP
dbWaitSleep = 2

stopSignalReceived = False  # Tracks whether a sigkill has been received or not

def fetchUUIDFromPath(path):
    # find UUID on end of SIP path
    uuidLen = -36
    if isUUID(path[uuidLen - 1:-1]):
        return path[uuidLen - 1:-1]

def findOrCreateSipInDB(path, waitSleep=dbWaitSleep, unit_type='SIP'):
    """Matches a directory to a database sip by it's appended UUID, or path. If it doesn't find one, it will create one"""
    path = path.replace(django_settings.SHARED_DIRECTORY, "%sharedPath%", 1)

    query = Q(currentpath=path)

    # Find UUID on end of SIP path
    UUID = fetchUUIDFromPath(path)
    sip = None
    if UUID:
        query = query | Q(uuid=UUID)

    sips = SIP.objects.filter(query)
    count = sips.count()
    if count > 1:
        # This might have happened because the UUID at the end of the directory
        # name corresponds to a different SIP in the database.
        # Try refiltering the queryset on path alone, and see if that brought us
        # down to a single SIP.
        sips = sips.filter(currentpath=path)
        count = sips.count()

        # Darn: we must have multiple SIPs with the same path in the database.
        # We have no reasonable way to recover from this condition.
        if count > 1:
            logger.error('More than one SIP for path %s and/or UUID %s, using first result', path, UUID)
    if count > 0:
        sip = sips[0]
        UUID = sip.uuid
        logger.info('Using existing SIP %s at %s', UUID, path)
    else:
        logger.info('Not using existing SIP %s at %s', UUID, path)

    if sip is None:
        # Create it
        # Note that if UUID is None here, a new UUID will be generated
        # and returned by the function; otherwise it returns the
        # value that was passed in.
        UUID = databaseFunctions.createSIP(path, UUID=UUID)
        logger.info('Creating SIP %s at %s', UUID, path)
    else:
        current_path = sip.currentpath
        if current_path != path and unit_type == 'SIP':
            # Ensure path provided matches path in DB
            sip.currentpath = path
            sip.save()

    return UUID

@log_exceptions
@auto_close_db
def createUnitAndJobChain(path, config, terminate=False):
    path = unicodeToStr(path)
    if os.path.isdir(path):
        path = path + "/"
    logger.debug('Creating unit and job chain for %s with %s', path, config)
    unit = None
    if os.path.isdir(path):
        if config[3] == "SIP":
            UUID = findOrCreateSipInDB(path)
            unit = unitSIP(path, UUID)
        elif config[3] == "DIP":
            UUID = findOrCreateSipInDB(path, unit_type='DIP')
            unit = unitDIP(path, UUID)
        elif config[3] == "Transfer":
            unit = unitTransfer(path)
    elif os.path.isfile(path):
        if config[3] == "Transfer":
            unit = unitTransfer(path)
        else:
            return
            UUID = uuid.uuid4()
            unit = unitFile(path, UUID)
    else:
        return
    jobChain(unit, config[1])

    if terminate:
        exit(0)

def createUnitAndJobChainThreaded(path, config, terminate=True):
    global countOfCreateUnitAndJobChainThreaded
    try:
        logger.debug('Watching path %s', path)
        t = threading.Thread(target=createUnitAndJobChain, args=(path, config), kwargs={"terminate": terminate})
        t.daemon = True
        countOfCreateUnitAndJobChainThreaded += 1
        while(django_settings.LIMIT_TASK_THREADS <= threading.activeCount() + django_settings.RESERVED_AS_TASK_PROCESSING_THREADS):
            if stopSignalReceived:
                logger.info('Signal was received; stopping createUnitAndJobChainThreaded(path, config)')
                exit(0)
            logger.debug('Active thread count: %s', threading.activeCount())
            time.sleep(.5)
        countOfCreateUnitAndJobChainThreaded -= 1
        t.start()
    except Exception:
        logger.exception('Error creating threads to watch directories')

def watchDirectories():
    """Start watching the watched directories defined in the WatchedDirectories table in the database."""
    watched_dir_path = django_settings.WATCH_DIRECTORY
    interval = django_settings.WATCH_DIRECTORY_INTERVAL

    watched_directories = WatchedDirectory.objects.all()

    for watched_directory in watched_directories:
        directory = watched_directory.watched_directory_path.replace("%watchDirectoryPath%", watched_dir_path, 1)

        # Tuple of variables that may be used by a callback
        row = (watched_directory.watched_directory_path, watched_directory.chain_id, watched_directory.only_act_on_directories, watched_directory.expected_type.description)

        if not os.path.isdir(directory):
            os.makedirs(directory)
        for item in os.listdir(directory):
            if item == ".gitignore":
                continue
            item = item.decode("utf-8")
            path = os.path.join(unicode(directory), item)
            while(django_settings.LIMIT_TASK_THREADS <= threading.activeCount() + django_settings.RESERVED_AS_TASK_PROCESSING_THREADS):
                time.sleep(1)
            createUnitAndJobChainThreaded(path, row, terminate=False)
        actOnFiles = True
        if watched_directory.only_act_on_directories:
            actOnFiles = False
        watchDirectory.archivematicaWatchDirectory(
            directory,
            variablesAdded=row,
            callBackFunctionAdded=createUnitAndJobChainThreaded,
            alertOnFiles=actOnFiles,
            interval=interval,
        )

def signal_handler(signalReceived, frame):
    """Used to handle the stop/kill command signals (SIGKILL)"""
    logger.info('Recieved signal %s in frame %s', signalReceived, frame)
    global stopSignalReceived
    stopSignalReceived = True
    threads = threading.enumerate()
    for thread in threads:
        logger.warning('Not stopping %s %s', type(thread), thread)
    sys.stdout.flush()
    sys.stderr.flush()
    sys.exit(0)
    exit(0)

@log_exceptions
@auto_close_db
def debugMonitor():
    """Periodically prints out status of MCP, including whether the database lock is locked, thread count, etc."""
    global countOfCreateUnitAndJobChainThreaded
    while True:
        logger.debug('Debug monitor: datetime: %s', databaseFunctions.getUTCDate())
        logger.debug('Debug monitor: thread count: %s', threading.activeCount())
        logger.debug('Debug monitor: created job chain threaded: %s', countOfCreateUnitAndJobChainThreaded)
        time.sleep(3600)

@log_exceptions
@auto_close_db
def flushOutputs():
    while True:
        sys.stdout.flush()
        sys.stderr.flush()
        time.sleep(5)

def cleanupOldDbEntriesOnNewRun():
    Job.objects.filter(currentstep=Job.STATUS_AWAITING_DECISION).delete()
    Job.objects.filter(currentstep=Job.STATUS_EXECUTING_COMMANDS).update(currentstep=Job.STATUS_FAILED)
    Task.objects.filter(exitcode=None).update(exitcode=-1, stderror="MCP shut down while processing.")

def created_shared_directory_structure():
    dirs = (
        "arrange",
        "completed",
        "completed/transfers",
        "currentlyProcessing",
        "DIPbackups",
        "failed",
        "rejected",
        "sharedMicroServiceTasksConfigs",
        "sharedMicroServiceTasksConfigs/createXmlEventsAssist",
        "sharedMicroServiceTasksConfigs/generateAIP",
        "sharedMicroServiceTasksConfigs/generateAIP/bagit",
        "sharedMicroServiceTasksConfigs/processingMCPConfigs",
        "sharedMicroServiceTasksConfigs/transcoder",
        "sharedMicroServiceTasksConfigs/transcoder/defaultIcons",
        "SIPbackups",
        "tmp",
        "watchedDirectories",
        "watchedDirectories/activeTransfers",
        "watchedDirectories/activeTransfers/baggitDirectory",
        "watchedDirectories/activeTransfers/baggitZippedDirectory",
        "watchedDirectories/activeTransfers/Dspace",
        "watchedDirectories/activeTransfers/maildir",
        "watchedDirectories/activeTransfers/standardTransfer",
        "watchedDirectories/activeTransfers/TRIM",
        "watchedDirectories/approveNormalization",
        "watchedDirectories/approveSubmissionDocumentationIngest",
        "watchedDirectories/quarantined",
        "watchedDirectories/SIPCreation",
        "watchedDirectories/SIPCreation/completedTransfers",
        "watchedDirectories/SIPCreation/SIPsUnderConstruction",
        "watchedDirectories/storeAIP",
        "watchedDirectories/system",
        "watchedDirectories/system/autoProcessSIP",
        "watchedDirectories/system/autoRestructureForCompliance",
        "watchedDirectories/system/createAIC",
        "watchedDirectories/system/reingestAIP",
        "watchedDirectories/uploadDIP",
        "watchedDirectories/uploadedDIPs",
        "watchedDirectories/workFlowDecisions",
        "watchedDirectories/workFlowDecisions/compressionAIPDecisions",
        "watchedDirectories/workFlowDecisions/createDip",
        "watchedDirectories/workFlowDecisions/createTree",
        "watchedDirectories/workFlowDecisions/examineContentsChoice",
        "watchedDirectories/workFlowDecisions/extractPackagesChoice",
        "watchedDirectories/workFlowDecisions/metadataReminder",
        "watchedDirectories/workFlowDecisions/quarantineTransfer",
        "watchedDirectories/workFlowDecisions/selectFormatIDToolIngest",
        "watchedDirectories/workFlowDecisions/selectFormatIDToolTransfer",
        "www",
        "www/AIPsStore",
        "www/AIPsStore/transferBacklog",
        "www/AIPsStore/transferBacklog/arrange",
        "www/AIPsStore/transferBacklog/originals",
        "www/DIPsStore"
    )
    for d in dirs:
        d = os.path.join(django_settings.SHARED_DIRECTORY, d)
        if os.path.isdir(d):
            continue
        logger.info('Creating directory: %s', d)
        os.makedirs(d, mode=770)

def _except_hook_log_everything(exc_type, exc_value, exc_traceback):
    """
    Replacement for default exception handler that logs exceptions.
    """
    # Reference http://stackoverflow.com/a/16993115/2475775
    logger.error("Uncaught exception", exc_info=(exc_type, exc_value, exc_traceback))
    sys.__excepthook__(exc_type, exc_value, exc_traceback)

if __name__ == '__main__':
    logger = logging.getLogger("archivematica.mcp.server")

    # Replace exception handler with one that logs exceptions
    sys.excepthook = _except_hook_log_everything

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    logger.info('This PID: %s', os.getpid())
    logger.info('User: %s', getpass.getuser())

    dicts.setup(
        shared_directory=django_settings.SHARED_DIRECTORY,
        processing_directory=django_settings.PROCESSING_DIRECTORY,
        watch_directory=django_settings.WATCH_DIRECTORY,
        rejected_directory=django_settings.REJECTED_DIRECTORY,
    )

    created_shared_directory_structure()

    t = threading.Thread(target=debugMonitor)
    t.daemon = True
    t.start()

    t = threading.Thread(target=flushOutputs)
    t.daemon = True
    t.start()
    cleanupOldDbEntriesOnNewRun()
    watchDirectories()

    # This is blocking the main thread with the worker loop
    RPCServer.startRPCServer()
<commit_msg>MCPServer: populate defaultProcessingMCP.xml, see #8<commit_after>#!/usr/bin/env python2

# This file is part of Archivematica.
#
# Copyright 2010-2013 Artefactual Systems Inc. <http://artefactual.com>
#
# Archivematica is free software: you can redistribute it and/or modify
# it under the terms of the GNU Affero General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# Archivematica is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with Archivematica.  If not, see <http://www.gnu.org/licenses/>.

# @package Archivematica
# @subpackage MCPServer
# @author Joseph Perry <joseph@artefactual.com>

# ~DOC~
#
# --- This is the MCP (master control program) ---
# The intention of this program is to provide a centralized automated distributed system for performing an arbitrary set of tasks on a directory.
# Distributed in that the work can be performed on more than one physical computer simultaneously.
# Centralized in that there is one centre point for configuring flow through the system.
# Automated in that the tasks performed will be based on the config files and instantiated for each of the targets.
#
# It loads configurations from the database.
#
# stdlib, alphabetical by import source
import logging
import logging.config
import getpass
import os
import signal
import sys
import threading
import time
import uuid

import django
django.setup()

from django.conf import settings as django_settings
from django.db.models import Q

# This project, alphabetical by import source
import watchDirectory
from utils import log_exceptions

from jobChain import jobChain
from unitSIP import unitSIP
from unitDIP import unitDIP
from unitFile import unitFile
from unitTransfer import unitTransfer
from utils import isUUID
import RPCServer

from django_mysqlpool import auto_close_db
import databaseFunctions
from archivematicaFunctions import unicodeToStr
import dicts

from main.models import Job, SIP, Task, WatchedDirectory

countOfCreateUnitAndJobChainThreaded = 0

# time to sleep to allow db to be updated with the new location of a SIP
dbWaitSleep = 2

stopSignalReceived = False  # Tracks whether a sigkill has been received or not

DEFAULT_PROCESSING_CONFIG = u"""<processingMCP>
  <preconfiguredChoices>
    <!-- Send to quarantine? -->
    <preconfiguredChoice>
      <appliesTo>755b4177-c587-41a7-8c52-015277568302</appliesTo>
      <goToChain>d4404ab1-dc7f-4e9e-b1f8-aa861e766b8e</goToChain>
    </preconfiguredChoice>
    <!-- Display metadata reminder -->
    <preconfiguredChoice>
      <appliesTo>eeb23509-57e2-4529-8857-9d62525db048</appliesTo>
      <goToChain>5727faac-88af-40e8-8c10-268644b0142d</goToChain>
    </preconfiguredChoice>
    <!-- Remove from quarantine -->
    <preconfiguredChoice>
      <appliesTo>19adb668-b19a-4fcb-8938-f49d7485eaf3</appliesTo>
      <goToChain>333643b7-122a-4019-8bef-996443f3ecc5</goToChain>
      <delay unitCtime="yes">2419200.0</delay>
    </preconfiguredChoice>
    <!-- Extract packages -->
    <preconfiguredChoice>
      <appliesTo>dec97e3c-5598-4b99-b26e-f87a435a6b7f</appliesTo>
      <goToChain>01d80b27-4ad1-4bd1-8f8d-f819f18bf685</goToChain>
    </preconfiguredChoice>
    <!-- Delete extracted packages -->
    <preconfiguredChoice>
      <appliesTo>f19926dd-8fb5-4c79-8ade-c83f61f55b40</appliesTo>
      <goToChain>85b1e45d-8f98-4cae-8336-72f40e12cbef</goToChain>
    </preconfiguredChoice>
    <!-- Select pre-normalize file format identification command -->
    <preconfiguredChoice>
      <appliesTo>7a024896-c4f7-4808-a240-44c87c762bc5</appliesTo>
      <goToChain>3c1faec7-7e1e-4cdd-b3bd-e2f05f4baa9b</goToChain>
    </preconfiguredChoice>
    <!-- Select compression algorithm -->
    <preconfiguredChoice>
      <appliesTo>01d64f58-8295-4b7b-9cab-8f1b153a504f</appliesTo>
      <goToChain>9475447c-9889-430c-9477-6287a9574c5b</goToChain>
    </preconfiguredChoice>
    <!-- Select compression level -->
    <preconfiguredChoice>
      <appliesTo>01c651cb-c174-4ba4-b985-1d87a44d6754</appliesTo>
      <goToChain>414da421-b83f-4648-895f-a34840e3c3f5</goToChain>
    </preconfiguredChoice>
    <!-- Examine contents -->
    <preconfiguredChoice>
      <appliesTo>accea2bf-ba74-4a3a-bb97-614775c74459</appliesTo>
      <goToChain>e0a39199-c62a-4a2f-98de-e9d1116460a8</goToChain>
    </preconfiguredChoice>
    <!-- Transcribe file -->
    <preconfiguredChoice>
      <appliesTo>7079be6d-3a25-41e6-a481-cee5f352fe6e</appliesTo>
      <goToChain>1170e555-cd4e-4b2f-a3d6-bfb09e8fcc53</goToChain>
    </preconfiguredChoice>
    <!-- Transfer tree diagram -->
    <preconfiguredChoice>
      <appliesTo>56eebd45-5600-4768-a8c2-ec0114555a3d</appliesTo>
      <goToChain>e9eaef1e-c2e0-4e3b-b942-bfb537162795</goToChain>
    </preconfiguredChoice>
  </preconfiguredChoices>
</processingMCP>
"""

def fetchUUIDFromPath(path):
    # find UUID on end of SIP path
    uuidLen = -36
    if isUUID(path[uuidLen - 1:-1]):
        return path[uuidLen - 1:-1]

def findOrCreateSipInDB(path, waitSleep=dbWaitSleep, unit_type='SIP'):
    """Matches a directory to a database sip by it's appended UUID, or path. If it doesn't find one, it will create one"""
    path = path.replace(django_settings.SHARED_DIRECTORY, "%sharedPath%", 1)

    query = Q(currentpath=path)

    # Find UUID on end of SIP path
    UUID = fetchUUIDFromPath(path)
    sip = None
    if UUID:
        query = query | Q(uuid=UUID)

    sips = SIP.objects.filter(query)
    count = sips.count()
    if count > 1:
        # This might have happened because the UUID at the end of the directory
        # name corresponds to a different SIP in the database.
        # Try refiltering the queryset on path alone, and see if that brought us
        # down to a single SIP.
        sips = sips.filter(currentpath=path)
        count = sips.count()

        # Darn: we must have multiple SIPs with the same path in the database.
        # We have no reasonable way to recover from this condition.
        if count > 1:
            logger.error('More than one SIP for path %s and/or UUID %s, using first result', path, UUID)
    if count > 0:
        sip = sips[0]
        UUID = sip.uuid
        logger.info('Using existing SIP %s at %s', UUID, path)
    else:
        logger.info('Not using existing SIP %s at %s', UUID, path)

    if sip is None:
        # Create it
        # Note that if UUID is None here, a new UUID will be generated
        # and returned by the function; otherwise it returns the
        # value that was passed in.
        UUID = databaseFunctions.createSIP(path, UUID=UUID)
        logger.info('Creating SIP %s at %s', UUID, path)
    else:
        current_path = sip.currentpath
        if current_path != path and unit_type == 'SIP':
            # Ensure path provided matches path in DB
            sip.currentpath = path
            sip.save()

    return UUID

@log_exceptions
@auto_close_db
def createUnitAndJobChain(path, config, terminate=False):
    path = unicodeToStr(path)
    if os.path.isdir(path):
        path = path + "/"
    logger.debug('Creating unit and job chain for %s with %s', path, config)
    unit = None
    if os.path.isdir(path):
        if config[3] == "SIP":
            UUID = findOrCreateSipInDB(path)
            unit = unitSIP(path, UUID)
        elif config[3] == "DIP":
            UUID = findOrCreateSipInDB(path, unit_type='DIP')
            unit = unitDIP(path, UUID)
        elif config[3] == "Transfer":
            unit = unitTransfer(path)
    elif os.path.isfile(path):
        if config[3] == "Transfer":
            unit = unitTransfer(path)
        else:
            return
            UUID = uuid.uuid4()
            unit = unitFile(path, UUID)
    else:
        return
    jobChain(unit, config[1])

    if terminate:
        exit(0)

def createUnitAndJobChainThreaded(path, config, terminate=True):
    global countOfCreateUnitAndJobChainThreaded
    try:
        logger.debug('Watching path %s', path)
        t = threading.Thread(target=createUnitAndJobChain, args=(path, config), kwargs={"terminate": terminate})
        t.daemon = True
        countOfCreateUnitAndJobChainThreaded += 1
        while(django_settings.LIMIT_TASK_THREADS <= threading.activeCount() + django_settings.RESERVED_AS_TASK_PROCESSING_THREADS):
            if stopSignalReceived:
                logger.info('Signal was received; stopping createUnitAndJobChainThreaded(path, config)')
                exit(0)
            logger.debug('Active thread count: %s', threading.activeCount())
            time.sleep(.5)
        countOfCreateUnitAndJobChainThreaded -= 1
        t.start()
    except Exception:
        logger.exception('Error creating threads to watch directories')

def watchDirectories():
    """Start watching the watched directories defined in the WatchedDirectories table in the database."""
    watched_dir_path = django_settings.WATCH_DIRECTORY
    interval = django_settings.WATCH_DIRECTORY_INTERVAL

    watched_directories = WatchedDirectory.objects.all()

    for watched_directory in watched_directories:
        directory = watched_directory.watched_directory_path.replace("%watchDirectoryPath%", watched_dir_path, 1)

        # Tuple of variables that may be used by a callback
        row = (watched_directory.watched_directory_path, watched_directory.chain_id, watched_directory.only_act_on_directories, watched_directory.expected_type.description)

        if not os.path.isdir(directory):
            os.makedirs(directory)
        for item in os.listdir(directory):
            if item == ".gitignore":
                continue
            item = item.decode("utf-8")
            path = os.path.join(unicode(directory), item)
            while(django_settings.LIMIT_TASK_THREADS <= threading.activeCount() + django_settings.RESERVED_AS_TASK_PROCESSING_THREADS):
                time.sleep(1)
            createUnitAndJobChainThreaded(path, row, terminate=False)
        actOnFiles = True
        if watched_directory.only_act_on_directories:
            actOnFiles = False
        watchDirectory.archivematicaWatchDirectory(
            directory,
            variablesAdded=row,
            callBackFunctionAdded=createUnitAndJobChainThreaded,
            alertOnFiles=actOnFiles,
            interval=interval,
        )

def signal_handler(signalReceived, frame):
    """Used to handle the stop/kill command signals (SIGKILL)"""
    logger.info('Recieved signal %s in frame %s', signalReceived, frame)
    global stopSignalReceived
    stopSignalReceived = True
    threads = threading.enumerate()
    for thread in threads:
        logger.warning('Not stopping %s %s', type(thread), thread)
    sys.stdout.flush()
    sys.stderr.flush()
    sys.exit(0)
    exit(0)

@log_exceptions
@auto_close_db
def debugMonitor():
    """Periodically prints out status of MCP, including whether the database lock is locked, thread count, etc."""
    global countOfCreateUnitAndJobChainThreaded
    while True:
        logger.debug('Debug monitor: datetime: %s', databaseFunctions.getUTCDate())
        logger.debug('Debug monitor: thread count: %s', threading.activeCount())
        logger.debug('Debug monitor: created job chain threaded: %s', countOfCreateUnitAndJobChainThreaded)
        time.sleep(3600)

@log_exceptions
@auto_close_db
def flushOutputs():
    while True:
        sys.stdout.flush()
        sys.stderr.flush()
        time.sleep(5)

def cleanupOldDbEntriesOnNewRun():
    Job.objects.filter(currentstep=Job.STATUS_AWAITING_DECISION).delete()
    Job.objects.filter(currentstep=Job.STATUS_EXECUTING_COMMANDS).update(currentstep=Job.STATUS_FAILED)
    Task.objects.filter(exitcode=None).update(exitcode=-1, stderror="MCP shut down while processing.")

def created_shared_directory_structure():
    dirs = (
        "arrange",
        "completed",
        "completed/transfers",
        "currentlyProcessing",
        "DIPbackups",
        "failed",
        "rejected",
        "sharedMicroServiceTasksConfigs",
        "sharedMicroServiceTasksConfigs/createXmlEventsAssist",
        "sharedMicroServiceTasksConfigs/generateAIP",
        "sharedMicroServiceTasksConfigs/generateAIP/bagit",
        "sharedMicroServiceTasksConfigs/processingMCPConfigs",
        "sharedMicroServiceTasksConfigs/transcoder",
        "sharedMicroServiceTasksConfigs/transcoder/defaultIcons",
        "SIPbackups",
        "tmp",
        "watchedDirectories",
        "watchedDirectories/activeTransfers",
        "watchedDirectories/activeTransfers/baggitDirectory",
        "watchedDirectories/activeTransfers/baggitZippedDirectory",
        "watchedDirectories/activeTransfers/Dspace",
        "watchedDirectories/activeTransfers/maildir",
        "watchedDirectories/activeTransfers/standardTransfer",
        "watchedDirectories/activeTransfers/TRIM",
        "watchedDirectories/approveNormalization",
        "watchedDirectories/approveSubmissionDocumentationIngest",
        "watchedDirectories/quarantined",
        "watchedDirectories/SIPCreation",
        "watchedDirectories/SIPCreation/completedTransfers",
        "watchedDirectories/SIPCreation/SIPsUnderConstruction",
        "watchedDirectories/storeAIP",
        "watchedDirectories/system",
        "watchedDirectories/system/autoProcessSIP",
        "watchedDirectories/system/autoRestructureForCompliance",
        "watchedDirectories/system/createAIC",
        "watchedDirectories/system/reingestAIP",
        "watchedDirectories/uploadDIP",
        "watchedDirectories/uploadedDIPs",
        "watchedDirectories/workFlowDecisions",
        "watchedDirectories/workFlowDecisions/compressionAIPDecisions",
        "watchedDirectories/workFlowDecisions/createDip",
        "watchedDirectories/workFlowDecisions/createTree",
        "watchedDirectories/workFlowDecisions/examineContentsChoice",
        "watchedDirectories/workFlowDecisions/extractPackagesChoice",
        "watchedDirectories/workFlowDecisions/metadataReminder",
        "watchedDirectories/workFlowDecisions/quarantineTransfer",
        "watchedDirectories/workFlowDecisions/selectFormatIDToolIngest",
        "watchedDirectories/workFlowDecisions/selectFormatIDToolTransfer",
        "www",
        "www/AIPsStore",
        "www/AIPsStore/transferBacklog",
        "www/AIPsStore/transferBacklog/arrange",
        "www/AIPsStore/transferBacklog/originals",
        "www/DIPsStore"
    )
    for d in dirs:
        d = os.path.join(django_settings.SHARED_DIRECTORY, d)
        if os.path.isdir(d):
            continue
        logger.info('Creating directory: %s', d)
        os.makedirs(d, mode=770)

    # Populate defaultProcessingMCP.xml (DEFAULT_PROCESSING_CONFIG)
    default_proc_config = os.path.join(django_settings.SHARED_DIRECTORY, 'sharedMicroServiceTasksConfigs/processingMCPConfigs/defaultProcessingMCP.xml')
    if os.path.isfile(default_proc_config):
        return
    with open(default_proc_config, 'w') as fd:
        fd.write(DEFAULT_PROCESSING_CONFIG.encode('utf-8'))

def _except_hook_log_everything(exc_type, exc_value, exc_traceback):
    """
    Replacement for default exception handler that logs exceptions.
    """
    # Reference http://stackoverflow.com/a/16993115/2475775
    logger.error("Uncaught exception", exc_info=(exc_type, exc_value, exc_traceback))
    sys.__excepthook__(exc_type, exc_value, exc_traceback)

if __name__ == '__main__':
    logger = logging.getLogger("archivematica.mcp.server")

    # Replace exception handler with one that logs exceptions
    sys.excepthook = _except_hook_log_everything

    signal.signal(signal.SIGINT, signal_handler)
    signal.signal(signal.SIGTERM, signal_handler)

    logger.info('This PID: %s', os.getpid())
    logger.info('User: %s', getpass.getuser())

    dicts.setup(
        shared_directory=django_settings.SHARED_DIRECTORY,
        processing_directory=django_settings.PROCESSING_DIRECTORY,
        watch_directory=django_settings.WATCH_DIRECTORY,
        rejected_directory=django_settings.REJECTED_DIRECTORY,
    )

    created_shared_directory_structure()

    t = threading.Thread(target=debugMonitor)
    t.daemon = True
    t.start()

    t = threading.Thread(target=flushOutputs)
    t.daemon = True
    t.start()
    cleanupOldDbEntriesOnNewRun()
    watchDirectories()

    # This is blocking the main thread with the worker loop
    RPCServer.startRPCServer()
<|endoftext|>


<commit_before>from django.conf import settings
from django.db import transaction
from django.utils.translation import ugettext as _
from rest_framework import serializers
from rest_framework.exceptions import PermissionDenied

from karrot.history.models import History, HistoryTypus
from karrot.utils.misc import find_changed
from karrot.places.models import Place as PlaceModel, PlaceStatus, PlaceSubscription

class PlaceHistorySerializer(serializers.ModelSerializer):
    class Meta:
        model = PlaceModel
        fields = '__all__'

class PlaceSerializer(serializers.ModelSerializer):
    class Meta:
        model = PlaceModel
        fields = [
            'id',
            'name',
            'description',
            'group',
            'address',
            'latitude',
            'longitude',
            'weeks_in_advance',
            'status',
            'is_subscribed',
        ]
        extra_kwargs = {
            'name': {
                'min_length': 3,
            },
            'description': {
                'trim_whitespace': False,
                'max_length': settings.DESCRIPTION_MAX_LENGTH,
            },
        }
        read_only_fields = [
            'id',
        ]

    status = serializers.ChoiceField(
        choices=[status.value for status in PlaceStatus], default=PlaceModel.DEFAULT_STATUS
    )
    is_subscribed = serializers.SerializerMethodField()

    def get_is_subscribed(self, place):
        return place.placesubscription_set.filter(user=self.context['request'].user).exists()

    def save(self, **kwargs):
        return super().save(last_changed_by=self.context['request'].user)

    def create(self, validated_data):
        place = super().create(validated_data)

        # TODO move into receiver
        History.objects.create(
            typus=HistoryTypus.STORE_CREATE,
            group=place.group,
            place=place,
            users=[
                self.context['request'].user,
            ],
            payload=self.initial_data,
            after=PlaceHistorySerializer(place).data,
        )
        place.group.refresh_active_status()
        return place

    def validate_group(self, group):
        if not group.is_member(self.context['request'].user):
            raise PermissionDenied(_('You are not a member of this group.'))
        if not group.is_editor(self.context['request'].user):
            raise PermissionDenied(_('You need to be a group editor'))
        return group

    def validate_weeks_in_advance(self, w):
        if w < 1:
            raise serializers.ValidationError(_('Set at least one week in advance'))
        if w > settings.STORE_MAX_WEEKS_IN_ADVANCE:
            raise serializers.ValidationError(
                _('Do not set more than %(count)s weeks in advance') % {'count': settings.STORE_MAX_WEEKS_IN_ADVANCE}
            )
        return w

class PlaceUpdateSerializer(PlaceSerializer):
    class Meta:
        model = PlaceModel
        fields = PlaceSerializer.Meta.fields
        read_only_fields = PlaceSerializer.Meta.read_only_fields
        extra_kwargs = PlaceSerializer.Meta.extra_kwargs

    def save(self, **kwargs):
        self._validated_data = find_changed(self.instance, self.validated_data)
        skip_update = len(self.validated_data.keys()) == 0
        if skip_update:
            return self.instance
        return super().save(**kwargs)

    @transaction.atomic()
    def update(self, place, validated_data):
        before_data = PlaceHistorySerializer(place).data
        place = super().update(place, validated_data)
        after_data = PlaceHistorySerializer(place).data

        if before_data != after_data:
            History.objects.create(
                typus=HistoryTypus.STORE_MODIFY,
                group=place.group,
                place=place,
                users=[self.context['request'].user],
                payload={k: self.initial_data.get(k)
                         for k in validated_data.keys()},
                before=before_data,
                after=after_data,
            )
        place.group.refresh_active_status()
        return place

class PlaceSubscriptionSerializer(serializers.ModelSerializer):
    class Meta:
        model = PlaceSubscription
        fields = [
            'place',
        ]

    def save(self, **kwargs):
        return super().save(user=self.context['request'].user)

    def validate_place(self, place):
        if place.placesubscription_set.filter(user=self.context['request'].user).exists():
            raise serializers.ValidationError(_('You are already subscribed to this place'))
        return place

    def create(self, validated_data):
        return PlaceSubscription.objects.create(**validated_data)
<commit_msg>add subscribers list to place API<commit_after>from django.conf import settings
from django.db import transaction
from django.utils.translation import ugettext as _
from rest_framework import serializers
from rest_framework.exceptions import PermissionDenied

from karrot.history.models import History, HistoryTypus
from karrot.utils.misc import find_changed
from karrot.places.models import Place as PlaceModel, PlaceStatus, PlaceSubscription

class PlaceHistorySerializer(serializers.ModelSerializer):
    class Meta:
        model = PlaceModel
        fields = '__all__'

class PlaceSerializer(serializers.ModelSerializer):
    class Meta:
        model = PlaceModel
        fields = [
            'id',
            'name',
            'description',
            'group',
            'address',
            'latitude',
            'longitude',
            'weeks_in_advance',
            'status',
            'is_subscribed',
            'subscribers',
        ]
        extra_kwargs = {
            'name': {
                'min_length': 3,
            },
            'description': {
                'trim_whitespace': False,
                'max_length': settings.DESCRIPTION_MAX_LENGTH,
            },
        }
        read_only_fields = [
            'id',
            'subscribers',
        ]

    status = serializers.ChoiceField(
        choices=[status.value for status in PlaceStatus], default=PlaceModel.DEFAULT_STATUS
    )
    is_subscribed = serializers.SerializerMethodField()

    def get_is_subscribed(self, place):
        return place.placesubscription_set.filter(user=self.context['request'].user).exists()

    def save(self, **kwargs):
        return super().save(last_changed_by=self.context['request'].user)

    def create(self, validated_data):
        place = super().create(validated_data)

        # TODO move into receiver
        History.objects.create(
            typus=HistoryTypus.STORE_CREATE,
            group=place.group,
            place=place,
            users=[
                self.context['request'].user,
            ],
            payload=self.initial_data,
            after=PlaceHistorySerializer(place).data,
        )
        place.group.refresh_active_status()
        return place

    def validate_group(self, group):
        if not group.is_member(self.context['request'].user):
            raise PermissionDenied(_('You are not a member of this group.'))
        if not group.is_editor(self.context['request'].user):
            raise PermissionDenied(_('You need to be a group editor'))
        return group

    def validate_weeks_in_advance(self, w):
        if w < 1:
            raise serializers.ValidationError(_('Set at least one week in advance'))
        if w > settings.STORE_MAX_WEEKS_IN_ADVANCE:
            raise serializers.ValidationError(
                _('Do not set more than %(count)s weeks in advance') % {'count': settings.STORE_MAX_WEEKS_IN_ADVANCE}
            )
        return w

class PlaceUpdateSerializer(PlaceSerializer):
    class Meta:
        model = PlaceModel
        fields = PlaceSerializer.Meta.fields
        read_only_fields = PlaceSerializer.Meta.read_only_fields
        extra_kwargs = PlaceSerializer.Meta.extra_kwargs

    def save(self, **kwargs):
        self._validated_data = find_changed(self.instance, self.validated_data)
        skip_update = len(self.validated_data.keys()) == 0
        if skip_update:
            return self.instance
        return super().save(**kwargs)

    @transaction.atomic()
    def update(self, place, validated_data):
        before_data = PlaceHistorySerializer(place).data
        place = super().update(place, validated_data)
        after_data = PlaceHistorySerializer(place).data

        if before_data != after_data:
            History.objects.create(
                typus=HistoryTypus.STORE_MODIFY,
                group=place.group,
                place=place,
                users=[self.context['request'].user],
                payload={k: self.initial_data.get(k)
                         for k in validated_data.keys()},
                before=before_data,
                after=after_data,
            )
        place.group.refresh_active_status()
        return place

class PlaceSubscriptionSerializer(serializers.ModelSerializer):
    class Meta:
        model = PlaceSubscription
        fields = [
            'place',
        ]

    def save(self, **kwargs):
        return super().save(user=self.context['request'].user)

    def validate_place(self, place):
        if place.placesubscription_set.filter(user=self.context['request'].user).exists():
            raise serializers.ValidationError(_('You are already subscribed to this place'))
        return place

    def create(self, validated_data):
        return PlaceSubscription.objects.create(**validated_data)
<|endoftext|>