Python Data Selection Function for 'dBZ', 'KDP', and 'ZDR' Directories

This function selects data from three directories: 'dBZ/3.0km/', 'KDP/3.0km/', and 'ZDR/3.0km/'. It then appends the selected data to the file 'select.txt'. The function iterates through each directory and loads the data into a list. It checks if the length of the data list is equal to 20, and if so, it checks if the sum of the data list equals 0. If the sum is less than 0.8 * 20 * 256 * 256, it writes the data to the file and prints the directory and frame. If the sum is greater than or equal to that value, it resets the index and the data list.

def select_data():
    dbz_dir_path = 'data/NJU_CPOL_update2308/dBZ/3.0km/'
    kdp_dir_path = 'data/NJU_CPOL_update2308/KDP/3.0km/'
    zdr_dir_path = 'data/NJU_CPOL_update2308/ZDR/3.0km/'
    data_select = 'data/select.txt'
    data_select_file = open(data_select, mode='a')
    dir_list = os.listdir(dbz_dir_path)
    for dir in dir_list:
        dbz_data_path = dbz_dir_path + dir
        dbz_list = os.listdir(dbz_data_path)
        i = 0
        data_list = []
        for i in range(len(dbz_list)):
            data_list.append(np.load(dbz_data_path + '/' + dbz_list[i]))
            if (len(dbz_list) - i) < (20 - len(data_list)):
                break
            if len(data_list) == 20:
                data_list = np.array(data_list)
                if (np.sum(data_list == 0)) < 0.8 * 20 * 256 * 256:
                    # print(np.sum(data_list==0))
                    data_select_file.write('
' + '3.0km' + ',' + dir + ',' + dbz_list[i + 1 - 20])
                    print('dir:', dir, 'frame:', dbz_list[i + 1 - 20])
                    data_list = []
                    continue
                else:
                    i = i - 19
                    data_list = []
                    continue