Compress a density grid with lossy + lossless compression

compress.py

#!/usr/bin/env python3
'''
Compress a density grid.

Usage: ./compress.py --help
'''

from pathlib import Path

import click
import numpy as np
import asdf

@click.command()
@click.argument('fn')
@click.option('--truncbits', default=10, help='Number of low bits to truncate in the density field')
def compress(fn, truncbits):
    '''
    Read the density file in FN (a npz file) and write a compressed version.

    This uses the blosc ASDF compression, supplied by abacusutils on pip.
    '''
    fn = Path(fn)
    d = np.load(fn)['arr_0']

    d = d.astype(np.float32)
    d = (d.view(np.uint32) & ~np.uint32((1<<truncbits)-1)).view(np.float32)

    with asdf.AsdfFile(tree=dict(density=d)) as af:
        af.write_to(fn.parent / (fn.stem + '_compressed.asdf'), all_array_compression='blsc',
            compression_kwargs=dict(shuffle='shuffle')
            )

if __name__ == '__main__':
    compress()