text:
#!/usr/bin/env python3

from tensorboardX import SummaryWriter
from itertools import count
import torch
import argparse
import numpy as np
import os
import time
from time import gmtime, strftime

import threading as th
import glob
from datetime import datetime
import glob
import pickle

"""
class PriorInfo(object):
    #docstring for PriorInfo
    def __init__(self, dataset_path="/home/alex/catkin_ws/datasets/datasets_v1"):
        super(PriorInfo, self).__init__()
        self.dataset_path = dataset_path
        self.joint_angles = []
        self.joint_next_angles = []
        self.joint_actions = []

    def load_demonstration(self):
        #dataset_path = "/home/alex/catkin_ws/src/cobot_controllers/datasets/without_extra_grabber/npys/"
        #dataset_path = "/home/alex/catkin_ws/datasets/datasets_v0"
        dataset_path = self.dataset_path

        traj_id = 0
        exp_trajs = []
        traj = None
        sample_ratio = 5

        filename_list = sorted(glob.glob(dataset_path+'/*.npy'))

        for filename in filename_list:
            if filename.endswith(".npy"):# and mode in filename:
                exp_traj = np.load(os.path.join(dataset_path, filename))[:-10][::sample_ratio]
                exp_trajs.append(exp_traj)

        for exp_t in exp_trajs:
            actions = []
            for tid in range(len(exp_t)-1):
                delta = exp_t[tid+1] - exp_t[tid]
                actions.append(delta)

                self.joint_angles.append(exp_t[tid])
                self.joint_actions.append(delta)
                self.joint_next_angles.append(exp_t[tid+1])

        self.joint_angles = np.array(self.joint_angles)
        self.joint_next_angles = np.array(self.joint_next_angles)
        self.joint_actions = np.array(self.joint_actions)

        return exp_trajs[0:]

    def retrieve_closest_prior(self, joints):
        norms = np.linalg.norm(self.joint_angles - joints, axis=1)
        closest_id = np.argmin(norms)

        return self.joint_angles[closest_id], self.joint_actions[closest_id], self.joint_next_angles[closest_id]
"""

class PriorInfo(object):
    """docstring for PriorInfo"""
    def __init__(self, dataset_path="/home/alex/Documents/franka_panda_insertion_logs/experts/"):
        super(PriorInfo, self).__init__()
        self.dataset_path = dataset_path
        self.joint_angles = []
        self.joint_next_angles = []
        self.cartesian_poses = []
        self.cartesian_next_poses = []
        self.joint_actions = []
        self.sample_ratio = 10

    def load_demonstration(self, load_ids=[1]):
        dataset_path = self.dataset_path

        exp_joint_trajs = []
        traj = None
        sample_ratio = self.sample_ratio

        exp_cartesian_trajs = np.load(self.dataset_path+'expert_cartesian_poses.npy', allow_pickle=True)
        exp_joint_trajs = np.load(self.dataset_path+'expert_joints_poses.npy', allow_pickle=True)

        exp_subsample_c_trajs = exp_cartesian_trajs[load_ids].copy()
        exp_subsample_j_trajs = exp_joint_trajs[load_ids].copy()

        for eid, exp_t in enumerate(exp_subsample_j_trajs):
            actions = []
            exp_c_t = exp_cartesian_trajs[eid]

            exp_subsample_j_trajs[eid] = exp_t[::sample_ratio]
            #exp_subsample_c_trajs[eid] = exp_c_t[::sample_ratio]

            for tid in range(len(exp_t)-sample_ratio):
                delta = exp_t[tid+sample_ratio] - exp_t[tid]

                self.joint_angles.append(exp_t[tid])
                self.joint_next_angles.append(exp_t[tid+sample_ratio])
                self.joint_actions.append(delta)
                #self.cartesian_poses.append(exp_c_t[tid])
                #self.cartesian_next_poses.append(exp_c_t[tid+sample_ratio])

        self.joint_angles = np.array(self.joint_angles)
        self.joint_next_angles = np.array(self.joint_next_angles)
        self.joint_actions = np.array(self.joint_actions)
        #self.cartesian_poses = np.array(self.cartesian_poses)
        #self.cartesian_next_poses = np.array(self.cartesian_next_poses)

        return exp_subsample_j_trajs#, exp_subsample_c_trajs

    def retrieve_closest_prior(self, inputs,
                                dis_type='cartesian'):
        if dis_type=='joints':
            norms = np.linalg.norm(self.joint_angles - inputs, axis=1)
            closest_id = np.argmin(norms)
        elif dis_type=='cartesian':
            norms = np.linalg.norm(self.cartesian_poses - inputs, axis=1)
            closest_id = np.argmin(norms)

        return closest_id, self.joint_angles[closest_id], self.joint_actions[closest_id], self.joint_next_angles[closest_id]#, self.cartesian_next_poses[closest_id]

all_acts = np.load('/home/alex/Documents/franka_panda_insertion_logs/expert_action_under_prior.npy', allow_pickle=True)

# initiate the expert prior class, which contains the expert state-action-next_state in joint angle space
prior = PriorInfo()
# load expert joint trajectory
demos = prior.load_demonstration()

for d in demos:
    print(len(d))

for exp_ep in range(len(demos)):
    demo_traj = demos[exp_ep]
    for ts in range(len(demo_traj)-1):
        curr_joints = demo_traj[ts]
        cid_j, _, _, prior_target_joints = prior.retrieve_closest_prior(curr_joints, dis_type='joints')

        target_joint = demo_traj[ts+1]
        print(cid_j)


text:
 ## Import libraries
import os
from PIL import Image
import numpy as np
import matplotlib.pyplot as plt
import matplotlib.image as mpimg
import skimage
import scipy.ndimage.interpolation as sci_int
import scipy.ndimage.morphology as sci_morph

from skimage import data, color, img_as_uint
from skimage.transform import rescale, resize, downscale_local_mean
from skimage.io import imread, imshow
from skimage.feature import blob_dog, blob_log, blob_doh
from skimage.color import rgb2gray
from skimage.filters import threshold_otsu

from skimage import segmentation
from skimage.morphology import watershed
from skimage.feature import peak_local_max
from scipy import ndimage
from skimage import morphology
from skimage.filters import gaussian
from skimage.measure import regionprops
from skimage.restoration import (denoise_tv_chambolle)

import sys

from subprocess import call, run
import subprocess

from math import sqrt

import SimpleITK as sitk

from shutil import move

import math

import seaborn as sns

#### Import custom libraries
import sys
sys.path.insert(0, 'INSERT FOLDER WHERE toolbox.py IS SAVED')

import importlib
import toolbox as bf
importlib.reload(bf)
####

# Paths
study_folder_server = r'INSERT PATH TO THE OUTPUT FOLDER '

# Selected reference brain from study
reference_brain_id = r'INSERT REFERENCE SAMPLE ID'

# Define path and sample names
sample_id = []
full_brain_udrive = []
sample_flip = []

#### SAMPLES
# # # # reference brain
sample_id.append(r'INSERT REFERENCE SAMPLE ID')
full_brain_udrive.append(r'INSERT FULL PATH TO REFERENCE SAMPLE FOLDER')
sample_flip.append('0') # 0 for no flipping, 1 for flipping

# # # # other brain samples
sample_id.append(r'INSERT SAMPLE ID')
sample_flip.append('0')
full_brain_udrive.append(r'INSERT FULL PATH TO SAMPLE FOLDER')

#### NOTE: Working directory and paramter files for registrations
# ### CREATE A FOLDER "elastix" IN THE "study_folder_server"
# ### ADD A FOLDER "workingDir" along with affine and Bspline registration parameter files in the "elastix" folder

# Analyze samples
for i in range(len(sample_id)):
    if sample_flip[i] == '1' or sample_flip[i] == '0':

        print(sample_id[i])

        ### Read auto and spec of full Brain
        auto = []
        file_names = os.listdir(full_brain_udrive[i])
        file_names.sort()

        nb_files = len(file_names)
        file_loop = range(0,nb_files)

        # read all files
        for idx,val in enumerate(file_loop):
            if os.path.isfile(os.path.join(full_brain_udrive[i], file_names[idx])):
                if os.path.join(full_brain_udrive[i], file_names[idx]).split('.')[-1]=='tif':
                    if 'C00' in file_names[idx]:
                        #print(idx)
                        # read image with scikit
                        im = imread(os.path.join(full_brain_udrive[i], file_names[idx]),as_grey=False, plugin='pil')
                        # resize image to allow efficient registration
                        image_rescaled = rescale(im, 4.794/20) # think abut interpolation order and anti-aliasing
                        image_rescaled_uint = img_as_uint(image_rescaled)

                        auto.append(image_rescaled_uint)

        # Convert list to array
        auto = np.array(auto)
        print(auto.shape)

        auto = downscale_local_mean(auto, (2,1,1))

        # Pre-processing
        #CREATE CORSE TISSUE MASK
        cutoff = np.percentile(auto,5)
        print(cutoff)

        mask = np.zeros(auto.shape,'uint8')
        mask[auto>cutoff] = 1

        bf.saveNifti(auto, study_folder_server+r'/nifti/auto.nii.gz')
        bf.saveNifti(mask, study_folder_server+r'/nifti/mask.nii.gz')

        #### BIAS FIELD CORRECTION #####
        program = r'INSERT PATH TO ../run_Bias_field_correction.sh '
        runtime_path = r'INSERT PATH TO ../MATLAB/MATLAB_Runtime/v93 '
        vol_in =  study_folder_server+r'/nifti/auto.nii.gz '
        mask_in = study_folder_server+r'/nifti/mask.nii.gz '
        cor_out = study_folder_server+r'/nifti/auto_biascor.nii.gz '
        origin_out = study_folder_server+r'/nifti/origin.nii.gz '

        print(program + runtime_path + vol_in  + mask_in + cor_out + origin_out)
        os.system(program + runtime_path + vol_in  + mask_in + cor_out + origin_out)

        corr = bf.readNifti(study_folder_server+r'/nifti/auto_biascor.nii.gz')

        # RESCALE TO 8 BIT
        scale_limit = np.percentile(corr, (99.999))
        corr = skimage.exposure.rescale_intensity(corr, in_range=(0,scale_limit), out_range='uint8')
        auto = np.copy(corr)

        # rescale intensity based on mean/median of tissue
        auto_temp = np.copy(auto)
        scale_thres = skimage.filters.threshold_otsu(auto_temp)
        auto_temp[auto_temp<scale_thres] = 0
        nz_mean = np.mean(auto_temp[auto_temp>0].flatten())
        print(nz_mean)
        scale_fact = 120/nz_mean
        auto = auto*scale_fact
        print(np.max(auto))
        auto[auto>255] = 255
        auto= auto.astype('uint8')

        # CLAHE IN THREE ORTHOGONAL PLANES
        auto_temp_hor = np.copy(auto)
        auto_temp_cor = np.copy(auto)
        auto_temp_sag = np.copy(auto)

        for h in range(auto_temp_cor.shape[1]):
            temp_img = auto_temp_cor[:,h,:]

            temp_img[0:2,0:2] = 250
            temp_img[3:5,3:5] = 0

            clahe_im = skimage.exposure.equalize_adapthist(temp_img, kernel_size=(int(temp_img.shape[0]/3),int(temp_img.shape[1]/6)), clip_limit=0.01, nbins=255)
            clahe_im[0:2,0:2] = 0

            clahe_im = clahe_im*255
            clahe_im[clahe_im<0] = 0
            clahe_im = np.uint8(clahe_im)
            auto_temp_cor[:,h,:] = clahe_im

        for h in range(auto_temp_hor.shape[0]):
            temp_img = auto_temp_hor[h,:,:]

            temp_img[0:2,0:2] = 250
            temp_img[3:5,3:5] = 0

            clahe_im = skimage.exposure.equalize_adapthist(temp_img, kernel_size=(int(temp_img.shape[0]/3),int(temp_img.shape[1]/6)), clip_limit=0.01, nbins=255)
            clahe_im[0:2,0:2] = 0

            clahe_im = clahe_im*255
            clahe_im[clahe_im<0] = 0
            clahe_im = np.uint8(clahe_im)
            auto_temp_hor[h,:,:] = clahe_im

        for h in range(auto_temp_sag.shape[2]):
            temp_img = auto_temp_sag[:,:,h]

            temp_img[0:2,0:2] = 250
            temp_img[3:5,3:5] = 0

            clahe_im = skimage.exposure.equalize_adapthist(temp_img, kernel_size=(int(temp_img.shape[0]/3),int(temp_img.shape[1]/6)), clip_limit=0.01, nbins=255)
            clahe_im[0:2,0:2] = 0

            clahe_im = clahe_im*255
            clahe_im[clahe_im<0] = 0
            clahe_im = np.uint8(clahe_im)
            auto_temp_sag[:,:,h] = clahe_im

        # combine angles
        clahe_all = np.zeros((auto_temp_sag.shape[0],auto_temp_sag.shape[1],auto_temp_sag.shape[2],3))
        clahe_all[:,:,:,0] = auto_temp_hor
        clahe_all[:,:,:,1] = auto_temp_cor
        clahe_all[:,:,:,2] = auto_temp_sag
        clahe_all_mean = np.mean(clahe_all,3)

        # combine with original volume
        clahe_final = np.zeros((auto_temp_sag.shape[0],auto_temp_sag.shape[1],auto_temp_sag.shape[2],2))
        clahe_final[:,:,:,0] = clahe_all_mean
        clahe_final[:,:,:,1] = auto

        clahe_final = np.mean(clahe_final,3)

        # # HISTOGRAM MATCHING
        reference = bf.readNifti(os.path.join(study_folder_server+r'/nifti',reference_brain_id+'_regi.nii.gz'))
        clahe_final = skimage.transform.match_histograms(clahe_final.astype('float'), reference.astype('float'), multichannel=False)

        # flip if needed
        if sample_flip[i] == '1':
            # flip brain
            clahe_final = np.flip(clahe_final,2)
            print('Flipping sample...')

        # save volume as nifty file
        clahe_final = np.uint8(clahe_final)
        final_sitk = sitk.GetImageFromArray(clahe_final)
        final_sitk.SetOrigin((round(final_sitk.GetWidth()/2), round(final_sitk.GetHeight()/2), -round(final_sitk.GetDepth()/2)))
        sitk.WriteImage(final_sitk,os.path.join(study_folder_server+r'/nifti',sample_id[i]+'_regi.nii.gz'))
        print('Brain ready for registration')

        # #### UNCOMMENT IF CORTEX OR OLFACTORY VOLUME
        #
        # ## AFFINE
        # program = r'elastix -threads 16 '; #elastix is added to PATH
        # fixed_name = r'-f ' + os.path.join(study_folder_server+r'/nifti',reference_brain_id+'_regi_orig.nii.gz ');
        # moving_name = r'-m ' + os.path.join(study_folder_server+r'/nifti',sample_id[i]+'_regi_orig.nii.gz ');
        # outdir = r'-out ' + study_folder_server + r'/elastix/workingDir ';
        # params = r'-p ' + study_folder_server + r'/elastix/Affine_Gubra_June2019.txt ';
        #
        # os.system(program + fixed_name + moving_name  + outdir + params)
        # move(study_folder_server + r'/elastix/workingDir/TransformParameters.0.txt', os.path.join(study_folder_server+r'/nifti',sample_id[i]+'_affine.txt'))
        #
        #
        # # BSPLINE
        # params = r'-p ' + study_folder_server + r'/elastix/Bspline_Gubra_June2019.txt ';
        # t0 = r'-t0 ' + os.path.join(study_folder_server+r'/nifti',sample_id[i]+'_affine.txt')
        #
        # os.system(program + fixed_name + moving_name  + outdir + params + t0)
        #
        # move(study_folder_server + r'/elastix/workingDir/result.0.nii.gz', os.path.join(study_folder_server+r'/nifti',sample_id[i]+'_bspline_orig.nii.gz'))
        # move(study_folder_server + r'/elastix/workingDir/TransformParameters.0.txt', os.path.join(study_folder_server+r'/nifti',sample_id[i]+'_bspline.txt'))
        #
        # ## temp hack as the nifti file from elastix could not be opened in ITK-SNAP
        # temp = sitk.ReadImage(os.path.join(study_folder_server+r'/nifti',sample_id[i]+'_bspline_orig.nii.gz'))
        # temp = sitk.GetArrayFromImage(temp)
        #
        # final = temp
        # final = final.astype('uint8')
        # final_sitk = sitk.GetImageFromArray(final)
        # final_sitk.SetOrigin((round(final_sitk.GetWidth()/2), round(final_sitk.GetHeight()/2), -round(final_sitk.GetDepth()/2)))
        # sitk.WriteImage(final_sitk,os.path.join(study_folder_server+r'/nifti',sample_id[i]+'_bspline.nii.gz'))

ext: py

max_stars_repo_name: Gubra-ApS/LSFM-atlas

max_issues_repo_name: Gubra-ApS/LSFM-atlas

max_forks_repo_name: Gubra-ApS/LSFM-atlas

reason: import numpy,import scipy,from scipy


text:
import ctypes
import os
from six.moves import urllib
import zipfile
import stat
import logging
import json

logging.basicConfig(level=logging.DEBUG)
print "logging"

def response(status_code, response_body):
    return {
                'statusCode': status_code,
                'body': json.dumps(response_body) if response_body else json.dumps({}),
                'headers': {
                    'Content-Type': 'application/json',
                },
            }

def vectorize_sequences(sequences, dimension):
    results = np.zeros((len(sequences), dimension))
    for i, sequence in enumerate(sequences):
       results[i, sequence] = 1.
    return results

for d, _, files in os.walk('lib'):
    for f in files:
        if f.endswith('.a') or f.endswith('.settings'):
            continue
        print('loading %s...' % f)
        ctypes.cdll.LoadLibrary(os.path.join(d, f))

import keras
from keras.preprocessing.text import one_hot
from numpy import array
import numpy as np

model = keras.models.load_model(os.environ['MODEL_NAME'])

def handler(event, context):
    print event
    #params = event['queryStringParameters']

    sms = event['body']

    print sms

    sentence = one_hot(str(sms), 87413)
    print sentence
    sentences = [sentence]

    vector = vectorize_sequences(sentences,87413)

    result = model.predict(vector)
    print "Verdict: ", result
    print "Shape: ", result.shape
    print "Type: ", type(result)
    result = float(np.array2string(result)[2:-2])
    print(result)

    return response(200, result)

max_stars_repo_name: diegonat/smsantispam

max_issues_repo_name: diegonat/smsantispam

max_forks_repo_name: diegonat/smsantispam

reason: import numpy,from numpy


text:
import numpy as np
import matplotlib.pyplot as plt
import cmocean.cm as cmo
import os
import sys
import xarray as xr

# add the 'src/' directory as one where we can import modules
src_dir = os.path.join(os.environ.get('projdir'),'src')
sys.path.append(src_dir)

from features.rotate_vector_roms import rotate_vector_roms

def plot_uv(ds,grd,block=2):

    # Radius of the Earth in metres
    r = 6.371e6
    # Degrees to radians conversion factor
    deg2rad = np.pi/180
    # Side length of blocks to average vectors over (can't plot vector at every
    # single point or the plot will be way too crowded)

    lon = grd.lon_rho.values
    lat = grd.lat_rho.values

    u = ds.ubar.values
    v = ds.vbar.values

    angle = np.zeros(np.shape(lon))
    u_rho,v_rho = rotate_vector_roms(u, v, angle)

    speed = np.sqrt(np.square(u_rho)+np.square(v_rho))

    #print('initialize and fill up the arrays')
    numy = np.size(lon,0) #530
    numx = np.size(lon,1) #630
    x = np.arange(numx)
    y = np.arange(numy)
    xmesh,ymesh = np.meshgrid(x,y)
    #print(numx,numy,x,y)

    # Average x, y, u_circ, and v_circ over block x block intervals
    # Calculate number of blocks
    size0 = int(np.ceil(numy/float(block)))
    size1 = int(np.ceil(numx/float(block)))

    # Set up arrays for averaged fields
    x_block = np.ma.empty([size0, size1])
    y_block = np.ma.empty([size0, size1])
    u_block = np.ma.empty([size0, size1])
    v_block = np.ma.empty([size0, size1])
    # Set up arrays containing boundary indices
    posn0 = list(np.arange(0, numy, block))
    posn0.append(numy)
    posn1 = list(np.arange(0, numx, block))
    posn1.append(numx)
    #print(posn0,posn1)
    # Double loop to average each block (can't find a more efficient way to do
    # this)
    for j in np.arange(size0):
        for i in np.arange(size1):
            start0 = posn0[j]
            end0 = posn0[j+1]
            start1 = posn1[i]
            end1 = posn1[i+1]
            x_block[j,i] = np.mean(xmesh[start0:end0, start1:end1])
            y_block[j,i] = np.mean(ymesh[start0:end0, start1:end1])
            u_block[j,i] = np.mean(u_rho[start0:end0, start1:end1])
            v_block[j,i] = np.mean(v_rho[start0:end0, start1:end1])

    # Make the plot
    fig,ax0 = plt.subplots(1,figsize=(15,10))

    speedP = ax0.pcolormesh(xmesh,ymesh,speed*100, vmin=0,vmax=10, cmap=cmo.speed)
    plt.colorbar(speedP,ax=ax0)
    # Add vectors for each block
    quiverP = ax0.quiver(x_block, y_block, u_block, v_block,pivot="mid", color='black',scale_units='xy',scale=0.01)
    plt.quiverkey(quiverP, 0.75, 0.99, 0.1, r'$10 \frac{cm}{s}$', labelpos='E',
                       coordinates='figure')
    ax0.set_title('Mean barotropic velocity (cm/s)', fontsize=16)
    ax0.set_aspect('equal')
    ax0.axis('off')

    print('test')
    return fig

reason: import numpy


text:
# Copyright (c) 2021 PaddlePaddle Authors. 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.

import paddle.fluid.core as core
import paddle.fluid.eager.eager_tensor_patch_methods as eager_tensor_patch_methods
import paddle
import numpy as np
from paddle.fluid.framework import _test_eager_guard
from paddle.fluid.data_feeder import convert_dtype
import unittest

class EagerScaleTestCase(unittest.TestCase):
    def test_scale_base(self):
        with _test_eager_guard():
            paddle.set_device("cpu")
            arr = np.ones([4, 16, 16, 32]).astype('float32')
            tensor = paddle.to_tensor(arr, 'float32', core.CPUPlace())
            print(tensor)
            tensor = core.eager.scale(tensor, 2.0, 0.9, True, False)
            for i in range(0, 100):
                tensor = core.eager.scale(tensor, 2.0, 0.9, True, False)
            print(tensor)
            self.assertEqual(tensor.shape, [4, 16, 16, 32])
            self.assertEqual(tensor.stop_gradient, True)

    def test_retain_grad_and_run_backward(self):
        with _test_eager_guard():
            paddle.set_device("cpu")

            input_data = np.ones([4, 16, 16, 32]).astype('float32')
            data_eager = paddle.to_tensor(input_data, 'float32',
                                          core.CPUPlace(), False)

            grad_data = np.ones([4, 16, 16, 32]).astype('float32')
            grad_eager = paddle.to_tensor(grad_data, 'float32', core.CPUPlace())

            core.eager.retain_grad_for_tensor(data_eager)

            out_eager = core.eager.scale(data_eager, 1.0, 0.9, True, True)
            self.assertFalse(data_eager.grad._is_initialized())
            core.eager.run_backward([out_eager], [grad_eager], False)
            self.assertTrue(data_eager.grad._is_initialized())
            self.assertTrue(np.array_equal(data_eager.grad.numpy(), input_data))

class EagerDtypeTestCase(unittest.TestCase):
    def check_to_tesnsor_and_numpy(self, dtype, proto_dtype):
        with _test_eager_guard():
            arr = np.random.random([4, 16, 16, 32]).astype(dtype)
            tensor = paddle.to_tensor(arr, dtype)
            self.assertEqual(tensor.dtype, proto_dtype)
            self.assertTrue(np.array_equal(arr, tensor.numpy()))

    def test_dtype_base(self):
        print("Test_dtype")
        self.check_to_tesnsor_and_numpy('bool', core.VarDesc.VarType.BOOL)
        self.check_to_tesnsor_and_numpy('int8', core.VarDesc.VarType.INT8)
        self.check_to_tesnsor_and_numpy('uint8', core.VarDesc.VarType.UINT8)
        self.check_to_tesnsor_and_numpy('int16', core.VarDesc.VarType.INT16)
        self.check_to_tesnsor_and_numpy('int32', core.VarDesc.VarType.INT32)
        self.check_to_tesnsor_and_numpy('int64', core.VarDesc.VarType.INT64)
        self.check_to_tesnsor_and_numpy('float16', core.VarDesc.VarType.FP16)
        self.check_to_tesnsor_and_numpy('float32', core.VarDesc.VarType.FP32)
        self.check_to_tesnsor_and_numpy('float64', core.VarDesc.VarType.FP64)
        self.check_to_tesnsor_and_numpy('complex64',
                                        core.VarDesc.VarType.COMPLEX64)
        self.check_to_tesnsor_and_numpy('complex128',
                                        core.VarDesc.VarType.COMPLEX128)

class EagerTensorPropertiesTestCase(unittest.TestCase):
    def test_properties(self):
        print("Test_properties")
        with _test_eager_guard():
            paddle.set_device("cpu")
            arr = np.ones([4, 16, 16, 32]).astype('float32')
            tensor = paddle.to_tensor(arr, core.VarDesc.VarType.FP32,
                                      core.CPUPlace())
            self.assertEqual(tensor.shape, [4, 16, 16, 32])
            tensor.name = 'tensor_name_test'
            self.assertEqual(tensor.name, 'tensor_name_test')
            self.assertEqual(tensor.persistable, False)
            tensor.persistable = True
            self.assertEqual(tensor.persistable, True)
            tensor.persistable = False
            self.assertEqual(tensor.persistable, False)
            self.assertTrue(tensor.place.is_cpu_place())
            self.assertEqual(tensor._place_str, 'CPUPlace')
            self.assertEqual(tensor.stop_gradient, True)
            tensor.stop_gradient = False
            self.assertEqual(tensor.stop_gradient, False)
            tensor.stop_gradient = True
            self.assertEqual(tensor.stop_gradient, True)

    def test_global_properties(self):
        print("Test_global_properties")
        self.assertFalse(core._in_eager_mode())
        with _test_eager_guard():
            self.assertTrue(core._in_eager_mode())
        self.assertFalse(core._in_eager_mode())

    def test_place_guard(self):
        core._enable_eager_mode()
        if core.is_compiled_with_cuda():
            paddle.set_device("gpu:0")
            with paddle.fluid.framework._dygraph_place_guard(core.CPUPlace()):
                self.assertTrue(core.eager._get_expected_place().is_cpu_place())
        else:
            paddle.set_device("cpu")
            with paddle.fluid.framework._dygraph_place_guard(core.CPUPlace()):
                self.assertTrue(core.eager._get_expected_place().is_cpu_place())
        core._disable_eager_mode()

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


text:
import math
import numpy as np
import random

class ExchangeMutation(object):
    def __init__(self, mutation_rate):
        self.mutation_rate = mutation_rate

    def __call__(self, sequence):
        if random.random() < self.mutation_rate:
            sequence_length = len(sequence)
            first           = random.randrange(sequence_length)
            second          = random.randrange(sequence_length)
            sequence[first], sequence[second] = sequence[second], sequence[first]
        return sequence

class OrderedCrossover(object):
    def __call__(self, parent_sequence_a, parent_sequence_b):
        sequence_length  = len(parent_sequence_a)
        child_sequence_a = [None] * sequence_length
        child_sequence_b = [None] * sequence_length

        first  = random.randrange(sequence_length)
        second = random.randrange(sequence_length)

        left, right = min(first, second), max(first, second)

        for m in range(left, right + 1):
            child_sequence_a[m] = parent_sequence_a[m]
            child_sequence_b[m] = parent_sequence_b[m]

        m   = (right + 1) % sequence_length
        n_a = m
        n_b = m

        while m != left:
            while parent_sequence_a[n_a] in child_sequence_b:
                n_a = (n_a + 1) % sequence_length

            while parent_sequence_b[n_b] in child_sequence_a:
                n_b = (n_b + 1) % sequence_length

            child_sequence_b[m] = parent_sequence_a[n_a]
            child_sequence_a[m] = parent_sequence_b[n_b]

            m = (m + 1) % sequence_length

        return child_sequence_a, child_sequence_b

class SequenceCreator(object):
    def __init__(self, length):
        self.length = length

    def __call__(self):
        sequence = list(range(self.length))
        random.shuffle(sequence)
        return sequence

def diversity(population):
    population_size = len(population)
    sequence_length = len(population[0].genotype)

    counts = np.zeros((sequence_length, sequence_length))

    for individual in population:
        for index, gene in enumerate(individual.genotype):
            counts[index, gene] += 1

    diversity = 0.0

    for index in range(sequence_length):
        for character in range(sequence_length):
            frequency  = counts[index, character] / population_size
            if frequency != 0 and frequency != 1:
                diversity -= frequency * math.log(frequency) + (1 - frequency) * math.log(1 - frequency)

    return diversity


text:
#!/usr/bin/env python3
#
# numpy_to_adjacency_matrix.py: converts an `.npy` file to a (textual!)
# adjacency matrix.

import argparse
import sys

import numpy as np

if __name__ == '__main__':
    parser = argparse.ArgumentParser()
    parser.add_argument('INPUT', type=str, help='Input file')
    parser.add_argument('-o', '--output',
        type=str,
        help='Output file (optional)'
    )

    args = parser.parse_args()

    M = np.load(args.INPUT)

    if not args.output or args.output == '-':
        output = sys.stdout
    else:
        output = args.output

    # TODO: make output format configurable?
    np.savetxt(output, M)

max_stars_repo_name: eudoxos/Aleph

max_issues_repo_name: eudoxos/Aleph

max_forks_repo_name: eudoxos/Aleph


text:
# -*- coding: utf-8 -*-
# ---
# jupyter:
#   jupytext:
#     formats: ipynb,py
#     text_representation:
#       extension: .py
#       format_name: light
#       format_version: '1.5'
#       jupytext_version: 1.9.1+dev
#   kernelspec:
#     display_name: Python [conda env:core_acc] *
#     language: python
#     name: conda-env-core_acc-py
# ---

# # KEGG enrichment of stable genes
#
# This notebooks looks at the group of most and least stable genes and performs a KEGG enrichment analysis to determine if there are any KEGG pathways that are significantly over-represented in our most or least stable gene sets.

# %load_ext autoreload
# %autoreload 2
# %matplotlib inline
import os
import pandas as pd
import numpy as np
import seaborn as sns
import matplotlib.pyplot as plt
import scipy.stats
import statsmodels.stats.multitest
from scripts import paths, utils, annotations

# Load KEGG pathway data
pao1_pathway_filename = "https://raw.githubusercontent.com/greenelab/adage/7a4eda39d360b224268921dc1f2c14b32788ab16/Node_interpretation/pseudomonas_KEGG_terms.txt"

pao1_pathways = annotations.load_format_KEGG(pao1_pathway_filename)
print(pao1_pathways.shape)
pao1_pathways.head()

# +
# Load transcriptional similarity df
# These are the subset of genes that we will consider
pao1_similarity_scores_filename = "pao1_similarity_scores_spell.tsv"

pao1_similarity_scores = pd.read_csv(
    pao1_similarity_scores_filename, sep="\t", header=0, index_col=0
)
# -

pao1_similarity_scores.head()

# Get most and least stable core genes based on label
pao1_most_stable_genes = list(
    pao1_similarity_scores[pao1_similarity_scores["label"] == "most stable"].index
)
pao1_least_stable_genes = list(
    pao1_similarity_scores[pao1_similarity_scores["label"] == "least stable"].index
)

# For each KEGG pathway, perform stat test, save p-values to get corrected p-values, report stats per pathway
def KEGG_enrichment_of_stable_genes(similarity_score_df, gene_list, kegg_df):
    """
    This function performs a KEGG enrichment using most or least stable genes,
    provided in `gene_list`
    """

    all_genes = set(similarity_score_df.index)
    module_genes = set(gene_list)
    not_module_genes = all_genes.difference(module_genes)

    rows = []
    # Find the KEGG pathway with significant over-representation
    for kegg_name in kegg_df.index:
        num_kegg_genes = kegg_df.loc[kegg_name, 1]
        kegg_genes = set(kegg_df.loc[kegg_name, 2])
        not_kegg_genes = all_genes.difference(kegg_genes)

        # Make contingency table
        # ---------------------| most stable  | not most stable
        # in KEGG pathway      | # genes      | # genes
        # not in KEGG pathway  | # genes     | # genes
        module_kegg_genes = module_genes.intersection(kegg_genes)
        not_module_kegg_genes = not_module_genes.intersection(kegg_genes)
        module_not_kegg_genes = module_genes.intersection(not_kegg_genes)
        not_module_not_kegg_genes = not_module_genes.intersection(not_kegg_genes)

        observed_contingency_table = np.array(
            [
                [len(module_kegg_genes), len(not_module_kegg_genes)],
                [len(module_not_kegg_genes), len(not_module_not_kegg_genes)],
            ]
        )
        # Fisher's exact test
        oddsr, pval = scipy.stats.fisher_exact(
            observed_contingency_table, alternative="greater"
        )
        # chi2 test will not accept 0 counts for the contingency table
        # chi2, pval, dof, expected_counts = scipy.stats.chi2_contingency(
        #    observed_contingency_table
        # )
        # print(oddsr, pval)

        rows.append(
            {
                "enriched KEGG pathway": kegg_name,
                "odds ratio": oddsr,
                "p-value": pval,
                "num shared genes": len(module_kegg_genes),
                "size gene set": len(module_genes),
                "size KEGG pathway": num_kegg_genes,
            }
        )

    enrichment_df = pd.DataFrame(rows)

    # Get corrected pvalues
    (
        reject_,
        pvals_corrected_,
        alphacSidak,
        alphacBonf,
    ) = statsmodels.stats.multitest.multipletests(
        enrichment_df["p-value"].values,
        alpha=0.05,
        method="fdr_bh",
        is_sorted=False,
    )

    enrichment_df["corrected p-value"] = pvals_corrected_
    enrichment_df["-log10 adj p-value"] = -np.log10(enrichment_df["corrected p-value"])

    return enrichment_df

pao1_most_stable_enrichment = KEGG_enrichment_of_stable_genes(
    pao1_similarity_scores, pao1_most_stable_genes, pao1_pathways
)

pao1_least_stable_enrichment = KEGG_enrichment_of_stable_genes(
    pao1_similarity_scores, pao1_least_stable_genes, pao1_pathways
)

print(pao1_most_stable_enrichment.shape)
pao1_most_stable_enrichment.sort_values(by="corrected p-value").head()

print(pao1_least_stable_enrichment.shape)
pao1_least_stable_enrichment.sort_values(by="corrected p-value").head()

# ## Plot

# +
# Create a shared scale to use for both plots
# np.linspace(3.759057e-22, 1)

# +
# Plot top most stable
pao1_most_stable_enrichment_top = pao1_most_stable_enrichment.sort_values(
    by="corrected p-value"
).head(10)
norm = plt.Normalize(0, 22)
cmap = sns.cubehelix_palette(as_cmap=True)

plt.figure(figsize=(8, 6))
f = plt.scatter(
    data=pao1_most_stable_enrichment_top,
    x="odds ratio",
    y="enriched KEGG pathway",
    s=300,
    c="-log10 adj p-value",
    norm=norm,
    cmap=cmap,
)
sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
cf = plt.colorbar(sm)
cf.ax.set_title("-log10 adj p-value", fontsize=12)

plt.title("Enrichment of most stable core genes", fontsize=16, y=1.05)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.xlabel("Odds ratio", fontsize=16)
plt.ylabel("")

# +
# Plot top most stable
pao1_least_stable_enrichment_top = pao1_least_stable_enrichment.sort_values(
    by="corrected p-value"
).head(10)
plt.figure(figsize=(8, 6))
norm = plt.Normalize(0.0, 22)
cmap = sns.cubehelix_palette(as_cmap=True)

g = plt.scatter(
    data=pao1_least_stable_enrichment_top,
    x="odds ratio",
    y="enriched KEGG pathway",
    s=300,
    c="-log10 adj p-value",
    norm=norm,
    cmap=cmap,
)

sm = plt.cm.ScalarMappable(cmap=cmap, norm=norm)
sm.set_array([])
cg = plt.colorbar(sm)
cg.ax.set_title("-log10 adj p-value", fontsize=12)

plt.title("Enrichment of least stable core genes", fontsize=16, y=1.05)
plt.xticks(fontsize=14)
plt.yticks(fontsize=14)
plt.xlabel("Odds ratio", fontsize=16)
plt.ylabel("")

# TO DO
# Update colorbar text to include label for -log10 p-value
# -

# Save
pao1_most_stable_enrichment.to_csv("pao1_most_stable_enrichment_spell.tsv", sep="\t")
pao1_least_stable_enrichment.to_csv("pao1_least_stable_enrichment_spell.tsv", sep="\t")

# **Takeaway:**
#
# KEGG enrichment analysis found that stable genes were significantly associated with essential functions: ribosome, lipopolysaccharide biosynthesis, citrate cycle. However, there are also pathways like the secretion systems, which allow for inter-strain warfare, that we’d expect to vary across strains but were found to be conserved (T3SS, T6SS) - looks like most of the genes annotated as T3/6SS are related to the secretion machinery, which is conserved across strains.
#
# There does NOT appear to be any significantly enriched KEGG pathways in the least stable core genes. However the least stable genes with the top enrichment score seem to be related to stress response
#
# KEGG enrichment: https://docs.google.com/spreadsheets/d/1lXZZXXjZSOuQ-cMOZ9I5llIJ2OBOQBur0Spso2WN4oY/edit#gid=0

max_stars_repo_name: greenelab/core-accessory-interactome

max_issues_repo_name: greenelab/core-accessory-interactome

max_forks_repo_name: greenelab/core-accessory-interactome

reason: import numpy,import scipy,import statsmodels


text:
import numpy as np
from collections import deque

class HistoryBuffer():
    def __init__(self,preprocess_fn,image_shape,frames_for_state) :
        self.buf = deque(maxlen=frames_for_state)
        self.preprocess_fn = preprocess_fn
        self.image_shape = image_shape
        self.clear()

    def clear(self) :
        for i in range(self.buf.maxlen):
            self.buf.append(np.zeros(self.image_shape,np.float32))

    def add(self,o) :
        #assert( list(o.shape) == self.image_shape ),'%s, %s'%(o.shape,self.image_shape)
        self.buf.append(self.preprocess_fn(o))
        state = np.concatenate([img for img in self.buf], axis=2)
        return state


text:
#!/usr/bin/env python2.7
import matplotlib
matplotlib.rc('xtick', labelsize=8)
matplotlib.use('Agg')
import matplotlib.pyplot as plt
import numpy as np
import os
import sys

def plot(xs, ys, data_dir):
    fig, ax = plt.subplots()
    ax.set_xticks(np.arange(0, 0.1, step=0.01))
    for tick in ax.get_xticklabels():
        tick.set_rotation(45)
    ax.set_xlim(0, 0.1)
    ax.plot(xs, ys)

    ax.set(xlabel='Delta (%)', ylabel='Number of NSS pexecs',
                  title='Impact of varying absolute segment tolerance')
    ax.grid()
    fig.savefig(os.path.join(data_dir, "plot.pdf"))
    plt.show()

def main(data_dir):
    data = []
    rel = []
    with open(os.path.join(data_dir, "nss_counts.csv")) as fh:
        for line in fh:
            line = line.strip()
            delta, count = line.split(",")

            if delta.endswith("%"):
                rel.append(True)
                data.append((delta[:-1], count))
            else:
                rel.append(False)
                data.append((delta, count))

    # There shouldn't be a mix of absolute and relative deltas
    assert rel.count(rel[0]) == len(rel)

    data = sorted(data, key=lambda t: t[0])
    xs = [t[0] for t in data]
    ys = [t[1] for t in data]
    plot(xs, ys, data_dir)

def usage():
    print("usage: plot.py <data-dir>")
    sys.exit(1)

if __name__ == "__main__":
    try:
        data_dir = sys.argv[1]
    except IndexError:
        usage()

    main(data_dir)


text:
from argparse import ArgumentParser
import os
import random
import string
import sys

import numpy as np

from logger import get_logger

logger = get_logger(__name__)

###
# file system
###

def make_dirs(path, empty=False):
    """
    create dir in path and clear dir if required
    """
    dir_path = os.path.dirname(path)
    os.makedirs(dir_path, exist_ok=True)

    if empty:
        files = [os.path.join(dir_path, item) for item in os.listdir(dir_path)]
        for item in files:
            if os.path.isfile(item):
                os.remove(item)

    return dir_path

def path_join(*paths, empty=False):
    """
    join paths and create dir
    """
    path = os.path.abspath(os.path.join(*paths))
    make_dirs(os.path.dirname(path), empty)
    return path

###
# data processing
###

def create_dictionary():
    """
    create char2id, id2char and vocab_size
    from printable ascii characters.
    """
    chars = sorted(ch for ch in string.printable if ch not in ("\x0b", "\x0c", "\r"))
    char2id = dict((ch, i + 1) for i, ch in enumerate(chars))
    char2id.update({"": 0})
    id2char = dict((char2id[ch], ch) for ch in char2id)
    vocab_size = len(char2id)
    return char2id, id2char, vocab_size

CHAR2ID, ID2CHAR, VOCAB_SIZE = create_dictionary()

def encode_text(text, char2id=CHAR2ID):
    """
    encode text to array of integers with CHAR2ID
    """
    return np.fromiter((char2id.get(ch, 0) for ch in text), int)

def decode_text(int_array, id2char=ID2CHAR):
    """
    decode array of integers to text with ID2CHAR
    """
    return "".join((id2char[ch] for ch in int_array))

def one_hot_encode(indices, num_classes):
    """
    one-hot encoding
    """
    return np.eye(num_classes)[indices]

def batch_generator(sequence, batch_size=64, seq_len=64, one_hot_features=False, one_hot_labels=False):
    """
    batch generator for sequence
    ensures that batches generated are continuous along axis 1
    so that hidden states can be kept across batches and epochs
    """
    # calculate effective length of text to use
    num_batches = (len(sequence) - 1) // (batch_size * seq_len)
    if num_batches == 0:
        raise ValueError("No batches created. Use smaller batch size or sequence length.")
    logger.info("number of batches: %s.", num_batches)
    rounded_len = num_batches * batch_size * seq_len
    logger.info("effective text length: %s.", rounded_len)

    x = np.reshape(sequence[: rounded_len], [batch_size, num_batches * seq_len])
    if one_hot_features:
        x = one_hot_encode(x, VOCAB_SIZE)
    logger.info("x shape: %s.", x.shape)

    y = np.reshape(sequence[1: rounded_len + 1], [batch_size, num_batches * seq_len])
    if one_hot_labels:
        y = one_hot_encode(y, VOCAB_SIZE)
    logger.info("y shape: %s.", y.shape)

    epoch = 0
    while True:
        # roll so that no need to reset rnn states over epochs
        x_epoch = np.split(np.roll(x, -epoch, axis=0), num_batches, axis=1)
        y_epoch = np.split(np.roll(y, -epoch, axis=0), num_batches, axis=1)
        for batch in range(num_batches):
            yield x_epoch[batch], y_epoch[batch]
        epoch += 1

###
# text generation
###

def generate_seed(text, seq_lens=(2, 4, 8, 16, 32)):
    """
    select subsequence randomly from input text
    """
    # randomly choose sequence length
    seq_len = random.choice(seq_lens)
    # randomly choose start index
    start_index = random.randint(0, len(text) - seq_len - 1)
    seed = text[start_index: start_index + seq_len]
    return seed

def sample_from_probs(probs, top_n=10):
    """
    truncated weighted random choice.
    """
    # need 64 floating point precision
    probs = np.array(probs, dtype=np.float64)
    # set probabilities after top_n to 0
    probs[np.argsort(probs)[:-top_n]] = 0
    # renormalise probabilities
    probs /= np.sum(probs)
    sampled_index = np.random.choice(len(probs), p=probs)
    return sampled_index

###
# main
###

def main(framework, train_main, generate_main):
    arg_parser = ArgumentParser(
        description="{} character embeddings LSTM text generation model.".format(framework))
    subparsers = arg_parser.add_subparsers(title="subcommands")

    # train args
    train_parser = subparsers.add_parser("train", help="train model on text file")
    train_parser.add_argument("--checkpoint-path", required=True,
                              help="path to save or load model checkpoints (required)")
    train_parser.add_argument("--text-path", required=True,
                              help="path of text file for training (required)")
    train_parser.add_argument("--restore", nargs="?", default=False, const=True,
                              help="whether to restore from checkpoint_path "
                                   "or from another path if specified")
    train_parser.add_argument("--seq-len", type=int, default=64,
                              help="sequence length of inputs and outputs (default: %(default)s)")
    train_parser.add_argument("--embedding-size", type=int, default=32,
                              help="character embedding size (default: %(default)s)")
    train_parser.add_argument("--rnn-size", type=int, default=128,
                              help="size of rnn cell (default: %(default)s)")
    train_parser.add_argument("--num-layers", type=int, default=2,
                              help="number of rnn layers (default: %(default)s)")
    train_parser.add_argument("--drop-rate", type=float, default=0.,
                              help="dropout rate for rnn layers (default: %(default)s)")
    train_parser.add_argument("--learning-rate", type=float, default=0.001,
                              help="learning rate (default: %(default)s)")
    train_parser.add_argument("--clip-norm", type=float, default=5.,
                              help="max norm to clip gradient (default: %(default)s)")
    train_parser.add_argument("--batch-size", type=int, default=64,
                              help="training batch size (default: %(default)s)")
    train_parser.add_argument("--num-epochs", type=int, default=32,
                              help="number of epochs for training (default: %(default)s)")
    train_parser.add_argument("--log-path", default=os.path.join(os.path.dirname(__file__), "main.log"),
                              help="path of log file (default: %(default)s)")
    train_parser.set_defaults(main=train_main)

    # generate args
    generate_parser = subparsers.add_parser("generate", help="generate text from trained model")
    generate_parser.add_argument("--checkpoint-path", required=True,
                                 help="path to load model checkpoints (required)")
    group = generate_parser.add_mutually_exclusive_group(required=True)
    group.add_argument("--text-path", help="path of text file to generate seed")
    group.add_argument("--seed", default=None, help="seed character sequence")
    generate_parser.add_argument("--length", type=int, default=1024,
                                 help="length of character sequence to generate (default: %(default)s)")
    generate_parser.add_argument("--top-n", type=int, default=3,
                                 help="number of top choices to sample (default: %(default)s)")
    generate_parser.add_argument("--log-path", default=os.path.join(os.path.dirname(__file__), "main.log"),
                                 help="path of log file (default: %(default)s)")
    generate_parser.set_defaults(main=generate_main)

    args = arg_parser.parse_args()
    get_logger("__main__", log_path=args.log_path, console=True)
    logger = get_logger(__name__, log_path=args.log_path, console=True)
    logger.debug("call: %s", " ".join(sys.argv))
    logger.debug("ArgumentParser: %s", args)

    try:
        args.main(args)
    except Exception as e:
        logger.exception(e)

max_stars_repo_name: yxtay/char-rnn-text-generation

max_issues_repo_name: fluffygravy/char-rnn-text-generation

max_forks_repo_name: fluffygravy/char-rnn-text-generation


text:
import numpy as np
import imageio
# val = [0, 3, 4, 23, 24, 40, 41, 42, 43, 45, 46, 48, 58, 59, 60]
val =[0,1,2,9,10,11,12,13,14,16,17,24,25,27,28]

for i in range(15):
    img_file = './data/nerf_synthetic/clevr_c_d4/train/r_{:d}.png'.format(val[i]+1)
    img=imageio.imread(img_file)

    mask_file1 = 'results/testing_2/segmentation/val_000200_r_{:d}_slot3.png'.format(i)
    mask1 = imageio.imread(mask_file1)

    mask1 = mask1 /255.
    mask1 = mask1[..., None]

    mask_file2 = 'results/testing_2/segmentation/val_000200_r_{:d}_slot2.png'.format(i)
    mask2 = imageio.imread(mask_file2)

    mask2 = mask2 /255.
    mask2 = mask2[..., None]

    output = img * (mask1+mask2)

    output_file = 'data/nerf_synthetic/clevr_c_d4/training/r_{:d}.png'.format(val[i]+1)
    imageio.imwrite(output_file, output)


text:
from os.path import join
import numpy as np
from parameters import *
from sklearn.model_selection import train_test_split

class DatasetLoader:
    def load_from_save(self):
        images = np.load(join(SAVE_DIRECTORY, SAVE_DATASET_IMAGES_FILENAME))
        images = images.reshape([-1, SIZE_FACE, SIZE_FACE, 1])
        labels = np.load(join(SAVE_DIRECTORY, SAVE_DATASET_LABELS_FILENAME)).\
            reshape([-1, len(EMOTIONS)])
        self._images, self._images_test, self._labels, self._labels_test = train_test_split\
            (images, labels, test_size=0.10, random_state=42)

    @property
    def images(self):
        return self._images

    @property
    def labels(self):
        return self._labels

    @property
    def images_test(self):
        return self._images_test

    @property
    def labels_test(self):
        return self._labels_test


text:
from __future__ import print_function, division, absolute_import

import sys
import functools

from numba import config

def allow_interpreter_mode(fn):
    """Temporarily re-enable intepreter mode
    """
    @functools.wraps(fn)
    def _core(*args, **kws):
        config.COMPATIBILITY_MODE = True
        try:
            fn(*args, **kws)
        finally:
            config.COMPATIBILITY_MODE = False
    return _core

def discover_tests(startdir):
    """Discover test under a directory
    """
    # Avoid importing unittest
    from numba import unittest_support as unittest
    loader = unittest.TestLoader()
    suite = loader.discover(startdir)
    return suite

def run_tests(suite, xmloutput=None, verbosity=1, nomultiproc=False):
    """
    args
    ----
    - suite [TestSuite]
        A suite of all tests to run
    - xmloutput [str or None]
        Path of XML output directory (optional)
    - verbosity [int]
        Verbosity level of tests output

    Returns the TestResult object after running the test *suite*.
    """
    from numba.tests import NumbaTestProgram

    if xmloutput is not None:
        import xmlrunner
        runner = xmlrunner.XMLTestRunner(output=xmloutput)
    else:
        runner = None
    prog = NumbaTestProgram(suite=suite, testRunner=runner, exit=False,
                            verbosity=verbosity,
                            nomultiproc=nomultiproc)
    return prog.result

def test(**kwargs):
    """
    Run all tests under ``numba.tests``.

    kwargs
    ------
    - descriptions
    - verbosity
    - buffer
    - failfast
    - xmloutput [str]
        Path of XML output directory
    """
    from numba import cuda

    suite = discover_tests("numba.tests")
    ok = run_tests(suite, **kwargs).wasSuccessful()
    if not ok:
        return ok

    # Test no cuda tests
    from numba.cuda.tests.nocuda.runtests import test as nocuda_test
    ok = nocuda_test()

    # Test CUDA
    if ok:
        if cuda.is_available():
            gpus = cuda.list_devices()
            if gpus and gpus[0].compute_capability >= (2, 0):
                print("== Run CUDA tests ==")
                ok = cuda.test()
            else:
                print("== Skipped CUDA tests because GPU CC < 2.0 ==")
        else:
            print("== Skipped CUDA tests ==")

    return ok

if __name__ == "__main__":
    sys.exit(0 if test() else 1)

max_stars_repo_name: gdementen/numba

max_issues_repo_name: gdementen/numba

max_forks_repo_name: gdementen/numba

reason: from numba


text:
# -*- coding: utf-8 -*-
# Copyright (C) 2020-2021 by SCICO Developers
# All rights reserved. BSD 3-clause License.
# This file is part of the SPORCO package. Details of the copyright
# and user license can be found in the 'LICENSE.txt' file distributed
# with the package.

r"""Extensions of numpy ndarray class.

 .. testsetup::

   >>> import scico.numpy as snp
   >>> from scico.blockarray import BlockArray
   >>> import numpy as np
   >>> import jax.numpy

The class :class:`.BlockArray` is a `jagged array
<https://en.wikipedia.org/wiki/Jagged_array>`_ that aims to mimic the :class:`numpy.ndarray`
interface where appropriate.

A :class:`.BlockArray` object consists of a tuple of `DeviceArray` objects that share their memory
buffers with non-overlapping, contiguous regions of a common one-dimensional `DeviceArray`.
A :class:`.BlockArray` contains the following size attributes:

* `shape`:  A tuple of tuples containing component dimensions.
* `size`: The sum of the size of each component block; this is the length of the underlying
  one-dimensional `DeviceArray`.
* `num_blocks`: The number of components (blocks) that comprise the :class:`.BlockArray`.

Motivating Example
------------------

Consider a two dimensional array :math:`\mb{x} \in \mathbb{R}^{n \times m}`.

We compute the discrete differences of :math:`\mb{x}` in the horizontal and vertical directions,
generating two new arrays: :math:`\mb{x}_h \in \mathbb{R}^{n \times (m-1)}` and :math:`\mb{x}_v \in
\mathbb{R}^{(n-1) \times m}`.

As these arrays are of different sizes, we cannot combine them into a single `ndarray`. Instead,
we might vectorize each array and concatenate the resulting vectors, leading to :math:`\mb{\bar{x}}
\in \mathbb{R}^{n(m-1) + m(n-1)}`, which can be stored as a one-dimensional `ndarray`.
Unfortunately, this makes it hard to access the individual components :math:`\mb{x}_h` and
:math:`\mb{x}_v`.

Instead, we can form a :class:`.BlockArray`: :math:`\mb{x}_B = [\mb{x}_h, \mb{x}_v]`

  ::

    >>> n = 32
    >>> m = 16
    >>> x_h = np.random.randn(n, m-1)
    >>> x_v = np.random.randn(n-1, m)

    # Form the blockarray
    >>> x_B = BlockArray.array([x_h, x_v])

    # The blockarray shape is a tuple of tuples
    >>> x_B.shape
    ((32, 15), (31, 16))

    # Each block component can be easily accessed
    >>> x_B[0].shape
    (32, 15)
    >>> x_B[1].shape
    (31, 16)

Constructing a BlockArray
-------------------------

Construct from a tuple of arrays (either `ndarray` or `DeviceArray`)
####################################################################
  .. doctest::

     >>> from scico.blockarray import BlockArray
     >>> import numpy as np
     >>> x0 = np.random.randn(32, 32)
     >>> x1 = np.random.randn(16)
     >>> X = BlockArray.array((x0, x1))
     >>> X.shape
     ((32, 32), (16,))
     >>> X.size
     1040
     >>> X.num_blocks
     2

| While :func:`.BlockArray.array` will accept either `ndarray` or `DeviceArray` as input, the
  resulting :class:`.BlockArray` will be backed by a `DeviceArray` memory buffer.

| **Note**: constructing a :class:`.BlockArray` always involves a copy to a new `DeviceArray` memory buffer.

| **Note**: by default, the resulting :class:`.BlockArray` cast to single precision and will have
  dtype `float32` or `complex64`.

Construct from a single vector and tuple of shapes
##################################################
  ::

     >>> x_flat = np.random.randn(1040)
     >>> shape_tuple = ((32, 32), (16,))
     >>> X = BlockArray.array_from_flattened(x_flat, shape_tuple=shape_tuple)
     >>> X.shape
     ((32, 32), (16,))

Operating on a BlockArray
-------------------------

Indexed Updating
----------------

BlockArrays support the JAX DeviceArray `indexed update syntax
<https://jax.readthedocs.io/en/latest/jax.ops.html#indexed-update-operators>`_

The index must be of the form [ibk] or [ibk,idx],
where `ibk` is the index of the block to be updated, and `idx` is a
general index of the elements to be updated in that block.  In particular, `ibk`
cannot be a `slice`.  The general index `idx` can be omitted, in which case
an entire block is updated.

==============================   ==============================================
Alternate syntax                 Equivalent in-place expression
==============================   ==============================================
``x.at[ibk, idx].set(y)``        ``x[ibk, idx] = y``
``x.at[ibk, idx].add(y)``        ``x[ibk, idx] += y``
``x.at[ibk, idx].multiply(y)``   ``x[ibk, idx] *= y``
``x.at[ibk, idx].divide(y)``     ``x[ibk, idx] /= y``
``x.at[ibk, idx].power(y)``      ``x[ibk, idx] **= y``
``x.at[ibk, idx].min(y)``        ``x[ibk, idx] = np.minimum(x[idx], y)``
``x.at[ibk, idx].max(y)``        ``x[ibk, idx] = np.maximum(x[idx], y)``
==============================   ==============================================

Arithmetic and Broadcasting
###########################

Suppose :math:`\mb{x}` is a BlockArray with shape :math:`((n, n), (m,))`.

  ::

    >>> x1 = np.random.randn(4, 4)
    >>> x2 = np.random.randn(5)
    >>> x = BlockArray.array( (x1, x2) )
    >>> x.shape
    ((4, 4), (5,))
    >>> x.num_blocks
    2
    >>> x.size  # 4*4 + 5
    21

Illustrated for the operation ``+``, but equally valid for operators
``+, -, *, /, //, **, <, <=, >, >=, ==``

Operations with BlockArrays with same number of blocks
******************************************************

Let :math:`\mb{y}` be a BlockArray with the same number of blocks as :math:`\mb{x}`.

  .. math::
     \mb{x} + \mb{y}
     =
     \begin{bmatrix}
       \mb{x}[0] + \mb{y}[0] \\
       \mb{x}[1] + \mb{y}[1] \\
     \end{bmatrix}

This operation depends on pair of blocks from :math:`\mb{x}` and :math:`\mb{y}`
being broadcastable against each other.

Operations with a scalar
************************

The scalar is added to each element of the :class:`.BlockArray`:

  .. math::
     \mb{x} + 1
     =
     \begin{bmatrix}
       \mb{x}[0] + 1 \\
       \mb{x}[1] + 1\\
     \end{bmatrix}

  ::

     >>> y = x + 1
     >>> np.testing.assert_allclose(y[0], x[0] + 1)
     >>> np.testing.assert_allclose(y[1], x[1] + 1)

Operations with a 1D `ndarray` of size equal to `num_blocks`
************************************************************

The *i*\th scalar is added to the *i*\th  block of the :class:`.BlockArray`:

  .. math::
   \mb{x}
     +
   \begin{bmatrix}
     1 \\
     2
   \end{bmatrix}
   =
   \begin{bmatrix}
     \mb{x}[0] + 1 \\
     \mb{x}[1] + 2\\
   \end{bmatrix}

  ::

     >>> y = x + np.array([1, 2])
     >>> np.testing.assert_allclose(y[0], x[0] + 1)
     >>> np.testing.assert_allclose(y[1], x[1] + 2)

Operations with an ndarray of `size` equal to :class:`.BlockArray` size
***********************************************************************

We first cast the `ndarray` to a BlockArray with same shape as :math:`\mb{x}`, then apply the
operation on the resulting BlockArrays.  With ``y.size = x.size``, we have:

  .. math::
   \mb{x}
   +
   \mb{y}
   =
   \begin{bmatrix}
     \mb{x}[0] + \mb{y}[0] \\
     \mb{x}[1] + \mb{y}[1]\\
   \end{bmatrix}

Equivalently, the BlockArray is first flattened, then added to the flattened `ndarray`, and the
result is reformed into a BlockArray with the same shape as :math:`\mb{x}`

MatMul
######

Between two BlockArrays
***********************

The matmul is computed between each block of the two BlockArrays.

The BlockArrays must have the same number of blocks, and each pair of blocks
must be broadcastable.

  .. math::
   \mb{x} @ \mb{y}
   =
   \begin{bmatrix}
     \mb{x}[0] @ \mb{y}[0] \\
     \mb{x}[1] @ \mb{y}[1]\\
   \end{bmatrix}

Between BlockArray and Ndarray/DeviceArray
******************************************

This operation is not defined.

Between BlockArray and :class:`.LinearOperator`
***********************************************

.. todo::
   Improve this

The :class:`.Operator` and :class:`.LinearOperator` classes are designed to work on :class:`.BlockArray`.  The shapes must conform:

   ::

      x = BlockArray.array((np.random.randn(32, 32), np.random.randn(16)))
      A # LinearOperator with A.shape[1] == x.shape
      A @ x

NumPy ufuncs
############

`NumPy universal functions (ufuncs) <https://numpy.org/doc/stable/reference/ufuncs.html>`_ are
functions that operate on an `ndarray` on an element-by-element fashion and support array
broadcasting. Examples of ufuncs are ``abs``, ``sign``, ``conj``, and ``exp``.

The JAX library implements most NumPy ufuncs in the :mod:`jax.numpy` module.
However, as JAX does not support subclassing of `DeviceArray`, the JAX ufuncs
cannot be used on :class:`.BlockArray`. As a workaround, we have wrapped several
JAX ufuncs for use on :class:`.BlockArray`; these are located in the
:mod:`scico.numpy` module.

Reductions
##########

Reductions are functions that take an array-like as an input and return an array of lower
dimension. Examples include ``mean``, ``sum``, ``norm``. BlockArray reductions are located in the
:mod:`scico.numpy` module

:class:`.BlockArray` tries to mirror `ndarray` reduction semantics where possible, but
cannot provide a one-to-one match as the block components may be of different size.

Consider the example BlockArray

  .. math::
   \mb{x} = \begin{bmatrix}
   \begin{bmatrix}
    1 & 1 \\
     1 & 1
   \end{bmatrix} \\
   \begin{bmatrix}
    2 \\
    2
   \end{bmatrix}
   \end{bmatrix}.

We have

  .. doctest::

    >>> import scico.numpy as snp
    >>> x = BlockArray.array((np.ones((2,2)), 2*np.ones((2))))
    >>> x.shape
    ((2, 2), (2,))
    >>> x.size
    6
    >>> x.num_blocks
    2

  If no axis is specified, the reduction is applied to the flattened array:

  .. doctest::

    >>> snp.sum(x, axis=None).item()
    8.0

  Reducing along the 0-th axis crushes the `BlockArray` down into a single `DeviceArray`
  and requires all blocks to have the same shape otherwise, an error is raised.

  .. doctest::

    >>> snp.sum(x, axis=0)
    Traceback (most recent call last):
    ValueError: Evaluating sum of BlockArray along axis=0 requires all blocks to be same shape; got ((2, 2), (2,))

    >>> y = BlockArray.array((np.ones((2,2)), 2*np.ones((2, 2))))
    >>> snp.sum(y, axis=0)
    DeviceArray([[3., 3.],
                 [3., 3.]], dtype=float32)

  Reducing along axis :math:`n` is equivalent to reducing each component along axis :math:`n-1`:

  .. math::
   \text{sum}(x, axis=1) = \begin{bmatrix}
   \begin{bmatrix}
    2 \\
     2
   \end{bmatrix} \\
   \begin{bmatrix}
    4 \\
   \end{bmatrix}
   \end{bmatrix}

  If a component does not have axis :math:`n-1`, the reduction is not applied to that component. In this example,
  ``x[1].ndim == 1``, so no reduction is applied to block ``x[1]``.

  .. math::
   \text{sum}(x, axis=2) = \begin{bmatrix}
   \begin{bmatrix}
    2 \\
     2
   \end{bmatrix} \\
   \begin{bmatrix}
    2 \\
    2
   \end{bmatrix}
   \end{bmatrix}

Code version

  .. doctest::

    >>> snp.sum(x, axis=1)  # doctest: +SKIP
    BlockArray([[2, 2],
                [4,] ])

    >>> snp.sum(x, axis=2)  # doctest: +SKIP
    BlockArray([ [2, 2],
                 [2,] ])

"""

from __future__ import annotations

from functools import wraps
from typing import Iterator, List, Optional, Tuple, Union

import numpy as np

import jax
import jax.numpy as jnp
from jax import core
from jax.interpreters import xla
from jax.interpreters.xla import DeviceArray
from jax.tree_util import register_pytree_node, tree_flatten

from jaxlib.xla_extension import Buffer

from scico.typing import Axes, BlockShape, DType, JaxArray, Shape
from scico.util import is_nested

_arraylikes = (Buffer, DeviceArray, np.ndarray)

__author__ = """\n""".join(["Luke Pfister <pfister@lanl.gov>", "Brendt Wohlberg <brendt@ieee.org>"])

def atleast_1d(*arys):
    """Convert inputs to arrays with at least one dimension.

    BlockArrays are returned unmodified.

    LAX-backend implementation of :func:`atleast_1d`.

    The JAX version of this function will return a copy rather than a view of the input.

    *Original docstring below.*

    Scalar inputs are converted to 1-dimensional arrays, whilst
        higher-dimensional inputs are preserved.

    Args:
        arys1, arys2, ... : One or more input arrays (array_like).

    Returns:
        An array, or list of arrays, each with ``a.ndim >= 1``. Copies are made only if
        necessary.
    """

    if len(arys) == 1:
        arr = arys[0]
        return arr if isinstance(arr, BlockArray) else jnp.atleast_1d(arr)
    else:
        out = []
        for arr in arys:
            if isinstance(arr, BlockArray):
                out.append(arr)
            else:
                out.append(jnp.atleast_1d(arr))
        return out

def reshape(
    a: Union[JaxArray, BlockArray], newshape: Union[Shape, BlockShape]
) -> Union[JaxArray, BlockArray]:
    """Gives a new shape to an array without changing its data.

    Args:
        a : Array to be reshaped.
        newshape:  The new shape should be compatible with the original shape. If an integer,
            then the result will be a 1-D array of that length.  One shape dimension can be -1.
            In this case, the value is inferred from the length of the array and remaining
            dimensions.  If a tuple of tuple of ints, a :class:`.BlockArray` is returned.

    Returns:
        The reshaped array.  Unlike :func:`numpy.reshape`, a copy is always returned.
    """

    if is_nested(newshape):
        # x is a blockarray
        return BlockArray.array_from_flattened(a, newshape)
    else:
        return jnp.reshape(a, newshape)

def block_sizes(shape: Union[Shape, BlockShape]) -> Axes:
    r"""Computes the 'sizes' of (possibly nested) block shapes.

    This function computes ``block_sizes(z.shape) == (_.size for _ in z)``

    Args:
       shape:  A shape tuple; possibly containing nested tuples.

    Examples:

    .. doctest::
        >>> import scico.numpy as snp

        >>> x = BlockArray.ones( ( (4, 4), (2,)))
        >>> x.size
        18

        >>> y = snp.ones((3, 3))
        >>> y.size
        9

        >>> z = BlockArray.array([x, y])
        >>> block_sizes(z.shape)
        (18, 9)

        >>> zz = BlockArray.array([z, z])
        >>> block_sizes(zz.shape)
        (27, 27)
    """

    if isinstance(shape, BlockArray):
        raise TypeError(
            "Expected a `shape` (possibly nested tuple of ints); got :class:`.BlockArray`."
        )

    out = []
    if is_nested(shape):
        # shape is nested -> at least one element came from a blockarray
        for y in shape:
            if is_nested(y):
                # recursively calculate the block size until we arrive at
                # a tuple (shape of a non-block array)
                while is_nested(y):
                    y = block_sizes(y)
                out.append(np.sum(y))  # adjacent block sizes are added together
            else:
                # this is a tuple; size given by product of elements
                out.append(np.prod(y))
        return tuple(out)
    else:
        # shape is a non-nested tuple; return the product
        return np.prod(shape)

def _flatten_blockarrays(inp, *args, **kwargs):
    """Flattens any blockarrays present in inp, args, or kwargs"""

    def _flatten_if_blockarray(inp):
        if isinstance(inp, BlockArray):
            return inp._data
        else:
            return inp

    inp_ = _flatten_if_blockarray(inp)
    args_ = (_flatten_if_blockarray(_) for _ in args)
    kwargs_ = {key: _flatten_if_blockarray(val) for key, val in kwargs.items()}
    return inp_, args_, kwargs_

def _block_array_ufunc_wrapper(func):
    """Wraps a "ufunc" to allow for joint operation on `DeviceArray` and `BlockArray`"""

    @wraps(func)
    def wrapper(inp, *args, **kwargs):
        all_args = (inp,) + args + tuple(kwargs.items())
        if any([isinstance(_, BlockArray) for _ in all_args]):
            # If 'inp' is a BlockArray, call func on inp._data
            # Then return a BlockArray of the same shape as inp

            inp_, args_, kwargs_ = _flatten_blockarrays(inp, *args, **kwargs)
            flat_out = func(inp_, *args_, **kwargs_)
            return BlockArray.array_from_flattened(flat_out, inp.shape)
        else:
            # Otherwise call the function normally
            return func(inp, *args, **kwargs)

    if not hasattr(func, "__doc__") or func.__doc__ is None:
        return wrapper
    else:
        wrapper.__doc__ = (
            f":func:`{func.__name__}` wrapped to operate on :class:`BlockArray`"
            + "\n\n"
            + func.__doc__
        )
    return wrapper

def _block_array_reduction_wrapper(func):
    """Wraps a reduction (eg sum, norm) to allow for joint operation on `DeviceArray` and
    `BlockArray`"""

    @wraps(func)
    def wrapper(inp, *args, axis=None, **kwargs):

        all_args = (inp,) + args + tuple(kwargs.items())
        if any([isinstance(_, BlockArray) for _ in all_args]):
            if axis is None:
                # Treat as a single long vector
                inp_, args_, kwargs_ = _flatten_blockarrays(inp, *args, **kwargs)
                return func(inp_, *args_, **kwargs_)
            elif type(axis) == tuple:
                raise Exception(
                    f"""Evaluating {func.__name__} on a BlockArray with a tuple argument to
                    axis is not currently supported"""
                )
            else:
                if axis == 0:  # reduction along block axis
                    # reduction along axis=0 only makes sense if all blocks are the same shape
                    # so we can convert to a standard DeviceArray of shape (inp.num_blocks, ...)
                    # and reduce along axis = 0
                    if all([bk_shape == inp.shape[0] for bk_shape in inp.shape]):
                        view_shape = (inp.num_blocks,) + inp.shape[0]
                        return func(inp._data.reshape(view_shape), *args, axis=0, **kwargs)
                    else:
                        raise ValueError(
                            f"Evaluating {func.__name__} of BlockArray along axis=0 requires "
                            f"all blocks to be same shape; got {inp.shape}"
                        )
                else:
                    # Reduce each block individually along axis-1
                    out = []
                    for bk in inp:
                        if isinstance(bk, BlockArray):
                            # This block is itself a blockarray, so call this wrapped reduction
                            # on axis-1
                            tmp = _block_array_reduction_wrapper(func)(
                                bk, *args, axis=axis - 1, **kwargs
                            )
                        else:
                            if axis - 1 >= bk.ndim:
                                # Trying to reduce along a dim that doesn't exist for this block,
                                # so just return the block.
                                # ie broadcast to shape (..., 1) and reduce along axis=-1
                                tmp = bk
                            else:
                                tmp = func(bk, *args, axis=axis - 1, **kwargs)
                        out.append(atleast_1d(tmp))
                    return BlockArray.array(out)

        elif axis is None:
            # 'axis' might not be a valid kwarg (eg dot, vdot), so don't pass it
            return func(inp, *args, **kwargs)
        else:
            return func(inp, *args, axis=axis, **kwargs)

    if not hasattr(func, "__doc__") or func.__doc__ is None:
        return wrapper
    else:
        wrapper.__doc__ = (
            f":func:`{func.__name__}` wrapped to operate on :class:`BlockArray`"
            + "\n\n"
            + func.__doc__
        )
    return wrapper

def _block_array_matmul_wrapper(func):
    @wraps(func)
    def wrapper(self, other):
        if isinstance(self, BlockArray):
            if isinstance(other, BlockArray):
                # Both blockarrays, work block by block
                return BlockArray.array([func(x, y) for x, y in zip(self, other)])
            else:
                raise TypeError(
                    f"Operation {func.__name__} not implemented between {type(self)} and {type(other)}"
                )
        else:
            return func(self, other)

    if not hasattr(func, "__doc__") or func.__doc__ is None:
        return wrapper
    else:
        wrapper.__doc__ = (
            f":func:`{func.__name__}` wrapped to operate on :class:`BlockArray`"
            + "\n\n"
            + func.__doc__
        )
    return wrapper

def _block_array_binary_op_wrapper(func):
    """Returns a decorator that performs type and shape checking for :class:`.BlockArray`
    arithmetic
    """

    @wraps(func)
    def wrapper(self, other):
        if isinstance(other, BlockArray):
            if other.shape == self.shape:
                # Same shape blocks, can operate on flattened arrays
                return BlockArray.array_from_flattened(func(self._data, other._data), self.shape)
            elif other.num_blocks == self.num_blocks:
                # Will work as long as the shapes are broadcastable
                return BlockArray.array([func(x, y) for x, y in zip(self, other)])
            else:
                raise ValueError(
                    f"operation not valid on operands with shapes {self.shape}  {other.shape}"
                )
        elif any([isinstance(other, _) for _ in _arraylikes]):
            if other.size == 1:
                # Same as operating on a scalar
                return BlockArray.array_from_flattened(func(self._data, other), self.shape)
            elif other.size == self.size:
                # A little fast and loose, treat the block array as a length self.size vector
                return BlockArray.array_from_flattened(func(self._data, other), self.shape)
            elif other.size == self.num_blocks:
                return BlockArray.array([func(blk, other_) for blk, other_ in zip(self, other)])
            else:
                raise ValueError(
                    f"operation not valid on operands with shapes {self.shape}  {other.shape}"
                )
        elif jnp.isscalar(other) or isinstance(other, core.Tracer):
            return BlockArray.array_from_flattened(func(self._data, other), self.shape)
        else:
            raise TypeError(
                f"Operation {func.__name__} not implemented between {type(self)} and {type(other)}"
            )

    if not hasattr(func, "__doc__") or func.__doc__ is None:
        return wrapper
    else:
        wrapper.__doc__ = (
            f":func:`{func.__name__}` wrapped to operate on :class:`BlockArray`"
            + "\n\n"
            + func.__doc__
        )
    return wrapper

class _AbstractBlockArray(core.ShapedArray):
    """Abstract BlockArray class for JAX tracing.

    See https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primitives_work.html

    """

    array_abstraction_level = 0  # Same as jax.core.ConcreteArray

    def __init__(self, shapes, dtype):

        sizes = block_sizes(shapes)
        size = np.sum(sizes)

        super(_AbstractBlockArray, self).__init__((size,), dtype)

        #: Abstract data value
        self._data_aval: core.ShapedArray = core.ShapedArray((size,), dtype)

        #: Array dtype
        self.dtype: DType = dtype

        #: Shape of each block
        self.shapes: BlockShape = shapes

        #: Size of each block
        self.sizes: Shape = sizes

        #: Array specifying boundaries of components as indices in base array
        self.bndpos: np.ndarray = np.r_[0, np.cumsum(sizes)]

    @core.aval_property
    def data(self):
        return bk_data_p.bind(self)

# Register BlockArray._data as a primitive
bk_data_p = core.Primitive("bk_data")

@bk_data_p.def_impl
def _bk_data_impl(mat):
    return mat._data

# The Jax class is heavily inspired by SparseArray/AbstractSparseArray here:
# https://github.com/google/jax/blob/7724322d1c08c13008815bfb52759a29c2a6823b/tests/custom_object_test.py
class BlockArray:
    """A tuple of :class:`jax.interpreters.xla.DeviceArray` objects.

    A tuple of `DeviceArray` objects that all share their memory buffers with
    non-overlapping, contiguous regions of a common one-dimensional `DeviceArray`.
    It can be used as the common one-dimensional array via the :func:`BlockArray.ravel`
    method, or individual component arrays can be accessed individually.
    """

    # Ensure we use BlockArray.__radd__,__rmul__, etc for binary operations of the form
    # op(np.ndarray, BlockArray)
    # See https://docs.scipy.org/doc/numpy-1.10.1/user/c-info.beyond-basics.html#ndarray.__array_priority__
    __array_priority__ = 1

    def __init__(self, aval: _AbstractBlockArray, data: JaxArray):
        """BlockArray init method.

        Args:
            aval: `Abstract value`_  associated to this array (shape+dtype+weak_type)
            data: The underlying contiguous, flattened `DeviceArray`.

        .. _Abstract value:  https://jax.readthedocs.io/en/latest/notebooks/How_JAX_primitives_work.html

        """
        self._aval = aval
        self._data = data

    def __repr__(self):
        return "scico.blockarray.BlockArray: \n" + self._data.__repr__()

    def __getitem__(self, idx: Union[int, Ellipsis]) -> JaxArray:
        if isinstance(idx, slice):
            raise TypeError(f"Slicing not supported on block index")
        elif idx == Ellipsis:
            return reshape(self._data, self.shape)
        elif idx < 0:
            idx = self.num_blocks + idx
        return reshape(self._data[self.bndpos[idx] : self.bndpos[idx + 1]], self.shape[idx])

    @_block_array_matmul_wrapper
    def __matmul__(self, other: Union[np.ndarray, BlockArray, JaxArray]) -> JaxArray:
        return self @ other

    @_block_array_matmul_wrapper
    def __rmatmul__(self, other: Union[np.ndarray, BlockArray, JaxArray]) -> JaxArray:
        return other @ self

    @_block_array_binary_op_wrapper
    def __sub__(a, b):
        return a - b

    @_block_array_binary_op_wrapper
    def __rsub__(a, b):
        return b - a

    @_block_array_binary_op_wrapper
    def __mul__(a, b):
        return a * b

    @_block_array_binary_op_wrapper
    def __rmul__(a, b):
        return a * b

    @_block_array_binary_op_wrapper
    def __add__(a, b):
        return a + b

    @_block_array_binary_op_wrapper
    def __radd__(a, b):
        return a + b

    @_block_array_binary_op_wrapper
    def __truediv__(a, b):
        return a / b

    @_block_array_binary_op_wrapper
    def __rtruediv__(a, b):
        return b / a

    @_block_array_binary_op_wrapper
    def __floordiv__(a, b):
        return a // b

    @_block_array_binary_op_wrapper
    def __rfloordiv__(a, b):
        return b // a

    @_block_array_binary_op_wrapper
    def __pow__(a, b):
        return a ** b

    @_block_array_binary_op_wrapper
    def __rpow__(a, b):
        return b ** a

    @_block_array_binary_op_wrapper
    def __gt__(a, b):
        return a > b

    @_block_array_binary_op_wrapper
    def __ge__(a, b):
        return a >= b

    @_block_array_binary_op_wrapper
    def __lt__(a, b):
        return a < b

    @_block_array_binary_op_wrapper
    def __le__(a, b):
        return a <= b

    @_block_array_binary_op_wrapper
    def __eq__(a, b):
        return a == b

    @_block_array_binary_op_wrapper
    def __ne__(a, b):
        return a != b

    def __iter__(self) -> Iterator[int]:
        for i in range(self.num_blocks):
            yield self[i]

    @property
    def blocks(self) -> Iterator[int]:
        """Returns an iterator yielding component blocks."""
        return self.__iter__()

    @property
    def bndpos(self) -> np.ndarray:
        """Array specifying boundaries of components as indices in base array."""
        return self._aval.bndpos

    @property
    def dtype(self) -> DType:
        """Array dtype"""
        return self._data.dtype

    @property
    def device_buffer(self) -> Buffer:
        """The :class:`jaxlib.xla_extension.Buffer` that backs the underlying data array"""
        return self._data.device_buffer

    @property
    def size(self) -> int:
        """Total number of elements in the array."""
        return self._aval.size

    @property
    def num_blocks(self) -> int:
        """Number of :class:`.BlockArray` components."""

        return len(self.shape)

    @property
    def ndim(self) -> Shape:
        """Tuple of component ndims."""

        return tuple(len(c) for c in self.shape)

    @property
    def shape(self) -> BlockShape:
        """Tuple of component shapes."""
        return self._aval.shapes

    @property
    def split(self) -> Tuple[JaxArray, ...]:
        """Tuple of component arrays."""

        return tuple(self[k] for k in range(self.num_blocks))

    def conj(self) -> BlockArray:
        """Returns a :class:`.BlockArray` with complex-conjugated elements."""

        # Much faster than BlockArray.array([_.conj() for _ in self.blocks])
        return BlockArray.array_from_flattened(self.ravel().conj(), self.shape)

    @property
    def real(self) -> BlockArray:
        """Returns a :class:`.BlockArray` with the real part of this array."""
        return BlockArray.array_from_flattened(self.ravel().real, self.shape)

    @property
    def imag(self) -> BlockArray:
        """Returns a :class:`.BlockArray` with the imaginary part of this array."""
        return BlockArray.array_from_flattened(self.ravel().imag, self.shape)

    @classmethod
    def array(
        cls, alst: List[Union[np.ndarray, JaxArray]], dtype: Optional[np.dtype] = None
    ) -> BlockArray:
        """Construct a :class:`.BlockArray` from a list or tuple of existing array-like.

        Args:
            alst : Initializers for array components.
                Can be :class:`numpy.ndarray` or :class:`jax.interpreters.xla.DeviceArray`
            dtype : Data type of array. If none, dtype is derived from dtype of initializers

        Returns:
            :class:`.BlockArray` initialized from `alst` tuple
        """

        if isinstance(alst, (tuple, list)) is False:
            raise TypeError("Input to `array` must be a list or tuple of existing arrays")

        if dtype is None:
            present_types = jax.tree_flatten(jax.tree_map(lambda x: x.dtype, alst))[0]
            dtype = np.find_common_type(present_types, [])

        # alst can be a list/tuple of arrays, or a list/tuple containing list/tuples of arrays
        # consider alst to be a tree where leaves are arrays (possibly abstract arrays)
        # use tree_map to find the shape of each leaf
        # `shapes` will be a tuple of ints and tuples containing ints (possibly nested further)

        # ensure any scalar leaves are converted to (1,) arrays
        def shape_atleast_1d(x):
            return x.shape if x.shape != () else (1,)

        shapes = tuple(
            jax.tree_map(shape_atleast_1d, alst, is_leaf=lambda x: not isinstance(x, (list, tuple)))
        )

        _aval = _AbstractBlockArray(shapes, dtype)
        data_ravel = jnp.hstack(jax.tree_map(lambda x: x.ravel(), jax.tree_flatten(alst)[0]))
        return cls(_aval, data_ravel)

    @classmethod
    def array_from_flattened(
        cls, data_ravel: Union[np.ndarray, JaxArray], shape_tuple: BlockShape
    ) -> BlockArray:
        """Construct a :class:`.BlockArray` from a flattened array and tuple of shapes.

        Args:
            data_ravel: Flattened data array
            shape_tuple: Tuple of tuples containing desired block shapes.

        Returns:
            :class:`.BlockArray` initialized from `data_ravel` and `shape_tuple`
        """

        if not isinstance(data_ravel, DeviceArray):
            data_ravel = jax.device_put(data_ravel)

        shape_tuple_size = np.sum(block_sizes(shape_tuple))

        if shape_tuple_size != data_ravel.size:
            raise ValueError(
                f"""The specified shape_tuple is incompatible with provided data_ravel
    shape_tuple = {shape_tuple}
    shape_tuple_size = {shape_tuple_size}
    len(data_ravel) = {len(data_ravel)}
            """
            )

        _aval = _AbstractBlockArray(shape_tuple, dtype=data_ravel.dtype)
        return cls(_aval, data_ravel)

    @classmethod
    def ones(cls, shape_tuple: BlockShape, dtype: DType = np.float32) -> BlockArray:
        """
        Return a new :class:`.BlockArray` with given block shapes and type, filled with ones.

        Args:
            shape_tuple: Tuple of shapes for component blocks
            dtype: Desired data-type for the :class:`.BlockArray`. Default is `numpy.float32`.

        Returns:
            :class:`.BlockArray` of ones with the given component shapes and dtype
        """
        _aval = _AbstractBlockArray(shape_tuple, dtype=dtype)
        data_ravel = jnp.ones(_aval.size, dtype=dtype)
        return cls(_aval, data_ravel)

    @classmethod
    def zeros(cls, shape_tuple: BlockShape, dtype: DType = np.float32) -> BlockArray:
        """
        Return a new :class:`.BlockArray` with given block shapes and type, filled with zeros.

        Args:
            shape_tuple: Tuple of shapes for component blocks
            dtype: Desired data-type for the :class:`.BlockArray`. Default is `numpy.float32`.

        Returns:
            :class:`.BlockArray` of zeros with the given component shapes and dtype
        """
        _aval = _AbstractBlockArray(shape_tuple, dtype=dtype)
        data_ravel = jnp.zeros(_aval.size, dtype=dtype)
        return cls(_aval, data_ravel)

    @classmethod
    def empty(cls, shape_tuple: BlockShape, dtype: DType = np.float32) -> BlockArray:
        """
        Return a new :class:`.BlockArray` with given block shapes and type, filled with zeros.

        Note: like :func:`jax.numpy.empty`, this does not return an uninitalized array.

        Args:
            shape_tuple: Tuple of shapes for component blocks
            dtype: Desired data-type for the :class:`.BlockArray`. Default is `numpy.float32`.

        Returns:
            :class:`.BlockArray` of zeros with the given component shapes and dtype.
        """
        _aval = _AbstractBlockArray(shape_tuple, dtype=dtype)
        data_ravel = jnp.empty(_aval.size, dtype=dtype)
        return cls(_aval, data_ravel)

    @classmethod
    def full(
        cls,
        shape_tuple: BlockShape,
        fill_value: Union[float, complex, int],
        dtype: DType = np.float32,
    ) -> BlockArray:
        """
        Return a new :class:`.BlockArray` with given block shapes and type, filled with
        `fill_value`.

        Args:
            shape_tuple: Tuple of shapes for component blocks.
            fill_value:  Fill value
            dtype: Desired data-type for the BlockArray. The default,
              None, means `np.array(fill_value).dtype`.

        Returns:
            :class:`.BlockArray` with the given component shapes and dtype and all entries
            equal to `fill_value`.
        """
        if dtype is None:
            dtype = np.asarray(fill_value).dtype

        _aval = _AbstractBlockArray(shape_tuple, dtype=dtype)
        data_ravel = jnp.full(_aval.size, fill_value=fill_value, dtype=dtype)
        return cls(_aval, data_ravel)

    def copy(self) -> BlockArray:
        """Returns a copy of this :class:`.BlockArray`.

        This method is not implemented for BlockArray.

        See Also:
            :meth:`.to_numpy`: Converts a :class:`.BlockArray` into a flattened NumPy array.
        """
        # jax DeviceArray copies return a NumPy ndarray.  This blockarray class must be backed
        # by a DeviceArray, so cannot be converted to a NumPy-backed BlockArray.  The BlockArray
        # .to_numpy() method returns a flattened ndarray.
        #
        # This method may be implemented in the future if jax DeviceArray .copy() is modified to
        # returns another DeviceArray.
        raise NotImplementedError

    def to_numpy(self) -> np.ndarray:
        """Returns a :class:`numpy.ndarray` containing the flattened form of this
        :class:`.BlockArray`."""

        if isinstance(self._data, DeviceArray):
            host_arr = jax.device_get(self._data.copy())
        else:
            host_arr = self._data.copy()
        return host_arr

    def blockidx(self, idx: int) -> jax._src.ops.scatter._Indexable:
        """Returns :class:`jax.ops.index` for a given component block.

        Args:
            idx: Desired block index

        Returns:
            :class:`jax.ops.index` pointing to desired block
        """
        return jax.ops.index[self.bndpos[idx] : self.bndpos[idx + 1]]

    def ravel(self) -> JaxArray:
        """Return a copy of ``self._data`` as a contiguous, flattened `DeviceArray`.

        Note that a copy, rather than a view, of the underlying array is returned.
        This is consistent with :func:`jax.numpy.ravel`.

        Returns:
            Copy of underlying flattened array

        """
        return self._data[:]

    def flatten(self) -> JaxArray:
        """Return a copy of ``self._data`` as a contiguous, flattened `DeviceArray`.

        Note that a copy, rather than a view, of the underlying array is returned.
        This is consistent with :func:`jax.numpy.ravel`.

        Returns:
            Copy of underlying flattened array

        """
        return self._data[:]

    def sum(self, axis=None, keepdims=False):
        """Return the sum of the blockarray elements over the given axis.

        Refer to :func:`scico.numpy.sum` for full documentation.
        """
        # Can't just call scico.numpy.sum due to pesky circular import...
        return _block_array_reduction_wrapper(jnp.sum)(self, axis=axis, keepdims=keepdims)

## Register BlockArray as a Jax type
# Our BlockArray is just a single large vector with some extra sugar
class _ConcreteBlockArray(_AbstractBlockArray):
    pass

def _block_array_result_handler(device, _aval):
    def build_block_array(data_buf):
        data = xla.DeviceArray(_aval._data_aval, device, None, data_buf)
        return BlockArray(_aval, data)

    return build_block_array

def _block_array_shape_handler(a):
    return (xla.xc.Shape.array_shape(a._data_aval.dtype, a._data_aval.shape),)

def _block_array_device_put_handler(a, device):
    return (xla.xb.get_device_backend(device).buffer_from_pyval(a._data, device),)

core.pytype_aval_mappings[BlockArray] = lambda x: x._aval
core.raise_to_shaped_mappings[_AbstractBlockArray] = lambda _aval, _: _aval
xla.pytype_aval_mappings[BlockArray] = lambda x: x._aval
xla.canonicalize_dtype_handlers[BlockArray] = lambda x: x
xla.device_put_handlers[BlockArray] = _block_array_device_put_handler
xla.xla_result_handlers[_AbstractBlockArray] = _block_array_result_handler
xla.xla_shape_handlers[_AbstractBlockArray] = _block_array_shape_handler

## Handlers to use jax.device_put on BlockArray
def _block_array_tree_flatten(block_arr):
    """Flattens a :class:`.BlockArray` pytree.

    See :func:`jax.tree_util.tree_flatten`.

    Args:
        block_arr (:class:`.BlockArray`): :class:`.BlockArray` to flatten

    Returns:
        children (tuple): :class:`.BlockArray` leaves
        aux_data (tuple): Extra metadata used to reconstruct BlockArray.
    """

    data_children, data_aux_data = tree_flatten(block_arr._data)
    return (data_children, block_arr._aval)

def _block_array_tree_unflatten(aux_data, children):
    """Constructs a :class:`.BlockArray` from a flattened pytree.

    See jax.tree_utils.tree_unflatten

    Args:
        aux_data (tuple): Metadata needed to construct block array
        children (tuple): Contains block array elements

    Returns:
        block_arr: Constructed :class:`.BlockArray`
    """
    return BlockArray(aux_data, children[0])

register_pytree_node(BlockArray, _block_array_tree_flatten, _block_array_tree_unflatten)

# Syntactic sugar for the .at operations
# see https://github.com/google/jax/blob/56e9f7cb92e3a099adaaca161cc14153f024047c/jax/_src/numpy/lax_numpy.py#L5900
class _BlockArrayIndexUpdateHelper:
    """The helper class for the `at` property to call indexed update functions.

    The `at` property is syntactic sugar for calling the indexed update
    functions as is done in jax. The index must be of the form [ibk] or [ibk,idx],
    where `ibk` is the index of the block to be updated, and `idx` is a
    general index of the elements to be updated in that block.

    In particular:
    - ``x = x.at[ibk].set(y)`` is an equivalent of ``x[ibk] = y``.
    - ``x = x.at[ibk,idx].set(y)`` is an equivalent of ``x[ibk,idx] = y``.

    The methods ``set, add, multiply, divide, power, maximum, minimum`` are supported.

    """

    __slots__ = ("_block_array",)

    def __init__(self, block_array):
        self._block_array = block_array

    def __getitem__(self, index):
        if isinstance(index, tuple):
            if isinstance(index[0], slice):
                raise TypeError(f"Slicing not supported along block index")
        return _BlockArrayIndexUpdateRef(self._block_array, index)

    def __repr__(self):
        print(f"_BlockArrayIndexUpdateHelper({repr(self._block_array)})")

class _BlockArrayIndexUpdateRef:
    """Helper object to call indexed update functions for an (advanced) index.

    This object references a source block array and a specific indexer, with the
    first integer specifying the block being updated, and rest being the indexer
    into the array of that block.
    Methods on this object return copies of the source block array that have
    been modified at the positions specified by the indexer in the given block.

    """

    __slots__ = ("_block_array", "bk_index", "index")

    def __init__(self, block_array, index):
        self._block_array = block_array
        if isinstance(index, int):
            self.bk_index = index
            self.index = Ellipsis
        elif index == Ellipsis:
            self.bk_index = Ellipsis
            self.index = Ellipsis
        else:
            self.bk_index = index[0]
            self.index = index[1:]

    def __repr__(self):
        return f"_BlockArrayIndexUpdateRef({repr(self._block_array)}, {repr(self.bk_index)}, {repr(self.index)})"

    def _index_wrapper(self, func_str, values):
        bk_index = self.bk_index
        index = self.index
        arr_tuple = self._block_array.split
        if bk_index == Ellipsis:
            # TODO does this result in multiple copies? One per sub-blockarray,
            # then one to combine into a nested BA?
            retval = BlockArray.array([getattr(_.at[index], func_str)(values) for _ in arr_tuple])
        else:
            retval = BlockArray.array(
                arr_tuple[:bk_index]
                + (getattr(arr_tuple[bk_index].at[index], func_str)(values),)
                + arr_tuple[bk_index + 1 :]
            )
        return retval

    def set(self, values):
        """Pure equivalent of ``x[idx] = y``.

        Returns the value of ``x`` that would result from the NumPy-style
        :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] = y``.

        See :mod:`jax.ops` for details.
        """
        return self._index_wrapper("set", values)

    def add(self, values):
        """Pure equivalent of ``x[idx] += y``.

        Returns the value of ``x`` that would result from the NumPy-style
        :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] += y``.

        See :mod:`jax.ops` for details.
        """
        return self._index_wrapper("add", values)

    def multiply(self, values):
        """Pure equivalent of ``x[idx] *= y``.

        Returns the value of ``x`` that would result from the NumPy-style
        :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] *= y``.

        See :mod:`jax.ops` for details.
        """
        return self._index_wrapper("multiply", values)

    def divide(self, values):
        """Pure equivalent of ``x[idx] /= y``.

        Returns the value of ``x`` that would result from the NumPy-style
        :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] /= y``.

        See :mod:`jax.ops` for details.
        """
        return self._index_wrapper("divide", values)

    def power(self, values):
        """Pure equivalent of ``x[idx] **= y``.

        Returns the value of ``x`` that would result from the NumPy-style
        :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] **= y``.

        See :mod:`jax.ops` for details.
        """
        return self._index_wrapper("power", values)

    def min(self, values):
        """Pure equivalent of ``x[idx] = minimum(x[idx], y)``.

        Returns the value of ``x`` that would result from the NumPy-style
        :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] = minimum(x[idx], y)``.

        See :mod:`jax.ops` for details.
        """
        return self._index_wrapper("min", values)

    def max(self, values):
        """Pure equivalent of ``x[idx] = maximum(x[idx], y)``.

        Returns the value of ``x`` that would result from the NumPy-style
        :mod:indexed assignment <numpy.doc.indexing>` ``x[idx] = maximum(x[idx], y)``.

        See :mod:`jax.ops` for details.
        """
        return self._index_wrapper("max", values)

setattr(BlockArray, "at", property(_BlockArrayIndexUpdateHelper))

max_stars_repo_name: lukepfister/scico

max_issues_repo_name: lukepfister/scico

max_forks_repo_name: lukepfister/scico

reason: import numpy,import jax,from jax


text:
import pandas as pd
import numpy as np
from sklearn.linear_model import LogisticRegression

lr=LogisticRegression()
data=pd.read_csv("final_data.csv")
data=np.array(data)
data=data[:,1:]
msk1 = np.random.rand(len(data)) < 0.7
train = data[msk1]
test = data[~msk1]
x_train=train[:,1:]
y_train=train[:,0]
x_test=test[:,1:]
y_test=test[:,0]
lr.fit(x_train,y_train)
def prediction():
    #learn the regressor
    #print(lr.predict(data))
    count=0
    for i in range(x_test.shape[0]):
        if(lr.predict([x_test[i]])==y_test[i]):
            count+=1
    print("efficiency"+str((count/x_test.shape[0])*100))
def prediction1(data):
    return ((lr.predict(data)))

max_stars_repo_name: anchittandon/Disaster-prediction

max_issues_repo_name: anchittandon/Disaster-prediction

max_forks_repo_name: anchittandon/Disaster-prediction