text:
## @ingroup Methods-Performance
# V_n_diagram.py
#
# Created:  Nov 2018, S. Karpuk
# Modified:

# ----------------------------------------------------------------------
#  Imports
# ----------------------------------------------------------------------

# SUave Imports
import SUAVE
from SUAVE.Core import Data
from SUAVE.Core import Units

from SUAVE.Analyses.Mission.Segments.Conditions import Aerodynamics,Numerics
from SUAVE.Methods.Aerodynamics.Fidelity_Zero.Lift import compute_max_lift_coeff
from SUAVE.Methods.Flight_Dynamics.Static_Stability.Approximations import datcom

# package imports
import numpy as np
import matplotlib.pyplot as plt

# ----------------------------------------------------------------------
#  Compute a V-n diagram
# ----------------------------------------------------------------------

## @ingroup Methods-Performance
def V_n_diagram(vehicle,analyses,weight,altitude,delta_ISA):

    """ Computes a V-n diagram for a given aircraft and given regulations for ISA conditions

    Source:
    S. Gudmundsson "General Aviation Aircraft Design: Applied Methods and Procedures", Butterworth-Heinemann; 1 edition
    CFR FAR Part 23: https://www.ecfr.gov/cgi-bin/text-idx?SID=0e6a13c7c1de7f501d0eb0a4d71418bd&mc=true&tpl=/ecfrbrowse/Title14/14cfr23_main_02.tpl
    CFR FAR Part 25: https://www.ecfr.gov/cgi-bin/text-idx?tpl=/ecfrbrowse/Title14/14cfr25_main_02.tpl

    Inputs:
    analyses.base.atmosphere               [SUAVE data type]
    vehicle.
      reference_area                       [m^2]
      maximum_lift_coefficient             [Unitless]
      minimum_lift_coefficient             [Unitless]
      chords.mean_aerodynamic              [m]
      envelope.FARpart_number              [Unitless]
        limit_loads.positive               [Unitless]
        limit_loads.negative               [Unitless]
        cruise_mach                        [Unitless]
    weight                                 [kg]
    altitude                               [m]
    delta_ISA                              [deg C]

    Outputs:
    V_n_data

    Properties Used:
    N/A

    Description:
    The script creates an aircraft V-n diagram based on the input parameters specified by the user.
    Depending on the certification flag, an appropriate diagram, output and log files are created.
    """

    #------------------------
    # Create a log - file
    #------------------------
    flog = open("V_n_diagram_" + vehicle.tag + ".log","w")

    print('Running the V-n diagram calculation...')
    flog.write('Running the V-n diagram calculation...\n')
    flog.write('Aircraft: ' + vehicle.tag + '\n')
    flog.write('Category: ' + vehicle.envelope.category + '\n')
    flog.write('FAR certification: Part ' + str(vehicle.envelope.FAR_part_number) + '\n\n')

    # ----------------------------------------------
    # Unpack
    # ----------------------------------------------
    flog.write('Unpacking the input and calculating required inputs...\n')
    FAR_part_number = vehicle.envelope.FAR_part_number
    atmo            = analyses.atmosphere
    Mc              = vehicle.envelope.cruise_mach

    for wing in vehicle.wings:
        reference_area  = vehicle.reference_area
        Cmac            = wing.chords.mean_aerodynamic

    for envelope in vehicle:
        pos_limit_load  = vehicle.envelope.limit_loads.positive
        neg_limit_load  = vehicle.envelope.limit_loads.negative

    category_tag = vehicle.envelope.category

    # ----------------------------------------------
    # Computing atmospheric conditions
    # ----------------------------------------------
    atmo_values       = atmo.compute_values(altitude,delta_ISA)
    SL_atmo_values    = atmo.compute_values(0,delta_ISA)
    conditions        = SUAVE.Analyses.Mission.Segments.Conditions.Aerodynamics()

    rho               = atmo_values.density
    sea_level_rho     = SL_atmo_values.density
    sea_level_gravity = atmo.planet.sea_level_gravity
    Vc                = Mc * (1.4 * 287 * atmo_values.temperature) ** 0.5

    # ------------------------------
    # Computing lift-curve slope
    # ------------------------------
    CLa = datcom(vehicle.wings.main_wing, np.array([Mc]))
    CLa = CLa[0]

    # -----------------------------------------------------------
    # Determining vehicle minimum and maximum lift coefficients
    # -----------------------------------------------------------
    try:   # aircraft maximum lift informed by user
        maximum_lift_coefficient = vehicle.maximum_lift_coefficient
    except:
        # Condition to CLmax calculation: 0.333 * Vc @ specified altitude, ISA
        conditions.freestream                   = Data()
        conditions.freestream.density           = atmo_values.density
        conditions.freestream.dynamic_viscosity = atmo_values.dynamic_viscosity
        conditions.freestream.velocity          = 0.333 * Vc
        try:
            max_lift_coefficient, induced_drag_high_lift \
                                             = compute_max_lift_coeff(vehicle,conditions)
            maximum_lift_coefficient         = max_lift_coefficient[0][0]
            vehicle.maximum_lift_coefficient = maximum_lift_coefficient

        except:
            raise ValueError("Maximum lift coefficient calculation error. Please, check inputs")

    try:    # aircraft minimum lift informed by user
        minimum_lift_coefficient = vehicle.minimum_lift_coefficient
    except:
        raise ValueError("The value not found. Specify minimum lift coefficient")

    # -----------------------------------------------------------------------------
    # Convert all terms to English (Used for FAR) and remove elements from arrays
    # -----------------------------------------------------------------------------
    altitude          = altitude / Units.ft
    rho               = rho[0,0] / Units['slug/ft**3']
    sea_level_rho     = sea_level_rho[0,0] / Units['slug/ft**3']
    density_ratio     = (rho/sea_level_rho)**0.5
    sea_level_gravity = sea_level_gravity / Units['ft/s**2']
    weight            = weight / Units['slug'] * sea_level_gravity
    reference_area    = reference_area / Units['ft**2']
    Cmac              = Cmac / Units.ft
    wing_loading      = weight / reference_area
    Vc                = Vc / Units['ft/s']

    load_factors_pos    = np.zeros(shape=(5));
    load_factors_neg    = np.zeros(shape=(5));
    load_factors_pos[1] =  1;
    load_factors_neg[1] = -1;
    airspeeds_pos       = np.zeros(shape=(5));
    airspeeds_neg       = np.zeros(shape=(5))

    # --------------------------------------------------
    # Establish limit maneuver load factors n+ and n-
    # --------------------------------------------------
    flog.write('Establish limit maneuver load factors n+ and n-...\n')
    # CFR Part 25
    # Positive and negative limits
    if FAR_part_number == 25:
        # Positive limit
        flog.write('    Estimating n+...\n')
        load_factors_pos[2] = 2.1 + 24000 / (weight + 10000)
        if load_factors_pos[2] < 2.5:
            flog.write('        Defined positive limit load factor < 2.5. Setting 2.5...\n')
            load_factors_pos[2] = 2.5
        elif load_factors_pos[2] < pos_limit_load:
            load_factors_pos[2] = pos_limit_load
        elif load_factors_pos[2] > 3.8:
            flog.write('        Defined positive limit load factor > 3.8. Setting 3.8...\n')
            load_factors_pos[2] = 3.8

        # Negative limit
        flog.write('    Estimating n-...\n')
        load_factors_neg[2] = neg_limit_load
        if load_factors_neg[2] > -1:
            flog.write('        Negative limit load factor magnitude is too small. Setting to -1...\n')
            load_factors_neg[2] = -1

    elif FAR_part_number == 23:
        if category_tag == 'normal' or category_tag == 'commuter':
            # Positive limit
            flog.write('    Estimating n+...\n')
            load_factors_pos[2] = 2.1 + 24000 / (weight + 10000)
            if load_factors_pos[2] < 2.5:
                flog.write('        Defined positive limit load factor < 2.5. Setting 2.5...\n')
                load_factors_pos[2] = 2.5
            elif load_factors_pos[2] < pos_limit_load:
                load_factors_pos[2] = pos_limit_load
            elif load_factors_pos[2] > 3.8:
                flog.write('        Defined positive limit load factor > 3.8. Setting 3.8...\n')
                load_factors_pos[2] = 3.8

            # Negative limit
            flog.write('    Estimating n-...\n')
            load_factors_neg[2] = -0.4 * load_factors_pos[2]

        elif category_tag == 'utility':
            # Positive limit
            flog.write('    Estimating n+...\n')
            load_factors_pos[2] =  pos_limit_load
            if load_factors_pos[2] < 4.4:
                 flog.write('        Defined positive limit load factor < 4.4. Setting 4.4...\n')
                 load_factors_pos[2] = 4.4

            # Negative limit
            flog.write('    Estimating n-...\n')
            load_factors_neg[2] = -0.4 * load_factors_pos[2]

        elif category_tag == 'acrobatic':
            # Positive limit
            load_factors_pos[2] = pos_limit_load
            if load_factors_pos[2] < 6.0:
                flog.write('        Defined positive limit load factor < 6.0. Setting 6.0...\n')
                load_factors_pos[2] = 6.0

            # Negative limit
            load_factors_neg[2] = -0.5 * load_factors_pos[2]

        else:
            raise ValueError("Check the category_tag input. The parameter was not found")

    else:
        raise ValueError("Check the FARflag input. The parameter was not found")

    # Check input of the limit load
    if abs(neg_limit_load) > abs(load_factors_neg[2]):
        load_factors_neg[2] = neg_limit_load

    #----------------------------------------
    # Generate a V-n diagram data structure
    #----------------------------------------
    V_n_data                          = Data()
    V_n_data.limit_loads              = Data()
    V_n_data.limit_loads.dive         = Data()
    V_n_data.load_factors             = Data()
    V_n_data.gust_load_factors        = Data()
    V_n_data.Vb_load_factor           = Data()
    V_n_data.airspeeds                = Data()
    V_n_data.Vs1                      = Data()
    V_n_data.Va                       = Data()
    V_n_data.Vb                       = Data()
    V_n_data.load_factors.positive    = load_factors_pos
    V_n_data.load_factors.negative    = load_factors_neg
    V_n_data.airspeeds.positive       = airspeeds_pos
    V_n_data.airspeeds.negative       = airspeeds_neg
    V_n_data.Vc                       = Vc
    V_n_data.weight                   = weight
    V_n_data.wing_loading             = wing_loading
    V_n_data.altitude                 = altitude
    V_n_data.density                  = rho
    V_n_data.density_ratio            = density_ratio
    V_n_data.reference_area           = reference_area
    V_n_data.maximum_lift_coefficient = maximum_lift_coefficient
    V_n_data.minimum_lift_coefficient = minimum_lift_coefficient
    V_n_data.limit_loads.positive     = load_factors_pos[2]
    V_n_data.limit_loads.negative     = load_factors_neg[2]

    # --------------------------------------------------
    # Computing critical speeds (Va, Vc, Vb, Vd, Vs1)
    # --------------------------------------------------
    flog.write('Computing critical speeds (Va, Vc, Vb, Vd, Vs1)...\n')

    # Calculate Stall and Maneuver speeds
    flog.write('    Computing Vs1 and Va...\n')
    stall_maneuver_speeds(V_n_data)

    # convert speeds to KEAS for future calculations
    convert_keas(V_n_data)

    # unpack modified airspeeds
    airspeeds_pos = V_n_data.airspeeds.positive
    airspeeds_neg = V_n_data.airspeeds.negative
    Vc            = V_n_data.Vc
    Va_pos        = V_n_data.Va.positive
    Va_neg        = V_n_data.Va.negative
    Va_pos        = V_n_data.Va.positive
    Va_neg        = V_n_data.Va.negative

    flog.write('    Checking Vc wrt Va...\n')
    if Va_neg > Vc and Va_neg > Va_pos:
        flog.write('        Negative Va > Vc. Setting Vc = 1.15 * Va...\n')
        Vc = 1.15 * Va_neg
    elif Va_pos > Vc and Va_neg < Va_pos:
        flog.write('        Positive Va > Vc. Setting Vc = 1.15 * Va...\n')
        Vc = 1.15 * Va_pos

    # Gust speeds between Vb and Vc (EAS) and minimum Vc
    miu = 2 * wing_loading / (rho * Cmac * CLa * sea_level_gravity)
    Kg  = 0.88 * miu / (5.3 + miu)

    if FAR_part_number == 25:
        if altitude < 15000:
            Uref_cruise = (-0.0008 * altitude + 56)
            Uref_rough  = Uref_cruise
            Uref_dive   = 0.5 * Uref_cruise
        else:
            Uref_cruise = (-0.0005142 * altitude + 51.7133)
            Uref_rough  = Uref_cruise
            Uref_dive   = 0.5 * Uref_cruise

        # Minimum Cruise speed Vc_min
        coefs = [1, -Uref_cruise * (2.64 + (Kg * CLa * airspeeds_pos[1]**2)/(498 * wing_loading)), 1.72424 * Uref_cruise**2 - airspeeds_pos[1]**2]
        Vc1   = max(np.roots(coefs))

    elif FAR_part_number == 23:
        if altitude < 20000:
            Uref_cruise = 50;
            Uref_dive   = 25;
        else:
            Uref_cruise = (-0.0008333 * altitude + 66.67);
            Uref_dive = (-0.0004167 * altitude + 33.334);

        if category_tag == 'commuter':
            if altitude < 20000:
                Uref_rough = 66
            else:
                Uref_rough = -0.000933 * altitude + 84.667
        else:
            Uref_rough = Uref_cruise

        # Minimum Cruise speed Vc_min
        if category_tag == 'acrobatic':
            Vc1 = 36 * wing_loading**0.5
            if wing_loading >= 20:
                Vc1 = (-0.0925 * wing_loading + 37.85) * wing_loading **0.5
        else:
            Vc1 = 33 * wing_loading**0.5
            if wing_loading >= 20:
                Vc1 = (-0.055 * wing_loading + 34.1) * wing_loading **0.5

    # checking input Cruise speed
    flog.write('    Checking Vc wrt Vcmin... \n')
    if Vc1 > Vc:
        flog.write('        Specified cruise speed is less than the minimum required. Setting th minimum required value...\n')
        Vc = Vc1

    # Dive speed
    flog.write('    Computing Vd...\n')
    if FAR_part_number == 25:
        airspeeds_pos[4] = 1.25 * Vc
    elif FAR_part_number == 23:
        if category_tag == 'acrobatic':
            airspeeds_pos[4] = 1.55 * Vc1
            if wing_loading > 20:
                airspeeds_pos[4] = (-0.0025 * wing_loading + 1.6) * Vc1
        elif category_tag == 'utility':
            airspeeds_pos[4] = 1.5 * Vc1
            if wing_loading > 20:
                airspeeds_pos[4] = (-0.001875 * wing_loading + 1.5375) * Vc1
        else:
            airspeeds_pos[4] = 1.4 * Vc1
            if wing_loading > 20:
                airspeeds_pos[4] = (-0.000625 * wing_loading + 1.4125) * Vc1

        if airspeeds_pos[4] < 1.15 * Vc:
            flog.write('        Based on min Vc, Vd is too close Vc. Setting Vd = 1.15 Vc...\n')
            airspeeds_pos[4] = 1.15 * Vc

    Vd               = airspeeds_pos[4]
    airspeeds_pos[3] = airspeeds_pos[4]
    airspeeds_neg[3] = Vc
    airspeeds_neg[4] = airspeeds_pos[4]

    # complete initial load factors
    load_factors_pos[4] = 0
    load_factors_neg[4] = 0
    if category_tag == 'acrobatic' or category_tag == 'utility':
        load_factors_neg[4] = -1

    load_factors_pos[3] = load_factors_pos[2]
    load_factors_neg[3] = load_factors_neg[2]

    # add parameters to the data structure
    V_n_data.load_factors.positive    = load_factors_pos
    V_n_data.load_factors.negative    = load_factors_neg
    V_n_data.airspeeds.positive       = airspeeds_pos
    V_n_data.airspeeds.negative       = airspeeds_neg
    V_n_data.Vd                       = Vd
    V_n_data.Vc                       = Vc

    #------------------------
    # Create Stall lines
    #------------------------
    flog.write('Creating stall lines...\n')
    Num_of_points = 20                                            # number of points for the stall line
    upper_bound = 2;
    lower_bound = 1;
    stall_line(V_n_data, upper_bound, lower_bound, Num_of_points, 1)
    stall_line(V_n_data, upper_bound, lower_bound, Num_of_points, 2)

    # ----------------------------------------------
    # Determine Gust loads
    # ----------------------------------------------
    flog.write('Calculating gust loads...\n')
    V_n_data.gust_data           = Data()
    V_n_data.gust_data.airspeeds = Data()

    V_n_data.gust_data.airspeeds.rough_gust  = Uref_rough
    V_n_data.gust_data.airspeeds.cruise_gust = Uref_cruise
    V_n_data.gust_data.airspeeds.dive_gust   = Uref_dive

    gust_loads(category_tag, V_n_data, Kg, CLa, Num_of_points, FAR_part_number, 1)
    gust_loads(category_tag, V_n_data, Kg, CLa, Num_of_points, FAR_part_number, 2)

    #----------------------------------------------------------------
    # Finalize the load factors for acrobatic and utility aircraft
    #----------------------------------------------------------------
    if category_tag == 'acrobatic' or category_tag == 'utility':
        V_n_data.airspeeds.negative    = np.append(V_n_data.airspeeds.negative, Vd)
        V_n_data.load_factors.negative = np.append(V_n_data.load_factors.negative, 0)

    # ----------------------------------------------
    # Post-processing the V-n diagram
    # ----------------------------------------------
    flog.write('Post-Processing...\n')
    V_n_data.limit_loads.positive = max(V_n_data.load_factors.positive)
    V_n_data.limit_loads.negative = min(V_n_data.load_factors.negative)

    post_processing(category_tag, Uref_rough, Uref_cruise, Uref_dive, V_n_data, vehicle)

    print('Calculation complete...')
    flog.write('Calculation complete\n')
    flog.close()

    return V_n_data

#------------------------------------------------------------------------------------------------------
# USEFUL FUNCTIONS
#------------------------------------------------------------------------------------------------------

def stall_maneuver_speeds(V_n_data):

    """ Computes stall and maneuver speeds for positive and negative halves of the the V-n diagram

    Source:
    S. Gudmundsson "General Aviation Aircraft Design: Applied Methods and Procedures", Butterworth-Heinemann; 1 edition

    Inputs:
    V_n_data.
        airspeeds.positive                      [kts]
            negative                            [kts]
        weight                                  [lb]
        density                                 [slug/ft**3]
        density_ratio                           [Unitless]
        reference_area                          [ft**2]
        maximum_lift_coefficient                [Unitless]
        minimum_lift_coefficient                [Unitless]
        load_factors.positive                   [Unitless]
        load_factors.negative                   [Unitless]

    Outputs:
    V_n_data.
        airspeeds.positive                      [kts]
            negative                            [kts]
        Vs1.positive                            [kts]
            negative                            [kts]
        Va.positive                             [kts]
            negative                            [kts]
        load_factors.positive                   [Unitless]
            negative                            [Unitless]

    Properties Used:
    N/A

    Description:
    """

    # Unpack
    airspeeds_pos    = V_n_data.airspeeds.positive
    airspeeds_neg    = V_n_data.airspeeds.negative
    weight           = V_n_data.weight
    rho              = V_n_data.density
    reference_area   = V_n_data.reference_area
    max_lift_coef    = V_n_data.maximum_lift_coefficient
    min_lift_coef    = V_n_data.minimum_lift_coefficient
    load_factors_pos = V_n_data.load_factors.positive
    load_factors_neg = V_n_data.load_factors.negative

    # Stall speeds
    airspeeds_pos[1] = (2 * weight / (rho * reference_area * max_lift_coef)) ** 0.5
    airspeeds_neg[1] = (2 * weight / (rho * reference_area * abs(min_lift_coef))) ** 0.5
    airspeeds_pos[0] = airspeeds_pos[1]
    airspeeds_neg[0] = airspeeds_neg[1]

    # Maneuver speeds
    airspeeds_pos[2] = airspeeds_pos[1] * load_factors_pos[2] ** 0.5
    airspeeds_neg[2] = (2 * weight * abs(load_factors_neg[2]) / (rho * reference_area * \
                                                                 abs(min_lift_coef))) ** 0.5

    # Pack
    V_n_data.airspeeds.positive           = airspeeds_pos
    V_n_data.airspeeds.negative           = airspeeds_neg
    V_n_data.load_factors.positive        = load_factors_pos
    V_n_data.load_factors.negative        = load_factors_neg
    V_n_data.Vs1.positive                 = airspeeds_pos[1]
    V_n_data.Vs1.negative                 = airspeeds_neg[1]
    V_n_data.Va.positive                  = airspeeds_pos[2]
    V_n_data.Va.negative                  = airspeeds_neg[2]

#------------------------------------------------------------------------------------------------------------

def stall_line(V_n_data, upper_bound, lower_bound, Num_of_points, sign_flag):

    """ Calculates Stall lines of positive and negative halves of the V-n diagram

    Source:
    S. Gudmundsson "General Aviation Aircraft Design: Applied Methods and Procedures", Butterworth-Heinemann; 1 edition

    Inputs:
    V_n_data.
        airspeeds.positive                  [kts]
            negative                        [kts]
        load_factors.positive               [Unitless]
            negative                        [Unitless]
        weight                              [lb]
        density                             [slug/ft**3]
        density_ratio                       [Unitless]
        lift_coefficient                    [Unitless]
        reference_area                      [ft**2]
    lower_bound                             [Unitless]
    Num_of_points                           [Unitless]
    sign_flag                               [Unitless]

    Outputs:
    V_n_data.
        load_factors.positive               [Unitless]
            negative                        [Unitless]
        airspeeds.positive                  [kts]
            negative                        [kts]

    Properties Used:
    N/A

    Description:
    """
    # Unpack
    weight          = V_n_data.weight
    reference_area  = V_n_data.reference_area
    density_ratio   = V_n_data.density_ratio
    rho             = V_n_data.density

    if sign_flag == 1:
        load_factors     = V_n_data.load_factors.positive
        airspeeds        = V_n_data.airspeeds.positive
        lift_coefficient = V_n_data.maximum_lift_coefficient

    elif sign_flag == 2:
        load_factors     = V_n_data.load_factors.negative
        airspeeds        = V_n_data.airspeeds.negative
        lift_coefficient = V_n_data.minimum_lift_coefficient

    delta = (airspeeds[upper_bound] - airspeeds[lower_bound]) / (Num_of_points + 1)     # Step size
    for i in range(Num_of_points):
        coef      = lower_bound + i + 1
        airspeeds = np.concatenate((airspeeds[:coef], [airspeeds[lower_bound] + (i + 1) * delta], airspeeds[coef:]))
        Vtas      = airspeeds[coef] / density_ratio * Units.knots / Units['ft/s']
        if load_factors[1] > 0:
            nl = 0.5 * rho * Vtas**2 * reference_area * lift_coefficient / weight
        else:
            nl = -0.5 * rho * Vtas**2 * reference_area * abs(lift_coefficient) / weight

        load_factors = np.concatenate((load_factors[:coef], [nl], load_factors[coef:]))

    # Pack
    if sign_flag == 1:
        V_n_data.load_factors.positive     = load_factors
        V_n_data.airspeeds.positive        = airspeeds

    elif sign_flag == 2:
        V_n_data.load_factors.negative     = load_factors
        V_n_data.airspeeds.negative        = airspeeds

    return
#--------------------------------------------------------------------------------------------------------------

def gust_loads(category_tag, V_n_data, Kg, CLa, Num_of_points, FAR_part_number, sign_flag):

    """ Calculates airspeeds and load factors for gust loads and modifies the V-n diagram

    Source:
    S. Gudmundsson "General Aviation Aircraft Design: Applied Methods and Procedures", Butterworth-Heinemann; 1 edition

    Inputs:
    V_n_data.
        airspeeds.positive                      [kts]
            negative                            [kts]
        load_factors.positive                   [Unitless]
            negative                            [Unitless]
        limit_loads.positive                    [Unitless]
            negative                            [Unitless]
        weight                                  [lb]
        wing_loading                            [lb/ft**2]
        reference_area                          [ft**2]
        density                                 [slug/ft**3]
        density_ratio                           [Unitless]
        lift_coefficient                        [Unitless]
        Vc                                      [kts]
        Vd                                      [kts]
    Uref_rough                                  [ft/s]
    Uref_cruise                                 [ft/s]
    Uref_dive                                   [ft/s]
    Num_of_points                               [Unitless]
    sign_flag                                   [Unitless]

    Outputs:
    V_n_data.
        airspeeds.positive                      [kts]
            negative                            [kts]
        load_factors.positive                   [Unitless]
            .negative                           [Unitless]
        gust_load_factors.positive              [Unitless]
            negative                            [Unitless]

    Properties Used:
    N/A

    Description:
    The function calculates the gust-related load factors at critical speeds (Va, Vc, Vd). Then, if the load factors exceed
    the standart diagram limits, the diagram is modified to include the new limit loads.
    For more details, refer to S. Gudmundsson "General Aviation Aircraft Design: Applied Methods and Procedures"
    """

    # Unpack
    weight          = V_n_data.weight
    wing_loading    = V_n_data.wing_loading
    reference_area  = V_n_data.reference_area
    density         = V_n_data.density
    density_ratio   = V_n_data.density_ratio
    Vc              = V_n_data.Vc
    Vd              = V_n_data.Vd

    if sign_flag == 1:
        airspeeds        = V_n_data.airspeeds.positive
        load_factors     = V_n_data.load_factors.positive
        lift_coefficient = V_n_data.maximum_lift_coefficient
        limit_load       = V_n_data.limit_loads.positive
        Uref_rough       = V_n_data.gust_data.airspeeds.rough_gust
        Uref_cruise      = V_n_data.gust_data.airspeeds.cruise_gust
        Uref_dive        = V_n_data.gust_data.airspeeds.dive_gust

    elif sign_flag == 2:
        airspeeds        = V_n_data.airspeeds.negative
        load_factors     = V_n_data.load_factors.negative
        lift_coefficient = V_n_data.minimum_lift_coefficient
        limit_load       = V_n_data.limit_loads.negative
        Uref_rough       = -V_n_data.gust_data.airspeeds.rough_gust
        Uref_cruise      = -V_n_data.gust_data.airspeeds.cruise_gust
        Uref_dive        = -V_n_data.gust_data.airspeeds.dive_gust

    gust_load_factors    = np.zeros(shape=(4));
    gust_load_factors[0] = 1;

    # Cruise speed Gust loads at Va and Vc
    gust_load_factors[1] = 1 + Kg * CLa * airspeeds[Num_of_points+2] * Uref_rough / (498 * wing_loading)
    gust_load_factors[2] = 1 + Kg * CLa * Vc * Uref_cruise/(498 * wing_loading)

    # Intersection between cruise gust load and Va
    if abs(gust_load_factors[1]) > abs(limit_load):
        sea_level_rho      = density / density_ratio**2
        coefs              = [709.486 * sea_level_rho * lift_coefficient, -Kg * Uref_rough * CLa, -498 * wing_loading]
        V_inters           = max(np.roots(coefs))
        load_factor_inters = 1 + Kg * CLa * V_inters * Uref_rough / (498 * wing_loading)

        airspeeds    = np.concatenate((airspeeds[:(Num_of_points + 3)], [V_inters], airspeeds[(Num_of_points + 3):]))
        load_factors = np.concatenate((load_factors[:(Num_of_points + 3)], [load_factor_inters], load_factors[(Num_of_points + 3):]))

        # Pack
        if sign_flag == 1:
            V_n_data.airspeeds.positive          = airspeeds
            V_n_data.load_factors.positive       = load_factors
            V_n_data.gust_load_factors.positive  = gust_load_factors
            V_n_data.Vb_load_factor.positive     = load_factor_inters
            V_n_data.Vb.positive                 = V_inters

        if sign_flag == 2:
            V_n_data.airspeeds.negative          = airspeeds
            V_n_data.load_factors.negative       = load_factors
            V_n_data.gust_load_factors.negative  = gust_load_factors
            V_n_data.Vb_load_factor.negative     = load_factor_inters
            V_n_data.Vb.negative                 = V_inters

        # continue stall lines
        Num_of_points_ext       = 5;
        upper_bound             = Num_of_points+3;
        lower_bound             = Num_of_points+2;
        stall_line(V_n_data, upper_bound, lower_bound, Num_of_points_ext, sign_flag)
        Num_of_points = Num_of_points + Num_of_points_ext + 1

        # Unpack
        if sign_flag == 1:
            airspeeds        = V_n_data.airspeeds.positive
            load_factors     = V_n_data.load_factors.positive
            lift_coefficient = V_n_data.maximum_lift_coefficient

        elif sign_flag == 2:
            airspeeds        = V_n_data.airspeeds.negative
            load_factors     = V_n_data.load_factors.negative
            lift_coefficient = V_n_data.minimum_lift_coefficient

        # insert the cruise speed Vc in the positive load factor line
        if load_factors[1] > 0:
            airspeeds    = np.concatenate((airspeeds[:(Num_of_points + 3)], [Vc], airspeeds[(Num_of_points + 3):]))
            load_factors = np.concatenate((load_factors[:(Num_of_points + 3)], [gust_load_factors[2]], \
                                           load_factors[(Num_of_points + 3):]))

    # Intersection between cruise gust load and maximum load at Vc
    elif abs(gust_load_factors[2]) > abs(limit_load):
        V_inters      = 498 * (load_factors[Num_of_points + 2] - 1) * wing_loading / (Kg * Uref_cruise * CLa)
        airspeeds     = np.concatenate((airspeeds[:(Num_of_points + 3)], [V_inters], airspeeds[(Num_of_points + 3):]))
        load_factors  = np.concatenate((load_factors[:(Num_of_points + 3)], [limit_load], \
                                       load_factors[(Num_of_points + 3):]))
        Num_of_points = Num_of_points + 1

        if load_factors[1] > 0:
            airspeeds    = np.concatenate((airspeeds[:(Num_of_points + 3)], [Vc], airspeeds[(Num_of_points + 3):]))
            load_factors = np.concatenate((load_factors[:(Num_of_points + 3)], [gust_load_factors[2]], \
                                           load_factors[(Num_of_points + 3):]))
        else:
            load_factors[len(airspeeds)-2] = gust_load_factors[2]

    # Dive speed Gust loads Vd
    gust_load_factors[3] = 1 + Kg * CLa * Vd * Uref_dive / (498 * wing_loading)

    # Resolve the upper half of the dive section
    if load_factors[1] > 0:
        if abs(gust_load_factors[2]) > abs(limit_load):
            if abs(gust_load_factors[3]) > abs(limit_load):
                load_factors[len(load_factors) - 2] = gust_load_factors[3]
            else:
                airspeeds, load_factors = gust_dive_speed_intersection(category_tag, load_factors, gust_load_factors, airspeeds, \
                                                                       len(load_factors)-2, Vc, Vd)

    # Resolve the lower half of the dive section
    else:
        if gust_load_factors[3] < load_factors[len(load_factors) - 1]:
            airspeeds    = np.concatenate((airspeeds[:(len(load_factors) - 1)], [airspeeds[(len(load_factors) - 1)]], \
                                           airspeeds[(len(load_factors) - 1):]))
            load_factors = np.concatenate((load_factors[:(len(load_factors) - 1)], [gust_load_factors[3]], \
                                           load_factors[(len(load_factors) - 1):]))

            if abs(gust_load_factors[2]) > abs(limit_load):
                load_factors[len(load_factors) - 2] = gust_load_factors[3]
            else:
                airspeeds, load_factors = gust_dive_speed_intersection(category_tag, load_factors, gust_load_factors, airspeeds, \
                                                                       len(load_factors)-2, Vc, Vd)

            V_n_data.limit_loads.dive.negative = load_factors[len(load_factors) - 2]
        else:
            V_n_data.limit_loads.dive.negative = load_factors[len(load_factors) - 1]

    # gusts load extension for gust lines at Vd
    gust_load_factors = np.append(gust_load_factors, 1 + Kg * CLa * (1.05 * Vd) * Uref_cruise/(498 * wing_loading))
    gust_load_factors = np.append(gust_load_factors, 1 + Kg * CLa * (1.05 * Vd) * Uref_dive/(498 * wing_loading))

    # guts load extension for gust lines at Vb and Vc for a rough gust
    if category_tag == 'commuter':
        gust_load_factors = np.append(gust_load_factors, 1 + Kg * CLa * (1.05 * Vd) * Uref_rough/(498 * wing_loading))
    else:
        gust_load_factors = np.append(gust_load_factors, 0)

    # Pack
    if sign_flag == 1:
        V_n_data.airspeeds.positive          = airspeeds
        V_n_data.load_factors.positive       = load_factors
        V_n_data.gust_load_factors.positive  = gust_load_factors
        V_n_data.limit_loads.dive.positive   = load_factors[len(load_factors) - 2]

    if sign_flag == 2:
        V_n_data.airspeeds.negative          = airspeeds
        V_n_data.load_factors.negative       = load_factors
        V_n_data.gust_load_factors.negative  = gust_load_factors

    return
#------------------------------------------------------------------------------------------------------------------------

def gust_dive_speed_intersection(category_tag, load_factors, gust_load_factors, airspeeds, element_num, Vc, Vd):

    """ Calculates intersection between the general V-n diagram and the gust load for Vd

    Source:
    S. Gudmundsson "General Aviation Aircraft Design: Applied Methods and Procedures", Butterworth-Heinemann; 1 edition

    Inputs:
    load_factors                                [Unitless]
    gust_load_factors                           [Unitless]
    airspeeds                                   [kts]
    Vc                                          [kts]
    Vd                                          [kts]
    element_num                                 [Unitless]

    Outputs:
    airspeeds                                    [kts]
    load_factors                                 [Unitless]

    Properties Used:
    N/A

    Description:
    A specific function for CFR FAR Part 25 regulations. For negative loads, the general curve must go linear from a specific
    load factor at Vc to zero at Vd. Gust loads may put a non-zero load factor at Vd, so an intersection between the old and
    the new curves is required
    """

    if load_factors[1] > 0:
        V_inters = (load_factors[element_num] - gust_load_factors[2]) * (Vd - Vc) \
                    / (gust_load_factors[3] - gust_load_factors[2]) + Vc
        load     = load_factors[element_num]
    else:
        if category_tag == 'acrobatic':
            V_inters = ((gust_load_factors[3] + 1) * Vc + (min(load_factors) - gust_load_factors[2]) * Vd) / \
                       (gust_load_factors[3] - gust_load_factors[2] + min(load_factors) + 1)
            load     = (min(load_factors) + 1) / (Vc - Vd) * V_inters - ((min(load_factors) + 1) / (Vc - Vd) * Vd + 1)
        else:
            V_inters = (gust_load_factors[3] * Vc + (min(load_factors) - gust_load_factors[2]) * Vd) / \
                       (gust_load_factors[3] - gust_load_factors[2] + min(load_factors))
            load     = min(load_factors) * (V_inters - Vd) / (Vc - Vd)

    airspeeds    = np.concatenate((airspeeds[:(element_num)], [V_inters], airspeeds[(element_num):]))
    load_factors = np.concatenate((load_factors[:(element_num)], [load], \
                                        load_factors[(element_num):]))

    return airspeeds, load_factors
#--------------------------------------------------------------------------------------------------------------------------

def post_processing(category_tag, Uref_rough, Uref_cruise, Uref_dive, V_n_data, vehicle):

    """ Plot graph, save the final figure, and create results output file

    Source:

    Inputs:
    V_n_data.
        airspeeds.positive                  [kts]
        airspeeds.negative                  [kts]
        Vc                                  [kts]
        Vd                                  [kts]
        Vs1.positive                        [kts]
            negative                        [kts]
        Va.positive                         [kts]
            negative                        [kts]
        load_factors.positive               [Unitless]
            negative                        [Unitless]
        gust_load_factors.positive          [Unitless]
            negative                        [Unitless]
        weight                              [lb]
        altitude                            [ft]
    vehicle._base.tag                       [Unitless]
    Uref_rough                              [ft/s]
    Uref_cruise                             [ft/s]
    Uref_dive                               [ft/s]

    Outputs:

    Properties Used:
    N/A

    Description:
    """

    # Unpack
    load_factors_pos        = V_n_data.load_factors.positive
    load_factors_neg        = V_n_data.load_factors.negative
    airspeeds_pos           = V_n_data.airspeeds.positive
    airspeeds_neg           = V_n_data.airspeeds.negative
    Vc                      = V_n_data.Vc
    Vd                      = V_n_data.Vd
    Vs1_pos                 = V_n_data.Vs1.positive
    Vs1_neg                 = V_n_data.Vs1.negative
    Va_pos                  = V_n_data.Va.positive
    Va_neg                  = V_n_data.Va.negative
    gust_load_factors_pos   = V_n_data.gust_load_factors.positive
    gust_load_factors_neg   = V_n_data.gust_load_factors.negative
    weight                  = V_n_data.weight
    altitude                = V_n_data.altitude

    #-----------------------------
    # Plotting the V-n diagram
    #-----------------------------
    fig, ax = plt.subplots()
    ax.fill(airspeeds_pos, load_factors_pos, c='b', alpha=0.3)
    ax.fill(airspeeds_neg, load_factors_neg, c='b', alpha=0.3)
    ax.plot(airspeeds_pos, load_factors_pos, c='b')
    ax.plot(airspeeds_neg, load_factors_neg, c='b')

    # Plotting gust lines
    ax.plot([0, Vc,1.05*Vd],[1,gust_load_factors_pos[2],gust_load_factors_pos[len(gust_load_factors_pos)-3]],'--', c='r', label = ('Gust ' + str(round(Uref_cruise)) + 'fps'))
    ax.plot([0, Vd,1.05*Vd],[1,gust_load_factors_pos[3],gust_load_factors_pos[len(gust_load_factors_pos)-2]],'--', c='g', label = ('Gust ' + str(round(Uref_dive)) + 'fps'))
    ax.plot([0, Vc,1.05*Vd],[1,gust_load_factors_neg[2],gust_load_factors_neg[len(gust_load_factors_neg)-3]],'--', c='r')
    ax.plot([0, Vd,1.05*Vd],[1,gust_load_factors_neg[3],gust_load_factors_neg[len(gust_load_factors_neg)-2]],'--', c='g')

    if category_tag == 'commuter':
        ax.plot([0, 1.05*Vd],[1,gust_load_factors_pos[len(gust_load_factors_pos)-1]],'--', c='m', label = ('Gust ' + str(round(Uref_rough)) + 'fps'))
        ax.plot([0, 1.05*Vd],[1,gust_load_factors_neg[len(gust_load_factors_neg)-1]],'--', c='m')

    # Formating the plot
    ax.set_xlabel('Airspeed, KEAS')
    ax.set_ylabel('Load Factor')
    ax.set_title(vehicle.tag + '  Weight=' + str(round(weight)) + 'lb  ' + ' Altitude=' + str(round(altitude)) + 'ft ')
    ax.legend()
    ax.grid()
    plt.savefig('Vn_diagram_'+ vehicle.tag + '.png')

    #---------------------------------
    # Creating results output file
    #---------------------------------
    fres = open("V_n_diagram_results_" + vehicle.tag +".dat","w")
    fres.write('V-n diagram summary\n')
    fres.write('-------------------\n')
    fres.write('Aircraft: ' + vehicle.tag + '\n')
    fres.write('category: ' + vehicle.envelope.category + '\n')
    fres.write('FAR certification: Part ' + str(vehicle.envelope.FAR_part_number) + '\n')
    fres.write('Weight = ' + str(round(weight)) + ' lb\n')
    fres.write('Altitude = ' + str(round(altitude)) + ' ft\n')
    fres.write('---------------------------------------------------------------\n\n')
    fres.write('Airspeeds: \n')
    fres.write('    Positive stall speed (Vs1)   = ' + str(round(Vs1_pos,1)) + ' KEAS\n')
    fres.write('    Negative stall speed (Vs1)   = ' + str(round(Vs1_neg,1)) + ' KEAS\n')
    fres.write('    Positive maneuver speed (Va) = ' + str(round(Va_pos,1))  + ' KEAS\n')
    fres.write('    Negative maneuver speed (Va) = ' + str(round(Va_neg,1))  + ' KEAS\n')
    fres.write('    Cruise speed (Vc)            = ' + str(round(Vc,1))      + ' KEAS\n')
    fres.write('    Dive speed (Vd)              = ' + str(round(Vd,1))      + ' KEAS\n')
    fres.write('Load factors: \n')
    fres.write('    Positive limit load factor (n+) = ' + str(round(max(load_factors_pos),2)) + '\n')
    fres.write('    Negative limit load factor (n-) = ' + str(round(min(load_factors_neg),2)) + '\n')
    fres.write('    Positive load factor at Vd      = ' + str(round(V_n_data.limit_loads.dive.positive,2)) + '\n')
    fres.write('    Negative load factor at Vd      = ' + str(round(V_n_data.limit_loads.dive.negative,2)) + '\n')

    return
#------------------------------------------------------------------------------------------------------------------------

def convert_keas(V_n_data):

    """ Convert speed to KEAS

    Source:

    Inputs:
    V_n_data.
        airspeeds.positive              [ft/s]
            negative                    [ft/s]
        Vc                              [ft/s]
        Va.positive                     [ft/s]
        Va.negative                     [ft/s]
        Vs1.negative                    [ft/s]
        Vs1.positive                    [ft/s]
        density_ratio                   [Unitless]

    Outputs:
    V_n_data.
        airspeeds.positive              [kts]
            negative                    [kts]
        Vc                              [kts]
        Va.positive                     [kts]
        Va.negative                     [kts]
        Vs1.negative                    [kts]
        Vs1.positive                    [kts]

    Properties Used:
    N/A

    Description:
    """

    # Unpack
    airspeeds_pos = V_n_data.airspeeds.positive
    airspeeds_neg = V_n_data.airspeeds.negative
    density_ratio = V_n_data.density_ratio
    Vc            = V_n_data.Vc
    Va_pos        = V_n_data.Va.positive
    Va_neg        = V_n_data.Va.negative
    Vs1_neg       = V_n_data.Vs1.negative
    Vs1_pos       = V_n_data.Vs1.positive

    airspeeds_pos = airspeeds_pos * Units['ft/s'] / Units.knots * density_ratio
    airspeeds_neg = airspeeds_neg * Units['ft/s'] / Units.knots * density_ratio

    Vs1_pos       = Vs1_pos * Units['ft/s'] / Units.knots * density_ratio
    Vs1_neg       = Vs1_neg * Units['ft/s'] / Units.knots * density_ratio
    Va_pos        = Va_pos * Units['ft/s'] / Units.knots * density_ratio
    Va_neg        = Va_neg * Units['ft/s'] / Units.knots * density_ratio
    Vc            = Vc[0] * Units['ft/s'] / Units.knots * density_ratio
    Vc            = Vc[0]

    # Pack
    V_n_data.airspeeds.positive = airspeeds_pos
    V_n_data.airspeeds.negative = airspeeds_neg
    V_n_data.Vc                 = Vc
    V_n_data.Va.positive        = Va_pos
    V_n_data.Va.negative        = Va_neg
    V_n_data.Vs1.negative       = Vs1_neg
    V_n_data.Vs1.positive       = Vs1_pos


text:
#!/usr/bin/python3
# default python version used : 3.9
# follow this link to change your python version if you have
# a different version installed
# https://stackoverflow.com/questions/52584907/how-to-downgrade-python-from-3-7-to-3-6

# ######################################
# #############  Notes #################
# ######################################

# Adapted for Python by SAUpadhyayula
# Originally written by TRHayes for Matlab

# 1.0.0 2021-07-23 : SAUpadhyayula wrote it
# 2.1.0 2020-01-17 : OSF release updated for TRHayes's matlab code

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

import numpy as np
from math import hypot

def patch_stitch(dat, scene, scale):
    """
    Compiles the patch ratings and stitches them together.

    This takes in the patch ratings for each patch across all participants,
    computes the average rating, and stores the average ratings for all such
    patches into an array of shape [width x height]. It does it for each of
    the specified scales.

    Args:
        dat (dict): The input dict file is a dictionary of dicts.
        It should containg the following fields: {scene}.
        Scene field contains  a dict {name of the image:scale} key value pair.
        Scale contains another dictionary with key
        value pairs as {scale_name: patch_ids}. Scale name is:
        ["coarse", "fine"].
        patch_ids correspond to the number in the grid that was used to
        generate the patches.
        Refer to patch_cut.py and create_scene_patches.py for more
        information.

        scene (string): Name of the image file to be processed.

        scale (string): Scale at which the image should be processed.

    Returns:
        dict: Returns a dictionary of dictionaries. It would contain the
        following fields: {scene}. Scene field contains another dict with
        {name of the image: scale} key value pair. Scale contains another
        dictionary with {scale_name: rating_array} as a key value pair.
        Rating array is a numpy array with dimensions [width X height]

    Raises:
        Exception: description

    """

    # %% 010: Define image related parameters
    patch = {}
    patch["img_size"] = [768, 1024]  # resolution of the input image
    # Patch diameter - Fine patches have smaller sizes
    patch["diameter"] = [87, 205]
    patch["density"] = [300, 108]  # Number of pixels per patch
    patch["scale"] = ["fine", "coarse"]

    # %% 020: Create the image grid

    # Define the grid
    x_grid = {}
    y_grid = {}

    # Illustrate how the grid templates look
    for i, d in enumerate(patch["scale"]):
        density = patch["density"][i]
        # Determine pixel frequency to create a grid of patches
        freq = np.round(
            np.sqrt(patch["img_size"][0] * patch["img_size"][1]) / np.sqrt(density)
        )
        [y, x] = np.meshgrid(
            np.arange(freq, patch["img_size"][1], freq),
            np.arange(freq, patch["img_size"][0], freq),
        )

        # Center grid on images
        offset_x = freq - (patch["img_size"][0] - np.max(x) + freq) / 2
        offset_y = freq - (patch["img_size"][1] - np.max(y) + freq) / 2

        # Draw circle patch grid on the gray image
        # th = np.arange(0, 2 * np.pi + 1, np.pi / 50)
        x_grid[d] = x - offset_x
        y_grid[d] = y - offset_y

        x_grid[d] = x_grid[d].flatten()
        y_grid[d] = y_grid[d].flatten()

    # %% 030: Create the output rating_array
    xdim = patch["img_size"][0]
    ydim = patch["img_size"][1]
    rating_array = [[[] for i in range(ydim)] for i in range(xdim)]

    # %% 040: Fill in the array with patch ratings
    # Get the patch diameter
    if scale == "fine":
        d = patch["diameter"][0]
    elif scale == "coarse":
        d = patch["diameter"][1]

    for en, (x, y) in enumerate(zip(x_grid[scale], y_grid[scale])):
        # Get the corresponding pixels of the patch location in the image
        lx = max(int(np.round(x)) - int(d / 2), 0)
        rx = min(int(np.round(x)) + int(d / 2), xdim)
        ly = max(int(np.round(y)) - int(d / 2), 0)
        ry = min(int(np.round(y)) + int(d / 2), ydim)

        for i in range(lx, rx):
            for j in range(ly, ry):
                if hypot(i - x, j - y) <= int(d / 2):
                    rating_array[i][j].extend(dat[scene][scale][en + 1])

    return rating_array

ext: py

reason: import numpy


text:
from typing import List

import numpy as np

def save_np_ndarray(string_format: str,
                    string_args: List[str],
                    destination_dir: str,
                    array: np.ndarray
                    ) -> None:
    """Saves numpy arrays with the specified format name to the output dir.

    Args:
        string_format (str): filename format the file should have
        string_args (List[str]): list of parameters to be inserted to the format
        destination_dir (str): directory where the file should be saved
        array (np.ndarray): the array that should be saved
    """
    np.save((destination_dir + string_format).format(*string_args), array)

def cut_np_2D_ndarray(array: np.ndarray,
                      start_offset: int,
                      end_offset: int
                      ) -> np.ndarray:
    #TODO: add documentation
    return array[:, max(0, start_offset) : min(end_offset, array.shape[1])]

def cut_np_1D_ndarray(array: np.ndarray,
                   start_offset: int,
                   end_offset: int
                   ) -> np.ndarray:
    #TODO: add documentation
    return array[max(0, start_offset) : min(end_offset, len(array))]

max_stars_repo_name: MatePaulinovic/thesis

max_issues_repo_name: MatePaulinovic/thesis

max_forks_repo_name: MatePaulinovic/thesis


text:
# --------------------------------------------------------
# Fast R-CNN
# Copyright (c) 2015 Microsoft
# Licensed under The MIT License [see LICENSE for details]
# Written by Ross Girshick
# --------------------------------------------------------

"""Factory method for easily getting imdbs by name."""

__sets = {}

from datasets.pascal_voc import pascal_voc
from datasets.coco import coco
from datasets.icub import icub_dataset
import numpy as np

#icub transformation dataset
icub_manual_path = '/home/IIT.LOCAL/emaiettini/workspace/Repos/SSM/data/iCubWorld-Transformations_manual'
# icub_devkit_path = '/home/elisa/Repos/py-faster_icubworld/data/iCubWorld-Translation_devkit'

for t in ['train', 'test']:
    for task in ['TASK2']:
        for numObjs in ['30', '21']:
            for supervision in ['supervised', 'unsupervised']:
                split = '{}_{}_{}objs_{}'.format(t, task, numObjs, supervision)
                name = '{}_{}_{}objs_{}'.format(t, task, numObjs, supervision)
                __sets[name] = (lambda split=split: icub_dataset(split))

for t in ['train']:
    for task in ['TASK2', 'TASK2_EAL']:
        for numObjs in ['21']:
            for color in ['white', 'pois_odd', 'green', 'full']:
                split = '{}_{}_{}objs_{}'.format(t, task, numObjs, color)
                name = '{}_{}_{}objs_{}'.format(t, task, numObjs, color)
                __sets[name] = (lambda split=split: icub_dataset(split))

for t in ['test']:
    for task in ['TASK2', 'TASK2_EAL']:
        for numObjs in ['21']:
            for color in ['white', 'pois_odd', 'green']:
                split = '{}_{}_{}objs_{}'.format(t, task, numObjs, color)
                name = '{}_{}_{}objs_{}'.format(t, task, numObjs, color)
                __sets[name] = (lambda split=split: icub_dataset(split, devkit_path=icub_manual_path))

for t in ['train', 'test']:
    for task in ['TASK2']:
        for numObjs in ['21']:
                split = '{}_{}_{}objs'.format(t, task, numObjs, supervision)
                name = '{}_{}_{}objs'.format(t, task, numObjs, supervision)
                __sets[name] = (lambda split=split: icub_dataset(split))

for t in ['train', 'test']:
    for task in ['TASK2']:
        for numObjs in ['21']:
                for supervision in ['acquisition1', 'acquisition2']:
                    split = '{}_{}_{}objs_{}'.format(t, task, numObjs, supervision)
                    name = '{}_{}_{}objs_{}'.format(t, task, numObjs, supervision)
                    __sets[name] = (lambda split=split: icub_dataset(split, devkit_path=icub_manual_path))

for t in ['train']:
    for task in ['TASK2', 'TASK2_sequence']:
        for numObjs in ['30', '21', '5', '10']:
                for supervision in ['TRA_supervised', 'NoTRA_unsupervised', 'NoTRA_unsupervised_ordered']:
                    split = '{}_{}_{}objs_{}'.format(t, task, numObjs, supervision)
                    name = '{}_{}_{}objs_{}'.format(t, task, numObjs, supervision)
                    __sets[name] = (lambda split=split: icub_dataset(split))

for t in ['train']:
    for task in ['TASK2']:
        for numObjs in ['30']:
                for supervision in ['supervised', 'unsupervised']:
                    split = '{}_{}_{}objs_1over8_{}_fake'.format(t, task, numObjs, supervision)
                    name = '{}_{}_{}objs_1over8_{}_fake'.format(t, task, numObjs, supervision)
                    __sets[name] = (lambda split=split: icub_dataset(split))

for t in ['test']: # split sarebbe l'imageset
    for task in ['TASK2']:
        for numObjs in ['30', '5', '10']:
            split = '{}_{}_{}objs_manual'.format(t, task, numObjs)
            name = '{}_{}_{}objs_manual'.format(t, task, numObjs)
            __sets[name] = (lambda split=split: icub_dataset(split, devkit_path=icub_manual_path))

# Set up voc_<year>_<split> using selective search "fast" mode
for year in ['2007', '2012', '0712']:
    for split in ['train', 'val', 'trainval', 'test']:
        name = 'voc_{}_{}'.format(year, split)
        __sets[name] = (lambda split=split, year=year: pascal_voc(split, year))

# Set up coco_2014_<split>
for year in ['2014']:
    for split in ['train', 'val', 'minival', 'valminusminival']:
        name = 'coco_{}_{}'.format(year, split)
        __sets[name] = (lambda split=split, year=year: coco(split, year))

# Set up coco_2015_<split>
for year in ['2015']:
    for split in ['test', 'test-dev']:
        name = 'coco_{}_{}'.format(year, split)
        __sets[name] = (lambda split=split, year=year: coco(split, year))

def get_imdb(name):
  """Get an imdb (image database) by name."""
  if name not in __sets:
    raise KeyError('Unknown dataset: {}'.format(name))
  return __sets[name]()

def list_imdbs():
    """List all registered imdbs."""
    return __sets.keys()


text:
from mnist2 import load_mnist
from numpy import *

def read_dataset(size_training, size_testing):
    digits = [0,1,2,3,4,5,6,7,8,9]
    images_train, labels_train = load_mnist('training',digits)
    images_test, labels_test = load_mnist('testing',digits)

    total_training = len(labels_train)
    if(size_training > total_training):
        size_training = total_training
    total_testing = len(labels_test)
    if(size_testing > total_testing):
        size_testing = total_testing

    random_training_instances = range(total_training)
    random.shuffle(random_training_instances)
    random_testing_instances = range(total_testing)
    random.shuffle(random_training_instances)

    images_train = images_train.astype(float64)
    images_test  = images_test.astype(float64)
    images_train = images_train[random_training_instances[0:(size_training)],:]
    labels_train = labels_train[random_training_instances[0:(size_training)]]
    images_test  = images_test[random_testing_instances[0:(size_testing)],:]
    labels_test  = labels_test[random_testing_instances[0:(size_testing)]]

    return images_train, labels_train, images_test, labels_test

max_stars_repo_path: Assignment 1 - Neural networks/read_dataset.py

max_stars_repo_name: marioyc/Hand-written-Digit-Recognition

max_issues_repo_name: marioyc/Hand-written-Digit-Recognition

max_forks_repo_name: marioyc/Hand-written-Digit-Recognition

reason: from numpy


text:
import gym
import numpy as np
import pytest
import torch as th

from stable_baselines3 import A2C, PPO, SAC
from stable_baselines3.common.callbacks import BaseCallback
from stable_baselines3.common.policies import ActorCriticPolicy

class CustomEnv(gym.Env):
    def __init__(self, max_steps=8):
        super(CustomEnv, self).__init__()
        self.observation_space = gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
        self.action_space = gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
        self.max_steps = max_steps
        self.n_steps = 0

    def seed(self, seed):
        self.observation_space.seed(seed)

    def reset(self):
        self.n_steps = 0
        return self.observation_space.sample()

    def step(self, action):
        self.n_steps += 1

        done = False
        reward = 0.0
        if self.n_steps >= self.max_steps:
            reward = 1.0
            done = True

        return self.observation_space.sample(), reward, done, {}

class InfiniteHorizonEnv(gym.Env):
    def __init__(self, n_states=4):
        super().__init__()
        self.n_states = n_states
        self.observation_space = gym.spaces.Discrete(n_states)
        self.action_space = gym.spaces.Box(low=-1, high=1, shape=(2,), dtype=np.float32)
        self.current_state = 0

    def reset(self):
        self.current_state = 0
        return self.current_state

    def step(self, action):
        self.current_state = (self.current_state + 1) % self.n_states
        return self.current_state, 1.0, False, {}

class CheckGAECallback(BaseCallback):
    def __init__(self):
        super(CheckGAECallback, self).__init__(verbose=0)

    def _on_rollout_end(self):
        buffer = self.model.rollout_buffer
        rollout_size = buffer.size()

        max_steps = self.training_env.envs[0].max_steps
        gamma = self.model.gamma
        gae_lambda = self.model.gae_lambda
        value = self.model.policy.constant_value
        # We know in advance that the agent will get a single
        # reward at the very last timestep of the episode,
        # so we can pre-compute the lambda-return and advantage
        deltas = np.zeros((rollout_size,))
        advantages = np.zeros((rollout_size,))
        # Reward should be 1.0 on final timestep of episode
        rewards = np.zeros((rollout_size,))
        rewards[max_steps - 1 :: max_steps] = 1.0
        # Note that these are episode starts (+1 timestep from done)
        episode_starts = np.zeros((rollout_size,))
        episode_starts[::max_steps] = 1.0

        # Final step is always terminal (next would episode_start = 1)
        deltas[-1] = rewards[-1] - value
        advantages[-1] = deltas[-1]
        for n in reversed(range(rollout_size - 1)):
            # Values are constants
            episode_start_mask = 1.0 - episode_starts[n + 1]
            deltas[n] = rewards[n] + gamma * value * episode_start_mask - value
            advantages[n] = deltas[n] + gamma * gae_lambda * advantages[n + 1] * episode_start_mask

        # TD(lambda) estimate, see Github PR #375
        lambda_returns = advantages + value

        assert np.allclose(buffer.advantages.flatten(), advantages)
        assert np.allclose(buffer.returns.flatten(), lambda_returns)

    def _on_step(self):
        return True

class CustomPolicy(ActorCriticPolicy):
    """Custom Policy with a constant value function"""

    def __init__(self, *args, **kwargs):
        super(CustomPolicy, self).__init__(*args, **kwargs)
        self.constant_value = 0.0

    def forward(self, obs, deterministic=False):
        actions, values, log_prob = super().forward(obs, deterministic)
        # Overwrite values with ones
        values = th.ones_like(values) * self.constant_value
        return actions, values, log_prob

@pytest.mark.parametrize("model_class", [A2C, PPO])
@pytest.mark.parametrize("gae_lambda", [1.0, 0.9])
@pytest.mark.parametrize("gamma", [1.0, 0.99])
@pytest.mark.parametrize("num_episodes", [1, 3])
def test_gae_computation(model_class, gae_lambda, gamma, num_episodes):
    env = CustomEnv(max_steps=64)
    rollout_size = 64 * num_episodes
    model = model_class(
        CustomPolicy,
        env,
        seed=1,
        gamma=gamma,
        n_steps=rollout_size,
        gae_lambda=gae_lambda,
    )
    model.learn(rollout_size, callback=CheckGAECallback())

    # Change constant value so advantage != returns
    model.policy.constant_value = 1.0
    model.learn(rollout_size, callback=CheckGAECallback())

@pytest.mark.parametrize("model_class", [A2C, SAC])
@pytest.mark.parametrize("handle_timeout_termination", [False, True])
def test_infinite_horizon(model_class, handle_timeout_termination):
    max_steps = 8
    gamma = 0.98
    env = gym.wrappers.TimeLimit(InfiniteHorizonEnv(n_states=4), max_steps)
    kwargs = {}
    if model_class == SAC:
        policy_kwargs = dict(net_arch=[64], n_critics=1)
        kwargs = dict(
            replay_buffer_kwargs=dict(handle_timeout_termination=handle_timeout_termination),
            tau=0.5,
            learning_rate=0.005,
        )
    else:
        policy_kwargs = dict(net_arch=[64])
        kwargs = dict(learning_rate=0.002)
        # A2C always handle timeouts
        if not handle_timeout_termination:
            return

    model = model_class("MlpPolicy", env, gamma=gamma, seed=1, policy_kwargs=policy_kwargs, **kwargs)
    model.learn(1500)
    # Value of the initial state
    obs_tensor = model.policy.obs_to_tensor(0)[0]
    if model_class == A2C:
        value = model.policy.predict_values(obs_tensor).item()
    else:
        value = model.critic(obs_tensor, model.actor(obs_tensor))[0].item()
    # True value (geometric series with a reward of one at each step)
    infinite_horizon_value = 1 / (1 - gamma)

    if handle_timeout_termination:
        # true value +/- 1
        assert abs(infinite_horizon_value - value) < 1.0
    else:
        # wrong estimation
        assert abs(infinite_horizon_value - value) > 1.0


text:
# Refer: https://en.wikipedia.org/wiki/Test_functions_for_optimization

import numpy as np
import matplotlib.pyplot as plt
import pso_solver

pso_solver.setSeed(1)

lower = xlower = ylower = zlower = -50
upper = xupper = yupper = zupper = 50
particleList = pso_solver.createRandomParticleList(3, numParticles=20, lower=lower, upper=upper)

# Testing on sphere function
f = lambda x, y, z: x**2 + y**2 + z**2

pso_solver.psoVisualizeTrivariate(particleList, f, xlower, xupper, ylower, yupper, zlower, zupper, c1=1, c2=1, W=0.5, numIters=20, maxFlag=False, sleepTime=0.5, accuracy=2, verbose=False)


text:
import os
import sys

import numpy as np
import tensorflow as tf
import tensors_saver

tensors_saver.set_out_path(sys.argv[1])

a = np.array([
    [0.1, 1.2, 4.3],
    [4.1, 0.2, 7.3],
    [0.06, 2.01, 0.23],
    [5.6, 2.3, 1.18]
])

b = np.array([
    [2.1, -5, 4.2],
    [4.1, 0.6, 8.2],
    [102.76, 342.91, -2.23],
    [-50.6, 56.23, 15.18]
])

a_node = tf.Variable(a, dtype=tf.float32)
b_node = tf.Variable(b, dtype=tf.float32)
y_node = a_node + b_node
sess = tf.Session()
init = tf.global_variables_initializer()
sess.run(init)

tf_y = sess.run(y_node)
tensors_saver.add(tf_y)


text:
#!/usr/bin/env python3
# -*-coding: utf-8 -*-
#
help = 'imgfuncのテスト用コード'
#

import cv2
import unittest
import numpy as np

import Tools.imgfunc as IMG

lenna_path = './Tools/Tests/Lenna.bmp'
mandrill_path = './Tools/Tests/Mandrill.bmp'

class TestImgFunc(unittest.TestCase):

    def test_getCh(self):
        self.assertEqual(IMG.getCh(-1), cv2.IMREAD_UNCHANGED)
        self.assertEqual(IMG.getCh(0), cv2.IMREAD_UNCHANGED)
        self.assertEqual(IMG.getCh(1), cv2.IMREAD_GRAYSCALE)
        self.assertEqual(IMG.getCh(2), cv2.IMREAD_UNCHANGED)
        self.assertEqual(IMG.getCh(3), cv2.IMREAD_COLOR)
        self.assertEqual(IMG.getCh(4), cv2.IMREAD_UNCHANGED)
        self.assertEqual(IMG.getCh(2.5), cv2.IMREAD_UNCHANGED)

    def test_blank(self):
        img = IMG.blank((128, 128, 3), 0)
        self.assertEqual(img.shape, (128, 128, 3))
        self.assertEqual(np.sum(img), 0)
        with self.assertRaises(SystemExit):
            img = IMG.blank((128, 128, -1), 0)

        img = IMG.blank((128, 128, 3), -1)
        self.assertEqual(img.shape, (128, 128, 3))
        self.assertEqual(np.sum(img), 0)
        img = IMG.blank((128, 128, 1), -1)
        self.assertEqual(img.shape, (128, 128, 1))
        self.assertEqual(np.sum(img), 0)
        img = IMG.blank((128, 128), -1)
        self.assertEqual(img.shape, (128, 128, 1))
        self.assertEqual(np.sum(img), 0)
        img = IMG.blank((128, 128, 3), (255, 255, 255))
        self.assertEqual(img.shape, (128, 128, 3))
        self.assertEqual(np.sum(img), 255 * 128 * 128 * 3)
        img = IMG.blank((128, 128), (255, 255, 255))
        self.assertEqual(img.shape, (128, 128, 3))
        self.assertEqual(np.sum(img), 255 * 128 * 128 * 3)

    def test_isImgPath(self):
        self.assertTrue(IMG.isImgPath(lenna_path))
        self.assertFalse(IMG.isImgPath('./Tools/Tests/Lenno.bmp'))
        self.assertFalse(IMG.isImgPath(None))
        self.assertFalse(IMG.isImgPath(0))

    def test_encodeDecode(self):
        l = cv2.imread(lenna_path)
        m = cv2.imread(mandrill_path)
        self.assertEqual(len(IMG.encodeDecodeN([l, m], 3)), 2)
        self.assertEqual(len(IMG.encodeDecodeN([l, m], 1)), 2)
        l = cv2.imread(lenna_path, IMG.getCh(1))
        m = cv2.imread(mandrill_path, IMG.getCh(1))
        self.assertEqual(len(IMG.encodeDecodeN([l, m], 3)), 2)
        self.assertEqual(len(IMG.encodeDecodeN([l, m], 1)), 2)

    def test_cut(self):
        l = cv2.imread(lenna_path)
        m = cv2.imread(mandrill_path)
        self.assertEqual(IMG.cutN([l, m], 64).shape, (2, 64, 64, 3))
        lm16 = [l, l, l, l, m, m, m, m, l, l, l, l, m, m, m, m]
        self.assertEqual(IMG.cutN(lm16, 64, 100).shape, (16, 64, 64, 3))
        self.assertEqual(IMG.cutN(lm16, 64, 10).shape, (10, 64, 64, 3))
        self.assertEqual(IMG.cutN([l, m], 1).shape, (2, 256, 256, 3))
        self.assertEqual(IMG.cutN([l, m], 0).shape, (2, 256, 256, 3))
        self.assertEqual(IMG.cutN([l, m], -1).shape, (2, 256, 256, 3))

    def test_split(self):
        l = cv2.imread(lenna_path)
        m = cv2.imread(mandrill_path)
        imgs, split = IMG.splitSQN([l, m], 32)
        self.assertEqual(imgs.shape, (128, 32, 32, 3))
        self.assertEqual(split, (8, 8))

        imgs, split = IMG.splitSQN([l, m], 0)
        self.assertEqual(imgs.shape, (512, 2, 2))
        self.assertEqual(split, (1, 1))

        imgs, split = IMG.splitSQN([l, m], 32, 10)
        self.assertEqual(imgs.shape, (120, 32, 32, 3))
        self.assertEqual(split, (8, 8))

        imgs, split = IMG.splitSQN([l, m], 32, 100)
        self.assertEqual(imgs.shape, (100, 32, 32, 3))
        self.assertEqual(split, (8, 8))

        imgs, split = IMG.splitSQN([l, m], 32, 1000)
        self.assertEqual(imgs.shape, (128, 32, 32, 3))
        self.assertEqual(split, (8, 8))

        print(l.shape, m.shape)
        imgs, split = IMG.splitSQN([l, m], 1024)
        self.assertEqual(imgs.shape, (2, 256, 256, 3))
        self.assertEqual(split, (1, 1))

        bk = IMG.blank((100, 120, 3), 255)
        imgs, split = IMG.splitSQN([bk], 1024)
        self.assertEqual(imgs.shape, (1, 100, 100, 3))
        self.assertEqual(split, (1, 1))

        l = cv2.imread(lenna_path, IMG.getCh(1))
        m = cv2.imread(mandrill_path, IMG.getCh(1))
        imgs, split = IMG.splitSQN([l, m], 32)
        self.assertEqual(imgs.shape, (128, 32, 32))
        self.assertEqual(split, (8, 8))

    def test_rotateRN(self):
        l = cv2.imread(lenna_path)
        m = cv2.imread(mandrill_path)
        imgs, angle = IMG.rotateRN([l, m], 3)
        self.assertEqual(imgs.shape, (6, 256, 256, 3))
        self.assertEqual(angle.shape, (6,))

    def test_flip(self):
        l = cv2.imread(lenna_path)
        m = cv2.imread(mandrill_path)
        self.assertEqual(IMG.flipN([l, m]).shape, (6, 256, 256, 3))
        self.assertEqual(IMG.flipN([l, m], -1).shape, (2, 256, 256, 3))
        self.assertEqual(IMG.flipN([l, m], 0).shape,  (2, 256, 256, 3))
        self.assertEqual(IMG.flipN([l, m], 1).shape,  (4, 256, 256, 3))
        self.assertEqual(IMG.flipN([l, m], 2).shape,  (6, 256, 256, 3))
        self.assertEqual(IMG.flipN([l, m], 3).shape,  (8, 256, 256, 3))

    def test_resize(self):
        l = cv2.imread(lenna_path)
        m = cv2.imread(mandrill_path)
        imgs = IMG.size2x([l, m])
        self.assertEqual(imgs[0].shape, (512, 512, 3))
        self.assertEqual(imgs[1].shape, (512, 512, 3))
        self.assertEqual(IMG.resize(l, -1).shape, (256, 256, 3))
        self.assertEqual(IMG.resize(l, 2).shape, (512, 512, 3))
        self.assertEqual(IMG.resize(l, 1.5).shape, (384, 384, 3))
        self.assertEqual(IMG.resize(l, 0.5).shape, (128, 128, 3))

    def test_imgs2arr(self):
        l = cv2.imread(lenna_path)
        m = cv2.imread(mandrill_path)
        self.assertEqual(IMG.imgs2arr([l, m]).shape, (2, 3, 256, 256))

        l = cv2.imread(lenna_path, IMG.getCh(1))
        m = cv2.imread(mandrill_path, IMG.getCh(1))
        self.assertEqual(IMG.imgs2arr([l, m]).shape, (2, 1, 256, 256))

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


text:
import pytest
import gym
import ray
from ray import tune

from ray.rllib.examples.models.shared_weights_model import TF2SharedWeightsModel
from ray.rllib.models import ModelCatalog
from ray.rllib.policy.policy import PolicySpec

import gym_pcgrl
from gym_pcgrl.envs.multi_pcgrl_env import MAPcgrlEnv
from gym_pcgrl.parallel_multiagent_wrappers import MARL_CroppedImagePCGRLWrapper
from ray.rllib.agents.ppo import PPOTrainer
from ray.rllib.env.wrappers.pettingzoo_env import ParallelPettingZooEnv

@pytest.fixture
def env():
    env = gym.make('MAPcgrl-binary-narrow-v0', num_agents=None, binary_actions=True)
    return env

"""
test to ensure that a new position is returned every time by the environment
"""
def test_default_agent_initialization():
    env = gym.make('MAPcgrl-binary-narrow-v0', num_agents=None, binary_actions=True)
    init_obs = env.reset()
    assert isinstance(init_obs, dict)
    assert len(init_obs) == 2

def test_custom_agent_initialization():
    num_agents = 3
    env = gym.make('MAPcgrl-binary-narrow-v0', num_agents=num_agents, binary_actions=False)
    init_obs = env.reset()
    assert isinstance(init_obs, dict)
    assert len(init_obs) == num_agents

def test_apply_action_binary_actions():
    env = gym.make('MAPcgrl-binary-narrow-v0', num_agents=None, binary_actions=True)
    init_obs = env.reset()
    empty_pos = init_obs['empty']['pos']
    solid_pos = init_obs['solid']['pos']
    # check that the map state observed by both agents is the same
    assert (init_obs['empty']['map'] == init_obs['solid']['map']).all()
    # set actions to add tile
    actions = {'empty': 1, 'solid': 1}
    obs, rews, dones, infos = env.step(actions)
    # if the initial positions of the agents are not the same, then check to find the expected value
    assert (obs['empty']['map'] == obs['solid']['map']).all()
    if not (empty_pos == solid_pos).all():
        empty_x, empty_y = empty_pos
        assert obs['empty']['map'][empty_y][empty_x] == 0
        solid_x, solid_y = solid_pos
        assert obs['solid']['map'][solid_y][solid_x] == 1

def test_apply_action():
    env = gym.make('MAPcgrl-binary-narrow-v0', num_agents=2, binary_actions=False)
    init_obs = env.reset()
    empty_pos = init_obs[0]['pos']
    solid_pos = init_obs[1]['pos']
    # check that the map state observed by both agents is the same
    assert (init_obs[0]['map'] == init_obs[1]['map']).all()
    # set actions to add tile
    actions = {0: 1, 1: 1}
    obs, rews, dones, infos = env.step(actions)
    # If the initial positions of the agents are not the same, then check to
    # find the expected value.
    # It is possible for the stored positions to not have the expected values
    # if the positions are the same.
    assert (obs[0]['map'] == obs[1]['map']).all()
    if not (empty_pos == solid_pos).all():
        empty_x, empty_y = empty_pos
        assert obs[0]['map'][empty_y][empty_x] == 0
        solid_x, solid_y = solid_pos
        assert obs[1]['map'][solid_y][solid_x] == 0

def test_set_map(env):
    from gym_pcgrl.envs.reps.representation import Representation
    from gym_pcgrl.envs.probs import PROBLEMS
    from gym_pcgrl.envs.helper import get_int_prob, get_string_map
    import numpy as np

    prob_name = 'binary'
    rep = Representation()
    prob = PROBLEMS[prob_name]()
    tile_probs = get_int_prob(prob._prob, prob.get_tile_types())
    rep.reset(prob._width, prob._height, tile_probs)
    level = rep._map

    init_obs = env.reset(initial_level=level)
    np.testing.assert_equal(level, env.get_map())

def test_human_actions(env):
    init_obs = env.reset()
    tile = 'empty'
    action = 1
    #action = env.get_human_action(tile, action)
    action = env.get_human_action(tile, action)
    assert action == f'place {tile}'
    print(action)
#"""
#There are three possible conditions that can lead to a finished environment:
#    changes >= max_changes
#    iterations => max_iterations
#    level reaches satisfiable quality
#"""
#def test_done_conditions():
#    pass
#
#
#def test_reward():
#    env = gym.make('MAPcgrl-binary-narrow-v0', num_agents=None, binary_actions=True)
#    init_obs = env.reset()
#    empty_pos = init_obs['empty']['pos']
#    solid_pos = init_obs['solid']['pos']
#    # check that the map state observed by both agents is the same
#    assert (init_obs['empty']['map'] == init_obs['solid']['map']).all()
#
#    # check that the agents rewards
#
#def test_update_heatmap():
#    pass
#
#
#def test_agent_action():
#    pass

max_stars_repo_name: rohin-dasari/gym-pcgrl

max_issues_repo_name: rohin-dasari/gym-pcgrl

max_forks_repo_name: rohin-dasari/gym-pcgrl


text:
# vim: fdm=marker
'''
author:     Fabio Zanini
date:       13/01/14
content:    Check the quality score along reads for read1 and read2.
            This rapid check is useful at the beginning, before even demultiplexing.
'''
# Modules
import os
import sys
import argparse
import gzip
import numpy as np
from operator import itemgetter
from Bio import SeqIO
from itertools import izip
from Bio.SeqIO.QualityIO import FastqGeneralIterator as FGI
import pysam

from hivwholeseq.sequencing.filenames import get_read_filenames
from hivwholeseq.cluster.fork_cluster import fork_quality_along_read as fork_self
from hivwholeseq.utils.mapping import extract_mapped_reads_subsample_open
from hivwholeseq.sequencing.samples import load_sequencing_run

# Functions
def quality_score_along_reads(read_len, reads_filenames,
                              skipreads=0,
                              randomreads=False,
                              maxreads=-1, VERBOSE=0):
    '''Calculate the quality score along the reads'''

    quality = [[[] for j in xrange(read_len)] for i in xrange(2)]

    # Precompute conversion table
    SANGER_SCORE_OFFSET = ord("!")
    q_mapping = dict()
    for letter in range(0, 255):
        q_mapping[chr(letter)] = letter - SANGER_SCORE_OFFSET

    if reads_filenames[0][-3:] == '.gz':
        openf = gzip.open
        file_readmode = 'rb'
    else:
        openf = open
        file_readmode = 'r'

    # Iterate over all reads (using fast iterators)
    with openf(reads_filenames[0], file_readmode) as fh1, \
         openf(reads_filenames[1], file_readmode) as fh2:

        if randomreads:
            if VERBOSE:
                print 'Getting number of reads',
                sys.stdout.flush()
            n_reads = sum(1 for read in FGI(fh1))
            fh1.rewind()
            if VERBOSE:
                print n_reads

            inds = np.arange(skipreads, n_reads)
            np.random.shuffle(inds)
            inds = inds[:maxreads]
            inds.sort()
            indi = 0
            if VERBOSE:
                print 'Random indices from ', inds[0], 'to', inds[-1]

            for (i, reads) in enumerate(izip(FGI(fh1), FGI(fh2))):
                if VERBOSE and (not ((i + 1) % 10000)):
                    print i + 1

                if (i != inds[indi]):
                    continue

                for ip, read in enumerate(reads):
                    for j, qletter in enumerate(read[2]):
                        quality[ip][j].append(q_mapping[qletter])

                indi += 1
                if indi == maxreads:
                    if VERBOSE:
                        print 'Maximal number of read pairs reached:', maxreads
                    break

        else:

            for (i, reads) in enumerate(izip(FGI(fh1), FGI(fh2))):
                if i < skipreads:
                    continue

                if i == skipreads + maxreads:
                    if VERBOSE:
                        print 'Maximal number of read pairs reached:', maxreads
                    break

                if VERBOSE and (not ((i + 1) % 10000)):
                    print i + 1

                for ip, read in enumerate(reads):
                    for j, qletter in enumerate(read[2]):
                        quality[ip][j].append(q_mapping[qletter])

    for qual in quality:
        for qpos in qual:
            qpos.sort()

    return quality

def plot_quality_along_reads(data_folder, adaID, title, quality, VERBOSE=0, savefig=False):
    '''Plot the results of the quality scores along reads'''

    import matplotlib.pyplot as plt
    from matplotlib import cm
    fig, axs = plt.subplots(1, 2, figsize=(16, 9))
    for i, (ax, qual) in enumerate(izip(axs, quality)):
        for j, qpos in enumerate(qual):
            x = qpos
            y = np.linspace(0, 1, len(x))[::-1]
            ax.plot(x, y, color=cm.jet(int(255.0 * j / len(qual))),
                    alpha=0.5,
                    lw=2)
        ax.set_xlabel('Phred quality', fontsize=14)
        ax.set_ylabel('Fraction of bases above quality x', fontsize=14)
        ax.set_title('Read'+str(i+1), fontsize=16)
        ax.text(2, 0.03, 'blue to red: 0 to '+str(len(qual))+' base', fontsize=18)

    fig.suptitle(title, fontsize=20)

    if savefig:
        from hivwholeseq.utils.generic import mkdirs
        from hivwholeseq.sequencing.filenames import get_figure_folder, \
                get_quality_along_reads_filename
        fig_folder = get_figure_folder(data_folder, adaID)
        fig_filename = get_quality_along_reads_filename(data_folder, adaID)
        mkdirs(fig_folder)
        fig.savefig(fig_filename)

    else:
        plt.tight_layout()
        plt.ion()
        plt.show()

def plot_cuts_quality_along_reads(data_folder, adaID, quality, title='',
                                  VERBOSE=0, savefig=False):
    '''Plot some cuts of the quality along the read'''
    from scipy.stats import percentileofscore as pof
    import matplotlib.pyplot as plt
    from matplotlib import cm
    fig, axs = plt.subplots(1, 2, figsize=(14, 8))
    qthreshs = [10, 20, 30, 35]
    for i, (ax, qual) in enumerate(izip(axs, quality)):
        for j, qthresh in enumerate(qthreshs):
            x = np.arange(len(qual))
            y = np.array([100 - pof(qual[k], qthresh) for k in xrange(len(qual))])
            ax.plot(x, y, color=cm.jet(int(255.0 * j / len(qthreshs))),
                    alpha=0.8,
                    lw=2,
                    label='Q = '+str(qthresh))
        ax.set_xlabel('Position [bp]', fontsize=14)
        ax.set_ylabel('Percentage of bases above quality x', fontsize=14)
        ax.set_title('Read'+str(i+1), fontsize=16)
        ax.set_ylim(-1, 101)
        ax.set_xlim(-1, len(qual) + 1)
        ax.legend(loc='best')

    if title:
        fig.suptitle(title, fontsize=20)

    if savefig:
        from hivwholeseq.utils.generic import mkdirs
        if savefig == True:
            from hivwholeseq.sequencing.filenames import get_figure_folder, \
                    get_quality_along_reads_filename
            fig_folder = get_figure_folder(data_folder, adaID)
            fig_filename = get_quality_along_reads_filename(data_folder, adaID, simple=True)
        elif isinstance(savefig, basestring):
            import os
            fig_filename = savefig
            fig_folder = os.path.dirname(fig_filename)

        else:
            raise ValueError('savefig must be a bool or a figure filename (string)')

        mkdirs(fig_folder)
        fig.savefig(fig_filename)

    else:
        plt.tight_layout()
        plt.ion()
        plt.show()

# Script
if __name__ == '__main__':

    # Parse input args
    parser = argparse.ArgumentParser(description='Check quality along reads',
                                     formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument('--run', required=True,
                        help='Seq run to analyze (e.g. Tue28, test_tiny)')
    parser.add_argument('--adaID', required=True,
                        help='Adapter ID to analyze (e.g. TS2)')
    parser.add_argument('--verbose', type=int, default=0,
                        help='Verbosity level [0-3]')
    parser.add_argument('--maxreads', type=int, default=-1,
                        help='Maximal number of reads to analyze')
    parser.add_argument('--submit', action='store_true', default=False,
                        help='Fork the job to the cluster via qsub')
    parser.add_argument('--no-savefig', action='store_false', dest='savefig',
                        help='Show figure instead of saving it')

    args = parser.parse_args()
    seq_run = args.run
    VERBOSE = args.verbose
    submit = args.submit
    maxreads = args.maxreads
    adaID = args.adaID
    savefig = args.savefig

    if submit:
        fork_self(seq_run, VERBOSE=VERBOSE, maxreads=maxreads, savefig=savefig)
        sys.exit()

    dataset = load_sequencing_run(seq_run)
    data_folder = dataset.folder
    read_len = dataset.cycles // 2

    reads_filenames = get_read_filenames(data_folder, adaID, gzip=True)
    if not os.path.isfile(reads_filenames[0]):
        reads_filenames = get_read_filenames(data_folder, adaID, gzip=False)
    title = seq_run+', '+adaID

    quality = quality_score_along_reads(read_len, reads_filenames,
                                        randomreads=(maxreads >= 1),
                                        maxreads=maxreads, VERBOSE=VERBOSE)

    plot_cuts_quality_along_reads(data_folder, adaID,
                                  quality,
                                  title=title,
                                  VERBOSE=VERBOSE,
                                  savefig=savefig)

    #if plotfull:
    #    plot_quality_along_reads(data_folder, adaID, title,
    #                             quality, VERBOSE=VERBOSE,
    #                             savefig=savefig)

max_stars_repo_name: iosonofabio/hivwholeseq

max_issues_repo_name: iosonofabio/hivwholeseq

max_forks_repo_name: iosonofabio/hivwholeseq

reason: import numpy,from scipy