text:
import os
import sys
import pandas as pd
import numpy as np
import cPickle
import h5py

import sqtl
import utils
import config

import matplotlib; matplotlib.use('Agg')
import matplotlib.pyplot as plt
import seaborn as sns
from matplotlib.colors import ListedColormap

import matplotlib.gridspec as gridspec
from adjustText import adjust_text

sys.path.append('/cluster/home/starks/git/tools-python/viz/'); import gwas

# Location of trans/cis dataframes
#_DATA_DIR              = '/cluster/work/grlab/projects/TCGA/PanCancer/QTL_Analysis_results/analysis_10kset/2dmanhattan_information_v1'
_DATA_DIR              = '/cluster/work/grlab/projects/TCGA/PanCancer/QTL_Analysis_results/analysis_10kset/2dmanhattan_information_v3'

_GENOTYPE_H5_PATH      = '/cluster/work/grlab/projects/TCGA/PanCancer/hdf5_10Kset/mergedFiles_clean.hdf5'

# Genetic "constants"
_CANCER_GENE_PATH      = '/cluster/work/grlab/projects/TCGA/PanCancer/annotation/gencodeV14.v7.pancan_subset.ensembleID.list'
_CELL_CYCLE_GENE_PATH  = '/cluster/work/grlab/projects/TCGA/PanCancer/annotation/biomart_go_cell_cycle_strictv4.txt'
_RNA_BINDING_GENE_PATH = '/cluster/work/grlab/projects/TCGA/PanCancer/annotation/biomart_go_rna_binding_mRNA.txt'
_CHRM_LEN_PATH         = '/cluster/work/grlab/projects/TCGA/PanCancer/genome/hg19_hs37d5/chr_sizes'
_ENSG_TRANSLATE_PATH   = '/cluster/work/grlab/projects/ICGC/annotation/gencode.v19.annotation.hs37d5_chr.gtf.lookup.pickle'

# Load data fxns
def _hotfix_trans_colnames(df):
    df.index.name = df.columns[0].replace('# ', '')
    df.columns = df.columns.tolist()[1:-1] + df.columns[-1].split()
    return

def old_load_data():
    trans_df = pd.read_csv(os.path.join(_DATA_DIR, 'transdataresults.tsv'), sep='\t')
    _hotfix_trans_colnames(trans_df)
    cis_df = pd.read_csv(os.path.join(_DATA_DIR, 'cisdataresults.tsv'), sep='\t')
    cis_df.columns = trans_df.reset_index().columns
    cis_df = cis_df.set_index(cis_df.columns[0])
    assert cis_df.columns.equals(trans_df.columns)
    data = pd.concat((trans_df, cis_df))
    data['is_trans'] = data.loc[:, 'snp_chrm'] != data.loc[:, 'event-chrm']
    return data

def load_data():
    trans_df = pd.read_csv(os.path.join(_DATA_DIR, 'transdataresults_test_v3.tsv'), sep='\t')
    if trans_df.columns[0].startswith('#'):
        cols = trans_df.columns.tolist()
        cols[0] = cols[0].replace('# ', '')
        trans_df.columns = cols
    trans_df = trans_df.set_index('gene-id(event)')
    trans_df['is_trans'] = True

    cis_df = pd.read_csv(os.path.join(_DATA_DIR, 'cisdataresults_test_v3.tsv'), sep='\t')
    if cis_df.columns[0].startswith('#'):
        cols = cis_df.columns.tolist()
        cols[0] = cols[0].replace('# ', '')
        cis_df.columns = cols
    cis_df = cis_df.set_index('gene-id(event)')
    cis_df['is_trans'] = False

    assert cis_df.columns.equals(trans_df.columns)
    data = pd.concat((trans_df, cis_df))
    return data

def _load_chrm_lens():
    '''Dictionary mapping chrm 1-22 to lens (keys & values are floats)
    '''
    df = pd.read_csv(_CHRM_LEN_PATH, sep='\t', header=None, index_col=0)
    df = df.loc[np.arange(1,23).astype(str)]
    df.index = df.index.astype(float)
    df.iloc[:, 0] = df.iloc[:, 0].astype(float)
    return df.iloc[:, 0].to_dict()

def _load_ensg_to_gene():
    data = cPickle.load(open(_ENSG_TRANSLATE_PATH, 'r'))
    return data['ensembl2genename']

# Format plot fxns
def _format_ax(ax):
    ax.set_xticklabels([])
    ax.set_yticklabels([])
    return

def _increase_margin(ax, factor=.01):
    ylim = list(ax.get_ylim())
    ylength = ylim[1] - ylim[0]
    ylim[0] = ylim[0] - ylength * factor
    ylim[1] = ylim[1] + ylength * factor
    ax.set_ylim(ylim)

    xlim = list(ax.get_xlim())
    ylength = xlim[1] - xlim[0]
    xlim[0] = xlim[0] - ylength * factor
    xlim[1] = xlim[1] + ylength * factor
    ax.set_xlim(xlim)
    return

def _add_barplot_ensg_labels(ax, x_data, y_data, text_data):
    texts = [ax.text(x,y,t) for x,y,t in zip(x_data, y_data, text_data)]
    adjust_text(texts, expand_text=(1.2, .5),
                autoalign='xy', force_text=0.9, ax=ax,
                arrowprops=dict(arrowstyle="->", color='r', lw=0.5))
    return

def plot_sqtl_hist(ax, data, thold=100, ax_text=None):
    col_data = data.loc[data['is_trans'], :].set_index('gene-id(mutation)').loc[:, ['pos_gt']]
    col_data.index = col_data.index.map(lambda x: str(x).split('.')[0])
    unq, cnt = np.unique(col_data['pos_gt'], return_counts=True)
    map_pos_gt_to_count = dict(zip(unq, cnt))
    col_data['pos_gt_count'] = col_data['pos_gt'].map(lambda x: map_pos_gt_to_count[x])
    col_data = col_data.drop_duplicates()

    # Stagger small bars
    small_mask = col_data['pos_gt_count'] < 200
    bars = ax.bar(col_data.loc[~small_mask, 'pos_gt'],
                  col_data.loc[~small_mask, 'pos_gt_count'],
                  width=ax.get_xlim()[1] / (col_data.shape[0]),
                  edgecolor='black', linewidth=1)

    small_bars = ax.bar(col_data.loc[small_mask, 'pos_gt'],
                        col_data.loc[small_mask, 'pos_gt_count'],
                        width=ax.get_xlim()[1] / (col_data.shape[0]),
                        color=bars[0].get_facecolor(),
                        edgecolor='black', linewidth=1)

    ax.set_ylim(top=ax.get_ylim()[1] * 1.25)
    to_write = col_data.loc[col_data['pos_gt_count'] > thold].reset_index()[['pos_gt', 'pos_gt_count', 'gene-id(mutation)']]
    if ax_text is None:
        ax_text = ax
    else:
        ax_text.set_ylim(ax.get_ylim())
        ax_text.set_xlim(ax.get_xlim())
        ax_text.scatter(to_write['pos_gt'].values, to_write['pos_gt_count'].values, s=1)
    _add_barplot_ensg_labels(ax_text, to_write['pos_gt'].values,
                                      to_write['pos_gt_count'].values,
                                      to_write['gene-id(mutation)'].map(lambda x: _ENSG_TO_GENE[x]).values)

    ax.set_ymargin(0.2)
    ax.set_ylabel('# sQTL', size = 15)
    ax.set_xticks([])
    ax.yaxis.set_ticks_position('left')
    yticks = [0, int(50 * round(float(cnt.max())/50))]
    ax.set_yticks(yticks)
    ax.set_yticklabels(yticks)
    ax.spines['top'].set_visible(False)
    ax.spines['right'].set_visible(False)
    return bars, to_write

def plot_events_hist(ax, data):
    isTrans = data.loc[:, 'snp_chrm'] != data.loc[:, 'event-chrm']
#    ax.hist(data.loc[isTrans, 'pos_pt'].values, bins = data.shape[0] , orientation='horizontal')
    unq, cnt = np.unique(data.loc[data['is_trans'], 'pos_pt'], return_counts=True)

    ax.hist(data.loc[data['is_trans'], 'pos_pt'], bins=data.shape[0]/40 , orientation='horizontal')
#    ax.bar(unq, cnt, width=ax.get_ylim()[1] / unq.size, orientation='horizontal')
    ax.set_xlabel('# Events', size = 15)
    ax.set_yticks([])

    ax.spines['right'].set_visible(False)
    ax.spines['top'].set_visible(False)
    ax.xaxis.set_ticks_position('bottom')
    mylims = ax.get_xlim()
    ax.set_xticks(mylims)
    return

def plot_manhattan(ax, data):
    pos_gt, pos_pt = gwas.makeManhattanPlot2D(pv=data.loc[:, 'nominal p-value'].astype(float).values,
                                              pos_gt=data.loc[:, ['snp_chrm', 'snp_pos']].astype(float).values,
                                              pos_pt=data.loc[:, ['event-chrm', 'event-pos1']].astype(float).values,
                                              marker_color='b', chrm_lens=_CHRM_LENS,
                                              color_chrms=True, return_pos=True,
                                              ax=ax) #, tag = genesettag)#### better evpo

    ax.set_xticklabels(ax.get_xticklabels(), rotation = 90)
    ax.set_ylabel('Splice Event Position', size = 15)
    ax.set_xlabel(ax.get_xlabel(), size = 15)
    return pos_gt, pos_pt

def plot_gene_tags(ax, data):
    ax.grid(False)
    ax.set_ymargin(0.2)
    ax.set_facecolor('white')

    cancer_ensg = np.loadtxt(_CANCER_GENE_PATH, dtype=str)
    cell_cycle_ensg = np.loadtxt(_CELL_CYCLE_GENE_PATH, usecols=[0], dtype=str, skiprows=1)
    rna_bind_ensg = np.loadtxt(_RNA_BINDING_GENE_PATH, usecols=[0], dtype=str, skiprows=1)

    col_data = data.set_index('gene-id(mutation)').loc[:, 'pos_gt']
    col_data.index = col_data.index.map(lambda x: str(x).split('.')[0])

    pos_cancer_tags = col_data.loc[col_data.index.intersection(cancer_ensg)].unique()
    pos_cell_cycle_tags = col_data.loc[col_data.index.intersection(cell_cycle_ensg)].unique()
    pos_rna_bind_tags = col_data.loc[col_data.index.intersection(rna_bind_ensg)].unique()

    ax.scatter(pos_cancer_tags,     [0]*pos_cancer_tags.size, label='Cancer')
    ax.scatter(pos_cell_cycle_tags, [1]*pos_cell_cycle_tags.size, label='Cell Cycle')
    ax.scatter(pos_rna_bind_tags,   [2]*pos_rna_bind_tags.size, label='RNA Bind')
    ax.axis('off')
    return

def plot_event_tags(ax, data):
    ax.grid(False)
    ax.set_xmargin(0.2)
    ax.set_facecolor('white')

    cancer_ensg = np.loadtxt(_CANCER_GENE_PATH, dtype=str)
    cell_cycle_ensg = np.loadtxt(_CELL_CYCLE_GENE_PATH, usecols=[0], dtype=str, skiprows=1)
    rna_bind_ensg = np.loadtxt(_RNA_BINDING_GENE_PATH, usecols=[0], dtype=str, skiprows=1)

    row_data = data.loc[:, 'pos_pt']
    row_data.index = row_data.index.map(lambda x: str(x).split('.')[0])

    pos_cancer_tags = row_data.loc[row_data.index.intersection(cancer_ensg)].unique()
    pos_cell_cycle_tags = row_data.loc[row_data.index.intersection(cell_cycle_ensg)].unique()
    pos_rna_bind_tags = row_data.loc[row_data.index.intersection(rna_bind_ensg)].unique()

    ax.scatter([0]*pos_cancer_tags.size, pos_cancer_tags)
    ax.scatter([1]*pos_cell_cycle_tags.size,pos_cell_cycle_tags)
    ax.scatter([2]*pos_rna_bind_tags.size, pos_rna_bind_tags)
    ax.axis('off')
    return

_CHRM_LENS = _load_chrm_lens()
_ENSG_TO_GENE = _load_ensg_to_gene()

def run_plot_manhattan(data, write=True):
    fig = plt.figure(figsize=(13,13))
    gs = gridspec.GridSpec(4, 3,
                           width_ratios=[9, 1, .5],
                           height_ratios=[.5, 1, 1, 9],
                           hspace=0.05, wspace=0.05)

    ax_tdots = plt.subplot(gs[0,0]); _format_ax(ax_tdots)
    ax_text = plt.subplot(gs[1,0]); _format_ax(ax_text)
    ax_top = plt.subplot(gs[-2,0]); _format_ax(ax_top)
    ax_mh = plt.subplot(gs[-1,0]); _format_ax(ax_mh)
    ax_right = plt.subplot(gs[-1,-2]); _format_ax(ax_right)
    ax_rdots = plt.subplot(gs[-1,-1]); _format_ax(ax_rdots)
    ax_legend = plt.subplot(gs[0, -1]); _format_ax(ax_legend)

    plt.suptitle('Taken from: %s' %_DATA_DIR)
    pos_gt, pos_pt = plot_manhattan(ax_mh, data)
    assert pos_gt.size == data.shape[0]
    assert pos_pt.size == data.shape[0]
    data['pos_gt'] = pos_gt
    data['pos_pt'] = pos_pt
    _increase_margin(ax_mh, factor=.005)

    ax_top.set_xlim(ax_mh.get_xlim())
    ax_tdots.set_xlim(ax_mh.get_xlim())

    ax_right.set_ylim(ax_mh.get_ylim())
    ax_rdots.set_ylim(ax_mh.get_ylim())

    ax_text.axis('off')
    thold = 80
    _, labeled_gene_df = plot_sqtl_hist(ax_top, data, thold=thold, ax_text=ax_text)

    if write:
        plot_gene_tags(ax_tdots, data)

        plot_events_hist(ax_right, data)
        plot_event_tags(ax_rdots, data)

        # Add legend -- cancer, cell cycle, rna_bind
        dot_handles, dot_labels = ax_tdots.get_legend_handles_labels()
        ax_legend.legend(dot_handles, dot_labels)
        ax_legend.axis('off')

        fn_out = os.path.join('/cluster/work/grlab/projects/TCGA/PanCancer/rerun_plots', 'starks_sqtl', 'manhattan_sqtl_events_thold_%d_v3.png'%thold)
        if not os.path.exists(os.path.dirname(fn_out)): os.makedirs(os.path.dirname(fn_out))
        plt.savefig(fn_out, layout='tight', dpi=300)
        print "Writing to %s" %fn_out
        pdf_out = fn_out.replace('.png', '.pdf')
        print "Writing to %s" %pdf_out
        plt.savefig(pdf_out, layout='tight')
        plt.close()
        np.save(fn_out.replace('.png', '') + '_labeled_ensg.npy', labeled_gene_df['gene-id(mutation)'].values)
    return labeled_gene_df['gene-id(mutation)'].values

def run_plot_heatmap_labeled_ensg(ensg_labeled, gtdf):
    pass

def _sort_by_cnc(mut_status):
    tss_index = mut_status.index.map(lambda x: x.split('-')[1])
    return mut_status.iloc[np.argsort(tss_index)]

if __name__ == '__main__':
    data = load_data()
    run_plot_manhattan(data, write=True)


text:
import logging

import numpy as np
import hypothesis.strategies as hst
from hypothesis import HealthCheck, given, example, settings

from qcodes.dataset.measurements import Measurement

@given(n_points=hst.integers(min_value=1, max_value=100))
@example(n_points=5)
@settings(deadline=None, suppress_health_check=(HealthCheck.function_scoped_fixture,))
def test_datasaver_1d(experiment, DAC, DMM, caplog,
                      n_points):
    meas = Measurement()
    meas.register_parameter(DAC.ch1)
    meas.register_parameter(DMM.v1, setpoints=(DAC.ch1,))

    n_points_expected = 5

    meas.set_shapes({DMM.v1.full_name: (n_points_expected,)})

    with meas.run() as datasaver:

        for set_v in np.linspace(0, 1, n_points):
            DAC.ch1()
            datasaver.add_result((DAC.ch1, set_v),
                                 (DMM.v1, DMM.v1()))

    ds = datasaver.dataset
    caplog.clear()
    data = ds.get_parameter_data()

    for dataarray in data[DMM.v1.full_name].values():
        assert dataarray.shape == (n_points,)

    if n_points == n_points_expected:
        assert len(caplog.record_tuples) == 0
    elif n_points > n_points_expected:
        assert len(caplog.record_tuples) == 2
        exp_module = "qcodes.dataset.sqlite.queries"
        exp_level = logging.WARNING
        exp_msg = ("Tried to set data shape for {} in "
                   "dataset {} "
                   "from metadata when loading "
                   "but found inconsistent lengths {} and {}")
        assert caplog.record_tuples[0] == (exp_module,
                                           exp_level,
                                           exp_msg.format(DMM.v1.full_name,
                                                          DMM.v1.full_name,
                                                          n_points,
                                                          n_points_expected))
        assert caplog.record_tuples[1] == (exp_module,
                                           exp_level,
                                           exp_msg.format(DAC.ch1.full_name,
                                                          DMM.v1.full_name,
                                                          n_points,
                                                          n_points_expected))

@settings(deadline=None, suppress_health_check=(HealthCheck.function_scoped_fixture,))
@given(n_points_1=hst.integers(min_value=1, max_value=50),
       n_points_2=hst.integers(min_value=1, max_value=50))
@example(n_points_1=5, n_points_2=10)
def test_datasaver_2d(experiment, DAC, DMM, caplog,
                      n_points_1, n_points_2):
    meas = Measurement()
    meas.register_parameter(DAC.ch1)
    meas.register_parameter(DAC.ch2)
    meas.register_parameter(DMM.v1, setpoints=(DAC.ch1,
                                               DAC.ch2))

    n_points_expected_1 = 5
    n_points_expected_2 = 10

    meas.set_shapes({DMM.v1.full_name: (n_points_expected_1,
                                        n_points_expected_2,)})

    with meas.run() as datasaver:

        for set_v_1 in np.linspace(0, 1, n_points_1):
            for set_v_2 in np.linspace(0, 1, n_points_2):
                datasaver.add_result((DAC.ch1, set_v_1),
                                     (DAC.ch2, set_v_2),
                                     (DMM.v1, DMM.v1()))

    ds = datasaver.dataset
    caplog.clear()
    data = ds.get_parameter_data()

    if n_points_1*n_points_2 == n_points_expected_1*n_points_expected_2:
        assert len(caplog.record_tuples) == 0
        for dataarray in data[DMM.v1.full_name].values():
            assert dataarray.shape == (n_points_expected_1, n_points_expected_2)
    elif n_points_1*n_points_2 > n_points_expected_1*n_points_expected_2:
        assert len(caplog.record_tuples) == 3
        exp_module = "qcodes.dataset.sqlite.queries"
        exp_level = logging.WARNING
        exp_msg = ("Tried to set data shape for {} in "
                   "dataset {} "
                   "from metadata when loading "
                   "but found inconsistent lengths {} and {}")
        assert caplog.record_tuples[0] == (
            exp_module,
            exp_level,
            exp_msg.format(
                DMM.v1.full_name,
                DMM.v1.full_name,
                n_points_1*n_points_2,
                n_points_expected_1*n_points_expected_2
            )
        )
        assert caplog.record_tuples[1] == (
            exp_module,
            exp_level,
            exp_msg.format(
                DAC.ch1.full_name,
                DMM.v1.full_name,
                n_points_1*n_points_2,
                n_points_expected_1*n_points_expected_2)
        )
        assert caplog.record_tuples[2] == (
            exp_module,
            exp_level,
            exp_msg.format(
                DAC.ch2.full_name,
                DMM.v1.full_name,
                n_points_1*n_points_2,
                n_points_expected_1*n_points_expected_2
            )
        )

ext: py

max_forks_repo_name: nikhartman/Qcodes

reason: import numpy


text:
"""
Example of using meshing_ats.py to generate a column mesh with layers and variable resolution.

"""

import sys,os
import numpy as np
from copy import deepcopy

# This is the standard path for ATS's source directory
try:
    import meshing_ats
except ImportError:
    #sys.path.append(os.path.join(os.environ['ATS_SRC_DIR'],'tools','meshing', 'meshing_ats'))
    sys.path.append(os.path.join('/Users/f9y/mygithub/ATS_REPOS/amanzi/src/physics/ats', \
                                 'tools','meshing', 'meshing_ats'))
    import meshing_ats

# set up the surface mesh, which is a single unit cell
# TODO - read-in latixy/longxy/topo from 'surfdata.nc' and do some lat/lon<->x/y conversion
x = np.array([0.0, 1.0],'d')
elv = np.array([0.0, 0.0], 'd')

# using from_Transect extrudes the x,elv line in the y-direction to
# create 1 cell in y.  This results in a single cell.
m2 = meshing_ats.Mesh2D.from_Transect(x,elv)

# layer extrusion
# -- data structures needed for extrusion
layer_types = []
layer_data = []
layer_ncells = []
layer_mat_ids = []
z = 0.0

# -- standard soil layers from ELM's 15-layer column--
#  variable layer thickness
#  15 layers
#  mat-id for each top 10 layer and 1 for rest 5 layers (called bedrock in ELM)

ncells = 15
jidx = np.array(range(ncells))+1
zsoi = 0.025*(np.exp(0.5*(jidx-0.5))-1.0)       #ELM soil layer node depths - somewhere inside a layer but not centroid
dzsoi= np.zeros_like(zsoi)
dzsoi[0] = 0.5*(zsoi[0]+zsoi[1])                #thickness b/n two vertical interfaces (vertices)
for j in range(1,ncells-1):
    dzsoi[j]= 0.5*(zsoi[j+1]-zsoi[j-1])
dzsoi[ncells-1] = zsoi[ncells-1]-zsoi[ncells-2]

nlevsoi = 10
for j in range(ncells):
    i=ncells-j-1
    z = z+dzsoi[i]
    print('z, dz, z_centroid, z_node_elm: ', i, z, dzsoi[i], z-dzsoi[i]/2.0, zsoi[i])
    layer_types.append("snapped")
    layer_data.append(z)
    layer_ncells.append(1)
    if i<nlevsoi:
        layer_mat_ids.append(1001+i)
    else:
        layer_mat_ids.append(1001+nlevsoi)

# -- print out a summary --
#meshing_ats.summarize_extrusion(layer_types, layer_data, layer_ncells, layer_mat_ids)
dz = deepcopy(layer_data)
dz.insert(0,0.0)
dz = np.diff(dz)
meshing_ats.summarize_extrusion(layer_types, dz, layer_ncells, layer_mat_ids)   # here 'layer_data' shall be 'dz'

# Extrude the 3D model with this structure and write to file
m3 = meshing_ats.Mesh3D.extruded_Mesh2D(m2, layer_types,
                                        layer_data,                               # here 'layer_data' shall be 'z', depth of vertical vertices
                                        layer_ncells,
                                        layer_mat_ids)
if os.path.exists('soilcolumn_elm4ats.exo'):
    os.remove('soilcolumn_elm4ats.exo')
m3.write_exodus("soilcolumn_elm4ats.exo")


text:
import SimpleITK as sitk
import glob
import numpy as np
from PIL import Image
import cv2
import glob

import matplotlib.pyplot as plt # plt 用于显示图片

import os
res=[]
niipath = 'testct'
new_seg = 'testseg'
for filename in os.listdir(r'CT'):
    # print(filename+'/DICOM_anon')
    res.append(filename+'/Ground')
# print(res)
for i in res:
    file_path = os.path.join('CT',i) #png存放文件夹
    idx = i.split('/')[0]
    # print(idx,i)
    png_path = glob.glob('./CT/'+idx+'/Ground/*')
    png_path.sort(reverse=True)
    # print(png_path)
    niip = os.path.join(niipath,idx+'.nii.gz')
    # print(niip)
    ct = sitk.ReadImage(niip, sitk.sitkInt16)
    ct_array = sitk.GetArrayFromImage(ct)
    print(ct_array.shape)
    # for i in range(ct_array.shape[0]):
    #     plt.imshow(ct_array[i,:,:])
    #     # cv2.imwrite(target+'/'+'30'+str(i)+".png",ct_array[i,:,:])
    #     plt.show()
    seg_array = np.zeros(ct_array.shape, dtype='uint8')
    for i,d in enumerate(png_path):
        # print(i,d)
        img_as_img = Image.open(d)
        # img_as_np = np.asarray(img_as_img)
        # allImg[i, :, :] = img_as_np
        # img_as_img.show()
        seg_array[i,:,:] = img_as_img

    new_seg = sitk.GetImageFromArray(seg_array)
    new_seg.SetDirection(ct.GetDirection())
    new_seg.SetOrigin(ct.GetOrigin())
    new_seg.SetSpacing(ct.GetSpacing())
    sitk.WriteImage(new_seg, idx+'.nii.gz')


text:
from typing import List
import io
from os.path import isfile, basename
import os

import numpy as np

from .keyed_vectors import KeyedVectors
from .oov_error import OOVError
from .utils import download_data, extract_gz

def _load_text_file(path: str):
    """Load .vec file"""
    fin = io.open(path, 'rb')
    first_line = fin.readline().decode('utf8')
    vocab_size, embedding_size = map(int, first_line.split())

    # init vocab list
    vocab_list = ['0'] * vocab_size

    # record start position of each line in file
    byte_pos = [0] * (vocab_size + 1)
    byte_pos[0] = fin.tell()

    for idx in range(vocab_size):
        line = fin.readline()
        tokens = line.rstrip().split(b' ')
        vocab_list[idx] = tokens[0].decode('utf8')
        byte_pos[idx + 1] = fin.tell()
    fin.close()
    return embedding_size, vocab_size, vocab_list, byte_pos

def _load_bin_file(path: str):
    # load .bin file
    # Note that float in this file should be float32
    # float64 is not allowed
    fin = open(path, 'rb')
    header = fin.readline().decode('utf8')
    vocab_size, embedding_size = (int(x) for x in header.split())

    # init vocab list
    vocab_list = ['0'] * vocab_size

    # record start position of each line in file
    byte_pos = [0] * vocab_size

    binary_len = 4 * embedding_size  # 4: because of float32
    for idx in range(vocab_size):
        # mixed text and binary: read text first, then binary
        word = []
        while True:
            ch = fin.read(1)
            if ch == b' ':
                break
            if ch == b'':
                raise EOFError(
                    "unexpected end of input; is count incorrect or file otherwise damaged?")
            # ignore newlines in front of words (some binary files have)
            if ch != b'\n':
                word.append(ch)

        byte_pos[idx] = fin.tell()
        vocab_list[idx] = b''.join(word).decode('utf8')
        fin.read(binary_len)
    fin.close()
    return embedding_size, vocab_size, vocab_list, byte_pos

class KeyedVectorsLight(KeyedVectors):

    def __init__(self, path: str, binary: bool = False):
        super().__init__(path=path, binary=binary)

    def build(self):
        if self._is_built:
            pass

        if not isfile(self._path):
            # if data is not at self._path
            # download it through url in .env
            download_data(
                url=os.getenv(basename(self._path)),
                output_path=self._path + '.gz',
            )
            extract_gz(self._path + '.gz')

        (
            self._embedding_size,
            self._vocab_size,
            self._vocab_list,
            self._byte_pos,
        ) = self._load_data(
            path=self._path,
            binary=self._binary,
        )

        self._vloader = LowMemoryVecLoader(
            path=self._path,
            byte_pos=self._byte_pos,
            binary=self._binary,
        )
        self._is_built = True

    def _get_vector(self, index: int) -> np.ndarray:
        if (index >= 0) and (index < self._vocab_size):
            return self._vloader[index]
        else:
            raise OOVError

    @staticmethod
    def _load_data(path: str, binary: bool = False):
        if binary:
            return _load_bin_file(path=path)
        else:
            return _load_text_file(path=path)

    def __del__(self):
        if '_vloader' in self.__dict__:
            del self._vloader

class LowMemoryVecLoader:

    def __init__(
            self,
            path: str,
            byte_pos: List[int],
            binary: bool = False,
        ):
        self._byte_pos = byte_pos
        self._binary = binary

        self.fin = open(path, 'rb')
        header = self.fin.readline().decode('utf8')
        self._vocab_size, self._embedding_size = map(int, header.split())

    def _get_vector_from_text(self, index: int) -> np.ndarray:
        # get start end position from _byte_pos
        start_pos = self._byte_pos[index]
        end_pos = self._byte_pos[index + 1]

        #
        diff = end_pos - start_pos

        self.fin.seek(start_pos, 0)
        line = self.fin.read(diff).decode('utf8')

        tokens = line.rstrip().split(' ')
        vector = list(map(float, tokens[1:]))
        return np.array(vector).astype(np.float32)

    def _get_vector_from_bin(self, index: int) -> np.ndarray:
        binary_len = 4 * self._embedding_size  # 4: because of float32

        start_pos = self._byte_pos[index]

        self.fin.seek(start_pos, 0)

        vector = np.frombuffer(
            self.fin.read(binary_len),
            dtype='float32',
        )
        return vector

    def __getitem__(self, index: int) -> np.ndarray:
        if (index < 0) or (index > self._vocab_size):
            raise ValueError('Out of index')

        if self._binary:
            vector = self._get_vector_from_bin(index=index)
        else:
            vector = self._get_vector_from_text(index=index)
        return vector

    def __del__(self):
        self.fin.close()
        del self.fin

max_stars_repo_name: Yoctol/word-embedder

max_issues_repo_name: Yoctol/word-embedder

max_forks_repo_name: Yoctol/word-embedder


text:
#To plot los signal for 4 given los for a given offset, polystr and tau

import numpy as np
import matplotlib.pyplot as plt
import argparse
import sys

######################################################
#### Choose (x,y) for plotting los signal
los_array=[[34,54], [12,183], [234,145], [218,87]]
######################################################

parser = argparse.ArgumentParser()

parser.add_argument("-p", "--polystr", type=int, default=-1, help="Exponents in the terms of the fitting polynomial")
parser.add_argument("-o", "--offset", type=float, default=-1, help="Offset added to the original cube")
parser.add_argument("-t", "--tau", type=float, default=-1, help="Integration time (in hrs)")
args = parser.parse_args()

if any(i==-1 for i in [args.polystr,args.offset,args.tau]):
  print("Invalid Input.\nRun `python3 plot_los_signal.py --help` for more info")
  sys.exit()

def load_binary_data(filename, dtype=np.float64):
  f = open(filename, "rb")
  data = f.read()
  f.close()
  _data = np.fromstring(data, dtype)
  if sys.byteorder == 'big':
    _data = _data.byteswap()
  return _data

##############################################################
####Plotting starts here
##############################################################

freq=np.loadtxt("../ska_data/frequency.dat", unpack=True)
freq=[i/1e6 for i in freq]

#Plan:
#recovered/total[0]=eor_p_fg signal
#recovered/total[1]=eor_p_noise signal
#recovered/total[2]=eor_p_noise_p_fg signal

recovered=[None,None,None]
total=[None,None,None]

k=0
for i in ["eor_p_fg", "eor_p_noise", "eor_p_noise_p_fg"]:
  file1="../results/ska_cube/cube/%s_recovered_21cm_signal_XY_300_%.1f_hrs_Npol_%d_offset_%.6f.bin" % (i, args.tau, args.polystr, args.offset)
  file2="../ska_data/%s_pre_total_field_300_%.1f_hr-image.bin" % (i, args.tau)
  recovered[k]=load_binary_data(file1)
  total[k]=load_binary_data(file2)
  k+=1
fig=plt.figure()
ax=fig.add_subplot(1,1,1)

for los in los_array:

  y1=[]; y2=[]; y3=[]; y4=[]; y5=[]; y6=[]
  i=los[0]
  j=los[1]

  for k in range(21):
    y1.append(recovered[0][k+21*(j+300*i)])
    y2.append(recovered[1][k+21*(j+300*i)])
    y3.append(recovered[2][k+21*(j+300*i)])
    y4.append(total[0][k+21*(j+300*i)])
    y5.append(total[1][k+21*(j+300*i)])
    y6.append(total[2][k+21*(j+300*i)])

  ax.plot(freq,y1,"r--",label="eor+fg")
  ax.plot(freq,y2,"b:",label="eor+noise")
  ax.plot(freq,y3,"g-.",label="eor+noise+fg")
  plt.xlabel("Frequency [MHz]")
  plt.ylabel(r"$T_{21}$ [K]")
  plt.tight_layout()
  plt.legend(loc="best")
  fig.savefig("ska/los_1/los_x_%d_y_%d_recovered_all_tau%.1f_polystr%d_offset%.6f.png"%(i,j,args.tau,args.polystr,args.offset))
  plt.cla()

  ax.plot(freq,y4,"r--",label="eor+fg")
  ax.plot(freq,y5,"b:",label="eor+noise")
  ax.plot(freq,y6,"g-.",label="eor+noise+fg")
  plt.xlabel("Frequency [MHz]")
  plt.ylabel(r"$T_{21}$ [K]")
  plt.tight_layout()
  plt.legend(loc="best")
  fig.savefig("ska/los_1/los_x_%d_y_%d_total_all_tau%.1f_polystr%d_offset%.6f.png"%(i,j,args.tau,args.polystr,args.offset))
  plt.cla()

  #Setting average of all the fields to zero
  ave=sum(y1)/len(y1)
  y1=[i-ave for i in y1]

  ave=sum(y2)/len(y2)
  y2=[i-ave for i in y2]

  ave=sum(y3)/len(y3)
  y3=[i-ave for i in y3]

  ave=sum(y4)/len(y4)
  y4=[i-ave for i in y4]

  ave=sum(y5)/len(y5)
  y5=[i-ave for i in y5]

  ave=sum(y6)/len(y6)
  y6=[i-ave for i in y6]

  ax.plot(freq,y1,"r--",label="eor+fg, recovered")
  ax.plot(freq,y4,"b:",label="eor+fg, total")
  plt.xlabel("Frequency [MHz]")
  plt.ylabel(r"$T_{21}$ [K]")
  plt.tight_layout()
  plt.legend(loc="best")
  fig.savefig("ska/los_1/los_x_%d_y_%d_eor_p_fg_tau%.1f_polystr%d_offset%.6f.png"%(i,j,args.tau,args.polystr,args.offset))
  plt.cla()

  ax.plot(freq,y2,"r--",label="eor+noise, recovered")
  ax.plot(freq,y5,"b:",label="eor+noise, total")
  plt.xlabel("Frequency [MHz]")
  plt.ylabel(r"$T_{21}$ [K]")
  plt.tight_layout()
  plt.legend(loc="best")
  fig.savefig("ska/los_1/los_x_%d_y_%d_eor_p_noise_tau%.1f_polystr%d_offset%.6f.png"%(i,j,args.tau,args.polystr,args.offset))
  plt.cla()

  ax.plot(freq,y3,"r--",label="eor+noise+fg, recovered")
  ax.plot(freq,y6,"b:",label="eor+noise+fg, total")
  plt.xlabel("Frequency [MHz]")
  plt.ylabel(r"$T_{21}$ [K]")
  plt.tight_layout()
  plt.legend(loc="best")
  fig.savefig("ska/los_1/los_x_%d_y_%d_eor_p_noise_fg_tau%.1f_polystr%d_offset%.6f.png"%(i,j,args.tau,args.polystr,args.offset))
  plt.cla()


text:
__author__ = 'lucabasa'
__version__ = '1.1.1'
__status__ = 'development'

import pandas as pd
import numpy as np

from sklearn.base import BaseEstimator, TransformerMixin

class tr_numeric(BaseEstimator, TransformerMixin):
    def __init__(self, SF_room=True, bedroom=True, bath=True, lot=True, service=True):
        self.columns = []  # useful to well behave with FeatureUnion
        self.SF_room = SF_room
        self.bedroom = bedroom
        self.bath = bath
        self.lot = lot
        self.service = service

    def fit(self, X, y=None):
        return self

    def remove_skew(self, X, column):
        X[column] = np.log1p(X[column])
        return X

    def SF_per_room(self, X):
        if self.SF_room:
            X['sf_per_room'] = X['GrLivArea'] / X['TotRmsAbvGrd']
        return X

    def bedroom_prop(self, X):
        if self.bedroom:
            X['bedroom_prop'] = X['BedroomAbvGr'] / X['TotRmsAbvGrd']
            del X['BedroomAbvGr'] # the new feature makes it redundant and it is not important
        return X

    def total_bath(self, X):
        if self.bath:
            X['total_bath'] = (X[[col for col in X.columns if 'FullBath' in col]].sum(axis=1) +
                             0.5 * X[[col for col in X.columns if 'HalfBath' in col]].sum(axis=1))
            del X['FullBath']  # redundant

        del X['HalfBath']  # not useful anyway
        del X['BsmtHalfBath']
        del X['BsmtFullBath']
        return X

    def lot_prop(self, X):
        if self.lot:
            X['lot_prop'] = X['LotArea'] / X['GrLivArea']
        return X

    def service_area(self, X):
        if self.service:
            X['service_area'] = X['TotalBsmtSF'] + X['GarageArea']
            del X['TotalBsmtSF']
            del X['GarageArea']
        return X

    def transform(self, X, y=None):
        for col in ['GrLivArea', '1stFlrSF', 'LotArea']:
            X = self.remove_skew(X, col)

        X = self.SF_per_room(X)
        X = self.bedroom_prop(X)
        X = self.total_bath(X)
        X = self.lot_prop(X)
        X = self.service_area(X)

        self.columns = X.columns
        return X

    def get_feature_names(self):
        return self.columns


text:
import os
import math
import pandas as pd
import sklearn.metrics as metrics
from sklearn.ensemble import RandomForestClassifier
from sklearn.linear_model import LogisticRegression
from sklearn.metrics import classification_report, confusion_matrix, roc_auc_score, accuracy_score
from sklearn.metrics import average_precision_score
from numpyencoder import NumpyEncoder
from get_data import read_params
import argparse
import joblib
import json

def train_and_evaluate(config_path):
    """
    This function trains and evaluates a machine learning model on the dataset.

    :param config_path: the path of the config file to use
    """
    # Read configuration options
    config = read_params(config_path)
    model_type = config["train_and_evaluate"]["model"]
    test_data_path = config["split_data"]["test_path"]
    train_data_path = config["split_data"]["train_path"]
    random_state = config["base"]["random_state"]
    model_dir = config["model_dir"]
    target = [config["base"]["target_col"]]
    scores_file = config["reports"]["scores"]
    prc_file = config["reports"]["prc"]
    roc_file = config["reports"]["roc"]

    # Load training and validation datasets
    train = pd.read_csv(train_data_path, sep=",")
    test = pd.read_csv(test_data_path, sep=",")

    # Separate features (x) from label (y)
    train_y = train[target]
    test_y = test[target]
    train_x = train.drop(target, axis=1)
    test_x = test.drop(target, axis=1)

    if model_type == "logistic_regression":
        # Build logistic regression model
        model = LogisticRegression(solver='sag', random_state=random_state).fit(train_x, train_y)
    elif model_type == "random_forest":
        # Build random forest model
        model = RandomForestClassifier(n_estimators=50)
    else:
        return

    # Fit the model to the training data
    model.fit(train_x, train_y)

    # Report training set score
    train_score = model.score(train_x, train_y) * 100
    print(train_score)
    # Report test set score
    test_score = model.score(test_x, test_y) * 100
    print(test_score)

    # Predict output for observations in validation set
    predicted_val = model.predict(test_x)

    # Calculate performance metrics
    precision, recall, prc_thresholds = metrics.precision_recall_curve(test_y, predicted_val)
    fpr, tpr, roc_thresholds = metrics.roc_curve(test_y, predicted_val)
    nth_point = math.ceil(len(prc_thresholds) / 1000)
    prc_points = list(zip(precision, recall, prc_thresholds))[::nth_point]

    with open(prc_file, "w") as fd:
        prcs = {
            "prc": [
                {"precision": p, "recall": r, "threshold": t}
                for p, r, t in prc_points
            ]
        }
        json.dump(prcs, fd, indent=4, cls=NumpyEncoder)

    with open(roc_file, "w") as fd:
        rocs = {
            "roc": [
                {"fpr": fp, "tpr": tp, "threshold": t}
                for fp, tp, t in zip(fpr, tpr, roc_thresholds)
            ]
        }
        json.dump(rocs, fd, indent=4, cls=NumpyEncoder)

    # Print classification report
    print(classification_report(test_y, predicted_val))

    # Confusion Matrix and plot
    cm = confusion_matrix(test_y, predicted_val)
    print(cm)
    df1 = pd.DataFrame(predicted_val, columns=['Predicted'])
    df_cm = pd.concat([test_y, df1], axis=1)
    print(df_cm)
    df_cm.to_csv('cm.csv', index=False)

    # Receiver operating characteristic - area under curve
    roc_auc = roc_auc_score(test_y, model.predict_proba(test_x)[:, 1])
    print('ROC_AUC:{0:0.2f}'.format(roc_auc))

    # Model accuracy
    model_accuracy = accuracy_score(test_y, predicted_val)
    print('Model Accuracy:{0:0.2f}'.format(model_accuracy))

    # Average precision score
    average_precision = average_precision_score(test_y, predicted_val)
    print('Average precision-recall score: {0:0.2f}'.format(average_precision))

    with open(scores_file, "w") as f:
        scores = {
            "train_score": train_score,
            "test_score": test_score,
            "roc_auc": roc_auc,
            "Precision": list(precision),
            "Recall": list(recall),
            "Average precision": average_precision,
            "Model Accuracy": model_accuracy
        }
        json.dump(scores, f, indent=4)

    # Output model
    os.makedirs(model_dir, exist_ok=True)
    model_path = os.path.join(model_dir, "model.joblib")
    joblib.dump(model, model_path)

if __name__ == "__main__":
    # If the file is being run, parse command line arguments to get config filepath
    args = argparse.ArgumentParser()
    args.add_argument("--config", default="params.yaml")
    parsed_args = args.parse_args()

    # Train and evaluate the model
    train_and_evaluate(config_path=parsed_args.config)

max_stars_repo_name: RatnaDS/NCSU

reason: from numpy


text:
#!/usr/bin/env python2

import rospy
import numpy as np
# import threading
from timeit import default_timer as timer
import time
import json

class SIGNAL_ANALYZER:

    def __init__(self, signal_name="signal_analysis", th_abnormally_high=1.0):
        #
        self.name = signal_name

        # states
        self.value = None
        self.value_avg = None
        #
        self.stamp_start = timer()
        self.stamp_last = self.stamp_start

        # Parameters
        self.alpha = 2*np.pi*0.1 # 0.1 Hz
        self.th_abnormally_high = th_abnormally_high
        #
        self.event_str_dict = dict()
        self.event_str_dict[0] = "abnormally high"

        #
        self.event_publisher = None

    def event_2_json(self, event_code):
        """
        Through ROS std_msgs.String
        json string:
        {
            "module": "yyy"
            "status": "OK"/"WARN"/"ERROR"/"FATAL"/"UNKNOWN"
            "event_str": "xxx event of yyy module"
        }
        Output: json string
        """
        json_dict = dict()
        json_dict["module"] = self.name
        json_dict["status"] = "WARN"
        json_dict["event_str"] = self.event_str_dict[event_code]
        return json.dumps(json_dict)

    def update(self, value_in):
        """
        """
        stamp_now = timer()
        delta_T = stamp_now - self.stamp_last
        #---------------------------#
        if self.value is None:
            self.value = value_in
        if self.value_avg is None:
            self.value_avg = value_in
        #
        delta_value_short = value_in - self.value
        delta_value_long = value_in - self.value_avg
        # Check 1
        if delta_value_long > self.th_abnormally_high:
            event_json = self.event_2_json(0)
            print(event_json)
            if self.event_publisher:
                self.event_publisher.publish( event_json )
        #
        adT = self.alpha * delta_T
        if adT > 1.0:
            adT = 1.0
        elif adT < 0.0:
            adT = 0.0
        # Update the average
        self.value_avg += adT * (value_in - self.value_avg)
        #---------------------------#
        self.value = value_in
        self.stamp_last = stamp_now

if __name__ == "__main__":
    sig_analyzer = SIGNAL_ANALYZER()
    while True:
        sig_analyzer.update( np.random.rand() )
        time.sleep(0.3)


text:
#Procedimentos:
#1: Definir a origem e calibrar a posição do objeto na origem escolhida
#2: Calibrar a profundidade (z) utilizando regressão linear (medida real x medida obtida)
#3: Calibrar a relação entre a variável z e as coordenadas x e y. Relacionar os valores de (x,z) e (y,z).

import cv2 as cv
from threading import Thread
import numpy as np
from numpy.linalg import inv, det
import scipy as sci
from scipy.spatial.transform import Rotation as R
from imutils.video import WebcamVideoStream
import matplotlib.pyplot as plt
import matplotlib.animation as animation
import serial
import collections
import glob
import sys
import time
import math
from numpy import cos as c
from numpy import sin as s

from vpython import *

#Calibration position: In the beggining of the video, the script will get 40 early datas to compute the mean
#measurements and so find offset values. Having the offset, we can calculate the relative origin.
#Keep in mind that the marker's position was acquired with respect to camera frame.

isReceive = False
isRun = True
gx = 0.0
gy = 0.0
gz = 0.0
ax = 0.0
ay = 0.0
az = 0.0
cx = 0.0
cy = 0.0
cz = 0.0
dt = 0.1
g = 9.81
q_ant = np.array([1.0, 0.0, 0.0, 0.0])

x_ant_ori = np.array([[0], [0], [0], [0], [0], [0]])
P_ant_ori = np.eye(6)*1000

x_ant_pos = np.array([[0], [0], [0], [0], [0], [0]])
P_ant_pos = np.eye(6)*1000

def EKF_Accel_Camera(x_ant, P_ant, accel, y, dt):

    # Q ----> Process Noise Covariance
    # C ----> Control Noise Covariance
    # R  ---> Measurement Noise Covariance

    F = np.array([[1, 0, 0, dt, 0, 0],
                  [0, 1, 0, 0, dt, 0],
                  [0, 0, 1, 0, 0, dt],
                  [0, 0, 0, 1, 0, 0],
                  [0, 0, 0, 0, 1, 0],
                  [0, 0, 0, 0, 0, 1]])
    V = np.array([[0.5*dt**2, 0, 0],
                  [0, 0.5*dt**2, 0],
                  [0, 0, 0.5*dt**2],
                  [dt, 0, 0],
                  [0, dt, 0],
                  [0, 0, dt]])

    H = np.array([[1, 0, 0, 0, 0, 0],
                  [0, 1, 0, 0, 0, 0],
                  [0, 0, 1, 0, 0, 0]])

    Q = np.eye(6)
    R = np.eye(3)

    #Process Equations
    x = F@x_ant + V@accel

    #State covariance matrix
    P = F@P_ant@F.T + Q

    #Kalman Gain and Inovation
    K = P@H.T@inv(H@P@H.T + R)
    z = y - H@x

    #Update step
    x_est = x + K@z
    P_est = P - K@H@P

    return x_est, P_est

def KF_Orientation(x_ant, P_ant, y, dt):

    A = np.array([[1, 0, 0, dt, 0, 0],
                  [0, 1, 0, 0, dt, 0],
                  [0, 0, 1, 0, 0, dt],
                  [0, 0, 0, 1, 0, 0],
                  [0, 0, 0, 0, 1, 0],
                  [0, 0, 0, 0, 0, 1]])

    B = np.array([[]])

    H = np.eye(6)

    Q = np.array([[10, 0, 0, 0, 0, 0],
                  [0, 10, 0, 0, 0, 0],
                  [0, 0, 1, 0, 0, 0],
                  [0, 0, 0, 1, 0, 0],
                  [0, 0, 0, 0, 1, 0],
                  [0, 0, 0, 0, 0, 1]])

    R = np.array([[.1, 0, 0, 0, 0, 0],
                  [0, .1, 0, 0, 0, 0],
                  [0, 0, .01, 0, 0, 0],
                  [0, 0, 0, .01, 0, 0],
                  [0, 0, 0, 0, .01, 0],
                  [0, 0, 0, 0, 0, .01]])

    #Prediction
    x_p = A@x_ant
    P_p = A@P_ant@A.T + Q

    #Kalman Gain and Inovation
    K = P_p@H.T@inv(H@P_p@H.T + R)
    inov = y - H@x_p

    #Update
    x_est = x_p + K@inov
    P_est = P_p - K@H@P_p

    return x_est, P_est

def KF_Position(x_ant, P_ant, u, y, dt):

    A = np.array([[1, 0, 0, dt, 0, 0],
                  [0, 1, 0, 0, dt, 0],
                  [0, 0, 1, 0, 0, dt],
                  [0, 0, 0, 1, 0, 0],
                  [0, 0, 0, 0, 1, 0],
                  [0, 0, 0, 0, 0, 1]])

    B = np.array([[dt**2/2, 0, 0],
                  [0, dt**2/2, 0],
                  [0, 0, dt**2/2],
                  [dt, 0, 0],
                  [0, dt, 0],
                  [0, 0, dt]])

    H = np.array([[1, 0, 0, 0, 0, 0],
                  [0, 1, 0, 0, 0, 0],
                  [0, 0, 1, 0, 0, 0]])

    Q = np.array([[100, 0, 0, 0, 0, 0],
                  [0, 100, 0, 0, 0, 0],
                  [0, 0, 100, 0, 0, 0],
                  [0, 0, 0, 1, 0, 0],
                  [0, 0, 0, 0, 1, 0],
                  [0, 0, 0, 0, 0, 1]])

    R = np.array([[100, 0, 0],
                  [0, 100, 0],
                  [0, 0, 100]])

    #Prediction
    x_p = A@x_ant + B@u
    P_p = A@P_ant@A.T + Q

    #Kalman Gain and Inovation
    K = P_p@H.T@inv(H@P_p@H.T + R)
    inov = y - H@x_p

    #Update
    x_est = x_p + K@inov
    P_est = P_p - K@H@P_p

    return x_est, P_est

def quat_product(a, b):
    k1 = a[0]*b[0] - a[1]*b[1] - a[2]*b[2] - a[3]*b[3]
    k2 = a[0]*b[1] + a[1]*b[0] + a[2]*b[3] - a[3]*b[2]
    k3 = a[0]*b[2] - a[1]*b[3] + a[2]*b[0] + a[3]*b[1]
    k4 = a[0]*b[3] + a[1]*b[2] - a[2]*b[1] + a[3]*b[0]

    k = np.array([k1, k2, k3, k4])

    return k

def MadgwickUpdate_OnlyAccel(gx, gy, gz, ax, ay, az, cx, cy, cz, dt, q_ant):

    q0 = q_ant[0]
    q1 = q_ant[1]
    q2 = q_ant[2]
    q3 = q_ant[3]
    beta = 0.5

    # Rate of change of quaternion from gyroscope
    qDot1 = 0.5 * (-q1 * gx - q2 * gy - q3 * gz)
    qDot2 = 0.5 * (q0 * gx + q2 * gz - q3 * gy)
    qDot3 = 0.5 * (q0 * gy - q1 * gz + q3 * gx)
    qDot4 = 0.5 * (q0 * gz + q1 * gy - q2 * gx)

    # Compute feedback only if accelerometer measurement valid (avoids NaN in accelerometer normalisation)
    if ax !=0 and ay != 0 and az != 0:

        # Normalise accelerometer measurement
        recipNorm = 1/(math.sqrt(ax * ax + ay * ay + az * az))
        ax *= recipNorm
        ay *= recipNorm
        az *= recipNorm

        print('ax: {0}, ay: {1}, az: {2}'.format(ax, ay, az))

        # Auxiliary variables to avoid repeated arithmetic
        _2q0 = 2.0 * q0
        _2q1 = 2.0 * q1
        _2q2 = 2.0 * q2
        _2q3 = 2.0 * q3
        _4q0 = 4.0 * q0
        _4q1 = 4.0 * q1
        _4q2 = 4.0 * q2
        _4q3 = 4.0 * q3
        _8q1 = 8.0 * q1
        _8q2 = 8.0 * q2
        q0q0 = q0 * q0
        q1q1 = q1 * q1
        q2q2 = q2 * q2
        q3q3 = q3 * q3

        f1 = 2*(q1*q3 - q0*q2) - ax
        f2 = 2*(q0*q1 + q2*q3) - ay
        f3 = 2*(0.5 - q1q1 - q2q2) - az
        f4 = 2*(0.5 - q2q2 - q3q3) - cx
        f5 = 2*(q1*q2 - q0*q3) - cy
        f6 = 2*(q0*q2 + q1*q3) - cz

        print('cx: {0}, cy: {1}, cz: {2}'.format(cx, cy, cz))

        # Gradient decent algorithm corrective step
        s0 = -_2q2*f1 + _2q1*f2 - _2q3*f5 + _2q2*f6
        s1 = _2q3*f1 + _2q0*f2 - _4q1*f3 + _2q2*f5 + _2q3*f6
        s2 = -_2q0*f1 + _2q3*f2 - _4q2*f3 - _4q2*f4 + _2q1*f5 + _2q0*f6
        s3 = _2q1*f1 + _2q2*f2 - _4q3*f4 - _2q0*f5 + _2q1*f6
        recipNorm = 1/ (math.sqrt(s0 * s0 + s1 * s1 + s2 * s2 + s3 * s3)) ## normalise step magnitude
        s0 *= recipNorm
        s1 *= recipNorm
        s2 *= recipNorm
        s3 *= recipNorm

        # Apply feedback step
        qDot1 -= beta * s0
        qDot2 -= beta * s1
        qDot3 -= beta * s2
        qDot4 -= beta * s3

    # Integrate rate of change of quaternion to yield quaternion
    q0 += qDot1 * dt
    q1 += qDot2 * dt
    q2 += qDot3 * dt
    q3 += qDot4 * dt

    # Normalise quaternion
    recipNorm = 1 / (math.sqrt(q0 * q0 + q1 * q1 + q2 * q2 + q3 * q3))
    q0 *= recipNorm
    q1 *= recipNorm
    q2 *= recipNorm
    q3 *= recipNorm

    q = np.array([q0, q1, q2, q3])

    # print('q0: {0}, q1: {1}, q2: {2}, q3: {3}'.format(q[0], q[1], q[2], q[3]))p

    return q

def computeAngles(q0, q1, q2, q3):

    roll = 180*math.atan2(q0*q1 + q2*q3, 0.5 - q1*q1 - q2*q2)/math.pi
    pitch = 180*math.asin(-2.0 * (q1*q3 - q0*q2))/math.pi
    yaw = 180*math.atan2(q1*q2 + q0*q3, 0.5 - q2*q2 - q3*q3)/math.pi

    return roll, pitch, yaw

def getData():
    time.sleep(1.0)
    serialConnection.reset_input_buffer()

    while(isRun):

        global isReceive
        global roll, pitch, yaw
        global q0, q1, q2, q3
        global accel
        global ax, az, ay, gx, gy, gz

        data = serialConnection.readline()
        data = str(data, 'utf-8')
        splitdata = data.split(',')
        q0 = float(splitdata[0])
        q1 = float(splitdata[1])
        q2 = float(splitdata[2])
        q3 = float(splitdata[3])
        ax = float(splitdata[4])
        ay = float(splitdata[5])
        az = float(splitdata[6])
        gx = float(splitdata[7])
        gy = float(splitdata[8])
        gz = float(splitdata[9])

        accel = np.array([[ax], [ay], [az]])

        roll, pitch, yaw = computeAngles(q0, q1, q2, q3)

        isReceive = True

serialPort = 'COM5'
baudRate = 9600

try:
    serialConnection = serial.Serial(serialPort, baudRate)
except:
    print('Cannot connect to the port')

thread = Thread(target=getData)
thread.start()

#Load the predefinied dictionary
dictionary = cv.aruco.Dictionary_get(cv.aruco.DICT_4X4_50)

# Load previously saved calibration data
path = './camera_data/camera_calibration.npz'
npzfile = np.load(path)
#Camera Matrix
mtx = npzfile[npzfile.files[0]]
#Distortion Matrix
dist = npzfile[npzfile.files[1]]
print(mtx ,dist)

#Font setup
font = cv.FONT_HERSHEY_PLAIN

start_time = time.time()

#Create position data.txt
pos_data = open("pos_data.txt",'w')
pos_data.close()

#Create attitude data.txt
angle_data = open("angle_data.txt", 'w')
angle_data.close()

#Camera instance with thread
cap = WebcamVideoStream(src=0).start()
# cap = cv.VideoCapture(0)

#Create board object
# board = cv.aruco.GridBoard_create(2, 2, 0.037, 0.005, dictionary)

start_clock = time.perf_counter()

#Declare somes important variables
rvecs = None
tvecs = None

pos_calib_obj = True

offset_xo = 0
offset_yo = 0
offset_zo = 0

frame_id = 0

roll_o = []
pitch_o = []
yaw_o = []

roll_c = []
pitch_c = []
yaw_c = []

pos_xo = []
pos_yo = []
pos_zo = []

roll_obj = 0
pitch_obj = 0
yaw_obj = 0

alfa = 0
beta = 0
gama = 0

while True:
    img = cap.read()

    # img = cv.rotate(img, cv.ROTATE_180)
    #Convert frame to gray scale
    gray = cv.cvtColor(img, cv.COLOR_BGR2GRAY)
    #Frame count increment
    frame_id += 1

    #Set parameters for the marker tracking
    parameters = cv.aruco.DetectorParameters_create()
    parameters.minMarkerPerimeterRate = 0.1
    # parameters.minDistanceToBorder = 6
    # parameters.adaptiveThreshConstant = 25
    parameters.cornerRefinementWinSize = 3
    parameters.cornerRefinementMethod = cv.aruco.CORNER_REFINE_CONTOUR
    # parameters.cornerRefinementMaxIterations = 10
    parameters.cornerRefinementMinAccuracy = 0

    #Detect the markers in the image
    markerCorners, markerIDs, rejectedCandidates = cv.aruco.detectMarkers(gray, dictionary, parameters = parameters)

    #Refine detection
    # markerCorners, markerIDs, rejectedCandidates, recoveredIds = cv.aruco.refineDetectedMarkers(img, board, markerCorners,
                                                                                # markerIDs, rejectedCandidates, mtx, dist)
    #Open position file to write data position
    pos_data = open("pos_data.txt", "a+")

    #Open attitude file to write estimation
    angle_data = open("angle_data.txt", "a+")

    # print('ID:', markerIDs)
    #Verify if there is some marker detected
    if markerIDs is not None and len(markerIDs)==2:

        for i in range(0, len(markerCorners)):
            # print("Marker Detected")
            #Compute board's pose

            #Reference marker
            if markerIDs[i]==10:
                # pose, rvecs, tvecs = cv.aruco.estimatePoseBoard(markerCorners, markerIDs, board, mtx, dist, rvecs, tvecs)
                rvec_ref, tvec_ref, _ = cv.aruco.estimatePoseSingleMarkers(markerCorners[i], 0.066, mtx, dist)
                rvec_ref = np.reshape(rvec_ref, (3,1))
                tvec_ref = np.reshape(tvec_ref, (3,1))

                #Use Rodrigues formula to transform rotation vector into matrix
                #Pose marker w.r.t camera reference frame
                R_rc, _ = cv.Rodrigues(rvec_ref)

                #Homogeneous Transformation Fixed Frame to Camera Frame
                last_col = np.array([[0, 0, 0, 1]])
                T_rc = np.concatenate((R_rc, tvec_ref), axis=1)
                T_rc = np.concatenate((T_rc, last_col), axis=0)
                #Homegeneous Transformation Camera Frame to Fixed Frame
                T_cr = np.linalg.inv(T_rc)

                #Euler angles
                # r_ref = sci.spatial.transform.Rotation.from_matrix(T_cr[0:3, 0:3])
                # euler_ref = r_ref.as_euler('ZYX', degrees=True)
                # phi_ref = euler_ref[0]
                # theta_ref = euler_ref[1]
                # psi_ref = euler_ref[2]

                # print("phi:", phi_ref)
                # print("theta:", theta_ref)
                # print("psi:", psi_ref)

                #Attitude estimation
                # phi = 180*math.atan2(R_marker[2,1], R_marker[2,2])/math.pi
                # theta = 180*math.atan2(-R_matrix[2, 0], math.sqrt(R_matrix[2,1]**2 + R_matrix[2,2]**2))/math.pi
                # psi = 180*math.atan2(R_matrix[1,0], R_matrix[0,0])/math.pi

                cv.aruco.drawAxis(img, mtx, dist, rvec_ref, tvec_ref, 0.033)
                cv.aruco.drawDetectedMarkers(img, markerCorners)

            #Object marker
            if markerIDs[i]==4:

                # pose, rvecs, tvecs = cv.aruco.estimatePoseBoard(markerCorners, markerIDs, board, mtx, dist, rvecs, tvecs)
                rvec_obj, tvec_obj, _ = cv.aruco.estimatePoseSingleMarkers(markerCorners[i], 0.066, mtx, dist)
                rvec_obj = np.reshape(rvec_obj, (3,1))
                tvec_obj = np.reshape(tvec_obj, (3,1))

                #Use Rodrigues formula to transform rotation vector into matrix
                R_dc, _ = cv.Rodrigues(rvec_obj)

                #Homogeneous Transformation Object Frame to Camera Frame
                T_dc = np.concatenate((R_dc, tvec_obj), axis=1)
                T_dc = np.concatenate((T_dc, last_col), axis=0)

                #Homogeneous Transformation Object Frame to Fixed Frame
                T_dr = T_cr@T_dc

                #Getting quaternions from rotation matrix
                r_obj = sci.spatial.transform.Rotation.from_matrix(T_dr[0:3, 0:3])
                q_obj = r_obj.as_quat()
                # # print('Quaternion:')
                # # print(q)

                roll_obj, pitch_obj, yaw_obj = computeAngles(q_obj[3], q_obj[1], q_obj[0], -q_obj[2])
                roll_obj -= 8.2
                pitch_obj += 7.2

                # euler_obj = r_obj.as_euler('ZYX', degrees=True)
                # phi_obj = euler_obj[0]
                # theta_obj = euler_obj[1]
                # psi_obj = euler_obj[2]

                # print("phi:", phi_obj)
                # print("theta:", theta_obj)
                # print("psi:", psi_obj)

                #Position correction

                # xf_obj = 0.755*float(T_dr[0,3]) - 0.193
                # yf_obj = 1.06*float(T_dr[1,3]) + 0.018
                # zf_obj = 0.709*float(T_dr[2,3]) + 0.0687

                xf_obj = 0.774*float(T_dr[0,3]) - 0.129
                yf_obj = -0.349*float(T_dr[1,3])**2 + 1.09*float(T_dr[1,3]) + 0.129
                zf_obj = 0.767*float(T_dr[2,3]) - 0.0694

                # if pos_calib_obj:
                #     #Position correction
                #     pos_xo.append(xf_obj)
                #     pos_yo.append(yf_obj)
                #     pos_zo.append(zf_obj)

                #     if len(pos_xo) == 100 and len(pos_yo) == 100 and len(pos_zo)==100:

                #         # print('X: {0}, Y: {1}, Z: {2}'.format(np.average(pos_xo), np.average(pos_yo), np.average(pos_zo)))

                #         pos_xo=[]
                #         pos_yo=[]
                #         pos_zo=[]

                cv.aruco.drawAxis(img, mtx, dist, rvec_obj, tvec_obj, 0.033)
                print(tvec_obj)
                cv.aruco.drawDetectedMarkers(img, markerCorners)

                modulo = np.sqrt(q_obj[0]**2 + q_obj[1]**2 + q_obj[2]**2 + q_obj[3]**2)
                #Print position values in frame
                cv.putText(img, "X:"+str(np.round(float(xf_obj), 4)), (80,600), font, 1, (0,0,0), 2)
                cv.putText(img, "Y:"+str(np.round(float(yf_obj), 4)), (180,600), font, 1, (0,0,0), 2)
                cv.putText(img, "Z:"+str(np.round(float(zf_obj), 4)), (280,600), font, 1, (0,0,0), 2)

                cv.putText(img, "Orientacao Estimada por Camera:", (10, 200), font, 1, (255, 255, 255), 2)
                cv.putText(img, "Phi:"+str(np.round(float(roll_obj), 2)), (10,220), font, 1, (0,0,255), 2)
                cv.putText(img, "Theta:"+str(np.round(float(pitch_obj), 2)), (10,240), font, 1, (0,255,0), 2)
                cv.putText(img, "Psi:"+str(np.round(float(yaw_obj), 2)), (10,260), font, 1, (255,0,0), 2)
                # cv.putText(img, "q0:"+str(np.round(float(q_obj[3]), 3)), (500,120), font, 1, (255,255,255), 2)
                # cv.putText(img, "q1:"+str(np.round(float(q_obj[0]), 3)), (500,140), font, 1, (255,255,255), 2)
                # cv.putText(img, "q2:"+str(np.round(float(q_obj[1]), 3)), (500,160), font, 1, (255,255,255), 2)
                # cv.putText(img, "q3:"+str(np.round(float(q_obj[2]), 3)), (500,180), font, 1, (255,255,255), 2)
                # cv.putText(img, "Modulo:" + str(np.round(float(modulo), 3)), (500, 210), font, 1, (255, 0, 0),2)

            #Data sensor
            if isReceive:

                # cv.putText(img, "q0:"+str(np.round(float(q0), 3)), (500,320), font, 1, (0,0,0), 2)
                # cv.putText(img, "q1:"+str(np.round(float(q1), 3)), (500,340), font, 1, (0,0,0), 2)
                # cv.putText(img, "q2:"+str(np.round(float(q2), 3)), (500,360), font, 1, (0,0,0), 2)
                # cv.putText(img, "q3:"+str(np.round(float(q3), 3)), (500,380), font, 1, (0,0,0), 2)

                cv.putText(img, "Orientacao Estimada por IMU:", (10, 300), font, 1, (255,255,255), 2)
                cv.putText(img, "Roll:"+str(np.round(float(roll), 3)), (10,320), font, 1, (0,0,255), 2)
                cv.putText(img, "Pitch:"+str(np.round(float(pitch), 3)), (10,340), font, 1, (0,255,0), 2)
                cv.putText(img, "Psi:"+str(np.round(float(yaw), 3)), (10,360), font, 1, (255,0,0), 2)

                #########
                roll_c.append(roll_obj)
                pitch_c.append(pitch_obj)
                yaw_c.append(yaw_obj)

                if len(roll_c) == 100 and len(pitch_c) == 100 and len(yaw_c)==100:

                    print('Roll_C: {0}, Pitch_C: {1}, Yaw_C: {2}'.format(np.average(roll_c), np.average(pitch_c), np.average(yaw_c)))

                    roll_c=[]
                    pitch_c=[]
                    yaw_c=[]

                ######
                roll_o.append(roll)
                pitch_o.append(pitch)
                yaw_o.append(yaw)

                if len(roll_o) == 100 and len(pitch_o) == 100 and len(yaw_o)==100:

                    print('Roll: {0}, Pitch: {1}, Yaw: {2}'.format(np.average(roll_o), np.average(pitch_o), np.average(yaw_o)))

                    roll_o=[]
                    pitch_o=[]
                    yaw_o=[]

                #Kalman Filter - Orientation Estimation

                y_ori = np.array([[roll],[pitch],[yaw_obj],[gx],[gy],[gz]])
                x_est, P_est = KF_Orientation(x_ant_ori, P_ant_ori, y_ori, dt)
                x_ant_ori = x_est
                P_ant_ori = P_est

                roll_est = x_est[0]
                pitch_est = x_est[1]
                yaw_est = x_est[2]

                cv.putText(img, "Orientacao Estimada por Filtro:", (10, 400), font, 1, (255,255,255), 2)
                cv.putText(img, "Roll:"+str(np.round(float(roll_est), 3)), (10,420), font, 1, (0,0,255), 2)
                cv.putText(img, "Pitch:"+str(np.round(float(pitch_est), 3)), (10,440), font, 1, (0,255,0), 2)
                cv.putText(img, "Psi:"+str(np.round(float(yaw_est), 3)), (10,460), font, 1, (255,0,0), 2)

                #Kalman Filter - Position Estimation
                R_IMU = R.from_euler('zyx',[[float(yaw_est), float(pitch_est), float(roll_est)]], degrees=True).as_matrix()

                y_pos = np.array([[xf_obj],[yf_obj],[zf_obj]])
                u_pos = R_IMU@np.array([[ax],[ay],[az]]) - np.array([[0],[0],[9.81]])
                # print(u_pos)
                pos_est, P_pos = KF_Position(x_ant_pos, P_ant_pos, u_pos.reshape(3,1), y_pos, dt)
                x_ant_pos = pos_est
                P_ant_pos = P_pos

                x_kf = pos_est[0]
                y_kf = pos_est[1]
                z_kf = pos_est[2]

                cv.putText(img, "X:"+str(np.round(float(x_kf), 4)), (80,700), font, 1, (0,0,0), 2)
                cv.putText(img, "Y:"+str(np.round(float(y_kf), 4)), (180,700), font, 1, (0,0,0), 2)
                cv.putText(img, "Z:"+str(np.round(float(z_kf), 4)), (280,700), font, 1, (0,0,0), 2)

                #Matriz de rotação dinâmica referencia
                if elapsed_time > 10:
                    if gama < np.pi/4:

                        gama += np.pi/800

                    elif beta < np.pi/6:

                        beta += np.pi/800

                R_ref = np.array([[c(alfa)*c(gama) - s(alfa)*c(beta)*s(gama), s(alfa)*c(gama) + c(alfa)*c(beta)*s(gama), s(beta)*s(gama)],
                                [-c(alfa)*s(gama) - s(alfa)*c(beta)*c(gama), -s(alfa)*s(gama)+c(alfa)*c(beta)*c(gama), s(beta)*c(gama)],
                                [s(alfa)*s(beta), -c(alfa)*s(beta), c(beta)]])

                R_ref = T_rc[0:3, 0:3]@R_ref

                rvec_ref,_ = cv.Rodrigues(R_ref)

                #Getting quaternions from rotation matrix
                r_refe = sci.spatial.transform.Rotation.from_matrix(T_cr[0:3,0:3]@R_ref)
                q_ref = r_refe.as_quat()
                # # print('Quaternion:')
                # # print(q)

                roll_ref, pitch_ref, yaw_ref = computeAngles(q_ref[3], q_ref[1], q_ref[0], -q_ref[2])

                cv.drawFrameAxes(img, mtx, dist, rvec_ref, np.array([[-0.2],[-0.1],[1.6]]), 0.066, )

                cv.putText(img, "Orientacao Dinamica de Referencia:", (500, 200), font, 1, (255, 255, 255), 2)
                cv.putText(img, "Phi:"+str(np.round(float(roll_ref), 2)), (500,220), font, 1, (0,0,255), 2)
                cv.putText(img, "Theta:"+str(np.round(float(pitch_ref), 2)), (500,240), font, 1, (0,255,0), 2)
                cv.putText(img, "Psi:"+str(np.round(float(yaw_ref), 2)), (500,260), font, 1, (255,0,0), 2)

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

                #Write position data in pos_data.txt
                pos_data.write("{:.4f} , ".format(float(xf_obj)) + "{:.4f} , ".format(float(yf_obj)) + "{:.4f} , ".format(float(zf_obj)) + str(time.perf_counter()-start_clock) + "\n")
                #Write attitude data in file
                angle_data.write("{:.2f} , ".format(float(roll_ref)) + "{:.2f} , ".format(float(pitch_ref)) + "{:.2f} , ".format(float(yaw_ref))
                                + "{:.2f} , ".format(float(roll_est)) + "{:.2f} , ".format(float(pitch_est)) + "{:.2f} , ".format(float(yaw_est)) +  str(time.perf_counter()-start_clock) + "\n")

    #Compute FPS
    elapsed_time = time.perf_counter() - start_clock
    print("Start time: {0} \n Elapsed Time: {1}".format(start_clock, elapsed_time))
    fps = int(frame_id / elapsed_time)
    #Print FPS on the screen
    cv.putText(img, "FPS:" + str(fps), (10, 80), font, 1, (255,255,255), 1)
    cv.putText(img, "Time Elapsed:" + str(round(elapsed_time,1)), (100, 80), font, 1, (255,255,255), 1)
    cv.imshow('img', img)

    key = cv.waitKey(10)
    if key == ord('n') or key == ord('p'):
        break

pos_data.close()
angle_data.close()
cap.stop()
# cap.release()
cv.destroyAllWindows()

isRun = False
thread.join()
serialConnection.close()

plt.style.use("fivethirtyeight")
fig, (ax, ax2, ax3) = plt.subplots(3, 1, figsize=(10,10), sharex=True)
rr_list, pr_list, yr_list, re_list, pe_list, ye_list, time = [], [], [], [], [], [], []

data = open("angle_data.txt", "r").read()
lines = data.split('\n')

for line in lines:
    if len(line)>1:
        rr, pr, yr, re, pe, ye, times = line.split(' , ')
        rr_list.append(float(rr))
        pr_list.append(float(pr))
        yr_list.append(float(yr))
        re_list.append(float(re))
        pe_list.append(float(pe))
        ye_list.append(float(ye))
        time.append(float(times))

ax.plot(time, rr_list, 'r--', alpha=0.6, label = r'$\phi_{ref} (t)$', linewidth = 1.5)
ax.plot(time, re_list, 'r', label=r'$\phi_{est}(t)$', linewidth=1)
ax2.plot(time, pr_list, 'g--', alpha=0.6, label = r'$\theta_{ref} (t)$', linewidth=1.5)
ax2.plot(time, pe_list, 'g', label=r'$\theta_{est}(t)$', linewidth=1)
ax3.plot(time, yr_list, 'b--', alpha=0.6, label = r'$\psi_{ref} (t)$', linewidth=1.5)
ax3.plot(time, ye_list, 'b', label=r'$\psi_{est}(t)$', linewidth=1)
ax3.set_xlabel('Time (s)')

ax.set_ylabel(r'$\phi$ (°)')
ax2.set_ylabel(r'$\theta$ (°)')
ax3.set_ylabel(r'$\psi$ (°)')
ax.set_title('Orientação')

ax.set_ylim([-90, 90])
ax2.set_ylim([-90, 90])
ax3.set_ylim([-90, 90])
ax.legend()
ax2.legend()
ax3.legend()

plt.tight_layout()
plt.show()

max_stars_repo_path: ComputerVision/ArUco detection.py

max_stars_repo_name: mateusribs/DissertacaoMestrado

max_issues_repo_path: ComputerVision/ArUco detection.py

max_issues_repo_name: mateusribs/DissertacaoMestrado

max_forks_repo_path: ComputerVision/ArUco detection.py

max_forks_repo_name: mateusribs/DissertacaoMestrado

reason: import numpy,from numpy,import scipy,from scipy


text:
#!/usr/bin/env python
# encoding: utf-8

"""
@Author: yangwenhao
@Contact: 874681044@qq.com
@Software: PyCharm
@File: output_extract.py
@Time: 2020/3/21 5:57 PM
@Overview:
"""
from __future__ import print_function

import argparse
import json
import os
import pickle
import random
import time
from collections import OrderedDict
import gradients.constants as c

import numpy as np
import torch
import torch._utils
import torch.backends.cudnn as cudnn
import torch.nn as nn
import torchvision.transforms as transforms
from kaldi_io import read_mat
from torch.autograd import Variable
from torch.utils.data import DataLoader
from tqdm import tqdm

from dataset.Datasets import ScriptTrainDataset, ScriptValidDataset
from dataset.audio_processing import ConcateOrgInput, ConcateVarInput
from common_func import create_model

# Version conflict

try:
    torch._utils._rebuild_tensor_v2
except AttributeError:
    def _rebuild_tensor_v2(storage, storage_offset, size, stride, requires_grad, backward_hooks):
        tensor = torch._utils._rebuild_tensor(storage, storage_offset, size, stride)
        tensor.requires_grad = requires_grad
        tensor._backward_hooks = backward_hooks
        return tensor

    torch._utils._rebuild_tensor_v2 = _rebuild_tensor_v2
import warnings

warnings.filterwarnings("ignore")

# Training settings
parser = argparse.ArgumentParser(description='PyTorch Speaker Recognition: Class Activation Mapping')
# Data options
parser.add_argument('--train-dir', type=str, help='path to dataset')
parser.add_argument('--test-dir', type=str, help='path to voxceleb1 test dataset')
parser.add_argument('--train-set-name', type=str, required=True, help='path to voxceleb1 test dataset')
parser.add_argument('--test-set-name', type=str, required=True, help='path to voxceleb1 test dataset')
parser.add_argument('--sitw-dir', type=str, help='path to voxceleb1 test dataset')
parser.add_argument('--sample-utt', type=int, default=120, metavar='SU', help='Dimensionality of the embedding')
parser.add_argument('--test-only', action='store_true', default=False, help='using Cosine similarity')
parser.add_argument('--check-path', help='folder to output model checkpoints')
parser.add_argument('--extract-path', help='folder to output model grads, etc')

parser.add_argument('--start-epochs', type=int, default=36, metavar='E', help='number of epochs to train (default: 10)')
parser.add_argument('--epochs', type=int, default=36, metavar='E', help='number of epochs to train (default: 10)')

# Data options
parser.add_argument('--feat-dim', default=64, type=int, metavar='N', help='acoustic feature dimension')
parser.add_argument('--input-dim', default=257, type=int, metavar='N', help='acoustic feature dimension')

parser.add_argument('--revert', action='store_true', default=False, help='using Cosine similarity')
parser.add_argument('--input-length', choices=['var', 'fix'], default='var',
                    help='choose the acoustic features type.')
parser.add_argument('--remove-vad', action='store_true', default=False, help='using Cosine similarity')
parser.add_argument('--mvnorm', action='store_true', default=False,
                    help='using Cosine similarity')
# Model options
parser.add_argument('--model', type=str, help='path to voxceleb1 test dataset')
parser.add_argument('--cam', type=str, default='gradient', help='path to voxceleb1 test dataset')
parser.add_argument('--cam-layers',
                    default=['conv1', 'layer1.0.conv2', 'conv2', 'layer2.0.conv2', 'conv3', 'layer3.0.conv2'],
                    type=list, metavar='CAML', help='The channels of convs layers)')
parser.add_argument('--resnet-size', default=8, type=int, metavar='RES', help='The channels of convs layers)')

parser.add_argument('--filter', type=str, default='None', help='replace batchnorm with instance norm')
parser.add_argument('--input-norm', type=str, default='Mean', help='batchnorm with instance norm')
parser.add_argument('--vad', action='store_true', default=False, help='vad layers')
parser.add_argument('--inception', action='store_true', default=False, help='multi size conv layer')
parser.add_argument('--inst-norm', action='store_true', default=False, help='batchnorm with instance norm')

parser.add_argument('--mask-layer', type=str, default='None', help='time or freq masking layers')
parser.add_argument('--mask-len', type=int, default=20, help='maximum length of time or freq masking layers')
parser.add_argument('--block-type', type=str, default='None', help='replace batchnorm with instance norm')
parser.add_argument('--relu-type', type=str, default='relu', help='replace batchnorm with instance norm')
parser.add_argument('--encoder-type', type=str, help='path to voxceleb1 test dataset')
parser.add_argument('--transform', type=str, default="None", help='add a transform layer after embedding layer')

parser.add_argument('--channels', default='64,128,256', type=str,
                    metavar='CHA', help='The channels of convs layers)')
parser.add_argument('--fast', type=str, default='None', help='max pooling for fast')

parser.add_argument('--kernel-size', default='5,5', type=str, metavar='KE',
                    help='kernel size of conv filters')
parser.add_argument('--padding', default='', type=str, metavar='KE', help='padding size of conv filters')
parser.add_argument('--stride', default='2', type=str, metavar='ST', help='stride size of conv filters')
parser.add_argument('--time-dim', default=1, type=int, metavar='FEAT', help='acoustic feature dimension')
parser.add_argument('--avg-size', type=int, default=4, metavar='ES', help='Dimensionality of the embedding')

parser.add_argument('--loss-type', type=str, default='soft', help='path to voxceleb1 test dataset')
parser.add_argument('--dropout-p', type=float, default=0., metavar='BST',
                    help='input batch size for testing (default: 64)')

# args for additive margin-softmax
parser.add_argument('--margin', type=float, default=0.3, metavar='MARGIN',
                    help='the margin value for the angualr softmax loss function (default: 3.0')
parser.add_argument('--s', type=float, default=15, metavar='S',
                    help='the margin value for the angualr softmax loss function (default: 3.0')
# args for a-softmax
parser.add_argument('--m', type=int, default=3, metavar='M',
                    help='the margin value for the angualr softmax loss function (default: 3.0')
parser.add_argument('--lambda-min', type=int, default=5, metavar='S',
                    help='random seed (default: 0)')
parser.add_argument('--lambda-max', type=float, default=0.05, metavar='S',
                    help='random seed (default: 0)')

parser.add_argument('--alpha', default=0, type=float,
                    metavar='l2 length', help='acoustic feature dimension')
parser.add_argument('--cos-sim', action='store_true', default=True, help='using Cosine similarity')
parser.add_argument('--embedding-size', type=int, metavar='ES', help='Dimensionality of the embedding')

parser.add_argument('--nj', default=12, type=int, metavar='NJOB', help='num of job')
parser.add_argument('--batch-size', type=int, default=1, metavar='BS',
                    help='input batch size for training (default: 128)')
parser.add_argument('--test-batch-size', type=int, default=1, metavar='BST',
                    help='input batch size for testing (default: 64)')
parser.add_argument('--input-per-spks', type=int, default=192, metavar='IPFT',
                    help='input sample per file for testing (default: 8)')
parser.add_argument('--test-input-per-file', type=int, default=1, metavar='IPFT',
                    help='input sample per file for testing (default: 8)')

# Device options
parser.add_argument('--no-cuda', action='store_true', default=False,
                    help='enables CUDA training')
parser.add_argument('--gpu-id', default='1', type=str,
                    help='id(s) for CUDA_VISIBLE_DEVICES')
parser.add_argument('--seed', type=int, default=123456, metavar='S',
                    help='random seed (default: 0)')
parser.add_argument('--log-interval', type=int, default=1, metavar='LI',
                    help='how many batches to wait before logging training status')

parser.add_argument('--acoustic-feature', choices=['fbank', 'spectrogram', 'mfcc'], default='fbank',
                    help='choose the acoustic features type.')
parser.add_argument('--makemfb', action='store_true', default=False,
                    help='need to make mfb file')
parser.add_argument('--makespec', action='store_true', default=False,
                    help='need to make spectrograms file')

args = parser.parse_args()

# Set the device to use by setting CUDA_VISIBLE_DEVICES env variable in
# order to prevent any memory allocation on unused GPUs
os.environ['CUDA_VISIBLE_DEVICES'] = args.gpu_id

args.cuda = not args.no_cuda and torch.cuda.is_available()
random.seed(args.seed)
np.random.seed(args.seed)
torch.manual_seed(args.seed)
torch.multiprocessing.set_sharing_strategy('file_system')

if args.cuda:
    cudnn.benchmark = True

# Define visulaize SummaryWriter instance
kwargs = {'num_workers': args.nj, 'pin_memory': False} if args.cuda else {}
l2_dist = nn.CosineSimilarity(dim=1, eps=1e-6) if args.cos_sim else nn.PairwiseDistance(p=2)

if args.input_length == 'var':
    transform = transforms.Compose([
        ConcateOrgInput(remove_vad=args.remove_vad),
    ])
elif args.input_length == 'fix':
    transform = transforms.Compose([
        ConcateVarInput(remove_vad=args.remove_vad),
    ])

# if args.mvnorm:
#     transform.transforms.append(mvnormal())

file_loader = read_mat

train_dir = ScriptTrainDataset(dir=args.train_dir, samples_per_speaker=args.input_per_spks,
                               loader=file_loader, transform=transform, return_uid=True)
indices = list(range(len(train_dir)))
random.shuffle(indices)
indices = indices[:args.sample_utt]
train_part = torch.utils.data.Subset(train_dir, indices)

valid_dir = ScriptValidDataset(valid_set=train_dir.valid_set, spk_to_idx=train_dir.spk_to_idx,
                               valid_uid2feat=train_dir.valid_uid2feat, valid_utt2spk_dict=train_dir.valid_utt2spk_dict,
                               loader=file_loader, transform=transform, return_uid=True)
indices = list(range(len(valid_dir)))
random.shuffle(indices)
indices = indices[:args.sample_utt]
valid_part = torch.utils.data.Subset(valid_dir, indices)

cam_layers = args.cam_layers

out_feature_grads = []
in_feature_grads = []
in_layer_feat = []
out_layer_feat = []
in_layer_grad = []
out_layer_grad = []

bias_layers = []
biases = []
handlers = []

def extract_layer_bias(module):
    # extract bias of each layer
    if isinstance(module, torch.nn.BatchNorm2d):
        #         print('bn2')
        b = - (module.running_mean * module.weight
               / torch.sqrt(module.running_var + module.eps)) + module.bias
        return b.data
    elif module.bias is None:
        return None
    else:
        return module.bias.data

def _extract_layer_grads(module, in_grad, out_grad):
    # function to collect the gradient outputs
    # from each layer
#     print(module._get_name())
#     print('Input_grad shape:', in_grad[0].shape)
#     print('Output_grad shape:', out_grad[0].shape)
    global in_feature_grads
    global out_feature_grads
    if not module.bias is None:
        in_feature_grads.append(in_grad[0])
        out_feature_grads.append(out_grad[0])

def _extract_layer_feat(module, input, output):
    # function to collect the gradient outputs from each layer
    #     print(module._get_name())
    #     print('Input_grad shape:', in_grad[0].shape)
    #     print('Output_grad shape:', out_grad[0].shape)
    #     if not module.bias is None:
    global in_layer_feat
    global in_layer_feat
    in_layer_feat.append(input[0])
    out_layer_feat.append(output[0])

def _extract_layer_grad(module, in_grad, out_grad):
    # function to collect the gradient outputs from each layer
    #     print(module._get_name())
    #     print('Input_grad shape:', in_grad[0].shape)
    #     print('Output_grad shape:', out_grad[0].shape)
    global in_layer_grad
    global out_layer_grad
    in_layer_grad.append(in_grad[0])
    out_layer_grad.append(out_grad[0])

def train_extract(train_loader, model, file_dir, set_name, save_per_num=2500):
    # switch to evaluate mode
    model.eval()

    input_grads = []
    inputs_uids = []
    pbar = tqdm(enumerate(train_loader))
    global out_feature_grads
    global in_feature_grads
    global in_layer_feat
    global out_layer_feat
    global in_layer_grad
    global out_layer_grad

    for batch_idx, (data, label, uid) in pbar:

        # orig = data.detach().numpy().squeeze().astype(np.float32)

        if data.shape[2] > 5 * c.NUM_FRAMES_SPECT:
            num_half = int(data.shape[2] / (4 * c.NUM_FRAMES_SPECT))
            rest_frame = data.shape[2] % num_half
            if rest_frame > 0:
                x = data[:, :, :-rest_frame, :].chunk(num_half, dim=2)
            else:
                x = data.chunk(num_half, dim=2)
            data = torch.cat(x, dim=0)

        data = Variable(data.cuda(), requires_grad=True)
        ups = torch.nn.UpsamplingBilinear2d(size=data.shape[-2:])

        if len(data) == 1:
            logit, _ = model(data)

            if args.loss_type == 'asoft':
                classifed, _ = logit
            else:
                classifed = logit

            try:
                classifed[0][label.long()].backward()
            except Exception as e:
                print(data.shape, ",", uid)
                raise e

            if args.cam == 'gradient':
                grad = data.grad  # .cpu().numpy().squeeze().astype(np.float32)

            elif args.cam == 'grad_cam':
                grad = torch.zeros_like(data)
                assert len(out_layer_grad) == len(cam_layers), print(len(out_layer_grad), " is not equal to ", len(cam_layers))

                last_grad = out_layer_grad[0]
                last_feat = out_layer_feat[-1]

                weight = last_grad.mean(dim=(2, 3), keepdim=True)
                weight /= weight.sum()
                feat = last_feat  # .copy()

                T = (feat * weight).clamp_min(0).sum(dim=1, keepdim=True)  # .clamp_min(0)
                grad += ups(T).abs()

            elif args.cam == 'grad_cam_pp':
                # grad cam ++ last
                last_grad = out_layer_grad[0].cpu()
                last_feat = out_layer_feat[-1].cpu()
                first_derivative = classifed[0][label.long()].exp().cpu() * last_grad
                alpha = last_grad.pow(2) / (
                        2 * last_grad.pow(2) + (last_grad.pow(3) * last_feat).sum(dim=(2, 3), keepdim=True))
                # weight = alpha * (first_derivative.clamp_min_(0))
                weight = alpha * (first_derivative.abs())

                weight = weight.mean(dim=(2, 3), keepdim=True)
                weight /= weight.sum()

                grad = (last_feat * weight).sum(dim=1, keepdim=True)
                grad = ups(grad)

                # grad_cam_pp -= grad_cam_pp.min()
                grad = grad.abs()
                grad /= grad.max()

            elif args.cam == 'fullgrad':
                # full grad
                input_gradient = (data.grad * data)
                full_grad = input_gradient.clone().clamp_min(0)

                L = len(bias_layers)

                for i, l in enumerate(bias_layers):
                    bias = biases[L - i - 1]
                    if len(bias.shape) == 1:
                        bias = bias.reshape(1, -1, 1, 1)
                    bias = bias.expand_as(out_feature_grads[i])

                    #     bias_grad = (out_feature_grads[i]*bias).sum(dim=1, keepdim=True)
                    #     bias_grad = (out_feature_grads[i]*bias).mean(dim=1, keepdim=True)
                    bias_grad = (out_feature_grads[i] * bias).mean(dim=1, keepdim=True).clamp_min(0)
                    bias_grad /= bias_grad.max()
                    full_grad += ups(bias_grad)

                # full_grad -= full_grad.min()
                # full_grad = full_grad.abs()
                full_grad /= full_grad.max()
                grad = full_grad.detach().cpu()

        else:
            grad = []
            all_data = []
            for i in range(len(data)):
                out_feature_grads = []
                in_feature_grads = []
                in_layer_feat = []
                out_layer_feat = []
                in_layer_grad = []
                out_layer_grad = []

                data_a = data[i].unsqueeze(0)
                data_a = Variable(data_a.cuda(), requires_grad=True)

                logit, _ = model(data_a)

                if args.loss_type == 'asoft':
                    classifed, _ = logit
                else:
                    classifed = logit

                classifed[0][label.long()].backward()

                if args.cam == 'gradient':
                    grad_a = data_a.grad  # .cpu().numpy().squeeze().astype(np.float32)
                elif args.cam == 'grad_cam':
                    grad_a = torch.zeros_like(data_a)
                    L = len(cam_layers)
                    assert len(out_layer_grad) == L, print(len(out_layer_grad))

                    last_grad = out_layer_grad[0]
                    last_feat = out_layer_feat[-1]

                    weight = last_grad.mean(dim=(2, 3), keepdim=True)
                    weight /= weight.sum()
                    feat = last_feat  # .copy()

                    T = (feat * weight).clamp_min(0).sum(dim=1, keepdim=True)  # .clamp_min(0)

                    grad_a += ups(T).abs()

                elif args.cam == 'grad_cam_pp':
                    # grad cam ++ last
                    last_grad = out_layer_grad[0]
                    last_feat = out_layer_feat[-1]
                    first_derivative = classifed[0][label.long()].exp() * last_grad
                    alpha = last_grad.pow(2) / (
                            2 * last_grad.pow(2) + (last_grad.pow(3) * last_feat).sum(dim=(2, 3), keepdim=True))
                    # weight = alpha * (first_derivative.clamp_min_(0))
                    weight = alpha * (first_derivative.abs())

                    weight = weight.mean(dim=(2, 3), keepdim=True)
                    weight /= weight.sum()

                    grad_a = (last_feat * weight).sum(dim=1, keepdim=True)
                    grad_a = ups(grad_a)

                    # grad_cam_pp -= grad_cam_pp.min()
                    grad_a = grad_a.abs()
                    grad_a /= grad_a.max()

                elif args.cam == 'fullgrad':
                    # full grad
                    input_gradient = (data_a.grad * data_a)
                    full_grad = input_gradient.clone().clamp_min(0)

                    L = len(bias_layers)

                    for i, l in enumerate(bias_layers):
                        bias = biases[L - i - 1]
                        if len(bias.shape) == 1:
                            bias = bias.reshape(1, -1, 1, 1)
                        bias = bias.expand_as(out_feature_grads[i])

                        #     bias_grad = (out_feature_grads[i]*bias).sum(dim=1, keepdim=True)
                        #     bias_grad = (out_feature_grads[i]*bias).mean(dim=1, keepdim=True)
                        bias_grad = (out_feature_grads[i] * bias).mean(dim=1, keepdim=True).clamp_min(0)
                        bias_grad /= bias_grad.max()
                        full_grad += ups(bias_grad)

                    # full_grad -= full_grad.min()
                    # full_grad = full_grad.abs()
                    full_grad /= full_grad.max()
                    grad_a = full_grad

                grad.append(grad_a.detach().cpu().squeeze())
                all_data.append(data_a.detach().squeeze())

            grad = torch.cat(grad, dim=0)
            data = torch.cat(all_data, dim=0)

        out_feature_grads = []
        in_feature_grads = []
        in_layer_feat = []
        out_layer_feat = []
        in_layer_grad = []
        out_layer_grad = []

        grad = grad.cpu().detach().numpy().squeeze().astype(np.float32)
        data = data.data.cpu().numpy().squeeze().astype(np.float32)

        if args.revert:
            grad = grad.transpose()
            data = data.transpose()

        input_grads.append([data, grad])
        inputs_uids.append(uid)

        model.zero_grad()

        if batch_idx % args.log_interval == 0:
            pbar.set_description('Saving {} : [{:8d}/{:8d} ({:3.0f}%)] '.format(
                uid,
                batch_idx + 1,
                len(train_loader.dataset),
                100. * batch_idx / len(train_loader)))

        if (batch_idx + 1) % save_per_num == 0 or (batch_idx + 1) == len(train_loader.dataset):
            num = batch_idx // save_per_num if batch_idx + 1 % save_per_num == 0 else batch_idx // save_per_num + 1
            # checkpoint_dir / extract / < dataset > / < set >.*.bin

            filename = file_dir + '/%s.%d.bin' % (set_name, num)
            with open(filename, 'wb') as f:
                pickle.dump(input_grads, f)

            with open(file_dir + '/inputs.%s.%d.json' % (set_name, num), 'w') as f:
                json.dump(inputs_uids, f)

            input_grads = []
            inputs_uids = []

    print('Saving pairs in %s.\n' % file_dir)
    torch.cuda.empty_cache()

def main():
    print('\nNumber of Speakers: {}.'.format(train_dir.num_spks))
    # print the experiment configuration
    print('Current time is \33[91m{}\33[0m.'.format(str(time.asctime())))
    print('Parsed options: {}'.format(vars(args)))

    # instantiate model and initialize weights
    kernel_size = args.kernel_size.split(',')
    kernel_size = [int(x) for x in kernel_size]
    if args.padding == '':
        padding = [int((x - 1) / 2) for x in kernel_size]
    else:
        padding = args.padding.split(',')
        padding = [int(x) for x in padding]

    kernel_size = tuple(kernel_size)
    padding = tuple(padding)
    stride = args.stride.split(',')
    stride = [int(x) for x in stride]

    channels = args.channels.split(',')
    channels = [int(x) for x in channels]

    model_kwargs = {'input_dim': args.input_dim, 'feat_dim': args.feat_dim, 'kernel_size': kernel_size,
                    'mask': args.mask_layer, 'mask_len': args.mask_len, 'block_type': args.block_type,
                    'filter': args.filter, 'inst_norm': args.inst_norm, 'input_norm': args.input_norm,
                    'stride': stride, 'fast': args.fast, 'avg_size': args.avg_size, 'time_dim': args.time_dim,
                    'padding': padding, 'encoder_type': args.encoder_type, 'vad': args.vad,
                    'transform': args.transform, 'embedding_size': args.embedding_size, 'ince': args.inception,
                    'resnet_size': args.resnet_size, 'num_classes': train_dir.num_spks,
                    'channels': channels, 'alpha': args.alpha, 'dropout_p': args.dropout_p,
                    'loss_type': args.loss_type, 'm': args.m, 'margin': args.margin, 's': args.s,}

    print('Model options: {}'.format(model_kwargs))

    model = create_model(args.model, **model_kwargs)

    train_loader = DataLoader(train_part, batch_size=args.batch_size, shuffle=False, **kwargs)
    valid_loader = DataLoader(valid_part, batch_size=args.batch_size, shuffle=False, **kwargs)

    resume_path = args.check_path + '/checkpoint_{}.pth'
    print('=> Saving output in {}\n'.format(args.extract_path))
    epochs = np.arange(args.start_epochs, args.epochs + 1)

    for ep in epochs:
        # Load model from Checkpoint file
        if os.path.isfile(resume_path.format(ep)):
            print('=> loading checkpoint {}'.format(resume_path.format(ep)))
            checkpoint = torch.load(resume_path.format(ep))
            checkpoint_state_dict = checkpoint['state_dict']
            if isinstance(checkpoint_state_dict, tuple):
                checkpoint_state_dict = checkpoint_state_dict[0]

            filtered = {k: v for k, v in checkpoint_state_dict.items() if 'num_batches_tracked' not in k}
            if list(filtered.keys())[0].startswith('module'):
                new_state_dict = OrderedDict()
                for k, v in filtered.items():
                    name = k[7:]  # remove `module.`，表面从第7个key值字符取到最后一个字符，去掉module.
                    new_state_dict[name] = v  # 新字典的key值对应的value为一一对应的值。

                model.load_state_dict(new_state_dict)
            else:
                model_dict = model.state_dict()
                model_dict.update(filtered)
                model.load_state_dict(model_dict)

        else:
            print('=> no checkpoint found at %s' % resume_path.format(ep))
            continue
        model.cuda()

        # cam_layers = ['conv1', 'layer1.2.conv2']

        global bias_layers
        global biases
        global handlers

        for name, m in model.named_modules():
            try:
                if name in cam_layers:
                    handlers.append(m.register_forward_hook(_extract_layer_feat))
                    handlers.append(m.register_backward_hook(_extract_layer_grad))
                #         if isinstance(m, nn.Conv2d) or isinstance(m, nn.Linear) or isinstance(m, nn.BatchNorm2d):
                if not ('fc' in name or 'classifier' in name or 'CBAM' in name):
                    #             print(m)
                    b = extract_layer_bias(m)
                    if (b is not None):
                        biases.append(b)
                        bias_layers.append(name)
                        #                 biases.append(_extract_layer_bias(m))
                        #                 print("bias:", _extract_layer_bias(m))
                        m.register_backward_hook(_extract_layer_grads)
            #                     m.register_forward_hook(_extract_layer_feat)
            except Exception as e:
                continue
        print("The number of layers with biases: ", len(biases))

        file_dir = args.extract_path + '/epoch_%d' % ep
        if not os.path.exists(file_dir):
            os.makedirs(file_dir)

        if not args.test_only:

            train_extract(train_loader, model, file_dir, '%s_train' % args.train_set_name)
            train_extract(valid_loader, model, file_dir, '%s_valid' % args.train_set_name)

if __name__ == '__main__':
    main()

max_stars_repo_name: Wenhao-Yang/GradientFrequencyAttention

max_issues_repo_name: Wenhao-Yang/GradientFrequencyAttention

max_forks_repo_name: Wenhao-Yang/GradientFrequencyAttention


text:
import numpy as np
import pytest

import pfh.glidersim as gsim

@pytest.fixture
def harness():
    return gsim.paraglider_harness.Spherical(
        mass=1,
        z_riser=1,
        S=2,
        CD=1,
        kappa_w=1,
    )

def test_1(harness):
    # No drag, no weight
    f, g = harness.resultant_force(
        delta_w=0,
        v_W2h=[0, 0, 0],
        rho_air=0,
        g=[0, 0, 0],  # No gravity
        r_R2RM=[0, 0, 0],  # R = RM
    )
    assert np.allclose(f, [0, 0, 0])
    assert np.allclose(g, [0, 0, 0])

def test_2(harness):
    # Drag straight backwards, no weight
    f, g = harness.resultant_force(
        delta_w=0,
        v_W2h=[-1, 0, 0],
        rho_air=1,
        g=[0, 0, 0],  # No gravity
        r_R2RM=[0, 0, 0],  # R = RM
    )
    assert np.allclose(f, [-1, 0, 0])
    assert np.allclose(g, [0, -1, 0])

def test_3(harness):
    f, g = harness.resultant_force(
        delta_w=0,
        v_W2h=[0, 0, 0],
        rho_air=1,
        g=[0, 0, 1],
        r_R2RM=[0, 0, 0],
    )
    assert np.allclose(f, [0, 0, 1])
    assert np.allclose(g, [0, 0, 0])

def test_4(harness):
    f, g = harness.resultant_force(
        delta_w=1,
        v_W2h=[0, 0, 0],
        rho_air=1,
        g=[0, 0, 1],
        r_R2RM=[0, 0, 0],
    )
    assert np.allclose(f, [0, 0, 1])
    assert np.allclose(g, [1, 0, 0])

max_stars_repo_name: pfheatwole/glidersim

max_issues_repo_name: pfheatwole/glidersim

max_forks_repo_name: pfheatwole/glidersim


text:
'''
Units tests for the cpptraj wrappers for nasqm
'''
import os
import pytest
import numpy as np
from pynasqm.nmr.trajdistance import TrajDistance

def setup_module(module):
    '''
    Switch to test directory
    '''
    os.chdir("tests/trajDistance")

def teardown_module(module):
    '''
    Return to main directory
    '''
    os.chdir("../..")

def test_trajdistance():
    parmtop = 'm1.prmtop'
    solute_mask = ':1'
    trajin = '1/ground_snap.1'
    solvent_mask = ':400'
    trajdist = TrajDistance(parmtop, solute_mask)
    distance = trajdist(trajin, solvent_mask)
    assert distance == 5.4782

def test_trajdistance2():
    parmtop = 'm1.prmtop'
    solute_mask = ':1'
    trajin = '2/ground_snap.2'
    solvent_mask = ':400'
    trajdist = TrajDistance(parmtop, solute_mask)
    distance = trajdist(trajin, solvent_mask, 1)
    assert distance == 5.8772

max_stars_repo_name: PotentialParadox/pynasqm

max_issues_repo_name: PotentialParadox/pynasqm

max_forks_repo_name: PotentialParadox/pynasqm