# -*- coding: utf-8 -*- """ Created on Wed Oct 12 15:34:02 2016 @author: rmcleod """ import blosc import numpy as np import time import mrcz import tempfile, os, os.path, glob print( "==== Generating simulated Poisson stacks ====" ) print( "Note: will write files to: %s. (they will not be deleted as they can be re-used!)" % tempfile.gettempdir() ) stackShape = [50,4096,4096] doseLevelList = [ 0.01, 0.1, 0.25, 0.5, 1.0, 1.5, 2.0, 4.0 ] doseFilenames = [] poissonDataArray = np.zeros( [len(doseLevelList), stackShape[0],stackShape[1],stackShape[2]] ) for I, dose in enumerate(doseLevelList): poissonName = os.path.join(tempfile.gettempdir(), "poisson%f.mrcz" % dose ) doseFilenames.append( poissonName ) if not os.path.isfile( poissonName ): print( "Generating Poisson random numbers and writing %s" % poissonName ) poissonDataArray[I,...] = np.random.poisson( lam=dose, size=stackShape ).astype('int8') mrcz.MRCExport( poissonDataArray[I,...], poissonName, compressor='zstd' ) else: # Load from disk print( "Loading from disk: %s" % poissonName ) poissonDataArray[I,...] = mrcz.MRCImport( poissonName ) print( "==== Starting compression benchmark ====" ) def doCompression( dataStack, compressor='zstd', blocksize=2**20, n_threads=16, shuffle=blosc.BITSHUFFLE, clevel=5 ): blosc.set_blocksize( blocksize ) blosc.set_nthreads( n_threads ) typeSize = dataStack.dtype.itemsize packedDataList = [None] * dataStack.shape[0] for J in np.arange(dataStack.shape[0]): packedDataList[J] = blosc.compress( dataStack[J,:,:], typesize=typeSize, clevel=clevel, shuffle=shuffle, cname=compressor ) return packedDataList def doDecompression( packedDataList, shape, n_threads ): blosc.set_nthreads( n_threads ) dataList = [None] * len(packedDataList) for J in np.arange(len(packedDataList) ): # dataStack[J,:,:] = np.reshape( # np.frombuffer( blosc.decompress( packedDataList[J] ), dtype='uint8' ), # shape[1:] ) # Something here Numpy-side is very slow, so let's not include that in our # benchmark. dataList[J] = blosc.decompress( packedDataList[J] ) return dataList #t_half0 = time.time() #halfimage = dm4image_8bit[:,:,::2] + np.left_shift(dm4image_8bit[:,:,1::2],4) #t_half1 = time.time() #restoreimage = np.empty( header['dimensions'], dtype='uint8' ) ##image[0::2] = np.left_shift(interlaced_image,4)/16 ##image[1::2] = np.right_shift(interlaced_image,4) ## Different interlace option ## TODO: AND array with 15 instead? #restoreimage[:,:,::2] = (np.left_shift( halfimage, 4 ) & 15 ) #restoreimage[:,:,1::2] = np.right_shift( halfimage, 4 ) #t_half2 = time.time() # #print( "4-byte encoding time (s): %f" % (t_half1 - t_half0) ) #print( "4-byte DEcoding time (s): %f" % (t_half2 - t_half1) ) originalBytes = 50*4096*4096*4 # Guess of size if we saved as float-32 # Something is wrong with lz4hc, it's really not competitive # All codecs slow down if you use bloscLZ for some reason. blocksize adjustment? # Because our data is bit-limited by shot noise generally we will always # want to use the BITSHUFFLE filter when working with a counting electron detector # Probably BITSHUFFLE is not so great for floating-point data, however. SHUFFLE = blosc.BITSHUFFLE nThreadsList = np.arange( 12, 48+1, 6 ) compressorList = np.array( ['lz4', 'zlib', 'zstd' ] ) clevelList = np.arange(1,7+1) testRuns = 1 blockSizeList = np.array( [2**15, 2**16, 2**17, 2**18, 2**19, 2**20, 2**21, 2**22], dtype='int64' ) optiShape = [len(doseLevelList), len(compressorList), len(clevelList), len(nThreadsList), len(blockSizeList), testRuns] cCompress = np.zeros( optiShape ) cBytes = np.zeros( optiShape ) cRatio = np.zeros( optiShape ) cDecompress = np.zeros( optiShape ) for I, doseLevel in enumerate( doseLevelList ): for J, compressor in enumerate(compressorList): for K, clevel in enumerate( clevelList ): for N, N_THREADS in enumerate( nThreadsList ): for B, blocksize in enumerate( blockSizeList ): print( "Testing compressor %s level %d with %d threads and blocksize %d at dose level %f counts/pix" % (compressor, clevel, N_THREADS, blocksize, doseLevel) ) for L in np.arange(testRuns): t0 = time.time() packedData = doCompression( poissonDataArray[I, ...], compressor=compressor, n_threads=N_THREADS, shuffle=SHUFFLE, clevel=clevel, blocksize=blocksize ) t1 = time.time() cCompress[I,J,K,N,B,L] = t1 - t0 for index in np.arange( len(packedData) ): cBytes[I,J,K,N,B,L] += len( packedData[index] ) cRatio[I,J,K,N,B,L] = 100.0 * originalBytes / cBytes[I,J,K,N,B,L] t2 = time.time() # Will this work on all frames? No.... result = doDecompression( packedData, len(packedData), N_THREADS ) t3 = time.time() cDecompress[I,J,K,N,B,L] = t3 - t2 np.save( "cCompress.npy", cCompress ) np.save( "cBytes.npy", cBytes ) np.save( "cRatio.npy", cRatio ) np.save( "cDecompress.npy", cDecompress ) #print( "4-byte encoding time (s): %f" % (t_half1 - t_half0) ) #print( "4-byte DEncoding time (s): %f" % (t_half2 - t_half1) ) print( "==== Compressor ====" ) print( compressorList ) print( "==== Compression Time (s) ====" ) print( cCompress ) print( "==== DEcompression Time (s) ====" ) print( cDecompress ) print( "==== Compressed Size (MB) ====" ) print( cBytes / 2**20 ) print( "==== Compression Ratio (%) ====" ) print( cRatio ) print( "TODO: test decompression times" ) print( "TODO: test read/write times" ) #print( "Compression time: %f s" %(t1-t0) ) #print( "Original DM4 size: %d MB" % (originalBytes /(2**20)) ) #uncompressedBytes = np.prod( dm4image.shape ) #print( "UNcompressed size: %d MB" % (uncompressedBytes/(2**20)) ) #print( "Compressed size: %d MB" % (header['compressedBytes']/(2**20)) ) #print( "Effective compression ratio: %d %% " % (100.0 * originalBytes/header['compressedBytes']) )