text:
import os
import re
import time
import glob
import datetime
import random
import requests
import json
import pymysql
import xml.etree.ElementTree as et
import pandas as pd
import numpy as np
import socket
socket.setdefaulttimeout(8)

host = "localhost"
user = "root"
passwd = "root"
db = "quotation"
port = 3306

def _get_conn():
    return pymysql.connect(
        host,
        user,
        passwd,
        db,
        port = port)

def _load_trading_day():
    trading_day = pd.read_csv("trading_day.csv")
    trading_day['dt'] = pd.to_datetime(trading_day['dt']).apply(lambda x: x.date())
    return trading_day

def _create_tb(conn):
    sql = """create table if not exists future_day_web (
    exchange char(3),
    code char(2),
    contract char(6),
    date date,
    open float,
    high float,
    low float,
    close float,
    settle float,
    presettle float,
    vol int,
    opn int,
    turnover float,
    primary key(contract, date))
    """
    cur = conn.cursor()
    cur.execute(sql)
    conn.commit()
    cur.close()

def download(day):
    def _requests(address, data=None):
        TRYTIME = 5
        trytime = 0
        req = None
        while True:
            try:
                trytime += 1
                if data:
                    req = requests.post(address, data)
                else:
                    req = requests.get(address)
                break
            except:
                time.sleep(1)
                if trytime > TRYTIME:
                    raise Exception("request get trytime > 5 with address: %s" % address)
        if req is None or req.status_code != 200 or req.url.split("/")[-1] == "error_404.html":
            raise Exception("help_requests failed with address ", address)
        return req

    def _download(filename, address, post_data=None, startd=datetime.date(2010,1,1)):
        if not os.path.exists(filename) and day >= startd:
            req = _requests(address, post_data)
            with open(filename, 'w') as outfile:
                outfile.write(req.text)

    # cfe
    cfe_filename = "origin/cfe/cfe_%s.csv" % format(day, "%Y%m%d")
    cfe_address = "http://www.cffex.com.cn/sj/hqsj/rtj/%s/index.xml" % format(day, "%Y%m/%d")
    _download(cfe_filename, cfe_address, startd=datetime.date(2010,4,16))

    # shf
    shf_filename = "origin/shf/shf_%s.csv" % format(day, "%Y%m%d")
    shf_address = "http://www.shfe.com.cn/data/dailydata/kx/kx%s.dat" % format(day, "%Y%m%d")
    _download(shf_filename, shf_address)

    # dce
    dce_filename = "origin/dce/dce_%s.csv" % format(day, "%Y%m%d")
    dce_address = "http://www.dce.com.cn/publicweb/quotesdata/exportDayQuotesChData.html"
    dce_post_dat = {
        'dayQuotes.trade_type': 0,
        'dayQuotes.variety': 'all',
        'exportFlag': 'txt',
        'year': int(format(day, "%Y")),
        'month': int(format(day, "%m")) - 1,
        'day': int(format(day, "%d"))}
    _download(dce_filename, dce_address, dce_post_dat)

    # czc
    czc_filename = "origin/czc/czc_%s.csv" %format(day, "%Y%m%d")
    czc_address = "http://www.czce.com.cn/cn/DFSStaticFiles/Future/%s/%s/FutureDataDaily.txt" % (
            format(day, "%Y"), format(day, "%Y%m%d"))
    _download(czc_filename, czc_address, startd=datetime.date(2018,1,1))

def standard(filename):
    """
    标准化：交易所，品种，合约，ohlc，结算，昨结算，成交量，持仓量，成交金额
    """

    def _standard_record(record):
        for i in range(len(record)):
            for j in range(4, len(record[i])):
                if record[i][j] == 'null' or record[i][j] == '':
                    record[i][j] = None
                if isinstance(record[i][j], str):
                    record[i][j] = float(record[i][j].replace(',', ''))
            for k in range(2, 6):
                if record[i][k] == 0:
                    record[i][k] = None
        return record

    def _cfe_parse(text, day):
        root = et.fromstring(text.lower())
        record = [
            ["cfe",
             c.find("productid").text.strip(),
             c.find("instrumentid").text.strip(),
             c.find("tradingday").text.strip(),
             c.find("openprice").text,
             c.find("highestprice").text,
             c.find("lowestprice").text,
             c.find("closeprice").text,
             c.find("settlementprice").text,
             c.find("presettlementprice").text,
             c.find("volume").text,
             c.find("openinterest").text,
             c.find("turnover").text
             ] for c in root if c.tag == "dailydata"
        ]
        record = _standard_record(record)
        return record

    def _shf_parse(text, day):
        data = json.loads(text.lower())
        pat = re.compile("^[0-9]{4}$")
        record = []
        for v in data['o_curinstrument']:
            if not pat.match(v['deliverymonth']):
                continue
            contract = v['productid'].strip().split("_")[0] + v['deliverymonth']
            record.append([
                'shf',
                v['productid'].strip().split("_")[0],
                contract,
                day,
                v['openprice'],
                v['highestprice'],
                v['lowestprice'],
                v['closeprice'],
                v['settlementprice'],
                v['presettlementprice'],
                v['volume'],
                v['openinterest'],
                None
            ])
        record = _standard_record(record)
        return record

    def _dce_parse(text, day):
        codes = {
            '豆一': 'a',
            '豆二': 'b',
            '豆粕': 'm',
            '豆油': 'y',
            '玉米': 'c',
            '玉米淀粉': 'cs',
            '鸡蛋': 'jd',
            '棕榈油': 'p',
            '聚乙烯': 'l',
            '聚氯乙烯': 'v',
            '胶合板': 'bb',
            '纤维板': 'fb',
            '铁矿石': 'i',
            '焦炭': 'j',
            '焦煤': 'jm',
            '聚丙烯': 'pp',
            '乙二醇': 'eg'
        }
        data = text.lower()
        record = []
        for v in [t.strip() for t in data.split("\n") if len(t.strip()) > 0]:
            v = [t.strip() for t in v.split("\t") if len(t.strip()) > 0]
            if v[0] != '商品名称' and v[0].find('小计') == -1 and v[0] != '总计':
                contract = codes[v[0]] + v[1]
                record.append([
                    'dce',
                    codes[v[0]],
                    contract,
                    day,
                    v[2],
                    v[3],
                    v[4],
                    v[5],
                    v[7],
                    v[6],
                    v[10],
                    v[11],
                    v[13]
                ])
        record = _standard_record(record)
        return record

    def _czc_parse(text, day):
        contract_pat = re.compile("^[a-z]{1,2}[0-9]{3}$")
        data = [v.strip() for v in text.lower().split("\n") if len(v.strip()) > 0]
        record = []
        for i in range(1, len(data)):
            v = [t.strip() for t in data[i].split("|")]
            contract = v[0].lower()
            if contract_pat.match(contract):
                contract = contract[:-3] + '1' + contract[-3:]
                record.append([
                    'czc',
                    contract[:-4],
                    contract,
                    day,
                    v[2],
                    v[3],
                    v[4],
                    v[5],
                    v[6],
                    v[1],
                    v[9],
                    v[10],
                    float(v[12].replace(",", '')) * 10000
                ])
        record = _standard_record(record)
        return record

    def _czc_history_parse(text, day):
        contract_pat = re.compile("^[a-z]{1,2}[0-9]{3}$")
        data = [v.strip() for v in text.lower().split("\n") if len(v.strip()) > 0]
        record = []
        for i in range(1, len(data)):
            v = [t.strip() for t in data[i].split("|")]
            contract = v[1].lower()
            if contract_pat.match(contract):
                contract = contract[:-3] + '1' + contract[-3:]
                record.append([
                    'czc',
                    contract[:-4],
                    contract,
                    v[0],
                    v[3],
                    v[4],
                    v[5],
                    v[6],
                    v[7],
                    v[2],
                    v[10],
                    v[11],
                    float(v[13].replace(",", '')) * 10000
                ])
        record = _standard_record(record)
        return record

    _f = os.path.split(filename)[-1]
    with open(filename, 'r') as infile:
        text = infile.read()

    day = None
    if _f.find("_") == -1:
        func = _czc_history_parse
    else:
        exchange, day = _f[:-4].split("_")
        func = eval("_%s_parse" % exchange)

    record = func(text, day)
    return record

def storage(record):
    conn = _get_conn()
    cur = conn.cursor()
    for v in record:
        v[0] = "'%s'"%v[0]
        v[1] = "'%s'"%v[1]
        v[2] = "'%s'"%v[2]
        v[3] = "'%s'"%v[3]
        sql = "insert into future_day_web value (%s);" %(
            ','.join([str(item) for item in v]))
        sql = sql.replace("None", "null")
        try:
            cur.execute(sql)
        except pymysql.err.IntegrityError:
            return
        except:
            print(sql)
            raise("error")
    conn.commit()
    cur.close()
    conn.close()

def check():
    conn = _get_conn()

    def _check_day():
        # 日期完整性检验
        sql = "select exchange, date, count(*) from future_day_web group by exchange, date order by date"
        data = pd.read_sql(sql, conn)
        trading_day = _load_trading_day()
        for v in trading_day['dt']:
            if v <= data['date'].max():
                if v < datetime.date(2010,4,16):
                    assert len(data[data['date'] == v]) == 3
                else:
                    assert len(data[data['date'] == v]) == 4

    _check_day()
    print("日期完整性检验通过")

    try:
        from WindPy import w
        w.start()
    except ModuleNotFoundError:
        return

    def _wsd(contract, stroption, startd, endd):
        data = w.wsd(contract, stroption, startd, endd)
        data = {'Fields': data.Fields, 'Times': data.Times, 'Data': data.Data}
        names = ["contract", "date"] + [v.lower() for v in data['Fields']]
        record = []
        for i in range(len(data['Times'])):
            record.append([contract, data['Times'][i]] +
                          [data['Data'][j][i] for j in range(len(data['Data']))])
        return pd.DataFrame(record, columns=names)

    def _check_contract(contract):
        print("\tcheck", contract, end=' ')
        sql = """select exchange,contract,date,open,high,low,close,settle,
        presettle as pre_settle, vol as volume, opn as oi
        from future_day_web where contract='%s'""" % contract
        data = pd.read_sql(sql, conn)

        startd = data['date'].iloc[0]
        endd = data['date'].iloc[-1]
        ncontract = contract+'.' + data['exchange'][0]

        ndata = _wsd(ncontract, "open,high,low,close,settle,pre_settle,volume,oi", startd, endd)
        if data['exchange'][0]=='czc' and ndata['close'].isnull().all():
            ncontract = ncontract[:-8] + ncontract[-7:]
            ndata = _wsd(ncontract, "open,high,low,close,settle,pre_settle,volume,oi", startd, endd)

        if contract[:-4] == 'wr':
            # wr 2018年9月交易所网站没有数据
            ndata = ndata[ndata['volume'].notnull()]
            ndata.index = range(len(ndata))

        assert len(data) == len(ndata)

        for v in ['date', 'open', 'high', 'low', 'close', 'settle', 'volume', 'oi']:
            # pre_settle处理有区别
            data = data.fillna(0.0)
            ndata = ndata.fillna(0.0)
            assert (data[v] == ndata[v]).all() == True

        return True

    maxdate = pd.read_sql("select max(date) as d from future_day_web", conn)['d'][0]
    contracts = list(pd.read_sql("select contract from future_day_web where date='%s'" % maxdate, conn)['contract'])
    contracts = random.sample(contracts, 10)

    for v in contracts:
        print(_check_contract(v))
    print("抽样检查通过")

def main():
    # 确定startd和endd
    startd = datetime.date(2010, 1, 1)
    if int(time.strftime("%H", time.localtime(time.time()))) >= 17:
        endd = datetime.date.today()
    else:
        endd = datetime.date.today() - datetime.timedelta(days=1)

    # 下载原始数据
    trading_day = _load_trading_day()
    trading_day = trading_day[np.logical_and(startd <= trading_day['dt'], trading_day['dt'] <= endd)]
    for v in trading_day['dt'].sort_values(ascending=False):
        download(v)

    # 更新mysql数据库
    conn = _get_conn()
    tables = list(pd.read_sql("show tables", conn)['Tables_in_quotation'])
    if 'future_day_web' not in tables:
        # 全部更新
        _create_tb(conn)
        for f in glob.glob("origin/*/*.csv"):
            record = standard(f)
            storage(record)
    else:
        # 增量更新
        maxdate = pd.read_sql("select max(date) as d from future_day_web", conn)['d'][0]
        for f in glob.glob("origin/*/*.csv"):
            _f = os.path.split(f)[-1]
            if _f.find("_") == -1:
                continue
            else:
                exchange, date = _f[:-4].split("_")
                if date > format(maxdate, "%Y%m%d"):
                    record = standard(f)
                    storage(record)

    check()

if __name__ == "__main__":
    main()


text:
import numpy as np
from sklearn.model_selection import train_test_split
from sklearn.decomposition import PCA
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import MinMaxScaler

# # -*- coding: utf-8 -*-
# import click
# import logging
# from pathlib import Path
# from dotenv import find_dotenv, load_dotenv

# @click.command()
# @click.argument('input_filepath', type=click.Path(exists=True))
# @click.argument('output_filepath', type=click.Path())
# def main(input_filepath, output_filepath):
#     """ Runs data processing scripts to turn raw data from (../raw) into
#         cleaned data ready to be analyzed (saved in ../processed).
#     """
#     logger = logging.getLogger(__name__)
#     logger.info('making final data set from raw data')

# if __name__ == '__main__':
#     log_fmt = '%(asctime)s - %(name)s - %(levelname)s - %(message)s'
#     logging.basicConfig(level=logging.INFO, format=log_fmt)

#     # not used in this stub but often useful for finding various files
#     project_dir = Path(__file__).resolve().parents[2]

#     # find .env automagically by walking up directories until it's found, then
#     # load up the .env entries as environment variables
#     load_dotenv(find_dotenv())

#     main()

X_raw = np.load("../data/external/MNIST_X_28x28.npy")
X = X_raw.flatten().reshape(70000, 784)  # Flatten the data for convenience
Y = np.load("../data/external/MNIST_y.npy")

def get_split_data(features=X, labels=Y, test_size=0.2):
    return train_test_split(
        features, labels, test_size=0.2, random_state=42, shuffle=True
    )

def getNormalizedProjectedData(components=2):
    steps = Pipeline(
        [
            ("minmax", MinMaxScaler()),
            ("pca", PCA(n_components=components, random_state=42)),
        ]
    )
    X_transformed = steps.fit_transform(X)
    return X_transformed

ext: py

reason: import numpy


text:
# class FrequencyDomain:

# def simple_resonance_detector()

from scipy.optimize import curve_fit, least_squares
import matplotlib.pyplot as plt
import numpy as np
from .helper_functions import *

def lorentzian_fit_func(f, f0, gamma, a, b):
    omega, omega0 = 2 * np.pi * f, 2 * np.pi * f0

    return a + b / np.pi * (gamma / 2) / ((omega - omega0) ** 2 + (gamma / 2) ** 2)

def analyze_lorentzian(f, sig, p0=None):
    sig_mag = sig
    if sig.dtype == complex:
        sig_mag = np.abs(sig) ** 2

    if p0 is None:
        if (np.max(sig_mag) - np.mean(sig_mag)) > (np.mean(sig_mag) - np.min(sig_mag)):
            # peak detected case
            f0 = f[np.argmax(sig_mag)]
            a = np.mean(sig_mag[np.argsort(sig_mag)[:int(len(sig_mag) // 10)]]) # baseline (average of smallest 10% samples)
            # linewidth is extracted from sample closest to half-max
            gamma = 2 * np.abs(f[np.argmin(np.abs(sig_mag - 0.5 * (np.max(sig_mag) + a)))] - f0)
            b = np.pi * gamma / 2 * (np.max(sig_mag) - a)

            p0 = [f0, gamma, a, b]
        elif (np.max(sig_mag) - np.mean(sig_mag)) < (np.mean(sig_mag) - np.min(sig_mag)):
            # valley detected case
            f0 = f[np.argmin(sig_mag)]
            a = np.mean(sig_mag[np.argsort(-sig_mag)[:int(len(sig) // 10)]]) # baseline (average of largest 10% samples)
            # linewidth is extracted from sample closest to half-max
            gamma = 2 * np.abs(f[np.argmin(np.abs(sig_mag - 0.5 * (np.min(sig_mag) + a)))] - f0)
            b = np.pi * gamma / 2 * (np.min(sig_mag) - a)

            p0 = [f0, gamma, a, b]

    fit = curve_fit(lorentzian_fit_func, f, sig_mag, p0=p0,
                    bounds=([p0[0] * 0.5, p0[1] * 0.5, 0, p0[3] * 0.1],
                            [p0[0] * 1.5, p0[1] * 1.5, np.inf, p0[3] * 10]))
    return fit

def gaussian_fit_func(f, f0, a, c, d):
    return a * np.exp(-(f - f0)**2 / (2 * c**2)) + d

def analyze_gaussian(f, sig, p0=None):
    sig_mag = sig
    if sig.dtype == complex:
        sig_mag = np.abs(sig) ** 2

    if p0 is None:
        if (np.max(sig_mag) - np.mean(sig_mag)) > (np.mean(sig_mag) - np.min(sig_mag)):
            # peak detected case
            f0 = f[np.argmax(sig_mag)]
            d = np.mean(sig_mag[np.argsort(sig_mag)[:int(len(sig_mag) // 10)]]) # baseline (average of smallest 10% samples)
            # linewidth is extracted from sample closest to half-max
            c = 1 / np.sqrt(2) * np.abs(f[np.argmin(np.abs(sig_mag - ((np.max(sig_mag) - d) / np.exp(1) + d)))] - f0)
            a = (np.max(sig_mag) - d)

            p0 = [f0, a, c, d]
        elif (np.max(sig_mag) - np.mean(sig_mag)) < (np.mean(sig_mag) - np.min(sig_mag)):
            # valley detected case
            f0 = f[np.argmin(sig_mag)]
            d = np.mean(sig_mag[np.argsort(-sig_mag)[:int(len(sig) // 10)]]) # baseline (average of largest 10% samples)
            # linewidth is extracted from sample closest to half-max
            c = 1 / np.sqrt(2) * np.abs(f[np.argmin(np.abs(sig_mag - ((np.min(sig_mag) - d) / np.exp(1) + d)))] - f0)
            a = (np.min(sig_mag) - d)

            p0 = [f0, a, c, d]

        fit = curve_fit(gaussian_fit_func, f, sig_mag, p0=p0,
                        bounds=([p0[0] * 0.5, p0[1] * 0.5, p0[2] * 0.1, 0],
                                [p0[0] * 1.5, p0[1] * 1.5, p0[2] * 10, np.inf]))
    return fit

# class DispersiveShift:

class FrequencyDomain:
    def __init__(self, freq, signal):
        # initialize parameters
        self.frequency = freq
        self.signal = signal
        self.n_pts = len(self.signal)
        self.is_analyzed = False
        self.p0 = None
        self.popt = None
        self.pcov = None

    def _guess_init_params(self):
        """
        Guess initial parameters from data. Will be overwritten in subclass
        """

    def _set_init_params(self, p0):
        if p0 is None:
            self._guess_init_params()
        else:
            self.p0 = p0
    def _save_fit_results(self, popt, pcov):
        self.popt = popt
        self.pcov = pcov

    def analyze(self, p0=None, plot=True, **kwargs):
        """
        Analyze the data with initial parameter `p0`.
        """
        # set initial fit parameters
        self._set_init_params(p0)
        # perform fitting
        popt, pcov = curve_fit(self.fit_func, self.frequency, self.signal,
                               p0=self.p0, **kwargs)
        self.is_analyzed = True

        # save fit results
        self._save_fit_results(popt, pcov)

        if plot:
            self.plot_result()

    def _plot_base(self):
        fig = plt.figure()

        # plot data
        _, self.frequency_prefix = number_with_si_prefix(np.max(np.abs(self.frequency)))
        self.frequency_scaler = si_prefix_to_scaler(self.frequency_prefix)

        plt.plot(self.frequency / self.frequency_scaler,
                 self.signal, '.', label="Data", color="black")
        plt.xlabel("Frequency (" + self.frequency_prefix + "Hz)")
        plt.ylabel("Signal")
        plt.legend(loc=0, fontsize=14)

        fig.tight_layout()
        return fig

    def plot_result(self):
        """
        Will be overwritten in subclass
        """
        if not self.is_analyzed:
            raise ValueError("The data must be analyzed before plotting")

    def _get_const_baseline(self, baseline_portion=0.2,
                            baseline_ref='symmetric'):

        samples = self.signal

        N = len(samples)
        if baseline_ref == 'left':
            bs = np.mean(samples[:int(baseline_portion * N)])
        elif baseline_ref == 'right':
            bs = np.mean(samples[int(-baseline_portion * N):])
        elif baseline_ref == 'symmetric':
            bs_left = np.mean(samples[int(-baseline_portion * N / 2):])
            bs_right = np.mean(samples[:int(baseline_portion * N / 2)])
            bs = np.mean([bs_left, bs_right])

        return bs

class LorentzianFit(FrequencyDomain):
    """

    """
    def fit_func(self, f, f0, df, a, b):
        """
        Lorentzian fit function

        Parameters
        ----------
        f : TYPE
            DESCRIPTION.
        f0 : TYPE
            Resonant frequency.
        df : TYPE
            Full-width half-maximum linewidth.
        a : TYPE
            DESCRIPTION.
        b : TYPE
            DESCRIPTION.

        Returns
        -------
        TYPE
            DESCRIPTION.

        """

        return a / ((f - f0) ** 2 + (df / 2) ** 2) + b

    def _guess_init_params(self):
        """
        Guess initial parameters from data.
        """
        signal = self.signal
        f = self.frequency

        b0 = self._get_const_baseline()

        peak_A0, dip_A0 = np.max(signal) - b0, np.min(signal) - b0
        if peak_A0 > - dip_A0: # peak detected case
            A0 = peak_A0
            f0 = f[np.argmax(signal)]
        else: # valley detected case
            A0 = dip_A0
            f0 = f[np.argmin(signal)]

        # linewidth is extracted from sample closest to half-max(arg1, arg2, _args)
        df0 = 2 * np.abs(f[np.argmin(np.abs(signal - (0.5 * A0 + b0)))] - f0)
        a0 = A0 * (df0 / 2) ** 2

        self.p0 = [f0, df0, a0, b0]

    def _save_fit_results(self, popt, pcov):
        super()._save_fit_results(popt, pcov)

        self.f0 = popt[0]
        self.f0_sigma_err = np.sqrt(pcov[0, 0])
        self.df = popt[1]
        self.df_sigma_err = np.sqrt(pcov[1, 1])

    def plot_result(self, fit_n_pts=1000):
        super().plot_result()

        # get most of the plotting done
        fig = self._plot_base()

        freq_fit = np.linspace(self.frequency[0],
                               self.frequency[-1], fit_n_pts)

        plt.plot(freq_fit / self.frequency_scaler,
                 self.fit_func(freq_fit, *(self.p0)),
                 label="Fit (Init. Param.)", ls='--', lw=2, color="orange")

        plt.plot(freq_fit / self.frequency_scaler,
                 self.fit_func(freq_fit, *(self.popt)),
                 label="Fit (Opt. Param.)", lw=2, color="red")

        f0_string = (r"$f_0$ = %.4f $\pm$ %.4f " %
                     (self.f0 / self.frequency_scaler,
                      2 * self.f0_sigma_err / self.frequency_scaler) +
                     self.frequency_prefix + 'Hz')

        _, df_prefix = number_with_si_prefix(np.max(np.abs(self.df)))
        df_scaler = si_prefix_to_scaler(df_prefix)

        df_string = (r"$\Delta f$ = %.4f $\pm$ %.4f " %
                     (self.df / df_scaler, 2 * self.df_sigma_err / df_scaler) +
                     df_prefix + 'Hz')

        plt.title(f0_string + ', ' + df_string)
        plt.legend(loc=0, fontsize='x-small')
        fig.tight_layout()
        plt.show()

        return fig

class GaussianFit(FrequencyDomain):

    def fit_func(self, f, f0, sigma_f, a, b):

        return a * np.exp(-(f - f0) ** 2 / (2 * sigma_f ** 2)) + b

    def _guess_init_params(self):
        """
        Guess initial parameters from data.
        """
        signal = self.signal
        f = self.frequency

        b0 = self._get_const_baseline()

        peak_a0, dip_a0 = np.max(signal) - b0, np.min(signal) - b0
        if peak_a0 > - dip_a0: # peak detected case
            a0 = peak_a0
            f0 = f[np.argmax(signal)]
        else: # valley detected case
            a0 = dip_a0
            f0 = f[np.argmin(signal)]

        # sigma linewidth is extracted from sample closest to 1/2 of max
        sigma_f0 = np.sqrt(np.log(2) / 2) * np.abs(f[np.argmin(np.abs(signal - (0.5 * a0 + b0)))] - f0)

        self.p0 = [f0, sigma_f0, a0, b0]

    def _save_fit_results(self, popt, pcov):
        super()._save_fit_results(popt, pcov)

        self.f0 = popt[0]
        self.f0_sigma_err = np.sqrt(pcov[0, 0])
        self.sigma_f = popt[1]
        self.sigma_f_sigma_err = np.sqrt(pcov[1, 1])

    def plot_result(self, fit_n_pts=1000):
        super().plot_result()

        # get most of the plotting done
        fig = self._plot_base()

        freq_fit = np.linspace(self.frequency[0],
                               self.frequency[-1], fit_n_pts)

        plt.plot(freq_fit / self.frequency_scaler,
                 self.fit_func(freq_fit, *(self.p0)),
                 label="Fit (Init. Param.)", ls='--', lw=2, color="orange")

        plt.plot(freq_fit / self.frequency_scaler,
                 self.fit_func(freq_fit, *(self.popt)),
                 label="Fit (Opt. Param.)", lw=2, color="red")

        f0_string = (r"$f_0$ = %.4f $\pm$ %.4f " %
                     (self.f0 / self.frequency_scaler,
                      2 * self.f0_sigma_err / self.frequency_scaler) +
                     self.frequency_prefix + 'Hz')

        _, sigma_f_prefix = number_with_si_prefix(np.max(np.abs(self.sigma_f)))
        sigma_f_scaler = si_prefix_to_scaler(sigma_f_prefix)

        sigma_f_string = (r"$\sigma_f$ = %.4f $\pm$ %.4f " %
                     (self.sigma_f / sigma_f_scaler,
                      2 * self.sigma_f_sigma_err / sigma_f_scaler) +
                     sigma_f_prefix + 'Hz')

        plt.title(f0_string + ', ' + sigma_f_string)
        plt.legend(loc=0, fontsize='x-small')

        fig.tight_layout()
        plt.show()

        return fig

# class SingleSidedS11Fit(FrequencyDomain):
#     """
#     Class for implementing fitting of lineshape in reflection measurement
#     """
#     def __init__(self, freq, data, fit_type='complex'):
#         super().__init__(freq, data, df=df, fit_mag_dB=fit_mag_dB, plot_mag_dB=plot_mag_dB)

#         self.fit_type = fit_type

#     def fit_func(self, *args):
#         """
#         Reflection fit for resonator single-sided coupled to waveguide
#         according to relation
#         S11 = a * exp(i(ϕ + 2pi * (f - f[0]) * τ)) .* ...
#              (1 - 2 * QoverQe * (1 + 2i * δf / f0) / (1 + 2i * Q * (f - f0) / f0))

#         Note: If df = 0 (no background) this gives
#         S11 = 1 - (2 kappa_e) / (kappa_i + kappa_e + 2i (omega - omega0))
#             = (kappa_i - kappa_e + 2i (omega - omega0)) / (kappa_i + kappa_e + 2i (omega - omega0))
#         See Aspelmeyer et al, "Cavity Optomechanics", Rev. Mod. Phys. 86, 1391 (2014).

#         If the fit type is complex, the arguments of `fit_func` are:
#             f, f0, Q, QoverQe, df, a, phi, tau
#         If the fit type is magnitude, the arguments of `fit_func` are:
#             f, f0, Q, QoverQe, df, a
#         If the fit type is phase, the arguments of `fit_func` are:
#             f, f0, Q, QoverQe, df, phi, tau
#         """
#         if self.fit_type == 'complex':
#             f, f0, Q, QoverQe, df, a, phi, tau = args

#             return (a * np.exp(1j * (phi + 2 * np.pi * (f - f[0]) * tau)) *
#                     (1 - 2 * QoverQe * (1 + 2j * df / f0) /
#                     (1 + 2j * Q * (f - f0) / f0)))
#         elif self.fit_type == 'magnitude':
#             f, f0, Q, QoverQe, df, a = args

#             return np.abs(a * (1 - 2 * QoverQe * (1 + 2j * df / f0) /
#                                (1 + 2j * Q * (f - f0) / f0)))
#         elif self.fit_type == 'phase':
#             f, f0, Q, QoverQe, df, phi, tau = args

#             return np.angle(np.exp(1j * (phi + 2 * np.pi * (f - f[0]) * tau)) *
#                             (1 - 2 * QoverQe * (1 + 2j * df / f0) /
#                              (1 + 2j * Q * (f - f0) / f0)))

#     def _estimate_f0_FWHM(self):
#         f = self.frequency
#         mag2 = np.abs(self.data) ** 2

#         # background in magnitude estimated by linear interpolation of first and last point
#         mag2_bg = (mag2[-1] - mag2[0]) / (f[-1] - f[0]) * (f - f[0]) + mag2[0]

#         mag2_subtracted = mag2 - mag2_bg
#         f0 = f[mag2_subtracted.argmin()]
#         smin, smax = np.min(mag2_subtracted), np.max(mag2_subtracted)

#         # data in frequency < f0 or frequency >= f0
#         f_l, s_l = f[f < f0], mag2_subtracted[f < f0]
#         f_r, s_r = f[f >= f0], mag2_subtracted[f >= f0]

#         # find frequencies closest to the mean of smin and smax
#         f1 = f_l[np.abs(s_l - 0.5 * (smin + smax)).argmin()]
#         f2 = f_r[np.abs(s_r - 0.5 * (smin + smax)).argmin()]

#         # numerically find full width half max from magnitude squared
#         Δf = f2 - f1

#         return f0, Δf

#     def _guess_init_params(self, df):

#         # magnitude data
#         _mag = np.abs(self.data)
#         # phase data
#         _ang = np.angle(self.data)
#         # unwrapped phase data
#         _angU = np.unwrap(_ang)

#         f = self.frequency

#         a0 = np.sqrt((_mag[0] ** 2 + _mag[-1] ** 2) / 2)
#         phi0 = _angU[0]
#         tau0 = 0.0
#         if (np.max(_angU) - np.min(_angU)) > 2.1 * np.pi:
#             # if phase data at start and stop frequencies differ by more than 2pi,
#             # perform phase subtraction associated with delay
#             tau0 = (_angU[-1] - _angU[0]) / ((f[-1] - f[0]))/ (2 * np.pi)

#         # Estimate total Q from the FWHM in |mag|^2
#         f0, Δf = self._estimate_f0_FWHM()
#         amin = _mag[np.abs(f - f0).argmin()]
#         QoverQe0 = 0.5 * (1 - amin / a0)
#         Q0 = f0 / Δf

#         if self.fit_type == 'complex':
#             self.p0 = [f0, Q0, QoverQe0, 0.0, a0, phi0, tau0]
#         elif self.fit_type == 'magnitude':
#             self.p0 = [f0, Q0, QoverQe0, 0.0, a0]
#         elif self.fit_type == 'phase':
#             self.p0 = [f0, Q0, QoverQe0, 0.0, phi0, tau0]

# class acStarkShift:
#     """
#     Class to implement analysis of ac Stark shift measurement with varying
#     power applied to the readout resonator.

#     Parameters
#     ----------
#     freq : `numpy.ndarray`
#         A 1D numpy array that specifies the range of frequencies (in units of
#         Hz) for qubit spectroscopy.
#     power_dBm : `numpy.ndarray`
#         A 1D numpy array that specifies the range of power (in units of dBm)
#         for readout resonator.
#     signal : `numpy.ndarray`
#         A 2D numpy array that stores the complex signal obtained from qubit
#         spectroscopy with taken with a range of readout power. The row indices
#         and column indices correspond to indices of `power_dBm` and `freq`,
#         respectively.
#     disp_freq_shift : `float`
#         Dispersive frequency shift of readout resonator 2 * chi / (2 * pi) in
#         units of Hz.

#     Attributes
#     ----------

#     Methods
#     -------

#     """
#     def __init__(self, freq, power_dBm, signal, disp_freq_shift):
#         self.frequency = freq
#         # complex signal obtained from spectroscopy
#         self.signal = signal
#         # readout power in units of dBm
#         self.power_dBm = power_dBm
#         # dispersive frequency shift (2*chi / 2*pi) in units of Hz
#         self.disp_freq_shift = disp_freq_shift # 2 * chi / (2 * pi)

#         self.p0 = None
#         self.popt = None
#         self.pcov = None

#     def analyze(self, plot=True, p0=None):

#         self.res_frequency = np.zeros(len(self.power_dBm))
#         for idx in range(len(self.power_dBm)):
#             lorentzian_fit = analyze_lorentzian(self.frequency, self.signal[idx])
#             self.res_frequency[idx] = lorentzian_fit[0][0]

#         self._set_init_params(p0)

#         self.popt, self.pcov = curve_fit(self.fit_func, self.power_dBm,
#                                          self.res_frequency, p0=self.p0)
#         a, f0 = self.popt

#         self.single_photon_power_dBm = watt_to_dBm(self.disp_freq_shift / a)

#         if plot:
#             self.plot_result()

#     def _set_init_params(self, p0):
#         if p0 is None:
#             dp = dBm_to_watt(np.max(self.power_dBm))
#             maxind, minind = np.argmax(self.power_dBm), np.argmin(self.power_dBm)
#             df = self.res_frequency[maxind] - self.res_frequency[minind]
#             self.p0 = [df / dp, self.res_frequency[np.argmin(self.power_dBm)]]
#         else:
#             self.p0 = p0

#     def fit_func(self, power_dBm, a, f0):

#         power_Watt = dBm_to_watt(power_dBm)
#         return a * power_Watt + f0

#     def plot_result(self):
#         fig = plt.figure()
#         plt.pcolormesh(self.frequency / 1e9, self.power_dBm, np.abs(self.signal), shading='auto')
#         plt.plot(self.res_frequency / 1e9, self.power_dBm, '.', color='k', label='Res. Freq.')
#         plt.plot(self.fit_func(self.power_dBm, *self.popt) / 1e9,
#                  self.power_dBm, '--', color='white', label='Fit')
#         plt.title(r"Single-photon power = $%.2f$ dBm" % self.single_photon_power_dBm)
#         plt.legend()
#         plt.xlabel("Frequency (GHz)")
#         plt.ylabel("Power (dBm)")

# class WaveguideCoupledS21Fit(FrequencyDomain):
#     """
#     Class for imlementing fitting of lineshape in reflection measurement
#     """
#     def __init__(self, freq, data, df=0, fit_mag_dB=False, plot_mag_dB=False):
#         super().__init__(freq, data, df=df, fit_mag_dB=fit_mag_dB, plot_mag_dB=plot_mag_dB)

#         self.fit_type = "WaveguideCoupledTransmission"

#     def fit_func(self, f, f0_MHz, Q, QoverQe, δf, a, ϕ, τ_ns):
#         """
#         Reflection fit for resonator single-sided coupled to waveguide
#         according to relation
#         S11 = a0 * exp(i(ϕ0 + 2pi * (f - f[0]) * τ0)) .* ...
#              (1 - QoverQe * (1 + 2i * δf / f0) / (1 + 2i * Q * (f - f0) / f0))

#         Note: If df = 0 (no background) this gives
#         S21 = 1 - (kappa_e) / (kappa_i + kappa_e + 2i (omega - omega0))
#             = (kappa_i + 2i (omega - omega0)) / (kappa_i + kappa_e + 2i (omega - omega0))
#         See Khalil et al, "An analysis method for asymmetric resonator
#         transmission applied to superconducting devices",
#         J. Appl. Phys. 111, 054510 (2012).
#         """
#         return (a * np.exp(1j * (ϕ + 2 * np.pi * (f - f[0]) * τ_ns * 1e-9)) *
#                 (1 - QoverQe * (1 + 2j * δf / f0_MHz) /
#                 (1 + 2j * (Q) * (f/1e6 - f0_MHz) / (f0_MHz))))

#     def _init_fit_params(self, df):
#         # magnitude data
#         _mag = np.abs(self.data)
#         # phase data
#         _ang = np.angle(self.data)
#         # unwrapped phase data
#         _angU = np.unwrap(_ang)

#         f = self.frequency

#         a0 = _mag[0]
#         ϕ0 = _angU[0]
#         τ0 = 0.0
#         if (np.max(_angU) - np.min(_angU)) > 2.1 * np.pi:
#             # if phase data at start and stop frequencies differ more than 2pi,
#             # perform phase subtraction associated with delay
#             τ0 = (_angU[-1] - _angU[0]) / ((f[-1] - f[0]))/ (2 * np.pi)

#         # Estimate total Q from the FWHM in |mag|^2
#         f0, Δf = self._estimate_f0_FWHM()
#         QoverQe0 = (1 - np.min(_mag) / a0)
#         Q0 = f0 / Δf
#         p0_mag, p0_ang = self._prepare_fit_params(f0, Q0, QoverQe0,
#                                                   df, a0, ϕ0, τ0)

#         self.p0 = p0_mag + p0_ang
#         return p0_mag, p0_ang

max_stars_repo_name: eunjongkim/Qanalysis

max_issues_repo_name: eunjongkim/Qanalysis

max_forks_repo_name: eunjongkim/Qanalysis

reason: import numpy,from scipy


text:
"""Tests for the PyMapper class."""
import os
import time
import pytest
import numpy as np
import tensorflow as tf

from ramjet.py_mapper import PyMapper, map_py_function_to_dataset

class TestPyMapper:
    """Tests for the PyMapper class."""

    @pytest.fixture
    def dataset(self) -> tf.data.Dataset:
        """
        Sets up the dataset for use in a test.

        :return: The dataset.
        """
        return tf.data.Dataset.from_tensor_slices([0, 10, 20, 30])

    @pytest.mark.slow
    def test_py_map_runs_function_on_multiple_processes(self, dataset: tf.data.Dataset):
        py_mapper = PyMapper(sleep_and_get_pid, number_of_parallel_calls=4)
        map_dataset = py_mapper.map_to_dataset(dataset)
        batch_dataset = map_dataset.batch(batch_size=4)
        batch = next(iter(batch_dataset))
        batch_array = batch.numpy()
        unique_pids = set(batch_array)
        assert len(unique_pids) > 2

    def test_py_map_correctly_applies_map_function(self, dataset: tf.data.Dataset):
        py_mapper = PyMapper(add_one, number_of_parallel_calls=4)
        map_dataset = py_mapper.map_to_dataset(dataset)
        batch_dataset = map_dataset.batch(batch_size=4)
        batch = next(iter(batch_dataset))
        batch_array = batch.numpy()
        assert np.array_equal(batch_array, np.array([1, 11, 21, 31]))

    def test_py_map_correctly_applies_map_function_with_two_outputs(self, dataset: tf.data.Dataset):
        py_mapper = PyMapper(add_one_and_add_two, number_of_parallel_calls=4)
        map_dataset = py_mapper.map_to_dataset(dataset, output_types=[tf.float32, tf.float32])
        batch_dataset = map_dataset.batch(batch_size=4)
        batch = next(iter(batch_dataset))
        plus_one_array = batch[0].numpy()
        plus_two_array = batch[1].numpy()
        assert np.array_equal(plus_one_array, np.array([1, 11, 21, 31]))
        assert np.array_equal(plus_two_array, np.array([2, 12, 22, 32]))

    def test_py_map_correctly_applies_map_function_with_two_inputs(self, dataset: tf.data.Dataset):
        py_mapper = PyMapper(add_tensors, number_of_parallel_calls=4)
        zipped_dataset = tf.data.Dataset.zip((dataset, dataset))
        map_dataset = py_mapper.map_to_dataset(zipped_dataset)
        batch_dataset = map_dataset.batch(batch_size=4)
        batch = next(iter(batch_dataset))
        batch_array = batch.numpy()
        assert np.array_equal(batch_array, np.array([0, 20, 40, 60]))

    def test_single_function_wrapper(self, dataset):
        mapped_dataset = map_py_function_to_dataset(dataset=dataset, map_function=add_one, number_of_parallel_calls=4,
                                                    output_types=tf.float32)
        batch_dataset = mapped_dataset.batch(batch_size=4)
        batch = next(iter(batch_dataset))
        batch_array = batch.numpy()
        assert np.array_equal(batch_array, np.array([1, 11, 21, 31]))

    def test_py_map_can_be_applied_as_flat_map(self):
        dataset = tf.data.Dataset.from_tensor_slices([[[0, 0], [10, 10]], [[20, 20], [30, 30]]])
        py_mapper = PyMapper(add_one, number_of_parallel_calls=4)
        map_dataset = py_mapper.map_to_dataset(dataset, flat_map=True)
        batch_dataset = map_dataset.batch(batch_size=4)
        batch = next(iter(batch_dataset))
        batch_array = batch.numpy()
        assert np.array_equal(batch_array, np.array([[1, 1], [11, 11], [21, 21], [31, 31]]))

    def test_flat_map_can_be_used_from_single_function_wrapper(self):
        dataset = tf.data.Dataset.from_tensor_slices([[[0, 0], [10, 10]], [[20, 20], [30, 30]]])
        mapped_dataset = map_py_function_to_dataset(dataset=dataset, map_function=add_one, number_of_parallel_calls=4,
                                                    output_types=tf.float32, flat_map=True)
        batch_dataset = mapped_dataset.batch(batch_size=4)
        batch = next(iter(batch_dataset))
        batch_array = batch.numpy()
        assert np.array_equal(batch_array, np.array([[1, 1], [11, 11], [21, 21], [31, 31]]))

    def test_py_map_returns_specified_shape(self):
        dataset = tf.data.Dataset.from_tensor_slices([[0, 0, 0], [1, 1, 1]])
        py_mapper = PyMapper(add_one, number_of_parallel_calls=4)
        map_dataset = py_mapper.map_to_dataset(dataset, output_shapes=(3,))
        assert map_dataset.element_spec.shape == (3,)

    def test_py_map_returns_specified_shape_for_multiple_elements(self):
        dataset0 = tf.data.Dataset.from_tensor_slices([[0, 0, 0], [1, 1, 1]])
        dataset1 = tf.data.Dataset.from_tensor_slices([[2, 2], [3, 3]])
        zipped_dataset = tf.data.Dataset.zip((dataset0, dataset1))
        py_mapper = PyMapper(add_one, number_of_parallel_calls=4)
        map_dataset = py_mapper.map_to_dataset(zipped_dataset, output_types=[tf.float32, tf.float32],
                                               output_shapes=[(3,), (2,)])
        assert map_dataset.element_spec[0].shape == (3,)
        assert map_dataset.element_spec[1].shape == (2,)

    def test_py_map_returns_specified_shape_when_shape_passed_in_wrapper(self):
        dataset = tf.data.Dataset.from_tensor_slices([[0, 0, 0], [1, 1, 1]])
        map_dataset = map_py_function_to_dataset(dataset=dataset, map_function=add_one, number_of_parallel_calls=4,
                                                 output_shapes=(3,))
        assert map_dataset.element_spec.shape == (3,)

    def test_flat_map_with_output_shapes_are_applied_in_the_correct_order(self):
        dataset = tf.data.Dataset.from_tensor_slices([[[0, 0], [10, 10]], [[20, 20], [30, 30]]])
        mapped_dataset = map_py_function_to_dataset(dataset=dataset, map_function=add_one, number_of_parallel_calls=4,
                                                    output_types=tf.float32, flat_map=True, output_shapes=(2,))
        batch_dataset = mapped_dataset.batch(batch_size=4)
        batch = next(iter(batch_dataset))
        batch_array = batch.numpy()
        assert np.array_equal(batch_array, np.array([[1, 1], [11, 11], [21, 21], [31, 31]]))

def sleep_and_get_pid(element_tensor: tf.Tensor) -> int:
    """
    A simple sleep and get pid function to test multiprocessing.

    :param element_tensor: Input value.
    :return: The pid of the process that ran this function.
    """
    time.sleep(0.1)
    return os.getpid()

# noinspection PyPackageRequirements
def add_tensors(element_tensor0: tf.Tensor, element_tensor1: tf.Tensor) -> float:
    """
    Adds two elements together.

    :param element_tensor0: First input value.
    :param element_tensor1: Second input value.
    :return: The added inputs.
    """
    element0 = element_tensor0.numpy()
    element1 = element_tensor1.numpy()
    return element0 + element1

def add_one(element_tensor: tf.Tensor) -> float:
    """
    Adds 1

    :param element_tensor: Input value.
    :return: Input plus 1.
    """
    element = element_tensor.numpy()
    return element + 1

def add_one_and_add_two(element_tensor: tf.Tensor) -> (float, float):
    """
    Adds 1 and adds 2 returning both.

    :param element_tensor: Input value.
    :return: Input plus 1 and input plus 2.
    """
    element = element_tensor.numpy()
    return element + 1, element + 2

max_stars_repo_name: golmschenk/ramjet

max_issues_repo_name: golmschenk/ramjet

max_forks_repo_name: golmschenk/ramjet


text:
"""
A module containing a bunch of objects for testing.

See also: https://github.com/i2mint/py2http/blob/master/py2http/tests/objects_for_testing.py
"""

###############################################################################
# Empty callables, meant to give an example of different kinds of callables ###
# Note that x can be applied to a number, string, or list #####################
def f(x):
    """Note that x can be applied to a number, string, or list"""
    return x * 2

class C:
    @classmethod
    def c(cls, x):
        return f(x)

    @staticmethod
    def s(x):
        return f(x)

    # instance method
    def i(self, x):
        return f(x)

    # making an instance itself callable
    def __call__(self, x):
        return f(x)

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

def times_2(x):
    return x * 2

def plus_1(z):
    return z + 1

def add(a, b: float = 0.0) -> float:
    """Adds numbers"""
    return a + b

def mult(x: float, y=1):
    return x * y

def formula1(w, /, x: float, y=1, *, z: int = 1):
    """Has 3 kinds & every of the 4 combinations of (default y/n, annotated y/n)"""
    return ((w + x) * y) ** z

def formula2(w, /, x: float, y=2, *, z: int = 3):
    """Has 3 kinds & every of the 4 combinations of (default y/n, annotated y/n)

    >>> formula2(0, 1, 2)
    '(w:=0) + (x:=1) * (y:=2) ** (z:=3) == 8'
    """
    return f'(w:={w}) + (x:={x}) * (y:={y}) ** (z:={z}) == {w + x * y ** z}'

# Only *args and **kwargs. Also has a doc with a doc test
def sum_of_args(*args, **kwargs):
    """Sums the args, and returns a copy of kwargs with the sum of args added to each value.

    >>> sum_of_args(1,2,3, a=0, b=10)
    {'a': 6, 'b': 16}
    """
    t = sum(args)
    return {k: t + v for k, v in kwargs.items()}

def test_func_1(
    any_var,
    a_list: list,
    a_dict: dict,
    an_int: int,
    a_float: float = 3.14,
    a_bool=True,
    a_str: str = None,
) -> str:
    """
    This is the first line,
    continued on the next.

    This function is meant to demo several combinations of type declaration (in signature or/and in docs),
    as well as with or without defaults.

    :param any_var: no type declared in either signature of doc
    :param list a_list: type declared in signature AND doc
    :param a_dict: type declared in signature but NOT in doc
    :param int an_int: type declared in doc but NOT in signature
    :param float a_float: type declared in signature AND doc, and has a default of the desired type
    :param a_bool: type declared in signature but NOT in doc, and has a default of the desired type
    :param str a_str: type declared in signature AND doc, but has a default that is not of that type
    line description
    :return str: Will just return the 'test_func_1 returned' string
    :tags foo, bar

    >>> # here are sometests
    >>> test_func_1('anything', [1, 2], {'this': 0, 'that': 'foo'}, 1, 1.1, a_bool=True, a_str='hello')
    'test_func_1 returned'
    """
    # This is the first comment
    return 'test_func_1 returned'  # this is a comment on the same line as some code

class SomeClass(object):
    pass

class SomeOtherClass(object):
    pass

def test_func_2(
    any_var,
    an_int: int,
    a_float=0.1,
    a_tuple: tuple = (),
    an_obj: SomeClass = SomeClass(),
    another_obj=SomeOtherClass(),
):
    """
    This is the first line,
    continued on the next.

    And this is some more information about the function:
    Blah
    Blah
    Bloo

    :param any_var: foo
    :param an_int: you again?!
    :param a_float:
    :param a_tuple: For real
    :param an_obj: Does nothing
    :param another_obj: Multi-
    line
    :return: Just pi
    :tags any, old, tag
    :keyword one_keyword
    :keyword keyword1, keyword2
    """

    """
    This is the first line,
    continued on the next.

    And this is some more information about the function:
    Blah
    Blah
    Bloo

    :param any_var: this can be of any type
    :param an_int: an int
    :param a_float: a float
    :param a_bool: a bool
    :param a_str: a str
    :param a_list:
    :param a_dict: An example of a continued
    line description
    :return:

    >>> # here are sometests
    >>> test_func_1('anything', 1, 1.1, True, 'hello', [1, 2], {'this': 0, 'that': 'foo'})
    'returned'
    """
    return 3.14

class AClass(object):
    def __init__(self, a=1, b=0):
        """The doc of the __init__"""
        self.a = a
        self.b = b

    def __call__(self, x: float) -> float:
        """The actual callable"""
        return self.a * x + self.b

class TestObj:
    constarg: str = ''

    def __init__(self, constarg: str):
        print(constarg)
        self.constarg = constarg

    def methodnum(self, methodarg1: int) -> int:
        print(methodarg1)
        return methodarg1 + 1

    def methodstr(self, methodarg2: str = 'hi') -> str:
        print(methodarg2)
        return methodarg2 + ' test ' + self.constarg

from contextlib import suppress

with suppress(ModuleNotFoundError, ImportError):
    from numpy import ndarray, matrix
    from pandas import DataFrame, Series

    def a_func_with_df_in_and_sr_out(df: DataFrame) -> Series:
        return df.sum()

    def a_func_with_np_and_df(
        df: DataFrame, sr: Series, arr: ndarray, mat: matrix
    ) -> matrix:
        return (df.sum().sum() + sr.sum() + arr.sum()) * mat

reason: from numpy


text:
import numpy as np
import pandas as pd
from assistant import cvs_to_arr

def get_metrics(pred_arr):
    """
    Get metrics from predictions

    arr : prediction_array

    return -> metric values; accuracy[0], sensitivity[1], specificity[2]
    """
    arr = pred_arr.transpose(1, 0)
    pred = arr[1]                               # prediction value
    corr_label = arr[0]                         # true value
    tp_tn = pred == corr_label                  # if pred == corr_label, it returns True else False
    correct_pred = tp_tn.astype(int)
    tp, tn, fp, fn = 0, 0, 0, 0
    # count = 0
    for index in range(0, len(arr[0])):
        if correct_pred[index] == 1:
            # count += 1
            if corr_label[index] == 1:
                tp += 1
            elif corr_label[index] == 0:
                tn += 1
        elif correct_pred[index] == 0:
            # count += 1
            if corr_label[index] == 0:
                fp += 1
            elif corr_label[index] == 1:
                fn += 1

    # print('tp',tp)
    # print('np',tn)
    # print('fp',fp)
    # print('fn',fn)

    accuracy = (tp + tn) / (tp + tn + fp + fn)
    sensitivity = tp / (tp + fn)                # recall = sensitivity
    specificity = tn / (tn + fp)

    # print("accuracy: %.3f" % accuracy +"\n" + "sensitivity: %.3f" % sensitivity + "\n" + "specificity: %.3f" %
    # specificity)
    return accuracy, sensitivity, specificity

def cvs_to_metric(cvs_path, requirement):
    """
    Read prediction.csv files and make metrics

    requirement: accuracy or sensitivity or specificity

    return -> wanted metric array
    """
    requirements = []
    for item in cvs_path:
        arr = cvs_to_arr(item)
        metric = get_metrics(arr)
        if requirement == 'Accuracy':
            requirements.append(metric[0])
        if requirement == 'Sensitivity':
            requirements.append(metric[1])
        if requirement == 'Specificity':
            requirements.append(metric[2])
    return requirements

if __name__ == "__main__":
    arr = cvs_to_arr(cvs_file='/example_results/train_data/ep_0_predictions.csv')
    metrics = get_metrics(arr)
    print(metrics[0])


text:
# coding=utf-8
# /usr/bin/env python

'''
Author: wenqiangw
Email: wenqiangw@opera.com
Date: 2020-04-30 14:42
Desc:
    标准评分卡(lr)
    log(odds)=w*x
    score = self.score_offset + self.score_factor * np.log(odds)
'''

from pynovice.score_card.src.score_card import ScoreCardModel
from sklearn.linear_model import LogisticRegression
from sklearn.model_selection import train_test_split
import pandas as pd
import numpy as np

class LRModel(ScoreCardModel):
    '''
        用于预测某一条预处理过的特征向量得分的方法
    :return:
        score_label: 最终预测的评分
        score_feature: 各个特征的得分字典
    '''

    def __init__(self,**kwargs):
        super().__init__()
        self.model_ = None
        self.woe_score = None
        self.kwargs = kwargs

    def train(self,train_feature, train_label, test_feature=pd.DataFrame, test_label=pd.DataFrame,eval=True,eval_plot=False):
        if (test_feature.empty or test_label.empty):
            train_feature, test_feature, train_label, test_label = \
                train_test_split(train_feature, train_label, test_size=0.2, random_state=0)

        self.model_ = LogisticRegression(**self.kwargs)
        self.model_.fit(train_feature, train_label)
        if eval:
            self.woe_score = self.get_card_info(train_feature)
            # 模型评估
            print("training eval:")
            df_pre, _ = self.predict(train_feature)
            auc, _ = self.get_auc(df_pre, train_label)
            ks, _ = self.get_ks(df_pre, train_label)
            lift, _ = self.get_lift(df_pre, train_label)
            print("auc={}, ks={}, lift={}".format(auc, ks['gap'].values[0], lift))

            print('testing eval:')
            df_pre, _ = self.predict(test_feature)
            auc, _ = self.get_auc(df_pre, test_label)
            ks, _ = self.get_ks(df_pre, test_label)
            lift, _ = self.get_lift(df_pre, test_label)
            print("auc={}, ks={}, lift={}".format(auc, ks['gap'].values[0], lift))

        if eval_plot:
            df_pre, _ = self.predict(test_feature)
            self.plot_roc(df_pre, test_label, pre_target=1, save_path='.')
            self.plot_ks(df_pre, test_label, pre_target=1, save_path='.')
            self.plot_lift(df_pre, test_label, pre_target=1, save_path='.')

    def predict(self,x):
        '''
        :param x: list/datafame
        :return: array([proba]),array([score_label])
        '''
        proba = self.model_.predict_proba(x)[:,1]
        score_label = np.array([self.probability_to_score(bad_prob=proba_i) for proba_i in proba])
        return proba,score_label

    def predict_detail(self,x):
        '''

        :param x: list/datafame
        :return: base_score,array([field_score]),array([final_score])
        '''
        coef_ = self.model_.coef_[0]
        intercept_ = self.model_.intercept_[0]

        base_score = round(self.score_offset + self.score_factor * intercept_,2)  # 基础分
        field_score = (np.array(x)*coef_*self.score_factor).round(decimals=2)  # 特征的分值
        final_score = (base_score + field_score.sum(axis=1)).astype(int)  #综合评分
        return base_score, field_score, final_score

    def get_card_info(self,df_field):
        '''
        计算评分卡模型详细分数信息
        :param df_field:
        :return:
        '''
        _out = []
        base_score = None
        for i,field in enumerate(df_field.columns):
            df_idx = df_field[field].drop_duplicates().index
            df_eval = df_field.loc[df_idx,:]
            base_score, field_score, _ = self.predict_detail(x=df_eval)
            # 变量评分
            df_iter = pd.DataFrame(df_eval[field].values,columns=['value'])
            df_iter['score'] = field_score[:,i]
            df_iter.insert(loc=0, column='fields', value=field)
            _out.append(df_iter.sort_values('value'))
        df_card = pd.concat(_out, axis=0)
        df_card.loc['__'] = ['_base',None, base_score]
        df_card = df_card.sort_values('fields').reset_index(drop=True)
        return df_card

if __name__ == '__main__':
    df = pd.DataFrame([[1, 2, 3, 4, 5, 5, 5, 6, 8, 3, 2, 1, 5, 7],[10, 2, 3, 42, 534, 5, 53, 6, 83, 3, 42, 21, 25, 7], [1, 1, 0, 0, 0, 0, 0, 1, 1, 0, 1, 1, 1, 1]]).T
    df.columns = ['field1', 'field2', 'label']
    df_fields, df_label = df[['field1','field2']], df['label']
    lr = LRModel(class_weight='balanced')
    lr.train(df_fields,df_label,eval=True)
    print(lr.predict_detail([[2,8]]))
    print(lr.predict_detail(df_fields.head()))

    print(lr.predict([[2,8]]))
    print(lr.predict(df_fields.head()))

max_stars_repo_name: wqwangchn/novice

max_issues_repo_name: wqwangchn/novice

max_forks_repo_name: wqwangchn/novice


text:
import numpy as np
import matplotlib.pyplot as plt
import husimi_sp as hu
import scipy.fftpack as sp
import utils as u_

N = 256
L = 1.
dx = L/N
x = dx*(.5+np.arange(-N/2, N/2))
kx = 2*np.pi*sp.fftfreq(N, d = L/N)
u = np.asarray(sp.fftshift(2*np.pi*sp.fftfreq(N, d = dx)))
hbar = 1.0e-2
u = hbar*np.asarray(sp.fftshift(2*np.pi*sp.fftfreq(N, d = dx)))

dt = 1e-4
Tfinal = 5.

sig_x = dx*5
K = hu.K_H_(x, x.copy(), u, hbar, sig_x, L)

A = 100.
B = 73.
C = B*1e-3
mu = L/4.

V = A*x**2 + C*x**4
def F(x):
    return -2*A*x- 4*B*C**3
V_ = V + B*x**4
def F_(x):
    return -2*A*x - 4*B*x**3

def GetPQ(psi):
    q = np.sum(x*np.abs(psi)**2)*dx
    q2 = np.sum(x*x*np.abs(psi)**2)*dx
    psi_k = sp.fft(psi)
    psi_k /= np.sqrt(np.sum(np.abs(psi_k)**2))
    p = np.sum(kx*np.abs(psi_k)**2)
    p2 = np.sum(kx*kx*np.abs(psi_k)**2)
    Varq = q2 - q**2
    Varp = hbar**2*(p2 - p**2)
    a2 = .5*(p+q*1j)
    return q, p*hbar, Varq / mu**2

def makeFig(psi, psiT, psiT_ , psiT2, psiT2_,  r,v,rT,vT, rT2, vT2  ,rT_,vT_, rT2_, vT2_):
    fig, axs = plt.subplots(2,3,figsize = (16,12))

    H = hu.f_H(psi, K, dx)
    HT = hu.f_H(psiT, K, dx)
    HT_ = hu.f_H(psiT_, K, dx)
    HT2 = hu.f_H(psiT2, K, dx)
    HT2_ = hu.f_H(psiT2_, K, dx)

    cmin = 0.
    cmax = np.max(H)*.7

    im = axs[0][0].imshow(H, interpolation = 'none', extent=[-L/2.01, L/2.01, np.min(u),np.max(u)], aspect = 'auto', origin = 'lower',cmap = 'binary')
    im.set_clim(vmin=cmin, vmax=cmax)
    axs[0][0].text(.8,.9,r'$T = 0$'  , ha='center', va='center', transform= axs[0][0].transAxes, bbox = {'facecolor': 'white', 'pad': 5})
    axs[0][0].plot(r,v, 'ro')
    p,q, Q = GetPQ(psi)
    axs[0][0].plot(p,q, 'c^', linewidth = 5)
    axs[0][0].text(.2,.9,r'$Q_q = %.2f$' %(Q) , ha='center', va='center', transform= axs[0][0].transAxes, bbox = {'facecolor': 'white', 'pad': 5})

    im = axs[1][0].imshow(H, interpolation = 'none', extent=[-L/2.01, L/2.01, np.min(u),np.max(u)], aspect = 'auto', origin = 'lower',cmap = 'binary')
    im.set_clim(vmin=cmin, vmax=cmax)
    axs[1][0].plot(r,v, 'ro')
    p,q, Q = GetPQ(psi)
    axs[1][0].plot(p,q, 'c^', linewidth = 5)
    axs[1][0].text(.2,.9,r'$Q_q = %.2f$' %(Q) , ha='center', va='center', transform= axs[1][0].transAxes, bbox = {'facecolor': 'white', 'pad': 5})

    im = axs[0][1].imshow(HT, interpolation = 'none', extent=[-L/2.01, L/2.01, np.min(u),np.max(u)], aspect = 'auto', origin = 'lower',cmap = 'binary')
    im.set_clim(vmin=cmin, vmax=cmax)
    axs[0][1].text(.8,.9,r'$T = 5$'  , ha='center', va='center', transform= axs[0][1].transAxes, bbox = {'facecolor': 'white', 'pad': 5})
    axs[0][1].plot(rT,vT, 'ro')
    p,q, Q = GetPQ(psiT)
    axs[0][1].plot(p,q, 'c^', linewidth = 5)
    axs[0][1].text(.2,.9,r'$Q_q = %.2f$' %(Q) , ha='center', va='center', transform= axs[0][1].transAxes, bbox = {'facecolor': 'white', 'pad': 5})

    im = axs[1][1].imshow(HT_, interpolation = 'none', extent=[-L/2.01, L/2.01, np.min(u),np.max(u)], aspect = 'auto', origin = 'lower',cmap = 'binary')
    im.set_clim(vmin=cmin, vmax=cmax)
    axs[1][1].plot(rT_,vT_, 'ro')
    p,q, Q = GetPQ(psiT_)
    axs[1][1].plot(p,q, 'c^', linewidth = 5)
    axs[1][1].text(.2,.9,r'$Q_q = %.2f$' %(Q) , ha='center', va='center', transform= axs[1][1].transAxes, bbox = {'facecolor': 'white', 'pad': 5})

    im = axs[0][2].imshow(HT2, interpolation = 'none', extent=[-L/2.01, L/2.01, np.min(u),np.max(u)], aspect = 'auto', origin = 'lower',cmap = 'binary')
    im.set_clim(vmin=cmin, vmax=cmax)
    axs[0][2].text(.8,.9,r'$T = 10$'  , ha='center', va='center', transform= axs[0][2].transAxes, bbox = {'facecolor': 'white', 'pad': 5})
    axs[0][2].plot(rT2,vT2, 'ro')
    p,q, Q = GetPQ(psiT2)
    axs[0][2].plot(p,q, 'c^', linewidth = 5)
    axs[0][2].text(.2,.9,r'$Q_q = %.2f$' %(Q) , ha='center', va='center', transform= axs[0][2].transAxes, bbox = {'facecolor': 'white', 'pad': 5})

    im = axs[1][2].imshow(HT2_, interpolation = 'none', extent=[-L/2.01, L/2.01, np.min(u),np.max(u)], aspect = 'auto', origin = 'lower',cmap = 'binary')
    im.set_clim(vmin=cmin, vmax=cmax)
    axs[1][2].plot(rT2_,vT2_, 'ro')
    p,q, Q = GetPQ(psiT2_)
    axs[1][2].plot(p,q, 'c^', linewidth = 5)
    axs[1][2].text(.2,.9,r'$Q_q = %.2f$' %(Q) , ha='center', va='center', transform= axs[1][2].transAxes, bbox = {'facecolor': 'white', 'pad': 5})

    axs[0][0].set_xticklabels([])
    axs[0][1].set_yticklabels([])
    axs[0][1].set_xticklabels([])
    axs[1][1].set_yticklabels([])
    axs[0][2].set_yticklabels([])
    axs[0][2].set_xticklabels([])
    axs[1][2].set_yticklabels([])

    axs[1][0].set_xlabel(r'$q \, [L]$')
    axs[1][1].set_xlabel(r'$q \, [L]$')
    axs[1][2].set_xlabel(r'$q \, [L]$')

    axs[0][0].set_ylabel(r'$p \, [\hbar / L]$')
    axs[1][0].set_ylabel(r'$p \, [\hbar / L]$')

    x0 = axs[0][0].get_position().x0
    y0 = axs[0][0].get_position().y1
    h = .015
    w = axs[0][2].get_position().x1 - x0
    cbar_ax = fig.add_axes([x0,y0, w, h])
    cb = fig.colorbar(im, cax=cbar_ax, orientation = "horizontal")
    cbar_ax.xaxis.set_ticks_position("top")

    plt.subplots_adjust(wspace=0, hspace=0)

    return fig

def EvolveWaveFuncs(psi_, Phi, r_, v_, a):
    T = 0
    psi = psi_.copy()

    r = r_
    v = v_
    while(T < Tfinal):

        r += v*dt/2.
        v += a(r) * dt
        r += v*dt/2.

        # update position half-step
        psi_k = sp.fft(psi)
        psi_k *= np.exp(-1j*dt*hbar*(kx**2)/(4.))
        psi = sp.ifft(psi_k)

        # update momentum full-step
        psi *= np.exp(-1j*dt*Phi/hbar)

        # update position half-step
        psi_k = sp.fft(psi)
        psi_k *= np.exp(-1j*dt*hbar*(kx**2)/(4.))
        psi = sp.ifft(psi_k)

        T += dt

    return psi, r, v

def GetICs():
    sig = L/20.

    arg = (x-mu)**2 / sig**2
    psi = np.exp(-arg)
    psi /= np.sqrt(np.sum(np.abs(psi)**2)*dx )

    return psi, mu, 0.

if __name__ == "__main__":
    u_.orientPlot()

    psi, r, v = GetICs()
    psiT, rT, vT = EvolveWaveFuncs(psi, V, r, v, F)
    psiT_, rT_, vT_ = EvolveWaveFuncs(psi, V_, r, v, F_)
    psiT2, rT2, vT2 = EvolveWaveFuncs(psiT, V, rT, vT, F)
    psiT2_, rT2_, vT2_ = EvolveWaveFuncs(psiT_, V_, rT_, vT_, F_)

    fig = makeFig(psi, psiT, psiT_ , psiT2, psiT2_,  r,v,rT,vT, rT2, vT2 ,  rT_,vT_, rT2_, vT2_)
    fig.savefig("../Figs/" + "QuantumErrors.pdf",bbox_inches = 'tight')

max_stars_repo_name: andillio/CHiMES

max_issues_repo_name: andillio/CHiMES

max_forks_repo_name: andillio/CHiMES

reason: import numpy,import scipy


text:
# Copyright (c) Facebook, Inc. and its affiliates.

import copy
import numpy as np

from fairmotion.core import motion as motion_class
from fairmotion.ops import conversions, math, quaternion

def append(motion1, motion2):
    """
    Combines two motion sequences into one. motion2 is appended to motion1.
    The operation is not done in place.

    Note that the operation places the sequences next to each other without
    attempting to blend between the poses. To interpolate between the end of
    motion1 and start of motion2, use the `append_and_blend` operation.

    Args:
        motion1, motion2: Motion sequences to be combined.
    """
    assert isinstance(motion1, motion_class.Motion)
    assert isinstance(motion2, motion_class.Motion)
    assert motion1.skel.num_joints() == motion2.skel.num_joints()

    combined_motion = copy.deepcopy(motion1)
    combined_motion.name = f"{motion1.name}+{motion2.name}"
    combined_motion.poses.extend(motion2.poses)

    return combined_motion

def append_and_blend(motion1, motion2, blend_length=0):
    assert isinstance(motion1, motion_class.Motion)
    assert isinstance(motion2, motion_class.Motion)
    assert motion1.skel.num_joints() == motion2.skel.num_joints()

    combined_motion = copy.deepcopy(motion1)
    combined_motion.name = f"{motion1.name}+{motion2.name}"

    if motion1.num_frames() == 0:
        for i in range(motion2.num_frames()):
            combined_motion.poses.append(motion2.poses[i])
            if hasattr(motion1, "velocities"):
                combined_motion.velocities.append(motion2.velocities[i])
        return combined_motion

    frame_target = motion2.time_to_frame(blend_length)
    frame_source = motion1.time_to_frame(motion1.length() - blend_length)

    # Translate and rotate motion2 to location of frame_source
    pose1 = motion1.get_pose_by_frame(frame_source)
    pose2 = motion2.get_pose_by_frame(0)

    R1, p1 = conversions.T2Rp(pose1.get_root_transform())
    R2, p2 = conversions.T2Rp(pose2.get_root_transform())

    # Translation to be applied
    dp = p1 - p2
    dp = dp - math.projectionOnVector(dp, motion1.skel.v_up_env)
    axis = motion1.skel.v_up_env

    # Rotation to be applied
    Q1 = conversions.R2Q(R1)
    Q2 = conversions.R2Q(R2)
    _, theta = quaternion.Q_closest(Q1, Q2, axis)
    dR = conversions.A2R(axis * theta)

    motion2 = translate(motion2, dp)
    motion2 = rotate(motion2, dR)

    t_total = motion1.fps * frame_source
    t_processed = 0.0
    poses_new = []
    for i in range(motion2.num_frames()):
        dt = 1 / motion2.fps
        t_total += dt
        t_processed += dt
        pose_target = motion2.get_pose_by_frame(i)
        # Blend pose for a moment
        if t_processed <= blend_length:
            alpha = t_processed / float(blend_length)
            pose_source = motion1.get_pose_by_time(t_total)
            pose_target = blend(pose_source, pose_target, alpha)
        poses_new.append(pose_target)

    del combined_motion.poses[frame_source + 1 :]
    for i in range(len(poses_new)):
        combined_motion.add_one_frame(0, copy.deepcopy(poses_new[i].data))

    return combined_motion

def blend(pose1, pose2, alpha=0.5):
    assert 0.0 <= alpha <= 1.0
    pose_new = copy.deepcopy(pose1)
    for j in range(pose1.skel.num_joints()):
        R0, p0 = conversions.T2Rp(pose1.get_transform(j, local=True))
        R1, p1 = conversions.T2Rp(pose2.get_transform(j, local=True))
        R, p = math.slerp(R0, R1, alpha), math.lerp(p0, p1, alpha)
        pose_new.set_transform(j, conversions.Rp2T(R, p), local=True)
    return pose_new

def transform(motion, T, local=False):
    """
    Apply transform to all poses of a motion sequence. The operation is done
    in-place.

    Args:
        motion: Motion sequence to be transformed
        T: Transformation matrix of shape (4, 4) to be applied to poses of
            motion
        local: Optional; Set local=True if the transformations are to be
            applied locally, relative to parent of each joint.
    """
    for pose_id in range(len(motion.poses)):
        R0, p0 = conversions.T2Rp(motion.poses[pose_id].get_root_transform())
        R1, p1 = conversions.T2Rp(T)
        if local:
            R, p = np.dot(R0, R1), p0 + np.dot(R0, p1)
        else:
            R, p = np.dot(R1, R0), p0 + p1
        motion.poses[pose_id].set_root_transform(
            conversions.Rp2T(R, p), local=False,
        )
    return motion

def translate(motion, v, local=False):
    """
    Apply translation to motion sequence.

    Args:
        motion: Motion sequence to be translated
        v: Array of shape (3,) indicating translation vector to be applied to
            all poses of motion sequence
        local: Optional; Set local=True if the translation is to be applied
            locally, relative to root position.
    """
    return transform(motion, conversions.p2T(v), local)

def rotate(motion, R, local=False):
    return transform(motion, conversions.R2T(R), local)

def cut(motion, frame_start, frame_end):
    """
    Returns motion object with poses from [frame_start, frame_end) only. The
    operation is not done in-place.

    Args:
        motion: Motion sequence to be cut
        frame_start, frame_end: Frame number range that defines the boundary of
            motion to be cut. Pose at frame_start is included, and pose at
            frame_end is excluded in the returned motion object
    """
    cut_motion = copy.deepcopy(motion)
    cut_motion.name = f"{motion.name}_{frame_start}_{frame_end}"
    cut_motion.poses = motion.poses[frame_start:frame_end]

    return cut_motion

def resample(motion, fps):
    """
    Upsample/downsample frame rate of motion object to `fps` Hz. For
    upsampling, poses are interpolated using `Pose.interpolate` method to
    fill in the gaps.

    Args:
        motion: Motion sequence to be resampled
        fps: Frequency of motion desired
    """
    poses_new = []

    dt = 1.0 / fps
    t = 0
    while t < motion.fps * len(motion.poses):
        pose = motion.get_pose_by_time(t)
        pose.skel = motion.skel
        poses_new.append(pose)
        t += dt

    motion.poses = poses_new
    motion.fps = fps
    return motion

def position_wrt_root(motion):
    """
    Returns position of joints with respect to the root, for all poses in the
    motion sequence.
    """
    matrix = motion.to_matrix(local=False)
    # Extract positions
    matrix = matrix[:, :, :3, 3]
    # Subtract root position from all joint positions
    matrix = matrix - matrix[:, np.newaxis, 0]
    return matrix

max_stars_repo_name: CristianNajeraL/fairmotion

max_issues_repo_name: CristianNajeraL/fairmotion

max_forks_repo_name: CristianNajeraL/fairmotion


text:
# ***************************************************************
# Copyright (c) 2021 Jittor. All Rights Reserved.
# Maintainers:
#     Zheng-Ning Liu <lzhengning@gmail.com>
#     Dun Liang <randonlang@gmail.com>.
#
# This file is subject to the terms and conditions defined in
# file 'LICENSE.txt', which is part of this source code package.
# ***************************************************************

import unittest
import jittor as jt
import numpy as np

def topk(input, k, dim=None, largest=True, sorted=True):
    if dim is None:
        dim = -1
    if dim < 0:
        dim += input.ndim

    transpose_dims = [i for i in range(input.ndim)]
    transpose_dims[0] = dim
    transpose_dims[dim] = 0
    input = input.transpose(transpose_dims)
    index, values = jt.argsort(input, dim=0, descending=largest)
    indices = index[:k]
    values = values[:k]
    indices = indices.transpose(transpose_dims)
    values = values.transpose(transpose_dims)
    return [values, indices]

def knn(x, k):
    inner = -2 * jt.nn.bmm(x.transpose(0, 2, 1), x)
    xx = jt.sum(x ** 2, dim=1, keepdims=True)
    distance = -xx - inner - xx.transpose(0, 2, 1)
    return topk(distance, k=k, dim=-1)

class TestKnnOp(unittest.TestCase):
    def test_knn(self):
        jt_a = jt.randn(32,512,3)
        a1, b1 = jt.misc.knn(jt_a, jt_a, 16)
        a2, b2 = knn(jt_a.transpose(0,2,1), 16)
        a2 *= -1
        np.testing.assert_allclose(a1.data, a2.data, atol=1e-4)

        if jt.has_cuda:
            with jt.flag_scope(use_cuda=1):
                jt_a = jt.randn(32,512,3)
                a1, b1 = jt.misc.knn(jt_a, jt_a, 16)
                a2, b2 = knn(jt_a.transpose(0,2,1), 16)
                a2 *= -1
                np.testing.assert_allclose(a1.data, a2.data, atol=1e-4)

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

max_stars_repo_name: Jittor/Jittor

max_issues_repo_name: Jittor/Jittor

max_forks_repo_name: Jittor/Jittor


text:
import torch.nn.functional as F
from torch.nn import Linear
from torch.nn import BCELoss, Module
from torch.optim import Adam
import pandas as pd
from torch import Tensor, save
import numpy as np
import pickle
from sklearn.model_selection import train_test_split
from sklearn.preprocessing import MinMaxScaler
from sklearn.metrics import f1_score
from Model import BreastCancerPrognosisModel

model = BreastCancerPrognosisModel()
criterion = BCELoss()
optimizer = Adam(model.parameters())

# Headers for the dataset
columns = ['ID', 'CLUMP_THICKNNESS', 'UNIFORMITY_OF_CELL_SIZE', 'UNIFORMITY_OF_CELL_SHAPE', 'MARGINAL_ADHESION',
		   'SINGLE_EPITHELIAL_CELL_SIZE', 'BARE_NUCLEI', 'BLAND_CHROMATIN', 'NORMAL_NUCLEI', 'MITOSIS', 'TARGET_CLASS']

raw_data = pd.read_csv('breast-cancer-wisconsin.data', header=None)
raw_data.columns = columns

# Fill missing values
raw_data = raw_data.replace('?', np.nan)
raw_data = raw_data.fillna(raw_data.median())

# Map benign to 0, malignant to 1
raw_data['TARGET_CLASS'] = raw_data['TARGET_CLASS'].map({2: 0, 4: 1})
print(raw_data.tail())
# Drop the ID column
raw_data = raw_data.loc[:, raw_data.columns != 'ID']
columns.remove('ID')

# Scale the values between 0-1
scaler = MinMaxScaler()
y = raw_data['TARGET_CLASS']
raw_data = scaler.fit_transform(raw_data.drop(['TARGET_CLASS'], axis=1))
columns.remove('TARGET_CLASS')
raw_data = pd.DataFrame(raw_data, columns=columns)
raw_data = pd.concat([raw_data, y], axis=1)
columns.append('TARGET_CLASS')
train, test = train_test_split(raw_data, test_size=0.15)
train = pd.DataFrame(data=train, columns=columns)
test = pd.DataFrame(data=test, columns=columns)

# Model training and test data
X_train = train.loc[:, train.columns != 'TARGET_CLASS']
Y_train = train['TARGET_CLASS']

X_test = test.loc[:, test.columns != 'TARGET_CLASS']
Y_test = test['TARGET_CLASS']
print('Shape={}'.format(X_train.values.shape))
n_features_tr = X_train.values.shape[1]
n_samples_tr = X_train.values.shape[0]

X_train = Tensor(X_train.values)
Y_train = Tensor(Y_train.values).reshape(-1, 1)
X_test = Tensor(X_test.values)
Y_test = Tensor(Y_test.values).reshape(-1, 1)

N_epochs = 500

for i in range(1, N_epochs + 1):
	optimizer.zero_grad()

	Y_hat = model.forward(X_train)

	loss = criterion(Y_hat, Y_train)
	loss.backward()
	optimizer.step()
	print('Loss = {}, epoch = {}'.format(loss.item(), i))

# Testing
model.eval()
print(X_test.shape)
Y_hat = np.round(model.forward(X_test).detach().numpy())
accuracy = f1_score(y_pred=Y_hat, y_true=Y_test)
print('Accuracy: {}'.format(accuracy))

# Saving model and scaler
save(model.state_dict(), 'breast_cancer_prognosis.pt')
with open('scaler.pkl', 'wb') as f:
	pickle.dump(scaler, f)


text:
import os
import pickle
from PIL import Image
import torch
import torch.utils.data as data
import numpy as np
from PIL import ImageFile
ImageFile.LOAD_TRUNCATED_IMAGES = True

def conversion(img):
    #img: PIL format
    img=np.array(img)
    h,w,c=img.shape
    if h>w:
        img2=np.ones((h,h,c),dtype='uint8')*255
        start=(h-w)//2
        img2[:,start:start+w,:]=img
    elif w>h:
        img2=np.ones((w,w,c),dtype='uint8')*255
        start=(w-h)//2
        img2[start:start+h,:,:]=img
    else:
        img2=img
    img2=Image.fromarray(img2)
    return img2

class retrieval_dataset(data.Dataset):
    def __init__(self,root_path,transform = None,crop_aug=False):
        #import pdb; pdb.set_trace()
        self.root=root_path
        self.image_list=os.listdir(self.root)
        self.image_list.sort()
        self.transform=transform

    def __getitem__(self,idx):
        #try:
        image=Image.open(os.path.join(self.root,self.image_list[idx])).convert('RGB')
        #except:
            #continue
        image=conversion(image)
        output_imgs=self.transform(image)
        return output_imgs, self.image_list[idx]

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


text:
import matplotlib.pyplot as plt
import numpy as np
import pandas as pd
from tabulate import tabulate
from mpl_toolkits.mplot3d import Axes3D
from sklearn.preprocessing import StandardScaler

"""
The dataset contains many values which do not make any sense.
As an example, a blood pressure that is >1000 or <=0 is unrealistic!
"""
def clean_data(df):
    return df[ (df.BloodPressure < 1000)
                & (df.BloodPressure > 0)
                & (df.SkinThickness > 0)
                & (df.Glucose > 0)
                & (df.BMI > 0)
                & (df.Insulin > 0) ]

def plot_3D(df, label1, label2, title1, title2):
    #---- Parameters which are the same for both figures
    fig = plt.figure(figsize=(20,8))
    x = np.array(df['Age'])
    y = np.array(df['BloodPressure'])
    z = np.array(df['BMI'])

    #---- First subplot
    ax = fig.add_subplot(121, projection='3d')

    ax.scatter(x, y, z, marker="s", c=label1, s=25, cmap="RdBu")
    ax.set_title("%s" % title1, fontsize = 16)
    ax.set_xlabel('Age')
    ax.set_ylabel('BloodPressure')
    ax.set_zlabel('BMI')
    ax.view_init(15) # rotation of the 3D plot

    #---- Second subplot
    ax = fig.add_subplot(122, projection='3d')

    ax.scatter(x, y, z, marker="s", c=label2, s=25, cmap="RdBu")
    ax.set_title("%s" % title2, fontsize = 16)
    ax.set_xlabel('Age')
    ax.set_ylabel('BloodPressure')
    ax.set_zlabel('BMI')
    ax.view_init(15) # rotation of the 3D plot

    plt.show()

def plot_TSNE(df, label1, label2, title1, title2):
    fig, ax = plt.subplots(2, sharex=True, figsize=(8,10))

    ax[0].scatter(df[:,0], df[:,1], c = label1, cmap="RdBu")
    ax[0].set_title('t-SNE using labels from %s Clustering' % title1, fontsize = 16)

    ax[1].scatter(df[:,0], df[:,1], c = label2, cmap="RdBu")
    ax[1].set_title('t-SNE using labels from %s Clustering' % title2, fontsize = 16)

    fig.subplots_adjust(hspace=0.2)
    plt.show()

def plot_overview(df):
    # make a figure twice as wide as it is tall
    plt.figure(figsize=plt.figaspect(0.5))

    # save the 3 different plots created as 'figs'
    figs = plt.plot(df)
    # iterate over 'figs' and create a label for each one from the column names of t1_t2_clean
    plt.legend(iter(figs), list(df))

    plt.show()

def plot_boxplot(df, label):
    # Create a figure instance
    fig = plt.figure(figsize=(14,7))

    # Create the subplots, sharing the y axis
    ax1 = plt.subplot(121)
    plt.boxplot(df[label == 0].values)
    plt.setp(ax1, xticklabels=list(df))
    plt.title('Class I', fontsize=18)

    ax2 = plt.subplot(122, sharey=ax1)
    plt.boxplot(df[label == 1].values)
    plt.setp(ax2, xticklabels=list(df))
    plt.title('Class II', fontsize=18)

    plt.show()

def print_tab(df, label):
    # Use the labels of KMeans, as they gave a better separation.
    mean1 = df[label == 0].mean(axis=0)
    std1 = df[label == 0].std(axis=0)

    mean2 = df[label == 1].mean(axis=0)
    std2 = df[label == 1].std(axis=0)

    # Print the values in a table
    print( tabulate([['Class I', "%i +/- %i"%(mean1[0],std1[0]), "%i +/- %i"%(mean1[1],std1[1]), "%i +/- %i"%(mean1[2],std1[2])],
                     ['Class II', "%i +/- %i"%(mean2[0],std2[0]), "%i +/- %i"%(mean2[1],std2[1]), "%i +/- %i"%(mean2[2],std2[2])]],
                    headers=list(df), tablefmt="fancy_grid") )


text:
import numpy as np
import os
import matplotlib.pyplot as plt
import sys

try:
    plt.style.use('presentation')
except:
    pass

class SSP(object):
    def __init__(self, filename):
        self.base, self.name = os.path.split(filename)
        self.data = read_ssp_output(filename)
        self.load_best()

    def load_best(self):
        self.best = self.data[np.argmin(self.data['fit'])]

    def marginalize(self, marg_attr, value='best'):
        vals = ['Av', 'dmod', 'lage', 'logZ']
        assert marg_attr in vals, 'must choose between {}'.format(vals)

        if value == 'best':
            value = self.best[marg_attr]

        inds, = np.nonzero(self.data[marg_attr]==value)

        new_attr = '{}_inds'.format(marg_attr)
        self.__setattr__(new_attr, inds)
        return inds

    def age_plot(self):
        [self.marginalize(v) for v in ['Av', 'dmod', 'logZ']]
        inds = list(set(self.Av_inds) & set(self.dmod_inds) & set(self.logZ_inds))

        mfit = np.max(self.data['fit'])
        fit = 1 - self.data['fit'] / mfit
        fig, ax = plt.subplots()

        ax.set_yticks([])
        ax.plot(self.data['lage'], fit, 'o', alpha=0.3, color='gray', mec='none')
        bestage = self.data['lage'][inds]
        bestfit = fit[inds]
        isort = np.argsort(bestage)
        ax.plot(bestage[isort], bestfit[isort], 'o', ms=8, color='k')
        ax.set_title(translate('Av={}\ dmod={}\ logZ={}'.format(self.best['Av'],
                                                             self.best['dmod'],
                                                             self.best['logZ'])))
        ax.set_xlabel(translate('lage'))
        return fig, ax

    def ssp_plot0(self, marg_attr):
        # not sure this was worth doing...
        vals = ['Av', 'dmod', 'lage', 'logZ', 'sfr']

        try:
            vals.pop(vals.index(marg_attr))
        except ValueError:
            print('must choose between {}'.format(vals))
            return

        # marginalize
        inds = list(set(self.marginalize(marg_attr)) & set(self.marginalize('Av')))
        data = self.data[inds]

        fig, axs = plot_config()
        for i, j in itertools.product(range(4), range(4)):
            if j > i:
                continue
            y = vals[i]
            x = vals[j]
            ax = axs[i, j]
            ax.set_ylabel(translate(vals[i]))
            ax.set_xlabel(translate(vals[j]))
            if i == j:
                ax.plot(data[x], data['fit'])
            # plot x vs y
            ax.scatter(data[x], data[y], c=data['fit'], marker='s', s=40,
                       cmap=plt.cm.Spectral, alpha=.5, edgecolors='none')
            ax.axhline(self.best[y])
            ax.axvline(self.best[x])

        axs[0, -1].annotate(translate('{}={}'.format(marg_attr, self.best[marg_attr])),
                            (1, 1), ha='right', fontsize=24, color='k')

        return fig, axs

def translate(string):
    if 'log' in string:
        string = string.replace('log', '\log\ ')
    if 'lage' in string:
        string = string.replace('la', '\log\ A')
    if 'dmod' in string:
        string = string.replace('dmod', '\mu')
    if 'Av' in string:
        string = string.replace('Av', 'A_V')
    if 'sfr' in string:
        string = string.replace('sfr', 'SFR')
    return r'${}$'.format(string)

def read_ssp_output(filename):
    colnames = ['Av', 'IMF', 'dmod', 'lage', 'logZ', 'fit', 'sfr', 'sfrperr',
                'sfrmerr']
    return np.genfromtxt(filename, skip_header=10, skip_footer=1, names=colnames)

def plot_config():
    """
    Make an 4x4 figure with top corner frames off and ticks on the left
    and bottom edges
    """
    fig, axs = plt.subplots(4, 4, figsize=(9.7, 9.7))
    fig.subplots_adjust(left=0.1, bottom=0.1, right=0.96, top=0.96,
                        wspace=0.05, hspace=0.05)
    for i in range(4):
        for j in range(4):
            ax = axs[i, j]
            # only have y ticks on left axes (axs[:, 0])
            ax.set_yticks([])
            # turn off top right frames
            if j > i:
                ax.set_frame_on(False)
            # only have x ticks on bottom axes (axs[-1, 0])
            if i != 4:
                ax.set_xticks([])
    return fig, axs

if __name__ == '__main__':
    ssp = SSP(sys.argv[1])
    fig, ax = ssp.age_plot()
    plt.savefig(sys.argv[1] + '.png')

max_stars_repo_name: kmcquinn/match

max_issues_repo_name: kmcquinn/match

max_forks_repo_name: kmcquinn/match


text:
"""
A Keras data generator, suitable for use with keras.model.fit_generator()

This generator assumes the default folder structure of unzipped Inria data:

Inria/train/images/*.tif
           /gt/*.tif
     /val/images/*.tif
         /gt/*.tif
     /test/images/*.tif
"""

import os
from glob import glob
import rasterio
import numpy as np
from keras.utils import to_categorical

class InriaGenerator(object):
    """
    Data generator for the Inria Aerial image data set
    """

    def __init__(self, datadir, tile_size=(512, 512)):

        self.datadir = datadir
        self.tile_size = np.array(tile_size)
        self.images = [os.path.basename(imgpath) for imgpath in glob(os.path.join(self.datadir, 'images', '*.tif'))]

    def _random_image(self):
        return np.random.choice(self.images)

    def _random_tile(self, image_name):
        """
        Choose a random tile fully contained within the given image / label
        """
        with rasterio.open(os.path.join(self.datadir, 'images', image_name)) as img_src:
            img_shape = img_src.shape
            r_start = np.random.randint(0, img_shape[1] - self.tile_size[1])
            c_start = np.random.randint(0, img_shape[0] - self.tile_size[0])
            window = ((r_start, r_start+self.tile_size[1]), (c_start, c_start+self.tile_size[0]))
            img = np.rollaxis(img_src.read((1,2,3), window=window) / 255, 0, 3)
            with rasterio.open(os.path.join(self.datadir, 'gt', image_name)) as lab_src:
                lab = lab_src.read((1, ), window=window)[0] // 255
            lab_one_hot = to_categorical(lab.ravel(), 2)
            return img[None], lab_one_hot[None]

    def _fully_tile(self, image_name):
        with rasterio.open(os.path.join(self.datadir, 'images', image_name)) as img_src:
            with rasterio.open(os.path.join(self.datadir, 'gt', image_name)) as lab_src:
                img_shape = img_src.shape
                for r_start in range(0, img_shape[1] - self.tile_size[1], self.tile_size[1]):
                    for c_start in range(0, img_shape[0] - self.tile_size[0], self.tile_size[0]):
                        window = ((r_start, r_start+self.tile_size[1]), (c_start, c_start+self.tile_size[0]))
                        img = np.rollaxis(img_src.read((1,2,3), window=window) / 255, 0, 3)
                        lab = lab_src.read((1, ), window=window)[0] // 255
                        lab_one_hot = to_categorical(lab.ravel(), 2)
                        yield img[None], lab_one_hot[None]

    def random_generator(self, batch_size=5):
        while True:
            img_batch = []
            lab_batch = []
            for _ in range(batch_size):
                img, lab = self._random_tile(self._random_image())
                img_batch.append(img)
                lab_batch.append(lab)
            yield np.concatenate(img_batch), np.concatenate(lab_batch)


text:
import dolfin as df
import pickle
import numpy as np
import finmag.energies.demag.fk_demag as fk
import matplotlib.pyplot as plt
from table_printer import TablePrinter
from finmag.util.helpers import vector_valued_function
from finmag.native.llg import compute_bem_fk
from aeon import default_timer

def _prepare_demag_object(S3, m, Ms, unit_length, s_param, solver, p_param, preconditioner, bem, boundary_to_global):
    demag = fk.FKDemag()
    if bem is not None:
        # Use the precomputed BEM if it was provided.
        demag.precomputed_bem(bem, boundary_to_global)
    demag.parameters[s_param] = solver
    demag.parameters[p_param] = preconditioner
    demag.setup(S3, m, Ms, unit_length)
    return demag

def _loaded_results_table(timed_method_name, results, solvers, preconditioners):
    print "Time shown for {} in s.".format(timed_method_name)
    print "Results loaded from file.\n"

    table = TablePrinter(preconditioners, xlabels_width_min=10)
    for solver in solvers:
        table.new_row(solver)
        for prec in preconditioners:
            try:
                time = results[solver][prec]
            except KeyError:
                time = "-"
            table.new_entry(time)

    default = results['default']['default']
    fastest = fastest_run(results, solvers)

    print "\n\nDefault combination ran in {:.3} s.".format(default)
    print "Fastest combination {}/{} ran in {:.3} s.".format(fastest[0], fastest[1], fastest[2])
    print "That is an {:.1%} improvement.\n".format(1 - fastest[2] / default)

class ToleranceNotReachedException(Exception):
    pass

def create_measurement_runner(S3, m, Ms, unit_length, H_expected=None, tol=1e-3, repeats=10, bem=None, boundary_to_global=None):
    def runner(timed_method_name, s_param, solvers, p_param, preconditioners, skip=[], full_timings_log="timings.txt", results_cache=None):
        r = np.ma.zeros((len(solvers), len(preconditioners)))
        results, failed = {}, []

        try:
            with open(str(results_cache), "r") as f:
                results = pickle.load(f)
                _loaded_results_table(timed_method_name, results, solvers, preconditioners)
                return results, failed
        except IOError:
            print "No recorded results found. Will run benchmarks."
        print "Time shown for {} in s.\n".format(timed_method_name)

        for timer in (default_timer, fk.fk_timer):
            timer.reset()
        log = open(full_timings_log, "w")

        table = TablePrinter(preconditioners)
        for solver in solvers:
            table.new_row(solver)
            #print "\nsolving with {}\n".format(solver)
            results_for_this_solver = {}
            # Compute the demagnetising field with the current solver and each of the preconditioners.
            for prec in preconditioners:
                #print "\nprec {} entry {}\n".format(prec, table.entry_counter)
                if (solver, prec) in skip:
                    table.new_entry("s")
                    continue
                demag = _prepare_demag_object(S3, m, Ms, unit_length, s_param, solver, p_param, prec, bem, boundary_to_global)
                checked_result = False
                try:
                    for _ in xrange(repeats):  # Repeat to average out little fluctuations.
                        H = demag.compute_field()  # This can fail with some method/preconditioner combinations.
                        if not checked_result:
                            # By checking here, we can avoid computing the
                            # field ten times if it's wrong.
                            checked_result = True
                            max_diff = np.max(np.abs(H - H_expected))
                            if max_diff > tol:
                                # We need to raise an exception, otherwise the 'else' clause below is
                                # executed and the measurement will be recorded twice.
                                raise ToleranceNotReachedException
                except ToleranceNotReachedException:
                    table.new_entry("x")
                    failed.append({'solver': solver, 'preconditioner': prec, 'message': "error {:.3} higher than allowed {:.3}"})
                except RuntimeError as e:
                    default_timer.get("compute_field", "FKDemag").stop()
                    table.new_entry("x")
                    failed.append({'solver': solver, 'preconditioner': prec, 'message': e.message})
                else:
                    measured_time = fk.fk_timer.time_per_call(timed_method_name, "FKDemag")
                    table.new_entry(measured_time)
                    results_for_this_solver[prec] = measured_time
                log.write("\nTimings for {} with {}.\n".format(solver, prec))
                log.write(fk.fk_timer.report() + "\n")
                default_timer.reset()
                fk.fk_timer.reset()
                del(demag)
            results[solver] = results_for_this_solver

        if results_cache is not None:
            with open(results_cache, "w") as f:
                pickle.dump(results, f)

        log.close()

        default = results['default']['default']
        fastest = fastest_run(results, solvers)

        print "\nDefault combination ran in {:.3} s.".format(default)
        print "Fastest combination {}/{} ran in {:.3} s.".format(fastest[0], fastest[1], fastest[2])
        print "That is an {:.1%} improvement.".format(1 - fastest[2] / default)

        return results, failed
    return runner

def run_measurements(m, mesh, unit_length, tol, repetitions=10, H_expected=None, name="", skip=[]):
    S3 = df.VectorFunctionSpace(mesh, "CG", 1)
    m = vector_valued_function(m, S3)
    Ms = 1
    bem, boundary_to_global = compute_bem_fk(df.BoundaryMesh(mesh, 'exterior', False))

    if H_expected is not None:
        H = vector_valued_function(H_expected, S3)
        H_expected = H.vector().array()
    else:
        # use default/default as reference then.
        demag = fk.FKDemag()
        demag.precomputed_bem(bem, boundary_to_global)
        demag.setup(S3, m, Ms, unit_length)
        H_expected = demag.compute_field()
        del(demag)

    if name == "":
        pass
    else:
        name = name + "_"

    runner = create_measurement_runner(S3, m, Ms, unit_length,
                                       H_expected, tol, repetitions,
                                       bem, boundary_to_global)

    solvers = [s[0] for s in df.krylov_solver_methods()]
    preconditioners = [p[0] for p in df.krylov_solver_preconditioners()]

    results_1, failed_1 = runner("first linear solve",
                                 "phi_1_solver", solvers,
                                 "phi_1_preconditioner", preconditioners,
                                 skip,
                                 "{}timings_log_1.txt".format(name),
                                 "{}results_1.pickled".format(name))

    results_2, failed_2 = runner("second linear solve",
                                 "phi_2_solver", solvers,
                                 "phi_2_preconditioner", preconditioners,
                                 skip,
                                 "{}timings_log_2.txt".format(name),
                                 "{}results_2.pickled".format(name))

    return solvers, preconditioners, results_1, failed_1, results_2, failed_2

def fastest_run(results, solvers):
    return _extremum(results, solvers, min)

def fastest_runs(results, solvers):
    return _extremum(results, solvers, min, per_solver=True)

def slowest_run(results, solvers):
    return _extremum(results, solvers, max)

def slowest_runs(results, solvers):
    return _extremum(results, solvers, max, per_solver=True)

def _extremum(results, solvers, func, per_solver=False):
    fastest = []
    for solver in solvers:
        p = func(results[solver].keys(), key=lambda x: results[solver][x])
        fastest.append((solver, p, results[solver][p]))
    if per_solver:
        return fastest
    return func(fastest, key=lambda x: x[2])

def _runs_with_preconditioner(results, solvers, preconditioner):
    indices = []
    times = []
    for i, solver in enumerate(solvers):
        try:
            times.append(results[solver][preconditioner])
        except KeyError:
            pass  # Run with this combination didn't succeed.
        else:
            indices.append(i)  # Will allow us to skip bars in the column chart.
    return times, np.array(indices)

def column_chart(results, solvers, offset=None):
    times_slowest = [r[2] for r in slowest_runs(results, solvers)]
    times_fastest, names_fastest = zip(* [(r[2], r[1]) for r in fastest_runs(results, solvers)])
    times_without_prec, iwop = _runs_with_preconditioner(results, solvers, "none")
    times_with_default_prec, idp = _runs_with_preconditioner(results, solvers, "default")

    fig = plt.figure()
    ax = fig.add_subplot(111)

    width = 0.2
    ind = np.arange(len(results))

    ax.axhline(results['default']['default'], color=(0.3, 0.3, 0.3), ls=":", zorder=0)
    ax.bar(width + idp, times_with_default_prec, width, color=(0.3, 0.3, 0.3), label="default prec.")
    ax.bar(2 * width + iwop, times_without_prec, width, color=(0.8, 0.8, 0.8), label="no prec.")
    ax.bar(3 * width + ind, times_slowest, width, color="red", label="slowest prec.")
    rects_fastest = ax.bar(ind, times_fastest, width, color="green", label="fastest prec.")
    for i, rect in enumerate(rects_fastest):
        height = rect.get_height()
        offset = offset if offset is not None else 1.05 * height
        ax.text(rect.get_x() + rect.get_width() / 2.0, height + offset, names_fastest[i], ha='center', va='bottom', rotation=90)

    ax.set_xlabel("Solver")
    ax.set_ylabel("Time (s)")
    ax.legend()
    ax.set_xticks(ind + 2 * width)
    xtickNames = plt.setp(ax, xticklabels=solvers)
    plt.setp(xtickNames, rotation=45)
    return fig

max_stars_repo_name: davidcortesortuno/finmag

max_issues_repo_name: davidcortesortuno/finmag

max_forks_repo_name: davidcortesortuno/finmag