text:
#from fluids import *
from fluids.numerics import IS_PYPY
from fluids.constants import *
#IS_PYPY = True

if not IS_PYPY:
    import fluids.numba
    import numba
from datetime import datetime
import numpy as np
import pytz
import inspect

def also_numba(f):
    if not IS_PYPY:
        f.duplicate_with_numba = True
    return f

from fluids.atmosphere import ATMOSPHERE_1976, ATMOSPHERE_NRLMSISE00, airmass, solar_position, earthsun_distance, sunrise_sunset, solar_irradiation

if not IS_PYPY:
    ATMOSPHERE_1976_numba = fluids.numba.ATMOSPHERE_1976
    airmass_numba = fluids.numba.atmosphere.airmass

    @numba.njit
    def numba_int_airmass(Z):
        return fluids.numba.atmosphere.ATMOSPHERE_1976(Z, 0).rho

class BaseTimeSuite(object):
    def setup(self):
        if not IS_PYPY:
            for k in dir(self.__class__):
                if 'time' in k and 'numba' in k:
                    c = getattr(self, k)
                    c()

class TimeAtmosphereSuite(BaseTimeSuite):
    def setup(self):
        if not IS_PYPY:
            self.time_ATMOSPHERE_1976_numba()
        self.date_test_es = datetime(2020, 6, 6, 10, 0, 0, 0)
        self.tz_dt = pytz.timezone('Australia/Perth').localize(datetime(2020, 6, 6, 7, 10, 57))
        self.tz_dt2 = pytz.timezone('America/Edmonton').localize(datetime(2018, 4, 15, 13, 43, 5))

    def time_ATMOSPHERE_1976(self):
        ATMOSPHERE_1976(5000.0)

    def time_ATMOSPHERE_1976_numba(self):
        ATMOSPHERE_1976_numba(5000.0)

    def time_ATMOSPHERE_1976_pressure_integral(self):
        ATMOSPHERE_1976.pressure_integral(288.6, 84100.0, 147.0)

    def time_ATMOSPHERE_NRLMSISE00(self):
        ATMOSPHERE_NRLMSISE00(1E3, 45, 45, 150)

    def time_airmass(self):
        airmass(lambda Z : ATMOSPHERE_1976(Z).rho, 90.0)

    def time_airmass_numba(self):
        airmass_numba(numba_int_airmass, 90.0)

    def time_earthsun_distance(self):
        earthsun_distance(self.date_test_es)

    def time_solar_position(self):
        solar_position(self.tz_dt, -31.95265, 115.85742)

    def time_sunrise_sunset(self):
        sunrise_sunset(self.tz_dt, 51.0486, -114.07)

    def time_solar_irradiation(self):
        solar_irradiation(Z=1100.0, latitude=51.0486, longitude=-114.07, linke_turbidity=3, moment=self.tz_dt2, surface_tilt=41.0, surface_azimuth=180.0)

from fluids import isothermal_gas, isentropic_work_compression, isentropic_efficiency, P_isothermal_critical_flow
if not IS_PYPY:
    isentropic_work_compression_numba = fluids.numba.isentropic_work_compression
    isentropic_efficiency_numba = fluids.numba.isentropic_efficiency
    P_isothermal_critical_flow_numba = fluids.numba.P_isothermal_critical_flow
    isothermal_gas_numba = fluids.numba.isothermal_gas

class TimeCompressibleSuite(BaseTimeSuite):
    def setup(self):
        if not IS_PYPY:
            self.time_isentropic_work_compression_numba()
            self.time_isentropic_efficiency_numba()
            #self.time_P_isothermal_critical_flow_numba()
            self.time_compressible_D_numba()
            #self.time_compressible_P1_numba()
            #self.time_compressible_P2_numba()

    def time_compressible_D(self):
        isothermal_gas(rho=11.3, fd=0.00185, P1=1E6, P2=9E5, L=1000, m=145.48475726, D=None)

    def time_compressible_D_numba(self):
        isothermal_gas_numba(rho=11.3, fd=0.00185, P1=1E6, P2=9E5, L=1000, m=145.48475726, D=None)

    def time_compressible_P1(self):
        isothermal_gas(rho=11.3, fd=0.00185, P2=9E5, L=1000, m=145.48475726, D=0.5)

    #def time_compressible_P1_numba(self):
        #isothermal_gas_numba(rho=11.3, fd=0.00185, P2=9E5, L=1000, m=145.48475726, D=0.5)

    def time_compressible_P2(self):
        isothermal_gas(rho=11.3, fd=0.00185, P1=1E6, L=1000, m=145.48475726, D=0.5)

    #def time_compressible_P2_numba(self):
        #isothermal_gas_numba(rho=11.3, fd=0.00185, P1=1E6, L=1000, m=145.48475726, D=0.5)

    def time_isentropic_work_compression(self):
        isentropic_work_compression(P1=1E5, P2=1E6, T1=300.0, k=1.4, eta=0.78)

    def time_isentropic_work_compression_numba(self):
        isentropic_work_compression_numba(P1=1E5, P2=1E6, T1=300.0, k=1.4, eta=0.78)

    def time_isentropic_efficiency(self):
        isentropic_efficiency(1E5, 1E6, 1.4, eta_p=0.78)

    def time_isentropic_efficiency_numba(self):
        isentropic_efficiency_numba(1E5, 1E6, 1.4, eta_p=0.78)

    def time_P_isothermal_critical_flow(self):
        P_isothermal_critical_flow(P=1E6, fd=0.00185, L=1000., D=0.5)

    #def time_P_isothermal_critical_flow_numba(self):
        #P_isothermal_critical_flow_numba(P=1E6, fd=0.00185, L=1000., D=0.5)

from fluids import control_valve_noise_g_2011, control_valve_noise_l_2015, size_control_valve_l, size_control_valve_g
if not IS_PYPY:
    control_valve_noise_l_2015_numba = fluids.numba.control_valve_noise_l_2015
    control_valve_noise_g_2011_numba = fluids.numba.control_valve_noise_g_2011

class TimeControlValveSuite(BaseTimeSuite):
    def time_size_control_valve_g(self):
        size_control_valve_g(T=433., MW=44.01, mu=1.4665E-4, gamma=1.30,  Z=0.988, P1=680E3, P2=310E3, Q=38/36., D1=0.08, D2=0.1, d=0.05, FL=0.85, Fd=0.42, xT=0.60)

    def time_size_control_valve_l(self):
        size_control_valve_l(rho=965.4, Psat=70.1E3, Pc=22120E3, mu=3.1472E-4, P1=680E3, P2=220E3, Q=0.1, D1=0.15, D2=0.15, d=0.15, FL=0.9, Fd=0.46)

    def time_control_valve_noise_l_2015(self):
        control_valve_noise_l_2015(m=40, P1=1E6, P2=6.5E5, Psat=2.32E3, rho=997, c=1400, Kv=77.848, d=0.1, Di=0.1071, FL=0.92, Fd=0.42, t_pipe=0.0036, rho_pipe=7800.0, c_pipe=5000.0,rho_air=1.293, c_air=343.0, An=-4.6)

    def time_control_valve_noise_l_2015_numba(self):
        control_valve_noise_l_2015_numba(m=40, P1=1E6, P2=6.5E5, Psat=2.32E3, rho=997, c=1400, Kv=77.848, d=0.1, Di=0.1071, FL=0.92, Fd=0.42, t_pipe=0.0036, rho_pipe=7800.0, c_pipe=5000.0,rho_air=1.293, c_air=343.0, An=-4.6)

    def time_control_valve_noise_g_2011(self):
        control_valve_noise_g_2011(m=2.22, P1=1E6, P2=7.2E5, T1=450, rho=5.3, gamma=1.22, MW=19.8, Kv=77.85,  d=0.1, Di=0.2031, FL=None, FLP=0.792, FP=0.98, Fd=0.296, t_pipe=0.008, rho_pipe=8000.0, c_pipe=5000.0, rho_air=1.293, c_air=343.0, An=-3.8, Stp=0.2)

    def time_control_valve_noise_g_2011_numba(self):
        control_valve_noise_g_2011_numba(m=2.22, P1=1E6, P2=7.2E5, T1=450, rho=5.3, gamma=1.22, MW=19.8, Kv=77.85,  d=0.1, Di=0.2031, FL=None, FLP=0.792, FP=0.98, Fd=0.296, t_pipe=0.008, rho_pipe=8000.0, c_pipe=5000.0, rho_air=1.293, c_air=343.0, An=-3.8, Stp=0.2)

from fluids import drag_sphere, v_terminal, integrate_drag_sphere

if not IS_PYPY:
    drag_sphere_numba = fluids.numba.drag_sphere
    v_terminal_numba = fluids.numba.v_terminal
#integrate_drag_sphere_numba = fluids.numba.integrate_drag_sphere

class TimeDragSuite(BaseTimeSuite):

    def time_drag_sphere(self):
        drag_sphere(20000.0, 'Barati_high')

    def time_drag_sphere_numba(self):
        drag_sphere_numba(20000.0, 'Barati_high')

    def time_v_terminal(self):
        v_terminal(D=70E-6, rhop=2600., rho=1000., mu=1E-3)

    def time_v_terminal_numba(self):
        v_terminal_numba(D=70E-6, rhop=2600., rho=1000., mu=1E-3)

    def time_integrate_drag_sphere(self):
        integrate_drag_sphere(D=0.001, rhop=2200., rho=1.2, mu=1.78E-5, t=0.5, V=30, distance=True)

from fluids import change_K_basis, entrance_distance, Darby3K, Hooper2K, K_angle_valve_Crane, v_lift_valve_Crane, K_branch_converging_Crane
if not IS_PYPY:

    change_K_basis_numba = fluids.numba.change_K_basis
    entrance_distance_numba = fluids.numba.entrance_distance
    Darby3K_numba = fluids.numba.Darby3K
    Hooper2K_numba = fluids.numba.Hooper2K
    K_angle_valve_Crane_numba = fluids.numba.K_angle_valve_Crane
    v_lift_valve_Crane_numba = fluids.numba.v_lift_valve_Crane
    K_branch_converging_Crane_numba = fluids.numba.K_branch_converging_Crane

class TimeFittingsSuite(BaseTimeSuite):
    #def setup(self):
        #pass
        #self.time_change_K_basis_numba()
        #self.time_entrance_distance_idelchik_numba()
        #self.time_entrance_distance_harris_numba()
        #self.time_Darby3K_numba()
        #self.time_Hooper2K_numba()

    def time_change_K_basis(self):
        change_K_basis(K1=32.68875692997804, D1=.01, D2=.02)

    def time_change_K_basis_numba(self):
        change_K_basis_numba(K1=32.68875692997804, D1=.01, D2=.02)

    def time_entrance_distance_idelchik(self):
        entrance_distance(Di=0.1, t=0.0005, l=.02, method='Idelchik')

    def time_entrance_distance_idelchik_numba(self):
        entrance_distance_numba(Di=0.1, t=0.0005, l=.02, method='Idelchik')

    def time_entrance_distance_harris(self):
        entrance_distance(Di=0.1, t=0.0005, l=.02, method='Harris')

    def time_entrance_distance_harris_numba(self):
        entrance_distance_numba(Di=0.1, t=0.0005, l=.02, method='Harris')

    def time_Darby3K(self):
        Darby3K(NPS=2., Re=10000., name='Valve, Angle valve, 45°, full line size, β = 1')

    def time_Darby3K_numba(self):
        Darby3K_numba(NPS=2., Re=10000., name='Valve, Angle valve, 45°, full line size, β = 1')

    def time_Hooper2K(self):
         Hooper2K(Di=2., Re=10000., name='Valve, Globe, Standard')

    def time_Hooper2K_numba(self):
         Hooper2K_numba(Di=2., Re=10000., name='Valve, Globe, Standard')

    def time_K_angle_valve_Crane(self):
         K_angle_valve_Crane(.01, .02)

    def time_K_angle_valve_Crane_numba(self):
         K_angle_valve_Crane_numba(.01, .02)

    def time_v_lift_valve_Crane(self):
        v_lift_valve_Crane(rho=998.2, D1=0.0627, D2=0.0779, style='lift check straight')

    def time_v_lift_valve_Crane_numba(self):
        v_lift_valve_Crane_numba(rho=998.2, D1=0.0627, D2=0.0779, style='lift check straight')

    def time_K_branch_converging_Crane(self):
        K_branch_converging_Crane(0.1023, 0.1023, 0.018917, 0.00633)

    def time_K_branch_converging_Crane_numba(self):
        K_branch_converging_Crane_numba(0.1023, 0.1023, 0.018917, 0.00633)

from fluids import C_Reader_Harris_Gallagher, differential_pressure_meter_solver, dP_venturi_tube
if not IS_PYPY:
    C_Reader_Harris_Gallagher_numba = fluids.numba.C_Reader_Harris_Gallagher
    differential_pressure_meter_solver_numba = fluids.numba.differential_pressure_meter_solver
    dP_venturi_tube_numba = fluids.numba.dP_venturi_tube

class TimeFlowMeterSuite(BaseTimeSuite):

    def time_C_Reader_Harris_Gallagher(self):
        C_Reader_Harris_Gallagher(D=0.07391, Do=0.0222, rho=1.165, mu=1.85E-5, m=0.12, taps='flange')

    def time_C_Reader_Harris_Gallagher_numba(self):
        C_Reader_Harris_Gallagher_numba(D=0.07391, Do=0.0222, rho=1.165, mu=1.85E-5, m=0.12, taps='flange')

    def time_dP_venturi_tube(self):
        dP_venturi_tube(D=0.07366, Do=0.05, P1=200000.0, P2=183000.0)

    def time_dP_venturi_tube_numba(self):
        dP_venturi_tube_numba(D=0.07366, Do=0.05, P1=200000.0, P2=183000.0)

    def time_differential_pressure_meter_solver_m(self):
        differential_pressure_meter_solver(D=0.07366, D2=0.05, P1=200000.0, P2=183000.0, rho=999.1, mu=0.0011, k=1.33, meter_type='ISO 5167 orifice', taps='D')

    def time_differential_pressure_meter_solver_numba_m(self):
        differential_pressure_meter_solver_numba(D=0.07366, D2=0.05, P1=200000.0, P2=183000.0, rho=999.1, mu=0.0011, k=1.33, meter_type='ISO 5167 orifice', taps='D')

    def time_differential_pressure_meter_solver_P2(self):
        differential_pressure_meter_solver(D=0.07366, D2=0.05, P1=200000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='ISO 5167 orifice', taps='D')

    def time_differential_pressure_meter_solver_numba_P2(self):
        differential_pressure_meter_solver_numba(D=0.07366, D2=0.05, P1=200000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='ISO 5167 orifice', taps='D')

    def time_differential_pressure_meter_solver_P1(self):
        differential_pressure_meter_solver(D=0.07366, D2=0.05, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='ISO 5167 orifice', taps='D')

    def time_differential_pressure_meter_solver_numba_P1(self):
        differential_pressure_meter_solver_numba(D=0.07366, D2=0.05, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='ISO 5167 orifice', taps='D')

    def time_differential_pressure_meter_solver_D2(self):
        differential_pressure_meter_solver(D=0.07366, P1=200000.0, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='ISO 5167 orifice', taps='D')

    def time_differential_pressure_meter_solver_numba_D2(self):
        differential_pressure_meter_solver_numba(D=0.07366, P1=200000.0, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='ISO 5167 orifice', taps='D')

    def time_differential_pressure_meter_solver_m_Hollingshead(self):
        differential_pressure_meter_solver(D=0.07366, D2=0.05, P1=200000.0, P2=183000.0, rho=999.1, mu=0.0011, k=1.33, meter_type='Hollingshead orifice', taps='D')

    def time_differential_pressure_meter_solver_numba_m_Hollingshead(self):
        differential_pressure_meter_solver_numba(D=0.07366, D2=0.05, P1=200000.0, P2=183000.0, rho=999.1, mu=0.0011, k=1.33, meter_type='Hollingshead orifice', taps='D')

    def time_differential_pressure_meter_solver_P2_Hollingshead(self):
        differential_pressure_meter_solver(D=0.07366, D2=0.05, P1=200000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Hollingshead orifice', taps='D')

    def time_differential_pressure_meter_solver_numba_P2_Hollingshead(self):
        differential_pressure_meter_solver_numba(D=0.07366, D2=0.05, P1=200000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Hollingshead orifice', taps='D')

    def time_differential_pressure_meter_solver_P1_Hollingshead(self):
        differential_pressure_meter_solver(D=0.07366, D2=0.05, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Hollingshead orifice', taps='D')

    def time_differential_pressure_meter_solver_numba_P1_Hollingshead(self):
        differential_pressure_meter_solver_numba(D=0.07366, D2=0.05, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Hollingshead orifice', taps='D')

    def time_differential_pressure_meter_solver_D2_Hollingshead(self):
        differential_pressure_meter_solver(D=0.07366, P1=200000.0, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Hollingshead orifice', taps='D')

    def time_differential_pressure_meter_solver_numba_D2_Hollingshead(self):
        differential_pressure_meter_solver_numba(D=0.07366, P1=200000.0, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Hollingshead orifice', taps='D')

    def time_differential_pressure_meter_solver_m_Miller_orifice(self):
        differential_pressure_meter_solver(D=0.07366, D2=0.05, P1=200000.0, P2=183000.0, rho=999.1, mu=0.0011, k=1.33, meter_type='Miller orifice', taps='corner')

    def time_differential_pressure_meter_solver_numba_m_Miller_orifice(self):
        differential_pressure_meter_solver_numba(D=0.07366, D2=0.05, P1=200000.0, P2=183000.0, rho=999.1, mu=0.0011, k=1.33, meter_type='Miller orifice', taps='corner')

    def time_differential_pressure_meter_solver_P2_Miller_orifice(self):
        differential_pressure_meter_solver(D=0.07366, D2=0.05, P1=200000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Miller orifice', taps='corner')

    def time_differential_pressure_meter_solver_numba_P2_Miller_orifice(self):
        differential_pressure_meter_solver_numba(D=0.07366, D2=0.05, P1=200000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Miller orifice', taps='corner')

    def time_differential_pressure_meter_solver_P1_Miller_orifice(self):
        differential_pressure_meter_solver(D=0.07366, D2=0.05, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Miller orifice', taps='corner')

    def time_differential_pressure_meter_solver_numba_P1_Miller_orifice(self):
        differential_pressure_meter_solver_numba(D=0.07366, D2=0.05, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Miller orifice', taps='corner')

    def time_differential_pressure_meter_solver_D2_Miller_orifice(self):
        differential_pressure_meter_solver(D=0.07366, P1=200000.0, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Miller orifice', taps='corner')

    def time_differential_pressure_meter_solver_numba_D2_Miller_orifice(self):
        differential_pressure_meter_solver_numba(D=0.07366, P1=200000.0, P2=183000.0, m=7.702338035732167, rho=999.1, mu=0.0011, k=1.33, meter_type='Miller orifice', taps='corner')

from fluids import ft_Crane, friction_factor, friction_factor_curved, friction_plate_Kumar, roughness_Farshad
if not IS_PYPY:
    ft_Crane_numba = fluids.numba.ft_Crane
    friction_factor_numba = fluids.numba.friction_factor
    friction_factor_curved_numba = fluids.numba.friction_factor_curved
    friction_plate_Kumar_numba = fluids.numba.friction_plate_Kumar
    roughness_Farshad_numba = fluids.numba.roughness_Farshad

class TimeFrictionSuite(BaseTimeSuite):
    def time_ft_Crane(self):
        ft_Crane(.1)

    def time_ft_Crane_numba(self):
        ft_Crane_numba(.1)

    def time_friction_factor(self):
        friction_factor(Re=1E5, eD=1E-4)

    def time_friction_factor_numba(self):
        friction_factor_numba(Re=1E5, eD=1E-4)

    def time_friction_factor_S2(self):
        friction_factor(Re=2.9E5, eD=1E-5, Method='Serghides_2')

    def time_friction_factor_S2_numba(self):
        friction_factor_numba(Re=2.9E5, eD=1E-5, Method='Serghides_2')

    def time_friction_factor_curved(self):
        friction_factor_curved(Re=1E5, Di=0.02, Dc=0.5)

    def time_friction_factor_curved_numba(self):
        friction_factor_curved_numba(Re=1E5, Di=0.02, Dc=0.5)

    def time_friction_plate_Kumar(self):
        friction_plate_Kumar(Re=2000.0, chevron_angle=30.0)

    def time_friction_plate_Kumar_numba(self):
        friction_plate_Kumar_numba(Re=2000.0, chevron_angle=30.0)

    def time_roughness_Farshad(self):
        roughness_Farshad('Cr13, bare', 0.05)

    def time_roughness_Farshad_numba(self):
        roughness_Farshad_numba('Cr13, bare', 0.05)

from fluids import V_horiz_conical, V_horiz_spherical, V_horiz_torispherical, V_vertical_torispherical_concave, SA_tank, V_tank, SA_partial_horiz_spherical_head, SA_partial_horiz_ellipsoidal_head, SA_partial_horiz_guppy_head, SA_partial_horiz_torispherical_head, V_from_h, SA_from_h
if not IS_PYPY:
    V_horiz_conical_numba = fluids.numba.V_horiz_conical
    V_horiz_spherical_numba = fluids.numba.V_horiz_spherical
    V_horiz_torispherical_numba = fluids.numba.V_horiz_torispherical
    V_vertical_torispherical_concave_numba = fluids.numba.V_vertical_torispherical_concave
    SA_tank_numba = fluids.numba.SA_tank
    V_tank_numba = fluids.numba.V_tank
    SA_partial_horiz_spherical_head_numba = fluids.numba.SA_partial_horiz_spherical_head
    SA_partial_horiz_ellipsoidal_head_numba = fluids.numba.SA_partial_horiz_ellipsoidal_head
    SA_partial_horiz_guppy_head_numba = fluids.numba.SA_partial_horiz_guppy_head
    SA_partial_horiz_torispherical_head_numba = fluids.numba.SA_partial_horiz_torispherical_head
    V_from_h_numba = fluids.numba.V_from_h
    SA_from_h_numba = fluids.numba.SA_from_h

class TimeGeometrySuite(BaseTimeSuite):
    def time_V_horiz_conical(self):
        V_horiz_conical(D=108., L=156., a=42., h=36)

    def time_V_horiz_conical_numba(self):
        V_horiz_conical_numba(D=108., L=156., a=42., h=36)

    def time_V_horiz_spherical(self):
        V_horiz_spherical(D=108., L=156., a=0.1, h=36.0)

    def time_V_horiz_spherical_numba(self):
        V_horiz_spherical_numba(D=108., L=156., a=0.1, h=36.0)

    def time_V_horiz_torispherical_1(self):
        V_horiz_torispherical(D=108., L=156., f=1., k=0.06, h=1.0)

    def time_V_horiz_torispherical_1_numba(self):
        V_horiz_torispherical_numba(D=108., L=156., f=1., k=0.06, h=1.0)

    def time_V_horiz_torispherical_2(self):
        V_horiz_torispherical(D=108., L=156., f=1., k=0.06, h=84.0)

    def time_V_horiz_torispherical_2_numba(self):
        V_horiz_torispherical_numba(D=108., L=156., f=1., k=0.06, h=84.0)

    def time_V_horiz_torispherical_3(self):
        V_horiz_torispherical(D=108., L=156., f=1., k=0.06, h=106.0)

    def time_V_horiz_torispherical_3_numba(self):
        V_horiz_torispherical_numba(D=108., L=156., f=1., k=0.06, h=106.0)

    def time_V_vertical_torispherical_concave(self):
        V_vertical_torispherical_concave(D=113., f=0.71, k=0.081, h=15.0)

    def time_V_vertical_torispherical_concave_numba(self):
        V_vertical_torispherical_concave_numba(D=113., f=0.71, k=0.081, h=15.0)

    def time_SA_tank_1(self):
        SA_tank(D=2.54, L=5, sideA='torispherical', sideB='torispherical', sideA_f=1.039370079, sideA_k=0.062362205, sideB_f=1.039370079, sideB_k=0.062362205)

    def time_SA_tank_1_numba(self):
        SA_tank_numba(D=2.54, L=5, sideA='torispherical', sideB='torispherical', sideA_f=1.039370079, sideA_k=0.062362205, sideB_f=1.039370079, sideB_k=0.062362205)

    def time_V_tank_1(self):
        V_tank(D=1.5, L=5., horizontal=False, sideA='conical', sideB='conical', sideA_a=2., sideB_a=1.)

    def time_V_tank_1_numba(self):
        V_tank_numba(D=1.5, L=5., horizontal=False, sideA='conical', sideB='conical', sideA_a=2., sideB_a=1.)

    def time_SA_partial_horiz_spherical_head(self):
        SA_partial_horiz_spherical_head(D=72., a=48.0, h=24.0)

    def time_SA_partial_horiz_spherical_head_numba(self):
        SA_partial_horiz_spherical_head_numba(D=72., a=48.0, h=24.0)

    def time_SA_partial_horiz_ellipsoidal_head(self):
        SA_partial_horiz_ellipsoidal_head(D=72., a=48.0, h=24.0)

    def time_SA_partial_horiz_ellipsoidal_head_numba(self):
        SA_partial_horiz_ellipsoidal_head_numba(D=72., a=48.0, h=24.0)

    def time_SA_partial_horiz_guppy_head(self):
        SA_partial_horiz_guppy_head(D=72., a=48.0, h=24.0)

    def time_SA_partial_horiz_guppy_head_numba(self):
        SA_partial_horiz_guppy_head_numba(D=72., a=48.0, h=24.0)

    def time_SA_partial_horiz_torispherical_head_1(self):
        SA_partial_horiz_torispherical_head(D=72., f=1, k=.06, h=1.0)

    def time_SA_partial_horiz_torispherical_head_1_numba(self):
        SA_partial_horiz_torispherical_head_numba(D=72., f=1, k=.06, h=1.0)

    def time_SA_partial_horiz_torispherical_head_2(self):
        SA_partial_horiz_torispherical_head(D=72., f=1, k=.06, h=20.0)

    def time_SA_partial_horiz_torispherical_head_2_numba(self):
        SA_partial_horiz_torispherical_head_numba(D=72., f=1, k=.06, h=20.0)

    def time_SA_partial_horiz_torispherical_head_3(self):
        SA_partial_horiz_torispherical_head(D=72., f=1, k=.06, h=66.0)

    def time_SA_partial_horiz_torispherical_head_3_numba(self):
        SA_partial_horiz_torispherical_head_numba(D=72., f=1, k=.06, h=66.0)

    def time_V_from_h(self):
        V_from_h(h=7, D=1.5, L=5., horizontal=False, sideA='conical', sideB='conical', sideA_a=2., sideB_a=1.)

    def time_V_from_h_numba(self):
        V_from_h_numba(h=7, D=1.5, L=5., horizontal=False, sideA='conical', sideB='conical', sideA_a=2., sideB_a=1.)

    def time_SA_from_h(self):
        SA_from_h(h=7, D=1.5, L=5., horizontal=False, sideA='conical', sideB='conical', sideA_a=2., sideB_a=1.)

    def time_SA_from_h_numba(self):
        SA_from_h_numba(h=7, D=1.5, L=5., horizontal=False, sideA='conical', sideB='conical', sideA_a=2., sideB_a=1.)

    # TODO test classes in geometry

from fluids import Q_weir_V_Shen
if not IS_PYPY:
    Q_weir_V_Shen_numba = fluids.numba.Q_weir_V_Shen

class TimeOpenFlowSuite(BaseTimeSuite):
    def time_Q_weir_V_Shen(self):
        Q_weir_V_Shen(0.6, angle=45)

    def time_Q_weir_V_Shen_numba(self):
        Q_weir_V_Shen_numba(0.6, angle=45)

from fluids import dP_packed_bed
if not IS_PYPY:
    dP_packed_bed_numba = fluids.numba.dP_packed_bed

class TimePackedBedSuite(BaseTimeSuite):
    def time_dP_packed_bed(self):
        dP_packed_bed(dp=0.05, voidage=0.492, vs=0.1, rho=1E3, mu=1E-3, Dt=0.015, Method='Guo, Sun, Zhang, Ding & Liu')

    def time_dP_packed_bed_numba(self):
        dP_packed_bed_numba(dp=0.05, voidage=0.492, vs=0.1, rho=1E3, mu=1E-3, Dt=0.015, Method='Guo, Sun, Zhang, Ding & Liu')

from fluids import Stichlmair_wet, Stichlmair_flood
if not IS_PYPY:
    Stichlmair_wet_numba = fluids.numba.Stichlmair_wet
    Stichlmair_flood_numba = fluids.numba.Stichlmair_flood

class TimePackedTowerSuite(BaseTimeSuite):
    def time_Stichlmair_wet(self):
        Stichlmair_wet(Vg=0.4, Vl = 5E-3, rhog=5., rhol=1200., mug=5E-5, voidage=0.68, specific_area=260., C1=32., C2=7., C3=1.)

    def time_Stichlmair_wet_numba(self):
        Stichlmair_wet_numba(Vg=0.4, Vl = 5E-3, rhog=5., rhol=1200., mug=5E-5, voidage=0.68, specific_area=260., C1=32., C2=7., C3=1.)

    def time_Stichlmair_flood(self):
        Stichlmair_flood(Vl = 5E-3, rhog=5., rhol=1200., mug=5E-5, voidage=0.68, specific_area=260., C1=32., C2=7., C3=1.)

    def time_Stichlmair_flood_numba(self):
        Stichlmair_flood_numba(Vl = 5E-3, rhog=5., rhol=1200., mug=5E-5, voidage=0.68, specific_area=260., C1=32., C2=7., C3=1.)

from fluids import ParticleSizeDistribution
if not IS_PYPY:
    ParticleSizeDistribution_numba = fluids.numba.ParticleSizeDistribution

class TimeParticleSizeDistributionSuite(BaseTimeSuite):

    # TODO optimize these; maybe add numba support in the future - or increase support for numpy inputs
    #
    def time_ParticleSizeDistribution_init(self):
        ds = 1E-6*np.array([240, 360, 450, 562.5, 703, 878, 1097, 1371, 1713, 2141, 2676, 3345, 4181, 5226, 6532])
        numbers = [65, 119, 232, 410, 629, 849, 990, 981, 825, 579, 297, 111, 21, 1]
        psd = ParticleSizeDistribution(ds=ds, fractions=numbers, order=0)

from fluids import nearest_pipe, gauge_from_t, t_from_gauge
if not IS_PYPY:
    nearest_pipe_numba = fluids.numba.nearest_pipe

class TimePipingSuite(BaseTimeSuite):
    def time_nearest_pipe(self):
        nearest_pipe(Di=0.021)

    def time_gauge_from_t(self):
        gauge_from_t(.5, SI=False, schedule='BWG')

    def time_t_from_gauge(self):
        t_from_gauge(.2, False, 'BWG')

from fluids import VFD_efficiency, CSA_motor_efficiency
if not IS_PYPY:
    VFD_efficiency_numba = fluids.numba.VFD_efficiency
    CSA_motor_efficiency_numba = fluids.numba.CSA_motor_efficiency

class TimePumpSuite(BaseTimeSuite):
    def time_VFD_efficiency(self):
        VFD_efficiency(100*hp, load=0.2)

    def time_VFD_efficiency_numba(self):
        VFD_efficiency_numba(100*hp, load=0.2)

    def time_CSA_motor_efficiency(self):
        CSA_motor_efficiency(100*hp, closed=True, poles=6, high_efficiency=True)

    def time_CSA_motor_efficiency_numba(self):
        CSA_motor_efficiency_numba(100*hp, closed=True, poles=6, high_efficiency=True)

from fluids import API520_A_g, API520_B, API520_SH
if not IS_PYPY:
    API520_A_g_numba = fluids.numba.API520_A_g
    API520_B_numba = fluids.numba.API520_B
    API520_SH_numba = fluids.numba.API520_SH

class TimeSafetyValveSuite(BaseTimeSuite):
    def time_API520_A_g(self):
        API520_A_g(m=24270/3600., T=348., Z=0.90, MW=51., k=1.11, P1=670E3, Kb=1, Kc=1)

    def time_API520_A_g_numba(self):
        API520_A_g_numba(m=24270/3600., T=348., Z=0.90, MW=51., k=1.11, P1=670E3, Kb=1, Kc=1)

    def time_API520_B(self):
        API520_B(1E6, 5E5)

    def time_API520_B_numba(self):
        API520_B_numba(1E6, 5E5)

    def time_API520_SH(self):
        API520_SH(593+273.15, 1066.325E3)

    def time_API520_SH_numba(self):
        API520_SH_numba(593+273.15, 1066.325E3)

from fluids import K_separator_Watkins
if not IS_PYPY:
    K_separator_Watkins_numba = fluids.numba.K_separator_Watkins

class TimeSeparatorSuite(BaseTimeSuite):
    def time_K_separator_Watkins(self):
        K_separator_Watkins(0.88, 985.4, 1.3, horizontal=True)

    def time_K_separator_Watkins_numba(self):
        K_separator_Watkins_numba(0.88, 985.4, 1.3, horizontal=True)

from fluids import liquid_gas_voidage, gas_liquid_viscosity
if not IS_PYPY:
    liquid_gas_voidage_numba = fluids.numba.liquid_gas_voidage
    gas_liquid_viscosity_numba = fluids.numba.gas_liquid_viscosity

class TimeTwoPhaseVoidageSuite(BaseTimeSuite):

    def time_liquid_gas_voidage(self):
        liquid_gas_voidage(m=0.6, x=0.1, rhol=915., rhog=2.67, mul=180E-6, mug=14E-6, sigma=0.0487, D=0.05)

    def time_liquid_gas_voidage_numba(self):
        liquid_gas_voidage_numba(m=0.6, x=0.1, rhol=915., rhog=2.67, mul=180E-6, mug=14E-6, sigma=0.0487, D=0.05)

    def time_gas_liquid_viscosity(self):
        gas_liquid_viscosity(x=0.4, mul=1E-3, mug=1E-5, rhol=850, rhog=1.2, Method='Duckler')

    def time_gas_liquid_viscosity_numba(self):
        gas_liquid_viscosity_numba(x=0.4, mul=1E-3, mug=1E-5, rhol=850, rhog=1.2, Method='Duckler')

from fluids import two_phase_dP, Taitel_Dukler_regime, Mandhane_Gregory_Aziz_regime
if not IS_PYPY:
    two_phase_dP_numba = fluids.numba.two_phase_dP
    Taitel_Dukler_regime_numba = fluids.numba.Taitel_Dukler_regime
    Mandhane_Gregory_Aziz_regime_numba = fluids.numba.Mandhane_Gregory_Aziz_regime

class TimeTwoPhaseSuite(BaseTimeSuite):
    def time_two_phase_dP(self):
        two_phase_dP(m=0.6, x=0.1, rhol=915., rhog=2.67, mul=180E-6, mug=14E-6, sigma=0.0487, D=0.05, L=1.0)

    def time_two_phase_dP_numba(self):
        two_phase_dP_numba(m=0.6, x=0.1, rhol=915., rhog=2.67, mul=180E-6, mug=14E-6, sigma=0.0487, D=0.05, L=1.0)

    def time_Taitel_Dukler_regime(self):
        Taitel_Dukler_regime(m=0.6, x=0.112, rhol=915.12, rhog=2.67, mul=180E-6, mug=14E-6, D=0.05, roughness=0.0, angle=0)

    def time_Taitel_Dukler_regime_numba(self):
        Taitel_Dukler_regime_numba(m=0.6, x=0.112, rhol=915.12, rhog=2.67, mul=180E-6, mug=14E-6, D=0.05, roughness=0.0, angle=0)

    def time_Mandhane_Gregory_Aziz_regime(self):
        Mandhane_Gregory_Aziz_regime(m=0.6, x=0.112, rhol=915.12, rhog=2.67, mul=180E-6, mug=14E-6, sigma=0.065, D=0.05)

    def time_Mandhane_Gregory_Aziz_regime_numba(self):
        Mandhane_Gregory_Aziz_regime_numba(m=0.6, x=0.112, rhol=915.12, rhog=2.67, mul=180E-6, mug=14E-6, sigma=0.065, D=0.05)

suites = [TimeAtmosphereSuite, TimeCompressibleSuite, TimeControlValveSuite,
          TimeDragSuite, TimeFittingsSuite, TimeFlowMeterSuite,
          TimeFrictionSuite, TimeGeometrySuite, TimeOpenFlowSuite,
          TimePackedBedSuite, TimePackedTowerSuite, TimeParticleSizeDistributionSuite,
          TimePipingSuite, TimePumpSuite, TimeSafetyValveSuite,
          TimeSeparatorSuite, TimeTwoPhaseVoidageSuite, TimeTwoPhaseSuite,
          ]

for suite in suites:
    continue
    # asv requires inspect to work :(
    # Do I want to write a file that writes this benchmark file?
    glbs, lcls = {}, {}
    for k in dir(suite):
        if 'time' in k:
            f = getattr(suite, k)
            if hasattr(f, 'duplicate_with_numba'):
                source = inspect.getsource(f)
                source = '\n'.join([s[4:] for s in source.split('\n')[1:]])
                orig_function = k.replace('time_', '')
                numba_function = orig_function + '_numba'
                new_function_name = k + '_numba'
                new_source = source.replace(orig_function, numba_function)
                exec(new_source, glbs, lcls)
                setattr(suite, new_function_name, lcls[new_function_name])

if IS_PYPY:
    for s in suites:
        for k in dir(s):
            if 'time' in k and 'numba' in k:
                delattr(s, k)

max_stars_repo_name: yoelcortes/fluids

max_issues_repo_name: yoelcortes/fluids

max_forks_repo_name: yoelcortes/fluids

reason: import numpy,import numba


text:
# -*- encoding:UTF-8 -*-
import numpy as np
import pandas as pd
import time

# np.random.seed(2)

# 一维世界的长度
N_STATES = 8

# 可执行的动作
ACTIONS = ['left', 'right']

# 贪心策略执行的概率
EPSILON = 0.9

# 学习率
ALPHA = 0.1

# 折扣因子
LAMBDA = 0.9

# 最大回合数(最大运行次数)
MAX_EPISODES = N_STATES * 2 + 1

# 走一步需要多少秒(多少秒刷新一次)
FRESH_TIME = 0.3

def build_q_table(n_states, actions):
    table = pd.DataFrame(
        np.zeros((n_states+1, len(actions))),  # n_states行len(actions)列
        columns=actions,
    )
    return table

def choose_action(state, q_table):
    # 根据Q表选择一个操作
    state_actions = q_table.iloc[state, :]
    if np.random.uniform() > EPSILON or state_actions.all() == 0:
        # 探索操作
        action_name = np.random.choice(ACTIONS)
    else:
        # 利用操作
        action_name = state_actions.argmax()
    return action_name

def get_env_feedback(S, A):
    # 获取环境反馈，下一个状态和获得的奖赏
    if A == 'right':
        if S == N_STATES - 2:  # 进行向右操作以后就会到达终点
            S_ = 'terminal'
            R = 1
        else:
            S_ = S + 1
            R = 0
    else:
        R = 0
        if S == 0:  # 撞到左边的墙
            S_ = S
        else:
            S_ = S - 1
    return S_, R

def update_env(S, episode, step_counter):
    env_list = ['-'] * (N_STATES - 1) + ['T']
    if S == 'terminal':
        interaction = 'Episode %s: total_steps = %s' % (episode + 1, step_counter)
        print('\r{}'.format(interaction))
        time.sleep(2)
    else:
        env_list[S] = 'o'
        interaction = ''.join(env_list)
        print('\r{}'.format(interaction))
        time.sleep(FRESH_TIME)

def rl():
    q_table = build_q_table(N_STATES, ACTIONS)
    for episode in range(MAX_EPISODES):
        step_counter = 0
        S = np.random.randint(0, N_STATES-1)
        is_terminated = False
        update_env(S, episode, step_counter)
        while not is_terminated:
            A = choose_action(S, q_table)  # 根据当前的状态和Q表选择一个执行动作
            S_, R = get_env_feedback(S, A)  # 在当前状态上执行动作，获取下一个状态和执行动作得到的奖赏
            q_predict = q_table.ix[S, A]  # 根据以往经验，获取在状态S上执行动作A，对往后累积奖励的期望

            # 更新经验值(Q表)
            if S_ != 'terminal':
                q_target = R + LAMBDA * q_table.iloc[S_, :].max()
            else:
                q_target = R
                is_terminated = True
            q_table.ix[S, A] += ALPHA * (q_target - q_predict)
            S = S_  # 更新状态

            update_env(S, episode, step_counter + 1)
            step_counter += 1
    return q_table

if __name__ == '__main__':
    rl()

ext: py

max_stars_repo_name: sineatos/ReinforcementLearningDemo

max_issues_repo_name: sineatos/ReinforcementLearningDemo

max_forks_repo_name: sineatos/ReinforcementLearningDemo

reason: import numpy


text:
import torch
from torch_geometric.data import InMemoryDataset
from torch_geometric.data.data import Data
from torch_geometric.utils import to_networkx
from torch_geometric.utils import to_undirected
import numpy as np
import networkx as nx
import pickle
import os
import scipy.io as sio
from math import comb

def get_n_params(model):
    pp=0
    for p in list(model.parameters()):
        nn=1
        for s in list(p.size()):
            nn = nn*s
        pp += nn
    return pp

class PtcDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):

        super(PtcDataset, self).__init__(root, transform, pre_transform)

        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["ptc.mat"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        a=sio.loadmat(self.raw_paths[0]) #'subgraphcount/randomgraph.mat')
        # list of adjacency matrix
        A=a['A'][0]
        F=a['F'][0]
        #Y=a['y'].astype(np.float32) #(a['y']+1)//2
        Y=a['Y'].astype(np.int)
        Y=Y[:,0]

        data_list = []
        for i in range(len(A)):
            E=np.where(A[i]>0)
            edge_index=torch.Tensor(np.vstack((E[0],E[1]))).type(torch.int64)
            x=torch.tensor(F[i]).type(torch.float32)
            y=torch.tensor([Y[i]]) #.type(torch.float32)
            data_list.append(Data(edge_index=edge_index, x=x, y=y))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class ProteinsDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None,contfeat=False):
        self.contfeat=contfeat
        super(ProteinsDataset, self).__init__(root, transform, pre_transform)

        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["proteins.mat"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        a=sio.loadmat(self.raw_paths[0]) #'subgraphcount/randomgraph.mat')
        # list of adjacency matrix
        A=a['A'][0]
        F=a['F'][0]
        #Y=a['y'].astype(np.float32) #(a['y']+1)//2
        Y=a['Y'].astype(np.int)
        Y=Y[:,0]

        data_list = []
        for i in range(len(A)):
            E=np.where(A[i]>0)
            edge_index=torch.Tensor(np.vstack((E[0],E[1]))).type(torch.int64)
            if self.contfeat:
                #ind=list(set(range(0,F[i].shape[1]))-set([3,4]))
                tmp=F[i]#[:,ind]
                x=torch.tensor(tmp).type(torch.float32)
            else:
                x=torch.tensor(F[i][:,0:3]).type(torch.float32)
            y=torch.tensor([Y[i]]) #.type(torch.float32)
            data_list.append(Data(edge_index=edge_index, x=x, y=y))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class EnzymesDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None,contfeat=False):
        self.contfeat=contfeat
        super(EnzymesDataset, self).__init__(root, transform, pre_transform)

        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["enzymes.mat"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        a=sio.loadmat(self.raw_paths[0]) #'subgraphcount/randomgraph.mat')
        # list of adjacency matrix
        A=a['A'][0]
        F=a['F'][0]
        #Y=a['y'].astype(np.float32) #(a['y']+1)//2
        Y=a['Y'][0].astype(np.int)

        data_list = []
        for i in range(len(A)):
            E=np.where(A[i]>0)
            edge_index=torch.Tensor(np.vstack((E[0],E[1]))).type(torch.int64)
            if self.contfeat:
                #ind=list(set(range(0,F[i].shape[1]))-set([3,4]))
                tmp=F[i]#[:,ind]
                x=torch.tensor(tmp).type(torch.float32)
            else:
                x=torch.tensor(F[i][:,0:3]).type(torch.float32)
            y=torch.tensor([Y[i]]) #.type(torch.float32)
            data_list.append(Data(edge_index=edge_index, x=x, y=y))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class MutagDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(MutagDataset, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["mutag.mat"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        a=sio.loadmat(self.raw_paths[0]) #'subgraphcount/randomgraph.mat')
        # list of adjacency matrix
        A=a['A'][0]
        TA=a['TA'][0]
        F=a['F'][0]
        #Y=a['y'].astype(np.float32) #(a['y']+1)//2
        Y=((a['y']+1)//2).astype(np.float32)

        data_list = []
        for i in range(len(A)):
            E=np.where(A[i]>0)
            edge_index=torch.Tensor(np.vstack((E[0],E[1]))).type(torch.int64)
            x=torch.tensor(F[i]).type(torch.float32)
            y=torch.tensor(Y[i]).type(torch.float32)
            data_list.append(Data(edge_index=edge_index, x=x, y=y))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class Zinc12KDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(Zinc12KDataset, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["Zinc.mat"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        a=sio.loadmat(self.raw_paths[0])
        # list of adjacency matrix
        F=a['F'][0]
        A=a['E'][0]
        Y=a['Y']
        nmax=37
        ntype=21
        maxdeg=4

        data_list = []
        for i in range(len(A)):
            E=np.where(A[i]>0)
            edge_index=torch.Tensor(np.vstack((E[0],E[1]))).type(torch.int64)
            x=torch.zeros(A[i].shape[0],ntype+maxdeg)
            deg=(A[i]>0).sum(1)
            for j in range(F[i][0].shape[0]):
                # put atom code
                x[j,F[i][0][j]]=1
                # put degree code
                x[j,-int(deg[j])]=1
            y=torch.tensor(Y[i,:])
            data_list.append(Data(edge_index=edge_index, x=x, y=y))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class BandClassDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(BandClassDataset, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["bandclass.mat"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        a=sio.loadmat(self.raw_paths[0]) #'subgraphcount/randomgraph.mat')
        # list of adjacency matrix
        A=a['A']
        F=a['F']
        Y=a['Y']
        F=np.expand_dims(F,2)

        data_list = []
        for i in range(len(A)):
            E=np.where(A[i]>0)
            edge_index=torch.Tensor(np.vstack((E[0],E[1]))).type(torch.int64)
            x=torch.tensor(F[i,:,:])
            y=torch.tensor(Y[i,:])
            data_list.append(Data(edge_index=edge_index, x=x, y=y))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class TwoDGrid30(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(TwoDGrid30, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["TwoDGrid30.mat"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        a=sio.loadmat(self.raw_paths[0]) #'subgraphcount/randomgraph.mat')
        # list of adjacency matrix
        A=a['A']
        # list of output
        F=a['F']
        F=F.astype(np.float32)

        data_list = []
        E=np.where(A>0)
        edge_index=torch.Tensor(np.vstack((E[0],E[1]))).type(torch.int64)
        x=torch.tensor(F[:,0:1])
        y=torch.tensor(F[:,1:4])
        mask=torch.tensor(F[:,12:13])
        data_list.append(Data(edge_index=edge_index, x=x, y=y,mask=mask))
        x=torch.tensor(F[:,4:5])
        y=torch.tensor(F[:,5:8])
        data_list.append(Data(edge_index=edge_index, x=x, y=y,mask=mask))
        x=torch.tensor(F[:,8:9])
        y=torch.tensor(F[:,9:12])
        data_list.append(Data(edge_index=edge_index, x=x, y=y,mask=mask))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class GraphCountDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(GraphCountDataset, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["randomgraph.mat"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        a=sio.loadmat(self.raw_paths[0]) #'subgraphcount/randomgraph.mat')
        # list of adjacency matrix
        A=a['A'][0]
        # list of output
        Y=a['F']

        data_list = []
        for i in range(len(A)):
            a=A[i]
            A2=a.dot(a)
            A3=A2.dot(a)
            tri=np.trace(A3)/6
            tailed=((np.diag(A3)/2)*(a.sum(0)-2)).sum()
            cyc4=1/8*(np.trace(A3.dot(a))+np.trace(A2)-2*A2.sum())
            cus= a.dot(np.diag(np.exp(-a.dot(a).sum(1)))).dot(a).sum()

            deg=a.sum(0)
            star=0
            for j in range(a.shape[0]):
                star+=comb(int(deg[j]),3)

            expy=torch.tensor([[tri,tailed,star,cyc4,cus]])

            E=np.where(A[i]>0)
            edge_index=torch.Tensor(np.vstack((E[0],E[1]))).type(torch.int64)
            x=torch.ones(A[i].shape[0],1)
            #y=torch.tensor(Y[i:i+1,:])
            data_list.append(Data(edge_index=edge_index, x=x, y=expy))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class PlanarSATPairsDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None, pre_filter=None):
        super(PlanarSATPairsDataset, self).__init__(root, transform, pre_transform, pre_filter)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["GRAPHSAT.pkl"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        pass

    def process(self):
        # Read data into huge `Data` list.
        data_list = pickle.load(open(os.path.join(self.root, "raw/GRAPHSAT.pkl"), "rb"))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class Grapg8cDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(Grapg8cDataset, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["graph8c.g6"]

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        dataset = nx.read_graph6(self.raw_paths[0])
        data_list = []
        for i,datum in enumerate(dataset):
            x = torch.ones(datum.number_of_nodes(),1)
            edge_index = to_undirected(torch.tensor(list(datum.edges())).transpose(1,0))
            data_list.append(Data(edge_index=edge_index, x=x, y=0))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class SRDataset(InMemoryDataset):
    def __init__(self, root, transform=None, pre_transform=None):
        super(SRDataset, self).__init__(root, transform, pre_transform)
        self.data, self.slices = torch.load(self.processed_paths[0])

    @property
    def raw_file_names(self):
        return ["sr251256.g6"]  #sr251256  sr351668

    @property
    def processed_file_names(self):
        return 'data.pt'

    def download(self):
        # Download to `self.raw_dir`.
        pass

    def process(self):
        # Read data into huge `Data` list.
        b=self.processed_paths[0]
        dataset = nx.read_graph6(self.raw_paths[0])
        data_list = []
        for i,datum in enumerate(dataset):
            x = torch.ones(datum.number_of_nodes(),1)
            edge_index = to_undirected(torch.tensor(list(datum.edges())).transpose(1,0))
            data_list.append(Data(edge_index=edge_index, x=x, y=0))

        if self.pre_filter is not None:
            data_list = [data for data in data_list if self.pre_filter(data)]

        if self.pre_transform is not None:
            data_list = [self.pre_transform(data) for data in data_list]

        data, slices = self.collate(data_list)
        torch.save((data, slices), self.processed_paths[0])

class SpectralDesign(object):

    def __init__(self,nmax=0,recfield=1,dv=5,nfreq=5,adddegree=False,laplacien=True,addadj=False,vmax=None):
        # receptive field. 0: adj, 1; adj+I, n: n-hop area
        self.recfield=recfield
        # b parameter
        self.dv=dv
        # number of sampled point of spectrum
        self.nfreq=  nfreq
        # if degree is added to node feature
        self.adddegree=adddegree
        # use laplacian or adjacency for spectrum
        self.laplacien=laplacien
        # add adjacecny as edge feature
        self.addadj=addadj
        # use given max eigenvalue
        self.vmax=vmax

        # max node for PPGN algorithm, set 0 if you do not use PPGN
        self.nmax=nmax

    def __call__(self, data):

        n =data.x.shape[0]
        nf=data.x.shape[1]

        data.x=data.x.type(torch.float32)

        nsup=self.nfreq+1
        if self.addadj:
            nsup+=1

        A=np.zeros((n,n),dtype=np.float32)
        SP=np.zeros((nsup,n,n),dtype=np.float32)
        A[data.edge_index[0],data.edge_index[1]]=1

        if self.adddegree:
            data.x=torch.cat([data.x,torch.tensor(A.sum(0)).unsqueeze(-1)],1)

        # calculate receptive field. 0: adj, 1; adj+I, n: n-hop area
        if self.recfield==0:
            M=A
        else:
            M=(A+np.eye(n))
            for i in range(1,self.recfield):
                M=M.dot(M)

        M=(M>0)

        d = A.sum(axis=0)
        # normalized Laplacian matrix.
        dis=1/np.sqrt(d)
        dis[np.isinf(dis)]=0
        dis[np.isnan(dis)]=0
        D=np.diag(dis)
        nL=np.eye(D.shape[0])-(A.dot(D)).T.dot(D)
        V,U = np.linalg.eigh(nL)
        V[V<0]=0
        # keep maximum eigenvalue for Chebnet if it is needed
        data.lmax=V.max().astype(np.float32)

        if not self.laplacien:
            V,U = np.linalg.eigh(A)

        # design convolution supports
        vmax=self.vmax
        if vmax is None:
            vmax=V.max()

        freqcenter=np.linspace(V.min(),vmax,self.nfreq)

        # design convolution supports (aka edge features)
        for i in range(0,len(freqcenter)):
            SP[i,:,:]=M* (U.dot(np.diag(np.exp(-(self.dv*(V-freqcenter[i])**2))).dot(U.T)))
        # add identity
        SP[len(freqcenter),:,:]=np.eye(n)
        # add adjacency if it is desired
        if self.addadj:
            SP[len(freqcenter)+1,:,:]=A

        # set convolution support weigths as an edge feature
        E=np.where(M>0)
        data.edge_index2=torch.Tensor(np.vstack((E[0],E[1]))).type(torch.int64)
        data.edge_attr2 = torch.Tensor(SP[:,E[0],E[1]].T).type(torch.float32)

        # set tensor for Maron's PPGN
        if self.nmax>0:
            H=torch.zeros(1,nf+2,self.nmax,self.nmax)
            H[0,0,data.edge_index[0],data.edge_index[1]]=1
            H[0,1,0:n,0:n]=torch.diag(torch.ones(data.x.shape[0]))
            for j in range(0,nf):
                H[0,j+2,0:n,0:n]=torch.diag(data.x[:,j])
            data.X2= H
            M=torch.zeros(1,2,self.nmax,self.nmax)
            for i in range(0,n):
                M[0,0,i,i]=1
            M[0,1,0:n,0:n]=1-M[0,0,0:n,0:n]
            data.M= M

        return data

class PPGNAddDegree(object):

    def __init__(self,nmax=0,adddegree=True,):

        # if degree is added to node feature
        self.adddegree=adddegree

        # max node for PPGN algorithm, set 0 if you do not use PPGN
        self.nmax=nmax

    def __call__(self, data):

        n =data.x.shape[0]
        nf=data.x.shape[1]

        data.x=data.x.type(torch.float32)
        A=np.zeros((n,n),dtype=np.float32)
        A[data.edge_index[0],data.edge_index[1]]=1

        if self.adddegree:
            data.x=torch.cat([data.x,torch.tensor(A.sum(0)).unsqueeze(-1)],1)

        # set tensor for Maron's PPGN
        if self.nmax>0:
            H=torch.zeros(1,nf+2,self.nmax,self.nmax)
            H[0,0,data.edge_index[0],data.edge_index[1]]=1
            H[0,1,0:n,0:n]=torch.diag(torch.ones(data.x.shape[0]))
            for j in range(0,nf):
                H[0,j+2,0:n,0:n]=torch.diag(data.x[:,j])
            data.X2= H
            M=torch.zeros(1,2,self.nmax,self.nmax)
            for i in range(0,n):
                M[0,0,i,i]=1
            M[0,1,0:n,0:n]=1-M[0,0,0:n,0:n]
            data.M= M

        return data

class DegreeMaxEigTransform(object):

    def __init__(self,adddegree=True,maxdeg=40,addposition=False,addmaxeig=True):
        self.adddegree=adddegree
        self.maxdeg=maxdeg
        self.addposition=addposition
        self.addmaxeig=addmaxeig

    def __call__(self, data):

        n=data.x.shape[0]
        A=np.zeros((n,n),dtype=np.float32)
        A[data.edge_index[0],data.edge_index[1]]=1
        if self.adddegree:
            data.x=torch.cat([data.x,torch.tensor(1/self.maxdeg*A.sum(0)).unsqueeze(-1)],1)
        if self.addposition:
            data.x=torch.cat([data.x,data.pos],1)

        if self.addmaxeig:
            d = A.sum(axis=0)
            # normalized Laplacian matrix.
            dis=1/np.sqrt(d)
            dis[np.isinf(dis)]=0
            dis[np.isnan(dis)]=0
            D=np.diag(dis)
            nL=np.eye(D.shape[0])-(A.dot(D)).T.dot(D)
            V,U = np.linalg.eigh(nL)
            vmax=np.abs(V).max()
            # keep maximum eigenvalue for Chebnet if it is needed
            data.lmax=vmax.astype(np.float32)
        return data

reason: import numpy,import scipy,import networkx


text:
"""Class managing a Stepmania "chart"

'Stepfiles' for Stepmania are organized into 'charts': lists of annotations for by an annotator for a song with some difficulty. Many charts can point to one song so we do not want to store song features for every chart. Instead, we have this helper class that will point to song features but store chart features. This way, our training and eval sets can be list of charts, instead of some mess with list of songs attached to multiple charts.
"""

from collections import Counter
import math
import random

import numpy as np

from beatcalc import BeatCalc
from util import make_onset_feature_context, np_pad

class Chart(object):
    def __init__(self, song_metadata, metadata, annotations):
        assert len(annotations) >= 2
        self.song_metadata = song_metadata
        self.metadata = metadata

        self.annotations = annotations
        self.label_counts = Counter()

        self.beat_calc = BeatCalc(song_metadata['offset'], song_metadata['bpms'], song_metadata['stops'])

        self.first_annotation_time = self.annotations[0][2]
        self.last_annotation_time = self.annotations[-1][2]

        assert self.last_annotation_time - self.first_annotation_time > 0.0

        self.time_annotated = self.last_annotation_time - self.first_annotation_time
        self.annotations_per_second = float(len(self.annotations)) / self.time_annotated

    def get_song_metadata(self):
        return self.song_metadata

    def get_coarse_difficulty(self):
        return self.metadata[0]

    def get_foot_difficulty(self):
        return self.metadata[1]

    def get_type(self):
        return self.metadata[2]

    def get_freetext(self):
        return self.metadata[3]

    def get_nannotations(self):
        return len(self.annotations)

    def get_time_annotated(self):
        return self.time_annotated

    def get_annotations_per_second(self):
        return self.annotations_per_second

class OnsetChart(Chart):
    def __init__(self, song_metadata, song_features, frame_rate, metadata, annotations):
        super(OnsetChart, self).__init__(song_metadata, metadata, annotations)

        self.song_features = song_features
        self.nframes = song_features.shape[0]

        self.dt = dt = 1.0 / frame_rate
        self.onsets = set(int(round(t / dt)) for _, _, t, _ in self.annotations)
        self.onsets = set(filter(lambda x: x >= 0, self.onsets))
        self.onsets = set(filter(lambda x: x < self.nframes, self.onsets))
        # TODO: filter greater than song_features len?

        self.first_onset = min(self.onsets)
        self.last_onset = max(self.onsets)
        self.nframes_annotated = (self.last_onset - self.first_onset) + 1

        assert self.first_onset >= 0
        assert self.last_onset < self.nframes
        assert self.nframes > 0
        assert self.nframes_annotated > 0
        assert self.nframes >= self.nframes_annotated
        assert len(self.onsets) > 0

        self.blanks = set(range(self.nframes)) - self.onsets
        self._blanks_memoized = {}

    def get_nframes(self):
        return self.nframes

    def get_nframes_annotated(self):
        return self.nframes_annotated

    def get_nonsets(self):
        return len(self.onsets)

    def get_onsets(self):
        return self.onsets

    def get_first_onset(self):
        return self.first_onset

    def get_last_onset(self):
        return self.last_onset

    def get_example(self,
                    frame_idx,
                    dtype,
                    time_context_radius=1,
                    diff_feet_to_id=None,
                    diff_coarse_to_id=None,
                    diff_dipstick=False,
                    freetext_to_id=None,
                    beat_phase=False,
                    beat_phase_cos=False):
        feats_audio = make_onset_feature_context(self.song_features, frame_idx, time_context_radius)
        feats_other = [np.zeros(0, dtype=dtype)]

        if diff_feet_to_id:
            diff_feet = diff_feet_to_id[str(self.get_foot_difficulty())]
            diff_feet_onehot = np.zeros(max(diff_feet_to_id.values()) + 1, dtype=dtype)
            if diff_dipstick:
                diff_feet_onehot[:diff_feet + 1] = 1.0
            else:
                diff_feet_onehot[diff_feet] = 1.0
            feats_other.append(diff_feet_onehot)
        if diff_coarse_to_id:
            diff_coarse = diff_coarse_to_id[str(self.get_coarse_difficulty())]
            diff_coarse_onehot = np.zeros(max(diff_coarse_to_id.values()) + 1, dtype=dtype)
            if diff_dipstick:
                diff_coarse_onehot[:diff_coarse + 1] = 1.0
            else:
                diff_coarse_onehot[diff_coarse] = 1.0
            feats_other.append(diff_coarse_onehot)
        if freetext_to_id:
            freetext = self.get_freetext()
            if freetext not in freetext_to_id:
                freetext = None
            freetext_id = freetext_to_id[freetext]
            freetext_onehot = np.zeros(max(freetext_to_id.values()) + 1, dtype=dtype)
            freetext_onehot[freetext_id] = 1.0
            feats_other.append(freetext_onehot)
        if beat_phase:
            beat = self.beat_calc.time_to_beat(frame_idx * self.dt)
            beat_phase = beat - int(beat)
            feats_other.append(np.array([beat_phase], dtype=dtype))
        if beat_phase_cos:
            beat = self.beat_calc.time_to_beat(frame_idx * self.dt)
            beat_phase = beat - int(beat)
            beat_phase_cos = np.cos(beat * 2.0 * np.pi)
            beat_phase_cos += 1.0
            beat_phase_cos *= 0.5
            feats_other.append(np.array([beat_phase_cos], dtype=dtype))

        y = dtype(frame_idx in self.onsets)
        return np.array(feats_audio, dtype=dtype), np.concatenate(feats_other), y

    def sample(self, n, exclude_onset_neighbors=0, nunroll=0):
        if self._blanks_memoized:
            valid = self._blanks_memoized
        else:
            valid = set(range(self.get_first_onset(), self.get_last_onset() + 1))
            if exclude_onset_neighbors > 0:
                onset_neighbors = [set([x + i for x in self.onsets]) | set([x - i for x in self.onsets]) for i in xrange(1, exclude_onset_neighbors + 1)]
                onset_neighbors = reduce(lambda x, y: x | y, onset_neighbors)
                valid -= onset_neighbors
            if nunroll > 0:
                valid -= set(range(self.get_first_onset(), self.get_first_onset() + nunroll))
            self._blanks_memoized = valid

        assert n <= len(valid)
        return random.sample(valid, n)

    def sample_onsets(self, n):
        assert n <= len(self.onsets)
        return random.sample(self.onsets, n)

    def sample_blanks(self, n, exclude_onset_neighbors=0, exclude_pre_onsets=True, exclude_post_onsets=True, include_onsets=False):
        exclusion_params = (exclude_onset_neighbors, exclude_pre_onsets, exclude_post_onsets, include_onsets)

        if exclusion_params in self._blanks_memoized:
            blanks = self._blanks_memoized[exclusion_params]
        else:
            blanks = self.blanks
            if exclude_onset_neighbors > 0:
                onset_neighbors = [set([x + i for x in self.onsets]) | set([x - i for x in self.onsets]) for i in xrange(1, exclude_onset_neighbors + 1)]
                onset_neighbors = reduce(lambda x, y: x | y, onset_neighbors)
                blanks -= onset_neighbors
            if exclude_pre_onsets:
                blanks -= set(range(self.get_first_onset()))
            if exclude_post_onsets:
                blanks -= set(range(self.get_last_onset(), self.nframes))
            if include_onsets:
                blanks |= self.onsets
            self._blanks_memoized[exclusion_params] = blanks

        assert n <= len(blanks)
        return random.sample(blanks, n)

    def get_subsequence(self,
            subseq_start,
            subseq_len,
            dtype,
            zack_hack_div_2=0,
            **feat_kwargs):
        seq_feats_audio = []
        seq_feats_other = []
        seq_y = []
        for i in xrange(subseq_start - zack_hack_div_2, subseq_start + subseq_len + zack_hack_div_2):
            feats_audio, feats_other, y = self.get_example(i, dtype=dtype, **feat_kwargs)
            seq_feats_audio.append(feats_audio)
            seq_feats_other.append(feats_other)
            seq_y.append(y)
        zhmin = zack_hack_div_2
        zhmax = zack_hack_div_2 + subseq_len

        return np.array(seq_feats_audio, dtype=dtype), np.array(seq_feats_other, dtype=dtype)[zhmin:zhmax], np.array(seq_y, dtype=dtype)[zhmin:zhmax]

class SymbolicChart(Chart):
    def __init__(self, song_metadata, song_features, frame_rate, metadata, annotations, pre=1, post=False):
        super(SymbolicChart, self).__init__(song_metadata, metadata, annotations)

        self.song_features = song_features
        dt = 1.0 / frame_rate

        self.sequence = []

        self.seq_beat_phase = []
        self.seq_beat_diff = []
        self.seq_beat_abs = []
        self.seq_time_diff = []
        self.seq_time_abs = []
        self.seq_prog_diff = []
        self.seq_prog_abs = []
        self.seq_meas_phase = []
        self.seq_meas_wraps = []
        self.seq_beat_wraps = []
        self.seq_frame_idxs = []

        prepend = ['<-{}>'.format(i + 1) for i in range(pre)[::-1]]
        prepend_annotations = [self.annotations[0][:3] + [p] for p in prepend]
        annotations = prepend_annotations + self.annotations
        prog_last = 0.0
        for i in xrange(len(annotations) - 1):
            pulse_last, beat_last, time_last, label_last = annotations[i]
            pulse, beat, time, label = annotations[i + 1]
            prog = (time - self.first_annotation_time) / self.time_annotated

            assert beat >= beat_last
            assert time >= time_last

            self.sequence.append(label_last)

            self.seq_beat_phase.append(beat - math.floor(beat))
            self.seq_beat_diff.append(beat - beat_last)
            self.seq_beat_abs.append(beat)
            self.seq_time_diff.append(time - time_last)
            self.seq_time_abs.append(time)
            self.seq_prog_abs.append(prog)
            self.seq_prog_diff.append(prog - prog_last)
            self.seq_meas_phase.append(float(pulse[2]) / pulse[1])
            self.seq_meas_wraps.append(pulse[0] - pulse_last[0])
            self.seq_beat_wraps.append(int(beat) - int(beat_last))
            self.seq_frame_idxs.append(int(round(time / dt)))

            prog_last = prog

        self.sequence.append(label)

    def get_subsequence(self,
            subseq_start,
            subseq_len,
            meas_phase_cos=False,
            meas_phase_sin=False,
            meas_phase=False,
            beat_phase_cos=False,
            beat_phase_sin=False,
            beat_phase=False,
            beat_diff=False,
            beat_diff_next=False,
            beat_abs=False,
            time_diff=False,
            time_diff_next=False,
            time_abs=False,
            prog_diff=False,
            prog_abs=False,
            beat_phase_nquant=0,
            beat_phase_max_nwraps=0,
            meas_phase_nquant=0,
            meas_phase_max_nwraps=0,
            diff_feet_to_id=None,
            diff_coarse_to_id=None,
            diff_dipstick=False,
            freetext_to_id=None,
            audio_time_context_radius=0,
            audio_deviation_max=0,
            bucket_beat_diff_n=None,
            bucket_beat_diff_max=None,
            bucket_time_diff_n=None,
            bucket_time_diff_max=None,
            dtype=np.float32):
        syms = self.sequence[subseq_start:subseq_start + subseq_len + 1]
        nvalid = len(syms) - 1

        feats = np.zeros((nvalid, 0), dtype=dtype)
        if meas_phase_cos:
            feats = np.append(feats, np.cos(2. * np.pi * np.array(self.seq_meas_phase[subseq_start:subseq_start + nvalid], dtype=dtype))[:, np.newaxis], axis=1)
        if meas_phase_sin:
            feats = np.append(feats, np.sin(2. * np.pi * np.array(self.seq_meas_phase[subseq_start:subseq_start + nvalid], dtype=dtype))[:, np.newaxis], axis=1)
        if meas_phase:
            feats = np.append(feats, np.array(self.seq_meas_phase[subseq_start:subseq_start + nvalid], dtype=dtype)[:, np.newaxis], axis=1)
        if beat_phase:
            feats = np.append(feats, np.array(self.seq_beat_phase[subseq_start:subseq_start + nvalid], dtype=dtype)[:, np.newaxis], axis=1)
        if beat_phase_cos:
            feats = np.append(feats, np.cos(2. * np.pi * np.array(self.seq_beat_phase[subseq_start:subseq_start + nvalid], dtype=dtype))[:, np.newaxis], axis=1)
        if beat_phase_sin:
            feats = np.append(feats, np.sin(2. * np.pi * np.array(self.seq_beat_phase[subseq_start:subseq_start + nvalid], dtype=dtype))[:, np.newaxis], axis=1)
        if beat_diff:
            feats = np.append(feats, np.array(self.seq_beat_diff[subseq_start:subseq_start + nvalid], dtype=dtype)[:, np.newaxis], axis=1)
        if beat_diff_next:
            feat = np.array(self.seq_beat_diff[subseq_start + 1:subseq_start + nvalid + 1], dtype=dtype)
            feat = np_pad(feat, nvalid, axis=0)
            feats = np.append(feats, feat[:, np.newaxis], axis=1)
        if beat_abs:
            feats = np.append(feats, np.array(self.seq_beat_abs[subseq_start:subseq_start + nvalid], dtype=dtype)[:, np.newaxis], axis=1)
        if time_diff:
            feats = np.append(feats, np.array(self.seq_time_diff[subseq_start:subseq_start + nvalid], dtype=dtype)[:, np.newaxis], axis=1)
        if time_diff_next:
            feat = np.array(self.seq_time_diff[subseq_start + 1:subseq_start + nvalid + 1], dtype=dtype)
            feat = np_pad(feat, nvalid, axis=0)
            feats = np.append(feats, feat[:, np.newaxis], axis=1)
        if time_abs:
            feats = np.append(feats, np.array(self.seq_time_abs[subseq_start:subseq_start + nvalid], dtype=dtype)[:, np.newaxis], axis=1)
        if prog_diff:
            feats = np.append(feats, np.array(self.seq_prog_diff[subseq_start:subseq_start + nvalid], dtype=dtype)[:, np.newaxis], axis=1)
        if prog_abs:
            feats = np.append(feats, np.array(self.seq_prog_abs[subseq_start:subseq_start + nvalid], dtype=dtype)[:, np.newaxis], axis=1)
        if beat_phase_nquant > 0:
            beat_phase = np.array(self.seq_beat_phase[subseq_start:subseq_start + nvalid])
            beat_phase_quant = np.rint(beat_phase_nquant * beat_phase).astype(np.int)
            beat_phase_quant = np.mod(beat_phase_quant, beat_phase_nquant)
            beat_phase_onehot = np.zeros((nvalid, beat_phase_nquant), dtype=dtype)
            beat_phase_onehot[np.arange(nvalid), beat_phase_quant] = 1.0
            feats = np.append(feats, beat_phase_onehot, axis=1)
        if beat_phase_max_nwraps > 0:
            beat_wraps = np.array(self.seq_beat_wraps[subseq_start:subseq_start + nvalid], dtype=np.int)
            beat_wraps = np.minimum(beat_wraps, beat_phase_max_nwraps)
            beat_wraps_onehot = np.zeros((nvalid, beat_phase_max_nwraps + 1), dtype=dtype)
            beat_wraps_onehot[np.arange(nvalid), beat_wraps] = 1.0
            feats = np.append(feats, beat_wraps_onehot, axis=1)
        if meas_phase_nquant > 0:
            meas_phase = np.array(self.seq_meas_phase[subseq_start:subseq_start + nvalid])
            meas_phase_quant = np.rint(meas_phase_nquant * meas_phase).astype(np.int)
            meas_phase_quant = np.mod(meas_phase_quant, meas_phase_nquant)
            meas_phase_onehot = np.zeros((nvalid, meas_phase_nquant), dtype=dtype)
            meas_phase_onehot[np.arange(nvalid), meas_phase_quant] = 1.0
            feats = np.append(feats, meas_phase_onehot, axis=1)
        if meas_phase_max_nwraps > 0:
            meas_wraps = np.array(self.seq_meas_wraps[subseq_start:subseq_start + nvalid], dtype=np.int)
            meas_wraps = np.minimum(meas_wraps, meas_phase_max_nwraps)
            meas_wraps_onehot = np.zeros((nvalid, meas_phase_max_nwraps + 1), dtype=dtype)
            meas_wraps_onehot[np.arange(nvalid), meas_wraps] = 1.0
            feats = np.append(feats, meas_wraps_onehot, axis=1)
        if diff_feet_to_id:
            diff_feet = diff_feet_to_id[str(self.get_foot_difficulty())]
            diff_feet_onehot = np.zeros((nvalid, max(diff_feet_to_id.values()) + 1), dtype=dtype)
            if diff_dipstick:
                diff_feet_onehot[np.arange(nvalid), :diff_feet + 1] = 1.0
            else:
                diff_feet_onehot[np.arange(nvalid), diff_feet] = 1.0
            feats = np.append(feats, diff_feet_onehot, axis=1)
        if diff_coarse_to_id:
            diff_coarse = diff_coarse_to_id[self.get_coarse_difficulty()]
            diff_coarse_onehot = np.zeros((nvalid, max(diff_coarse_to_id.values()) + 1), dtype=dtype)
            if diff_dipstick:
                diff_coarse_onehot[np.arange(nvalid), :diff_coarse + 1] = 1.0
            else:
                diff_coarse_onehot[np.arange(nvalid), diff_coarse] = 1.0
            feats = np.append(feats, diff_coarse_onehot, axis=1)
        if freetext_to_id:
            freetext = self.get_freetext()
            if freetext not in freetext_to_id:
                freetext = None
            freetext_id = freetext_to_id[freetext]
            freetext_onehot = np.zeros((nvalid, max(freetext_to_id.values()) + 1), dtype=dtype)
            freetext_onehot[np.arange(nvalid), freetext_id] = 1.0
            feats = np.append(feats, freetext_onehot, axis=1)
        if bucket_beat_diff_n and bucket_beat_diff_max:
            beat_diff = np.array(self.seq_beat_diff[subseq_start:subseq_start + nvalid])

            bucket_inc = bucket_beat_diff_max / bucket_beat_diff_n
            bucket_beat_diff = beat_diff / bucket_inc
            bucket_beat_diff = bucket_beat_diff.astype(np.int32)
            bucket_beat_diff[np.where(bucket_beat_diff > bucket_beat_diff_n)] = bucket_beat_diff_n

            buckets = np.zeros((nvalid, bucket_beat_diff_n + 1), dtype=dtype)
            buckets[np.arange(nvalid), bucket_beat_diff] = 1.0

            feats = np.append(feats, buckets, axis=1)
        if bucket_time_diff_n and bucket_time_diff_max:
            time_diff = np.array(self.seq_time_diff[subseq_start:subseq_start + nvalid])

            bucket_inc = bucket_time_diff_max / bucket_time_diff_n
            bucket_time_diff = time_diff / bucket_inc
            bucket_time_diff = bucket_time_diff.astype(np.int32)
            bucket_time_diff[np.where(bucket_time_diff > bucket_time_diff_n)] = bucket_time_diff_n

            buckets = np.zeros((nvalid, bucket_time_diff_n + 1), dtype=dtype)
            buckets[np.arange(nvalid), bucket_time_diff] = 1.0

            feats = np.append(feats, buckets, axis=1)

        feats_audio = np.zeros((nvalid, 0, 0, 0), dtype=dtype)
        if audio_time_context_radius >= 0 and self.song_features is not None:
            audio_time = audio_time_context_radius * 2 + 1
            feats_audio = np.zeros((nvalid, audio_time) + self.song_features.shape[1:], dtype=dtype)
            frame_idxs = self.seq_frame_idxs[subseq_start:subseq_start + nvalid]
            for i, frame_idx in enumerate(frame_idxs):
                if audio_deviation_max > 0:
                    frame_idx += random.randint(-audio_deviation_max, audio_deviation_max)
                feats_audio[i] = make_onset_feature_context(self.song_features, frame_idx, audio_time_context_radius)

        return syms, feats, feats_audio

    def get_random_subsequence(self, subseq_len, **feat_kwargs):
        assert subseq_len <= self.get_nannotations()

        max_idx = self.get_nannotations() - subseq_len
        i = random.randint(0, max_idx)
        #TODO: first sequence incredibly unlikely to appear, balance this

        return self.get_subsequence(i, subseq_len, **feat_kwargs)

    def get_sequence(self, **feat_kwargs):
        return self.get_subsequence(0, self.get_nannotations(), **feat_kwargs)


text:
import PIL.Image as PIL
import numpy as np

image_frames = []

days = np.arange(1, 15)

for k in days:
    new_frame = PIL.open(r'C:/Users/M1/Pictures/SCEC IMAGES/Slide Show' + '/Figure_' + str(k) + '.png')
    image_frames.append(new_frame)

# Saves to working directory
image_frames[0].save('movement_timelapse.gif', format= 'GIF', append_images=image_frames[1:],
                     save_all=True, duration=800, loop=0)  # set as an infinite loop


text:
## COMPUTE TRISPEC
from __future__ import print_function
path = '/home/mkloewer/python/swm/'
import os; os.chdir(path) # change working directory
import numpy as np
from scipy import sparse
import time as tictoc
from netCDF4 import Dataset
import glob
import matplotlib.pyplot as plt

# OPTIONS
runfolder = [3]
print('Calculating 3D spectrogramms from run ' + str(runfolder))

##
def trispec(a,dt,dx,dy):
    """ Computes a wavenumber-frequency plot for 3D (t,x,y) data via radial (k = sqrt(kx**2 + ky**2)) integration. TODO: correct normalisation, so that the integral in normal space corresponds to the integral in Fourier space.
    """

    nt,ny,nx = np.shape(a)
    kx = (1/(dx))*np.hstack((np.arange(0,(nx+1)/2.),np.arange(-nx/2.+1,0)))/float(nx)
    ky = (1/(dy))*np.hstack((np.arange(0,(ny+1)/2.),np.arange(-ny/2.+1,0)))/float(ny)
    f = (1/(dt))*np.hstack((np.arange(0,(nt+1)/2.),np.arange(-nt/2.+1,0)))/float(nt)

    kxx,kyy = np.meshgrid(kx,ky)
    # radial distance from kx,ky = 0
    kk = np.sqrt(kxx**2 + kyy**2)

    if nx >= ny: #kill negative wavenumbers
        k  = kx[:int(nx/2)+1]
    else:
        k  = ky[:int(ny/2)+1]

    f = f[:int(nt/2)+1] #kill negative frequencies
    dk = k[1] - k[0]

    # create radial coordinates, associated with k[i]
    # nearest point interpolation to get points within the -.5,.5 annulus
    rcoords = []
    for i in range(len(k)):
        rcoords.append(np.where((kk>(k[i]-.5*dk))*(kk<=(k[i]+.5*dk))))

    # 3D FFT
    p = np.fft.fftn(a)
    p = np.real(p * np.conj(p))

    # mulitply by dk to have the corresponding integral
    spec = np.zeros((len(k),len(f)))
    for i in range(len(k)):
        spec[i,:] = np.sum(p[:int(nt/2)+1,rcoords[i][0],rcoords[i][1]],axis=1)*dk

    return k,f,spec

## read data
for r,i in zip(runfolder,range(len(runfolder))):
    runpath = path+'data/run%04i' % r

    if i == 0:
        #u = np.load(runpath+'/u_sub.npy')
        #v = np.load(runpath+'/v_sub.npy')
        h = np.load(runpath+'/h_sub.npy')
        time = np.load(runpath+'/t_sub.npy')
        print('run %i read.' % r)

    else:
        #u = np.concatenate((u,np.load(runpath+'/u_sub.npy')))
        #v = np.concatenate((v,np.load(runpath+'/v_sub.npy')))
        h = np.concatenate((h,np.load(runpath+'/h_sub.npy')))
        time = np.hstack((time,np.load(runpath+'/t_sub.npy')))
        print('run %i read.' % r)

t = time / 3600. / 24.  # in days
tlen = len(time)
dt = time[1] - time[0]

## read param
global param
param = np.load(runpath+'/param.npy').all()

##
k,f,spec3d = trispec(h,dt,param['dx'],param['dy'])

## rossby radius
Ro_min = param['c_phase']/f_T.max()
Ro_max = param['c_phase']/f_T.min()
# dispersion relation
Ro_disp_min = param['beta']*k/(k**2 + 1./Ro_min**2)
Ro_disp_max = param['beta']*k/(k**2 + 1./Ro_max**2)

## plotting 3D
sd = 3600.*24.

fig,ax = plt.subplots(1,1)
c1 = ax.contourf(k[1:]*1e3,f[1:]*sd,np.log10(spec3d[1:,1:].T),64,cmap='viridis')
ax.loglog()

plt.colorbar(c1,ax=ax)

ax.plot(Ro_min*np.ones(2)*1e3,[1./f[-1]/sd,1/f[1]/sd],'w')
ax.plot(Ro_max*np.ones(2)*1e3,[1./f[-1]/sd,1/f[1]/sd],'w')

ax.plot(1./k[-1]*np.ones(2)/1e3,[1./f_T.min()/3600./24.,1./f_T.min()/3600./24.],'k')
ax.plot(1./k[-1]*np.ones(2)/1e3,[1./f_T.max()/3600./24.,1./f_T.max()/3600./24.],'k')

ax.plot(1./k[1:]/1e3,1./Ro_disp_min[1:]/3600./24.,'k')
ax.plot(1./k[1:]/1e3,1./Ro_disp_max[1:]/3600./24.,'k')

ax.set_xlim(1./k[-1]/1e3,1./k[1]/1e3)
ax.set_ylim(1./f[-1],1./f[1])

ax.set_xlabel('wavelength [km]')
ax.set_ylabel('period [days]')
ax.set_title('3D power spectrum of $\eta$')
plt.tight_layout()
plt.show()

reason: import numpy,from scipy


text:
import numpy as np
import cv2

import logging

logger = logging.getLogger(__name__)

from ipso_phen.ipapi.base.ipt_abstract import IptBase
from ipso_phen.ipapi.base import ip_common as ipc

class IptCleanHorizontalNoise(IptBase):
    def build_params(self):
        self.add_enabled_checkbox()
        self.add_spin_box(
            name="min_line_size",
            desc="Minimal size for a line to be considered for removal",
            default_value=11,
            minimum=2,
            maximum=4000,
        )
        self.add_spin_box(
            name="max_iter",
            desc="Maximum number of iterations to perform",
            default_value=100,
            minimum=1,
            maximum=1000,
        )
        self.add_checkbox(
            name="fully_isolated", desc="Remove blocks of lines", default_value=1
        )

    def process_wrapper(self, **kwargs):
        """
        Clean horizontal noise:
        Removes noise in the form of horizontal lines from masks.
        Used with light barriers
        Real time: False

        Keyword Arguments (in parentheses, argument name):
            * Activate tool (enabled): Toggle whether or not tool is active
            * Minimal size for a line to be considered for removal (min_line_size):
            * Maximum number of iterations to perform (max_iter):
            * Remove blocks of lines (fully_isolated):
        """

        wrapper = self.init_wrapper(**kwargs)
        if wrapper is None:
            return False

        res = False
        try:
            if self.get_value_of("enabled") == 1:
                img = wrapper.current_image
                mask = self.get_mask()
                if mask is None:
                    logger.error(f"FAIL {self.name}: mask must be initialized")
                    return

                nz_pixels = np.count_nonzero(mask)
                min_line_size = self.get_value_of("min_line_size")
                max_iter = self.get_value_of("max_iter")
                fully_isolated = self.get_value_of("fully_isolated")
                stable_ = False
                iter_ = 0
                all_lines = []

                while not stable_ and (iter_ < max_iter):
                    stable_ = True
                    iter_ += 1
                    msk_data = ipc.MaskData(mask=mask)
                    for l in msk_data.lines_data:
                        if (l.solidity >= 0.99) and (
                            l.nz_span >= msk_data.mask_width - 4
                        ):
                            ld_up, ld_down = msk_data.find_top_bottom_non_full_lines(
                                l.height_pos
                            )
                            l.merge_or(ld_up, ld_down)
                            all_lines.append([(l.height_pos, 0, msk_data.mask_width)])
                            mask[l.height_pos] = 0
                            for i in l.nz_pos:
                                mask[l.height_pos][i] = 255
                        else:
                            lines = msk_data.horizontal_lines_at(
                                l.height_pos, min_line_size, fully_isolated
                            )
                            if not lines:
                                continue
                            all_lines.append(lines)
                            for i, line in enumerate(lines):
                                stable_ = False
                                cv2.line(
                                    mask, (line[1], line[0]), (line[2], line[0]), 0, 1
                                )
                    wrapper.store_image(mask, f"cleaned_image_iter_{iter_}")

                lines_removed_ = list(set([line[0][0] for line in all_lines]))
                if lines_removed_:
                    if "hor_lines_removed" in wrapper.data_output:
                        wrapper.data_output["hor_lines_removed"].extend(lines_removed_)
                    else:
                        wrapper.data_output["hor_lines_removed"] = lines_removed_
                wrapper.data_output["hor_pixels_removed"] = nz_pixels - np.count_nonzero(
                    mask
                )
                wrapper.data_output["hor_lines_removed_hit_plant"] = 0

                self.result = mask

                wrapper.store_image(self.result, "line_cleaned_mask")
                res = True
            else:
                wrapper.store_image(wrapper.current_image, "current_image")
                res = True
        except Exception as e:
            res = False
            logger.error(f'Failed to process {self. name}: "{repr(e)}"')
        else:
            pass
        finally:
            return res

    @property
    def name(self):
        return "Clean horizontal noise"

    @property
    def package(self):
        return "Heliasen"

    @property
    def real_time(self):
        return False

    @property
    def result_name(self):
        return "mask"

    @property
    def output_kind(self):
        return "mask"

    @property
    def use_case(self):
        return [ipc.ToolFamily.MASK_CLEANUP]

    @property
    def description(self):
        return "Removes noise in the form of horizontal lines from masks.\nUsed with light barriers"

max_stars_repo_name: tpmp-inra/ipapi

max_issues_repo_name: tpmp-inra/ipapi

max_forks_repo_name: tpmp-inra/ipapi