diff --git a/abacusnbody/data/asdf.py b/abacusnbody/data/asdf.py index 6493da3..70fbd89 100644 --- a/abacusnbody/data/asdf.py +++ b/abacusnbody/data/asdf.py @@ -17,11 +17,6 @@ from asdf.extension import Compressor, Extension -def _monkey_patch(*args,**kwargs): - raise Exception("Please use abacusnbody.data.asdf.set_nthreads(nthreads)") - -asdf.compression.set_decompression_options = _monkey_patch - def set_nthreads(nthreads): blosc.set_nthreads(nthreads) diff --git a/abacusnbody/data/compaso_halo_catalog.py b/abacusnbody/data/compaso_halo_catalog.py index 174eebd..1e35484 100644 --- a/abacusnbody/data/compaso_halo_catalog.py +++ b/abacusnbody/data/compaso_halo_catalog.py @@ -20,17 +20,21 @@ from pathlib import PurePath import asdf -import asdf.compression import astropy.table import numba as nb import numpy as np from astropy.table import Table +try: + import asdf._compression as asdf_compression +except ImportError: + import asdf.compression as asdf_compression + from . import asdf as _asdf from . import bitpacked try: - asdf.compression.validate('blsc') + asdf_compression.validate('blsc') except Exception as e: raise Exception("Abacus ASDF extension not properly loaded! Try reinstalling abacusutils, or updating ASDF: `pip install asdf>=2.8`") from e diff --git a/abacusnbody/data/pipe_asdf.py b/abacusnbody/data/pipe_asdf.py index 5abb791..50980c6 100755 --- a/abacusnbody/data/pipe_asdf.py +++ b/abacusnbody/data/pipe_asdf.py @@ -86,11 +86,15 @@ from timeit import default_timer as timer import asdf -import asdf.compression import numpy as np try: - asdf.compression.validate('blsc') + import asdf._compression as asdf_compression +except ImportError: + import asdf.compression as asdf_compression + +try: + asdf_compression.validate('blsc') except Exception as e: raise Exception("Abacus ASDF extension not properly loaded! Try reinstalling abacusutils, or updating ASDF: `pip install asdf>=2.8`") from e diff --git a/abacusnbody/data/read_abacus.py b/abacusnbody/data/read_abacus.py index bc22276..db82908 100644 --- a/abacusnbody/data/read_abacus.py +++ b/abacusnbody/data/read_abacus.py @@ -71,9 +71,14 @@ def read_asdf(fn, load=None, colname=None, dtype=np.float32, verbose=True, **kwa ''' import asdf - import asdf.compression + + try: + import asdf._compression as asdf_compression + except ImportError: + import asdf.compression as asdf_compression + try: - asdf.compression.validate('blsc') + asdf_compression.validate('blsc') except Exception as e: raise Exception("Abacus ASDF extension not properly loaded! \ Try reinstalling abacusutils: `pip install 'abacusutils>=1'`, \