Fix ASDF error in ZCV module, add ASDF 3.1.0 support

CHANGES.rst

@@ -1,6 +1,19 @@
 Changelog
 =========
 
+2.0.1 (2024-03-01)
+------------------
+This is a bugfix release primarily to add support for ASDF 3.1.0.
+Several other bugs in the ZCV and HOD modules are also fixed.
+
+Fixes
+~~~~~
+- Fix ASDF error in ZCV module, add ASDF 3.1.0 support [#130]
+- bug fix keys() [#126]
+- Sandydev: fixed reseeding bug [#127]
+- fix reseed bug [#128]
+- backward compatible fix for velocity bias [#121]
+
 2.0.0 (2023-11-15)
 ------------------
 

abacusnbody/data/asdf.py

@@ -11,17 +11,11 @@
 import struct
 import time
 
-import asdf
 import blosc
 import numpy as np
 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)
 

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
 

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
 

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'`, \

abacusnbody/hod/zcv/advect_fields.py

@@ -262,8 +262,8 @@ def main(path2config, want_rsd=False, alt_simname=None, save_3D_power=False, onl
             print("Computing cross-correlation of", keynames[i], keynames[j])
 
             # load field
-            field_fft_i = asdf.open(fields_fft_fn[i])['data']
-            field_fft_j = asdf.open(fields_fft_fn[j])['data']
+            field_fft_i = asdf.open(fields_fft_fn[i], lazy_load=False)['data']
+            field_fft_j = asdf.open(fields_fft_fn[j], lazy_load=False)['data']
 
             if save_3D_power:
                 power_ij_fn = Path(save_z_dir) / f"power{rsd_str}_{keynames[i]}_{keynames[j]}_nmesh{nmesh:d}.asdf"
@@ -295,7 +295,7 @@ def main(path2config, want_rsd=False, alt_simname=None, save_3D_power=False, onl
                 P['binned_poles'] *= field_D[i]*field_D[j]
                 pk_auto.append(P['power'])
                 pk_ij_dict[f'P_kmu_{keynames[i]}_{keynames[j]}'] = P['power']
-                pk_ij_dict[f'N_kmu_{keynames[i]}_{keynames[j]}'] = P['N_modes']
+                pk_ij_dict[f'N_kmu_{keynames[i]}_{keynames[j]}'] = P['N_mode']
                 pk_ij_dict[f'P_ell_{keynames[i]}_{keynames[j]}'] = P['binned_poles']
                 pk_ij_dict[f'N_ell_{keynames[i]}_{keynames[j]}'] = P['N_mode_poles']
                 del field_fft_i, field_fft_j; gc.collect() # noqa: E702