text:
'''               ----  SQream Native Python API  ----              '''

import socket, json, ssl, logging, time, traceback, asyncio, sys, array, _thread as thread, threading
from struct import pack, pack_into, unpack, error as struct_error
from datetime import datetime, date, time as t
import multiprocessing as mp
from mmap import mmap
from functools import reduce
from collections import deque
from queue import Queue, Empty
from decimal import Decimal, getcontext
from math import floor, ceil
import functools
import operator
import re
from packaging import version
try:
    import cython
    CYTHON = True
except:
    CYTHON = False

try:
    import pyarrow as pa
    from pyarrow import csv
    import numpy as np
    ARROW = True
except:
    ARROW = False
else:
    sqream_to_pa = {
        'ftBool':     pa.bool_(),
        'ftUByte':    pa.uint8(),
        'ftShort':    pa.int16(),
        'ftInt':      pa.int32(),
        'ftLong':     pa.int64(),
        'ftFloat':    pa.float32(),
        'ftDouble':   pa.float64(),
        'ftDate':     pa.timestamp('ns'),
        'ftDateTime': pa.timestamp('ns'),
        'ftVarchar':  pa.string(),
        'ftBlob':     pa.utf8(),
        'ftNumeric':  pa.decimal128(38, 11)
    }

__version__ = '3.1.5'

WIN = True if sys.platform in ('win32', 'cygwin') else False
PROTOCOL_VERSION = 8
SUPPORTED_PROTOCOLS = 6, 7, 8
BUFFER_SIZE = 100 * int(1e6)  # For setting auto-flushing on netrwork insert
ROWS_PER_FLUSH = 100000
DEFAULT_CHUNKSIZE = 0  # Dummy variable for some jsons
FETCH_MANY_DEFAULT = 1  # default parameter for fetchmany()
VARCHAR_ENCODING = 'ascii'

clean_sqream_errors = True
support_pandas = False

# For encoding data to be sent to SQream using struct.pack() and for type checking by _set_val()
type_to_letter = {
    'ftBool': '?',
    'ftUByte': 'B',
    'ftShort': 'h',
    'ftInt': 'i',
    'ftLong': 'q',
    'ftFloat': 'f',
    'ftDouble': 'd',
    'ftDate': 'i',
    'ftDateTime': 'q',
    'ftVarchar': 's',
    'ftBlob': 's',
    'ftNumeric': '4i'
}

## Setup Logging and debug prints
## ------------------------------
dbg = False

def printdbg(*debug_print):
    if dbg:
        print(*debug_print)

class SQreamDbapiException(Exception):
    pass

logger  = logging.getLogger("dbapi_logger")
logger.setLevel(logging.DEBUG)
logger.disabled = True

def start_logging(log_path=None):

    log_path = log_path or '/tmp/sqream_dbapi.log'
    # logging.disable(logging.NOTSET)
    logger.disabled = False
    try:
        handler = logging.FileHandler(log_path)
    except Exception as e:
        raise Exception("Bad log path was given, please verify path is valid and no forbidden characters were used")

    handler.setLevel(logging.DEBUG)
    handler.setFormatter(logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s'))
    logger.addHandler(handler)

    return logger

def stop_logging():

    # logging.disable(logging.CRITICAL)
    logger.handlers = []
    logger.disabled = True

def log_and_raise(exception_type, error_msg):

    if logger.isEnabledFor(logging.ERROR):
        logger.error(error_msg, exc_info=True)

    raise exception_type(error_msg)

## --- To allow adaptive ROWS_PER_FLUSH ---
##

def get_ram_linux():

    vmstat, err = Popen('vmstat -s'.split(), stdout=PIPE, stderr=PIPE).communicate()

    return int(vmstat.splitlines()[0].split()[0])

def get_ram_windows():

    pass

get_ram = get_ram_windows if WIN else get_ram_linux

## Date and Datetime conversion functions
#  --------------------------------------
'''
  SQream uses a dedicated algorithm to store dates as ints and datetimes as longs. The algorithm is depicted here:
  https://alcor.concordia.ca/~gpkatch/gdate-algorithm.html
  The logic behind it is explained here:
  https:#alcor.concordia.ca/~gpkatch/gdate-method.html
'''

def pad_dates(num):
    return ('0' if num < 10 else '') + str(num)

def sq_date_to_py_date(sqream_date, date_convert_func=date):

    if sqream_date is None or sqream_date == 0:
        return None

    year = (10000 * sqream_date + 14780) // 3652425
    intermed_1 = 365 * year + year // 4 - year // 100 + year // 400
    intermed_2 = sqream_date - intermed_1
    if intermed_2 < 0:
        year = year - 1
        intermed_2 = sqream_date - (365 * year + year // 4 - year // 100 +
                                    year // 400)
    intermed_3 = (100 * intermed_2 + 52) // 3060

    year = year + (intermed_3 + 2) // 12
    month = int((intermed_3 + 2) % 12) + 1
    day = int(intermed_2 - (intermed_3 * 306 + 5) // 10 + 1)

    return date_convert_func(year, month, day)

def sq_datetime_to_py_datetime(sqream_datetime, dt_convert_func=datetime):
    ''' Getting the datetime items involves breaking the long into the date int and time it holds
        The date is extracted in the above, while the time is extracted here  '''

    if sqream_datetime is None or sqream_datetime == 0:
        return None

    date_part = sqream_datetime >> 32
    time_part = sqream_datetime & 0xffffffff
    date_part = sq_date_to_py_date(date_part)

    if date_part is None:
        return None

    msec = time_part % 1000
    sec = (time_part // 1000) % 60
    mins = (time_part // 1000 // 60) % 60
    hour = time_part // 1000 // 60 // 60

    return dt_convert_func(date_part.year, date_part.month, date_part.day,
                           hour, mins, sec, msec)

def date_to_int(d: date) -> int:

    year, month, day = d.timetuple()[:3]
    mth: int         = (month + 9) % 12
    yr: int          = year - mth // 10

    return 365 * yr + yr // 4 - yr // 100 + yr // 400 + (mth * 306 + 5) // 10 + (day - 1)

def datetime_to_long(dt: datetime) -> int:
    ''' self contained to avoid function calling overhead '''

    year, month, day, hour, minute, second = dt.timetuple()[:6]
    msecond = dt.microsecond

    mth: int = (month + 9) % 12
    yr: int = year - mth // 10
    date_int: int = 365 * yr + yr // 4 - yr // 100 + yr // 400 + (
        mth * 306 + 5) // 10 + (day - 1)
    time_int: int = hour * 3600 * 1000 + minute * 60 * 1000 + second * 1000 + msecond // 1000

    return (date_int << 32) + time_int

tenth = Decimal("0.1")
if getcontext().prec < 38:
    getcontext().prec = 38
def sq_numeric_to_decimal(bigint: int, scale: int) -> Decimal:
    if getcontext().prec < 38:
        getcontext().prec = 38
    return Decimal(bigint) * (tenth ** scale)

def decimal_to_sq_numeric(dec: Decimal, scale: int) -> int: # returns bigint
    if getcontext().prec < 38:
        getcontext().prec = 38
    res = dec * (10 ** scale)
    return ceil(res) if res > 0 else floor(res)

def bytes_to_bigint(bytes) -> int:
    c = unpack('4i', bytes)
    res = ((c[3] << 96) + ((c[2] & 0xffffffff) << 64) + ((c[1] & 0xffffffff) << 32) + (c[0] & 0xffffffff))
    return res

try:
    from cythonized import date_to_int as pydate_to_int, datetime_to_long as pydt_to_long, sq_date_to_py_date as date_to_py, sq_datetime_to_py_datetime as dt_to_py
except:
    if CYTHON:
        try:
            import pyximport; pyximport.install(pyimport=True, language_level=3, inplace=True)
            from cythonized import date_to_int as pydate_to_int, datetime_to_long as pydt_to_long, sq_date_to_py_date as date_to_py, sq_datetime_to_py_datetime as dt_to_py
        except:
            CYTHON = False
else:
    CYTHON = True

def lengths_to_pairs(nvarc_lengths):
    ''' Accumulative sum generator, used for parsing nvarchar columns '''

    idx = new_idx = 0
    for length in nvarc_lengths:
        new_idx += length
        yield idx, new_idx
        idx = new_idx

def numpy_datetime_str_to_tup(numpy_dt):
    ''' '1970-01-01T00:00:00.699148800' '''

    numpy_dt = repr(numpy_dt).split("'")[1]
    date_part, time_part = numpy_dt.split('T')
    year, month, day = date_part.split('-')
    hms, ns = time_part.split('.')
    hour, mins, sec = hms.split(':')

    return year, month, day, hour, mins, sec, ns

def numpy_datetime_str_to_tup2(numpy_dt):
    ''' '1970-01-01T00:00:00.699148800' '''

    ts = (numpy_dt - np.datetime64('1970-01-01T00:00:00Z')) / np.timedelta64(1, 's')
    dt = datetime.utcfromtimestamp(ts)

    return dt.year, dt.month, dt.day

    return year, month, day, hour, mins, sec, ns

## Version compare
def version_compare(v1, v2) :
    if (v2 is None or v1 is None):
        return None
    r1 = re.search("\\d{4}(\\.\\d+)+", v1)
    r2 = re.search("\\d{4}(\\.\\d+)+", v2)
    if (r2 is None or r1 is None):
        return None
    v1 = version.parse(r1.group(0))
    v2 = version.parse(r2.group(0))
    return -1 if v1 < v2 else 1 if v1 > v2 else 0

## Socket related
#  --------------

class SQSocket:
    ''' Extended socket class with some'''

    def __init__(self, ip, port, use_ssl=False):
        self.ip, self.port, self.use_ssl = ip, port, use_ssl
        self._setup_socket(ip, port)

    def _setup_socket(self, ip, port):

        self.s = socket.socket()
        if self.use_ssl:
            self.s = ssl.wrap_socket(self.s)

        try:
            self.timeout(10)
            self.s.connect((ip, port))
        except ConnectionRefusedError as e:
            log_and_raise(ConnectionRefusedError, "Connection refused, perhaps wrong IP?")
        except ConnectionResetError:
            log_and_raise(Exception, 'Trying to connect to an SSL port with use_ssl = False')
        except Exception as e:
            if 'timeout' in repr(e):
                log_and_raise(Exception, "Timeout when connecting to SQream, perhaps wrong IP?")
            elif '[SSL: UNKNOWN_PROTOCOL] unknown protocol' in repr(e):
                 log_and_raise(Exception, 'Using use_ssl=True but connected to non ssl sqreamd port')
            elif 'EOF occurred in violation of protocol (_ssl.c:' in repr(e):
                 log_and_raise(Exception, 'Using use_ssl=True but connected to non ssl sqreamd port')
            else:
                log_and_raise(Exception, e)
        else:
            self.timeout(None)

    # General socket / tls socket functionality
    #

    def _check_server_up(self, ip = None, port = None, use_ssl = None):

        try:
            SQSocket(ip or self.ip, port or self.port, use_ssl or self.use_ssl)
        except ConnectionRefusedError:
            log_and_raise(ConnectionRefusedError, f"Connection to SQream interrupted")

    def send(self, data):

        # print ("sending: ", data)
        #try:
        return self.s.send(data)

        #except BrokenPipeError:
        #    raise BrokenPipeError('No connection to SQream. Try reconnecting')

    def close(self):

        return self.s.close()

    def timeout(self, timeout = 'not passed'):

        if timeout == 'not passed':
            return self.s.gettimeout()

        self.s.settimeout(timeout)

    # Extended functionality
    #

    def reconnect(self, ip = None, port = None):

        self.s.close()
        self._setup_socket(ip or self.ip, port or self.port)

    def receive(self, byte_num, timeout=None):
        ''' Read a specific amount of bytes from a given socket '''

        data = bytearray(byte_num)
        view = memoryview(data)
        total = 0

        if timeout:
            self.s.settimeout(timeout)

        while view:
            # Get whatever the socket gives and put it inside the bytearray
            received = self.s.recv_into(view)
            if received == 0:
                log_and_raise(ConnectionRefusedError, f'SQreamd connection interrupted - 0 returned by socket')
            view = view[received:]
            total += received

        if timeout:
            self.s.settimeout(None)

        return data

    def get_response(self, is_text_msg=True):
        ''' Get answer JSON string from SQream after sending a relevant message '''

        # Getting 10-byte response header back
        header = self.receive(10)
        server_protocol = header[0]
        if server_protocol not in SUPPORTED_PROTOCOLS:
            log_and_raise(Exception, f'Protocol mismatch, client version - {PROTOCOL_VERSION}, server version - {server_protocol}')
        # bytes_or_text =  header[1]
        message_len = unpack('q', header[2:10])[0]

        return self.receive(message_len).decode(
            'utf8') if is_text_msg else self.receive(message_len)

    # Non socket aux. functionality
    #

    def generate_message_header(self, data_length, is_text_msg=True, protocol_version=PROTOCOL_VERSION):
        ''' Generate SQream's 10 byte header prepended to any message '''

        return pack('bb', protocol_version, 1 if is_text_msg else 2) + pack(
            'q', data_length)

    def validate_response(self, response, expected):

        if expected not in response:
            # Color first line of SQream error (before the haskell thingy starts) in Red
            response = '\033[31m' + (response.split('\\n')[0] if clean_sqream_errors else response) + '\033[0m'
            log_and_raise(Exception, f'\nexpected response {expected} but got:\n\n {response}')

## Buffer setup and functionality
#  ------------------------------

buf_maps, buf_views = [], []

def init_lock(l):
    ''' To pass a lock to mp.Pool() '''

    global lock
    lock = l

class ColumnBuffer:
    ''' Buffer holding packed columns to be sent to SQream '''

    def __init__(self, size=BUFFER_SIZE):
        global buf_maps, buf_views

    def clear(self):
        if buf_maps:
            [buf_map.close() for buf_map in buf_maps[0]]

    def init_buffers(self, col_sizes, col_nul):
        if not WIN:
            try:
                self.pool.close()
                self.pool.join()
            except Exception as e:
                pass

            l = mp.Lock()
            self.pool = mp.Pool(initializer=init_lock, initargs=(l,))
        self.clear()
        buf_maps = [mmap(-1, ((1 if col_nul else 0)+(size if size!=0 else 104)) * ROWS_PER_FLUSH) for size in col_sizes]
        buf_views = [memoryview(buf_map) for buf_map in buf_maps]

    def pack_columns(self, cols, capacity, col_types, col_sizes, col_nul, col_tvc, col_scales):
        ''' Packs the buffer starting a given index with the column.
            Returns number of bytes packed '''

        pool_params = zip(cols, range(len(col_types)), col_types,
                          col_sizes, col_nul, col_tvc, col_scales)
        if WIN:
            packed_cols = []
            for param_tup in pool_params:
                packed_cols.append(_pack_column(param_tup))

        else:
            # self.pool = mp.Pool()
            # To use multiprocess type packing, we call a top level function with a single tuple parameter
            try:
                packed_cols = self.pool.map(_pack_column, pool_params, chunksize = 2)  # buf_end_indices
            except Exception as e:
                printdbg("Original error from pool.map: ", e)
                if logger.isEnabledFor(logging.ERROR):
                    logger.error("Original error from pool.map: ", e)
                log_and_raise(ProgrammingError, "Error packing columns. Check that all types match the respective column types")

        return list(packed_cols)

    def close(self):
        self.clear()
        try:
            self.pool.close()
            self.pool.join()
        except Exception as e:
            # print (f'testing pool closing, got: {e}')
            pass # no pool was initiated

## A top level packing function for Python's MP compatibility
def _pack_column(col_tup, return_actual_data = True):
    ''' Packs the buffer starting a given index with the column.
        Returns number of bytes packed '''

    global CYTHON
    col, col_idx, col_type, size, nullable, tvc, scale = col_tup
    col = list(col)
    capacity = len(col)
    buf_idx = 0
    buf_map =  mmap(-1, ((1 if nullable else 0)+(size if size!=0 else 104)) * ROWS_PER_FLUSH)
    buf_view = memoryview(buf_map)

    def pack_exception(e):
        ''' Allowing to return traceback info from parent process when using mp.Pool on _pack_column
            [add link]
        '''

        e.traceback = traceback.format_exc()
        error_msg =  f'Trying to insert unsuitable types to column number {col_idx + 1} of type {col_type}'
        with lock:
            logger.error(error_msg, exc_info=True)
        raise ProgrammingError(error_msg)

    # Numpy array for column
    if ARROW and isinstance(col, np.ndarray):
        # Pack null column if applicable
        if nullable:
            pack_into(f'{capacity}b', buf_view, buf_idx,
                      *[1 if item in (np.nan, b'') else 0 for item in col])
            buf_idx += capacity

            # Replace Nones with appropriate placeholder
            if 'S' in repr(col.dtype):    # already b''?
                pass
            elif 'U' in repr(col.dtype):
                pass
            else:
                # Swap out the nans
                col[col == np.nan] = 0

        # Pack nvarchar length column if applicable
        if tvc:
            buf_map.seek(buf_idx)
            lengths_as_bytes = np.vectorize(len)(col).astype('int32').tobytes()
            buf_map.write(lengths_as_bytes)
            buf_idx += len(lengths_as_bytes)

        # Pack the actual data
        if 'U' in repr(col.dtype):
            packed_strings = ''.join(col).encode('utf8')
            buf_map.seek(buf_idx)
            buf_map.write(packed_strings)
            buf_idx += len(packed_strings)
            print (f'unicode strings: {packed_strings}')
        else:
            packed_np = col.tobytes()
            buf_map.seek(buf_idx)
            buf_map.write(packed_np)
            buf_idx += len(packed_np)

        return buf_map[0:buf_idx] if return_actual_data else (0, buf_idx)

    # Pack null column if applicable
    type_code = type_to_letter[col_type]

    # Pack null column and replace None with appropriate placeholder
    col_placeholder = {
        'ftBool': 0,
        'ftUByte': 0,
        'ftShort': 0,
        'ftInt': 0,
        'ftLong': 0,
        'ftFloat': 0,
        'ftDouble': 0,
        'ftDate': None,     #updated separately
        'ftDateTime': None, #updated separately
        'ftVarchar': ''.ljust(size, ' '),
        'ftBlob': '',
        'ftNumeric': 0
    }

    if nullable:
        idx = -1
        while True:
            try:
                idx = col.index(None, idx+1)
            except ValueError:
                break
            else:
                buf_map.seek(buf_idx + idx)
                buf_map.write(b'\x01')
                col[idx] = col_placeholder[col_type]
        # buf_map.seek(buf_idx)
        # buf_map.write(nulls)
        buf_idx += capacity

    # If a text column, replace and pack in adavnce to see if the buffer is sufficient
    if col_type == 'ftBlob':
        try:
            encoded_col = [strn.encode('utf8') for strn in col]
        except AttributeError as e:  # Non strings will not have .encode()
            pack_exception(e)
        else:
            packed_strings = b''.join(encoded_col)

        needed_buf_size = len(packed_strings) + 5* capacity

        # Resize the buffer if not enough space for the current strings
        if needed_buf_size > len(buf_map):
            buf_view.release()
            buf_map.resize(needed_buf_size)
            buf_view = memoryview(buf_map)

        # Pack nvarchar length column
        pack_into(f'{capacity}i', buf_view, buf_idx, *[len(string) for string in encoded_col])

        buf_idx += 4 * capacity

    # Replace Nones with appropriate placeholder - this affects the data itself
    if col_type == 'ftBlob':
        pass     # Handled preemptively due to allow possible buffer resizing

    elif col_type == 'ftVarchar':
        try:
            col = (strn.encode(VARCHAR_ENCODING)[:size].ljust(size, b' ') for strn in col)
        except AttributeError as e:  # Non strings will not have .encode()
            pack_exception(e)
        else:
            packed_strings = b''.join(col)

    elif col_type == 'ftDate':
        # date_tuple_to_int(1900, 1, 1) = 693901
        pass
        # '''
        try:
            col = (date_to_int(deit) if deit is not None else 693901 for deit in col)
        except AttributeError as e:  # Non date/times will not have .timetuple()
            pack_exception(e)
        # '''

    elif col_type == 'ftDateTime':
        # datetime_tuple_to_long(1900, 1, 1, 0, 0, 0) = 2980282101661696
        try:
            col = (datetime_to_long(dt) if dt is not None else 2980282101661696 for dt in col)
        except AttributeError as e:
            pack_exception(e)

    elif col_type == 'ftNumeric':
        try:
            col = (decimal_to_sq_numeric(Decimal(num), scale) for num in col)
        except AttributeError as e:
            pack_exception(e)

    elif col_type in ('ftBool', 'ftUByte', 'ftShort', 'ftInt', 'ftLong','ftFloat', 'ftDouble'):
        pass
    else:
        error_msg = f'Bad column type passed: {col_type}'
        with lock:
            logger.error(error_msg, exc_info=True)
        raise ProgrammingError(error_msg)

    CYTHON = False
    # Done preceding column handling, pack the actual data
    if col_type in ('ftVarchar', 'ftBlob'):
        buf_map.seek(buf_idx)
        buf_map.write(packed_strings)
        buf_idx += len(packed_strings)
    elif col_type == 'ftNumeric':
        buf_map.seek(buf_idx)
        all = functools.reduce(operator.iconcat, (num.to_bytes(16, byteorder='little', signed=True) for num in col), [])
        buf_map.write(bytearray(all))
        buf_idx += len(all)
    else:
        try:
            if CYTHON:
                buf_map.seek(buf_idx)
                type_packer[col_type](col, size, buf_map, buf_idx)
            else:
                pack_into(f'{capacity}{type_code}', buf_view, buf_idx, *col)
        except struct_error as e:
            pack_exception(e)

        buf_idx += capacity * size

    return buf_map[0:buf_idx] if return_actual_data else (0, buf_idx)

class PingLoop(threading.Thread):
    def __init__(self, conn):
        self.conn = conn
        super(PingLoop, self).__init__()
        self.done = False

    def run(self):
        json_cmd = '{"ping":"ping"}'
        binary = self.conn.s.generate_message_header(len(json_cmd)) + json_cmd.encode('utf8');
        while self.sleep():
            conn = self.conn
            try:
                conn.s.send(binary)
            except:
                self.done = True

    def halt(self):
        self.done = True

    def sleep(self):
        if self.done:
            return False
        count = 0
        while (count < 100):
            count = count + 1
            time.sleep(.1)
            if self.done:
                return False
        return True

def _start_ping_loop(self):
    self.ping_loop = PingLoop(self)
    self.ping_loop.start()

def _end_ping_loop(self):
    if (self.ping_loop is not None):
        self.ping_loop.halt()
        self.ping_loop.join()
    self.ping_loop = None

class Connection:
    ''' Connection class used to interact with SQream '''

    base_conn_open = [False]

    def __init__(self, ip, port, clustered, use_ssl=False, log = False, base_connection=True, reconnect_attempts=3, reconnect_interval=10):

        self.buffer = ColumnBuffer(BUFFER_SIZE)  # flushing buffer every BUFFER_SIZE bytes
        self.row_size = 0
        self.rows_per_flush = 0
        self.version = None
        self.stmt_id = None  # For error handling when called out of order
        self.statement_type = None
        self.open_statement = False
        self.closed = False
        self.orig_ip, self.orig_port, self.clustered, self.use_ssl = ip, port, clustered, use_ssl
        self.reconnect_attempts, self.reconnect_interval = reconnect_attempts, reconnect_interval
        self.base_connection = base_connection
        self.ping_loop = None

        self._open_connection(clustered, use_ssl)
        self.arraysize = FETCH_MANY_DEFAULT
        self.rowcount = -1    # DB-API property
        self.more_to_fetch = False
        self.parsed_rows = []

        if self.base_connection:
            self.cursors = []

        self.lastrowid = None
        self.unpack_q = Queue()
        if CYTHON:
            # To allow hot swapping for testing
            date_to_int, datetime_to_long, sq_date_to_py_date, sq_datetime_to_py_datetime = pydate_to_int, pydt_to_long, date_to_py, dt_to_py

        if log is not False:
            start_logging(None if log is True else log)
        # Thread for unpacking fetched socket data
        # thread.start_new_thread(self._parse_fetched_cols, (self.unpack_q,))

    ## SQream mechanisms
    #  -----------------

    def _open_connection(self, clustered, use_ssl):
        ''' Get proper ip and port from picker if needed and connect socket. Used at __init__() '''
        if clustered is True:

            # Create non SSL socket for picker communication
            picker_socket = SQSocket(self.orig_ip, self.orig_port, False)
            # Parse picker response to get ip and port
            # Read 4 bytes to find length of how much to read
            picker_socket.timeout(5)
            try:
                read_len = unpack('i', picker_socket.receive(4))[0]
            except socket.timeout:

                log_and_raise(ProgrammingError, f'Connected with clustered=True, but apparently not a server picker port')
            picker_socket.timeout(None)

            # Read the number of bytes, which is the IP in string format
            # Using a nonblocking socket in case clustered = True was passed but not connected to picker
            self.ip = picker_socket.receive(read_len)

            # Now read port
            self.port = unpack('i', picker_socket.receive(4))[0]
            picker_socket.close()
        else:
            self.ip, self.port = self.orig_ip, self.orig_port

        # Create socket and connect to actual SQreamd server
        self.s = SQSocket(self.ip, self.port, use_ssl)
        if self.base_connection:
            self.base_conn_open[0] = True

    def _send_string(self, json_cmd, get_response=True, is_text_msg=True, sock=None):
        ''' Encode a JSON string and send to SQream. Optionally get response '''

        # Generating the message header, and sending both over the socket
        printdbg(f'string sent: {json_cmd}')
        self.s.send(self.s.generate_message_header(len(json_cmd)) + json_cmd.encode('utf8'))

        if get_response:
            return self.s.get_response(is_text_msg)

    def connect_database(self, database, username, password, service='sqream'):
        ''' Handle connection to database, with or without server picker '''

        self.database, self.username, self.password, self.service = database, username, password, service
        res = self._send_string(
            f'{{"username":"{username}", "password":"{password}", "connectDatabase":"{database}", "service":"{service}"}}'
        )
        res = json.loads(res)
        try:
            self.connection_id = res['connectionId']
            if 'version' in res:
                self.version = res['version']
        except KeyError as e:
            log_and_raise(ProgrammingError, f"Error connecting to database: {res['error']}")

        self.varchar_enc = res.get('varcharEncoding', 'ascii')

        if logger.isEnabledFor(logging.INFO):
            logger.info(f'Connection opened to database {database}. Connection ID: {self.connection_id}')

    def _attempt_reconnect(self):

        for attempt in range(self.reconnect_attempts):
            time.sleep(self.reconnect_interval)
            try:
                self._open_connection(self.clustered, self.use_ssl)
                self.connect_database(self.database, self.username, self.password, self.service)
            except ConnectionRefusedError as e:
                print(f'Connection lost, retry attempt no. {attempt+1} failed. Original error:\n{e}')
            else:
                return

        # Attempts failed
        log_and_raise(ConnectionRefusedError, 'Reconnection attempts to sqreamd failed')

    def execute_sqream_statement(self, stmt):
        ''' '''

        self.latest_stmt = stmt

        if self.open_statement:
            self.close_statement()
        self.open_statement = True

        self.more_to_fetch = True

        self.stmt_id = json.loads(self._send_string('{"getStatementId" : "getStatementId"}'))["statementId"]
        comp = version_compare(self.version, "2020.3.1")
        if (comp is not None and comp > -1):
            _start_ping_loop(self)
        stmt_json = json.dumps({"prepareStatement": stmt, "chunkSize": DEFAULT_CHUNKSIZE})
        res = self._send_string(stmt_json)

        self.s.validate_response(res, "statementPrepared")
        self.lb_params = json.loads(res)
        if self.lb_params.get('reconnect'):  # Reconnect exists and issued, otherwise False / None

            # Close socket, open a new socket with new port/ip sent be the reconnect response
            self.s.reconnect(
                self.lb_params['ip'], self.lb_params['port_ssl']
                if self.use_ssl else self.lb_params['port'])

            # Send reconnect and reconstruct messages
            reconnect_str = '{{"service": "{}", "reconnectDatabase":"{}", "connectionId":{}, "listenerId":{},"username":"{}", "password":"{}"}}'.format(
                self.service, self.database, self.connection_id,
                self.lb_params['listener_id'], self.username, self.password)
            self._send_string(reconnect_str)
            self._send_string('{{"reconstructStatement": {}}}'.format(
                self.stmt_id))

        # Reconnected/reconstructed if needed,  send  execute command
        self.s.validate_response(self._send_string('{"execute" : "execute"}'), 'executed')

        # Send queryType message/s
        res = json.loads(self._send_string('{"queryTypeIn": "queryTypeIn"}'))
        self.column_list = res.get('queryType', '')

        if not self.column_list:
            res = json.loads(
                self._send_string('{"queryTypeOut" : "queryTypeOut"}'))
            self.column_list = res.get('queryTypeNamed', '')
            if not self.column_list:
                self.statement_type = 'DML'
                self.close_statement()
                return

            self.statement_type = 'SELECT' if self.column_list else 'DML'
            self.result_rows = []
            self.unparsed_row_amount = 0
            self.data_columns = []
        else:
            self.statement_type = 'INSERT'

        # {"isTrueVarChar":false,"nullable":true,"type":["ftInt",4,0]}
        self.col_names, self.col_tvc, self.col_nul, self.col_type_tups = \
            list(zip(*[(col.get("name", ""), col["isTrueVarChar"], col["nullable"], col["type"]) for col in self.column_list]))
        self.col_names_map = {
            name: idx
            for idx, name in enumerate(self.col_names)
        }

        if self.statement_type == 'INSERT':
            self.col_types = [type_tup[0] for type_tup in self.col_type_tups]
            self.col_sizes = [type_tup[1] for type_tup in self.col_type_tups]
            self.col_scales = [type_tup[2] for type_tup in self.col_type_tups]
            self.row_size = sum(self.col_sizes) + sum(
                self.col_nul) + 4 * sum(self.col_tvc)
            self.rows_per_flush = ROWS_PER_FLUSH
            self.buffer.init_buffers(self.col_sizes, self.col_nul)

        # if self.statement_type == 'SELECT':
        self.parsed_rows = []
        self.parsed_row_amount = 0

        if logger.isEnabledFor(logging.INFO):
            logger.info(f'Executing statement over connection {self.connection_id} with statement id {self.stmt_id}:\n{stmt}')

    ## Select

    def _fetch(self, sock=None):

        sock = sock or self.s
        # JSON correspondence
        res = self._send_string('{"fetch" : "fetch"}')
        self.s.validate_response(res, "colSzs")
        fetch_meta = json.loads(res)
        num_rows_fetched, column_sizes = fetch_meta['rows'], fetch_meta['colSzs']
        if num_rows_fetched == 0:
            self.close_statement()
            return num_rows_fetched

        # Get preceding header
        self.s.receive(10)

        # Get data as memoryviews of bytearrays.
        unsorted_data_columns = [
            memoryview(self.s.receive(size))
            for idx, size in enumerate(column_sizes)
        ]

        # Sort by columns, taking a memoryview and casting to the proper type
        self.data_columns = []

        for type_tup, nullable, tvc in zip(self.col_type_tups, self.col_nul,
                                           self.col_tvc):
            column = []
            if nullable:
                column.append(unsorted_data_columns.pop(0))
            if tvc:
                column.append(unsorted_data_columns.pop(0).cast('i'))

            column.append(unsorted_data_columns.pop(0))
            if type_tup[0] not in ('ftVarchar', 'ftBlob', 'ftNumeric'):
                column[-1] = column[-1].cast(type_to_letter[type_tup[0]])
            else:
                column[-1] = column[-1].tobytes()
            self.data_columns.append(column)

        self.unparsed_row_amount = num_rows_fetched

        return num_rows_fetched

    def _parse_fetched_cols(self, queue = None):
        ''' Used by _fetch_and_parse ()  '''

        self.extracted_cols = []

        if not self.data_columns:
            return self.extracted_cols

        for idx, raw_col_data in enumerate(self.data_columns):
            # Extract data according to column type
            if self.col_tvc[idx]:  # nvarchar
                nvarc_sizes = raw_col_data[1 if self.col_nul[idx] else 0]
                col = [
                    raw_col_data[-1][start:end].decode('utf8')
                    for (start, end) in lengths_to_pairs(nvarc_sizes)
                ]
            elif self.col_type_tups[idx][0] == "ftVarchar":
                varchar_size = self.col_type_tups[idx][1]
                col = [
                    raw_col_data[-1][idx:idx + varchar_size].decode(
                        self.varchar_enc, "ignore").replace('\x00', '').rstrip()
                    for idx in range(0, len(raw_col_data[-1]), varchar_size)
                ]
            elif self.col_type_tups[idx][0] == "ftDate":
                col = [sq_date_to_py_date(d) for d in raw_col_data[-1]]
            elif self.col_type_tups[idx][0] == "ftDateTime":
                col = [sq_datetime_to_py_datetime(d) for d in raw_col_data[-1]]
            elif self.col_type_tups[idx][0] == "ftNumeric":
                scale = self.col_type_tups[idx][2]
                col = [
                    sq_numeric_to_decimal(bytes_to_bigint(raw_col_data[-1][idx:idx + 16]), scale)
                    for idx in range(0, len(raw_col_data[-1]), 16)
                ]

            else:
                col = raw_col_data[-1]

            # Fill Nones if / where needed
            if self.col_nul[idx]:
                nulls = raw_col_data[0]  # .tolist()
                col = [
                    item if not null else None
                    for item, null in zip(col, nulls)
                ]
            else:
                pass

            self.extracted_cols.append(col)

        # Done with the raw data buffers
        self.unparsed_row_amount = 0
        self.data_columns = []

        return self.extracted_cols

    def _fetch_and_parse(self, requested_row_amount, data_as='rows'):
        ''' See if this amount of data is available or a fetch from sqream is required
            -1 - fetch all available data. Used by fetchmany() '''

        if data_as == 'rows':
            while (requested_row_amount > len(self.parsed_rows)
                   or requested_row_amount == -1) and self.more_to_fetch:
                self.more_to_fetch = bool(self._fetch())  # _fetch() updates self.unparsed_row_amount

                self.parsed_rows.extend(zip(*self._parse_fetched_cols()))

    ## Insert

    def _send_columns(self, cols=None, capacity=None):
        ''' Perform network insert - "put" json, header, binarized columns. Used by executemany() '''

        cols = cols or self.cols
        cols = cols if isinstance(cols, (list, tuple, set, dict)) else list(cols)

        capacity = capacity or self.capacity

        # Send columns and metadata to be packed into our buffer
        packed_cols = self.buffer.pack_columns(cols, capacity, self.col_types,
                                               self.col_sizes, self.col_nul,
                                               self.col_tvc, self.col_scales)
        del cols
        byte_count = sum(len(packed_col) for packed_col in packed_cols)

        # Sending put message and binary header
        self._send_string(f'{{"put":{capacity}}}', False)
        self.s.send((self.s.generate_message_header(byte_count, False)))

        # Sending packed data (binary buffer)
        for packed_col in packed_cols:
            self.s.send((packed_col))

        self.s.validate_response(self.s.get_response(), '{"putted":"putted"}')
        del packed_cols

    ## Closing

    def close_statement(self, sock=None):

        if self.open_statement:
            sock = sock or self.s
            self._send_string('{"closeStatement": "closeStatement"}')
            self.open_statement = False
            self.buffer.close()
            _end_ping_loop(self)

        if logger.isEnabledFor(logging.INFO):
            logger.info(f'Done executing statement {self.stmt_id} over connection {self.connection_id}')

    def close_connection(self, sock=None):

        if self.closed:
            log_and_raise(ProgrammingError, "Trying to close a connection that's already closed")

        self._send_string('{"closeConnection":  "closeConnection"}')
        self.s.close()
        self.buffer.close()
        _end_ping_loop(self)
        self.closed = True
        self.base_conn_open[0] = False if self.base_connection else True

        if logger.isEnabledFor(logging.INFO):
            logger.info(f'Connection closed to database {self.database}. Connection ID: {self.connection_id}')

    # '''
    def csv_to_table(self, csv_path, table_name, read = None, parse = None, convert = None, con = None, auto_infer = False):
        ' Pyarrow CSV reader documentation: https://arrow.apache.org/docs/python/generated/pyarrow.csv.read_csv.html '

        if not ARROW:
            return "Optional pyarrow dependency not found. To install: pip3 install pyarrow"

        sqream_to_pa = {
            'ftBool':     pa.bool_(),
            'ftUByte':    pa.uint8(),
            'ftShort':    pa.int16(),
            'ftInt':      pa.int32(),
            'ftLong':     pa.int64(),
            'ftFloat':    pa.float32(),
            'ftDouble':   pa.float64(),
            'ftDate':     pa.timestamp('ns'),
            'ftDateTime': pa.timestamp('ns'),
            'ftVarchar':  pa.string(),
            'ftBlob':     pa.utf8(),
            'ftNumeric':  pa.decimal128(28, 11)
        }

        start = time.time()
        # Get table metadata
        con = con or self
        con.execute(f'select * from {table_name} where 1=0')

        # Map column names to pyarrow types and set Arrow's CSV parameters
        sqream_col_types = [col_type[0] for col_type in con.col_type_tups]
        column_types = zip(con.col_names, [sqream_to_pa[col_type[0]] for col_type in con.col_type_tups])
        read = read or csv.ReadOptions(column_names=con.col_names)
        parse = parse or csv.ParseOptions(delimiter='|')
        convert = convert or csv.ConvertOptions(column_types = None if auto_infer else column_types)

        # Read CSV to in-memory arrow format
        csv_arrow = csv.read_csv(csv_path, read_options=read, parse_options=parse, convert_options=convert).combine_chunks()
        num_chunks = len(csv_arrow[0].chunks)
        numpy_cols = []

        # For each column, get the numpy representation for quick packing
        for col_type, col in zip(sqream_col_types, csv_arrow):
            # Only one chunk after combine_chunks()
            col = col.chunks[0]
            if col_type in  ('ftVarchar', 'ftBlob', 'ftDate', 'ftDateTime', 'ftNumeric'):
                col = col.to_pandas()
            else:
                col = col.to_numpy()

            numpy_cols.append(col)

        print (f'total loading csv: {time.time()-start}')
        start = time.time()

        # Insert columns into SQream
        col_num = csv_arrow.shape[1]
        con.executemany(f'insert into {table_name} values ({"?,"*(col_num-1)}?)', numpy_cols)
        print (f'total inserting csv: {time.time()-start}')

    # '''

    '''  -- Metadata  --
         ---------------   '''
    def get_statement_type(self):

        return self.statement_type

    def get_statement_id(self):

        return self.stmt_id

    ## DB-API API
    #  ----------

    def _verify_open(self):

        if not self.base_conn_open[0]:
            log_and_raise(ProgrammingError, 'Connection has been closed')

        if self.closed:
            log_and_raise(ProgrammingError, 'Cursor has been closed')

    def _verify_query_type(self, query_type):

        if not self.open_statement:
            log_and_raise(ProgrammingError, 'No open statement while attempting fetch operation')

    def _fill_description(self):
        '''Getting parameters for the cursor's 'description' attribute, even for
           a query that returns no rows. For each column, this includes:
           (name, type_code, display_size, internal_size, precision, scale, null_ok) '''

        if self.statement_type != 'SELECT':
            self.description = None
            return self.description

        self.description = []
        for col_name, col_nullalbe, col_type_tup in zip(
                self.col_names, self.col_nul, self.col_type_tups):
            type_code = typecodes[
                col_type_tup[0]]  # Convert SQream type to DBAPI identifier
            display_size = internal_size = col_type_tup[
                1]  # Check if other size is available from API
            precision = 38
            scale = col_type_tup[2]

            self.description.append(
                (col_name, type_code, display_size, internal_size, precision,
                 scale, col_nullalbe))

        return self.description

    def execute(self, query, params=None):
        ''' Execute a statement. Parameters are not supported '''

        self._verify_open()
        if params:

            log_and_raise(ProgrammingError, "Parametered queries not supported. \
                If this is an insert query, use executemany() with the data rows as the parameter")

        else:
            self.execute_sqream_statement(query)

        self._fill_description()
        self.rows_fetched = 0
        self.rows_returned = 0

        return self

    def executemany(self, query, rows_or_cols=None, data_as='rows', amount=None):
        ''' Execute a statement, including parametered data insert '''

        self._verify_open()
        self.execute(query)

        if rows_or_cols is None:
            return self

        if data_as =='alchemy_flat_list':
            # Unflatten SQLalchemy data list
            row_len = len(self.column_list)
            rows_or_cols = [rows_or_cols[i:i+row_len] for i in range(0, len(rows_or_cols), row_len)]
            data_as = 'rows'

        if 'numpy' in repr(type(rows_or_cols[0])):
            data_as = 'numpy'

        # Network insert starts here if data was passed
        column_lengths = [len(row_or_col) for row_or_col in rows_or_cols]
        if column_lengths.count(column_lengths[0]) != len(column_lengths):
            log_and_raise(ProgrammingError,
                "Incosistent data sequences passed for inserting. Please use rows/columns of consistent length"
            )
        if data_as == 'rows':
            self.capacity = amount or len(rows_or_cols)
            self.cols = list(zip(*rows_or_cols))
        else:
            self.cols = rows_or_cols
            self.capacity = len(self.cols)

        # Slice a chunk of columns and pass to _send_columns()
        start_idx = 0
        while self.cols != [()]:
            col_chunk = [col[start_idx:start_idx + self.rows_per_flush] for col in self.cols]
            chunk_len = len(col_chunk[0])
            if chunk_len == 0:
                break
            self._send_columns(col_chunk, chunk_len)
            start_idx += self.rows_per_flush
            del col_chunk

            if logger.isEnabledFor(logging.INFO):
                logger.info(f'Sent {chunk_len} rows of data')

        self.close_statement()

        return self

    def fetchmany(self, size=None, data_as='rows', fetchone=False):
        ''' Fetch an amount of result rows '''

        size = size or self.arraysize
        self._verify_open()
        if self.statement_type not in (None, 'SELECT'):
            log_and_raise(ProgrammingError,'No open statement while attempting fetch operation')

        if self.more_to_fetch is False:
            # All data from server for this select statement was fetched
            if len(self.parsed_rows) == 0:
                # Nothing
                return [] if not fetchone else None
        else:
            self._fetch_and_parse(size, data_as)

        # Get relevant part of parsed rows and reduce storage and counter
        if data_as == 'rows':
            res = self.parsed_rows[0:size if size != -1 else None]
            del self.parsed_rows[:size if size != -1 else None]

        if logger.isEnabledFor(logging.INFO):
            logger.info(f'Fetched {size} rows')

        return (res if res else []) if not fetchone else (res[0] if res else None)

    def fetchone(self, data_as='rows'):
        ''' Fetch one result row '''

        if data_as not in ('rows',):
            log_and_raise(ProgrammingError, "Bad argument to fetchone()")

        return self.fetchmany(1, data_as, fetchone=True)

    def fetchall(self, data_as='rows'):
        ''' Fetch all result rows '''

        if data_as not in ('rows',):
            log_and_raise(ProgrammingError, "Bad argument to fetchall()")

        return self.fetchmany(-1, data_as)

    def cursor(self):
        ''' Return a new connection with the same parameters.
            We use a connection as the equivalent of a 'cursor' '''

        cur = Connection(
            self.orig_ip if self.clustered is True else self.ip,
            self.orig_port if self.clustered is True else self.port,
            self.clustered,
            self.use_ssl,
            base_connection=False
        )  # self is the calling connection instance, so cursor can trace back to pysqream
        cur.connect_database(self.database, self.username, self.password, self.service)

        self.cursors.append(cur)

        return cur

    def commit(self):
        self._verify_open()

    def rollback(self):
        pass

    def close(self):

        if self.closed:
            log_and_raise(ProgrammingError, "Trying to close a connection that's already closed")

        self.close_statement()
        self.close_connection()
        self.closed = True

    def nextset(self):
        ''' No multiple result sets so currently always returns None '''

        return None

    # DB-API Do nothing (for now) methods
    # -----------------------------------

    def setinputsizes(self, sizes):

        self._verify_open()

    def setoutputsize(self, size, column=None):

        self._verify_open()

    # Internal Methods
    # ----------------
    ''' Include: __enter__(), __exit__() - for using "with",
                 __iter__ for use in for-in clause  '''

    def __enter__(self):
        return self

    def __exit__(self, exc_type, exc_value, exc_traceback):
        self.close()

    def __iter__(self):
        for item in self.fetchall():
            yield item

## Top level functionality
#  -----------------------

def connect(host, port, database, username, password, clustered = False, use_ssl = False, service='sqream', log=False, reconnect_attempts=3, reconnect_interval=10):
    ''' Connect to SQream database '''
    if not isinstance(reconnect_attempts, int) or reconnect_attempts < 0:
        log_and_raise(Exception, f'reconnect attempts should be a positive integer, got : {reconnect_attempts}')
    if not isinstance(reconnect_interval, int) or reconnect_attempts < 0:
        log_and_raise(Exception, f'reconnect interval should be a positive integer, got : {reconnect_interval}')

    conn = Connection(host, port, clustered, use_ssl, log=log, base_connection=True, reconnect_attempts=reconnect_attempts, reconnect_interval=reconnect_interval)
    conn.connect_database(database, username, password, service)

    return conn

## DBapi compatibility
#  -------------------
''' To fully comply to Python's DB-API 2.0 database standard. Ignore when using internally '''

# Type objects and constructors required by the DB-API 2.0 standard
Binary = memoryview
Date = date
Time = t
Timestamp = datetime

class Error(Exception):
    pass

class Warning(Exception):
    pass

class InterfaceError(Error):
    pass

class DatabaseError(Error):
    pass

class DataError(DatabaseError):
    pass

class OperationalError(DatabaseError):
    pass

class IntegrityError(DatabaseError):
    pass

class InternalError(DatabaseError):
    pass

class ProgrammingError(DatabaseError):
    pass

class NotSupportedError(DatabaseError):
    pass

class _DBAPITypeObject:
    """DB-API type object which compares equal to all values passed to the constructor.
        https://www.python.org/dev/peps/pep-0249/#implementation-hints-for-module-authors
    """
    def __init__(self, *values):
        self.values = values

    def __eq__(self, other):
        return other in self.values

STRING = "STRING"
BINARY = _DBAPITypeObject("BYTES", "RECORD", "STRUCT")
NUMBER = _DBAPITypeObject("INTEGER", "INT64", "FLOAT", "FLOAT64", "NUMERIC",
                          "BOOLEAN", "BOOL")
DATETIME = _DBAPITypeObject("TIMESTAMP", "DATE", "TIME", "DATETIME")
ROWID = "ROWID"

typecodes = {
    'ftBool': 'NUMBER',
    'ftUByte': 'NUMBER',
    'ftInt': 'NUMBER',
    'ftShort': 'NUMBER',
    'ftLong': 'NUMBER',
    'ftDouble': 'NUMBER',
    'ftFloat': 'NUMBER',
    'ftDate': 'DATETIME',
    'ftDateTime': 'DATETIME',
    'ftVarchar': 'STRING',
    'ftBlob': 'STRING',
    'ftNumeric': 'NUMBER'
}

def DateFromTicks(ticks):
    return Date.fromtimestamp(ticks)

def TimeFromTicks(ticks):
    return Time(
        *time.localtime(ticks)[3:6]
    )  # localtime() returns a namedtuple, fields 3-5 are hr/min/sec

def TimestampFromTicks(ticks):
    return Timestamp.fromtimestamp(ticks)

# DB-API global parameters
apilevel = '2.0'

threadsafety = 1 # Threads can share the module but not a connection

paramstyle = 'qmark'

if __name__ == '__main__':

    print('PySqream DB-API connector, version ', __version__)

max_stars_repo_name: SQream/SQream-Python-Driver-Connector

max_issues_repo_name: SQream/SQream-Python-Driver

max_forks_repo_name: SQream/SQream-Python-Driver


text:
import matplotlib as mpl
#mpl.use('Agg')
mpl.rcParams['agg.path.chunksize'] = 2_000_000
from time import time as t
import numpy as np

# Create prm
from PRM_predict import create_prm
prm = create_prm(month_prediction=True)

from downscale.data_source.MNT import MNT
from downscale.data_source.NWP import NWP
from downscale.data_source.observation import Observation
from downscale.operators.helbig import DwnscHelbig
from downscale.utils.context_managers import timer_context
from downscale.visu.visualization import Visualization

IGN = MNT(prm=prm)
AROME = NWP(prm["selected_path"], begin=prm["begin"], end=prm["end"], prm=prm)
AROME.data_xr = AROME.data_xr.isel(time=0)
AROME.convert_to_mnt_format(extract_wind=True)
AROME_interpolated = AROME.data_xr.interp_like(IGN.data_xr, method="nearest")
BDclim = Observation(prm["BDclim_stations_path"], prm["BDclim_data_path"], prm=prm)

#IGN.data_xr = IGN.data_xr.sel(x=slice(900_000, 950_000), y=slice(6_450_000, 6_500_000))

"""
mnt = IGN.data_xr.data[0, :, :]#[0, :1000, :2000]

d = DwnscHelbig()
# mu_mnt = d.mu_helbig_map(mnt, dx=25)
# laplacian_mnt = d.laplacian_map(mnt, dx=25)

with timer_context("downscale Helbig et al. 2017", level="", unit="minute", verbose=True):
    output_downscale = d.downscale_helbig(mnt, dx=30, library_subgrid="tensorflow", library_downscale="numba")

result = AROME_interpolated.data * output_downscale.reshape((1, output_downscale.shape[0], output_downscale.shape[1]))
print(result.shape)
"""

"""
import matplotlib.pyplot as plt

mnt_small = IGN.data[330:12_300, 420:7_200]
x_min = np.nanmin(IGN.data_xr.x.data[420:7_200])
y_max = np.nanmax(IGN.data_xr.y.data[330:12_300])

stations = BDclim.stations
resolution_x = 25
resolution_y = 25
idx_x, idx_y = IGN.find_nearest_MNT_index(stations["X"], stations["Y"])

mu_mnt_idx = downscale_1.mu_helbig_idx(mnt, dx=25, idx_x=idx_x, idx_y=idx_y)
laplacian_mnt_idx = downscale_1.laplacian_idx(mnt, dx=25, idx_x=idx_x, idx_y=idx_y)
x_dsc_topo_idx = downscale_1.x_dsc_topo_helbig(mnt, dx=25, idx_x=idx_x, idx_y=idx_y, type_input="indexes")

downscale_1 = DwnscHelbig()

mu_mnt = downscale_1.mu_helbig_map(mnt_small, dx=25)
laplacian_mnt = downscale_1.laplacian_map(mnt_small, dx=25)
x_dsc_topo = downscale_1.x_dsc_topo_helbig(mnt_small)

plt.figure()
plt.imshow(mnt_small)
plt.colorbar()
plt.figure()
plt.imshow(laplacian_mnt, vmin=-0.5, vmax=0.5)
plt.colorbar()
plt.title("Laplacian")
plt.figure()
plt.imshow(mu_mnt)
plt.colorbar()
plt.title("Mu")
plt.figure()
plt.imshow(x_dsc_topo)
plt.colorbar()
plt.title("x_dsc_topo")

plt.figure()
plt.scatter(laplacian_mnt, mu_mnt, c=x_dsc_topo, s=3)
plt.scatter(laplacian_mnt_idx, mu_mnt_idx, c="red", s=3)
plt.xlabel("Laplacian")
plt.ylabel("Mu")
plt.colorbar()
plt.savefig("laplacian_mu2.png")
"""

ext: py

reason: import numpy


text:
import numpy as np

def get_pose_count(data_dict):
    """
    Calculates the number of poses in the dataset. A pose is one frame of a sequence, independent of cameras.
    In other words, each pose was taken from 4 cameras.
    """
    pose_lens = {}
    for key in data_dict.keys():
        subject, action, seq, camera = key
        if (subject, action, seq) in pose_lens:
            assert pose_lens[(subject, action, seq)] == len(data_dict[key]), \
                'Inconsistent number of poses at ' + str(key)
        else:
            pose_lens[(subject, action, seq)] = len(data_dict[key])

    return sum([x for x in pose_lens.values()])

def pose_index(data_dict):
    """
    Calculates a pose index for the given input dict.

    :returns: (poselist, keys) tuple, where poselist has a row for each pose and two columns.
        The first column is the frame number, the second an index into keys. keys is a list of tuples of 3 elements:
        (subject, action, sequence). Note that camera information is not present.
    """
    pose_lens = {}
    for key in data_dict.keys():
        subject, action, seq, camera = key
        if (subject, action, seq) in pose_lens:
            assert pose_lens[(subject, action, seq)] == len(data_dict[key])
        else:
            pose_lens[(subject, action, seq)] = len(data_dict[key])

    N = sum([x for x in pose_lens.values()])
    poselist = np.zeros((N, 2), 'int64')
    keys = sorted(pose_lens.keys())

    start = 0
    for i, k in enumerate(keys):
        n = pose_lens[k]
        poselist[start:start + n, 0] = np.arange(n)
        poselist[start:start + n, 1] = i

        start += n

    assert start == len(poselist)
    for i in range(len(keys) - 1):
        assert keys[i] < keys[i + 1]

    for i in range(len(poselist) - 1):
        assert poselist[i, 1] <= poselist[i + 1, 1]

    return poselist, keys

def keys_to_stack(x):
    """
    Creates an index that can be used to merge the data that is in baseline format.
    The output is deterministic.

    :returns: ``(framelist, keys)`` tuple, ``keys`` is an array of video sequence keys, ``framelist`` is a numpy array,
              each row contains a frame number and a key id which can be looked up in keys.
    """
    keys = sorted(x.keys())
    key2id = {k: i for i, k in enumerate(keys)}

    x_len = sum([v.shape[0] for v in x.values()])
    framelist = np.empty((x_len, 2), dtype='int32')

    start = 0
    for k in keys:
        cnt = x[k].shape[0]
        framelist[start:start + cnt, 1] = key2id[k]
        framelist[start:start + cnt, 0] = np.arange(cnt)
        start += cnt

    assert start == len(framelist)
    for i in range(len(keys) - 1):
        assert keys[i] < keys[i + 1]

    for i in range(len(framelist) - 1):
        assert framelist[i, 1] <= framelist[i + 1, 1]

    return framelist, keys

# stack y data according to framelist
def stack_by_keys(data_dict, framelist, keys):
    """ Stacks a pose dictionary according to ``framelist``. ``keys`` provides the keys to the video name. """
    shape = list(data_dict[keys[0]].shape)
    shape[0] = len(framelist)
    out = np.empty(shape)

    for i, key in enumerate(keys):
        out[framelist[:, 1] == i, :] = data_dict[key]

    return out

max_stars_repo_name: vegesm/siamese-pose-estimation

max_issues_repo_name: vegesm/siamese-pose-estimation

max_forks_repo_name: vegesm/siamese-pose-estimation


text:
"""Synthetic datasets for quantile learning."""

from scipy.stats import norm
from numpy import sort, sin, array, pi

from sklearn.utils import check_random_state
from sklearn.base import BaseEstimator

class SinePattern(BaseEstimator):
    """"Sine pattern.

    The inputs X is drawn uniformly in the 1D interval [min, max].  The target
    y is computed as a sine wave with some period modulated by a sine wave
    envelope with another period and shifted by a constant.

    Attributes
    ----------
    inputs_bound : :rtype: (float, float)
        The bounds (min, max) of the inputs.

    sine_period : float
        The period of the sine wave.

    enveloppe : :rtype: (float, float)
        The (shift, period) of the enveloppe of the enveloppe.

    random_state : int, RandomState instance or None, optional (default=None)
        If int, random_state is the seed used by the random number generator;
        If RandomState instance, random_state is the random number generator;
        If None, the random number generator is the RandomState instance used
        by `np.random`.

    References
    ----------
    * Maxime Sangnier, Olivier Fercoq, Florence D'Alche-Buc
    "Joint quantile regression in vector-valued RKHSs." In Advances in Neural
    Information Processing Systems (2016, pp. 3693-3701).
    """

    def __init__(self, inputs_bound=(0., 1.5), sine_period=1.,
                 enveloppe=(1. / 3., 1.), random_state=None):
        """Initialize the sine pattern.

        Parameters
        ----------
        inputs_bound : :rtype: (float, float), optional (default = (0., 1.5))
            The bounds (min, max) of the inputs

        sine_period : float, optional (default = 1.0)
            The period of the sine wave.

        enveloppe : :rtype: (float, float), optional (default = (1. / 3., 1.))
            The (shift, period) of the enveloppe of the sine wave

        random_state : int, RandomState instance or None,
                       optional (default=None)
            If int, random_state is the seed used by the random number
            generator; If RandomState instance, random_state is the random
            number generator; If None, the random number generator is the
            RandomState instance used by `np.random`.
        """
        self.inputs_bound = inputs_bound
        self.sine_period = sine_period
        self.enveloppe = enveloppe
        self.random_state = random_state

    def __call__(self, n_samples):
        """Generate the data.

        Parameters
        ----------

        n_points : int
            The number of points in the generated dataset.

        Returns
        -------

        inputs : array, shape = [n_samples, 1]
            Returns the generated inputs.

        targets : array, shape = [n_samples, 1]
            Returns the generated targets.
        """
        random_state = check_random_state(self.random_state)
        envelope_period, envelope_shift = self.enveloppe
        inputs = (random_state.rand(n_samples) *
                  (self.inputs_bound[1] - self.inputs_bound[0]) +
                  self.inputs_bound[0])
        # Pattern of the signal
        inputs = sort(inputs)
        pattern = -sin(2 * pi * inputs * self.sine_period)

        # Enveloppe of the signal
        enveloppe = envelope_shift + sin(2 * pi * inputs * envelope_period)

        return (inputs.reshape((n_samples, 1)),
                (pattern * enveloppe).reshape((n_samples, 1)))

def toy_data_quantile(n_samples=50, probs=0.5, pattern=SinePattern(),
                      noise=(1., .2, 0., 1.5), random_state=None):
    """Sine wave toy dataset.

    The target y is computed as a sine curve at of modulated by a sine envelope
    with some period (default 1/3Hz) and mean (default 1).  Moreover, this
    pattern is distorted with a random Gaussian noise with mean 0 and a
    linearly decreasing standard deviation (default from 1.2 at X = 0 to 0.2 at
    X = 1 .  5).

    Parameters
    ----------
    n_samples : int
        Number of samples to generate.

    probs : list or float, shape = [n_quantiles], default=0.5
        Probabilities (quantiles levels).

    pattern : callable, default = SinePattern()
        Callable which generates n_sample 1D data (inputs and targets).

    noise : :rtype: (float, float, float, float)
        Noise parameters (variance, shift, support_min, support_max).

    Returns
    -------
    X : array, shape = [n_samples, 1]
        Input data.

    y : array, shape = [n_sample, 1]
        Targets.

    quantiles : array, shape = [n_samples x n_quantiles]
        True conditional quantiles.
    """
    probs = array(probs, ndmin=1)
    noise_var, noise_shift, noise_min, noise_max = noise

    random_state = check_random_state(random_state)
    pattern.random_state = random_state
    inputs, targets = pattern(n_samples)

    # Noise decreasing std (from noise+0.2 to 0.2)
    noise_std = noise_shift + (noise_var * (noise_max - inputs) /
                               (noise_max - noise_min))
    # Gaussian noise with decreasing std
    add_noise = noise_std * random_state.randn(n_samples, 1)
    observations = targets + add_noise
    quantiles = [targets + array([norm.ppf(p, loc=0., scale=abs(noise_c))
                                  for noise_c in noise_std]) for p in probs]
    return inputs, observations.ravel(), quantiles

max_stars_repo_name: vishalbelsare/operalib

max_issues_repo_name: vishalbelsare/operalib

max_forks_repo_name: RomainBrault/operalib

reason: from numpy,from scipy


text:
import gym
import dso.task.control
import numpy as np

import click

@click.command()
@click.option('--seed', default=0, help='Value for seeding environments.')
@click.option('--rollout_length', default=10, help='How long of episodes to do rollouts for comparing environments.')
def setup_envs_run_rollouts(seed: int, rollout_length: int):
    continuous_envs = {"LunarLanderContinuous-v2": {}, "LunarLanderCustom-v0": dict(reward_shaping_coef=1, continuous=True)}
    discrete_envs = {"LunarLander-v2": {}, "LunarLanderCustom-v0": dict(reward_shaping_coef=1, continuous=False)}
    all_envs = [discrete_envs, continuous_envs]

    results = {}
    for i, envs in enumerate(all_envs):
        for env_name, kwargs in envs.items():
            env = gym.make(env_name, **kwargs)
            env.seed(seed)
            env.action_space.seed(seed)
            env.observation_space.seed(seed)

            results[env_name] = {}

            obs = env.reset()
            r, done = 0, False
            for i in range(rollout_length):
                action = env.action_space.sample()

                results[env_name]["state"] = obs
                results[env_name]["reward"] = r
                results[env_name]["done"] = done

                obs, r, done, _ = env.step(action)

            results[env_name]["state"] = obs
            results[env_name]["reward"] = r
            results[env_name]["done"] = done

        names = list(envs.keys())
        print(f"Comparing: {names}...")
        assert np.array_equal(results[names[0]]['state'], results[names[1]]['state']), "States do not match between above environments."
        assert np.array_equal(results[names[0]]['reward'], results[names[1]]['reward']), "Rewards do not match between above environments."
        assert np.array_equal(results[names[0]]['done'], results[names[1]]['done']), "Episode dones do not match between above environments."
        print(f"Tests passed.")

    return 0

if __name__ == '__main__':
    setup_envs_run_rollouts()

max_stars_repo_name: ryok/deep-symbolic-optimization

max_issues_repo_name: ryok/deep-symbolic-optimization

max_forks_repo_name: ryok/deep-symbolic-optimization


text:
import cv2
import numpy as np
from PIL import Image, ImageFont, ImageDraw, ImageEnhance
import pytesseract
import mechanize
import nltk
from bs4 import BeautifulSoup
from html2text import html2text
import re
from googlesearch import search

pytesseract.pytesseract.tesseract_cmd = 'C://Users//Mr-Robot//Desktop//KBC-Hack//Tesseract-OCR/tesseract'

def word_count(str,A,B,C,D):
    cnt_a = str.count(A)
    cnt_b = str.count(B)
    cnt_c = str.count(C)
    cnt_d = str.count(D)
    max= cnt_a+cnt_b+cnt_c+cnt_d
    print("-----------The answer probabilty--------")
    print("A : ",A,int((cnt_a/max)*100),"%")
    print("B : ",B,int((cnt_b/max)*100),'%')
    print("C : ",C,int((cnt_c/max)*100),'%')
    print("D : ",D,int((cnt_d/max)*100),'%')

def clean_html(html):
    # First we remove inline JavaScript/CSS:
    cleaned = re.sub(r"(?is)<(script|style).*?>.*?(</\1>)", "", html.strip())
    # Then we remove html comments. This has to be done before removing regular
    # tags since comments can contain '>' characters.
    cleaned = re.sub(r"(?s)<!--(.*?)-->[\n]?", "", cleaned)
    # Next we can remove the remaining tags:
    cleaned = re.sub(r"(?s)<.*?>", " ", cleaned)
    # Finally, we deal with whitespace
    cleaned = re.sub(r" ", " ", cleaned)
    cleaned = re.sub(r"  ", " ", cleaned)
    cleaned = re.sub(r"  ", " ", cleaned)
    return cleaned.strip()

def text_proceesor(text):
        string = ''
        text = text.lower()
        l = text.splitlines()
        for i in l:
            string = string + i + ' '
        return string

img = Image.open('kbc image 1.jpeg')
draw = ImageDraw.Draw(img)
draw.rectangle(((60,437), (475,263)),fill=None)#question
draw.rectangle(((54,551), (482,473)),fill=None)#A
draw.rectangle(((54,662), (482,588)),fill=None)#B
draw.rectangle(((54,777), (482,699)),fill=None)#C
draw.rectangle(((54,890), (482,814)),fill=None)#D
#img.save("with_boundaries.png")
question = img.crop((60,263,475,437))#question_cropped
#question.save("question.png")
A = img.crop((54,473,475,551))#option_A_cropped
#A.save("option_A.png")
B = img.crop((60,588,475,662))#option_B_cropped
#B.save("option_B.png")
C = img.crop((60,699,475,777))#option_C_cropped
#C.save("option_C.png")
D = img.crop((60,814,475,890))#option_D_cropped
#D.save("option_D.png")
# A.show()
# B.show()
# C.show()
# D.show()
text = pytesseract.image_to_string(question)
question = text_proceesor(text)
A=pytesseract.image_to_string(A)
A = text_proceesor(A)
B=pytesseract.image_to_string(B)
B = text_proceesor(B)
C=pytesseract.image_to_string(C)
C = text_proceesor(C)
D=pytesseract.image_to_string(D)
D = text_proceesor(D)
# print(A,B,C,D,sep="\n")
x = search(question, tld="com", num=10, stop=1, pause=2)
for i in x:
    url = i
    br = mechanize.Browser()
    br.set_handle_robots(False)
    br.addheaders = [('User-agent', 'Firefox')]
    html = br.open(url).read().decode('utf-8')
    cleanhtml = clean_html(html)
    text = html2text(cleanhtml)
    text = text_proceesor(text)
    soup = BeautifulSoup(html,features="lxml")
    text2 = soup.get_text()
    #I haven't used BeatifulSoup for text as its text is not clean but its a better method
    word_count(text,A,B,C,D)
    #print(text)

max_stars_repo_name: kartikeyangupta/KBC-Hack

max_issues_repo_name: kartikeyangupta/KBC-Hack

max_forks_repo_name: kartikeyangupta/KBC-Hack


text:
#!/usr/bin/env python3

# The MIT License (MIT)
# =====================
#
# Copyright © 2020 Azavea
#
# Permission is hereby granted, free of charge, to any person
# obtaining a copy of this software and associated documentation
# files (the “Software”), to deal in the Software without
# restriction, including without limitation the rights to use,
# copy, modify, merge, publish, distribute, sublicense, and/or sell
# copies of the Software, and to permit persons to whom the
# Software is furnished to do so, subject to the following
# conditions:
#
# The above copyright notice and this permission notice shall be
# included in all copies or substantial portions of the Software.
#
# THE SOFTWARE IS PROVIDED “AS IS”, WITHOUT WARRANTY OF ANY KIND,
# EXPRESS OR IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES
# OF MERCHANTABILITY, FITNESS FOR A PARTICULAR PURPOSE AND
# NONINFRINGEMENT. IN NO EVENT SHALL THE AUTHORS OR COPYRIGHT
# HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER LIABILITY,
# WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING
# FROM, OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR
# OTHER DEALINGS IN THE SOFTWARE.

import argparse
import copy
import csv
import glob
import math
import re

import numpy as np
import scipy.ndimage

import rasterio as rio
import rasterio.transform

def cli_parser() -> argparse.ArgumentParser:
    parser = argparse.ArgumentParser()
    parser.add_argument('--imagery-dir', required=True, type=str)
    parser.add_argument('--elevation-dir', required=False, type=str)
    parser.add_argument('--flood-label-dir', required=True, type=str)
    parser.add_argument('--perm-label-dir', required=False, type=str)
    parser.add_argument('--test-set-csv', required=False, type=str)
    parser.add_argument('--imagery-file', required=True, type=str)
    parser.add_argument('--labels-file', required=True, type=str)
    parser.add_argument('--regexp', required=False,
                        type=str, default='[Bb]olivia')
    parser.add_argument('--mode', choices=['train', 'test'], default='train')
    parser.add_argument('--chip-size', required=False, type=int, default=512)
    return parser

if __name__ == '__main__':
    args = cli_parser().parse_args()
    regexp = re.compile(args.regexp)
    imagery = sorted(glob.glob('{dir}/*.*'.format(dir=args.imagery_dir)))
    water = sorted(glob.glob('{dir}/*.*'.format(dir=args.flood_label_dir)))
    if args.perm_label_dir:
        perm = sorted(glob.glob('{dir}/*.*'.format(dir=args.perm_label_dir)))
    else:
        perm = None

    if not len(imagery) == len(water):
        print('WARNING: different number of imagery and water chips')
        water_set = set(map(lambda t: t.split(
            '/')[-1].replace('_NoQC', ''), water))
        imagery = list(filter(lambda t: t.split(
            '/')[-1].replace('_S1', '') in water_set, imagery))
    assert(len(imagery) == len(water))
    if perm is not None:
        assert(len(water) == len(perm))

    if perm is not None:
        filenames = list(zip(imagery, water, perm))
    else:
        filenames = list(zip(imagery, water))

    test_set = set()
    if args.test_set_csv is not None:
        with open(args.test_set_csv) as csv_file:
            for row in csv.reader(csv_file):
                test_set |= set(map(lambda r: r.split('/')[-1], row))
    if (args.mode == 'test'):
        filenames = list(filter(
            lambda t: t[0].split('/')[-1] in test_set, filenames))
    elif (args.mode == 'train'):
        filenames = list(filter(
            lambda t: not regexp.search(t[0]), filenames))
        filenames = list(filter(
            lambda t: t[0].split('/')[-1] not in test_set, filenames))

    # Dimensions
    sqrt_chip_n = int(math.ceil(math.sqrt(len(filenames))))
    width = sqrt_chip_n * args.chip_size
    height = sqrt_chip_n * args.chip_size

    # Profiles
    with rio.open(imagery[0], 'r') as ds:
        images_profile = copy.deepcopy(ds.profile)
    images_profile.update(
        compress='deflate',
        predictor=2,
        driver='GTiff',
        width=width,
        height=height,
        tiled=True,
        bigtiff=True,
        transform=rasterio.transform.from_bounds(
            31.132830, 29.978150, 31.135448, 29.980260, width, height),
        crs='epsg:4326',
        dtype=np.float32,
        nodata=np.nan,
    )
    if args.elevation_dir is not None:
        images_profile.update(count=images_profile.get('count') + 1)
    labels_profile = copy.deepcopy(images_profile)
    labels_profile.update(
        count=1,
        dtype=np.uint8,
        nodata=0,
    )

    # Arrays
    images = np.zeros((images_profile.get('count'), width,
                       height), dtype=images_profile.get('dtype'))
    labels = np.zeros((1, width, height), dtype=labels_profile.get('dtype'))

    # Data
    i = 0
    for t in filenames:
        if (i % 107) == 0:
            print('.')
        x = (i // sqrt_chip_n) * args.chip_size
        y = (i % sqrt_chip_n) * args.chip_size
        with rio.open(t[0], 'r') as ds:
            chip = ds.read()
            if args.elevation_dir:
                dem_filename = t[0].replace(
                    args.imagery_dir, args.elevation_dir)
                with rio.open(dem_filename, 'r') as ds_dem:
                    dem_chip = ds_dem.read().astype(np.float32)
                nodata_mask = (dem_chip == -32768)
                dem_chip -= np.min(np.extract(nodata_mask != True, dem_chip))
                dem_chip[nodata_mask] = np.nan
                chip = np.concatenate([chip, dem_chip], axis=0)
            ratiox = float(args.chip_size) / chip.shape[1]
            ratioy = float(args.chip_size) / chip.shape[2]
            if ratiox != 1.0 or ratioy != 1.0:
                chip = scipy.ndimage.zoom(
                    chip, [1.0, ratiox, ratioy], order=0, prefilter=False)
            images[:, x:(x+args.chip_size), y:(y+args.chip_size)] = chip
        with rio.open(t[1], 'r') as ds:
            chip = ds.read() + 1
            ratiox = float(args.chip_size) / chip.shape[1]
            ratioy = float(args.chip_size) / chip.shape[2]
            if ratiox != 1.0 or ratioy != 1.0:
                chip = scipy.ndimage.zoom(
                    chip, [1.0, ratiox, ratioy], order=0, prefilter=False)
            labels[:, x:(x+args.chip_size), y:(y+args.chip_size)] = chip
        if len(t) == 3:
            with rio.open(t[2], 'r') as ds:
                chip = ds.read()
                ratiox = float(args.chip_size) / chip.shape[1]
                ratioy = float(args.chip_size) / chip.shape[2]
                if ratiox != 1.0 or ratioy != 1.0:
                    chip = scipy.ndimage.zoom(
                        chip, [1.0, ratiox, ratioy], order=0, prefilter=False)
                labels[:, x:(x+args.chip_size),
                       y:(y+args.chip_size)] += chip
        i += 1

    print('Writing ...')
    with rio.open(args.imagery_file, 'w', **images_profile) as ds:
        ds.write(images)
    with rio.open(args.labels_file, 'w', **labels_profile) as ds:
        ds.write(labels)

max_stars_repo_name: geotrellis/deeplab-nlcd

max_issues_repo_name: geotrellis/deeplab-nlcd

max_forks_repo_name: geotrellis/deeplab-nlcd

reason: import numpy,import scipy


text:
# Copyright 2021 NREL

# Licensed under the Apache License, Version 2.0 (the "License"); you may not
# use this file except in compliance with the License. You may obtain a copy of
# the License at http://www.apache.org/licenses/LICENSE-2.0

# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS, WITHOUT
# WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the
# License for the specific language governing permissions and limitations under
# the License.

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

from SALib.sample import saltelli
from SALib.analyze import sobol

from pandas.core.base import DataError

from .. import floris_tools as ftools

class floris_sobol_analysis():
    def __init__(self, fi, problem, calc_second_order=False):
        self.fi = fi

        # Default parameters
        self.param_dict = {
                      'ad': 0.0,
                      'alpha': 0.58,
                      'bd': 0.0,
                      'beta': 0.077,
                      'eps_gain': 0.20,
                      'ka': 0.38,
                      'kb': 0.004,
                      'ti_ai': 0.80,
                      'ti_constant': 0.50,
                      'ti_downstream': -0.32,
                      'ti_initial': 0.10
                      }

        self.problem = problem
        self.Si = None
        self.calc_second_order = calc_second_order
        self.samples_x = None
        self.samples_y = None
        self.N = None

    def _get_model_params_dict(self, id):
        new_params_dict = {}
        n_params = len(self.problem['names'])
        for i in range(n_params):
            var_name = self.problem['names'][i]
            value = self.samples_x[id, i]
            new_params_dict[var_name] = value

        keys = list(new_params_dict.keys())
        values = list(new_params_dict.values())
        for i in range(len(keys)):
            self.param_dict[keys[i]] = values[i]

        params = {
            "Wake Deflection Parameters":
                {"ad": self.param_dict['ad'],
                 "alpha": self.param_dict['alpha'],
                 "bd": self.param_dict['bd'],
                 "beta": self.param_dict['beta'],
                 "eps_gain": self.param_dict['eps_gain'],
                 "ka": self.param_dict['ka'],
                 "kb": self.param_dict['kb'],
                 "use_secondary_steering": True},
            "Wake Turbulence Parameters":
                {"ti_ai": self.param_dict['ti_ai'],
                 "ti_constant": self.param_dict['ti_constant'],
                 "ti_downstream": self.param_dict['ti_downstream'],
                 "ti_initial": self.param_dict['ti_initial']},
            "Wake Velocity Parameters":
                {"alpha": self.param_dict['alpha'],
                 "beta": self.param_dict['beta'],
                 "eps_gain": self.param_dict['eps_gain'],
                 "ka": self.param_dict['ka'],
                 "kb": self.param_dict['kb'],
                 "use_yaw_added_recovery": True}
        }

        return params
        # self.fi.set_model_parameters(params=params, verbose=False)

    def _create_evals_dataframe(self):
        Nt = self.samples_x.shape[0]
        params_array = [[] for _ in range(Nt)]
        for id in range(Nt):
            params_array[id] = self._get_model_params_dict(id)
        df = pd.DataFrame({'model_params_dict': params_array})
        self.df_eval = df

    # Step 1: generating samples for a particular problem
    def generate_samples(self, N, problem=None, calc_second_order=None):
        if problem is None:
            problem = self.problem

        if calc_second_order is None:
            calc_second_order = self.calc_second_order
        self.problem = problem
        self.calc_second_order = calc_second_order

        Ns = N
        self.samples_x = saltelli.sample(
            problem, Ns, calc_second_order=calc_second_order)

        self.N = self.samples_x.shape[0]
        self.samples_y = np.zeros(self.N)
        self._create_evals_dataframe()

    def calculate_wfpower_for_samples(self, num_threads=1):
        if self.samples_x is None:
            raise DataError('Please run generate_samples first.')

        # Copy and write wd and ws to dataframe
        # Nt = self.df_eval.shape[0]
        df = self.df_eval
        df['wd'] = self.fi.floris.farm.wind_direction[0]
        df['ws'] = self.fi.floris.farm.wind_speed[0]

        # Calculate floris predictions
        df_out = ftools.calc_floris(df, self.fi, num_threads=10, num_workers=2)
        pow_cols = ['pow_%03d' % ti for ti in range(len(self.fi.layout_x))]
        self.samples_y = np.array(df_out[pow_cols].sum(axis=1), dtype=float)

        return self.samples_y

    # def calculate_aep_for_samples(self, wd, ws, freq):
    #     if self.samples_x is None:
    #         raise DataError('Please run generate_samples first.')

    #     print('Calculating AEP for %d samples.' % self.N)
    #     for i in range(self.N):
    #         self._set_fi_by_sample_id(i)
    #         aep = self.fi.get_farm_AEP(wd=wd, ws=ws, freq=freq)
    #         self.samples_y[i] = aep

    #     return self.samples_y

    def get_sobol_sensitivity_indices(self, verbose=False):
        self.Si = sobol.analyze(self.problem, self.samples_y,
                                print_to_console=verbose,
                                calc_second_order=self.calc_second_order)
        return self.Si

    def plot_sobol_results(self, save_path=None, fig_format='png', fig_dpi=200):
        if self.Si is None:
            raise DataError('No Sobol results to show. ' +
                            'Have you run get_sobol_sensitivity_indices()?')
        problem = self.problem
        if self.calc_second_order:
            fig, ax = plt.subplots(2)
        else:
            fig, ax = plt.subplots()
            ax = [ax]
        width = 0.30
        Nv = problem['num_vars']

        # Plot first order sensitivity plus uncertainties
        bar1 = ax[0].bar(np.arange(Nv) - width/2, self.Si['S1'], width=0.3,
                         align='center', label='First order sensitivity',
                         color='deepskyblue', edgecolor='black', hatch='//')
        # ax[0].bar_label(bar1, padding=7, fmt='%.1e')

        for i in range(Nv):
            ymean = self.Si['S1'][i]
            ystd = self.Si['S1_conf'][i]
            ax[0].plot(np.repeat(i - width/2, 3),
                       ymean + np.array([-ystd, 0., ystd]),
                       color='black')
            ax[0].plot(i - width/2, ymean, 'o', color='black')
            ax[0].plot([i - 0.75*width, i-0.25*width],
                       [ymean-ystd, ymean-ystd], color='black')
            ax[0].plot([i - 0.75*width, i-0.25*width],
                       [ymean+ystd, ymean+ystd], color='black')

        # Plot total sensitivity including uncertainties
        bar2 = ax[0].bar(np.arange(Nv) + width/2, self.Si['ST'], width=0.3,
                         align='center', label='Total sensitivity',
                         color='orangered', edgecolor='black', hatch='.')
        # ax[0].bar_label(bar2, padding=7, fmt='%.1e')

        for i in range(Nv):
            ymean = self.Si['ST'][i]
            ystd = self.Si['ST_conf'][i]
            ax[0].plot(np.repeat(i + width/2, 3),
                       ymean + np.array([-ystd, 0., ystd]),
                       color='black')
            ax[0].plot(i + width/2, ymean, 'o', color='black')
            ax[0].plot([i + 0.25*width, i+0.75*width],
                       [ymean-ystd, ymean-ystd], color='black')
            ax[0].plot([i + 0.25*width, i+0.75*width],
                       [ymean+ystd, ymean+ystd], color='black')

        # Plot settings
        ax[0].set_xticks(range(Nv))
        ax[0].set_xticklabels(problem['names'])
        ax[0].set_title('First order effects')
        ax[0].set_ylabel('Sensitivity (-)')
        ax[0].legend()
        ax[0].grid('minor')
        ax[0].set_ylim([0, 1.1])

        # Second order tabular/imshow plot, if calculated
        if self.calc_second_order:
            im = ax[1].imshow(self.Si['S2'][:-1,:])
            ax[1].set_xticks(range(Nv))
            ax[1].set_yticks(range(Nv-1))
            ax[1].set_xticklabels(self.problem['names'])
            ax[1].set_yticklabels(self.problem['names'][:-1])
            ax[1].set_title('Second order effects')
            for ii in range(Nv-1):
                for jj in range(Nv):
                    ax[1].text(y=ii, x=jj,
                               s='%.3e' % self.Si['S2'][ii, jj],
                               ha='center', color='white')
            plt.copper()
            plt.colorbar(im, ax=ax[1])
            fig.tight_layout()

            if save_path is not None:
                plt.savefig(save_path + '/Sobol_sensitivity_indices.%s'
                            % fig_format, dpi=fig_dpi)

        return fig, ax

    def plot_convergence(self, save_path=None, fig_format='png', fig_dpi=200):
        print('Analyzing convergence...')

        # Create copies of original results
        samples_x_full = self.samples_x
        samples_y_full = self.samples_y
        Si_full = self.Si
        N = self.N

        # Create Sobol outputs for data subsets
        Si_list = []
        self.generate_samples(N=10)
        dN = self.samples_x.shape[0]
        N_array = np.arange(dN, N, dN, dtype='int')
        if N_array[-1] < N:
            N_array = np.append(N_array, N)

        for n in N_array:
            self.samples_y = samples_y_full[0:n]
            self.get_sobol_sensitivity_indices(verbose=False)
            Si_list.append(self.Si)

        # Restore original results
        self.samples_x = samples_x_full
        self.samples_y = samples_y_full
        self.Si = Si_full
        self.N = N

        # Plot convergence for S1
        Nv = self.problem['num_vars']
        fig, ax = plt.subplots(nrows=Nv, sharex=True)
        for i in range(Nv):
            ax[i].errorbar(x=N_array, y=[S['S1'][i] for S in Si_list],
                           yerr=[S['S1_conf'][i] for S in Si_list],
                           label='First order sensitivity (S1)')
            ax[i].errorbar(x=N_array, y=[S['ST'][i] for S in Si_list],
                           yerr=[S['ST_conf'][i] for S in Si_list],
                           label='Total sensitivity (ST)')
            ax[i].plot([0, self.N], [self.Si['S1'][i], self.Si['S1'][i]],
                       '--', color='gray', label=None)
            ax[i].plot([0, self.N], [self.Si['ST'][i], self.Si['ST'][i]],
                       '--', color='gray', label=None)
            ax[i].legend()
            ax[i].grid('minor')
            ax[i].set_ylabel(self.problem['names'][i])
            ax[i].set_xlabel('Number of iterations')

        if save_path is not None:
            plt.savefig(save_path + '/Sobol_convergence_order1.%s'
                        % fig_format, dpi=fig_dpi)

        if self.calc_second_order:
            for i in range(Nv-1):
                for j in range(i+1, Nv):
                    fig, ax = plt.subplots()
                    ax.errorbar(
                        x=N_array,
                        y=[S['S2'][i, j] for S in Si_list],
                        yerr=[S['S2_conf'][i, j] for S in Si_list],
                        label=('Second order sensitivity with %s'
                               % self.problem['names'][j])
                        )
                    ax.plot(
                        [0, self.N],
                        [self.Si['S2'][i, j], self.Si['S2'][i, j]],
                        '--', color='black', label=None
                        )
                    ax.legend()
                    ax.set_ylabel(self.problem['names'][i])
                    ax.set_xlabel('Number of iterations')
                    ax.grid('minor')
                    if save_path is not None:
                        plt.savefig(
                            save_path + '/Sobol_convergence_order2_%s.%s'
                            % (self.problem['names'][j], fig_format),
                            dpi=fig_dpi)

        return fig, ax

max_stars_repo_name: NREL/flasc

max_issues_repo_name: NREL/flasc

max_forks_repo_name: NREL/flasc


text:
import os
from optparse import OptionParser

import numpy as np
import pandas as pd

import file_handling as fh

"""
Randomly split a jsonlist (and optionally, a labels and/or covaraites dataframe) into train and test.
"""

def main():
    usage = "%prog input.jsonlist ouput_dir [labels.csv covariates.csv ...]"
    parser = OptionParser(usage=usage)
    parser.add_option('--test_prop', dest='test_prop', default=0.2,
                      help='proportion of documents to use for test data: default=%default')
    parser.add_option('--train', dest='train', default='train',
                      help='output prefix for training data: default=%default')
    parser.add_option('--test', dest='test', default='test',
                      help='output prefix for test data: default=%default')

    (options, args) = parser.parse_args()
    infile = args[0]
    output_dir = args[1]
    if len(args) > 2:
        csv_files = args[2:]
    else:
        csv_files = []

    test_prop = float(options.test_prop)
    train_prefix = options.train
    test_prefix = options.test

    print("Reading", infile)
    items = fh.read_jsonlist(infile)
    n_items = len(items)

    n_test = int(n_items * test_prop)
    print("Creating random test set of %d items" % n_test)
    n_train = n_items - n_test
    train_indices = np.random.choice(np.arange(n_items), size=n_train, replace=False)
    test_indices = list(set(range(n_items)) - set(train_indices))

    train_items = [items[i] for i in train_indices]
    test_items = [items[i] for i in test_indices]

    fh.write_jsonlist(train_items, os.path.join(output_dir, train_prefix + '.jsonlist'))
    fh.write_jsonlist(test_items, os.path.join(output_dir, test_prefix + '.jsonlist'))

    for file in csv_files:
        print(file)
        basename = os.path.basename(file)
        df = pd.read_csv(file, header=0, index_col=0)
        train_df_index = [df.index[i] for i in train_indices]
        train_df = df.loc[train_df_index]
        train_df.to_csv(os.path.join(output_dir, train_prefix + '.' + basename))

        test_df_index = [df.index[i] for i in test_indices]
        test_df = df.loc[test_df_index]
        test_df.to_csv(os.path.join(output_dir, test_prefix + '.' + basename))

if __name__ == '__main__':
    main()

max_stars_repo_name: WillSkywalker/scholar

max_issues_repo_name: WillSkywalker/scholar

max_forks_repo_name: WillSkywalker/scholar


text:
from typing import Callable, Union

import numpy as np
import tensorflow as tf

from ..config import default_float
from ..inducing_variables import InducingPoints, InducingVariables
from .model import BayesianModel
from .training_mixins import Data, ExternalDataTrainingLossMixin

def inducingpoint_wrapper(
    inducing_variable: Union[InducingVariables, tf.Tensor, np.ndarray]
) -> InducingVariables:
    """
    This wrapper allows transparently passing either an InducingVariables
    object or an array specifying InducingPoints positions.
    """
    if not isinstance(inducing_variable, InducingVariables):
        inducing_variable = InducingPoints(inducing_variable)
    return inducing_variable

def _assert_equal_data(data1, data2):
    if isinstance(data1, tf.Tensor) and isinstance(data2, tf.Tensor):
        tf.debugging.assert_equal(data1, data2)
    else:
        for v1, v2 in zip(data1, data2):
            tf.debugging.assert_equal(v1, v2)

def training_loss_closure(
    model: BayesianModel, data: Data, **closure_kwargs
) -> Callable[[], tf.Tensor]:
    if isinstance(model, ExternalDataTrainingLossMixin):
        return model.training_loss_closure(data, **closure_kwargs)
    else:
        _assert_equal_data(model.data, data)
        return model.training_loss_closure(**closure_kwargs)

def training_loss(model: BayesianModel, data: Data) -> tf.Tensor:
    if isinstance(model, ExternalDataTrainingLossMixin):
        return model.training_loss(data)
    else:
        _assert_equal_data(model.data, data)
        return model.training_loss()

def maximum_log_likelihood_objective(model: BayesianModel, data: Data) -> tf.Tensor:
    if isinstance(model, ExternalDataTrainingLossMixin):
        return model.maximum_log_likelihood_objective(data)
    else:
        _assert_equal_data(model.data, data)
        return model.maximum_log_likelihood_objective()

def data_input_to_tensor(structure):
    """
    Converts non-tensor elements of a structure to TensorFlow tensors retaining the structure itself.
    The function doesn't keep original element's dtype and forcefully converts
    them to GPflow's default float type.
    """

    def convert_to_tensor(elem):
        if tf.is_tensor(elem):
            return elem
        elif isinstance(elem, np.ndarray):
            return tf.convert_to_tensor(elem)
        return tf.convert_to_tensor(elem, dtype=default_float())

    return tf.nest.map_structure(convert_to_tensor, structure)

max_stars_repo_name: yiwei-prowler/GPflow

max_issues_repo_name: yiwei-prowler/GPflow

max_forks_repo_name: yiwei-prowler/GPflow


text:
import sys
import os
import ntpath
import logging

from pathlib import Path
from astropy.io import fits
from PyQt5 import QtWidgets
from PyQt5 import QtGui, QtCore

'''
FITS Header parser
'''
class FitsHeaderTab():
    logger = logging.getLogger(__name__)

    def __init__(self, mainW, app):
        super().__init__()
        self.mainW = mainW
        self.app = app

    def readHeaders(self):
        filename, _ = QtWidgets.QFileDialog.getOpenFileName(
            self.mainW,
            'Select a FITS file to open…',
            QtCore.QDir.homePath(),
            'FITS Files (*.fits *.fit *.FIT *.FITS) ;; All Files (*)'
        )
        if filename:
            with open(filename) as fh:
                self.mainW.ui.plainTextEditFitsHeader.clear()
                self.mainW.ui.plainTextEditMsg.clear()
                self.mainW.ui.lineEditFitsHeader.setText(filename)
                self.logger.info(f"Opened {filename}")

                hdu = fits.open(filename)
                hdu.verify('silentfix')
                hdr = hdu[0].header
                for i,key in enumerate(hdr):
                    sval = str(hdr[key])[0:70]
                    if key  in self.app.filterDictToList('fitsHeader'):
                        color = "green"
                        if not hdr[key]:
                            sval= "Empty value: set a default value in Settings before importing"
                    else:
                        color = "#ccc"
                    s = "
"+str(key)+"\t"+sval+"
"
                    self.mainW.ui.plainTextEditFitsHeader.appendHtml(s)

                l = ('FILE', 'HASH', 'ALT', 'AZ')
                requestedH = set(self.app.filterDictToList('fitsHeader')) - set(l)
                missingH = set(requestedH) - set(hdr)- set(self.app.filterDictToList('fitsDefault', 'fitsHeader'))
                matchedH =  set(requestedH).intersection(set(hdr))
                defaultedH = set(self.app.filterDictToList('fitsDefault', 'fitsHeader')) -set(hdr)
                msg=""
                if len(matchedH)> 0:
                    msg +="Matched keywords: "+','.join(matchedH)+ "
"
                if len(defaultedH)> 0:
                    msg +="Keywords with a default: "+','.join(defaultedH)+ "
"
                if len(missingH)> 0:
                    msg +="Missing keywords: "+','.join(missingH)+ ""
                else:
                    msg +="This file can be imported in Astrodom."

                self.mainW.ui.plainTextEditMsg.appendHtml(msg)

max_stars_repo_name: fenriques/AstroDom

max_issues_repo_name: fenriques/AstroDom

max_forks_repo_name: fenriques/AstroDom

reason: from astropy


text:
import logging

import numpy as np

from metadrive.component.vehicle_module.PID_controller import PIDController
from metadrive.policy.base_policy import BasePolicy
from metadrive.policy.manual_control_policy import ManualControlPolicy
from metadrive.utils.math_utils import not_zero, wrap_to_pi, point_distance
from metadrive.utils.scene_utils import is_same_lane_index, is_following_lane_index

from metadrive.engine.core.manual_controller import KeyboardController, SteeringWheelController
from metadrive.utils import clip
from metadrive.examples import expert
from metadrive.policy.env_input_policy import EnvInputPolicy
from direct.controls.InputState import InputState
from metadrive.engine.engine_utils import get_global_config

import gym
from gym import Wrapper
from gym import spaces
# from metadrive.envs.base_env import BaseEnv
# from metadrive.envs.metadrive_env import MetaDriveEnv
from typing import Callable

class ActionType(object):
    #def __init__(self, env: 'MetaDriveEnv', **kwargs) -> None:
    def __init__(self, env, **kwargs) -> None:
        self.env = env
        self.__controlled_vehicle = None

    def space(self) -> spaces.Space:
        raise NotImplementedError

    @property
    def vehicle_class(self) -> Callable:
        raise NotImplementedError

    def act(self, action) -> None:
        raise NotImplementedError

    @property
    def controlled_vehicle(self):
        return self.__controlled_vehicle or self.env.vehicle

    @controlled_vehicle.setter
    def controlled_vehicle(self, vehicle):
        self.__controlled_vehicle = vehicle

class DiscreteMetaAction(ActionType):
    ACTIONS_ALL = {0: 'LANE_LEFT', 1: 'IDLE', 2: 'LANE_RIGHT', 3: 'FASTER', 4: 'SLOWER', 5: 'Holdon'}

    def __init__(self, **kwargs):
        self.actions = self.ACTIONS_ALL
        self.actions_indexes = {v: k for k, v in self.actions.items()}

    def space(self) -> spaces.Space:
        return spaces.Discrete(5)
        #return spaces.Discrete(len(self.actions))

    # @property
    # def vehicle_class(self) -> Callable:
    #     return MDPVehicle

    def act(self, action: int) -> None:
        self.controlled_vehicle.act(self.actions[action])