<commit_before>###############################################################################
# plot_distanceintegral_smallpatch.py: make a useful plot of the distance
#                                      integration in the effective volume
#                                      for a small patch of the sky with better
#                                      distance resolution
###############################################################################
import sys
import os, os.path
import pickle
import numpy
import h5py
from scipy import signal, integrate, interpolate
import healpy
import matplotlib
matplotlib.use('Agg')
from galpy.util import save_pickles, bovy_plot, multi
import multiprocessing
import densprofiles
import dust
import gaia_rc
# magnitude limits for the survey
_GMIN= 3.
_GMAX= 20.
_NGSAMPLES= 100
_HIRESGREEN15DISTMODS= numpy.linspace(4.,19.,61)
_HIRESGREEN15DISTS= dust.distmod2dist(_HIRESGREEN15DISTMODS)
_NSIDE_HIRES= 1024
def read_patch():
    with h5py.File(os.path.join(os.getenv('DATADIR'),'bovy','apogee-maps',
                                'l60_b03.h5'),'r') as greendata:
            pix_info= greendata['/pixel_info'][:]
            best_fit= greendata['/best_fit'][:]
    # Upgrade the Nside=512 pixels to Nside=1024
    healpix_index_out= numpy.zeros(len(pix_info)+3*numpy.sum(pix_info['nside'] == 512),
                                   dtype='int64')
    best_fit_out= numpy.zeros((len(pix_info)+3*numpy.sum(pix_info['nside'] == 512),61))
    healpix_index_out[:len(pix_info)-numpy.sum(pix_info['nside'] == 512)]=\
        pix_info['healpix_index'][pix_info['nside'] == 1024]
    best_fit_out[:len(pix_info)-numpy.sum(pix_info['nside'] == 512)]=\
        best_fit[pix_info['nside'] == 1024]
    for ii in range(numpy.sum(pix_info['nside'] == 512)):
        healpix_index_out[len(pix_info)-numpy.sum(pix_info['nside'] == 512)+4*ii:len(pix_info)-numpy.sum(pix_info['nside'] == 512)+4*ii+4]=\
            pix_info['healpix_index'][pix_info['nside'] == 512][ii]*4+numpy.arange(4)
        best_fit_out[len(pix_info)-numpy.sum(pix_info['nside'] == 512)+4*ii:len(pix_info)-numpy.sum(pix_info['nside'] == 512)+4*ii+4]=\
            best_fit[pix_info['nside'] == 512][ii]
    return (healpix_index_out,best_fit_out*0.460)

def plot_distanceintegral_smallpatch(savename,plotname):
    if os.path.exists(savename):
        with open(savename,'rb') as savefile:
            area= pickle.load(savefile)
    else:
        # Load the patch
        hpIndx, dmap= read_patch()
        # For samping over the absolute magnitude distribution
        iso= gaia_rc.load_iso()
        Gsamples= gaia_rc.sample_Gdist(iso,n=_NGSAMPLES)
        # l and b of the pixels
        theta, phi= healpy.pixelfunc.pix2ang(_NSIDE_HIRES,hpIndx,nest=True)
        cosb= numpy.sin(theta)
        area= multi.parallel_map(lambda x: distanceIntegrandHires(\
                _HIRESGREEN15DISTS[x],theta,phi,cosb,Gsamples),
                                 range(len(_HIRESGREEN15DISTS)),
                                 numcores=numpy.amin([16,
                                                      len(_HIRESGREEN15DISTS),
                                                      multiprocessing.cpu_count()]))

        save_pickles(savename,area)
    # Plot the power spectrum
    if True:
        psdx, psd= signal.periodogram(area*_HIRESGREEN15DISTS**3./numpy.sum(area*_HIRESGREEN15DISTS**3.),
                                      fs=(_HIRESGREEN15DISTMODS[1]-_HIRESGREEN15DISTMODS[0]),
                                      detrend=lambda x: x,scaling='spectrum')
        bovy_plot.bovy_print(fig_height=3.)
        matplotlib.rcParams['text.latex.preamble']=[r"\usepackage{yfonts}"]
        bovy_plot.bovy_plot(psdx[1:]*2.*numpy.pi,psd[1:],
                            'k-',loglog=True,
                            xlabel=r'$2\pi\,k_\mu\,(\mathrm{mag}^{-1})$',
                            ylabel=r'$P_k$',
                            xrange=[0.04,4.])
        bovy_plot.bovy_text(r'$\mathrm{normalized}\ D^3\,\nu_*(\mu|\theta)\,\textswab{S}(\mu)$',
                            bottom_left=True,size=16.)
        bovy_plot.bovy_end_print(plotname)
    else:
        bovy_plot.bovy_print(fig_height=3.)
        matplotlib.rcParams['text.latex.preamble']=[r"\usepackage{yfonts}"]
        bovy_plot.bovy_plot(_HIRESGREEN15DISTMODS,
                            area*_HIRESGREEN15DISTS**3.,
                            'k-',
                            xlabel=r'$\mu\,(\mathrm{mag}^{-1})$',
                            ylabel=r'$D^3\,\nu_*(\mu|\theta)\,\textswab{S}(\mu)$')
        bovy_plot.bovy_end_print(plotname)
    spl= interpolate.InterpolatedUnivariateSpline(_HIRESGREEN15DISTMODS,
                                                  area*_HIRESGREEN15DISTS**3.,
                                                  k=5)
    fthder= [spl.derivatives(dm)[4] for dm in _HIRESGREEN15DISTMODS]
    print "Simpson error= ", 0.5**4./180.*numpy.mean(numpy.fabs(fthder))/integrate.simps(area*_HIRESGREEN15DISTS**3.,dx=0.5)
    return None

def distanceIntegrandHires(dist,theta,phi,cosb,Gsamples,hpIndx,dmap):
    # Calculate the density
    densmap= densprofiles.expdist(phi,numpy.pi/2.-theta,
                                  dist*numpy.ones(len(theta)),glon=True,
                                  params=[1./3.,1./0.3])
    # Sample over the distribution of MG
    combinedmask= numpy.zeros_like(dmap)
    G0= 0.68+dust.dist2distmod(dist)
    #if dust.dist2distmod(dist) == 9.5 or dust.dist2distmod(dist) == 10.5:
    #    print numpy.sum(numpy.isnan(combinedmap))
    for jj in range(_NGSAMPLES):
        combinedmask+= ((dmap > (_GMIN-G0-Gsamples[jj]+0.68))\
                            *(dmap < (_GMAX-G0-Gsamples[jj]+0.68))).astype('float')
    combinedmask/= _NGSAMPLES
    #if dust.dist2distmod(dist) == 9.5 or dust.dist2distmod(dist) == 10.5:
    #    print numpy.sum(combinedmask)
    #    print dust.dist2distmod(dist), numpy.sum(cosb*densmap*combinedmask)
    # Compute cross correlation
    return float(numpy.sum(cosb*densmap*combinedmask))

if __name__ == '__main__':
    plot_distanceintegral_smallpatch(sys.argv[1], # savefilename
                                     sys.argv[2]) # plotfilename

<commit_msg>apogee-maps: update small-patch code<commit_after>###############################################################################
# plot_distanceintegral_smallpatch.py: make a useful plot of the distance
#                                      integration in the effective volume
#                                      for a small patch of the sky with better
#                                      distance resolution
###############################################################################
import sys
import os, os.path
import pickle
import numpy
import h5py
from scipy import signal, integrate, interpolate
import healpy
import matplotlib
matplotlib.use('Agg')
from galpy.util import save_pickles, bovy_plot, multi
import multiprocessing
import densprofiles
import dust
import gaia_rc
# magnitude limits for the survey
_GMIN= 3.
_GMAX= 20.
_NGSAMPLES= 100
_HIRESGREEN15DISTMODS= numpy.linspace(4.,19.,61)
_HIRESGREEN15DISTS= dust.distmod2dist(_HIRESGREEN15DISTMODS)
_NSIDE_HIRES= 1024
def read_patch():
    with h5py.File(os.path.join(os.getenv('DATADIR'),'bovy','apogee-maps',
                                'l60_b03.h5'),'r') as greendata:
            pix_info= greendata['/pixel_info'][:]
            best_fit= greendata['/best_fit'][:]
    # Upgrade the Nside=512 pixels to Nside=1024
    healpix_index_out= numpy.zeros(len(pix_info)+3*numpy.sum(pix_info['nside'] == 512),
                                   dtype='int64')
    best_fit_out= numpy.zeros((len(pix_info)+3*numpy.sum(pix_info['nside'] == 512),61))
    healpix_index_out[:len(pix_info)-numpy.sum(pix_info['nside'] == 512)]=\
        pix_info['healpix_index'][pix_info['nside'] == 1024]
    best_fit_out[:len(pix_info)-numpy.sum(pix_info['nside'] == 512)]=\
        best_fit[pix_info['nside'] == 1024]
    for ii in range(numpy.sum(pix_info['nside'] == 512)):
        healpix_index_out[len(pix_info)-numpy.sum(pix_info['nside'] == 512)+4*ii:len(pix_info)-numpy.sum(pix_info['nside'] == 512)+4*ii+4]=\
            pix_info['healpix_index'][pix_info['nside'] == 512][ii]*4+numpy.arange(4)
        best_fit_out[len(pix_info)-numpy.sum(pix_info['nside'] == 512)+4*ii:len(pix_info)-numpy.sum(pix_info['nside'] == 512)+4*ii+4]=\
            best_fit[pix_info['nside'] == 512][ii]
    return (healpix_index_out,best_fit_out*0.460)

def plot_distanceintegral_smallpatch(savename,plotname):
    if os.path.exists(savename):
        with open(savename,'rb') as savefile:
            area= pickle.load(savefile)
    else:
        # Load the patch
        hpIndx, dmap= read_patch()
        # For samping over the absolute magnitude distribution
        iso= gaia_rc.load_iso()
        Gsamples= gaia_rc.sample_Gdist(iso,n=_NGSAMPLES)
        # l and b of the pixels
        theta, phi= healpy.pixelfunc.pix2ang(_NSIDE_HIRES,hpIndx,nest=True)
        cosb= numpy.sin(theta)
        area= multi.parallel_map(lambda x: distanceIntegrandHires(\
                _HIRESGREEN15DISTS[x],theta,phi,cosb,Gsamples,dmap),
                                 range(len(_HIRESGREEN15DISTS)),
                                 numcores=numpy.amin([16,
                                                      len(_HIRESGREEN15DISTS),
                                                      multiprocessing.cpu_count()]))

        save_pickles(savename,area)
    # Plot the power spectrum
    if True:
        psdx, psd= signal.periodogram(area*_HIRESGREEN15DISTS**3./numpy.sum(area*_HIRESGREEN15DISTS**3.),
                                      fs=(_HIRESGREEN15DISTMODS[1]-_HIRESGREEN15DISTMODS[0]),
                                      detrend=lambda x: x,scaling='spectrum')
        bovy_plot.bovy_print(fig_height=3.)
        matplotlib.rcParams['text.latex.preamble']=[r"\usepackage{yfonts}"]
        bovy_plot.bovy_plot(psdx[1:]*2.*numpy.pi,psd[1:],
                            'k-',loglog=True,
                            xlabel=r'$2\pi\,k_\mu\,(\mathrm{mag}^{-1})$',
                            ylabel=r'$P_k$',
                            xrange=[0.04,4.])
        bovy_plot.bovy_text(r'$\mathrm{normalized}\ D^3\,\nu_*(\mu|\theta)\,\textswab{S}(\mu)$',
                            bottom_left=True,size=16.)
        bovy_plot.bovy_end_print(plotname)
    else:
        bovy_plot.bovy_print(fig_height=3.)
        matplotlib.rcParams['text.latex.preamble']=[r"\usepackage{yfonts}"]
        bovy_plot.bovy_plot(_HIRESGREEN15DISTMODS,
                            area*_HIRESGREEN15DISTS**3.,
                            'k-',
                            xlabel=r'$\mu\,(\mathrm{mag}^{-1})$',
                            ylabel=r'$D^3\,\nu_*(\mu|\theta)\,\textswab{S}(\mu)$')
        bovy_plot.bovy_end_print(plotname)
    spl= interpolate.InterpolatedUnivariateSpline(_HIRESGREEN15DISTMODS,
                                                  area*_HIRESGREEN15DISTS**3.,
                                                  k=5)
    fthder= [spl.derivatives(dm)[4] for dm in _HIRESGREEN15DISTMODS]
    print "Simpson error= ", 0.5**4./180.*numpy.mean(numpy.fabs(fthder))/integrate.simps(area*_HIRESGREEN15DISTS**3.,dx=0.5)
    return None

def distanceIntegrandHires(dist,theta,phi,cosb,Gsamples,dmap):
    # Calculate the density
    densmap= densprofiles.expdisk(phi,numpy.pi/2.-theta,
                                  dist*numpy.ones(len(theta)),glon=True,
                                  params=[1./3.,1./0.3])
    # Distance pixel
    tpix= numpy.argmin(numpy.fabs(dist-_HIRESGREEN15DISTS))
    dmap= dmap[:,tpix]
    # Sample over the distribution of MG
    combinedmask= numpy.zeros_like(dmap)
    G0= 0.68+dust.dist2distmod(dist)
    #if dust.dist2distmod(dist) == 9.5 or dust.dist2distmod(dist) == 10.5:
    #    print numpy.sum(numpy.isnan(combinedmap))
    for jj in range(_NGSAMPLES):
        combinedmask+= ((dmap > (_GMIN-G0-Gsamples[jj]+0.68))\
                            *(dmap < (_GMAX-G0-Gsamples[jj]+0.68))).astype('float')
    combinedmask/= _NGSAMPLES
    #if dust.dist2distmod(dist) == 9.5 or dust.dist2distmod(dist) == 10.5:
    #    print numpy.sum(combinedmask)
    #    print dust.dist2distmod(dist), numpy.sum(cosb*densmap*combinedmask)
    # Compute cross correlation
    return float(numpy.sum(cosb*densmap*combinedmask))

if __name__ == '__main__':
    plot_distanceintegral_smallpatch(sys.argv[1], # savefilename
                                     sys.argv[2]) # plotfilename

<|endoftext|>


<commit_before>#!/usr/bin/env python3 -tt
# Copyright 2010 Google Inc.
# Licensed under the Apache License, Version 2.0
# http://www.apache.org/licenses/LICENSE-2.0

# Google's Python Class
# http://code.google.com/edu/languages/google-python-class/

import sys
import re
import baby_parser

"""
Baby Names exercise

Define the extract_names() function below and change main()
to call it.

Suggested milestones for incremental development
-Extract the year and print it
-Extract the names and rank numbers and just print them
-Get the names data into a dict and print it
-Build the [year, 'name rank', ... ] list and print it
-Fix main() to use the extract_names list
"""

def year_from_babyfile_string(babyfile_string):
    """
    return year as a string
    """
    print('year_from_babyfile_string')
    match = re.search(r'Popularity in \w\w\w\w', babyfile_string)
    if match:
        print('found {}', match.group())
        year_int = match.group()[-4:]
        year = str(year_int)
    else:
        year = None
    print(year)
    return year

def baby_names_table_rows_from_babyfile_string(babyfile_string):
    """
    babyfile_string sample excerpt with lines of html

1	Michael	Jessica

2	Christopher	Ashley

3	Matthew	Brittany

    return list of dictionaries
    row = {'rank' : rank, 'name_boy' : name_boy, 'name_girl' : name_girl}
    """
    print('baby_names_table_rows_from_babyfile_string')

    table_rows = []
    # findall with regular expression with () groups returns a list of tuples.
    baby_tuples = re.findall(r'
(\d+)	(\w+)	(\w+)	', babyfile_string)
    for baby_tuple in baby_tuples:
        rank = baby_tuple[0]
        name_boy = baby_tuple[1]
        name_girl = baby_tuple[2]
        row = {'rank' : rank, 'name_boy' : name_boy, 'name_girl' : name_girl}
        table_rows.append(row)

    #print(table_rows)
    return table_rows

def baby_names_list_from_table_rows(a_table_rows):
    """
    a_table_rows is a list of dictionaries
    return list of lists, discarding gender information contained in a_table_rows
    element [name, rank]
    list is sorted by name
    """
    print('baby_names_list_from_table_rows')

    baby_names_list = []
    for row in a_table_rows:
        baby_names_list.append([row['name_boy'], row['rank']])
        baby_names_list.append([row['name_girl'], row['rank']])

    def sort_key_from_array(an_array):
        """
        return key for sorted()
        """
        return an_array[0]

    baby_names_list = sorted(baby_names_list, key=sort_key_from_array)

    #print(baby_names_list)
    return baby_names_list

def baby_names_collapsed_from_list(a_baby_names_list):
    """
    a_baby_names_list is a list of lists, each element [name, rank]
    Collapse list element to a string
    """
    print('baby_names_collapsed_from_list')

    baby_names_collapsed = []
    for baby_element in a_baby_names_list:
        baby_names_collapsed.append('{} {}'.format(baby_element[0], baby_element[1]))

    #print(baby_names_collapsed)
    return baby_names_collapsed

def baby_names(babyfile_string):
    parser = baby_parser.BabyParser()
    parser.feed(babyfile_string)

def extract_names(filename):
    """
    Given a file name for baby.html, returns a list starting with the year string
    followed by the name-rank strings in alphabetical order.
    ['2006', 'Aaliyah 91', Aaron 57', 'Abagail 895', ' ...]
    """
    # +++your code here+++
    print('extract_names')
    print(filename)
    babyfile = open(filename, 'r')
    babyfile_string = babyfile.read()
    babyfile.close()
    year = year_from_babyfile_string(babyfile_string)
    #baby_names(babyfile_string)
    baby_names_table_rows = baby_names_table_rows_from_babyfile_string(babyfile_string)
    baby_names_list = baby_names_list_from_table_rows(baby_names_table_rows)
    baby_names_collapsed = baby_names_collapsed_from_list(baby_names_list)
    baby_names_collapsed.insert(0, year)
    #print(baby_names_collapsed)
    return baby_names_collapsed

def main():
    # This command-line parsing code is provided.
    # Make a list of command line arguments, omitting the [0] element
    # which is the script itself.
    args = sys.argv[1:]

    print(args)

    if not args:
        print('usage: [--summaryfile] file [file ...]')
        sys.exit(1)

    # Notice the summary flag and remove it from args if it is present.
    summary = False
    if args[0] == '--summaryfile':
        summary = True
        del args[0]

    # +++your code here+++
    # For each filename, get the names, then either print the text output
    # or write it to a summary file
    for filename in args:
        baby_names_extracted = extract_names(filename)
        print(baby_names_extracted)

if __name__ == '__main__':
    main()
<commit_msg>In babynames.py add function summary_filename()<commit_after>#!/usr/bin/env python3 -tt
# Copyright 2010 Google Inc.
# Licensed under the Apache License, Version 2.0
# http://www.apache.org/licenses/LICENSE-2.0

# Google's Python Class
# http://code.google.com/edu/languages/google-python-class/

import sys
import re
import baby_parser

"""
Baby Names exercise

Define the extract_names() function below and change main()
to call it.

Suggested milestones for incremental development
-Extract the year and print it
-Extract the names and rank numbers and just print them
-Get the names data into a dict and print it
-Build the [year, 'name rank', ... ] list and print it
-Fix main() to use the extract_names list
"""

def year_from_babyfile_string(babyfile_string):
    """
    return year as a string
    """
    print('year_from_babyfile_string')
    match = re.search(r'Popularity in \w\w\w\w', babyfile_string)
    if match:
        print('found {}', match.group())
        year_int = match.group()[-4:]
        year = str(year_int)
    else:
        year = None
    print(year)
    return year

def baby_names_table_rows_from_babyfile_string(babyfile_string):
    """
    babyfile_string sample excerpt with lines of html

1	Michael	Jessica

2	Christopher	Ashley

3	Matthew	Brittany

    return list of dictionaries
    row = {'rank' : rank, 'name_boy' : name_boy, 'name_girl' : name_girl}
    """
    print('baby_names_table_rows_from_babyfile_string')

    table_rows = []
    # findall with regular expression with () groups returns a list of tuples.
    baby_tuples = re.findall(r'
(\d+)	(\w+)	(\w+)	', babyfile_string)
    for baby_tuple in baby_tuples:
        rank = baby_tuple[0]
        name_boy = baby_tuple[1]
        name_girl = baby_tuple[2]
        row = {'rank' : rank, 'name_boy' : name_boy, 'name_girl' : name_girl}
        table_rows.append(row)

    #print(table_rows)
    return table_rows

def baby_names_list_from_table_rows(a_table_rows):
    """
    a_table_rows is a list of dictionaries
    return list of lists, discarding gender information contained in a_table_rows
    element [name, rank]
    list is sorted by name
    """
    print('baby_names_list_from_table_rows')

    baby_names_list = []
    for row in a_table_rows:
        baby_names_list.append([row['name_boy'], row['rank']])
        baby_names_list.append([row['name_girl'], row['rank']])

    def sort_key_from_array(an_array):
        """
        return key for sorted()
        """
        return an_array[0]

    baby_names_list = sorted(baby_names_list, key=sort_key_from_array)

    #print(baby_names_list)
    return baby_names_list

def baby_names_collapsed_from_list(a_baby_names_list):
    """
    a_baby_names_list is a list of lists, each element [name, rank]
    Collapse list element to a string
    """
    print('baby_names_collapsed_from_list')

    baby_names_collapsed = []
    for baby_element in a_baby_names_list:
        baby_names_collapsed.append('{} {}'.format(baby_element[0], baby_element[1]))

    #print(baby_names_collapsed)
    return baby_names_collapsed

def baby_names(babyfile_string):
    parser = baby_parser.BabyParser()
    parser.feed(babyfile_string)

def extract_names(filename):
    """
    Given a file name for baby.html, returns a list starting with the year string
    followed by the name-rank strings in alphabetical order.
    ['2006', 'Aaliyah 91', Aaron 57', 'Abagail 895', ' ...]
    """
    # +++your code here+++
    print('extract_names')
    print(filename)
    babyfile = open(filename, 'r')
    babyfile_string = babyfile.read()
    babyfile.close()
    year = year_from_babyfile_string(babyfile_string)
    #baby_names(babyfile_string)
    baby_names_table_rows = baby_names_table_rows_from_babyfile_string(babyfile_string)
    baby_names_list = baby_names_list_from_table_rows(baby_names_table_rows)
    baby_names_collapsed = baby_names_collapsed_from_list(baby_names_list)
    baby_names_collapsed.insert(0, year)
    #print(baby_names_collapsed)
    return baby_names_collapsed

def summary_filename(filename):
        """
        for filename baby1990.html return baby1990_summary.txt
        """
        filename_start = filename[:-5]
        summary_name = '{}{}'.format(filename_start, '_summary.txt')
        return summary_name

def main():
    # This command-line parsing code is provided.
    # Make a list of command line arguments, omitting the [0] element
    # which is the script itself.
    args = sys.argv[1:]

    print(args)

    if not args:
        print('usage: [--summaryfile] file [file ...]')
        sys.exit(1)

    # Notice the summary flag and remove it from args if it is present.
    summary = False
    if args[0] == '--summaryfile':
        summary = True
        del args[0]

    # +++your code here+++
    # For each filename, get the names, then either print the text output
    # or write it to a summary file
    for filename in args:
        baby_names_extracted = extract_names(filename)
        print(baby_names_extracted)

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


<commit_before># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import logging

from monty.json import MSONable

from pymatgen.core.structure import Molecule
from pymatgen.analysis.graphs import build_MoleculeGraph
from pymatgen.analysis.local_env import OpenBabelNN
import networkx as nx

__author__ = "Samuel Blau"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Samuel Blau"
__email__ = "samblau1@gmail.com"
__status__ = "Alpha"
__date__ = "7/26/18"

logger = logging.getLogger(__name__)

class BondDissociationEnergies(MSONable):

    def __init__(self, molecule_entry, fragment_entries):
        """
        Standard constructor for bond dissociation energies.

        Args:

        Get list of all bonds in target molecule
        Build molecule graph of target molecule
        Loop through bonds aka edges and attempt to split into disconnected subgraphs
            On success:
                Search through fragments: find entries with initial OR final molecules which are isomorphic
                Should be three entires per fragment: original charge, OC+1, OC-1
                If there are too few entries:
                    If one of the fragments is an H atom, you should only find two. H+ should be set (to what value???)
                    If fragment matches initial, but not final structure:
                        Just accept the rearrangement as part of the energetic cost of dissociation?
                        At the very least, flag this structure change in the eventual entry
                    If fragment matches final, but not initial structure:
                        Ignore unless no other fragment matches either?
                    If fragment is unstable:
                        Find its subfragments? Is that even sensical?
                        Just pass?
                    If fragment truly cannot be found, aka has not been calculated or calculate failed to complete:
                        Tell user to calculate it and pass
                If there are too many entries:
                    Get rid of duplicates
                        Going to be annoying given time stamps, rounding errors, etc
                    Remove those whose initial rems don't conform to current defaults - Happens before passing
                    If still more than three, for now, isolate differences, report to user, and grab most recent
                    Just take more negative?
                Grab energies of each fragment at all three charges
                If target molecule is neutral:
                    Calculate E_frag1(0)  + E_frag2(0)
                    Calculate E_frag1(+1) + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(+1)
                If target molecule is +1:
                    Calculate E_frag1(0)  + E_frag2(+1)
                    Calculate E_frag1(+1) + E_frag2(0)
                    Calculate E_frag1(+2) + E_frag2(-1)?
                    Calculate E_frag1(-1) + E_frag2(+2)?
                If target molecule is -1:
                    Calculate E_frag1(0)  + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(0)
                    Calculate E_frag1(-2) + E_frag2(+1)?
                    Calculate E_frag1(+1) + E_frag2(-2)?
                Don't limit - just make final charge add to original charge
                Save E_molecule - calculated values associated with bond indices & charges in some data structure
            On failure:
                Track bonds that do not break to yield subgraphs as ring bonds
        Make all possible pairs of ring bonds
        Loop through ring bond pairs and attempt to split into disconnected subgraphs
            Same success/failure procedure as above. Should be wrapped in a function.
        Resolve individual rings
        Attempt to solve for single ring bond dissociation energies, careful with charge

        """

        self.molecule_entry = molecule_entry
        self.filter_fragment_entries(fragment_entries)
        print(str(len(self.filtered_entries)) + " filtered entries")
        self.bond_dissociation_energies = []
        self.done_frag_pairs = []
        self.ring_bonds = []
        self.bad_pairs = []
        self.mol_graph = build_MoleculeGraph(Molecule.from_dict(molecule_entry["output"]["optimized_molecule"]),
                                             strategy=OpenBabelNN,
                                             reorder=False,
                                             extend_structure=False)
        for bond in self.mol_graph.graph.edges:
            bonds = [(bond[0],bond[1])]
            self.fragment_and_process(bonds)
        self.bond_pairs = []
        for ii,bond in enumerate(self.ring_bonds):
            for jj in range(ii+1,len(self.ring_bonds)):
                bond_pair = [bond, self.ring_bonds[jj]]
                self.bond_pairs += [bond_pair]
        if len(self.bond_pairs) > 0:
            print("Ring bonds detected!")
        for bond_pair in self.bond_pairs:
            self.fragment_and_process(bond_pair)
        self.solve_ring_bonds()

    def fragment_and_process(self, bonds):
        try:
            frags = self.mol_graph.split_molecule_subgraphs(bonds,allow_reverse=True)
            frag_success = True
        except RuntimeError:
            if len(bonds) == 1:
                self.ring_bonds += bonds
            elif len(bonds) == 2:
                self.bad_pairs += bonds
            else:
                print('No reason to try and break more than two bonds at once! Exiting...')
                raise ValueError
            frag_success = False
        if frag_success:
            frags_done = False
            for frag_pair in self.done_frag_pairs:
                if frag_pair[0].isomorphic_to(frags[0]):
                    if frag_pair[1].isomorphic_to(frags[1]):
                        frags_done = True
                        break
                elif frag_pair[1].isomorphic_to(frags[0]):
                    if frag_pair[0].isomorphic_to(frags[1]):
                        frags_done = True
                        break
            if not frags_done:
                self.done_frag_pairs += [frags]
                frag1_entries = self.search_fragment_entries(frags[0])
                frag2_entries = self.search_fragment_entries(frags[1])
                for frag1 in frag1_entries[0]:
                    for frag2 in frag2_entries[0]:
                        if frag1["calcs_reversed"][-1]["input"]["rem"]["method"] == frag2["calcs_reversed"][-1]["input"]["rem"]["method"] and frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag2["calcs_reversed"][-1]["input"]["rem"]["basis"]: # Add PCM check?
                            if frag1["output"]["initial_molecule"]["charge"] + frag2["output"]["initial_molecule"]["charge"] == self.molecule_entry["output"]["optimized_molecule"]["charge"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["method"] == frag1["calcs_reversed"][-1]["input"]["rem"]["method"]:
                                coords = nx.get_node_attributes(self.mol_graph.graph, "coords")
                                specie = nx.get_node_attributes(self.mol_graph.graph, "specie")
                                frag1_charge = frag1["output"]["initial_molecule"]["charge"]
                                frag2_charge = frag2["output"]["initial_molecule"]["charge"]
                                if frag1["output"]["final_energy"] == "unstable" or frag2["output"]["final_energy"] == "unstable":
                                    new_entry = ["unstable", bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], frag1["smiles"], frag1_charge, frag1["output"]["final_energy"], frag2["smiles"], frag2_charge, frag2["output"]["final_energy"]]
                                else:
                                    new_entry = [self.molecule_entry["output"]["final_energy"] - (frag1["output"]["final_energy"] + frag2["output"]["final_energy"]), bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], frag1["smiles"], frag1_charge, frag1["output"]["final_energy"], frag2["smiles"], frag2_charge, frag2["output"]["final_energy"]]
                                self.bond_dissociation_energies += [new_entry]

    def search_fragment_entries(self, frag):
        entries = []
        initial_entries = []
        final_entries = []
        for entry in self.filtered_entries:
            if frag.isomorphic_to(entry["initial"]) and frag.isomorphic_to(entry["final"]):
                entries += [entry["doc"]]
            elif frag.isomorphic_to(entry["initial"]):
                entries += [entry["doc"]]
                initial_entries += [entry["doc"]]
            elif frag.isomorphic_to(entry["final"]):
                final_entries += [entry["doc"]]
        return [entries, initial_entries, final_entries]

    def filter_fragment_entries(self,fragment_entries):
        self.filtered_entries = []
        for entry in fragment_entries:
            print(len(self.filtered_entries))
            this_dict = {}
            this_dict["doc"] = entry
            this_dict["initial"] = build_MoleculeGraph(Molecule.from_dict(entry["input"]["initial_molecule"]),
                                          strategy=OpenBabelNN,
                                          reorder=False,
                                          extend_structure=False)
            this_dict["final"] = build_MoleculeGraph(Molecule.from_dict(entry["output"]["initial_molecule"]),
                                        strategy=OpenBabelNN,
                                        reorder=False,
                                        extend_structure=False)
            found_similar_entry = False
            for ii,filtered_entry in enumerate(self.filtered_entries):
                # print(filtered_entry)
                # print(filtered_entry["initial"])
                # print(this_dict["initial"])
                # print(filtered_entry["initial"].isomorphic_to(this_dict["initial"]))
                # print(filtered_entry["final"])
                # print(this_dict["final"])
                # print(filtered_entry["final"].isomorphic_to(this_dict["final"]))
                # print(filtered_entry["doc"]["input"]["initial_molecule"]["charge"])
                # print(entry["input"]["initial_molecule"]["charge"])
                if filtered_entry["initial"].isomorphic_to(this_dict["initial"]) and filtered_entry["final"].isomorphic_to(this_dict["final"]) and filtered_entry["doc"]["input"]["initial_molecule"]["charge"] == entry["input"]["initial_molecule"]["charge"]:
                    found_similar_entry = True
                    if entry["output"]["final_energy"] < filtered_entry["doc"]["output"]["final_energy"]:
                        self.filtered_entries[ii] = this_dict
                    break
            if not found_similar_entry:
                self.filtered_entries += [this_dict]

    def solve_ring_bonds(self):
        pass
<commit_msg>Check smiles first<commit_after># coding: utf-8
# Copyright (c) Pymatgen Development Team.
# Distributed under the terms of the MIT License.

from __future__ import division, unicode_literals

import logging

from monty.json import MSONable

from pymatgen.core.structure import Molecule
from pymatgen.analysis.graphs import build_MoleculeGraph
from pymatgen.analysis.local_env import OpenBabelNN
import networkx as nx

__author__ = "Samuel Blau"
__copyright__ = "Copyright 2018, The Materials Project"
__version__ = "1.0"
__maintainer__ = "Samuel Blau"
__email__ = "samblau1@gmail.com"
__status__ = "Alpha"
__date__ = "7/26/18"

logger = logging.getLogger(__name__)

class BondDissociationEnergies(MSONable):

    def __init__(self, molecule_entry, fragment_entries):
        """
        Standard constructor for bond dissociation energies.

        Args:

        Get list of all bonds in target molecule
        Build molecule graph of target molecule
        Loop through bonds aka edges and attempt to split into disconnected subgraphs
            On success:
                Search through fragments: find entries with initial OR final molecules which are isomorphic
                Should be three entires per fragment: original charge, OC+1, OC-1
                If there are too few entries:
                    If one of the fragments is an H atom, you should only find two. H+ should be set (to what value???)
                    If fragment matches initial, but not final structure:
                        Just accept the rearrangement as part of the energetic cost of dissociation?
                        At the very least, flag this structure change in the eventual entry
                    If fragment matches final, but not initial structure:
                        Ignore unless no other fragment matches either?
                    If fragment is unstable:
                        Find its subfragments? Is that even sensical?
                        Just pass?
                    If fragment truly cannot be found, aka has not been calculated or calculate failed to complete:
                        Tell user to calculate it and pass
                If there are too many entries:
                    Get rid of duplicates
                        Going to be annoying given time stamps, rounding errors, etc
                    Remove those whose initial rems don't conform to current defaults - Happens before passing
                    If still more than three, for now, isolate differences, report to user, and grab most recent
                    Just take more negative?
                Grab energies of each fragment at all three charges
                If target molecule is neutral:
                    Calculate E_frag1(0)  + E_frag2(0)
                    Calculate E_frag1(+1) + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(+1)
                If target molecule is +1:
                    Calculate E_frag1(0)  + E_frag2(+1)
                    Calculate E_frag1(+1) + E_frag2(0)
                    Calculate E_frag1(+2) + E_frag2(-1)?
                    Calculate E_frag1(-1) + E_frag2(+2)?
                If target molecule is -1:
                    Calculate E_frag1(0)  + E_frag2(-1)
                    Calculate E_frag1(-1) + E_frag2(0)
                    Calculate E_frag1(-2) + E_frag2(+1)?
                    Calculate E_frag1(+1) + E_frag2(-2)?
                Don't limit - just make final charge add to original charge
                Save E_molecule - calculated values associated with bond indices & charges in some data structure
            On failure:
                Track bonds that do not break to yield subgraphs as ring bonds
        Make all possible pairs of ring bonds
        Loop through ring bond pairs and attempt to split into disconnected subgraphs
            Same success/failure procedure as above. Should be wrapped in a function.
        Resolve individual rings
        Attempt to solve for single ring bond dissociation energies, careful with charge

        """

        self.molecule_entry = molecule_entry
        self.filter_fragment_entries(fragment_entries)
        print(str(len(self.filtered_entries)) + " filtered entries")
        self.bond_dissociation_energies = []
        self.done_frag_pairs = []
        self.ring_bonds = []
        self.bad_pairs = []
        self.mol_graph = build_MoleculeGraph(Molecule.from_dict(molecule_entry["output"]["optimized_molecule"]),
                                             strategy=OpenBabelNN,
                                             reorder=False,
                                             extend_structure=False)
        for bond in self.mol_graph.graph.edges:
            bonds = [(bond[0],bond[1])]
            self.fragment_and_process(bonds)
        self.bond_pairs = []
        for ii,bond in enumerate(self.ring_bonds):
            for jj in range(ii+1,len(self.ring_bonds)):
                bond_pair = [bond, self.ring_bonds[jj]]
                self.bond_pairs += [bond_pair]
        if len(self.bond_pairs) > 0:
            print("Ring bonds detected!")
        for bond_pair in self.bond_pairs:
            self.fragment_and_process(bond_pair)
        self.solve_ring_bonds()

    def fragment_and_process(self, bonds):
        try:
            frags = self.mol_graph.split_molecule_subgraphs(bonds,allow_reverse=True)
            frag_success = True
        except RuntimeError:
            if len(bonds) == 1:
                self.ring_bonds += bonds
            elif len(bonds) == 2:
                self.bad_pairs += bonds
            else:
                print('No reason to try and break more than two bonds at once! Exiting...')
                raise ValueError
            frag_success = False
        if frag_success:
            frags_done = False
            for frag_pair in self.done_frag_pairs:
                if frag_pair[0].isomorphic_to(frags[0]):
                    if frag_pair[1].isomorphic_to(frags[1]):
                        frags_done = True
                        break
                elif frag_pair[1].isomorphic_to(frags[0]):
                    if frag_pair[0].isomorphic_to(frags[1]):
                        frags_done = True
                        break
            if not frags_done:
                self.done_frag_pairs += [frags]
                frag1_entries = self.search_fragment_entries(frags[0])
                frag2_entries = self.search_fragment_entries(frags[1])
                for frag1 in frag1_entries[0]:
                    for frag2 in frag2_entries[0]:
                        if frag1["calcs_reversed"][-1]["input"]["rem"]["method"] == frag2["calcs_reversed"][-1]["input"]["rem"]["method"] and frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag2["calcs_reversed"][-1]["input"]["rem"]["basis"]: # Add PCM check?
                            if frag1["output"]["initial_molecule"]["charge"] + frag2["output"]["initial_molecule"]["charge"] == self.molecule_entry["output"]["optimized_molecule"]["charge"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["basis"] == frag1["calcs_reversed"][-1]["input"]["rem"]["basis"] and self.molecule_entry["calcs_reversed"][-1]["input"]["rem"]["method"] == frag1["calcs_reversed"][-1]["input"]["rem"]["method"]:
                                coords = nx.get_node_attributes(self.mol_graph.graph, "coords")
                                specie = nx.get_node_attributes(self.mol_graph.graph, "specie")
                                frag1_charge = frag1["output"]["initial_molecule"]["charge"]
                                frag2_charge = frag2["output"]["initial_molecule"]["charge"]
                                if frag1["output"]["final_energy"] == "unstable" or frag2["output"]["final_energy"] == "unstable":
                                    new_entry = ["unstable", bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], frag1["smiles"], frag1_charge, frag1["output"]["final_energy"], frag2["smiles"], frag2_charge, frag2["output"]["final_energy"]]
                                else:
                                    new_entry = [self.molecule_entry["output"]["final_energy"] - (frag1["output"]["final_energy"] + frag2["output"]["final_energy"]), bonds, specie[bonds[0][0]], specie[bonds[0][1]], coords[bonds[0][0]], coords[bonds[0][1]], frag1["smiles"], frag1_charge, frag1["output"]["final_energy"], frag2["smiles"], frag2_charge, frag2["output"]["final_energy"]]
                                self.bond_dissociation_energies += [new_entry]

    def search_fragment_entries(self, frag):
        entries = []
        initial_entries = []
        final_entries = []
        for entry in self.filtered_entries:
            if frag.isomorphic_to(entry["initial"]) and frag.isomorphic_to(entry["final"]):
                entries += [entry["doc"]]
            elif frag.isomorphic_to(entry["initial"]):
                entries += [entry["doc"]]
                initial_entries += [entry["doc"]]
            elif frag.isomorphic_to(entry["final"]):
                final_entries += [entry["doc"]]
        return [entries, initial_entries, final_entries]

    def filter_fragment_entries(self,fragment_entries):
        self.filtered_entries = []
        for entry in fragment_entries:
            print(len(self.filtered_entries))
            this_dict = {}
            this_dict["doc"] = entry
            this_dict["initial"] = build_MoleculeGraph(Molecule.from_dict(entry["input"]["initial_molecule"]),
                                          strategy=OpenBabelNN,
                                          reorder=False,
                                          extend_structure=False)
            this_dict["final"] = build_MoleculeGraph(Molecule.from_dict(entry["output"]["initial_molecule"]),
                                        strategy=OpenBabelNN,
                                        reorder=False,
                                        extend_structure=False)
            found_similar_entry = False
            for ii,filtered_entry in enumerate(self.filtered_entries):
                # print(filtered_entry)
                # print(filtered_entry["initial"])
                # print(this_dict["initial"])
                # print(filtered_entry["initial"].isomorphic_to(this_dict["initial"]))
                # print(filtered_entry["final"])
                # print(this_dict["final"])
                # print(filtered_entry["final"].isomorphic_to(this_dict["final"]))
                # print(filtered_entry["doc"]["input"]["initial_molecule"]["charge"])
                # print(entry["input"]["initial_molecule"]["charge"])
                if filtered_entry["doc"]["smiles"] == entry["smiles"]:
                    if filtered_entry["initial"].isomorphic_to(this_dict["initial"]) and filtered_entry["final"].isomorphic_to(this_dict["final"]) and filtered_entry["doc"]["input"]["initial_molecule"]["charge"] == entry["input"]["initial_molecule"]["charge"]:
                        found_similar_entry = True
                        if entry["output"]["final_energy"] < filtered_entry["doc"]["output"]["final_energy"]:
                            self.filtered_entries[ii] = this_dict
                        break
            if not found_similar_entry:
                self.filtered_entries += [this_dict]

    def solve_ring_bonds(self):
        pass
<|endoftext|>


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

__version__ = '$Revision$'

import os
import re
import string
import sys

class Node:
    __rmjunk = re.compile("<#\d+#>")

    continuation = 0

    def __init__(self, link, str, seqno):
        self.links = [link]
        self.seqno = seqno
        # remove <#\d+#> left in by moving the data out of LaTeX2HTML
        str = self.__rmjunk.sub('', str)
        # build up the text
        self.text = split_entry_text(str)
        self.key = split_entry_key(str)

    def __cmp__(self, other):
        """Comparison operator includes sequence number, for use with
        list.sort()."""
        return self.cmp_entry(other) or cmp(self.seqno, other.seqno)

    def cmp_entry(self, other):
        """Comparison 'operator' that ignores sequence number."""
        c = 0
        for i in range(min(len(self.key), len(other.key))):
            c = (cmp_part(self.key[i], other.key[i])
                 or cmp_part(self.text[i], other.text[i]))
            if c:
                break
        return c or cmp(self.key, other.key) or cmp(self.text, other.text)

    def __repr__(self):
        return "<Node for %s (%s)>" % (string.join(self.text, '!'), self.seqno)

    def __str__(self):
        return string.join(self.key, '!')

    def dump(self):
        return "%s\1%s###%s\n" \
               % (string.join(self.links, "\1"),
                  string.join(self.text, '!'),
                  self.seqno)

def cmp_part(s1, s2):
    result = cmp(s1, s2)
    if result == 0:
        return 0
    l1 = string.lower(s1)
    l2 = string.lower(s2)
    minlen = min(len(s1), len(s2))
    if len(s1) < len(s2) and l1 == l2[:len(s1)]:
        result = -1
    elif len(s2) < len(s1) and l2 == l1[:len(s2)]:
        result = 1
    else:
        result = cmp(l1, l2) or cmp(s1, s2)
    return result

def split_entry(str, which):
    stuff = []
    parts = string.split(str, '!')
    parts = map(string.split, parts, ['@'] * len(parts))
    for entry in parts:
        if len(entry) != 1:
            key = entry[which]
        else:
            key = entry[0]
        stuff.append(key)
    return stuff

_rmtt = re.compile(r"""(.*)<tt(?: class=['"][a-z0-9]+["'])?>(.*)</tt>(.*)$""",
                   re.IGNORECASE)
_rmparens = re.compile(r"\(\)")

def split_entry_key(str):
    parts = split_entry(str, 1)
    for i in range(len(parts)):
        m = _rmtt.match(parts[i])
        if m:
            parts[i] = string.join(m.group(1, 2, 3), '')
        else:
            parts[i] = string.lower(parts[i])
        # remove '()' from the key:
        parts[i] = _rmparens.sub('', parts[i])
    return map(trim_ignored_letters, parts)

def split_entry_text(str):
    if '<' in str:
        m = _rmtt.match(str)
        if m:
            str = string.join(m.group(1, 2, 3), '')
    return split_entry(str, 1)

def load(fp):
    nodes = []
    rx = re.compile("(.*)\1(.*)###(.*)$")
    while 1:
        line = fp.readline()
        if not line:
            break
        m = rx.match(line)
        if m:
            link, str, seqno = m.group(1, 2, 3)
            nodes.append(Node(link, str, seqno))
    return nodes

def trim_ignored_letters(s):
    # ignore $ to keep environment variables with the
    # leading letter from the name
    s = string.lower(s)
    if s[0] == "$":
        return s[1:]
    else:
        return s

def get_first_letter(s):
    return string.lower(trim_ignored_letters(s)[0])

def split_letters(nodes):
    letter_groups = []
    if nodes:
        group = []
        append = group.append
        letter = get_first_letter(nodes[0].text[0])
        letter_groups.append((letter, group))
        for node in nodes:
            nletter = get_first_letter(node.text[0])
            if letter != nletter:
                letter = nletter
                group = []
                letter_groups.append((letter, group))
                append = group.append
            append(node)
    return letter_groups

# need a function to separate the nodes into columns...
def split_columns(nodes, columns=1):
    if columns <= 1:
        return [nodes]
    # This is a rough height; we may have to increase to avoid breaks before
    # a subitem.
    colheight = len(nodes) / columns
    numlong = len(nodes) % columns
    if numlong:
        colheight = colheight + 1
    else:
        numlong = columns
    cols = []
    for i in range(numlong):
        start = i * colheight
        end = start + colheight
        cols.append(nodes[start:end])
    del nodes[:end]
    colheight = colheight - 1
    try:
        numshort = len(nodes) / colheight
    except ZeroDivisionError:
        cols = cols + (columns - len(cols)) * [[]]
    else:
        for i in range(numshort):
            start = i * colheight
            end = start + colheight
            cols.append(nodes[start:end])
    #
    # If items continue across columns, make sure they are marked
    # as continuations so the user knows to look at the previous column.
    #
    for i in range(len(cols) - 1):
        try:
            prev = cols[i][-1]
            next = cols[i + 1][0]
        except IndexError:
            return cols
        else:
            n = min(len(prev.key), len(next.key))
            for j in range(n):
                if prev.key[j] != next.key[j]:
                    break
                next.continuation = j + 1
    return cols

DL_LEVEL_INDENT = "  "

def format_column(nodes):
    strings = ["
"]
    append = strings.append
    level = 0
    previous = []
    for node in nodes:
        current = node.text
        count = 0
        for i in range(min(len(current), len(previous))):
            if previous[i] != current[i]:
                break
            count = i + 1
        if count > level:
            append("
" * (count - level) + "\n")
            level = count
        elif level > count:
            append("\n")
            append(level * DL_LEVEL_INDENT)
            append("
" * (level - count))
            level = count
        # else: level == count
        for i in range(count, len(current) - 1):
            term = node.text[i]
            level = level + 1
            if node.continuation > i:
                extra = " (continued)"
            else:
                extra = ""
            append("\n
%s%s\n
\n%s
"
                   % (term, extra, level * DL_LEVEL_INDENT))
        append("\n%s
%s%s"
               % (level * DL_LEVEL_INDENT, node.links[0], node.text[-1]))
        for link in node.links[1:]:
            append(",\n%s    %s[Link]" % (level * DL_LEVEL_INDENT, link))
        previous = current
    append("\n")
    append("
" * (level + 1))
    return string.join(strings, '')

def format_nodes(nodes, columns=1):
    strings = []
    append = strings.append
    if columns > 1:
        colnos = range(columns)
        colheight = len(nodes) / columns
        if len(nodes) % columns:
            colheight = colheight + 1
        colwidth = 100 / columns
        append('

')
        for col in split_columns(nodes, columns):
            append('\n' % colwidth)
            append(format_column(col))
            append("\n	")
        append("\n

")
    else:
        append(format_column(nodes))
    append("\n
\n")
    return string.join(strings, '')

def format_letter(letter):
    if letter == '.':
        lettername = ". (dot)"
    elif letter == '_':
        lettername = "_ (underscore)"
    else:
        lettername = string.upper(letter)
    return "\n
\n
%s
\n\n" \
           % (letter, lettername)

def format_html_letters(nodes, columns=1):
    letter_groups = split_letters(nodes)
    items = []
    for letter, nodes in letter_groups:
        s = "%s" % (letter, letter)
        items.append(s)
    s = ["
\n%s\n" % string.join(items, " |\n")]
    for letter, nodes in letter_groups:
        s.append(format_letter(letter))
        s.append(format_nodes(nodes, columns))
    return string.join(s, '')

def format_html(nodes, columns):
    return format_nodes(nodes, columns)

def collapse(nodes):
    """Collapse sequences of nodes with matching keys into a single node.
    Destructive."""
    if len(nodes) < 2:
        return
    prev = nodes[0]
    i = 1
    while i < len(nodes):
        node = nodes[i]
        if not node.cmp_entry(prev):
            prev.links.append(node.links[0])
            del nodes[i]
        else:
            i = i + 1
            prev = node

def dump(nodes, fp):
    for node in nodes:
        fp.write(node.dump())

def process_nodes(nodes, columns, letters):
    nodes.sort()
    collapse(nodes)
    if letters:
        return format_html_letters(nodes, columns)
    else:
        return format_html(nodes, columns)

def main():
    import getopt
    ifn = "-"
    ofn = "-"
    columns = 1
    letters = 0
    opts, args = getopt.getopt(sys.argv[1:], "c:lo:",
                               ["columns=", "letters", "output="])
    for opt, val in opts:
        if opt in ("-o", "--output"):
            ofn = val
        elif opt in ("-c", "--columns"):
            columns = string.atoi(val)
        elif opt in ("-l", "--letters"):
            letters = 1
    if not args:
        args = [ifn]
    nodes = []
    for fn in args:
        nodes = nodes + load(open(fn))
    num_nodes = len(nodes)
    html = process_nodes(nodes, columns, letters)
    program = os.path.basename(sys.argv[0])
    if ofn == "-":
        sys.stdout.write(html)
        sys.stderr.write("\n%s: %d index nodes" % (program, num_nodes))
    else:
        open(ofn, "w").write(html)
        print
        print "%s: %d index nodes" % (program, num_nodes)

if __name__ == "__main__":
    main()
<commit_msg>Make this do the right thing with entries which start with the percent sign, in response to Skip's comments in SF bug #487165.<commit_after>#! /usr/bin/env python

__version__ = '$Revision$'

import os
import re
import string
import sys

bang_join = "!".join
null_join = "".join

class Node:
    __rmjunk = re.compile("<#\d+#>")

    continuation = 0

    def __init__(self, link, str, seqno):
        self.links = [link]
        self.seqno = seqno
        # remove <#\d+#> left in by moving the data out of LaTeX2HTML
        str = self.__rmjunk.sub('', str)
        # build up the text
        self.text = split_entry_text(str)
        self.key = split_entry_key(str)

    def __cmp__(self, other):
        """Comparison operator includes sequence number, for use with
        list.sort()."""
        return self.cmp_entry(other) or cmp(self.seqno, other.seqno)

    def cmp_entry(self, other):
        """Comparison 'operator' that ignores sequence number."""
        c = 0
        for i in range(min(len(self.key), len(other.key))):
            c = (cmp_part(self.key[i], other.key[i])
                 or cmp_part(self.text[i], other.text[i]))
            if c:
                break
        return c or cmp(self.key, other.key) or cmp(self.text, other.text)

    def __repr__(self):
        return "<Node for %s (%s)>" % (bang_join(self.text), self.seqno)

    def __str__(self):
        return bang_join(self.key)

    def dump(self):
        return "%s\1%s###%s\n" \
               % (string.join(self.links, "\1"),
                  bang_join(self.text),
                  self.seqno)

def cmp_part(s1, s2):
    result = cmp(s1, s2)
    if result == 0:
        return 0
    l1 = s1.lower()
    l2 = s2.lower()
    minlen = min(len(s1), len(s2))
    if len(s1) < len(s2) and l1 == l2[:len(s1)]:
        result = -1
    elif len(s2) < len(s1) and l2 == l1[:len(s2)]:
        result = 1
    else:
        result = cmp(l1, l2) or cmp(s1, s2)
    return result

def split_entry(str, which):
    stuff = []
    parts = str.split('!')
    parts = [part.split('@') for part in parts]
    for entry in parts:
        if len(entry) != 1:
            key = entry[which]
        else:
            key = entry[0]
        stuff.append(key)
    return stuff

_rmtt = re.compile(r"""(.*)<tt(?: class=['"][a-z0-9]+["'])?>(.*)</tt>(.*)$""",
                   re.IGNORECASE)
_rmparens = re.compile(r"\(\)")

def split_entry_key(str):
    parts = split_entry(str, 1)
    for i in range(len(parts)):
        m = _rmtt.match(parts[i])
        if m:
            parts[i] = null_join(m.group(1, 2, 3))
        else:
            parts[i] = parts[i].lower()
        # remove '()' from the key:
        parts[i] = _rmparens.sub('', parts[i])
    return map(trim_ignored_letters, parts)

def split_entry_text(str):
    if '<' in str:
        m = _rmtt.match(str)
        if m:
            str = null_join(m.group(1, 2, 3))
    return split_entry(str, 1)

def load(fp):
    nodes = []
    rx = re.compile("(.*)\1(.*)###(.*)$")
    while 1:
        line = fp.readline()
        if not line:
            break
        m = rx.match(line)
        if m:
            link, str, seqno = m.group(1, 2, 3)
            nodes.append(Node(link, str, seqno))
    return nodes

def trim_ignored_letters(s):
    # ignore $ to keep environment variables with the
    # leading letter from the name
    if s.startswith("$"):
        return s[1:].lower()
    else:
        return s.lower()

def get_first_letter(s):
    if s.startswith("<tex2html_percent_mark>"):
        return "%"
    else:
        return trim_ignored_letters(s)[0]

def split_letters(nodes):
    letter_groups = []
    if nodes:
        group = []
        append = group.append
        letter = get_first_letter(nodes[0].text[0])
        letter_groups.append((letter, group))
        for node in nodes:
            nletter = get_first_letter(node.text[0])
            if letter != nletter:
                letter = nletter
                group = []
                letter_groups.append((letter, group))
                append = group.append
            append(node)
    return letter_groups

def group_symbols(groups):
    entries = []
    ident_letters = string.ascii_letters + "_"
    while groups[0][0] not in ident_letters:
        entries += groups[0][1]
        del groups[0]
    if entries:
        groups.insert(0, ("Symbols", entries))

# need a function to separate the nodes into columns...
def split_columns(nodes, columns=1):
    if columns <= 1:
        return [nodes]
    # This is a rough height; we may have to increase to avoid breaks before
    # a subitem.
    colheight = int(len(nodes) / columns)
    numlong = int(len(nodes) % columns)
    if numlong:
        colheight = colheight + 1
    else:
        numlong = columns
    cols = []
    for i in range(numlong):
        start = i * colheight
        end = start + colheight
        cols.append(nodes[start:end])
    del nodes[:end]
    colheight = colheight - 1
    try:
        numshort = int(len(nodes) / colheight)
    except ZeroDivisionError:
        cols = cols + (columns - len(cols)) * [[]]
    else:
        for i in range(numshort):
            start = i * colheight
            end = start + colheight
            cols.append(nodes[start:end])
    #
    # If items continue across columns, make sure they are marked
    # as continuations so the user knows to look at the previous column.
    #
    for i in range(len(cols) - 1):
        try:
            prev = cols[i][-1]
            next = cols[i + 1][0]
        except IndexError:
            return cols
        else:
            n = min(len(prev.key), len(next.key))
            for j in range(n):
                if prev.key[j] != next.key[j]:
                    break
                next.continuation = j + 1
    return cols

DL_LEVEL_INDENT = "  "

def format_column(nodes):
    strings = ["
"]
    append = strings.append
    level = 0
    previous = []
    for node in nodes:
        current = node.text
        count = 0
        for i in range(min(len(current), len(previous))):
            if previous[i] != current[i]:
                break
            count = i + 1
        if count > level:
            append("
" * (count - level) + "\n")
            level = count
        elif level > count:
            append("\n")
            append(level * DL_LEVEL_INDENT)
            append("
" * (level - count))
            level = count
        # else: level == count
        for i in range(count, len(current) - 1):
            term = node.text[i]
            level = level + 1
            if node.continuation > i:
                extra = " (continued)"
            else:
                extra = ""
            append("\n
%s%s\n
\n%s
"
                   % (term, extra, level * DL_LEVEL_INDENT))
        append("\n%s
%s%s"
               % (level * DL_LEVEL_INDENT, node.links[0], node.text[-1]))
        for link in node.links[1:]:
            append(",\n%s    %s[Link]" % (level * DL_LEVEL_INDENT, link))
        previous = current
    append("\n")
    append("
" * (level + 1))
    return null_join(strings)

def format_nodes(nodes, columns=1):
    strings = []
    append = strings.append
    if columns > 1:
        colnos = range(columns)
        colheight = int(len(nodes) / columns)
        if len(nodes) % columns:
            colheight = colheight + 1
        colwidth = int(100 / columns)
        append('

')
        for col in split_columns(nodes, columns):
            append('\n' % colwidth)
            append(format_column(col))
            append("\n	")
        append("\n

")
    else:
        append(format_column(nodes))
    append("\n
\n")
    return null_join(strings)

def format_letter(letter):
    if letter == '.':
        lettername = ". (dot)"
    elif letter == '_':
        lettername = "_ (underscore)"
    else:
        lettername = letter.capitalize()
    return "\n
\n
%s
\n\n" \
           % (letter, lettername)

def format_html_letters(nodes, columns, group_symbol_nodes):
    letter_groups = split_letters(nodes)
    if group_symbol_nodes:
        group_symbols(letter_groups)
    items = []
    for letter, nodes in letter_groups:
        s = "%s" % (letter, letter)
        items.append(s)
    s = ["
\n%s\n" % string.join(items, " |\n")]
    for letter, nodes in letter_groups:
        s.append(format_letter(letter))
        s.append(format_nodes(nodes, columns))
    return null_join(s)

def format_html(nodes, columns):
    return format_nodes(nodes, columns)

def collapse(nodes):
    """Collapse sequences of nodes with matching keys into a single node.
    Destructive."""
    if len(nodes) < 2:
        return
    prev = nodes[0]
    i = 1
    while i < len(nodes):
        node = nodes[i]
        if not node.cmp_entry(prev):
            prev.links.append(node.links[0])
            del nodes[i]
        else:
            i = i + 1
            prev = node

def dump(nodes, fp):
    for node in nodes:
        fp.write(node.dump())

def process_nodes(nodes, columns, letters=0, group_symbol_nodes=0):
    nodes.sort()
    collapse(nodes)
    if letters:
        return format_html_letters(nodes, columns, group_symbol_nodes)
    else:
        return format_html(nodes, columns)

def main():
    import getopt
    ifn = "-"
    ofn = "-"
    columns = 1
    letters = 0
    group_symbol_nodes = 1
    opts, args = getopt.getopt(sys.argv[1:], "c:lo:",
                               ["columns=", "dont-group-symbols",
                                "group-symbols", "letters", "output="])
    for opt, val in opts:
        if opt in ("-o", "--output"):
            ofn = val
        elif opt in ("-c", "--columns"):
            columns = int(val, 10)
        elif opt in ("-l", "--letters"):
            letters = 1
        elif opt == "--group-symbols":
            group_symbol_nodes = 1
        elif opt == "--dont-group-symbols":
            group_symbol_nodes = 0
    if not args:
        args = [ifn]
    nodes = []
    for fn in args:
        nodes = nodes + load(open(fn))
    num_nodes = len(nodes)
    html = process_nodes(nodes, columns, letters, group_symbol_nodes)
    program = os.path.basename(sys.argv[0])
    if ofn == "-":
        sys.stdout.write(html)
        sys.stderr.write("\n%s: %d index nodes" % (program, num_nodes))
    else:
        open(ofn, "w").write(html)
        print
        print "%s: %d index nodes" % (program, num_nodes)

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


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

"""
 Counts words in UTF8 encoded, '\n' delimited text directly received from Kafka in every 2 seconds.
 Usage: direct_kafka_wordcount.py <broker_list> <topic>

 To run this on your local machine, you need to setup Kafka and create a producer first, see
 http://kafka.apache.org/documentation.html#quickstart

 and then run the example
    `$ bin/spark-submit --jars \
      external/kafka-assembly/target/scala-*/spark-streaming-kafka-assembly-*.jar \
      examples/src/main/python/streaming/distributed_route_streaming.py \
      localhost:9092 test`
"""

import sys

from pyspark import SparkContext
from pyspark.streaming import StreamingContext
from pyspark.streaming.kafka import KafkaUtils

if __name__ == "__main__":
    if len(sys.argv) != 3:
        #print("Usage: direct_kafka_wordcount.py <broker_list> <topic>", file=sys.stderr)
        exit(-1)

    sc = SparkContext(appName="PythonStreamingDistributedRoute")
    ssc = StreamingContext(sc, 2)

    brokers, topic = sys.argv[1:]
    kvs = KafkaUtils.createDirectStream(ssc, [topic], {"metadata.broker.list": brokers})

    # Splits up data based on the "geohash1_geohash2 incidentType~1234" outline given by team 1
    def handleRDD(rdd):
        dat = rdd.split(" ")
        cost = calcCost(dat[1])
	geo = dat[0].split('_')
	# Insert into database here

    # Function where our cost calculation formula will manipulate the values attached to the different incidents (cars per hour, et cetera)
    def calcCost(dat):
	datSplit = dat.split('~')
	return datSplit[1]

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

"""
 Counts words in UTF8 encoded, '\n' delimited text directly received from Kafka in every 2 seconds.
 Usage: direct_kafka_wordcount.py <broker_list> <topic>

 To run this on your local machine, you need to setup Kafka and create a producer first, see
 http://kafka.apache.org/documentation.html#quickstart

 and then run the example
    `$ bin/spark-submit --jars \
      external/kafka-assembly/target/scala-*/spark-streaming-kafka-assembly-*.jar \
      examples/src/main/python/streaming/distributed_route_streaming.py \
      localhost:9092 test`
"""

import sys

from pyspark import SparkContext
from pyspark.streaming import StreamingContext
from pyspark.streaming.kafka import KafkaUtils

if __name__ == "__main__":
    if len(sys.argv) != 3:
        #print("Usage: direct_kafka_wordcount.py <broker_list> <topic>", file=sys.stderr)
        exit(-1)

    sc = SparkContext(appName="PythonStreamingDistributedRoute")
    ssc = StreamingContext(sc, 2)
    brokers, topic = sys.argv[1:]
    kvs = KafkaUtils.createDirectStream(ssc, [topic], {"metadata.broker.list": brokers})
    lines = kvs.map(lambda x: x[1])
    lines.pprint()
    #data = lines.split(" ")
    #locationData = data[0].split('_')
    #incidentData = data[1].split('~')
    #cost = incidentData[1]
    #incidentType = incidentData[0]
    #print cost
    #print incidentType
    ssc.start()
    ssc.awaitTermination()
<|endoftext|>


<commit_before>##########################################################################
#
#  Copyright (c) 2011-2013, Image Engine Design Inc. All rights reserved.
#  Copyright (c) 2012-2013, John Haddon. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#      * Redistributions of source code must retain the above
#        copyright notice, this list of conditions and the following
#        disclaimer.
#
#      * Redistributions in binary form must reproduce the above
#        copyright notice, this list of conditions and the following
#        disclaimer in the documentation and/or other materials provided with
#        the distribution.
#
#      * Neither the name of John Haddon nor the names of
#        any other contributors to this software may be used to endorse or
#        promote products derived from this software without specific prior
#        written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

from __future__ import with_statement

import Gaffer
import GafferUI

## Supported plug metadata :
#
# - "boolPlugValueWidget:displayMode", with a value of "checkBox" or "switch"
class BoolPlugValueWidget( GafferUI.PlugValueWidget ) :

	def __init__( self, plug, displayMode=GafferUI.BoolWidget.DisplayMode.CheckBox, **kw ) :

		self.__boolWidget = GafferUI.BoolWidget( displayMode=displayMode )

		GafferUI.PlugValueWidget.__init__( self, self.__boolWidget, plug, **kw )

		self._addPopupMenu( self.__boolWidget )

		self.__stateChangedConnection = self.__boolWidget.stateChangedSignal().connect( Gaffer.WeakMethod( self.__stateChanged ) )

		self._updateFromPlug()

	def setHighlighted( self, highlighted ) :

		GafferUI.PlugValueWidget.setHighlighted( self, highlighted )
		self.__boolWidget.setHighlighted( highlighted )

	def _updateFromPlug( self ) :

		if self.getPlug() is not None :

			with self.getContext() :
				with Gaffer.BlockedConnection( self.__stateChangedConnection ) :
					self.__boolWidget.setState( self.getPlug().getValue() )

			displayMode = Gaffer.Metadata.plugValue( self.getPlug(), "boolPlugValueWidget:displayMode" )
			if displayMode is not None :
				self.__boolWidget.setDisplayMode( self.__boolWidget.DisplayMode.Switch if displayMode == "switch" else self.__boolWidget.DisplayMode.CheckBox )

		self.__boolWidget.setEnabled( self._editable( canEditAnimation = True ) )

	def __stateChanged( self, widget ) :

		self.__setPlugValue()

		return False

	def __setPlugValue( self ) :

		with Gaffer.UndoContext( self.getPlug().ancestor( Gaffer.ScriptNode ) ) :

			if Gaffer.Animation.isAnimated( self.getPlug() ) :
				curve = Gaffer.Animation.acquire( self.getPlug() )
				curve.addKey(
					Gaffer.Animation.Key(
						time = self.getContext().getTime(),
						value = self.__boolWidget.getState(),
						type = Gaffer.Animation.Type.Step
					)
				)
			else :
				self.getPlug().setValue( self.__boolWidget.getState() )

GafferUI.PlugValueWidget.registerType( Gaffer.BoolPlug, BoolPlugValueWidget )
<commit_msg>BoolPlugValueWidget : Harden against computation errors.<commit_after>##########################################################################
#
#  Copyright (c) 2011-2013, Image Engine Design Inc. All rights reserved.
#  Copyright (c) 2012-2013, John Haddon. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#      * Redistributions of source code must retain the above
#        copyright notice, this list of conditions and the following
#        disclaimer.
#
#      * Redistributions in binary form must reproduce the above
#        copyright notice, this list of conditions and the following
#        disclaimer in the documentation and/or other materials provided with
#        the distribution.
#
#      * Neither the name of John Haddon nor the names of
#        any other contributors to this software may be used to endorse or
#        promote products derived from this software without specific prior
#        written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

from __future__ import with_statement

import IECore

import Gaffer
import GafferUI

## Supported plug metadata :
#
# - "boolPlugValueWidget:displayMode", with a value of "checkBox" or "switch"
class BoolPlugValueWidget( GafferUI.PlugValueWidget ) :

	def __init__( self, plug, displayMode=GafferUI.BoolWidget.DisplayMode.CheckBox, **kw ) :

		self.__boolWidget = GafferUI.BoolWidget( displayMode=displayMode )

		GafferUI.PlugValueWidget.__init__( self, self.__boolWidget, plug, **kw )

		self._addPopupMenu( self.__boolWidget )

		self.__stateChangedConnection = self.__boolWidget.stateChangedSignal().connect( Gaffer.WeakMethod( self.__stateChanged ) )

		self._updateFromPlug()

	def setHighlighted( self, highlighted ) :

		GafferUI.PlugValueWidget.setHighlighted( self, highlighted )
		self.__boolWidget.setHighlighted( highlighted )

	def _updateFromPlug( self ) :

		if self.getPlug() is not None :

			value = None
			with self.getContext() :
				# Since BoolWidget doesn't yet have a way of
				# displaying errors, we just ignore exceptions
				# and leave UI components like GraphGadget to
				# display them via Node.errorSignal().
				with IECore.IgnoredExceptions( Exception ) :
					value = self.getPlug().getValue()

			if value is not None :
				with Gaffer.BlockedConnection( self.__stateChangedConnection ) :
					self.__boolWidget.setState( value )

			displayMode = Gaffer.Metadata.plugValue( self.getPlug(), "boolPlugValueWidget:displayMode" )
			if displayMode is not None :
				self.__boolWidget.setDisplayMode( self.__boolWidget.DisplayMode.Switch if displayMode == "switch" else self.__boolWidget.DisplayMode.CheckBox )

		self.__boolWidget.setEnabled( self._editable( canEditAnimation = True ) )

	def __stateChanged( self, widget ) :

		self.__setPlugValue()

		return False

	def __setPlugValue( self ) :

		with Gaffer.UndoContext( self.getPlug().ancestor( Gaffer.ScriptNode ) ) :

			if Gaffer.Animation.isAnimated( self.getPlug() ) :
				curve = Gaffer.Animation.acquire( self.getPlug() )
				curve.addKey(
					Gaffer.Animation.Key(
						time = self.getContext().getTime(),
						value = self.__boolWidget.getState(),
						type = Gaffer.Animation.Type.Step
					)
				)
			else :
				self.getPlug().setValue( self.__boolWidget.getState() )

GafferUI.PlugValueWidget.registerType( Gaffer.BoolPlug, BoolPlugValueWidget )
<|endoftext|>


<commit_before>##########################################################################
#
#  Copyright (c) 2013, Image Engine Design Inc. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#      * Redistributions of source code must retain the above
#        copyright notice, this list of conditions and the following
#        disclaimer.
#
#      * Redistributions in binary form must reproduce the above
#        copyright notice, this list of conditions and the following
#        disclaimer in the documentation and/or other materials provided with
#        the distribution.
#
#      * Neither the name of John Haddon nor the names of
#        any other contributors to this software may be used to endorse or
#        promote products derived from this software without specific prior
#        written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

import Gaffer
import GafferUI
import imath
import IECore

from Qt import QtCore
from Qt import QtGui

class RampPlugValueWidget( GafferUI.PlugValueWidget ) :

	def __init__( self, plug, **kw ) :

		column = GafferUI.ListContainer( GafferUI.ListContainer.Orientation.Vertical, spacing = 4 )

		GafferUI.PlugValueWidget.__init__( self, column, plug, **kw )

		with column :
			with GafferUI.ListContainer( GafferUI.ListContainer.Orientation.Horizontal, spacing = 4 ) :
				GafferUI.Label( "Display Mode" )
				drawModeWidget = GafferUI.MultiSelectionMenu( allowMultipleSelection = False, allowEmptySelection = False )
				drawModeWidget.append( "Ramp" )
				drawModeWidget.append( "Curves" )
				drawModeWidget.setSelection( "Ramp" )
				self.__drawModeChangedConnection = drawModeWidget.selectionChangedSignal().connect( Gaffer.WeakMethod( self.__drawModeChanged ) )

				GafferUI.Spacer( imath.V2i( 0 ), parenting = { "expand" : True } )

				# TODO: Since we don't have a good way to register metadata on child plugs, we just write the
				# metadata on this child plug right before constructing a widget for it.  There should probably
				# be some way to do this genericly during initialization
				Gaffer.Metadata.registerValue( plug['interpolation'],
					"plugValueWidget:type", "GafferUI.PresetsPlugValueWidget", persistent=False )
				for name, value in sorted( Gaffer.SplineDefinitionInterpolation.names.items() ):
					Gaffer.Metadata.registerValue( plug['interpolation'], "preset:" + name, value, persistent=False )

				GafferUI.PlugWidget( GafferUI.PlugValueWidget.create( plug["interpolation"] ) )

			self.__splineWidget = GafferUI.SplineWidget()
			self.__splineWidget.setDrawMode( self.__splineWidget.DrawMode.Ramp )
			self.__splineWidget._qtWidget().setMinimumHeight( 50 )

			self.__slider = GafferUI.Slider()
			self.__slider.setSizeEditable( True )
			self.__slider.setMinimumSize( 2 )
			self.__positionsChangedConnection = self.__slider.positionChangedSignal().connect( Gaffer.WeakMethod( self.__positionsChanged ) )
			self.__indexRemovedConnection = self.__slider.indexRemovedSignal().connect( Gaffer.WeakMethod( self.__indexRemoved ) )
			self.__selectedIndexChangedConnection = self.__slider.selectedIndexChangedSignal().connect( Gaffer.WeakMethod( self.__selectedIndexChanged ) )

			self.__lastPositionChangedReason = None
			self.__positionsMergeGroupId = 0

			with GafferUI.ListContainer( GafferUI.ListContainer.Orientation.Horizontal, spacing = 4 ) :

				self.__positionLabel = GafferUI.LabelPlugValueWidget( plug.pointXPlug( 0 ) )
				self.__positionField = GafferUI.NumericPlugValueWidget( plug.pointXPlug( 0 ) )

				self.__valueLabel = GafferUI.LabelPlugValueWidget( plug.pointYPlug( 0 ) )
				if isinstance( plug.pointYPlug( 0 ), Gaffer.FloatPlug ):
					self.__valueField = GafferUI.NumericPlugValueWidget( plug.pointYPlug( 0 ) )
				else:
					self.__valueField = GafferUI.ColorPlugValueWidget( plug.pointYPlug( 0 ) )

		self.setPlug( plug )

	def __drawModeChanged( self, drawModeWidget ) :
		name = drawModeWidget.getSelection()[0]
		self.__splineWidget.setDrawMode( self.__splineWidget.DrawMode.Ramp if name == "Ramp" else self.__splineWidget.DrawMode.Splines )

	def setPlug( self, plug ) :

		GafferUI.PlugValueWidget.setPlug( self, plug )

		self.__slider.setSelectedIndex( 0 )

	def getToolTip( self ) :

		result = GafferUI.PlugValueWidget.getToolTip( self )

		if self.getPlug() is not None :
			result += "
"
			result += "
Click empty space in slider to add handle"
			result += "
Click handle to select"
			result += "
Delete to remove selected handle"
			result += "
Cursor left/right to nudge selected handle"
			result += "
"

		return result

	def _updateFromPlug( self ) :

		plug = self.getPlug()
		with self.getContext() :

			self.__splineWidget.setSpline( plug.getValue().spline() )

			positions = []
			for i in range( 0, plug.numPoints() ) :
				positions.append( plug.pointXPlug( i ).getValue() )

			with Gaffer.BlockedConnection( self.__positionsChangedConnection ) :
				self.__slider.setPositions( positions )

	def __positionsChanged( self, slider, reason ) :

		if not slider.changesShouldBeMerged( self.__lastPositionChangedReason, reason ) :
			self.__positionsMergeGroupId += 1
		self.__lastPositionChangedReason = reason

		plug = self.getPlug()
		with Gaffer.UndoScope(
			plug.ancestor( Gaffer.ScriptNode ),
			mergeGroup = "RampPlugValudWidget%d%d" % ( id( self, ), self.__positionsMergeGroupId )
		) :

			if len( slider.getPositions() ) == plug.numPoints() :
				# the user has moved an existing point on the slider
				for index, position in enumerate( slider.getPositions() ) :
					plug.pointXPlug( index ).setValue( position )
			else :
				# a new position was added on the end by the user clicking
				# on an empty area of the slider.
				numPoints = plug.numPoints()
				assert( len( slider.getPositions() ) == numPoints + 1 )
				spline = plug.getValue().spline()
				position = slider.getPositions()[numPoints]
				plug.addPoint()
				plug.pointXPlug( numPoints ).setValue( position )
				plug.pointYPlug( numPoints ).setValue( spline( position ) )

	def __indexRemoved( self, slider, index ) :

		with Gaffer.UndoScope( self.getPlug().ancestor( Gaffer.ScriptNode ) ) :
			self.getPlug().removePoint( index )

	def __selectedIndexChanged( self, slider ) :

		index = slider.getSelectedIndex()
		if self.getPlug() is not None and index is not None :
			pointPlug = self.getPlug().pointPlug( index )
			self.__positionLabel.setPlug( pointPlug["x"] )
			self.__positionField.setPlug( pointPlug["x"] )
			self.__valueLabel.setPlug( pointPlug["y"] )
			self.__valueField.setPlug( pointPlug["y"] )

		else :
			self.__positionLabel.setPlug( None )
			self.__positionField.setPlug( None )
			self.__valueLabel.setPlug( None )
			self.__valueField.setPlug( None )

		self.__positionLabel.label().setText( "Position" )
		self.__valueLabel.label().setText( "Value" )

# we don't register this automatically for any plugs, as it takes up a lot of room
# in the node editor. this means the SplinePlugValueWidget will be used instead, and
# that will launch a dialogue containing a RampPlugValueWidget when appropriate. for
# nodes which want a large editor directly in the node editor, the RampPlugValueWidget
# can be registered directly for specific plugs.
<commit_msg>RampPlugValueWidget : Default to Spline draw mode for float splines<commit_after>##########################################################################
#
#  Copyright (c) 2013, Image Engine Design Inc. All rights reserved.
#
#  Redistribution and use in source and binary forms, with or without
#  modification, are permitted provided that the following conditions are
#  met:
#
#      * Redistributions of source code must retain the above
#        copyright notice, this list of conditions and the following
#        disclaimer.
#
#      * Redistributions in binary form must reproduce the above
#        copyright notice, this list of conditions and the following
#        disclaimer in the documentation and/or other materials provided with
#        the distribution.
#
#      * Neither the name of John Haddon nor the names of
#        any other contributors to this software may be used to endorse or
#        promote products derived from this software without specific prior
#        written permission.
#
#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS
#  IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO,
#  THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
#  PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR
#  CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
#  EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
#  PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
#  PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF
#  LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING
#  NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
#
##########################################################################

import Gaffer
import GafferUI
import imath
import IECore

from Qt import QtCore
from Qt import QtGui

class RampPlugValueWidget( GafferUI.PlugValueWidget ) :

	def __init__( self, plug, **kw ) :

		column = GafferUI.ListContainer( GafferUI.ListContainer.Orientation.Vertical, spacing = 4 )

		GafferUI.PlugValueWidget.__init__( self, column, plug, **kw )

		with column :
			with GafferUI.ListContainer( GafferUI.ListContainer.Orientation.Horizontal, spacing = 4 ) :
				GafferUI.Label( "Display Mode" )
				drawModeWidget = GafferUI.MultiSelectionMenu( allowMultipleSelection = False, allowEmptySelection = False )
				drawModeWidget.append( "Ramp" )
				drawModeWidget.append( "Curves" )
				drawModeWidget.setSelection( "Ramp" )
				self.__drawModeChangedConnection = drawModeWidget.selectionChangedSignal().connect( Gaffer.WeakMethod( self.__drawModeChanged ) )

				GafferUI.Spacer( imath.V2i( 0 ), parenting = { "expand" : True } )

				# TODO: Since we don't have a good way to register metadata on child plugs, we just write the
				# metadata on this child plug right before constructing a widget for it.  There should probably
				# be some way to do this genericly during initialization
				Gaffer.Metadata.registerValue( plug['interpolation'],
					"plugValueWidget:type", "GafferUI.PresetsPlugValueWidget", persistent=False )
				for name, value in sorted( Gaffer.SplineDefinitionInterpolation.names.items() ):
					Gaffer.Metadata.registerValue( plug['interpolation'], "preset:" + name, value, persistent=False )

				GafferUI.PlugWidget( GafferUI.PlugValueWidget.create( plug["interpolation"] ) )

			self.__splineWidget = GafferUI.SplineWidget()
			if isinstance( plug, ( Gaffer.SplinefColor3fPlug, Gaffer.SplinefColor4fPlug ) ) :
				self.__splineWidget.setDrawMode( self.__splineWidget.DrawMode.Ramp )
			else:
				self.__splineWidget.setDrawMode( self.__splineWidget.DrawMode.Splines )
			self.__splineWidget._qtWidget().setMinimumHeight( 50 )

			self.__slider = GafferUI.Slider()
			self.__slider.setSizeEditable( True )
			self.__slider.setMinimumSize( 2 )
			self.__positionsChangedConnection = self.__slider.positionChangedSignal().connect( Gaffer.WeakMethod( self.__positionsChanged ) )
			self.__indexRemovedConnection = self.__slider.indexRemovedSignal().connect( Gaffer.WeakMethod( self.__indexRemoved ) )
			self.__selectedIndexChangedConnection = self.__slider.selectedIndexChangedSignal().connect( Gaffer.WeakMethod( self.__selectedIndexChanged ) )

			self.__lastPositionChangedReason = None
			self.__positionsMergeGroupId = 0

			with GafferUI.ListContainer( GafferUI.ListContainer.Orientation.Horizontal, spacing = 4 ) :

				self.__positionLabel = GafferUI.LabelPlugValueWidget( plug.pointXPlug( 0 ) )
				self.__positionField = GafferUI.NumericPlugValueWidget( plug.pointXPlug( 0 ) )

				self.__valueLabel = GafferUI.LabelPlugValueWidget( plug.pointYPlug( 0 ) )
				if isinstance( plug.pointYPlug( 0 ), Gaffer.FloatPlug ):
					self.__valueField = GafferUI.NumericPlugValueWidget( plug.pointYPlug( 0 ) )
				else:
					self.__valueField = GafferUI.ColorPlugValueWidget( plug.pointYPlug( 0 ) )

		self.setPlug( plug )

	def __drawModeChanged( self, drawModeWidget ) :
		name = drawModeWidget.getSelection()[0]
		self.__splineWidget.setDrawMode( self.__splineWidget.DrawMode.Ramp if name == "Ramp" else self.__splineWidget.DrawMode.Splines )

	def setPlug( self, plug ) :

		GafferUI.PlugValueWidget.setPlug( self, plug )

		self.__slider.setSelectedIndex( 0 )

	def getToolTip( self ) :

		result = GafferUI.PlugValueWidget.getToolTip( self )

		if self.getPlug() is not None :
			result += "
"
			result += "
Click empty space in slider to add handle"
			result += "
Click handle to select"
			result += "
Delete to remove selected handle"
			result += "
Cursor left/right to nudge selected handle"
			result += "
"

		return result

	def _updateFromPlug( self ) :

		plug = self.getPlug()
		with self.getContext() :

			self.__splineWidget.setSpline( plug.getValue().spline() )

			positions = []
			for i in range( 0, plug.numPoints() ) :
				positions.append( plug.pointXPlug( i ).getValue() )

			with Gaffer.BlockedConnection( self.__positionsChangedConnection ) :
				self.__slider.setPositions( positions )

	def __positionsChanged( self, slider, reason ) :

		if not slider.changesShouldBeMerged( self.__lastPositionChangedReason, reason ) :
			self.__positionsMergeGroupId += 1
		self.__lastPositionChangedReason = reason

		plug = self.getPlug()
		with Gaffer.UndoScope(
			plug.ancestor( Gaffer.ScriptNode ),
			mergeGroup = "RampPlugValudWidget%d%d" % ( id( self, ), self.__positionsMergeGroupId )
		) :

			if len( slider.getPositions() ) == plug.numPoints() :
				# the user has moved an existing point on the slider
				for index, position in enumerate( slider.getPositions() ) :
					plug.pointXPlug( index ).setValue( position )
			else :
				# a new position was added on the end by the user clicking
				# on an empty area of the slider.
				numPoints = plug.numPoints()
				assert( len( slider.getPositions() ) == numPoints + 1 )
				spline = plug.getValue().spline()
				position = slider.getPositions()[numPoints]
				plug.addPoint()
				plug.pointXPlug( numPoints ).setValue( position )
				plug.pointYPlug( numPoints ).setValue( spline( position ) )

	def __indexRemoved( self, slider, index ) :

		with Gaffer.UndoScope( self.getPlug().ancestor( Gaffer.ScriptNode ) ) :
			self.getPlug().removePoint( index )

	def __selectedIndexChanged( self, slider ) :

		index = slider.getSelectedIndex()
		if self.getPlug() is not None and index is not None :
			pointPlug = self.getPlug().pointPlug( index )
			self.__positionLabel.setPlug( pointPlug["x"] )
			self.__positionField.setPlug( pointPlug["x"] )
			self.__valueLabel.setPlug( pointPlug["y"] )
			self.__valueField.setPlug( pointPlug["y"] )

		else :
			self.__positionLabel.setPlug( None )
			self.__positionField.setPlug( None )
			self.__valueLabel.setPlug( None )
			self.__valueField.setPlug( None )

		self.__positionLabel.label().setText( "Position" )
		self.__valueLabel.label().setText( "Value" )

# we don't register this automatically for any plugs, as it takes up a lot of room
# in the node editor. this means the SplinePlugValueWidget will be used instead, and
# that will launch a dialogue containing a RampPlugValueWidget when appropriate. for
# nodes which want a large editor directly in the node editor, the RampPlugValueWidget
# can be registered directly for specific plugs.
<|endoftext|>