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

# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""
Top level script to make a castro plot in mass / sigmav space
"""
from __future__ import absolute_import, division, print_function

import os
import numpy as np

from astropy.table import Table

from fermipy.utils import init_matplotlib_backend, load_yaml
from fermipy.jobs.utils import is_not_null
from fermipy.jobs.link import Link
from fermipy.jobs.scatter_gather import ScatterGather
from fermipy.jobs.slac_impl import make_nfs_path

from dmpipe.dm_spectral import DMCastroData, DMSpecTable
from dmpipe.dm_plotting_utils import plot_dm_castro
from dmpipe.dm_plotting_utils import plot_dm_spectra_by_mass, plot_dm_spectra_by_channel
from dmpipe.dm_plotting_utils import plot_limits_from_arrays, plot_mc_truth

from dmpipe.name_policy import NameFactory
from dmpipe import defaults

init_matplotlib_backend()
NAME_FACTORY = NameFactory(basedir='.')

def get_ul_bands(table, prefix):
    """ Get the Upper limit bands out of a `astropy.Table` objects """
    o = dict(q02=np.squeeze(table["%s_q02" % prefix]),
             q16=np.squeeze(table["%s_q16" % prefix]),
             q84=np.squeeze(table["%s_q84" % prefix]),
             q97=np.squeeze(table["%s_q97" % prefix]),
             median=np.squeeze(table["%s_median" % prefix]))
    return o

class PlotDMSpectra(Link):
    """Small class wrap an analysis script.

    This is useful for parallelizing analysis using the fermipy.jobs module.
    """
    appname = 'dmpipe-plot-dm-spectra'
    linkname_default = 'plot-dm-spectra'
    usage = '%s [options]' % (appname)
    description = "Plot the DM spectra stored in pre-computed tables"

    default_options = dict(infile=defaults.generic['infile'],
                           outfile=defaults.generic['outfile'],
                           chan=defaults.common['chan'],
                           mass=defaults.common['mass'],
                           spec_type=defaults.common['spec_type'])

    def run_analysis(self, argv):
        """Run this analysis"""
        args = self._parser.parse_args(argv)

        dm_spec_table = DMSpecTable.create_from_fits(args.infile)
        dm_plot_by_mass = plot_dm_spectra_by_mass(
            dm_spec_table, chan=args.chan, spec_type=args.spec_type)
        dm_plot_by_chan = plot_dm_spectra_by_channel(
            dm_spec_table, mass=args.mass, spec_type=args.spec_type)

        if args.outfile:
            dm_plot_by_mass[0].savefig(
                args.outfile.replace(
                    '.png', '_%s.png' %
                    args.chan))
            dm_plot_by_chan[0].savefig(
                args.outfile.replace(
                    '.png', '_%1.FGeV.png' %
                    args.mass))

class PlotLimits(Link):
    """Small class wrap an analysis script.

    This is useful for parallelizing analysis using the fermipy.jobs module.
    """
    appname = 'dmpipe-plot-limits'
    linkname_default = 'plot-limits'
    usage = '%s [options]' % (appname)
    description = "Plot DM limits on <sigma v> versus mass"

    default_options = dict(infile=defaults.generic['infile'],
                           outfile=defaults.generic['outfile'],
                           chan=defaults.common['chan'],
                           bands=defaults.collect['bands'],
                           sim=defaults.sims['sim'])

    def run_analysis(self, argv):
        """Run this analysis"""
        args = self._parser.parse_args(argv)

        if is_not_null(args.infile):
            tab_m = Table.read(args.infile, hdu="MASSES")
            tab_s = Table.read(args.infile, hdu=args.chan)

            xvals = tab_m['MASSES'][0]
            yvals = tab_s['UL_0.95'][0]
            ldict = dict(limits=(xvals, yvals))
        else:
            ldict = {}

        if is_not_null(args.bands):
            tab_b = Table.read(args.bands, hdu=args.chan)
            tab_bm = Table.read(args.bands, hdu="MASSES")
            bands = get_ul_bands(tab_b, 'UL_0.95')
            bands['MASSES'] = tab_bm['MASSES'][0]
        else:
            bands = None

        if is_not_null(args.sim):
            sim_srcs = load_yaml(args.sim)
            injected_src = sim_srcs.get('injected_source', None)
        else:
            injected_src = None

        xlims = (1e1, 1e4)
        ylims = (1e-28, 1e-22)

        dm_plot = plot_limits_from_arrays(ldict, xlims, ylims, bands)

        if injected_src is not None:
            mc_model = injected_src['source_model']
            plot_mc_truth(dm_plot[1], mc_model)

        if args.outfile:
            dm_plot[0].savefig(args.outfile)
            return None
        return dm_plot

class PlotDM(Link):
    """Small class wrap an analysis script.

    This is useful for parallelizing analysis using the fermipy.jobs module.
    """
    appname = 'dmpipe-plot-dm'
    linkname_default = 'plot-dm'
    usage = "%s [options]" % (appname)
    description = "Plot the likelihood vs <sigma v> and DM particle mass"

    default_options = dict(infile=defaults.generic['infile'],
                           outfile=defaults.generic['outfile'],
                           chan=defaults.common['chan'])

    def run_analysis(self, argv):
        """Run this analysis"""
        args = self._parser.parse_args(argv)
        tab_m = Table.read(args.infile, hdu="MASSES")
        tab_s = Table.read(args.infile, hdu=args.chan)
        dm_castro = DMCastroData.create_from_tables(tab_s, tab_m)

        dm_plot = plot_dm_castro(dm_castro)
        if args.outfile:
            dm_plot[0].savefig(args.outfile)
            return None
        return dm_plot

class PlotLimits_SG(ScatterGather):
    """Small class to generate configurations for this script

    This adds the following arguments:
    """
    appname = 'dmpipe-plot-limits-sg'
    usage = "%s [options]" % (appname)
    description = "Make castro plots for set of targets"
    clientclass = PlotLimits

    job_time = 60

    default_options = dict(ttype=defaults.common['ttype'],
                           targetlist=defaults.common['targetlist'],
                           bands=defaults.collect['bands'],
                           channels=defaults.common['channels'],
                           jpriors=defaults.common['jpriors'],
                           dry_run=defaults.common['dry_run'])

    def build_job_configs(self, args):
        """Hook to build job configurations
        """
        job_configs = {}

        ttype = args['ttype']
        (targets_yaml, sim) = NAME_FACTORY.resolve_targetfile(args)
        if targets_yaml is None:
            return job_configs

        jpriors = args['jpriors']
        channels = args['channels']

        targets = load_yaml(targets_yaml)
        for target_name, target_list in targets.items():
            for targ_prof in target_list:
                for jprior in jpriors:
                    name_keys = dict(target_type=ttype,
                                     target_name=target_name,
                                     profile=targ_prof,
                                     jprior=jprior,
                                     fullpath=True)
                    input_path = NAME_FACTORY.dmlimitsfile(**name_keys)
                    for chan in channels:
                        targ_key = "%s:%s:%s:%s" % (
                            target_name, targ_prof, jprior, chan)

                        output_path = input_path.replace(
                            '.fits', '_%s.png' % chan)
                        logfile = make_nfs_path(
                            output_path.replace('.png', '.log'))
                        job_config = dict(infile=input_path,
                                          outfile=output_path,
                                          jprior=jprior,
                                          logfile=logfile,
                                          chan=chan)
                        job_configs[targ_key] = job_config

        return job_configs

class PlotStackedLimits_SG(ScatterGather):
    """Small class to generate configurations for this script

    This adds the following arguments:
    """
    appname = 'dmpipe-plot-stacked-limits-sg'
    usage = "%s [options]" % (appname)
    description = "Make castro plots for set of targets"
    clientclass = PlotLimits

    job_time = 60

    default_options = dict(ttype=defaults.common['ttype'],
                           rosterlist=defaults.common['targetlist'],
                           bands=defaults.collect['bands'],
                           channels=defaults.common['channels'],
                           jpriors=defaults.common['jpriors'],
                           sim=defaults.sims['sim'],
                           nsims=defaults.sims['nsims'],
                           seed=defaults.sims['seed'],
                           dry_run=defaults.common['dry_run'])

    def build_job_configs(self, args):
        """Hook to build job configurations
        """
        job_configs = {}

        ttype = args['ttype']
        (roster_yaml, sim) = NAME_FACTORY.resolve_rosterfile(args)
        if roster_yaml is None:
            return job_configs

        roster_dict = load_yaml(roster_yaml)

        jpriors = args['jpriors']
        channels = args['channels']

        for roster_name in roster_dict.keys():
            for jprior in jpriors:
                name_keys = dict(target_type=ttype,
                                 roster_name=roster_name,
                                 jprior=jprior,
                                 sim_name=sim,
                                 fullpath=True)
                for chan in channels:
                    targ_key = "%s:%s:%s" % (roster_name, jprior, chan)
                    if sim is not None:
                        seedlist = range(
                            args['seed'], args['seed'] + args['nsims'])
                        sim_path = os.path.join('config', 'sim_%s.yaml' % sim)
                    else:
                        seedlist = [None]
                        sim_path = None

                    for seed in seedlist:
                        if seed is not None:
                            name_keys['seed'] = "%06i" % seed
                            input_path = NAME_FACTORY.sim_stackedlimitsfile(
                                **name_keys)
                            full_targ_key = "%s_%06i" % (targ_key, seed)
                        else:
                            input_path = NAME_FACTORY.stackedlimitsfile(
                                **name_keys)
                            full_targ_key = targ_key

                        output_path = input_path.replace(
                            '.fits', '_%s.png' % chan)
                        logfile = make_nfs_path(
                            output_path.replace('.png', '.log'))
                        job_config = dict(infile=input_path,
                                          outfile=output_path,
                                          jprior=jprior,
                                          logfile=logfile,
                                          sim=sim_path,
                                          chan=chan)
                        job_configs[full_targ_key] = job_config

        return job_configs

class PlotDM_SG(ScatterGather):
    """Small class to generate configurations for this script

    This adds the following arguments:
    """
    appname = 'dmpipe-plot-dm-sg'
    usage = "%s [options]" % (appname)
    description = "Make castro plots for set of targets"
    clientclass = PlotDM

    job_time = 60

    default_options = dict(ttype=defaults.common['ttype'],
                           targetlist=defaults.common['targetlist'],
                           channels=defaults.common['channels'],
                           jpriors=defaults.common['jpriors'],
                           dry_run=defaults.common['dry_run'])

    def build_job_configs(self, args):
        """Hook to build job configurations
        """
        job_configs = {}

        ttype = args['ttype']
        (targets_yaml, sim) = NAME_FACTORY.resolve_targetfile(args)
        if targets_yaml is None:
            return job_configs

        targets = load_yaml(targets_yaml)

        jpriors = args['jpriors']
        channels = args['channels']

        for target_name, target_list in targets.items():
            for targ_prof in target_list:
                for jprior in jpriors:
                    name_keys = dict(target_type=ttype,
                                     target_name=target_name,
                                     profile=targ_prof,
                                     jprior=jprior,
                                     fullpath=True)
                    input_path = NAME_FACTORY.dmlikefile(**name_keys)
                    for chan in channels:
                        targ_key = "%s:%s:%s:%s" % (
                            target_name, targ_prof, jprior, chan)
                        output_path = input_path.replace(
                            '.fits', '_%s.png' % chan)
                        logfile = make_nfs_path(
                            output_path.replace('.png', '.log'))
                        job_config = dict(infile=input_path,
                                          outfile=output_path,
                                          jprior=jprior,
                                          logfile=logfile,
                                          chan=chan)
                        job_configs[targ_key] = job_config

        return job_configs

class PlotStackedDM_SG(ScatterGather):
    """Small class to generate configurations for this script

    This adds the following arguments:
    """
    appname = 'dmpipe-plot-stacked-dm-sg'
    usage = "%s [options]" % (appname)
    description = "Make castro plots for set of targets"
    clientclass = PlotDM

    job_time = 60

    default_options = dict(ttype=defaults.common['ttype'],
                           rosterlist=defaults.common['targetlist'],
                           channels=defaults.common['channels'],
                           jpriors=defaults.common['jpriors'],
                           sim=defaults.sims['sim'],
                           nsims=defaults.sims['nsims'],
                           seed=defaults.sims['seed'],
                           dry_run=defaults.common['dry_run'])

    def build_job_configs(self, args):
        """Hook to build job configurations
        """
        job_configs = {}

        ttype = args['ttype']
        (roster_yaml, sim) = NAME_FACTORY.resolve_rosterfile(args)
        if roster_yaml is None:
            return job_configs

        roster_dict = load_yaml(roster_yaml)

        jpriors = args['jpriors']
        channels = args['channels']

        for roster_name in roster_dict.keys():
            for jprior in jpriors:
                name_keys = dict(target_type=ttype,
                                 roster_name=roster_name,
                                 jprior=jprior,
                                 sim_name=sim,
                                 fullpath=True)

                for chan in channels:
                    targ_key = "%s:%s:%s" % (roster_name, jprior, chan)

                    if sim is not None:
                        seedlist = range(
                            args['seed'], args['seed'] + args['nsims'])
                    else:
                        seedlist = [None]

                    for seed in seedlist:
                        if seed is not None:
                            name_keys['seed'] = "%06i" % seed
                            input_path = NAME_FACTORY.sim_resultsfile(
                                **name_keys)
                            full_targ_key = "%s_%06i" % (targ_key, seed)
                        else:
                            input_path = NAME_FACTORY.resultsfile(**name_keys)
                            full_targ_key = targ_key

                        output_path = input_path.replace(
                            '.fits', '_%s.png' % chan)
                        logfile = make_nfs_path(
                            output_path.replace('.png', '.log'))
                        job_config = dict(infile=input_path,
                                          outfile=output_path,
                                          jprior=jprior,
                                          logfile=logfile,
                                          chan=chan)
                        job_configs[full_targ_key] = job_config

        return job_configs

def register_classes():
    """Register these classes with the `LinkFactory` """
    PlotDMSpectra.register_class()
    PlotLimits.register_class()
    PlotLimits_SG.register_class()
    PlotDM.register_class()
    PlotDM_SG.register_class()
    PlotStackedDM_SG.register_class()
    PlotStackedLimits_SG.register_class()
<commit_msg>added a couple of option to pass along in PlotLimits_SG<commit_after>#!/usr/bin/env python
#

# Licensed under a 3-clause BSD style license - see LICENSE.rst

"""
Top level script to make a castro plot in mass / sigmav space
"""
from __future__ import absolute_import, division, print_function

import os
import numpy as np

from astropy.table import Table

from fermipy.utils import init_matplotlib_backend, load_yaml
from fermipy.jobs.utils import is_not_null
from fermipy.jobs.link import Link
from fermipy.jobs.scatter_gather import ScatterGather
from fermipy.jobs.slac_impl import make_nfs_path

from dmpipe.dm_spectral import DMCastroData, DMSpecTable
from dmpipe.dm_plotting_utils import plot_dm_castro
from dmpipe.dm_plotting_utils import plot_dm_spectra_by_mass, plot_dm_spectra_by_channel
from dmpipe.dm_plotting_utils import plot_limits_from_arrays, plot_mc_truth

from dmpipe.name_policy import NameFactory
from dmpipe import defaults

init_matplotlib_backend()
NAME_FACTORY = NameFactory(basedir='.')

def get_ul_bands(table, prefix):
    """ Get the Upper limit bands out of a `astropy.Table` objects """
    o = dict(q02=np.squeeze(table["%s_q02" % prefix]),
             q16=np.squeeze(table["%s_q16" % prefix]),
             q84=np.squeeze(table["%s_q84" % prefix]),
             q97=np.squeeze(table["%s_q97" % prefix]),
             median=np.squeeze(table["%s_median" % prefix]))
    return o

class PlotDMSpectra(Link):
    """Small class wrap an analysis script.

    This is useful for parallelizing analysis using the fermipy.jobs module.
    """
    appname = 'dmpipe-plot-dm-spectra'
    linkname_default = 'plot-dm-spectra'
    usage = '%s [options]' % (appname)
    description = "Plot the DM spectra stored in pre-computed tables"

    default_options = dict(infile=defaults.generic['infile'],
                           outfile=defaults.generic['outfile'],
                           chan=defaults.common['chan'],
                           mass=defaults.common['mass'],
                           spec_type=defaults.common['spec_type'])

    def run_analysis(self, argv):
        """Run this analysis"""
        args = self._parser.parse_args(argv)

        dm_spec_table = DMSpecTable.create_from_fits(args.infile)
        dm_plot_by_mass = plot_dm_spectra_by_mass(
            dm_spec_table, chan=args.chan, spec_type=args.spec_type)
        dm_plot_by_chan = plot_dm_spectra_by_channel(
            dm_spec_table, mass=args.mass, spec_type=args.spec_type)

        if args.outfile:
            dm_plot_by_mass[0].savefig(
                args.outfile.replace(
                    '.png', '_%s.png' %
                    args.chan))
            dm_plot_by_chan[0].savefig(
                args.outfile.replace(
                    '.png', '_%1.FGeV.png' %
                    args.mass))

class PlotLimits(Link):
    """Small class wrap an analysis script.

    This is useful for parallelizing analysis using the fermipy.jobs module.
    """
    appname = 'dmpipe-plot-limits'
    linkname_default = 'plot-limits'
    usage = '%s [options]' % (appname)
    description = "Plot DM limits on <sigma v> versus mass"

    default_options = dict(infile=defaults.generic['infile'],
                           outfile=defaults.generic['outfile'],
                           chan=defaults.common['chan'],
                           bands=defaults.collect['bands'],
                           sim=defaults.sims['sim'])

    def run_analysis(self, argv):
        """Run this analysis"""
        args = self._parser.parse_args(argv)

        if is_not_null(args.infile):
            tab_m = Table.read(args.infile, hdu="MASSES")
            tab_s = Table.read(args.infile, hdu=args.chan)

            xvals = tab_m['MASSES'][0]
            yvals = tab_s['UL_0.95'][0]
            ldict = dict(limits=(xvals, yvals))
        else:
            ldict = {}

        if is_not_null(args.bands):
            tab_b = Table.read(args.bands, hdu=args.chan)
            tab_bm = Table.read(args.bands, hdu="MASSES")
            bands = get_ul_bands(tab_b, 'UL_0.95')
            bands['MASSES'] = tab_bm['MASSES'][0]
        else:
            bands = None

        if is_not_null(args.sim):
            sim_srcs = load_yaml(args.sim)
            injected_src = sim_srcs.get('injected_source', None)
        else:
            injected_src = None

        xlims = (1e1, 1e4)
        ylims = (1e-28, 1e-22)

        dm_plot = plot_limits_from_arrays(ldict, xlims, ylims, bands)

        if injected_src is not None:
            mc_model = injected_src['source_model']
            plot_mc_truth(dm_plot[1], mc_model)

        if args.outfile:
            dm_plot[0].savefig(args.outfile)
            return None
        return dm_plot

class PlotDM(Link):
    """Small class wrap an analysis script.

    This is useful for parallelizing analysis using the fermipy.jobs module.
    """
    appname = 'dmpipe-plot-dm'
    linkname_default = 'plot-dm'
    usage = "%s [options]" % (appname)
    description = "Plot the likelihood vs <sigma v> and DM particle mass"

    default_options = dict(infile=defaults.generic['infile'],
                           outfile=defaults.generic['outfile'],
                           chan=defaults.common['chan'])

    def run_analysis(self, argv):
        """Run this analysis"""
        args = self._parser.parse_args(argv)
        tab_m = Table.read(args.infile, hdu="MASSES")
        tab_s = Table.read(args.infile, hdu=args.chan)
        dm_castro = DMCastroData.create_from_tables(tab_s, tab_m)

        dm_plot = plot_dm_castro(dm_castro)
        if args.outfile:
            dm_plot[0].savefig(args.outfile)
            return None
        return dm_plot

class PlotLimits_SG(ScatterGather):
    """Small class to generate configurations for this script

    This adds the following arguments:
    """
    appname = 'dmpipe-plot-limits-sg'
    usage = "%s [options]" % (appname)
    description = "Make castro plots for set of targets"
    clientclass = PlotLimits

    job_time = 60

    default_options = dict(ttype=defaults.common['ttype'],
                           targetlist=defaults.common['targetlist'],
                           channels=defaults.common['channels'],
                           jpriors=defaults.common['jpriors'],
                           dry_run=defaults.common['dry_run'])

    def build_job_configs(self, args):
        """Hook to build job configurations
        """
        job_configs = {}

        ttype = args['ttype']
        (targets_yaml, sim) = NAME_FACTORY.resolve_targetfile(args)
        if targets_yaml is None:
            return job_configs

        jpriors = args['jpriors']
        channels = args['channels']

        targets = load_yaml(targets_yaml)
        for target_name, target_list in targets.items():
            for targ_prof in target_list:
                for jprior in jpriors:
                    name_keys = dict(target_type=ttype,
                                     target_name=target_name,
                                     profile=targ_prof,
                                     jprior=jprior,
                                     fullpath=True)
                    input_path = NAME_FACTORY.dmlimitsfile(**name_keys)
                    for chan in channels:
                        targ_key = "%s:%s:%s:%s" % (
                            target_name, targ_prof, jprior, chan)

                        output_path = input_path.replace(
                            '.fits', '_%s.png' % chan)
                        logfile = make_nfs_path(
                            output_path.replace('.png', '.log'))
                        job_config = dict(infile=input_path,
                                          outfile=output_path,
                                          jprior=jprior,
                                          bands=None,
                                          sim=sim,
                                          logfile=logfile,
                                          chan=chan)
                        job_configs[targ_key] = job_config

        return job_configs

class PlotStackedLimits_SG(ScatterGather):
    """Small class to generate configurations for this script

    This adds the following arguments:
    """
    appname = 'dmpipe-plot-stacked-limits-sg'
    usage = "%s [options]" % (appname)
    description = "Make castro plots for set of targets"
    clientclass = PlotLimits

    job_time = 60

    default_options = dict(ttype=defaults.common['ttype'],
                           rosterlist=defaults.common['targetlist'],
                           bands=defaults.collect['bands'],
                           channels=defaults.common['channels'],
                           jpriors=defaults.common['jpriors'],
                           sim=defaults.sims['sim'],
                           nsims=defaults.sims['nsims'],
                           seed=defaults.sims['seed'],
                           dry_run=defaults.common['dry_run'])

    def build_job_configs(self, args):
        """Hook to build job configurations
        """
        job_configs = {}

        ttype = args['ttype']
        (roster_yaml, sim) = NAME_FACTORY.resolve_rosterfile(args)
        if roster_yaml is None:
            return job_configs

        roster_dict = load_yaml(roster_yaml)

        jpriors = args['jpriors']
        channels = args['channels']

        for roster_name in roster_dict.keys():
            for jprior in jpriors:
                name_keys = dict(target_type=ttype,
                                 roster_name=roster_name,
                                 jprior=jprior,
                                 sim_name=sim,
                                 fullpath=True)
                for chan in channels:
                    targ_key = "%s:%s:%s" % (roster_name, jprior, chan)
                    if sim is not None:
                        seedlist = range(
                            args['seed'], args['seed'] + args['nsims'])
                        sim_path = os.path.join('config', 'sim_%s.yaml' % sim)
                    else:
                        seedlist = [None]
                        sim_path = None

                    for seed in seedlist:
                        if seed is not None:
                            name_keys['seed'] = "%06i" % seed
                            input_path = NAME_FACTORY.sim_stackedlimitsfile(
                                **name_keys)
                            full_targ_key = "%s_%06i" % (targ_key, seed)
                        else:
                            input_path = NAME_FACTORY.stackedlimitsfile(
                                **name_keys)
                            full_targ_key = targ_key

                        output_path = input_path.replace(
                            '.fits', '_%s.png' % chan)
                        logfile = make_nfs_path(
                            output_path.replace('.png', '.log'))
                        job_config = dict(infile=input_path,
                                          outfile=output_path,
                                          jprior=jprior,
                                          logfile=logfile,
                                          sim=sim_path,
                                          chan=chan)
                        job_configs[full_targ_key] = job_config

        return job_configs

class PlotDM_SG(ScatterGather):
    """Small class to generate configurations for this script

    This adds the following arguments:
    """
    appname = 'dmpipe-plot-dm-sg'
    usage = "%s [options]" % (appname)
    description = "Make castro plots for set of targets"
    clientclass = PlotDM

    job_time = 60

    default_options = dict(ttype=defaults.common['ttype'],
                           targetlist=defaults.common['targetlist'],
                           channels=defaults.common['channels'],
                           jpriors=defaults.common['jpriors'],
                           dry_run=defaults.common['dry_run'])

    def build_job_configs(self, args):
        """Hook to build job configurations
        """
        job_configs = {}

        ttype = args['ttype']
        (targets_yaml, sim) = NAME_FACTORY.resolve_targetfile(args)
        if targets_yaml is None:
            return job_configs

        targets = load_yaml(targets_yaml)

        jpriors = args['jpriors']
        channels = args['channels']

        for target_name, target_list in targets.items():
            for targ_prof in target_list:
                for jprior in jpriors:
                    name_keys = dict(target_type=ttype,
                                     target_name=target_name,
                                     profile=targ_prof,
                                     jprior=jprior,
                                     fullpath=True)
                    input_path = NAME_FACTORY.dmlikefile(**name_keys)
                    for chan in channels:
                        targ_key = "%s:%s:%s:%s" % (
                            target_name, targ_prof, jprior, chan)
                        output_path = input_path.replace(
                            '.fits', '_%s.png' % chan)
                        logfile = make_nfs_path(
                            output_path.replace('.png', '.log'))
                        job_config = dict(infile=input_path,
                                          outfile=output_path,
                                          jprior=jprior,
                                          logfile=logfile,
                                          chan=chan)
                        job_configs[targ_key] = job_config

        return job_configs

class PlotStackedDM_SG(ScatterGather):
    """Small class to generate configurations for this script

    This adds the following arguments:
    """
    appname = 'dmpipe-plot-stacked-dm-sg'
    usage = "%s [options]" % (appname)
    description = "Make castro plots for set of targets"
    clientclass = PlotDM

    job_time = 60

    default_options = dict(ttype=defaults.common['ttype'],
                           rosterlist=defaults.common['targetlist'],
                           channels=defaults.common['channels'],
                           jpriors=defaults.common['jpriors'],
                           sim=defaults.sims['sim'],
                           nsims=defaults.sims['nsims'],
                           seed=defaults.sims['seed'],
                           dry_run=defaults.common['dry_run'])

    def build_job_configs(self, args):
        """Hook to build job configurations
        """
        job_configs = {}

        ttype = args['ttype']
        (roster_yaml, sim) = NAME_FACTORY.resolve_rosterfile(args)
        if roster_yaml is None:
            return job_configs

        roster_dict = load_yaml(roster_yaml)

        jpriors = args['jpriors']
        channels = args['channels']

        for roster_name in roster_dict.keys():
            for jprior in jpriors:
                name_keys = dict(target_type=ttype,
                                 roster_name=roster_name,
                                 jprior=jprior,
                                 sim_name=sim,
                                 fullpath=True)

                for chan in channels:
                    targ_key = "%s:%s:%s" % (roster_name, jprior, chan)

                    if sim is not None:
                        seedlist = range(
                            args['seed'], args['seed'] + args['nsims'])
                    else:
                        seedlist = [None]

                    for seed in seedlist:
                        if seed is not None:
                            name_keys['seed'] = "%06i" % seed
                            input_path = NAME_FACTORY.sim_resultsfile(
                                **name_keys)
                            full_targ_key = "%s_%06i" % (targ_key, seed)
                        else:
                            input_path = NAME_FACTORY.resultsfile(**name_keys)
                            full_targ_key = targ_key

                        output_path = input_path.replace(
                            '.fits', '_%s.png' % chan)
                        logfile = make_nfs_path(
                            output_path.replace('.png', '.log'))
                        job_config = dict(infile=input_path,
                                          outfile=output_path,
                                          jprior=jprior,
                                          logfile=logfile,
                                          chan=chan)
                        job_configs[full_targ_key] = job_config

        return job_configs

def register_classes():
    """Register these classes with the `LinkFactory` """
    PlotDMSpectra.register_class()
    PlotLimits.register_class()
    PlotLimits_SG.register_class()
    PlotDM.register_class()
    PlotDM_SG.register_class()
    PlotStackedDM_SG.register_class()
    PlotStackedLimits_SG.register_class()
<|endoftext|>


<commit_before>'''
File: sites.py
Author: Ivan Gonzalez
Description: The class for single sites
'''
from numpy import zeros, eye
from dmrg_exceptions import DMRGException

class Site(object):
    """A general single site

    You use this class to create a single site. The site comes empty,
    and you should add operators.

    Attributes:
        dim: An integer with the size of the Hilbert space
        operators: A dictionary with the operators for the site
    """
    def __init__(self, dim):
    	"""Creates an empty site of dimension dim

    	The dimension must be at least 1. A site of dim=1 represents the
	vaccum.

	Raises:
	    DMRGException: if dim < 1.

	Postcond:
	    The identity operator (ones in the diagonal, zeros elsewhere)
	    is added to the self.operators dictionary.
    	"""
    	if dim < 1:
    		raise DMRGException("Site dim msut be at least 1")
    	super(Site, self).__init__()
    	self.dim = dim
	self.operators={ "id" : eye(self.dim, self.dim) }

    def add_operator(self, operator_name):
    	"""Adds an operator to the site

    	Attributes:
    	    operator_name: a string with the operator name.

	Raises:
	    DMRGException: if operator_name is already in the dict.

    	Postcond:
    	    o self.operators has one item more, and
   	    o the newly created operator is a square matrix of
    	      self.dim times self.dim full of zeros.
            """
	if str(operator_name) in self.operators.keys():
    	    raise DMRGException("Operator name exists already")
    	else:
    	    self.operators[str(operator_name)]=zeros(self.dim, self.dim)

class SpinOneHalfSite(Site):
    """A site for spin 1/2 models

    You use this site for models where the single sites are spin
    one-half sites. The Hilbert space is ordered such as the first state
    is the spin down, and the second state is the spin up. Therefore e.g.
    you have the following relation between operator matrix elements:

    .. math:
        \langle \down_arrow| A | \up_arrow\rangle=A_{0,1}

    Postcond:
        The site has already built-in the spin operators for s_z, s_p,
	s_m.
    """
    def __init__(self):
        super(SpinOneHalfSite, self).__init__(2)
	# add the operators
        super(SpinOneHalfSite, self).operators.add["s_z"]
        super(SpinOneHalfSite, self).operators.add["s_p"]
        super(SpinOneHalfSite, self).operators.add["s_m"]
	# set the matrix elements different from zero to the right values
        self.operators["s_z"](0,0)=-1.0
        self.operators["s_z"](1,1)=1.0
        self.operators["s_p"](1,0)=1.0
        self.operators["s_m"](0,1)=1.0

class ElectronicSite(Site):
    """A site for electronic models

    You use this site for models where the single sites are electron
    sites. The Hilbert space is ordered such as:
        - the first state, labelled 0,  is the empty site,
	- the second, labelled 1, is spin down,
	- the third, labelled 2, is spin up, and
	- the fourth, labelled 3, is double occupancy.

    Postcond:
	The site has already built-in the spin operators for:
	    - c_up, destroys an spin up electron
	    - c_up_dag, creates an spin up electron
	    - c_down, destroys an spin down electron
	    - c_down_dag, creates an spin down electron
	    - s_z, component z of spin
	    - s_p, raises the component z of spin
	    - s_m, lowers the component z of spin
	    - n_up, number of electrons with spin up
	    - n_down, number of electrons with spin down
	    - n, number of electrons, i.e. n_up+n_down
	    - u, number of double occupancies, i.e. n_up*n_down
    """
    def __init__(self):
        super(ElectronicSiteSpinOneHalfSite, self).__init__(4)
	# add the operators
        super(ElectronicSite, self).operators.add["c_up"]
        super(ElectronicSite, self).operators.add["c_up_dag"]
        super(ElectronicSite, self).operators.add["c_down"]
        super(ElectronicSite, self).operators.add["c_down_dag"]
        super(ElectronicSite, self).operators.add["s_z"]
        super(ElectronicSite, self).operators.add["s_p"]
        super(ElectronicSite, self).operators.add["s_m"]
        super(ElectronicSite, self).operators.add["n_up"]
        super(ElectronicSite, self).operators.add["n_down"]
        super(ElectronicSite, self).operators.add["n"]
        super(ElectronicSite, self).operators.add["u"]
	# set the matrix elements different from zero to the right values
        self.operators["c_up"](0,2)=1.0
        self.operators["c_up"](1,3)=1.0
        self.operators["c_up_dag"](2,0)=1.0
        self.operators["c_up_dag"](3,1)=1.0
        self.operators["c_down"](0,1)=1.0
        self.operators["c_down"](2,3)=1.0
        self.operators["c_down_dag"](1,0)=1.0
        self.operators["c_down_dag"](3,2)=1.0
        self.operators["s_z"](1,1)=-1.0
        self.operators["s_z"](2,2)=1.0
	# todo: missing s_p, s_m
        self.operators["n_up"](2,2)=1.0
        self.operators["n_up"](3,3)=1.0
        self.operators["n_down"](1,1)=1.0
        self.operators["n_down"](3,3)=1.0
        self.operators["n"](1,1)=1.0
        self.operators["n"](2,2)=1.0
        self.operators["n"](3,3)=2.0
        self.operators["u"](3,3)=1.0
<commit_msg>Fixing the docs and some other problems with the sites.<commit_after>#
# File: sites.py
# Author: Ivan Gonzalez
#
""" A module for single sites.
"""
from numpy import zeros, eye
from dmrg_exceptions import DMRGException

class Site(object):
    """A general single site

    You use this class to create a single site. The site comes empty (i.e.
    with no operators included), but for th identity operator. You should
    add operators you need to make you site up.

    Parameters
    ----------
    dim : an int
	Size of the Hilbert space. The dimension must be at least 1. A site of
        dim = 1  represents the vaccum (or something strange like that, it's
        used for demo purposes mostly.)
    operators : dictionary of string and ndarray (with ndim = 2)
	Operators for the site.

    Examples
    --------
    >>> from dmrg101.core.sites import Site
    >>> brand_new_site = Site(2)
    >>> # the Hilbert space has dimension 2
    >>> print brand_new_site.dim
    2
    >>> # the only operator is the identity
    >>> print brand_new_site.operators
    {'id': array([[ 1.,  0.],
           [ 0.,  1.]])}
    """
    def __init__(self, dim):
    	"""Creates an empty site of dimension dim.

	Raises
	------
	DMRGException
	     if `dim` < 1.

	Notes
	-----
	Postcond : The identity operator (ones in the diagonal, zeros elsewhere)
	is added to the `self.operators` dictionary.
    	"""
    	if dim < 1:
    	    raise DMRGException("Site dim must be at least 1")
    	super(Site, self).__init__()
    	self.dim = dim
	self.operators={ "id" : eye(self.dim, self.dim) }

    def add_operator(self, operator_name):
    	"""Adds an operator to the site

    	Parameters
	----------
    	operator_name : string
	    The operator name.

	Raises
	------
	DMRGException
	    if `operator_name` is already in the dict.

	Notes
	-----
	Postcond:
    	    - `self.operators` has one item more, and
   	    - the newly created operator is a (`self.dim`, `self.dim`)
	      matrix of full of zeros.

	Examples
	--------
	>>> from dmrg101.core.sites import Site
	>>> new_site = Site(2)
	>>> print new_site.operators.keys()
	['id']
	>>> new_site.add_operator('s_z')
	>>> print new_site.operators.keys()
	['s_z', 'id']
	>>> # note that the newly created op has all zeros
	>>> print new_site.operators['s_z']
	[[ 0.  0.]
	 [ 0.  0.]]
        """
	if str(operator_name) in self.operators.keys():
    	    raise DMRGException("Operator name exists already")
    	else:
    	    self.operators[str(operator_name)] = zeros((self.dim, self.dim))

class SpinOneHalfSite(Site):
    """A site for spin 1/2 models.

    You use this site for models where the single sites are spin
    one-half sites. The Hilbert space is ordered such as the first state
    is the spin down, and the second state is the spin up. Therefore e.g.
    you have the following relation between operator matrix elements:

    .. math::
        \langle \downarrow \left| A \\right|\uparrow \\rangle = A_{0,1}

    Notes
    -----
        Postcond : The site has already built-in the spin operators for s_z, s_p, s_m.

    Examples
    --------
    >>> from dmrg101.core.sites import SpinOneHalfSite
    >>> spin_one_half_site = SpinOneHalfSite()
    >>> # check all it's what you expected
    >>> print spin_one_half_site.dim
    2
    >>> print spin_one_half_site.operators.keys()
    ['s_p', 's_z', 's_m', 'id']
    >>> print spin_one_half_site.operators['s_z']
    [[-1.  0.]
     [ 0.  1.]]
    >>> print spin_one_half_site.operators['s_p']
    [[ 0.  0.]
     [ 1.  0.]]
    >>> print spin_one_half_site.operators['s_m']
    [[ 0.  1.]
     [ 0.  0.]]
    """
    def __init__(self):
	"""Creates the spin one-half site.

	Notes
	-----
	    Postcond : the dimension is set to 2, and the Pauli matrices
	    are added as operators.
	"""
        super(SpinOneHalfSite, self).__init__(2)
	# add the operators
        self.add_operator("s_z")
        self.add_operator("s_p")
        self.add_operator("s_m")
	# for clarity
        s_z = self.operators["s_z"]
        s_p = self.operators["s_p"]
        s_m = self.operators["s_m"]
	# set the matrix elements different from zero to the right values
        s_z[0, 0] = -1.0
        s_z[1, 1] = 1.0
        s_p[1, 0] = 1.0
        s_m[0, 1] = 1.0

class ElectronicSite(Site):
    """A site for electronic models

    You use this site for models where the single sites are electron
    sites. The Hilbert space is ordered such as:
        - the first state, labelled 0,  is the empty site,
	- the second, labelled 1, is spin down,
	- the third, labelled 2, is spin up, and
	- the fourth, labelled 3, is double occupancy.

    Notes
    -----
    Postcond: The site has already built-in the spin operators for:
        - c_up, destroys an spin up electron,
        - c_up_dag, creates an spin up electron,
        - c_down, destroys an spin down electron,
        - c_down_dag, creates an spin down electron,
        - s_z, component z of spin,
        - s_p, raises the component z of spin,
        - s_m, lowers the component z of spin,
        - n_up, number of electrons with spin up,
        - n_down, number of electrons with spin down,
        - n, number of electrons, i.e. n_up+n_down, and
        - u, number of double occupancies, i.e. n_up*n_down.

    Examples
    --------
    >>> from dmrg101.core.sites import ElectronicSite
    >>> hubbard_site = ElectronicSite()
    >>> # check all it's what you expected
    >>> print hubbard_site.dim
    4
    >>> print hubbard_site.operators.keys() # doctest: +ELLIPSIS
    ['s_p', ...]
    >>> print hubbard_site.operators['n_down']
    [[ 0.  0.  0.  0.]
     [ 0.  1.  0.  0.]
     [ 0.  0.  0.  0.]
     [ 0.  0.  0.  1.]]
    >>> print hubbard_site.operators['n_up']
    [[ 0.  0.  0.  0.]
     [ 0.  0.  0.  0.]
     [ 0.  0.  1.  0.]
     [ 0.  0.  0.  1.]]
    >>> print hubbard_site.operators['u']
    [[ 0.  0.  0.  0.]
     [ 0.  0.  0.  0.]
     [ 0.  0.  0.  0.]
     [ 0.  0.  0.  1.]]
    """
    def __init__(self):
        super(ElectronicSite, self).__init__(4)
	# add the operators
        self.add_operator("c_up")
        self.add_operator("c_up_dag")
        self.add_operator("c_down")
        self.add_operator("c_down_dag")
        self.add_operator("s_z")
        self.add_operator("s_p")
        self.add_operator("s_m")
        self.add_operator("n_up")
        self.add_operator("n_down")
        self.add_operator("n")
        self.add_operator("u")
	# for clarity
        c_up = self.operators["c_up"]
        c_up_dag = self.operators["c_up_dag"]
        c_down = self.operators["c_down"]
        c_down_dag = self.operators["c_down_dag"]
        s_z = self.operators["s_z"]
        s_p = self.operators["s_p"]
        s_m = self.operators["s_m"]
        n_up = self.operators["n_up"]
        n_down = self.operators["n_down"]
        n = self.operators["n"]
        u = self.operators["u"]
	# set the matrix elements different from zero to the right values
	# TODO: missing s_p, s_m
        c_up[0,2] = 1.0
        c_up[1,3] = 1.0
        c_up_dag[2,0] = 1.0
        c_up_dag[3,1] = 1.0
        c_down[0,1] = 1.0
        c_down[2,3] = 1.0
        c_down_dag[1,0] = 1.0
        c_down_dag[3,2] = 1.0
        s_z[1,1] = -1.0
        s_z[2,2] = 1.0
        n_up[2,2] = 1.0
        n_up[3,3] = 1.0
        n_down[1,1] = 1.0
        n_down[3,3] = 1.0
        n[1,1] = 1.0
        n[2,2] = 1.0
        n[3,3] = 2.0
        u[3,3] = 1.0
<|endoftext|>


<commit_before>"""
Load Motor Imagery example data
"""
import numpy as np
from os.path import abspath, dirname, join
from scot.matfiles import loadmat
from eegtopo.eegpos3d import positions as eeg_locations

def __load():
    matfile = loadmat(join(abspath(dirname(__file__)), 'motorimagery.mat'))['s0']

    class Data: pass

    Data.description = """
                       The Data set contains a continuous 45 channel EEG recording of a motor imagery
                       experiment. The Data was preprocessed to reduce eye movement artifacts and
                       resampled to a sampling rate of 100 Hz.
                       With a visual cue the subject was instructed to perform either hand of foot
                       motor imagery. The the trigger time points of the cues are stored in 'tr', and
                       'cl' contains the class labels (hand: 1, foot: -1). Duration of the motor
                       imagery period was approximately 6 seconds.
                       """

    Data.samplerate = matfile['fs'] # Sampling rate
    Data.num_trials = matfile['T']  # Number of trials
    Data.triggers = matfile['tr']   # Trigger locations
    Data.eeg = matfile['eeg']       # EEG Data

    # Class labels
    cltrans = { -1 : 'hand', 1 : 'foot' }
    Data.classes = np.array([cltrans[c] for c in matfile['cl']])

    # Unfortunately, the EEG channel labels are not stored in the file, so we set them manually.
    Data.labels = ['AF7', 'AFz', 'AF8', 'F3', 'F1',
                   'Fz', 'F2', 'F4', 'FT7', 'FC5',
                   'FC3', 'FC1', 'FCz', 'FC2', 'FC4',
                   'FC6', 'FT8', 'C5', 'C3', 'C1',
                   'Cz', 'C2', 'C4', 'C6', 'CP5',
                   'CP3', 'CP1', 'CPz', 'CP2', 'CP4',
                   'CP6', 'P7', 'P3', 'Pz', 'P4',
                   'P8', 'PO3', 'POz', 'PO4', 'O1',
                   'Oz', 'O2', 'O9', 'Iz', 'O10']

    # Obtain default electrode locations corresponding to the channels
    Data.locations = [[v for v in eeg_locations[l].vector] for l in Data.labels]

    return Data

data = __load()

<commit_msg>Fixed example class labels<commit_after>"""
Load Motor Imagery example data
"""
import numpy as np
from os.path import abspath, dirname, join
from scot.matfiles import loadmat
from eegtopo.eegpos3d import positions as eeg_locations

def __load():
    matfile = loadmat(join(abspath(dirname(__file__)), 'motorimagery.mat'))['s0']

    class Data: pass

    Data.description = """
                       The Data set contains a continuous 45 channel EEG recording of a motor imagery
                       experiment. The Data was preprocessed to reduce eye movement artifacts and
                       resampled to a sampling rate of 100 Hz.
                       With a visual cue the subject was instructed to perform either hand of foot
                       motor imagery. The the trigger time points of the cues are stored in 'tr', and
                       'cl' contains the class labels (hand: 1, foot: -1). Duration of the motor
                       imagery period was approximately 6 seconds.
                       """

    Data.samplerate = matfile['fs'] # Sampling rate
    Data.num_trials = matfile['T']  # Number of trials
    Data.triggers = matfile['tr']   # Trigger locations
    Data.eeg = matfile['eeg']       # EEG Data

    # Class labels
    cltrans = { 1 : 'hand', -1 : 'foot' }
    Data.classes = np.array([cltrans[c] for c in matfile['cl']])

    # Unfortunately, the EEG channel labels are not stored in the file, so we set them manually.
    Data.labels = ['AF7', 'AFz', 'AF8', 'F3', 'F1',
                   'Fz', 'F2', 'F4', 'FT7', 'FC5',
                   'FC3', 'FC1', 'FCz', 'FC2', 'FC4',
                   'FC6', 'FT8', 'C5', 'C3', 'C1',
                   'Cz', 'C2', 'C4', 'C6', 'CP5',
                   'CP3', 'CP1', 'CPz', 'CP2', 'CP4',
                   'CP6', 'P7', 'P3', 'Pz', 'P4',
                   'P8', 'PO3', 'POz', 'PO4', 'O1',
                   'Oz', 'O2', 'O9', 'Iz', 'O10']

    # Obtain default electrode locations corresponding to the channels
    Data.locations = [[v for v in eeg_locations[l].vector] for l in Data.labels]

    return Data

data = __load()

<|endoftext|>


<commit_before>from copy import copy
from collections import defaultdict

from . import _hoomd
from hoomd.box import Box
from hoomd.snapshot import Snapshot
from hoomd.data import LocalSnapshot, LocalSnapshotGPU

def _create_domain_decomposition(device, box):
    """ Create a default domain decomposition.

    This method is a quick hack to get basic MPI simulations working with
    the new API. We will need to consider designing an appropriate user-facing
    API to set the domain decomposition.
    """
    if not _hoomd.is_MPI_available():
        return None

    # if we are only running on one processor, we use optimized code paths
    # for single-GPU execution
    if device.comm.num_ranks == 1:
        return None

    # create a default domain decomposition
    result = _hoomd.DomainDecomposition(device.cpp_exec_conf,
                                        box.getL(),
                                        0,
                                        0,
                                        0,
                                        False)

    return result

class State:
    R""" Simulation state.

    Args:
        snapshot
    """

    def __init__(self, simulation, snapshot):
        self._simulation = simulation
        self._in_context_manager = False
        snapshot._broadcast_box()
        domain_decomp = _create_domain_decomposition(
            simulation.device,
            snapshot._cpp_obj._global_box)

        if domain_decomp is not None:
            self._cpp_sys_def = _hoomd.SystemDefinition(
                snapshot._cpp_obj, simulation.device.cpp_exec_conf,
                domain_decomp)
        else:
            self._cpp_sys_def = _hoomd.SystemDefinition(
                snapshot._cpp_obj, simulation.device.cpp_exec_conf)
        self._groups = defaultdict(dict)

    @property
    def snapshot(self):
        cpp_snapshot = self._cpp_sys_def.takeSnapshot_double()
        return Snapshot._from_cpp_snapshot(cpp_snapshot,
                                           self._simulation.device.comm)

    @snapshot.setter
    def snapshot(self, snapshot):
        R""" Re-initializes the system from a snapshot.

        Args:
            snapshot:. The snapshot to initialize the system from.

        Snapshots temporarily store system data. Snapshots contain the complete
        simulation state in a single object. They can be used to restart a
        simulation.

        Example use cases in which a simulation may be restarted from a snapshot
        include python-script-level Monte-Carlo schemes, where the system state
        is stored after a move has been accepted (according to some criterion),
        and where the system is re-initialized from that same state in the case
        when a move is not accepted.

        Example::

            system = init.read_xml("some_file.xml")

            ... run a simulation ...

            snapshot = system.take_snapshot(all=True)
            ...
            system.restore_snapshot(snapshot)

        Warning:
                restore_snapshot() may invalidate force coefficients,
                neighborlist r_cut values, and other per type quantities if
                called within a callback during a run(). You can restore a
                snapshot during a run only if the snapshot is of a previous
                state of the currently running system. Otherwise, you need to
                use restore_snapshot() between run() commands to ensure that all
                per type coefficients are updated properly.

        """

        if self._simulation.device.comm.rank == 0:
            if len(snapshot.particles.types) != len(self.particle_types):
                raise RuntimeError(
                    "Number of particle types must remain the same")
            if len(snapshot.bonds.types) != len(self.bond_types):
                raise RuntimeError("Number of bond types must remain the same")
            if len(snapshot.angles.types) != len(self.angle_types):
                raise RuntimeError(
                    "Number of angle types must remain the same")
            if len(snapshot.dihedrals.types) != len(self.dihedral_types):
                raise RuntimeError(
                    "Number of dihedral types must remain the same")
            if len(snapshot.impropers.types) != len(self.improper_types):
                raise RuntimeError(
                    "Number of dihedral types must remain the same")
            if len(snapshot.pairs.types) != len(self.special_pair_types):
                raise RuntimeError("Number of pair types must remain the same")

        self._cpp_sys_def.initializeFromSnapshot(snapshot._cpp_obj)

    @property
    def types(self):
        return dict(particle_types=self.particle_types,
                    bond_types=self.bond_types,
                    angle_types=self.angles_types,
                    dihedral_types=self.dihedral_types,
                    improper_types=self.improper_types,
                    special_pair_types=self.special_pair_types
                    )

    @property
    def particle_types(self):
        return self._cpp_sys_def.getParticleData().getTypes()

    @property
    def bond_types(self):
        return self._cpp_sys_def.getBondData().getTypes()

    @property
    def angle_types(self):
        return self._cpp_sys_def.getAngleData().getTypes()

    @property
    def dihedral_types(self):
        return self._cpp_sys_def.getDihedralData().getTypes()

    @property
    def improper_types(self):
        return self._cpp_sys_def.getImproperData().getTypes()

    @property
    def special_pair_types(self):
        return self._cpp_sys_def.getPairData().getTypes()

    @property
    def box(self):
        """The state's box (a :py:class:`hoomd.Box` object).

        Ediing the box directly is not allowed.  For example
        ``state.box.scale(1.1)`` would not scale the state's box. To set the
        state's box to a new box ``state.box = new_box`` must be used.
        """
        b = Box._from_cpp(self._cpp_sys_def.getParticleData().getGlobalBox())
        return Box.from_box(b)

    @box.setter
    def box(self, value):
        if self._in_context_manager:
            raise RuntimeError(
                "Cannot set system box within local snapshot context manager.")
        try:
            value = Box.from_box(value)
        except Exception:
            raise ValueError('{} is not convertable to hoomd.Box using '
                             'hoomd.Box.from_box'.format(value))

        if value.dimensions != self._cpp_sys_def.getNDimensions():
            self._simulation.device.cpp_msg.warning(
                "Box changing dimensions from {} to {}."
                "".format(self._cpp_sys_def.getNDimensions(),
                          value.dimensions))
            self._cpp_sys_def.setNDimensions(value.dimensions)
        self._cpp_sys_def.getParticleData().setGlobalBox(value._cpp_obj)

    def replicate(self):
        raise NotImplementedError

    def scale_system(self):
        raise NotImplementedError

    def get_group(self, filter_):
        cls = filter_.__class__
        if filter_ in self._groups[cls]:
            return self._groups[cls][filter_]
        else:
            group = _hoomd.ParticleGroup(self._cpp_sys_def, filter_)
            self._groups[cls][filter_] = group
            return group

    def update_group_dof(self):
        """Update the number of degrees of freedom in each group.

        The groups of particles selected by filters each need to know the number
        of degrees of freedom given to that group by the simulation's
        Integrator. This method is called automatically when:

        * The Integrator is attached to the simulation

        Call it manually to force an update.
        """
        integrator = self._simulation.operations.integrator

        for groups in self._groups.values():
            for group in groups.values():
                if integrator is not None:
                    if not integrator.is_attached:
                        raise RuntimeError(
                            "Call update_group_dof after attaching")

                    integrator._cpp_obj.updateGroupDOF(group)
                else:
                    group.setTranslationalDOF(0)
                    group.setRotationalDOF(0)

    @property
    def cpu_local_snapshot(self):
        """hoomd.data.LocalSnapshot: Directly expose HOOMD data buffers
        on the CPU.

        Provides access directly to the system state's particle, bond, angle,
        dihedral, improper, constaint, and pair data through a context manager.
        The `hoomd.data.LocalSnapshot` object is only usable within a
        context manager (i.e. ``with sim.state.cpu_local_snapshot as data:``)
        The interface is similar to that of the `hoomd.Snapshot`. Data is local
        to a given MPI rank. The returned arrays are `hoomd.array.HOOMDArray`
        objects. Through this interface zero-copy access is available (access is
        guarenteed to be zero-copy when running on CPU, may be zero copy if
        running on GPU).

        Changing the data in the buffers exposed by the local snapshot will
        change the data across the HOOMD-blue simulation. For a trivial example,
        this example would set all particle z-axis positions to 0.

        .. code-block:: python

            with sim.state.cpu_local_snapshot as data:
                data.particles.position[:, 2] = 0

        Note:
            The state's box and the number of particles, bonds, angles,
            dihedrals, impropers, constaints, and pairs cannot change within the
            context manager.
        """
        if self._in_context_manager:
            raise RuntimeError(
                "Cannot enter cpu_local_snapshot context manager inside "
                "another local_snapshot context manager.")
        return LocalSnapshot(self)

    @property
    def gpu_local_snapshot(self):
        """hoomd.data.LocalSnapshotGPU: Directly expose HOOMD data
        buffers on the GPU.

        Provides access directly to the system state's particle, bond, angle,
        dihedral, improper, constaint, and pair data through a context manager.
        The `hoomd.data.LocalSnapshotGPU` object is only usable
        within a context manager (i.e. ``with sim.state.gpu_local_snapshot as
        data:``) The interface is similar to that of the `hoomd.Snapshot`. Data
        is local to a given MPI rank. The returned arrays are
        `hoomd.array.HOOMDGPUArray` objects. Through this interface potential
        zero-copy access is available (access cannot be guarenteed to be
        zero-copy, but will be if the most recent copy of the data is on the
        GPU).

        Changing the data in the buffers exposed by the local snapshot will
        change the data across the HOOMD-blue simulation. For a trivial example,
        this example would set all particle z-axis positions to 0.

        .. code-block:: python

            with sim.state.gpu_local_snapshot as data:
                data.particles.position[:, 2] = 0

        Note:
            The state's box and the number of particles, bonds, angles,
            dihedrals, impropers, constaints, and pairs cannot change within the
            context manager.

        Note:
            This property is only available when running on a GPU(s).
        """
        if self._simulation.device.mode != 'gpu':
            raise RuntimeError(
                "Cannot access gpu_snapshot with a non GPU device.")
        elif self._in_context_manager:
            raise RuntimeError(
                "Cannot enter gpu_local_snapshot context manager inside "
                "another local_snapshot context manager.")
        else:
            return LocalSnapshotGPU(self)
<commit_msg>Add N_particles and related methods to State<commit_after>from copy import copy
from collections import defaultdict

from . import _hoomd
from hoomd.box import Box
from hoomd.snapshot import Snapshot
from hoomd.data import LocalSnapshot, LocalSnapshotGPU

def _create_domain_decomposition(device, box):
    """ Create a default domain decomposition.

    This method is a quick hack to get basic MPI simulations working with
    the new API. We will need to consider designing an appropriate user-facing
    API to set the domain decomposition.
    """
    if not _hoomd.is_MPI_available():
        return None

    # if we are only running on one processor, we use optimized code paths
    # for single-GPU execution
    if device.comm.num_ranks == 1:
        return None

    # create a default domain decomposition
    result = _hoomd.DomainDecomposition(device.cpp_exec_conf,
                                        box.getL(),
                                        0,
                                        0,
                                        0,
                                        False)

    return result

class State:
    R""" Simulation state.

    Args:
        snapshot
    """

    def __init__(self, simulation, snapshot):
        self._simulation = simulation
        self._in_context_manager = False
        snapshot._broadcast_box()
        domain_decomp = _create_domain_decomposition(
            simulation.device,
            snapshot._cpp_obj._global_box)

        if domain_decomp is not None:
            self._cpp_sys_def = _hoomd.SystemDefinition(
                snapshot._cpp_obj, simulation.device.cpp_exec_conf,
                domain_decomp)
        else:
            self._cpp_sys_def = _hoomd.SystemDefinition(
                snapshot._cpp_obj, simulation.device.cpp_exec_conf)
        self._groups = defaultdict(dict)

    @property
    def snapshot(self):
        cpp_snapshot = self._cpp_sys_def.takeSnapshot_double()
        return Snapshot._from_cpp_snapshot(cpp_snapshot,
                                           self._simulation.device.comm)

    @snapshot.setter
    def snapshot(self, snapshot):
        R""" Re-initializes the system from a snapshot.

        Args:
            snapshot:. The snapshot to initialize the system from.

        Snapshots temporarily store system data. Snapshots contain the complete
        simulation state in a single object. They can be used to restart a
        simulation.

        Example use cases in which a simulation may be restarted from a snapshot
        include python-script-level Monte-Carlo schemes, where the system state
        is stored after a move has been accepted (according to some criterion),
        and where the system is re-initialized from that same state in the case
        when a move is not accepted.

        Example::

            system = init.read_xml("some_file.xml")

            ... run a simulation ...

            snapshot = system.take_snapshot(all=True)
            ...
            system.restore_snapshot(snapshot)

        Warning:
                restore_snapshot() may invalidate force coefficients,
                neighborlist r_cut values, and other per type quantities if
                called within a callback during a run(). You can restore a
                snapshot during a run only if the snapshot is of a previous
                state of the currently running system. Otherwise, you need to
                use restore_snapshot() between run() commands to ensure that all
                per type coefficients are updated properly.

        """

        if self._simulation.device.comm.rank == 0:
            if len(snapshot.particles.types) != len(self.particle_types):
                raise RuntimeError(
                    "Number of particle types must remain the same")
            if len(snapshot.bonds.types) != len(self.bond_types):
                raise RuntimeError("Number of bond types must remain the same")
            if len(snapshot.angles.types) != len(self.angle_types):
                raise RuntimeError(
                    "Number of angle types must remain the same")
            if len(snapshot.dihedrals.types) != len(self.dihedral_types):
                raise RuntimeError(
                    "Number of dihedral types must remain the same")
            if len(snapshot.impropers.types) != len(self.improper_types):
                raise RuntimeError(
                    "Number of dihedral types must remain the same")
            if len(snapshot.pairs.types) != len(self.special_pair_types):
                raise RuntimeError("Number of pair types must remain the same")

        self._cpp_sys_def.initializeFromSnapshot(snapshot._cpp_obj)

    @property
    def types(self):
        return dict(particle_types=self.particle_types,
                    bond_types=self.bond_types,
                    angle_types=self.angles_types,
                    dihedral_types=self.dihedral_types,
                    improper_types=self.improper_types,
                    special_pair_types=self.special_pair_types
                    )

    @property
    def N_particles(self):
        """int: The number of particles in the simulation."""
        return self._cpp_sys_def.getParticleData().getNGlobal()

    @property
    def N_bonds(self):
        """int: The number of bonds in the simulation."""
        return self._cpp_sys_def.getBondData().getNGlobal()

    @property
    def N_angles(self):
        """int: The number of angles in the simulation."""
        return self._cpp_sys_def.getAngleData().getNGlobal()

    @property
    def N_impropers(self):
        """int: The number of impropers in the simulation."""
        return self._cpp_sys_def.getImproperData().getNGlobal()

    @property
    def N_special_pairs(self):
        """int: The number of special pairs in the simulation."""
        return self._cpp_sys_def.getPairData().getNGlobal()

    @property
    def N_dihedrals(self):
        """int: The number of dihedrals in the simulation."""
        return self._cpp_sys_def.getDihedralData().getNGlobal()

    @property
    def particle_types(self):
        return self._cpp_sys_def.getParticleData().getTypes()

    @property
    def bond_types(self):
        return self._cpp_sys_def.getBondData().getTypes()

    @property
    def angle_types(self):
        return self._cpp_sys_def.getAngleData().getTypes()

    @property
    def dihedral_types(self):
        return self._cpp_sys_def.getDihedralData().getTypes()

    @property
    def improper_types(self):
        return self._cpp_sys_def.getImproperData().getTypes()

    @property
    def special_pair_types(self):
        return self._cpp_sys_def.getPairData().getTypes()

    @property
    def box(self):
        """The state's box (a :py:class:`hoomd.Box` object).

        Ediing the box directly is not allowed.  For example
        ``state.box.scale(1.1)`` would not scale the state's box. To set the
        state's box to a new box ``state.box = new_box`` must be used.
        """
        b = Box._from_cpp(self._cpp_sys_def.getParticleData().getGlobalBox())
        return Box.from_box(b)

    @box.setter
    def box(self, value):
        if self._in_context_manager:
            raise RuntimeError(
                "Cannot set system box within local snapshot context manager.")
        try:
            value = Box.from_box(value)
        except Exception:
            raise ValueError('{} is not convertable to hoomd.Box using '
                             'hoomd.Box.from_box'.format(value))

        if value.dimensions != self._cpp_sys_def.getNDimensions():
            self._simulation.device.cpp_msg.warning(
                "Box changing dimensions from {} to {}."
                "".format(self._cpp_sys_def.getNDimensions(),
                          value.dimensions))
            self._cpp_sys_def.setNDimensions(value.dimensions)
        self._cpp_sys_def.getParticleData().setGlobalBox(value._cpp_obj)

    def replicate(self):
        raise NotImplementedError

    def scale_system(self):
        raise NotImplementedError

    def get_group(self, filter_):
        cls = filter_.__class__
        if filter_ in self._groups[cls]:
            return self._groups[cls][filter_]
        else:
            group = _hoomd.ParticleGroup(self._cpp_sys_def, filter_)
            self._groups[cls][filter_] = group
            return group

    def update_group_dof(self):
        """Update the number of degrees of freedom in each group.

        The groups of particles selected by filters each need to know the number
        of degrees of freedom given to that group by the simulation's
        Integrator. This method is called automatically when:

        * The Integrator is attached to the simulation

        Call it manually to force an update.
        """
        integrator = self._simulation.operations.integrator

        for groups in self._groups.values():
            for group in groups.values():
                if integrator is not None:
                    if not integrator.is_attached:
                        raise RuntimeError(
                            "Call update_group_dof after attaching")

                    integrator._cpp_obj.updateGroupDOF(group)
                else:
                    group.setTranslationalDOF(0)
                    group.setRotationalDOF(0)

    @property
    def cpu_local_snapshot(self):
        """hoomd.data.LocalSnapshot: Directly expose HOOMD data buffers
        on the CPU.

        Provides access directly to the system state's particle, bond, angle,
        dihedral, improper, constaint, and pair data through a context manager.
        The `hoomd.data.LocalSnapshot` object is only usable within a
        context manager (i.e. ``with sim.state.cpu_local_snapshot as data:``)
        The interface is similar to that of the `hoomd.Snapshot`. Data is local
        to a given MPI rank. The returned arrays are `hoomd.array.HOOMDArray`
        objects. Through this interface zero-copy access is available (access is
        guarenteed to be zero-copy when running on CPU, may be zero copy if
        running on GPU).

        Changing the data in the buffers exposed by the local snapshot will
        change the data across the HOOMD-blue simulation. For a trivial example,
        this example would set all particle z-axis positions to 0.

        .. code-block:: python

            with sim.state.cpu_local_snapshot as data:
                data.particles.position[:, 2] = 0

        Note:
            The state's box and the number of particles, bonds, angles,
            dihedrals, impropers, constaints, and pairs cannot change within the
            context manager.
        """
        if self._in_context_manager:
            raise RuntimeError(
                "Cannot enter cpu_local_snapshot context manager inside "
                "another local_snapshot context manager.")
        return LocalSnapshot(self)

    @property
    def gpu_local_snapshot(self):
        """hoomd.data.LocalSnapshotGPU: Directly expose HOOMD data
        buffers on the GPU.

        Provides access directly to the system state's particle, bond, angle,
        dihedral, improper, constaint, and pair data through a context manager.
        The `hoomd.data.LocalSnapshotGPU` object is only usable
        within a context manager (i.e. ``with sim.state.gpu_local_snapshot as
        data:``) The interface is similar to that of the `hoomd.Snapshot`. Data
        is local to a given MPI rank. The returned arrays are
        `hoomd.array.HOOMDGPUArray` objects. Through this interface potential
        zero-copy access is available (access cannot be guarenteed to be
        zero-copy, but will be if the most recent copy of the data is on the
        GPU).

        Changing the data in the buffers exposed by the local snapshot will
        change the data across the HOOMD-blue simulation. For a trivial example,
        this example would set all particle z-axis positions to 0.

        .. code-block:: python

            with sim.state.gpu_local_snapshot as data:
                data.particles.position[:, 2] = 0

        Note:
            The state's box and the number of particles, bonds, angles,
            dihedrals, impropers, constaints, and pairs cannot change within the
            context manager.

        Note:
            This property is only available when running on a GPU(s).
        """
        if self._simulation.device.mode != 'gpu':
            raise RuntimeError(
                "Cannot access gpu_snapshot with a non GPU device.")
        elif self._in_context_manager:
            raise RuntimeError(
                "Cannot enter gpu_local_snapshot context manager inside "
                "another local_snapshot context manager.")
        else:
            return LocalSnapshotGPU(self)
<|endoftext|>


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

from config import AWS_KEY, SECRET_KEY
from amazonproduct import API

#: a list of root nodes retrieved from
#: http://docs.amazonwebservices.com/AWSECommerceService/2009-11-01/DG/index.html?BrowseNodeIDs.html
NODE_IDS = {
    'ca' : {
        'Books' : 927726,
        'Classical' : 962454,
        'DVD' : 14113311,
        'Electronics' : 677211011,
        'ForeignBooks' : 927726,
        'Kitchen' : 2206275011,
        'Music' : 962454,
        'Software' : 3234171,
        'SoftwareVideoGames' : 3323751,
        'VHS' : 962072,
        'Video' : 962454,
        'VideoGames' : 110218011,
    },
    'de' : {
        'Apparel' : 78689031,
        'Automotive' : 78194031,
        'Baby' : 357577011,
        'Beauty' : 64257031,
        'Books' : 541686,
        'Classical' : 542676,
        'DVD' : 547664,
        'Electronics' : 569604,
        'ForeignBooks' : 54071011,
        'HealthPersonalCare' : 64257031,
        'HomeGarden' : 10925241,
        'Jewelry' : 327473011,
        'Kitchen' : 3169011,
        'Lighting' : 213083031,
        'Magazines' : 1198526,
        'MP3Downloads' : 77256031,
        'Music' : 542676,
        'OfficeProducts' : 16291311,
        'OutdoorLiving' : 10925051,
        'PCHardware' : 569604,
        'Photo' : 569604,
        'Software' : 542064,
        'SoftwareVideoGames' : 541708,
        'SportingGoods' : 16435121,
        'Toys' : 12950661,
        'VHS' : 547082,
        'Video' : 547664,
        'VideoGames' : 541708,
        'Watches' : 193708031,
    },
    'fr' : {
        'Baby' : 206617031,
        'Beauty' : 197858031,
        'Books' : 468256,
        'Classical' : 537366,
        'DVD' : 578608,
        'Electronics' : 1058082,
        'ForeignBooks' : 69633011,
        'HealthPersonalCare' : 197861031,
        'Jewelry' : 193711031,
        'Kitchen' : 57686031,
        'Lighting' : 213080031,
        'MP3Downloads' : 206442031,
        'Music' : 537366,
        'OfficeProducts' : 192420031,
        'Shoes' : 215934031,
        'Software' : 548012,
        'SoftwareVideoGames' : 548014,
        'Toys' : 548014,
        'VHS' : 578610,
        'Video' : 578608,
        'VideoGames' : 548014,
        'Watches' : 60937031,
    },
    'jp' : {
        'Apparel' : 361299011,
        'Automotive' : 2017304051,
        'Baby' : 13331821,
        'Beauty' : 52391051,
        'Books' : 465610,
        'Classical' : 562032,
        'DVD' : 562002,
        'Electronics' : 3210991,
        'ForeignBooks' : 388316011,
        'Grocery' : 57240051,
        'HealthPersonalCare' : 161669011,
        'Hobbies' : 13331821,
        'HomeImprovement' : 2016929051,
        'Jewelry' : 85896051,
        'Kitchen' : 3839151,
        'Music' : 562032,
        'Shoes' : 2016926051,
        'Software' : 637630,
        'SportingGoods' : 14315361,
        'Toys' : 13331821,
        'VHS' : 561972,
        'Video' : 561972,
        'VideoGames' : 637872,
        'Watches' : 14315361,
    },
    'uk' : {
        'Apparel' : 83451031,
        'Automotive' : 248877031,
        'Baby' : 60032031,
        'Beauty' : 66280031,
        'Books' : 1025612,
        'Classical' : 505510,
        'DVD' : 283926,
        'Electronics' : 560800,
        'HealthPersonalCare' : 66280031,
        'HomeGarden' : 11052591,
        'Jewelry' : 193717031,
        'Kitchen' : 11052591,
        'Lighting' : 213077031,
        'MP3Downloads' : 77198031,
        'Music' : 505510,
        'OfficeProducts' : 560800,
        'OutdoorLiving' : 11052591,
        'Software' : 1025614,
        'SoftwareVideoGames' : 1025616,
        'SportingGoods' : 319530011,
        'Tools' : 11052591,
        'Toys' : 712832,
        'VHS' : 283926,
        'Video' : 283926,
        'VideoGames' : 1025616,
        'Watches' : 595312,
    },
    'us' : {
        'Apparel' : 1036682,
        'Automotive' : 15690151,
        'Baby' : 1036682,
        'Beauty' : 11055981,
        'Books' : 1000,
        'Classical' : 301668,
        'DigitalMusic' : 195208011,
        'DVD' : 130,
        'Electronics' : 493964,
        'GourmetFood' : 3580501,
        'Grocery' : 3760931,
        'HealthPersonalCare' : 3760931,
        'HomeGarden' : 285080,
        'Industrial' : 228239,
        'Jewelry' : 3880591,
        'Kitchen' : 1063498,
        'Magazines' : 599872,
        'Merchants' : 493964,
        'Miscellaneous' : 10304191,
        'MP3Downloads' : 195211011,
        'Music' : 301668,
        'MusicalInstruments' : 11965861,
        'OfficeProducts' : 1084128,
        'OutdoorLiving' : 1063498,
        'PCHardware' : 493964,
        'PetSupplies' : 1063498,
        'Photo' : 493964,
        'Software' : 409488,
        'SportingGoods' : 1079730,
        'Tools' : 468240,
        'Toys' : 493964,
        'VHS' : 404272,
        'Video' : 130,
        'VideoGames' : 493964,
        'Watches' : 1079730,
        'Wireless' : 508494,
        'WirelessAccessories' : 13900851,
    },
}

class BrowseNodeExplorer (gtk.Window):

    """
    Gtk explorer for Amazon BrowseNodes.
    """

    def on_delete(self, widget, event, data=None):
        # closes the window and quits.
        gtk.main_quit()
        return False

    def on_tree_click(self, widget, event, data=None):
        # if double click
        if event.type == gtk.gdk._2BUTTON_PRESS:

            # get data from highlighted selection
            treeselection = self.treeview.get_selection()
            model, iter = treeselection.get_selected()
            name_of_data = self.treestore.get_value(iter, 0)

            # and fetch selected node
            self.fetch_nodes(name_of_data)

    def __init__(self, locale='de'):

        gtk.Window.__init__(self, gtk.WINDOW_TOPLEVEL)

        self.set_title("BrowseNode Explorer")
        self.set_size_request(400, 200)
        self.connect("delete_event", self.on_delete)

        self.locale = locale
        self.api = API(AWS_KEY, SECRET_KEY, self.locale)

        # create a TreeStore with one string column to use as the model
        self.treestore = gtk.TreeStore(int, str)

        # create the TreeView using treestore
        self.treeview = gtk.TreeView(self.treestore)

        # add column id
        renderer = gtk.CellRendererText()
        column = gtk.TreeViewColumn('id', renderer, text=0)
        self.treeview.append_column(column)

        # add column name
        renderer = gtk.CellRendererText()
        column = gtk.TreeViewColumn('name', renderer, text=1)
        column.set_sort_column_id(1) # Allow sorting on the column
        self.treeview.append_column(column)

        # make it clickable
        self.treeview.add_events(gtk.gdk.BUTTON_PRESS_MASK)
        self.treeview.connect('button_press_event', self.on_tree_click)

        scrolled = gtk.ScrolledWindow()
        scrolled.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
        scrolled.add(self.treeview)

        self.add(scrolled)
        self.show_all()

        # populate with root nodes
        for name, id in NODE_IDS[self.locale].items():
            self.treestore.append(None, [id, name])

    def fetch_nodes(self, node_id):
        """
        Fetches a BrowseNode from Amazon.
        """
        self.treestore.clear()

        root = self.api.browse_node_lookup(node_id)
        node = root.BrowseNodes.BrowseNode
        id = node.BrowseNodeId.pyval
        try:
            is_root = node.IsCategoryRoot.pyval == 1
        except AttributeError:
            is_root = False
        name = node.Name.pyval

        try:
            parents = dict((parent.BrowseNodeId.pyval, parent.Name.pyval)
                        for parent in node.Ancestors.BrowseNode)
        except AttributeError:
            parents = {}

        try:
            children = dict((child.BrowseNodeId.pyval, child.Name.pyval)
                        for child in node.Children.BrowseNode)
        except AttributeError:
            children = {}

        piter = None
        for parent_id, parent_name in parents.items():
            piter = self.treestore.append(None, [parent_id, parent_name])

        iter = self.treestore.append(piter, [id, name])

        for child_id, child_name in children.items():
            self.treestore.append(iter, [child_id, child_name])

        self.treeview.expand_all()

    def main(self):
        gtk.main()

if __name__ == "__main__":
    explorer = BrowseNodeExplorer()
    explorer.main()
<commit_msg>Fetched data is appended to nodes.<commit_after>#!/usr/bin/env python
import gtk

from config import AWS_KEY, SECRET_KEY
from amazonproduct import API

#: a list of root nodes retrieved from
#: http://docs.amazonwebservices.com/AWSECommerceService/2009-11-01/DG/index.html?BrowseNodeIDs.html
NODE_IDS = {
    'ca' : {
        'Books' : 927726,
        'Classical' : 962454,
        'DVD' : 14113311,
        'Electronics' : 677211011,
        'ForeignBooks' : 927726,
        'Kitchen' : 2206275011,
        'Music' : 962454,
        'Software' : 3234171,
        'SoftwareVideoGames' : 3323751,
        'VHS' : 962072,
        'Video' : 962454,
        'VideoGames' : 110218011,
    },
    'de' : {
        'Apparel' : 78689031,
        'Automotive' : 78194031,
        'Baby' : 357577011,
        'Beauty' : 64257031,
        'Books' : 541686,
        'Classical' : 542676,
        'DVD' : 547664,
        'Electronics' : 569604,
        'ForeignBooks' : 54071011,
        'HealthPersonalCare' : 64257031,
        'HomeGarden' : 10925241,
        'Jewelry' : 327473011,
        'Kitchen' : 3169011,
        'Lighting' : 213083031,
        'Magazines' : 1198526,
        'MP3Downloads' : 77256031,
        'Music' : 542676,
        'OfficeProducts' : 16291311,
        'OutdoorLiving' : 10925051,
        'PCHardware' : 569604,
        'Photo' : 569604,
        'Software' : 542064,
        'SoftwareVideoGames' : 541708,
        'SportingGoods' : 16435121,
        'Toys' : 12950661,
        'VHS' : 547082,
        'Video' : 547664,
        'VideoGames' : 541708,
        'Watches' : 193708031,
    },
    'fr' : {
        'Baby' : 206617031,
        'Beauty' : 197858031,
        'Books' : 468256,
        'Classical' : 537366,
        'DVD' : 578608,
        'Electronics' : 1058082,
        'ForeignBooks' : 69633011,
        'HealthPersonalCare' : 197861031,
        'Jewelry' : 193711031,
        'Kitchen' : 57686031,
        'Lighting' : 213080031,
        'MP3Downloads' : 206442031,
        'Music' : 537366,
        'OfficeProducts' : 192420031,
        'Shoes' : 215934031,
        'Software' : 548012,
        'SoftwareVideoGames' : 548014,
        'Toys' : 548014,
        'VHS' : 578610,
        'Video' : 578608,
        'VideoGames' : 548014,
        'Watches' : 60937031,
    },
    'jp' : {
        'Apparel' : 361299011,
        'Automotive' : 2017304051,
        'Baby' : 13331821,
        'Beauty' : 52391051,
        'Books' : 465610,
        'Classical' : 562032,
        'DVD' : 562002,
        'Electronics' : 3210991,
        'ForeignBooks' : 388316011,
        'Grocery' : 57240051,
        'HealthPersonalCare' : 161669011,
        'Hobbies' : 13331821,
        'HomeImprovement' : 2016929051,
        'Jewelry' : 85896051,
        'Kitchen' : 3839151,
        'Music' : 562032,
        'Shoes' : 2016926051,
        'Software' : 637630,
        'SportingGoods' : 14315361,
        'Toys' : 13331821,
        'VHS' : 561972,
        'Video' : 561972,
        'VideoGames' : 637872,
        'Watches' : 14315361,
    },
    'uk' : {
        'Apparel' : 83451031,
        'Automotive' : 248877031,
        'Baby' : 60032031,
        'Beauty' : 66280031,
        'Books' : 1025612,
        'Classical' : 505510,
        'DVD' : 283926,
        'Electronics' : 560800,
        'HealthPersonalCare' : 66280031,
        'HomeGarden' : 11052591,
        'Jewelry' : 193717031,
        'Kitchen' : 11052591,
        'Lighting' : 213077031,
        'MP3Downloads' : 77198031,
        'Music' : 505510,
        'OfficeProducts' : 560800,
        'OutdoorLiving' : 11052591,
        'Software' : 1025614,
        'SoftwareVideoGames' : 1025616,
        'SportingGoods' : 319530011,
        'Tools' : 11052591,
        'Toys' : 712832,
        'VHS' : 283926,
        'Video' : 283926,
        'VideoGames' : 1025616,
        'Watches' : 595312,
    },
    'us' : {
        'Apparel' : 1036682,
        'Automotive' : 15690151,
        'Baby' : 1036682,
        'Beauty' : 11055981,
        'Books' : 1000,
        'Classical' : 301668,
        'DigitalMusic' : 195208011,
        'DVD' : 130,
        'Electronics' : 493964,
        'GourmetFood' : 3580501,
        'Grocery' : 3760931,
        'HealthPersonalCare' : 3760931,
        'HomeGarden' : 285080,
        'Industrial' : 228239,
        'Jewelry' : 3880591,
        'Kitchen' : 1063498,
        'Magazines' : 599872,
        'Merchants' : 493964,
        'Miscellaneous' : 10304191,
        'MP3Downloads' : 195211011,
        'Music' : 301668,
        'MusicalInstruments' : 11965861,
        'OfficeProducts' : 1084128,
        'OutdoorLiving' : 1063498,
        'PCHardware' : 493964,
        'PetSupplies' : 1063498,
        'Photo' : 493964,
        'Software' : 409488,
        'SportingGoods' : 1079730,
        'Tools' : 468240,
        'Toys' : 493964,
        'VHS' : 404272,
        'Video' : 130,
        'VideoGames' : 493964,
        'Watches' : 1079730,
        'Wireless' : 508494,
        'WirelessAccessories' : 13900851,
    },
}

class BrowseNodeExplorer (gtk.Window):

    """
    Gtk explorer for Amazon BrowseNodes.
    """

    def on_delete(self, widget, event, data=None):
        # closes the window and quits.
        gtk.main_quit()
        return False

    def on_tree_click(self, widget, event, data=None):
        # if double click
        if event.type == gtk.gdk._2BUTTON_PRESS:

            # get data from highlighted selection
            treeselection = self.treeview.get_selection()
            model, iter = treeselection.get_selected()
            name_of_data = self.treestore.get_value(iter, 0)

            # and fetch selected node
            self.fetch_nodes(name_of_data)

    def __init__(self, locale='de'):

        gtk.Window.__init__(self, gtk.WINDOW_TOPLEVEL)

        self.set_title("BrowseNode Explorer")
        self.set_size_request(400, 200)
        self.connect("delete_event", self.on_delete)

        self.locale = locale
        self.api = API(AWS_KEY, SECRET_KEY, self.locale)

        # create a TreeStore with one string column to use as the model
        self.treestore = gtk.TreeStore(int, str)

        # create the TreeView using treestore
        self.treeview = gtk.TreeView(self.treestore)

        # add column id
        renderer = gtk.CellRendererText()
        column = gtk.TreeViewColumn('id', renderer, text=0)
        self.treeview.append_column(column)

        # add column name
        renderer = gtk.CellRendererText()
        column = gtk.TreeViewColumn('name', renderer, text=1)
        column.set_sort_column_id(1) # Allow sorting on the column
        self.treeview.append_column(column)

        # make it clickable
        self.treeview.add_events(gtk.gdk.BUTTON_PRESS_MASK)
        self.treeview.connect('button_press_event', self.on_tree_click)

        scrolled = gtk.ScrolledWindow()
        scrolled.set_policy(gtk.POLICY_AUTOMATIC, gtk.POLICY_AUTOMATIC)
        scrolled.add(self.treeview)

        self.add(scrolled)
        self.show_all()

        # populate with root nodes
        # but avoid duplicated node ids
        node_ids = set(NODE_IDS[self.locale].values())
        for name, id in NODE_IDS[self.locale].items():
            if id in node_ids:
                self.treestore.append(None, [id, name])
                node_ids.remove(id)

    def _find_row(self, node_id):
        def match_func(row, data):
            # data is a tuple containing column number, key
            column, key = data
            return row[column] == key
        def search(rows, func, data):
            if not rows: return None
            for row in rows:
                if func(row, data):
                    return row
                result = search(row.iterchildren(), func, data)
                if result: return result
            return None
        return search(self.treestore, match_func, (0, node_id))

    def fetch_nodes(self, node_id):
        """
        Fetches a BrowseNode from Amazon.
        """
        # fetch matching row from treestore
        row = self._find_row(node_id)

        # fetch Amazon data
        node = self.api.browse_node_lookup(node_id).BrowseNodes.BrowseNode
        id = node.BrowseNodeId.pyval
        name = node.Name.pyval
        is_root = hasattr(node, 'IsCategoryRoot') and node.IsCategoryRoot.pyval == 1

        #~ from lxml import etree
        #~ print etree.tostring(node, pretty_print=True)

        #try:
        #    parents = dict((parent.BrowseNodeId.pyval, parent.Name.pyval)
        #                for parent in node.Ancestors.BrowseNode)
        #except AttributeError:
        #    parents = {}
        #
        #piter = None
        #for parent_id, parent_name in parents.items():
        #    piter = self.treestore.append(None, [parent_id, parent_name])
        #
        #iter = self.treestore.append(piter, [id, name])

        # replace node name
        row[1] = name

        try:
            children = dict((child.BrowseNodeId.pyval, child.Name.pyval)
                        for child in node.Children.BrowseNode)
        except AttributeError:
            children = {}

        for child_id, child_name in children.items():
            self.treestore.append(row.iter, [child_id, child_name])

        # expand nodes of just added
        self.treeview.expand_row(tuple(row.path), True)

    def main(self):
        gtk.main()

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