<commit_before>import logging
import os
import yaml

import pytest

from kubeflow.testing import util
from kubeflow.kfctl.testing.util import deploy_utils
from kubeflow.kfctl.testing.util import aws_util as kfctl_aws_util

def set_logging():
    logging.basicConfig(level=logging.INFO,
                        format=('%(levelname)s|%(asctime)s'
                                '|%(pathname)s|%(lineno)d| %(message)s'),
                        datefmt='%Y-%m-%dT%H:%M:%S',
                        )
    logging.getLogger().setLevel(logging.INFO)

def get_platform_app_name(app_path):
    with open(os.path.join(app_path, "tmp.yaml")) as f:
        kfdef = yaml.safe_load(f)
    app_name = kfdef["metadata"]["name"]
    platform = ""
    apiVersion = kfdef["apiVersion"].strip().split("/")
    if len(apiVersion) != 2:
        raise RuntimeError("Invalid apiVersion: " + kfdef["apiVersion"].strip())
    if apiVersion[-1] == "v1alpha1":
        platform = kfdef["spec"]["platform"]
    elif apiVersion[-1] in ["v1beta1", "v1"]:
        for plugin in kfdef["spec"].get("plugins", []):
            if plugin.get("kind", "") == "KfGcpPlugin":
                platform = "gcp"
            elif plugin.get("kind", "") == "KfAwsPlugin":
                platform = "aws"
            elif plugin.get("kind", "") == "KfExistingArriktoPlugin":
                platform = "existing_arrikto"
            else:
                # Indicate agnostic Kubeflow Platform
                platform = "agnostic"
    else:
        raise RuntimeError("Unknown version: " + apiVersion[-1])
    return platform, app_name

def check_deployments_ready(record_xml_attribute, namespace, name, deployments, cluster_name):
    """Test that Kubeflow deployments are successfully deployed.

  Args:
    namespace: The namespace Kubeflow is deployed to.
  """
    set_logging()
    util.set_pytest_junit(record_xml_attribute, name)

    kfctl_aws_util.aws_auth_load_kubeconfig(cluster_name)

    api_client = deploy_utils.create_k8s_client()

    for deployment_name in deployments:
        logging.info("Verifying that deployment %s started...", deployment_name)
        util.wait_for_deployment(api_client, namespace, deployment_name, 10)

def test_admission_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "admission-webhook-deployment"
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_admission_is_ready", deployment_names,
                            cluster_name)

def test_katib_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "katib-controller",
        "katib-mysql",
        "katib-db-manager",
        "katib-ui",
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_katib_is_ready", deployment_names,
                            cluster_name)

def test_metadata_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "metadata-deployment",
        "metadata-grpc-deployment",
        "metadata-db",
        "metadata-ui",
        "metadata-envoy-deployment",
        "metadata-writer",
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_metadata_is_ready", deployment_names,
                            cluster_name)

def test_pipeline_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "argo-ui",
        "cache-deployer-deployment",
        "cache-server",
        "kubeflow-pipelines-profile-controller",
        "minio",
        "ml-pipeline",
        "ml-pipeline-persistenceagent",
        "ml-pipeline-scheduledworkflow",
        "ml-pipeline-ui",
        "ml-pipeline-viewer-crd",
        "ml-pipeline-visualizationserver",
        "mysql",
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_pipeline_is_ready", deployment_names,
                            cluster_name)

def test_notebook_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "jupyter-web-app-deployment",
        "notebook-controller-deployment",
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_notebook_is_ready", deployment_names,
                            cluster_name)

def test_centraldashboard_is_ready(record_xml_attribute, namespace, cluster_name):
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_centraldashboard_is_ready", ["centraldashboard"],
                            cluster_name)

def test_profiles_is_ready(record_xml_attribute, namespace, cluster_name):
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_profile_is_ready", ["profiles-deployment"],
                            cluster_name)

def test_seldon_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "seldon-controller-manager"
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_seldon_is_ready", deployment_names,
                            cluster_name)

def test_spark_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "spark-operatorsparkoperator"
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_spark_is_ready", deployment_names,
                            cluster_name)

def test_training_operators_are_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "mpi-operator",
        "mxnet-operator",
        "pytorch-operator",
        "tf-job-operator",
    ]

    check_deployments_ready(record_xml_attribute, namespace,
                            "test_training_is_ready", deployment_names,
                            cluster_name)

def test_workflow_controller_is_ready(record_xml_attribute, namespace, cluster_name):
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_workflow_controller_is_ready", ["workflow-controller"],
                            cluster_name)

def test_kf_is_ready(record_xml_attribute, namespace, use_basic_auth,
                     app_path, cluster_name):
    """Test that Kubeflow was successfully deployed.

  Args:
    namespace: The namespace Kubeflow is deployed to.
  """
    set_logging()
    util.set_pytest_junit(record_xml_attribute, "test_kf_is_ready")

    kfctl_aws_util.aws_auth_load_kubeconfig(cluster_name)

    api_client = deploy_utils.create_k8s_client()

    # Verify that components are actually deployed.
    deployment_names = []

    stateful_set_names = []

    platform, _ = get_platform_app_name(app_path)

    # TODO(PatrickXYS): not sure why istio-galley can't found
    ingress_related_deployments = [
        "cluster-local-gateway",
        "istio-citadel",
        "istio-ingressgateway",
        "istio-pilot",
        "istio-policy",
        "istio-sidecar-injector",
        "istio-telemetry",
        "prometheus",
    ]
    ingress_related_stateful_sets = []

    knative_namespace = "knative-serving"
    knative_related_deployments = [
        "activator",
        "autoscaler",
        "autoscaler-hpa",
        "controller",
        "networking-istio",
        "webhook",
    ]

    if platform == "gcp":
        deployment_names.extend(["cloud-endpoints-controller"])
        stateful_set_names.extend(["kfserving-controller-manager"])
        if use_basic_auth:
            deployment_names.extend(["basic-auth-login"])
            ingress_related_stateful_sets.extend(["backend-updater"])
        else:
            ingress_related_deployments.extend(["iap-enabler"])
            ingress_related_stateful_sets.extend(["backend-updater"])
    elif platform == "existing_arrikto":
        deployment_names.extend(["dex"])
        ingress_related_deployments.extend(["authservice"])
        knative_related_deployments = []
    elif platform == "aws":
        # TODO(PatrickXYS): Extend List with AWS Deployment
        deployment_names.extend(["alb-ingress-controller"])

    # TODO(jlewi): Might want to parallelize this.
    for deployment_name in deployment_names:
        logging.info("Verifying that deployment %s started...", deployment_name)
        util.wait_for_deployment(api_client, namespace, deployment_name, 10)

    ingress_namespace = "istio-system"
    for deployment_name in ingress_related_deployments:
        logging.info("Verifying that deployment %s started...", deployment_name)
        util.wait_for_deployment(api_client, ingress_namespace, deployment_name, 10)

    all_stateful_sets = [(namespace, name) for name in stateful_set_names]
    all_stateful_sets.extend([(ingress_namespace, name) for name in ingress_related_stateful_sets])

    for ss_namespace, name in all_stateful_sets:
        logging.info("Verifying that stateful set %s.%s started...", ss_namespace, name)
        try:
            util.wait_for_statefulset(api_client, ss_namespace, name)
        except:
            # Collect debug information by running describe
            util.run(["kubectl", "-n", ss_namespace, "describe", "statefulsets", name])
            raise

    ingress_names = ["istio-ingress"]
    # Check if Ingress is Ready and Healthy
    if platform in ["aws"]:
        for ingress_name in ingress_names:
            logging.info("Verifying that ingress %s started...", ingress_name)
            util.wait_for_ingress(api_client, ingress_namespace, ingress_name, 10)

    for deployment_name in knative_related_deployments:
        logging.info("Verifying that deployment %s started...", deployment_name)
        util.wait_for_deployment(api_client, knative_namespace, deployment_name, 10)

    # Check if Dex is Ready and Healthy
    dex_deployment_names = ["dex"]
    dex_namespace = "auth"
    for dex_deployment_name in dex_deployment_names:
        logging.info("Verifying that deployment %s started...", dex_deployment_name)
        util.wait_for_deployment(api_client, dex_namespace, dex_deployment_name, 10)

    # Check if Cert-Manager is Ready and Healthy
    cert_manager_deployment_names = [
        "cert-manager",
        "cert-manager-cainjector",
        "cert-manager-webhook",
    ]
    cert_manager_namespace = "cert-manager"
    for cert_manager_deployment_name in cert_manager_deployment_names:
        logging.info("Verifying that deployment %s started...", cert_manager_deployment_name)
        util.wait_for_deployment(api_client, cert_manager_namespace, cert_manager_deployment_name, 10)

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO,
                        format=('%(levelname)s|%(asctime)s'
                                '|%(pathname)s|%(lineno)d| %(message)s'),
                        datefmt='%Y-%m-%dT%H:%M:%S',
                        )
    logging.getLogger().setLevel(logging.INFO)
    pytest.main()<commit_msg>Remove autoscaler-hpa knative deployment check (#439)<commit_after>import logging
import os
import yaml

import pytest

from kubeflow.testing import util
from kubeflow.kfctl.testing.util import deploy_utils
from kubeflow.kfctl.testing.util import aws_util as kfctl_aws_util

def set_logging():
    logging.basicConfig(level=logging.INFO,
                        format=('%(levelname)s|%(asctime)s'
                                '|%(pathname)s|%(lineno)d| %(message)s'),
                        datefmt='%Y-%m-%dT%H:%M:%S',
                        )
    logging.getLogger().setLevel(logging.INFO)

def get_platform_app_name(app_path):
    with open(os.path.join(app_path, "tmp.yaml")) as f:
        kfdef = yaml.safe_load(f)
    app_name = kfdef["metadata"]["name"]
    platform = ""
    apiVersion = kfdef["apiVersion"].strip().split("/")
    if len(apiVersion) != 2:
        raise RuntimeError("Invalid apiVersion: " + kfdef["apiVersion"].strip())
    if apiVersion[-1] == "v1alpha1":
        platform = kfdef["spec"]["platform"]
    elif apiVersion[-1] in ["v1beta1", "v1"]:
        for plugin in kfdef["spec"].get("plugins", []):
            if plugin.get("kind", "") == "KfGcpPlugin":
                platform = "gcp"
            elif plugin.get("kind", "") == "KfAwsPlugin":
                platform = "aws"
            elif plugin.get("kind", "") == "KfExistingArriktoPlugin":
                platform = "existing_arrikto"
            else:
                # Indicate agnostic Kubeflow Platform
                platform = "agnostic"
    else:
        raise RuntimeError("Unknown version: " + apiVersion[-1])
    return platform, app_name

def check_deployments_ready(record_xml_attribute, namespace, name, deployments, cluster_name):
    """Test that Kubeflow deployments are successfully deployed.

  Args:
    namespace: The namespace Kubeflow is deployed to.
  """
    set_logging()
    util.set_pytest_junit(record_xml_attribute, name)

    kfctl_aws_util.aws_auth_load_kubeconfig(cluster_name)

    api_client = deploy_utils.create_k8s_client()

    for deployment_name in deployments:
        logging.info("Verifying that deployment %s started...", deployment_name)
        util.wait_for_deployment(api_client, namespace, deployment_name, 10)

def test_admission_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "admission-webhook-deployment"
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_admission_is_ready", deployment_names,
                            cluster_name)

def test_katib_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "katib-controller",
        "katib-mysql",
        "katib-db-manager",
        "katib-ui",
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_katib_is_ready", deployment_names,
                            cluster_name)

def test_metadata_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "metadata-deployment",
        "metadata-grpc-deployment",
        "metadata-db",
        "metadata-ui",
        "metadata-envoy-deployment",
        "metadata-writer",
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_metadata_is_ready", deployment_names,
                            cluster_name)

def test_pipeline_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "argo-ui",
        "cache-deployer-deployment",
        "cache-server",
        "kubeflow-pipelines-profile-controller",
        "minio",
        "ml-pipeline",
        "ml-pipeline-persistenceagent",
        "ml-pipeline-scheduledworkflow",
        "ml-pipeline-ui",
        "ml-pipeline-viewer-crd",
        "ml-pipeline-visualizationserver",
        "mysql",
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_pipeline_is_ready", deployment_names,
                            cluster_name)

def test_notebook_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "jupyter-web-app-deployment",
        "notebook-controller-deployment",
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_notebook_is_ready", deployment_names,
                            cluster_name)

def test_centraldashboard_is_ready(record_xml_attribute, namespace, cluster_name):
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_centraldashboard_is_ready", ["centraldashboard"],
                            cluster_name)

def test_profiles_is_ready(record_xml_attribute, namespace, cluster_name):
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_profile_is_ready", ["profiles-deployment"],
                            cluster_name)

def test_seldon_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "seldon-controller-manager"
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_seldon_is_ready", deployment_names,
                            cluster_name)

def test_spark_is_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "spark-operatorsparkoperator"
    ]
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_spark_is_ready", deployment_names,
                            cluster_name)

def test_training_operators_are_ready(record_xml_attribute, namespace, cluster_name):
    deployment_names = [
        "mpi-operator",
        "mxnet-operator",
        "pytorch-operator",
        "tf-job-operator",
    ]

    check_deployments_ready(record_xml_attribute, namespace,
                            "test_training_is_ready", deployment_names,
                            cluster_name)

def test_workflow_controller_is_ready(record_xml_attribute, namespace, cluster_name):
    check_deployments_ready(record_xml_attribute, namespace,
                            "test_workflow_controller_is_ready", ["workflow-controller"],
                            cluster_name)

def test_kf_is_ready(record_xml_attribute, namespace, use_basic_auth,
                     app_path, cluster_name):
    """Test that Kubeflow was successfully deployed.

  Args:
    namespace: The namespace Kubeflow is deployed to.
  """
    set_logging()
    util.set_pytest_junit(record_xml_attribute, "test_kf_is_ready")

    kfctl_aws_util.aws_auth_load_kubeconfig(cluster_name)

    api_client = deploy_utils.create_k8s_client()

    # Verify that components are actually deployed.
    deployment_names = []

    stateful_set_names = []

    platform, _ = get_platform_app_name(app_path)

    # TODO(PatrickXYS): not sure why istio-galley can't found
    ingress_related_deployments = [
        "cluster-local-gateway",
        "istio-citadel",
        "istio-ingressgateway",
        "istio-pilot",
        "istio-policy",
        "istio-sidecar-injector",
        "istio-telemetry",
        "prometheus",
    ]
    ingress_related_stateful_sets = []

    knative_namespace = "knative-serving"
    knative_related_deployments = [
        "activator",
        "autoscaler",
        "controller",
        "networking-istio",
        "webhook",
    ]

    if platform == "gcp":
        deployment_names.extend(["cloud-endpoints-controller"])
        stateful_set_names.extend(["kfserving-controller-manager"])
        if use_basic_auth:
            deployment_names.extend(["basic-auth-login"])
            ingress_related_stateful_sets.extend(["backend-updater"])
        else:
            ingress_related_deployments.extend(["iap-enabler"])
            ingress_related_stateful_sets.extend(["backend-updater"])
    elif platform == "existing_arrikto":
        deployment_names.extend(["dex"])
        ingress_related_deployments.extend(["authservice"])
        knative_related_deployments = []
    elif platform == "aws":
        # TODO(PatrickXYS): Extend List with AWS Deployment
        deployment_names.extend(["alb-ingress-controller"])

    # TODO(jlewi): Might want to parallelize this.
    for deployment_name in deployment_names:
        logging.info("Verifying that deployment %s started...", deployment_name)
        util.wait_for_deployment(api_client, namespace, deployment_name, 10)

    ingress_namespace = "istio-system"
    for deployment_name in ingress_related_deployments:
        logging.info("Verifying that deployment %s started...", deployment_name)
        util.wait_for_deployment(api_client, ingress_namespace, deployment_name, 10)

    all_stateful_sets = [(namespace, name) for name in stateful_set_names]
    all_stateful_sets.extend([(ingress_namespace, name) for name in ingress_related_stateful_sets])

    for ss_namespace, name in all_stateful_sets:
        logging.info("Verifying that stateful set %s.%s started...", ss_namespace, name)
        try:
            util.wait_for_statefulset(api_client, ss_namespace, name)
        except:
            # Collect debug information by running describe
            util.run(["kubectl", "-n", ss_namespace, "describe", "statefulsets", name])
            raise

    ingress_names = ["istio-ingress"]
    # Check if Ingress is Ready and Healthy
    if platform in ["aws"]:
        for ingress_name in ingress_names:
            logging.info("Verifying that ingress %s started...", ingress_name)
            util.wait_for_ingress(api_client, ingress_namespace, ingress_name, 10)

    for deployment_name in knative_related_deployments:
        logging.info("Verifying that deployment %s started...", deployment_name)
        util.wait_for_deployment(api_client, knative_namespace, deployment_name, 10)

    # Check if Dex is Ready and Healthy
    dex_deployment_names = ["dex"]
    dex_namespace = "auth"
    for dex_deployment_name in dex_deployment_names:
        logging.info("Verifying that deployment %s started...", dex_deployment_name)
        util.wait_for_deployment(api_client, dex_namespace, dex_deployment_name, 10)

    # Check if Cert-Manager is Ready and Healthy
    cert_manager_deployment_names = [
        "cert-manager",
        "cert-manager-cainjector",
        "cert-manager-webhook",
    ]
    cert_manager_namespace = "cert-manager"
    for cert_manager_deployment_name in cert_manager_deployment_names:
        logging.info("Verifying that deployment %s started...", cert_manager_deployment_name)
        util.wait_for_deployment(api_client, cert_manager_namespace, cert_manager_deployment_name, 10)

if __name__ == "__main__":
    logging.basicConfig(level=logging.INFO,
                        format=('%(levelname)s|%(asctime)s'
                                '|%(pathname)s|%(lineno)d| %(message)s'),
                        datefmt='%Y-%m-%dT%H:%M:%S',
                        )
    logging.getLogger().setLevel(logging.INFO)
    pytest.main()
<|endoftext|>


<commit_before>import unittest
import random, sys, time, re
sys.path.extend(['.','..','py'])
import h2o_browse as h2b, h2o_gbm

import h2o, h2o_cmd, h2o_hosts, h2o_browse as h2b, h2o_import as h2i, h2o_glm, h2o_util, h2o_rf, h2o_jobs as h2j
class Basic(unittest.TestCase):
    def tearDown(self):
        h2o.check_sandbox_for_errors()

    @classmethod
    def setUpClass(cls):
        global SEED, localhost
        SEED = h2o.setup_random_seed()
        global localhost
        localhost = h2o.decide_if_localhost()
        if (localhost):
            h2o.build_cloud(3, java_heap_GB=4)
        else:
            h2o_hosts.build_cloud_with_hosts()

    @classmethod
    def tearDownClass(cls):
        h2o.tear_down_cloud()

    def test_GBM_manyfiles_train_test(self):
        bucket = 'home-0xdiag-datasets'
        modelKey = 'GBMModelKey'
        if localhost:
            files = [
                # None forces num_cols to be used. assumes you set it from Inspect
                # problems with categoricals not in the train data set? (warnings in h2o stdout)
                ## ('manyfiles-nflx-gz', 'file_1.dat.gz', 'file_1.hex', 1800, None, 'file_11.dat.gz', 'file_1_test.hex')
                # just use matching
                ('manyfiles-nflx-gz', 'file_1.dat.gz', 'file_1.hex', 1800, None, 'file_1.dat.gz', 'file_1_test.hex')
                ]
        else:
            files = [
                # None forces num_cols to be used. assumes you set it from Inspect
                ('manyfiles-nflx-gz', 'file_[0-9].dat.gz', 'file_10.hex', 1800, None, 'file_1[0-9].dat.gz', 'file_10_test.hex')
                ]

        # if I got to hdfs, it's here
        # hdfs://192.168.1.176/datasets/manyfiles-nflx-gz/file_99.dat.gz

        # h2b.browseTheCloud()
        for (importFolderPath, trainFilename, trainKey, timeoutSecs, response, testFilename, testKey) in files:
            h2o.beta_features = False #turn off beta_features
            # PARSE train****************************************
            start = time.time()
            xList = []
            eList = []
            fList = []

            # Parse (train)****************************************
            if h2o.beta_features:
                print "Parsing to fvec directly! Have to noPoll=true!, and doSummary=False!"
            csvPathname = importFolderPath + "/" + trainFilename
            parseTrainResult = h2i.import_parse(bucket=bucket, path=csvPathname, schema='local',
                hex_key=trainKey, timeoutSecs=timeoutSecs, noPoll=h2o.beta_features, doSummary=False)
            # hack
            if h2o.beta_features:
                h2j.pollWaitJobs(timeoutSecs=timeoutSecs, pollTimeoutSecs=timeoutSecs)
                print "Filling in the parseTrainResult['destination_key'] for h2o"
                parseTrainResult['destination_key'] = trainKey

            elapsed = time.time() - start
            print "train parse end on ", trainFilename, 'took', elapsed, 'seconds',\
                "%d pct. of timeout" % ((elapsed*100)/timeoutSecs)
            print "train parse result:", parseTrainResult['destination_key']

            ### h2o_cmd.runSummary(key=parsTraineResult['destination_key'])

            # if you set beta_features here, the fvec translate will happen with the Inspect not the GBM
            # h2o.beta_features = True
            inspect = h2o_cmd.runInspect(key=parseTrainResult['destination_key'])
            print "\n" + csvPathname, \
                "    num_rows:", "{:,}".format(inspect['num_rows']), \
                "    num_cols:", "{:,}".format(inspect['num_cols'])
            num_rows = inspect['num_rows']
            num_cols = inspect['num_cols']

            # Make col 378 it something we can do binomial regression on!
            execExpr = '%s=colSwap(%s,378,(%s[378]>15 ? 1 : 0))' % (trainKey, trainKey, trainKey)
            resultExec = h2o_cmd.runExec(expression=execExpr)

            # Parse (test)****************************************
            if h2o.beta_features:
                print "Parsing to fvec directly! Have to noPoll=true!, and doSummary=False!"

            parseTestResult = h2i.import_parse(bucket=bucket, path=importFolderPath + "/" + testFilename, schema='local',
                hex_key=testKey, timeoutSecs=timeoutSecs, noPoll=h2o.beta_features, doSummary=False)
            # hack
            if h2o.beta_features:
                h2j.pollWaitJobs(timeoutSecs=timeoutSecs, pollTimeoutSecs=timeoutSecs)
                print "Filling in the parseTestResult['destination_key'] for h2o"
                parseTestResult['destination_key'] = testKey

            elapsed = time.time() - start
            print "test parse end on ", testFilename, 'took', elapsed, 'seconds',\
                "%d pct. of timeout" % ((elapsed*100)/timeoutSecs)
            print "test parse result:", parseTestResult['destination_key']

            # Make col 378 it something we can do binomial regression on!
            execExpr = '%s=colSwap(%s,378,(%s[378]>15 ? 1 : 0))' % (testKey, testKey, testKey)
            resultExec = h2o_cmd.runExec(expression=execExpr)

            # Note ..no inspect of test data here..so translate happens later?

            # GBM (train iterate)****************************************
            # if not response:
            #     response = num_cols - 1
            response = 378

            # randomly ignore a bunch of cols, just to make it go faster
            x = range(num_cols)
            del x[response]
            ignored_cols_by_name = random.sample(x, 300)

            print "Using the same response %s for train and test (which should have a output value too)" % response

            ntrees = 10
            # ignore 200 random cols (not the response)
            for max_depth in [5, 40]:
                params = {
                    'learn_rate': .2,
                    'nbins': 1024,
                    'ntrees': ntrees,
                    'max_depth': max_depth,
                    'min_rows': 10,
                    'response': response,
                    'ignored_cols_by_name': ignored_cols_by_name,
                }
                print "Using these parameters for GBM: ", params
                kwargs = params.copy()
                h2o.beta_features = True

                # GBM train****************************************
                trainStart = time.time()
                gbmTrainResult = h2o_cmd.runGBM(parseResult=parseTrainResult,
                    noPoll=True, timeoutSecs=timeoutSecs, destination_key=modelKey, **kwargs)
                # hack
                if h2o.beta_features:
                    h2j.pollWaitJobs(timeoutSecs=timeoutSecs, pollTimeoutSecs=timeoutSecs)
                trainElapsed = time.time() - trainStart
                print "GBM training completed in", trainElapsed, "seconds. On dataset: ", trainFilename

                gbmTrainView = h2o_cmd.runGBMView(model_key=modelKey)
                # errrs from end of list? is that the last tree?
                errsLast = gbmTrainView['gbm_model']['errs'][-1]
                print "GBM 'errsLast'", errsLast

                cm = gbmTrainView['gbm_model']['cm']
                pctWrongTrain = h2o_gbm.pp_cm_summary(cm);
                print "Last line of this cm might be NAs, not CM"
                print "\nTrain\n==========\n"
                print h2o_gbm.pp_cm(cm)

                # GBM test****************************************
                predictKey = 'Predict.hex'
                ### h2o_cmd.runInspect(key=parseTestResult['destination_key'])
                start = time.time()
                gbmTestResult = h2o_cmd.runPredict(
                    data_key=parseTestResult['destination_key'],
                    model_key=modelKey,
                    destination_key=predictKey,
                    timeoutSecs=timeoutSecs)
                # hack
                if h2o.beta_features:
                    h2j.pollWaitJobs(timeoutSecs=timeoutSecs, pollTimeoutSecs=timeoutSecs)
                elapsed = time.time() - start
                print "GBM predict completed in", elapsed, "seconds. On dataset: ", testFilename

                print "This is crazy!"
                gbmPredictCMResult =h2o.nodes[0].predict_confusion_matrix(
                    actual=parseTestResult['destination_key'],
                    vactual=response,
                    predict=predictKey,
                    vpredict='predict', # choices are 0 and 'predict'
                    )

                # errrs from end of list? is that the last tree?
                # all we get is cm
                cm = gbmPredictCMResult['cm']

                # These will move into the h2o_gbm.py
                pctWrong = h2o_gbm.pp_cm_summary(cm);
                print "Last line of this cm is really NAs, not CM"
                print "\nTest\n==========\n"
                print h2o_gbm.pp_cm(cm)

                # xList.append(ntrees)
                xList.append(max_depth)
                eList.append(pctWrong)
                fList.append(trainElapsed)

            h2o.beta_features = False
            xLabel = 'max_depth'
            eLabel = 'pctWrong'
            fLabel = 'trainElapsed'
            eListTitle = ""
            fListTitle = ""
            h2o_gbm.plotLists(xList, xLabel, eListTitle, eList, eLabel, fListTitle, fList, fLabel)

if __name__ == '__main__':
    h2o.unit_main()
<commit_msg>params forced to be static. https://0xdata.atlassian.net/browse/HEX-120'0<commit_after>import unittest
import random, sys, time, re
sys.path.extend(['.','..','py'])
import h2o_browse as h2b, h2o_gbm

FORCE_FAIL_CASE = True

import h2o, h2o_cmd, h2o_hosts, h2o_browse as h2b, h2o_import as h2i, h2o_glm, h2o_util, h2o_rf, h2o_jobs as h2j
class Basic(unittest.TestCase):
    def tearDown(self):
        h2o.check_sandbox_for_errors()

    @classmethod
    def setUpClass(cls):
        global SEED, localhost
        SEED = h2o.setup_random_seed()
        global localhost
        localhost = h2o.decide_if_localhost()
        if (localhost):
            h2o.build_cloud(3, java_heap_GB=4)
        else:
            h2o_hosts.build_cloud_with_hosts()

    @classmethod
    def tearDownClass(cls):
        h2o.tear_down_cloud()

    def test_GBM_manyfiles_train_test(self):
        bucket = 'home-0xdiag-datasets'
        modelKey = 'GBMModelKey'
        if localhost:
            files = [
                # None forces num_cols to be used. assumes you set it from Inspect
                # problems with categoricals not in the train data set? (warnings in h2o stdout)
                ## ('manyfiles-nflx-gz', 'file_1.dat.gz', 'file_1.hex', 1800, None, 'file_11.dat.gz', 'file_1_test.hex')
                # just use matching
                ('manyfiles-nflx-gz', 'file_1.dat.gz', 'file_1.hex', 1800, None, 'file_1.dat.gz', 'file_1_test.hex')
                ]
        else:
            files = [
                # None forces num_cols to be used. assumes you set it from Inspect
                ('manyfiles-nflx-gz', 'file_[0-9].dat.gz', 'file_10.hex', 1800, None, 'file_1[0-9].dat.gz', 'file_10_test.hex')
                ]

        # if I got to hdfs, it's here
        # hdfs://192.168.1.176/datasets/manyfiles-nflx-gz/file_99.dat.gz

        # h2b.browseTheCloud()
        for (importFolderPath, trainFilename, trainKey, timeoutSecs, response, testFilename, testKey) in files:
            h2o.beta_features = False #turn off beta_features
            # PARSE train****************************************
            start = time.time()
            xList = []
            eList = []
            fList = []

            # Parse (train)****************************************
            if h2o.beta_features:
                print "Parsing to fvec directly! Have to noPoll=true!, and doSummary=False!"
            csvPathname = importFolderPath + "/" + trainFilename
            parseTrainResult = h2i.import_parse(bucket=bucket, path=csvPathname, schema='local',
                hex_key=trainKey, timeoutSecs=timeoutSecs, noPoll=h2o.beta_features, doSummary=False)
            # hack
            if h2o.beta_features:
                h2j.pollWaitJobs(timeoutSecs=timeoutSecs, pollTimeoutSecs=timeoutSecs)
                print "Filling in the parseTrainResult['destination_key'] for h2o"
                parseTrainResult['destination_key'] = trainKey

            elapsed = time.time() - start
            print "train parse end on ", trainFilename, 'took', elapsed, 'seconds',\
                "%d pct. of timeout" % ((elapsed*100)/timeoutSecs)
            print "train parse result:", parseTrainResult['destination_key']

            ### h2o_cmd.runSummary(key=parsTraineResult['destination_key'])

            # if you set beta_features here, the fvec translate will happen with the Inspect not the GBM
            # h2o.beta_features = True
            inspect = h2o_cmd.runInspect(key=parseTrainResult['destination_key'])
            print "\n" + csvPathname, \
                "    num_rows:", "{:,}".format(inspect['num_rows']), \
                "    num_cols:", "{:,}".format(inspect['num_cols'])
            num_rows = inspect['num_rows']
            num_cols = inspect['num_cols']

            # Make col 378 it something we can do binomial regression on!
            execExpr = '%s=colSwap(%s,378,(%s[378]>15 ? 1 : 0))' % (trainKey, trainKey, trainKey)
            resultExec = h2o_cmd.runExec(expression=execExpr)

            # Parse (test)****************************************
            if h2o.beta_features:
                print "Parsing to fvec directly! Have to noPoll=true!, and doSummary=False!"

            parseTestResult = h2i.import_parse(bucket=bucket, path=importFolderPath + "/" + testFilename, schema='local',
                hex_key=testKey, timeoutSecs=timeoutSecs, noPoll=h2o.beta_features, doSummary=False)
            # hack
            if h2o.beta_features:
                h2j.pollWaitJobs(timeoutSecs=timeoutSecs, pollTimeoutSecs=timeoutSecs)
                print "Filling in the parseTestResult['destination_key'] for h2o"
                parseTestResult['destination_key'] = testKey

            elapsed = time.time() - start
            print "test parse end on ", testFilename, 'took', elapsed, 'seconds',\
                "%d pct. of timeout" % ((elapsed*100)/timeoutSecs)
            print "test parse result:", parseTestResult['destination_key']

            # Make col 378 it something we can do binomial regression on!
            execExpr = '%s=colSwap(%s,378,(%s[378]>15 ? 1 : 0))' % (testKey, testKey, testKey)
            resultExec = h2o_cmd.runExec(expression=execExpr)

            # Note ..no inspect of test data here..so translate happens later?

            # GBM (train iterate)****************************************
            # if not response:
            #     response = num_cols - 1
            response = 378

            # randomly ignore a bunch of cols, just to make it go faster
            x = range(num_cols)
            del x[response]
            ignored_cols_by_name = random.sample(x, 300)

            print "Using the same response %s for train and test (which should have a output value too)" % response

            ntrees = 10
            # ignore 200 random cols (not the response)
            for max_depth in [5, 40]:
                params = {
                    'learn_rate': .2,
                    'nbins': 1024,
                    'ntrees': ntrees,
                    'max_depth': max_depth,
                    'min_rows': 10,
                    'response': response,
                    'ignored_cols_by_name': ignored_cols_by_name,
                }

                if FORCE_FAIL_CASE:
                    params = {'learn_rate': 0.2, 'classification': None, 'min_rows': 10, 'ntrees': 10, 'response': 378, 'nbins': 1024, 'ignored_cols_by_name': [256, 382, 399, 50, 176, 407, 375, 113, 170, 313, 364, 33, 361, 426, 121, 371, 232, 327, 480, 75, 37, 312, 225, 195, 244, 406, 268, 230, 321, 257, 274, 197, 35, 501, 360, 72, 213, 79, 1, 466, 362, 160, 444, 437, 5, 59, 108, 454, 73, 374, 509, 337, 183, 252, 21, 314, 100, 200, 159, 379, 405, 367, 432, 181, 8, 420, 118, 284, 281, 465, 456, 359, 291, 330, 258, 523, 243, 487, 408, 392, 15, 231, 482, 481, 70, 171, 182, 31, 409, 492, 471, 53, 45, 448, 83, 527, 452, 350, 423, 93, 447, 130, 126, 54, 354, 169, 253, 49, 42, 431, 305, 498, 216, 189, 508, 122, 308, 228, 190, 293, 451, 63, 133, 304, 397, 425, 333, 19, 158, 391, 153, 282, 112, 64, 502, 7, 16, 469, 163, 136, 40, 99, 302, 264, 325, 434, 187, 311, 286, 278, 179, 109, 348, 287, 467, 400, 164, 384, 422, 43, 117, 91, 276, 211, 175, 329, 541, 438, 145, 534, 218, 177, 317, 222, 210, 162, 402, 98, 299, 245, 385, 233, 188, 516, 143, 13, 532, 429, 172, 455, 470, 518, 236, 296, 388, 468, 110, 395, 185, 25, 489, 196, 120, 435, 165, 168, 271, 74, 510, 36, 76, 208, 223, 270, 515, 421, 87, 66, 473, 220, 46, 486, 102, 38, 156, 48, 132, 331, 51, 403, 234, 23, 449, 341, 303, 410, 479, 203, 413, 512, 513, 9, 446, 511, 55, 6, 339, 418, 476, 178, 266, 22, 141, 259, 349, 86, 144, 34, 290, 326, 318, 519, 424, 127, 174, 472, 116, 17, 152, 280, 215, 514, 103, 377, 537, 373, 238, 47, 353, 428, 94, 214, 61, 123, 386, 351, 246, 411, 101, 249, 240, 520, 307, 288, 199, 147, 436, 77, 464, 414], 'source': u'file_1.hex', 'validation': u'file_1.hex', 'max_depth': 5}

                ### print "Using these parameters for GBM: ", params
                kwargs = params.copy()
                h2o.beta_features = True

                # GBM train****************************************
                trainStart = time.time()
                gbmTrainResult = h2o_cmd.runGBM(parseResult=parseTrainResult,
                    noPoll=True, timeoutSecs=timeoutSecs, destination_key=modelKey, **kwargs)
                # hack
                if h2o.beta_features:
                    h2j.pollWaitJobs(timeoutSecs=timeoutSecs, pollTimeoutSecs=timeoutSecs)
                trainElapsed = time.time() - trainStart
                print "GBM training completed in", trainElapsed, "seconds. On dataset: ", trainFilename

                gbmTrainView = h2o_cmd.runGBMView(model_key=modelKey)
                # errrs from end of list? is that the last tree?
                errsLast = gbmTrainView['gbm_model']['errs'][-1]
                print "GBM 'errsLast'", errsLast

                cm = gbmTrainView['gbm_model']['cm']
                pctWrongTrain = h2o_gbm.pp_cm_summary(cm);
                print "Last line of this cm might be NAs, not CM"
                print "\nTrain\n==========\n"
                print h2o_gbm.pp_cm(cm)

                # GBM test****************************************
                predictKey = 'Predict.hex'
                ### h2o_cmd.runInspect(key=parseTestResult['destination_key'])
                start = time.time()
                gbmTestResult = h2o_cmd.runPredict(
                    data_key=parseTestResult['destination_key'],
                    model_key=modelKey,
                    destination_key=predictKey,
                    timeoutSecs=timeoutSecs)
                # hack
                if h2o.beta_features:
                    h2j.pollWaitJobs(timeoutSecs=timeoutSecs, pollTimeoutSecs=timeoutSecs)
                elapsed = time.time() - start
                print "GBM predict completed in", elapsed, "seconds. On dataset: ", testFilename

                print "This is crazy!"
                gbmPredictCMResult =h2o.nodes[0].predict_confusion_matrix(
                    actual=parseTestResult['destination_key'],
                    vactual=response,
                    predict=predictKey,
                    vpredict='predict', # choices are 0 and 'predict'
                    )

                # errrs from end of list? is that the last tree?
                # all we get is cm
                cm = gbmPredictCMResult['cm']

                # These will move into the h2o_gbm.py
                pctWrong = h2o_gbm.pp_cm_summary(cm);
                print "Last line of this cm is really NAs, not CM"
                print "\nTest\n==========\n"
                print h2o_gbm.pp_cm(cm)

                # xList.append(ntrees)
                xList.append(max_depth)
                eList.append(pctWrong)
                fList.append(trainElapsed)

            h2o.beta_features = False
            xLabel = 'max_depth'
            eLabel = 'pctWrong'
            fLabel = 'trainElapsed'
            eListTitle = ""
            fListTitle = ""
            h2o_gbm.plotLists(xList, xLabel, eListTitle, eList, eLabel, fListTitle, fList, fLabel)

if __name__ == '__main__':
    h2o.unit_main()
<|endoftext|>


<commit_before>#
# Copyright 2018 Analytics Zoo Authors.
#
# 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.
#

from bigdl.nn.keras.layer import KerasLayer
from bigdl.nn.layer import Node

from zoo.pipeline.api.keras.utils import *

if sys.version >= '3':
    long = int
    unicode = str

class ZooKerasCreator(JavaValue):
    def jvm_class_constructor(self):
        name = "createZooKeras" + self.__class__.__name__
        print("creating: " + name)
        return name

class ZooKerasLayer(ZooKerasCreator, KerasLayer):
    @classmethod
    def of(cls, jvalue, bigdl_type="float"):
        return KerasLayer(jvalue, bigdl_type)

class KerasNet(ZooKerasLayer):
    def compile(self, optimizer, loss, metrics=None):
        """
        Configure the learning process. It MUST be called before fit or evaluate.

        # Arguments
        optimizer: Optimization method to be used. One can alternatively pass in the corresponding
                   string representation, such as 'sgd'.
        loss: Criterion to be used. One can alternatively pass in the corresponding string
              representation, such as 'mse'.
        metrics: List of validation methods to be used. Default is None if no validation is needed.
                 One can alternatively use ['accuracy'].
        """
        if isinstance(optimizer, six.string_types):
            optimizer = to_bigdl_optim_method(optimizer)
        if isinstance(loss, six.string_types):
            loss = to_bigdl_criterion(loss)
        if callable(loss):
            from zoo.pipeline.api.autograd import CustomLoss
            loss = CustomLoss(loss, self.get_output_shape()[1:])
        if metrics and all(isinstance(metric, six.string_types) for metric in metrics):
            metrics = to_bigdl_metrics(metrics)
        callBigDlFunc(self.bigdl_type, "zooCompile",
                      self.value,
                      optimizer,
                      loss,
                      metrics)

    def set_tensorboard(self, log_dir, app_name):
        """
        Set summary information during the training process for visualization purposes.
        Saved summary can be viewed via TensorBoard.
        In order to take effect, it needs to be called before fit.

        Training summary will be saved to 'log_dir/app_name/train'
        and validation summary (if any) will be saved to 'log_dir/app_name/validation'.

        # Arguments
        log_dir: The base directory path to store training and validation logs.
        app_name: The name of the application.
        """
        callBigDlFunc(self.bigdl_type, "zooSetTensorBoard",
                      self.value,
                      log_dir,
                      app_name)

    def set_checkpoint(self, path, over_write=True):
        """
        Configure checkpoint settings to write snapshots every epoch during the training process.
        In order to take effect, it needs to be called before fit.

        # Arguments
        path: The path to save snapshots. Make sure this path exists beforehand.
        over_write: Whether to overwrite existing snapshots in the given path. Default is True.
        """
        callBigDlFunc(self.bigdl_type, "zooSetCheckpoint",
                      self.value,
                      path,
                      over_write)

    def clear_gradient_clipping(self):
        """
        Clear gradient clipping parameters. In this case, gradient clipping will not be applied.
        In order to take effect, it needs to be called before fit.
        """
        callBigDlFunc(self.bigdl_type, "zooClearGradientClipping",
                      self.value)

    def set_constant_gradient_clipping(self, min, max):
        """
        Set constant gradient clipping during the training process.
        In order to take effect, it needs to be called before fit.

        # Arguments
        min: The minimum value to clip by. Float.
        max: The maximum value to clip by. Float.
        """
        callBigDlFunc(self.bigdl_type, "zooSetConstantGradientClipping",
                      self.value,
                      float(min),
                      float(max))

    def set_gradient_clipping_by_l2_norm(self, clip_norm):
        """
        Clip gradient to a maximum L2-Norm during the training process.
        In order to take effect, it needs to be called before fit.

        # Arguments
        clip_norm: Gradient L2-Norm threshold. Float.
        """
        callBigDlFunc(self.bigdl_type, "zooSetGradientClippingByL2Norm",
                      self.value,
                      float(clip_norm))

    def fit(self, x, y=None, batch_size=32, nb_epoch=10, validation_data=None, distributed=True):
        """
        Train a model for a fixed number of epochs on a dataset.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample or Image DataSet.
        y: Labels. A Numpy array. Default is None if x is already RDD of Sample or Image DataSet.
        batch_size: Number of samples per gradient update.
        nb_epoch: Number of iterations to train.
        validation_data: Tuple (x_val, y_val) where x_val and y_val are both Numpy arrays.
                         Or RDD of Sample. Default is None if no validation is involved.
        distributed: Boolean. Whether to train the model in distributed mode or local mode.
                     Default is True. In local mode, x and y must both be Numpy arrays.
        """
        if distributed:
            if isinstance(x, np.ndarray) and isinstance(y, np.ndarray):
                training_data = to_sample_rdd(x, y)
                if validation_data:
                    validation_data = to_sample_rdd(*validation_data)
            elif (isinstance(x, RDD) or isinstance(x, DataSet)) and not y:
                training_data = x
            else:
                raise TypeError("Unsupported training data type: %s" % type(x))
            callBigDlFunc(self.bigdl_type, "zooFit",
                          self.value,
                          training_data,
                          batch_size,
                          nb_epoch,
                          validation_data)
        else:
            if validation_data:
                val_x = [JTensor.from_ndarray(x) for x in to_list(validation_data[0])]
                val_y = JTensor.from_ndarray(validation_data[1])
            else:
                val_x, val_y = None, None
            callBigDlFunc(self.bigdl_type, "zooFit",
                          self.value,
                          [JTensor.from_ndarray(x) for x in to_list(x)],
                          JTensor.from_ndarray(y),
                          batch_size,
                          nb_epoch,
                          val_x,
                          val_y)

    def evaluate(self, x, y=None, batch_size=32):
        """
        Evaluate a model on a given dataset in distributed mode.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample.
        y: Labels. A Numpy array. Default is None if x is already RDD of Sample.
        batch_size: Number of samples per gradient update.
        """
        if isinstance(x, np.ndarray) and isinstance(y, np.ndarray):
            evaluation_data = to_sample_rdd(x, y)
        elif isinstance(x, RDD) and not y:
            evaluation_data = x
        else:
            raise TypeError("Unsupported evaluation data type: %s" % type(x))
        return callBigDlFunc(self.bigdl_type, "zooEvaluate",
                             self.value,
                             evaluation_data,
                             batch_size)

    def predict(self, x, distributed=True):
        """
        Use a model to do prediction.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample.
        distributed: Boolean. Whether to do prediction in distributed mode or local mode.
                     Default is True. In local mode, x must be a Numpy array.
        """
        if is_distributed:
            if isinstance(x, np.ndarray):
                features = to_sample_rdd(x, np.zeros([x.shape[0]]))
            elif isinstance(x, RDD):
                features = x
            else:
                raise TypeError("Unsupported prediction data type: %s" % type(x))
            return self.predict_distributed(features)
        else:
            if isinstance(x, np.ndarray):
                return self.predict_local(x)
            else:
                raise TypeError("Unsupported prediction data type: %s" % type(x))

    def to_model(self):
        from zoo.pipeline.api.keras.models import Model
        return Model.from_jvalue(callBigDlFunc(self.bigdl_type, "kerasNetToModel", self.value))

class Input(ZooKerasCreator, Node):
    """
    Used to instantiate an input node.

    # Arguments
    shape: A shape tuple, not including batch.
    name: String to set the name of the layer.
          If not specified, its name will by default to be a generated string.

    >>> input = Input(name="input1", shape=(3, 5))
    creating: createZooKerasInput
    """
    def __init__(self, shape=None, name=None, bigdl_type="float"):
        super(Input, self).__init__(None, bigdl_type,
                                    name,
                                    list(shape) if shape else None)

class InputLayer(ZooKerasLayer):
    """
    Used as an entry point into a model.

    # Arguments
    input_shape: A shape tuple, not including batch.
    name: String to set the name of the layer.
          If not specified, its name will by default to be a generated string.

    >>> inputlayer = InputLayer(input_shape=(3, 5), name="input1")
    creating: createZooKerasInputLayer
    """
    def __init__(self, input_shape=None, **kwargs):
        super(InputLayer, self).__init__(None,
                                         list(input_shape) if input_shape else None,
                                         **kwargs)

class Merge(ZooKerasLayer):
    """
    Used to merge a list of inputs into a single output, following some merge mode.
    Merge must have at least two input layers.

    When using this layer as the first layer in a model, you need to provide the argument
    input_shape for input layers (a list of shape tuples, does not include the batch dimension).

    # Arguments
    layers: A list of layer instances. Must be more than one layer.
    mode: Merge mode. String, must be one of: 'sum', 'mul', 'concat', 'ave', 'cos',
          'dot', 'max', 'sub', 'div', 'min'. Default is 'sum'.
    concat_axis: Int, axis to use when concatenating layers.
                 Only specify this when merge mode is 'concat'.
                 Default is -1, meaning the last axis of the input.
    input_shape: A list of shape tuples, each not including batch.
    name: String to set the name of the layer.
          If not specified, its name will by default to be a generated string.

    >>> l1 = InputLayer(input_shape=(3, 5))
    creating: createZooKerasInputLayer
    >>> l2 = InputLayer(input_shape=(3, 5))
    creating: createZooKerasInputLayer
    >>> merge = Merge(layers=[l1, l2], mode='sum', name="merge1")
    creating: createZooKerasMerge
    """
    def __init__(self, layers=None, mode="sum", concat_axis=-1,
                 input_shape=None, **kwargs):
        super(Merge, self).__init__(None,
                                    list(layers) if layers else None,
                                    mode,
                                    concat_axis,
                                    input_shape,
                                    **kwargs)

def merge(inputs, mode="sum", concat_axis=-1, name=None):
    """
    Functional merge. Only use this method if you are defining a graph model.
    Used to merge a list of input nodes into a single output node (NOT layers!),
    following some merge mode.

    # Arguments
    inputs: A list of node instances. Must be more than one node.
    mode: Merge mode. String, must be one of: 'sum', 'mul', 'concat', 'ave', 'cos',
          'dot', 'max', 'sub', 'div', 'min'. Default is 'sum'.
    concat_axis: Int, axis to use when concatenating nodes.
                 Only specify this when merge mode is 'concat'.
                 Default is -1, meaning the last axis of the input.
    name: String to set the name of the functional merge.
          If not specified, its name will by default to be a generated string.
    """
    return Merge(mode=mode, concat_axis=concat_axis, name=name)(list(inputs))
<commit_msg>Fix GraphNet and kerasNet submodules bug (#221)<commit_after>#
# Copyright 2018 Analytics Zoo Authors.
#
# 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.
#

from bigdl.nn.keras.layer import KerasLayer
from bigdl.nn.layer import Node

from zoo.pipeline.api.keras.utils import *

if sys.version >= '3':
    long = int
    unicode = str

class ZooKerasCreator(JavaValue):
    def jvm_class_constructor(self):
        name = "createZooKeras" + self.__class__.__name__
        print("creating: " + name)
        return name

class ZooKerasLayer(ZooKerasCreator, KerasLayer):
    @classmethod
    def of(cls, jvalue, bigdl_type="float"):
        return KerasLayer(jvalue, bigdl_type)

class KerasNet(ZooKerasLayer):
    def compile(self, optimizer, loss, metrics=None):
        """
        Configure the learning process. It MUST be called before fit or evaluate.

        # Arguments
        optimizer: Optimization method to be used. One can alternatively pass in the corresponding
                   string representation, such as 'sgd'.
        loss: Criterion to be used. One can alternatively pass in the corresponding string
              representation, such as 'mse'.
        metrics: List of validation methods to be used. Default is None if no validation is needed.
                 One can alternatively use ['accuracy'].
        """
        if isinstance(optimizer, six.string_types):
            optimizer = to_bigdl_optim_method(optimizer)
        if isinstance(loss, six.string_types):
            loss = to_bigdl_criterion(loss)
        if callable(loss):
            from zoo.pipeline.api.autograd import CustomLoss
            loss = CustomLoss(loss, self.get_output_shape()[1:])
        if metrics and all(isinstance(metric, six.string_types) for metric in metrics):
            metrics = to_bigdl_metrics(metrics)
        callBigDlFunc(self.bigdl_type, "zooCompile",
                      self.value,
                      optimizer,
                      loss,
                      metrics)

    def set_tensorboard(self, log_dir, app_name):
        """
        Set summary information during the training process for visualization purposes.
        Saved summary can be viewed via TensorBoard.
        In order to take effect, it needs to be called before fit.

        Training summary will be saved to 'log_dir/app_name/train'
        and validation summary (if any) will be saved to 'log_dir/app_name/validation'.

        # Arguments
        log_dir: The base directory path to store training and validation logs.
        app_name: The name of the application.
        """
        callBigDlFunc(self.bigdl_type, "zooSetTensorBoard",
                      self.value,
                      log_dir,
                      app_name)

    def set_checkpoint(self, path, over_write=True):
        """
        Configure checkpoint settings to write snapshots every epoch during the training process.
        In order to take effect, it needs to be called before fit.

        # Arguments
        path: The path to save snapshots. Make sure this path exists beforehand.
        over_write: Whether to overwrite existing snapshots in the given path. Default is True.
        """
        callBigDlFunc(self.bigdl_type, "zooSetCheckpoint",
                      self.value,
                      path,
                      over_write)

    def clear_gradient_clipping(self):
        """
        Clear gradient clipping parameters. In this case, gradient clipping will not be applied.
        In order to take effect, it needs to be called before fit.
        """
        callBigDlFunc(self.bigdl_type, "zooClearGradientClipping",
                      self.value)

    def set_constant_gradient_clipping(self, min, max):
        """
        Set constant gradient clipping during the training process.
        In order to take effect, it needs to be called before fit.

        # Arguments
        min: The minimum value to clip by. Float.
        max: The maximum value to clip by. Float.
        """
        callBigDlFunc(self.bigdl_type, "zooSetConstantGradientClipping",
                      self.value,
                      float(min),
                      float(max))

    def set_gradient_clipping_by_l2_norm(self, clip_norm):
        """
        Clip gradient to a maximum L2-Norm during the training process.
        In order to take effect, it needs to be called before fit.

        # Arguments
        clip_norm: Gradient L2-Norm threshold. Float.
        """
        callBigDlFunc(self.bigdl_type, "zooSetGradientClippingByL2Norm",
                      self.value,
                      float(clip_norm))

    def fit(self, x, y=None, batch_size=32, nb_epoch=10, validation_data=None, distributed=True):
        """
        Train a model for a fixed number of epochs on a dataset.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample or Image DataSet.
        y: Labels. A Numpy array. Default is None if x is already RDD of Sample or Image DataSet.
        batch_size: Number of samples per gradient update.
        nb_epoch: Number of iterations to train.
        validation_data: Tuple (x_val, y_val) where x_val and y_val are both Numpy arrays.
                         Or RDD of Sample. Default is None if no validation is involved.
        distributed: Boolean. Whether to train the model in distributed mode or local mode.
                     Default is True. In local mode, x and y must both be Numpy arrays.
        """
        if distributed:
            if isinstance(x, np.ndarray) and isinstance(y, np.ndarray):
                training_data = to_sample_rdd(x, y)
                if validation_data:
                    validation_data = to_sample_rdd(*validation_data)
            elif (isinstance(x, RDD) or isinstance(x, DataSet)) and not y:
                training_data = x
            else:
                raise TypeError("Unsupported training data type: %s" % type(x))
            callBigDlFunc(self.bigdl_type, "zooFit",
                          self.value,
                          training_data,
                          batch_size,
                          nb_epoch,
                          validation_data)
        else:
            if validation_data:
                val_x = [JTensor.from_ndarray(x) for x in to_list(validation_data[0])]
                val_y = JTensor.from_ndarray(validation_data[1])
            else:
                val_x, val_y = None, None
            callBigDlFunc(self.bigdl_type, "zooFit",
                          self.value,
                          [JTensor.from_ndarray(x) for x in to_list(x)],
                          JTensor.from_ndarray(y),
                          batch_size,
                          nb_epoch,
                          val_x,
                          val_y)

    def evaluate(self, x, y=None, batch_size=32):
        """
        Evaluate a model on a given dataset in distributed mode.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample.
        y: Labels. A Numpy array. Default is None if x is already RDD of Sample.
        batch_size: Number of samples per gradient update.
        """
        if isinstance(x, np.ndarray) and isinstance(y, np.ndarray):
            evaluation_data = to_sample_rdd(x, y)
        elif isinstance(x, RDD) and not y:
            evaluation_data = x
        else:
            raise TypeError("Unsupported evaluation data type: %s" % type(x))
        return callBigDlFunc(self.bigdl_type, "zooEvaluate",
                             self.value,
                             evaluation_data,
                             batch_size)

    def predict(self, x, distributed=True):
        """
        Use a model to do prediction.

        # Arguments
        x: Input data. A Numpy array or RDD of Sample.
        distributed: Boolean. Whether to do prediction in distributed mode or local mode.
                     Default is True. In local mode, x must be a Numpy array.
        """
        if is_distributed:
            if isinstance(x, np.ndarray):
                features = to_sample_rdd(x, np.zeros([x.shape[0]]))
            elif isinstance(x, RDD):
                features = x
            else:
                raise TypeError("Unsupported prediction data type: %s" % type(x))
            return self.predict_distributed(features)
        else:
            if isinstance(x, np.ndarray):
                return self.predict_local(x)
            else:
                raise TypeError("Unsupported prediction data type: %s" % type(x))

    def to_model(self):
        from zoo.pipeline.api.keras.models import Model
        return Model.from_jvalue(callBigDlFunc(self.bigdl_type, "kerasNetToModel", self.value))

    @property
    def layers(self):
        jlayers = callBigDlFunc(self.bigdl_type, "getSubModules", self)
        layers = [Layer.of(jlayer) for jlayer in jlayers]
        return layers

    def flattened_layers(self, include_container=False):
        jlayers = callBigDlFunc(self.bigdl_type, "getFlattenSubModules", self, include_container)
        layers = [Layer.of(jlayer) for jlayer in jlayers]
        return layers

class Input(ZooKerasCreator, Node):
    """
    Used to instantiate an input node.

    # Arguments
    shape: A shape tuple, not including batch.
    name: String to set the name of the layer.
          If not specified, its name will by default to be a generated string.

    >>> input = Input(name="input1", shape=(3, 5))
    creating: createZooKerasInput
    """
    def __init__(self, shape=None, name=None, bigdl_type="float"):
        super(Input, self).__init__(None, bigdl_type,
                                    name,
                                    list(shape) if shape else None)

class InputLayer(ZooKerasLayer):
    """
    Used as an entry point into a model.

    # Arguments
    input_shape: A shape tuple, not including batch.
    name: String to set the name of the layer.
          If not specified, its name will by default to be a generated string.

    >>> inputlayer = InputLayer(input_shape=(3, 5), name="input1")
    creating: createZooKerasInputLayer
    """
    def __init__(self, input_shape=None, **kwargs):
        super(InputLayer, self).__init__(None,
                                         list(input_shape) if input_shape else None,
                                         **kwargs)

class Merge(ZooKerasLayer):
    """
    Used to merge a list of inputs into a single output, following some merge mode.
    Merge must have at least two input layers.

    When using this layer as the first layer in a model, you need to provide the argument
    input_shape for input layers (a list of shape tuples, does not include the batch dimension).

    # Arguments
    layers: A list of layer instances. Must be more than one layer.
    mode: Merge mode. String, must be one of: 'sum', 'mul', 'concat', 'ave', 'cos',
          'dot', 'max', 'sub', 'div', 'min'. Default is 'sum'.
    concat_axis: Int, axis to use when concatenating layers.
                 Only specify this when merge mode is 'concat'.
                 Default is -1, meaning the last axis of the input.
    input_shape: A list of shape tuples, each not including batch.
    name: String to set the name of the layer.
          If not specified, its name will by default to be a generated string.

    >>> l1 = InputLayer(input_shape=(3, 5))
    creating: createZooKerasInputLayer
    >>> l2 = InputLayer(input_shape=(3, 5))
    creating: createZooKerasInputLayer
    >>> merge = Merge(layers=[l1, l2], mode='sum', name="merge1")
    creating: createZooKerasMerge
    """
    def __init__(self, layers=None, mode="sum", concat_axis=-1,
                 input_shape=None, **kwargs):
        super(Merge, self).__init__(None,
                                    list(layers) if layers else None,
                                    mode,
                                    concat_axis,
                                    input_shape,
                                    **kwargs)

def merge(inputs, mode="sum", concat_axis=-1, name=None):
    """
    Functional merge. Only use this method if you are defining a graph model.
    Used to merge a list of input nodes into a single output node (NOT layers!),
    following some merge mode.

    # Arguments
    inputs: A list of node instances. Must be more than one node.
    mode: Merge mode. String, must be one of: 'sum', 'mul', 'concat', 'ave', 'cos',
          'dot', 'max', 'sub', 'div', 'min'. Default is 'sum'.
    concat_axis: Int, axis to use when concatenating nodes.
                 Only specify this when merge mode is 'concat'.
                 Default is -1, meaning the last axis of the input.
    name: String to set the name of the functional merge.
          If not specified, its name will by default to be a generated string.
    """
    return Merge(mode=mode, concat_axis=concat_axis, name=name)(list(inputs))
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
"""
Written by Daniel M. Aukes
Email: danaukes<at>gmail.com
Please see LICENSE for full license.
"""

import pynamics
#pynamics.script_mode = True
from pynamics.frame import Frame
from pynamics.variable_types import Differentiable,Constant
from pynamics.system import System
from pynamics.body import Body
from pynamics.dyadic import Dyadic
from pynamics.output import Output
from pynamics.particle import Particle

#import sympy
import numpy
import scipy.integrate
import matplotlib.pyplot as plt
plt.ion()
from sympy import pi
system = System()
tol = 1e-10
lA = Constant('lA',0.05699,system)

mA = Constant('mA',0.0054531675,system)

g = Constant('g',9.81,system)
b_air = Constant('b_air',0,system)
b_joint = Constant('b_joint',0.000023855338,system)
k = Constant('k',0.016705166667,system)

Ixx_A = Constant('Ixx_A',1,system)
Iyy_A = Constant('Iyy_A',1,system)
Izz_A = Constant('Izz_A',0.00000428720184,system)

tinitial = 0
tfinal = 5
tstep = .001
t = numpy.r_[tinitial:tfinal:tstep]

preload1 = Constant('preload1',0*pi/180,system)

qA,qA_d,qA_dd = Differentiable(system,'qA')

initialvalues = {}
initialvalues[qA]=10*pi/180
initialvalues[qA_d]=0*pi/180

statevariables = system.get_q(0)+system.get_q(1)
ini = [initialvalues[item] for item in statevariables]

N = Frame('N')
A = Frame('A')

system.set_newtonian(N)
A.rotate_fixed_axis_directed(N,[0,0,1],qA,system)

pNA=0*N.x
pAB=pNA+lA*A.x
vAB=pAB.time_derivative(N,system)

#ParticleA = Particle(system,pAB,mA,'ParticleA')
IA = Dyadic.build(A,Ixx_A,Iyy_A,Izz_A)
BodyA = Body('BodyA',A,pAB,mA,IA,system)

wNA = N.getw_(A)

lab2 = vAB.dot(vAB)
uab = vAB * (1/(lab2**.5+tol))

#squared term
#system.addforce(-b_air*lab2*uab,vAB)
#linear term
system.addforce(-b_air*vAB,vAB)
system.addforce(-b_joint*wNA,wNA)
system.addforcegravity(-g*N.x)
system.add_spring_force(k,(qA-preload1)*N.z,wNA)

#x1 = ParticleA.pCM.dot(N.x)
#y1 = ParticleA.pCM.dot(N.y)
x1 = BodyA.pCM.dot(N.x)
y1 = BodyA.pCM.dot(N.y)

KE = system.KE
PE = system.getPEGravity(pNA) - system.getPESprings()

pynamics.tic()
print('solving dynamics...')
f,ma = system.getdynamics()
print('creating second order function...')
func = system.state_space_post_invert(f,ma)
print('integrating...')
states=scipy.integrate.odeint(func,ini,t,rtol=1e-12,atol=1e-12,hmin=1e-14)
pynamics.toc()
print('calculating outputs..')
output = Output([x1,y1,KE-PE,qA],system)
y = output.calc(states)
pynamics.toc()

#plt.figure(1)
#plt.plot(y[:,0],y[:,1])
#plt.axis('equal')
plt.figure(1)
plt.plot(t,y[:,0])
plt.axis('equal')

plt.figure(2)
plt.plot(y[:,1]*180/pi)

plt.figure(3)
plt.plot(y[:,2]*180/pi)

plt.figure(4)
plt.plot(t,y[:,-1]*180/pi)
plt.show()

#
#import numpy.random
#
#f = f[0].simplify()
#ma = ma[0].simplify()
#
#q = y[:,-1].astype(float)
#q += numpy.random.rand(len(q))*1e-6
#q_d = (q[2:]-q[:-2])/(2*tstep)
#q_dd = (q_d[2:]-q_d[:-2])/(2*tstep)
#
#
#q = q[2:-2]
#t = t[2:-2]
#q_d = q_d[1:-1]
#
#plt.figure()
#plt.plot(t,q)
#plt.figure()
#plt.plot(t,q_d)
#plt.figure()
#plt.plot(t,q_dd)
#
#
#x = numpy.c_[q,numpy.cos(q),q_d]
#m = float((ma/qA_dd).subs(system.constants))
#y = m*q_dd
#
#C = numpy.linalg.solve(x.T.dot(x),x.T.dot(y))
#y2 = numpy.r_[[C]].dot(x.T).T
#
#plt.figure()
#plt.plot(t,y)
#plt.plot(t,y2)
<commit_msg>added file writing<commit_after># -*- coding: utf-8 -*-
"""
Written by Daniel M. Aukes
Email: danaukes<at>gmail.com
Please see LICENSE for full license.
"""

import pynamics
#pynamics.script_mode = True
from pynamics.frame import Frame
from pynamics.variable_types import Differentiable,Constant
from pynamics.system import System
from pynamics.body import Body
from pynamics.dyadic import Dyadic
from pynamics.output import Output
from pynamics.particle import Particle

#import sympy
import numpy
import scipy.integrate
import matplotlib.pyplot as plt
plt.ion()
from sympy import pi
system = System()
tol = 1e-10
lA = Constant('lA',0.05699,system)

mA = Constant('mA',0.0054531675,system)

g = Constant('g',9.81,system)
b_air = Constant('b_air',0,system)
b_joint = Constant('b_joint',0.000023855338,system)
k = Constant('k',0.016705166667,system)

Ixx_A = Constant('Ixx_A',1,system)
Iyy_A = Constant('Iyy_A',1,system)
Izz_A = Constant('Izz_A',0.00000428720184,system)

tinitial = 0
tfinal = 5
tstep = .001
t = numpy.r_[tinitial:tfinal:tstep]

preload1 = Constant('preload1',0*pi/180,system)

qA,qA_d,qA_dd = Differentiable(system,'qA')

initialvalues = {}
initialvalues[qA]=10*pi/180
initialvalues[qA_d]=0*pi/180

statevariables = system.get_q(0)+system.get_q(1)
ini = [initialvalues[item] for item in statevariables]

N = Frame('N')
A = Frame('A')

system.set_newtonian(N)
A.rotate_fixed_axis_directed(N,[0,0,1],qA,system)

pNA=0*N.x
pAB=pNA+lA*A.x
vAB=pAB.time_derivative(N,system)

#ParticleA = Particle(system,pAB,mA,'ParticleA')
IA = Dyadic.build(A,Ixx_A,Iyy_A,Izz_A)
BodyA = Body('BodyA',A,pAB,mA,IA,system)

wNA = N.getw_(A)

lab2 = vAB.dot(vAB)
uab = vAB * (1/(lab2**.5+tol))

#squared term
#system.addforce(-b_air*lab2*uab,vAB)
#linear term
system.addforce(-b_air*vAB,vAB)
system.addforce(-b_joint*wNA,wNA)
system.addforcegravity(-g*N.x)
system.add_spring_force(k,(qA-preload1)*N.z,wNA)

#x1 = ParticleA.pCM.dot(N.x)
#y1 = ParticleA.pCM.dot(N.y)
x1 = BodyA.pCM.dot(N.x)
y1 = BodyA.pCM.dot(N.y)

KE = system.KE
PE = system.getPEGravity(pNA) - system.getPESprings()

pynamics.tic()
print('solving dynamics...')
f,ma = system.getdynamics()
print('creating second order function...')
func = system.state_space_post_invert(f,ma)
print('integrating...')
states=scipy.integrate.odeint(func,ini,t,rtol=1e-12,atol=1e-12,hmin=1e-14)
pynamics.toc()
print('calculating outputs..')
output = Output([x1,y1,KE-PE,qA],system)
y = output.calc(states)
pynamics.toc()

#plt.figure(1)
#plt.plot(y[:,0],y[:,1])
#plt.axis('equal')
plt.figure(1)
plt.plot(t,y[:,0])
plt.axis('equal')

plt.figure(2)
plt.plot(y[:,1]*180/pi)

plt.figure(3)
plt.plot(y[:,2]*180/pi)

plt.figure(4)
plt.plot(t,y[:,-1]*180/pi)
plt.show()

#
#import numpy.random
#
#f = f[0].simplify()
#ma = ma[0].simplify()
#
q = y[:,-1].astype(float)
q += numpy.random.rand(len(q))*1e-6
q_d = (q[2:]-q[:-2])/(2*tstep)
q_dd = (q_d[2:]-q_d[:-2])/(2*tstep)
#
#
q = q[2:-2]
t = t[2:-2]
q_d = q_d[1:-1]
#
#plt.figure()
#plt.plot(t,q)
#plt.figure()
#plt.plot(t,q_d)
#plt.figure()
#plt.plot(t,q_dd)
#
#
#x = numpy.c_[q,numpy.cos(q),q_d]
#m = float((ma/qA_dd).subs(system.constants))
#y = m*q_dd
#
#C = numpy.linalg.solve(x.T.dot(x),x.T.dot(y))
#y2 = numpy.r_[[C]].dot(x.T).T
#
#plt.figure()
#plt.plot(t,y)
#plt.plot(t,y2)

lines = []
for item in zip(t,q,q_d,q_dd):
    lines.append('{0:0.5e},{1:0.5e},{2:0.5e},{3:0.5e}\n'.format(*item))

with open('output.csv','w') as f:
    f.writelines(lines)<|endoftext|>


<commit_before>import sys
import glob
import re

NEW_VERSION = sys.argv[1]

with open('VERSION') as f:
    VERSION=f.read()

print NEW_VERSION
print VERSION

def set_assemblyinfo_version(version):

    aLineRe = "AssemblyVersion|AssemblyFileVersion\(\"([\.0-9]+)\"\)"
    aVersionRe = "(\d\.\d\.\d)"

    print "Changing version numbers in AssemblyInfo.cs files"

    for name in glob.glob("./*/Properties/AssemblyInfo.cs"):
        print "Working on " + name

        new_file = []

        with open(name) as f:
            for line in f.readlines():
                if line.startswith("//") != True:
                    reProg = re.compile(aLineRe)
                    result = reProg.search(line)
                    if (result != None):
                        line = re.sub(aVersionRe, NEW_VERSION, line)
                new_file.append(line)

        with open(name, "w") as f:
            f.write("".join(new_file))

def set_version_file(version):
    print "Working on VERSION file"
    with open("VERSION", "w") as f:
        f.write(version)

set_verion_file(NEW_VERSION)
set_assemblyinfo_version(NEW_VERSION)<commit_msg>Revert "working on version file"<commit_after>import sys
import glob
import re

NEW_VERSION = sys.argv[1]

with open('VERSION') as f:
    VERSION=f.read()

print NEW_VERSION
print VERSION

def set_assemblyinfo_version(version):

    aLineRe = "AssemblyVersion|AssemblyFileVersion\(\"([\.0-9]+)\"\)"
    aVersionRe = "(\d\.\d\.\d)"

    print "Changing version numbers in AssemblyInfo.cs files"

    for name in glob.glob("./*/Properties/AssemblyInfo.cs"):
        print "Working on " + name

        new_file = []

        with open(name) as f:
            for line in f.readlines():
                if line.startswith("//") != True:
                    reProg = re.compile(aLineRe)
                    result = reProg.search(line)
                    if (result != None):
                        line = re.sub(aVersionRe, NEW_VERSION, line)
                new_file.append(line)

        with open(name, "w") as f:
            f.write("".join(new_file))

set_assemblyinfo_version(NEW_VERSION)<|endoftext|>


<commit_before># -*- coding: utf-8 -*-

"""
Copyright [2009-2020] EMBL-European Bioinformatics Institute
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 typing as ty
from collections import OrderedDict

import attr
from attr.validators import instance_of as is_a
from gffutils import Feature
from intervaltree import Interval

from rnacentral_pipeline.databases.data.regions import Exon
from rnacentral_pipeline.databases.data.databases import Database
from rnacentral_pipeline.rnacentral.ftp_export.coordinates.bed import BedEntry

from .extent import Extent
from .rna_type import RnaType

@attr.s(auto_attribs=True, frozen=True, slots=True)
class QaInfo:
    has_issue: bool
    incomplete_sequence: bool
    possible_contamination: bool
    missing_rfam_match: bool

    @classmethod
    def build(cls, raw):
        return cls(
            has_issue=raw["has_issue"],
            incomplete_sequence=raw["incomplete_sequence"],
            possible_contamination=raw["possible_contamination"],
            missing_rfam_match=raw["missing_rfam_match"],
        )

    def as_tuple(self):
        return (
            self.has_issue,
            self.incomplete_sequence,
            self.possible_contamination,
            self.missing_rfam_match,
        )

@attr.s(frozen=True, slots=True, hash=True)
class Count:
    mapped_count = attr.ib(validator=is_a(int))
    given_count = attr.ib(validator=is_a(int))
    total_count = attr.ib(validator=is_a(int))

    @classmethod
    def from_json(cls, raw: ty.Dict[str, int]):
        return cls(
            mapped_count=raw["mapped_count"],
            given_count=raw["given_count"],
            total_count=raw["total_count"],
        )

    def __attrs_post_init__(self):
        if self.total_count <= 0:
            raise ValueError("Must have >0 total counts")
        if self.mapped_count + self.given_count != self.total_count:
            raise ValueError("Counts do not sum")

@attr.s(frozen=True, slots=True, hash=True)
class LocationInfo:
    id: int = attr.ib(validator=is_a(int))
    urs_taxid: str = attr.ib(validator=is_a(str))
    extent: Extent = attr.ib(validator=is_a(Extent))
    exons: ty.Tuple[Exon] = attr.ib(validator=is_a(tuple))
    region_name: str = attr.ib(validator=is_a(str))
    rna_type: RnaType = attr.ib(validator=is_a(RnaType))
    qa: QaInfo = attr.ib(validator=is_a(QaInfo))
    providing_databases: ty.Tuple[Database] = attr.ib(validator=is_a(tuple))
    databases: ty.Tuple[Database] = attr.ib(validator=is_a(tuple))
    counts: Count = attr.ib(validator=is_a(Count))

    @classmethod
    def build(cls, context, raw: ty.Dict[str, ty.Any]):
        rna_type = RnaType.build(
            raw["insdc_rna_type"], raw["so_rna_type"], context.ontology
        )
        return cls(
            id=raw["region_id"],
            urs_taxid=raw["urs_taxid"],
            extent=Extent.build(raw),
            exons=tuple(sorted(Exon.from_dict(e) for e in raw["exons"])),
            region_name=raw["region_name"],
            rna_type=rna_type,
            qa=QaInfo.build(raw["qa"][0]),
            providing_databases=tuple(Database.build(d) for d in raw["providing_databases"]),
            databases=tuple(Database.build(d) for d in raw["databases"]),
            counts=context.count_for(raw["urs_taxid"]),
        )

    @property
    def start(self):
        return self.extent.start

    @property
    def stop(self):
        return self.extent.stop

    @property
    def name(self):
        return self.region_name

    def has_introns(self):
        return len(self.exons) > 1

    def is_rfam_only(self):
        inferred = {Database.genecards, Database.malacards, Database.pirbase}
        dbs = set(self.databases) - inferred
        return dbs == set([Database.rfam])

    def as_interval(self) -> Interval:
        return Interval(self.extent.start, self.extent.stop, self.id)

    def as_bed(self):
        region = self.extent.as_region()
        region = attr.assoc(region, exons=self.exons)
        yield BedEntry(
            rna_id=self.urs_taxid,
            rna_type=self.rna_type.so_term,
            databases=",".join(sorted(d.pretty() for d in self.providing_databases)),
            region=region,
        )

    def as_features(self):
        transcript_name = self.urs_taxid
        yield Feature(
            seqid=self.extent.chromosome,
            source="RNAcentral",
            featuretype="transcript",
            start=self.extent.start,
            end=self.extent.stop,
            strand=self.extent.string_strand(),
            frame=".",
            attributes=OrderedDict(
                [
                    ("Name", [transcript_name]),
                    ("type", [self.rna_type.so_term]),
                ]
            ),
        )

        for index, exon in enumerate(self.exons):
            exon_name = transcript_name + f":ncRNA_exon{index + 1}"
            yield Feature(
                seqid=self.extent.chromosome,
                source="RNAcentral",
                featuretype="noncoding_exon",
                start=exon.start,
                end=exon.stop,
                strand=self.extent.string_strand(),
                frame=".",
                attributes=OrderedDict(
                    [
                        ("Name", [exon_name]),
                        ("Parent", [transcript_name]),
                        ("type", [self.rna_type.so_term]),
                    ]
                ),
            )

    def as_writeable(self, status=None):
        yield [
            self.extent.assembly,
            self.id,
            self.urs_taxid,
            status,
        ]
<commit_msg>Improve gene naming for BED files<commit_after># -*- coding: utf-8 -*-

"""
Copyright [2009-2020] EMBL-European Bioinformatics Institute
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 re
import typing as ty
from collections import OrderedDict

import attr
from attr.validators import instance_of as is_a
from gffutils import Feature
from intervaltree import Interval

from rnacentral_pipeline.databases.data.regions import Exon
from rnacentral_pipeline.databases.data.databases import Database
from rnacentral_pipeline.rnacentral.ftp_export.coordinates.bed import BedEntry

from .extent import Extent
from .rna_type import RnaType

@attr.s(auto_attribs=True, frozen=True, slots=True)
class QaInfo:
    has_issue: bool
    incomplete_sequence: bool
    possible_contamination: bool
    missing_rfam_match: bool

    @classmethod
    def build(cls, raw):
        return cls(
            has_issue=raw["has_issue"],
            incomplete_sequence=raw["incomplete_sequence"],
            possible_contamination=raw["possible_contamination"],
            missing_rfam_match=raw["missing_rfam_match"],
        )

    def as_tuple(self):
        return (
            self.has_issue,
            self.incomplete_sequence,
            self.possible_contamination,
            self.missing_rfam_match,
        )

@attr.s(frozen=True, slots=True, hash=True)
class Count:
    mapped_count = attr.ib(validator=is_a(int))
    given_count = attr.ib(validator=is_a(int))
    total_count = attr.ib(validator=is_a(int))

    @classmethod
    def from_json(cls, raw: ty.Dict[str, int]):
        return cls(
            mapped_count=raw["mapped_count"],
            given_count=raw["given_count"],
            total_count=raw["total_count"],
        )

    def __attrs_post_init__(self):
        if self.total_count <= 0:
            raise ValueError("Must have >0 total counts")
        if self.mapped_count + self.given_count != self.total_count:
            raise ValueError("Counts do not sum")

@attr.s(frozen=True, slots=True, hash=True)
class LocationInfo:
    id: int = attr.ib(validator=is_a(int))
    urs_taxid: str = attr.ib(validator=is_a(str))
    extent: Extent = attr.ib(validator=is_a(Extent))
    exons: ty.Tuple[Exon] = attr.ib(validator=is_a(tuple))
    region_name: str = attr.ib(validator=is_a(str))
    rna_type: RnaType = attr.ib(validator=is_a(RnaType))
    qa: QaInfo = attr.ib(validator=is_a(QaInfo))
    providing_databases: ty.Tuple[Database] = attr.ib(validator=is_a(tuple))
    databases: ty.Tuple[Database] = attr.ib(validator=is_a(tuple))
    counts: Count = attr.ib(validator=is_a(Count))

    @classmethod
    def build(cls, context, raw: ty.Dict[str, ty.Any]):
        rna_type = RnaType.build(
            raw["insdc_rna_type"], raw["so_rna_type"], context.ontology
        )
        return cls(
            id=raw["region_id"],
            urs_taxid=raw["urs_taxid"],
            extent=Extent.build(raw),
            exons=tuple(sorted(Exon.from_dict(e) for e in raw["exons"])),
            region_name=raw["region_name"],
            rna_type=rna_type,
            qa=QaInfo.build(raw["qa"][0]),
            providing_databases=tuple(Database.build(d) for d in raw["providing_databases"]),
            databases=tuple(Database.build(d) for d in raw["databases"]),
            counts=context.count_for(raw["urs_taxid"]),
        )

    @property
    def start(self) -> int:
        return self.extent.start

    @property
    def stop(self) -> int:
        return self.extent.stop

    @property
    def name(self) -> str:
        return self.region_name

    @property
    def short_urs_suffix(self) -> str:
        urs, _taxid = self.urs_taxid.split("_")
        return re.sub("URS0*", "", urs)

    def has_introns(self):
        return len(self.exons) > 1

    def is_rfam_only(self):
        inferred = {Database.genecards, Database.malacards, Database.pirbase}
        dbs = set(self.databases) - inferred
        return dbs == set([Database.rfam])

    def as_interval(self) -> Interval:
        return Interval(self.extent.start, self.extent.stop, self.id)

    def as_bed(self) -> ty.Iterable[BedEntry]:
        region = self.extent.as_region()
        region = attr.assoc(region, exons=self.exons)
        yield BedEntry(
            rna_id=self.urs_taxid,
            rna_type=self.rna_type.so_term,
            databases=",".join(sorted(d.pretty() for d in self.providing_databases)),
            region=region,
        )

    def as_features(self) -> ty.Iterable[Feature]:
        transcript_name = self.urs_taxid
        yield Feature(
            seqid=self.extent.chromosome,
            source="RNAcentral",
            featuretype="transcript",
            start=self.extent.start,
            end=self.extent.stop,
            strand=self.extent.string_strand(),
            frame=".",
            attributes=OrderedDict(
                [
                    ("Name", [transcript_name]),
                    ("type", [self.rna_type.so_term]),
                ]
            ),
        )

        for index, exon in enumerate(self.exons):
            exon_name = transcript_name + f":ncRNA_exon{index + 1}"
            yield Feature(
                seqid=self.extent.chromosome,
                source="RNAcentral",
                featuretype="noncoding_exon",
                start=exon.start,
                end=exon.stop,
                strand=self.extent.string_strand(),
                frame=".",
                attributes=OrderedDict(
                    [
                        ("Name", [exon_name]),
                        ("Parent", [transcript_name]),
                        ("type", [self.rna_type.so_term]),
                    ]
                ),
            )

    def as_writeable(self, status=None) -> ty.Iterable[ty.List[str]]:
        yield [
            self.extent.assembly,
            self.id,
            self.urs_taxid,
            status,
        ]
<|endoftext|>