<commit_before>#!/usr/bin/env python
import os
import sys
import argparse
from phylotoast import util, biom_calc as bc, otu_calc as oc, graph_util as gu
errors = []
try:
    import matplotlib.pyplot as plt
    import matplotlib as mpl
except ImportError as ie:
    errors.append(ie)
try:
    import numpy as np
except ImportError as ie:
    errors.append(ie)
try:
    import pandas as pd
except ImportError as ie:
    errors.append(ie)
try:
    import biom
except ImportError as ie:
    errors.append(ie)
try:
    from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
except ImportError as ie:
    errors.append(ie)
if len(errors) != 0:
    for item in errors:
        print "Import Error:", item
    sys.exit()

def get_relative_abundance(biomfile):
    """
    Return relative abundance from a OTU table. OTUIDs are converted to their
    genus-species identifier.
    """
    biomf = biom.load_table(biomfile)
    norm_biomf = biomf.norm(inplace=False)
    rel_abd = {}
    for sid in norm_biomf.ids():
        rel_abd[sid] = {}
        for otuid in norm_biomf.ids("observation"):
            otuname = oc.otu_name(norm_biomf.metadata(otuid, axis="observation")["taxonomy"])
            abd = norm_biomf.get_value_by_ids(otuid, sid)
            rel_abd[sid][otuname] = abd
    ast_rel_abd = bc.arcsine_sqrt_transform(rel_abd)
    return ast_rel_abd

def plot_LDA(X_lda, y_lda, class_colors, exp_var, style, out_fp=""):
    """
    Plot transformed LDA data.
    """
    cats = class_colors.keys()
    group_lda = {c: [] for c in cats}
    fig = plt.figure(figsize=(15, 10))
    ax = fig.add_subplot(111)
    for i, target_name in zip(range(len(cats)), cats):
        cat_x = X_lda[:, 0][y_lda == target_name]
        if X_lda.shape[1] == 1:
            cat_y = np.ones((cat_x.shape[0], 1)) + i
        else:
            cat_y = X_lda[:, 1][y_lda == target_name]
        group_lda[target_name].append(cat_x)
        group_lda[target_name].append(cat_y)
        plt.scatter(x=cat_x, y=cat_y, label=target_name,
                    color=class_colors[target_name],
                    alpha=0.85, s=250, edgecolors="k")
    mpl.rc("font", family="Arial")  # define font for figure text
    mpl.rc('xtick', labelsize=12)  # increase X axis ticksize
    mpl.rc('ytick', labelsize=12)  # increase Y axis ticksize
    if X_lda.shape[1] == 1:
        plt.ylim((0.5, 2.5))
    plt.xlabel("LD1 (Percent Explained Variance: {:.3f}%)".format(exp_var[0]*100), fontsize=16)
    plt.ylabel("LD2 (Percent Explained Variance: {:.3f}%)".format(exp_var[1]*100), fontsize=16)
    leg = plt.legend(loc="best", frameon=True, framealpha=1, fontsize=16)
    leg.get_frame().set_edgecolor('k')
    if style:
        gu.ggplot2_style(ax)
        fc = "0.8"
    else:
        fc = "none"

    # save or display result
    if out_fp:
        plt.savefig(out_fp, facecolor=fc, edgecolor="none", dpi=300,
                    bbox_inches="tight", pad_inches=0.1)
    else:
        plt.show()

def run_LDA(df):
    """
    Run LinearDiscriminantAnalysis on input dataframe (df) and return
    transformed data, scalings and
    """
    # Prep variables for sklearn LDA
    X = df[range(1, df.shape[1])].values     # input data matrix
    y = df["Condition"].values               # data categories list

    # Calculate LDA
    sklearn_lda = LDA()
    X_lda_sklearn = sklearn_lda.fit_transform(X, y)
    exp_var = sklearn_lda.explained_variance_ratio_

    return X_lda_sklearn, y, exp_var

def handle_program_options():
    parser = argparse.ArgumentParser(description="Create an LDA plot from\
                                     sample-grouped OTU data. It is necessary\
                                     to remove the header cell '#OTU ID'\
                                     before running this program.")
    parser.add_argument("-i", "--input_data_type", required=True,
                        choices=["biom", "unifrac_dm"],
                        default="unifrac_dm",
                        help="Specify if the input file is biom file format OTU \
                              table or unifrac distance matrix. If biom file is \
                              provided, the arc-sine transformed relative abundances \
                              eill be used as input whereas, if unifrac distance matrix.\
                              is given, unifrac distances will be used as input to LDA.\
                              [REQUIRED]")
    parser.add_argument("-bf", "--biom_file", required=True,
                        help="Input biom file format. [REQUIRED]")
    parser.add_argument("-m", "--map_fp", required=True,
                        help="Metadata mapping file. [REQUIRED]")
    parser.add_argument("-uf", "--unifrac_file",
                        help="Input unifrac datdistance matrix file. This is the \
                              output from ")
    parser.add_argument("-g", "--group_by", required=True,
                        help="A column name in the mapping file containing\
                              categorical values that will be used to identify \
                              groups. Each sample ID must have a group entry. \
                              Default is no categories and all the data will be \
                              treated as a single group.")
    parser.add_argument("-c", "--color_by", required=True,
                        help="A column name in the mapping file containing\
                              hexadecimal (#FF0000) color values that will\
                              be used to color the groups. Each sample ID must\
                              have a color entry.")
    parser.add_argument("--bubble",
                        help="If set, provide a file with 1 OTU name per line \
                              for bubble plotting. OTU name must be condensed to \
                              genus-species identifier. Default parameter value \
                              will not plot bubble plots.")
    parser.add_argument("--save_lda_input",
                        help="Save a CSV-format file of the transposed LDA-input\
                              table to the file specifed by this option.")
    parser.add_argument("--plot_title", default=None,
                        help="Plot title. Default is no title.")
    parser.add_argument("-o", "--out_fp", default="",
                        help="The path and file name to save the LDA plot under.\
                              If specified, the figure will be saved directly\
                              instead of opening a window in which the plot \
                              can be viewed before saving")
    parser.add_argument("-od", "--output_dir", default=".",
                        help="The directory to save the LDA bubble plots to.")
    parser.add_argument("--scale_by", default=1000, type=float,
                        help="Species relative abundance is multiplied by this \
                              factor in order to make appropriate visible \
                              bubbles in the output plots. Default is 1000.")
    parser.add_argument("-s", "--save_as", default="svg",
                        help="The type of image file for LDA plots. By default,\
                              files will be saved in SVG format.")
    parser.add_argument("--ggplot2_style", action="store_true",
                        help="Apply ggplot2 styling to the figure.")
    return parser.parse_args()

def main():
    args = handle_program_options()

    try:
        with open(args.map_fp):
            pass
    except IOError as ioe:
        err_msg = "\nError in metadata mapping filepath (-m): {}\n"
        sys.exit(err_msg.format(ioe))

    # Parse and read mapping file and obtain group colors
    header, imap = util.parse_map_file(args.map_fp)
    class_colors = util.color_mapping(imap, header, args.group_by, args.color_by)

    if args.input_data_type == "unifrac_dm":
        try:
            with open(args.unifrac_file):
                pass
        except IOError as ioe:
            err_msg = "\nError with unifrac distance matrix file (-d): {}\n"
            sys.exit(err_msg.format(ioe))
        uf_data = pd.read_csv(args.unifrac_file, sep="\t", index_col=0)
        uf_data.insert(0, "Condition", [imap[sid][header.index(args.group_by)]
                                        for sid in uf_data.index])
        sampleids = uf_data.index
        if args.save_lda_input:
            uf_data.to_csv(args.save_lda_input, sep="\t")
        # Run LDA
        X_lda, y_lda, exp_var = run_LDA(uf_data)
        # Plot LDA
        plot_LDA(X_lda, y_lda, class_colors, exp_var, style=args.ggplot2_style,
                 out_fp=args.out_fp)
    else:
        # Load biom file and calculate relative abundance
        try:
            rel_abd = get_relative_abundance(args.biom_file)
        except ValueError as ve:
            err_msg = "\nError with biom format file (-d): {}\n"
            sys.exit(err_msg.format(ve))
        df_rel_abd = pd.DataFrame(rel_abd).T
        df_rel_abd.insert(0, "Condition", [imap[sid][header.index(args.group_by)]
                                           for sid in df_rel_abd.index])
        sampleids = df_rel_abd.index
        if args.save_lda_input:
            df_rel_abd.to_csv(args.save_lda_input, sep="\t")
        # Run LDA
        X_lda, y_lda, exp_var = run_LDA(df_rel_abd)
        # Plot LDA
        plot_LDA(X_lda, y_lda, class_colors, exp_var, style=args.ggplot2_style,
                 out_fp=args.out_fp)

    if args.bubble:
        # Get otus for LDA bubble plots
        try:
            with open(args.bubble) as hojiehr:
                for line in hojiehr.readlines():
                    bubble_otus = line.strip().split("\r")
        except IOError as ioe:
            err_msg = "\nError in OTU name list file (--bubble): {}\n"
            sys.exit(err_msg.format(ioe))

        # Load biom file and calculate relative abundance
        try:
            rel_abd = get_relative_abundance(args.biom_file)
        except ValueError as ve:
            err_msg = "\nError with biom format file (-d): {}\n"
            sys.exit(err_msg.format(ve))
        category_idx = header.index(args.group_by)

        # Calculate position and size of SampleIDs to plot for each OTU
        for otuname in bubble_otus:
            plot_data = {cat: {"x": [], "y": [], "size": [], "label": []}
                         for cat in class_colors.keys()}
            for sid, data in zip(sampleids, X_lda):
                category = plot_data[imap[sid][category_idx]]
                try:
                    size = rel_abd[sid][otuname] * args.scale_by
                except KeyError as ke:
                    print "{} not found in {} sample.".format(ke, sid)
                    continue
                category["x"].append(float(data[0]))
                category["y"].append(float(data[1]))
                category["size"].append(size)

            # Plot LDA bubble for each OTU
            fig = plt.figure(figsize=(12, 9))
            ax = fig.add_subplot(111)
            for i, cat in enumerate(plot_data):
                plt.scatter(plot_data[cat]["x"], plot_data[cat]["y"],
                            plot_data[cat]["size"], label=cat,
                            color=class_colors[cat],
                            alpha=0.85, marker="o", edgecolor="k")
            mpl.rc("font", family="Arial")  # define font for figure text
            mpl.rc("xtick", labelsize=12)  # increase X axis ticksize
            mpl.rc("ytick", labelsize=12)  # increase Y axis ticksize
            if X_lda.shape[1] == 1:
                plt.ylim((0.5, 2.5))
            plt.title(" ".join(otuname.split("_")), style="italic")
            plt.xlabel("LD1 (Percent Explained Variance: {:.3f}%)".format(exp_var[0]*100),
                       fontsize=12)
            plt.ylabel("LD2 (Percent Explained Variance: {:.3f}%)".format(exp_var[1]*100),
                       fontsize=12)
            lgnd = plt.legend(loc="best", scatterpoints=3, fontsize=12)
            # Change the legend marker size manually
            for i in range(len(class_colors.keys())):
                lgnd.legendHandles[i]._sizes = [75]

            # Set style for LDA bubble plots
            if args.ggplot2_style:
                gu.ggplot2_style(ax)
                fc = "0.8"
            else:
                fc = "none"

            # Save LDA bubble plots to output directory
            print "Saving chart for {}".format(" ".join(otuname.split("_")))
            fig.savefig(os.path.join(args.output_dir, "_".join(otuname.split())) + "." + args.save_as,
                        facecolor=fc, edgecolor="none", dpi=300,
                        bbox_inches="tight", pad_inches=0.2)
            plt.close(fig)

if __name__ == "__main__":
    sys.exit(main())
<commit_msg>Updated LDA.py<commit_after>#!/usr/bin/env python
"""
Abstract: This script calculates and returns LDA plots based on normalized relative
          abundances or distance matrices (for e.g. unifrac distance matrix).
"""

import sys
import argparse
from phylotoast import util, biom_calc as bc, graph_util as gu
errors = []
try:
    import matplotlib as mpl
    import matplotlib.pyplot as plt
    from mpl_toolkits.mplot3d import Axes3D
    mpl.rc("font", family="Arial")  # define font for figure text
except ImportError as ie:
    errors.append(ie)
try:
    import numpy as np
except ImportError as ie:
    errors.append(ie)
try:
    import pandas as pd
except ImportError as ie:
    errors.append(ie)
try:
    import biom
except ImportError as ie:
    errors.append(ie)
try:
    from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
except ImportError as ie:
    errors.append(ie)
if len(errors) != 0:
    for item in errors:
        print("Import Error:", item)
    sys.exit()

def plot_LDA(X_lda, y_lda, class_colors, exp_var, style, fig_size, label_pad, font_size,
             dim=2, zangles=None, out_fp=""):
    """
    Plot transformed LDA data.
    """
    cats = class_colors.keys()
    fig = plt.figure(figsize=fig_size)
    if dim == 3:
        ax = fig.add_subplot(111, projection="3d")
        ax.view_init(elev=zangles[1], azim=zangles[0])
        try:
            ax.set_zlabel("LD3 (Percent Explained Variance: {:.3f}%)".
                          format(exp_var[2]*100), fontsize=font_size, labelpad=label_pad)
        except:
            ax.set_zlabel("LD3", fontsize=font_size, labelpad=label_pad)
        for i, target_name in zip(range(len(cats)), cats):
            cat_x = X_lda[:, 0][y_lda == target_name]
            cat_y = X_lda[:, 1][y_lda == target_name]
            cat_z = X_lda[:, 2][y_lda == target_name]
            ax.scatter(xs=cat_x, ys=cat_y, zs=cat_z, label=target_name,
                       c=class_colors[target_name], alpha=0.85, s=250,
                       edgecolors="k", zdir="z")
    else:
        ax = fig.add_subplot(111)
        for i, target_name in zip(range(len(cats)), cats):
            cat_x = X_lda[:, 0][y_lda == target_name]
            if X_lda.shape[1] == 1:
                cat_y = np.ones((cat_x.shape[0], 1)) + i
            else:
                cat_y = X_lda[:, 1][y_lda == target_name]
            ax.scatter(x=cat_x, y=cat_y, label=target_name,
                       color=class_colors[target_name],
                       alpha=0.85, s=250, edgecolors="k")
    if X_lda.shape[1] == 1:
        plt.ylim((0.5, 2.5))
    try:
        ax.set_xlabel("LD1 (Percent Explained Variance: {:.3f}%)".
                      format(exp_var[0]*100), fontsize=font_size, labelpad=label_pad)
    except:
        ax.set_xlabel("LD1", fontsize=font_size, labelpad=label_pad)
    try:
        ax.set_ylabel("LD2 (Percent Explained Variance: {:.3f}%)".
                      format(exp_var[1]*100), fontsize=font_size, labelpad=label_pad)
    except:
        ax.set_ylabel("LD2", fontsize=font_size, labelpad=label_pad)

    leg = plt.legend(loc="best", scatterpoints=3, frameon=True, framealpha=1, fontsize=15)
    leg.get_frame().set_edgecolor('k')
    if dim == 2 and style:
        gu.ggplot2_style(ax)
        fc = "0.8"
    else:
        fc = "none"

    # save or display result
    if out_fp:
        plt.savefig(out_fp, facecolor=fc, edgecolor="none", dpi=300,
                    bbox_inches="tight", pad_inches=0.1)
    else:
        plt.show()

def run_LDA(df):
    """
    Run LinearDiscriminantAnalysis on input dataframe (df) and return
    transformed data, scalings and
    """
    # Prep variables for sklearn LDA
    X = df[range(1, df.shape[1])].values     # input data matrix
    y = df["Condition"].values               # data categories list

    # Calculate LDA
    sklearn_lda = LDA()
    X_lda_sklearn = sklearn_lda.fit_transform(X, y)
    try:
        exp_var = sklearn_lda.explained_variance_ratio_
    except AttributeError as ae:
        print("\n{}: explained variance cannot be computed.\nPlease check this GitHub PR:"
              " https://github.com/scikit-learn/scikit-learn/pull/6027".format(ae))
        return X_lda_sklearn, y, "NA"
    return X_lda_sklearn, y, exp_var

def handle_program_options():
    parser = argparse.ArgumentParser(description="Create an LDA plot from\
                                     sample-grouped OTU data. It is necessary\
                                     to remove the header cell '#OTU ID'\
                                     before running this program.")
    parser.add_argument("-i", "--otu_table", required=True,
                        help="Input biom file format OTU table. [REQUIRED]")
    parser.add_argument("-m", "--map_fp", required=True,
                        help="Metadata mapping file. [REQUIRED]")
    parser.add_argument("-g", "--group_by", required=True,
                        help="A column name in the mapping file containing\
                              categorical values that will be used to identify \
                              groups. Each sample ID must have a group entry. \
                              Default is no categories and all the data will be \
                              treated as a single group. [REQUIRED]")
    parser.add_argument("-c", "--color_by", required=True,
                        help="A column name in the mapping file containing\
                              hexadecimal (#FF0000) color values that will\
                              be used to color the groups. Each sample ID must\
                              have a color entry. [REQUIRED]")
    parser.add_argument("-dm", "--dist_matrix_file",
                        help="Input unifrac distance matrix file.")
    parser.add_argument("--save_lda_input",
                        help="Save a CSV-format file of the transposed LDA-input\
                              table to the file specifed by this option.")
    parser.add_argument("--plot_title", default=None,
                        help="Plot title. Default is no title.")
    parser.add_argument("-o", "--out_fp", default="",
                        help="The path and file name to save the LDA plot under.\
                              If specified, the figure will be saved directly\
                              instead of opening a window in which the plot \
                              can be viewed before saving")
    parser.add_argument("-d", "--dimensions", default=2, type=int, choices=[2, 3],
                        help="Choose whether to plot 2D or 3D.")
    parser.add_argument("--z_angles", type=float, nargs=2, default=[45., 30.],
                        help="Specify the azimuth and elevation angles for a 3D plot.")
    parser.add_argument("--figsize", default=[14, 8], type=int, nargs=2,
                        help="Specify the 'width height' in inches for LDA plots."
                             "By default, figure size is 14x8 inches.")
    parser.add_argument("--font_size", default=12, type=int,
                        help="Sets the font size for text elements in the plot.")
    parser.add_argument("--label_padding", default=15, type=int,
                        help="Sets the spacing in points between the each axis and its \
                             label.")
    parser.add_argument("--ggplot2_style", action="store_true",
                        help="Apply ggplot2 styling to the figure.")
    return parser.parse_args()

def main():
    args = handle_program_options()

    # Parse and read mapping file
    try:
        header, imap = util.parse_map_file(args.map_fp)
        category_idx = header.index(args.group_by)
    except IOError as ioe:
        err_msg = "\nError in metadata mapping filepath (-m): {}\n"
        sys.exit(err_msg.format(ioe))
    # Obtain group colors
    class_colors = util.color_mapping(imap, header, args.group_by, args.color_by)

    if args.dist_matrix_file:
        try:
            with open(args.dist_matrix_file):
                pass
        except IOError as ioe:
            err_msg = "\nError with unifrac distance matrix file (-d): {}\n"
            sys.exit(err_msg.format(ioe))
        uf_data = pd.read_csv(args.dist_matrix_file, sep="\t", index_col=0)
        uf_data.insert(0, "Condition", [imap[sid][category_idx] for sid in uf_data.index])
        if args.save_lda_input:
            uf_data.to_csv(args.save_lda_input, sep="\t")
        # Run LDA
        X_lda, y_lda, exp_var = run_LDA(uf_data)
        # Plot LDA
        if args.dimensions == 3:
            plot_LDA(X_lda, y_lda, class_colors, exp_var, style=args.ggplot2_style,
                     fig_size=args.figsize, label_pad=args.label_padding,
                     font_size=args.font_size, dim=3, zangles=args.z_angles,
                     out_fp=args.out_fp)
        else:
            plot_LDA(X_lda, y_lda, class_colors, exp_var, style=args.ggplot2_style,
                     fig_size=args.figsize, label_pad=args.label_padding,
                     font_size=args.font_size, out_fp=args.out_fp)
    else:
        # Load biom file and calculate relative abundance
        try:
            biomf = biom.load_table(args.otu_table)
        except IOError as ioe:
            err_msg = "\nError with biom format file (-d): {}\n"
            sys.exit(err_msg.format(ioe))
        # Get normalized relative abundances
        rel_abd = bc.relative_abundance(biomf)
        rel_abd = bc.arcsine_sqrt_transform(rel_abd)
        df_rel_abd = pd.DataFrame(rel_abd).T
        df_rel_abd.insert(0, "Condition", [imap[sid][category_idx] for sid in df_rel_abd.index])
        if args.save_lda_input:
            df_rel_abd.to_csv(args.save_lda_input, sep="\t")
        # Run LDA
        X_lda, y_lda, exp_var = run_LDA(df_rel_abd)
        # Plot LDA
        if args.dimensions == 3:
            plot_LDA(X_lda, y_lda, class_colors, exp_var, style=args.ggplot2_style,
                     fig_size=args.figsize, label_pad=args.label_padding,
                     font_size=args.font_size, dim=3, zangles=args.z_angles,
                     out_fp=args.out_fp)
        else:
            plot_LDA(X_lda, y_lda, class_colors, exp_var, style=args.ggplot2_style,
                     fig_size=args.figsize, label_pad=args.label_padding,
                     font_size=args.font_size, out_fp=args.out_fp)

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


<commit_before># Copyright 2019 Google LLC
#
# 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
#
#     https://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.

from functools import partial
import numpy as np
import unittest

from absl.testing import absltest
from absl.testing import parameterized

from jax import numpy as jnp
from jax import test_util as jtu, jit

from jax.config import config
config.parse_flags_with_absl()

all_dtypes = jtu.dtypes.floating + jtu.dtypes.integer + jtu.dtypes.complex

# TODO: these tests fail without fixed PRNG seeds.

class TestPolynomial(jtu.JaxTestCase):

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_dtype={}_leading={}_trailing={}".format(
       jtu.format_shape_dtype_string((length+leading+trailing,), dtype),
       leading, trailing),
     "dtype": dtype, "length": length, "leading": leading, "trailing": trailing}
    for dtype in all_dtypes
    for length in [0, 3, 9, 10, 17]
    for leading in [0, 1, 2, 3, 5, 7, 10]
    for trailing in [0, 1, 2, 3, 5, 7, 10]))
  # TODO(phawkins): no nonsymmetric eigendecomposition implementation on GPU.
  @jtu.skip_on_devices("gpu")
  def testRoots(self, dtype, length, leading, trailing):
    rng = jtu.rand_default(np.random.RandomState(0))

    def args_maker():
      p = rng((length,), dtype)
      return jnp.concatenate(
        [jnp.zeros(leading, p.dtype), p, jnp.zeros(trailing, p.dtype)]),

    jnp_fn = lambda arg: jnp.sort(jnp.roots(arg))
    np_fn = lambda arg: np.sort(np.roots(arg))
    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=False,
                            tol=3e-6)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_dtype={}_trailing={}".format(
        jtu.format_shape_dtype_string((length+trailing,), dtype), trailing),
     "dtype": dtype, "length": length, "trailing": trailing}
    for dtype in all_dtypes
    for length in [0, 1, 3, 10]
    for trailing in [0, 1, 3, 7]))
  # TODO(phawkins): no nonsymmetric eigendecomposition implementation on GPU.
  @jtu.skip_on_devices("gpu")
  def testRootsNostrip(self, length, dtype, trailing):
    rng = jtu.rand_default(np.random.RandomState(0))

    def args_maker():
      p = rng((length,), dtype)
      if length != 0:
        return jnp.concatenate([p, jnp.zeros(trailing, p.dtype)]),
      else:
        # adding trailing would make input invalid (start with zeros)
        return p,

    jnp_fn = lambda arg: jnp.sort(jnp.roots(arg, strip_zeros=False))
    np_fn = lambda arg: np.sort(np.roots(arg))
    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker,
                            check_dtypes=False, tol=1e-6)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_dtype={}_trailing={}".format(
      jtu.format_shape_dtype_string((length + trailing,), dtype), trailing),
      "dtype": dtype, "length": length, "trailing": trailing}
    for dtype in all_dtypes
    for length in [0, 1, 3, 10]
    for trailing in [0, 1, 3, 7]))
  # TODO: enable when there is an eigendecomposition implementation
  # for GPU/TPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testRootsJit(self, length, dtype, trailing):
    rng = jtu.rand_default(np.random.RandomState(0))

    def args_maker():
      p = rng((length,), dtype)
      if length != 0:
        return jnp.concatenate([p, jnp.zeros(trailing, p.dtype)]),
      else:
        # adding trailing would make input invalid (start with zeros)
        return p,

    roots_compiled = jit(partial(jnp.roots, strip_zeros=False))
    jnp_fn = lambda arg: jnp.sort(roots_compiled(arg))
    np_fn = lambda arg: np.sort(np.roots(arg))
    # Using strip_zeros=False makes the algorithm less efficient
    # and leads to slightly different values compared ot numpy
    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker,
                            check_dtypes=False, tol=1e-6)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_dtype={}_zeros={}_nonzeros={}".format(
        jtu.format_shape_dtype_string((zeros+nonzeros,), dtype),
        zeros, nonzeros),
     "zeros": zeros, "nonzeros": nonzeros, "dtype": dtype}
    for dtype in all_dtypes
    for zeros in [1, 2, 5]
    for nonzeros in [0, 3]))
  @jtu.skip_on_devices("gpu")
  def testRootsInvalid(self, zeros, nonzeros, dtype):
    raise unittest.SkipTest("getting segfaults on MKL")  # TODO(#3711)
    rng = jtu.rand_default(np.random.RandomState(0))

    # The polynomial coefficients here start with zero and would have to
    # be stripped before computing eigenvalues of the companion matrix.
    # Setting strip_zeros=False skips this check,
    # allowing jit transformation but yielding nan's for these inputs.
    p = jnp.concatenate([jnp.zeros(zeros, dtype), rng((nonzeros,), dtype)])

    if p.size == 1:
      # polynomial = const has no roots
      self.assertTrue(jnp.roots(p, strip_zeros=False).size == 0)
    else:
      self.assertTrue(jnp.any(jnp.isnan(jnp.roots(p, strip_zeros=False))))

if __name__ == "__main__":
  absltest.main(testLoader=jtu.JaxTestLoader())
<commit_msg>Skip unimplemented TPU cases in polynomial_test.py<commit_after># Copyright 2019 Google LLC
#
# 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
#
#     https://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.

from functools import partial
import numpy as np
import unittest

from absl.testing import absltest
from absl.testing import parameterized

from jax import numpy as jnp
from jax import test_util as jtu, jit

from jax.config import config
config.parse_flags_with_absl()

all_dtypes = jtu.dtypes.floating + jtu.dtypes.integer + jtu.dtypes.complex

# TODO: these tests fail without fixed PRNG seeds.

class TestPolynomial(jtu.JaxTestCase):

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_dtype={}_leading={}_trailing={}".format(
       jtu.format_shape_dtype_string((length+leading+trailing,), dtype),
       leading, trailing),
     "dtype": dtype, "length": length, "leading": leading, "trailing": trailing}
    for dtype in all_dtypes
    for length in [0, 3, 9, 10, 17]
    for leading in [0, 1, 2, 3, 5, 7, 10]
    for trailing in [0, 1, 2, 3, 5, 7, 10]))
  # TODO(phawkins): no nonsymmetric eigendecomposition implementation on GPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testRoots(self, dtype, length, leading, trailing):
    rng = jtu.rand_default(np.random.RandomState(0))

    def args_maker():
      p = rng((length,), dtype)
      return jnp.concatenate(
        [jnp.zeros(leading, p.dtype), p, jnp.zeros(trailing, p.dtype)]),

    jnp_fn = lambda arg: jnp.sort(jnp.roots(arg))
    np_fn = lambda arg: np.sort(np.roots(arg))
    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker, check_dtypes=False,
                            tol=3e-6)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_dtype={}_trailing={}".format(
        jtu.format_shape_dtype_string((length+trailing,), dtype), trailing),
     "dtype": dtype, "length": length, "trailing": trailing}
    for dtype in all_dtypes
    for length in [0, 1, 3, 10]
    for trailing in [0, 1, 3, 7]))
  # TODO(phawkins): no nonsymmetric eigendecomposition implementation on GPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testRootsNostrip(self, length, dtype, trailing):
    rng = jtu.rand_default(np.random.RandomState(0))

    def args_maker():
      p = rng((length,), dtype)
      if length != 0:
        return jnp.concatenate([p, jnp.zeros(trailing, p.dtype)]),
      else:
        # adding trailing would make input invalid (start with zeros)
        return p,

    jnp_fn = lambda arg: jnp.sort(jnp.roots(arg, strip_zeros=False))
    np_fn = lambda arg: np.sort(np.roots(arg))
    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker,
                            check_dtypes=False, tol=1e-6)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_dtype={}_trailing={}".format(
      jtu.format_shape_dtype_string((length + trailing,), dtype), trailing),
      "dtype": dtype, "length": length, "trailing": trailing}
    for dtype in all_dtypes
    for length in [0, 1, 3, 10]
    for trailing in [0, 1, 3, 7]))
  # TODO: enable when there is an eigendecomposition implementation
  # for GPU/TPU.
  @jtu.skip_on_devices("gpu", "tpu")
  def testRootsJit(self, length, dtype, trailing):
    rng = jtu.rand_default(np.random.RandomState(0))

    def args_maker():
      p = rng((length,), dtype)
      if length != 0:
        return jnp.concatenate([p, jnp.zeros(trailing, p.dtype)]),
      else:
        # adding trailing would make input invalid (start with zeros)
        return p,

    roots_compiled = jit(partial(jnp.roots, strip_zeros=False))
    jnp_fn = lambda arg: jnp.sort(roots_compiled(arg))
    np_fn = lambda arg: np.sort(np.roots(arg))
    # Using strip_zeros=False makes the algorithm less efficient
    # and leads to slightly different values compared ot numpy
    self._CheckAgainstNumpy(np_fn, jnp_fn, args_maker,
                            check_dtypes=False, tol=1e-6)

  @parameterized.named_parameters(jtu.cases_from_list(
    {"testcase_name": "_dtype={}_zeros={}_nonzeros={}".format(
        jtu.format_shape_dtype_string((zeros+nonzeros,), dtype),
        zeros, nonzeros),
     "zeros": zeros, "nonzeros": nonzeros, "dtype": dtype}
    for dtype in all_dtypes
    for zeros in [1, 2, 5]
    for nonzeros in [0, 3]))
  @jtu.skip_on_devices("gpu")
  def testRootsInvalid(self, zeros, nonzeros, dtype):
    raise unittest.SkipTest("getting segfaults on MKL")  # TODO(#3711)
    rng = jtu.rand_default(np.random.RandomState(0))

    # The polynomial coefficients here start with zero and would have to
    # be stripped before computing eigenvalues of the companion matrix.
    # Setting strip_zeros=False skips this check,
    # allowing jit transformation but yielding nan's for these inputs.
    p = jnp.concatenate([jnp.zeros(zeros, dtype), rng((nonzeros,), dtype)])

    if p.size == 1:
      # polynomial = const has no roots
      self.assertTrue(jnp.roots(p, strip_zeros=False).size == 0)
    else:
      self.assertTrue(jnp.any(jnp.isnan(jnp.roots(p, strip_zeros=False))))

if __name__ == "__main__":
  absltest.main(testLoader=jtu.JaxTestLoader())
<|endoftext|>


<commit_before>#!/usr/bin/env python
import argparse
import sys

errors = []
try:
  import matplotlib.pyplot as plt
except ImportError as ie:
    errors.append(ie)
try:
  import numpy as np
except ImportError as ie:
    errors.append(ie)
try:
  import pandas as pd
except ImportError as ie:
    errors.append(ie)
try:
  from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
except ImportError as ie:
    errors.append(ie)

if len(errors) != 0:
    for item in errors:
        print 'Import Error:', item
    sys.exit()

from phylotoast import util

def merge_dicts(*dict_args):
    """
    Given any number of dicts, shallow copy and merge into a new dict,
    precedence goes to key value pairs in latter dicts.
    """
    result = {}
    for dictionary in dict_args:
        result.update(dictionary)
    return result

def plot_LDA(X_lda, y_lda, class_colors, out_fp="", dpi=150, title=None):
    cats = class_colors.keys()
    group_lda = {c: [] for c in cats}

    for i, target_name in zip(range(len(cats)), cats):
        cat_x = X_lda[:, 0][y_lda == target_name]
        if X_lda.shape[1] == 1:
            cat_y = np.ones((cat_x.shape[0], 1)) + i
        else:
            cat_y = X_lda[:, 1][y_lda == target_name]
        group_lda[target_name].append(cat_x)
        group_lda[target_name].append(cat_y)
        plt.scatter(x=cat_x, y=cat_y, label=target_name,
                    color=class_colors[target_name],
                    alpha=0.85, s=80, edgecolors='k')

    if X_lda.shape[1] == 1:
        plt.ylim((0.5, 2.5))
    plt.xlabel('LD1')
    plt.ylabel('LD2')
    plt.legend(loc='best')
    if title:
        plt.title(title)

    # save or display result
    if out_fp:
        plt.savefig(out_fp, facecolor='white',
                    edgecolor='none', dpi=dpi,
                    bbox_inches='tight', pad_inches=0.2)
    else:
        plt.show()

def run_LDA(df):
    # Prep variables for sklearn LDA
    X = df[range(2, df.shape[1])].values     # input data matrix
    y = df['Condition'].values               # data categories list

    # Calculate LDA
    sklearn_lda = LDA(n_components=2)
    X_lda_sklearn = sklearn_lda.fit_transform(X, y)

    # Quality Test - can be ignored
#     print len(X_lda_sklearn)
#     print sklearn_lda.predict_proba(X)
#     print(sklearn_lda.score(X, y))

    return X_lda_sklearn, y

def handle_program_options():
    parser = argparse.ArgumentParser(description="Create an LDA plot from\
                                     sample-grouped OTU data. It is necessary\
                                     to remove the header cell '#OTU ID'\
                                     before running this program.")
    parser.add_argument('-i', '--biom_tsv', required=True,
                        help="Sample-OTU abundance table in TSV format with the\
                        arcsin sqrt transform already applied. [REQUIRED]")
    parser.add_argument('-m', '--map_fp', required=True,
                        help="Metadata mapping file. [REQUIRED]")
    parser.add_argument('-g', '--group_by', required=True, nargs='+',
                        help="Any mapping categories, such as treatment type, \
                              that will be used to group the data in the \
                              output iTol table. For example, one category \
                              with three types will result in three data \
                              columns in the final output. Two categories with\
                              three types each will result in six data \
                              columns. Default is no categories and all the \
                              data will be treated as a single group.")
    parser.add_argument('-c', '--color_by', required=True,
                        help="A column name in the mapping file containing\
                              hexadecimal (#FF0000) color values that will\
                              be used to color the groups. Each sample ID must\
                              have a color entry.")
    parser.add_argument('--dpi', default=200, type=int,
                        help="Set plot quality in Dots Per Inch (DPI). Larger\
                              DPI will result in larger file size.")
    parser.add_argument('--save_lda_input',
                        help="Save a CSV-format file of the transposed LDA-input\
                              table to the file specifed by this option.")
    parser.add_argument('--plot_title', default=None,
                        help="Plot title. Default is no title.")
    parser.add_argument('-o', '--out_fp', default="",
                        help="The path and file name to save the plot under.\
                              If specified, the figure will be saved directly\
                              instead of opening a window in which the plot \
                              can be viewed before saving")

    return parser.parse_args()

def main():
    args = handle_program_options()

    map_header, imap = util.parse_map_file(args.map_fp)

    df = pd.read_csv(args.biom_tsv, sep='\t', index_col=0).T
    # exclude Sample IDs not in the mapping file
    df = df.loc[imap.keys()]

    cat_gather = util.gather_categories(imap, map_header, args.group_by)
    if len(cat_gather) < 2:
        sys.stderr.write("ERROR: Only one category value found. Linear \
        Discriminant Analysis requires at least two categories to compare.")
        return

    color_gather = util.gather_categories(imap, map_header, [args.color_by])

    class_map = merge_dicts(*[{sid: cat for sid in cat_gather[cat].sids}
                              for cat in cat_gather])
    class_colors = merge_dicts(*[{class_map[sid]: color
                                  for sid in color_gather[color].sids}
                                 for color in color_gather])

    df.insert(0, "Condition", [class_map[entry] for entry in df.index])

    if args.save_lda_input:
        df.to_csv(args.save_lda_input)

    X_lda, y_lda = run_LDA(df)

    plot_LDA(X_lda, y_lda, class_colors, out_fp=args.out_fp, dpi=args.dpi,
             title=args.plot_title)

if __name__ == "__main__":
    main()
<commit_msg>Updated LDA.py. Added figsize parameter and added mpl import for setting figure font size. Increased point and font size for scatter plot. Updated axis labels. Deleted quality testing for LDA calculations.<commit_after>#!/usr/bin/env python
import argparse
import sys

errors = []
try:
  import matplotlib.pyplot as plt
  import matplotlib as mpl
except ImportError as ie:
    errors.append(ie)
try:
  import numpy as np
except ImportError as ie:
    errors.append(ie)
try:
  import pandas as pd
except ImportError as ie:
    errors.append(ie)
try:
  from sklearn.discriminant_analysis import LinearDiscriminantAnalysis as LDA
except ImportError as ie:
    errors.append(ie)

if len(errors) != 0:
    for item in errors:
        print "Import Error:", item
    sys.exit()

from phylotoast import util

def merge_dicts(*dict_args):
    """
    Given any number of dicts, shallow copy and merge into a new dict,
    precedence goes to key value pairs in latter dicts.
    """
    result = {}
    for dictionary in dict_args:
        result.update(dictionary)
    return result

def plot_LDA(X_lda, y_lda, class_colors, out_fp="", dpi=300, title=None):
    cats = class_colors.keys()
    group_lda = {c: [] for c in cats}
    fig = plt.figure(figsize=(15, 10))
    for i, target_name in zip(range(len(cats)), cats):
        cat_x = X_lda[:, 0][y_lda == target_name]
        if X_lda.shape[1] == 1:
            cat_y = np.ones((cat_x.shape[0], 1)) + i
        else:
            cat_y = X_lda[:, 1][y_lda == target_name]
        group_lda[target_name].append(cat_x)
        group_lda[target_name].append(cat_y)
        plt.scatter(x=cat_x, y=cat_y, label=target_name,
                    color=class_colors[target_name],
                    alpha=0.85, s=250,
                    edgecolors="k")
    font = {"size": 16}
    mpl.rc("font", **font)
    if X_lda.shape[1] == 1:
        plt.ylim((0.5, 2.5))
    plt.xlabel("Linear Discriminant 1")
    plt.ylabel("Linear Discriminant 2")
    plt.legend(loc="best")
    if title:
        plt.title(title)

    # save or display result
    if out_fp:
        plt.savefig(out_fp, facecolor="white",
                    edgecolor="none", dpi=dpi,
                    bbox_inches="tight", pad_inches=0.1)
    else:
        plt.show()

def run_LDA(df):
    # Prep variables for sklearn LDA
    X = df[range(2, df.shape[1])].values     # input data matrix
    y = df["Condition"].values               # data categories list

    # Calculate LDA
    sklearn_lda = LDA(n_components=2)
    X_lda_sklearn = sklearn_lda.fit_transform(X, y)

    return X_lda_sklearn, y

def handle_program_options():
    parser = argparse.ArgumentParser(description="Create an LDA plot from\
                                     sample-grouped OTU data. It is necessary\
                                     to remove the header cell '#OTU ID'\
                                     before running this program.")
    parser.add_argument("-i", "--biom_tsv", required=True,
                        help="Sample-OTU abundance table in TSV format with the\
                        arcsin sqrt transform already applied. [REQUIRED]")
    parser.add_argument("-m", "--map_fp", required=True,
                        help="Metadata mapping file. [REQUIRED]")
    parser.add_argument("-g", "--group_by", required=True, nargs="+",
                        help="Any mapping categories, such as treatment type, \
                              that will be used to group the data in the \
                              output iTol table. For example, one category \
                              with three types will result in three data \
                              columns in the final output. Two categories with\
                              three types each will result in six data \
                              columns. Default is no categories and all the \
                              data will be treated as a single group.")
    parser.add_argument("-c", "--color_by", required=True,
                        help="A column name in the mapping file containing\
                              hexadecimal (#FF0000) color values that will\
                              be used to color the groups. Each sample ID must\
                              have a color entry.")
    parser.add_argument("--dpi", default=300, type=int,
                        help="Set plot quality in Dots Per Inch (DPI). Larger\
                              DPI will result in larger file size. Default value\
                              is 300.")
    parser.add_argument("--save_lda_input",
                        help="Save a CSV-format file of the transposed LDA-input\
                              table to the file specifed by this option.")
    parser.add_argument("--plot_title", default=None,
                        help="Plot title. Default is no title.")
    parser.add_argument("-o", "--out_fp", default="",
                        help="The path and file name to save the plot under.\
                              If specified, the figure will be saved directly\
                              instead of opening a window in which the plot \
                              can be viewed before saving")

    return parser.parse_args()

def main():
    args = handle_program_options()

    map_header, imap = util.parse_map_file(args.map_fp)

    df = pd.read_csv(args.biom_tsv, sep="\t", index_col=0).T
    # exclude Sample IDs not in the mapping file
    df = df.loc[imap.keys()]

    cat_gather = util.gather_categories(imap, map_header, args.group_by)
    if len(cat_gather) < 2:
        sys.stderr.write("ERROR: Only one category value found. Linear \
        Discriminant Analysis requires at least two categories to compare.")
        return

    color_gather = util.gather_categories(imap, map_header, [args.color_by])

    class_map = merge_dicts(*[{sid: cat for sid in cat_gather[cat].sids}
                              for cat in cat_gather])
    class_colors = merge_dicts(*[{class_map[sid]: color
                                  for sid in color_gather[color].sids}
                                 for color in color_gather])

    df.insert(0, "Condition", [class_map[entry] for entry in df.index])

    if args.save_lda_input:
        df.to_csv(args.save_lda_input)

    X_lda, y_lda = run_LDA(df)

    plot_LDA(X_lda, y_lda, class_colors, out_fp=args.out_fp, dpi=args.dpi,
             title=args.plot_title)

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


<commit_before>#!/usr/bin/env python
# vim: set encoding=utf-8
# pylint: disable=wrong-import-position,wrong-import-order

"""
Main server program.

Configuration parameters:

    path.internal.malformed
    path.internal.static
    path.internal.templates
    path.log.main
    path.log.queries
"""

from __future__ import print_function

import sys
if sys.version_info[0] < 3:
    reload(sys)
    sys.setdefaultencoding('utf8')

import sys
import logging
import os
import requests
import jinja2
from flask import Flask, request, send_from_directory, redirect, Response

sys.path.append(os.path.abspath(os.path.join(__file__, "..", "..", "lib")))
from config import CONFIG
from limits import Limits
from cheat_wrapper import cheat_wrapper
from post import process_post_request
from options import parse_args

from stateful_queries import save_query, last_query

if not os.path.exists(os.path.dirname(CONFIG["path.log.main"])):
    os.makedirs(os.path.dirname(CONFIG["path.log.main"]))
logging.basicConfig(
    filename=CONFIG["path.log.main"],
    level=logging.DEBUG,
    format='%(asctime)s %(message)s')
# Fix Flask "exception and request logging" to `stderr`.
#
# When Flask's werkzeug detects that logging is already set, it
# doesn't add its own logger that prints exceptions.
stderr_handler = logging.StreamHandler()
logging.getLogger().addHandler(stderr_handler)
#
# Alter log format to disting log lines from everything else
stderr_handler.setFormatter(logging.Formatter('%(filename)s:%(lineno)s: %(message)s'))
#
# Sometimes werkzeug starts logging before an app is imported
# (https://github.com/pallets/werkzeug/issues/1969)
# resulting in duplicating lines. In that case we need root
# stderr handler to skip lines from werkzeug.
class SkipFlaskLogger(object):
    def filter(self, record):
        if record.name != 'werkzeug':
            return True
if logging.getLogger('werkzeug').handlers:
    stderr_handler.addFilter(SkipFlaskLogger())

app = Flask(__name__) # pylint: disable=invalid-name
app.jinja_loader = jinja2.ChoiceLoader([
    app.jinja_loader,
    jinja2.FileSystemLoader(CONFIG["path.internal.templates"])])

LIMITS = Limits()

def is_html_needed(user_agent):
    """
    Basing on `user_agent`, return whether it needs HTML or ANSI
    """
    plaintext_clients = [
        'curl', 'wget', 'fetch', 'httpie', 'lwp-request', 'openbsd ftp', 'python-requests']
    return all([x not in user_agent for x in plaintext_clients])

def is_result_a_script(query):
    return query in [':cht.sh']

@app.route('/files/<path:path>')
def send_static(path):
    """
    Return static file `path`.
    Can be served by the HTTP frontend.
    """
    return send_from_directory(CONFIG["path.internal.static"], path)

@app.route('/favicon.ico')
def send_favicon():
    """
    Return static file `favicon.ico`.
    Can be served by the HTTP frontend.
    """
    return send_from_directory(CONFIG["path.internal.static"], 'favicon.ico')

@app.route('/malformed-response.html')
def send_malformed():
    """
    Return static file `malformed-response.html`.
    Can be served by the HTTP frontend.
    """
    dirname, filename = os.path.split(CONFIG["path.internal.malformed"])
    return send_from_directory(dirname, filename)

def log_query(ip_addr, found, topic, user_agent):
    """
    Log processed query and some internal data
    """
    log_entry = "%s %s %s %s\n" % (ip_addr, found, topic, user_agent)
    with open(CONFIG["path.log.queries"], 'ab') as my_file:
        my_file.write(log_entry.encode('utf-8'))

def get_request_ip(req):
    """
    Extract IP address from `request`
    """

    if req.headers.getlist("X-Forwarded-For"):
        ip_addr = req.headers.getlist("X-Forwarded-For")[0]
        if ip_addr.startswith('::ffff:'):
            ip_addr = ip_addr[7:]
    else:
        ip_addr = req.remote_addr
    if req.headers.getlist("X-Forwarded-For"):
        ip_addr = req.headers.getlist("X-Forwarded-For")[0]
        if ip_addr.startswith('::ffff:'):
            ip_addr = ip_addr[7:]
    else:
        ip_addr = req.remote_addr

    return ip_addr

def get_answer_language(request):
    """
    Return preferred answer language based on
    domain name, query arguments and headers
    """

    def _parse_accept_language(accept_language):
        languages = accept_language.split(",")
        locale_q_pairs = []

        for language in languages:
            try:
                if language.split(";")[0] == language:
                    # no q => q = 1
                    locale_q_pairs.append((language.strip(), "1"))
                else:
                    locale = language.split(";")[0].strip()
                    weight = language.split(";")[1].split("=")[1]
                    locale_q_pairs.append((locale, weight))
            except IndexError:
                pass

        return locale_q_pairs

    def _find_supported_language(accepted_languages):
        for lang_tuple in accepted_languages:
            lang = lang_tuple[0]
            if '-' in lang:
                lang = lang.split('-', 1)[0]
            return lang
        return None

    lang = None
    hostname = request.headers['Host']
    if hostname.endswith('.cheat.sh'):
        lang = hostname[:-9]

    if 'lang' in request.args:
        lang = request.args.get('lang')

    header_accept_language = request.headers.get('Accept-Language', '')
    if lang is None and header_accept_language:
        lang = _find_supported_language(
            _parse_accept_language(header_accept_language))

    return lang

def _proxy(*args, **kwargs):
    # print "method=", request.method,
    # print "url=", request.url.replace('/:shell-x/', ':3000/')
    # print "headers=", {key: value for (key, value) in request.headers if key != 'Host'}
    # print "data=", request.get_data()
    # print "cookies=", request.cookies
    # print "allow_redirects=", False

    url_before, url_after = request.url.split('/:shell-x/', 1)
    url = url_before + ':3000/'

    if 'q' in request.args:
        url_after = '?' + "&".join("arg=%s" % x for x in request.args['q'].split())

    url += url_after
    print(url)
    print(request.get_data())
    resp = requests.request(
        method=request.method,
        url=url,
        headers={key: value for (key, value) in request.headers if key != 'Host'},
        data=request.get_data(),
        cookies=request.cookies,
        allow_redirects=False)

    excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
    headers = [(name, value) for (name, value) in resp.raw.headers.items()
               if name.lower() not in excluded_headers]

    response = Response(resp.content, resp.status_code, headers)
    return response

@app.route("/", methods=['GET', 'POST'])
@app.route("/<path:topic>", methods=["GET", "POST"])
def answer(topic=None):
    """
    Main rendering function, it processes incoming weather queries.
    Depending on user agent it returns output in HTML or ANSI format.

    Incoming data:
        request.args
        request.headers
        request.remote_addr
        request.referrer
        request.query_string
    """

    user_agent = request.headers.get('User-Agent', '').lower()
    html_needed = is_html_needed(user_agent)
    options = parse_args(request.args)

    if topic in ['apple-touch-icon-precomposed.png', 'apple-touch-icon.png', 'apple-touch-icon-120x120-precomposed.png'] \
        or (topic is not None and any(topic.endswith('/'+x) for x in ['favicon.ico'])):
        return ''

    request_id = request.cookies.get('id')
    if topic is not None and topic.lstrip('/') == ':last':
        if request_id:
            topic = last_query(request_id)
        else:
            return "ERROR: you have to set id for your requests to use /:last\n"
    else:
        if request_id:
            save_query(request_id, topic)

    if request.method == 'POST':
        process_post_request(request, html_needed)
        if html_needed:
            return redirect("/")
        return "OK\n"

    if 'topic' in request.args:
        return redirect("/%s" % request.args.get('topic'))

    if topic is None:
        topic = ":firstpage"

    if topic.startswith(':shell-x/'):
        return _proxy()
        #return requests.get('http://127.0.0.1:3000'+topic[8:]).text

    lang = get_answer_language(request)
    if lang:
        options['lang'] = lang

    ip_address = get_request_ip(request)
    if '+' in topic:
        not_allowed = LIMITS.check_ip(ip_address)
        if not_allowed:
            return "429 %s\n" % not_allowed, 429

    html_is_needed = is_html_needed(user_agent) and not is_result_a_script(topic)
    if html_is_needed:
        output_format='html'
    else:
        output_format='ansi'
    result, found = cheat_wrapper(topic, request_options=options, output_format=output_format)
    if 'Please come back in several hours' in result and html_is_needed:
        malformed_response = open(os.path.join(CONFIG["path.internal.malformed"])).read()
        return malformed_response

    log_query(ip_address, found, topic, user_agent)
    if html_is_needed:
        return result
    return Response(result, mimetype='text/plain')
<commit_msg>Add powershell to PLAIN_TEXT_AGENTS (#273)<commit_after>#!/usr/bin/env python
# vim: set encoding=utf-8
# pylint: disable=wrong-import-position,wrong-import-order

"""
Main server program.

Configuration parameters:

    path.internal.malformed
    path.internal.static
    path.internal.templates
    path.log.main
    path.log.queries
"""

from __future__ import print_function

import sys
if sys.version_info[0] < 3:
    reload(sys)
    sys.setdefaultencoding('utf8')

import sys
import logging
import os
import requests
import jinja2
from flask import Flask, request, send_from_directory, redirect, Response

sys.path.append(os.path.abspath(os.path.join(__file__, "..", "..", "lib")))
from config import CONFIG
from limits import Limits
from cheat_wrapper import cheat_wrapper
from post import process_post_request
from options import parse_args

from stateful_queries import save_query, last_query

if not os.path.exists(os.path.dirname(CONFIG["path.log.main"])):
    os.makedirs(os.path.dirname(CONFIG["path.log.main"]))
logging.basicConfig(
    filename=CONFIG["path.log.main"],
    level=logging.DEBUG,
    format='%(asctime)s %(message)s')
# Fix Flask "exception and request logging" to `stderr`.
#
# When Flask's werkzeug detects that logging is already set, it
# doesn't add its own logger that prints exceptions.
stderr_handler = logging.StreamHandler()
logging.getLogger().addHandler(stderr_handler)
#
# Alter log format to disting log lines from everything else
stderr_handler.setFormatter(logging.Formatter('%(filename)s:%(lineno)s: %(message)s'))
#
# Sometimes werkzeug starts logging before an app is imported
# (https://github.com/pallets/werkzeug/issues/1969)
# resulting in duplicating lines. In that case we need root
# stderr handler to skip lines from werkzeug.
class SkipFlaskLogger(object):
    def filter(self, record):
        if record.name != 'werkzeug':
            return True
if logging.getLogger('werkzeug').handlers:
    stderr_handler.addFilter(SkipFlaskLogger())

app = Flask(__name__) # pylint: disable=invalid-name
app.jinja_loader = jinja2.ChoiceLoader([
    app.jinja_loader,
    jinja2.FileSystemLoader(CONFIG["path.internal.templates"])])

LIMITS = Limits()

PLAIN_TEXT_AGENTS = [
    "curl",
    "httpie",
    "lwp-request",
    "wget",
    "python-requests",
    "openbsd ftp",
    "powershell",
    "fetch",
]

def _is_html_needed(user_agent):
    """
    Basing on `user_agent`, return whether it needs HTML or ANSI
    """
    return all([x not in user_agent for x in PLAIN_TEXT_AGENTS])

def is_result_a_script(query):
    return query in [':cht.sh']

@app.route('/files/<path:path>')
def send_static(path):
    """
    Return static file `path`.
    Can be served by the HTTP frontend.
    """
    return send_from_directory(CONFIG["path.internal.static"], path)

@app.route('/favicon.ico')
def send_favicon():
    """
    Return static file `favicon.ico`.
    Can be served by the HTTP frontend.
    """
    return send_from_directory(CONFIG["path.internal.static"], 'favicon.ico')

@app.route('/malformed-response.html')
def send_malformed():
    """
    Return static file `malformed-response.html`.
    Can be served by the HTTP frontend.
    """
    dirname, filename = os.path.split(CONFIG["path.internal.malformed"])
    return send_from_directory(dirname, filename)

def log_query(ip_addr, found, topic, user_agent):
    """
    Log processed query and some internal data
    """
    log_entry = "%s %s %s %s\n" % (ip_addr, found, topic, user_agent)
    with open(CONFIG["path.log.queries"], 'ab') as my_file:
        my_file.write(log_entry.encode('utf-8'))

def get_request_ip(req):
    """
    Extract IP address from `request`
    """

    if req.headers.getlist("X-Forwarded-For"):
        ip_addr = req.headers.getlist("X-Forwarded-For")[0]
        if ip_addr.startswith('::ffff:'):
            ip_addr = ip_addr[7:]
    else:
        ip_addr = req.remote_addr
    if req.headers.getlist("X-Forwarded-For"):
        ip_addr = req.headers.getlist("X-Forwarded-For")[0]
        if ip_addr.startswith('::ffff:'):
            ip_addr = ip_addr[7:]
    else:
        ip_addr = req.remote_addr

    return ip_addr

def get_answer_language(request):
    """
    Return preferred answer language based on
    domain name, query arguments and headers
    """

    def _parse_accept_language(accept_language):
        languages = accept_language.split(",")
        locale_q_pairs = []

        for language in languages:
            try:
                if language.split(";")[0] == language:
                    # no q => q = 1
                    locale_q_pairs.append((language.strip(), "1"))
                else:
                    locale = language.split(";")[0].strip()
                    weight = language.split(";")[1].split("=")[1]
                    locale_q_pairs.append((locale, weight))
            except IndexError:
                pass

        return locale_q_pairs

    def _find_supported_language(accepted_languages):
        for lang_tuple in accepted_languages:
            lang = lang_tuple[0]
            if '-' in lang:
                lang = lang.split('-', 1)[0]
            return lang
        return None

    lang = None
    hostname = request.headers['Host']
    if hostname.endswith('.cheat.sh'):
        lang = hostname[:-9]

    if 'lang' in request.args:
        lang = request.args.get('lang')

    header_accept_language = request.headers.get('Accept-Language', '')
    if lang is None and header_accept_language:
        lang = _find_supported_language(
            _parse_accept_language(header_accept_language))

    return lang

def _proxy(*args, **kwargs):
    # print "method=", request.method,
    # print "url=", request.url.replace('/:shell-x/', ':3000/')
    # print "headers=", {key: value for (key, value) in request.headers if key != 'Host'}
    # print "data=", request.get_data()
    # print "cookies=", request.cookies
    # print "allow_redirects=", False

    url_before, url_after = request.url.split('/:shell-x/', 1)
    url = url_before + ':3000/'

    if 'q' in request.args:
        url_after = '?' + "&".join("arg=%s" % x for x in request.args['q'].split())

    url += url_after
    print(url)
    print(request.get_data())
    resp = requests.request(
        method=request.method,
        url=url,
        headers={key: value for (key, value) in request.headers if key != 'Host'},
        data=request.get_data(),
        cookies=request.cookies,
        allow_redirects=False)

    excluded_headers = ['content-encoding', 'content-length', 'transfer-encoding', 'connection']
    headers = [(name, value) for (name, value) in resp.raw.headers.items()
               if name.lower() not in excluded_headers]

    response = Response(resp.content, resp.status_code, headers)
    return response

@app.route("/", methods=['GET', 'POST'])
@app.route("/<path:topic>", methods=["GET", "POST"])
def answer(topic=None):
    """
    Main rendering function, it processes incoming weather queries.
    Depending on user agent it returns output in HTML or ANSI format.

    Incoming data:
        request.args
        request.headers
        request.remote_addr
        request.referrer
        request.query_string
    """

    user_agent = request.headers.get('User-Agent', '').lower()
    html_needed = _is_html_needed(user_agent)
    options = parse_args(request.args)

    if topic in ['apple-touch-icon-precomposed.png', 'apple-touch-icon.png', 'apple-touch-icon-120x120-precomposed.png'] \
        or (topic is not None and any(topic.endswith('/'+x) for x in ['favicon.ico'])):
        return ''

    request_id = request.cookies.get('id')
    if topic is not None and topic.lstrip('/') == ':last':
        if request_id:
            topic = last_query(request_id)
        else:
            return "ERROR: you have to set id for your requests to use /:last\n"
    else:
        if request_id:
            save_query(request_id, topic)

    if request.method == 'POST':
        process_post_request(request, html_needed)
        if html_needed:
            return redirect("/")
        return "OK\n"

    if 'topic' in request.args:
        return redirect("/%s" % request.args.get('topic'))

    if topic is None:
        topic = ":firstpage"

    if topic.startswith(':shell-x/'):
        return _proxy()
        #return requests.get('http://127.0.0.1:3000'+topic[8:]).text

    lang = get_answer_language(request)
    if lang:
        options['lang'] = lang

    ip_address = get_request_ip(request)
    if '+' in topic:
        not_allowed = LIMITS.check_ip(ip_address)
        if not_allowed:
            return "429 %s\n" % not_allowed, 429

    html_is_needed = _is_html_needed(user_agent) and not is_result_a_script(topic)
    if html_is_needed:
        output_format='html'
    else:
        output_format='ansi'
    result, found = cheat_wrapper(topic, request_options=options, output_format=output_format)
    if 'Please come back in several hours' in result and html_is_needed:
        malformed_response = open(os.path.join(CONFIG["path.internal.malformed"])).read()
        return malformed_response

    log_query(ip_address, found, topic, user_agent)
    if html_is_needed:
        return result
    return Response(result, mimetype='text/plain')
<|endoftext|>


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

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

import rospy
import roslib; roslib.load_manifest('rsync_ros')
import actionlib
from rsync import Rsync
from rsync_ros.msg import RsyncAction, RsyncResult, RsyncFeedback

class RsyncActionServer:

    def __init__(self, name):
        self._action_name = name
        self.server = actionlib.SimpleActionServer(self._action_name, RsyncAction, self.execute, False)
        self.server.start()
        rospy.loginfo("Ready to sync files.")

    def progress_update_cb(self, line, percent_complete, transfer_rate):
        #This is run everytime the progress is published to stdout
        #rospy.loginfo('Total transfer percentage: {}'.format(percent_complete))

        self.feedback.percent_complete = percent_complete
        self.feedback.transfer_rate = transfer_rate
        self.server.publish_feedback(self.feedback)

        if line:
            rospy.loginfo(line)

        # check if preempt (cancel action) has been requested by the client
        if self.server.is_preempt_requested():
            rospy.loginfo('%s: Preempted' % self._action_name)
            self.rsync = None #Delete the instance of the Rsync class
            self.server.set_preempted() #TO-DO, fix logic error changing states upon preempt request

    def execute(self, goal):
        self.result = RsyncResult()
        self.feedback = RsyncFeedback()

        rospy.loginfo("Executing rsync command '%s %s %s'", 'rsync ' + ' '.join(goal.rsync_args) + ' --progress --outbuf=L', goal.source_path, goal.destination_path)

        self.rsync = Rsync(goal.rsync_args, goal.source_path, goal.destination_path, progress_callback=self.progress_update_cb)

        self.result.sync_success = self.rsync.sync()

        if not self.server.is_preempt_requested():

            if self.rsync.stderr_block:
                rospy.logerr('\n{}'.format(self.rsync.stderr_block))

            rospy.loginfo("Rsync command result '%s'", self.result.sync_success)
            self.server.set_succeeded(self.result)

if __name__ == "__main__":
    try:
        rospy.init_node('rsync_ros')
        RsyncActionServer(rospy.get_name())
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
<commit_msg>Fixed preempting by calling a sigterm to the rsync process. If that fails; call sigkill with the rsync pid<commit_after>#!/usr/bin/env python

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

import rospy
import roslib; roslib.load_manifest('rsync_ros')
import actionlib
import os
from rsync import Rsync
from rsync_ros.msg import RsyncAction, RsyncResult, RsyncFeedback

class RsyncActionServer:

    def __init__(self, name):
        self._action_name = name
        self.server = actionlib.SimpleActionServer(self._action_name, RsyncAction, self.execute, False)
        self.server.start()
        rospy.loginfo("Ready to sync files.")

    def progress_update_cb(self, line, percent_complete, transfer_rate):
        #This is run everytime the progress is published to stdout
        #rospy.loginfo('Total transfer percentage: {}'.format(percent_complete))

        self.feedback.percent_complete = percent_complete
        self.feedback.transfer_rate = transfer_rate
        self.server.publish_feedback(self.feedback)

        if line:
            rospy.loginfo(line)

        # check if preempt (cancel action) has been requested by the client
        if self.server.is_preempt_requested():
            # Get the process id & try to terminate it gracefuly
            pid = self.rsync.p.pid
            self.rsync.p.terminate()

            # Check if the process has really terminated & force kill if not.
            try:
                os.kill(pid, 0)
                self.rsync.p.kill()
                print "Forced kill"
            except OSError, e:
                print "Terminated gracefully"

            rospy.loginfo('%s: Preempted' % self._action_name)
            self.server.set_preempted() #TO-DO, fix logic error changing states upon preempt request

    def execute(self, goal):
        self.result = RsyncResult()
        self.feedback = RsyncFeedback()

        rospy.loginfo("Executing rsync command '%s %s %s'", 'rsync ' + ' '.join(goal.rsync_args) + ' --progress --outbuf=L', goal.source_path, goal.destination_path)

        self.rsync = Rsync(goal.rsync_args, goal.source_path, goal.destination_path, progress_callback=self.progress_update_cb)

        self.result.sync_success = self.rsync.sync()

        if not self.server.is_preempt_requested():

            if self.rsync.stderr_block:
                rospy.logerr('\n{}'.format(self.rsync.stderr_block))

            rospy.loginfo("Rsync command result '%s'", self.result.sync_success)
            self.server.set_succeeded(self.result)

if __name__ == "__main__":
    try:
        rospy.init_node('rsync_ros')
        RsyncActionServer(rospy.get_name())
        rospy.spin()
    except rospy.ROSInterruptException:
        pass
<|endoftext|>


<commit_before>import json
import StringIO

import geojson
import requests
from shapely.wkt import loads

from girder.api import access
from girder.api.describe import Description
from girder.api.rest import Resource
from girder.utility.model_importer import ModelImporter
from girder.plugins.minerva.rest.geojson_dataset import GeojsonDataset
from girder.plugins.minerva.utility.minerva_utility import findDatasetFolder

class TwoFishes(Resource):
    def __init__(self):
        self.resourceName = 'minerva_geocoder'
        self.route('GET', ('autocomplete',), self.autocomplete)
        self.route('GET', ('geojson',), self.getGeojson)
        self.route('POST', ('geojson',), self.postGeojson)

    @staticmethod
    def getWktFromTwoFishes(twofishes, location):
        """Gets wkt from twofishes for a given location"""
        r = requests.get(twofishes,
                         params={'query': location,
                                 'responseIncludes': 'WKT_GEOMETRY'})
        wkt = r.json()['interpretations'][0]['feature']['geometry']['wktGeometry']

        return wkt

    @staticmethod
    def createGeometryFromWkt(wkt):
        """Creates a shapely geometry from wkt"""
        return loads(wkt)

    @staticmethod
    def createGeojson(twofishes, locations):
        """Create geojson for given locations and twofishes url"""
        geoms = []

        for i in locations:
            wkt = TwoFishes.getWktFromTwoFishes(twofishes, i)
            geom = TwoFishes.createGeometryFromWkt(wkt)
            for g in geom:
                geoms.append(geojson.Feature(geometry=g,
                                             properties={'location': i}))

        multiPoly = geojson.FeatureCollection(geoms)

        return multiPoly

    def createMinervaDataset(self, geojsonString, name):
        """Creates a dataset from a geojson string"""
        output = StringIO.StringIO(json.dumps(geojsonString))
        outputSize = output.len
        user = self.getCurrentUser()
        datasetFolder = findDatasetFolder(user, user)
        itemModel = ModelImporter.model('item')
        uploadModel = ModelImporter.model('upload')
        item = itemModel.createItem(name, user, datasetFolder)
        geojsonFile = uploadModel.uploadFromFile(output, outputSize, name,
                                                 'item', item, user)
        GeojsonDataset().createGeojsonDataset(itemId=geojsonFile['itemId'],
                                              params={})
        return geojsonFile

    @access.public
    def autocomplete(self, params):
        r = requests.get(params['twofishes'],
                         params={'autocomplete': True,
                                 'query': params['location'],
                                 'maxInterpretations': 10,
                                 'autocompleteBias': None})

        return [i['feature']['matchedName'] for i in r.json()['interpretations']]

    autocomplete.description = (
        Description('Autocomplete result for a given location name')
        .param('twofishes', 'Twofishes url')
        .param('location', 'Location name to autocomplete')
    )

    @access.public
    def getGeojson(self, params):
        locations = json.loads(params['locations'])
        geojson = TwoFishes.createGeojson(params['twofishes'], locations)
        return geojson

    getGeojson.description = (
        Description('Create a geojson string from multiple locations')
        .param('twofishes', 'Twofishes url')
        .param('locations', 'List of locations', dataType='list')
    )

    @access.public
    def postGeojson(self, params):
        twofishes = params['twofishes']
        try:
            locationInfo = json.loads(params['locations'])
            geojson = TwoFishes.createGeojson(twofishes, locationInfo)
        except ValueError:
            locationInfo = params['locations']
            geojson = TwoFishes.createGeojson(twofishes, locationInfo)

        minervaDataset = self.createMinervaDataset(geojson, params['name'])
        return minervaDataset

    postGeojson.description = (
        Description('Create a minerva dataset from the search result/results')
        .param('twofishes', 'Twofishes url')
        .param('locations', 'Location name or list of locations to get a geojson')
        .param('name', 'Name for the geojson dataset')
    )
<commit_msg>Improve type checking and add docstring<commit_after>#!/usr/bin/env python
# -*- coding: utf-8 -*-

###############################################################################
#  Copyright Kitware Inc.
#
#  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 ast
import json
import StringIO

import geojson
import requests
from shapely.wkt import loads

from girder.api import access
from girder.api.describe import Description
from girder.api.rest import Resource
from girder.utility.model_importer import ModelImporter
from girder.plugins.minerva.rest.geojson_dataset import GeojsonDataset
from girder.plugins.minerva.utility.minerva_utility import findDatasetFolder

class TwoFishes(Resource):
    """Resource that handles geocoding related operations"""
    def __init__(self):
        self.resourceName = 'minerva_geocoder'
        self.route('GET', ('autocomplete',), self.autocomplete)
        self.route('GET', ('geojson',), self.getGeojson)
        self.route('POST', ('geojson',), self.postGeojson)

    @staticmethod
    def getWktFromTwoFishes(twofishes, location):
        """Gets wkt from twofishes for a given location"""
        r = requests.get(twofishes,
                         params={'query': location,
                                 'responseIncludes': 'WKT_GEOMETRY'})
        wkt = r.json()['interpretations'][0]['feature']['geometry']['wktGeometry']

        return wkt

    @staticmethod
    def createGeometryFromWkt(wkt):
        """Creates a shapely geometry from wkt"""
        return loads(wkt)

    @staticmethod
    def createGeojson(twofishes, locations):
        """Create geojson for given locations and twofishes url"""

        try:
            locations = ast.literal_eval(locations)
        except ValueError:
            pass

        geoms = []

        for i in locations:
            wkt = TwoFishes.getWktFromTwoFishes(twofishes, i)
            geom = TwoFishes.createGeometryFromWkt(wkt)
            for g in geom:
                geoms.append(geojson.Feature(geometry=g,
                                             properties={'location': i}))

        multiPoly = geojson.FeatureCollection(geoms)

        return multiPoly

    def createMinervaDataset(self, geojsonString, name):
        """Creates a dataset from a geojson string"""
        output = StringIO.StringIO(json.dumps(geojsonString))
        outputSize = output.len
        user = self.getCurrentUser()
        datasetFolder = findDatasetFolder(user, user, create=True)
        itemModel = ModelImporter.model('item')
        uploadModel = ModelImporter.model('upload')
        item = itemModel.createItem(name, user, datasetFolder)
        geojsonFile = uploadModel.uploadFromFile(output, outputSize, name,
                                                 'item', item, user)
        GeojsonDataset().createGeojsonDataset(itemId=geojsonFile['itemId'],
                                              params={})
        return geojsonFile

    @access.public
    def autocomplete(self, params):
        r = requests.get(params['twofishes'],
                         params={'autocomplete': True,
                                 'query': params['location'],
                                 'maxInterpretations': 10,
                                 'autocompleteBias': None})

        return [i['feature']['matchedName'] for i in r.json()['interpretations']]

    autocomplete.description = (
        Description('Autocomplete result for a given location name')
        .param('twofishes', 'Twofishes url')
        .param('location', 'Location name to autocomplete')
    )

    @access.public
    def getGeojson(self, params):
        locations = json.loads(params['locations'])
        geojson = TwoFishes.createGeojson(params['twofishes'], locations)
        return geojson

    getGeojson.description = (
        Description('Create a geojson string from multiple locations')
        .param('twofishes', 'Twofishes url')
        .param('locations', 'List of locations', dataType='list')
    )

    @access.public
    def postGeojson(self, params):
        twofishes = params['twofishes']
        try:
            locationInfo = json.loads(params['locations'])
            geojson = TwoFishes.createGeojson(twofishes, locationInfo)
        except ValueError:
            locationInfo = params['locations']
            geojson = TwoFishes.createGeojson(twofishes, locationInfo)

        minervaDataset = self.createMinervaDataset(geojson, params['name'])
        return minervaDataset

    postGeojson.description = (
        Description('Create a minerva dataset from the search result/results')
        .param('twofishes', 'Twofishes url')
        .param('locations', 'Location name or list of locations to get a geojson')
        .param('name', 'Name for the geojson dataset')
    )
<|endoftext|>


<commit_before>from samplics.datasets.datasets import (
    Birth,
    CountyCrop,
    CountyCropMeans,
    Nhanes2,
    Nhanes2brr,
    Nhanes2jk,
    Nmihs,
    PSUFrame,
    PSUSample,
    SSUSample,
)
from samplics.estimation import ReplicateEstimator, TaylorEstimator
from samplics.sae import EblupAreaModel, EblupUnitModel, EbUnitModel, EllUnitModel
from samplics.sampling import OneMeanSampleSize, SampleSelection, SampleSize, allocate
from samplics.utils.basic_functions import transform
from samplics.utils.formats import array_to_dict
from samplics.weighting import ReplicateWeight, SampleWeight

__all__ = [
    "allocate",
    "array_to_dict",
    "Birth",
    "CountyCrop",
    "CountyCropMeans",
    "EblupAreaModel",
    "EblupUnitModel",
    "EbUnitModel",
    "EllUnitModel",
    "Nhanes2",
    "Nhanes2brr",
    "Nmihs",
    "OneMeanSampleSize",
    "PSUFrame",
    "PSUSample",
    "SampleSelection",
    "SampleSize",
    "SampleWeight",
    "SSUSample",
    "ReplicateWeight",
    "ReplicateEstimator",
    "TaylorEstimator",
    "transform",
]

__version__ = "0.3.2"
<commit_msg>adding datasets module<commit_after>from samplics.datasets.datasets import (
    Birth,
    CountyCrop,
    CountyCropMeans,
    ExpenditureMilk,
    Nhanes2,
    Nhanes2brr,
    Nhanes2jk,
    Nmihs,
    PSUFrame,
    PSUSample,
    SSUSample,
)
from samplics.estimation import ReplicateEstimator, TaylorEstimator
from samplics.sae import EblupAreaModel, EblupUnitModel, EbUnitModel, EllUnitModel
from samplics.sampling import OneMeanSampleSize, SampleSelection, SampleSize, allocate
from samplics.utils.basic_functions import transform
from samplics.utils.formats import array_to_dict
from samplics.weighting import ReplicateWeight, SampleWeight

__all__ = [
    "allocate",
    "array_to_dict",
    "Birth",
    "CountyCrop",
    "CountyCropMeans",
    "EblupAreaModel",
    "EblupUnitModel",
    "EbUnitModel",
    "EllUnitModel",
    "ExpenditureMilk",
    "Nhanes2",
    "Nhanes2brr",
    "Nmihs",
    "OneMeanSampleSize",
    "PSUFrame",
    "PSUSample",
    "SampleSelection",
    "SampleSize",
    "SampleWeight",
    "SSUSample",
    "ReplicateWeight",
    "ReplicateEstimator",
    "TaylorEstimator",
    "transform",
]

__version__ = "0.3.2"
<|endoftext|>