Added option to compress network file

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

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

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)