experiment_log_tracking_arg.py

import os
import pickle
from datetime import datetime
from typing import Dict, List, OrderedDict, Callable, Any, Optional

import gurobipy as gp
import numpy as np
import ray
import torch
import tqdm
from gurobipy import GRB
from multiprocess import Process, Manager

import local_setting
import utils
from algorithm_helper import update_policy_weights, save_data, generate_rff, least_squares, print_debug
from cobs import Model
from icnn import ICNN


def warm_start(
    available_zones: List[str],
    fg_model: ICNN,
    pretrained_models: Dict[str, List[OrderedDict[str, torch.Tensor]]],
    rff_functions: List[Callable],
    log_path: str,
    save_path: str,
    save_result: bool = True
) -> Dict[str, Any]:
    """
    Warm start the OCFT algorithm by generating ICNN inference results from raw log data.

    Parameters
    ----------
    available_zones : List[str]
        List of available zones in the building
    fg_model : ICNN
        The ICNN model to be used for inference
    pretrained_models : Dict[str, List[OrderedDict[str, torch.Tensor]]]
        Pretrained ICNN model weights for F and G functions
    rff_functions : List[Callable]
        Random Fourier Features functions for G function
    log_path : str
        Path to the raw log data
    save_path : str
        Path to save the warm start result
    save_result : bool
        Whether to save the warm start result or not

    Returns
    -------
    Dict[str, Dict[str, List]]
        Dictionary containing the ICNN outputs and objectives for each zone and each function (F and G)
        Also return the last state of the log data
    """
    print(f"[{datetime.now()}] ---------- Warm Start: Begin ----------")

    # Check if warm start is needed
    if save_result and os.path.isfile(save_path):
        print(f"[{datetime.now()}] Warm start checkpoint found in {save_path} ...")
        with open(save_path, "rb") as pklfile:
            data = pickle.load(pklfile)
        print(f"[{datetime.now()}] ---------- Warm Start: Finish ----------")
        return data

    with open(log_path, "rb") as pklfile:
        data = pickle.load(pklfile)
    # Convert log data back to list of dictionaries
    data = [dict(zip(data, row)) for row in zip(*data.values())]

    icnn_outputs = {zone: {'f': list(), 'g': list()} for zone in available_zones}
    objectives = {zone: {'f': list(), 'g': list()} for zone in available_zones}

    # Do inference with each (prev_state, state) pair
    for i in tqdm.tqdm(range(1, len(data))):
        prev_state = data[i - 1]
        state = data[i]

        for zone in available_zones:
            current_icnn_outputs, current_objectives = utils.icnn_inference(
                zone=zone,
                state=state,
                prev_state=prev_state,
                fg_model=fg_model,
                pretrained_models=pretrained_models,
                rff_functions=rff_functions
            )

            for key in "fg":
                icnn_outputs[zone][key].append(current_icnn_outputs[key])
                objectives[zone][key].append(current_objectives[key])

    prev_state = data[-1]
    data = {
        "objectives": objectives,
        "icnn_outputs": icnn_outputs,
        "last_state": prev_state
    }

    # Store warm start result to skip next time
    if save_result:
        print(f"[{datetime.now()}] Saving CBC constraints from raw data ...")
        save_data(
            save_dir=os.path.dirname(save_path),
            name=os.path.basename(save_path.replace(".pkl", '')),
            data=data
        )

    print(f"[{datetime.now()}] ---------- Warm Start: Finish ----------")
    return data


@ray.remote(num_cpus=2)
def main(savedir, file_name, icnn_idxs):
    import socket
    host_name = str(socket.gethostname())
    # Seed everything
    np.random.seed(3)
    torch.manual_seed(3)
    env = gp.Env(params=local_setting.options[host_name])

    Model.set_energyplus_folder(local_setting.energyplus_location)

    valid_daytypes = tuple(range(11))
    fg_model = ICNN(dim=model_design[0], dimout=1, softplus=False).to(device)

    available_zones = ["Core_bottom", "Core_mid", "Core_top",
                       "Perimeter_bot_ZN_1", "Perimeter_bot_ZN_2", "Perimeter_bot_ZN_3", "Perimeter_bot_ZN_4",
                       "Perimeter_mid_ZN_1", "Perimeter_mid_ZN_2", "Perimeter_mid_ZN_3", "Perimeter_mid_ZN_4",
                       "Perimeter_top_ZN_1", "Perimeter_top_ZN_2", "Perimeter_top_ZN_3", "Perimeter_top_ZN_4"]

    # load RL policy for action prediction
    rl_controller = utils.RLController(
        available_zones,
        log_period=args.log_days,
        climate=args.experiment_climate,
        season=args.experiment_season,
        selection_pickle=f"{args.dir_prefix}/log_data/rl_selections.pkl",
    )

    model = utils.set_model(available_zones, args)

    # ======================== Generate Log Data ========================
    if args.get_log:
        pickle_name = (f"{args.dir_prefix}/log_data/"
                       f"{args.month:02d}{args.day:02d}-{args.year}-{args.experiment_climate}.pkl")

        state = model.reset()
        all_state = [utils.extract_state(available_zones, state)]

        while not model.is_terminate():
            state = model.step()
            all_state.append(utils.extract_state(available_zones, state))

        with open(pickle_name, "wb") as pklfile:
            # Convert to dictionary of list to save storage
            all_state = {key: [state[key] for state in all_state] for key in all_state[0]}
            pickle.dump(all_state, pklfile)
        print(f"[{datetime.now()}] {os.path.basename(pickle_name)} generated.")
        return

    # ======================== Load ICNNs ========================
    pretrained_models = {
        'f': list(),
        "f_scale": list(),
        'g': list(),
        "g_scale": list(),
    }

    # Load all f and gs
    print(f"[{datetime.now()}] Start loading ICNNs (F and Gs) ...")
    for key in "fg":
        folder = "models/ICNN_energy_relu" if key == 'g' else "models/ICNN_temperature_relu"
        folder_names = sorted(os.listdir(folder))
        for i in icnn_idxs:
            pretrained_models[key].append(torch.load(f"{folder}/{folder_names[i]}"))

            # Get the scale of the ICNN output
            min_val, max_val = folder_names[i].split('-')[:2]
            pretrained_models[f"{key}_scale"].append((int(min_val), int(max_val)))
    print(f"[{datetime.now()}] F ang Gs loaded.")

    # ======================== OCFT Parameters ========================
    # Time step variables
    performed_action = list()
    zone_temperature = list()
    energy_costs = list()
    rl_action_sequence = dict()
    occupancy = list()

    # STEP 2 variables
    weights = dict()
    disturbances = dict()
    diameters = dict()

    # STEP 3 variables
    infeasibility = dict()
    delta = dict()

    rff_params, rff_functions = generate_rff(num_rff, 12)

    print(f"[{datetime.now()}] Initializing CBC ...")
    for i, zone in enumerate(available_zones):
        disturbances[zone] = {key: [disturbance_settings[key]] for key in "fg"}
        weights[zone] = {key: [np.zeros(len(icnn_idxs))] for key in "fg"}
        diameters[zone] = {key: [100] for key in "fg"}

        infeasibility[zone] = list()
        delta[zone] = list()
        rl_action_sequence[zone] = list()

    # ======================== OCFT Warm Start ========================
    #  Load raw log data to generate and store ICNN inference results
    result = warm_start(
        available_zones=available_zones,
        fg_model=fg_model,
        pretrained_models=pretrained_models,
        rff_functions=rff_functions,
        log_path=f"{args.dir_prefix}/log_data/{args.logfile_name}.pkl",
        save_path=f"{args.dir_prefix}/log_data/{args.logfile_name}-warm_start.pkl",
        save_result=len(icnn_idxs) == 60,
    )

    icnn_outputs: Dict[str, Dict[str, List[np.ndarray]]] = result["icnn_outputs"]
    objectives: Dict[str, Dict[str, List[float]]] = result["objectives"]
    prev_state: Dict[str, Any] = result["last_state"]

    # ======================== OCFT Algorithm ========================
    state = model.reset()
    state = utils.extract_state(available_zones, state)
    energy_costs.append(state["total hvac electricity"])
    step = 0
    pbar = tqdm.tqdm(desc="Experiment progress (STEP)", total=args.experiment_duration * 24 * args.experiment_interval)

    # ----- Multiprocessing -----
    threads = dict()
    manager = Manager()
    result_dict = manager.dict()

    print(f"[{datetime.now()}] Algorithm initialized. Starting experiment ...")
    while not model.is_terminate():
        pbar.update(1)
        # Dry run only test the code without any action taken, just to see if the code runs
        if args.dry_run:
            model.step()
            save_data(savedir, file_name, {"Random": icnn_idxs})
            continue

        rl_actions = rl_controller.get_actions(state)
        cobs_actions = list()

        assert int(state["day type"]) in valid_daytypes, f"Invalid day type {state['day type']} on step {step}"

        # Define occupancy state based on time
        occupancy_status = "occupied" if 7 <= state["hour"] <= 21 else "unoccupied"
        occupancy.append(int(occupancy_status == "occupied"))
        print_debug(
            f"[{datetime.now()}] ------------------- {occupancy_status.upper()} -------------------",
            verbose=args.verbose
        )
        comfort_lower, comfort_upper = comfort_constraints[occupancy_status]

        # Run OCFT for each zone
        def ocft_per_zone(
            zone_idx: int,
            zone_name: str,
        ) -> Optional[Dict[str, Any]]:
            """
            Run OCFT algorithm for each zone

            Parameters
            ----------
            zone_idx : int
                Index of the zone
            zone_name : str
                Name of the zone

            Returns
            -------
            Optional[Dict[str, Any]]
                Dictionary containing the OCFT results
            """
            zone_icnn_outputs = icnn_outputs[zone_name]
            zone_objectives = objectives[zone_name]
            zone_weights = weights[zone_name]
            zone_disturbances = disturbances[zone_name]
            zone_diameters = diameters[zone_name]
            zone_rl_actions = rl_actions[zone_name]

            fg_model = ICNN(dim=model_design[0], dimout=1, softplus=False).to(device)
            zone_temp = state[f"{zone_name} thermostat temperature"]

            return_result = {
                "icnn_outputs": {'f': None, 'g': None},
                "objectives": {'f': None, 'g': None},
                "weights": {'f': None, 'g': None},
                "disturbances": {'f': None, 'g': None},
                "diameters": {'f': None, 'g': None},
                "infeasibility": None,
                "delta": None,
                "action": None
            }

            # Evaluate the comfort constraints
            if not comfort_lower - padding < zone_temp < comfort_upper + padding:
                print_debug(
                    f"[{datetime.now()}] ---- Zone {zone_idx} temperature ({zone_temp}) outside bound "
                    f"{comfort_lower - padding} ~ {comfort_upper + padding} -----",
                    verbose=1
                )

            # Adjust f and g parameters with OCFT. We find theta* to satisfy (s_t-1, a_t-1) -> (s_t, r_t)
            # --------------- STEP 1: Observe ICNN outputs and objectives ---------------
            current_icnn_outputs, current_objectives = utils.icnn_inference(
                zone=zone_name,
                state=state,
                prev_state=prev_state,
                fg_model=fg_model,
                pretrained_models=pretrained_models,
                rff_functions=rff_functions
            )

            # --------------- STEP 2: Update weights based on the observed trajectory ---------------
            for key in 'f':
                # Add the current ICNN outputs and objectives to the list
                return_result["icnn_outputs"][key] = current_icnn_outputs[key]
                return_result["objectives"][key] = current_objectives[key]

                # Convert values to numpy for compute
                icnn_output_matrix = np.vstack(zone_icnn_outputs[key] + [current_icnn_outputs[key]])
                weight_vector = zone_weights[key][-1]
                objective_vector = np.array(zone_objectives[key] + [current_objectives[key]])
                disturbance_lower, disturbance_upper = zone_disturbances[key][-1]
                lower_bounds = objective_vector - disturbance_lower
                upper_bounds = objective_vector + disturbance_upper
                weighted_outputs = icnn_output_matrix.dot(weight_vector)
                status = "CONSISTENT"

                new_weights = zone_weights[key][-1]
                new_disturbance = zone_disturbances[key][-1]
                new_diameter = zone_diameters[key][-1]

                # Run baseline if needed
                if args.lse or args.cls:
                    new_weights = least_squares(
                        targets=objective_vector,
                        icnn_outputs=icnn_output_matrix,
                        lower_bounds=None if args.lse else lower_bounds,
                        upper_bounds=None if args.lse else upper_bounds,
                        regularization_coefficient=args.coef_ls,
                        gurobi_env=env
                    )

                # Check if the previous weight is still valid. If the previous weight is invalid, update the weight
                elif (weighted_outputs > upper_bounds + c_margin).any() or \
                        (weighted_outputs < lower_bounds - c_margin).any():
                    status = "INVALID"
                    new_weights, disturbance_lower_delta, disturbance_upper_delta, new_diameter = update_policy_weights(
                        icnn_outputs=icnn_output_matrix,
                        objectives=objective_vector,
                        lower_bounds=lower_bounds,
                        upper_bounds=upper_bounds,
                        disturbance_deltas=disturbance_settings[f"{key}_delta"],
                        box_constraints=box_constraints[key],
                        method=method[key],
                        use_simplex=False,
                        last_n_samples=last_n_samples,
                        select_n_samples=select_n_samples,
                        do_subsampling=False,
                        prev_weights=zone_weights[key][-1],
                        gurobi_env=env,
                        verbose=args.verbose
                    )

                    if disturbance_lower_delta != 0 and disturbance_upper_delta != 0:
                        new_disturbance = (
                            zone_disturbances[key][-1][0] + disturbance_lower_delta,
                            zone_disturbances[key][-1][1] + disturbance_upper_delta
                        )

                        print_debug(
                            f"[{datetime.now()}] ---------- Zone {zone_idx} Function {key} set empty! "
                            f"increasing to W = {new_disturbance} ----------",
                            verbose=args.verbose
                        )

                # Save new weights and disturbances
                return_result["weights"][key] = new_weights
                return_result["disturbances"][key] = new_disturbance
                return_result["diameters"][key] = new_diameter

                print_debug(
                    f"[{datetime.now()}] ---------- function = {key} , "
                    f"previous value = {np.linalg.norm(return_result['weights'][key])} , "
                    f"zone = {zone_idx}, {status}, "
                    f"with prediction error {max(abs(weighted_outputs - objective_vector))} ----------",
                    verbose=args.verbose
                )

            # --------------- STEP 3: Generate actions based on the latest adapted f and g ---------------
            # Define variables
            m = gp.Model(env=env)
            m.Params.LogToConsole = 0
            final_action = m.addMVar(1, lb=0.1, ub=1, vtype='C')

            # Prefill the input vector with s_t and var
            predictor_input = np.array([
                prev_state["outdoor temperature"],
                prev_state["outdoor wind"],
                prev_state["outdoor solar"],
                prev_state["hour"],
                prev_state[f"{zone_name} occupancy"],
                prev_state[f"{zone_name} thermostat heating setpoint"],
                prev_state[f"{zone_name} thermostat cooling setpoint"],
                prev_state[f"{zone_name} thermostat temperature"],
                prev_state[f"{zone_name} VAV inlet temp"],
                abs(prev_state[f"{zone_name} predict load"]),
                prev_state[f"{zone_name} total solar"]
            ]).astype("float32")
            z = m.addMVar((len(predictor_input) + 1,), lb=-1e9, ub=1e9, vtype='C', name='z')
            m.addConstr(z[:-1] == predictor_input, name="c0")
            m.addConstr(z[-1] == final_action, name="c1")

            # Formulate the output of f and g
            fg_output = dict()
            for key in 'f':
                fg_output[key] = 0
                for model_index, params in enumerate(pretrained_models[key]):
                    if current_icnn_outputs[key][model_index] == 0:
                        continue
                    for ii in range(1, len(model_design)):  # per layer
                        zn_raw = m.addMVar((model_design[ii],), lb=-1e7, ub=1e7, vtype='C', name=f"z_raw_{ii}")
                        if ii == 1:
                            m.addConstr(zn_raw == z @ params[f"Wzs.{ii - 1}.weight"].detach().numpy().T +
                                        params[f"Wzs.{ii - 1}.bias"].detach().numpy(), name=f"layer_{ii}")
                        else:
                            if f"Wxs.{ii - 2}.bias" in params.keys():
                                m.addConstr(zn_raw == zn @ params[f"Wzs.{ii - 1}.weight"].detach().numpy().T +
                                            z @ params[f"Wxs.{ii - 2}.weight"].detach().numpy().T +
                                            params[f"Wxs.{ii - 2}.bias"].detach().numpy(),
                                            name=f"middle layer_{ii}")
                            else:
                                m.addConstr(zn_raw == zn @ params[f"Wzs.{ii - 1}.weight"].detach().numpy().T +
                                            z @ params[f"Wxs.{ii - 2}.weight"].detach().numpy().T,
                                            name=f"final_layer_{ii}")

                        if ii != len(model_design) - 1:
                            zn = m.addMVar((model_design[ii],), lb=0, ub=1e7, vtype='C', name="zn")
                            for ij in range(model_design[ii]):
                                m.addGenConstrMax(zn[ij], [zn_raw[ij], 0], name="maximize")

                    scale_min, scale_max = pretrained_models[f"{key}_scale"][model_index]
                    current_weights = return_result["weights"][key][model_index]
                    fg_output[key] += (zn_raw * (scale_max - scale_min) + scale_min) * current_weights

            # add slack variable
            delta1 = m.addVar(vtype='C', lb=0, ub=10000, name="d1")
            delta2 = m.addVar(vtype='C', lb=0, ub=10000, name="d2")

            # Set temperature constraints, assume box constraints between 18 and 26
            m.addConstr(fg_output['f'] >= comfort_lower + padding + return_result["disturbances"]['f'][0] - delta1)
            m.addConstr(fg_output['f'] <= comfort_upper - padding - return_result["disturbances"]['f'][1] + delta2)

            # Use g in the objectives function
            m.setObjective(
                args.rl_lambda * (z[-1] - zone_rl_actions) * (z[-1] - zone_rl_actions) +
                eta_1 * delta1 * delta1 + eta_2 * delta2 * delta2,
                GRB.MINIMIZE
            )

            # Optimize and edge case handling
            # m.Params.FeasibilityTol = 1e-4
            m.optimize()
            if m.Status != GRB.OPTIMAL:
                out_action = zone_rl_actions
                print_debug(f"[{datetime.now()}] OPTIMIZATION NOT FEASIBLE, USING MEAN", verbose=args.verbose)
                new_infeasibility = 0.0
                new_delta = np.array([np.nan, np.nan])
            else:
                out_action = final_action.x
                new_infeasibility = 1.0
                new_delta = np.array([delta1.x, delta2.x])

            return_result["infeasibility"] = new_infeasibility
            return_result["delta"] = new_delta
            return_result["action"] = out_action

            if args.multiprocessing:
                result_dict[zone_name] = return_result
            else:
                return return_result

        # Run multiprocessing for each zone
        if args.multiprocessing:
            result_dict.clear()
            for i, zone in enumerate(available_zones):
                threads[zone] = Process(target=ocft_per_zone, args=(i, zone))
                threads[zone].start()

            for zone in available_zones:
                threads[zone].join()

        # Save experiment results
        for i, zone in enumerate(available_zones):
            # Get the result
            if args.multiprocessing:
                result = result_dict[zone]
            else:
                result = ocft_per_zone(i, zone)

            # Save the result
            for var in result:
                if var == "action":
                    continue
                if isinstance(result[var], dict):
                    for key in result[var]:
                        eval(var)[zone][key].append(result[var][key])
                else:
                    eval(var)[zone].append(result[var])

            # --------------- STEP 4: Perform action ---------------
            cobs_actions.append(
                {
                    "priority": 0,
                    "component_type": "Schedule:Constant",
                    "control_type": "Schedule Value",
                    "actuator_key": f"{zone} VAV Customized Schedule",
                    "value": result["action"],
                    "start_time": state["timestep"]
                }
            )
            rl_action_sequence[zone].append(rl_actions[zone])

        prev_state = state
        state = model.step(cobs_actions)
        state = utils.extract_state(available_zones, state)

        # BOOK KEEPING
        energy_costs.append(state["total hvac electricity"])
        zone_temperature.append([state[f"{zone} thermostat temperature"] for zone in available_zones])
        performed_action.append([action["value"] for action in cobs_actions])
        step += 1

        print_debug(
            f"[{datetime.now()}] ---------- Current energy: {energy_costs[-1]} ----------"
            f"[{datetime.now()}] action is: {performed_action[-1]}",
            verbose=args.verbose
        )

        if (step + 1) % (args.experiment_interval * 24) == 0:
            print(f"[{datetime.now()}] ---------- SAVING ----------")
            specs = {
                "comfort_constraints": comfort_constraints,
                "box_constraints": box_constraints,
                "disturbance_settings": disturbance_settings,
                "latest_sample": last_n_samples,
                "past_sample": select_n_samples,
                "c_margin": c_margin,
                "padding": padding,
                "rl_lambda": args.rl_lambda,
                "num_rff": num_rff,
                "method": method,
                "eta": [eta_1, eta_2],
                "run_duration": args.experiment_duration,
                "run_year": args.year,
                "run_month": args.month,
                "run_date": args.day,
                "run_interval": args.experiment_interval,
            }
            data = {
                "specs": specs,
                "performed_action": performed_action,
                "zone_temperature": zone_temperature,
                "energy_costs": energy_costs,
                "rl_action_sequences": rl_action_sequence,
                "occupancy": occupancy,
                "icnn_outputs": icnn_outputs,
                "objectives": objectives,
                "weights": weights,
                "disturbances": disturbances,
                "diameters": diameters,
                "infeasibility": infeasibility,
                "delta": delta,
                "icnns": icnn_idxs
            }

            save_data(savedir, file_name, data)

    total_energy_cost = sum(energy_costs)
    print(f"[{datetime.now()}] Total energy cost is: {total_energy_cost:.2f}, in MWh: {total_energy_cost / 4e6:.2f}")
    pbar.close()


if __name__ == "__main__":
    # file saving & parser
    ray.init()
    parser = utils.set_cli()
    args = parser.parse_args()

    # Overwrite experiment period if season is given
    if args.experiment_season:
        args.experiment_duration = args.log_days if args.get_log else 21
        args.year = 2000
        args.day = 15 - args.log_days if args.get_log else 15
        args.month = 7 if args.experiment_season == "summer" else 1
        print(f"[{datetime.now()}] Experiment period changed to "
              f"{args.year}-{args.month:02d}-{args.day:02d} for {args.experiment_duration} days ...")

        args.logfile_name = f"{args.month:02d}{15 - args.log_days:02d}-{args.year}-{args.experiment_climate}"

        print(f"[{datetime.now()}] Experiment log data changed to {args.logfile_name} ...")

    # Folder names for saving data and log data location
    experiment_type = "ocft"
    if args.lse:
        experiment_type = "lse"
    elif args.cls:
        experiment_type = "cls"

    savedir_global = (
        f"{args.dir_prefix}/cbc/"
        f"{experiment_type}-final-{args.experiment_building}"
        f"-{args.month:02d}{args.day:02d}"
        f"-{args.experiment_duration}days"
        f"-climate={args.experiment_climate}"
        f"-log={args.log_days}"
        f"-W_f={args.w}"
        f"-rl_lambda={args.rl_lambda}"
        f"-icnn_var={args.icnn_size}v{args.icnn_random}"
        f"{args.savedir_suffix}"
    )

    # Hyperparameters
    comfort_constraints = {
        "occupied": (
            23 if args.month in range(5, 11) else 20,
            26 if args.month in range(5, 11) else 23.5
        ),
        "unoccupied": (
            13,
            35
        ),
    }
    box_constraints = {
        'f': (-1e6, 1e6),
        'g': (-1e6, 1e6),
    }

    disturbance_settings = {
        'f': (args.w, args.w),
        "f_delta": (0.1, 0.1),
        'g': (8e3, 8e3),
        "g_delta": (1e3, 1e3)
    }

    # algorithm parameter
    last_n_samples = 1300
    select_n_samples = 150
    c_margin = 1e-1

    padding = 0.1
    num_rff = 0
    method_f = "projection"  # ["projection", "CBC", "one_step_LSE"]
    method_g = "projection"  # ["projection", "CBC", "one_step_LSE"]
    method = {'f': method_f, 'g': method_g}
    eta_1 = 2e4 if args.month in range(5, 11) else 1e3
    eta_2 = 1e3 if args.month in range(5, 11) else 1e4

    os.makedirs(savedir_global, exist_ok=True)
    file_name = "tracking"

    device = torch.device("cpu")
    model_design = (12, 100, 100, 1)

    print(f"[{datetime.now()}] Job started. savedir: {savedir_global}, file_name {file_name}")

    # ---------- RUN PROGRAM ----------
    futures = list()
    icnn_idxs = tuple(range(60))
    if args.icnn_random != 0:
        icnn_idxs = tuple(sorted(np.random.choice(np.arange(60), args.icnn_size, replace=False)))
    futures.append(main.remote(savedir_global, file_name, icnn_idxs))
    ray.get(futures)