Skip to content

Instantly share code, notes, and snippets.

@etale-cohomology
Created March 11, 2016 00:53
Show Gist options
  • Star 0 You must be signed in to star a gist
  • Fork 0 You must be signed in to fork a gist
  • Save etale-cohomology/16b67480ce510e2b63ab to your computer and use it in GitHub Desktop.
Save etale-cohomology/16b67480ce510e2b63ab to your computer and use it in GitHub Desktop.
Save a NumPy array to disk with optional compression using 7z
def save_ndarray(ndarray, filename, compress=False):
"""Store a single NumPy array on an .npy binary file. Tries to keep allow_pickle at False.
Optional multithreaded compression using 7-zip (needs to be in the PATH).
Args:
ndarray (ndarray): Array to be stored on disk.
filename (str): The file's name on disk. No need for .npy extension
compress (bool, optional): Set to True to save array as .npy and then compress it to .7z
Returns:
None
"""
if filename.endswith('.npy'):
file_sans = filename[:-4]
else:
file_sans = filename # String assignments are always copies!
filename += '.npy'
print('Saving {0!s}-array ({1:,} bytes) to disk...' .format(ndarray.shape, ndarray.nbytes))
if compress:
status_compression = 'compressed'
if shutil.which('7z') is None:
raise FileNotFoundError('7z not found on the PATH!')
try:
np.save(filename, ndarray, allow_pickle=False)
status_pickle = 'unpickled'
except ValueError:
np.save(filename, ndarray, allow_pickle=True)
status_pickle = 'pickled'
subprocess.Popen('7z a ' + file_sans + '.7z ' + file_sans + '.npy -mmt', shell=True).wait()
os.remove(filename)
filename = file_sans + '.7z '
print()
else:
status_compression = 'uncompressed'
try:
np.save(filename, ndarray, allow_pickle=False)
status_pickle = 'unpickled'
except ValueError:
np.save(filename, ndarray, allow_pickle=True)
status_pickle = 'pickled'
print('Successfully saved {0!s}-array ({1:,} bytes) as [{2!s}] in [{3}], [{4}] form!'
.format(ndarray.shape, ndarray.nbytes, filename, status_compression, status_pickle))
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment