<commit_before>#
# This is stress longevity test that runs light weight transactions in parallel with differnt node operations:
# disruptive and not disruptive
#
# Schema definition yaml: data_dir/c-s_lwt_basic.yaml
#
# After the test is finished will be performed the data validation.
#
# 2 kinds of LWT updated will be run:
#
# update one column:
#    set lwt_indicator=30000000,
#    condition: for all rows where if lwt_indicator < 0
#
# update two columns:
#    set lwt_indicator=20000000 and author='text',
#    condition: for all rows where if lwt_indicator > 0 and lwt_indicator <= 1000000 and author != 'text'
#
# Additional expected that all rows where lwt_indicator > 1000000 and lwt_indicator < 20000000 won't be updated
#
# Based on this 3 types of validation will be performed:
#
# - Rows which are expected to stay intact, are not updated
# - Rows for a certain row range, lwt_indicator is changed to 30000000
# - Rows for another row range lwt_indicator is changed to 20000000
#
#
# ***Validation implementation***
#
# 1. Rows which are expected to stay intact.
#    To be able to validate that, I added a MV called blogposts_not_updated_lwt_indicator
#
#     - create MATERIALIZED VIEW blogposts_not_updated_lwt_indicator as select lwt_indicator, author from blogposts
#       where domain is not null and published_date is not null and lwt_indicator > 1000000 and lwt_indicator < 20000000
#       PRIMARY KEY(lwt_indicator, domain, published_date);
#     This MV hold all rows that shouldn't be updated.
#
#     Once the prepare_write_cmd part will be completed, all data from the view
#     blogposts_not_updated_lwt_indicator will be copied to a side table called
#     blogposts_not_updated_lwt_indicator_expect (created by test). This table uses as the expected data for this
#     validation.
#
#     When test is finished, the test checks that data in the blogposts_not_updated_lwt_indicator and
#     blogposts_not_updated_lwt_indicator_expect will be same.
#
# 2. For the rows where lwt_indicator was changed to "30000000", the data validation behaves as follow:
#
#     Two more Materialized View are added.
#     First one holds rows that candidates for this update (before the update):
#
#     create MATERIALIZED VIEW blogposts_update_one_column_lwt_indicator as select domain, lwt_indicator, author
#     from blogposts where domain is not null and published_date is not null and lwt_indicator < 0
#     PRIMARY KEY(lwt_indicator, domain, published_date);
#     Second one holds rows with lwt_indicator=30000000 (means - after update):
#
#     create MATERIALIZED VIEW blogposts_update_one_column_lwt_indicator_upd as select lwt_indicator, author
#     from blogposts where domain is not null and published_date is not null and lwt_indicator = 30000000
#     PRIMARY KEY(lwt_indicator, domain, published_date);
#     Ideally, The view blogposts_update_one_column_lwt_indicator should be empty after the update and the view
#     blogposts_update_one_column_lwt_indicator_upd will hold same amount of data as
#     blogposts_update_one_column_lwt_indicator had before the update.
#
#     However, due to the fact we cannot control the c-s workload to go over all the rows (should visit all the rows
#     with lwt_indicator negative value), for now it's just validated that blogposts_update_one_column_lwt_indicator_upd
#     has some records.
#     (It's unreasonable that the workload won't visit any row with negative value)
#
# 3. For the rows where lwt_indicator is changed to "20000000", the data validation the data validation behaves as
#    follow:
#
#     Two more Materialized View are added.
#     First one holds rows that are candidates for this update (before update):
#
#      create MATERIALIZED VIEW blogposts_update_two_columns_lwt_indicator as select lwt_indicator, author
#      from blogposts where domain is not null and published_date is not null and lwt_indicator > 0 and
#      lwt_indicator <= 1000000 PRIMARY KEY(lwt_indicator, domain, published_date);
#     Second one holds rows with lwt_indicator=20000000 (means - after the update):
#
#     create MATERIALIZED VIEW blogposts_update_two_columns_lwt_indicator_upd as select lwt_indicator, author
#     from blogposts where domain is not null and published_date is not null and lwt_indicator = 20000000
#     PRIMARY KEY(lwt_indicator, domain, published_date);

#     Ideally, the view blogposts_update_two_columns_lwt_indicator should be empty after the update and the view
#     blogposts_update_two_columns_lwt_indicator_upd will hold the same amount of data that as the view
#     blogposts_update_two_columns_lwt_indicator had before the update.
#
#     However, due to the same reason as above, the current validation is just that
#     blogposts_update_two_columns_lwt_indicator_upd has some records.
#

import re

import time
from longevity_test import LongevityTest

from sdcm.utils.common import get_profile_content

class LWTLongevityTest(LongevityTest):
    def __init__(self, *args):
        super(LWTLongevityTest, self).__init__(*args)
        self._keyspace_name = None
        self._mv_for_not_updated_data = None
        self._expected_data_table = None
        self._validate_not_updated_data = True

    @property
    def keyspace_name(self):
        if not self._keyspace_name:
            prepare_write_cmd = self.params.get('prepare_write_cmd', default=[])
            profiles = [cmd for cmd in prepare_write_cmd if 'c-s_lwt' in cmd]
            if not profiles:
                self._validate_not_updated_data = False
                self.log.warning('Keyspace is not recognized. '
                                 'Data validation of not updated rows won\' be performed')
                return self._keyspace_name
            cs_profile = profiles[0]
            _, profile = get_profile_content(cs_profile)
            self._keyspace_name = profile['keyspace']

        return self._keyspace_name

    @property
    def mv_for_not_updated_data(self):
        """
        Get MV name that holds rows which are expected to stay intact
        """
        if not self._mv_for_not_updated_data:
            name_substr = '_not_updated'
            self._mv_for_not_updated_data = self.get_mv_name_from_profile(name_substr)
            return self._mv_for_not_updated_data

        return self._mv_for_not_updated_data

    @property
    def expected_data_table_name(self):
        """
        Table name for expected data (needs for validate rows which are expected to stay intact - implemented in
         validate_range_not_expected_to_change function)
        """
        if not self._expected_data_table:
            self._expected_data_table = '%s_expect' % self.mv_for_not_updated_data
        return self._expected_data_table

    def get_mv_name_from_profile(self, name_substr, cs_profile=None):
        mv_name = None

        if not cs_profile:
            prepare_write_cmd = self.params.get('prepare_write_cmd', default=[])
            cs_profile = [cmd for cmd in prepare_write_cmd if 'c-s_lwt' in cmd]

        if not cs_profile:
            return mv_name

        if not isinstance(cs_profile, list):
            cs_profile = [cs_profile]

        for profile in cs_profile:
            _, profile_content = get_profile_content(profile)
            mv_create_cmd = self.get_mv_cmd_from_profile(profile_content, name_substr)
            if mv_create_cmd:
                mv_name = self.get_mv_name_from_stress_cmd(mv_create_cmd, name_substr)
                break

        return mv_name

    @staticmethod
    def get_mv_cmd_from_profile(profile_content, name_substr):
        all_mvs = profile_content['extra_definitions']
        mv_cmd = [cmd for cmd in all_mvs if name_substr in cmd]

        mv_cmd = mv_cmd[0] if mv_cmd else ''
        return mv_cmd

    @staticmethod
    def get_mv_name_from_stress_cmd(mv_create_cmd, name_substr):
        find_mv_name = re.search(r'materialized view (.*%s.*) as' % name_substr, mv_create_cmd, re.I)
        return find_mv_name.group(1)

    def copy_expected_data(self):
        if self._validate_not_updated_data:
            if not self.copy_view(src_keyspace=self.keyspace_name, src_view=self.mv_for_not_updated_data,
                                  dest_keyspace=self.keyspace_name, dest_table=self.expected_data_table_name,
                                  copy_data=True):
                self._validate_not_updated_data = False
                self.log.warning('Problem during copying expected data. '
                                 'Data validation of not updated rows won\' be performed')

    def run_prepare_write_cmd(self):
        super(LWTLongevityTest, self).run_prepare_write_cmd()

        # TODO: Temporary print. Will be removed
        self.log.debug('Get rows count in {} MV before sleep'.format(self.mv_for_not_updated_data))
        self.get_rows_count(self.db_cluster.nodes[0])

        # Wait for MVs data will be fully inserted (running on background)
        time.sleep(300)

        # Run repair on all nodes. We need it to fix the case, when part of data wasn't inserted because of timeouts
        self.stop_nemesis_and_repair_cluster()

        self.copy_expected_data()

        # Run nemesis during stress as it was stopped before copy expected data
        if self.params.get('nemesis_during_prepare'):
            self.start_nemesis()

    def start_nemesis(self):
        self.db_cluster.termination_event.clear()
        self.db_cluster.start_nemesis()

    def stop_nemesis_and_repair_cluster(self):
        # Stop nemesis. Repair on all nodes will be run before data validation, so all nodes should be up
        # and also prevent case to run repair from nemesis in parallel
        if self.db_cluster.nemesis_threads:
            self.db_cluster.stop_nemesis()

        # TODO: Temporary print. Will be removed
        self.log.debug('Get rows count in {} MV after sleep and before repair'.format(self.mv_for_not_updated_data))
        self.get_rows_count(self.db_cluster.nodes[0])

        for node in self.db_cluster.nodes:
            self.log.debug('Start nodetool repair on {} node'.format(node.name))
            node.run_nodetool("repair")

        # TODO: Temporary print. Will be removed
        self.log.debug('Get rows count in {} MV after repair'.format(self.mv_for_not_updated_data))
        self.get_rows_count(self.db_cluster.nodes[0])

    def test_lwt_longevity(self):
        self.test_custom_time()

        # Run repair on all nodes. We need it to fix the case, when part of data wasn't inserted because of timeouts
        self.stop_nemesis_and_repair_cluster()
        self.validate_data()

    def validate_data(self):
        node = self.db_cluster.nodes[0]
        with self.cql_connection_patient(node, keyspace='cqlstress_lwt_example') as session:
            self.validate_range_not_expected_to_change(session=session)
            self.validate_updated_data(session=session)

    def validate_range_not_expected_to_change(self, session):
        """
        Part of data in the user profile table shouldn't be updated using LWT.
        This data will be saved in the materialized view with "not_updated" substring in the name
        After prepare write all data from this materialized view will be saved in the separate table as expected result.
        During stress (after prepare) LWT update statements will be run for a few hours.
        When updates will be finished this function verifies that data in "not_updated" MV and expected result table
        is same
        """
        if self._validate_not_updated_data and self.mv_for_not_updated_data and self.expected_data_table_name:
            self.log.debug('Verify not updated rows')
            # Get all rows, not use pagination
            session.default_fetch_size = 0

            actual_result = self.rows_to_list(session.execute("SELECT * FROM %s" % self.mv_for_not_updated_data))
            expected_result = self.rows_to_list(session.execute("SELECT * FROM %s" % self.expected_data_table_name))

            self.assertEqual(len(actual_result), len(expected_result),
                             'One or more rows are not as expected, suspected LWT wrong update. '
                             'Actual dataset length: {}, Expected dataset length: {}'.format(len(actual_result),
                                                                                len(expected_result)))

            self.assertEqual(actual_result, expected_result,
                             'One or more rows are not as expected, suspected LWT wrong update')

    def validate_updated_data(self, session):
        """
        In user profile 'data_dir/c-s_lwt_basic.yaml' LWT updates the lwt_indicator and author columns with hard coded
        values. Validate, that we can find these values
        Two more materialized views are added. The first one holds rows that are candidates for the update
        (i.e. all rows before the update).
        The second one holds rows with lwt_indicator=30000000 (i.e. only the updated rows)
        """
        stress_cmd = self.params.get('stress_cmd', default=[])
        stress_read_cmd = self.params.get('stress_read_cmd', default=[])

        one_column_run_update = [cmd for cmd in stress_cmd+stress_read_cmd if 'lwt_update_one_column' in cmd]
        two_columns_run_update = [cmd for cmd in stress_cmd+stress_read_cmd if 'lwt_update_two_columns' in cmd]

        if one_column_run_update or two_columns_run_update:
            self.log.info('Verify that rows were updated')

            if one_column_run_update:
                expected_data_mv_name = self.get_mv_name_from_profile(name_substr='one_column_lwt_indicator_upd',
                                                                      cs_profile=one_column_run_update)
                result = session.execute('SELECT * from {name}'.format(name=expected_data_mv_name))

                self.assertTrue(result, 'Data in the column "lwt_indicator" was not updated by light weight '
                                        'transaction')

            if two_columns_run_update:
                expected_data_mv_name = self.get_mv_name_from_profile(name_substr='two_columns_lwt_indicator_upd',
                                                                      cs_profile=two_columns_run_update)
                result = session.execute('SELECT * from {name}'.format(name=expected_data_mv_name))

                self.assertTrue(result, 'Data in the columns "lwt_indicator" and "author" were not updated by light '
                                'weight transaction')
<commit_msg>fix(lightweight trunsaction): not fail the test when more immutable rows<commit_after>#
# This is stress longevity test that runs light weight transactions in parallel with differnt node operations:
# disruptive and not disruptive
#
# Schema definition yaml: data_dir/c-s_lwt_basic.yaml
#
# After the test is finished will be performed the data validation.
#
# 2 kinds of LWT updated will be run:
#
# update one column:
#    set lwt_indicator=30000000,
#    condition: for all rows where if lwt_indicator < 0
#
# update two columns:
#    set lwt_indicator=20000000 and author='text',
#    condition: for all rows where if lwt_indicator > 0 and lwt_indicator <= 1000000 and author != 'text'
#
# Additional expected that all rows where lwt_indicator > 1000000 and lwt_indicator < 20000000 won't be updated
#
# Based on this 3 types of validation will be performed:
#
# - Rows which are expected to stay intact, are not updated
# - Rows for a certain row range, lwt_indicator is changed to 30000000
# - Rows for another row range lwt_indicator is changed to 20000000
#
#
# ***Validation implementation***
#
# 1. Rows which are expected to stay intact.
#    To be able to validate that, I added a MV called blogposts_not_updated_lwt_indicator
#
#     - create MATERIALIZED VIEW blogposts_not_updated_lwt_indicator as select lwt_indicator, author from blogposts
#       where domain is not null and published_date is not null and lwt_indicator > 1000000 and lwt_indicator < 20000000
#       PRIMARY KEY(lwt_indicator, domain, published_date);
#     This MV hold all rows that shouldn't be updated.
#
#     Once the prepare_write_cmd part will be completed, all data from the view
#     blogposts_not_updated_lwt_indicator will be copied to a side table called
#     blogposts_not_updated_lwt_indicator_expect (created by test). This table uses as the expected data for this
#     validation.
#
#     When test is finished, the test checks that data in the blogposts_not_updated_lwt_indicator and
#     blogposts_not_updated_lwt_indicator_expect will be same.
#
# 2. For the rows where lwt_indicator was changed to "30000000", the data validation behaves as follow:
#
#     Two more Materialized View are added.
#     First one holds rows that candidates for this update (before the update):
#
#     create MATERIALIZED VIEW blogposts_update_one_column_lwt_indicator as select domain, lwt_indicator, author
#     from blogposts where domain is not null and published_date is not null and lwt_indicator < 0
#     PRIMARY KEY(lwt_indicator, domain, published_date);
#     Second one holds rows with lwt_indicator=30000000 (means - after update):
#
#     create MATERIALIZED VIEW blogposts_update_one_column_lwt_indicator_upd as select lwt_indicator, author
#     from blogposts where domain is not null and published_date is not null and lwt_indicator = 30000000
#     PRIMARY KEY(lwt_indicator, domain, published_date);
#     Ideally, The view blogposts_update_one_column_lwt_indicator should be empty after the update and the view
#     blogposts_update_one_column_lwt_indicator_upd will hold same amount of data as
#     blogposts_update_one_column_lwt_indicator had before the update.
#
#     However, due to the fact we cannot control the c-s workload to go over all the rows (should visit all the rows
#     with lwt_indicator negative value), for now it's just validated that blogposts_update_one_column_lwt_indicator_upd
#     has some records.
#     (It's unreasonable that the workload won't visit any row with negative value)
#
# 3. For the rows where lwt_indicator is changed to "20000000", the data validation the data validation behaves as
#    follow:
#
#     Two more Materialized View are added.
#     First one holds rows that are candidates for this update (before update):
#
#      create MATERIALIZED VIEW blogposts_update_two_columns_lwt_indicator as select lwt_indicator, author
#      from blogposts where domain is not null and published_date is not null and lwt_indicator > 0 and
#      lwt_indicator <= 1000000 PRIMARY KEY(lwt_indicator, domain, published_date);
#     Second one holds rows with lwt_indicator=20000000 (means - after the update):
#
#     create MATERIALIZED VIEW blogposts_update_two_columns_lwt_indicator_upd as select lwt_indicator, author
#     from blogposts where domain is not null and published_date is not null and lwt_indicator = 20000000
#     PRIMARY KEY(lwt_indicator, domain, published_date);

#     Ideally, the view blogposts_update_two_columns_lwt_indicator should be empty after the update and the view
#     blogposts_update_two_columns_lwt_indicator_upd will hold the same amount of data that as the view
#     blogposts_update_two_columns_lwt_indicator had before the update.
#
#     However, due to the same reason as above, the current validation is just that
#     blogposts_update_two_columns_lwt_indicator_upd has some records.
#

import re

import time
from longevity_test import LongevityTest

from sdcm.utils.common import get_profile_content

class LWTLongevityTest(LongevityTest):
    def __init__(self, *args):
        super(LWTLongevityTest, self).__init__(*args)
        self._keyspace_name = None
        self._mv_for_not_updated_data = None
        self._expected_data_table = None
        self._validate_not_updated_data = True

    @property
    def keyspace_name(self):
        if not self._keyspace_name:
            prepare_write_cmd = self.params.get('prepare_write_cmd', default=[])
            profiles = [cmd for cmd in prepare_write_cmd if 'c-s_lwt' in cmd]
            if not profiles:
                self._validate_not_updated_data = False
                self.log.warning('Keyspace is not recognized. '
                                 'Data validation of not updated rows won\' be performed')
                return self._keyspace_name
            cs_profile = profiles[0]
            _, profile = get_profile_content(cs_profile)
            self._keyspace_name = profile['keyspace']

        return self._keyspace_name

    @property
    def mv_for_not_updated_data(self):
        """
        Get MV name that holds rows which are expected to stay intact
        """
        if not self._mv_for_not_updated_data:
            name_substr = '_not_updated'
            self._mv_for_not_updated_data = self.get_mv_name_from_profile(name_substr)
            return self._mv_for_not_updated_data

        return self._mv_for_not_updated_data

    @property
    def expected_data_table_name(self):
        """
        Table name for expected data (needs for validate rows which are expected to stay intact - implemented in
         validate_range_not_expected_to_change function)
        """
        if not self._expected_data_table:
            self._expected_data_table = '%s_expect' % self.mv_for_not_updated_data
        return self._expected_data_table

    def get_mv_name_from_profile(self, name_substr, cs_profile=None):
        mv_name = None

        if not cs_profile:
            prepare_write_cmd = self.params.get('prepare_write_cmd', default=[])
            cs_profile = [cmd for cmd in prepare_write_cmd if 'c-s_lwt' in cmd]

        if not cs_profile:
            return mv_name

        if not isinstance(cs_profile, list):
            cs_profile = [cs_profile]

        for profile in cs_profile:
            _, profile_content = get_profile_content(profile)
            mv_create_cmd = self.get_mv_cmd_from_profile(profile_content, name_substr)
            if mv_create_cmd:
                mv_name = self.get_mv_name_from_stress_cmd(mv_create_cmd, name_substr)
                break

        return mv_name

    @staticmethod
    def get_mv_cmd_from_profile(profile_content, name_substr):
        all_mvs = profile_content['extra_definitions']
        mv_cmd = [cmd for cmd in all_mvs if name_substr in cmd]

        mv_cmd = mv_cmd[0] if mv_cmd else ''
        return mv_cmd

    @staticmethod
    def get_mv_name_from_stress_cmd(mv_create_cmd, name_substr):
        find_mv_name = re.search(r'materialized view (.*%s.*) as' % name_substr, mv_create_cmd, re.I)
        return find_mv_name.group(1)

    def copy_expected_data(self):
        if self._validate_not_updated_data:
            if not self.copy_view(src_keyspace=self.keyspace_name, src_view=self.mv_for_not_updated_data,
                                  dest_keyspace=self.keyspace_name, dest_table=self.expected_data_table_name,
                                  copy_data=True):
                self._validate_not_updated_data = False
                self.log.warning('Problem during copying expected data. '
                                 'Data validation of not updated rows won\' be performed')

    def run_prepare_write_cmd(self):
        super(LWTLongevityTest, self).run_prepare_write_cmd()

        # TODO: Temporary print. Will be removed
        self.log.debug('Get rows count in {} MV before sleep'.format(self.mv_for_not_updated_data))
        self.get_rows_count(self.db_cluster.nodes[0])

        # Wait for MVs data will be fully inserted (running on background)
        time.sleep(300)

        # Run repair on all nodes. We need it to fix the case, when part of data wasn't inserted because of timeouts
        self.stop_nemesis_and_repair_cluster()

        self.copy_expected_data()

        # Run nemesis during stress as it was stopped before copy expected data
        if self.params.get('nemesis_during_prepare'):
            self.start_nemesis()

    def start_nemesis(self):
        self.db_cluster.termination_event.clear()
        self.db_cluster.start_nemesis()

    def stop_nemesis_and_repair_cluster(self):
        # Stop nemesis. Repair on all nodes will be run before data validation, so all nodes should be up
        # and also prevent case to run repair from nemesis in parallel
        if self.db_cluster.nemesis_threads:
            self.db_cluster.stop_nemesis()

        # TODO: Temporary print. Will be removed
        self.log.debug('Get rows count in {} MV after sleep and before repair'.format(self.mv_for_not_updated_data))
        self.get_rows_count(self.db_cluster.nodes[0])

        for node in self.db_cluster.nodes:
            self.log.debug('Start nodetool repair on {} node'.format(node.name))
            node.run_nodetool("repair")

        # TODO: Temporary print. Will be removed
        self.log.debug('Get rows count in {} MV after repair'.format(self.mv_for_not_updated_data))
        self.get_rows_count(self.db_cluster.nodes[0])

    def test_lwt_longevity(self):
        self.test_custom_time()

        # Run repair on all nodes. We need it to fix the case, when part of data wasn't inserted because of timeouts
        self.stop_nemesis_and_repair_cluster()
        self.validate_data()

    def validate_data(self):
        node = self.db_cluster.nodes[0]
        with self.cql_connection_patient(node, keyspace='cqlstress_lwt_example') as session:
            self.validate_range_not_expected_to_change(session=session)
            self.validate_updated_data(session=session)

    def validate_range_not_expected_to_change(self, session):
        """
        Part of data in the user profile table shouldn't be updated using LWT.
        This data will be saved in the materialized view with "not_updated" substring in the name
        After prepare write all data from this materialized view will be saved in the separate table as expected result.
        During stress (after prepare) LWT update statements will be run for a few hours.
        When updates will be finished this function verifies that data in "not_updated" MV and expected result table
        is same
        """
        if self._validate_not_updated_data and self.mv_for_not_updated_data and self.expected_data_table_name:
            self.log.debug('Verify not updated rows')
            # Get all rows, not use pagination
            session.default_fetch_size = 0

            actual_result = self.rows_to_list(session.execute("SELECT * FROM %s" % self.mv_for_not_updated_data))
            expected_result = self.rows_to_list(session.execute("SELECT * FROM %s" % self.expected_data_table_name))

            # Issue https://github.com/scylladb/scylla/issues/6181
            # Not fail the test if unexpected additional rows where found in actual result table
            if len(actual_result) > len(expected_result):
                self.log.warning('Actual dataset length {} more '
                                 'then expected dataset length: {}. '
                                 'Issue #6181'.format(len(actual_result),
                                                      len(expected_result)))
            else:

                self.assertEqual(len(actual_result), len(expected_result),
                                 'One or more rows are not as expected, suspected LWT wrong update. '
                                 'Actual dataset length: {}, Expected dataset length: {}'.format(len(actual_result),
                                                                                len(expected_result)))

                self.assertEqual(actual_result, expected_result,
                                 'One or more rows are not as expected, suspected LWT wrong update')

    def validate_updated_data(self, session):
        """
        In user profile 'data_dir/c-s_lwt_basic.yaml' LWT updates the lwt_indicator and author columns with hard coded
        values. Validate, that we can find these values
        Two more materialized views are added. The first one holds rows that are candidates for the update
        (i.e. all rows before the update).
        The second one holds rows with lwt_indicator=30000000 (i.e. only the updated rows)
        """
        stress_cmd = self.params.get('stress_cmd', default=[])
        stress_read_cmd = self.params.get('stress_read_cmd', default=[])

        one_column_run_update = [cmd for cmd in stress_cmd+stress_read_cmd if 'lwt_update_one_column' in cmd]
        two_columns_run_update = [cmd for cmd in stress_cmd+stress_read_cmd if 'lwt_update_two_columns' in cmd]

        if one_column_run_update or two_columns_run_update:
            self.log.info('Verify that rows were updated')

            if one_column_run_update:
                expected_data_mv_name = self.get_mv_name_from_profile(name_substr='one_column_lwt_indicator_upd',
                                                                      cs_profile=one_column_run_update)
                result = session.execute('SELECT * from {name}'.format(name=expected_data_mv_name))

                self.assertTrue(result, 'Data in the column "lwt_indicator" was not updated by light weight '
                                        'transaction')

            if two_columns_run_update:
                expected_data_mv_name = self.get_mv_name_from_profile(name_substr='two_columns_lwt_indicator_upd',
                                                                      cs_profile=two_columns_run_update)
                result = session.execute('SELECT * from {name}'.format(name=expected_data_mv_name))

                self.assertTrue(result, 'Data in the columns "lwt_indicator" and "author" were not updated by light '
                                'weight transaction')
<|endoftext|>


<commit_before>from django.core.management.base import BaseCommand
from django.utils.six.moves import input

from optparse import make_option

from cms.models import Post, Category, Localisation
from cms.git.models import GitPage, GitCategory

from html2text import html2text

class Command(BaseCommand):
    help = 'import all the current content from a github repository'

    option_list = BaseCommand.option_list + (
        make_option(
            '--quiet',
            action='store_true',
            dest='quiet',
            default=False,
            help='imports the data using default arguments'),
    )

    def handle(self, *args, **options):
        quiet = options.get('quiet')

        if not quiet:
            must_delete = self.get_input_data(
                'Do you want to delete existing data? Y/n: ', 'y')
        else:
            must_delete = 'y'

        if must_delete.lower() == 'y':
            print 'deleting existing content..'
            Post.objects.all().delete()
            Category.objects.all().delete()

        print 'creating categories..'
        categories = list(GitCategory.all())

        for instance in categories:
            Category.objects.create(
                slug=instance.slug,
                title=instance.title,
                subtitle=instance.subtitle,
                localisation=Localisation._for(instance.language),
                featured_in_navbar=instance.featured_in_navbar,
                uuid=instance.uuid
            )

        # second pass to add related fields
        for instance in categories:
            if instance.source:
                p = Category.objects.get(uuid=instance.uuid)
                p.source = Category.objects.get(uuid=instance.source.uuid)
                p.save()

        print 'creating pages..'
        pages = list(GitPage.all())

        for instance in pages:
            primary_category = None
            if instance.primary_category:
                primary_category = Category.objects.get(
                    uuid=instance.primary_category.uuid)
            Post.objects.create(
                title=instance.title,
                subtitle=instance.subtitle,
                slug=instance.slug,
                description=instance.description,
                content=html2text(instance.content),
                created_at=instance.created_at,
                modified_at=instance.modified_at,
                featured_in_category=instance.featured_in_category,
                featured=instance.featured,
                localisation=Localisation._for(instance.language),
                primary_category=primary_category,
                uuid=instance.uuid
            )

        # second pass to add related fields
        for instance in pages:
            if instance.source:
                p = Post.objects.get(uuid=instance.uuid)
                p.source = Post.objects.get(uuid=instance.source.uuid)
                p.save()

            if instance.linked_pages:
                p = Post.objects.get(uuid=instance.uuid)
                p.related_posts.add(*list(
                    Post.objects.filter(uuid__in=instance.linked_pages)))
                p.save()

        print 'done.'

    def get_input_data(self, message, default=None):
        raw_value = input(message)
        if default and raw_value == '':
            raw_value = default

        return raw_value.lower()
<commit_msg>add some error handling when importing<commit_after>from django.core.management.base import BaseCommand
from django.utils.six.moves import input

from optparse import make_option

from cms.models import Post, Category, Localisation
from cms.git.models import GitPage, GitCategory

from html2text import html2text

class Command(BaseCommand):
    help = 'import all the current content from a github repository'

    option_list = BaseCommand.option_list + (
        make_option(
            '--quiet',
            action='store_true',
            dest='quiet',
            default=False,
            help='imports the data using default arguments'),
    )

    def handle(self, *args, **options):
        quiet = options.get('quiet')

        if not quiet:
            must_delete = self.get_input_data(
                'Do you want to delete existing data? Y/n: ', 'y')
        else:
            must_delete = 'y'

        if must_delete.lower() == 'y':
            print 'deleting existing content..'
            Post.objects.all().delete()
            Category.objects.all().delete()

        print 'creating categories..'
        categories = list(GitCategory.all())

        for instance in categories:
            Category.objects.create(
                slug=instance.slug,
                title=instance.title,
                subtitle=instance.subtitle,
                localisation=Localisation._for(instance.language),
                featured_in_navbar=instance.featured_in_navbar,
                uuid=instance.uuid
            )

        # second pass to add related fields
        for instance in categories:
            if instance.source:
                p = Category.objects.get(uuid=instance.uuid)
                p.source = Category.objects.get(uuid=instance.source.uuid)
                p.save()

        print 'creating pages..'
        pages = list(GitPage.all())

        for instance in pages:
            primary_category = None
            if instance.primary_category:
                primary_category = Category.objects.get(
                    uuid=instance.primary_category.uuid)
            try:
                Post.objects.create(
                    title=instance.title,
                    subtitle=instance.subtitle,
                    slug=instance.slug,
                    description=instance.description,
                    content=html2text(instance.content),
                    created_at=instance.created_at,
                    modified_at=instance.modified_at,
                    featured_in_category=instance.featured_in_category,
                    featured=instance.featured,
                    localisation=Localisation._for(instance.language),
                    primary_category=primary_category,
                    uuid=instance.uuid
                )
            except Exception, e:
                print 'An error occured with: %s(%s)' % (
                    instance.title, instance.uuid)
                print e

        # second pass to add related fields
        for instance in pages:
            if instance.source:
                p = Post.objects.get(uuid=instance.uuid)
                p.source = Post.objects.get(uuid=instance.source.uuid)
                p.save()

            if instance.linked_pages:
                p = Post.objects.get(uuid=instance.uuid)
                p.related_posts.add(*list(
                    Post.objects.filter(uuid__in=instance.linked_pages)))
                p.save()

        print 'done.'

    def get_input_data(self, message, default=None):
        raw_value = input(message)
        if default and raw_value == '':
            raw_value = default

        return raw_value.lower()
<|endoftext|>


<commit_before># -*- coding: utf-8 -*-
import json
import math
import numpy as np
from flask import make_response

from py3dtiles.feature_table import (FeatureTableHeader, FeatureTableBody, FeatureTable)
from py3dtiles.tile import TileBody, TileHeader, Tile

from .utils import (
    read_uncompressed_patch, boundingbox_to_polygon, list_from_str, patch_nbpoints_unc
)
from .conf import Config
from .database import Session

LOD_MIN = 0
LOD_MAX = 5
LOD_LEN = LOD_MAX + 1 - LOD_MIN

def ThreeDTilesInfo(table, column):

    session = Session(table, column)
    # bounding box
    box = session.boundingbox

    # number of points for the first patch
    npoints = session.approx_row_count * session.patch_size

    # srs
    srs = session.srs

    # build json
    return {
        "bounds": [box['xmin'], box['ymin'], box['zmin'],
                   box['xmax'], box['ymax'], box['zmax']],
        "numPoints": npoints,
        "srs": srs
    }

def ThreeDTilesRead(table, column, bounds, lod):

    session = Session(table, column)
    # offsets = [round(off, 2) for off in list_from_str(offsets)]
    box = list_from_str(bounds)
    # requested = [scales, offsets]
    stored_patches = session.lopocstable.filter_stored_output()
    schema = stored_patches['point_schema']
    pcid = stored_patches['pcid']
    # scales = [scale] * 3
    scales = stored_patches['scales']
    offsets = stored_patches['offsets']

    [tile, npoints] = get_points(session, box, lod, offsets, pcid, scales, schema)

    if Config.DEBUG:
        tile.sync()
        print("NPOINTS: ", npoints)

    # build flask response
    response = make_response(tile.to_array().tostring())
    response.headers['content-type'] = 'application/octet-stream'
    return response

# FIXME: change pdt type in case of quantized values
cdt = np.dtype([('Red', np.uint8), ('Green', np.uint8), ('Blue', np.uint8)])
pdt = np.dtype([('X', np.float32), ('Y', np.float32), ('Z', np.float32)])

def get_points(session, box, lod, offsets, pcid, scales, schema):
    sql = sql_query(session, box, pcid, lod)
    if Config.DEBUG:
        print(sql)

    pcpatch_wkb = session.query(sql)[0][0]
    points, npoints = read_uncompressed_patch(pcpatch_wkb, schema)

    if max(points['Red']) > 255:
        # normalize
        rgb_reduced = np.c_[points['Red'] % 255, points['Green'] % 255, points['Blue'] % 255]
        rgb = np.array(np.core.records.fromarrays(rgb_reduced.T, dtype=cdt))
    else:
        rgb = points[['Red', 'Green', 'Blue']].astype(cdt)

    quantized_points_r = np.c_[
        points['X'] * scales[0],
        points['Y'] * scales[1],
        points['Z'] * scales[2]
    ]

    quantized_points = np.array(np.core.records.fromarrays(quantized_points_r.T, dtype=pdt))

    fth = FeatureTableHeader.from_dtype(
        quantized_points.dtype, rgb.dtype, npoints
    )
    ftb = FeatureTableBody()
    ftb.positions_itemsize = fth.positions_dtype.itemsize
    ftb.colors_itemsize = fth.colors_dtype.itemsize
    ftb.positions_arr = quantized_points.view(np.uint8)
    ftb.colors_arr = rgb.view(np.uint8)

    ft = FeatureTable()
    ft.header = fth
    ft.body = ftb

    # tile
    tb = TileBody()
    tb.feature_table = ft
    th = TileHeader()
    tile = Tile()
    tile.body = tb
    tile.header = th
    tile.body.feature_table.header.rtc = offsets

    return [tile, npoints]

def sql_query(session, box, pcid, lod):
    poly = boundingbox_to_polygon(box)

    maxppp = session.lopocstable.max_points_per_patch
    patch_size = session.patch_size

    if maxppp:
        range_min = 1
        range_max = maxppp
    else:
        # FIXME: may skip some points if patch_size/lod_len is decimal
        # we need to fix either here or at loading with the patch_size and lod bounds
        range_min = lod * int(patch_size / LOD_LEN) + 1
        range_max = (lod + 1) * int(patch_size / LOD_LEN)

    # build the sql query
    sql_limit = ""
    maxppq = session.lopocstable.max_patches_per_query
    if maxppq:
        sql_limit = " limit {0} ".format(maxppq)

    if Config.USE_MORTON:
        sql = ("select pc_union("
               "pc_filterbetween( "
               "pc_range({0}, {4}, {5}), 'Z', {6}, {7} )) from "
               "(select {0} from {1} "
               "where pc_intersects({0}, st_geomfromtext('polygon (("
               "{2}))',{3})) order by morton {8})_;"
               .format(session.column, session.table,
                       poly, session.srsid, range_min, range_max,
                       box[2] - 0.1, box[5] + 0.1, sql_limit,
                       pcid))
    else:
        sql = ("select pc_compress(pc_transform(pc_union("
               "pc_filterbetween( "
               "pc_range({0}, {4}, {5}), 'Z', {6}, {7} )), {9}), 'laz') from "
               "(select {0} from {1} where pc_intersects({0}, "
               "st_geomfromtext('polygon (({2}))',{3})) {8})_;"
               .format(session.column, session.table,
                       poly, session.srsid, range_min, range_max,
                       box[2], box[5], sql_limit,
                       pcid))

    return sql

def buildbox(bbox):
    width = bbox[3] - bbox[0]
    depth = bbox[4] - bbox[1]
    height = bbox[5] - bbox[2]
    midx = bbox[0] + width / 2
    midy = bbox[1] + depth / 2
    midz = bbox[2] + height / 2

    box = [midx, midy, midz]
    box.append(width / 2.0)
    box.append(0.0)
    box.append(0.0)
    box.append(0.0)
    box.append(depth / 2.0)
    box.append(0.0)
    box.append(0.0)
    box.append(0.0)
    box.append(height / 2.0)
    return box

def build_hierarchy_from_pg(session, baseurl, bbox):

    stored_patches = session.lopocstable.filter_stored_output()
    pcid = stored_patches['pcid']
    offsets = stored_patches['offsets']
    tileset = {}
    tileset["asset"] = {"version": "0.0"}
    tileset["geometricError"] = math.sqrt(
        (bbox[3] - bbox[0]) ** 2 + (bbox[4] - bbox[1]) ** 2 + (bbox[5] - bbox[2]) ** 2
    )
    if Config.DEBUG:
        print('tileset geometricErroc', tileset["geometricError"])

    bvol = {}
    bvol["box"] = buildbox(bbox)

    lod_str = "lod={0}".format(LOD_MIN)
    bounds = ("bounds=[{0},{1},{2},{3},{4},{5}]"
              .format(bbox[0], bbox[1], bbox[2], bbox[3], bbox[4], bbox[5]))
    resource = "{}.{}".format(session.table, session.column)

    base_url = "{0}/3dtiles/{1}/read.pnts".format(baseurl, resource)
    url = (
        "{0}?{1}&{2}"
        .format(base_url, lod_str, bounds)
    )

    GEOMETRIC_ERROR = tileset["geometricError"]

    root = {}
    root["refine"] = "add"
    root["boundingVolume"] = bvol
    root["geometricError"] = GEOMETRIC_ERROR / 20
    root["content"] = {"url": url}

    lod = 1
    children_list = []
    for bb in split_bbox(bbox):
        json_children = children(
            session, baseurl, offsets, bb, lod, pcid, GEOMETRIC_ERROR / 40
        )
        if len(json_children):
            children_list.append(json_children)

    if len(children_list):
        root["children"] = children_list

    tileset["root"] = root

    return json.dumps(tileset, indent=2, separators=(',', ': '))

def build_children_section(session, baseurl, offsets, bbox, err, lod):

    cjson = {}

    lod = "lod={0}".format(lod)
    bounds = ("bounds=[{0},{1},{2},{3},{4},{5}]"
              .format(bbox[0], bbox[1], bbox[2], bbox[3], bbox[4], bbox[5]))

    resource = "{}.{}".format(session.table, session.column)
    baseurl = "{0}/3dtiles/{1}/read.pnts".format(baseurl, resource)
    url = "{0}?{1}&{2}".format(baseurl, lod, bounds)

    bvol = {}
    bvol["box"] = buildbox(bbox)

    cjson["boundingVolume"] = bvol
    cjson["geometricError"] = err
    cjson["content"] = {"url": url}

    return cjson

def split_bbox(bbox):
    width = bbox[3] - bbox[0]
    length = bbox[4] - bbox[1]
    height = bbox[5] - bbox[2]

    up = bbox[5]
    middle = up - height / 2
    down = bbox[2]

    x = bbox[0]
    y = bbox[1]

    bbox_nwd = [x, y + length / 2, down, x + width / 2, y + length, middle]
    bbox_nwu = [x, y + length / 2, middle, x + width / 2, y + length, up]
    bbox_ned = [x + width / 2, y + length / 2, down, x + width, y + length, middle]
    bbox_neu = [x + width / 2, y + length / 2, middle, x + width, y + length, up]
    bbox_swd = [x, y, down, x + width / 2, y + length / 2, middle]
    bbox_swu = [x, y, middle, x + width / 2, y + length / 2, up]
    bbox_sed = [x + width / 2, y, down, x + width, y + length / 2, middle]
    bbox_seu = [x + width / 2, y, middle, x + width, y + length / 2, up]

    return [bbox_nwd, bbox_nwu, bbox_ned, bbox_neu, bbox_swd, bbox_swu,
            bbox_sed, bbox_seu]

def children(session, baseurl, offsets, bbox, lod, pcid, err):

    # run sql
    sql = sql_query(session, bbox, pcid, lod)
    pcpatch_wkb = session.query(sql)[0][0]

    json_me = {}
    if lod <= LOD_MAX and pcpatch_wkb:
        npoints = patch_nbpoints_unc(pcpatch_wkb)
        if npoints > 0:
            json_me = build_children_section(session, baseurl, offsets, bbox, err, lod)

        lod += 1

        children_list = []
        if lod <= LOD_MAX:
            for bb in split_bbox(bbox):
                json_children = children(
                    session, baseurl, offsets, bb, lod, pcid, err / 2
                )

                if len(json_children):
                    children_list.append(json_children)

        if len(children_list):
            json_me["children"] = children_list

    return json_me
<commit_msg>don't crash if no RGB in data<commit_after># -*- coding: utf-8 -*-
import json
import math

import numpy as np
from flask import make_response

from py3dtiles.feature_table import (FeatureTableHeader, FeatureTableBody, FeatureTable)
from py3dtiles.tile import TileBody, TileHeader, Tile

from .utils import (
    read_uncompressed_patch, boundingbox_to_polygon, list_from_str, patch_nbpoints_unc
)
from .conf import Config
from .database import Session

LOD_MIN = 0
LOD_MAX = 5
LOD_LEN = LOD_MAX + 1 - LOD_MIN

def ThreeDTilesInfo(table, column):

    session = Session(table, column)
    # bounding box
    box = session.boundingbox

    # number of points for the first patch
    npoints = session.approx_row_count * session.patch_size

    # srs
    srs = session.srs

    # build json
    return {
        "bounds": [box['xmin'], box['ymin'], box['zmin'],
                   box['xmax'], box['ymax'], box['zmax']],
        "numPoints": npoints,
        "srs": srs
    }

def ThreeDTilesRead(table, column, bounds, lod):

    session = Session(table, column)
    # offsets = [round(off, 2) for off in list_from_str(offsets)]
    box = list_from_str(bounds)
    # requested = [scales, offsets]
    stored_patches = session.lopocstable.filter_stored_output()
    schema = stored_patches['point_schema']
    pcid = stored_patches['pcid']
    # scales = [scale] * 3
    scales = stored_patches['scales']
    offsets = stored_patches['offsets']

    [tile, npoints] = get_points(session, box, lod, offsets, pcid, scales, schema)

    if Config.DEBUG:
        tile.sync()
        print("NPOINTS: ", npoints)

    # build flask response
    response = make_response(tile.to_array().tostring())
    response.headers['content-type'] = 'application/octet-stream'
    return response

# FIXME: change pdt type in case of quantized values
cdt = np.dtype([('Red', np.uint8), ('Green', np.uint8), ('Blue', np.uint8)])
pdt = np.dtype([('X', np.float32), ('Y', np.float32), ('Z', np.float32)])

def get_points(session, box, lod, offsets, pcid, scales, schema):
    sql = sql_query(session, box, pcid, lod)
    if Config.DEBUG:
        print(sql)

    pcpatch_wkb = session.query(sql)[0][0]
    points, npoints = read_uncompressed_patch(pcpatch_wkb, schema)

    if 'Red' in points:
        if max(points['Red']) > 255:
            # normalize
            rgb_reduced = np.c_[points['Red'] % 255, points['Green'] % 255, points['Blue'] % 255]
            rgb = np.array(np.core.records.fromarrays(rgb_reduced.T, dtype=cdt))
        else:
            rgb = points[['Red', 'Green', 'Blue']].astype(cdt)
    else:
        # No colors
        # FIXME: compute color gradient based on elevation or use classification
        rgb_reduced = np.c_[
            np.array([0] * npoints),
            np.array([0] * npoints),
            np.array([0] * npoints)
        ]
        rgb = np.array(np.core.records.fromarrays(rgb_reduced.T, dtype=cdt))

    quantized_points_r = np.c_[
        points['X'] * scales[0],
        points['Y'] * scales[1],
        points['Z'] * scales[2]
    ]

    quantized_points = np.array(np.core.records.fromarrays(quantized_points_r.T, dtype=pdt))

    fth = FeatureTableHeader.from_dtype(
        quantized_points.dtype, rgb.dtype, npoints
    )
    ftb = FeatureTableBody()
    ftb.positions_itemsize = fth.positions_dtype.itemsize
    ftb.colors_itemsize = fth.colors_dtype.itemsize
    ftb.positions_arr = quantized_points.view(np.uint8)
    ftb.colors_arr = rgb.view(np.uint8)

    ft = FeatureTable()
    ft.header = fth
    ft.body = ftb

    # tile
    tb = TileBody()
    tb.feature_table = ft
    th = TileHeader()
    tile = Tile()
    tile.body = tb
    tile.header = th
    tile.body.feature_table.header.rtc = offsets

    return [tile, npoints]

def sql_query(session, box, pcid, lod):
    poly = boundingbox_to_polygon(box)

    maxppp = session.lopocstable.max_points_per_patch
    patch_size = session.patch_size

    if maxppp:
        range_min = 1
        range_max = maxppp
    else:
        # FIXME: may skip some points if patch_size/lod_len is decimal
        # we need to fix either here or at loading with the patch_size and lod bounds
        range_min = lod * int(patch_size / LOD_LEN) + 1
        range_max = (lod + 1) * int(patch_size / LOD_LEN)

    # build the sql query
    sql_limit = ""
    maxppq = session.lopocstable.max_patches_per_query
    if maxppq:
        sql_limit = " limit {0} ".format(maxppq)

    if Config.USE_MORTON:
        sql = ("select pc_union("
               "pc_filterbetween( "
               "pc_range({0}, {4}, {5}), 'Z', {6}, {7} )) from "
               "(select {0} from {1} "
               "where pc_intersects({0}, st_geomfromtext('polygon (("
               "{2}))',{3})) order by morton {8})_;"
               .format(session.column, session.table,
                       poly, session.srsid, range_min, range_max,
                       box[2] - 0.1, box[5] + 0.1, sql_limit,
                       pcid))
    else:
        sql = ("select pc_compress(pc_transform(pc_union("
               "pc_filterbetween( "
               "pc_range({0}, {4}, {5}), 'Z', {6}, {7} )), {9}), 'laz') from "
               "(select {0} from {1} where pc_intersects({0}, "
               "st_geomfromtext('polygon (({2}))',{3})) {8})_;"
               .format(session.column, session.table,
                       poly, session.srsid, range_min, range_max,
                       box[2], box[5], sql_limit,
                       pcid))

    return sql

def buildbox(bbox):
    width = bbox[3] - bbox[0]
    depth = bbox[4] - bbox[1]
    height = bbox[5] - bbox[2]
    midx = bbox[0] + width / 2
    midy = bbox[1] + depth / 2
    midz = bbox[2] + height / 2

    box = [midx, midy, midz]
    box.append(width / 2.0)
    box.append(0.0)
    box.append(0.0)
    box.append(0.0)
    box.append(depth / 2.0)
    box.append(0.0)
    box.append(0.0)
    box.append(0.0)
    box.append(height / 2.0)
    return box

def build_hierarchy_from_pg(session, baseurl, bbox):

    stored_patches = session.lopocstable.filter_stored_output()
    pcid = stored_patches['pcid']
    offsets = stored_patches['offsets']
    tileset = {}
    tileset["asset"] = {"version": "0.0"}
    tileset["geometricError"] = math.sqrt(
        (bbox[3] - bbox[0]) ** 2 + (bbox[4] - bbox[1]) ** 2 + (bbox[5] - bbox[2]) ** 2
    )
    if Config.DEBUG:
        print('tileset geometricErroc', tileset["geometricError"])

    bvol = {}
    bvol["box"] = buildbox(bbox)

    lod_str = "lod={0}".format(LOD_MIN)
    bounds = ("bounds=[{0},{1},{2},{3},{4},{5}]"
              .format(bbox[0], bbox[1], bbox[2], bbox[3], bbox[4], bbox[5]))
    resource = "{}.{}".format(session.table, session.column)

    base_url = "{0}/3dtiles/{1}/read.pnts".format(baseurl, resource)
    url = (
        "{0}?{1}&{2}"
        .format(base_url, lod_str, bounds)
    )

    GEOMETRIC_ERROR = tileset["geometricError"]

    root = {}
    root["refine"] = "add"
    root["boundingVolume"] = bvol
    root["geometricError"] = GEOMETRIC_ERROR / 20
    root["content"] = {"url": url}

    lod = 1
    children_list = []
    for bb in split_bbox(bbox):
        json_children = children(
            session, baseurl, offsets, bb, lod, pcid, GEOMETRIC_ERROR / 40
        )
        if len(json_children):
            children_list.append(json_children)

    if len(children_list):
        root["children"] = children_list

    tileset["root"] = root

    return json.dumps(tileset, indent=2, separators=(',', ': '))

def build_children_section(session, baseurl, offsets, bbox, err, lod):

    cjson = {}

    lod = "lod={0}".format(lod)
    bounds = ("bounds=[{0},{1},{2},{3},{4},{5}]"
              .format(bbox[0], bbox[1], bbox[2], bbox[3], bbox[4], bbox[5]))

    resource = "{}.{}".format(session.table, session.column)
    baseurl = "{0}/3dtiles/{1}/read.pnts".format(baseurl, resource)
    url = "{0}?{1}&{2}".format(baseurl, lod, bounds)

    bvol = {}
    bvol["box"] = buildbox(bbox)

    cjson["boundingVolume"] = bvol
    cjson["geometricError"] = err
    cjson["content"] = {"url": url}

    return cjson

def split_bbox(bbox):
    width = bbox[3] - bbox[0]
    length = bbox[4] - bbox[1]
    height = bbox[5] - bbox[2]

    up = bbox[5]
    middle = up - height / 2
    down = bbox[2]

    x = bbox[0]
    y = bbox[1]

    bbox_nwd = [x, y + length / 2, down, x + width / 2, y + length, middle]
    bbox_nwu = [x, y + length / 2, middle, x + width / 2, y + length, up]
    bbox_ned = [x + width / 2, y + length / 2, down, x + width, y + length, middle]
    bbox_neu = [x + width / 2, y + length / 2, middle, x + width, y + length, up]
    bbox_swd = [x, y, down, x + width / 2, y + length / 2, middle]
    bbox_swu = [x, y, middle, x + width / 2, y + length / 2, up]
    bbox_sed = [x + width / 2, y, down, x + width, y + length / 2, middle]
    bbox_seu = [x + width / 2, y, middle, x + width, y + length / 2, up]

    return [bbox_nwd, bbox_nwu, bbox_ned, bbox_neu, bbox_swd, bbox_swu,
            bbox_sed, bbox_seu]

def children(session, baseurl, offsets, bbox, lod, pcid, err):

    # run sql
    sql = sql_query(session, bbox, pcid, lod)
    pcpatch_wkb = session.query(sql)[0][0]

    json_me = {}
    if lod <= LOD_MAX and pcpatch_wkb:
        npoints = patch_nbpoints_unc(pcpatch_wkb)
        if npoints > 0:
            json_me = build_children_section(session, baseurl, offsets, bbox, err, lod)

        lod += 1

        children_list = []
        if lod <= LOD_MAX:
            for bb in split_bbox(bbox):
                json_children = children(
                    session, baseurl, offsets, bb, lod, pcid, err / 2
                )

                if len(json_children):
                    children_list.append(json_children)

        if len(children_list):
            json_me["children"] = children_list

    return json_me
<|endoftext|>


<commit_before>from numpy import linspace, zeros, array, meshgrid, abs, empty, arange, \
    int32, unravel_index
from multiprocessing import Pool

from cobra.solvers import solver_dict, get_solver_name

# attempt to import plotting libraries
try:
    from matplotlib import pyplot
    from mpl_toolkits.mplot3d import axes3d
except ImportError:
    pyplot = None
    axes3d = None
mlab = None  # mayavi may crash python
try:  # for prettier colors
    from brewer2mpl import get_map
except:
    get_map = None

class phenotypePhasePlaneData:
    """class to hold results of a phenotype phase plane analysis"""
    def __init__(self,
            reaction1_name, reaction2_name,
            reaction1_range_max, reaction2_range_max,
            reaction1_npoints, reaction2_npoints):
        self.reaction1_name = reaction1_name
        self.reaction2_name = reaction2_name
        self.reaction1_range_max = reaction1_range_max
        self.reaction2_range_max = reaction2_range_max
        self.reaction1_npoints = reaction1_npoints
        self.reaction2_npoints = reaction2_npoints
        self.reaction1_fluxes = linspace(0, reaction1_range_max,
            reaction1_npoints)
        self.reaction2_fluxes = linspace(0, reaction2_range_max,
            reaction2_npoints)
        self.growth_rates = zeros((reaction1_npoints, reaction2_npoints))
        self.shadow_prices1 = zeros((reaction1_npoints, reaction2_npoints))
        self.shadow_prices2 = zeros((reaction1_npoints, reaction2_npoints))
        self.segments = zeros(self.growth_rates.shape, dtype=int32)
        self.phases = []

    def plot(self):
        """plot the phenotype phase plane in 3D using any available backend"""
        if pyplot is not None:
            self.plot_matplotlib()
        elif mlab is not None:
            self.plot_mayavi()
        else:
            raise ImportError("No suitable 3D plotting package found")

    def plot_matplotlib(self, theme="Paired", scale_grid=False):
        """Use matplotlib to plot a phenotype phase plane in 3D.

        theme: color theme to use (requires brewer2mpl)

        returns: maptlotlib 3d subplot object"""
        if pyplot is None:
            raise ImportError("Error importing matplotlib 3D plotting")
        colors = empty(self.growth_rates.shape, dtype="|S7")
        n_segments = self.segments.max()
        # pick colors
        if get_map is None:
            color_list = ['#A6CEE3', '#1F78B4', '#B2DF8A', '#33A02C',
                          '#FB9A99', '#E31A1C', '#FDBF6F', '#FF7F00',
                          '#CAB2D6', '#6A3D9A', '#FFFF99', '#B15928']
        else:
            color_list = get_map(theme, 'Qualitative', n_segments).hex_colors
        if n_segments > len(color_list):
            from warnings import warn
            warn("not enough colors to color all detected phases")
        if n_segments > 0 and n_segments <= len(color_list):
            for i in range(n_segments):
                colors[data.segments == (i + 1)] = color_list[i]
        else:
            colors[:, :] = 'b'
        if scale_grid:
            # grid wires should not have more than ~20 points
            xgrid_scale = int(self.reaction1_npoints / 20)
            ygrid_scale = int(self.reaction2_npoints / 20)
        else:
            xgrid_scale, ygrid_scale = (1, 1)
        figure = pyplot.figure()
        xgrid, ygrid = meshgrid(self.reaction1_fluxes, self.reaction2_fluxes)
        axes = figure.add_subplot(111, projection="3d")
        xgrid = xgrid.transpose()
        ygrid = ygrid.transpose()
        axes.plot_surface(xgrid, ygrid, self.growth_rates, rstride=1,
            cstride=1, facecolors=colors, linewidth=0, antialiased=False)
        axes.plot_wireframe(xgrid, ygrid, self.growth_rates, color="black", rstride=xgrid_scale, cstride=ygrid_scale)
        axes.set_xlabel(self.reaction1_name)
        axes.set_ylabel(self.reaction2_name)
        axes.set_zlabel("Growth rates")
        axes.view_init(elev=30, azim=-135)
        return axes

    def plot_mayavi(self):
        """Use mayavi to plot a phenotype phase plane in 3D.
        The resulting figure will be quick to interact with in real time,
        but might be difficult to save as a vector figure.
        returns: mlab figure object"""
        from mayavi import mlab
        figure = mlab.figure(bgcolor=(1, 1, 1), fgcolor=(0, 0, 0))
        figure.name = "Phenotype Phase Plane"
        max = 10.0
        xmax = self.reaction1_fluxes.max()
        ymax = self.reaction2_fluxes.max()
        zmax = self.growth_rates.max()
        xgrid, ygrid = meshgrid(self.reaction1_fluxes, self.reaction2_fluxes)
        xgrid = xgrid.transpose()
        ygrid = ygrid.transpose()
        xscale = max / xmax
        yscale = max / ymax
        zscale = max / zmax
        mlab.surf(xgrid * xscale, ygrid * yscale, self.growth_rates * zscale,
            representation="wireframe", color=(0, 0, 0), figure=figure)
        mlab.mesh(xgrid * xscale, ygrid * yscale, self.growth_rates * zscale,
            scalars=self.shadow_prices1 + self.shadow_prices2, resolution=1,
            representation="surface", opacity=0.75, figure=figure)
        # draw axes
        mlab.outline(extent=(0, max, 0, max, 0, max))
        mlab.axes(opacity=0, ranges=[0, xmax, 0, ymax, 0, zmax])
        mlab.xlabel(self.reaction1_name)
        mlab.ylabel(self.reaction2_name)
        mlab.zlabel("Growth rates")
        return figure

    def segment(self, threshold=0.01):
        """attempt to segment the data and identify the various phases"""
        self.segments *= 0
        # each entry in phases will consist of the following tuple
        # ((x, y), shadow_price1, shadow_price2)
        self.phases = []
        # initialize the area to be all False
        covered_area = (self.growth_rates * 0 == 1)
        # as long as part of the area has not been covered
        segment_id = 0
        while self.segments.min() == 0:
            segment_id += 1
            # i and j are indices for a current point which has not been
            # assigned a segment yet

            i, j = unravel_index(self.segments.argmin(), self.segments.shape)
            # update the segment id for any point with a similar shadow price
            # to the current point
            self.segments[
                (abs(self.shadow_prices1 - self.shadow_prices1[i, j])
                    < threshold) *
                (abs(self.shadow_prices2 - self.shadow_prices2[i, j])
                    < threshold)] += segment_id
            # add the current point as one of the phases
            self.phases.append((
                (self.reaction1_fluxes[i], self.reaction2_fluxes[j]),
                self.shadow_prices1[i, j], self.shadow_prices2[i, j]))

def _calculate_subset(arguments):
    """Calculate a subset of the phenotype phase plane data.
    Store each result tuple as:
    (i, j, growth_rate, shadow_price1, shadow_price2)"""

    model = arguments["model"]
    reaction1_fluxes = arguments["reaction1_fluxes"]
    reaction2_fluxes = arguments["reaction2_fluxes"]
    metabolite1_name = arguments["metabolite1_name"]
    metabolite2_name = arguments["metabolite2_name"]
    index1 = arguments["index1"]
    index2 = arguments["index2"]
    i_list = arguments["i_list"]
    j_list = arguments["j_list"]
    tolerance = arguments["tolerance"]
    solver = solver_dict[arguments["solver"]]

    results = []
    reaction1 = model.reactions[index1]
    reaction2 = model.reactions[index2]
    problem = solver.create_problem(model)
    solver.solve_problem(problem)
    for a, flux1 in enumerate(reaction1_fluxes):
        i = i_list[a]
        # flux is actually negative for uptake. Also some solvers require
        # float instead of numpy.float64
        flux1 = float(-1 * flux1)
        # change bounds on reaction 1
        solver.change_variable_bounds(problem, index1, flux1 - tolerance, flux1 + tolerance)
        for b, flux2 in enumerate(reaction2_fluxes):
            j = j_list[b]
            flux2 = float(-1 * flux2)  # same story as flux1
            # change bounds on reaction 2
            solver.change_variable_bounds(problem, index2, flux2 - tolerance, flux2 + tolerance)
            # solve the problem and save results
            solver.solve_problem(problem)
            solution = solver.format_solution(problem, model)
            if solution is not None and solution.status == "optimal":
                results.append((i, j, solution.f,
                    solution.y_dict[metabolite1_name],
                    solution.y_dict[metabolite2_name]))
            else:
                results.append((i, j, 0, 0, 0))
            # reset reaction 2 bounds
            solver.change_variable_bounds(problem, index2, float(reaction2.lower_bound),
                                          float(reaction2.upper_bound))
        # reset reaction 1 bounds
        solver.change_variable_bounds(problem, index1, float(reaction1.lower_bound),
                                      float(reaction1.upper_bound))
    return results

def calculate_phenotype_phase_plane(model,
        reaction1_name, reaction2_name,
        reaction1_range_max=20, reaction2_range_max=20,
        reaction1_npoints=50, reaction2_npoints=50,
        solver=None, n_processes=1, tolerance=1e-6):
    """calculates the growth rates while varying the uptake rates for two
    reactions.

    returns: an object containing the growth rates for the uptake rates.
    To plot the result, call the plot function of the returned object.

    Example:
    data = calculate_phenotype_phase_plane(my_model, "EX_foo", "EX_bar")
    data.plot()
    """
    if solver is None:
        get_solver_name()
    data = phenotypePhasePlaneData(
            str(reaction1_name), str(reaction2_name),
            reaction1_range_max, reaction2_range_max,
            reaction1_npoints, reaction2_npoints)
    # find the objects for the reactions and metabolites
    index1 = model.reactions.index(model.reactions.get_by_id(data.reaction1_name))
    index2 = model.reactions.index(model.reactions.get_by_id(data.reaction2_name))
    metabolite1_name = str(model.reactions.get_by_id( \
        reaction1_name)._metabolites.keys()[0])
    metabolite2_name = str(model.reactions.get_by_id( \
        reaction2_name)._metabolites.keys()[0])
    if n_processes > reaction1_npoints:  # limit the number of processes
        n_processes = reaction1_npoints
    range_add = reaction1_npoints / n_processes
    # prepare the list of arguments for each _calculate_subset call
    arguments_list = []
    i = arange(reaction1_npoints)
    j = arange(reaction2_npoints)
    for n in range(n_processes):
        start = n * range_add
        if n != n_processes - 1:
            r1_range = data.reaction1_fluxes[start:start + range_add]
            i_list = i[start:start + range_add]
        else:
            r1_range = data.reaction1_fluxes[start:]
            i_list = i[start:]
        arguments_list.append({"model": model,
            "index1": index1, "index2": index2,
            "metabolite1_name": metabolite1_name,
            "metabolite2_name": metabolite2_name,
            "reaction1_fluxes": r1_range,
            "reaction2_fluxes": data.reaction2_fluxes.copy(),
            "i_list": i_list, "j_list": j.copy(),
            "tolerance": tolerance, "solver": solver})
    if n_processes > 1:
        p = Pool(n_processes)
        results = list(p.map(_calculate_subset, arguments_list))
    else:
        results = [_calculate_subset(arguments_list[0])]
    for result_list in results:
        for result in result_list:
            i = result[0]
            j = result[1]
            data.growth_rates[i, j] = result[2]
            data.shadow_prices1[i, j] = result[3]
            data.shadow_prices2[i, j] = result[4]
    data.segment()
    return data

if __name__ == "__main__":
    from time import time

    import cobra
    from cobra.test import ecoli_pickle, create_test_model

    n1 = 40
    n2 = 40

    model = create_test_model(ecoli_pickle)
    start_time = time()
    data = calculate_phenotype_phase_plane(model, "EX_glc_e",
        'EX_o2_e', reaction1_npoints=n1, reaction2_npoints=n2,
        n_processes=1, solver="cglpk")
    print "took %.2f seconds with 1 process" % (time() - start_time)
    #start_time = time()
    #data = calculate_phenotype_phase_plane(model, 'EX_glc_e',
    #    'EX_o2_e', reaction1_npoints=n1, reaction2_npoints=n2,
    #    n_processes=2, solver="cglpk")
    #print "took %.2f seconds with 2 processes" % (time() - start_time)
    #data.plot_matplotlib("Pastel1")
    #data.plot_matplotlib("Dark2")
    data.plot_matplotlib()
<commit_msg>relative import for phenotypePhasePlane.py<commit_after>from numpy import linspace, zeros, array, meshgrid, abs, empty, arange, \
    int32, unravel_index
from multiprocessing import Pool

from ..solvers import solver_dict, get_solver_name

# attempt to import plotting libraries
try:
    from matplotlib import pyplot
    from mpl_toolkits.mplot3d import axes3d
except ImportError:
    pyplot = None
    axes3d = None
mlab = None  # mayavi may crash python
try:  # for prettier colors
    from brewer2mpl import get_map
except:
    get_map = None

class phenotypePhasePlaneData:
    """class to hold results of a phenotype phase plane analysis"""
    def __init__(self,
            reaction1_name, reaction2_name,
            reaction1_range_max, reaction2_range_max,
            reaction1_npoints, reaction2_npoints):
        self.reaction1_name = reaction1_name
        self.reaction2_name = reaction2_name
        self.reaction1_range_max = reaction1_range_max
        self.reaction2_range_max = reaction2_range_max
        self.reaction1_npoints = reaction1_npoints
        self.reaction2_npoints = reaction2_npoints
        self.reaction1_fluxes = linspace(0, reaction1_range_max,
            reaction1_npoints)
        self.reaction2_fluxes = linspace(0, reaction2_range_max,
            reaction2_npoints)
        self.growth_rates = zeros((reaction1_npoints, reaction2_npoints))
        self.shadow_prices1 = zeros((reaction1_npoints, reaction2_npoints))
        self.shadow_prices2 = zeros((reaction1_npoints, reaction2_npoints))
        self.segments = zeros(self.growth_rates.shape, dtype=int32)
        self.phases = []

    def plot(self):
        """plot the phenotype phase plane in 3D using any available backend"""
        if pyplot is not None:
            self.plot_matplotlib()
        elif mlab is not None:
            self.plot_mayavi()
        else:
            raise ImportError("No suitable 3D plotting package found")

    def plot_matplotlib(self, theme="Paired", scale_grid=False):
        """Use matplotlib to plot a phenotype phase plane in 3D.

        theme: color theme to use (requires brewer2mpl)

        returns: maptlotlib 3d subplot object"""
        if pyplot is None:
            raise ImportError("Error importing matplotlib 3D plotting")
        colors = empty(self.growth_rates.shape, dtype="|S7")
        n_segments = self.segments.max()
        # pick colors
        if get_map is None:
            color_list = ['#A6CEE3', '#1F78B4', '#B2DF8A', '#33A02C',
                          '#FB9A99', '#E31A1C', '#FDBF6F', '#FF7F00',
                          '#CAB2D6', '#6A3D9A', '#FFFF99', '#B15928']
        else:
            color_list = get_map(theme, 'Qualitative', n_segments).hex_colors
        if n_segments > len(color_list):
            from warnings import warn
            warn("not enough colors to color all detected phases")
        if n_segments > 0 and n_segments <= len(color_list):
            for i in range(n_segments):
                colors[data.segments == (i + 1)] = color_list[i]
        else:
            colors[:, :] = 'b'
        if scale_grid:
            # grid wires should not have more than ~20 points
            xgrid_scale = int(self.reaction1_npoints / 20)
            ygrid_scale = int(self.reaction2_npoints / 20)
        else:
            xgrid_scale, ygrid_scale = (1, 1)
        figure = pyplot.figure()
        xgrid, ygrid = meshgrid(self.reaction1_fluxes, self.reaction2_fluxes)
        axes = figure.add_subplot(111, projection="3d")
        xgrid = xgrid.transpose()
        ygrid = ygrid.transpose()
        axes.plot_surface(xgrid, ygrid, self.growth_rates, rstride=1,
            cstride=1, facecolors=colors, linewidth=0, antialiased=False)
        axes.plot_wireframe(xgrid, ygrid, self.growth_rates, color="black", rstride=xgrid_scale, cstride=ygrid_scale)
        axes.set_xlabel(self.reaction1_name)
        axes.set_ylabel(self.reaction2_name)
        axes.set_zlabel("Growth rates")
        axes.view_init(elev=30, azim=-135)
        return axes

    def plot_mayavi(self):
        """Use mayavi to plot a phenotype phase plane in 3D.
        The resulting figure will be quick to interact with in real time,
        but might be difficult to save as a vector figure.
        returns: mlab figure object"""
        from mayavi import mlab
        figure = mlab.figure(bgcolor=(1, 1, 1), fgcolor=(0, 0, 0))
        figure.name = "Phenotype Phase Plane"
        max = 10.0
        xmax = self.reaction1_fluxes.max()
        ymax = self.reaction2_fluxes.max()
        zmax = self.growth_rates.max()
        xgrid, ygrid = meshgrid(self.reaction1_fluxes, self.reaction2_fluxes)
        xgrid = xgrid.transpose()
        ygrid = ygrid.transpose()
        xscale = max / xmax
        yscale = max / ymax
        zscale = max / zmax
        mlab.surf(xgrid * xscale, ygrid * yscale, self.growth_rates * zscale,
            representation="wireframe", color=(0, 0, 0), figure=figure)
        mlab.mesh(xgrid * xscale, ygrid * yscale, self.growth_rates * zscale,
            scalars=self.shadow_prices1 + self.shadow_prices2, resolution=1,
            representation="surface", opacity=0.75, figure=figure)
        # draw axes
        mlab.outline(extent=(0, max, 0, max, 0, max))
        mlab.axes(opacity=0, ranges=[0, xmax, 0, ymax, 0, zmax])
        mlab.xlabel(self.reaction1_name)
        mlab.ylabel(self.reaction2_name)
        mlab.zlabel("Growth rates")
        return figure

    def segment(self, threshold=0.01):
        """attempt to segment the data and identify the various phases"""
        self.segments *= 0
        # each entry in phases will consist of the following tuple
        # ((x, y), shadow_price1, shadow_price2)
        self.phases = []
        # initialize the area to be all False
        covered_area = (self.growth_rates * 0 == 1)
        # as long as part of the area has not been covered
        segment_id = 0
        while self.segments.min() == 0:
            segment_id += 1
            # i and j are indices for a current point which has not been
            # assigned a segment yet

            i, j = unravel_index(self.segments.argmin(), self.segments.shape)
            # update the segment id for any point with a similar shadow price
            # to the current point
            self.segments[
                (abs(self.shadow_prices1 - self.shadow_prices1[i, j])
                    < threshold) *
                (abs(self.shadow_prices2 - self.shadow_prices2[i, j])
                    < threshold)] += segment_id
            # add the current point as one of the phases
            self.phases.append((
                (self.reaction1_fluxes[i], self.reaction2_fluxes[j]),
                self.shadow_prices1[i, j], self.shadow_prices2[i, j]))

def _calculate_subset(arguments):
    """Calculate a subset of the phenotype phase plane data.
    Store each result tuple as:
    (i, j, growth_rate, shadow_price1, shadow_price2)"""

    model = arguments["model"]
    reaction1_fluxes = arguments["reaction1_fluxes"]
    reaction2_fluxes = arguments["reaction2_fluxes"]
    metabolite1_name = arguments["metabolite1_name"]
    metabolite2_name = arguments["metabolite2_name"]
    index1 = arguments["index1"]
    index2 = arguments["index2"]
    i_list = arguments["i_list"]
    j_list = arguments["j_list"]
    tolerance = arguments["tolerance"]
    solver = solver_dict[arguments["solver"]]

    results = []
    reaction1 = model.reactions[index1]
    reaction2 = model.reactions[index2]
    problem = solver.create_problem(model)
    solver.solve_problem(problem)
    for a, flux1 in enumerate(reaction1_fluxes):
        i = i_list[a]
        # flux is actually negative for uptake. Also some solvers require
        # float instead of numpy.float64
        flux1 = float(-1 * flux1)
        # change bounds on reaction 1
        solver.change_variable_bounds(problem, index1, flux1 - tolerance, flux1 + tolerance)
        for b, flux2 in enumerate(reaction2_fluxes):
            j = j_list[b]
            flux2 = float(-1 * flux2)  # same story as flux1
            # change bounds on reaction 2
            solver.change_variable_bounds(problem, index2, flux2 - tolerance, flux2 + tolerance)
            # solve the problem and save results
            solver.solve_problem(problem)
            solution = solver.format_solution(problem, model)
            if solution is not None and solution.status == "optimal":
                results.append((i, j, solution.f,
                    solution.y_dict[metabolite1_name],
                    solution.y_dict[metabolite2_name]))
            else:
                results.append((i, j, 0, 0, 0))
            # reset reaction 2 bounds
            solver.change_variable_bounds(problem, index2, float(reaction2.lower_bound),
                                          float(reaction2.upper_bound))
        # reset reaction 1 bounds
        solver.change_variable_bounds(problem, index1, float(reaction1.lower_bound),
                                      float(reaction1.upper_bound))
    return results

def calculate_phenotype_phase_plane(model,
        reaction1_name, reaction2_name,
        reaction1_range_max=20, reaction2_range_max=20,
        reaction1_npoints=50, reaction2_npoints=50,
        solver=None, n_processes=1, tolerance=1e-6):
    """calculates the growth rates while varying the uptake rates for two
    reactions.

    returns: an object containing the growth rates for the uptake rates.
    To plot the result, call the plot function of the returned object.

    Example:
    data = calculate_phenotype_phase_plane(my_model, "EX_foo", "EX_bar")
    data.plot()
    """
    if solver is None:
        get_solver_name()
    data = phenotypePhasePlaneData(
            str(reaction1_name), str(reaction2_name),
            reaction1_range_max, reaction2_range_max,
            reaction1_npoints, reaction2_npoints)
    # find the objects for the reactions and metabolites
    index1 = model.reactions.index(model.reactions.get_by_id(data.reaction1_name))
    index2 = model.reactions.index(model.reactions.get_by_id(data.reaction2_name))
    metabolite1_name = str(model.reactions.get_by_id( \
        reaction1_name)._metabolites.keys()[0])
    metabolite2_name = str(model.reactions.get_by_id( \
        reaction2_name)._metabolites.keys()[0])
    if n_processes > reaction1_npoints:  # limit the number of processes
        n_processes = reaction1_npoints
    range_add = reaction1_npoints / n_processes
    # prepare the list of arguments for each _calculate_subset call
    arguments_list = []
    i = arange(reaction1_npoints)
    j = arange(reaction2_npoints)
    for n in range(n_processes):
        start = n * range_add
        if n != n_processes - 1:
            r1_range = data.reaction1_fluxes[start:start + range_add]
            i_list = i[start:start + range_add]
        else:
            r1_range = data.reaction1_fluxes[start:]
            i_list = i[start:]
        arguments_list.append({"model": model,
            "index1": index1, "index2": index2,
            "metabolite1_name": metabolite1_name,
            "metabolite2_name": metabolite2_name,
            "reaction1_fluxes": r1_range,
            "reaction2_fluxes": data.reaction2_fluxes.copy(),
            "i_list": i_list, "j_list": j.copy(),
            "tolerance": tolerance, "solver": solver})
    if n_processes > 1:
        p = Pool(n_processes)
        results = list(p.map(_calculate_subset, arguments_list))
    else:
        results = [_calculate_subset(arguments_list[0])]
    for result_list in results:
        for result in result_list:
            i = result[0]
            j = result[1]
            data.growth_rates[i, j] = result[2]
            data.shadow_prices1[i, j] = result[3]
            data.shadow_prices2[i, j] = result[4]
    data.segment()
    return data

if __name__ == "__main__":
    from time import time

    import cobra
    from cobra.test import ecoli_pickle, create_test_model

    n1 = 40
    n2 = 40

    model = create_test_model(ecoli_pickle)
    start_time = time()
    data = calculate_phenotype_phase_plane(model, "EX_glc_e",
        'EX_o2_e', reaction1_npoints=n1, reaction2_npoints=n2,
        n_processes=1, solver="cglpk")
    print "took %.2f seconds with 1 process" % (time() - start_time)
    #start_time = time()
    #data = calculate_phenotype_phase_plane(model, 'EX_glc_e',
    #    'EX_o2_e', reaction1_npoints=n1, reaction2_npoints=n2,
    #    n_processes=2, solver="cglpk")
    #print "took %.2f seconds with 2 processes" % (time() - start_time)
    #data.plot_matplotlib("Pastel1")
    #data.plot_matplotlib("Dark2")
    data.plot_matplotlib()
<|endoftext|>


<commit_before>from lsst_project.src.model import Classifier
from lsst_project.data import utils

import json

import pandas as pd
import numpy as np
import torch
from torch.nn import functional as f

from lsst_project.data import config

def log(message):
    m = "[INFO] {message}".format(message=message)
    print(m)

def pre_process(data, meta_data):
    assert isinstance(data, pd.DataFrame)
    data["unix_time"] = utils.mjd_to_unix_time(data.mjd)
    data = data.set_index('object_id').join(meta_data.set_index('object_id'))

    data["distmod"].fillna(0, inplace=True)
    return data

def scale_data(data, columns):
    features = data[columns]
    mean, variance = features.mean(), features.std()

    # m_dict, v_dict = mean.to_dict(), variance.to_dict()
    # todo store this lookup for retrieval at inference time
    # lookup = {col: {"mean": avg, "variance": var} for (col, avg), (_, var) in zip(m_dict.items(), v_dict.items())}

    data[columns] -= mean
    data[columns] /= variance
    return data, (mean, variance)

def split_train_val_sets(object_ids, val_size=512):
    np.random.seed(0)
    ids = np.random.permutation(object_ids)
    return ids[:val_size], ids[val_size:]

def shuffle_sample(object_ids, batch_size, num_batches, seed=0):
    np.random.seed(seed)
    ids = np.random.permutation(object_ids)

    for batch_num in range(num_batches):
        yield ids[batch_num * batch_size: (batch_num + 1) * batch_size]

def batch_generator(data, batch_ids, max_length, columns, one_hot_lookup):
    data = data[data.index.isin(batch_ids)]

    # need to sort everything in descending by sequence lengths
    # because that is the format that sequence packing method expects
    group = data.groupby(data.index)
    lengths = group.size()[batch_ids].sort_values(ascending=False)
    sorted_batch_ids = lengths.index

    target = group.target.unique()[sorted_batch_ids].astype(float).values
    target = utils.encode_targets(target_array=target, col_lookup=one_hot_lookup)

    seq_tensor = torch.autograd.Variable(torch.zeros((len(sorted_batch_ids), max_length, len(columns)))).float()
    for i in range(len(sorted_batch_ids)):
        object_id = sorted_batch_ids[i]
        seq_len = lengths[object_id]
        seq = data[data.index == object_id][columns].values
        seq_tensor[i, :seq_len] = torch.tensor(seq, dtype=torch.float32)

    return seq_tensor, torch.tensor(target, dtype=torch.long), lengths.values

def compute_f1(predictions, expectations, num_classes):
    y = np.zeros((len(expectations), num_classes), dtype=np.float32)
    y_hat = np.zeros((len(predictions), num_classes), dtype=np.float32)

    y[np.arange(y.shape[0]), expectations] = 1
    y_hat[np.arange(y_hat.shape[0]), predictions] = 1

    with np.errstate(divide='ignore', invalid='ignore'):
        tp = np.sum(y_hat * y, axis=0)
        fp = np.sum(y_hat * (1 - y), axis=0)
        fn = np.sum((1 - y_hat) * y, axis=0)

        precision = tp / (tp + fp)
        recall = tp / (tp + fn)

        precision[np.isnan(precision)] = 0
        recall[np.isnan(recall)] = 0

        f1 = 2 * precision * recall / (precision + recall)
        f1[np.isnan(f1)] = 0
        return f1, (precision, recall)

def run(hidden_dims=32, num_hidden_layers=1, lstm_dropout=0., dense_dropout=0.,
        num_epochs=2, batch_size=32, num_batches=50):
    log("Loading data")
    meta = utils.load_meta_data(training=True)
    data = utils.load_data(training=True)
    data = pre_process(data=data, meta_data=meta)

    log("Splitting training and cross-validation sets")
    val_ids, train_ids = split_train_val_sets(object_ids=meta.object_id.values)
    data_cv = data[data.index.isin(val_ids)]
    data = data[data.index.isin(train_ids)]

    log("Getting longest sequence for padding information")
    lengths = data.groupby(data.index).size().values

    columns_to_exclude = {
        "object_id_main",
        "object_id_meta",
        "mjd",
        "unix_time",
        "detected",
        "hostgal_photoz",
        "target"
    }
    feature_columns = [col for col in data.columns if col not in columns_to_exclude]

    log("Scaling data")
    data, (mean, variance) = scale_data(data=data, columns=feature_columns)
    data_cv[feature_columns] = (data_cv[feature_columns] - mean) / variance

    max_length = lengths.max()
    classes = sorted(meta.target.unique()) + [99]
    num_targets = len(classes)
    one_hot_lookup = {v: idx for idx, v in enumerate(classes)}
    # class_lookup = {v: k for k, v in one_hot_lookup.items()}

    with open(config.MODEL_PATH + "model.json", "w") as jf:
        d = dict(
            num_features=len(feature_columns),
            num_classes=num_targets,
            max_sequence_length=float(max_length),
            lstm_dim=hidden_dims,
            num_lstm_layers=num_hidden_layers,
            lstm_dropout=lstm_dropout,
            dense_dropout=dense_dropout
        )

        json.dump(d, jf, indent=2)

    model = Classifier(
        num_features=len(feature_columns),
        num_classes=num_targets,
        max_sequence_length=max_length,
        lstm_dim=hidden_dims,
        num_lstm_layers=num_hidden_layers,
        lstm_dropout=lstm_dropout,
        dense_dropout=dense_dropout
    )
    loss = torch.nn.CrossEntropyLoss()
    opt = torch.optim.Adam(model.parameters(), lr=1e-2)

    for epoch in range(num_epochs):

        log("\tEpoch: {epoch}".format(epoch=epoch + 1))
        running_loss = 0
        # running_f1 = 0
        i = 0

        for batch_ids in shuffle_sample(train_ids, batch_size, num_batches, seed=0):
            seq_tensor, target, train_lengths = batch_generator(
                data=data,
                batch_ids=batch_ids,
                max_length=max_length,
                columns=feature_columns,
                one_hot_lookup=one_hot_lookup
            )

            if seq_tensor.shape[0] == 0:
                continue

            opt.zero_grad()
            outputs = model.forward(
                x=seq_tensor,
                sequence_lengths=train_lengths,
                max_sequence_length=max_length
            )
            cost = loss(outputs, target)
            cost.backward()
            opt.step()

            running_loss += cost.item()

            if (i + 1) % 10 == 0:
                print("\t\tloss: {loss}".format(loss=running_loss / 10))
                running_loss = 0
            i += 1

        # evaluate the cross-validation set
        seq_tensor, target, cv_lengths = batch_generator(
            data=data_cv,
            batch_ids=val_ids,
            max_length=max_length,
            columns=feature_columns,
            one_hot_lookup=one_hot_lookup
        )
        outputs = model.forward(
            x=seq_tensor,
            sequence_lengths=cv_lengths,
            max_sequence_length=max_length
        )
        _, index = torch.max(f.softmax(outputs, dim=1), 1)
        f1, _ = compute_f1(predictions=index, expectations=target, num_classes=num_targets)
        print("\tcross-validation class-averaged F1: {f1}".format(f1=f1.mean()))

        torch.save(model.state_dict(), config.MODEL_PATH + "model.pth")

if __name__ == '__main__':
    run(
        hidden_dims=32,
        num_hidden_layers=2,
        lstm_dropout=0.3,
        dense_dropout=0.2,
        num_epochs=10,
        batch_size=64,
        num_batches=50
    )
<commit_msg>Add TensorBoard logging for training scalars.<commit_after>from lsst_project.src.model import Classifier
from lsst_project.data import utils

import json

import pandas as pd
import numpy as np
import torch
from torch.nn import functional as f

import tensorflow as tf

from lsst_project.data import config

def log(message):
    m = "[INFO] {message}".format(message=message)
    print(m)

def pre_process(data, meta_data):
    assert isinstance(data, pd.DataFrame)
    data["unix_time"] = utils.mjd_to_unix_time(data.mjd)
    data = data.set_index('object_id').join(meta_data.set_index('object_id'))

    data["distmod"].fillna(0, inplace=True)
    return data

def scale_data(data, columns):
    features = data[columns]
    mean, variance = features.mean(), features.std()

    # m_dict, v_dict = mean.to_dict(), variance.to_dict()
    # todo store this lookup for retrieval at inference time
    # lookup = {col: {"mean": avg, "variance": var} for (col, avg), (_, var) in zip(m_dict.items(), v_dict.items())}

    data[columns] -= mean
    data[columns] /= variance
    return data, (mean, variance)

def split_train_val_sets(object_ids, val_size=512):
    np.random.seed(0)
    ids = np.random.permutation(object_ids)
    return ids[:val_size], ids[val_size:]

def shuffle_sample(object_ids, batch_size, num_batches, seed=0):
    np.random.seed(seed)
    ids = np.random.permutation(object_ids)

    for batch_num in range(num_batches):
        yield ids[batch_num * batch_size: (batch_num + 1) * batch_size]

def batch_generator(data, batch_ids, max_length, columns, one_hot_lookup):
    data = data[data.index.isin(batch_ids)]

    # need to sort everything in descending by sequence lengths
    # because that is the format that sequence packing method expects
    group = data.groupby(data.index)
    lengths = group.size()[batch_ids].sort_values(ascending=False)
    sorted_batch_ids = lengths.index

    target = group.target.unique()[sorted_batch_ids].astype(float).values
    target = utils.encode_targets(target_array=target, col_lookup=one_hot_lookup)

    seq_tensor = torch.autograd.Variable(torch.zeros((len(sorted_batch_ids), max_length, len(columns)))).float()
    for i in range(len(sorted_batch_ids)):
        object_id = sorted_batch_ids[i]
        seq_len = lengths[object_id]
        seq = data[data.index == object_id][columns].values
        seq_tensor[i, :seq_len] = torch.tensor(seq, dtype=torch.float32)

    return seq_tensor, torch.tensor(target, dtype=torch.long), lengths.values

def compute_f1(predictions, expectations, num_classes):
    y = np.zeros((len(expectations), num_classes), dtype=np.float32)
    y_hat = np.zeros((len(predictions), num_classes), dtype=np.float32)

    y[np.arange(y.shape[0]), expectations] = 1
    y_hat[np.arange(y_hat.shape[0]), predictions] = 1

    with np.errstate(divide='ignore', invalid='ignore'):
        tp = np.sum(y_hat * y, axis=0)
        fp = np.sum(y_hat * (1 - y), axis=0)
        fn = np.sum((1 - y_hat) * y, axis=0)

        precision = tp / (tp + fp)
        recall = tp / (tp + fn)

        precision[np.isnan(precision)] = 0
        recall[np.isnan(recall)] = 0

        f1 = 2 * precision * recall / (precision + recall)
        f1[np.isnan(f1)] = 0
        return f1, (precision, recall)

def run(hidden_dims=32, num_hidden_layers=1, lstm_dropout=0., dense_dropout=0.,
        num_epochs=2, batch_size=32, num_batches=50):
    log("Loading data")
    meta = utils.load_meta_data(training=True)
    data = utils.load_data(training=True)
    data = pre_process(data=data, meta_data=meta)

    log("Splitting training and cross-validation sets")
    val_ids, train_ids = split_train_val_sets(object_ids=meta.object_id.values)
    data_cv = data[data.index.isin(val_ids)]
    data = data[data.index.isin(train_ids)]

    log("Getting longest sequence for padding information")
    lengths = data.groupby(data.index).size().values

    columns_to_exclude = {
        "object_id_main",
        "object_id_meta",
        "mjd",
        "unix_time",
        "detected",
        "hostgal_photoz",
        "target"
    }
    feature_columns = [col for col in data.columns if col not in columns_to_exclude]

    log("Scaling data")
    data, (mean, variance) = scale_data(data=data, columns=feature_columns)
    data_cv[feature_columns] = (data_cv[feature_columns] - mean) / variance

    max_length = lengths.max()
    classes = sorted(meta.target.unique()) + [99]
    num_targets = len(classes)
    one_hot_lookup = {v: idx for idx, v in enumerate(classes)}
    class_lookup = {v: k for k, v in one_hot_lookup.items()}

    with open(config.MODEL_PATH + "model.json", "w") as jf:
        d = dict(
            num_features=len(feature_columns),
            num_classes=num_targets,
            max_sequence_length=float(max_length),
            lstm_dim=hidden_dims,
            num_lstm_layers=num_hidden_layers,
            lstm_dropout=lstm_dropout,
            dense_dropout=dense_dropout
        )

        json.dump(d, jf, indent=2)

    model = Classifier(
        num_features=len(feature_columns),
        num_classes=num_targets,
        max_sequence_length=max_length,
        lstm_dim=hidden_dims,
        num_lstm_layers=num_hidden_layers,
        lstm_dropout=lstm_dropout,
        dense_dropout=dense_dropout
    )
    loss = torch.nn.CrossEntropyLoss()
    opt = torch.optim.Adam(model.parameters(), lr=1e-2)

    summary_writer_train = tf.summary.FileWriter(config.MODEL_PATH + 'training')
    summary_writer_dev = tf.summary.FileWriter(config.MODEL_PATH + 'validation')

    for epoch in range(num_epochs):

        log("\tEpoch: {epoch}".format(epoch=epoch + 1))
        running_loss = 0
        train_summary = tf.Summary()
        i = 0

        for batch_ids in shuffle_sample(train_ids, batch_size, num_batches, seed=0):
            seq_tensor, target, train_lengths = batch_generator(
                data=data,
                batch_ids=batch_ids,
                max_length=max_length,
                columns=feature_columns,
                one_hot_lookup=one_hot_lookup
            )

            if seq_tensor.shape[0] == 0:
                continue

            opt.zero_grad()
            outputs = model.forward(
                x=seq_tensor,
                sequence_lengths=train_lengths,
                max_sequence_length=max_length
            )
            cost = loss(outputs, target)
            cost.backward()
            opt.step()

            train_cost = cost.item()
            running_loss += train_cost
            train_summary.value.add(tag="cost", simple_value=train_cost)
            summary_writer_train.add_summary(train_summary, epoch * num_batches + i)

            if (i + 1) % 10 == 0:
                print("\t\tloss: {loss}".format(loss=running_loss / 10))
                running_loss = 0
            i += 1

        summary_writer_train.flush()
        # evaluate the cross-validation set
        seq_tensor, target, cv_lengths = batch_generator(
            data=data_cv,
            batch_ids=val_ids,
            max_length=max_length,
            columns=feature_columns,
            one_hot_lookup=one_hot_lookup
        )
        outputs = model.forward(
            x=seq_tensor,
            sequence_lengths=cv_lengths,
            max_sequence_length=max_length
        )
        _, index = torch.max(f.softmax(outputs, dim=1), 1)
        cv_cost = loss(outputs, target)
        f1, _ = compute_f1(predictions=index, expectations=target, num_classes=num_targets)

        dev_summary = tf.Summary()
        dev_summary.value.add(tag="cost", simple_value=cv_cost)
        for i in range(len(f1)):
            dev_summary.value.add(tag="F1-score/class-{label}".format(label=class_lookup[i]), simple_value=f1[i])

        summary_writer_dev.add_summary(dev_summary, epoch * num_batches + i)
        summary_writer_dev.flush()
        print("\tcross-validation class-averaged F1: {f1}".format(f1=f1.mean()))

        torch.save(model.state_dict(), config.MODEL_PATH + "model.pth")

if __name__ == '__main__':
    run(
        hidden_dims=32,
        num_hidden_layers=2,
        lstm_dropout=0.3,
        dense_dropout=0.2,
        num_epochs=100,
        batch_size=64,
        num_batches=50
    )
<|endoftext|>