<commit_before>from pylons import c

import ew

from pyforge.lib import validators as V
from pyforge.lib import helpers as h
from pyforge.lib.widgets import form_fields as ffw
from pyforge import model as M

# Discussion forms
class ModerateThread(ew.SimpleForm):
    class buttons(ew.WidgetsList):
        delete=ew.SubmitButton(label='Delete Thread')
    submit_text=None

class ModeratePost(ew.SimpleForm):
    template='genshi:pyforge.lib.widgets.templates.moderate_post'
    submit_text=None

class FlagPost(ew.SimpleForm):
    template='genshi:pyforge.lib.widgets.templates.flag_post'
    submit_text=None

class AttachPost(ew.SimpleForm):
    submit_text='Attach File'
    enctype='multipart/form-data'
    class fields(ew.WidgetsList):
        file_info=ew.InputField(field_type='file')

class ModeratePosts(ew.SimpleForm):
    template='genshi:pyforge.lib.widgets.templates.moderate_posts'
    submit_text=None
    def resources(self):
        for r in super(ModeratePosts, self).resources(): yield r
        yield ew.JSScript('''
      (function($){
          var tbl = $('form table');
          var checkboxes = $('input[type=checkbox]', tbl);
          $('a[href=#]', tbl).click(function() {
              checkboxes.each(function() {
                  if(this.checked) { this.checked = false; }
                  else { this.checked = true; }
              });
              return false;
          });
      }(jQuery));''')

class PostFilter(ew.SimpleForm):
    submit_text=None
    method='GET'
    fields = [
        ew.FieldSet(label='Post Filter', fields=[
                ew.SingleSelectField(
                    name='status',
                    label='Show posts with status',
                    options=[
                        ew.Option(py_value='-', label='Any'),
                        ew.Option(py_value='spam', label='Spam'),
                        ew.Option(py_value='pending', label='Pending moderation'),
                        ew.Option(py_value='ok', label='Ok')],
                    if_missing='-'),
                ew.IntField(name='flag',
                            label='Show posts with at least "n" flags',
                            css_class='text',
                            if_missing=0),
                ew.SubmitButton(label='Filter Posts')
                ])
        ]

class EditPost(ew.SimpleForm):
    show_subject=False

    # this ickiness is to override the default submit button
    def __call__(self, **kw):
        result = super(EditPost, self).__call__(**kw)
        submit_button = ffw.SubmitButton(label=result['submit_text'])
        result['extra_fields'] = [submit_button]
        result['buttons'] = [submit_button]
        return result

    @property
    def fields(self):
        def _():
            if hasattr(c, 'widget'):
                if c.widget.response.get('show_subject', self.show_subject):
                    yield ew.TextField(name='subject')
            else:
                yield ew.TextField(name='subject', if_missing=None)
            yield ffw.MarkdownEdit(name='text')
        return _()

    def resources(self):
        for r in ew.TextField(name='subject').resources(): yield r
        for r in ffw.MarkdownEdit(name='text').resources(): yield r

class _ThreadsTable(ew.SimpleForm):
    template='pyforge.lib.widgets.templates.threads_table'
    class hidden_fields(ew.WidgetsList):
        _id=ew.HiddenField(validator=V.Ming(M.Thread))
    class fields(ew.WidgetsList):
        num_replies=ew.HTMLField(show_label=True, label='Num Posts')
        num_views=ew.HTMLField(show_label=True)
        last_post=ew.HTMLField(text="${value and value.summary()}", show_label=True)
        subscription=ew.Checkbox(suppress_label=True, show_label=True)
    fields.insert(0, ew.LinkField(
            label='Subject', text="${value['subject']}",
            href="${value['url']()}", show_label=True))

class SubscriptionForm(ew.SimpleForm):
    template='pyforge.lib.widgets.templates.subscription_form'
    value=None
    show_discussion_email=False
    show_subject=False
    allow_create_thread=False
    params=['value',
            'show_discussion_email', 'show_subject', 'allow_create_thread']
    class fields(ew.WidgetsList):
        threads=_ThreadsTable()
        new_topic = EditPost(submit_text='New Topic', if_missing=None)
    submit_text='Update Subscriptions'
    def resources(self):
        for r in super(SubscriptionForm, self).resources(): yield r
        yield ew.JSScript('''
        $(window).load(function() {
            $('tbody').children(':even').addClass('even');
            $('.discussion_subscription_form').each(function(){
                var discussion = this;
                $('.submit', discussion).button();
                if($('.new_topic', discussion)){
                    $('.new_topic', discussion).click(function(ele){
                        $('.new_topic_form', discussion).show();
                        return false;
                    });
                }
                $('.follow', discussion).click(function(ele){
                    $('.follow_form', discussion).submit();
                    return false;
                });
            });
        });''')

# Widgets
class HierWidget(ew.Widget):
    widgets = {}

    def __call__(self, **kw):
        response = super(HierWidget, self).__call__(**kw)
        response['widgets'] = self.widgets
        return response

    def resources(self):
        for w in self.widgets.itervalues():
            for r in w.resources():
                yield r

class Attachment(ew.Widget):
    template='genshi:pyforge.lib.widgets.templates.attachment'
    params=['value', 'post']
    value=None
    post=None

class DiscussionHeader(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.discussion_header'
    params=['value']
    value=None
    widgets=dict(
        edit_post=EditPost(submit_text='New Thread'))

class ThreadHeader(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.thread_header'
    params=['value', 'offset', 'pagesize', 'total', 'show_moderate']
    value=None
    offset=None
    pagesize=None
    total=None
    show_moderate=False
    widgets=dict(
        moderate_thread=ModerateThread())

class PostHeader(ew.Widget):
    template='genshi:pyforge.lib.widgets.templates.post_header'
    params=['value']
    value=None

class PostThread(ew.Widget):
    template='genshi:pyforge.lib.widgets.templates.post_thread'
    params=['value']
    value=None

class Post(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.post'
    params=['value', 'show_subject', 'indent', 'supress_promote']
    value=None
    indent=0
    show_subject=False
    supress_promote=False
    widgets=dict(
        flag_post=FlagPost(),
        moderate_post=ModeratePost(),
        edit_post=EditPost(submit_text='Edit Post'),
        attach_post=AttachPost(submit_text='Attach'),
        attachment=Attachment())
    def resources(self):
        for r in super(Post, self).resources(): yield r
        for w in self.widgets.itervalues():
            for r in w.resources():
                yield r
        yield ew.JSScript('''
        (function(){
            $('.discussion-post').each(function(){
                var post = this;
                $('.submit', post).button();
                $('.flag_post, .delete_post', post).click(function(ele){
                    this.parentNode.submit();
                    return false;
                });
                if($('.edit_post', post)){
                    $('.edit_post', post).click(function(ele){
                        $('.edit_post_form', post).show();
                        return false;
                    });
                }
                if($('.reply_post', post)){
                    $('.reply_post', post).click(function(ele){
                        $('.reply_post_form', post).show();
                        $('.reply_post_form textarea', post).focus()
                        return false;
                    });
                    $('.reply_post', post).button();
                }
                if($('.add_attachment', post)){
                    $('.add_attachment', post).click(function(ele){
                        $('.add_attachment_form', post).show();
                        return false;
                    });
                }
                if($('.promote_to_thread', post)){
                    $('.promote_to_thread', post).click(function(ele){
                        $('.promote_to_thread_form', post).show();
                        return false;
                    });
                }
            });
        })();
        ''')

class Thread(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.thread'
    name='thread'
    params=['value', 'offset', 'pagesize', 'total', 'show_subject']
    value=None
    offset=None
    pagesize=None
    total=None
    show_subject=False
    widgets=dict(
        thread_header=ThreadHeader(),
        post_thread=PostThread(),
        post=Post(),
        edit_post=EditPost(submit_text='New Post'))

class Discussion(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.discussion'
    params=['value',
            'show_discussion_email', 'show_subject', 'allow_create_thread']
    value=None
    show_discussion_email=False
    show_subject=False
    allow_create_thread=False
    widgets=dict(
        discussion_header=DiscussionHeader(),
        edit_post=EditPost(submit_text='New Topic'),
        subscription_form=SubscriptionForm())

    def resources(self):
        for r in super(Discussion, self).resources(): yield r
<commit_msg>Fix for weird case where c.widget is set by SOMEONE to an empty string<commit_after>from pylons import c

import ew

from pyforge.lib import validators as V
from pyforge.lib import helpers as h
from pyforge.lib.widgets import form_fields as ffw
from pyforge import model as M

# Discussion forms
class ModerateThread(ew.SimpleForm):
    class buttons(ew.WidgetsList):
        delete=ew.SubmitButton(label='Delete Thread')
    submit_text=None

class ModeratePost(ew.SimpleForm):
    template='genshi:pyforge.lib.widgets.templates.moderate_post'
    submit_text=None

class FlagPost(ew.SimpleForm):
    template='genshi:pyforge.lib.widgets.templates.flag_post'
    submit_text=None

class AttachPost(ew.SimpleForm):
    submit_text='Attach File'
    enctype='multipart/form-data'
    class fields(ew.WidgetsList):
        file_info=ew.InputField(field_type='file')

class ModeratePosts(ew.SimpleForm):
    template='genshi:pyforge.lib.widgets.templates.moderate_posts'
    submit_text=None
    def resources(self):
        for r in super(ModeratePosts, self).resources(): yield r
        yield ew.JSScript('''
      (function($){
          var tbl = $('form table');
          var checkboxes = $('input[type=checkbox]', tbl);
          $('a[href=#]', tbl).click(function() {
              checkboxes.each(function() {
                  if(this.checked) { this.checked = false; }
                  else { this.checked = true; }
              });
              return false;
          });
      }(jQuery));''')

class PostFilter(ew.SimpleForm):
    submit_text=None
    method='GET'
    fields = [
        ew.FieldSet(label='Post Filter', fields=[
                ew.SingleSelectField(
                    name='status',
                    label='Show posts with status',
                    options=[
                        ew.Option(py_value='-', label='Any'),
                        ew.Option(py_value='spam', label='Spam'),
                        ew.Option(py_value='pending', label='Pending moderation'),
                        ew.Option(py_value='ok', label='Ok')],
                    if_missing='-'),
                ew.IntField(name='flag',
                            label='Show posts with at least "n" flags',
                            css_class='text',
                            if_missing=0),
                ew.SubmitButton(label='Filter Posts')
                ])
        ]

class EditPost(ew.SimpleForm):
    show_subject=False

    # this ickiness is to override the default submit button
    def __call__(self, **kw):
        result = super(EditPost, self).__call__(**kw)
        submit_button = ffw.SubmitButton(label=result['submit_text'])
        result['extra_fields'] = [submit_button]
        result['buttons'] = [submit_button]
        return result

    @property
    def fields(self):
        def _():
            if getattr(c, 'widget', '') != '':
                if c.widget.response.get('show_subject', self.show_subject):
                    yield ew.TextField(name='subject')
            else:
                yield ew.TextField(name='subject', if_missing=None)
            yield ffw.MarkdownEdit(name='text')
        return _()

    def resources(self):
        for r in ew.TextField(name='subject').resources(): yield r
        for r in ffw.MarkdownEdit(name='text').resources(): yield r

class _ThreadsTable(ew.SimpleForm):
    template='pyforge.lib.widgets.templates.threads_table'
    class hidden_fields(ew.WidgetsList):
        _id=ew.HiddenField(validator=V.Ming(M.Thread))
    class fields(ew.WidgetsList):
        num_replies=ew.HTMLField(show_label=True, label='Num Posts')
        num_views=ew.HTMLField(show_label=True)
        last_post=ew.HTMLField(text="${value and value.summary()}", show_label=True)
        subscription=ew.Checkbox(suppress_label=True, show_label=True)
    fields.insert(0, ew.LinkField(
            label='Subject', text="${value['subject']}",
            href="${value['url']()}", show_label=True))

class SubscriptionForm(ew.SimpleForm):
    template='pyforge.lib.widgets.templates.subscription_form'
    value=None
    show_discussion_email=False
    show_subject=False
    allow_create_thread=False
    params=['value',
            'show_discussion_email', 'show_subject', 'allow_create_thread']
    class fields(ew.WidgetsList):
        threads=_ThreadsTable()
        new_topic = EditPost(submit_text='New Topic', if_missing=None)
    submit_text='Update Subscriptions'
    def resources(self):
        for r in super(SubscriptionForm, self).resources(): yield r
        yield ew.JSScript('''
        $(window).load(function() {
            $('tbody').children(':even').addClass('even');
            $('.discussion_subscription_form').each(function(){
                var discussion = this;
                $('.submit', discussion).button();
                if($('.new_topic', discussion)){
                    $('.new_topic', discussion).click(function(ele){
                        $('.new_topic_form', discussion).show();
                        return false;
                    });
                }
                $('.follow', discussion).click(function(ele){
                    $('.follow_form', discussion).submit();
                    return false;
                });
            });
        });''')

# Widgets
class HierWidget(ew.Widget):
    widgets = {}

    def __call__(self, **kw):
        response = super(HierWidget, self).__call__(**kw)
        response['widgets'] = self.widgets
        return response

    def resources(self):
        for w in self.widgets.itervalues():
            for r in w.resources():
                yield r

class Attachment(ew.Widget):
    template='genshi:pyforge.lib.widgets.templates.attachment'
    params=['value', 'post']
    value=None
    post=None

class DiscussionHeader(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.discussion_header'
    params=['value']
    value=None
    widgets=dict(
        edit_post=EditPost(submit_text='New Thread'))

class ThreadHeader(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.thread_header'
    params=['value', 'offset', 'pagesize', 'total', 'show_moderate']
    value=None
    offset=None
    pagesize=None
    total=None
    show_moderate=False
    widgets=dict(
        moderate_thread=ModerateThread())

class PostHeader(ew.Widget):
    template='genshi:pyforge.lib.widgets.templates.post_header'
    params=['value']
    value=None

class PostThread(ew.Widget):
    template='genshi:pyforge.lib.widgets.templates.post_thread'
    params=['value']
    value=None

class Post(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.post'
    params=['value', 'show_subject', 'indent', 'supress_promote']
    value=None
    indent=0
    show_subject=False
    supress_promote=False
    widgets=dict(
        flag_post=FlagPost(),
        moderate_post=ModeratePost(),
        edit_post=EditPost(submit_text='Edit Post'),
        attach_post=AttachPost(submit_text='Attach'),
        attachment=Attachment())
    def resources(self):
        for r in super(Post, self).resources(): yield r
        for w in self.widgets.itervalues():
            for r in w.resources():
                yield r
        yield ew.JSScript('''
        (function(){
            $('.discussion-post').each(function(){
                var post = this;
                $('.submit', post).button();
                $('.flag_post, .delete_post', post).click(function(ele){
                    this.parentNode.submit();
                    return false;
                });
                if($('.edit_post', post)){
                    $('.edit_post', post).click(function(ele){
                        $('.edit_post_form', post).show();
                        return false;
                    });
                }
                if($('.reply_post', post)){
                    $('.reply_post', post).click(function(ele){
                        $('.reply_post_form', post).show();
                        $('.reply_post_form textarea', post).focus()
                        return false;
                    });
                    $('.reply_post', post).button();
                }
                if($('.add_attachment', post)){
                    $('.add_attachment', post).click(function(ele){
                        $('.add_attachment_form', post).show();
                        return false;
                    });
                }
                if($('.promote_to_thread', post)){
                    $('.promote_to_thread', post).click(function(ele){
                        $('.promote_to_thread_form', post).show();
                        return false;
                    });
                }
            });
        })();
        ''')

class Thread(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.thread'
    name='thread'
    params=['value', 'offset', 'pagesize', 'total', 'show_subject']
    value=None
    offset=None
    pagesize=None
    total=None
    show_subject=False
    widgets=dict(
        thread_header=ThreadHeader(),
        post_thread=PostThread(),
        post=Post(),
        edit_post=EditPost(submit_text='New Post'))

class Discussion(HierWidget):
    template='genshi:pyforge.lib.widgets.templates.discussion'
    params=['value',
            'show_discussion_email', 'show_subject', 'allow_create_thread']
    value=None
    show_discussion_email=False
    show_subject=False
    allow_create_thread=False
    widgets=dict(
        discussion_header=DiscussionHeader(),
        edit_post=EditPost(submit_text='New Topic'),
        subscription_form=SubscriptionForm())

    def resources(self):
        for r in super(Discussion, self).resources(): yield r
<|endoftext|>


<commit_before>import transitfeed
import time
from pytz import timezone
import pytz
import sys
import datetime
sys.path.append("../../..")
from graphserver.core import Graph, Street, ServicePeriod, TripHopSchedule, ServiceCalendar, State, Timezone, Wait, TripHop, Headway, TimezonePeriod
from graphserver.util import TimeHelpers
import csv
import calendar
import os
from datetime import date, timedelta

if [int(x) for x in transitfeed.__version__.split('.')] >= [1,1,6]:
    raise Exception("transitfeed.__version__ < 1.1.6")

def parse_date(date):
    return (int(date[0:4]), int(date[4:6]), int(date[6:8]))

def get_service_ids(sched, date):
    if type(date)==datetime.date or type(date)==datetime.datetime:
        date = "%0.4d%0.2d%0.2d"%(date.year,date.month,date.day)

    return [sp.service_id for sp in filter(lambda x:x.IsActiveOn(date), sched.GetServicePeriodList())]

def timezone_from_agency(sched, agency_id):
    #try getting agency by agency_id. If that fails, try getting by agency name
    try:
        agency = sched.GetAgency(agency_id)
    except KeyError:
        agency = filter(lambda x:x.agency_name==agency_id, sched.GetAgencyList())[0]

    return pytz.timezone( agency.agency_timezone )

def day_bounds_from_sched(sched):
    sid_start = min( [trip.GetStartTime() for trip in sched.GetTripList()] )
    sid_end   = max( [trip.GetEndTime() for trip in sched.GetTripList()] )

    return (sid_start, sid_end)

def schedule_to_service_calendar(sched, agency_id):
    timezone = timezone_from_agency(sched, agency_id)

    day_start, day_end = day_bounds_from_sched(sched)

    startdate, enddate = [ datetime.datetime( *parse_date(x) ) for x in sched.GetDateRange() ]

    cal = ServiceCalendar()

    for currdate in iter_dates(startdate, enddate):
        local_dt = timezone.localize(currdate)

        service_ids = [x.encode("ascii") for x in get_service_ids( sched, currdate )]

        this_day_begins = timezone.normalize( local_dt + timedelta(seconds=day_start) )
        this_day_ends = timezone.normalize( local_dt + timedelta(seconds=day_end)  )

        cal.add_period( TimeHelpers.datetime_to_unix(this_day_begins), TimeHelpers.datetime_to_unix(this_day_ends), service_ids )

    return cal

def group(ary, groupfunc=lambda x:x):
    clumper = {}
    for item in ary:
        grouplabel = groupfunc(item)
        if grouplabel not in clumper:
            clumper[grouplabel] = []
        clumper[grouplabel].append( item )

    return clumper.values()

def iter_dates(startdate, enddate):
    currdate = startdate
    while currdate <= enddate:
        yield currdate
        currdate += timedelta(1)

def stoptimes_from_stoptimetrips(stts):
    """return transitfeed StopTime objects from the tuple list returned from stop.GetStopTimeTrips()"""

    for time,(trip,trindex),is_tp in stts:
        yield trip.GetStopTimes()[trindex]

def eventtimes_from_stoptimetrips(stts):
    """return all events (arrivals and departures) that occur at this stop, given this stop's stop.GetStopTimeTrips() list of tuples"""

    for stoptime in stoptimes_from_stoptimetrips(stts):
        yield stoptime.arrival_secs
        yield stoptime.departure_secs

def distinct(ary):
    """eliminate duplicate items from list"""
    return list(set(ary))

def unique_eventtimes_from_stoptimetrips(stts):
    """return a sorted list of all departure and arrival times occurring at this stop, given the stops's stop.GetStopTimeTrips() list of tuples"""
    eventtimes = distinct(eventtimes_from_stoptimetrips(stts))
    eventtimes.sort()
    return eventtimes

def cons(ary):
    """return all pairs of consecutive elements"""
    for x,y in zip(ary[:-1],ary[1:]):
        yield (x,y)

def route_for_trip(trip):
    return trip._schedule.routes[ trip.route_id ]

class GTFSLoadable:
    def _triphops_from_stop(self, sched, stop, agency):
        ret = []

        for time, (trip,ix), is_timepoint in stop.GetStopTimeTrips():
            route = sched.routes[trip.route_id]

            stoptimes = trip.GetStopTimes()
            if ix != len(stoptimes)-1 and route.agency_id == agency.agency_id:
                ret.append( (trip, stoptimes[ix], stoptimes[ix+1], trip.route_id ) )

        return ret

    def _group_triphops(self, raw_triphops):
        """Takes a bunch of tuples of (trip,stoptime,stoptime). Returns them grouped by trip,fromv,tov,routeid"""

        #Get dictionary with keys corresponding to a destination stop_id,service_id pair, and initialize to []
        trip_ids = dict.fromkeys( ["%s+%s+%s"%(tov.stop.stop_id,trip.service_id,routeid) for trip,fromv,tov,routeid in raw_triphops] )
        for k in trip_ids.keys():
            trip_ids[k] = []

        #Sort the triphops into their buckets
        for trip,fromv,tov,routeid in raw_triphops:
            trip_ids["%s+%s+%s"%(tov.stop.stop_id,trip.service_id,routeid)].append( (trip,fromv,tov) )

        #Convert dict to list
        thss = trip_ids.values()

        #Sort each triphopschedule by trihop departure time
        for ths in thss:
            ths.sort(lambda x,y: cmp(x[1].departure_time,y[1].departure_time))

        return thss

    def _raw_triphopschedules_from_stop(self, sched, stop, agency):
        triphops = self._triphops_from_stop(sched, stop, agency)
        thss = self._group_triphops(triphops)

        return thss

    def load_gtfs(self, sched_or_datadir, prefix="gtfs"):

        if type(sched_or_datadir)==str:
            sched = transitfeed.Loader(sched_or_datadir).Load()
        else:
            sched = sched_or_datadir

        #add all vertices
        print "adding stop vertices"
        for stop in sched.GetStopList():
            self.add_vertex(prefix+stop.stop_id)

        self.numagencies = len(sched.GetAgencyList())

        for agency, i in zip( sched.GetAgencyList(), range(len(sched.GetAgencyList())) ):
            self._load_agency(sched, agency, i, prefix)

    def _load_agency(self, sched, agency, agency_int, prefix):
        cal = schedule_to_service_calendar(sched, agency.agency_id)

        gs_tz = Timezone.generate(agency.agency_timezone)

        for stop in sched.GetStopList():
            print "loading stop %s for agency %s"%(stop.stop_id, agency.agency_id)
            rawtriphopschedules = self._raw_triphopschedules_from_stop(sched, stop, agency)

            for rawtriphopschedule in rawtriphopschedules:
                #hops = [(fromv.departure_secs, tov.arrival_secs, trip.trip_id.encode("ascii")) for trip,fromv,tov in rawtriphopschedule]
                hops =[]
                service_id = rawtriphopschedule[0][0].service_id.encode("ascii")
                for trip,fromv,tov in rawtriphopschedule:
                  trip_id = trip.trip_id.encode("ascii")
                  if len(trip.GetHeadwayPeriodTuples())!=0:
                    for start_time, end_time, headway_secs in trip.GetHeadwayPeriodTuples():
                      hw = Headway( start_time, end_time, headway_secs, (tov.arrival_secs - fromv.departure_secs), trip_id, cal, gs_tz, agency_int, service_id )
                      e = self.add_edge( prefix+fromv.stop_id, prefix+tov.stop_id, hw )
                  else:
                    hops.append( (fromv.departure_secs, tov.arrival_secs, trip_id) )

                if hops!=[]:
                  ths = TripHopSchedule( hops, service_id, cal, gs_tz, agency_int )
                  e = self.add_edge( prefix+fromv.stop_id, prefix+tov.stop_id, ths )

    def load_gtfs_dag(self, sched_or_datadir, stops_timezone_name, prefix="gtfs"):

        if type(sched_or_datadir)==str:
            sched = transitfeed.Loader(sched_or_datadir).Load()
        else:
            sched = sched_or_datadir

        stz = Timezone.generate( stops_timezone_name )

        tzs = dict( [(agency.agency_id, Timezone.generate(agency.agency_timezone)) for agency in sched.GetAgencyList()] )
        scs = dict( [(agency.agency_id, schedule_to_service_calendar(sched, agency.agency_id)) for agency in sched.GetAgencyList()] )
        agints = dict( zip([x.agency_id for x in sched.GetAgencyList()], range(len(sched.GetAgencyList())) ))

        stop_times = dict( [(x.stop_id, set()) for x in sched.GetStopList()] )

        for route in sched.GetRouteList():
            print "adding all triphops for route %s"%route.route_id

            for trip in route.trips:
                agency_id = route_for_trip(trip).agency_id
                stops = trip.GetTimeInterpolatedStops()

                for (fromst_time,fromst_st,fromst_timed), (tost_time,tost_st,tost_timed) in cons(stops):
                    fromst_name = "%s@%s"%(fromst_st.stop.stop_id,fromst_time)
                    tost_name = "%s@%s"%(tost_st.stop.stop_id,tost_time)

                    self.add_vertex(fromst_name)
                    self.add_vertex(tost_name)

                    sc = scs[agency_id]
                    th = TripHop(fromst_time, tost_time, trip.trip_id, sc, tzs[agency_id], agints[agency_id], trip.service_id.encode("ascii"))

                    self.add_edge(fromst_name, tost_name, th)

                    stop_times[fromst_st.stop.stop_id].add(fromst_time)
                    stop_times[tost_st.stop.stop_id].add(tost_time)

        for stop_id, times in stop_times.iteritems():
            print "laying down beanstock for %s"%stop_id
            times = list(times)
            times.sort()

            self.add_vertex(stop_id)

            for time in times:
                self.add_edge( stop_id, "%s@%s"%(stop_id,time), Wait(time, stz) )
                self.add_edge( "%s@%s"%(stop_id,time), stop_id, Wait(time, stz) )

            for fromtime, totime in cons(times):
                self.add_edge( "%s@%s"%(stop_id,fromtime), "%s@%s"%(stop_id,totime), Wait(totime, stz) )

<commit_msg>wrong comparison operator<commit_after>import transitfeed
import time
from pytz import timezone
import pytz
import sys
import datetime
sys.path.append("../../..")
from graphserver.core import Graph, Street, ServicePeriod, TripHopSchedule, ServiceCalendar, State, Timezone, Wait, TripHop, Headway, TimezonePeriod
from graphserver.util import TimeHelpers
import csv
import calendar
import os
from datetime import date, timedelta

if [int(x) for x in transitfeed.__version__.split('.')] < [1,1,6]:
    raise Exception("transitfeed.__version__ < 1.1.6")

def parse_date(date):
    return (int(date[0:4]), int(date[4:6]), int(date[6:8]))

def get_service_ids(sched, date):
    if type(date)==datetime.date or type(date)==datetime.datetime:
        date = "%0.4d%0.2d%0.2d"%(date.year,date.month,date.day)

    return [sp.service_id for sp in filter(lambda x:x.IsActiveOn(date), sched.GetServicePeriodList())]

def timezone_from_agency(sched, agency_id):
    #try getting agency by agency_id. If that fails, try getting by agency name
    try:
        agency = sched.GetAgency(agency_id)
    except KeyError:
        agency = filter(lambda x:x.agency_name==agency_id, sched.GetAgencyList())[0]

    return pytz.timezone( agency.agency_timezone )

def day_bounds_from_sched(sched):
    sid_start = min( [trip.GetStartTime() for trip in sched.GetTripList()] )
    sid_end   = max( [trip.GetEndTime() for trip in sched.GetTripList()] )

    return (sid_start, sid_end)

def schedule_to_service_calendar(sched, agency_id):
    timezone = timezone_from_agency(sched, agency_id)

    day_start, day_end = day_bounds_from_sched(sched)

    startdate, enddate = [ datetime.datetime( *parse_date(x) ) for x in sched.GetDateRange() ]

    cal = ServiceCalendar()

    for currdate in iter_dates(startdate, enddate):
        local_dt = timezone.localize(currdate)

        service_ids = [x.encode("ascii") for x in get_service_ids( sched, currdate )]

        this_day_begins = timezone.normalize( local_dt + timedelta(seconds=day_start) )
        this_day_ends = timezone.normalize( local_dt + timedelta(seconds=day_end)  )

        cal.add_period( TimeHelpers.datetime_to_unix(this_day_begins), TimeHelpers.datetime_to_unix(this_day_ends), service_ids )

    return cal

def group(ary, groupfunc=lambda x:x):
    clumper = {}
    for item in ary:
        grouplabel = groupfunc(item)
        if grouplabel not in clumper:
            clumper[grouplabel] = []
        clumper[grouplabel].append( item )

    return clumper.values()

def iter_dates(startdate, enddate):
    currdate = startdate
    while currdate <= enddate:
        yield currdate
        currdate += timedelta(1)

def stoptimes_from_stoptimetrips(stts):
    """return transitfeed StopTime objects from the tuple list returned from stop.GetStopTimeTrips()"""

    for time,(trip,trindex),is_tp in stts:
        yield trip.GetStopTimes()[trindex]

def eventtimes_from_stoptimetrips(stts):
    """return all events (arrivals and departures) that occur at this stop, given this stop's stop.GetStopTimeTrips() list of tuples"""

    for stoptime in stoptimes_from_stoptimetrips(stts):
        yield stoptime.arrival_secs
        yield stoptime.departure_secs

def distinct(ary):
    """eliminate duplicate items from list"""
    return list(set(ary))

def unique_eventtimes_from_stoptimetrips(stts):
    """return a sorted list of all departure and arrival times occurring at this stop, given the stops's stop.GetStopTimeTrips() list of tuples"""
    eventtimes = distinct(eventtimes_from_stoptimetrips(stts))
    eventtimes.sort()
    return eventtimes

def cons(ary):
    """return all pairs of consecutive elements"""
    for x,y in zip(ary[:-1],ary[1:]):
        yield (x,y)

def route_for_trip(trip):
    return trip._schedule.routes[ trip.route_id ]

class GTFSLoadable:
    def _triphops_from_stop(self, sched, stop, agency):
        ret = []

        for time, (trip,ix), is_timepoint in stop.GetStopTimeTrips():
            route = sched.routes[trip.route_id]

            stoptimes = trip.GetStopTimes()
            if ix != len(stoptimes)-1 and route.agency_id == agency.agency_id:
                ret.append( (trip, stoptimes[ix], stoptimes[ix+1], trip.route_id ) )

        return ret

    def _group_triphops(self, raw_triphops):
        """Takes a bunch of tuples of (trip,stoptime,stoptime). Returns them grouped by trip,fromv,tov,routeid"""

        #Get dictionary with keys corresponding to a destination stop_id,service_id pair, and initialize to []
        trip_ids = dict.fromkeys( ["%s+%s+%s"%(tov.stop.stop_id,trip.service_id,routeid) for trip,fromv,tov,routeid in raw_triphops] )
        for k in trip_ids.keys():
            trip_ids[k] = []

        #Sort the triphops into their buckets
        for trip,fromv,tov,routeid in raw_triphops:
            trip_ids["%s+%s+%s"%(tov.stop.stop_id,trip.service_id,routeid)].append( (trip,fromv,tov) )

        #Convert dict to list
        thss = trip_ids.values()

        #Sort each triphopschedule by trihop departure time
        for ths in thss:
            ths.sort(lambda x,y: cmp(x[1].departure_time,y[1].departure_time))

        return thss

    def _raw_triphopschedules_from_stop(self, sched, stop, agency):
        triphops = self._triphops_from_stop(sched, stop, agency)
        thss = self._group_triphops(triphops)

        return thss

    def load_gtfs(self, sched_or_datadir, prefix="gtfs"):

        if type(sched_or_datadir)==str:
            sched = transitfeed.Loader(sched_or_datadir).Load()
        else:
            sched = sched_or_datadir

        #add all vertices
        print "adding stop vertices"
        for stop in sched.GetStopList():
            self.add_vertex(prefix+stop.stop_id)

        self.numagencies = len(sched.GetAgencyList())

        for agency, i in zip( sched.GetAgencyList(), range(len(sched.GetAgencyList())) ):
            self._load_agency(sched, agency, i, prefix)

    def _load_agency(self, sched, agency, agency_int, prefix):
        cal = schedule_to_service_calendar(sched, agency.agency_id)

        gs_tz = Timezone.generate(agency.agency_timezone)

        for stop in sched.GetStopList():
            print "loading stop %s for agency %s"%(stop.stop_id, agency.agency_id)
            rawtriphopschedules = self._raw_triphopschedules_from_stop(sched, stop, agency)

            for rawtriphopschedule in rawtriphopschedules:
                #hops = [(fromv.departure_secs, tov.arrival_secs, trip.trip_id.encode("ascii")) for trip,fromv,tov in rawtriphopschedule]
                hops =[]
                service_id = rawtriphopschedule[0][0].service_id.encode("ascii")
                for trip,fromv,tov in rawtriphopschedule:
                  trip_id = trip.trip_id.encode("ascii")
                  if len(trip.GetHeadwayPeriodTuples())!=0:
                    for start_time, end_time, headway_secs in trip.GetHeadwayPeriodTuples():
                      hw = Headway( start_time, end_time, headway_secs, (tov.arrival_secs - fromv.departure_secs), trip_id, cal, gs_tz, agency_int, service_id )
                      e = self.add_edge( prefix+fromv.stop_id, prefix+tov.stop_id, hw )
                  else:
                    hops.append( (fromv.departure_secs, tov.arrival_secs, trip_id) )

                if hops!=[]:
                  ths = TripHopSchedule( hops, service_id, cal, gs_tz, agency_int )
                  e = self.add_edge( prefix+fromv.stop_id, prefix+tov.stop_id, ths )

    def load_gtfs_dag(self, sched_or_datadir, stops_timezone_name, prefix="gtfs"):

        if type(sched_or_datadir)==str:
            sched = transitfeed.Loader(sched_or_datadir).Load()
        else:
            sched = sched_or_datadir

        stz = Timezone.generate( stops_timezone_name )

        tzs = dict( [(agency.agency_id, Timezone.generate(agency.agency_timezone)) for agency in sched.GetAgencyList()] )
        scs = dict( [(agency.agency_id, schedule_to_service_calendar(sched, agency.agency_id)) for agency in sched.GetAgencyList()] )
        agints = dict( zip([x.agency_id for x in sched.GetAgencyList()], range(len(sched.GetAgencyList())) ))

        stop_times = dict( [(x.stop_id, set()) for x in sched.GetStopList()] )

        for route in sched.GetRouteList():
            print "adding all triphops for route %s"%route.route_id

            for trip in route.trips:
                agency_id = route_for_trip(trip).agency_id
                stops = trip.GetTimeInterpolatedStops()

                for (fromst_time,fromst_st,fromst_timed), (tost_time,tost_st,tost_timed) in cons(stops):
                    fromst_name = "%s@%s"%(fromst_st.stop.stop_id,fromst_time)
                    tost_name = "%s@%s"%(tost_st.stop.stop_id,tost_time)

                    self.add_vertex(fromst_name)
                    self.add_vertex(tost_name)

                    sc = scs[agency_id]
                    th = TripHop(fromst_time, tost_time, trip.trip_id, sc, tzs[agency_id], agints[agency_id], trip.service_id.encode("ascii"))

                    self.add_edge(fromst_name, tost_name, th)

                    stop_times[fromst_st.stop.stop_id].add(fromst_time)
                    stop_times[tost_st.stop.stop_id].add(tost_time)

        for stop_id, times in stop_times.iteritems():
            print "laying down beanstock for %s"%stop_id
            times = list(times)
            times.sort()

            self.add_vertex(stop_id)

            for time in times:
                self.add_edge( stop_id, "%s@%s"%(stop_id,time), Wait(time, stz) )
                self.add_edge( "%s@%s"%(stop_id,time), stop_id, Wait(time, stz) )

            for fromtime, totime in cons(times):
                self.add_edge( "%s@%s"%(stop_id,fromtime), "%s@%s"%(stop_id,totime), Wait(totime, stz) )

<|endoftext|>


<commit_before># coding: utf-8

from __future__ import unicode_literals

from pymatgen.alchemy.filters import ContainsSpecieFilter, \
    SpecieProximityFilter, RemoveDuplicatesFilter
from pymatgen.core.lattice import Lattice
from pymatgen.core.structure import Structure
from pymatgen.core.periodic_table import Specie
from pymatgen.alchemy.transmuters import StandardTransmuter
from pymatgen.analysis.structure_matcher import StructureMatcher
from pymatgen.util.testing import PymatgenTest

from monty.json import MontyDecoder

import os
import json
import unittest

test_dir = os.path.join(os.path.dirname(__file__), "..", "..", "..",
                        'test_files')

class ContainsSpecieFilterTest(PymatgenTest):

    def test_filtering(self):
        coords = [[0, 0, 0], [0.75, 0.75, 0.75], [0.5, 0.5, 0.5],
                  [0.25, 0.25, 0.25]]
        lattice = Lattice([[3.0, 0.0, 0.0],
                           [1.0, 3.0, 0.00],
                           [0.00, -2.0, 3.0]])
        s = Structure(lattice,
                      [{"Si4+": 0.5, "O2-": 0.25, "P5+": 0.25},
                       {"Si4+": 0.5, "O2-": 0.25, "P5+": 0.25},
                       {"Si4+": 0.5, "O2-": 0.25, "P5+": 0.25},
                       {"Si4+": 0.5, "O2-": 0.25, "P5+": 0.25}], coords)

        species1 = [Specie('Si', 5), Specie('Mg', 2)]
        f1 = ContainsSpecieFilter(species1, strict_compare=True, AND=False)
        self.assertFalse(f1.test(s), 'Incorrect filter')
        f2 = ContainsSpecieFilter(species1, strict_compare=False, AND=False)
        self.assertTrue(f2.test(s), 'Incorrect filter')
        species2 = [Specie('Si', 4), Specie('Mg', 2)]
        f3 = ContainsSpecieFilter(species2, strict_compare=True, AND=False)
        self.assertTrue(f3.test(s), 'Incorrect filter')
        f4 = ContainsSpecieFilter(species2, strict_compare=False, AND=False)
        self.assertTrue(f4.test(s), 'Incorrect filter')

        species3 = [Specie('Si', 5), Specie('O', -2)]
        f5 = ContainsSpecieFilter(species3, strict_compare=True, AND=True)
        self.assertFalse(f5.test(s), 'Incorrect filter')
        f6 = ContainsSpecieFilter(species3, strict_compare=False, AND=True)
        self.assertTrue(f6.test(s), 'Incorrect filter')
        species4 = [Specie('Si', 4), Specie('Mg', 2)]
        f7 = ContainsSpecieFilter(species4, strict_compare=True, AND=True)
        self.assertFalse(f7.test(s), 'Incorrect filter')
        f8 = ContainsSpecieFilter(species4, strict_compare=False, AND=True)
        self.assertFalse(f8.test(s), 'Incorrect filter')

    def test_to_from_dict(self):
        species1 = ['Si5+', 'Mg2+']
        f1 = ContainsSpecieFilter(species1, strict_compare=True, AND=False)
        d = f1.as_dict()
        self.assertIsInstance(ContainsSpecieFilter.from_dict(d),
                              ContainsSpecieFilter)

class SpecieProximityFilterTest(PymatgenTest):

    def test_filter(self):
        s = self.get_structure("Li10GeP2S12")
        sf = SpecieProximityFilter({"Li": 1})
        self.assertTrue(sf.test(s))
        sf = SpecieProximityFilter({"Li": 2})
        self.assertFalse(sf.test(s))
        sf = SpecieProximityFilter({"P": 1})
        self.assertTrue(sf.test(s))
        sf = SpecieProximityFilter({"P": 5})
        self.assertFalse(sf.test(s))

    def test_to_from_dict(self):
        sf = SpecieProximityFilter({"Li": 1})
        d = sf.as_dict()
        self.assertIsInstance(SpecieProximityFilter.from_dict(d),
                              SpecieProximityFilter)

class RemoveDuplicatesFilterTest(unittest.TestCase):

    def setUp(self):
        with open(os.path.join(test_dir, "TiO2_entries.json"), 'r') as fp:
            entries = json.load(fp, cls=MontyDecoder)
        self._struct_list = [e.structure for e in entries]
        self._sm = StructureMatcher()

    def test_filter(self):
        transmuter = StandardTransmuter.from_structures(self._struct_list)
        fil = RemoveDuplicatesFilter()
        transmuter.apply_filter(fil)
        self.assertEqual(len(transmuter.transformed_structures), 11)

    def test_to_from_dict(self):
        fil = RemoveDuplicatesFilter()
        d = fil.as_dict()
        self.assertIsInstance(RemoveDuplicatesFilter().from_dict(d),
                              RemoveDuplicatesFilter)

if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.main()
<commit_msg>Add testcase for RemoveExistingFilter<commit_after># coding: utf-8

from __future__ import unicode_literals

from pymatgen.alchemy.filters import ContainsSpecieFilter, \
    SpecieProximityFilter, RemoveDuplicatesFilter, RemoveExistingFilter
from pymatgen.core.lattice import Lattice
from pymatgen.core.structure import Structure
from pymatgen.core.periodic_table import Specie
from pymatgen.alchemy.transmuters import StandardTransmuter
from pymatgen.analysis.structure_matcher import StructureMatcher
from pymatgen.util.testing import PymatgenTest

from monty.json import MontyDecoder

import os
import json
import unittest

test_dir = os.path.join(os.path.dirname(__file__), "..", "..", "..",
                        'test_files')

class ContainsSpecieFilterTest(PymatgenTest):

    def test_filtering(self):
        coords = [[0, 0, 0], [0.75, 0.75, 0.75], [0.5, 0.5, 0.5],
                  [0.25, 0.25, 0.25]]
        lattice = Lattice([[3.0, 0.0, 0.0],
                           [1.0, 3.0, 0.00],
                           [0.00, -2.0, 3.0]])
        s = Structure(lattice,
                      [{"Si4+": 0.5, "O2-": 0.25, "P5+": 0.25},
                       {"Si4+": 0.5, "O2-": 0.25, "P5+": 0.25},
                       {"Si4+": 0.5, "O2-": 0.25, "P5+": 0.25},
                       {"Si4+": 0.5, "O2-": 0.25, "P5+": 0.25}], coords)

        species1 = [Specie('Si', 5), Specie('Mg', 2)]
        f1 = ContainsSpecieFilter(species1, strict_compare=True, AND=False)
        self.assertFalse(f1.test(s), 'Incorrect filter')
        f2 = ContainsSpecieFilter(species1, strict_compare=False, AND=False)
        self.assertTrue(f2.test(s), 'Incorrect filter')
        species2 = [Specie('Si', 4), Specie('Mg', 2)]
        f3 = ContainsSpecieFilter(species2, strict_compare=True, AND=False)
        self.assertTrue(f3.test(s), 'Incorrect filter')
        f4 = ContainsSpecieFilter(species2, strict_compare=False, AND=False)
        self.assertTrue(f4.test(s), 'Incorrect filter')

        species3 = [Specie('Si', 5), Specie('O', -2)]
        f5 = ContainsSpecieFilter(species3, strict_compare=True, AND=True)
        self.assertFalse(f5.test(s), 'Incorrect filter')
        f6 = ContainsSpecieFilter(species3, strict_compare=False, AND=True)
        self.assertTrue(f6.test(s), 'Incorrect filter')
        species4 = [Specie('Si', 4), Specie('Mg', 2)]
        f7 = ContainsSpecieFilter(species4, strict_compare=True, AND=True)
        self.assertFalse(f7.test(s), 'Incorrect filter')
        f8 = ContainsSpecieFilter(species4, strict_compare=False, AND=True)
        self.assertFalse(f8.test(s), 'Incorrect filter')

    def test_to_from_dict(self):
        species1 = ['Si5+', 'Mg2+']
        f1 = ContainsSpecieFilter(species1, strict_compare=True, AND=False)
        d = f1.as_dict()
        self.assertIsInstance(ContainsSpecieFilter.from_dict(d),
                              ContainsSpecieFilter)

class SpecieProximityFilterTest(PymatgenTest):

    def test_filter(self):
        s = self.get_structure("Li10GeP2S12")
        sf = SpecieProximityFilter({"Li": 1})
        self.assertTrue(sf.test(s))
        sf = SpecieProximityFilter({"Li": 2})
        self.assertFalse(sf.test(s))
        sf = SpecieProximityFilter({"P": 1})
        self.assertTrue(sf.test(s))
        sf = SpecieProximityFilter({"P": 5})
        self.assertFalse(sf.test(s))

    def test_to_from_dict(self):
        sf = SpecieProximityFilter({"Li": 1})
        d = sf.as_dict()
        self.assertIsInstance(SpecieProximityFilter.from_dict(d),
                              SpecieProximityFilter)

class RemoveDuplicatesFilterTest(unittest.TestCase):

    def setUp(self):
        with open(os.path.join(test_dir, "TiO2_entries.json"), 'r') as fp:
            entries = json.load(fp, cls=MontyDecoder)
        self._struct_list = [e.structure for e in entries]
        self._sm = StructureMatcher()

    def test_filter(self):
        transmuter = StandardTransmuter.from_structures(self._struct_list)
        fil = RemoveDuplicatesFilter()
        transmuter.apply_filter(fil)
        self.assertEqual(len(transmuter.transformed_structures), 11)

    def test_to_from_dict(self):
        fil = RemoveDuplicatesFilter()
        d = fil.as_dict()
        self.assertIsInstance(RemoveDuplicatesFilter().from_dict(d),
                              RemoveDuplicatesFilter)

class RemoveExistingFilterTest(unittest.TestCase):

    def setUp(self):
        with open(os.path.join(test_dir, "TiO2_entries.json"), 'r') as fp:
            entries = json.load(fp, cls=MontyDecoder)
        self._struct_list = [e.structure for e in entries]
        self._sm = StructureMatcher()
        self._exisiting_structures = self._struct_list[:-1]

    def test_filter(self):
        fil = RemoveExistingFilter(self._exisiting_structures)
        transmuter = StandardTransmuter.from_structures(self._struct_list)
        transmuter.apply_filter(fil)
        self.assertEqual(len(transmuter.transformed_structures), 1)
        self.assertTrue(
            self._sm.fit(self._struct_list[-1],
                         transmuter.transformed_structures[-1].final_structure))

if __name__ == "__main__":
    #import sys;sys.argv = ['', 'Test.testName']
    unittest.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 re
import csv
import collections
import itertools
from io import open
import math
from six.moves import zip
import logging

from monty.json import MSONable, MontyDecoder

import numpy as np

from pymatgen.core.structure import Molecule
from pymatgen.io.qchem.outputs import edges_from_babel, build_MoleculeGraph, is_isomorphic
import networkx as nx

# have_babel = True
# try:
#     from pymatgen.io.babel import BabelMolAdaptor
#     import openbabel as ob
# except ImportError:
#     print("Cannot find OpenBabel! Thus, bonds must be provided by the user.")
#     have_babel = False

"""
This module defines tools to calculate bond dissocation energies.
"""

__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.fragment_entries = fragment_entries
        self.bond_dissociation_energies = []
        # self.all_bonds = edges_from_babel(target_molecule)
        self.ring_bonds = []
        self.bad_pairs = []
        self.mol_graph = build_MoleculeGraph(Molecule.from_dict(molecule_entry["output"]["optimized_molecule"]))#, self.all_bonds)
        for bond in self.mol_graph.graph.edges: #switch to all_bonds?
            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]
        for bond_pair in self.bond_pairs:
            self.fragment_and_process(bond_pair)
        self.solve_ring_bonds()
        # return self.bond_dissociation_energies

    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:
            frag1_entries = self.search_fragment_entries(frags[0])[0]
            frag2_entries = self.search_fragment_entries(frags[1])[0]
            for frag1 in frag1_entries:
                for frag2 in frag2_entries:
                    if frag1["input"]["rem"]["method"] == frag2["input"]["rem"]["method"] and frag1["input"]["rem"]["basis"] == frag2["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["input"]["rem"]["basis"] == frag1["input"]["rem"]["basis"] and self.molecule_entry["input"]["rem"]["method"] == frag1["input"]["rem"]["method"]:
                            new_entry = [bonds, self.molecule_entry["output"]["final_energy"] - (frag1["output"]["final_energy"] + frag2["output"]["final_energy"]), frag1["output"]["final_energy"], frag1["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"], frag2["output"]["initial_molecule"]["charge"]]
                            self.bond_dissociation_energies += [new_entry]

    def search_fragment_entries(self, frag):
        entries = []
        initial_entries = []
        final_entries = []
        for entry in self.fragment_entries:
            initial_graph = build_MoleculeGraph(Molecule.from_dict(entry["input"]["initial_molecule"])).graph
            final_graph = build_MoleculeGraph(Molecule.from_dict(entry["output"]["initial_molecule"])).graph
            if is_isomorphic(frag.graph, initial_graph) and is_isomorphic(frag.graph, final_graph):
                entries += [entry]
            elif is_isomorphic(frag.graph, initial_graph):
                initial_entries += [entry]
            elif is_isomorphic(frag.graph, final_graph):
                final_entries += [entry]
        return [entries, initial_entries, final_entries]

    def solve_ring_bonds(self):
        pass
<commit_msg>bug fix<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 re
import csv
import collections
import itertools
from io import open
import math
from six.moves import zip
import logging

from monty.json import MSONable, MontyDecoder

import numpy as np

from pymatgen.core.structure import Molecule
from pymatgen.io.qchem.outputs import edges_from_babel, build_MoleculeGraph, is_isomorphic
import networkx as nx

# have_babel = True
# try:
#     from pymatgen.io.babel import BabelMolAdaptor
#     import openbabel as ob
# except ImportError:
#     print("Cannot find OpenBabel! Thus, bonds must be provided by the user.")
#     have_babel = False

"""
This module defines tools to calculate bond dissocation energies.
"""

__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.fragment_entries = fragment_entries
        self.bond_dissociation_energies = []
        # self.all_bonds = edges_from_babel(target_molecule)
        self.ring_bonds = []
        self.bad_pairs = []
        self.mol_graph = build_MoleculeGraph(Molecule.from_dict(molecule_entry["output"]["optimized_molecule"]))#, self.all_bonds)
        for bond in self.mol_graph.graph.edges: #switch to all_bonds?
            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]
        for bond_pair in self.bond_pairs:
            self.fragment_and_process(bond_pair)
        self.solve_ring_bonds()
        # return self.bond_dissociation_energies

    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:
            frag1_entries = self.search_fragment_entries(frags[0])[0]
            frag2_entries = self.search_fragment_entries(frags[1])[0]
            for frag1 in frag1_entries:
                for frag2 in frag2_entries:
                    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"]:
                            new_entry = [bonds, self.molecule_entry["output"]["final_energy"] - (frag1["output"]["final_energy"] + frag2["output"]["final_energy"]), frag1["output"]["final_energy"], frag1["output"]["initial_molecule"]["charge"], frag2["output"]["final_energy"], frag2["output"]["initial_molecule"]["charge"]]
                            self.bond_dissociation_energies += [new_entry]

    def search_fragment_entries(self, frag):
        entries = []
        initial_entries = []
        final_entries = []
        for entry in self.fragment_entries:
            initial_graph = build_MoleculeGraph(Molecule.from_dict(entry["input"]["initial_molecule"])).graph
            final_graph = build_MoleculeGraph(Molecule.from_dict(entry["output"]["initial_molecule"])).graph
            if is_isomorphic(frag.graph, initial_graph) and is_isomorphic(frag.graph, final_graph):
                entries += [entry]
            elif is_isomorphic(frag.graph, initial_graph):
                initial_entries += [entry]
            elif is_isomorphic(frag.graph, final_graph):
                final_entries += [entry]
        return [entries, initial_entries, final_entries]

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


<commit_before># -*- coding: utf-8 -*-
import codecs
from unittest import TestCase

from nose.tools import assert_equal  # @UnresolvedImport

from pyoracc.atf.atffile import AtfFile
from pyoracc.test.fixtures import belsunu, output_filepath

from ...atf.atflex import AtfLexer
from ...atf.atfyacc import AtfParser
from pyoracc.model.line import Line

class testSerializer(TestCase):

    def setUp(self):
        """
        Initialize lexer and parser.
        """
        self.lexer = AtfLexer().lexer
        self.parser =  AtfParser().parser

    def parse(self, any_str):
        """
        Parse input string, could be just a line or a whole file content.
        """
        parsed = AtfFile(any_str)
        return parsed

    def serialize(self, any_object):
        """
        Serialize input object, from a simple lemma to a whole AtfFile object.
        """
        serialized = any_object.serialize()
        return serialized

    def parse_then_serialize(self, any_str):
        """
        Shorthand for testing serialization.
        """
        return self.serialize(self.parse(any_str))

    def open_file(self, filename):
        """
        Open serialized file and output contents
        """
        return codecs.open(filename, "r", "utf-8").read()

    def save_file(self, content, filename):
        """
        Write serialized file on disk
        """
        serialized_file = codecs.open(filename, "w", "utf-8")
        serialized_file.write(content)
        serialized_file.close()

    def test_belsunu_serializer(self):
        """
        Parse belsunu.atf, then serialize, parse again, serialize again, compare.
        Comparing serialized output with input file would bring up differences that might not be significant (white spaces, newlines, etc).
        The solution is to parse again the serialized file, serialize again, then compare the two serializations.
        """
        serialized_1 = self.parse_then_serialize(belsunu())
        self.save_file(serialized_1, output_filepath("belsunu.atf"))
        serialized_2 = self.parse_then_serialize(serialized_1)
        assert_equal(serialized_1, serialized_2)

    def test_line_word(self):
        """
        Get a sample word with unicode chars and check serialization is correct.
        """
        line=Line("1")
        line.words.append(u"\u2086")
        line_ser = line.serialize()
        assert_equal(line_ser,"1. " + u"\u2086 ")

    def test_line_words(self):
        """
        Get a sample word with unicode chars and check serialization is correct.
        1. [MU] 1.03-KAM {iti}AB GE₆ U₄ 2-KAM
        """
        atf_file = AtfFile(belsunu())
        uline = atf_file.text.children[0].children[0].children[0]
        uwords = uline.words
        gold = [u'[MU]', u'1.03-KAM', u'{iti}AB', u'GE\u2086', u'U\u2084', u'2-KAM']
        assert_equal(uwords, gold)

    def test_line_lemmas(self):
        """
        Get a sample word with unicode chars and check serialization is correct.
        šatti[year]N; n; Ṭebetu[1]MN; mūša[at night]AV; ūm[day]N; n
        """
        atf_file = AtfFile(belsunu())
        uline = atf_file.text.children[0].children[0].children[0]
        ulemmas = uline.lemmas
        gold = [u'\u0161atti[year]N', u'n', u'\u1e6cebetu[1]MN', u'm\u016b\u0161a[at night]AV', u'\u016bm[day]N', u'n']
        assert_equal(ulemmas, gold)

#TODO: Build list of atf files for testing and make a test to go through the list of test and try serializing each of them.

#     def test_text_code_and_description(self):
#         """
#         Check if serializing works for the code/description case - first line of ATF texts.
#         Note the parser always returns an AtfFile object, even when it's not ATF-compliant.
#         """
#         atf = self.parse("&X001001 = JCS 48, 089\n")
#         serialized = self.serialize(atf)
#         assert_equal(serialized.strip()+"\n", "&X001001 = JCS 48, 089\n")

#     def test_text_project(self):
#         """
#         Check if serializing works for the project lines.
#         Note the parser always returns an AtfFile object, even when it's not ATF-compliant.
#         """
#         serialized = self.parse_then_serialize("#project: cams/gkab\n")
#         assert_equal(serialized.strip()+"\n", "#project: cams/gkab\n")
#
#     def test_text_language(self):
#
#     def test_text_protocols(self):

        <commit_msg>Test word comment out since it needs to be compared to a word object instead of a line to work smoothly, and we don't have that.<commit_after># -*- coding: utf-8 -*-
import codecs
from unittest import TestCase

from nose.tools import assert_equal  # @UnresolvedImport

from pyoracc.atf.atffile import AtfFile
from pyoracc.test.fixtures import belsunu, output_filepath

from ...atf.atflex import AtfLexer
from ...atf.atfyacc import AtfParser
from pyoracc.model.line import Line

class testSerializer(TestCase):

    def setUp(self):
        """
        Initialize lexer and parser.
        """
        self.lexer = AtfLexer().lexer
        self.parser =  AtfParser().parser

    def parse(self, any_str):
        """
        Parse input string, could be just a line or a whole file content.
        """
        parsed = AtfFile(any_str)
        return parsed

    def serialize(self, any_object):
        """
        Serialize input object, from a simple lemma to a whole AtfFile object.
        """
        serialized = any_object.serialize()
        return serialized

    def parse_then_serialize(self, any_str):
        """
        Shorthand for testing serialization.
        """
        return self.serialize(self.parse(any_str))

    def open_file(self, filename):
        """
        Open serialized file and output contents
        """
        return codecs.open(filename, "r", "utf-8").read()

    def save_file(self, content, filename):
        """
        Write serialized file on disk
        """
        serialized_file = codecs.open(filename, "w", "utf-8")
        serialized_file.write(content)
        serialized_file.close()

    def test_belsunu_serializer(self):
        """
        Parse belsunu.atf, then serialize, parse again, serialize again, compare.
        Comparing serialized output with input file would bring up differences that might not be significant (white spaces, newlines, etc).
        The solution is to parse again the serialized file, serialize again, then compare the two serializations.
        """
        serialized_1 = self.parse_then_serialize(belsunu())
        self.save_file(serialized_1, output_filepath("belsunu.atf"))
        serialized_2 = self.parse_then_serialize(serialized_1)
        assert_equal(serialized_1, serialized_2)

#     def test_line_word(self):
#         """
#         Get a sample word with unicode chars and check serialization is correct.
#         """
#         line=Line("1")
#         line.words.append(u"\u2086")
#         line_ser = line.serialize()
#         assert_equal(line_ser,"1.\t" + u"\u2086")

    def test_line_words(self):
        """
        Get a sample word with unicode chars and check serialization is correct.
        1. [MU] 1.03-KAM {iti}AB GE₆ U₄ 2-KAM
        """
        atf_file = AtfFile(belsunu())
        uline = atf_file.text.children[0].children[0].children[0]
        uwords = uline.words
        gold = [u'[MU]', u'1.03-KAM', u'{iti}AB', u'GE\u2086', u'U\u2084', u'2-KAM']
        assert_equal(uwords, gold)

    def test_line_lemmas(self):
        """
        Get a sample word with unicode chars and check serialization is correct.
        šatti[year]N; n; Ṭebetu[1]MN; mūša[at night]AV; ūm[day]N; n
        """
        atf_file = AtfFile(belsunu())
        uline = atf_file.text.children[0].children[0].children[0]
        ulemmas = uline.lemmas
        gold = [u' \u0161atti[year]N', u'n', u'\u1e6cebetu[1]MN', u'm\u016b\u0161a[at night]AV', u'\u016bm[day]N', u'n']
        assert_equal(ulemmas, gold)

#TODO: Build list of atf files for testing and make a test to go through the list of test and try serializing each of them.

#     def test_text_code_and_description(self):
#         """
#         Check if serializing works for the code/description case - first line of ATF texts.
#         Note the parser always returns an AtfFile object, even when it's not ATF-compliant.
#         """
#         atf = self.parse("&X001001 = JCS 48, 089\n")
#         serialized = self.serialize(atf)
#         assert_equal(serialized.strip()+"\n", "&X001001 = JCS 48, 089\n")

#     def test_text_project(self):
#         """
#         Check if serializing works for the project lines.
#         Note the parser always returns an AtfFile object, even when it's not ATF-compliant.
#         """
#         serialized = self.parse_then_serialize("#project: cams/gkab\n")
#         assert_equal(serialized.strip()+"\n", "#project: cams/gkab\n")
#
#     def test_text_language(self):
#
#     def test_text_protocols(self):

        <|endoftext|>


<commit_before>import time, requests, json
from pyspider.database.base.projectdb import ProjectDB as BaseProjectDB

class ProjectDB(BaseProjectDB):
    __collection_name__ = 'projectdb'

    def __init__(self, url, database='projectdb'):
        self.base_url = url
        self.url = url + database + "/"
        self.database = database
        self.insert('', {})

    def _default_fields(self, each):
        if each is None:
            return each
        each.setdefault('group', None)
        each.setdefault('status', 'TODO')
        each.setdefault('script', '')
        each.setdefault('comments', None)
        each.setdefault('rate', 0)
        each.setdefault('burst', 0)
        each.setdefault('updatetime', 0)
        return each

    def insert(self, name, obj={}):
        url = self.url + name
        obj = dict(obj)
        obj['name'] = name
        obj['updatetime'] = time.time()
        res = requests.put(url, data = json.dumps(obj), headers = {"Content-Type": "application/json"}).json()
        print('[couchdb projectdb insert] - url: {} data: {} res: {}'.format(url, json.dumps(obj), res))
        return res

    def update(self, name, obj={}, **kwargs):
        if self.get(name) is None:
            return None
        obj = dict(obj)
        obj.update(kwargs)
        self.insert(name, obj)

    def get_all(self, fields=None):
        if fields is None:
            fields = []
        payload = {
            "selector": {},
            "fields": fields
        }
        url = self.url + "_find"
        res = requests.post(url, data=json.dumps(payload), headers={"Content-Type": "application/json"}).json()
        print('[couchdb projectdb get_all] - url: {} res: {}'.format(url, res))
        for doc in res['docs']:
            yield self._default_fields(doc)

    def get(self, name, fields=None):
        if fields is None:
            fields = []
        payload = {
            "selector": {"name": name},
            "fields": fields,
            "limit": 1
        }
        url = self.url + "_find"
        res = requests.post(url, data=json.dumps(payload), headers={"Content-Type": "application/json"}).json()
        print('[couchdb projectdb get] - url: {} res: {}'.format(url, res))
        if len(res['docs']) == 0:
            return None
        return self._default_fields(res['docs'][0])

    def check_update(self, timestamp, fields=None):
        if fields is None:
            fields = []
        for project in self.get_all(fields=('updatetime', 'name')):
            if project['updatetime'] > timestamp:
                project = self.get(project['name'], fields)
                yield self._default_fields(project)

    def drop(self, name):
        doc = self.get(name)
        url = self.url + name + "/" + doc["_id"]
        res = requests.delete(url, headers={"Content-Type": "application/json"}).json()
        print('[couchdb projectdb drop] - url: {} res: {}'.format(url, res))
        return res

    def drop_database(self):
        res = requests.delete(self.url, headers={"Content-Type": "application/json"}).json()
        print('[couchdb projectdb drop_database] - url: {} res: {}'.format(self.url, res))
        return res

<commit_msg>fixed update bug<commit_after>import time, requests, json
from pyspider.database.base.projectdb import ProjectDB as BaseProjectDB

class ProjectDB(BaseProjectDB):
    __collection_name__ = 'projectdb'

    def __init__(self, url, database='projectdb'):
        self.base_url = url
        self.url = url + database + "/"
        self.database = database
        self.insert('', {})

    def _default_fields(self, each):
        if each is None:
            return each
        each.setdefault('group', None)
        each.setdefault('status', 'TODO')
        each.setdefault('script', '')
        each.setdefault('comments', None)
        each.setdefault('rate', 0)
        each.setdefault('burst', 0)
        each.setdefault('updatetime', 0)
        return each

    def insert(self, name, obj={}):
        url = self.url + name
        obj = dict(obj)
        obj['name'] = name
        obj['updatetime'] = time.time()
        res = requests.put(url, data = json.dumps(obj), headers = {"Content-Type": "application/json"}).json()
        print('[couchdb projectdb insert] - url: {} data: {} res: {}'.format(url, json.dumps(obj), res))
        return res

    def update(self, name, obj={}, **kwargs):
        if self.get(name) is None:
            return None
        obj = dict(obj)
        obj.update(kwargs)
        obj['updatetime'] = time.time()
        self.insert(name, obj)

    def get_all(self, fields=None):
        if fields is None:
            fields = []
        payload = {
            "selector": {},
            "fields": fields
        }
        url = self.url + "_find"
        res = requests.post(url, data=json.dumps(payload), headers={"Content-Type": "application/json"}).json()
        print('[couchdb projectdb get_all] - url: {} res: {}'.format(url, res))
        for doc in res['docs']:
            yield self._default_fields(doc)

    def get(self, name, fields=None):
        if fields is None:
            fields = []
        payload = {
            "selector": {"name": name},
            "fields": fields,
            "limit": 1
        }
        url = self.url + "_find"
        res = requests.post(url, data=json.dumps(payload), headers={"Content-Type": "application/json"}).json()
        print('[couchdb projectdb get] - url: {} res: {}'.format(url, res))
        if len(res['docs']) == 0:
            return None
        return self._default_fields(res['docs'][0])

    def check_update(self, timestamp, fields=None):
        if fields is None:
            fields = []
        for project in self.get_all():
            # save an extra request
            if project['updatetime'] > timestamp:
                yield project

    def drop(self, name):
        doc = self.get(name)
        url = self.url + name + "/" + doc["_id"]
        res = requests.delete(url, headers={"Content-Type": "application/json"}).json()
        print('[couchdb projectdb drop] - url: {} res: {}'.format(url, res))
        return res

    def drop_database(self):
        res = requests.delete(self.url, headers={"Content-Type": "application/json"}).json()
        print('[couchdb projectdb drop_database] - url: {} res: {}'.format(self.url, res))
        return res

<|endoftext|>


<commit_before>"""Methods for matching the empirical probability distribution of two data sets."""

import numpy as np
from scipy import interpolate as sip
from scipy import optimize as sop

def compute_empirical_cdf(bin_edges, hist):
    """Compute an empirical cumulative distribution function from the given
    histogram.

    Parameters
    ----------
    bin_edges : array_like
        Coordinates of left edges of the histogram bins.
    hist : array_like
        Histogram counts for each bin.

    Returns
    -------
    out : ndarray
        CDF values corresponding to the bin edges.

    """
    cdf = []
    xs = 0.0

    for x,h in zip(zip(bin_edges[:-1], bin_edges[1:]), hist):
        cdf.append(xs)
        xs += (x[1] - x[0]) * h

    cdf.append(xs)
    cdf = np.array(cdf) / xs

    return cdf

def nonparam_match_empirical_cdf(R, R_trg):
    """Matches the empirical CDF of the initial array with the empirical CDF
    of a target array. Initial ranks are conserved, but empirical distribution
    matches the target one. Zero-pixels (i.e. pixels having the minimum value)
    in the initial array are conserved.

    Parameters
    ----------
    R : array_like
        The initial array whose CDF is to be matched with the target.
    R_trg : array_like
        The target array.

    Returns
    -------
    out : array_like
        The matched array.

    """
    if R.size != R_trg.size:
        raise ValueError("the input arrays must have the same size")
    if np.any(~np.isfinite(R)):
        raise ValueError("initial array contains non-finite values")
    if np.any(~np.isfinite(R_trg)):
        raise ValueError("target array contains non-finite values")

    # zeros in initial image
    zvalue = R.min()
    idxzeros = R == zvalue

    # zeros in target image
    zvalue_trg = R_trg.min()

    if np.sum(R_trg > zvalue_trg) > np.sum(R > zvalue):
        # adjust the fraction of rain in target distribution if the number of zeros
        # is greater than in the initial array
        # TODO: this needs more testing
        war = np.sum(R > zvalue)/R.size
        p = np.percentile(R_trg, 100*(1 - war))
        R_trg = R_trg.copy()
        R_trg[R_trg < p] = zvalue_trg

    # flatten the arrays
    arrayshape = R.shape
    R_trg = R_trg.flatten()
    R = R.flatten()

    # rank target values
    order = R_trg.argsort()
    ranked = R_trg[order]

    # rank initial values order
    orderin = R.argsort()
    ranks = np.empty(len(R), int)
    ranks[orderin] = np.arange(len(R))

    # get ranked values from target and rearrange with inital order
    R = ranked[ranks]

    # reshape as original array
    R = R.reshape(arrayshape)

    # readding original zeros
    R[idxzeros] = zvalue_trg

    return R

# TODO: A more detailed explanation of the PMM method + references.
def pmm_init(bin_edges_1, cdf_1, bin_edges_2, cdf_2):
    """Initialize a probability matching method (PMM) object from binned
    cumulative distribution functions (CDF).

    Parameters
    ----------
    bin_edges_1 : array_like
        Coordinates of the left bin edges of the source cdf.
    cdf_1 : array_like
        Values of the source CDF at the bin edges.
    bin_edges_2 : array_like
        Coordinates of the left bin edges of the target cdf.
    cdf_2 : array_like
        Values of the target CDF at the bin edges.

    """
    pmm = {}

    pmm["bin_edges_1"]      = bin_edges_1.copy()
    pmm["cdf_1"]            = cdf_1.copy()
    pmm["bin_edges_2"]      = bin_edges_2.copy()
    pmm["cdf_2"]            = cdf_2.copy()
    pmm["cdf_interpolator"] = sip.interp1d(bin_edges_1, cdf_1, kind="linear")

    return pmm

def pmm_compute(pmm, x):
    """For a given PMM object and x-coordinate, compute the probability matched
    value (i.e. the x-coordinate for which the target CDF has the same value as
    the source CDF).

    Parameters
    ----------
    pmm : dict
        A PMM object returned by pmm_init.
    x : float
        The coordinate for which to compute the probability matched value.

    """
    mask = np.logical_and(x >= pmm["bin_edges_1"][0], x <= pmm["bin_edges_1"][-1])
    p = pmm["cdf_interpolator"](x[mask])

    result = np.ones(len(mask)) * np.nan
    result[mask] = _invfunc(p, pmm["bin_edges_2"], pmm["cdf_2"])

    return result

def shift_scale(R, f, rain_fraction_trg, second_moment_trg, **kwargs):
    """Find shift and scale that is needed to return the required second_moment
    and rain area. The optimization is performed with the Nelder-Mead algorithm
    available in scipy.
    It ssumes a forward transformation ln_rain = ln(rain)-ln(min_rain) if
    rain > min_rain, else 0.

    Parameters
    ----------
    R : array_like
        The initial array to be shift and scaled.
    f : function
        The inverse transformation that is applied after the shift and scale.
    rain_fraction_trg : float
        The required rain fraction to be matched by shifting.
    second_moment_trg : float
        The required second moment to be matched by scaling.
        The second_moment is defined as second_moment = var + mean^2.

    Other Parameters
    ----------------
    scale : float
        Optional initial value of the scale parameter for the Nelder-Mead optimisation.
        Typically, this would be the scale parameter estimated the previous time step.
        Default : 1.
    max_iterations : int
        Maximum allowed number of iterations and function evaluations.
        More details: https://docs.scipy.org/doc/scipy/reference/optimize.minimize-neldermead.html
        Deafult: 100.
    tol : float
        Tolerance for termination.
        More details: https://docs.scipy.org/doc/scipy/reference/optimize.minimize-neldermead.html
        Default: 0.05*second_moment_trg, i.e. terminate the search if the error
        is less than 5% since the second moment is a bit unstable.

    Returns
    -------
    shift : float
        The shift value that produces the required rain fraction.
    scale : float
        The scale value that produces the required second_moment.
    R : array_like
        The shifted, scaled and back-transformed array.
    """

    shape = R.shape
    R = R.flatten()

    # defaults
    scale = kwargs.get("scale", 1.)
    max_iterations = kwargs.get("max_iterations", 100)
    tol = kwargs.get("tol", 0.05*second_moment_trg)

    # calculate the shift parameter based on the required rain fraction
    shift = np.percentile(R, 100*(1 - rain_fraction_trg))
    idx_wet = R > shift

    # define objective function
    def _get_error(scale):
        R_ = np.zeros_like(R)
        R_[idx_wet]  = f((R[idx_wet] - shift)*scale)
        R_[~idx_wet] = 0
        second_moment = np.nanstd(R_)**2 + np.nanmean(R_)**2
        return np.abs(second_moment - second_moment_trg)

    # Nelder-Mead optimisation
    nm_scale = sop.minimize(_get_error, scale, method="Nelder-Mead", tol=tol,
                options={"disp":False,"maxiter":max_iterations})
    scale = nm_scale["x"][0]

    R[idx_wet]  = f((R[idx_wet] - shift)*scale)
    R[~idx_wet] = 0

    return shift, scale, R.reshape(shape)

def _invfunc(y, fx, fy):
  if len(y) == 0:
      return np.array([])

  b = np.digitize(y, fy)
  mask = np.logical_and(b > 0, b < len(fy))
  c = (y[mask] - fy[b[mask]-1]) / (fy[b[mask]] - fy[b[mask]-1])

  result = np.ones(len(y)) * np.nan
  result[mask] = c * fx[b[mask]] + (1.0-c) * fx[b[mask]-1]

  return result
<commit_msg>Corrections to comment lines<commit_after>"""Methods for matching the empirical probability distribution of two data sets."""

import numpy as np
from scipy import interpolate as sip
from scipy import optimize as sop

def compute_empirical_cdf(bin_edges, hist):
    """Compute an empirical cumulative distribution function from the given
    histogram.

    Parameters
    ----------
    bin_edges : array_like
        Coordinates of left edges of the histogram bins.
    hist : array_like
        Histogram counts for each bin.

    Returns
    -------
    out : ndarray
        CDF values corresponding to the bin edges.

    """
    cdf = []
    xs = 0.0

    for x,h in zip(zip(bin_edges[:-1], bin_edges[1:]), hist):
        cdf.append(xs)
        xs += (x[1] - x[0]) * h

    cdf.append(xs)
    cdf = np.array(cdf) / xs

    return cdf

def nonparam_match_empirical_cdf(R, R_trg):
    """Matches the empirical CDF of the initial array with the empirical CDF
    of a target array. Initial ranks are conserved, but empirical distribution
    matches the target one. Zero-pixels (i.e. pixels having the minimum value)
    in the initial array are conserved.

    Parameters
    ----------
    R : array_like
        The initial array whose CDF is to be matched with the target.
    R_trg : array_like
        The target array.

    Returns
    -------
    out : array_like
        The matched array.

    """
    if R.size != R_trg.size:
        raise ValueError("the input arrays must have the same size")
    if np.any(~np.isfinite(R)):
        raise ValueError("initial array contains non-finite values")
    if np.any(~np.isfinite(R_trg)):
        raise ValueError("target array contains non-finite values")

    # zeros in initial array
    zvalue = R.min()
    idxzeros = R == zvalue

    # zeros in the target array
    zvalue_trg = R_trg.min()

    # adjust the fraction of rain in target distribution if the number of zeros
    # is greater than in the initial array
    if np.sum(R_trg > zvalue_trg) > np.sum(R > zvalue):
        war = np.sum(R > zvalue)/R.size
        p = np.percentile(R_trg, 100*(1 - war))
        R_trg = R_trg.copy()
        R_trg[R_trg < p] = zvalue_trg

    # flatten the arrays
    arrayshape = R.shape
    R_trg = R_trg.flatten()
    R = R.flatten()

    # rank target values
    order = R_trg.argsort()
    ranked = R_trg[order]

    # rank initial values order
    orderin = R.argsort()
    ranks = np.empty(len(R), int)
    ranks[orderin] = np.arange(len(R))

    # get ranked values from target and rearrange with the initial order
    R = ranked[ranks]

    # reshape to the original array dimensions
    R = R.reshape(arrayshape)

    # readd original zeros
    R[idxzeros] = zvalue_trg

    return R

# TODO: A more detailed explanation of the PMM method + references.
def pmm_init(bin_edges_1, cdf_1, bin_edges_2, cdf_2):
    """Initialize a probability matching method (PMM) object from binned
    cumulative distribution functions (CDF).

    Parameters
    ----------
    bin_edges_1 : array_like
        Coordinates of the left bin edges of the source cdf.
    cdf_1 : array_like
        Values of the source CDF at the bin edges.
    bin_edges_2 : array_like
        Coordinates of the left bin edges of the target cdf.
    cdf_2 : array_like
        Values of the target CDF at the bin edges.

    """
    pmm = {}

    pmm["bin_edges_1"]      = bin_edges_1.copy()
    pmm["cdf_1"]            = cdf_1.copy()
    pmm["bin_edges_2"]      = bin_edges_2.copy()
    pmm["cdf_2"]            = cdf_2.copy()
    pmm["cdf_interpolator"] = sip.interp1d(bin_edges_1, cdf_1, kind="linear")

    return pmm

def pmm_compute(pmm, x):
    """For a given PMM object and x-coordinate, compute the probability matched
    value (i.e. the x-coordinate for which the target CDF has the same value as
    the source CDF).

    Parameters
    ----------
    pmm : dict
        A PMM object returned by pmm_init.
    x : float
        The coordinate for which to compute the probability matched value.

    """
    mask = np.logical_and(x >= pmm["bin_edges_1"][0], x <= pmm["bin_edges_1"][-1])
    p = pmm["cdf_interpolator"](x[mask])

    result = np.ones(len(mask)) * np.nan
    result[mask] = _invfunc(p, pmm["bin_edges_2"], pmm["cdf_2"])

    return result

def shift_scale(R, f, rain_fraction_trg, second_moment_trg, **kwargs):
    """Find shift and scale that is needed to return the required second_moment
    and rain area. The optimization is performed with the Nelder-Mead algorithm
    available in scipy.
    It ssumes a forward transformation ln_rain = ln(rain)-ln(min_rain) if
    rain > min_rain, else 0.

    Parameters
    ----------
    R : array_like
        The initial array to be shift and scaled.
    f : function
        The inverse transformation that is applied after the shift and scale.
    rain_fraction_trg : float
        The required rain fraction to be matched by shifting.
    second_moment_trg : float
        The required second moment to be matched by scaling.
        The second_moment is defined as second_moment = var + mean^2.

    Other Parameters
    ----------------
    scale : float
        Optional initial value of the scale parameter for the Nelder-Mead optimisation.
        Typically, this would be the scale parameter estimated the previous time step.
        Default : 1.
    max_iterations : int
        Maximum allowed number of iterations and function evaluations.
        More details: https://docs.scipy.org/doc/scipy/reference/optimize.minimize-neldermead.html
        Deafult: 100.
    tol : float
        Tolerance for termination.
        More details: https://docs.scipy.org/doc/scipy/reference/optimize.minimize-neldermead.html
        Default: 0.05*second_moment_trg, i.e. terminate the search if the error
        is less than 5% since the second moment is a bit unstable.

    Returns
    -------
    shift : float
        The shift value that produces the required rain fraction.
    scale : float
        The scale value that produces the required second_moment.
    R : array_like
        The shifted, scaled and back-transformed array.
    """

    shape = R.shape
    R = R.flatten()

    # defaults
    scale = kwargs.get("scale", 1.)
    max_iterations = kwargs.get("max_iterations", 100)
    tol = kwargs.get("tol", 0.05*second_moment_trg)

    # calculate the shift parameter based on the required rain fraction
    shift = np.percentile(R, 100*(1 - rain_fraction_trg))
    idx_wet = R > shift

    # define objective function
    def _get_error(scale):
        R_ = np.zeros_like(R)
        R_[idx_wet]  = f((R[idx_wet] - shift)*scale)
        R_[~idx_wet] = 0
        second_moment = np.nanstd(R_)**2 + np.nanmean(R_)**2
        return np.abs(second_moment - second_moment_trg)

    # Nelder-Mead optimisation
    nm_scale = sop.minimize(_get_error, scale, method="Nelder-Mead", tol=tol,
                options={"disp":False,"maxiter":max_iterations})
    scale = nm_scale["x"][0]

    R[idx_wet]  = f((R[idx_wet] - shift)*scale)
    R[~idx_wet] = 0

    return shift, scale, R.reshape(shape)

def _invfunc(y, fx, fy):
  if len(y) == 0:
      return np.array([])

  b = np.digitize(y, fy)
  mask = np.logical_and(b > 0, b < len(fy))
  c = (y[mask] - fy[b[mask]-1]) / (fy[b[mask]] - fy[b[mask]-1])

  result = np.ones(len(y)) * np.nan
  result[mask] = c * fx[b[mask]] + (1.0-c) * fx[b[mask]-1]

  return result
<|endoftext|>