Source code for NetworkSim.simulation.tools.load_save

__all__ = [
    "load_model",
    "save_model",
    "clear_env"
]
__author__ = ["Hongyi Yang"]

import pickle
import bz2
import os
from tqdm.auto import tqdm

from NetworkSim.simulation.simulator.base import BaseSimulator
from NetworkSim.simulation.simulator.parallel import ParallelSimulator


[docs]def load_model(fname=None, dir=None): """ Function to load previously saved pickle files as models. Parameters ---------- fname : str, optional File name of the model, by default ``None``, \ such that all models under the directory are loaded. dir : str, optional Directory of the models, by default ``None``, \ corresponding to the `models` directory. Returns ------- simulator : list A list of simulators loaded. """ dir_path = os.path.dirname(os.path.realpath(__file__)) if dir is None: load_dir = os.path.join(dir_path, '../../../', 'models') else: load_dir = dir os.chdir(load_dir) simulators = [] for filename in tqdm(os.listdir(load_dir), desc="Scanning files"): _load_file = False if fname is None or filename == fname: _load_file = True if _load_file: try: if os.path.getsize(filename) > 0: print("Load", filename) # with open(filename, 'rb') as input: # _simulator = pickle.load(input) input = bz2.BZ2File(filename, 'rb') _simulator = pickle.load(input) if isinstance(_simulator, BaseSimulator) or isinstance(_simulator, ParallelSimulator): simulators.append(_simulator) else: print(filename, "not recognised as a valid simulator object.") except pickle.PickleError: print(filename, "cannot be unpickled.") continue os.chdir(dir_path) return simulators
[docs]def clear_env(simulator, minimal): """ Function to clear all `simpy Environment` in the simulator. Parameters ---------- simulator : BaseSimulator The simulator whose `env` variable is to be cleared. minimal : bool, optional If the size is further reduced. Returns ------- simulator : BaseSimulator The simulator with all `env` variable cleared. """ simulator.env = None for ram_process in simulator.RAM: ram_process.env = None if simulator.bidirectional: for transmitter in simulator.upstream_transmitter: transmitter.env = None for transmitter in simulator.downstream_transmitter: transmitter.env = None for receiver in simulator.upstream_receiver: receiver.env = None for receiver in simulator.downstream_receiver: receiver.env = None if minimal: simulator.upstream_transmitter = None simulator.downstream_transmitter = None simulator.receiver = None simulator.model.data_rings = None simulator.model.control_ring = None simulator.model.reversed_data_rings = None for ram in simulator.RAM: ram.model = None ram.add_to_queue_record = None ram.pop_from_queue_record = None ram.generated_data_packet = None else: for transmitter in simulator.transmitter: transmitter.env = None for receiver in simulator.receiver: receiver.env = None if minimal: simulator.transmitter = None simulator.receiver = None simulator.model.data_rings = None simulator.model.control_ring = None for ram in simulator.RAM: ram.model = None ram.add_to_queue_record = None ram.pop_from_queue_record = None ram.generated_data_packet = None return simulator
[docs]def save_model(simulator, fname, dir=None, minimal=True): """ Function to save simulation models as pickle files. Note that all `env` variables present in the objects are cleared before saving. Parameters ---------- simulator : BaseSimulator or ParallelSimulator The simulator to be saved. fname : str The file name of the saved model. dir : str, optional The output directory of the saved model, by default ``None``, \ corresponding to the `models` directory. minimal : bool, optional If file size is further reduced by removing transmitter and receiver objects. """ if isinstance(simulator, BaseSimulator): simulator = clear_env(simulator, minimal=minimal) elif isinstance(simulator, ParallelSimulator): for base_simulator in simulator.simulator: base_simulator = clear_env(base_simulator, minimal=minimal) dir_path = os.path.dirname(os.path.realpath(__file__)) if dir is None: load_dir = os.path.join(dir_path, '../../../', 'models') else: load_dir = dir os.chdir(load_dir) # with open(fname, 'wb') as output: # pickle.dump(simulator, output) output = bz2.BZ2File(fname, 'w') pickle.dump(simulator, output) os.chdir(dir_path)