text:
# from imagepy.ui.widgets import CMapSelCtrl
# from imagepy.core.manager import ColorManager, ImageManager, WindowsManager, ToolsManager
from imagepy.core.engine import Macros
import numpy as np
import wx, os.path as osp

def make_bitmap(bmp):
    img = bmp.ConvertToImage()
    img.Resize((20, 20), (2, 2))
    return img.ConvertToBitmap()

# apply, paint, fill, width, slic
class Plugin ( wx.Panel ):
	title = 'Label Tool'
	def __init__( self, parent ):
		wx.Panel.__init__ ( self, parent, id = wx.ID_ANY, pos = wx.DefaultPosition, size = wx.Size(-1,-1), style = wx.TAB_TRAVERSAL )
		outsizer = wx.BoxSizer(wx.HORIZONTAL)
		sizer = wx.BoxSizer( wx.VERTICAL )
		sizer_color = wx.BoxSizer( wx.HORIZONTAL )
		self.btns = []
		self.btn_make =  wx.Button( self, wx.ID_ANY, 'New Mark', wx.DefaultPosition, wx.DefaultSize, 0 )
		sizer_color.Add(self.btn_make, 0, wx.ALL, 2)
		for i in range(11):
			btn = wx.Button( self, wx.ID_ANY, str(i), wx.DefaultPosition, wx.DefaultSize, 0 )
			btn.SetMaxSize( wx.Size( 30,-1 ) )
			self.btns.append(btn)
			sizer_color.Add( btn, 0, wx.ALL, 2 )
		self.spn_num = wx.SpinCtrl( self, wx.ID_ANY, wx.EmptyString, wx.DefaultPosition, wx.DefaultSize, wx.SP_ARROW_KEYS, 2, 15, 0 )
		self.spn_num.SetMaxSize(wx.Size(45, -1))
		sizer_color.Add( self.spn_num, 0, wx.ALL|wx.EXPAND, 2 )

		sizer.Add(sizer_color, 0, wx.ALL|wx.EXPAND, 0)

		sizer_other = wx.BoxSizer( wx.HORIZONTAL )

		self.btn_update =  wx.Button( self, wx.ID_ANY, 'Update', wx.DefaultPosition, wx.DefaultSize, 0 )
		sizer_other.Add( self.btn_update, 0, wx.ALL, 2)

		self.cmapsel = CMapSelCtrl(self)
		self.cmapsel.SetItems(ColorManager.luts)
		sizer_other.Add(self.cmapsel, 0, wx.ALL|wx.EXPAND, 2 )

		com_backChoices = [ u"No Background" ]
		self.com_back = wx.ComboBox( self, wx.ID_ANY, u"No Background", wx.DefaultPosition, wx.DefaultSize, com_backChoices, wx.CB_READONLY)
		self.com_back.SetSelection( 0 )
		sizer_other.Add( self.com_back, 1, wx.ALL|wx.EXPAND, 2 )

		com_modeChoices = [ u"None", u"Max", u"Min", u"Mask", u"2-8mix", u"4-6mix", u"5-5mix", u"6-4mix", u"8-2mix" ]
		self.com_mode = wx.ComboBox( self, wx.ID_ANY, u"Min", wx.DefaultPosition, wx.DefaultSize, com_modeChoices, wx.CB_READONLY )
		self.com_mode.SetSelection( 0 )
		sizer_other.Add( self.com_mode, 0, wx.ALL|wx.EXPAND, 2 )

		self.chk_hide = wx.CheckBox( self, wx.ID_ANY, u"Hide", wx.DefaultPosition, wx.DefaultSize, 0 )
		sizer_other.Add( self.chk_hide, 0, wx.ALIGN_CENTER|wx.ALL, 0 )

		sizer.Add(sizer_other, 0, wx.ALL|wx.EXPAND, 0)

		#sizer_tol = wx.GridSizer( 0, 3, 0, 0 )
		self.pens = []
		name = ['01.gif','03.gif','05.gif','10.gif','fill.gif']
		path = osp.abspath(osp.dirname(__file__))
		for i in (0,1,2,3,4):
			pen = wx.BitmapButton(self, wx.ID_ANY, make_bitmap(wx.Bitmap(osp.join(path, 'imgs', name[i]))),#make_bitmap(wx.Bitmap(data[1])),
            wx.DefaultPosition, (30, 30), wx.BU_AUTODRAW|wx.RAISED_BORDER )
            #wx.Button( self, wx.ID_ANY, str(i), wx.DefaultPosition, wx.DefaultSize, 0 )
			pen.SetMaxSize( wx.Size( 30,-1 ) )
			self.pens.append( pen )
			sizer_color.Add( pen, 0, wx.ALL, 2 )

		outsizer.AddStretchSpacer(prop=1)
		#outsizer.Add(sizer_tol, 0, wx.ALL, 0)
		outsizer.Add(sizer, 0, wx.ALL, 0)
		outsizer.AddStretchSpacer(prop=1)

		self.SetSizer( outsizer )
		self.Layout()

		# Connect Events
		self.spn_num.Bind( wx.EVT_SPINCTRL, self.on_cmapsel )
		self.com_back.Bind(wx.EVT_COMBOBOX_DROPDOWN, self.on_items )
		self.cmapsel.Bind(wx.EVT_COMBOBOX,  self.on_cmapsel)
		self.com_back.Bind( wx.EVT_COMBOBOX, self.on_setback)
		self.com_mode.Bind( wx.EVT_COMBOBOX, self.on_mode)
		self.chk_hide.Bind( wx.EVT_CHECKBOX, self.on_mode)
		self.pens[-1].Bind( wx.EVT_BUTTON, self.on_fill)
		for i in self.btns: i.Bind(wx.EVT_BUTTON, self.on_color)
		for i in range(4): self.pens[i].Bind(wx.EVT_BUTTON, \
			lambda e, x=(1,3,5,10)[i]: self.on_pen(x))
		self.btn_make.Bind( wx.EVT_BUTTON, self.on_make)

	def on_make(self, event):
		Macros(None, ['Build Mark Image>None']).start()

	def on_fill(self, event):
		tol = ToolsManager.get('Flood Fill')()
		tol.para['tor'] = 0
		tol.start()

	def on_pen(self, width):
		tol = ToolsManager.get('Pencil')()
		tol.para['width'] = width
		tol.start()

	def on_color(self, event):
		ColorManager.set_front(self.btns.index(event.GetEventObject()))

	def on_items(self, event):
		items = ['No Background Image']+ImageManager.get_titles()
		self.com_back.SetItems(items)
		if self.com_back.GetValue() in items:
			self.com_back.Select(items.index(self.com_back.GetValue()))
		else: self.com_back.Select(0)

	def on_cmapsel(self, event):
		key = self.cmapsel.GetValue()
		lut = ColorManager.get_lut(key)
		n = self.spn_num.GetValue()+1
		idx = np.linspace(0, 255, n).astype(int)
		cs = list(lut[idx]) + [(128,128,128)]*(16-n)
		for btn, c in zip(self.btns, cs):
			btn.SetBackgroundColour(c)

		ips = ImageManager.get()
		if ips is None: return
		newlut = lut*0
		newlut[:n] = lut[idx]
		ips.lut = newlut
		ips.update()

	def on_setback(self, event):
		name = self.com_back.GetValue()
		if name is None: return
		ImageManager.get().back = ImageManager.get(name)
		#curwin = WindowsManager.get()
		#curwin.set_back(ImageManager.get(name))
		ImageManager.get().update()

	def on_mode(self, event):
		ips = ImageManager.get()
		if ips is None: return
		if self.chk_hide.GetValue():
			ips.chan_mode = 0.0
			return ips.update()
		modes = ['set', 'max', 'min', 'msk', 0.2, 0.4, 0.5, 0.6, 0.8]
		ips.chan_mode = modes[self.com_mode.GetSelection()]
		ips.update()

	def __del__( self ):
		pass

max_stars_repo_name: pengguanjun/imagepy

max_issues_repo_name: cycleuser/imagepy

max_forks_repo_name: cycleuser/imagepy


text:
# Code to estimate mixture entropy using Monte Carlo sampling

import numpy as np
import scipy
from entropy import pairwise_distance2_np

from randomgen import RandomGenerator, MT19937
rnd = RandomGenerator(MT19937())

def get_mc_entropy(mx, var):
    n, d     = mx.shape
    mx2 = mx + rnd.standard_normal(mx.shape, dtype=np.float32)*np.sqrt(var)

    dist_norm = pairwise_distance2_np(mx, mx2) / (-2.0 * var)

    const     =  - 0.5  * d * np.log(2.0 * np.pi * var) - np.log(n)
    logprobs  = const + scipy.special.logsumexp(dist_norm , axis=0)

    return -np.mean(logprobs)

ext: py

max_stars_repo_name: artemyk/nonlinearIB

max_issues_repo_name: artemyk/nonlinearIB

max_forks_repo_name: artemyk/nonlinearIB

reason: import numpy,import scipy


text:
import numpy as np
import random
import matplotlib.pyplot as plt

class KMeans:
    def __init__(self):
        pass

    def KMeansDemo(self):
        centerList = [[0, 0], [30, 30], [30, -30], [-30, 30], [-30, -30]]
        sigmaList = []
        clustreNumList = []
        rangeList = []
        k = 5
        eleNum = 2
        for i in xrange(0, k):
            sigmaList.append(random.random() * 1.0)
            clustreNumList.append(np.int32(50.0 + random.random() * 50.0))
            rangeList.append(random.random() * 20.0)
        rawData = self._GenerateKClustre(2, centerList, sigmaList, clustreNumList, rangeList)
        mData = self._MergeData(rawData)
        clustre, centers, iterCentroid = self.Kmeans(k, mData)
        #plt.plot(mData[:, 0], mData[:, 1], 'rs')
        plt.plot(centers[:, 0], centers[:, 1], 'r^', markersize=20)
        print centers
        plt.plot(iterCentroid[:, 0], iterCentroid[:, 1], 'kx', markersize=20)
        # plt.plot(clustre[0][:, 0], clustre[0][:, 1], 'mo', markersize=5, fillstyle='none')
        # plt.plot(clustre[1][:, 0], clustre[1][:, 1], 'yo', markersize=5, fillstyle='none')
        # plt.plot(clustre[2][:, 0], clustre[2][:, 1], 'go', markersize=5, fillstyle='none')
        # plt.plot(clustre[3][:, 0], clustre[3][:, 1], 'bo', markersize=5, fillstyle='none')
        # plt.plot(clustre[4][:, 0], clustre[4][:, 1], 'co', markersize=5, fillstyle='none')
        plt.plot(clustre[0][:, 0], clustre[0][:, 1], 'mo', fillstyle='none')
        plt.plot(clustre[1][:, 0], clustre[1][:, 1], 'yo', fillstyle='none')
        plt.plot(clustre[2][:, 0], clustre[2][:, 1], 'go', fillstyle='none')
        plt.plot(clustre[3][:, 0], clustre[3][:, 1], 'bo', fillstyle='none')
        plt.plot(clustre[4][:, 0], clustre[4][:, 1], 'co',  fillstyle='none')
        plt.show()

    def Kmeans(self, k, data):
        centers = self._SelectSeeds(data, k)
        iterCentroid = centers.copy()
        iterCentroidOld = np.zeros((iterCentroid.shape))
        clustreNum = np.zeros((centers.shape[0],1))
        dataTag = np.zeros((data.shape[0],1), np.uint32)
        while np.linalg.norm(iterCentroid - iterCentroidOld) > 1:
            iterCentroidOld = iterCentroid.copy()
            clustreCenter = np.zeros((centers.shape[0], data.shape[1]))
            clustreNum = np.zeros((centers.shape[0], 1))
            for i in xrange(0, data.shape[0]):
                nearestCenter, nearestIndex = self._GetNearestClustreCenter(iterCentroid, data[i, :])
                dataTag[i] = nearestIndex
                clustreNum[nearestIndex] += 1
                clustreCenter[nearestIndex, :] += data[i, :]
            for i in xrange(0, iterCentroid.shape[0]):
                iterCentroid[i, :] = clustreCenter[i, :] / clustreNum[i]
            #print iterCentroid
        clustre = [np.zeros((0, data.shape[1])) for i in xrange(0,centers.shape[0])]
        for i in xrange(0, data.shape[0]):
            clustre[dataTag[i]] = np.vstack((clustre[dataTag[i]], data[i, :]))
        return clustre, centers, iterCentroid

    def _MergeData(self, rawData):
        if len(rawData) == 0:
            return None
        mData = rawData[0]
        for i in xrange(1, len(rawData)):
            mData = np.concatenate((mData, rawData[i]), axis=0)
        return mData

    def _GenerateKClustre(self, eleNum, centerList, sigmaList, clustreNumList, rangeList):
        rawData = []
        for i in xrange(0, len(clustreNumList)):
            tmpData = np.zeros((clustreNumList[i], eleNum))
            center = centerList[i]
            sigma = sigmaList[i]
            for j in xrange(0, clustreNumList[i]):
                for ij in xrange(0, eleNum):
                    tmpData[j, ij] = center[ij] + random.gauss(0, sigma) * rangeList[i]
            rawData.append(tmpData)
        return rawData

    def _GetNearestClustreCenter(self, centers, datum):
        distMatrix = np.zeros((centers.shape[0], 1))
        for i in xrange(0, centers.shape[0]):
            distMatrix[i] = np.linalg.norm(datum - centers[i])
        nearestIndex = np.argmin(distMatrix)
        nearestCenter = centers[nearestIndex, :]
        return nearestCenter, nearestIndex

    def _SelectSeeds(self, data, k):
        #row vector
        seed = np.zeros((k, data.shape[1]))
        distMatrix = np.zeros((data.shape[0], data.shape[0]))
        for i in xrange(0, data.shape[0]):
            for j in xrange(i, data.shape[0]):
                distMatrix[i,j] = np.linalg.norm(data[i, :] - data[j, :])
        farestPtInd = np.unravel_index(distMatrix.argmax(), distMatrix.shape)
        seed[0, :] = data[farestPtInd[0], :]
        seed[1, :] = data[farestPtInd[1], :]
        #seed[0,:] = data[np.int32(np.random.random() * data.shape[0]), :]
        for i in xrange(2, k):
            seed[i, :] = self._GetFarestData(data, seed[0:i, :])
        return seed

    def _GetFarestData(self, data, curSeeds):
        distMatrix = np.zeros((data.shape[0],1))
        for i in xrange(0, data.shape[0]):
            nearestCenter, nearestIndex = self._GetNearestClustreCenter(curSeeds, data[i, :])
            distMatrix[i] = np.linalg.norm(data[i, :] - nearestCenter)
        farestIndex = np.argmax(distMatrix)
        newSeed = data[farestIndex, :]
        return newSeed

max_stars_repo_name: artzers/MachineLearning

max_issues_repo_name: ZhidongYang/MachineLearning

max_forks_repo_name: ZhidongYang/MachineLearning

reason: import numpy


text:
import scipy.stats as st
from brent_search import minimize
from numpy import concatenate, empty, kron
from numpy.random import RandomState
from numpy.testing import assert_allclose
from numpy_sugar.linalg import economic_qs

from glimix_core._util import assert_interface, vec
from glimix_core.lmm import FastScanner, Kron2Sum, KronFastScanner

def test_lmm_kron_scan():
    random = RandomState(0)
    n = 20
    Y = random.randn(n, 3)
    A = random.randn(3, 3)
    A = A @ A.T
    F = random.randn(n, 2)
    G = random.randn(n, 6)
    lmm = Kron2Sum(Y, A, F, G, restricted=True)
    lmm.fit(verbose=False)
    scan = lmm.get_fast_scanner()

    m = lmm.mean()
    K = lmm.covariance()

    def func(scale):
        mv = st.multivariate_normal(m, scale * K)
        return -mv.logpdf(vec(Y))

    s = minimize(func, 1e-3, 5.0, 1e-5)[0]

    assert_allclose(scan.null_lml(), st.multivariate_normal(m, s * K).logpdf(vec(Y)))
    assert_allclose(kron(A, F) @ scan.null_beta, m)

    A1 = random.randn(3, 2)
    F1 = random.randn(n, 4)

    r = scan.scan(A1, F1)
    assert_allclose(r["scale"], 0.7365021111700154, rtol=1e-3)

    m = kron(A, F) @ vec(r["effsizes0"]) + kron(A1, F1) @ vec(r["effsizes1"])

    def func(scale):
        mv = st.multivariate_normal(m, scale * K)
        return -mv.logpdf(vec(Y))

    s = minimize(func, 1e-3, 5.0, 1e-5)[0]

    assert_allclose(r["lml"], st.multivariate_normal(m, s * K).logpdf(vec(Y)))

    r = scan.scan(empty((3, 0)), F1)
    assert_allclose(r["lml"], -85.36667704747371, rtol=1e-4)
    assert_allclose(r["scale"], 0.8999995537936586, rtol=1e-3)
    assert_allclose(
        r["effsizes0"],
        [
            [0.21489119796865844, 0.6412947101778663, -0.7176143380221816],
            [0.8866722740598517, -0.18731140321348416, -0.26118052682069],
        ],
        rtol=1e-2,
        atol=1e-2,
    )
    assert_allclose(r["effsizes1"], [])

def test_lmm_kron_scan_with_lmm():
    random = RandomState(0)
    n = 15
    Y = random.randn(n, 3)
    A = random.randn(3, 3)
    A = A @ A.T
    F = random.randn(n, 2)
    G = random.randn(n, 6)

    klmm = Kron2Sum(Y, A, F, G, restricted=True)
    klmm.fit(verbose=False)
    kscan = klmm.get_fast_scanner()

    K = klmm.covariance()

    X = kron(A, F)
    QS = economic_qs(K)
    scan = FastScanner(vec(Y), X, QS, 0.0)

    assert_allclose(klmm.covariance(), K)
    assert_allclose(kscan.null_scale, scan.null_scale)
    assert_allclose(kscan.null_beta, scan.null_beta)
    assert_allclose(kscan.null_lml(), scan.null_lml())
    assert_allclose(kscan.null_beta_covariance, scan.null_beta_covariance)

    A1 = random.randn(3, 2)
    F1 = random.randn(n, 2)
    M = kron(A1, F1)

    kr = kscan.scan(A1, F1)
    r = scan.scan(M)
    assert_allclose(kr["lml"], r["lml"])
    assert_allclose(kr["scale"], r["scale"])
    assert_allclose(vec(kr["effsizes0"]), r["effsizes0"])
    assert_allclose(vec(kr["effsizes1"]), r["effsizes1"])
    assert_allclose(vec(kr["effsizes0_se"]), r["effsizes0_se"])
    assert_allclose(vec(kr["effsizes1_se"]), r["effsizes1_se"])

def test_lmm_kron_scan_unrestricted():
    random = RandomState(0)
    n = 15
    Y = random.randn(n, 3)
    A = random.randn(3, 3)
    A = A @ A.T
    F = random.randn(n, 2)
    G = random.randn(n, 6)
    lmm = Kron2Sum(Y, A, F, G, restricted=False)
    lmm.fit(verbose=False)
    scan = lmm.get_fast_scanner()

    assert_allclose(scan.null_scale, 1.0, rtol=1e-3)
    assert_allclose(lmm.beta_covariance, scan.null_beta_covariance, rtol=1e-3)

def test_lmm_kron_scan_redundant():
    random = RandomState(0)
    n = 30
    Y = random.randn(n, 3)
    A = random.randn(3, 3)
    A = A @ A.T
    F = random.randn(n, 2)
    G = random.randn(n, 6)
    G = concatenate([G, G], axis=1)
    lmm = Kron2Sum(Y, A, F, G, restricted=True)
    lmm.fit(verbose=False)
    scan = lmm.get_fast_scanner()

    m = lmm.mean()
    K = lmm.covariance()

    def func(scale):
        mv = st.multivariate_normal(m, scale * K)
        return -mv.logpdf(vec(Y))

    s = minimize(func, 1e-3, 5.0, 1e-5)[0]

    assert_allclose(scan.null_lml(), st.multivariate_normal(m, s * K).logpdf(vec(Y)))
    assert_allclose(kron(A, F) @ scan.null_beta, m)

    A1 = random.randn(3, 2)
    F1 = random.randn(n, 4)
    F1 = concatenate([F1, F1], axis=1)

    r = scan.scan(A1, F1)
    assert_allclose(r["scale"], 0.8843540849467378, rtol=1e-3)

    m = kron(A, F) @ vec(r["effsizes0"]) + kron(A1, F1) @ vec(r["effsizes1"])

    def func(scale):
        mv = st.multivariate_normal(m, scale * K)
        return -mv.logpdf(vec(Y))

    s = minimize(func, 1e-3, 5.0, 1e-5)[0]

    assert_allclose(r["lml"], st.multivariate_normal(m, s * K).logpdf(vec(Y)))

def test_lmm_kron_scan_public_attrs():
    assert_interface(
        KronFastScanner,
        ["null_lml", "scan"],
        ["null_beta", "null_beta_covariance", "null_beta_se", "null_scale"],
    )

max_stars_repo_name: Horta/limix-inference

max_issues_repo_name: Horta/limix-inference

max_forks_repo_name: Horta/limix-inference

reason: from numpy,import scipy


text:
import datetime
import numpy as np
import pandas as pd
from shapely.geometry import Point
import seis_utils
from seis_settings import EPSG_PSD93
from seis_database import DbUtils

class SpsDb:
    table_sps_files = 'sps_files'
    table_sps = 'sps_records'

    @classmethod
    def create_database(cls):
        DbUtils().create_database()

    @classmethod
    @DbUtils.connect
    def delete_table_sps(cls, cursor):
        sql_string = f'DROP TABLE {cls.table_sps};'
        cursor.execute(sql_string)
        print(f'delete table {cls.table_sps}')

    @classmethod
    @DbUtils.connect
    def delete_table_sps_files(cls, cursor):
        sql_string = f'DROP TABLE {cls.table_sps_files};'
        cursor.execute(sql_string)
        print(f'delete table {cls.table_sps_files}')

    @classmethod
    @DbUtils.connect
    def create_table_sps_files(cls, cursor):
        sql_string = (
            f'CREATE TABLE {cls.table_sps_files} ('
            f'id INTEGER PRIMARY KEY, '
            f'file_name VARCHAR(100), '
            f'file_date TIMESTAMP, '
            f'block_name VARCHAR(10));'
        )
        cursor.execute(sql_string)
        print(f'create table {cls.table_sps_files}')

    @classmethod
    @DbUtils.connect
    def create_table_sps(cls, cursor):
        # create the table according to SPS format
        sql_string = (
            f'CREATE TABLE {cls.table_sps} ('
            f'id INTEGER PRIMARY KEY, '
            f'file_id INTEGER REFERENCES {cls.table_sps_files}(id) ON DELETE CASCADE, '
            f'sps_type VARCHAR(1), '
            f'line INTEGER, '
            f'point INTEGER, '
            f'point_index INTEGER, '
            f'source_type VARCHAR(5),'
            f'easting DOUBLE PRECISION, '
            f'northing DOUBLE PRECISION, '
            f'elevation REAL, '
            f'dpg_filename VARCHAR(30), '
            f'time_break TIMESTAMP, '
            f'vibrator);'
        )
        cursor.executescript(sql_string)

        # once table is created you can add the geomety column
        sql_string = (
            f'SELECT AddGeometryColumn("{cls.table_sps}", '
            f'"geom", {EPSG_PSD93}, "POINT", "XY");'
        )
        cursor.execute(sql_string)

        print(f'create table {cls.table_sps}')

    @classmethod
    @DbUtils.connect
    def update_sps_file(cls, sps_file, cursor):
        ''' method to to check if file_name exists in the database, if it does not then
            add the filename to the data base
            returns:
            -1, if file is found
            n, new file_id number if no file is found
        '''
        # check if file exists
        sql_string = (
            f'SELECT id FROM {cls.table_sps_files} WHERE '
            f'file_name like \'%{sps_file.file_name}\' ;'
        )
        cursor.execute(sql_string)
        try:
            # check if id exists
            _ = cursor.fetchone()[0]
            return -1

        except TypeError:
            # no id was found so go on to create one
            pass

        sql_string = (
            f'INSERT INTO {cls.table_sps_files} ('
            f'file_name, file_date, block_name) '
            f'VALUES (?, ?, ?); '
        )
        cursor.execute(sql_string, (
            sps_file.file_name, sps_file.file_date, sps_file.block_name)
        )
        return cursor.lastrowid

    @classmethod
    @DbUtils.connect
    def update_sps(cls, sps_records, cursor):
        progress_message = seis_utils.progress_message_generator(
            f'populate database for table: {cls.table_sps}                   ')

        sql_vp_record = (
            f'INSERT INTO {cls.table_sps} ('
            f'file_id, sps_type, line, point, point_index, source_type, '
            f'easting, northing, elevation, dpg_filename, time_break, '
            f'vibrator, geom) '
            f'VALUES ({", ".join(["?"]*12)}, MakePoint(?, ?, ?));'
        )

        for sps_record in sps_records:
            point = Point(sps_record.easting, sps_record.northing)

            cursor.execute(sql_vp_record, (
                sps_record.file_id,
                sps_record.sps_type,
                sps_record.line,
                sps_record.point,
                sps_record.point_index,
                sps_record.source_type,
                sps_record.easting,
                sps_record.northing,
                sps_record.elevation,
                sps_record.dpg_filename,
                sps_record.time_break,
                sps_record.vibrator,
                point.x, point.y, EPSG_PSD93
                )
            )
            next(progress_message)

    @classmethod
    def get_sps_data_by_time(cls, start_time, end_time):
        ''' retrieve vp data by time interval
            arguments:
              start_time: datetime object
              end_time: datetime object
            returns:
              pandas dataframe with all database attributes
        '''
        assert end_time >= start_time, "end time must be greater equal than start time"

        engine = DbUtils().get_db_engine()
        sql_string = (
            f'SELECT * FROM {cls.table_sps} WHERE '
            f'time_break BETWEEN \'{start_time}\' AND \'{end_time}\' '
            f'ORDER BY time_break;'
        )
        return pd.read_sql_query(sql_string, con=engine)

    @classmethod
    def get_sps_data_by_date(cls, production_date):
        ''' retrieve vp data by date
            arguments:
              database_table: 'VAPS' or 'VP'
              production_date: datetime object
            returns:
              pandas dataframe with vp attributes for production_date
        '''
        engine = DbUtils().get_db_engine()
        sql_string = (f'SELECT * FROM {cls.table_sps} WHERE '
                      f'DATE(time_break) = \'{production_date.strftime("%Y-%m-%d")}\';')
        return pd.read_sql_query(sql_string, con=engine)

    @classmethod
    def get_vp_data_by_line(cls, line):
        ''' retrieve vp data by line number
            arguments:
              line: integer
            returns:
              pandas dataframe with all database attributes
        '''
        engine = DbUtils().get_db_engine()
        sql_string = (
            f'SELECT * FROM {cls.table_sps} WHERE '
            f'line = {line} ORDER BY station;'
        )
        return pd.read_sql_query(sql_string, con=engine)

    @classmethod
    @DbUtils.connect
    def get_all_line_points(cls, block_name, cursor):
        ''' get all line, point of the table_sps in a pandas df
        '''
        sql_string = (
            f'CREATE VIEW line_points (line, point) '
            f'AS SELECT line, point FROM {cls.table_sps} as r '
            f'INNER JOIN {cls.table_sps_files} as f ON f.id = r.file_id '
            f'WHERE f.block_name = \"{block_name}\";'
        )
        cursor.execute(sql_string)

        engine = DbUtils().get_db_engine()
        sql_string = (
            f'SELECT line, point FROM line_points '
            f'ORDER BY line, point'
        )
        sps_df = pd.read_sql_query(sql_string, con=engine)
        lines = np.array(sps_df['line'].to_list()) * 10_000
        points = np.array(sps_df['point'].to_list())
        linepoints = np.add(lines, points)

        sql_string = (
            f'DROP VIEW IF EXISTS line_points;'
        )
        cursor.execute(sql_string)

        return linepoints

max_stars_repo_name: bvermeulen/Seistools

max_issues_repo_name: bvermeulen/Seistools

max_forks_repo_name: bvermeulen/Seistools


text:
import numpy as np

from recommendation_engines import RandomRecommender
from recommendation_engines import RuleBasedRecommender
from recommendation_engines import SoftmaxRecommender

areas_mapper = {
    0: 'Action',
    1: 'Comedy',
    2: 'Horror',
    3: 'Drama',
    4: 'Thriller'
}

products_mapper = {
    'Action': {
        0: 'Indiana Jones',
        1: 'Terminator',
        2: 'The Matrix'
    },

    'Comedy': {
        0: 'Scary Movie',
        1: 'Borat',
        2: 'Hot Fuzz'
    },

    'Horror': {
        0: 'The Evil Dead',
        1: 'The Thing',
        2: 'They Live'
    },

    'Drama': {
        0: 'Bicycle Thieves',
        1: 'Patch Adams',
        2: 'Awakenings'
    },

    'Thriller': {
        0: 'Zodiac',
        1: 'Silence of the Lamb',
        2: 'Shutter Island'
    }

}

# probabilities for the rule-based recommender
areas_ps = {}
for customer_id in range(100000):

    p = np.random.random(size=5)
    p = p / p.sum()
    areas_ps[customer_id] = p

products_ps = {}
for area_id in range(5):

    p = np.random.random(size=3)
    p = p / p.sum()
    products_ps[area_id] = p

# frequencies for the softmax recommender
areas_dict = {
    customer_id: np.random.randint(0, 100, size=5) for
    customer_id in range(100000)
}
products_dict = {
    area_id: np.random.randint(0, 100, size=(3)) for area_id in range(5)
}

areas_dict = {
    customer_id: np.random.randint(0, 100, size=5) for
    customer_id in range(100000)
}
products_dict = {
    area_id: np.random.randint(0, 100, size=(3)) for area_id in range(5)
}

random_users = np.random.randint(0, 100000, 15)

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

print('Random recommendations')
print('')

recommender = RandomRecommender(
    areas_mapper=areas_mapper,
    products_mapper=products_mapper
)
for user_id in random_users:

    area, product = recommender.recommend(
        query_id=user_id
    )
    print(f'User {user_id} might like a {area} movie like {product}')

print('')

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

print('Rule Based recommendations')
print('')

recommender = RuleBasedRecommender(
    areas_mapper=areas_mapper,
    products_mapper=products_mapper,
    areas_ps=areas_ps,
    products_ps=products_ps
)
for user_id in random_users:

    area, product = recommender.recommend(
        query_id=user_id
    )
    print(f'User {user_id} might like a {area} movie like {product}')

print('')

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

print('Softmax recommendations')
print('')

recommender = SoftmaxRecommender(
    areas_dict=areas_dict,
    products_dict=products_dict,
    areas_mapper=areas_mapper,
    products_mapper=products_mapper
)

for user_id in random_users:

    area, product = recommender.recommend(
        query_id=user_id,
        temperature_area=2,
        temperature_product=2
    )
    print(f'User {user_id} might like a {area} movie like {product}')

print('')

max_stars_repo_name: mtaung/bazaar

max_issues_repo_name: mtaung/bazaar

max_forks_repo_name: mtaung/bazaar


text:
import sys
import os
os.environ["path"] = os.path.dirname(sys.executable) + ";" + os.environ["path"]
import glob
import numpy
import json

###############################################################################
# GLOBAL CONSTANTS
###############################################################################

DATA_FOLDER = "data"
PRINT_LEVEL=2

###############################################################################
# UTILITIES
###############################################################################

def myprint(str, level=0):
	if (level >= PRINT_LEVEL):
		print(str)

def print_as_csv(table, level=0):
	if (level < PRINT_LEVEL):
		return

	print("------")
	for row in table:
		msg = ""
		for col in row:
			msg += "\"" + str(col) + "\","
		print(msg[0:-1])
	print("------")

def clean_path(path):
	path = path.replace("\\", "/")
	return path

###############################################################################
# MAIN
###############################################################################

def glob_json(params):
	root = os.path.join(params["glob_root"], "**", "*.json")
	myprint("Loading jsons in " + root, 3)
	all_json = glob.glob(root, recursive=True)
	loaded_jsons = []
	for json_file in all_json:
		myprint("processing " + json_file)
		with open(json_file, 'r') as jsonfile:
			data = json.load(jsonfile)
			data["src"] = clean_path(json_file)
		loaded_jsons.append(data)
	myprint("Loaded " + str(len(loaded_jsons)) + " jsons", 1)
	params["data"] = loaded_jsons

def crafting_report(params):

	#"recyclable": {
	#	"energy": 1200,
	#	"player_sale_range": [1900,2200],
	#	"player_buy_range": [2800,3500]
	#},
	#"disassembling": {
	#	"produce":"data/json/items/misc/spare_parts.json",
	#	"count":4,
	#	"energy_cost":1000
	#},

	#"converter": {
	#	"maximum_energy":40000,
	#	"recipes": [

	#{
	#	"name": "Laser Turret MK2",
	#	"requirements":[{"type":"energy", "amount":1000}, {"type":"spare_parts", "amount":4}],
	#	"produce":"data/json/items/weapons/laser_turret_mk2.json",
	#	"ap_cost":5.0,
	#	"amount":1
	#},

	result = [["name_id", "recyclable_energy", "sale_range", "buy_range", "disassemble_parts", "disassemble_cost", "ap_cost", "energy_cost", "spare_parts_cost"]]

	for object in params["data"]:
		name_id = object["name_id"]
		has_some_data = False
		myprint("processing " + name_id)
		recyclable_energy = "N/A"
		recyclable_sale = "N/A"
		recyclable_buy = "N/A"
		disassemble_count = "-1"
		disassemble_energy = "-1"
		if "recyclable" in object:
			has_some_data = True
			recyclable_data = object["recyclable"]
			recyclable_energy = str(recyclable_data["energy"])
			recyclable_sale = str(recyclable_data["player_sale_range"])
			recyclable_buy = str(recyclable_data["player_buy_range"])
		if "disassembling" in object:
			has_some_data = True
			disassembling_data = object["disassembling"]
			disassemble_count = str(disassembling_data["count"])
			disassemble_energy = str(disassembling_data["energy_cost"])

		ap_cost = -1
		energy_cost = -1
		spare_parts_cost = -1
		for converter in params["data"]:
			if "converter" in converter and "recipes" in converter["converter"]:
				recipes = converter["converter"]["recipes"]
				for recipe in recipes:
					produce = clean_path(recipe["produce"])
					if produce in object["src"]:
						has_some_data = True
						if ap_cost > 0 and ap_cost != recipe["ap_cost"]:
							myprint("************ERROR in {}, ap_cost seen twice ({},{})".format(name_id, str(ap_cost), str(recipe["ap_cost"])), 10)
						ap_cost = recipe["ap_cost"]
						for req in recipe["requirements"]:
							if "type" in req and req["type"] == "energy":
								if energy_cost > 0 and energy_cost != req["amount"]:
									myprint("************ERROR in {}, energy_cost seen twice ({},{})".format(name_id, str(energy_cost), str(req["amount"])), 10)
								energy_cost = req["amount"]
							if "type" in req and req["type"] == "spare_parts":
								if spare_parts_cost > 0 and spare_parts_cost != req["amount"]:
									myprint("************ERROR in {}, spare_parts_cost seen twice ({},{})".format(name_id, str(spare_parts_cost), str(req["amount"])), 10)
								spare_parts_cost = req["amount"]

		if has_some_data:
			result.append([
				name_id,
				str(recyclable_energy),
				str(recyclable_sale),
				str(recyclable_buy),
				str(disassemble_count),
				str(disassemble_energy),
				str(ap_cost),
				str(energy_cost),
				str(spare_parts_cost)
			])

	print_as_csv(result)

def weapon_energy_report(params):
	#"weapon_data": {
	#	"base_dam":8.0,
	#	"max_dam":12.0,
	#	"fire_range":3.0,
	#	"fire_pattern":"+",
	#	"fire_energy_cost":750.0,
	#	"fire_speed":1.0
	#},

	# for each weapon
	#1. how much energy to fire
	#2. if ammo, how much energy to produce ammo
	#3. ap_shot
	#4. min_dam, max_dam

	result = [["name_id", "fire_cost", "ammo_cost", "ap", "min_dam", "max_dam", "range"]]
	for object in params["data"]:
		name_id = object["name_id"]
		myprint("processing " + name_id)

		fire_cost = 0
		ammo_cost = 0
		ap_cost = 0
		min_dam = 0
		max_dam = 0
		range = 0
		has_some_data = False

		if "weapon_data" in object:
			weapon_data = object["weapon_data"]
			has_some_data = True
			min_dam = weapon_data["base_dam"]
			max_dam = weapon_data["max_dam"]
			range = weapon_data["fire_range"]
			ap_cost = weapon_data["fire_speed"]
			fire_cost = weapon_data["fire_energy_cost"]
			if "ammo" in weapon_data:
				for converter in params["data"]:
					if "converter" in converter and "recipes" in converter["converter"]:
						recipes = converter["converter"]["recipes"]
						for recipe in recipes:
							produce = clean_path(recipe["produce"])
							if produce in weapon_data["ammo"]:
								for req in recipe["requirements"]:
									if "type" in req and req["type"] == "energy":
										if ammo_cost > 0 and ammo_cost != req["amount"]:
											myprint("************ERROR in {}, ammo_cost seen twice ({},{})".format(name_id, str(ammo_cost), str(req["amount"])), 10)
										ammo_cost = req["amount"]

		if has_some_data:
			result.append([
				name_id,
				str(fire_cost),
				str(ammo_cost),
				str(ap_cost),
				str(min_dam),
				str(max_dam),
				str(range)
			])

	print_as_csv(result)

def converter_recipe_report(params):
	# for each json objects
		# collect converter & recipes
		# add column header
	# for each json objects
		# check if it is in any converter
			# if not print warning
			# if it is, add a row

	converter_list = []
	header = ["name_id"]
	for converter in params["data"]:
		name_id = converter["name_id"]
		if "converter" in converter and "recipes" in converter["converter"]:
			converter_list.append(converter)
			header.append(name_id)

	result = [header]
	for object in params["data"]:
		name_id = object["name_id"]
		print("processing " + object["src"])
		in_any_converter = False
		row_data = [name_id]
		for converter in converter_list:
			recipes = converter["converter"]["recipes"]
			in_this_converter = False
			for recipe in recipes:
				produce = clean_path(recipe["produce"])
				if produce in object["src"]:
					print("converter {} : produce {} matches {}".format(converter["name_id"], produce, object["src"]))
					in_any_converter = True
					in_this_converter = True
					row_data.append("O")
			if not in_this_converter:
				row_data.append("X")
		if in_any_converter:
			result.append(row_data)
		else:
			myprint("WARNING: {} not in any converter".format(name_id), 5)

	print_as_csv(result)

def count_mounts(params):
	for ship in params["data"]:
		name_id = ship["name_id"]
		if "mounts" in ship:
			if "mount_attributes" not in ship:
				myprint("WARNING: {} is missing mount_attributes".format(name_id))
				continue
			checklist = ["weapon", "shield", "scanner", "converter", "utility"]
			for m in checklist:
				if m in ship["mounts"]:
					if m not in ship["mount_attributes"]:
						myprint("WARNING: {} is missing {} in mount_attributes".format(name_id, m))
						if len(ship["mount_attributes"][m]) != len(ship["mounts"][m]):
							myprint("ERROR: {} {} count doesn't match in mount_attributes".format(name_id, m))

def search_invalid_filename(params):
	for f in params["data"]:
		search_dict(f)

def search_dict(root):
	for comp in root:
		comp_val = root[comp]
		if type(comp_val) is dict:
			search_dict(comp_val)
		elif type(comp_val) is list:
			search_list(comp_val)
		elif type(comp_val) is str:
			if "json" in comp_val or "tscn" in comp_val:
				if not os.path.exists(comp_val):
					myprint("ERROR: File not found {}".format(comp_val), 5)

def search_list(root):
	for comp_val in root:
		if type(comp_val) is dict:
			search_dict(comp_val)
		elif type(comp_val) is list:
			search_list(comp_val)
		elif type(comp_val) is str:
			if "json" in comp_val or "tscn" in comp_val:
				if not os.path.exists(comp_val):
					myprint("ERROR: File not found {}".format(comp_val), 5)

def list_types(params):
	types = {}
	for f in params["data"]:
		if "type" in f:
			if f["type"] not in types:
				types[f["type"]] = 1
			else:
				types[f["type"]] += 1

	print(types)

def do_actions(actions, params):
	if "glob_json" in actions:
		glob_json(params)
	if "crafting_report" in actions:
		crafting_report(params)
	if "weapon_energy_report" in actions:
		weapon_energy_report(params)
	if "converter_recipe_report" in actions:
		converter_recipe_report(params)
	if "count_mounts" in actions:
		count_mounts(params)
	if "search_invalid_filename" in actions:
		search_invalid_filename(params)
	if "list_types" in actions:
		list_types(params)

if __name__ == '__main__':
	actions = [
		"glob_json",
		"search_invalid_filename",
		#"count_mounts",
		#"crafting_report",
		#"weapon_energy_report",
		#"converter_recipe_report",
		#"list_types",
		"nothing" # just so I don't need to play with the last ,
	]
	params = {
		"glob_root":os.path.join(DATA_FOLDER, "json"),
		"nothing" : None # don't have to deal with last ,
	}
	do_actions(actions, params)

max_stars_repo_name: Ombarus/SolarRogue

max_issues_repo_name: Ombarus/SolarRogue

max_forks_repo_name: Ombarus/SolarRogue


text:
import pytest
import string
import numpy as np
import pandas as pd
from fast_plotter.postproc import functions as funcs

def _make_string(index):
    chars = string.printable
    start = index % len(chars)
    stop = (index + 28) % len(chars)
    return chars[start] + chars[stop]

@pytest.fixture
def binned_df():
    anInt = list(range(4))
    aCat = ["foo", "bar"]
    anInterval = pd.IntervalIndex.from_breaks(np.linspace(100, 104, 6))
    index = pd.MultiIndex.from_product([anInt, aCat, anInterval],
                                       names=["int", "cat", "interval"])
    data = dict(a=np.arange(len(index)), b=map(_make_string, range(len(index))))
    df = pd.DataFrame(data, index=index)
    return df

def test_query(binned_df):
    result = funcs.query(binned_df, "cat=='bar'")
    assert len(result) == 20

    result = funcs.query(binned_df, "cat=='bar' and int > 2")
    assert len(result) == 5

def test_rebin(binned_df):
    result = funcs.rebin(binned_df.copy(), rename="hruff", axis="int", mapping=dict(zip(range(4), [0, 2] * 2)))
    assert len(result) == 20
    assert list(result.index.unique("hruff")) == [0, 2]

    mapping = {0: dict(bar="foo"), 2: dict(foo="bar"), 3: dict(foo="BAZ", bar="BAZ")}
    result = funcs.rebin(binned_df.copy(), axis=["int", 'cat'], mapping=mapping)
    assert len(result) == 25
    assert set(result.index.unique("cat")) == {"bar", "BAZ", "foo"}

def test_keep_bins(binned_df):
    result = funcs.keep_bins(binned_df.copy(), "int", keep=[0, 2])
    assert len(result) == 20

    result = funcs.keep_bins(binned_df.copy(), "int", keep=binned_df.index.unique("int"))
    assert len(result) == 40

def test_keep_specific_bins(binned_df):
    keep = {"0": ["foo"], "1": ["bar"]}
    result = funcs.keep_specific_bins(binned_df, axis=["int", "cat"], keep=keep)
    assert len(result) == 10

    result = funcs.keep_specific_bins(binned_df, axis=["int", "cat"], keep=keep, expansions=[])
    assert len(result) == 10

    expansions = dict(one=["bar"])
    keep = {"0": ["foo"], "1": ["{one}"]}
    result = funcs.keep_specific_bins(binned_df, axis=["int", "cat"], keep=keep, expansions=expansions)
    assert len(result) == 10

    expansions = dict(one=["bar"], two=["foo", "bar"])
    keep = {"0": ["{two}"], "1": ["{one}"]}
    result = funcs.keep_specific_bins(binned_df, axis=["int", "cat"], keep=keep, expansions=expansions)
    assert len(result) == 15

def test_combine_cols_AND_split_dimension(binned_df):
    result = funcs.combine_cols(binned_df, {"a;b": "{a};{b}"})
    assert len(result.columns) == 3
    assert all(result.columns == ["a", "b", "a;b"])
    assert len(result) == 40

# def test_split_dimension(binned_df):
#     #def split_dimension(df, axis, delimeter=";"):
#     result = funcs.split_dimension(binned_df, ["interval"], ",")
#     pass

# def test_regex_split_dimension():
#     #def regex_split_dimension(df, axis, regex):
#     pass

# def test_rename_cols():
#     #def rename_cols(df, mapping):
#     pass

def test_rename_dim(binned_df):
    result = funcs.rename_dim(binned_df, {"int": "integers", "cat": "CATEGORICALS"})
    assert result.index.names == ["integers", "CATEGORICALS", "interval"]

def test_split(binned_df):
    results = funcs.split(binned_df, "cat", keep_split_dim=True)
    assert len(results) == 2
    assert all([r[0].index.nlevels == 3 for r in results])

    results = funcs.split(binned_df, "int", keep_split_dim=True)
    assert len(results) == 4
    assert all([r[0].index.nlevels == 3 for r in results])

def test_filter_cols(binned_df):
    df = binned_df.index.to_frame()

    result = funcs.filter_cols(df, items=["int"])
    assert len(result.columns) == 1
    assert result.columns[0] == "int"

    result = funcs.filter_cols(df, items=["int", "cat"])
    assert len(result.columns) == 2
    assert set(result.columns) == set(("int", "cat"))

    result = funcs.filter_cols(df, like="int")
    assert len(result.columns) == 2
    assert set(result.columns) == set(("int", "interval"))

    result = funcs.filter_cols(df, like=["int", "cat"])
    assert len(result.columns) == 3
    assert set(result.columns) == set(("int", "cat", "interval"))

    result = funcs.filter_cols(df, regex="^int.*")
    assert len(result.columns) == 2
    assert set(result.columns) == set(("int", "interval"))

# def test_reorder_dimensions():
#     #def reorder_dimensions(df, order):
#     pass

# def test_densify():
#     #def densify(df, known={}):
#     pass

# def test_stack_weights():
#     #def stack_weights(df, drop_n_col=False):
#     pass

# def test_to_datacard_inputs():
#     #def to_datacard_inputs(df, select_data, rename_syst_vars=False):
#     pass

# def test_assign_col():
#     #def assign_col(df, assignments={}, evals={}, drop_cols=[]):
#     pass

# def test_assign_dim():
#     #def assign_dim(df, assignments={}, evals={}, drop_cols=[]):
#     pass

# def test_merge():
#     #def merge(dfs):
#     pass

def test_multiply_values(binned_df):
    result = funcs.multiply_values(binned_df, constant=3)
    assert np.array_equal(result.a, np.arange(len(binned_df)) * 3)
    assert np.array_equal(result.b, binned_df.b)

    result = funcs.multiply_values(binned_df, apply_if="int % 2 == 0", constant=3)
    assert np.array_equal(result.a, np.arange(len(binned_df)) * np.repeat([3, 1, 3, 1], 10))
    assert np.array_equal(result.b, binned_df.b)

    result = funcs.multiply_values(binned_df, mapping={"int % 2 == 0": 3, "int % 2 == 1": 7.2})
    assert np.array_equal(result.a, np.arange(len(binned_df)) * np.repeat([3, 7.2, 3, 7.2], 10))
    assert np.array_equal(result.b, binned_df.b)

    result = funcs.multiply_values(binned_df, mapping={"cat=='foo'": 1.2, "cat=='bar'": 19})
    tiled_vals = np.tile(np.repeat([1.2, 19], 5), 4)
    assert np.array_equal(result.a, np.arange(len(binned_df)) * tiled_vals)
    assert np.array_equal(result.b, binned_df.b)

# def test_multiply_dataframe():
#     #def multiply_dataframe(df, multiply_df, use_column=None):
#     pass

# def test_normalise_group():
#     #def normalise_group(df, groupby_dimensions, apply_if=None, use_column=None):
#     pass

# def test_open_many():
#     #def open_many(file_list, return_meta=True):
#     pass

# def test_write_out():
#     #def write_out(df, meta, filename="tbl_{dims}--{name}.csv", out_dir=None):
#     pass

max_stars_repo_name: eshwen/fast-plotter

max_issues_repo_name: eshwen/fast-plotter

max_forks_repo_name: benkrikler/fast-plotter-github-test


text:
import numpy as np
import pandas as pd
import matplotlib.pyplot as plt
from matplotlib.ticker import MultipleLocator

# 使用numpy产生数据
x = np.arange(-5, 5, 0.1)
y = x * 3

# 创建窗口、子图
# 方法1：先创建窗口，再创建子图。（一定绘制）
fig = plt.figure(num=1, figsize=(15, 8), dpi=80)  # 开启一个窗口，同时设置大小，分辨率
ax1 = fig.add_subplot(2, 1, 1)  # 通过fig添加子图，参数：行数，列数，第几个。
ax2 = fig.add_subplot(2, 1, 2)  # 通过fig添加子图，参数：行数，列数，第几个。
print(fig, ax1, ax2)
# 方法2：一次性创建窗口和多个子图。（空白不绘制）
fig, axarr = plt.subplots(4, 1)  # 开一个新窗口，并添加4个子图，返回子图数组
ax1 = axarr[0]  # 通过子图数组获取一个子图
print(fig, ax1)
# 方法3：一次性创建窗口和一个子图。（空白不绘制）
ax1 = plt.subplot(1, 1, 1, facecolor='white')  # 开一个新窗口，创建1个子图。facecolor设置背景颜色
print(ax1)
# 获取对窗口的引用，适用于上面三种方法
# fig = plt.gcf()   #获得当前figure
# fig=ax1.figure   #获得指定子图所属窗口

# fig.subplots_adjust(left=0)                         #设置窗口左内边距为0，即左边留白为0。

# 设置子图的基本元素
ax1.set_title('python-drawing')  # 设置图体，plt.title
ax1.set_xlabel('x-name')  # 设置x轴名称,plt.xlabel
ax1.set_ylabel('y-name')  # 设置y轴名称,plt.ylabel
plt.axis([-6, 6, -10, 10])  # 设置横纵坐标轴范围，这个在子图中被分解为下面两个函数
ax1.set_xlim(-5, 5)  # 设置横轴范围，会覆盖上面的横坐标,plt.xlim
ax1.set_ylim(-10, 10)  # 设置纵轴范围，会覆盖上面的纵坐标,plt.ylim

xmajorLocator = MultipleLocator(2)  # 定义横向主刻度标签的刻度差为2的倍数。就是隔几个刻度才显示一个标签文本
ymajorLocator = MultipleLocator(3)  # 定义纵向主刻度标签的刻度差为3的倍数。就是隔几个刻度才显示一个标签文本

ax1.xaxis.set_major_locator(xmajorLocator)  # x轴 应用定义的横向主刻度格式。如果不应用将采用默认刻度格式
ax1.yaxis.set_major_locator(ymajorLocator)  # y轴 应用定义的纵向主刻度格式。如果不应用将采用默认刻度格式

ax1.xaxis.grid(True, which='major')  # x坐标轴的网格使用定义的主刻度格式
ax1.yaxis.grid(True, which='major')  # x坐标轴的网格使用定义的主刻度格式

ax1.set_xticks([])  # 去除坐标轴刻度
ax1.set_xticks((-5, -3, -1, 1, 3, 5))  # 设置坐标轴刻度
ax1.set_xticklabels(labels=['x1', 'x2', 'x3', 'x4', 'x5','x6'], rotation=-30,
                    fontsize='small')  # 设置刻度的显示文本，rotation旋转角度，fontsize字体大小

plot1 = ax1.plot(x, y, marker='o', color='g', label='legend1')  # 点图：marker图标
plot2 = ax1.plot(x, y, linestyle='--', alpha=0.5, color='r',
                 label='legend2')  # 线图：linestyle线性，alpha透明度，color颜色，label图例文本

ax1.legend(loc='upper left')  # 显示图例,plt.legend()
ax1.text(2.8, 7, r'y=3*x')  # 指定位置显示文字,plt.text()
ax1.annotate('important point', xy=(2, 6), xytext=(3, 1.5),  # 添加标注，参数：注释文本、指向点、文字位置、箭头属性
             arrowprops=dict(facecolor='black', shrink=0.05),
             )
# 显示网格。which参数的值为major(只绘制大刻度)、minor(只绘制小刻度)、both，默认值为major。axis为'x','y','both'
ax1.grid(b=True, which='major', axis='both', alpha=0.5, color='skyblue', linestyle='--', linewidth=2)

axes1 = plt.axes([.2, .3, .1, .1], facecolor='y')  # 在当前窗口添加一个子图，rect=[左, 下, 宽, 高]，是使用的绝对布局，不和以存在窗口挤占空间
axes1.plot(x, y)  # 在子图上画图
plt.savefig('aa.jpg', dpi=400, bbox_inches='tight')  # savefig保存图片，dpi分辨率，bbox_inches子图周边白色空间的大小
plt.show()  # 打开窗口，对于方法1创建在窗口一定绘制，对于方法2方法3创建的窗口，若坐标系全部空白，则不绘制


text:
#Author: Bogdan Bintu
#Date: 6/3/2019
# Fitting tools
import numpy as np
import pyfftw
from pyfftw.interfaces.numpy_fft import rfftn,irfftn
from scipy import fftpack

#gaussian functions
from scipy import signal
def gker(gaus=[3,3,3],exp=8):
    szgaus = [int(gaus_*exp) for gaus_ in gaus]
    gker_ = np.outer(np.outer(signal.gaussian(szgaus[0], gaus[0]), signal.gaussian(szgaus[1], gaus[1])),
             signal.gaussian(szgaus[2], gaus[2])).reshape(szgaus)
    gker_ = gker_/np.sum(gker_)
    return gker_.astype(np.float32)
def reflect(arr,sizes=[10,10,10]):
    arr_ = arr.copy()
    for axis,size in zip(list(range(len(arr.shape))),sizes):
        arr_ = arr_.swapaxes(axis,0)
        arr_ = np.concatenate([arr_[:size][::-1],arr_,arr_[-size:][::-1]])
        arr_ = arr_.swapaxes(axis,0)
    return arr_[:-1,:-1,:-1]
def fft_gaussian(arr,gaus=[3,3,3],exp=8):
    gk3d = gker(gaus,exp)
    szs = (np.array(gk3d.shape)/2).astype(int)
    arr_ = reflect(arr,sizes=szs)
    return signal.fftconvolve(arr_,gk3d,mode='valid')
def best_rotation(zns_zxy,zos_zxy):
    """Returns drift and rotation matrix. Use as: np.dot(R,zns_zxy)+t==zos_zxy"""
    PA = np.array(zns_zxy)
    PB = np.array(zos_zxy)
    cA = np.mean(PA,0)
    cB = np.mean(PB,0)
    H=np.dot((PA-cA).T,(PB-cB))
    U,S,V = np.linalg.svd(H)
    R = np.dot(U,V)
    if np.linalg.det(R)<0:
        R[:,-1]*=-1
    t = -np.dot(cA,R)+cB
    #np.dot(R,zns_zxy)+t==zos_zxy
    return t,R

def _centered(arr, newshape):
    # Return the center newshape portion of the array.
    newshape = np.asarray(newshape)
    currshape = np.array(arr.shape)
    startind = (currshape - newshape) // 2
    endind = startind + newshape
    myslice = [slice(startind[k], endind[k]) for k in range(len(endind))]
    return arr[tuple(myslice)]
def fftconvolvefast(in1,in2):
    s1 = np.array(in1.shape)
    s2 = np.array(in2.shape)

    shape = s1 + s2 - 1

    fshape = [fftpack.helper.next_fast_len(int(d)) for d in shape]
    fslice = tuple([slice(0, int(sz)) for sz in shape])

    sp1 = rfftn(in1, fshape)
    sp2 = rfftn(in2, fshape)
    ret = (irfftn(sp1 * sp2, fshape)[fslice].copy())
    return _centered(ret, s1 - s2 + 1)
def fft_gaussian_fast(arr,gaus=[3,3,3],exp=8):
    gk3d = gker(gaus,exp)
    szs = (np.array(gk3d.shape)/2).astype(int)
    arr_ = reflect(arr,sizes=szs)
    return fftconvolvefast(arr_,gk3d)
from scipy.ndimage.filters import maximum_filter
def get_seed_points_base(im_sm,gfilt_size=5,filt_size=3,th_seed=3.,max_num=None):
    """Get the seed points in an image.
    #1 perform a gaussian filter
    #2 find local maxima within a radius <filt_size> with <th_seed> above local maximum (in log space)
    """
    im_diff = np.log(im_sm)-np.log(fft_gaussian_fast(im_sm,gaus=[gfilt_size]*3))
    im_tomax = im_diff
    im_max = maximum_filter(im_tomax,filt_size)
    z,x,y = np.where(im_tomax == im_max)
    std_ = np.std(im_diff)
    h = im_diff[z,x,y]
    kp = h>th_seed*std_
    z,x,y,h = z[kp],x[kp],y[kp],h[kp]
    centers = np.array([z,x,y])
    #sort by relative brightness
    ind = np.argsort(h)[::-1]
    centers = np.array([z[ind],x[ind],y[ind],h[ind]])
    if max_num is not None:
        centers = centers[:,:max_num]
    return centers,std_

import cv2
def normalzie_im(im,sz=20):
    im_ = np.array(im,dtype=np.float32)
    im_blur = np.array([cv2.blur(im__,(sz,sz)) for im__ in im_],dtype=np.float32)
    im_ =im_-im_blur
    return im_

def get_seed_points_base_v2(im_sm,gfilt_size=5,filt_size=3,th_seed=3.,max_num=None):
    """Get the seed points in an image.
    #1 perform a gaussian filter
    #2 find local maxima within a radius <filt_size> with <th_seed> above local maximum (in log space)
    """
    if gfilt_size!=0: im_norm = normalzie_im(im_sm,gfilt_size)
    else: im_norm=im_sm
    std_ = np.std(im_norm)
    hcutoff = std_*th_seed
    z,x,y = np.where(im_norm>hcutoff)
    h_im = im_norm[z,x,y]
    sz,sx,sy = im_norm.shape
    keep = h_im>0
    pix = int(filt_size/2)
    deltas = list(range(-pix,pix+1))
    for deltax in deltas:
        for deltay in deltas:
            for deltaz in deltas:
                #if (np.abs(deltax)+np.abs(deltay)+np.abs(deltaz))>0:
                keep &= (h_im>=im_norm[(z+deltaz)%sz,(x+deltax)%sx,(y+deltay)%sy])
    zf,xf,yf = z[keep],x[keep],y[keep]
    hf = im_norm[zf,xf,yf]
    inds = np.argsort(hf)[::-1]
    centers_zxyh = np.array([zf[inds],xf[inds],yf[inds],hf[inds]])
    if max_num is not None:
        centers_zxyh = centers_zxyh[:,:max_num]
    return centers_zxyh,std_

def gauss_ker(sig_xyz=[2,2,2],sxyz=16,xyz_disp=[0,0,0]):
    dim = len(xyz_disp)
    xyz=np.indices([sxyz+1]*dim)
    for i in range(len(xyz.shape)-1):
        sig_xyz=np.expand_dims(sig_xyz,axis=-1)
        xyz_disp=np.expand_dims(xyz_disp,axis=-1)
    im_ker = np.exp(-np.sum(((xyz-xyz_disp-sxyz/2.)/sig_xyz)**2,axis=0)/2.)
    return im_ker
def add_source(im_,pos=[0,0,0],h=200,sig=[2,2,2]):
    im=np.array(im_,dtype=float)
    pos_int = np.array(pos,dtype=int)
    xyz_disp = -pos_int+pos
    im_ker = gauss_ker(sig_xyz=sig,sxyz=int(np.max(sig)*20),xyz_disp=xyz_disp)
    im_ker_sz = np.array(im_ker.shape,dtype=int)
    pos_min = (pos_int-im_ker_sz/2).astype(int)
    pos_max = pos_min+im_ker_sz
    im_shape = np.array(im.shape)
    def in_im(pos__):
        pos_=np.array(pos__,dtype=int)
        pos_[pos_>=im_shape]=im_shape[pos_>=im_shape]-1
        pos_[pos_<0]=0
        return pos_
    pos_min_ = in_im(pos_min)
    pos_max_ = in_im(pos_max)
    pos_min_ker = pos_min_-pos_min
    pos_max_ker = im_ker_sz+pos_max_-pos_max
    #print zip(pos_min_ker,pos_max_ker),zip(pos_min_,pos_max_),zip(pos_min,pos_max)
    slices_ker = [slice(pm,pM)for pm,pM in zip(pos_min_ker,pos_max_ker)]
    slices_im = [slice(pm,pM)for pm,pM in zip(pos_min_,pos_max_)]
    im[slices_im]+=im_ker[slices_ker]*h
    return im

from scipy.optimize import leastsq
class GaussianFit():
    def __init__(self,im,X,center=None,n_aprox=10,min_w=0.5,max_w=4.,delta_center=3.,
                 init_w=1.5):
        self.min_w = min_w*min_w
        self.max_w = max_w*max_w

        self.delta_center = delta_center
        self.im = np.array(im,dtype=np.float32)
        self.x,self.y,self.z = np.array(X,dtype=np.float32)
        #get estimates
        argsort_im = np.argsort(im)
        if center is None:
            center = np.median(X[:,argsort_im][:,-n_aprox:],-1)
        self.center_est = center
        sorted_im = im[argsort_im]
        eps =  np.exp(-10.)
        bk_guess = np.log(np.max([np.mean(sorted_im[:n_aprox]),eps]))
        h_guess = np.log(np.max([np.mean(sorted_im[-n_aprox:]),eps]))
        wsq = init_w**2
        wg = np.log((self.max_w - wsq)/(wsq-self.min_w))
        self.p_ = np.array([bk_guess,h_guess,0,0,0,wg,wg,wg,0,0],dtype=np.float32)
        self.to_natural_paramaters()
        self.success = False

    def to_center(self,c0_,c1_,c2_):
        """constrains via sigmoidal function close to local center"""
        delta = self.delta_center
        # c0
        if c0_ >= np.log(np.finfo(c0_.dtype).max):
            c0 = - delta + self.center_est[0]
        elif c0_ <= - np.log(np.finfo(c0_.dtype).max):
            c0 = delta + self.center_est[0]
        else:
            c0 = 2. *delta / (1.+np.exp(c0_)) - delta + self.center_est[0]
        # c1
        if c1_ >= np.log(np.finfo(c1_.dtype).max):
            c1 = - delta + self.center_est[1]
        elif c1_ <= - np.log(np.finfo(c1_.dtype).max):
            c1 = delta + self.center_est[1]
        else:
            c1 = 2. *delta / (1.+np.exp(c1_)) - delta + self.center_est[1]
        # c2
        if c2_ >= np.log(np.finfo(c2_.dtype).max):
            c2 = - delta + self.center_est[2]
        elif c2_ <= - np.log(np.finfo(c2_.dtype).max):
            c2 = delta + self.center_est[2]
        else:
            c2 = 2. *delta / (1.+np.exp(c2_)) - delta + self.center_est[2]

        #c0 = 2. *delta / (1.+np.exp(c0_)) - delta + self.center_est[0]
        #c1 = 2.*delta/(1.+np.exp(c1_))-delta+self.center_est[1]
        #c2 = 2.*delta/(1.+np.exp(c2_))-delta+self.center_est[2]
        return c0,c1,c2

    def to_sine(self, t_):
        """constrain sin-angles to -1,1"""
        #eps =  10E-5
        #self.sine_eps = eps
        #return 2.*(1-eps)/(1+np.exp(t_))-1.+eps
        if t_ >= np.log(np.finfo(t_.dtype).max):
            return -1
        elif t_ <= - np.log(np.finfo(t_.dtype).max):
            return 1
        else:
            return 2./(1+np.exp(t_))-1.

    def to_ws(self,w_):
        """constrain widths"""
        min_ws = self.min_w
        delta_ws = self.max_w - min_ws

        if w_ >= np.log(np.finfo(w_.dtype).max):
            ws = min_ws
        elif w_ <= - np.log(np.finfo(w_.dtype).max):
            ws = delta_ws + min_ws
        else:
            ws = delta_ws/(1.+np.exp(w_))+min_ws
        return ws

    def to_natural_paramaters(self,parms=None):
        """
        Convert from constrained paramaters to [hf,xc,yc,zc,bkf,w1f,w2f,w3f,t,p,eps]
        """
        if parms is None:
            parms = self.p_
        bk,h,xp,yp,zp,w1,w2,w3,pp,tp = parms
        bkf,hf=np.exp(bk),np.exp(h)
        t,p = self.to_sine(tp),self.to_sine(pp)
        w1f,w2f,w3f = np.sqrt(self.to_ws(w1)),np.sqrt(self.to_ws(w2)),np.sqrt(self.to_ws(w3))
        xc,yc,zc = self.to_center(xp,yp,zp)
        eps = self.calc_eps(parms)
        eps = np.mean(np.abs(eps))
        self.p = np.array([hf,xc,yc,zc,bkf,w1f,w2f,w3f,t,p,eps],dtype=np.float32)
        return self.p
    def calc_f(self,parms):
        self.p_ = parms

        bk,h,xp,yp,zp,w1,w2,w3,pp,tp = parms
        t,p = self.to_sine(tp),self.to_sine(pp)
        ws1,ws2,ws3 = self.to_ws(w1),self.to_ws(w2),self.to_ws(w3)
        xc,yc,zc = self.to_center(xp,yp,zp)
        xt,yt,zt = self.x-xc,self.y-yc,self.z-zc

        p2 = p*p
        t2 = t*t
        tc2 = 1-t2
        pc2 = 1-p2
        tc= np.sqrt(tc2)
        pc= np.sqrt(pc2)
        s1,s2,s3 = 1./ws1,1./ws2,1./ws3
        x2c = pc2*tc2*s1 + t2*s2 + p2*tc2*s3
        y2c = pc2*t2*s1 + tc2*s2 + p2*t2*s3
        z2c = p2*s1 + pc2*s3
        xyc = 2*tc*t*(pc2*s1 - s2 + p2*s3)
        xzc = 2*p*pc*tc*(s3 - s1)
        yzc = 2*p*pc*t*(s3 - s1)

        xsigmax = x2c*xt*xt+y2c*yt*yt+z2c*zt*zt+xyc*xt*yt+xzc*xt*zt+yzc*yt*zt
        self.f0 = np.exp(h-0.5*xsigmax)
        self.f = np.exp(bk)+self.f0

        return self.f

    def frac_conv(self,x1,x2):return 2*np.abs(x1-x2)/(x1+x2)<self.eps_frac
    def dist_conv(self,x1,x2):return np.abs(x1-x2)<self.eps_dist
    def angle_conv(self,x1,x2):return np.abs(x1-x2)<self.eps_angle
    def calc_eps(self,parms):
        """
        calculate the loss function
        """
        #Decided not to include this extra step of convergence
        """
        if self.p_old is not None:
            if np.any((parms-self.parms_old)!=0):
                p_new = self.to_natural_paramaters(parms)
                h1,x1,y1,z1,bk1,wx1,wy1,wz1,t1,p1 = self.p_old[:10]
                h2,x2,y2,z2,bk2,wx2,wy2,wz2,t2,p2 = p_new[:10]
                #print self.p_old,p_new
                self.p_old = p_new

                converged = self.frac_conv(h1,h2) and self.frac_conv(bk1,bk2)
                converged = converged and self.dist_conv(x1,x2) and self.dist_conv(y1,y2) and self.dist_conv(z1,z2)
                converged = converged and self.dist_conv(wx1,wx2) and self.dist_conv(wy1,wy2) and self.dist_conv(wz1,wz2)
                converged = converged and self.angle_conv(t1,t2) and self.angle_conv(p1,p2)
                if converged:
                    self.converged = True
                    return np.zeros(len(self.im),dtype=np.float32)
        else:
            self.p_old = self.to_natural_paramaters(parms)
            self.parms_old = parms
        """
        return self.calc_f(parms)-self.im
    def calc_jac(self,parms):
        bk,h,xp,yp,zp,w1,w2,w3,pp,tp = parms
        t,p = self.to_sine(tp),self.to_sine(pp)
        ws1,ws2,ws3 = self.to_ws(w1),self.to_ws(w2),self.to_ws(w3)
        xc,yc,zc = self.to_center(xp,yp,zp)
        xt,yt,zt = self.x-xc,self.y-yc,self.z-zc
        p2 = p*p
        t2 = t*t
        tc2 = 1-t2
        pc2 = 1-p2
        tc= np.sqrt(tc2)
        pc= np.sqrt(pc2)
        s1,s2,s3 = 1./ws1,1./ws2,1./ws3
        x2c = pc2*tc2*s1 + t2*s2 + p2*tc2*s3
        y2c = pc2*t2*s1 + tc2*s2 + p2*t2*s3
        z2c = p2*s1 + pc2*s3
        xyc = 2*tc*t*(pc2*s1 - s2 + p2*s3)
        xzc = 2*p*pc*tc*(s3 - s1)
        yzc = 2*p*pc*t*(s3 - s1)
        xt2,xtyt,xtzt,yt2,ytzt,zt2 = xt*xt,xt*yt,xt*zt,yt*yt,yt*zt,zt*zt
        xsigmax = x2c*xt2+y2c*yt2+z2c*zt2+xyc*xtyt+xzc*xtzt+yzc*ytzt

        d,minw,maxw = self.delta_center,self.min_w,self.max_w

        f2 = np.exp(h-0.5*xsigmax)
        f1 =  np.exp(bk)+np.zeros(len(f2))
        e_xp,e_yp,e_zp = np.exp(-np.abs(xp)),np.exp(-np.abs(yp)),np.exp(-np.abs(zp))
        norm_xp = -d*e_xp/((1 + e_xp)*(1 + e_xp))
        norm_yp = -d*e_yp/((1 + e_yp)*(1 + e_yp))
        norm_zp = -d*e_zp/((1 + e_zp)*(1 + e_zp))
        f3 = (f2*(2*x2c*xt + xyc*yt + xzc*zt))*norm_xp
        f4 = (f2*(xt*xyc + 2*y2c*yt + yzc*zt))*norm_yp
        f5 = (f2*(xt*xzc + yt*yzc + 2*z2c*zt))*norm_zp
        f6 = (f2*(-pc2*tc2*xt2 - 2*pc2*t*tc*xtyt - pc2*t2*yt2 + 2*p*pc*tc*xtzt + 2*p*pc*t*ytzt - p2*zt2))*self.norm_w(w1,minw,maxw)
        f7 = (f2*(-t2*xt2 + 2*t*tc*xtyt - tc2*yt2))*self.norm_w(w2,minw,maxw)
        f8 = (f2*(-p2*tc2*xt2 - 2*p2*t*tc*xtyt - p2*t2*yt2 - 2*p*pc*tc*xtzt - 2*p*pc*t*ytzt - pc2*zt2))*self.norm_w(w3,minw,maxw)
        e_p = np.exp(-np.abs(pp)/2)
        norm_p = e_p/(1+e_p*e_p)
        f9 = f2*(s3-s1)*((2*pc2-1.)*(tc*xtzt + t*ytzt) + p*pc*(tc2*xt2 + 2*t*tc*xtyt + t2*yt2 - zt2))*norm_p
        e_t = np.exp(-np.abs(tp)/2)
        norm_t = e_t/(1+e_t*e_t)
        f10 = f2*((pc2 *s1 - s2 +  p2*s3)*(t *tc*(yt2 - xt2) - (t2 - tc2)*xtyt) + p* pc *(s1 - s3)* (t *xtzt - tc*ytzt))*norm_t

        self.jac = np.array([f1,f2,f3,f4,f5,f6,f7,f8,f9,f10],np.float32).T

        return self.jac

    def norm_w(self,w,minw,maxw):
        if w>0:
            e_w = np.exp(-w)
            return 0.5*(maxw - minw)*e_w/(maxw*e_w+minw)**2
        else:
            e_w = np.exp(w)
            return 0.5*(maxw - minw)*e_w/(minw*e_w+maxw)**2

    def fit(self,eps_frac=10E-3,eps_dist=10E-3,eps_angle=10E-3):
        """
        This implements the Levenberg-Marquardt algorithm for 3D gaussian fitting.
        Stores the results in [height,x,y,z,background,width_1,width_2,width_3,sin_theta,sin_phi,error] = self.p
        """
        if len(self.p_)>len(self.im):
            self.success = False
        else:
            self.eps_frac,self.eps_dist,self.eps_angle = eps_frac,eps_dist,eps_angle
            parms0 = self.p_
            self.p_old = None
            parmsf,_ = leastsq(self.calc_eps,parms0,Dfun=self.calc_jac, maxfev=1000) # changed maxfev to block warning message
            #parmsf=parms0#####pisici
            self.p_ = parmsf
            self.to_natural_paramaters()
            self.center = self.p[1:4]
            self.success = True
    def get_im(self):
        self.calc_f(self.p_)
        return self.f0

def in_dim(x,y,z,xmax,ymax,zmax):
    keep = ((x>=0)&(x<xmax)&(y>=0)&(y<ymax)&(z>=0)&(z<zmax))>0
    return x[keep],y[keep],z[keep]
from scipy.spatial.distance import cdist
def closest(x,y,z,ic,all_centers):
    xyz = np.array([x,y,z]).T
    dists = cdist(xyz,all_centers)
    closest_point_id = np.argmin(dists,axis=-1)
    keep = closest_point_id==ic
    xyz_keep = xyz[keep,:]
    return xyz_keep.T
def closest_fast(x,y,z,ic,all_centers,rsearch =6):
    xyz = np.array([x,y,z]).T
    dists = cdist([all_centers[ic]],all_centers,'chebyshev')[0,:]
    close = dists<rsearch
    close_inds = np.arange(len(all_centers))[close]
    close_centers = all_centers[close]
    dists = cdist(xyz,close_centers)
    closest_point_id = np.argmin(dists,axis=-1)
    keep = close_inds[closest_point_id]==ic
    xyz_keep = xyz[keep,:]
    return xyz_keep.T
from scipy.spatial import cKDTree as KDTree
def closest_faster(xyz,ic,tree,rsearch = 6):
    dists_,nns_ = tree.query(xyz,distance_upper_bound=rsearch)
    return xyz[nns_==ic].T

def inv_sigma(sigma):
    [[a, d, e], [d, b, f], [e, f, c]] = sigma

    sigma_inv = np.array([[(b*c - f ** 2)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (-(c*d) + e*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (-(b*e) + d*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2)],
                      [(-(c*d) + e*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (a*c - e ** 2)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (d*e - a*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2)],
                      [(-(b*e) + d*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (d*e - a*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (a*b - d ** 2)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2)]])
    return sigma_inv
def gfit_fast(im_,X_,bk_f=0.1,reconstruct=False,plt_val=False,compare_with_fitting=False):
    pfinal = np.array([np.nan]*12)
    if len(im_)>0:
        #background, height and new weights
        nim = len(im_)
        bk = np.sort(im_)[int(nim*bk_f)]
        weights = im_-bk
        weights[weights<0]=0
        h = np.max(weights)
        weights = weights/np.sum(weights)

        #approximate position
        zxy_ = np.sum(X_*weights,-1)
        z_,x_,y_ = zxy_
        #approximate shape
        X_c = X_.T-zxy_
        Cov = np.array([[ X_c[:,i]*X_c[:,j]for i in range(3)] for j in range(3)])
        Cov = np.sum(Cov*weights,-1)
        [[a, d, e], [d, b, f], [e, f, c]] = Cov
        #compute reconstruction
        eps = np.nan
        if reconstruct:
            iCov = inv_sigma(Cov)
            im_fit = h*np.exp(-np.sum(np.dot(X_c,iCov)*X_c,-1)*0.5)+bk
            eps = np.mean(np.abs(im_-im_fit))
        pfinal = np.array([h,z_,x_,y_,bk,a,b,c,d,e,f,eps])
    if plt_val:
        import matplotlib.pyplot as plt
        zs,xs,ys = np.min(X_,1)

        X_im = X_.T-[zs,xs,ys]
        size_im = np.max(X_im.T,1)+1
        im__ = np.zeros(size_im)+np.nan
        im__[X_im[:,0],X_im[:,1],X_im[:,2]]=im_
        plt.figure()
        plt.imshow(np.nanmax(im__,0))
        plt.plot(y_-ys,x_-xs,'bx')
        # plot reconstructed image and fit
        im__ = np.zeros(size_im)+np.nan
        im__[X_im[:,0],X_im[:,1],X_im[:,2]]=im_fit
        plt.figure()
        plt.imshow(np.nanmax(im__,0))
        plt.plot(y_-ys,x_-xs,'bx')
        plt.figure()
        plt.plot(im_,im_fit,'o')
        if compare_with_fitting:
            center = X_.T[np.argmax(im_)]
            obj = GaussianFit(im_,X_,center=center,delta_center=1.5)
            obj.fit()
            im_fit2 = obj.get_im()+obj.p[4]
            im__ = np.zeros(size_im)+np.nan
            im__[X_im[:,0],X_im[:,1],X_im[:,2]]=im_fit2
            plt.figure()
            plt.imshow(np.nanmax(im__,0))
            plt.plot(y_-ys,x_-xs,'bx')
            plt.plot(obj.p[3]-ys,obj.p[2]-xs,'rx')
            #print obj.p[1:3]-pfinal[1:3]
    return pfinal
def in_dim(x,y,z,xmax,ymax,zmax):
    keep = ((x>=0)&(x<xmax)&(y>=0)&(y<ymax)&(z>=0)&(z<zmax))>0
    return x[keep],y[keep],z[keep]

from scipy.spatial import cKDTree as KDTree
def fast_fit_big_image(im,centers_zxy,radius_fit = 4,avoid_neigbors=True,recenter=False ,verbose = True,better_fit=False,troubleshoot=False):
    ps = []
    if len(centers_zxy)>0:
        if avoid_neigbors:
            centers_tree = KDTree(centers_zxy)
            list_inters = centers_tree.query_ball_tree(centers_tree,radius_fit*2)
        centers_ = centers_zxy
        if verbose:
            from tqdm import tqdm_notebook as tqdm
            centers_ = tqdm(centers_zxy)
        zb,xb,yb = np.reshape(np.indices([radius_fit*2]*3)-radius_fit,[3,-1]).astype(int)
        keep = zb*zb+xb*xb+yb*yb<=radius_fit**2
        zb,xb,yb = zb[keep],xb[keep],yb[keep]
        sz,sx,sy = im.shape
        X_c = np.array([zb,xb,yb]).T

        for ic,(zc,xc,yc) in enumerate(centers_):
            if avoid_neigbors:
                common_pts_ids = list_inters[ic]
                index_ic = list_inters[ic].index(ic)
                centers__ = centers_zxy[common_pts_ids]-[zc,xc,yc]
                nns_ = np.argmin(cdist(centers__,X_c),0)
                zb_,xb_,yb_ = X_c[nns_==index_ic].T
            else:
                zb_,xb_,yb_ = zb,xb,yb

            z_keep,x_keep,y_keep = int(zc)+zb_,int(xc)+xb_,int(yc)+yb_
            z_keep,x_keep,y_keep = in_dim(z_keep,x_keep,y_keep,sz,sx,sy)

            X_ = np.array([z_keep,x_keep,y_keep]).T
            im_ = im[z_keep,x_keep,y_keep]

            #recenter to max in the imge - generally indicated if using seeds from a different image
            if recenter:
                if len(im_)>0:
                    #while True:
                    im_c = np.argmax(im_)
                    zcn,xcn,ycn = z_keep[im_c],x_keep[im_c],y_keep[im_c]
                    #if zcn==zc and xcn==xc and ycn==yc:
                    #    break
                    zc,xc,yc = zcn,xcn,ycn
                    z_keep,x_keep,y_keep = int(zc)+zb_,int(xc)+xb_,int(yc)+yb_
                    z_keep,x_keep,y_keep = in_dim(z_keep,x_keep,y_keep,sz,sx,sy)

                    X_ = np.array([z_keep,x_keep,y_keep]).T
                    im_ = im[z_keep,x_keep,y_keep]

            if not better_fit:
                reconstruct,plt_val,compare_with_fitting=False,False,False
                if troubleshoot:
                    reconstruct,plt_val,compare_with_fitting=True,True,True
                p_ = gfit_fast(im_,X_.T,bk_f=0.1,
                                reconstruct=reconstruct,plt_val=plt_val,compare_with_fitting=compare_with_fitting)
            else:
                p_ = np.array([np.nan]*11)
                if len(im_)>0:
                    center = X_[np.argmax(im_)]
                    obj = GaussianFit(im_,X_.T,center=center,delta_center=2.5)
                    obj.fit()
                    p_  = obj.p
            ps.append(p_)
    ps = np.array(ps)
    return ps
class iter_fit_seed_points():
    def __init__(self,im,centers,radius_fit=5,min_delta_center=1.,max_delta_center=2.5,
                 n_max_iter = 10,max_dist_th=0.1,
                 min_w=0.5, max_w=4, init_w=1.5):
        """
        Given a set of seeds <centers> in a 3d image <im> iteratively 3d gaussian fit around the seeds (in order of brightness)
        and subtract the gaussian signal.
        Retruns a numpy array of size Nx(height, x, y, z, width_x, width_y,width_z,background) where N~len(centers). Bad fits are disregarded.
        Warning: Generally a bit slow. In practice, the faster version fast_local_fit is used.
        """

        #internalize
        self.im = im
        self.radius_fit = radius_fit
        self.n_max_iter = n_max_iter
        self.max_dist_th = max_dist_th
        self.min_delta_center = min_delta_center
        self.max_delta_center = max_delta_center

        self.centers = centers.T
        self.z,self.x,self.y = centers
        self.zb,self.xb,self.yb = np.reshape(np.indices([self.radius_fit*2]*3)-self.radius_fit,[3,-1])
        keep = self.zb*self.zb+self.xb*self.xb+self.yb*self.yb<=self.radius_fit**2
        self.zb,self.xb,self.yb = self.zb[keep],self.xb[keep],self.yb[keep]
        self.zxyb = np.array([self.zb,self.xb,self.yb]).T
        self.sz,self.sx,self.sy = im.shape

        self.min_w = min_w
        self.max_w = max_w
        self.init_w = init_w

    def firstfit(self):
        """
        Perform a first fit on the sample with the gaussian constrained close to the local maximum
        """
        if len(self.centers)>0:
            #fit the points in order of brightness and at each fit subtract the fitted signal
            self.ps = []
            self.ims_rec=[]
            self.im_subtr = np.array(self.im,dtype=float)
            self.centers_fit = []
            self.success=[]
            self.centers_tree = KDTree(self.centers)
            from tqdm import tqdm_notebook as tqdm
            #centers_ = tqdm(self.centers)
            centers_ = self.centers
            self.gparms = []
            for ic,(zc,xc,yc) in enumerate(centers_):
                z_keep,x_keep,y_keep = int(zc)+self.zb,int(xc)+self.xb,int(yc)+self.yb
                z_keep,x_keep,y_keep = in_dim(z_keep,x_keep,y_keep,self.sz,self.sx,self.sy)

                X_full = np.array([z_keep,x_keep,y_keep],dtype=int)
                center = [zc,xc,yc]
                z_keep,x_keep,y_keep = closest_faster(X_full.T,ic,self.centers_tree,rsearch = self.radius_fit*2)

                X = np.array([z_keep,x_keep,y_keep])
                im_ = self.im[z_keep,x_keep,y_keep]

                self.gparms

                obj = GaussianFit(im_,X,center=center,delta_center=self.min_delta_center,
                                  min_w=self.min_w, max_w=self.max_w, init_w=self.init_w)
                obj.fit()

                self.gparms.append([im_,X,center])
                n_p = len(obj.p)
                self.success.append(obj.success)
                if obj.success:
                    self.ps.append(obj.p)
                    self.centers_fit.append(obj.center)
                    z_keep,x_keep,y_keep = X_full
                    obj.x,obj.y,obj.z = X_full
                    im_rec = obj.get_im()
                    self.ims_rec.append(im_rec)
                    self.im_subtr[z_keep,x_keep,y_keep] -= im_rec
                else:
                    self.ims_rec.append(np.nan)
                    self.ps.append([np.nan]*n_p)
                    self.centers_fit.append([np.nan]*3)

        self.im_add = np.array(self.im_subtr)

    def repeatfit(self):
            self.n_iter = 0

            self.converged = np.zeros(len(self.centers),dtype=bool)
            self.dists = np.zeros(len(self.centers))+np.inf
            converged = np.all(self.converged)
            while not converged:

                self.success_old,self.centers_fit_old=np.array(self.success),np.array(self.centers_fit)
                #self.ps_old=np.array(self.ps)
                for ic,(zc,xc,yc) in enumerate(self.centers):
                    if not self.converged[ic]:
                        #get modified image positions
                        z_keep,x_keep,y_keep = int(zc)+self.zb,int(xc)+self.xb,int(yc)+self.yb
                        z_keep,x_keep,y_keep = in_dim(z_keep,x_keep,y_keep,self.sz,self.sx,self.sy)
                        X = np.array([z_keep,x_keep,y_keep])
                        #get image, adding back the fit
                        im_ = self.im_add[z_keep,x_keep,y_keep]

                        if self.success_old[ic]:
                            im_rec = self.ims_rec[ic]
                            im_=im_rec+im_ #add back the image

                        delta_center = self.max_delta_center
                        obj = GaussianFit(im_,X,center=[zc,xc,yc],delta_center=delta_center,
                                          min_w=self.min_w, max_w=self.max_w, init_w=self.init_w)
                        obj.fit()
                        self.success[ic] = obj.success
                        if obj.success:
                            im_rec = obj.get_im()
                            self.ps[ic]=obj.p
                            self.centers_fit[ic]=obj.center
                            self.ims_rec[ic]=im_rec
                            self.im_add[z_keep,x_keep,y_keep] = im_-im_rec

                keep = (np.array(self.success)&np.array(self.success_old))>0
                self.dists[~keep]=0
                self.dists[keep]=np.sum((np.array(self.centers_fit_old)[keep]-np.array(self.centers_fit)[keep])**2,axis=-1)
                self.converged = self.dists<self.max_dist_th**2
                converged = np.all(self.converged)
                self.n_iter+=1
                converged = converged or (self.n_iter>self.n_max_iter)

def to_sigmas(t,p,w_1,w_2,w_3):
    """Given t = sin(theta), p = sin(phi), and the stds this computes the covariance matrix and its inverse"""
    p2 = p*p
    t2 = t*t
    tc2 = 1-t2
    pc2 = 1-p2
    tc= np.sqrt(tc2)
    pc= np.sqrt(pc2)
    s1,s2,s3 = 1./(w_1*w_1),1./(w_2*w_2),1./(w_3*w_3)
    a = pc2*tc2*s1 + t2*s2 + p2*tc2*s3
    b = pc2*t2*s1 + tc2*s2 + p2*t2*s3
    c = p2*s1 + pc2*s3
    d = tc*t*(pc2*s1 - s2 + p2*s3)
    e = p*pc*tc*(s3 - s1)
    f = p*pc*t*(s3 - s1)
    sigma_inv = np.array([[a, d, e], [d, b, f], [e, f, c]])
    sigma = np.array([[(b*c - f ** 2)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (-(c*d) + e*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (-(b*e) + d*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2)],
                      [(-(c*d) + e*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (a*c - e ** 2)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (d*e - a*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2)],
                      [(-(b*e) + d*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (d*e - a*f)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2), (a*b - d ** 2)/(a*b*c - c*d ** 2 - b*e ** 2 + 2*d*e*f - a*f ** 2)]])
    return sigma,sigma_inv
def to_sigmas_abc(a,b,c,d,e,f):
    sigma = np.array([[a, d, e], [d, b, f], [e, f, c]])
    sigma_inv = np.array([[b*c - f*f,-(c*d) + e*f,-(b*e) + d*f],[-(c*d) + e*f,a*c - e*e,d*e - a*f],[-(b*e) + d*f,d*e - a*f,a*b - d*d]])
    det_sigma = a*b*c - c*d**2 - b*e**2 + 2*d*e*f - a*f**2
    sigma_inv/=det_sigma
    return sigma,sigma_inv,det_sigma
def KL_dif(p1,p2):
    """Given two gaussian paramters p1 and p2 this computes their KL divergence."""
    if len(p1)<11:
        return 0
    elif len(p1)==11:
        [height1,center_01,center_11,center_21,background1,widths_01,widths_11,widths_21,t_1,p_1,e1]=p1
        [height2,center_02,center_12,center_22,background2,widths_02,widths_12,widths_22,t_2,p_2,e2]=p2
        sigma1,sigma_inv1 = to_sigmas(*[t_1,p_1,widths_01,widths_11,widths_21])
        sigma2,sigma_inv2 = to_sigmas(*[t_2,p_2,widths_02,widths_12,widths_22])
        k = len(sigma1)
        inds = np.arange(k)
        return 0.5*(np.sum(np.dot(sigma_inv2,sigma1)[inds,inds])-k+np.log(np.linalg.det(sigma2))-np.log(np.linalg.det(sigma1)))
    else:
        sigma1,sigma_inv1,det1 =to_sigmas_abc(*p1[5:11])
        sigma2,sigma_inv2,det2 =to_sigmas_abc(*p2[5:11])
        k = len(sigma1)
        inds = np.arange(k)
        #
        #return 0.5*(np.sum(np.dot(sigma_inv2,sigma1)[inds,inds])-k+np.log(det2)-np.log(det1))
        return 0
##2d images
import cv2
def blurnorm2d(im,gb):
    """Normalize an input 2d image <im> by dividing by a cv2 gaussian filter of the image"""
    im_ = im.astype(np.float32)
    blurred = cv2.blur(im_,(gb,gb))
    return im_-blurred
def fft3d_from2d(im1,im2,gb=5,max_disp=150,plt_val=False,return_cor=False):
    """Given a refence 3d image <im1> and a target image <im2>
    this max-projects along the first (z) axis and finds the best tx,ty using fftalign_2d.
    Then it trims and max-projects along the last (y) axis and finds tz.
    Before applying fftalignment we normalize the images using blurnorm2d for stability."""
    im1_ = blurnorm2d(np.max(im1,0),gb)
    im2_ = blurnorm2d(np.max(im2,0),gb)
    tx,ty,cor_xy = fftalign_2d(im1_, im2_, center=[0, 0], max_disp=max_disp, plt_val=plt_val,return_cor=True)
    sx,sy = im1_.shape
    im1_t = blurnorm2d(np.max(im1[:,max(tx,0):sx+tx,max(ty,0):sy+ty],axis=-1),gb)
    im2_t = blurnorm2d(np.max(im2[:,max(-tx,0):sx-tx,max(-ty,0):sy-ty],axis=-1),gb)
    tz,_,cor_z = fftalign_2d(im1_t, im2_t, center=[0, 0], max_disp=max_disp, plt_val=plt_val,return_cor=True)
    if return_cor:
        return np.array([tz,tx,ty]),cor_xy,cor_z
    return np.array([tz,tx,ty])
import matplotlib.pyplot as plt
def fftalign_2d(im1,im2,center=[0,0],max_disp=50,plt_val=False,return_cor=False):
    """
    Inputs: 2 2D images <im1>, <im2>, the expected displacement <center>, the maximum displacement <max_disp> around the expected vector.
    This computes the cross-cor between im1 and im2 using fftconvolve (fast) and determines the maximum
    """
    from scipy.signal import fftconvolve
    im2_=np.array(im2[::-1,::-1],dtype=float)
    im2_-=np.mean(im2_)
    im2_/=np.std(im2_)
    im1_=np.array(im1,dtype=float)
    im1_-=np.mean(im1_)
    im1_/=np.std(im1_)
    im_cor = fftconvolve(im1_,im2_, mode='full')

    sx_cor,sy_cor = im_cor.shape
    center_ = np.array(center)+np.array([sx_cor,sy_cor])/2.

    x_min = int(min(max(center_[0]-max_disp,0),sx_cor))
    x_max = int(min(max(center_[0]+max_disp,0),sx_cor))
    y_min = int(min(max(center_[1]-max_disp,0),sy_cor))
    y_max = int(min(max(center_[1]+max_disp,0),sy_cor))

    im_cor0=np.zeros_like(im_cor)
    im_cor0[x_min:x_max,y_min:y_max]=1
    im_cor = im_cor*im_cor0

    y, x = np.unravel_index(np.argmax(im_cor), im_cor.shape)
    xt,yt=np.array([y,x],dtype=int)-im2_.shape+1
    if np.sum(im_cor>0)>0:
        im_cor[im_cor==0]=np.min(im_cor[im_cor>0])
    else:
        im_cor[im_cor==0]=0
    if plt_val:
        plt.figure()
        plt.plot([x],[y],'k+')
        plt.imshow(im_cor,interpolation='nearest')
        plt.show()
        std_2 = np.std(im2)
        std_1 = np.std(im1)
        im2t = translate(minmax(im2,min_=-2*std_2,max_=std_2*2),[-xt,-yt])
        im1t = minmax(im1,min_=-2*std_1,max_=std_1*2)
        #print im1.shape,im2.shape,im1t.shape,im2t.shape
        #print im_cor.shape,y, x
        szx,szy = np.min([im2t.shape,im1t.shape],axis=0)
        plt.figure()
        plt.imshow(im1t[:szx,:szy])
        plt.figure()
        plt.imshow(im2t[:szx,:szy])
        plt.figure()
        plt.imshow(np.dstack([im2t[:szx,:szy],im1t[:szx,:szy],im2t[:szx,:szy]]))
        plt.show()
    if return_cor:
        return xt,yt,np.max(im_cor)/np.prod(np.min([im1_.shape,im2_.shape],0))
    return xt,yt
def minmax(im,min_=None,max_=None):
    if min_ is None: min_=np.min(im)
    if max_ is None: max_=np.max(im)
    return (np.array(im,dtype=np.float32)-min_)/(max_-min_)
def translate(im,trans):
    """translates ND images"""
    shape_ = im.shape
    trans_=np.array(np.round(trans),dtype=int)
    im_slices = tuple([slice(max(t_,0),min(sh_,sh_+t_)) for t_,sh_ in zip(trans_,shape_)])
    im_slices0 = tuple([slice(max(-t_,0),min(sh_,sh_-t_)) for t_,sh_ in zip(trans_,shape_)])
    im_base_0 = np.zeros_like(im)+np.median(im)
    im_base_0[im_slices0]=im[im_slices]
    return im_base_0
from sklearn import linear_model
def calc_color_matrix(x,y,order=2,ransac=False):
    """This gives a quadratic color transformation (in matrix form)
    x is Nx3 vector of positions in the reference channel (typically cy5)
    y is the Nx3 vector of positions in another channel (i.e. cy7)
    return m_ a 3x7 matrix which when multipled with x,x**2,1 returns y-x
    This m_ is indended to be used with apply_colorcor
    """
    x_ = np.array(y)# ref zxy
    y_ = np.array(x)-x_# dif zxy
    # get a list of exponents
    exps = []
    for p in range(order+1):
        for i in range(p+1):
            for j in range(p+1):
                if i+j<=p:
                    exps.append([i,j,p-i-j])
    # construct A matrix
    A = np.zeros([len(x_),len(exps)])
    for iA,(ix,iy,iz) in enumerate(exps):
        s = (x_[:,0]**ix*x_[:,1]**iy*x_[:,2]**iz)
        A[:,iA]=s
    if ransac:
        ransac = linear_model.RANSACRegressor()
        ransac.base_estimator = linear_model.LinearRegression(fit_intercept=False)
        m_ = []
        for iy in range(len(x_[0])):
            ransac.fit(A, y_[:,iy])
            m_.append(ransac.estimator_.coef_)
        m_=np.array(m_)
    else:
        m_ = [np.linalg.lstsq(A, y_[:,iy])[0] for iy in range(len(x_[0]))]
    m_=np.array(m_)
    return m_
def apply_colorcor(x,m=None):
    """This applies chromatic abberation correction to order 2
    x is a Nx3 vector of positions (typically 750(-->647))
    m is a matrix computed by function calc_color_matrix
    y is the corrected vector in another channel"""
    if m is None:
        return x
    exps = []
    order_max=10
    for p in range(order_max+1):
        for i in range(p+1):
            for j in range(p+1):
                if i+j<=p:
                    exps.append([i,j,p-i-j])
    #find the order
    mx,my = m.shape
    order = int((my-1)/mx)
    assert(my<len(exps))
    x_ = np.array(x)
    # construct A matrix
    exps = exps[:my]
    A = np.zeros([len(x_),len(exps)])
    for iA,(ix,iy,iz) in enumerate(exps):
        s = (x_[:,0]**ix*x_[:,1]**iy*x_[:,2]**iz)
        A[:,iA]=s
    diff = [np.dot(A,m_) for m_ in m]
    return x_+np.array(diff).T

reason: import numpy,from scipy


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

from __future__ import print_function

import unittest
import numpy as np
from op_test import OpTest
import paddle.fluid as fluid
from paddle.fluid import compiler, Program, program_guard, core

class TestSplitOp(OpTest):
    def setUp(self):
        self._set_op_type()
        self.dtype = self.get_dtype()
        axis = 1
        x = np.random.random((4, 5, 6)).astype(self.dtype)
        out = np.split(x, [2, 3], axis)
        self.inputs = {'X': x}
        self.attrs = {'axis': axis, 'sections': [2, 1, 2]}
        self.outputs = {'Out': [('out%d' % i, out[i]) \
            for i in range(len(out))]}

    def get_dtype(self):
        return "float64"

    def _set_op_type(self):
        self.op_type = "split"

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], ['out0', 'out1', 'out2'])

# test with attr(num)
class TestSplitOp_2(OpTest):
    def setUp(self):
        self._set_op_type()
        self.dtype = self.get_dtype()
        self.init_data()
        self.inputs = {'X': self.x}
        self.attrs = {
            'axis': self.axis,
            'sections': self.sections,
            'num': self.num
        }

        out = np.split(self.x, self.indices_or_sections, self.axis)
        self.outputs = {'Out': [('out%d' % i, out[i]) \
                                for i in range(len(out))]}

    def init_data(self):
        self.x = np.random.random((4, 5, 6)).astype(self.dtype)
        self.axis = 2
        self.sections = []
        self.num = 3
        self.indices_or_sections = 3

    def get_dtype(self):
        return "float64"

    def _set_op_type(self):
        self.op_type = "split"

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], ['out0', 'out1', 'out2'])

# attr(axis) is Tensor
class TestSplitOp_AxisTensor(OpTest):
    def setUp(self):
        self._set_op_type()
        self.dtype = self.get_dtype()
        self.init_data()
        self.inputs = {
            'X': self.x,
            'AxisTensor': np.array([self.axis]).astype("int32")
        }
        self.attrs = {'sections': self.sections, 'num': self.num}

        out = np.split(self.x, self.indices_or_sections, self.axis)
        self.outputs = {'Out': [('out%d' % i, out[i]) \
                                for i in range(len(out))]}

    def init_data(self):
        self.x = np.random.random((4, 5, 6)).astype(self.dtype)
        self.axis = 2
        self.sections = []
        self.num = 3
        self.indices_or_sections = 3

    def get_dtype(self):
        return "float64"

    def _set_op_type(self):
        self.op_type = "split"

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], ['out0', 'out1', 'out2'])

# attr(sections) is list containing Tensor
class TestSplitOp_SectionsTensor(OpTest):
    def setUp(self):
        self._set_op_type()
        self.dtype = self.get_dtype()
        self.init_data()
        self.inputs = {'X': self.x}

        sections_tensor = []
        for index, ele in enumerate(self.sections):
            sections_tensor.append(("x" + str(index), np.ones(
                (1)).astype('int32') * ele))

        self.inputs['SectionsTensorList'] = sections_tensor

        self.attrs = {
            'axis': self.axis,
            'sections': self.sections_infer,
            'num': self.num
        }

        out = np.split(self.x, self.indices_or_sections, self.axis)
        self.outputs = {'Out': [('out%d' % i, out[i]) \
                                for i in range(len(out))]}

    def init_data(self):
        self.x = np.random.random((4, 5, 6)).astype(self.dtype)
        self.axis = 1
        self.sections = [2, 1, 2]
        self.sections_infer = [-1, -1, -1]
        self.num = 0
        self.indices_or_sections = [2, 3]

    def get_dtype(self):
        return "float64"

    def _set_op_type(self):
        self.op_type = "split"

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], ['out0', 'out1', 'out2'])

class TestSplitOp_unk_section(OpTest):
    def setUp(self):
        self._set_op_type()
        self.dtype = self.get_dtype()
        self.init_data()
        self.inputs = {'X': self.x}
        self.attrs = {
            'axis': self.axis,
            'sections': self.sections,
            'num': self.num
        }

        out = np.split(self.x, self.indices_or_sections, self.axis)
        self.outputs = {'Out': [('out%d' % i, out[i]) \
                                for i in range(len(out))]}

    def init_data(self):
        self.x = np.random.random((4, 5, 6)).astype(self.dtype)
        self.axis = 2
        self.sections = [2, 1, -1]
        self.num = 0
        self.indices_or_sections = [2, 3]

    def get_dtype(self):
        return "float64"

    def _set_op_type(self):
        self.op_type = "split"

    def test_check_output(self):
        self.check_output()

    def test_check_grad(self):
        self.check_grad(['X'], ['out0', 'out1', 'out2'])

class TestSplitByrefOp(OpTest):
    def _set_op_type(self):
        self.op_type = "split_byref"

#----------------Split Fp16----------------

def create_test_fp16(parent):
    @unittest.skipIf(not core.is_compiled_with_cuda(),
                     "core is not compiled with CUDA")
    class TestSplitFp16(parent):
        def get_dtype(self):
            return np.float16

        def test_check_grad(self):
            pass

    cls_name = "{0}_{1}".format(parent.__name__, "Fp16")
    TestSplitFp16.__name__ = cls_name
    globals()[cls_name] = TestSplitFp16

create_test_fp16(TestSplitOp)

class TestSplitAPI(unittest.TestCase):
    def test_api(self):
        input_1 = np.random.random([4, 5, 6]).astype("int32")
        positive_1_int32 = fluid.layers.fill_constant([1], "int32", 1)
        positive_1_int64 = fluid.layers.fill_constant([1], "int64", 1)
        positive_2_int64 = fluid.layers.fill_constant([1], "int64", 2)
        x_1 = fluid.data(shape=[4, 5, 6], dtype='int32', name='x_1')
        x_2 = fluid.data(shape=[4, 5, None], dtype='int32', name='x_2')

        out_0, out_1, out_2 = fluid.layers.split(
            input=x_1,
            num_or_sections=[positive_2_int64, positive_1_int32, -1],
            dim=positive_1_int64)

        out_3, out_4, out_5 = fluid.layers.split(
            input=x_1, num_or_sections=[2, 1, 2], dim=positive_1_int32)
        fluid.layers.split(input=x_2, num_or_sections=2, dim=2)

        exe = fluid.Executor(place=fluid.CPUPlace())
        [res_0, res_1, res_2, res_3, res_4, res_5] = exe.run(
            fluid.default_main_program(),
            feed={"x_1": input_1,
                  "x_2": input_1},
            fetch_list=[out_0, out_1, out_2, out_3, out_4, out_5])

        out = np.split(input_1, [2, 3], 1)
        assert np.array_equal(res_0, out[0])
        assert np.array_equal(res_1, out[1])
        assert np.array_equal(res_2, out[2])
        assert np.array_equal(res_3, out[0])
        assert np.array_equal(res_4, out[1])
        assert np.array_equal(res_5, out[2])

class TestSplitOpError(unittest.TestCase):
    def test_errors(self):
        with program_guard(Program(), Program()):
            # The type of axis in split_op should be int or Variable.
            def test_axis_type():
                x6 = fluid.layers.data(shape=[4], dtype='float16', name='x3')
                fluid.layers.split(input=x6, num_or_sections=2, dim=3.2)

            self.assertRaises(TypeError, test_axis_type)

            # The type of num_or_sections in split_op should be int, tuple or list.
            def test_num_or_sections_type():
                x6 = fluid.layers.data(shape=[4], dtype='float16', name='x4')
                fluid.layers.split(input=x6, num_or_sections=2.1, dim=3)

            self.assertRaises(TypeError, test_num_or_sections_type)

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


text:
import dash
import numpy as np
import plotly.graph_objs as go
import pandas as pd
import dash_html_components as html
import dash_core_components as dcc
from dash.dependencies import Input, Output
import os

class DaysUntilVacDataset:
    @staticmethod
    def load():
        dir_path = os.path.dirname(os.path.realpath(__file__)) + "/"
        return pd.read_csv(dir_path + "tempo-desde-17-por-uf-categorizado-agrupado.csv")

    @staticmethod
    def filter_by_uf(df, uf):
        if uf != "Todos":
            df = df.query(f"paciente_endereco_uf==\"{uf}\"")
        df = df.drop("paciente_endereco_uf", axis=1) \
            .sort_values(["Dias até vacinar", "count"]) \
            .groupby(["categoria", "Dias até vacinar"]) \
            .sum() \
            .reset_index()
        df["Dias até vacinar"] = df["Dias até vacinar"].astype(int)
        return df

    @staticmethod
    def sort_by_mean(df):
        mean_df = df[["categoria", "count", "Dias até vacinar"]]
        mean_df["Dias até vacinar"] = mean_df["count"] * mean_df["Dias até vacinar"].astype(int)
        mean_df = mean_df.groupby("categoria").sum()
        mean_df["Média dias"] = mean_df["Dias até vacinar"] / mean_df["count"]
        mean_df = mean_df.drop(["count", "Dias até vacinar"], axis=1)
        return df.join(mean_df, on="categoria").sort_values(["Média dias", "Dias até vacinar"]).drop(["Média dias"],
                                                                                axis=1)

    @staticmethod
    def calculate_whiskers(q1, q3, min_, max_):
        iqr = q3 - q1
        lowerfence = max(q1 - iqr, min_)
        upperfence = min(q3 + iqr, max_)
        return lowerfence, upperfence

    @staticmethod
    def detect_quantiles(df):
        df = df.set_index("categoria")
        df["quantile1_pos"] = df.reset_index()[["categoria", "count"]].groupby("categoria").sum()["count"].apply(
            lambda c: c // 4)
        df["median_pos"] = df.reset_index()[["categoria", "count"]].groupby("categoria").sum()["count"].apply(
            lambda c: c // 2)
        df["quantile3_pos"] = df.reset_index()[["categoria", "count"]].groupby("categoria").sum()["count"].apply(
            lambda c: (3 * c) // 4)
        summary = {}
        for index in df.index.unique():
            category_df = df.loc[index]
            quantile1_pos = category_df.iloc[0]["quantile1_pos"]
            median_pos = category_df.iloc[0]["median_pos"]
            quantile3_pos = category_df.iloc[0]["quantile3_pos"]
            count_series = category_df["count"]
            acu = 0
            targets_pos = [quantile1_pos, median_pos, quantile3_pos]
            targets_values = []
            for i, c in enumerate(count_series):
                acu += c
                if len(targets_pos) == 0:
                    break
                if acu > targets_pos[0]:
                    targets_pos.pop(0)
                    value = category_df.iloc[i]["Dias até vacinar"]
                    targets_values.append(value)
            q1, median, q3 = targets_values
            min_, max_ = category_df.iloc[0]["Dias até vacinar"], category_df.iloc[-1]["Dias até vacinar"]
            lowerfence, upperfence = DaysUntilVacDataset.calculate_whiskers(q1, q3, min_, max_)
            targets_values.insert(0, lowerfence)
            targets_values.append(upperfence)
            summary[index] = targets_values
        return summary

class DaysUntilVac:
    fig: go.Figure
    df: pd.DataFrame

    def show(self):
        self.fig.show()

    def __init__(self, df=DaysUntilVacDataset.load(), uf="Todos"):
        self.df = df
        self.sorted_df = DaysUntilVacDataset.sort_by_mean(DaysUntilVacDataset.filter_by_uf(self.df, uf))
        self.summary = np.array(list(DaysUntilVacDataset.detect_quantiles(self.sorted_df).values()))
        self.generate_figure()

    def generate_figure(self):
        self.fig = go.Figure()
        for i, cat in zip(self.summary, self.sorted_df.set_index("categoria").index.unique()):
            self.fig.add_trace(go.Box(x=[""], name=cat, lowerfence=[i[0]], q1=[i[1]], median=[i[2]],
                                      q3=[i[3]], upperfence=[i[4]]))
        self.fig.update_layout(height=650, width=1300, yaxis_title="Dias até vacinar", boxmode="group", font=dict(
            size=16))
        return self.fig

class DaysUntilVacDash(DaysUntilVac):
    def __init__(self, app: dash.Dash):
        super().__init__()
        uf_label = "paciente_endereco_uf"
        options = [{'label': i, 'value': i} for i in self.df.query(f'{uf_label}!="XX"')[uf_label].unique().tolist()]
        options.append({'label': "Todos", 'value': "Todos"})

        self.component_html = html.Div([
            html.Hr(),
            html.H4('População vacinada e imunizada por sexo e idade'),
            html.Label("Selecione o estado"),
            html.Div([
                dcc.Dropdown(
                    id='uf',
                    options=options,
                    value='Todos',
                )
            ],
                style={'width': '48%', 'display': 'inline-block'}),
            dcc.Graph(id='days-util-vac')
        ])
        app.callback(Output('days-util-vac', 'figure'), Input('uf', 'value'))(self.callback)

    def callback(self, uf: str):
        self.sorted_df = DaysUntilVacDataset.sort_by_mean(DaysUntilVacDataset.filter_by_uf(self.df, uf))
        self.summary = np.array(list(DaysUntilVacDataset.detect_quantiles(self.sorted_df).values()))
        return self.generate_figure()