diff --git a/bgq/submit_job.py b/bgq/submit_job.py index 05bfa64..d5cc230 100644 --- a/bgq/submit_job.py +++ b/bgq/submit_job.py @@ -43,6 +43,22 @@ network_file = args.network assert network_file, "Must supply a network file to run." +# Handle gz files +if path.splitext(network_file)[1] == '.gz': + new_name = path.splitext(network_file)[0] + + # Unzip file (writes out to new file) + import gzip + with gzip.open(network_file, 'rb') as f_in: + with open(new_name, 'wb') as f_out: + f_out.writelines(f_in) + + # Remove gzipped file + os.remove(network_file) + + # Use new file + network_file = new_name + network_name = path.split(network_file)[-1] network_name = path.splitext(network_name)[0] print "Network name: ", network_name diff --git a/nengo_mpi/model.py b/nengo_mpi/model.py index 7ed5d19..6d04837 100644 --- a/nengo_mpi/model.py +++ b/nengo_mpi/model.py @@ -368,13 +368,18 @@ class MpiModel(builder.Model): def __init__( self, n_components, assignments, dt=0.001, label=None, - decoder_cache=NoDecoderCache(), save_file="", free_memory=True): + decoder_cache=NoDecoderCache(), save_file="", + compress_save_file=False, free_memory=True): self.n_components = n_components self.assignments = assignments if save_file: - self.save_file = open(save_file, 'w') + if compress_save_file: + import gzip + self.save_file = gzip.open(save_file + '.gz', 'wb') + else: + self.save_file = open(save_file, 'w') self.save_file.write( "%s%s%s" % (n_components, MpiModel.outfile_delim, dt)) else: @@ -561,9 +566,12 @@ def finalize_build(self): probe, self.sig[probe]['in'], sample_every=probe.sample_every) - if not self.save_file: + if self.save_file: + self.save_file.close() + else: self.mpi_sim.finalize_build() + def add_ops_to_mpi(self, component): """ Adds to MPI all ops that are meant for the given component. Which ops diff --git a/nengo_mpi/simulator.py b/nengo_mpi/simulator.py index 6ddc246..89b6bbc 100644 --- a/nengo_mpi/simulator.py +++ b/nengo_mpi/simulator.py @@ -16,7 +16,8 @@ class Simulator(object): def __init__( self, network, dt=0.001, seed=None, model=None, - partitioner=None, assignments=None, save_file=""): + partitioner=None, assignments=None, save_file="", + compress_save_file=False): """ Creates a Simulator for a nengo network than can be executed in parallel using MPI. @@ -87,7 +88,8 @@ def __init__( self.n_components, self.assignments, dt=dt, label="%s, dt=%f" % (network, dt), decoder_cache=get_default_decoder_cache(), - save_file=save_file) + save_file=save_file, + compress_save_file=compress_save_file) MpiBuilder.build(self.model, network)