text:
import numpy as np
import pandas as pd
import h5py
from typing import Dict, Optional, Union, List
import pickle
import warnings
from pathlib import Path
from collections import OrderedDict
import json
import sparse
from fastcore.all import is_listy

from sklearn.pipeline import Pipeline

__all__ = ['FoldYielder', 'HEPAugFoldYielder']

class FoldYielder:
    r'''
    Interface class for accessing data from foldfiles created by :meth:`~lumin.data_processing.file_proc.df2foldfile`

    Arguments:
        foldfile: filename of hdf5 file or opened hdf5 file
        cont_feats: list of names of continuous features present in input data, not required if foldfile contains meta data already
        cat_feats: list of names of categorical features present in input data, not required if foldfile contains meta data already
        ignore_feats: optional list of input features which should be ignored
        input_pipe: optional Pipeline, or filename for pickled Pipeline, which was used for processing the inputs
        output_pipe: optional Pipeline, or filename for pickled Pipeline, which was used for processing the targets
        yield_matrix: whether to actually yield matrix data if present
        matrix_pipe: preprocessing pipe for matrix data

    Examples::
        >>> fy = FoldYielder('train.h5')
        >>>
        >>> fy = FoldYielder('train.h5', ignore_feats=['phi'], input_pipe='input_pipe.pkl')
        >>>
        >>> fy = FoldYielder('train.h5', input_pipe=input_pipe, matrix_pipe=matrix_pipe)
        >>>
        >>> fy = FoldYielder('train.h5', input_pipe=input_pipe, yield_matrix=False)
    '''

    # TODO: Matrix example

    def __init__(self, foldfile:Union[str,Path,h5py.File], cont_feats:Optional[List[str]]=None, cat_feats:Optional[List[str]]=None,
                 ignore_feats:Optional[List[str]]=None, input_pipe:Optional[Union[str,Pipeline,Path]]=None, output_pipe:Optional[Union[str,Pipeline,Path]]=None,
                 yield_matrix:bool=True, matrix_pipe:Optional[Union[str,Pipeline,Path]]=None):
        self.cont_feats,self.cat_feats,self.input_pipe,self.output_pipe = cont_feats,cat_feats,input_pipe,output_pipe
        self.yield_matrix,self.matrix_pipe = yield_matrix,matrix_pipe
        self.augmented,self.aug_mult,self.train_time_aug,self.test_time_aug = False,0,False,False
        self._set_foldfile(foldfile)
        self.input_feats = self.cont_feats + self.cat_feats
        self.orig_cont_feats,self.orig_cat_feat,self._ignore_feats = self.cont_feats,self.cat_feats,[]
        if isinstance(self.input_pipe, str) or isinstance(self.input_pipe, Path): self.add_input_pipe_from_file(self.input_pipe)
        if isinstance(self.output_pipe, str) or isinstance(self.output_pipe, Path): self.add_output_pipe_from_file(self.output_pipe)
        if isinstance(self.matrix_pipe, str) or isinstance(self.matrix_pipe, Path): self.add_matrix_pipe_from_file(self.matrix_pipe)
        if ignore_feats is not None: self.add_ignore(ignore_feats)

    def __repr__(self) -> str: return f'FoldYielder with {self.n_folds} folds, containing {self.columns()}'

    def __len__(self) -> int: return self.n_folds

    def __getitem__(self, idx:int) -> Dict[str,np.ndarray]: return self.get_fold(idx)

    def __iter__(self) -> Dict[str,np.ndarray]:
        for i in range(self.n_folds): yield self.get_fold(i)

    def columns(self) -> List[str]:
        r'''
        Returns list of columns present in foldfile

        Returns:
            list of columns present in foldfile
        '''

        return [k for k in self.foldfile["fold_0"].keys()]

    def add_ignore(self, feats:Union[str,List[str]]) -> None:
        r'''
        Add features to ignored features.

        Arguments:
            feats: list of feature names to ignore
        '''

        if not is_listy(feats): feats = [feats]
        self._ignore_feats += feats
        self.cont_feats = [f for f in self.cont_feats if f not in self._ignore_feats]
        self.cat_feats  = [f for f in self.cat_feats  if f not in self._ignore_feats]

    def get_ignore(self) -> List[str]:
        r'''
        Returns list of ignored features

        Returns:
            Features removed from training data
        '''

        return self._ignore_feats

    def get_use_cont_feats(self) -> List[str]:
        r'''
        Returns list of continuous features which will be present in training data, accounting for ignored features.

        Returns:
            List of continuous features
        '''

        return [f for f in self.cont_feats if f not in self._ignore_feats]

    def get_use_cat_feats(self) -> List[str]:
        r'''
        Returns list of categorical features which will be present in training data, accounting for ignored features.

        Returns:
            List of categorical features
        '''

        return [f for f in self.cat_feats if f not in self._ignore_feats]

    def _set_foldfile(self, foldfile:Union[str,Path,h5py.File]) -> None:
        r'''
        Sets the file from which to access data

        Arguments:
            foldfile: filename of h5py file or opened h5py file
        '''

        if not isinstance(foldfile,  h5py.File): foldfile = h5py.File(foldfile, "r+")
        self.foldfile, self.n_folds = foldfile, len([f for f in foldfile if 'fold_' in f])
        self.has_matrix = 'matrix_inputs' in self.columns()
        if 'meta_data' in self.foldfile: self._load_meta_data()
        self.fld_szs = {i:self.foldfile[f'fold_{i}/targets'].shape[0] for i in range(self.n_folds)}

    def get_data_count(self, idxs:Union[int,List[int]]) -> int:
        r'''
        Returns total number of data entries in requested folds

        Arguments:
            idxs: list of indices to check

        Returns:
            Total number of entries in the folds
        '''

        if not is_listy(idxs): idxs = [idxs]
        s = 0
        for i in idxs: s += self.fld_szs[i]
        return s

    def _load_meta_data(self) -> None:
        if self.cont_feats is not None:
            warnings.warn("Fold file contains meta data information, explicit passing of continuous and categorical feature lists is no longer required.")
        self.cont_feats = json.loads(self.foldfile['meta_data/cont_feats'][()])
        self.cat_feats  = json.loads(self.foldfile['meta_data/cat_feats'][()])
        self.targ_feats = json.loads(self.foldfile['meta_data/targ_feats'][()])
        if 'wgt_feat' in self.foldfile['meta_data']: self.wgt_feat = json.loads(self.foldfile['meta_data/wgt_feat'][()])
        if 'cat_maps' in self.foldfile['meta_data']: self.cat_maps = OrderedDict(json.loads(self.foldfile['meta_data/cat_maps'][()]))
        if self.has_matrix:
            self.matrix_feats = json.loads(self.foldfile['meta_data/matrix_feats'][()])
            self.matrix_feats['missing'] = np.array(self.matrix_feats['missing'], dtype=np.bool)
            self.matrix_is_sparse = self.matrix_feats['is_sparse'] if 'is_sparse' in self.matrix_feats else False
            self.matrix_shape = self.matrix_feats['shape'] if 'shape' in self.matrix_feats else False

    def _append_matrix(self, data, idx) -> Dict[str,np.ndarray]:
        data['inputs'] = (data['inputs'],self.get_column('matrix_inputs', n_folds=1, fold_idx=idx))
        return data

    def close(self) -> None:
        r'''
        Closes the foldfile
        '''

        self.foldfile.close()

    def add_input_pipe(self, input_pipe:Union[str,Pipeline]) -> None:
        r'''
        Adds an input pipe to the FoldYielder for use when deprocessing data

        Arguments:
            input_pipe: Pipeline which was used for preprocessing the input data or name of pkl file containing Pipeline
        '''

        if isinstance(input_pipe, str) or isinstance(input_pipe, Path): self.add_input_pipe_from_file(input_pipe)
        else:                                                           self.input_pipe = input_pipe

    def add_matrix_pipe(self, matrix_pipe:Union[str,Pipeline]) -> None:
        r'''
        Adds an matrix pipe to the FoldYielder for use when deprocessing data

        .. Warning:: Deprocessing matrix data is not yet implemented

        Arguments:
            matrix_pipe: Pipeline which was used for preprocessing the input data or name of pkl file containing Pipeline
        '''

        if isinstance(matrix_pipe, str) or isinstance(matrix_pipe, Path): self.add_matrix_pipe_from_file(matrix_pipe)
        else:                                                             self.matrix_pipe = matrix_pipe

    def add_output_pipe(self, output_pipe:Union[str,Pipeline]) -> None:
        r'''
        Adds an output pipe to the FoldYielder for use when deprocessing data

        Arguments:
            output_pipe: Pipeline which was used for preprocessing the target data or name of pkl file containing Pipeline
        '''

        if isinstance(output_pipe, str) or isinstance(output_pipe, Path): self.add_output_pipe_from_file(output_pipe)
        else:                                                             self.output_pipe = output_pipe

    def add_input_pipe_from_file(self, name:Union[str,Path]) -> None:
        r'''
        Adds an input pipe from a pkl file to the FoldYielder for use when deprocessing data

        Arguments:
            name: name of pkl file containing Pipeline which was used for preprocessing the input data
        '''

        with open(name, 'rb') as fin: self.input_pipe = pickle.load(fin)

    def add_matrix_pipe_from_file(self, name:str) -> None:
        r'''
        Adds an matrix pipe from a pkl file to the FoldYielder for use when deprocessing data

        Arguments:
            name: name of pkl file containing Pipeline which was used for preprocessing the matrix data
        '''

        with open(name, 'rb') as fin: self.matrix_pipe = pickle.load(fin)

    def add_output_pipe_from_file(self, name:Union[str,Path]) -> None:
        r'''
        Adds an output pipe from a pkl file to the FoldYielder for use when deprocessing data

        Arguments:
            name: name of pkl file containing Pipeline which was used for preprocessing the target data
        '''

        with open(name, 'rb') as fin: self.output_pipe = pickle.load(fin)

    def get_fold(self, idx:int) -> Dict[str,np.ndarray]:
        r'''
        Get data for single fold. Data consists of dictionary of inputs, targets, and weights.
        Accounts for ignored features.
        Inputs, except for matrix data, are passed through np.nan_to_num to deal with nans and infs.

        Arguments:
            idx: fold index to load

        Returns:
            tuple of inputs, targets, and weights as Numpy arrays
        '''

        data = self.get_data(n_folds=1, fold_idx=idx)
        if len(self._ignore_feats) == 0:
            return self._append_matrix(data, idx) if self.has_matrix and self.yield_matrix else data
        else:
            inputs = pd.DataFrame(data['inputs'], columns=self.input_feats)
            inputs = inputs[[f for f in self.input_feats if f not in self._ignore_feats]]  # TODO Improve this with preconfigured mask
            data['inputs'] = inputs.values
            return self._append_matrix(data, idx) if self.has_matrix and self.yield_matrix else data

    def get_column(self, column:str, n_folds:Optional[int]=None, fold_idx:Optional[int]=None, add_newaxis:bool=False) -> Union[np.ndarray, None]:
        r'''
        Load column (h5py group) from foldfile. Used for getting arbitrary data which isn't automatically grabbed by other methods.

        Arguments:
            column: name of h5py group to get
            n_folds: number of folds to get data from. Default all folds. Not compatable with fold_idx
            fold_idx: Only load group from a single, specified fold. Not compatable with n_folds
            add_newaxis: whether expand shape of returned data if data shape is ()

        Returns:
            Numpy array of column data
        '''

        if column not in self.columns(): return None

        if fold_idx is None:
            data = []
            for i, fold in enumerate([f for f in self.foldfile if 'fold_' in f]):
                if n_folds is not None and i >= n_folds: break
                tmp = self.foldfile[f'{fold}/{column}'][()]
                if column == 'matrix_inputs' and self.matrix_is_sparse:
                    c = tmp[1:].astype(int)
                    tmp = sparse.COO(coords=c, data=tmp[0], shape=[c[0][-1]+1]+self.matrix_shape).todense()
                data.append(tmp)
            data = np.concatenate(data)
        else:
            if f'fold_{fold_idx}' not in self.foldfile: raise IndexError(f"Fold {fold_idx} does not exist")
            data = self.foldfile[f'fold_{fold_idx}/{column}'][()]
            if column == 'matrix_inputs' and self.matrix_is_sparse:
                c = data[1:].astype(int)
                data = sparse.COO(coords=c, data=data[0], shape=[c[0][-1]+1]+self.matrix_shape).todense()
        return data[:, None] if data[0].shape == () and add_newaxis else data

    def get_data(self, n_folds:Optional[int]=None, fold_idx:Optional[int]=None) -> Dict[str,np.ndarray]:
        r'''
        Get data for single, specified fold or several of folds. Data consists of dictionary of inputs, targets, and weights.
        Does not account for ignored features.
        Inputs are passed through np.nan_to_num to deal with nans and infs.

        Arguments:
            n_folds: number of folds to get data from. Default all folds. Not compatable with fold_idx
            fold_idx: Only load group from a single, specified fold. Not compatable with n_folds

        Returns:
            tuple of inputs, targets, and weights as Numpy arrays
        '''

        return {'inputs':  np.nan_to_num(self.get_column('inputs',  n_folds=n_folds, fold_idx=fold_idx)),
                'targets':               self.get_column('targets', n_folds=n_folds, fold_idx=fold_idx, add_newaxis=True),
                'weights':               self.get_column('weights', n_folds=n_folds, fold_idx=fold_idx, add_newaxis=True)}

    def get_df(self, pred_name:str='pred', targ_name:str='targets', wgt_name:str='weights', n_folds:Optional[int]=None, fold_idx:Optional[int]=None,
               inc_inputs:bool=False, inc_ignore:bool=False, deprocess:bool=False, verbose:bool=True, suppress_warn:bool=False,
               nan_to_num:bool=False, inc_matrix:bool=False) -> pd.DataFrame:
        r'''
        Get a Pandas DataFrameof the data in the foldfile. Will add columns for inputs (if requested), targets, weights, and predictions (if present)

        Arguments:
            pred_name: name of prediction group
            targ_name: name of target group
            wgt_name: name of weight group
            n_folds: number of folds to get data from. Default all folds. Not compatable with fold_idx
            fold_idx: Only load group from a single, specified fold. Not compatable with n_folds
            inc_inputs: whether to include input data
            inc_ignore: whether to include ignored features
            deprocess: whether to deprocess inputs and targets if pipelines have been
            verbose: whether to print the number of datapoints loaded
            suppress_warn: whether to supress the warning about missing columns
            nan_to_num: whether to pass input data through `np.nan_to_num`
            inc_matrix: whether to include flattened matrix data in output, if present

        Returns:
            Pandas DataFrame with requested data
        '''

        # TODO Decide how to handle deprocessing matrix data: option for object by object, flattened out?

        if inc_inputs:
            inputs = self.get_column('inputs', n_folds=n_folds, fold_idx=fold_idx)
            if deprocess and self.input_pipe is not None:
                try:
                    inputs = np.hstack((self.input_pipe.inverse_transform(inputs[:,:len(self.orig_cont_feats)]), inputs[:,len(self.orig_cont_feats):]))
                except ValueError:
                    if self.has_matrix:
                        print('Deprocessing of flat data failed, possible due to the input_pipe expecting to also transform matrix data. Deprocessing of matrix'
                              'is not currently implemented, and deprocessing of flat data using an input_pipe which expects matrix data as well is difficult '
                              'due to loss of variable ordering. In future please use separate pipes to preprocess flat data and matrix data. Returning inputs '
                              'as processed.')
                    else:
                        print('Deprocessing of flat data failed, returning inputs as processed.')

            if nan_to_num: inputs = np.nan_to_num(inputs)
            data = pd.DataFrame(inputs, columns=self.input_feats)
            if len(self._ignore_feats) > 0 and not inc_ignore: data = data[[f for f in self.input_feats if f not in self._ignore_feats]]
            if self.has_matrix and inc_matrix:
                mat = self.get_column('matrix_inputs', n_folds=n_folds, fold_idx=fold_idx).reshape(len(inputs),np.multiply(*self.matrix_feats['shape']))
                mat = mat[:,np.logical_not(self.matrix_feats['missing'])]
                # if deprocess and self.matrix_pipe is not None: mat = self.matrix_pipe.inverse_transform(mat)
                if nan_to_num: mat = np.nan_to_num(mat)
                data = data.join(pd.DataFrame(mat, columns=self.matrix_feats['present_feats']))
        else:
            data = pd.DataFrame()

        targets = self.get_column(targ_name, n_folds=n_folds, fold_idx=fold_idx)
        if deprocess and self.output_pipe is not None: targets = self.output_pipe.inverse_transform(targets)
        if targets is not None and len(targets.shape) > 1:
            for t in range(targets.shape[-1]): data[f'gen_target_{t}'] = targets[:,t]
        elif targets is None and not suppress_warn:
            warnings.warn(f"{targ_name} not found in file")
        else:
            data['gen_target'] = targets

        weights = self.get_column(wgt_name, n_folds=n_folds, fold_idx=fold_idx)
        if weights is not None and weights is not None and len(weights.shape) > 1:
            for w in range(weights.shape[-1]): data[f'gen_weight_{w}'] = weights[:,w]
        elif weights is None and not suppress_warn:
            warnings.warn(f"{wgt_name} not found in file")
        else:
            data['gen_weight'] = weights

        preds = self.get_column(pred_name, n_folds=n_folds, fold_idx=fold_idx)
        if deprocess and self.output_pipe is not None: preds = self.output_pipe.inverse_transform(preds)
        if preds is not None and len(preds.shape) > 1:
            for p in range(preds.shape[-1]): data[f'pred_{p}'] = preds[:,p]
        elif preds is not None:
            data['pred'] = preds
        elif not suppress_warn:
            warnings.warn(f'{pred_name} not found in foldfile file')
        if verbose: print(f'{len(data)} datapoints loaded')
        return data

    def save_fold_pred(self, pred:np.ndarray, fold_idx:int, pred_name:str='pred') -> None:
        r'''
        Save predictions for given fold as a new column in the foldfile

        Arguments:
            pred: array of predictions in the same order as data appears in the file
            fold_idx: index for fold
            pred_name: name of column to save predictions under
        '''

        n = f'fold_{fold_idx}/{pred_name}'
        if n in self.foldfile: del self.foldfile[n]
        self.foldfile.create_dataset(n, shape=pred.shape, dtype='float32')
        self.foldfile[n][...] = pred

class HEPAugFoldYielder(FoldYielder):
    r'''
    Specialised version of :class:`~lumin.nn.data.fold_yielder.FoldYielder` providing HEP specific data augmetation at train and test time.

    Arguments:
        foldfile: filename of hdf5 file or opened hdf5 file
        cont_feats: list of names of continuous features present in input data, not required if foldfile contains meta data already
        cat_feats: list of names of categorical features present in input data, not required if foldfile contains meta data already
        ignore_feats: optional list of input features which should be ignored
        targ_feats: optional list of target features to also be transformed
        rot_mult: number of rotations of event in phi to make at test-time (currently must be even).
                  Greater than zero will also apply random rotations during train-time
        random_rot: whether test-time rotation angles should be random or in steps of 2pi/rot_mult
        reflect_x: whether to reflect events in x axis at train and test time
        reflect_y: whether to reflect events in y axis at train and test time
        reflect_z: whether to reflect events in z axis at train and test time
        train_time_aug: whether to apply augmentations at train time
        test_time_aug: whether to apply augmentations at test time
        input_pipe: optional Pipeline, or filename for pickled Pipeline, which was used for processing the inputs
        output_pipe: optional Pipeline, or filename for pickled Pipeline, which was used for processing the targets
        yield_matrix: whether to actually yield matrix data if present
        matrix_pipe: preprocessing pipe for matrix data

    Examples::
        >>> fy = HEPAugFoldYielder('train.h5',
        ...                        cont_feats=['pT','eta','phi','mass'],
        ...                        rot_mult=2, reflect_y=True, reflect_z=True,
        ...                        input_pipe='input_pipe.pkl')
    '''

    def __init__(self, foldfile:Union[str,Path,h5py.File], cont_feats:Optional[List[str]]=None, cat_feats:Optional[List[str]]=None,
                 ignore_feats:Optional[List[str]]=None, targ_feats:Optional[List[str]]=None,
                 rot_mult:int=2, random_rot:bool=False,
                 reflect_x:bool=False, reflect_y:bool=True, reflect_z:bool=True,
                 train_time_aug:bool=True, test_time_aug:bool=True,
                 input_pipe:Optional[Pipeline]=None, output_pipe:Optional[Pipeline]=None,
                 yield_matrix:bool=True, matrix_pipe:Optional[Union[str,Pipeline]]=None):
        super().__init__(foldfile=foldfile, cont_feats=cont_feats, cat_feats=cat_feats,
                         ignore_feats=ignore_feats, input_pipe=input_pipe, output_pipe=output_pipe,
                         yield_matrix=yield_matrix, matrix_pipe=matrix_pipe)

        if rot_mult > 0 and not random_rot and rot_mult % 2 != 0:
            warnings.warn('Warning: rot_mult must currently be even for fixed rotations, adding an extra rotation multiplicity')
            rot_mult += 1
        self.rot_mult,self.random_rot,self.reflect_x,self.reflect_y,self.reflect_z,self.train_time_aug,self.test_time_aug,self.targ_feats = \
            rot_mult,random_rot,reflect_x,reflect_y,reflect_z,train_time_aug,test_time_aug,targ_feats
        self.augmented,self.reflect_axes,self.aug_mult = True,[],1
        self.vectors = [x[:-3] for x in self.cont_feats if '_px' in x]
        if self.targ_feats is not None: self.targ_vectors = [x[:-3] for x in self.targ_feats if '_px' in x]

        if self.rot_mult:
            print("Augmenting via phi rotations")
            self.aug_mult = self.rot_mult
            if self.reflect_y:
                print("Augmenting via y flips")
                self.reflect_axes += ['_py']
                self.aug_mult *= 2
            if self.reflect_z:
                print("Augmenting via longitunidnal flips")
                self.reflect_axes += ['_pz']
                self.aug_mult *= 2
        else:
            if self.reflect_x:
                print("Augmenting via x flips")
                self.reflect_axes += ['_px']
                self.aug_mult *= 2
            if self.reflect_y:
                print("Augmenting via y flips")
                self.reflect_axes += ['_py']
                self.aug_mult *= 2
            if self.reflect_z:
                print("Augmenting via longitunidnal flips")
                self.reflect_axes += ['_pz']
                self.aug_mult *= 2
        print(f'Total augmentation multiplicity is {self.aug_mult}')

    def _rotate(self, df:pd.DataFrame, vecs:List[str]) -> None:
        for vec in vecs:
            df.loc[:, f'{vec}_pxtmp'] = df.loc[:, f'{vec}_px']*np.cos(df.loc[:, 'aug_angle'])-df.loc[:, f'{vec}_py']*np.sin(df.loc[:, 'aug_angle'])
            df.loc[:, f'{vec}_py']    = df.loc[:, f'{vec}_py']*np.cos(df.loc[:, 'aug_angle'])+df.loc[:, f'{vec}_px']*np.sin(df.loc[:, 'aug_angle'])
            df.loc[:, f'{vec}_px']    = df.loc[:, f'{vec}_pxtmp']

    def _reflect(self, df:pd.DataFrame, vectors:List[str]) -> None:
        for vector in vectors:
            for coord in self.reflect_axes:
                try:
                    cut = (df[f'aug{coord}'] == 1)
                    df.loc[cut, f'{vector}{coord}'] = -df.loc[cut, f'{vector}{coord}']
                except KeyError:
                    pass

    def get_fold(self, idx:int) -> Dict[str,np.ndarray]:
        r'''
        Get data for single fold applying random train-time data augmentaion. Data consists of dictionary of inputs, targets, and weights.
        Accounts for ignored features.
        Inputs, except for matrix data, are passed through np.nan_to_num to deal with nans and infs.

        Arguments:
            idx: fold index to load

        Returns:
            tuple of inputs, targets, and weights as Numpy arrays
        '''

        data = self.get_data(n_folds=1, fold_idx=idx)
        if not self.augmented: return data
        inputs = pd.DataFrame(self.foldfile[f'fold_{idx}/inputs'][()], columns=self.input_feats)
        if self.targ_feats is not None: targets = pd.DataFrame(self.foldfile[f'fold_{idx}/targets'][()], columns=self.targ_feats)

        if self.rot_mult:
            inputs['aug_angle'] = (2*np.pi*np.random.random(size=len(inputs)))-np.pi
            self._rotate(inputs, self.vectors)
            if self.targ_feats is not None:
                targets['aug_angle'] = inputs['aug_angle']
                self._rotate(targets, self.targ_vectors)

        for coord in self.reflect_axes:
            inputs[f'aug{coord}'] = np.random.randint(0, 2, size=len(inputs))
            if self.targ_feats is not None: targets[f'aug{coord}'] = inputs[f'aug{coord}']
        self._reflect(inputs, self.vectors)
        if self.targ_feats is not None: self._reflect(targets, self.targ_vectors)

        inputs = inputs[[f for f in self.input_feats if f not in self._ignore_feats]]
        data['inputs'] = np.nan_to_num(inputs.values)
        if self.targ_feats is not None:
            targets = targets[self.targ_feats]
            data['targets'] = np.nan_to_num(targets.values)
        return self._append_matrix(data, idx) if self.has_matrix and self.yield_matrix else data

    def _get_ref_idx(self, aug_idx:int) -> str:
        n_axes = len(self.reflect_axes)
        div = self.rot_mult if self.rot_mult else 1
        if   n_axes == 3: return '{0:03b}'.format(int(aug_idx/div))
        elif n_axes == 2: return '{0:02b}'.format(int(aug_idx/div))
        elif n_axes == 1: return '{0:01b}'.format(int(aug_idx/div))

    def get_test_fold(self, idx:int, aug_idx:int) -> Dict[str, np.ndarray]:
        r'''
        Get test data for single fold applying test-time data augmentaion. Data consists of dictionary of inputs, targets, and weights.
        Accounts for ignored features.
        Inputs, except for matrix data, are passed through np.nan_to_num to deal with nans and infs.

        Arguments:
            idx: fold index to load
            aug_idx: index for the test-time augmentaion (ignored if random test-time augmentation requested)

        Returns:
            tuple of inputs, targets, and weights as Numpy arrays
        '''

        if aug_idx >= self.aug_mult: raise ValueError(f"Invalid augmentation idx passed {aug_idx}")
        data = self.get_data(n_folds=1, fold_idx=idx)
        if not self.augmented: return data

        inputs = pd.DataFrame(self.foldfile[f'fold_{idx}/inputs'][()], columns=self.input_feats)
        if len(self.reflect_axes) > 0 and self.rot_mult > 0:
            rot_idx = aug_idx % self.rot_mult
            ref_idx = self._get_ref_idx(aug_idx)
            if self.random_rot: inputs['aug_angle'] = (2*np.pi*np.random.random(size=len(inputs)))-np.pi
            else:               inputs['aug_angle'] = np.linspace(0, 2*np.pi, (self.rot_mult)+1)[rot_idx]
            self._rotate(inputs, self.vectors)

            for i, coord in enumerate(self.reflect_axes): inputs[f'aug{coord}'] = int(ref_idx[i])
            self._reflect(inputs, self.vectors)

        elif len(self.reflect_axes) > 0:
            ref_idx = self._get_ref_idx(aug_idx)
            for i, coord in enumerate(self.reflect_axes): inputs[f'aug{coord}'] = int(ref_idx[i])
            self._reflect(inputs, self.vectors)

        elif self.rot_mult:
            if self.random_rot: inputs['aug_angle'] = (2*np.pi*np.random.random(size=len(inputs)))-np.pi
            else:               inputs['aug_angle'] = np.linspace(0, 2*np.pi, (self.rot_mult)+1)[aug_idx]
            self._rotate(inputs, self.vectors)

        inputs = inputs[[f for f in self.input_feats if f not in self._ignore_feats]]
        data['inputs'] = np.nan_to_num(inputs.values)

        if self.targ_feats is not None:
            targets = pd.DataFrame(self.foldfile[f'fold_{idx}/targets'][()], columns=self.targ_feats)
            if len(self.reflect_axes) > 0 and self.rot_mult > 0:
                rot_idx = aug_idx % self.rot_mult
                ref_idx = self._get_ref_idx(aug_idx)
                if self.random_rot: targets['aug_angle'] = (2*np.pi*np.random.random(size=len(targets)))-np.pi
                else:               targets['aug_angle'] = np.linspace(0, 2*np.pi, (self.rot_mult)+1)[rot_idx]
                self._rotate(targets, self.targ_vectors)

                for i, coord in enumerate(self.reflect_axes): targets[f'aug{coord}'] = int(ref_idx[i])
                self._reflect(targets, self.targ_vectors)

            elif len(self.reflect_axes) > 0:
                ref_idx = self._get_ref_idx(aug_idx)
                for i, coord in enumerate(self.reflect_axes): targets[f'aug{coord}'] = int(ref_idx[i])
                self._reflect(targets, self.targ_vectors)

            elif self.rot_mult:
                if self.random_rot: targets['aug_angle'] = (2*np.pi*np.random.random(size=len(targets)))-np.pi
                else:               targets['aug_angle'] = np.linspace(0, 2*np.pi, (self.rot_mult)+1)[aug_idx]
                self._rotate(targets, self.targ_vectors)

            targets = targets[self.targ_feats]
            data['targets'] = np.nan_to_num(targets.values)
        return self._append_matrix(data, idx) if self.has_matrix and self.yield_matrix else data

max_stars_repo_name: GilesStrong/lumin

max_issues_repo_name: GilesStrong/lumin

max_forks_repo_name: GilesStrong/lumin


text:
from keras import initializers, layers
import keras.backend as K
import tensorflow as tf
import numpy as np
from keras.utils.conv_utils import conv_output_length
from ops import update_routing

class ConvCapsuleLayer(layers.Layer):
    def __init__(self, kernel_size, num_caps, caps_dim, strides=1, padding='same', routings=3,
                 kernel_initializer='he_normal', **kwargs):
        super(ConvCapsuleLayer, self).__init__(**kwargs)
        self.kernel_size = kernel_size
        self.num_caps = num_caps
        self.caps_dim = caps_dim
        self.strides = strides
        self.padding = padding
        self.routings = routings
        self.kernel_initializer = initializers.get(kernel_initializer)

    def build(self, input_shape):
        assert len(input_shape) == 5, "The input Tensor should have shape=[None, input_height, input_width," \
                                      " input_num_capsule, input_num_atoms]"
        self.input_height = input_shape[1]
        self.input_width = input_shape[2]
        self.num_in_caps = input_shape[3]
        self.in_caps_dim = input_shape[4]

        # Transform matrix
        self.W = self.add_weight(shape=[self.kernel_size, self.kernel_size,
                                        self.in_caps_dim, self.num_caps * self.caps_dim],
                                 initializer=self.kernel_initializer,
                                 name='W')

        self.b = self.add_weight(shape=[1, 1, self.num_caps, self.caps_dim],
                                 initializer=initializers.constant(0.1),
                                 name='b')

        self.built = True

    def call(self, input_tensor, training=None):
        input_transposed = tf.transpose(input_tensor, [3, 0, 1, 2, 4])
        input_shape = K.shape(input_transposed)
        input_tensor_reshaped = K.reshape(input_transposed, [
            input_shape[0] * input_shape[1], self.input_height, self.input_width, self.in_caps_dim])
        input_tensor_reshaped.set_shape((None, self.input_height, self.input_width, self.in_caps_dim))

        conv = K.conv2d(input_tensor_reshaped, self.W, (self.strides, self.strides),
                        padding=self.padding, data_format='channels_last')

        votes_shape = K.shape(conv)
        _, conv_height, conv_width, _ = conv.get_shape()

        votes = K.reshape(conv, [input_shape[1], input_shape[0], votes_shape[1], votes_shape[2],
                                 self.num_caps, self.caps_dim])
        votes.set_shape((None, self.num_in_caps, conv_height.value, conv_width.value,
                         self.num_caps, self.caps_dim))

        logit_shape = K.stack([input_shape[1], input_shape[0], votes_shape[1], votes_shape[2], self.num_caps])
        biases_replicated = K.tile(self.b, [conv_height.value, conv_width.value, 1, 1])

        activations = update_routing(votes=votes,
                                     biases=biases_replicated,
                                     logit_shape=logit_shape,
                                     num_dims=6,
                                     input_dim=self.num_in_caps,
                                     output_dim=self.num_caps,
                                     num_routing=self.routings)

        return activations

    def compute_output_shape(self, input_shape):
        space = input_shape[1:-2]
        new_space = []
        for i in range(len(space)):
            new_dim = conv_output_length(
                space[i],
                self.kernel_size,
                padding=self.padding,
                stride=self.strides,
                dilation=1)
            new_space.append(new_dim)

        return (input_shape[0],) + tuple(new_space) + (self.num_caps, self.caps_dim)

    def get_config(self):
        config = {'kernel_size': self.kernel_size,
                  'num_capsule': self.num_caps,
                  'num_atoms': self.caps_dim,
                  'strides': self.strides,
                  'padding': self.padding,
                  'routings': self.routings,
                  'kernel_initializer': initializers.serialize(self.kernel_initializer)}
        base_config = super(ConvCapsuleLayer, self).get_config()
        return dict(list(base_config.items()) + list(config.items()))

ext: py

max_stars_repo_name: mawanda-jun/NoLabels

max_issues_repo_name: mawanda-jun/NoLabels

max_forks_repo_name: mawanda-jun/NoLabels

reason: import numpy


text:
import numpy as np
import pandas as pd
import pytest
from raga_pose_estimation.openpose_parts import OpenPoseParts
from raga_pose_estimation.reshaper import reshape_dataframes

def _get_dummy_dataframes():
    data = [
        {"x": [2, 15], "y": [8, 5], "confidence": [0.1, 0.8],},
        {"x": [3, 10], "y": [5, 4], "confidence": [0.5, 0.9],},
        {"x": [4, np.nan], "y": [4, np.nan], "confidence": [0.7, np.nan],},
    ]
    dataframes = []
    for d in data:
        dataframes.append(
            pd.DataFrame(
                d,
                columns=["x", "y", "confidence"],
                index=pd.Index(
                    [OpenPoseParts.R_EAR.value, OpenPoseParts.R_EYE.value,]
                ),
            )
        )
    return dataframes

@pytest.fixture
def dummy_dataframes():
    return _get_dummy_dataframes()

@pytest.fixture
def single_frame_person_df():
    return reshape_dataframes([_get_dummy_dataframes()[0]])

@pytest.fixture
def three_frame_person_dfs():
    return reshape_dataframes(_get_dummy_dataframes())


text:
import time
# only required to run python3 examples/cvt_arm.py
import torch
import torch.nn as nn
import torch.nn.functional as F
import torchvision
from torchvision import datasets, transforms
from torch.utils.data import Dataset

import os
import math
import numpy as np

device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
batch_size = 64 # batch size in every epoch

class CustomDataset(Dataset):
    def __init__(self, data, targets, transform=None, target_transform=None):
        self.data = data
        self.targets = targets
        self.transform = transform
        self.target_transform = target_transform

    def __len__(self):
        return len(self.targets)

    def __getitem__(self, idx):

        img, target = self.data[idx], self.targets[idx]

        if self.transform is not None:
            img = self.transform(img)

        if self.target_transform is not None:
            target = self.target_transform(target)

        return img, target

train_dataset = datasets.MNIST(root = 'data/', train=True, download=False,
                    transform=transforms.Compose([
                        transforms.ToTensor(),
                        transforms.Normalize((0.1307,), (0.3081,))
                    ]))

# training set
train_loader = torch.utils.data.DataLoader(train_dataset,
                                           batch_size=batch_size,
                                           shuffle=True,
                                           num_workers=4,
                                           drop_last=True)

test_dataset =  datasets.MNIST(root = 'data/', train=False, download=False,
                    transform=transforms.Compose([
                        transforms.ToTensor(),
                        transforms.Normalize((0.1307,), (0.3081,))
                    ]))

# test set
test_loader = torch.utils.data.DataLoader(test_dataset,
                                          batch_size=100,
                                          shuffle=True,
                                          num_workers=4,
                                          drop_last=False)

inv_adv_examples = np.load("invariance_examples/final_l0/inv_adv_examples.npy") # visualize this for sanity check
human_labels = np.load("invariance_examples/final_l0/human_labels.npy")

inv_eg_dataset = CustomDataset(data=inv_adv_examples,
                               targets=human_labels,
                               transform=transforms.Compose([
                        transforms.ToTensor(),
                        transforms.Normalize((0.1307,), (0.3081,))
                    ]))

inv_eg_dataloader = torch.utils.data.DataLoader(inv_eg_dataset,
                                                batch_size=10,
                                                shuffle=True,
                                                num_workers=4,
                                                drop_last=False)

class MLP(nn.Module):
    def __init__(self, input_size=784, output_size=10):
        super(MLP, self).__init__()
        self.fc1 = nn.Linear(input_size, 20)
        self.fc2 = nn.Linear(20, output_size)

    def forward(self, x):
        x = x.view(-1, 28 * 28)
        x = F.relu(self.fc1(x))
        x = F.relu(self.fc2(x))
        return F.log_softmax(x, dim=1)

model = MLP()
model = model.to(device)

epochs = 100

# specify loss function
criterion = nn.CrossEntropyLoss()

# specify optimizer
optimizer = torch.optim.SGD(model.parameters(), lr=0.1)
model.train() # prep model for training

standard_acc_arr = []
robust_acc_arr = []

for epoch in range(epochs):
    # monitor training loss
    train_loss = 0.0

    ###################
    # train the model #
    ###################

    for data, target in train_loader:
        data, target = data.to(device), target.to(device)

        # clear the gradients of all optimized variables
        optimizer.zero_grad()
        # forward pass: compute predicted outputs by passing inputs to the model
        output = model(data)
        # calculate the loss
        loss = criterion(output, target)
        # backward pass: compute gradient of the loss with respect to model parameters
        loss.backward()
        # perform a single optimization step (parameter update)
        optimizer.step()
        # update running training loss
        train_loss += loss.item()*data.size(0)

    #--------------------------------------------------------------------------------------
    # initialize lists to monitor test loss and accuracy
    class_correct = list(0. for i in range(10))
    class_total = list(0. for i in range(10))

    model.eval() # prep model for *evaluation*

    for data, target in test_loader:
        data, target = data.to(device), target.to(device)

        # forward pass: compute predicted outputs by passing inputs to the model
        output = model(data)

        # convert output probabilities to predicted class
        _, pred = torch.max(output, 1)

        # compare predictions to true label
        correct = np.squeeze(pred.eq(target.data.view_as(pred)))
        # calculate test accuracy for each object class
        for i in range(len(target)):
            label = target.data[i]
            class_correct[label] += correct[i].item()
            class_total[label] += 1

    test_accuracy = 100. * np.sum(class_correct) / np.sum(class_total)
    standard_acc_arr.append(test_accuracy)
    print('\n Standard test accuracy: %2d%% (%2d/%2d)' % (test_accuracy,
                                                          np.sum(class_correct),
                                                          np.sum(class_total)))
    #--------------------------------------------------------------------------------------

    #--------------------------------------------------------------------------------------
    # initialize lists to monitor test loss and accuracy
    class_correct = list(0. for i in range(10))
    class_total = list(0. for i in range(10))

    model.eval() # prep model for *evaluation*

    for data, target in inv_eg_dataloader:
        data, target = data.to(device), target.to(device)

        # forward pass: compute predicted outputs by passing inputs to the model
        output = model(data)

        # convert output probabilities to predicted class
        _, pred = torch.max(output, 1)

        # compare predictions to true label
        correct = np.squeeze(pred.eq(target.data.view_as(pred)))
        # calculate test accuracy for each object class
        for i in range(len(target)):
            label = target.data[i]
            class_correct[label] += correct[i].item()
            class_total[label] += 1

    test_accuracy = 100. * np.sum(class_correct) / np.sum(class_total)
    robust_acc_arr.append(test_accuracy)
    print('\n Robust test accuracy: %2d%% (%2d/%2d)' % (test_accuracy,
                                                        np.sum(class_correct),
                                                        np.sum(class_total)))
    #--------------------------------------------------------------------------------------

save_path = "saved_stuff/eval_mlp/"
if not os.path.exists(save_path):
    os.makedirs(save_path)

torch.save(standard_acc_arr, os.path.join(save_path, "standard_acc_arr.pt"))
torch.save(robust_acc_arr, os.path.join(save_path, "robust_acc_arr.pt"))

max_stars_repo_name: greninja/RobustNNSearch

max_issues_repo_name: greninja/RobustNNSearch

max_forks_repo_name: greninja/RobustNNSearch


text:
import numpy as np

def mean_average_precision(rs):
    """Score is mean average precision
    Relevance is binary (nonzero is relevant).
    >>> rs = [[1, 1, 0, 1, 0, 1, 0, 0, 0, 1]]
    >>> mean_average_precision(rs)
    0.78333333333333333
    >>> rs = [[1, 1, 0, 1, 0, 1, 0, 0, 0, 1], [0]]
    >>> mean_average_precision(rs)
    0.39166666666666666
    Args:
        rs: Iterator of relevance scores (list or numpy) in rank order
            (first element is the first item)
    Returns:
        Mean average precision
    """
    return np.mean([average_precision(r) for r in rs])

def average_precision(r):
    """Score is average precision (area under PR curve)
    Relevance is binary (nonzero is relevant).
    >>> r = [1, 1, 0, 1, 0, 1, 0, 0, 0, 1]
    >>> delta_r = 1. / sum(r)
    >>> sum([sum(r[:x + 1]) / (x + 1.) * delta_r for x, y in enumerate(r) if y])
    0.7833333333333333
    >>> average_precision(r)
    0.78333333333333333
    Args:
        r: Relevance scores (list or numpy) in rank order
            (first element is the first item)
    Returns:
        Average precision
    """
    r = np.asarray(r) != 0
    out = [precision_at_k(r, k + 1) for k in range(r.size) if r[k]]
    if not out:
        return 0.
    return np.mean(out)

def precision_at_k(r, k):
    """Score is precision @ k
    Relevance is binary (nonzero is relevant).
    >>> r = [0, 0, 1]
    >>> precision_at_k(r, 1)
    0.0
    >>> precision_at_k(r, 2)
    0.0
    >>> precision_at_k(r, 3)
    0.33333333333333331
    >>> precision_at_k(r, 4)
    Traceback (most recent call last):
        File "<stdin>", line 1, in ?
    ValueError: Relevance score length < k
    Args:
        r: Relevance scores (list or numpy) in rank order
            (first element is the first item)
    Returns:
        Precision @ k
    Raises:
        ValueError: len(r) must be >= k
    """
    assert k >= 1
    r = np.asarray(r)[:k] != 0
    if r.size != k:
        raise ValueError('Relevance score length < k')
    return np.mean(r)

def dcg_at_k(relevance, topics):

    scores = np.zeros(relevance.shape[0], dtype=np.float)
    for u, top in enumerate(relevance):
        tops = set([])

        for i, item in enumerate(top):
            wi = 1/(2**i)
            if topics[item] is not None:
                scores[u] += len(set(topics[item]).difference(tops))*wi
                tops.update(topics[item])
    return scores

def dcg_at_k_with_importance(relevance, topics, importance):

    scores = np.zeros(relevance.shape[0], dtype=np.float)
    for u, top in enumerate(relevance):
        tops = []
        for i, item in enumerate(top):
            wi = 1/(2**i)
            if topics[item] is not None:
                scores[u] += (importance[item]**tops.count(topics[item])) * wi
                tops.append(topics[item])
    return scores


text:
import numpy as np
import os
from sklearn.neighbors import NearestNeighbors

from pydrake.multibody.rigid_body import RigidBody
from pydrake.all import (
        AddFlatTerrainToWorld,
        AddModelInstancesFromSdfString,
        AddModelInstanceFromUrdfFile,
        FindResourceOrThrow,
        FloatingBaseType,
        InputPort,
        Isometry3,
        OutputPort,
        RgbdCamera,
        RigidBodyPlant,
        RigidBodyTree,
        RigidBodyFrame,
        RollPitchYaw,
        RollPitchYawFloatingJoint,
        RotationMatrix,
        Value,
        VisualElement,
    )

import meshcat
import meshcat.transformations as tf
import meshcat.geometry as g

# From
# https://www.opengl.org/discussion_boards/showthread.php/197893-View-and-Perspective-matrices
def normalize(x):
    return x / np.linalg.norm(x)

def save_pointcloud(pc, normals, path):
    joined = np.hstack([pc.T, normals.T])
    np.savetxt(path, joined)

def load_pointcloud(path):
    joined = np.loadtxt(path)
    return joined[:, 0:3].T, joined[:, 3:6].T

def translate(x):
    T = np.eye(4)
    T[0:3, 3] = x[:3]
    return T

def get_pose_error(tf_1, tf_2):
    rel_tf = transform_inverse(tf_1).dot(tf_2)
    if np.allclose(np.diag(rel_tf[0:3, 0:3]), [1., 1., 1.]):
        angle_dist = 0.
    else:
        # Angle from rotation matrix
        angle_dist = np.arccos(
            (np.sum(np.diag(rel_tf[0:3, 0:3])) - 1) / 2.)
    euclid_dist = np.linalg.norm(rel_tf[0:3, 3])
    return euclid_dist, angle_dist

# If misalignment_tol = None, returns the average
# distance between the model clouds when transformed
# by est_tf and gt_tf (using nearest-point lookups
# for each point in the gt-tf'd model cloud).
# If misalignment_tol is a number, it returns
# the percent of points that are misaligned by more
# than the misalignment error under the same distance
# metric.
def get_earth_movers_error(est_tf, gt_tf, model_cloud,
                           misalignment_tol=0.005):
    # Transform the model cloud into both frames
    est_model_cloud = transform_points(est_tf, model_cloud)
    gt_model_cloud = transform_points(gt_tf, model_cloud)
    # For every point in the model cloud, find the distance
    # to the closest point in the estimated model cloud,
    # as a way of finding the swept volume between the
    # models in those poses.
    neigh = NearestNeighbors(n_neighbors=1)
    neigh.fit(gt_model_cloud.T)
    dist, _ = neigh.kneighbors(
        est_model_cloud[0:3, :].T, return_distance=True)
    if misalignment_tol is None:
        return np.mean(dist)
    else:
        return np.mean(dist > misalignment_tol)

def draw_points(vis, vis_prefix, name, points,
                normals=None, colors=None, size=0.001,
                normals_length=0.01):
    vis[vis_prefix][name].set_object(
        g.PointCloud(position=points,
                     color=colors,
                     size=size))
    n_pts = points.shape[1]
    if normals is not None:
        # Drawing normals for debug
        lines = np.zeros([3, n_pts*2])
        inds = np.array(range(0, n_pts*2, 2))
        lines[:, inds] = points[0:3, :]
        lines[:, inds+1] = points[0:3, :] + \
            normals * normals_length
        vis[vis_prefix]["%s_normals" % name].set_object(
            meshcat.geometry.LineSegmentsGeometry(
                lines, None))

def transform_points(tf, pts):
    return ((tf[:3, :3].dot(pts).T) + tf[:3, 3]).T

def transform_inverse(tf):
    new_tf = np.eye(4)
    new_tf[:3, :3] = tf[:3, :3].T
    new_tf[:3, 3] = -new_tf[:3, :3].dot(tf[:3, 3])
    return new_tf

def lookat(eye, target, up):
    # For a camera with +x right, +y down, and +z forward.
    eye = np.array(eye)
    target = np.array(target)
    up = np.array(up)
    F = target[:3] - eye[:3]
    f = normalize(F)
    U = normalize(up[:3])
    s = np.cross(f, U)  # right
    u = np.cross(s, f)  # up
    M = np.eye(4)
    M[:3, :3] = np.vstack([s, -u, f]).T

    # OLD:
    # flip z -> x
    # -x -> y
    # -y -> z
    # CAMERA FORWARD is +x-axis
    # CAMERA RIGHT is -y axis
    # CAMERA UP is +z axis
    # Why does the Drake documentation lie to me???
    T = translate(eye)
    return T.dot(M)

def add_single_instance_to_rbt(
        rbt, config, instance_config, i,
        floating_base_type=FloatingBaseType.kRollPitchYaw):
    class_name = instance_config["class"]
    if class_name not in config["objects"].keys():
        raise ValueError("Class %s not in classes." % class_name)
    if len(instance_config["pose"]) != 6:
        raise ValueError("Class %s has pose size != 6. Use RPY plz" %
                         class_name)
    frame = RigidBodyFrame(
        "%s_%d" % (class_name, i), rbt.world(),
        instance_config["pose"][0:3],
        instance_config["pose"][3:6])
    model_path = config["objects"][class_name]["model_path"]
    _, extension = os.path.splitext(model_path)
    if extension == ".urdf":
        AddModelInstanceFromUrdfFile(
            model_path, floating_base_type, frame, rbt)
    elif extension == ".sdf":
        AddModelInstancesFromSdfString(
            open(model_path).read(), floating_base_type, frame, rbt)
    else:
        raise ValueError("Class %s has non-sdf and non-urdf model name." %
                         class_name)

def setup_scene(rbt, config):
    if config["with_ground"] is True:
        AddFlatTerrainToWorld(rbt)

    for i, instance_config in enumerate(config["instances"]):
        add_single_instance_to_rbt(rbt, config, instance_config, i,
                                   floating_base_type=FloatingBaseType.kFixed)
    # Add camera geometry!
    camera_link = RigidBody()
    camera_link.set_name("camera_link")
    # necessary so this last link isn't pruned by the rbt.compile() call
    camera_link.set_spatial_inertia(np.eye(6))
    camera_link.add_joint(
        rbt.world(),
        RollPitchYawFloatingJoint(
            "camera_floating_base",
            np.eye(4)))
    rbt.add_rigid_body(camera_link)

    # - Add frame for camera fixture.
    camera_frame = RigidBodyFrame(
        name="rgbd_camera_frame", body=camera_link,
        xyz=[0.0, 0., 0.], rpy=[0., 0., 0.])
    rbt.addFrame(camera_frame)
    rbt.compile()


text:
from django.shortcuts import render
from rest_framework.permissions import IsAuthenticated,AllowAny
from .serializers import WaterLevelSerializer, FetchWaterLevelSerializer
from rest_framework.generics import ListAPIView,CreateAPIView
from dashboard.models import WaterLevel
from rest_framework.views import APIView
from decimal import *
import uuid
from rest_framework.response import Response
from rest_framework import status

class PredictWaterLevel(APIView):
    permission_classes = (AllowAny,)
    serializer_class = WaterLevelSerializer

    def post(self, request):
        serializer = WaterLevelSerializer(data=request.data)
        serializer.is_valid(raise_exception=True)
        temp = request.data['temp']
        humidity = request.data['humidity']
        pressure = request.data['pressure']
        id=uuid.uuid4();
        serializer.save(id=id,waterlevel=self.trainmodel(temp,humidity,pressure))
        serializer.save(waterlevel=self.trainmodel(temp,humidity,pressure))
        return Response(serializer.data,status=status.HTTP_201_CREATED)

    def trainmodel(self,temp,humidity,pressure):

    # coding: utf-8

    # # Water Crisis

    # In[1]:

        import pandas as pd

        # In[2]:

        import numpy as np

        # In[3]:

        df=pd.read_csv("/home/codeforvision/weatherHistory.csv")

        # In[4]:

        a=np.random.randint(20,8000,96453)

        # In[5]:

        b=pd.DataFrame(a)

        # In[6]:

        df1=pd.concat([df,b],axis=1)

        # In[7]:

        df1.rename(columns={0: "water level quantity (cm^3)"})

        # In[8]:

        features=df1.iloc[:5000,[3,5,10]].values
        labels=df1.iloc[:5000,-1].values

        # In[9]:

        #Spliting the dataset into training and testing set

        from sklearn.model_selection import train_test_split
        x_train,x_test,y_train,y_test=train_test_split(features,labels,test_size=0.10,random_state=0)

        # In[10]:

        #Features scaling
        from sklearn.preprocessing import StandardScaler
        sc=StandardScaler()
        x_train=sc.fit_transform(x_train)
        x_test=sc.transform(x_test)

        # In[11]:

        #Fitting the multiple linear regression to training set
        from sklearn.linear_model import LinearRegression
        regressor=LinearRegression()
        regressor.fit(x_train,y_train)

        # In[12]:

        y=regressor.predict([[int(temp),int(humidity),int(pressure)]])
        return int(y[0])

class FetchWaterLevel(ListAPIView):
    permission_classes = (AllowAny,)
    serializer_class = FetchWaterLevelSerializer

    def get_queryset(self):

        return WaterLevel.objects.filter(id=self.request.POST.get('id'))


text:
from __future__ import print_function

import numpy
import irlib

# Compute basis functions very accurately

max_dim = 1000
cutoff = 1e-8
r_tol = 1e-8
verbose = True
prec = 64
n_local_poly = 20
n_gl_node = 24

## Construct basis
statis = 'F'
for Lambda in [10.0]:
    print("Computing basis functions... It may take some time")
    if statis == 'F':
        b = irlib.compute_basis(irlib.FERMIONIC, Lambda, max_dim, cutoff, "mp", r_tol, prec, n_local_poly, n_gl_node, verbose)
        irlib.savetxt("basis_f-mp-Lambda"+str(Lambda)+".txt", b)
    else:
        b = irlib.compute_basis(irlib.BOSONIC, Lambda, max_dim, cutoff, "mp", r_tol, prec, n_local_poly, n_gl_node, verbose)
        irlib.savetxt("basis_b-mp-Lambda"+str(Lambda)+".txt", b)
    print("Done!")

max_stars_repo_name: dombrno/irlib

max_issues_repo_name: dombrno/irlib

max_forks_repo_name: dombrno/irlib


text:
import numpy as np
import math

"""
 ==================================
 Problem 3: Neural Network Basics
 ==================================
    Generates a neural network with the following architecture:
        Fully connected neural network.
        Input vector takes in two features.
        One hidden layer with three neurons whose activation function is ReLU.
        One output neuron whose activation function is the identity function.
"""

def rectified_linear_unit(x):
    """ Returns the ReLU of x, or the maximum between 0 and x."""
    return x*(x > 0)

def rectified_linear_unit_derivative(x):
    """ Returns the derivative of ReLU."""
    ret = 1 if x>0 else 0
    return ret

def output_layer_activation(x):
    """ Linear function, returns input as is. """
    return x

def output_layer_activation_derivative(x):
    """ Returns the derivative of a linear function: 1. """
    return 1

class NeuralNetwork():
    """
        Contains the following functions:
            -train: tunes parameters of the neural network based on error obtained from forward propagation.
            -predict: predicts the label of a feature vector based on the class's parameters.
            -train_neural_network: trains a neural network over all the data points for the specified number of epochs during initialization of the class.
            -test_neural_network: uses the parameters specified at the time in order to test that the neural network classifies the points given in testing_points within a margin of error.
    """

    def __init__(self):

        # DO NOT CHANGE PARAMETERS
        self.input_to_hidden_weights = np.matrix('1 1; 1 1; 1 1')
        self.hidden_to_output_weights = np.matrix('1 1 1')
        self.biases = np.matrix('0; 0; 0')
        self.learning_rate = .001
        self.epochs_to_train = 1
        self.training_points = [((2,1), 10), ((3,3), 21), ((4,5), 32), ((6, 6), 42)]
        # self.training_points=[((-7, -9), -16), ((7, 7), 14), ((-2, -5), -7), ((0, 5), 5), ((3, 8), 11), ((-9, 2), -7), ((-4, -7), -11), ((9, -6), 3), ((6, 2), 8), ((-5, 7), 2)]
        self.testing_points = [(1,1), (2,2), (3,3), (5,5), (10,10)]

    def train(self, x1, x2, y):

        ### Forward propagation ###
        input_values = np.matrix([[x1],[x2]]) # 2 by 1
        x=input_values
        w1=self.input_to_hidden_weights
        w2=self.hidden_to_output_weights
        b=self.biases
        f=output_layer_activation
        z1=w1*x+b
        tmp=np.vectorize(rectified_linear_unit)
        a1=tmp(z1)
        z2=np.multiply(a1.T,w2)
        o=f(z2)
        fprim=np.vectorize(output_layer_activation_derivative)
        E0=np.multiply((y-o),fprim(z2))
        tmprim=np.vectorize(rectified_linear_unit_derivative)
        Eh=np.multiply(np.multiply(w2,E0),tmprim(z1).T)
        deltab=Eh
        deltaw2=np.multiply(E0,a1.T)
        deltaw1=np.multiply(Eh,x).sum(axis=0)

        bnew=b-np.multiply(self.learning_rate,deltab.T)
        w2new=w2-self.learning_rate*deltaw2
        w1new=w1-self.learning_rate*deltaw1.T
        self.biases = bnew
        self.input_to_hidden_weights = w1new
        self.hidden_to_output_weights=w2new

        # Calculate the input and activation of the hidden layer

# t=1
# x=3
# w1=0.01
# w2=-5
# b=-1
# z1=w1*x
# a1=ReLU(z1)
# z2=w2*a1+b
# y=sigmoid(z2)
# c=0.5*(float(y)-t)*(float(y)-t)
        hidden_layer_weighted_input = np.dot(self.input_to_hidden_weights,input_values)+ self.biases
        # print("hidden_weighted:%s"%str(hidden_layer_weighted_input.shape))
        # (3 by 1 matrix)
        gx=np.vectorize(rectified_linear_unit)
        hidden_layer_activation = gx(hidden_layer_weighted_input)
        # print("hidden_active:%s"%str(hidden_layer_activation.shape))

        output =   np.dot(self.hidden_to_output_weights,hidden_layer_activation)
        activated_output = output_layer_activation(output)
        # print("active_out:%s"%str(activated_output.shape))
        ### Backpropagation ###

        # Compute gradients
        # lossy=1/2*(activated_output-y).T*(activated_output-y)
        dLoss_Yh=y-activated_output
        g = np.vectorize(output_layer_activation_derivative)

# ========================================
        dLoss_Z2 = dLoss_Yh * g(activated_output)
        dLoss_A1 = np.dot(self.hidden_to_output_weights.T,dLoss_Z2)
        dLoss_W2 = 1./activated_output.shape[1] * np.dot(dLoss_Z2, activated_output.T)
        dLoss_b2 = 1./ activated_output.shape[1] * np.dot(dLoss_Z2, np.ones([dLoss_Z2.shape[1],1]))

        g2 = np.vectorize(rectified_linear_unit_derivative)
        dLoss_Z1 = np.dot(dLoss_A1.T ,g2(hidden_layer_weighted_input))
        # dLoss_A0 = np.dot(self.input_to_hidden_weights.T,dLoss_Z1)
        dLoss_W1 = 1./input_values.shape[1] * np.dot(dLoss_Z1,input_values.T)
        dLoss_b1 = 1./input_values.shape[1] * np.dot(dLoss_Z1, np.ones([dLoss_Z1.shape[1],1]))
        # print(str(dLoss_W1))
        # print(str(dLoss_b1))
        # print("okejko")

# ===============================================

        # dy_z2=g(activated_output)
        # print("dlossy:%s"%str(dlossy.shape))

        # dlossZ2=np.dot(dlossy,dy_z2)
        # print("dlossZ2:%s"%str(dlossZ2.shape))
        # dlossA1=np.dot(self.hidden_to_output_weights.T,dlossZ2)
        # dlossW2=1./hidden_layer_activation.shape[1]*np.dot(dlossZ2,hidden_layer_activation.T)
        # print("dlossW2:%s"%str(dlossW2.shape))

        # g2 = np.vectorize(rectified_linear_unit_derivative)
        # da1_z1=g2(hidden_layer_weighted_input)
        # dlossZ1=np.dot(dlossA1,da1_z1)

        # print("hidden_layer_weighted_input:%s"%str(hidden_layer_weighted_input.shape))
        # print("da1_z1:%s"%str(da1_z1.shape))
        # print("dlossA1:%s"%str(dlossA1.shape))

        # tstsas=1./input_values.shape[1] * np.dot(dlossZ1,input_values.T)
        # print("tstsas:%s"%str(tstsas.shape))
        # print("dlossZ1:%s"%str(dlossZ1.shape))
        # dlossA0=np.dot(self.input_to_hidden_weights.T,dlossZ1)
        # dz1_w1=input_values

        # dlossW1=np.dot(tstsas,input_values.T)
        # dlossB1=np.dot(tstsas,np.ones([dlossZ1.shape[1],1]))

        # Use gradients to adjust weights and biases using gradient descent
        self.biases = self.biases-self.learning_rate*dLoss_b1
        self.input_to_hidden_weights = self.input_to_hidden_weights - self.learning_rate*dLoss_W1
        self.hidden_to_output_weights = self.hidden_to_output_weights - self.learning_rate*dLoss_W2
        # print("passx")

    # def predict(self, x1, x2):

    #     input_values = np.matrix([[x1],[x2]])

    #     # Compute output for a single input(should be same as the forward propagation in training)
    #     hidden_layer_weighted_input = # TODO
    #     hidden_layer_activation = # TODO
    #     output = # TODO
    #     activated_output = # TODO

    #     return activated_output.item()

    # Run this to train your neural network once you complete the train method
    def train_neural_network(self):

        for epoch in range(self.epochs_to_train):
            for x,y in self.training_points:
                self.train(x[0], x[1], y)
            print("bias")
            print(self.biases)
            print("w2")
            print(self.hidden_to_output_weights)
            print("w1")
            print(self.input_to_hidden_weights)

    # Run this to test your neural network implementation for correctness after it is trained
    def test_neural_network(self):

        for point in self.testing_points:
            print("Point,", point, "Prediction,", self.predict(point[0], point[1]))
            if abs(self.predict(point[0], point[1]) - 7*point[0]) < 0.1:
                print("Test Passed")
            else:
                print("Point ", point[0], point[1], " failed to be predicted correctly.")
                return

x = NeuralNetwork()

x.train_neural_network()

# UNCOMMENT THE LINE BELOW TO TEST YOUR NEURAL NETWORK
# x.test_neural_network()

max_stars_repo_name: stepinski/machinelearning

max_issues_repo_name: stepinski/machinelearning

max_forks_repo_name: stepinski/machinelearning


text:
#!/usr/bin/python

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

from consts import *
import numpy as np
import scipy as sp
from calc_thd import *
import calc_delay

# calculate THD for dut_playback_thd case
# Input: host recording (mono),
#        frequency of sine in Hz (i64)
#        THD pass level in percentile (double)
# Output: THD device (double) in percentile

def playback_thd(inputData, inputTypes):
    output = []
    outputData = []
    outputTypes = []
    # basic sanity check
    inputError = False
    if (inputTypes[0] != TYPE_MONO):
        inputError = True
    if (inputTypes[1] != TYPE_I64):
        inputError = True
    if (inputTypes[2] != TYPE_DOUBLE):
        inputError = True
    if inputError:
        output.append(RESULT_ERROR)
        output.append(outputData)
        output.append(outputTypes)
        return output

    hostRecording = inputData[0]
    signalFrequency = inputData[1]
    thdPassPercentile = inputData[2]
    samplingRate = 44100

    thd = calc_thd(hostRecording, signalFrequency, samplingRate, 0.02) * 100
    print "THD %", thd, "Margain % ", thdPassPercentile
    if (thd < thdPassPercentile):
        output.append(RESULT_PASS)
    else:
        output.append(RESULT_OK)
    outputData.append(thd)
    outputTypes.append(TYPE_DOUBLE)
    output.append(outputData)
    output.append(outputTypes)
    return output

max_stars_repo_name: HelixOS/cts

max_issues_repo_name: HelixOS/cts

max_forks_repo_name: HelixOS/cts

reason: import numpy,import scipy


text:
"""
Model API.
"""

import importlib.resources

import numpy as np
import onnxruntime as rt

class Sonar:
    _map = {0: "hate_speech", 1: "offensive_language", 2: "neither"}

    def __init__(self):
        with importlib.resources.path("hatesonar.data", "pipeline.onnx") as pipeline_file:
            self.pipeline = rt.InferenceSession(str(pipeline_file.absolute()))

    def ping(self, text: str) -> dict:
        assert isinstance(text, str)

        proba = self.pipeline.run(None, {"input": [text]})[1][0]

        res = {
            "text": text,
            "top_class": Sonar._map[np.argmax(proba)],
            "classes": [
                {"class_name": Sonar._map[k], "confidence": proba[k]}
                for k in sorted(Sonar._map)
            ],
        }

        return res


text:
import numpy as np
import pandas as pd
import seaborn as sns
import math
import matplotlib as mpl
import matplotlib.pyplot as plt
from adjustText import adjust_text

import streamlit as st
from st_aggrid import GridOptionsBuilder, AgGrid, GridUpdateMode, DataReturnMode, JsCode

def display_options(df):
        # Get columns with specific dtypes
        all_columns = df.columns.values
        # Get columns with numeric data
        numeric_columns = []
        # Get columns with string data
        alphabetic_columns = []
        short_alpha_columns = []
        short_alpha_columns.append('None')

        for col in all_columns:
                if df[col].dtype == object:
                        alphabetic_columns.append(col)
                else:
                        numeric_columns.append(col)

        for alpha in alphabetic_columns:
                if df[alpha].nunique() < 10:
                        short_alpha_columns.append(alpha)

        # Side bar contents
        # Allow user to select rows for annotation
        st.sidebar.subheader("Selection options")
        selection_mode = st.sidebar.radio("Selection Mode", ['single','multiple'], index=1)

        use_checkbox = st.sidebar.checkbox("Use check box for selection", value=True)

        if ((selection_mode == 'multiple') & (not use_checkbox)):
                rowMultiSelectWithClick = st.sidebar.checkbox("Multiselect with click (instead of holding CTRL)", value=False)
                if not rowMultiSelectWithClick:
                    suppressRowDeselection = st.sidebar.checkbox("Suppress deselection (while holding CTRL)", value=False)
                else:
                    suppressRowDeselection=False
        st.sidebar.text("_____________________________")

        # Display df with desired options

        gb = GridOptionsBuilder.from_dataframe(df)

        # Customize gridOptions
        gb.configure_default_column(groupable=True, value=True, enableRowGroup=True, aggFunc='sum', editable=True)
        for col in numeric_columns:
                if df[col].dtype == float:
                        gb.configure_column(col,
                                            type=["numericColumn","numberColumnFilter","customNumericFormat"],
                                            precision = 2)

        gb.configure_selection(selection_mode)
        if use_checkbox:
                gb.configure_selection(selection_mode, use_checkbox=True, groupSelectsChildren=True, groupSelectsFiltered=True)
        if ((selection_mode == 'multiple') & (not use_checkbox)):
                gb.configure_selection(selection_mode, use_checkbox=False, rowMultiSelectWithClick=rowMultiSelectWithClick, suppressRowDeselection=suppressRowDeselection)

        gb.configure_grid_options(domLayout='normal')
        gridOptions = gb.build()

        # Allow user to select graph type
        st.sidebar.subheader('Plot options')

        plot_type_menu = ['scatter', 'bar', 'box', 'histogram']
        plot_type_choice = st.sidebar.selectbox('Select a plot type:', plot_type_menu)
        plot_size_slider = st.sidebar.slider('Plot size:', 5, 15)
        f, ax = plt.subplots(figsize = (plot_size_slider, 5))

        ###################################################################################################################################

        #Display the grid
        st.header("Streamlit Ag-Grid")

        grid_response = AgGrid(
                df,
                gridOptions=gridOptions,
                height=300,
                data_return_mode='FILTERED_AND_SORTED',
                update_mode='MODEL_CHANGED',
                fit_columns_on_grid_load=False,
                allow_unsafe_jscode=True, #Set it to True to allow jsfunction to be injected
                enable_enterprise_modules=True,
                )

        # Make a few color palettes
        pal_1 = 'Set1'
        pal_2 = 'Paired'
        pal_3 = 'Set2'
        pal_4 = 'Accent'
        pal_5 = 'Set1_r'
        pal_6 = 'Set2_r'

        if plot_type_choice == 'scatter':
                x_menu = all_columns
                y_menu = numeric_columns
                group_menu = short_alpha_columns
                pal_menu = [pal_1, pal_2, pal_3, pal_4, pal_5, pal_6]
                anno_menu = alphabetic_columns

                x_selection = st.sidebar.selectbox('x axis:', x_menu)
                y_selection = st.sidebar.selectbox('y axis:', y_menu)
                group_selection = st.sidebar.selectbox('Group by:', group_menu)
                palette_selection = st.sidebar.selectbox('Color palette:', pal_menu)
                anno_selection = st.sidebar.selectbox('Annotate by:', anno_menu)

                # Make scatter plot
                if group_selection != 'None':
                        to_plot = grid_response['data'].dropna()
                        sns.scatterplot(x = x_selection, y = y_selection, data = to_plot,
                                        ax = ax, s = 100,
                                        hue = group_selection,
                                        palette = sns.color_palette(palette = palette_selection,
                                                                    n_colors = to_plot[group_selection].nunique()))

                else:
                        to_plot = grid_response['data'].copy()
                        sns.scatterplot(x = x_selection, y = y_selection, data = df,
                                        ax = ax, s = 100)

                ymin = to_plot[y_selection].min()
                ymax = to_plot[y_selection].max()
                yrange = ymax - ymin

                xmin = to_plot[x_selection].min()
                xmax = to_plot[x_selection].max()
                xrange = xmax - xmin

                # Rotate xlabels if data type is string
                if to_plot[x_selection].dtypes == 'object':
                        plt.xticks(rotation = 90)

                # Set x and y limits
                else:
                        #ax.set_xlim([xmin - 10*math.ceil(abs(xmin)), xmax*1.1])
                        #ax.set_ylim([ymin - 10*math.ceil(abs(ymin)), ymax*1.1])
                        ax.set_xlim([xmin - 0.2*xrange, xmax + 0.2*xrange])
                        ax.set_ylim([ymin - 0.2*yrange, ymax + 0.2*yrange])

                if '%' in y_selection:
                        ax.set_ylim([0, 100])

                # Annotate data points
                anno_ls = []
                selected = grid_response['selected_rows']
                selected_df = pd.DataFrame(selected)
                potential_annos = to_plot[anno_selection].tolist()

                for k, v in selected_df.iterrows():
                        if v[x_selection] != 'nan' and v[anno_selection] in potential_annos:
                                anno = ax.text(x = v[x_selection], y = v[y_selection],
                                               s = v[anno_selection], size = 10)
                                anno_ls.append(anno)

                adjust_text(anno_ls, ax = ax)

                plt.legend(loc = 'upper left', bbox_to_anchor = (1.01, 1.02), title = group_selection)

                st.pyplot(f)

        if plot_type_choice == 'bar':
                x_menu = alphabetic_columns
                y_menu = numeric_columns
                group_menu = short_alpha_columns
                pal_menu = [pal_1, pal_2, pal_3, pal_4, pal_5, pal_6]

                x_selection = st.sidebar.selectbox('x axis:', x_menu)
                y_selection = st.sidebar.selectbox('y axis:', y_menu)
                group_selection = st.sidebar.selectbox('Group by:', group_menu)
                palette_selection = st.sidebar.selectbox('Color palette:', pal_menu)

                # Make a bar chart
                if group_selection != 'None':
                        sns.barplot(x = x_selection, y = y_selection, data = df,
                                    ax = ax,
                                    hue = group_selection,
                                    palette = sns.color_palette(palette = palette_selection,
                                                                n_colors = df[group_selection].nunique()))
                else:
                        sns.barplot(x = x_selection, y = y_selection, data = df,
                                    ax = ax)

                plt.xticks(rotation = 90)
                plt.legend(bbox_to_anchor = (1.01,1), title = group_selection)
                st.pyplot(f)

        if plot_type_choice == 'box':
                x_menu = alphabetic_columns
                y_menu = numeric_columns
                anno_menu = alphabetic_columns

                x_selection = st.sidebar.selectbox('x axis:', x_menu)
                y_selection = st.sidebar.selectbox('y axis:', y_menu)
                anno_selection = st.sidebar.selectbox('Annotate by:', anno_menu)

                # Make a boxplot
                to_plot = grid_response['data'].dropna()
                sns.boxplot(x = x_selection, y = y_selection, data = to_plot,
                                ax = ax, linewidth = 1, color = 'white', saturation = 1)
                sns.swarmplot(x = x_selection, y = y_selection, data = to_plot,
                                ax = ax, edgecolor = 'black', linewidth = 1, size = 8)

                plt.xticks(rotation = 90)
                tick_labels = ax.get_xticklabels()
                tick_dict = {}
                for i in tick_labels:
                        tick_dict[i.get_text()] = i.get_position()[0]

                # Annotate data points
                anno_ls = []
                selected = grid_response['selected_rows']
                selected_df = pd.DataFrame(selected)
                for k, v in selected_df.iterrows():
                        if v[x_selection] != 'nan':
                                anno = ax.text(x = tick_dict[v[x_selection]], y = v[y_selection],
                                               s = v[anno_selection], size = 10)
                                anno_ls.append(anno)

                adjust_text(anno_ls, ax = ax)

                st.pyplot(f)

        if plot_type_choice == 'histogram':
                x_menu = numeric_columns
                x_selection = st.sidebar.selectbox('x axis:', x_menu)
                bin_slider = st.slider('Number of bins:', 1, 35, 5)

                # Round values to get general range
                x_min = df[x_selection].min()
                x_max = df[x_selection].max()
                x_range = x_max - x_min
                if x_range >= 1000:
                        x_min = round(x_min, -2)
                        x_max = round(x_max, -2)
                        x_range = x_max-x_min
                else:
                        x_min = round(x_min)
                        x_max = round(x_max)
                        x_range = x_max-x_min

                # Make histogram
                bins = np.histogram_bin_edges(df[x_selection], bins=bin_slider)
                print(bins)
                sns.histplot(data = df, x = x_selection, bins = bin_slider)
                st.pyplot(f)

###################################################################################################################################
# Main function
###################################################################################################################################

def main():

        st.title('Plotting from pandas df')

        # Have user upload a file to plot
        data_file = st.file_uploader('Upload CSV or Excel file', type = ['csv', 'xlsx', 'xls'])

        if data_file is not None:
                # To see details
                file_details = {'filename': data_file.name,
                                'filetype': data_file.type,
                                'filesize': data_file.size}

                data_file_name = data_file.name.split('.')[0]
                if data_file.name.endswith('csv'):
                        df = pd.read_csv(data_file)
                elif data_file.name.endswith('xlsx') or data_file.name.endswith('xls'):
                        df = pd.ExcelFile(data_file).parse()

                display_options(df)

if __name__ == '__main__':
	main()


text:
from astropy import units as u
from perylune import orbit_tools
from math import pi

import pytest

# pairs of input, expected values
norm_test_cases = [ [1,1], [-pi, pi], [-4, 2*pi-4], [0, 0], [2*pi, 0], [7, 7-2*pi], [0*u.deg, 0*u.deg], [370*u.deg, 10*u.deg], \
                    [-10*u.deg, 350*u.deg]]

def test_norm_2pi():

    for inp, exp in norm_test_cases:
        assert exp == orbit_tools.normalize_2pi(inp)

def test_norm_pipi():

    for inp, exp in norm_test_cases:
        # Wrap it around by 180 degrees/pi rad
        if type(exp) == u.quantity.Quantity:
            if exp.unit == u.deg and exp.value > 180:
                exp -= 360 * u.deg
            if exp.unit == u.rad and exp.value > 2*pi:
                exp -= 2*pi*u.rad
        else:
            if exp >= pi:
                exp = exp - 2*pi

        assert exp == orbit_tools.normalize_pipi(inp)

max_stars_repo_name: tomaszmrugalski/astro

max_issues_repo_name: tomaszmrugalski/astro

max_forks_repo_name: tomaszmrugalski/astro

reason: from astropy