<commit_before>"""
  Cloud connected autonomous RC car.

  Copyright 2016 Visible Energy Inc. All Rights Reserved.
"""
__license__ = """
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import time
import subprocess as sp
import numpy as np
import cv2

# filename to fetch telemetry from -- updated atomically by the controller loop at 30 Hz
TELEMFNAME = '/tmpfs/meta.txt'

# filename where to store the last frame -- used by the application loop and as parameter for the CNN prediction
FRAMEFNAME = '/tmpfs/frame.yuv'

log=None
config=None
camera=None

# Vehicle object instantiated in the application module
car=None

streaming=False

def init(conf, logger):
  global config, log, camera

  config=conf
  log=logger

  videodevs = ["/dev/" + x for x in os.listdir("/dev/") if x.startswith("video") ]
  if len(videodevs) == 0:
    log("Fatal error. Cannot proceed without cameras connected.")
    return None

  print videodevs
  # create an empty telemetry file
  s = 'echo > ' + TELEMFNAME
  sp.check_call(s, shell=True)

  # use the first camera
  camera=videodevs[0]

  # set resolution and encoding (Logitech C920)
  s = 'v4l2-ctl --device=' + camera + ' --set-fmt-video=width=320,height=240,pixelformat=1'
  log("Setting camera: %s" % s)
  # execute and wait for completion
  sp.check_call(s, shell=True)

  return camera

def video_stop():
  """ Stop running video capture streamer """

  streaming = True
  return

  try:
    pname = config['video']['streamer']
  except:
    log("Error cannot stop the video streamer. Streamer not defined in config.json")
    return None

  s = 'killall ' + pname
  FNULL = open(os.devnull, 'w')
  try:
    # execute and wait for completion
    sp.check_call(s, shell=True, stderr=FNULL)
  except Exception, e:
    # fails when no ffmpeg is running
    if config['app_params']['verbose']:
      log("Error stopping streamer. %s" % e)
    else:
      pass
  return

def __video_start(telem):
  """ Start a video streamer """
  global config, log, camera

  # insure no streamer is running
  video_stop()
  time.sleep(1)

  try:
    pname = config['video']['streamer']
  except:
    log("Error cannot start the video streamer. Streamer not defined in config.json")
    return None

  # set video codec depending on platform
  vcodec = 'h264_omx' if 'raspberrypi' in os.uname() else 'h264'
  # to suppress output when running the streamer
  FNULL = open(os.devnull, 'w')
  if telem:
    # streaming video with embedded telemetry
    params = [pname, '-r','30', '-use_wallclock_as_timestamps', '1', '-thread_queue_size', '512', '-f', 'v4l2', '-i', camera, '-c:v ', vcodec, '-maxrate', '768k', '-bufsize', '960k']
    format = 'format=yuv444p,drawbox=y=ih-h:color=black@0.9:width=40:height=12:t=max,drawtext=fontfile=OpenSans-Regular.ttf:textfile=' + TELEMFNAME + ':reload=1:fontsize=10:fontcolor=white:x=0:y=(h-th-2),format=yuv420p'
    url = 'rtmp://' + config['video']['server'] + ':' + config['video']['port'] + '/src/' + config['video']['key']
    params = params + ['-vf', format, '-threads', '4', '-r', '30', '-g', '60', '-f', 'flv', url]
    # spawn a process and do not wait
    pid = sp.Popen(params, stderr=FNULL)
  else:
    # streaming video and saving the last frame for CNN prediction
    params = [pname, '-r','30', '-use_wallclock_as_timestamps', '1', '-thread_queue_size', '512', '-f', 'v4l2', '-i', camera, '-c:v ', vcodec, '-maxrate', '768k', '-bufsize', '960k']
    url = 'rtmp://' + config['video']['server'] + ':' + config['video']['port'] + '/src/' + config['video']['key']
    params = params + ['-threads', '4', '-r', '30', '-g', '60', '-f', 'flv', url]
    params = params + ['-vcodec', 'rawvideo', '-an', '-updatefirst', '1', '-y', '-f', 'image2', FRAMEFNAME]
    # to transcode use format YUYV422:
    # $ ffmpeg -vcodec rawvideo -s 320x240 -r 1 -pix_fmt  yuyv422  -i frame.yuv rawframe.jpg
    # spawn a process and do not wait
    pid = sp.Popen(params, stderr=FNULL)
  return pid

def video_start(telem):

  try:
    pname = config['video']['streamer']
  except:
    log("Error cannot start the video streamer. Streamer not defined in config.json")
    return None

  i_command = [ pname,
            '-r', '30',
            '-use_wallclock_as_timestamps', '1',
            '-f', 'v4l2',
            '-i', '/dev/video0',
            '-vb','1000k',
            '-f', 'image2pipe',
            '-pix_fmt', 'yuyv422',
            '-vcodec', 'rawvideo', '-']
  i_pipe = sp.Popen(i_command, stdout = sp.PIPE, bufsize=10**5)

  url = 'rtmp://' + config['video']['server'] + ':' + config['video']['port'] + '/src/' + config['video']['key']

  o_command = [ pname,
        '-f', 'rawvideo',
        '-vcodec','rawvideo',
        '-s', '320x240', # size of one frame
#         '-pix_fmt', 'rgb24',
        '-pix_fmt', 'rgb24', #'yuyv422', rgb24
        '-r', '30', # frames per second
        '-i', 'pipe:0', # The imput comes from a pipe
        '-an', # Tells FFMPEG not to expect any audio
        '-c:v','libx264',
        '-profile:v','main',
        '-preset','ultrafast',
        '-pix_fmt', 'yuv420p',
        '-b:v','1000k',
        '-bufsize','2000k',
        '-g','30',
        '-f','flv',
        url ]
  o_pipe = sp.Popen(o_command, stdin=sp.PIPE, stderr=sp.PIPE)

  width = 320 # 640
  height = 240 #480

  rows = height
  cols = width

  image_size = rows * cols * 2 # *3

  streaming = True
  while streaming:
    raw_image = i_pipe.stdout.read(image_size)

    if telem:
      msg = car.telemetry()
      car.com.send_data(json.dumps(msg))

    f = np.fromstring(raw_image, dtype=np.uint8)
    i_pipe.stdout.flush()

    img = f.reshape(rows, cols, 2)  # 2)

    # convert to RGB
    rgb_img =  cv2.cvtColor(img, cv2.COLOR_YUV2RGB_YUY2)  # working w camera format yuyv422

 #   rgb_img = img
 #   rgb_img =  cv2.cvtColor(img, cv2.COLOR_YUV420p2RGB)

    # draw a center rectangle
    cv2.rectangle(rgb_img,(130,100),(190,140),(255,0,0),2)
  #  M = cv2.getRotationMatrix2D((width/2,height/2),180,1)

    # rotate the image 90 degrees twice and bring back to normal
  #  dst = cv2.warpAffine(rgb_img,M,(width,height))

    rgb_img[0,0,0] = car.steering
    rgb_img[0,0,1] = car.throttle
    # output the image
    o_pipe.stdin.write(rgb_img.tostring())

<commit_msg>testing with embedding telemetry<commit_after>"""
  Cloud connected autonomous RC car.

  Copyright 2016 Visible Energy Inc. All Rights Reserved.
"""
__license__ = """
Licensed under the Apache License, Version 2.0 (the "License");
you may not use this file except in compliance with the License.
You may obtain a copy of the License at
    http://www.apache.org/licenses/LICENSE-2.0
Unless required by applicable law or agreed to in writing, software
distributed under the License is distributed on an "AS IS" BASIS,
WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
See the License for the specific language governing permissions and
limitations under the License.
"""

import os
import time
import subprocess as sp
import numpy as np
import cv2

# filename to fetch telemetry from -- updated atomically by the controller loop at 30 Hz
TELEMFNAME = '/tmpfs/meta.txt'

# filename where to store the last frame -- used by the application loop and as parameter for the CNN prediction
FRAMEFNAME = '/tmpfs/frame.yuv'

log=None
config=None
camera=None

# Vehicle object instantiated in the application module
car=None

streaming=False

def init(conf, logger):
  global config, log, camera

  config=conf
  log=logger

  videodevs = ["/dev/" + x for x in os.listdir("/dev/") if x.startswith("video") ]
  if len(videodevs) == 0:
    log("Fatal error. Cannot proceed without cameras connected.")
    return None

  print videodevs
  # create an empty telemetry file
  s = 'echo > ' + TELEMFNAME
  sp.check_call(s, shell=True)

  # use the first camera
  camera=videodevs[0]

  # set resolution and encoding (Logitech C920)
  s = 'v4l2-ctl --device=' + camera + ' --set-fmt-video=width=320,height=240,pixelformat=1'
  log("Setting camera: %s" % s)
  # execute and wait for completion
  sp.check_call(s, shell=True)

  return camera

def video_stop():
  """ Stop running video capture streamer """

  streaming = True
  return

  try:
    pname = config['video']['streamer']
  except:
    log("Error cannot stop the video streamer. Streamer not defined in config.json")
    return None

  s = 'killall ' + pname
  FNULL = open(os.devnull, 'w')
  try:
    # execute and wait for completion
    sp.check_call(s, shell=True, stderr=FNULL)
  except Exception, e:
    # fails when no ffmpeg is running
    if config['app_params']['verbose']:
      log("Error stopping streamer. %s" % e)
    else:
      pass
  return

def __video_start(telem):
  """ Start a video streamer """
  global config, log, camera

  # insure no streamer is running
  video_stop()
  time.sleep(1)

  try:
    pname = config['video']['streamer']
  except:
    log("Error cannot start the video streamer. Streamer not defined in config.json")
    return None

  # set video codec depending on platform
  vcodec = 'h264_omx' if 'raspberrypi' in os.uname() else 'h264'
  # to suppress output when running the streamer
  FNULL = open(os.devnull, 'w')
  if telem:
    # streaming video with embedded telemetry
    params = [pname, '-r','30', '-use_wallclock_as_timestamps', '1', '-thread_queue_size', '512', '-f', 'v4l2', '-i', camera, '-c:v ', vcodec, '-maxrate', '768k', '-bufsize', '960k']
    format = 'format=yuv444p,drawbox=y=ih-h:color=black@0.9:width=40:height=12:t=max,drawtext=fontfile=OpenSans-Regular.ttf:textfile=' + TELEMFNAME + ':reload=1:fontsize=10:fontcolor=white:x=0:y=(h-th-2),format=yuv420p'
    url = 'rtmp://' + config['video']['server'] + ':' + config['video']['port'] + '/src/' + config['video']['key']
    params = params + ['-vf', format, '-threads', '4', '-r', '30', '-g', '60', '-f', 'flv', url]
    # spawn a process and do not wait
    pid = sp.Popen(params, stderr=FNULL)
  else:
    # streaming video and saving the last frame for CNN prediction
    params = [pname, '-r','30', '-use_wallclock_as_timestamps', '1', '-thread_queue_size', '512', '-f', 'v4l2', '-i', camera, '-c:v ', vcodec, '-maxrate', '768k', '-bufsize', '960k']
    url = 'rtmp://' + config['video']['server'] + ':' + config['video']['port'] + '/src/' + config['video']['key']
    params = params + ['-threads', '4', '-r', '30', '-g', '60', '-f', 'flv', url]
    params = params + ['-vcodec', 'rawvideo', '-an', '-updatefirst', '1', '-y', '-f', 'image2', FRAMEFNAME]
    # to transcode use format YUYV422:
    # $ ffmpeg -vcodec rawvideo -s 320x240 -r 1 -pix_fmt  yuyv422  -i frame.yuv rawframe.jpg
    # spawn a process and do not wait
    pid = sp.Popen(params, stderr=FNULL)
  return pid

def video_start(telem):

  try:
    pname = config['video']['streamer']
  except:
    log("Error cannot start the video streamer. Streamer not defined in config.json")
    return None

  i_command = [ pname,
            '-r', '30',
            '-use_wallclock_as_timestamps', '1',
            '-f', 'v4l2',
            '-i', '/dev/video0',
            '-vb','1000k',
            '-f', 'image2pipe',
            '-pix_fmt', 'yuyv422',
            '-vcodec', 'rawvideo', '-']
  i_pipe = sp.Popen(i_command, stdout = sp.PIPE, bufsize=10**5)

  url = 'rtmp://' + config['video']['server'] + ':' + config['video']['port'] + '/src/' + config['video']['key']

  o_command = [ pname,
        '-f', 'rawvideo',
        '-vcodec','rawvideo',
        '-s', '320x240', # size of one frame
#         '-pix_fmt', 'rgb24',
        '-pix_fmt', 'yuyv422', #'yuyv422', rgb24
        '-r', '30', # frames per second
        '-i', 'pipe:0', # The imput comes from a pipe
        '-an', # Tells FFMPEG not to expect any audio
        '-c:v','libx264',
        '-profile:v','main',
        '-preset','ultrafast',
        '-pix_fmt', 'yuv420p',
        '-b:v','1000k',
        '-bufsize','2000k',
        '-g','30',
        '-f','flv',
        url ]
  o_pipe = sp.Popen(o_command, stdin=sp.PIPE, stderr=sp.PIPE)

  width = 320 # 640
  height = 240 #480

  rows = height
  cols = width

  image_size = rows * cols * 2 # *3

  streaming = True
  while streaming:
    raw_image = i_pipe.stdout.read(image_size)

    if telem:
      msg = car.telemetry()
      car.com.send_data(json.dumps(msg))

    f = np.fromstring(raw_image, dtype=np.uint8)
    i_pipe.stdout.flush()

    img = f.reshape(rows, cols, 2)  # 2)

    img[0,0] = car.steering
    img[0,1] = car.throttle
    o_pipe.stdin.write(img.tostring())
    continue

    # convert to RGB
    rgb_img =  cv2.cvtColor(img, cv2.COLOR_YUV2RGB_YUY2)  # working w camera format yuyv422

 #   rgb_img = img
 #   rgb_img =  cv2.cvtColor(img, cv2.COLOR_YUV420p2RGB)

    # draw a center rectangle
    cv2.rectangle(rgb_img,(130,100),(190,140),(255,0,0),2)
  #  M = cv2.getRotationMatrix2D((width/2,height/2),180,1)

    # rotate the image 90 degrees twice and bring back to normal
  #  dst = cv2.warpAffine(rgb_img,M,(width,height))

    rgb_img[0,0,0] = car.steering
    rgb_img[0,0,1] = car.throttle
    # output the image
    o_pipe.stdin.write(rgb_img.tostring())

<|endoftext|>


<commit_before>import numpy as np
import re
import sys

def get_total_peptide(node, segment='ha'):
    '''
    the concatenation of signal peptide, HA1, HA1
    '''
    if segment=='ha':
        return np.fromstring(node.translations['SigPep'] + node.translations['HA1'] + node.translations['HA2'], 'S1')
    elif segment=='na':
        return np.fromstring(node.translations['NA'], 'S1')

def read_masks(mask_file):
    ha_masks = {}
    with open(mask_file) as f:
        for line in f:
            (key, value) = line.strip().split()
            ha_masks[key] = np.fromstring(value, 'S1')=='1'
    return ha_masks

def calculate_metadata_scores(tree):
    """Calculate scores for each node in the given tree based on metadata already
    annotated on each node including age and sex of infected individuals.
    """
    root = tree.root

    for node in tree.get_terminals():
        node.tip_count=1.0
        if ('age' in node.attr) and type(node.attr['age']) in [str, unicode]:
            tmp_age = node.attr['age']
            if tmp_age[-1] in ['y', 'm']:
                node.attr['age'] = int(tmp_age[:-1])
                if tmp_age[-1]=='m':
                    node.attr['age']/=12.0
                elif node.attr['age']>120:
                    node.attr['age'] = 'unknown'
            else:
                node.attr['age'] = 'unknown'
        else:
            node.attr['age'] = 'unknown'
        node.tip_ages = [] if node.attr['age']=="unknown" else [node.attr['age']]

    # assign a list of ages to each internal node (will be used to calculate changes in age distribution)
    for node in tree.get_nonterminals(order='postorder'):
        valid_ages = [c.attr['age'] for c in node if c.attr['age'] is not "unknown"]
        node.tip_ages = sum( [c.tip_ages for c in node] , [])
        valid_tipcounts = [c.tip_count for c in node if c.attr['age'] is not "unknown"]
        node.tip_count = np.sum(valid_tipcounts)
        if len(valid_ages):
            node.attr['age'] = np.sum([a*t for a,t in zip(valid_ages, valid_tipcounts)])/np.sum(valid_tipcounts)
        else:
            node.attr['age'] = "unknown"
    root.outgroup_ages = []

    # assign avg_age. This is redundant in that deep nodes have average ages already
    # but terminal nodes or small clades don't have a smooth average. this assigns those
    # to the parent.
    # outgroup_ages: these are ages of all nodes NOT within the clade.
    root.attr['avg_age'] = root.attr['age']
    for node in tree.get_nonterminals(order='preorder'):
        for c1 in node:
            c1.outgroup_ages = node.outgroup_ages + sum([c2.tip_ages for c2 in node if c1!=c2], [])
            if c1.tip_count<20:
                c1.attr['avg_age'] = node.attr['avg_age']
            else:
                c1.attr['avg_age'] = c1.attr['age']

    # calculate the a score to spot changing age distributions (for what ever reason, i.e. geography)
    # these are calculated as a 2-sample KS test based on ages within and outside the clade in question
    from scipy import stats
    for node in tree.find_clades(order='preorder'):
        if node==root:
            node.attr['age_score'] = 0
            continue
        if len(node.tip_ages)>20 and len(node.outgroup_ages)>20:
            ks = stats.ks_2samp(node.outgroup_ages, node.tip_ages)
            node.attr['age_score'] = -np.log10(ks.pvalue)
        else:
            node.attr['age_score'] = node.up.attr['age_score']

        if np.isnan(node.attr['age_score']):
            node.attr['age_score']=0.0

    # gender seems not to be a relevant quantity as of now...
    for node in tree.get_terminals():
        node.tip_count=1.0
        if ('gender' in node.attr) and type(node.attr['gender']) in [str, unicode]:
            g = node.attr["gender"]
            node.attr["num_gender"] = -1 if g=='male' else (1 if g=='female' else 0)
        else:
            node.attr["num_gender"] = 0

    for node in tree.get_nonterminals(order='postorder'):
        node.tip_count = np.sum([c.tip_count for c in node])
        node.attr['num_gender'] = np.sum([c.attr['num_gender']*c.tip_count for c in node])/node.tip_count

def mask_sites(aa, mask):
    return aa[mask[:len(aa)]]

def mask_distance(aaA, aaB, mask):
    """Return distance of sequences aaA and aaB by comparing sites in the given binary mask."""
    sites_A = mask_sites(aaA, mask)
    sites_B = mask_sites(aaB, mask)
    distance = np.sum(sites_A != sites_B)
    return distance

def glycosylation_count(total_aa_seq, glyc_mask):
    # TODO: need to restrict to surface residues.
    total_aa_seq_masked = "".join([aa if mask else 'X'
                                   for (mask, aa) in zip(glyc_mask, total_aa_seq)])

    return len(re.findall('N[^P][ST][^P]', total_aa_seq_masked))

def calculate_sequence_scores(tree, mask_file, lineage, segment, epitope_mask_version, glyc_mask_version):
    """Calculate scores from the amino acid sequence of each node in the given tree.

    Sequence scores depend on lineage- and segment-specific amino acid site masks or named masks.
    """
    ha_masks = read_masks(mask_file)
    try:
        epitope_mask = ha_masks[epitope_mask_version]
    except KeyError as error:
        sys.stderr.write("ERROR: Could not find an epitope mask named '%s'.\n" % epitope_mask_version)
        raise error

    # Get amino acid sequence of root node.
    root = tree.root
    root_total_aa_seq = get_total_peptide(root, segment)

    # Confirm that mask and sequence lengths are equal.
    assert len(root_total_aa_seq) == len(epitope_mask), "Sequence and mask lengths are not equal: %s != %s." % (len(root_total_aa_seq), len(epitope_mask))

    # Setup receptor binding site mask.
    rbs_mask_name = "%s_%s_rbs" % (lineage, segment)
    rbs_mask = ha_masks.get(rbs_mask_name)

    # Setup glycosylation mask.
    if glyc_mask_version in ha_masks:
        glyc_mask = ha_masks[glyc_mask_version]
    else:
        sys.stderr.write("WARNING: Could not find a glycosylation mask named '%s'; using all positions instead.\n" % glyc_mask_version)
        glyc_mask = np.ones(len(root_total_aa_seq), dtype='bool')

    # Annotate scores to each node based on its amino acid sequence.
    for node in tree.find_clades():
        total_aa_seq = get_total_peptide(node, segment)
        node.attr['ep'] = mask_distance(total_aa_seq, root_total_aa_seq, epitope_mask)
        node.attr['ne'] = mask_distance(total_aa_seq, root_total_aa_seq, ~epitope_mask)
        node.attr['glyc'] = glycosylation_count(total_aa_seq, glyc_mask)

        if rbs_mask is not None:
            node.attr['rb'] = mask_distance(total_aa_seq, root_total_aa_seq, rbs_mask)

def select_nodes_in_season(tree, timepoint, time_window=0.6, **kwargs):
    """Annotate a boolean to each node in the tree if it is alive at the given
    timepoint or prior to the timepoint by the given time window preceding.

    This annotation is used by the LBI and epitope cross-immunity predictors.
    """
    for node in tree.find_clades(order="postorder"):
        if node.is_terminal():
            if node.attr['num_date'] <= timepoint and node.attr['num_date'] > timepoint - time_window:
                node.alive=True
            else:
                node.alive=False
        else:
            node.alive = any(ch.alive for ch in node.clades)

def calculate_LBI(tree, attr="lbi", tau=0.4, transform=lambda x:x, **kwargs):
    '''
    traverses the tree in postorder and preorder to calculate the
    up and downstream tree length exponentially weighted by distance.
    then adds them as LBI
    tree     -- biopython tree for whose node the LBI is being computed
    attr     -- the attribute name used to store the result
    '''
    # Calculate clock length.
    tree.root.clock_length = 0.0
    for node in tree.find_clades():
        for child in node.clades:
            child.clock_length = child.attr['num_date'] - node.attr['num_date']

    # traverse the tree in postorder (children first) to calculate msg to parents
    for node in tree.find_clades(order="postorder"):
        node.down_polarizer = 0
        node.up_polarizer = 0
        for child in node.clades:
            node.up_polarizer += child.up_polarizer
        bl =  node.clock_length / tau
        node.up_polarizer *= np.exp(-bl)
        if node.alive: node.up_polarizer += tau*(1-np.exp(-bl))

    # traverse the tree in preorder (parents first) to calculate msg to children
    for node in tree.get_nonterminals():
        for child1 in node.clades:
            child1.down_polarizer = node.down_polarizer
            for child2 in node.clades:
                if child1!=child2:
                    child1.down_polarizer += child2.up_polarizer

            bl =  child1.clock_length / tau
            child1.down_polarizer *= np.exp(-bl)
            if child1.alive: child1.down_polarizer += tau*(1-np.exp(-bl))

    # go over all nodes and calculate the LBI (can be done in any order)
    max_LBI = 0.0
    for node in tree.find_clades(order="postorder"):
        tmp_LBI = node.down_polarizer
        for child in node.clades:
            tmp_LBI += child.up_polarizer

        node.attr[attr] = transform(tmp_LBI)
        if node.attr[attr] > max_LBI:
            max_LBI = node.attr[attr]

    # Normalize LBI to range [0, 1].
    for node in tree.find_clades():
        node.attr[attr] /= max_LBI
        setattr(node, attr, node.attr[attr])

def calculate_phylogenetic_scores(tree, **kwargs):
    """Calculate scores based on a given phylogenetic tree and assign scores to each node.

    Scores include:

      - LBI
    """
    # Find maximum time point in the given tree.
    timepoint = max(node.attr["num_date"] for node in tree.find_clades())

    # Select nodes within time window for LBI.
    select_nodes_in_season(tree, timepoint, **kwargs)

    # Calculate LBI.
    calculate_LBI(tree, **kwargs)
<commit_msg>Exclude undetermined residues from amino acid mutation counts<commit_after>import numpy as np
import re
import sys

def get_total_peptide(node, segment='ha'):
    '''
    the concatenation of signal peptide, HA1, HA1
    '''
    if segment=='ha':
        return np.fromstring(node.translations['SigPep'] + node.translations['HA1'] + node.translations['HA2'], 'S1')
    elif segment=='na':
        return np.fromstring(node.translations['NA'], 'S1')

def read_masks(mask_file):
    ha_masks = {}
    with open(mask_file) as f:
        for line in f:
            (key, value) = line.strip().split()
            ha_masks[key] = np.fromstring(value, 'S1')=='1'
    return ha_masks

def calculate_metadata_scores(tree):
    """Calculate scores for each node in the given tree based on metadata already
    annotated on each node including age and sex of infected individuals.
    """
    root = tree.root

    for node in tree.get_terminals():
        node.tip_count=1.0
        if ('age' in node.attr) and type(node.attr['age']) in [str, unicode]:
            tmp_age = node.attr['age']
            if tmp_age[-1] in ['y', 'm']:
                node.attr['age'] = int(tmp_age[:-1])
                if tmp_age[-1]=='m':
                    node.attr['age']/=12.0
                elif node.attr['age']>120:
                    node.attr['age'] = 'unknown'
            else:
                node.attr['age'] = 'unknown'
        else:
            node.attr['age'] = 'unknown'
        node.tip_ages = [] if node.attr['age']=="unknown" else [node.attr['age']]

    # assign a list of ages to each internal node (will be used to calculate changes in age distribution)
    for node in tree.get_nonterminals(order='postorder'):
        valid_ages = [c.attr['age'] for c in node if c.attr['age'] is not "unknown"]
        node.tip_ages = sum( [c.tip_ages for c in node] , [])
        valid_tipcounts = [c.tip_count for c in node if c.attr['age'] is not "unknown"]
        node.tip_count = np.sum(valid_tipcounts)
        if len(valid_ages):
            node.attr['age'] = np.sum([a*t for a,t in zip(valid_ages, valid_tipcounts)])/np.sum(valid_tipcounts)
        else:
            node.attr['age'] = "unknown"
    root.outgroup_ages = []

    # assign avg_age. This is redundant in that deep nodes have average ages already
    # but terminal nodes or small clades don't have a smooth average. this assigns those
    # to the parent.
    # outgroup_ages: these are ages of all nodes NOT within the clade.
    root.attr['avg_age'] = root.attr['age']
    for node in tree.get_nonterminals(order='preorder'):
        for c1 in node:
            c1.outgroup_ages = node.outgroup_ages + sum([c2.tip_ages for c2 in node if c1!=c2], [])
            if c1.tip_count<20:
                c1.attr['avg_age'] = node.attr['avg_age']
            else:
                c1.attr['avg_age'] = c1.attr['age']

    # calculate the a score to spot changing age distributions (for what ever reason, i.e. geography)
    # these are calculated as a 2-sample KS test based on ages within and outside the clade in question
    from scipy import stats
    for node in tree.find_clades(order='preorder'):
        if node==root:
            node.attr['age_score'] = 0
            continue
        if len(node.tip_ages)>20 and len(node.outgroup_ages)>20:
            ks = stats.ks_2samp(node.outgroup_ages, node.tip_ages)
            node.attr['age_score'] = -np.log10(ks.pvalue)
        else:
            node.attr['age_score'] = node.up.attr['age_score']

        if np.isnan(node.attr['age_score']):
            node.attr['age_score']=0.0

    # gender seems not to be a relevant quantity as of now...
    for node in tree.get_terminals():
        node.tip_count=1.0
        if ('gender' in node.attr) and type(node.attr['gender']) in [str, unicode]:
            g = node.attr["gender"]
            node.attr["num_gender"] = -1 if g=='male' else (1 if g=='female' else 0)
        else:
            node.attr["num_gender"] = 0

    for node in tree.get_nonterminals(order='postorder'):
        node.tip_count = np.sum([c.tip_count for c in node])
        node.attr['num_gender'] = np.sum([c.attr['num_gender']*c.tip_count for c in node])/node.tip_count

def mask_sites(aa, mask):
    return aa[mask[:len(aa)]]

def mask_distance(aaA, aaB, mask):
    """Return distance of sequences aaA and aaB by comparing sites in the given binary mask.

    >>> aaA = np.array(["A", "B", "C"], dtype="S1")
    >>> aaB = np.array(["A", "B", "D"], dtype="S1")
    >>> mask = np.array([0, 1, 1], dtype=np.bool)
    >>> mask_distance(aaA, aaB, mask)
    1
    >>> aaB = np.array(["A", "B", "X"], dtype="S1")
    >>> mask_distance(aaA, aaB, mask)
    0
    """
    sites_A = mask_sites(aaA, mask)
    sites_B = mask_sites(aaB, mask)

    # Count sites that differ between sequences excluding undetermined residues.
    distance = np.sum((sites_A != sites_B) & (sites_A != "X") & (sites_B != "X"))

    return distance

def glycosylation_count(total_aa_seq, glyc_mask):
    # TODO: need to restrict to surface residues.
    total_aa_seq_masked = "".join([aa if mask else 'X'
                                   for (mask, aa) in zip(glyc_mask, total_aa_seq)])

    return len(re.findall('N[^P][ST][^P]', total_aa_seq_masked))

def calculate_sequence_scores(tree, mask_file, lineage, segment, epitope_mask_version, glyc_mask_version):
    """Calculate scores from the amino acid sequence of each node in the given tree.

    Sequence scores depend on lineage- and segment-specific amino acid site masks or named masks.
    """
    ha_masks = read_masks(mask_file)
    try:
        epitope_mask = ha_masks[epitope_mask_version]
    except KeyError as error:
        sys.stderr.write("ERROR: Could not find an epitope mask named '%s'.\n" % epitope_mask_version)
        raise error

    # Get amino acid sequence of root node.
    root = tree.root
    root_total_aa_seq = get_total_peptide(root, segment)

    # Confirm that mask and sequence lengths are equal.
    assert len(root_total_aa_seq) == len(epitope_mask), "Sequence and mask lengths are not equal: %s != %s." % (len(root_total_aa_seq), len(epitope_mask))

    # Setup receptor binding site mask.
    rbs_mask_name = "%s_%s_rbs" % (lineage, segment)
    rbs_mask = ha_masks.get(rbs_mask_name)

    # Setup glycosylation mask.
    if glyc_mask_version in ha_masks:
        glyc_mask = ha_masks[glyc_mask_version]
    else:
        sys.stderr.write("WARNING: Could not find a glycosylation mask named '%s'; using all positions instead.\n" % glyc_mask_version)
        glyc_mask = np.ones(len(root_total_aa_seq), dtype='bool')

    # Annotate scores to each node based on its amino acid sequence.
    for node in tree.find_clades():
        total_aa_seq = get_total_peptide(node, segment)
        node.attr['ep'] = mask_distance(total_aa_seq, root_total_aa_seq, epitope_mask)
        node.attr['ne'] = mask_distance(total_aa_seq, root_total_aa_seq, ~epitope_mask)
        node.attr['glyc'] = glycosylation_count(total_aa_seq, glyc_mask)

        if rbs_mask is not None:
            node.attr['rb'] = mask_distance(total_aa_seq, root_total_aa_seq, rbs_mask)

def select_nodes_in_season(tree, timepoint, time_window=0.6, **kwargs):
    """Annotate a boolean to each node in the tree if it is alive at the given
    timepoint or prior to the timepoint by the given time window preceding.

    This annotation is used by the LBI and epitope cross-immunity predictors.
    """
    for node in tree.find_clades(order="postorder"):
        if node.is_terminal():
            if node.attr['num_date'] <= timepoint and node.attr['num_date'] > timepoint - time_window:
                node.alive=True
            else:
                node.alive=False
        else:
            node.alive = any(ch.alive for ch in node.clades)

def calculate_LBI(tree, attr="lbi", tau=0.4, transform=lambda x:x, **kwargs):
    '''
    traverses the tree in postorder and preorder to calculate the
    up and downstream tree length exponentially weighted by distance.
    then adds them as LBI
    tree     -- biopython tree for whose node the LBI is being computed
    attr     -- the attribute name used to store the result
    '''
    # Calculate clock length.
    tree.root.clock_length = 0.0
    for node in tree.find_clades():
        for child in node.clades:
            child.clock_length = child.attr['num_date'] - node.attr['num_date']

    # traverse the tree in postorder (children first) to calculate msg to parents
    for node in tree.find_clades(order="postorder"):
        node.down_polarizer = 0
        node.up_polarizer = 0
        for child in node.clades:
            node.up_polarizer += child.up_polarizer
        bl =  node.clock_length / tau
        node.up_polarizer *= np.exp(-bl)
        if node.alive: node.up_polarizer += tau*(1-np.exp(-bl))

    # traverse the tree in preorder (parents first) to calculate msg to children
    for node in tree.get_nonterminals():
        for child1 in node.clades:
            child1.down_polarizer = node.down_polarizer
            for child2 in node.clades:
                if child1!=child2:
                    child1.down_polarizer += child2.up_polarizer

            bl =  child1.clock_length / tau
            child1.down_polarizer *= np.exp(-bl)
            if child1.alive: child1.down_polarizer += tau*(1-np.exp(-bl))

    # go over all nodes and calculate the LBI (can be done in any order)
    max_LBI = 0.0
    for node in tree.find_clades(order="postorder"):
        tmp_LBI = node.down_polarizer
        for child in node.clades:
            tmp_LBI += child.up_polarizer

        node.attr[attr] = transform(tmp_LBI)
        if node.attr[attr] > max_LBI:
            max_LBI = node.attr[attr]

    # Normalize LBI to range [0, 1].
    for node in tree.find_clades():
        node.attr[attr] /= max_LBI
        setattr(node, attr, node.attr[attr])

def calculate_phylogenetic_scores(tree, **kwargs):
    """Calculate scores based on a given phylogenetic tree and assign scores to each node.

    Scores include:

      - LBI
    """
    # Find maximum time point in the given tree.
    timepoint = max(node.attr["num_date"] for node in tree.find_clades())

    # Select nodes within time window for LBI.
    select_nodes_in_season(tree, timepoint, **kwargs)

    # Calculate LBI.
    calculate_LBI(tree, **kwargs)
<|endoftext|>


<commit_before>
from __future__ import print_function
from __future__ import absolute_import

import os
import numpy
import datetime

import storm

import batch.batch

days2seconds = lambda days: days * 60.0**2 * 24.0

class StormJob(batch.batch.job):
    r""""""

    def __init__(self, storm_num, base_path='./', storms_path='./'):

        super(StormJob, self).__init__()

        self.type = ""
        self.name = ""
        self.prefix = str(storm_num).zfill(5)
        self.storm_num = storm_num
        self.executable = "xgeoclaw"

        # Create base data object
        import setrun
        self.rundata = setrun.setrun()

        # Storm specific data
        self.storm_file_path = os.path.join(storms_path
                               os.path.abspath("./%s.storm" % storm_num))

        # Set storm file
        self.rundata.storm_data.storm_file = self.storm_file_path

    def __str__(self):
        output = super(StormJob, self).__str__()
        output += "\n  Storm Number: %s" % self.storm_num
        return output

    def write_data_objects(self):
        r""""""

        # Write out all data files
        super(StormJob, self).write_data_objects()
<commit_msg>Modify storm path"<commit_after>
from __future__ import print_function
from __future__ import absolute_import

import os
import numpy
import datetime

import storm

import batch.batch

days2seconds = lambda days: days * 60.0**2 * 24.0

class StormJob(batch.batch.job):
    r""""""

    def __init__(self, storm_num, base_path='./', storms_path='./'):

        super(StormJob, self).__init__()

        self.type = ""
        self.name = ""
        self.prefix = str(storm_num).zfill(5)
        self.storm_num = storm_num
        self.executable = "xgeoclaw"

        # Create base data object
        import setrun
        self.rundata = setrun.setrun()

        # Storm specific data
        self.storm_file_path = os.path.abspath(os.path.join(storms_path,
                                                        "%s.storm" % storm_num))

        # Set storm file
        self.rundata.storm_data.storm_file = self.storm_file_path

    def __str__(self):
        output = super(StormJob, self).__str__()
        output += "\n  Storm Number: %s" % self.storm_num
        return output

    def write_data_objects(self):
        r""""""

        # Write out all data files
        super(StormJob, self).write_data_objects()
<|endoftext|>


<commit_before>__version__ = '0.1.40'
__version__date__ = '2017/12/01'

# histórico
#
# - ?.?.? - 20??/??/??
#
# - 0.1.40 - 2017/12/01
#   . Pedido: grade de itens pedidos
#
# - 0.1.39 - 2017/11/30
#   . App RH
#   . Impressão de cartelas de lote melhorada
#   . Toda rotina de bipagem para inventário de rolos e geração de arquivo
#     de inventário para o Systêxtil
#
# - 0.1.38 - 2017/11/17
#   . Alteração de layout e algumas ULRs
#   . Adição de link de base de conhecimento
#
# - 0.1.37 - 2017/11/13
#   . Abrindo Systêxtil sem ajuda de página externa (funcionando também de fora
#     da Tussor)
#
# - 0.1.36 - 2017/11/09
#   . OP: Prevendo OP com mais de uma finha na informações básicas
#   . OP e OS: acertando quantidades apresentadas - considerando perda e etc.
#
# - 0.1.35 - 2017/11/08
#   . Pedido: Situação da venda
#   . OP e OS: Link para pedido
#   . Pedido: OPs
#
# - 0.1.34 - 2017/11/08
#   . OP pendente: filtro por coleção
#
# - 0.1.33 - 2017/11/07
#   . Conferência de pedido: melhorias
#   . Pedido: Inicio de desenvolvimento
#
# - 0.1.32 - 2017/11/06
#   . OP + número de pedido e pedido de cliente
#   . renomeando páginas da intranet
#   . Melhorias em várias telas
#
# - 0.1.31 - 2017/10/18
#   . Ordens pendentes por estágio
#   . Melhorias em várias telas
#
# - 0.1.30 - 2017/09/29
#   . Primeira rotina de impressão térmica
#     . Impressão de etiqueta de lote
#
# - 0.1.29 - 2017/09/26
#   . Produto: links para insumos
#   . Insumo: adicionada informação de "Usado em"
#     . links para produtos
#
# - 0.1.28 - 2017/09/26
#   . Layout: alterações
#   . Insumo: Início da APP
#
# - 0.1.26 - 2017/09/12
#   . OP: depósitos destino da produção
#
# - 0.1.25 - 2017/09/12
#   . gravando log de alterações na tabela RecordTracking
#
# - 0.1.24 - 2017/09/11
#   . criando log de alterações em registros de tabelas
#   . Criada tabela RecordTracking para guardar log de alterações
#   . OP: adicionando quat. de 2a. qualidade e perda ao resumo por estágio
#   . App Utils para Classes, Middlewares e outras coisas uteis para
#     vários Apps
#
# - 0.1.23 - 2017/08/30
#   . adição de varios links para referência
#   . criação de view de modelo
#
# - 0.1.22 - 2017/08/29
#   . ajustes de layout (finalizado)
#
# - 0.1.21 - 2017/08/28
#   . view de referências - continuando
#   . ajustes de layout (iniciando por produto / ref e contagem)
#
# - 0.1.20 - 2017/08/25
#   . view de referências (produtos confeccionados)
#   . Layout:
#     . sem bordas
#     . eliminando footer
#     . não imprime href link
#
# - 0.1.19 - 2017/08/08
#   . refactoring de lotes.models
#   . análise de produzido
#     . qtd por período e alternativa
#     . qtd por data de corte e alternativa
#
# - 0.1.18 - 2017/08/04
#   . OP: datas
#
# - 0.1.17 - 2017/08/04
#   . OS: datas, observação, insumos
#
# - 0.1.16 - 2017/08/03
#   . estatística de produto: prevendo PG
#   . OS: retornos, datas
#
# - 0.1.15 - 2017/08/01
#   . OS: lotes
#   . Posição de lote: OSs
#   . OP: Mais informações
#
# - 0.1.14 - 2017/07/31
#   . Posição de lote: Mais informações
#   . OP: Mais informações
#   . OS: Iniciando busca de informações de OS
#
# - 0.1.13 - 2017/07/20
#   . ficha de cliente: pequena correção e melhoria
#   . visualização de OP:
#     . links para lotes
#     . situação de OP
#     . OP relacionada
#   . posição de lote:
#     . GET
#     . apresentando como tipo (MD, PG, PA), ref., cor e tam.
#     . link para OP
#
# - 0.1.12 - 2017/07/12
#   . controle de data de saída:
#     . transportadora
#     . ajustes de tela
#
# - 0.1.11 - 2017/07/11
#   . melhorias na consulta de OP
#   . adicionados campos à tela de controle de data de saída
#
# - 0.1.10 - 2017/07/10
#   . Evoluindo controle de data de saída
#   . Evoluindo dados de pedido para conferencia
#     Era só infAdProd, adicionei EAN e Narrativa
#
# - 0.1.9 - 2017/07/07
#   . controle de data de saída: Mais informações
#
# - 0.1.8 - 2017/07/06
#   . iniciando app logística
#   . controle de data de saída
#
# - 0.1.7 - 2017/07/05
#   . ficha de cliente (consultando F1)
#
# - 0.1.6 - 2017/06/30
#   . consulta infAdProd por número de pedido
#
# - 0.1.5 - 2017/06/29
#   . corrigido bug em view de OP
#   . posição de lote reescrita - vira classe
#   . responsáveis aparece estágio sem responsável
#   . autofocus nos forms
#
# - 0.1.4 - 2017/06/23
#   . menus reestruturados
#   . posição de lote reescrita - tira ajax e utiliza models
#
# - 0.1.3 - 2017/06/16
#   . consulta de OP: lotes por estágios, por OS e individuais
#   . lista produtos nível 1 sem preço médio
#
# - 0.1.2 - 2017/06/14
#   . consulta de responsável por estágio
#   . mostra cadastro de estágios
#   . em Produtos, mostra tamanhos e cores
#   . em Produtos, nível 1, mostra alternativas e roteiros
#
# - 0.1.1 - 2017/06/12
#   . melhorias na visualização das informações do lote
#   . mostra cadastro de depósitos
#
# - 0.1.0 - 2017/06/07
#   . primeira versão divulgada
#   . estatística de produtos
<commit_msg>version<commit_after>__version__ = '0.1.50'
__version__date__ = '29/01/2018'

# histórico
#
# - ?.?.? - 20??/??/??
#
# - 0.1.50 - 2018/01/29
#   . Várias melhorias
#
# - 0.1.40 - 2017/12/01
#   . Pedido: grade de itens pedidos
#
# - 0.1.39 - 2017/11/30
#   . App RH
#   . Impressão de cartelas de lote melhorada
#   . Toda rotina de bipagem para inventário de rolos e geração de arquivo
#     de inventário para o Systêxtil
#
# - 0.1.38 - 2017/11/17
#   . Alteração de layout e algumas ULRs
#   . Adição de link de base de conhecimento
#
# - 0.1.37 - 2017/11/13
#   . Abrindo Systêxtil sem ajuda de página externa (funcionando também de fora
#     da Tussor)
#
# - 0.1.36 - 2017/11/09
#   . OP: Prevendo OP com mais de uma finha na informações básicas
#   . OP e OS: acertando quantidades apresentadas - considerando perda e etc.
#
# - 0.1.35 - 2017/11/08
#   . Pedido: Situação da venda
#   . OP e OS: Link para pedido
#   . Pedido: OPs
#
# - 0.1.34 - 2017/11/08
#   . OP pendente: filtro por coleção
#
# - 0.1.33 - 2017/11/07
#   . Conferência de pedido: melhorias
#   . Pedido: Inicio de desenvolvimento
#
# - 0.1.32 - 2017/11/06
#   . OP + número de pedido e pedido de cliente
#   . renomeando páginas da intranet
#   . Melhorias em várias telas
#
# - 0.1.31 - 2017/10/18
#   . Ordens pendentes por estágio
#   . Melhorias em várias telas
#
# - 0.1.30 - 2017/09/29
#   . Primeira rotina de impressão térmica
#     . Impressão de etiqueta de lote
#
# - 0.1.29 - 2017/09/26
#   . Produto: links para insumos
#   . Insumo: adicionada informação de "Usado em"
#     . links para produtos
#
# - 0.1.28 - 2017/09/26
#   . Layout: alterações
#   . Insumo: Início da APP
#
# - 0.1.26 - 2017/09/12
#   . OP: depósitos destino da produção
#
# - 0.1.25 - 2017/09/12
#   . gravando log de alterações na tabela RecordTracking
#
# - 0.1.24 - 2017/09/11
#   . criando log de alterações em registros de tabelas
#   . Criada tabela RecordTracking para guardar log de alterações
#   . OP: adicionando quat. de 2a. qualidade e perda ao resumo por estágio
#   . App Utils para Classes, Middlewares e outras coisas uteis para
#     vários Apps
#
# - 0.1.23 - 2017/08/30
#   . adição de varios links para referência
#   . criação de view de modelo
#
# - 0.1.22 - 2017/08/29
#   . ajustes de layout (finalizado)
#
# - 0.1.21 - 2017/08/28
#   . view de referências - continuando
#   . ajustes de layout (iniciando por produto / ref e contagem)
#
# - 0.1.20 - 2017/08/25
#   . view de referências (produtos confeccionados)
#   . Layout:
#     . sem bordas
#     . eliminando footer
#     . não imprime href link
#
# - 0.1.19 - 2017/08/08
#   . refactoring de lotes.models
#   . análise de produzido
#     . qtd por período e alternativa
#     . qtd por data de corte e alternativa
#
# - 0.1.18 - 2017/08/04
#   . OP: datas
#
# - 0.1.17 - 2017/08/04
#   . OS: datas, observação, insumos
#
# - 0.1.16 - 2017/08/03
#   . estatística de produto: prevendo PG
#   . OS: retornos, datas
#
# - 0.1.15 - 2017/08/01
#   . OS: lotes
#   . Posição de lote: OSs
#   . OP: Mais informações
#
# - 0.1.14 - 2017/07/31
#   . Posição de lote: Mais informações
#   . OP: Mais informações
#   . OS: Iniciando busca de informações de OS
#
# - 0.1.13 - 2017/07/20
#   . ficha de cliente: pequena correção e melhoria
#   . visualização de OP:
#     . links para lotes
#     . situação de OP
#     . OP relacionada
#   . posição de lote:
#     . GET
#     . apresentando como tipo (MD, PG, PA), ref., cor e tam.
#     . link para OP
#
# - 0.1.12 - 2017/07/12
#   . controle de data de saída:
#     . transportadora
#     . ajustes de tela
#
# - 0.1.11 - 2017/07/11
#   . melhorias na consulta de OP
#   . adicionados campos à tela de controle de data de saída
#
# - 0.1.10 - 2017/07/10
#   . Evoluindo controle de data de saída
#   . Evoluindo dados de pedido para conferencia
#     Era só infAdProd, adicionei EAN e Narrativa
#
# - 0.1.9 - 2017/07/07
#   . controle de data de saída: Mais informações
#
# - 0.1.8 - 2017/07/06
#   . iniciando app logística
#   . controle de data de saída
#
# - 0.1.7 - 2017/07/05
#   . ficha de cliente (consultando F1)
#
# - 0.1.6 - 2017/06/30
#   . consulta infAdProd por número de pedido
#
# - 0.1.5 - 2017/06/29
#   . corrigido bug em view de OP
#   . posição de lote reescrita - vira classe
#   . responsáveis aparece estágio sem responsável
#   . autofocus nos forms
#
# - 0.1.4 - 2017/06/23
#   . menus reestruturados
#   . posição de lote reescrita - tira ajax e utiliza models
#
# - 0.1.3 - 2017/06/16
#   . consulta de OP: lotes por estágios, por OS e individuais
#   . lista produtos nível 1 sem preço médio
#
# - 0.1.2 - 2017/06/14
#   . consulta de responsável por estágio
#   . mostra cadastro de estágios
#   . em Produtos, mostra tamanhos e cores
#   . em Produtos, nível 1, mostra alternativas e roteiros
#
# - 0.1.1 - 2017/06/12
#   . melhorias na visualização das informações do lote
#   . mostra cadastro de depósitos
#
# - 0.1.0 - 2017/06/07
#   . primeira versão divulgada
#   . estatística de produtos
<|endoftext|>


<commit_before>from gi.repository import Gtk

from forecastmgmt.ui.masterdata_mask import MasterdataMask
from forecastmgmt.ui.forecast_mask import ForecastMask

from forecastmgmt.ui.forecast.forecast_new_dialog import ForecastNewDialog

class MainWindow(Gtk.Window):

    def __init__(self):
        Gtk.Window.__init__(self, title="Forecaster")
        self.set_default_size(800,600)

        # The main area, grid
        self.grid = Gtk.Grid()
        self.grid.set_orientation(Gtk.Orientation.VERTICAL)
        self.add(self.grid)

        menubar=self.create_menubar()
        self.grid.add(menubar)

        toolbar=self.create_toolbar()
        self.grid.add(toolbar)

        self.working_area=Gtk.Box(orientation=Gtk.Orientation.VERTICAL)
        self.grid.add(self.working_area)
        self.set_working_area("forecast")

        self.create_status_bar()
        self.grid.add(self.statusbar)

    def create_status_bar(self):
        self.statusbar = Gtk.Statusbar()
        self.statusbar.add(Gtk.Label("statusbar"))

    def create_toolbar(self):
        toolbar=Gtk.Toolbar()

        toolbutton_forecast=Gtk.ToolButton(Gtk.STOCK_ABOUT)
        toolbutton_forecast.set_tooltip_text("forecast")
        toolbutton_forecast.connect("clicked", self.on_toolbutton_forecast)
        toolbar.add(toolbutton_forecast)

        toolbutton_master_data=Gtk.ToolButton(Gtk.STOCK_EXECUTE)
        toolbutton_master_data.set_tooltip_text("master data")
        toolbutton_master_data.connect("clicked", self.on_toolbutton_masterdata)
        toolbar.add(toolbutton_master_data)

        toolbutton_quit=Gtk.ToolButton(Gtk.STOCK_QUIT)
        toolbutton_quit.set_tooltip_text("quit")
        toolbutton_quit.connect("clicked", self.on_menu_file_quit)
        toolbar.add(toolbutton_quit)

        return toolbar

    def create_menubar(self):
        menubar = Gtk.MenuBar()

        file_menu_entry = Gtk.MenuItem("File")

        menu = Gtk.Menu()
        mitem_file_new=Gtk.MenuItem("New forecast...")

        mitem_quit = Gtk.MenuItem("Quit")
        mitem_quit.connect("activate", self.on_menu_file_quit)
        menu.insert(mitem_file_new, 0)
        menu.insert(mitem_quit, 1)

        file_menu_entry.set_submenu(menu)

        forecast_menu_entry = Gtk.MenuItem("Forecast")
        forecast_menu = Gtk.Menu()
        mitem_forecast_new=Gtk.MenuItem("New")
        mitem_forecast_new.connect("activate", self.on_new_forecast)
        forecast_menu.insert(mitem_forecast_new,0)

        forecast_menu_entry.set_submenu(forecast_menu)

        masterdata_menu_entry = Gtk.MenuItem("Master data")
        masterdata_menu = Gtk.Menu()
        mitem_masterdata_person=Gtk.MenuItem("Person")
        masterdata_menu.insert(mitem_masterdata_person,0)

        masterdata_menu_entry.set_submenu(masterdata_menu)

        #filemenu.append(file_menu_entry)

        menubar.append(file_menu_entry)
        menubar.append(forecast_menu_entry)
        menubar.append(masterdata_menu_entry)

        return menubar

    def clean_working_area(self):
        for child in self.working_area.get_children():
            self.working_area.remove(child)

    def set_working_area(self, action="masterdata"):
        self.clean_working_area()
        if action=="masterdata":
            self.working_area.pack_start(MasterdataMask(self), False, False, 0)
        elif action=="forecast":
            self.working_area.pack_start(ForecastMask(self), False, False, 0)
        else:
            print("unimplemented")

        self.working_area.show_all()

    def on_menu_file_quit(self, widget):
        Gtk.main_quit()

    def on_toolbutton_masterdata(self, widget):
        self.set_working_area("masterdata")

    def on_toolbutton_forecast(self, widget):
        self.set_working_area("forecast")

    def on_new_forecast(self, widget):
        new_forecast_dialog=ForecastNewDialog(self)
        response = new_forecast_dialog.run()

        if response==Gtk.ResponseType.OK:
            new_forecast_dialog.perform_insert()
            self.clean_working_area()
            self.set_working_area("forecast")

        elif response==Gtk.ResponseType.CANCEL:
            print("insert nothing")
        else:
            print("unknown action")

        new_forecast_dialog.destroy()

win = MainWindow()
win.connect("delete-event", Gtk.main_quit)
win.show_all()
Gtk.main()
<commit_msg>removed statusbar<commit_after>from gi.repository import Gtk

from forecastmgmt.ui.masterdata_mask import MasterdataMask
from forecastmgmt.ui.forecast_mask import ForecastMask

from forecastmgmt.ui.forecast.forecast_new_dialog import ForecastNewDialog

class MainWindow(Gtk.Window):

    def __init__(self):
        Gtk.Window.__init__(self, title="Forecaster")
        self.set_default_size(800,600)

        # The main area, grid
        self.grid = Gtk.Grid()
        self.grid.set_orientation(Gtk.Orientation.VERTICAL)
        self.add(self.grid)

        menubar=self.create_menubar()
        self.grid.add(menubar)

        toolbar=self.create_toolbar()
        self.grid.add(toolbar)

        self.working_area=Gtk.Box(orientation=Gtk.Orientation.VERTICAL)
        self.grid.add(self.working_area)
        self.set_working_area("forecast")

        self.create_status_bar()
        #self.grid.add(self.statusbar)

    def create_status_bar(self):
        self.statusbar = Gtk.Statusbar()
        self.statusbar.add(Gtk.Label("statusbar"))

    def create_toolbar(self):
        toolbar=Gtk.Toolbar()

        toolbutton_forecast=Gtk.ToolButton(Gtk.STOCK_ABOUT)
        toolbutton_forecast.set_tooltip_text("forecast")
        toolbutton_forecast.connect("clicked", self.on_toolbutton_forecast)
        toolbar.add(toolbutton_forecast)

        toolbutton_master_data=Gtk.ToolButton(Gtk.STOCK_EXECUTE)
        toolbutton_master_data.set_tooltip_text("master data")
        toolbutton_master_data.connect("clicked", self.on_toolbutton_masterdata)
        toolbar.add(toolbutton_master_data)

        toolbutton_quit=Gtk.ToolButton(Gtk.STOCK_QUIT)
        toolbutton_quit.set_tooltip_text("quit")
        toolbutton_quit.connect("clicked", self.on_menu_file_quit)
        toolbar.add(toolbutton_quit)

        return toolbar

    def create_menubar(self):
        menubar = Gtk.MenuBar()

        file_menu_entry = Gtk.MenuItem("File")

        menu = Gtk.Menu()
        mitem_file_new=Gtk.MenuItem("New forecast...")

        mitem_quit = Gtk.MenuItem("Quit")
        mitem_quit.connect("activate", self.on_menu_file_quit)
        menu.insert(mitem_file_new, 0)
        menu.insert(mitem_quit, 1)

        file_menu_entry.set_submenu(menu)

        menubar.append(file_menu_entry)

        return menubar

    def clean_working_area(self):
        for child in self.working_area.get_children():
            self.working_area.remove(child)

    def set_working_area(self, action="masterdata"):
        self.clean_working_area()
        if action=="masterdata":
            self.working_area.pack_start(MasterdataMask(self), False, False, 0)
        elif action=="forecast":
            self.working_area.pack_start(ForecastMask(self), False, False, 0)
        else:
            print("unimplemented")

        self.working_area.show_all()

    def on_menu_file_quit(self, widget):
        Gtk.main_quit()

    def on_toolbutton_masterdata(self, widget):
        self.set_working_area("masterdata")

    def on_toolbutton_forecast(self, widget):
        self.set_working_area("forecast")

    def on_new_forecast(self, widget):
        new_forecast_dialog=ForecastNewDialog(self)
        response = new_forecast_dialog.run()

        if response==Gtk.ResponseType.OK:
            new_forecast_dialog.perform_insert()
            self.clean_working_area()
            self.set_working_area("forecast")

        elif response==Gtk.ResponseType.CANCEL:
            print("insert nothing")
        else:
            print("unknown action")

        new_forecast_dialog.destroy()

win = MainWindow()
win.connect("delete-event", Gtk.main_quit)
win.show_all()
Gtk.main()
<|endoftext|>


<commit_before>#!/usr/bin/python
# -*- encoding: utf-8 -*-
"""

"""
from collections import defaultdict

__author__ = 'Martin Martimeo <martin@martimeo.de>'
__date__ = '02.07.14 - 17:04'

class BaseNode(object):
    tag = None
    ignore_whitespace = None
    disallow_text = None
    strip_before = None
    full_match = None

    def __init__(self, arg: str):
        self._parent = None
        self._children = []
        self._text = ""

        # Parse arg
        if "=" in arg:
            tag, args = arg.split("=", 1)
            self._args = args.split(",")
        else:
            self._args = []

    def close(self, arg: str):
        if self._text:
            self.appendText(self._text)
        return self._parent

    def appendText(self, text: str) -> bool:

        if self.ignore_whitespace and not text.strip():
            return False

        if self.disallow_text:
            raise TypeError("Node %s disallows in-between text" % type(self))

        self._children.append(text)
        return True

    def append(self, node) -> "BaseNode":

        if isinstance(node, str):
            self._text += node
            return node

        else:
            # Collect text until now as child
            if self._text:
                self.appendText(self._text)
                self._text = ""

            # Append node
            self._children.append(node)
            node._parent = self
            return node

    @classmethod
    def matches(cls, text: str):
        if cls.full_match:
            return cls.tag == text.strip()
        else:
            return text.startswith(cls.tag)

    def markdown(self):
        rtn = ""
        for child in self._children:
            if isinstance(child, str):
                # Escape markdown control sequences
                child = child.replace("*", "\\*").replace("_", "\\_").replace("·", "*")

                # Fold windows line endings together
                child = child.replace("\r\n", "\n").replace("\r", "\n")

                # Remove Magic characters
                child = child.replace('\x94', '').replace('\xa0', '')

                rtn += child
            else:
                if child.strip_before:
                    rtn = rtn.rstrip()
                rtn += child.markdown()
        return rtn

class RootNode(BaseNode):
    pass

class FontNode(BaseNode):
    tag = "font"

class LinkNode(BaseNode):
    tag = "url"

    @property
    def url(self):
        if len(self._args) >= 1:
            return self._args[0]

    def markdown(self):
        if not self.url:
            return "[%s]" % super().markdown()
        else:
            return "[%s](%s)" % (super().markdown(), self.url)

class ColorNode(BaseNode):
    tag = "color"

class FontSizeNode(BaseNode):
    tag = "size"

class BlockNode(BaseNode):
    tag = "block"

    def markdown(self):
        return "\n\n" + super().markdown() + "\n\n"

class UnderlineNode(BaseNode):
    tag = "u"
    full_match = True

class BoldNode(BaseNode):
    tag = "b"
    full_match = True

    def markdown(self):
        return "**" + super().markdown() + "**"

class ItalicNode(BaseNode):
    tag = "i"

    def markdown(self):
        return "_" + super().markdown() + "_"

class ListNode(BaseNode):
    tag = "list"
    ignore_whitespace = True
    disallow_text = True
    strip_before = True

    @property
    def list_type(self):

        try:
            if len(self._args) and int(self._args[0]):
                return "1."
            else:
                return " *"
        except ValueError:
            return " -"

    def markdown(self):
        """
            Lists must be preceeded by an empty line
        :return:
        """
        rtn = "\n\n"
        for child in self._children:

            if isinstance(child, LiNode):
                rtn += child.markdown()
            elif isinstance(child, ListNode):
                rtn = rtn.rstrip()
                rtn_child = child.markdown()
                while rtn_child.startswith("\n\n"):
                    rtn_child = rtn_child[1:]
                for line in rtn_child.split("\n"):
                    rtn += "    " + line + "\n"
            else:
                raise TypeError("Unknown child in list element: %s" % type(child))
        rtn = rtn.rstrip() + "\n\n"
        return rtn

class LiNode(BaseNode):
    tag = "*"

    def markdown(self):
        rtn = ""
        for line in super().markdown().split("\n"):
            rtn += "    " + line + "\n"
        rtn = " %s " % self._parent.list_type + rtn[4:]
        return rtn

class TableNode(BaseNode):
    tag = "table"
    ignore_whitespace = True
    disallow_text = True

    @property
    def width(self):
        if len(self._args) >= 2:
            return self._args[0]
        else:
            return None

    @property
    def height(self):
        if len(self._args) >= 2:
            return self._args[1]
        else:
            return None

    def markdown(self):

        # Convert all elements
        nm = {}
        widths = defaultdict(int)
        heights = defaultdict(int)
        for itr, tr in enumerate(self._children):
            icol = 0
            for itd, td in enumerate(tr._children):
                nm[itr, icol] = td.markdown().split("\n")
                widths[icol] = max(widths[icol], max(len(line) for line in nm[itr, icol]), 1)
                heights[itr] = max(heights[itr], len(nm[itr, icol]), 1)
                if td.rowspan > 1:
                    raise TypeError("Multirow table cells are not supported in markdown")
                icol += td.colspan

        # Print them
        rtn = ""

        # Print first line
        def print_separation():
            r = "+"
            for w in widths.values():
                r += "-" * w
                r += "+"
            r += "\n"
            return r

        rtn += print_separation()
        for itr, height in heights.items():
            for iline in range(0, height):
                rtn += "|"
                for itd, width in widths.items():
                    if not (itr, itd) in nm:
                        line = ""
                    elif len(nm[itr, itd]) > iline:
                        line = nm[itr, itd][iline]
                    else:
                        line = ""
                    rtn += line
                    rtn += (width - len(line)) * ' '
                    rtn += "|"
                rtn += "\n"
            rtn += print_separation()
        return rtn

class TrNode(BaseNode):
    tag = "tr"
    ignore_whitespace = True
    disallow_text = True

class TdNode(BaseNode):
    tag = "td"

    @property
    def colspan(self):
        if len(self._args) >= 4:
            return int(self._args[2]) or 1
        else:
            return 1

    @property
    def rowspan(self):
        if len(self._args) >= 4:
            return int(self._args[3]) or 1
        else:
            return 1
<commit_msg>Alway hard break line endings comming from bbcode<commit_after>#!/usr/bin/python
# -*- encoding: utf-8 -*-
"""

"""
from collections import defaultdict
import re

__author__ = 'Martin Martimeo <martin@martimeo.de>'
__date__ = '02.07.14 - 17:04'

class BaseNode(object):
    tag = None
    ignore_whitespace = None
    disallow_text = None
    strip_before = None
    full_match = None

    def __init__(self, arg: str):
        self._parent = None
        self._children = []
        self._text = ""

        # Parse arg
        if "=" in arg:
            tag, args = arg.split("=", 1)
            self._args = args.split(",")
        else:
            self._args = []

    def close(self, arg: str):
        if self._text:
            self.appendText(self._text)
        return self._parent

    def appendText(self, text: str) -> bool:

        if self.ignore_whitespace and not text.strip():
            return False

        if self.disallow_text:
            raise TypeError("Node %s disallows in-between text" % type(self))

        self._children.append(text)
        return True

    def append(self, node) -> "BaseNode":

        if isinstance(node, str):
            self._text += node
            return node

        else:
            # Collect text until now as child
            if self._text:
                self.appendText(self._text)
                self._text = ""

            # Append node
            self._children.append(node)
            node._parent = self
            return node

    @classmethod
    def matches(cls, text: str):
        if cls.full_match:
            return cls.tag == text.strip()
        else:
            return text.startswith(cls.tag)

    def markdown(self):
        rtn = ""
        for child in self._children:
            if isinstance(child, str):
                # Escape markdown control sequences
                child = child.replace("*", "\\*").replace("_", "\\_").replace("·", "*").replace("-", "\-")

                # Fold windows line endings together
                child = child.replace("\r\n", "\n").replace("\r", "\n")

                # Remove Magic characters
                child = child.replace('\x94', '').replace('\xa0', '')

                # Challenge linebreaks
                child = re.sub(r"[\n]+", "\n", child)
                child = child.replace("\n", "\n\n")

                rtn += child
            else:
                if child.strip_before:
                    rtn = rtn.rstrip()
                rtn += child.markdown()
        return rtn

class RootNode(BaseNode):
    pass

class FontNode(BaseNode):
    tag = "font"

class LinkNode(BaseNode):
    tag = "url"

    @property
    def url(self):
        if len(self._args) >= 1:
            return self._args[0]

    def markdown(self):
        if not self.url:
            return "[%s]" % super().markdown()
        else:
            return "[%s](%s)" % (super().markdown(), self.url)

class ColorNode(BaseNode):
    tag = "color"

class FontSizeNode(BaseNode):
    tag = "size"

class BlockNode(BaseNode):
    tag = "block"

    def markdown(self):
        return "\n\n" + super().markdown() + "\n\n"

class UnderlineNode(BaseNode):
    tag = "u"
    full_match = True

class BoldNode(BaseNode):
    tag = "b"
    full_match = True

    def markdown(self):
        rtn = ""
        for line in super().markdown().split("\n"):
            if line:
                rtn += "**" + line + "**\n"
            else:
                rtn += "\n"
        return rtn

class ItalicNode(BaseNode):
    tag = "i"

    def markdown(self):
        rtn = ""
        for line in super().markdown().split("\n"):
            if line:
                rtn += "_" + line + "_\n"
            else:
                rtn += "\n"
        return rtn

class ListNode(BaseNode):
    tag = "list"
    ignore_whitespace = True
    disallow_text = True
    strip_before = True

    @property
    def list_type(self):

        try:
            if len(self._args) and int(self._args[0]):
                return "1."
            else:
                return " *"
        except ValueError:
            return " -"

    def markdown(self):
        """
            Lists must be preceeded by an empty line
        :return:
        """
        rtn = "\n\n"
        for child in self._children:

            if isinstance(child, LiNode):
                rtn += child.markdown()
            elif isinstance(child, ListNode):
                rtn = rtn.rstrip()
                rtn_child = child.markdown()
                while rtn_child.startswith("\n\n"):
                    rtn_child = rtn_child[1:]
                for line in rtn_child.split("\n"):
                    rtn += "    " + line + "\n"
            else:
                raise TypeError("Unknown child in list element: %s" % type(child))
        rtn = rtn.rstrip() + "\n\n"
        return rtn

class LiNode(BaseNode):
    tag = "*"

    def markdown(self):
        rtn = ""
        for line in super().markdown().split("\n"):
            rtn += "    " + line + "\n"
        rtn = " %s " % self._parent.list_type + rtn[4:]
        return rtn

class TableNode(BaseNode):
    tag = "table"
    ignore_whitespace = True
    disallow_text = True

    @property
    def width(self):
        if len(self._args) >= 2:
            return self._args[0]
        else:
            return None

    @property
    def height(self):
        if len(self._args) >= 2:
            return self._args[1]
        else:
            return None

    def markdown(self):

        # Convert all elements
        nm = {}
        widths = defaultdict(int)
        heights = defaultdict(int)
        for itr, tr in enumerate(self._children):
            icol = 0
            for itd, td in enumerate(tr._children):
                nm[itr, icol] = td.markdown().split("\n")
                widths[icol] = max(widths[icol], max(len(line) for line in nm[itr, icol]), 1)
                heights[itr] = max(heights[itr], len(nm[itr, icol]), 1)
                if td.rowspan > 1:
                    raise TypeError("Multirow table cells are not supported in markdown")
                icol += td.colspan

        # Print them
        rtn = ""

        # Print first line
        def print_separation():
            r = "+"
            for w in widths.values():
                r += "-" * w
                r += "+"
            r += "\n"
            return r

        rtn += print_separation()
        for itr, height in heights.items():
            for iline in range(0, height):
                rtn += "|"
                for itd, width in widths.items():
                    if not (itr, itd) in nm:
                        line = ""
                    elif len(nm[itr, itd]) > iline:
                        line = nm[itr, itd][iline]
                    else:
                        line = ""
                    rtn += line
                    rtn += (width - len(line)) * ' '
                    rtn += "|"
                rtn += "\n"
            rtn += print_separation()
        return rtn

class TrNode(BaseNode):
    tag = "tr"
    ignore_whitespace = True
    disallow_text = True

class TdNode(BaseNode):
    tag = "td"

    @property
    def colspan(self):
        if len(self._args) >= 4:
            return int(self._args[2]) or 1
        else:
            return 1

    @property
    def rowspan(self):
        if len(self._args) >= 4:
            return int(self._args[3]) or 1
        else:
            return 1
<|endoftext|>


<commit_before># Lint as: python2, python3
# Copyright 2019 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for tfx.components.example_gen.utils."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import re
from typing import Text
# Standard Imports

import tensorflow as tf

from tfx.components.example_gen import utils
from tfx.orchestration import data_types
from tfx.proto import example_gen_pb2
from tfx.utils import io_utils
from tfx.utils import json_utils

class UtilsTest(tf.test.TestCase):

  def setUp(self):
    super(UtilsTest, self).setUp()
    # Create input splits.
    test_dir = os.path.join(
        os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
        self._testMethodName)
    self._input_base_path = os.path.join(test_dir, 'input_base')
    tf.io.gfile.makedirs(self._input_base_path)

  def testDictToExample(self):
    instance_dict = {
        'int': 10,
        'float': 5.0,
        'str': 'abc',
        'int_list': [1, 2],
        'float_list': [3.0],
        'str_list': ['ab', 'cd'],
        'none': None,
        'empty_list': [],
    }
    example = utils.dict_to_example(instance_dict)
    self.assertProtoEquals(
        """
        features {
          feature {
            key: "empty_list"
            value {
            }
          }
          feature {
            key: "float"
            value {
              float_list {
                value: 5.0
              }
            }
          }
          feature {
            key: "float_list"
            value {
              float_list {
                value: 3.0
              }
            }
          }
          feature {
            key: "int"
            value {
              int64_list {
                value: 10
              }
            }
          }
          feature {
            key: "int_list"
            value {
              int64_list {
                value: 1
                value: 2
              }
            }
          }
          feature {
            key: "none"
            value {
            }
          }
          feature {
            key: "str"
            value {
              bytes_list {
                value: "abc"
              }
            }
          }
          feature {
            key: "str_list"
            value {
              bytes_list {
                value: "ab"
                value: "cd"
              }
            }
          }
        }
        """, example)

  def testMakeOutputSplitNames(self):
    split_names = utils.generate_output_split_names(
        input_config=example_gen_pb2.Input(splits=[
            example_gen_pb2.Input.Split(name='train', pattern='train/*'),
            example_gen_pb2.Input.Split(name='eval', pattern='eval/*')
        ]),
        output_config=example_gen_pb2.Output())
    self.assertListEqual(['train', 'eval'], split_names)

    split_names = utils.generate_output_split_names(
        input_config=example_gen_pb2.Input(splits=[
            example_gen_pb2.Input.Split(name='single', pattern='single/*')
        ]),
        output_config=example_gen_pb2.Output(
            split_config=example_gen_pb2.SplitConfig(splits=[
                example_gen_pb2.SplitConfig.Split(name='train', hash_buckets=2),
                example_gen_pb2.SplitConfig.Split(name='eval', hash_buckets=1)
            ])))
    self.assertListEqual(['train', 'eval'], split_names)

  def testMakeDefaultOutputConfig(self):
    output_config = utils.make_default_output_config(
        utils.make_default_input_config())
    self.assertEqual(2, len(output_config.split_config.splits))

    output_config = utils.make_default_output_config(
        example_gen_pb2.Input(splits=[
            example_gen_pb2.Input.Split(name='train', pattern='train/*'),
            example_gen_pb2.Input.Split(name='eval', pattern='eval/*')
        ]))
    self.assertEqual(0, len(output_config.split_config.splits))

  def testMakeOutputSplitNamesWithParameter(self):
    split_name_param = data_types.RuntimeParameter(
        name='split-name', ptype=Text, default=u'train')
    split_names = utils.generate_output_split_names(
        input_config={
            'splits': [{
                'name': split_name_param,
                'pattern': 'train/*'
            }, {
                'name': 'eval',
                'pattern': 'eval/*'
            }]
        },
        output_config=example_gen_pb2.Output())
    # Assert the json serialized version because RuntimeParameters only get
    # serialized after that.
    self.assertEqual(
        json_utils.dumps([split_name_param, 'eval']),
        json_utils.dumps(split_names))

    split_names = utils.generate_output_split_names(
        input_config=example_gen_pb2.Input(splits=[
            example_gen_pb2.Input.Split(name='single', pattern='single/*')
        ]),
        output_config={
            'split_config': {
                'splits': [{
                    'name': split_name_param,
                    'hash_buckets': 2
                }, {
                    'name': 'eval',
                    'hash_buckets': 1
                }]
            }
        })
    # Assert the json serialized version because RuntimeParameters only get
    # serialized after that.
    self.assertEqual(
        json_utils.dumps([split_name_param, 'eval']),
        json_utils.dumps(split_names))

  def testMakeDefaultOutputConfigWithParameter(self):
    split_name_param = data_types.RuntimeParameter(
        name='split-name', ptype=Text, default=u'train')
    output_config = utils.make_default_output_config({
        'splits': [{
            'name': split_name_param,
            'pattern': 'train/*'
        }, {
            'name': 'eval',
            'pattern': 'eval/*'
        }]
    })
    self.assertEqual(0, len(output_config.split_config.splits))

  def testGlobToRegex(self):
    glob_pattern = 'a(b)c'
    self.assertEqual(1, re.compile(glob_pattern).groups)
    regex_pattern = utils._glob_to_regex(glob_pattern)
    self.assertEqual(0, re.compile(regex_pattern).groups)
    self.assertEqual(glob_pattern,
                     re.match(regex_pattern, glob_pattern).group())

  def testCalculateSplitsFingerprint(self):
    split1 = os.path.join(self._input_base_path, 'split1', 'data')
    io_utils.write_string_file(split1, 'testing')
    os.utime(split1, (0, 1))
    split2 = os.path.join(self._input_base_path, 'split2', 'data')
    io_utils.write_string_file(split2, 'testing2')
    os.utime(split2, (0, 3))

    splits = [
        example_gen_pb2.Input.Split(name='s1', pattern='split1/*'),
        example_gen_pb2.Input.Split(name='s2', pattern='split2/*')
    ]
    fingerprint, span, version = utils.calculate_splits_fingerprint_span_and_version(
        self._input_base_path, splits)
    self.assertEqual(
        fingerprint,
        'split:s1,num_files:1,total_bytes:7,xor_checksum:1,sum_checksum:1\n'
        'split:s2,num_files:1,total_bytes:8,xor_checksum:3,sum_checksum:3')
    self.assertEqual(span, '0')
    self.assertEqual(version, '0')

  def testSpanNoMatching(self):
    splits = [
        example_gen_pb2.Input.Split(name='s1', pattern='span{SPAN}/split1/*'),
        example_gen_pb2.Input.Split(name='s2', pattern='span{SPAN}/split2/*')
    ]
    with self.assertRaisesRegexp(ValueError,
                                 'Cannot find matching for split'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path,
                                                 splits)

  def testVersionNoMatching(self):
    splits = [
        example_gen_pb2.Input.Split(name='s1',
                                    pattern='version{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
                                    pattern='version{VERSION}/split2/*')
    ]
    with self.assertRaisesRegexp(ValueError,
                                 'Cannot find matching for split'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path,
                                                 splits)

  def testSpanWrongFormat(self):
    wrong_span = os.path.join(self._input_base_path, 'spanx', 'split1', 'data')
    io_utils.write_string_file(wrong_span, 'testing_wrong_span')

    splits = [
        example_gen_pb2.Input.Split(name='s1', pattern='span{SPAN}/split1/*'),
        example_gen_pb2.Input.Split(name='s2', pattern='span{SPAN}/split2/*')
    ]
    with self.assertRaisesRegexp(ValueError, 'Cannot find span number'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path,
                                                 splits)

  def testVersionWrongFormat(self):
    wrong_span = os.path.join(self._input_base_path, 'versionx', 'split1',
                              'data')
    io_utils.write_string_file(wrong_span, 'testing_wrong_version')

    splits = [
        example_gen_pb2.Input.Split(name='s1',
                                    pattern='version{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
                                    pattern='version{VERSION}/split2/*')
    ]
    with self.assertRaisesRegexp(ValueError, 'Cannot find version number'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path, splits)

  def testCalculateSplitsFpSpanAndVersion(self):
    # Test align of span and version numbers.
    span1_v1_split1 = os.path.join(self._input_base_path, 'span01', 'ver01',
                                   'split1', 'data')
    io_utils.write_string_file(span1_v1_split1, 'testing11')
    span1_v1_split2 = os.path.join(self._input_base_path, 'span01', 'ver01',
                                   'split2', 'data')
    io_utils.write_string_file(span1_v1_split2, 'testing12')
    span2_v1_split1 = os.path.join(self._input_base_path, 'span02', 'ver01',
                                   'split1', 'data')
    io_utils.write_string_file(span2_v1_split1, 'testing21')

    # Test if error raised when span does not align.
    splits = [
        example_gen_pb2.Input.Split(name='s1',
                                    pattern='span{SPAN}/ver{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
                                    pattern='span{SPAN}/ver{VERSION}/split2/*')
    ]
    with self.assertRaisesRegexp(
        ValueError, 'Latest span should be the same for each split'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path, splits)

    span2_v1_split2 = os.path.join(self._input_base_path, 'span02', 'ver01',
                                   'split2', 'data')
    io_utils.write_string_file(span2_v1_split2, 'testing22')
    span2_v2_split1 = os.path.join(self._input_base_path, 'span02', 'ver02',
                                   'split1', 'data')
    io_utils.write_string_file(span2_v2_split1, 'testing21')

    # Test if error raised when span aligns but version does not.
    splits = [
        example_gen_pb2.Input.Split(name='s1',
                                    pattern='span{SPAN}/ver{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
                                    pattern='span{SPAN}/ver{VERSION}/split2/*')
    ]
    with self.assertRaisesRegexp(
        ValueError, 'Latest version should be the same for each split'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path, splits)
    span2_v2_split2 = os.path.join(self._input_base_path, 'span02', 'ver02',
                                   'split2', 'data')
    io_utils.write_string_file(span2_v2_split2, 'testing22')

    # Test if latest span is selected when span aligns for each split.
    splits = [
        example_gen_pb2.Input.Split(name='s1',
                                    pattern='span{SPAN}/ver{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
                                    pattern='span{SPAN}/ver{VERSION}/split2/*')
    ]
    _, span, version = utils.calculate_splits_fingerprint_span_and_version(
        self._input_base_path, splits)
    self.assertEqual(span, '02')
    self.assertEqual(version, '02')

if __name__ == '__main__':
  tf.test.main()
<commit_msg>adding tests for updated utils<commit_after># Lint as: python2, python3
# Copyright 2019 Google LLC. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""Tests for tfx.components.example_gen.utils."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import re
from typing import Text
# Standard Imports

import tensorflow as tf

from tfx.components.example_gen import utils
from tfx.orchestration import data_types
from tfx.proto import example_gen_pb2
from tfx.utils import io_utils
from tfx.utils import json_utils

class UtilsTest(tf.test.TestCase):

  def setUp(self):
    super(UtilsTest, self).setUp()
    # Create input splits.
    test_dir = os.path.join(
        os.environ.get('TEST_UNDECLARED_OUTPUTS_DIR', self.get_temp_dir()),
        self._testMethodName)
    self._input_base_path = os.path.join(test_dir, 'input_base')
    tf.io.gfile.makedirs(self._input_base_path)

  def testDictToExample(self):
    instance_dict = {
        'int': 10,
        'float': 5.0,
        'str': 'abc',
        'int_list': [1, 2],
        'float_list': [3.0],
        'str_list': ['ab', 'cd'],
        'none': None,
        'empty_list': [],
    }
    example = utils.dict_to_example(instance_dict)
    self.assertProtoEquals(
        """
        features {
          feature {
            key: "empty_list"
            value {
            }
          }
          feature {
            key: "float"
            value {
              float_list {
                value: 5.0
              }
            }
          }
          feature {
            key: "float_list"
            value {
              float_list {
                value: 3.0
              }
            }
          }
          feature {
            key: "int"
            value {
              int64_list {
                value: 10
              }
            }
          }
          feature {
            key: "int_list"
            value {
              int64_list {
                value: 1
                value: 2
              }
            }
          }
          feature {
            key: "none"
            value {
            }
          }
          feature {
            key: "str"
            value {
              bytes_list {
                value: "abc"
              }
            }
          }
          feature {
            key: "str_list"
            value {
              bytes_list {
                value: "ab"
                value: "cd"
              }
            }
          }
        }
        """, example)

  def testMakeOutputSplitNames(self):
    split_names = utils.generate_output_split_names(
        input_config=example_gen_pb2.Input(splits=[
            example_gen_pb2.Input.Split(name='train', pattern='train/*'),
            example_gen_pb2.Input.Split(name='eval', pattern='eval/*')
        ]),
        output_config=example_gen_pb2.Output())
    self.assertListEqual(['train', 'eval'], split_names)

    split_names = utils.generate_output_split_names(
        input_config=example_gen_pb2.Input(splits=[
            example_gen_pb2.Input.Split(name='single', pattern='single/*')
        ]),
        output_config=example_gen_pb2.Output(
            split_config=example_gen_pb2.SplitConfig(splits=[
                example_gen_pb2.SplitConfig.Split(name='train', hash_buckets=2),
                example_gen_pb2.SplitConfig.Split(name='eval', hash_buckets=1)
            ])))
    self.assertListEqual(['train', 'eval'], split_names)

  def testMakeDefaultOutputConfig(self):
    output_config = utils.make_default_output_config(
        utils.make_default_input_config())
    self.assertEqual(2, len(output_config.split_config.splits))

    output_config = utils.make_default_output_config(
        example_gen_pb2.Input(splits=[
            example_gen_pb2.Input.Split(name='train', pattern='train/*'),
            example_gen_pb2.Input.Split(name='eval', pattern='eval/*')
        ]))
    self.assertEqual(0, len(output_config.split_config.splits))

  def testMakeOutputSplitNamesWithParameter(self):
    split_name_param = data_types.RuntimeParameter(
        name='split-name', ptype=Text, default=u'train')
    split_names = utils.generate_output_split_names(
        input_config={
            'splits': [{
                'name': split_name_param,
                'pattern': 'train/*'
            }, {
                'name': 'eval',
                'pattern': 'eval/*'
            }]
        },
        output_config=example_gen_pb2.Output())
    # Assert the json serialized version because RuntimeParameters only get
    # serialized after that.
    self.assertEqual(
        json_utils.dumps([split_name_param, 'eval']),
        json_utils.dumps(split_names))

    split_names = utils.generate_output_split_names(
        input_config=example_gen_pb2.Input(splits=[
            example_gen_pb2.Input.Split(name='single', pattern='single/*')
        ]),
        output_config={
            'split_config': {
                'splits': [{
                    'name': split_name_param,
                    'hash_buckets': 2
                }, {
                    'name': 'eval',
                    'hash_buckets': 1
                }]
            }
        })
    # Assert the json serialized version because RuntimeParameters only get
    # serialized after that.
    self.assertEqual(
        json_utils.dumps([split_name_param, 'eval']),
        json_utils.dumps(split_names))

  def testMakeDefaultOutputConfigWithParameter(self):
    split_name_param = data_types.RuntimeParameter(
        name='split-name', ptype=Text, default=u'train')
    output_config = utils.make_default_output_config({
        'splits': [{
            'name': split_name_param,
            'pattern': 'train/*'
        }, {
            'name': 'eval',
            'pattern': 'eval/*'
        }]
    })
    self.assertEqual(0, len(output_config.split_config.splits))

  def testGlobToRegex(self):
    glob_pattern = 'a(b)c'
    self.assertEqual(1, re.compile(glob_pattern).groups)
    regex_pattern = utils._glob_to_regex(glob_pattern)
    self.assertEqual(0, re.compile(regex_pattern).groups)
    self.assertEqual(glob_pattern,
                     re.match(regex_pattern, glob_pattern).group())

  def testCalculateSplitsFingerprint(self):
    split1 = os.path.join(self._input_base_path, 'split1', 'data')
    io_utils.write_string_file(split1, 'testing')
    os.utime(split1, (0, 1))
    split2 = os.path.join(self._input_base_path, 'split2', 'data')
    io_utils.write_string_file(split2, 'testing2')
    os.utime(split2, (0, 3))

    splits = [
        example_gen_pb2.Input.Split(name='s1', pattern='split1/*'),
        example_gen_pb2.Input.Split(name='s2', pattern='split2/*')
    ]
    fingerprint, span, version = utils.calculate_splits_fingerprint_span_and_version(
        self._input_base_path, splits)
    self.assertEqual(
        fingerprint,
        'split:s1,num_files:1,total_bytes:7,xor_checksum:1,sum_checksum:1\n'
        'split:s2,num_files:1,total_bytes:8,xor_checksum:3,sum_checksum:3')
    self.assertEqual(span, '0')
    self.assertEqual(version, '0')

  def testSpanNoMatching(self):
    splits = [
        example_gen_pb2.Input.Split(name='s1', pattern='span{SPAN}/split1/*'),
        example_gen_pb2.Input.Split(name='s2', pattern='span{SPAN}/split2/*')
    ]
    with self.assertRaisesRegexp(ValueError,
                                 'Cannot find matching for split'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path,
                                                 splits)

  def testVersionNoMatching(self):
    span_dir = os.path.join(self._input_base_path, 'span01', 'wrong')
    io_utils.write_string_file(span_dir, 'testing_version_no_matching')

    splits = [
        example_gen_pb2.Input.Split(name='s1',
            pattern='span{SPAN}/version{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
            pattern='span{SPAN}/version{VERSION}/split2/*')
    ]
    with self.assertRaisesRegexp(ValueError,
                                 'Cannot find matching for split'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path,
                                                 splits)

  def testSpanWrongFormat(self):
    wrong_span = os.path.join(self._input_base_path, 'spanx', 'split1', 'data')
    io_utils.write_string_file(wrong_span, 'testing_wrong_span')

    splits = [
        example_gen_pb2.Input.Split(name='s1', pattern='span{SPAN}/split1/*'),
        example_gen_pb2.Input.Split(name='s2', pattern='span{SPAN}/split2/*')
    ]
    with self.assertRaisesRegexp(ValueError, 'Cannot find properties'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path,
                                                 splits)

  def testVersionNoSpan(self):
    wrong_version = os.path.join(self._input_base_path, 'versionx', 'split1',
                              'data')
    io_utils.write_string_file(wrong_version, 'testing_wrong_version')

    splits = [
        example_gen_pb2.Input.Split(name='s1',
                                    pattern='version{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
                                    pattern='version{VERSION}/split2/*')
    ]
    with self.assertRaisesRegexp(ValueError,
        'Version spec provided, but Span spec is not present'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path,
                                                          splits)

  def testVersionWrongFormat(self):
    wrong_version = os.path.join(self._input_base_path, 'span01', 'versionx',
                                 'split1', 'data')
    io_utils.write_string_file(wrong_version, 'testing_wrong_version')

    splits = [
        example_gen_pb2.Input.Split(name='s1',
            pattern='span{SPAN}/version{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
            pattern='span{SPAN}/version{VERSION}/split2/*')
    ]
    with self.assertRaisesRegexp(ValueError, 'Cannot find properties'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path,
                                                          splits)

  def testCalculateSplitsFingerprintSpanAndVersion(self):
    # Test align of span and version numbers.
    span1_v1_split1 = os.path.join(self._input_base_path, 'span01', 'ver01',
                                   'split1', 'data')
    io_utils.write_string_file(span1_v1_split1, 'testing11')
    span1_v1_split2 = os.path.join(self._input_base_path, 'span01', 'ver01',
                                   'split2', 'data')
    io_utils.write_string_file(span1_v1_split2, 'testing12')
    span2_v1_split1 = os.path.join(self._input_base_path, 'span02', 'ver01',
                                   'split1', 'data')
    io_utils.write_string_file(span2_v1_split1, 'testing21')

    # Test if error raised when span does not align.
    splits = [
        example_gen_pb2.Input.Split(name='s1',
                                    pattern='span{SPAN}/ver{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
                                    pattern='span{SPAN}/ver{VERSION}/split2/*')
    ]
    with self.assertRaisesRegexp(
        ValueError, 'Latest span should be the same for each split'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path, splits)

    span2_v1_split2 = os.path.join(self._input_base_path, 'span02', 'ver01',
                                   'split2', 'data')
    io_utils.write_string_file(span2_v1_split2, 'testing22')
    span2_v2_split1 = os.path.join(self._input_base_path, 'span02', 'ver02',
                                   'split1', 'data')
    io_utils.write_string_file(span2_v2_split1, 'testing21')

    # Test if error raised when span aligns but version does not.
    splits = [
        example_gen_pb2.Input.Split(name='s1',
                                    pattern='span{SPAN}/ver{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
                                    pattern='span{SPAN}/ver{VERSION}/split2/*')
    ]
    with self.assertRaisesRegexp(
        ValueError, 'Latest version should be the same for each split'):
      utils.calculate_splits_fingerprint_span_and_version(self._input_base_path, splits)
    span2_v2_split2 = os.path.join(self._input_base_path, 'span02', 'ver02',
                                   'split2', 'data')
    io_utils.write_string_file(span2_v2_split2, 'testing22')

    # Test if latest span is selected when span aligns for each split.
    splits = [
        example_gen_pb2.Input.Split(name='s1',
                                    pattern='span{SPAN}/ver{VERSION}/split1/*'),
        example_gen_pb2.Input.Split(name='s2',
                                    pattern='span{SPAN}/ver{VERSION}/split2/*')
    ]
    _, span, version = utils.calculate_splits_fingerprint_span_and_version(
        self._input_base_path, splits)
    self.assertEqual(span, '02')
    self.assertEqual(version, '02')

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