text:
import re
import collections
import collections.abc

from copy import copy
from pathlib import Path
from numbers import Number
from operator import truediv, sub, gt
from abc import abstractmethod
from itertools import chain, repeat, accumulate, groupby, count
from typing import Any, Dict, List, Set, Tuple, Optional, Sequence, Hashable, Iterator, Union, Type, Callable, NamedTuple
from coba.backports import Literal

from coba.statistics import Mean, StandardDeviation, StandardErrorOfMean, BootstrapConfidenceInterval, BinomialConfidenceInterval, PointAndInterval
from coba.contexts import CobaContext
from coba.exceptions import CobaException
from coba.utilities import PackageChecker
from coba.pipes import Pipes, Sink, Source, JsonEncode, JsonDecode, DiskSource, DiskSink, ListSource, ListSink, Foreach

def exponential_moving_average(values:Sequence[float], span:int=None):
    #exponential moving average identical to Pandas df.ewm(span=span).mean()
    alpha = 2/(1+span)
    cumwindow  = list(accumulate(values          , lambda a,v: v + (1-alpha)*a))
    cumdivisor = list(accumulate([1.]*len(values), lambda a,v: v + (1-alpha)*a)) #type: ignore
    return list(map(truediv, cumwindow, cumdivisor))

def moving_average(values:Sequence[float], span:int=None):

    if span == 1:
        return values

    if span is None or span >= len(values):
        return [a/n for n,a in enumerate(accumulate(values),1)]

    window_sums  = accumulate(map(sub, values, chain(repeat(0,span),values)))
    window_sizes = chain(range(1,span), repeat(span))
    return list(map(truediv,window_sums,window_sizes))

class Table:
    """A container class for storing tabular data."""

    def __init__(self, name:str, primary_cols: Sequence[str], rows: Sequence[Dict[str,Any]], preferred_cols: Sequence[str] = []):
        """Instantiate a Table.

        Args:
            name: The name of the table. Used for display purposes.
            primary_cols: Table columns used to make each row's tuple "key".
            rows: The actual rows that should be stored in the table. Each row is required to contain the given primary_cols.
            preferred_cols: A list of columns that we prefer be displayed immediately after primary columns. All remaining
                columns (i.e., neither primary nor preferred) will be ordered alphabetically.
        """
        self._name    = name
        self._primary = primary_cols

        for row in rows:
            assert len(row.keys() & primary_cols) == len(primary_cols), 'A Table row was provided without a primary key.'

        all_columns: Set[str] = set()
        for row in rows:
            all_columns |= {'index'} if '_packed' in row else set()
            all_columns |= row.keys()-{'_packed'}
            all_columns |= all_columns.union(row.get('_packed',{}).keys())

        col_priority = list(chain(primary_cols + ['index'] + preferred_cols + sorted(all_columns)))

        self._columns = sorted(all_columns, key=col_priority.index)
        self._rows_keys: Dict[Hashable, None         ] = {}
        self._rows_flat: Dict[Hashable, Dict[str,Any]] = {}
        self._rows_pack: Dict[Hashable, Dict[str,Any]] = {}

        for row in rows:
            row_key  = row[primary_cols[0]] if len(primary_cols) == 1 else tuple(row[col] for col in primary_cols)
            row_pack = row.pop('_packed',{})
            row_flat = row

            if row_pack:
                row_pack['index'] = list(range(1,len(list(row_pack.values())[0])+1))

            self._rows_keys[row_key] = None
            self._rows_pack[row_key] = row_pack
            self._rows_flat[row_key] = row_flat

        self._rows_keys = collections.OrderedDict(zip(sorted(list(self._rows_keys.keys())),repeat(None)))

    @property
    def name(self) -> str:
        """The name of the table."""
        return self._name

    @property
    def keys(self) -> Sequence[Hashable]:
        """Keys for accessing data in the table."""
        return list(self._rows_keys.keys())

    @property
    def columns(self) -> Sequence[str]:
        """The columns in the table."""
        return self._columns

    @property
    def dtypes(self) -> Sequence[Type[Union[int,float,bool,object]]]:
        """The dtypes for the columns in the table."""
        flats = self._rows_flat
        packs = self._rows_pack

        columns_packed = [ any([ col in packs[key] for key in self.keys]) for col in self.columns ]
        columns_values = [ [flats[key].get(col, packs[key].get(col, self._default(col))) for key in self.keys] for col in self.columns ]

        return [ self._infer_type(column_packed, column_values) for column_packed, column_values in zip(columns_packed,columns_values)]

    def filter(self, row_pred:Callable[[Dict[str,Any]],bool] = None, **kwargs) -> 'Table':
        """Filter to specific rows.

        Args:
            pred: A predicate that returns true for row dictionaries that should be kept.
            kwargs: key value pairs where the key is the column and the value indicates what
                value a row should have in that column to be kept. Keeping logic depends on
                the row value type and the kwargs value type. If kwarg value == row value keep
                the row. If kwarg value is callable pass the row value to the predicate. If
                kwarg value is a collection keep the row if the row value is in the collection.
                If kwarg value is a string apply a regular expression match to the row value.
        """

        def satisifies_filter(col_filter,col_value):
            if col_filter == col_value:
                return True

            if isinstance(col_filter,Number) and isinstance(col_value,str):
                return re.search(f'(\D|^){col_filter}(\D|$)', col_value)

            if isinstance(col_filter,str) and isinstance(col_value,str):
                return re.search(col_filter, col_value)

            if callable(col_filter):
                return col_filter(col_value)

            return False

        def satisfies_all_filters(key):
            row = self[key]

            row_filter_results = [ row_pred is None or row_pred(row) ]
            col_filter_results = [ ]

            for col,col_filter in kwargs.items():

                if isinstance(col_filter,collections.abc.Container) and not isinstance(col_filter,str):
                    col_filter_results.append(row[col] in col_filter or any([satisifies_filter(cf,row[col]) for cf in col_filter]))
                else:
                    col_filter_results.append(satisifies_filter(col_filter,row.get(col,self._default(col)) ))

            return all(row_filter_results+col_filter_results)

        new_result = copy(self)
        new_result._rows_keys = collections.OrderedDict(zip(filter(satisfies_all_filters,self.keys), repeat(None)))

        return new_result

    def to_pandas(self) -> Any:
        """Turn the Table into a Pandas data frame."""

        PackageChecker.pandas("Table.to_pandas")
        import pandas as pd #type: ignore
        import numpy as np  #type: ignore #pandas installs numpy so if we have pandas we have numpy

        col_numpy = { col: np.empty(len(self), dtype=dtype) for col,dtype in zip(self.columns,self.dtypes)}

        row_index = 0

        for key in self.keys:

            flat = self._rows_flat[key]
            pack = self._rows_pack[key]

            pack_size = 1 if not pack else len(pack['index'])

            for col in self.columns:
                if col in pack:
                    val = pack[col]

                elif col in flat:
                    if isinstance(flat[col], (tuple,list)):
                        val = [flat[col]]
                    else:
                        val = flat[col]

                else:
                    val = self._default(col)

                col_numpy[col][row_index:(row_index+pack_size)] = val

            row_index += pack_size

        return pd.DataFrame(col_numpy, columns=self.columns)

    def to_dicts(self) -> Sequence[Dict[str,Any]]:
        """Turn the Table into a sequence of tuples."""

        dicts = []

        for key in self.keys:

            flat = self._rows_flat[key]
            pack = self._rows_pack[key]

            if not pack:
                dicts.append( { col:flat.get(col,self._default(col)) for col in self.columns } )
            else:
                tuples = list(zip(*[pack.get(col,repeat(flat.get(col,self._default(col)))) for col in self.columns]))
                dicts.extend([ dict(zip(self.columns,t)) for t in tuples ])

        return dicts

    def _default(self, column:str) -> Any:
        return [1] if column == "index" else None

    def _infer_type(self, is_packed: bool, values: Sequence[Any]) -> Type[Union[int,float,bool,object]]:

        types: List[Optional[Type[Any]]] = []

        to_type = lambda value: None if value is None else type(value)

        for value in values:
            if is_packed and isinstance(value, (list,tuple)):
                types.extend([to_type(v) for v in value])
            else:
                types.append(to_type(value))

        return self._resolve_types(types)

    def _resolve_types(self, types: Sequence[Optional[Type[Any]]]) -> Type[Union[int,float,bool,object]]:
        types = list(set(types))

        if len(types) == 1 and types[0] in [dict,str]:
            return object

        if len(types) == 1 and types[0] in [int,float,bool]:
            return types[0]

        if all(t in [None,int,float] for t in types):
            return float

        return object

    def __iter__(self) -> Iterator[Dict[str,Any]]:
        for key in self.keys:
            yield self[key]

    def __contains__(self, key: Union[Hashable, Sequence[Hashable]]) -> bool:
        return key in self.keys

    def __str__(self) -> str:
        return str({"Table": self.name, "Columns": self.columns, "Rows": len(self)})

    def _ipython_display_(self):
        #pretty print in jupyter notebook (https://ipython.readthedocs.io/en/stable/config/integrating.html)
        print(str(self))

    def __len__(self) -> int:
        return sum([ len(self._rows_pack[key].get('index',[None])) for key in self.keys ])

    def __getitem__(self, key: Union[Hashable, Sequence[Hashable]]) -> Dict[str,Any]:
        if key not in self._rows_keys: raise KeyError(key)
        return dict(**self._rows_flat[key], **self._rows_pack[key])

class InteractionsTable(Table):
    pass

class TransactionIO_V3(Source['Result'], Sink[Any]):

    def __init__(self, log_file: Optional[str] = None, minify:bool = True) -> None:

        self._log_file = log_file
        self._minify   = minify
        self._source   = DiskSource(log_file) if log_file else ListSource()
        self._sink     = DiskSink(log_file) if log_file else ListSink(self._source.items)

    def write(self, item: Any) -> None:
        if isinstance(self._sink, ListSink):
            self._sink.write(self._encode(item))
        else:
            if not Path(self._sink._filename).exists():self._sink.write('["version",3]')
            self._sink.write(JsonEncode(self._minify).filter(self._encode(item)))

    def read(self) -> 'Result':
        n_lrns   = None
        n_envs   = None
        lrn_rows = {}
        sim_rows = {}
        int_rows = {}

        if isinstance(self._source, ListSource):
            decoded_source = self._source
        else:
            decoded_source = Pipes.join(self._source, Foreach(JsonDecode()))

        for trx in decoded_source.read():

            if not trx: continue

            if trx[0] == "benchmark":
                n_lrns = trx[1]["n_learners"]
                n_envs = trx[1]["n_simulations"]

            if trx[0] == "S":
                sim_rows[trx[1]] = trx[2]

            if trx[0] == "L":
                lrn_rows[trx[1]] = trx[2]

            if trx[0] == "I":
                int_rows[tuple(trx[1])] = trx[2]

        return Result(sim_rows, lrn_rows, int_rows, {"n_learners":n_lrns,"n_environments":n_envs})

    def _encode(self,item):
        if item[0] == "T0":
            return ['benchmark', {"n_learners":item[1], "n_simulations":item[2]}]

        if item[0] == "T1":
            return ["L", item[1], item[2]]

        if item[0] == "T2":
            return ["S", item[1], item[2]]

        if item[0] == "T3":
            rows_T = collections.defaultdict(list)

            for row in item[2]:
                for col,val in row.items():
                    if col == "rewards" : col="reward"
                    if col == "reveals" : col="reveal"
                    rows_T[col].append(val)

            return ["I", item[1], { "_packed": rows_T }]

        return None

class TransactionIO_V4(Source['Result'], Sink[Any]):

    def __init__(self, log_file: Optional[str] = None, minify:bool=True) -> None:

        self._log_file = log_file
        self._minify   = minify
        self._source   = DiskSource(log_file) if log_file else ListSource()
        self._sink     = DiskSink(log_file)   if log_file else ListSink(self._source.items)

    def write(self, item: Any) -> None:
        if isinstance(self._sink, ListSink):
            self._sink.write(self._encode(item))
        else:
            if not Path(self._sink._filename).exists():self._sink.write('["version",4]')
            self._sink.write(JsonEncode(self._minify).filter(self._encode(item)))

    def read(self) -> 'Result':

        exp_dict = {}
        lrn_rows = {}
        env_rows = {}
        int_rows = {}

        if isinstance(self._source, ListSource):
            decoded_source = self._source
        else:
            decoded_source = Pipes.join(self._source, Foreach(JsonDecode()))

        for trx in decoded_source.read():

            if not trx: continue

            if trx[0] == "experiment":
                exp_dict = trx[1]

            if trx[0] == "E":
                env_rows[trx[1]] = trx[2]

            if trx[0] == "L":
                lrn_rows[trx[1]] = trx[2]

            if trx[0] == "I":
                int_rows[tuple(trx[1])] = trx[2]

        return Result(env_rows, lrn_rows, int_rows, exp_dict)

    def _encode(self,item):
        if item[0] == "T0":
            return ['experiment', {"n_learners":item[1], "n_environments":item[2], "description":None} if len(item)==3 else item[1] ]

        if item[0] == "T1":
            return ["L", item[1], item[2]]

        if item[0] == "T2":
            return ["E", item[1], item[2]]

        if item[0] == "T3":
            rows_T = collections.defaultdict(list)

            for row in item[2]:
                for col,val in row.items():
                    if col == "rewards" : col="reward"
                    if col == "reveals" : col="reveal"
                    rows_T[col].append(val)

            return ["I", item[1], { "_packed": rows_T }]

        return None

class TransactionIO(Source['Result'], Sink[Any]):

    def __init__(self, transaction_log: Optional[str] = None) -> None:

        if not transaction_log or not Path(transaction_log).exists():
            version = None
        else:
            version = JsonDecode().filter(next(DiskSource(transaction_log).read()))[1]

        if version == 3:
            self._transactionIO = TransactionIO_V3(transaction_log)

        elif version == 4:
            self._transactionIO = TransactionIO_V4(transaction_log)

        elif version is None:
            self._transactionIO = TransactionIO_V4(transaction_log)

        else:
            raise CobaException("We were unable to determine the appropriate Transaction reader for the file.")

    def write(self, transaction: Any) -> None:
        self._transactionIO.write(transaction)

    def read(self) -> 'Result':
        return self._transactionIO.read()

class Points(NamedTuple):
    X    : Sequence[Any]
    Y    : Sequence[float]
    XE   : Sequence[float]  = None
    YE   : Sequence[float]  = None
    color: Union[str,int]   = None
    alpha: float            = 1
    label: Optional[str]    = None
    style: Literal['-','.'] = "-"

class Plotter:
    @abstractmethod
    def plot(self,
        ax,
        lines: Sequence[Points],
        title: str,
        xlabel: str,
        ylabel: str,
        xlim: Optional[Tuple[Number,Number]],
        ylim: Optional[Tuple[Number,Number]],
        xticks: bool,
        yticks: bool,
        filename: Optional[str]) -> None:
        pass

class MatplotlibPlotter(Plotter):

    def plot(self,
        ax,
        lines: Sequence[Points],
        title: str,
        xlabel: str,
        ylabel: str,
        xlim: Optional[Tuple[Number,Number]],
        ylim: Optional[Tuple[Number,Number]],
        xticks: bool,
        yticks: bool,
        xrotation: Optional[float],
        yrotation: Optional[float],
        filename: Optional[str]
    ) -> None:

        PackageChecker.matplotlib('Result.plot_learners')
        import matplotlib.pyplot as plt #type: ignore

        color_cycle = plt.rcParams['axes.prop_cycle'].by_key()['color']

        ax: plt.Axes
        show = ax is None

        bad_xlim  = xlim and xlim[0] is not None and xlim[1] is not None and xlim[0] >= xlim[1]
        bad_ylim  = ylim and ylim[0] is not None and ylim[1] is not None and ylim[0] >= ylim[1]
        bad_lines = not lines

        if bad_xlim or bad_ylim or bad_lines:
            if bad_xlim:
                CobaContext.logger.log("The xlim end is less than the xlim start. Plotting is impossible.")
            if bad_ylim:
                CobaContext.logger.log("The ylim end is less than the ylim start. Plotting is impossible.")
            if bad_lines:
                CobaContext.logger.log(f"No data was found for plotting in the given results.")
        else:
            ax = ax or plt.figure(num="coba").add_subplot(111) #type: ignore

            in_lim = lambda v,lim: lim is None or ((lim[0] or -float('inf')) <= v and v <= (lim[1] or float('inf')))

            any_label = False
            num_coalesce = lambda x1,x2: x1 if isinstance(x1,(int,float)) else x2

            for X, Y, XE, YE, c, a, l, fmt in lines:

                if l: any_label = True

                # we remove values outside of the given lims because matplotlib won't correctly scale otherwise
                if not YE:
                    XY = [(x,y) for i,x,y in zip(count(),X,Y) if in_lim(num_coalesce(x,i),xlim) and in_lim(num_coalesce(y,i),ylim)]
                    X,Y = map(list,zip(*XY)) if XY else ([],[])
                else:
                    XYE = [(x,y,e) for i,x,y,e in zip(count(),X,Y,YE) if in_lim(num_coalesce(x,i),xlim) and in_lim(num_coalesce(y,i),ylim)]
                    X,Y,YE = map(list,zip(*XYE)) if XYE else ([],[],[])

                if isinstance(c,int): c = color_cycle[c]

                not_err_bar = lambda E: not E or all([e is None for e in E])

                if X and Y:
                    if all(map(not_err_bar,[XE,YE])):
                        ax.plot(X, Y, fmt,  color=c, alpha=a, label=l)
                    else:
                        XE = None if not_err_bar(XE) else list(zip(*XE)) if isinstance(XE[0],tuple) else XE
                        YE = None if not_err_bar(YE) else list(zip(*YE)) if isinstance(YE[0],tuple) else YE
                        error_every = max(int(len(X)*0.05),1) if fmt == "-" else 1
                        ax.errorbar(X, Y, YE, XE, fmt, elinewidth=0.5, errorevery=error_every, color=c, alpha=a, label=l)

            if xrotation is not None:
                plt.xticks(rotation=xrotation)

            if yrotation is not None:
                plt.yticks(rotation=yrotation)

            padding = .05
            ax.margins(0)
            ax.set_xticks([min(ax.get_xlim()[1], max(ax.get_xlim()[0],x)) for x in ax.get_xticks()])
            ax.margins(padding)

            if xlim:
                x_pad = padding*(xlim[1]-xlim[0])
                ax.set_xlim(xlim[0]-x_pad, xlim[1]+x_pad)

            if ylim:
                y_pad = padding*(ylim[1]-ylim[0])
                ax.set_ylim(ylim[0]-y_pad, ylim[1]+y_pad)

            ax.set_title(title, loc='left', pad=15)
            ax.set_ylabel(ylabel)
            ax.set_xlabel(xlabel)

            if ax.get_legend() is None: #pragma: no cover
                scale = 0.65
                box1 = ax.get_position()
                ax.set_position([box1.x0, box1.y0 + box1.height * (1-scale), box1.width, box1.height * scale])
            else:
                ax.get_legend().remove()

            if any_label:
                ax.legend(*ax.get_legend_handles_labels(), loc='upper left', bbox_to_anchor=(-.01, -.25), ncol=1, fontsize='medium') #type: ignore

            if not xticks:
                plt.xticks([])

            if not yticks:
                plt.yticks([])

            if show:
                plt.show()
                plt.close()

            if filename:
                plt.tight_layout()
                plt.savefig(filename, dpi=300)
                plt.close()

class FilterPlottingData:

    def filter(self, rows:Sequence[Dict[str,Any]], x:Sequence[str], y:str, learner_ids:Sequence[int] = None) -> Sequence[Dict[str,Any]]:

        if not rows: raise CobaException("This result doesn't contain any evaluation data to plot.")

        rows = [row for row in rows if not learner_ids or row["learner_id"] in learner_ids]

        learner_count   = len(set([row['learner_id'] for row in rows]))
        environ_counts  = collections.Counter([row['environment_id'] for row in rows])
        environ_lengths = set([len(row[y]) for row in rows])
        min_env_length  = min(environ_lengths)

        if max(environ_counts.values()) != learner_count:
            raise CobaException("This result does not contain an environment which has been finished for every learner. Plotting has been stopped.")

        if min(environ_counts.values()) != learner_count:
            CobaContext.logger.log("This result contains environments not present for all learners. Environments not present for all learners have been excluded. To supress this warning in the future call <result>.filter_fin() before plotting.")

        if len(environ_lengths) > 1 and x == ['index']:
            CobaContext.logger.log("This result contains environments of different lengths. The plot only includes interactions up to the shortest environment. To supress this warning in the future call <result>.filter_fin(n_interactions) before plotting.")

        complete    = lambda row: environ_counts[row['environment_id']] == learner_count
        subselect_y = lambda ys: ys[:min_env_length] if x == ['index'] else ys
        subselect_r = lambda row: {k:(row[k] if k != y else subselect_y(row[k])) for k in ['environment_id','learner_id',y]}

        return list(map(subselect_r,filter(complete,rows)))

class SmoothPlottingData:

    def filter(self, rows:Sequence[Dict[str,Any]], y:str, span:Optional[int]) -> Sequence[Dict[str,Any]]:

        e_key = "environment_id"
        l_key = "learner_id"

        return [{e_key:row[e_key], l_key:row[l_key], y:moving_average(row[y],span)} for row in rows]

class ContrastPlottingData:

    def filter(self, rows:Sequence[Dict[str,Any]], y:str, mode:str, learner_id1:int) -> Sequence[Dict[str,Any]]:

        sort_key  = lambda row: (row['environment_id'], 0 if row['learner_id']==learner_id1 else 1)
        group_key = lambda row: row['environment_id']

        if mode == "prob":
            contraster = lambda Y1,Y2: list(map(int,map(gt,Y1,Y2)))
        elif mode == "diff":
            contraster = lambda Y1,Y2: list(map(sub,Y1,Y2))
        else:
            contraster = lambda Y1,Y2: list(zip(Y1,Y2))

        contrast = lambda env_id, pair: {'environment_id': env_id, y: contraster(pair[0][y],pair[1][y]) }

        return [ contrast(env_id,list(pair)) for env_id, pair in groupby(sorted(rows,key=sort_key),key=group_key) ]

class TransformToXYE:

    def filter(self,
        rows: Sequence[Dict[str,Any]],
        envs: Dict[int,Dict[str,Any]],
        x:Sequence[str],
        y:str,
        err:Union[str,None,PointAndInterval]) -> Sequence[Tuple[Any,float,Union[None,float,Tuple[float,float]]]]:

        if err == 'sd':
            avg = Mean()
            std = StandardDeviation()
            YE = lambda z: (avg.calculate(z), std.calculate(z))
        elif err == 'se':
            YE = StandardErrorOfMean().calculate
        elif err == "bs":
            YE = BootstrapConfidenceInterval(.95, Mean().calculate).calculate
        elif err == "bi":
            YE = BinomialConfidenceInterval('wilson').calculate
        elif isinstance(err,PointAndInterval):
            YE = err.calculate
        else:
            avg = Mean()
            YE = lambda z: (avg.calculate(z) if len(z) > 1 else z[0], None)

        is_scatter = isinstance(rows[0][y][0],(list,tuple))

        if x == ['index']:
            Z = list(zip(*[row[y] for row in rows]))
            X = list(range(1,len(Z)+1))

            if not is_scatter:
                points = [(x,None)+YE(z) for x,z in zip(X,Z) ]
            else:
                points = [YE(z1)+YE(z2) for _,z in zip(X,Z) for z1,z2 in [tuple(zip(*z))]]
        else:
            XZ = collections.defaultdict(list)
            make_x = lambda env: env[x[0]] if len(x) == 1 else tuple(env[k] for k in x)
            for row in rows:
                XZ[str(make_x(envs[row["environment_id"]]))].append(row[y][-1])

            if not is_scatter:
                points = [(x,None)+YE(z) for x,z in XZ.items()]
            else:
                points = [YE(z1)+YE(z2) for _,z in XZ.items() for z1,z2 in [tuple(zip(*z))]]

        return [ (x,y,xe,ye) for x,xe,y,ye in points ]

class Result:
    """A class representing the result of an Experiment."""

    @staticmethod
    def from_file(filename: str) -> 'Result':
        """Create a Result from a transaction file."""

        if not Path(filename).exists():
            raise CobaException("We were unable to find the given Result file.")

        return TransactionIO(filename).read()

    def __init__(self,
        env_rows: Dict[int           ,Dict[str,Any]] = {},
        lrn_rows: Dict[int           ,Dict[str,Any]] = {},
        int_rows: Dict[Tuple[int,int],Dict[str,Any]] = {},
        exp_dict: Dict[str, Any]                     = {}) -> None:
        """Instantiate a Result class.

        This constructor should never be called directly. Instead a Result file should be created
        from an Experiment and the result file should be loaded via Result.from_file(filename).
        """

        env_flat = [ { "environment_id":k,                       **v } for k,v in env_rows.items() ]
        lrn_flat = [ {                        "learner_id" :k,   **v } for k,v in lrn_rows.items() ]
        int_flat = [ { "environment_id":k[0], "learner_id":k[1], **v } for k,v in int_rows.items() ]

        self.experiment = exp_dict

        self._environments = Table            ("Environments", ['environment_id'              ], env_flat, ["source"])
        self._learners     = Table            ("Learners"    , ['learner_id'                  ], lrn_flat, ["family","shuffle","take"])
        self._interactions = InteractionsTable("Interactions", ['environment_id', 'learner_id'], int_flat, ["index","reward"])

        self._plotter = MatplotlibPlotter()

    @property
    def learners(self) -> Table:
        """The collection of learners used in the Experiment.

        The primary key of this table is learner_id.
        """
        return self._learners

    @property
    def environments(self) -> Table:
        """The collection of environments used in the Experiment.

        The primary key of this table is environment_id.
        """
        return self._environments

    @property
    def interactions(self) -> InteractionsTable:
        """The collection of interactions that learners chose actions for in the Experiment.

        The primary key of this Table is (index, environment_id, learner_id). It should be noted that the InteractionTable
        has an additional property, to_progressive_pandas() which calculates the expanding moving average or exponential
        moving average for interactions ordered by index and grouped by environment_id and learner_id.
        """
        return self._interactions

    def set_plotter(self, plotter: Plotter) -> None:
        """Manually set the underlying plotting tool. By default matplotlib is used though this can be changed."""
        self._plotter = plotter

    def copy(self) -> 'Result':
        """Create a copy of Result."""

        result = Result()

        result._environments = copy(self._environments)
        result._learners     = copy(self._learners)
        result._interactions = copy(self._interactions)

        return result

    def filter_fin(self, n_interactions: int = None) -> 'Result':
        """Filter the result to only contain data about environments with all learners and interactions.

        Args:
            n_interactions: The number of interactions at which an environment is considered complete.
        """

        def has_all(env_id):
            return all((env_id, lrn_id) in self.interactions for lrn_id in self.learners.keys)

        def has_min(env_id):
            if not n_interactions:
                return True
            else:
                return all(len(self.interactions[(env_id, lrn_id)].get('index',[])) >= n_interactions for lrn_id in self.learners.keys)

        complete_ids = [env_id for env_id in self.environments.keys if has_all(env_id) and has_min(env_id)]

        new_result               = copy(self)
        new_result._environments = self.environments.filter(environment_id=complete_ids)
        new_result._interactions = self.interactions.filter(environment_id=complete_ids)

        if n_interactions:
            new_inter = copy(new_result.interactions)
            new_inter._rows_pack = { rk: { pk: pv[0:n_interactions] for pk, pv in rv.items() } for rk,rv in new_inter._rows_pack.items() }
            new_result._interactions = new_inter

        if len(new_result.environments) == 0:
            CobaContext.logger.log(f"There was no environment which was finished for every learner.")

        return new_result

    def filter_env(self, pred:Callable[[Dict[str,Any]],bool] = None, **kwargs: Any) -> 'Result':
        """Filter the result to only contain data about specific environments.

        Args:
            pred: A predicate that returns true for learner dictionaries that should be kept.
            **kwargs: key-value pairs to filter on. To see filtering options see Table.filter.
        """

        new_result = copy(self)
        new_result._environments = new_result.environments.filter(pred, **kwargs)
        new_result._interactions = new_result.interactions.filter(lambda row: row['environment_id'] in new_result.environments)

        if len(new_result.environments) == 0:
            CobaContext.logger.log(f"No environments matched the given filter.")

        return new_result

    def filter_lrn(self, pred:Callable[[Dict[str,Any]],bool] = None, **kwargs: Any) -> 'Result':
        """Filter the result to only contain data about specific learners.

        Args:
            pred: A predicate that returns true for learner dictionaries that should be kept.
            **kwargs: key-value pairs to filter on. To see filtering options see Table.filter.
        """

        new_result = copy(self)
        new_result._learners     = new_result.learners.filter(pred, **kwargs)
        new_result._interactions = new_result.interactions.filter(learner_id=new_result.learners)

        if len(new_result.learners) == 0:
            CobaContext.logger.log(f"No learners matched the given filter.")

        return new_result

    def plot_contrast(self,
        learner_id1: int,
        learner_id2: int,
        x          : Union[str, Sequence[str]] = "environment_id",
        y          : str = "reward",
        mode       : Literal["diff","prob",'scat'] = "diff",
        span       : int = None,
        err        : Union[Literal['se','sd','bs'], None, PointAndInterval] = None,
        labels     : Sequence[str] = None,
        colors     : Sequence[str] = None,
        xlim       : Tuple[Number,Number] = None,
        ylim       : Tuple[Number,Number] = None,
        xticks     : bool = True,
        yticks     : bool = True,
        filename   : str = None,
        ax = None) -> None:

        x = [x] if isinstance(x,str) else list(x)
        self._validate_parameters(x)

        rows = FilterPlottingData().filter(list(self.interactions), x, y, [learner_id1, learner_id2])
        rows = SmoothPlottingData().filter(rows, y, span)
        rows = ContrastPlottingData().filter(rows, y, mode, learner_id1)

        XYE = TransformToXYE().filter(rows, self.environments, x, y, err)

        if x != ['index']:
            XYE = sorted(XYE, key=lambda xye: xye[1])

        bound = .5 if mode == "prob" else 0

        win,tie,loss = [],[],[]
        for _x,_y,_xe,_ye in XYE:
            xl,xu = (0,0) if _xe is None else (_xe,_xe) if isinstance(_xe,Number) else _xe
            yl,yu = (0,0) if _ye is None else (_ye,_ye) if isinstance(_ye,Number) else _ye
            if mode != 'scat':
                (win if _y-yl > bound else loss if _y+yu < bound else tie).append((_x,_y,_xe,_ye))
            else:
                (win if _x-xl>_y+yu else loss if _y-yl>_x+xu else tie).append((_x,_y,_xe,_ye))

        colors = (colors or []) + [0,1,2]
        labels = (labels or []) + [self.learners[learner_id1]['full_name'], self.learners[learner_id2]['full_name']]

        fmt = "-" if x == ['index'] else "."

        plots = []
        if loss: plots.append(Points(*zip(*loss), 2, 1, labels[1] + " " + f"({len(loss)})", fmt))
        if tie : plots.append(Points(*zip(*tie) , 1, 1, 'Tie'     + " " + f"({len(tie )})", fmt))
        if win : plots.append(Points(*zip(*win) , 0, 1, labels[0] + " " + f"({len(win )})", fmt))

        if mode != 'scat':
            leftmost_x  = (loss+tie+win)[ 0][0]
            rightmost_x = (loss+tie+win)[-1][0]
            plots.append(Points((leftmost_x,rightmost_x),(bound,bound), None, None , "#888", 1, None, '-'))
        else:
            m    = max([p[0] for p in (loss+tie+win)]+[p[1] for p in (loss+tie+win)]+[1])
            plots.append(Points((0,m),(0,m), None, None , "#888", 1, None, '-'))

        xrotation = 90 if x != ['index'] and len(XYE)>5 else 0
        yrotation = 0

        if mode != "scat":
            xlabel = "Interactions" if x==['index'] else x[0] if len(x) == 1 else x
            ylabel = f"{labels[0]} - {labels[1]}" if mode=="diff" else f"P({labels[0]} > {labels[1]})"
        else:
            xlabel = y
            ylabel = y

        title  = f"{ylabel} ({len(rows) if x==['index'] else len(XYE)} Environments)"

        self._plotter.plot(ax, plots, title, xlabel, ylabel, xlim, ylim, xticks, yticks, xrotation, yrotation, filename)

    def plot_learners(self,
        x       : Union[str,Sequence[str]] = "index",
        y       : str = "reward",
        span    : int = None,
        err     : Union[Literal['se','sd','bs'], None, PointAndInterval] = None,
        labels  : Sequence[str] = None,
        colors  : Sequence[str] = None,
        xlim    : Tuple[Number,Number] = None,
        ylim    : Tuple[Number,Number] = None,
        xticks  : bool = True,
        yticks  : bool = True,
        filename: str = None,
        ax = None) -> None:
        """Plot the performance of multiple learners on multiple environments. It gives a sense of the expected
            performance for different learners across independent environments. This plot is valuable in gaining
            insight into how various learners perform in comparison to one another.

        Args:
            x: The value to plot on the x-axis. This can either be index or environment columns to group by.
            y: The value to plot on the y-axis.
            span: The number of y values to smooth together when reporting y. If this is None then the average of all y
                values up to current is shown otherwise a moving average with window size of span (the window will be
                smaller than span initially).
            err: This determines what kind of error bars to plot (if any). If `None` then no bars are plotted, if 'se'
                the standard error is shown, and if 'sd' the standard deviation is shown.
            xlim: Define the x-axis limits to plot. If `None` the x-axis limits will be inferred.
            ylim: Define the y-axis limits to plot. If `None` the y-axis limits will be inferred.
            labels: The legend labels to use in the plot. These should be in order of the actual legend labels.
            filename: Provide a filename to write plot image to disk.
            ax: Provide an optional axes that the plot will be drawn to. If not provided a new figure/axes is created.
        """

        x = [x] if isinstance(x,str) else list(x)
        self._validate_parameters(x)

        rows = FilterPlottingData().filter(list(self.interactions), x, y)
        rows = SmoothPlottingData().filter(rows, y, span)

        env_rows = self.environments
        n_envs   = len(set([row['environment_id'] for row in rows]))
        get_key  = lambda row: row['learner_id']
        lines    = []

        style = "-" if x == ['index'] else "."

        for i, (lrn_id, lrn_rows) in enumerate(groupby(sorted(rows, key=get_key),key=get_key)):
            lrn_rows = list(lrn_rows)
            XYE      = TransformToXYE().filter(lrn_rows, env_rows, x, y, err)
            color    = colors[i] if colors and i < len(colors) else i
            label    = labels[i] if labels and i < len(labels) else self.learners[lrn_id]['full_name']
            lines.append(Points(*zip(*XYE), color, 1, label, style))

        lines  = sorted(lines, key=lambda line: -line[1][-1])
        xlabel = "Interactions" if x==['index'] else x[0] if len(x) == 1 else x
        ylabel = y.capitalize().replace("_pct"," Percent")

        title = ("Instantaneous" if span == 1 else f"Span {span}" if span else "Progressive") + f" {ylabel}"
        title = title + f" ({len(lrn_rows) if x==['index'] else len(XYE)} Environments)"

        if x != ['index']:
            title = "Final " + title

        self._plotter.plot(ax, lines, title, xlabel, ylabel, xlim, ylim, xticks, yticks, 0, 0, filename)

    def __str__(self) -> str:
        return str({"Learners": len(self._learners), "Environments": len(self._environments), "Interactions": len(self._interactions) })

    def _ipython_display_(self):
        #pretty print in jupyter notebook (https://ipython.readthedocs.io/en/stable/config/integrating.html)
        print(str(self))

    def _validate_parameters(self, x:Sequence[str]):
        if 'index' in x and len(x) > 1:
            raise CobaException('The x-axis cannot contain both interaction index and environment features. Please choose one or the other.')

max_stars_repo_name: mrucker/banditbenchmark

max_issues_repo_name: mrucker/banditbenchmark

max_forks_repo_name: mrucker/banditbenchmark


text:
# coding: utf-8

# In[2]:

import numpy as np
import pandas as pd
from sklearn.neighbors import KNeighborsClassifier
from matplotlib import pyplot as plt

df_train = pd.read_csv("../data/mnist_train.csv").sample(n=2000)
df_train['intercept'] = 1
trainingData = df_train.drop("label", axis = 1).values
trainingResults = df_train["label"].values
df_test = pd.read_csv("../data/mnist_test.csv").sample(n=500)
df_test['intercept'] = 1
testData = df_test.drop("label", axis=1).values
testResults = df_test["label"].values

errors = []
for i in range(len(K_VALS)):
    k = K_VALS[i]
    # build the validation set
    start_index = i * len(trainingData)//len(K_VALS)
    end_index = len(trainingData)//len(K_VALS) * (i + 1)

    validation_data = trainingData[start_index:end_index]
    validation_classifications = trainingResults[start_index:end_index]

    # build the model
    model = np.concatenate((trainingData[:start_index], trainingData[end_index:]), axis=0)
    model_classifications = np.concatenate((trainingResults[:start_index], trainingResults[end_index:]), axis=0)

    classifier = KNeighborsClassifier(n_neighbors=k, weights='distance')
    classifier.fit(model, model_classifications)
    score = classifier.score(validation_data, validation_classifications)
    errors.append(1 - score)

plt.plot(K_VALS, errors)
plt.title("K vs. Classification Error")
plt.xlabel("k value")
plt.xscale('log')
plt.ylabel("classification error")
plt.savefig('k-nn-libraryimpl.png')
plt.show()

"""We saw that a k of 1 actually led to our lowest classification error. Thus, we'll use this k value and
run nearest neighbor regression on the test set using the SciKit implementation. This will give us
a baseline test error."""

# Get test error of 1-NN classification
df_train = pd.read_csv("../data/mnist_train.csv").sample(frac=0.25)
df_train['intercept'] = 1
trainingData = df_train.drop("label", axis = 1).values
trainingResults = df_train["label"].values
df_test = pd.read_csv("../data/mnist_test.csv")
df_test['intercept'] = 1
testData = df_test.drop("label", axis=1).values
testResults = df_test["label"].values

classifier = KNeighborsClassifier(n_neighbors=1, weights='distance')
classifier.fit(trainingData, trainingResults)
print("The error rate of the nearest neighbor regression is", 1 - classifier.score(testData, testResults))

ext: py

max_stars_repo_name: dericp/digit-recognition

max_issues_repo_name: dericp/digit-recognition

max_forks_repo_name: dericp/digit-recognition

reason: import numpy


text:
# SPDX-License-Identifier: Apache-2.0

import os
import sys
import unittest
import keras2onnx
import keras_contrib
import numpy as np
from keras2onnx import set_converter
from keras2onnx.proto import keras
from os.path import dirname, abspath
sys.path.insert(0, os.path.join(dirname(abspath(__file__)), '../../tests/'))
from test_utils import run_onnx_runtime, convert_InstanceNormalizationLayer

Activation = keras.layers.Activation
BatchNormalization = keras.layers.BatchNormalization
Concatenate = keras.layers.Concatenate
Conv2D = keras.layers.Conv2D
Dense = keras.layers.Dense
Dropout = keras.layers.Dropout
Embedding = keras.layers.Embedding
Flatten = keras.layers.Flatten
InstanceNormalization = keras_contrib.layers.InstanceNormalization
Input = keras.layers.Input
LeakyReLU = keras.layers.LeakyReLU
multiply = keras.layers.multiply
Reshape = keras.layers.Reshape
UpSampling2D = keras.layers.UpSampling2D
ZeroPadding2D = keras.layers.ZeroPadding2D

Sequential = keras.models.Sequential
Model = keras.models.Model

# From https://github.com/eriklindernoren/Keras-GAN/blob/master/discogan/discogan.py
class DiscoGAN():
    def __init__(self):
        # Input shape
        self.img_rows = 128
        self.img_cols = 128
        self.channels = 3
        self.img_shape = (self.img_rows, self.img_cols, self.channels)

        # Calculate output shape of D (PatchGAN)
        patch = int(self.img_rows / 2**4)
        self.disc_patch = (patch, patch, 1)

        # Number of filters in the first layer of G and D
        self.gf = 64
        self.df = 64

        # Build and compile the discriminators
        self.d_A = self.build_discriminator()
        self.d_B = self.build_discriminator()

        #-------------------------
        # Construct Computational
        #   Graph of Generators
        #-------------------------

        # Build the generators
        self.g_AB = self.build_generator()
        self.g_BA = self.build_generator()

        # Input images from both domains
        img_A = Input(shape=self.img_shape)
        img_B = Input(shape=self.img_shape)

        # Translate images to the other domain
        fake_B = self.g_AB(img_A)
        fake_A = self.g_BA(img_B)
        # Translate images back to original domain
        reconstr_A = self.g_BA(fake_B)
        reconstr_B = self.g_AB(fake_A)

        # For the combined model we will only train the generators
        self.d_A.trainable = False
        self.d_B.trainable = False

        # Discriminators determines validity of translated images
        valid_A = self.d_A(fake_A)
        valid_B = self.d_B(fake_B)

        # Objectives
        # + Adversarial: Fool domain discriminators
        # + Translation: Minimize MAE between e.g. fake B and true B
        # + Cycle-consistency: Minimize MAE between reconstructed images and original
        self.combined = Model(inputs=[img_A, img_B],
                              outputs=[ valid_A, valid_B,
                                        fake_B, fake_A,
                                        reconstr_A, reconstr_B ])

    def build_generator(self):
        """U-Net Generator"""

        def conv2d(layer_input, filters, f_size=4, normalize=True):
            """Layers used during downsampling"""
            d = Conv2D(filters, kernel_size=f_size, strides=2, padding='same')(layer_input)
            d = LeakyReLU(alpha=0.2)(d)
            if normalize:
                d = InstanceNormalization()(d)
            return d

        def deconv2d(layer_input, skip_input, filters, f_size=4, dropout_rate=0):
            """Layers used during upsampling"""
            u = UpSampling2D(size=2)(layer_input)
            u = Conv2D(filters, kernel_size=f_size, strides=1, padding='same', activation='relu')(u)
            if dropout_rate:
                u = Dropout(dropout_rate)(u)
            u = InstanceNormalization()(u)
            u = Concatenate()([u, skip_input])
            return u

        # Image input
        d0 = Input(shape=self.img_shape)

        # Downsampling
        d1 = conv2d(d0, self.gf, normalize=False)
        d2 = conv2d(d1, self.gf*2)
        d3 = conv2d(d2, self.gf*4)
        d4 = conv2d(d3, self.gf*8)
        d5 = conv2d(d4, self.gf*8)
        d6 = conv2d(d5, self.gf*8)
        d7 = conv2d(d6, self.gf*8)

        # Upsampling
        u1 = deconv2d(d7, d6, self.gf*8)
        u2 = deconv2d(u1, d5, self.gf*8)
        u3 = deconv2d(u2, d4, self.gf*8)
        u4 = deconv2d(u3, d3, self.gf*4)
        u5 = deconv2d(u4, d2, self.gf*2)
        u6 = deconv2d(u5, d1, self.gf)

        u7 = UpSampling2D(size=2)(u6)
        output_img = Conv2D(self.channels, kernel_size=4, strides=1,
                            padding='same', activation='tanh')(u7)

        return Model(d0, output_img)

    def build_discriminator(self):

        def d_layer(layer_input, filters, f_size=4, normalization=True):
            """Discriminator layer"""
            d = Conv2D(filters, kernel_size=f_size, strides=2, padding='same')(layer_input)
            d = LeakyReLU(alpha=0.2)(d)
            if normalization:
                d = InstanceNormalization()(d)
            return d

        img = Input(shape=self.img_shape)

        d1 = d_layer(img, self.df, normalization=False)
        d2 = d_layer(d1, self.df*2)
        d3 = d_layer(d2, self.df*4)
        d4 = d_layer(d3, self.df*8)

        validity = Conv2D(1, kernel_size=4, strides=1, padding='same')(d4)

        return Model(img, validity)

set_converter(keras_contrib.layers.InstanceNormalization, convert_InstanceNormalizationLayer)

class TestDiscoGAN(unittest.TestCase):

    def setUp(self):
        self.model_files = []

    def tearDown(self):
        for fl in self.model_files:
            os.remove(fl)

    def test_DiscoGAN(self):
        keras_model = DiscoGAN().combined
        batch = 5
        x = np.random.rand(batch, 128, 128, 3).astype(np.float32)
        y = np.random.rand(batch, 128, 128, 3).astype(np.float32)
        expected = keras_model.predict([x, y])
        onnx_model = keras2onnx.convert_keras(keras_model, keras_model.name)
        self.assertTrue(run_onnx_runtime(onnx_model.graph.name, onnx_model, {keras_model.input_names[0]: x, keras_model.input_names[1]: y}, expected, self.model_files, rtol=1.e-2, atol=1.e-2))

if __name__ == "__main__":
    unittest.main()

max_stars_repo_name: TomWildenhain-Microsoft/keras-onnx

max_issues_repo_name: TomWildenhain-Microsoft/keras-onnx

max_forks_repo_name: TomWildenhain-Microsoft/keras-onnx


text:
from numpy.core.numeric import outer
import tensorflow as tf
import os
import numpy as np
import kgcnn
from kgcnn.utils.adj import coordinates_to_distancematrix,define_adjacency_from_distance, distance_to_gauss_basis, get_angle_indices, sort_edge_indices
from kgcnn.utils.data import save_pickle_file,load_pickle_file, ragged_tensor_from_nested_numpy
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import StandardScaler
from ase.io import read
from utils import get_distance

def generate_input_data(input_dir, output_dir):
    """
    Generates nodes, edges, coordinates and distances lists from .pdb files
    Args:
        input_dir (str): direction for the .pdb files
        output_dir (str): output direction
    """

    search_path = input_dir
    filelist = [f for f in os.listdir(search_path) if os.path.isfile(os.path.join(search_path, f))]
    base_list = [x.split(".")[0].split("_")[0] + "_" + x.split(".")[0].split("_")[1] for x in filelist]
    base_set = sorted(set(base_list))

    all_nodes = []
    all_edges = []
    all_labels = []
    all_atoms = []
    for x in base_set:
        mol1 = read(search_path+x+"_1.pdb")
        an1 = mol1.get_atomic_numbers()
        n_1 = mol1.get_global_number_of_atoms()
        pos1 = mol1.positions
        mol2 = read(search_path+x+"_2.pdb")
        n_2 = mol2.get_global_number_of_atoms()
        an2 = mol2.get_atomic_numbers()
        pos2 = mol2.positions
        nodes1 = np.concatenate((np.array([0]),an1),axis=0)
        edges1 = np.concatenate((np.array([pos2[0]]),pos1),axis=0)
        nodes2 = np.concatenate((np.array([0]),an2),axis=0)
        edges2 = np.concatenate((np.array([pos1[0]]),pos2),axis=0)
        arr = np.load(search_path+x+".npz", allow_pickle=True)
        arr = np.expand_dims(arr[arr.files[0]], axis=0)[0]
        if 'e_max' in arr and 'e_00' in arr and 'e_10' in arr:
            lab1 = arr['e_max'] - arr['e_00']
            lab2 = arr['e_max'] - arr['e_10']
            all_labels.append(lab1)
            all_labels.append(lab2)
            all_nodes.append(nodes1)
            all_nodes.append(nodes2)
            all_edges.append(edges1)
            all_edges.append(edges2)
            all_atoms.append(n_1)
            all_atoms.append(n_2)

    dist = [coordinates_to_distancematrix(x) for x in all_edges]
    adj = [define_adjacency_from_distance(x,max_distance=8, max_neighbours=25)[0] for x in dist]
    edge_indices = [define_adjacency_from_distance(x,max_distance=8, max_neighbours=25)[1] for x in dist]
    edge_dist = [distance_to_gauss_basis(dist[i][adj[i]], bins=40, distance=8, sigma=0.4) for i in range(len(adj))]

    save_pickle_file(all_edges, output_dir + "cr_coord.pickle")
    save_pickle_file(all_nodes, output_dir + "cr_nodes.pickle",)
    save_pickle_file(edge_dist, output_dir + "cr_edges.pickle",)
    save_pickle_file(edge_indices, output_dir + "cr_indices.pickle")
    np.save(output_dir + "cr_atoms.npy", all_atoms)
    np.save(output_dir + "cr_labels.npy",all_labels)

def generate_input_data_cutoff(input_dir, output_dir, cutoff_dist):
    """
    Generates nodes, edges, coordinates and distances lists from .pdb files with a cutoff distance
    Args:
        input_dir (str): direction for the .pdb files
        output_dir (str): output direction
        cutoff_dist (int): distance for cutoff (in Angsroem)
    """
    search_path = input_dir
    filelist = [f for f in os.listdir(search_path) if os.path.isfile(os.path.join(search_path, f))]
    base_list = [x.split(".")[0].split("_")[0] + "_" + x.split(".")[0].split("_")[1] for x in filelist]
    base_set = sorted(set(base_list))

    threshold = cutoff_dist

    all_nodes = []
    all_edges = []
    all_labels = []
    all_atoms = []
    all_nodes_high = []
    all_edges_high = []
    all_labels_high = []
    all_atoms_high = []
    for x in base_set:
        mol1 = read(search_path+x+"_1.pdb")
        an1 = mol1.get_atomic_numbers()
        pos1 = mol1.positions
        n_1 = mol1.get_global_number_of_atoms()
        mol2 = read(search_path+x+"_2.pdb")
        an2 = mol2.get_atomic_numbers()
        pos2 = mol2.positions
        n_2 = mol2.get_global_number_of_atoms()
        nodes1 = np.concatenate((np.array([0]),an1),axis=0)
        edges1 = np.concatenate((np.array([pos2[0]]),pos1),axis=0)
        nodes2 = np.concatenate((np.array([0]),an2),axis=0)
        edges2 = np.concatenate((np.array([pos1[0]]),pos2),axis=0)
        dist = get_distance(edges1[0,:], edges2[0,:])

        arr = np.load(search_path+x+".npz", allow_pickle=True)
        arr = np.expand_dims(arr[arr.files[0]], axis=0)[0]
        if 'e_max' in arr and 'e_00' in arr and 'e_10' in arr:
            lab1 = arr['e_max'] - arr['e_00']
            lab2 = arr['e_max'] - arr['e_10']
            if dist < threshold:
                all_labels.append(lab1)
                all_labels.append(lab2)
                all_nodes.append(nodes1)
                all_nodes.append(nodes2)
                all_edges.append(edges1)
                all_edges.append(edges2)
                all_atoms.append(n_1)
                all_atoms.append(n_2)
            else:
                all_nodes_high.append(nodes1)
                all_nodes_high.append(nodes2)
                all_labels_high.append(lab1)
                all_labels_high.append(lab2)
                all_edges_high.append(edges1)
                all_edges_high.append(edges2)
                all_atoms_high.append(n_1)
                all_atoms_high.append(n_2)

    dist = [coordinates_to_distancematrix(x) for x in all_edges]
    adj = [define_adjacency_from_distance(x,max_distance=3, max_neighbours=25)[0] for x in dist]
    edge_indices = [define_adjacency_from_distance(x,max_distance=3, max_neighbours=25)[1] for x in dist]
    edge_dist = [distance_to_gauss_basis(dist[i][adj[i]], bins=40, distance=8, sigma=0.4) for i in range(len(adj))]

    dist = [coordinates_to_distancematrix(x) for x in all_edges_high]
    adj = [define_adjacency_from_distance(x,max_distance=3, max_neighbours=25)[0] for x in dist]
    edge_indices_high = [define_adjacency_from_distance(x,max_distance=3, max_neighbours=25)[1] for x in dist]
    edge_dist_high = [distance_to_gauss_basis(dist[i][adj[i]], bins=40, distance=8, sigma=0.4) for i in range(len(adj))]

    save_pickle_file(all_edges, output_dir+str(cutoff_dist)+"_low_cr_coord.pickle")
    save_pickle_file(all_nodes, output_dir+str(cutoff_dist)+"_low_cr_nodes.pickle",)
    save_pickle_file(edge_dist, output_dir+str(cutoff_dist)+"_low_cr_edges.pickle",)
    save_pickle_file(edge_indices, output_dir+str(cutoff_dist)+"_low_cr_indices.pickle")
    np.save(output_dir+str(cutoff_dist)+"_low_cr_labels.npy",all_labels)
    np.save(output_dir+str(cutoff_dist)+"_low_cr_atoms.npy",all_atoms)

    save_pickle_file(all_edges_high, output_dir+str(cutoff_dist)+"_high_cr_coord.pickle")
    save_pickle_file(all_nodes_high, output_dir+str(cutoff_dist)+"_high_cr_nodes.pickle",)
    save_pickle_file(edge_dist_high, output_dir+str(cutoff_dist)+"_high_cr_edges.pickle",)
    save_pickle_file(edge_indices_high, output_dir+str(cutoff_dist)+"_high_cr_indices.pickle")
    np.save(output_dir+str(cutoff_dist)+"_high_cr_labels.npy",all_labels_high)
    np.save(output_dir+str(cutoff_dist)+"_high_cr_atoms.npy",all_atoms_high)

def painn_dataset(input_dir, split_size):
    """Creates PaiNN graph input
    Args:
        input_dir (str): input direction with files created using generate_input_data function
        split_size (int): test size
    Return:
        xtrain (ragged tensor): training data
        xtest (ragged tensor): test data
        xval (ragged tensor): validation data
        ytrain (ragged tensor): training labels
        ytest (ragged tensor): test labels
        yval (ragged tensor): validation labels
        scaler: scaler config for later rescaling of the data
        """

    # load Dataset
    y_data = np.expand_dims(np.load(input_dir+"cr_labels.npy"),axis=-1)
    nodes = load_pickle_file(input_dir+"cr_nodes.pickle")
    coord = load_pickle_file(input_dir+"cr_coord.pickle")

    equivariant = [np.zeros((len(x), 128, 3)) for x in nodes]

    dist = [coordinates_to_distancematrix(x) for x in coord]
    edge_indices = [define_adjacency_from_distance(x, max_distance=5, max_neighbours=25)[1] for x in dist]
    edge_indices = [x if x[0,1]==1 else sort_edge_indices(np.concatenate([np.array([[0, 1], [1,0]]), x], axis=0)) for x in edge_indices]
    del dist  # Frees some memory

    # Should check again the data function to ensure that radical and hydrogen are on top an edge between them exists.
    radical_node_index = [np.array([[0, 1]]) for _ in edge_indices]
    radical_edge_index = [np.array([[0]]) for _ in radical_node_index]

    scaler = StandardScaler(copy=True)
    labels = scaler.fit_transform(y_data)

    # Train Test split
    labels_new, labels_val, nodes_new, nodes_val, equi_new, equi_val, coord_new, coord_val, \
        edge_indices_new, edge_indices_val, radical_node_index_new, radical_node_index_val, radical_edge_index_new, radical_edge_index_val= train_test_split(labels,
        nodes, equivariant, coord, edge_indices, radical_node_index, radical_edge_index, test_size=split_size)
    del labels, nodes, equivariant, coord, edge_indices, radical_node_index , radical_edge_index,

    # Train Test split
    labels_train, labels_test, nodes_train, nodes_test, equi_train, equi_test, coord_train, coord_test, \
        edge_indices_train, edge_indices_test, radical_node_index_train, radical_node_index_test, radical_edge_index_train, \
            radical_edge_index_test= train_test_split( labels_new,
        nodes_new, equi_new, coord_new, edge_indices_new, radical_node_index_new, radical_edge_index_new, test_size=split_size)
    del labels_new, nodes_new, equi_new, coord_new, edge_indices_new

    # Convert to tf.RaggedTensor or tf.tensor
    # a copy of the data is generated by ragged_tensor_from_nested_numpy()
    nodes_train, equi_train, coord_train, edge_indices_train, radical_node_index_train, radical_edge_index_train = ragged_tensor_from_nested_numpy(
        nodes_train), ragged_tensor_from_nested_numpy(equi_train), ragged_tensor_from_nested_numpy(coord_train), ragged_tensor_from_nested_numpy(
        edge_indices_train), ragged_tensor_from_nested_numpy(radical_node_index_train), ragged_tensor_from_nested_numpy(radical_edge_index_train)

    nodes_test, equi_test, coord_test, edge_indices_test, radical_node_index_test, radical_edge_index_test = ragged_tensor_from_nested_numpy(
        nodes_test),  ragged_tensor_from_nested_numpy(equi_test), ragged_tensor_from_nested_numpy(coord_test), ragged_tensor_from_nested_numpy(
        edge_indices_test), ragged_tensor_from_nested_numpy(radical_node_index_test), ragged_tensor_from_nested_numpy(radical_edge_index_test)

    nodes_val, equi_val, coord_val, edge_indices_val, radical_node_index_val, radical_edge_index_val = ragged_tensor_from_nested_numpy(
        nodes_val),  ragged_tensor_from_nested_numpy(equi_val), ragged_tensor_from_nested_numpy(coord_val), ragged_tensor_from_nested_numpy(
        edge_indices_val), ragged_tensor_from_nested_numpy(radical_node_index_val), ragged_tensor_from_nested_numpy(radical_edge_index_val)

    # Define input and output data
    xtrain = nodes_train, equi_train, coord_train, edge_indices_train, radical_node_index_train, radical_edge_index_train
    xtest = nodes_test, equi_test, coord_test, edge_indices_test, radical_node_index_test, radical_edge_index_test
    ytrain = labels_train
    ytest = labels_test
    xval = nodes_val, equi_val, coord_val, edge_indices_val, radical_node_index_val, radical_edge_index_val
    yval = labels_val
    return(xtrain, xtest, xval, ytrain, ytest, yval, scaler)

def schnet_dataset(input_dir, split_size):
    """
    Creates Schnet graph input
    Args:
        input_dir (str): input direction with files created using generate_input_data function
        split_size (int): test size
    Return:
        xtrain (ragged tensor): training data
        xtest (ragged tensor): test data
        xval (ragged tensor): validation data
        ytrain (ragged tensor): training labels
        ytest (ragged tensor): test labels
        yval (ragged tensor): validation labels
        scaler: scaler config for later rescaling of the data
    """

    y_data = np.expand_dims(np.load(input_dir+"cr_labels.npy"),axis=-1)
    nodes = load_pickle_file(input_dir+'cr_nodes.pickle')
    edges = load_pickle_file(input_dir+"cr_edges.pickle")
    edge_indices = load_pickle_file(input_dir+"cr_indices.pickle")

    radical_node_index = [np.expand_dims(x[0],axis=0) for x in edge_indices]
    radical_edge_index = [np.array([[0]]) for _ in radical_node_index]

    scaler = StandardScaler(copy=True)
    labels = scaler.fit_transform(y_data)

    # Train Test split
    labels_train, labels_test, nodes_train, nodes_test, edges_train, edges_test, edge_indices_train, edge_indices_test, \
    radical_node_index_train, radical_node_index_test, radical_edge_index_train, radical_edge_index_test = train_test_split(
        labels, nodes, edges, edge_indices, radical_node_index, radical_edge_index, test_size=split_size)
    del labels, nodes, edges, edge_indices, radical_node_index, radical_edge_index # Free memory after split, if possible

    # Convert to tf.RaggedTensor or tf.tensor
    # a copy of the data is generated by ragged_tensor_from_nested_numpy()
    nodes_train, edges_train, edge_indices_train, radical_node_index_train, radical_edge_index_train = ragged_tensor_from_nested_numpy(
        nodes_train), ragged_tensor_from_nested_numpy(edges_train), ragged_tensor_from_nested_numpy(
        edge_indices_train), ragged_tensor_from_nested_numpy(radical_node_index_train), ragged_tensor_from_nested_numpy(radical_edge_index_train)

    nodes_test, edges_test, edge_indices_test, radical_node_index_test, radical_edge_index_test = ragged_tensor_from_nested_numpy(
        nodes_test), ragged_tensor_from_nested_numpy(edges_test), ragged_tensor_from_nested_numpy(
        edge_indices_test), ragged_tensor_from_nested_numpy(radical_node_index_test), ragged_tensor_from_nested_numpy(radical_edge_index_test)

    # Define input and output data
    xtrain = nodes_train, edges_train, edge_indices_train, radical_node_index_train, radical_edge_index_train
    xval = nodes_test, edges_test, edge_indices_test, radical_node_index_test, radical_edge_index_test
    ytrain = labels_train
    yval = labels_test

    return(xtrain,xval,ytrain,yval,scaler)

def megnet_dataset(input_dir, split_size):
    """
    Creates MegNet graph input
    Args:
        input_dir (str): input direction with files created using generate_input_data function
        split_size (int): test size
    Return:
        xtrain (ragged tensor): training data
        xtest (ragged tensor): test data
        xval (ragged tensor): validation data
        ytrain (ragged tensor): training labels
        ytest (ragged tensor): test labels
        yval (ragged tensor): validation labels
        scaler: scaler config for later rescaling of the data
    """
    y_data = np.expand_dims(np.load(input_dir+"cr_labels.npy"),axis=-1)
    nodes = load_pickle_file(input_dir+"cr_nodes.pickle")
    edges = load_pickle_file(input_dir+"cr_edges.pickle")
    edge_indices = load_pickle_file(input_dir+"cr_indices.pickle")
    atoms = np.expand_dims(np.load(input_dir+"cr_atoms.npy"), axis=-1)

    radical_node_index = [np.array([[0, 1]]) for _ in edge_indices]
    radical_edge_index = [np.array([[0]]) for _ in radical_node_index]

    scaler = StandardScaler(copy=True)
    labels = scaler.fit_transform(y_data)

    # Train Test split
    labels_train, labels_test, nodes_train, nodes_test, edges_train, edges_test, edge_indices_train, edge_indices_test, \
    atoms_train, atoms_test, radical_node_index_train, radical_node_index_test, radical_edge_index_train, \
        radical_edge_index_test = train_test_split(
        labels, nodes, edges, edge_indices, atoms, radical_node_index, radical_edge_index, test_size=split_size)
    del labels, nodes, edges, edge_indices, atoms, radical_node_index, radical_edge_index  # Free memory after split, if possible

    # Convert to tf.RaggedTensor or tf.tensor
    # a copy of the data is generated by ragged_tensor_from_nested_numpy()
    nodes_train, edges_train, edge_indices_train, atoms_train, radical_node_index_train, radical_edge_index_train = ragged_tensor_from_nested_numpy(
        nodes_train), ragged_tensor_from_nested_numpy(edges_train), ragged_tensor_from_nested_numpy(
        edge_indices_train), tf.constant(atoms_train), ragged_tensor_from_nested_numpy(
        radical_node_index_train), ragged_tensor_from_nested_numpy(radical_edge_index_train)

    nodes_test, edges_test, edge_indices_test, atoms_test, radical_node_index_test, radical_edge_index_test = ragged_tensor_from_nested_numpy(
        nodes_test), ragged_tensor_from_nested_numpy(edges_test), ragged_tensor_from_nested_numpy(
        edge_indices_test), tf.constant(atoms_test), ragged_tensor_from_nested_numpy(
        radical_node_index_test), ragged_tensor_from_nested_numpy(radical_edge_index_test)

    # Define input and output data
    xtrain = nodes_train, edges_train, edge_indices_train, atoms_train, radical_node_index_train, radical_edge_index_train
    xval = nodes_test, edges_test, edge_indices_test, atoms_test, radical_node_index_test, radical_edge_index_test
    ytrain = labels_train
    yval = labels_test
    return(xtrain,xval,ytrain,yval,scaler)

def dimenet_dataset(input_dir, split_size):
    """
    Creates DimeNet++ graph input
    Args:
        input_dir (str): input direction with files created using generate_input_data function
        split_size (int): test size
    Return:
        xtrain (ragged tensor): training data
        xtest (ragged tensor): test data
        xval (ragged tensor): validation data
        ytrain (ragged tensor): training labels
        ytest (ragged tensor): test labels
        yval (ragged tensor): validation labels
        scaler: scaler config for later rescaling of the data
    """
    y_data = np.expand_dims(np.load(input_dir+"cr_labels.npy"), axis=-1)
    nodes = load_pickle_file(input_dir+"cr_nodes.pickle")
    coord = load_pickle_file(input_dir+"cr_coord.pickle")

    dist = [coordinates_to_distancematrix(x) for x in coord]
    edge_indices = [define_adjacency_from_distance(x, max_distance=5, max_neighbours=25)[1] for x in dist]
    edge_indices = [x if x[0,1]==1 else sort_edge_indices(np.concatenate([np.array([[0, 1], [1,0]]), x], axis=0)) for x in edge_indices]
    angle_indices = [get_angle_indices(x, is_sorted=True)[2] for x in edge_indices]
    del dist  # Frees some memory

    radical_node_index = [np.array([[0, 1]]) for _ in edge_indices]  # für Kohlenstoff
    radical_edge_index = [np.array([[0]]) for _ in radical_node_index]

    scaler = StandardScaler(copy=True)
    labels = scaler.fit_transform(y_data)

    # Train Test split
    labels_train, labels_test, nodes_train, nodes_test, coord_train, coord_test, edge_indices_train, edge_indices_test, \
    angle_indices_train, angle_indices_test, \
    radical_node_index_train, radical_node_index_test, radical_edge_index_train, radical_edge_index_test = train_test_split(
        labels, nodes, coord, edge_indices, angle_indices, radical_node_index, radical_edge_index, test_size=split_size)
    del labels, nodes, coord, edge_indices, angle_indices, radical_node_index, radical_edge_index  # Free memory after split, if possible

    # Convert to tf.RaggedTensor or tf.tensor
    # a copy of the data is generated by ragged_tensor_from_nested_numpy()
    nodes_train, coord_train, edge_indices_train, angle_indices_train, radical_node_index_train, radical_edge_index_train = ragged_tensor_from_nested_numpy(
        nodes_train), ragged_tensor_from_nested_numpy(coord_train), ragged_tensor_from_nested_numpy(
        edge_indices_train), ragged_tensor_from_nested_numpy(angle_indices_train), ragged_tensor_from_nested_numpy(
        radical_node_index_train), ragged_tensor_from_nested_numpy(radical_edge_index_train)

    nodes_test, coord_test, edge_indices_test, angle_indices_test, radical_node_index_test, radical_edge_index_test = ragged_tensor_from_nested_numpy(
        nodes_test), ragged_tensor_from_nested_numpy(coord_test), ragged_tensor_from_nested_numpy(
        edge_indices_test), ragged_tensor_from_nested_numpy(angle_indices_test), ragged_tensor_from_nested_numpy(
        radical_node_index_test), ragged_tensor_from_nested_numpy(radical_edge_index_test)

    # Define input and output data
    xtrain = nodes_train, coord_train, edge_indices_train, angle_indices_train, radical_node_index_train, radical_edge_index_train
    xval = nodes_test, coord_test, edge_indices_test, angle_indices_test, radical_node_index_test, radical_edge_index_test
    ytrain = labels_train
    yval = labels_test
    return(xtrain,xval,ytrain,yval,scaler)

reason: import numpy,from numpy


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

# This file implements GPU KMeans Clustering after collecting RDD's data
# using Spark's collect().

import sys
import numpy
import random
import time
from itertools import groupby
from operator import itemgetter

from pyspark import SparkContext

from pycuda import gpuarray
import pycuda.autoinit
from pycuda.compiler import SourceModule
import pycuda.driver as cuda

def parseVector(line):
    return numpy.array([float(x) for x in line.split(' ')])

def closestPoint(p, centers):
    bestIndex = 0
    closest = float("+inf")
    for i in range(len(centers)):
        tempDist = numpy.sum((p - centers[i]) ** 2)
        if tempDist < closest:
            closest = tempDist
            bestIndex = i
    return bestIndex

def reduceClusters(obj1, obj2): #(x[0], ((x[1],x[2],x[3]), 1))
    return ( obj1[0],
		[obj1[1][0]+obj2[1][0], obj1[1][1]+obj2[1][1], obj1[1][2]+obj2[1][2], obj1[1][3]+obj2[1][3]] )

def divideCluster(container, y):
    return [x/y for x in container]

if __name__ == "__main__":

    if len(sys.argv) != 4:
        print >> sys.stderr, "Usage: ./spark-submit kmeans_collect <file> <k> <convergeDist>"
        exit(-1)

    start = time.time()
    sc = SparkContext(appName="PythonKMeans")
    lines = sc.textFile(sys.argv[1], 2)
    rdddata = lines.map(parseVector).cache()
    data = numpy.asarray(rdddata.collect(), dtype = numpy.float32)

    k = int(sys.argv[2])
    convergeDist = float(sys.argv[3])
    kPoints = numpy.asarray(random.sample(data, k))
    datasize = len(data)
    centroids = numpy.empty((datasize), gpuarray.vec.float4)
    mod = SourceModule("""
	__global__ void assignToCentroid(float *data, int* datasize, int *k, float *kPoints, float4 * clusters)
	{
		//1d grid index
		unsigned idxInLeftBlocks = blockIdx.x * (blockDim.x * blockDim.y);
		unsigned idxInCurrBlock  = threadIdx.y * blockDim.x + threadIdx.x;
		unsigned idx = idxInLeftBlocks + idxInCurrBlock;
		int size = datasize[0];
		int K = k[0];
		if (idx < size*3){//consider xyz for each data point
			if (idx % 3 == 0){
				unsigned dataIdx = idx/3;
				clusters[dataIdx].x = 999999;
				clusters[dataIdx].y = data[idx];
				clusters[dataIdx].z = data[idx+1];
				clusters[dataIdx].w = data[idx+2];
			}
			float inf =  99999.9999;
			float closest = inf;

			float tempDist=0, tempDistX=0, tempDistY=0, tempDistZ=0;
			for (int i = 0; i < K; i++){//hardwired. TODO: make it an argument
				if( idx % 3 == 0 /*0,3,6*/)
					tempDistX = data[idx] - kPoints[i*3];
				else if(idx % 3 == 1 /*1,4,7*/)
					tempDistY = data[idx] - kPoints[i*3 + 1];
				else {/*2,5,8*/
					tempDistZ = data[idx] - kPoints[i*3 + 2];
					tempDist = pow(tempDistX,2) + pow(tempDistY,2) + pow(tempDistZ,2);
				}

				if(tempDist < closest){
					closest = tempDist;
					int dataIdx = idx/3;
					clusters[dataIdx].x = i;//coordinate belongs to cluster i
				}
			}
		}
	}

    """)

    gridNum = (datasize * 3)/256 + 1
    datasize = numpy.asarray(datasize)
    k = numpy.asarray(k)
    tempDist = 1.0
    while tempDist > convergeDist:
	func = mod.get_function("assignToCentroid")
        func(cuda.In(data), cuda.In(datasize), cuda.In(k), cuda.In(kPoints), cuda.Out(centroids),
		block=(16, 16, 1), grid=(gridNum, 1), shared=0)

        closest = [(x[0], [x[1],x[2],x[3], 1]) for x in  centroids]

        statPoints = [reduce(reduceClusters, group) for _, group in groupby(sorted(closest), key=itemgetter(0))]
        newPoints = map(lambda (x, y): (x, 	divideCluster(y[:3],y[3])), statPoints)

        tempDist = sum(numpy.sum((kPoints[int(x)] - y) ** 2) for (x, y) in newPoints)

   	for (x, y) in newPoints:                        # update centroids
             kPoints[x] = y

        print "Current distance: ", tempDist
    print "Final centers:\n"
    print newPoints

    sc.stop()
    stop = time.time()
    print "Completion time: ", (stop - start) ," seconds"

max_stars_repo_name: adobe-research/spark-gpu

max_issues_repo_name: adobe-research/spark-gpu

max_forks_repo_name: adobe-research/spark-gpu

reason: import numpy,import pycuda,from pycuda


text:
# -*- coding: utf-8 -*-

"""
@Author: Lyzhang
@Date:
@Description:
"""
import numpy as np

def evaluate_(batch_iter, model):
    """ To evaluate the model in EDU segmentation.
        因内存原因，我们依旧选择批量运算的方式
        不同的解码评测，不同的是，这个模式下只要一句一句进行，因为边界是未知
        的，所以模型要做完边界预测才知道下一步做什么。
    """
    c_b, g_b, h_b = 0., 0., 0.
    for n_batch, (inputs, target) in enumerate(batch_iter, start=1):
        word_ids, word_elmo_embeddings, pos_ids, graph, masks = inputs
        # (batch_size, boundary_num + 1),
        pred = model.decode_all(word_ids, word_elmo_embeddings, pos_ids, graph, masks)
        # (batch_size, boundary_num + 1, seq_len) --> (batch, boundary_num + 1)
        predict = pred.data.cpu().numpy()
        target = target.data.cpu().numpy()
        # 标准标签中标签值为1的下标对应的预测标签值数组中预测正确为1的个数。
        c_b += len(np.intersect1d(predict, target))
        # 人工标注的标签中值为1的标签整体
        g_b += len(target)
        # 预测结果中标签值为1的标签个数
        h_b += len(predict)
        # 因为均多包含了句子结束符，需要减去
        c_b -= 1
        g_b -= 1
        h_b -= 1
    # p r f
    p_b = 0. if h_b == 0 else c_b / h_b
    r_b = 0. if g_b == 0 else c_b / g_b
    f_b = 0. if (g_b + h_b) == 0 else (2 * c_b) / (g_b + h_b)
    return p_b, r_b, f_b


text:
import pytest
from sklearn import linear_model
import numpy
import chaospy

LINEAR_MODELS = {
    "none": None,
    "linear": linear_model.LinearRegression(fit_intercept=False),
    "elastic_net": linear_model.MultiTaskElasticNet(alpha=0.0001, fit_intercept=False),
    "lasso": linear_model.MultiTaskLasso(alpha=0.001, fit_intercept=False),
    "lasso_lars": linear_model.LassoLars(alpha=0.0001, fit_intercept=False),
    "lars": linear_model.Lars(n_nonzero_coefs=10, fit_intercept=False),
    "matching_pursuit": linear_model.OrthogonalMatchingPursuit(n_nonzero_coefs=10, fit_intercept=False),
    "ridge": linear_model.Ridge(alpha=0.1, fit_intercept=False),
}

@pytest.fixture
def samples(joint):
    return joint.sample(1000, rule="sobol")

@pytest.fixture
def evaluations(model_solver, samples):
    return numpy.array([model_solver(sample) for sample in samples.T])

@pytest.fixture(params=LINEAR_MODELS)
def linear_model(request, expansion_small, samples, evaluations):
    return chaospy.fit_regression(
        expansion_small, samples, evaluations, model=LINEAR_MODELS[request.param])

def test_regression_mean(linear_model, joint, true_mean):
    assert numpy.allclose(chaospy.E(linear_model, joint), true_mean, rtol=1e-2)

def test_regression_variance(linear_model, joint, true_variance):
    assert numpy.allclose(chaospy.Var(linear_model, joint), true_variance, rtol=3e-1)


text:
from argparse import ArgumentParser
from itertools import starmap

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from fyne import blackscholes, heston
from matplotlib.patches import Patch
from scipy.stats import gaussian_kde

import settings
from align_settings import STARTTIME, ENDTIME
from utils import resample

def safe_xs(*args, **kwargs):
    try:
        return pd.Series.xs(*args, **kwargs)
    except KeyError:
        return np.nan

def get_tick_size(quote):
    diffs = (quote['Ask'] + quote['Bid']).diff()
    diffs = diffs[diffs > 1e-6]
    return np.round(diffs.min(), 2)

def filter_tick_size(data, quote, size):
    tick_size = quote.groupby('Strike').apply(get_tick_size)
    return data.reindex(tick_size[tick_size == size].index, level='Strike')

def filter_trade_on_book(quote, trade):
    max_expiry = np.max(quote.index.get_level_values('Expiry'))
    trade = trade[trade.index.get_level_values('Expiry') <= max_expiry]

    quote_aligned = trade.groupby(['Class', 'Expiry', 'Strike']
                        ).apply(lambda o: resample(quote.xs(o.name),
                                                   o.xs(o.name).index))
    valid_trades = ((trade['Price'] == quote_aligned['Bid']) |
                    (trade['Price'] == quote_aligned['Ask']))

    filtered = trade[valid_trades]
    quote_aligned = quote_aligned.loc[valid_trades]
    filtered['Buy'] = filtered['Price'] == quote_aligned['Ask']
    filtered['Half-spread'] = (quote_aligned['Ask'] - quote_aligned['Bid']).round(2)/2

    return filtered

def compute_duration(quote):
    quote = quote.copy()
    quote['Half-spread'] = (quote['Ask'] - quote['Bid']).round(2)/2

    time = quote.reset_index('Time'
               ).set_index('Half-spread', append=True)[['Time']]
    time['Duration'] = time['Time'].groupby(['Class', 'Expiry', 'Strike']
                                  ).transform(lambda t: t.diff().shift(-1))
    time['Time'] += time['Duration']/2

    duration = time.set_index('Time', append=True)['Duration']
    duration /= pd.to_timedelta('1s')

    return duration

def compute_volume_duration(quote, trade):
    trade = filter_trade_on_book(quote, trade)
    volume = trade.set_index(['Half-spread', 'Buy'], append=True)['Volume']
    duration = compute_duration(quote)

    return volume, duration

def plot_arrival_rates_bubbles(volume, duration):
    volume = volume.groupby(['Class', 'Expiry', 'Strike', 'Half-spread', 'Buy']
                  ).sum()
    duration = duration.groupby(['Class', 'Expiry', 'Strike',
                                 'Half-spread']).sum()

    duration = duration[duration > 300]
    arrival_rate = volume.groupby(['Class', 'Expiry', 'Strike', 'Half-spread']
                        ).transform(lambda d: d.xs(d.name
                                              )/safe_xs(duration, d.name))
    arrival_rate.name = 'Arrival rate'

    fig, axes = plt.subplots(3, 2, sharey=True, sharex=True, figsize=(8, 10))
    patches = [Patch(color='b', alpha=.5, label='Call'),
               Patch(color='r', alpha=.5, label='Put')]
    axes[0, 1].legend(handles=patches)

    for row, (e, r_ex) in zip(axes, arrival_rate.groupby('Expiry')):
        for bs in ['Buy', 'Sell']:
            ax = row[0] if bs == 'Buy' else row[1]
            ax.set_title("Expiry: {}, {}".format(
                pd.to_datetime(e).strftime('%Y-%m-%d'), bs))

            for cp, cl in [('C', 'b'), ('P', 'r')]:
                r = r_ex.xs((cp, bs == 'Buy'), level=('Class', 'Buy'))
                r.reset_index(['Strike', 'Half-spread']).plot.scatter(
                    x='Strike', y='Half-spread', s=20*r/r_ex.mean(), ax=ax,
                    xlim=(325, 550), ylim=(0, None), alpha=.5, color=cl)

    return fig

def plot_arrival_rates(arrival_rate):
    depths = arrival_rate.index.get_level_values('Half-spread')
    arrival_rate = arrival_rate[depths > 0].dropna()
    bandwidth = 0.25
    levels = ['Class', 'Expiry', 'Buy']
    kernel = arrival_rate.groupby(levels).apply(
        lambda r: gaussian_kde(np.stack(r.xs(r.name, level=levels).index, axis=-1),
                               bandwidth, r.values))

    xlen, ylen = 200, 150
    xmin, xmax, ymin, ymax = -0.2, 0.15, 0.0, 0.3
    x = np.linspace(xmin, xmax, xlen)
    y = np.linspace(ymin, ymax, ylen)
    x_b, y_b = np.broadcast_arrays(x[:, None], y[None, :])

    fig, axes = plt.subplots(3, 2, sharex=True, sharey=True, figsize=(8, 10))
    patches = [Patch(color='tab:blue', label='Call'),
               Patch(color='tab:red', label='Put')]
    axes[0, 1].legend(handles=patches)
    for row, (e, k) in zip(axes, kernel.groupby('Expiry')):
        row[0].set_title("Expiry: {}, Buy".format(
            pd.to_datetime(e).strftime('%Y-%m-%d')))
        row[1].set_title("Expiry: {}, Sell".format(
            pd.to_datetime(e).strftime('%Y-%m-%d')))
        for cp, cm in [('C', plt.cm.Blues), ('P', plt.cm.Reds)]:
            z = k.xs((cp, e, True))(np.array([x_b.ravel(), y_b.ravel()]))
            z = np.rot90(np.reshape(z, x_b.shape))
            row[0].imshow(z, extent=[xmin, xmax, ymin, ymax], cmap=cm,
                          aspect='auto', alpha=.5)
            z = k.xs((cp, e, False))(np.array([x_b.ravel(), y_b.ravel()]))
            z = np.rot90(np.reshape(z, x_b.shape))
            row[1].imshow(z, extent=[xmin, xmax, ymin, ymax], cmap=cm,
                          aspect='auto', alpha=.5)

    return fig

def quote_slice(quote, start, end):
    quote = quote.copy()
    quote.loc[start] = np.nan
    quote.sort_index(inplace=True)
    quote.ffill(inplace=True)
    return quote.loc[start:end]

def plot_intraday(volume, duration):
    filtered_volume = volume.groupby(['Class',  'Expiry']).apply(lambda e: filter_tick_size(e.xs(e.name), quote.xs(e.name), 0.05))
    filtered_duration = duration.groupby(['Class', 'Expiry']).apply(lambda e: e.xs(e.name).reindex(np.unique(filtered_volume.xs(e.name).index.get_level_values('Strike')), level='Strike').dropna())
    filtered_duration = filtered_duration.groupby(['Class', 'Expiry', 'Strike']).apply(lambda o: o.xs(o.name).reindex(np.unique(filtered_volume.xs(o.name).index.get_level_values('Half-spread')), level='Half-spread').dropna())
    volume_by_depth = filtered_volume.groupby('Expiry').apply(lambda e: e.groupby('Half-spread').sum().loc[.025:.225])
    duration_by_depth = filtered_duration.groupby('Expiry').apply(lambda e: e.groupby('Half-spread').sum().loc[.025:.225])
    volume_kde = filtered_volume.groupby('Expiry').apply(lambda e: e.groupby(['Half-spread', 'Time']).sum().loc[.025:.225].groupby('Half-spread').apply(lambda d: gaussian_kde((d.xs(d.name).index - pd.to_datetime('2016-01-04'))/pd.to_timedelta('1s'), weights=d, bw_method=.25)))
    duration_kde = filtered_duration.groupby('Expiry').apply(lambda e: e.groupby(['Half-spread', 'Time']).sum().loc[.025:.225].groupby('Half-spread').apply(lambda d: gaussian_kde((d.xs(d.name).index - pd.to_datetime('2016-01-04'))/pd.to_timedelta('1s'), weights=d, bw_method=.05)))

    fig_volume, axes = plt.subplots(3, 1, sharex=True, **settings.PLOT)
    tmin, tmax = 8*3600, 16.5*3600
    t = np.linspace(tmin, tmax, 100)
    for ax, (e, v) in zip(axes, volume_kde.groupby('Expiry')):
        z = volume_by_depth.xs(e).values*np.array([k.xs(d)(t) for d, k in v.xs(e).groupby('Half-spread')]).T
        ax.imshow(np.rot90(z), extent=[tmin/3600, tmax/3600, 0.025, 0.225], aspect='auto')
        ax.set_ylabel('Half-spread (€)')
        ax.set_title('Expiry: {}'.format(pd.to_datetime(e).strftime('%Y-%m-%d')))
    ax.set_xlabel('Time (hours)')
    fig_volume.tight_layout()

    fig_duration, axes = plt.subplots(3, 1, sharex=True, **settings.PLOT)
    tmin, tmax = 8*3600, 16.5*3600
    t = np.linspace(tmin, tmax, 100)
    for ax, (e, du) in zip(axes, duration_kde.groupby('Expiry')):
        z = duration_by_depth.xs(e).values*np.array([k.xs(d)(t) for d, k in du.xs(e).groupby('Half-spread')]).T
        ax.imshow(np.rot90(z), extent=[tmin/3600, tmax/3600, 0.025, 0.225], aspect='auto')
        ax.set_ylabel('Half-spread (€)')
        ax.set_title('Expiry: {}'.format(pd.to_datetime(e).strftime('%Y-%m-%d')))
    ax.set_xlabel('Time (hours)')
    fig_duration.tight_layout()

    fig_arrival, axes = plt.subplots(3, 1, sharex=True, **settings.PLOT)
    tmin, tmax = 8*3600, 16.5*3600
    t = np.linspace(tmin, tmax, 100)
    for ax, (e, du), (_, v) in zip(axes, duration_kde.groupby('Expiry'), volume_kde.groupby('Expiry')):
        z_v = volume_by_depth.xs(e).values*np.array([k.xs(d)(t) for d, k in v.xs(e).groupby('Half-spread')]).T
        z_d = duration_by_depth.xs(e).values*np.array([k.xs(d)(t) for d, k in du.xs(e).groupby('Half-spread')]).T
        z = np.clip(z_v/z_d, 1e-3, 1e1)
        ax.imshow(np.rot90(np.log(z)), extent=[tmin/3600, tmax/3600, 0.025, 0.225], aspect='auto')
        ax.set_ylabel('Half-spread (€)')
        ax.set_title('Expiry: {}'.format(pd.to_datetime(e).strftime('%Y-%m-%d')))
    ax.set_xlabel('Time (hours)')
    fig_arrival.tight_layout()

    return fig_volume, fig_duration, fig_arrival

if __name__ == '__main__':
    cli = ArgumentParser()
    cli.add_argument('quote_filename')
    cli.add_argument('trade_filename')
    cli.add_argument('dest_bubbles_filename')
    cli.add_argument('dest_intraday_volume_filename')
    cli.add_argument('dest_intraday_duration_filename')
    cli.add_argument('dest_intraday_arrival_filename')
    args = cli.parse_args()

    quote = pd.read_parquet(args.quote_filename)
    trade = pd.read_parquet(args.trade_filename).xs('AEX')

    quote.sort_index(inplace=True)
    volume, duration = compute_volume_duration(quote, trade)
    fig = plot_arrival_rates_bubbles(volume, duration)
    fig.savefig(args.dest_bubbles_filename)

    figs = plot_intraday(volume, duration)
    figs[0].savefig(args.dest_intraday_volume_filename)
    figs[1].savefig(args.dest_intraday_duration_filename)
    figs[2].savefig(args.dest_intraday_arrival_filename)

max_stars_repo_name: dougmvieira/options-microstructure

max_issues_repo_name: dxcv/options-microstructure

max_forks_repo_name: dxcv/options-microstructure

reason: import numpy,from scipy


text:
#!/usr/bin/env python3

import numpy as np
import SimpleITK as sitk
import os
import pandas as pd
import sys
import seaborn as sns
import matplotlib.pyplot as plt

'''

***********
PLEASE READ
***********

Script to evaluated clot detection on benchmark dataset

Usage: python3 evaluate_seg.py path_to_roi path_to_ann save_results

Input arguments:

path_to_rois:
Full path to where ROIs are located
* User provides predictions in the form of roi labelled '<subject_id>_roi.nii.gz
* Subject IDs are in format sub-test*
* ROI images must have the same image properties (size, direction, origin, spacing) as the corresponding raw image:
    e.g., roi.GetSize() == annotation.GetSize()

path_to_annotations:
Full path to where the ground truth annotations are held
* Annotations are in the format '<subject_id>_seg.nii.gz
* A 3D ellipse that is [60, 60, 30] voxels wide is built around the ground truth annotations

save_results:
Full path to location for output ROC figure

'''

def get_score(path_to_roi, path_to_gt, subject, threshold):

    gt = sitk.ReadImage(path_to_gt)
    roi = sitk.ReadImage(path_to_roi)

    # connected component labelling
    gt_cc = sitk.ConnectedComponent(gt)

    # label statistics
    lsif = sitk.LabelShapeStatisticsImageFilter()
    lsif.Execute(gt_cc)

    # read in segmentation
    lsif_roi = sitk.LabelShapeStatisticsImageFilter()
    lsif_roi.Execute(roi)

    # if there is no clot
    if not lsif.GetLabels():
        # if you have provided an roi
        if not lsif_roi.GetLabels():
            return 'tn'
        # if you have provided an roi
        if lsif_roi.GetLabels():
            return 'fp'

    # if there is a clot
    else:
        # if you have provided an roi
        if lsif_roi.GetLabels():
            # get the location of the annotation centroid
            centroid = lsif.GetCentroid(1)
            # construct a 3D Gaussian sphere around it
            # sigma = [30, 30, 15]
            circle = sitk.GaussianSource(
                sitk.sitkInt16,
                size=gt_cc.GetSize(),
                sigma=[30, 30, 15],
                mean=gt_cc.TransformPhysicalPointToIndex(centroid))

            circle.CopyInformation(gt_cc)
            # calculate half max, to keep the circle diameter 2 x sigma
            max_val = np.max(sitk.GetArrayFromImage(circle))
            circle = circle > max_val /2

            # normalize labels
            circle_array = sitk.GetArrayFromImage(circle)
            circle_array = circle_array/np.max(circle_array)

            roi_array = sitk.GetArrayFromImage(roi)
            roi_array = roi_array/np.max(roi_array)

            # change background values so they are different from each other
            circle_array[circle_array == 0] = 10
            roi_array[roi_array == 0] = 5

            # subtract the prediction circle from the ground truth circle
            iou_array = circle_array - roi_array

            # count intersecting voxels
            intersection = len(np.where(iou_array == 0)[0])

            # count union voxels
            union = len(
                np.where(iou_array == -4)[0]) + len(
                np.where(iou_array == 9)[0]) + intersection

            # calculate Jaccard
            IoU = intersection/union

            # assign score
            if IoU > threshold:
                return 'tp'
            else:
                return 'fp'
        # if you have not provided an index
        if not lsif_roi.GetLabels():
            return 'fn'

# main part of the code
def main(path_to_rois, path_to_anns, save_results):
    # get list of roi files
    rois = [file for file in os.listdir(path_to_rois) if 'roi' in file]
    # subjects are in format sub-test*
    subs = [file.split('_')[0] for file in rois]

    # make an array of thresholds for ROC curve
    thresholds = np.arange(10, 100, 5)
    # make an empty DataFrame
    score_df = pd.DataFrame(
        index=subs,
        columns=[str(thresh) + '%' for thresh in thresholds])

    for roi, sub in zip(rois, subs):
        # get full paths to images
        path_to_roi = os.path.join(path_to_rois, roi)
        path_to_ann = os.path.join(path_to_anns, sub + '_seg.nii.gz')

        for thresh in thresholds:
            score = get_score(path_to_roi, path_to_ann, sub, thresh/10)
            column = str(thresh) + '%'
            score_df.loc[sub, column] = score

    # Calculate TPR and FPR over entire dataset
    tpr_row = pd.DataFrame(
        index=['tpr'],
        columns=[str(thresh) + '%' for thresh in thresholds]
    )
    fpr_row = pd.DataFrame(
        index=['fpr'],
        columns=[str(thresh) + '%' for thresh in thresholds]
    )

    for col in score_df:
        series = score_df[col]
        # get TPR and TPR for each column
        counts = series.value_counts()
        tpr = counts['tp'] / (counts['tp'] + counts['fn'])
        fpr = counts['fp'] / (counts['fp'] + counts['tn'])
        tpr_row.loc['tpr', col] = round(tpr, 3)
        fpr_row.loc['fpr', col] = round(fpr, 3)

    # append FPR and TPR for the bottom of the score card
    score_df = score_df.append(tpr_row)
    score_df = score_df.append(fpr_row)

    # make the ROC curve
    def create_triangle(tpr_0, tpr_1, fpr_0, fpr_1):
        plt.plot([tpr_0, tpr_1], [fpr_0, fpr_1], '-', lw=2, color='#4285F4')
        plt.plot([tpr_0, tpr_1], [fpr_1, fpr_1], '-', lw=2, color='#4285F4')
        plt.plot([tpr_0, tpr_0], [fpr_0, fpr_1], '-', lw=2, color='#4285F4')

    def create_rectangle(tpr_0, fpr_0):
        plt.plot([tpr_0, tpr_0], [fpr_0, 0], '-', lw=2, color='#4285F4')

    # create ROC image
    sns.set()
    plt.figure(figsize=(15, 7))
    plt.scatter(fpr_row, tpr_row, color='#0F9D58', s=100)

    plt.title('ROC Curve', fontsize=20)
    plt.xlabel('False Positive Rate', fontsize=16)
    plt.ylabel('True Positive Rate', fontsize=16)

    # pending testing (adding rectangles underneath)
    # plt.plot([1, 0], [0, 0], "-", lw=2, color="#4285F4")
    # for j in range(len(thresholds) - 1):
    #     create_rectangle(
    #         fpr_row.iloc[0, j],
    #         tpr_row.iloc[0, j])
    # for k in range(len(thresholds) - 1):
    #     create_triangle(
    #         fpr_row.iloc[0, k],
    #         fpr_row.iloc[0, k + 1],
    #         tpr_row.iloc[0, k],
    #         tpr_row.iloc[0, k + 1])

    # save figure and DataFrame
    plt.savefig(os.path.join(save_results, "roc_curve.png"), bbox_inches='tight')
    score_df.to_csv(os.path.join(save_results, 'scores.csv'))

    # construct the area under the ROC curve
    rectangle_auc = 0
    for k in range(len(thresholds)-1):
        rectangle_auc += (fpr_row.iloc[0, k + 1] - fpr_row.iloc[0, k]) * tpr_row.iloc[0, k]

    print('Rectangular AUC = {}'.format(rectangle_auc))

if __name__ == "__main__":
    if len(sys.argv) == 4:
        main(sys.argv[1], sys.argv[2], sys.argv[3])
    else:
        print('Insufficient arguments.')

max_stars_repo_name: MBC-Neuroimaging/ClotDetectEval

max_issues_repo_name: MBC-Neuroimaging/ClotDetectEval

max_forks_repo_name: MBC-Neuroimaging/ClotDetectEval


text:
"""Contains the 1D implicit_general solver.

Used to compute unidimensional thermal processes

"""

import numpy as np
import copy

def implicit_general(obj):
    """implicit_general solver.

    Used to compute one time step of 1D systems with fixed thermal
    conductivity.

    """
    # initializes the matrixes for the equation systems
    a = np.zeros((obj.num_points, obj.num_points))
    b = np.zeros(obj.num_points)

    # left boundary
    a[0][0] = 1
    if obj.boundaries[0] == 0:
        b[0] = obj.temperature[1][0]
    else:
        b[0] = obj.boundaries[0]

    # right boundary
    a[obj.num_points - 1][obj.num_points - 1] = 1
    if obj.boundaries[1] == 0:
        value = obj.temperature[obj.num_points - 2][0]
        b[obj.num_points - 1] = value
    else:
        b[obj.num_points - 1] = obj.boundaries[1]

    # creates the matrixes and solves the equation systems
    for i in range(1, obj.num_points - 1):
        beta = obj.k[i] * obj.dt / \
            (2 * obj.rho[i] * obj.Cp[i] * obj.dx * obj.dx)
        sigma = obj.dt / (obj.rho[i] * obj.Cp[i])

        a[i][i - 1] = -beta
        a[i][i] = 1 + 2 * beta - sigma * obj.Q[i]
        a[i][i + 1] = -beta
        b[i] = (1 - 2 * beta - sigma * obj.Q[i]) * \
            obj.temperature[i][0] + \
            beta * obj.temperature[i + 1][0] + \
            beta * obj.temperature[i - 1][0] + 2. * sigma * \
            (obj.Q0[i] - obj.Q[i] * obj.amb_temperature)

    x = np.linalg.solve(a, b)

    # latent heat
    lheat = copy.copy(obj.lheat)
    for i in range(1, obj.num_points - 1):
        j = 0
        for lh in obj.latent_heat[i]:
            temper = obj.temperature[i][0]
            if x[i] > lh[0] and temper <= lh[0] and lheat[i][j][1] != lh[1]:
                en = obj.Cp[i] * obj.rho[i] * (x[i] - obj.temperature[i][0])
                if en + lheat[i][j][1] >= lh[1]:
                    lheat[i][j][1] = lh[1]
                    energy_temp = lheat[i][j][1] + en - lh[1]
                    x[i] = obj.temperature[i][0] + \
                        energy_temp / (obj.Cp[i] * obj.rho[i])
                else:
                    lheat[i][j][1] += en
                    x[i] = obj.temperature[i][0]
            if x[i] < lh[0] and temper >= lh[0] and lheat[i][j][1] != 0:
                en = obj.Cp[i] * obj.rho[i] * (x[i] - obj.temperature[i][0])
                if en + lheat[i][j][1] <= 0.:
                    lheat[i][j][1] = 0.
                    energy_temp = (en + lheat[i][j][1])
                    x[i] = obj.temperature[i][0] + \
                        energy_temp / (obj.Cp[i] * obj.rho[i])
                else:
                    lheat[i][j][1] += en
                    x[i] = obj.temperature[i][0]
            j += 1

    y = copy.deepcopy(obj.temperature)

    # updates the temperature list
    for i in range(obj.num_points):
        y[i][1] = x[i]
        y[i][0] = x[i]

    return y, lheat

max_issues_repo_name: danieljosesilva/heatrapy