Fix dhdl2series and u_nk2series would not reattach the unit (#248)

* ensure that `@pass_attrs` is applied to all functions that handle our DataFrames
  fix: dhdl2series() and u_nk2series()
* update tests
* update CHANGELOG

Co-authored-by: Zhiyi Wu <zwu@exscientia.co.uk>

CHANGES

@@ -48,6 +48,7 @@ Fixes
   - All parsers now have a 'extract(file, T)' method that returns a dict with both
     "dHdl" and "u_nk" data (or None). THe AMBER parser when using this function will read the 
     file just once, extracting all data at once. (issue #222, PR #240)
+  - Fix dhdl2series and u_nk2series would not reattach the unit. (PR #248)
 
 07/22/2022 xiki-tempula, IAlibay, dotsdl, orbeckst, ptmerz
 

src/alchemlyb/__init__.py

@@ -1,4 +1,5 @@
 import pandas as pd
+from functools import wraps
 
 from ._version import get_versions
 __version__ = get_versions()['version']
@@ -11,6 +12,8 @@ def pass_attrs(func):
  
     .. versionadded:: 0.5.0
  '''
+
+    @wraps(func)
     def wrapper(input_dataframe, *args,**kwargs):
         dataframe = func(input_dataframe, *args,**kwargs)
         dataframe.attrs = input_dataframe.attrs

src/alchemlyb/preprocessing/subsampling.py

@@ -5,6 +5,7 @@
 from pymbar.timeseries import (statisticalInefficiency,
                                detectEquilibration,
                                subsampleCorrelatedData, )
+from .. import pass_attrs
 
 def decorrelate_u_nk(df, method='dhdl', drop_duplicates=True,
                      sort=True, remove_burnin=False, **kwargs):
@@ -106,6 +107,7 @@ def decorrelate_dhdl(df, drop_duplicates=True, sort=True,
     else:
         return statistical_inefficiency(df, series, **kwargs)
 
+@pass_attrs
 def u_nk2series(df, method='dhdl'):
     """Convert an u_nk DataFrame into a series based on the selected method
     for subsampling.
@@ -175,6 +177,7 @@ def u_nk2series(df, method='dhdl'):
     return series
 
 
+@pass_attrs
 def dhdl2series(df, method=''):
     """Convert a dhdl DataFrame to a series for subsampling.
 

src/alchemlyb/tests/test_units.py

@@ -4,8 +4,12 @@
 import alchemlyb
 from alchemlyb import pass_attrs
 from alchemtest.gmx import load_benzene
-from alchemlyb.parsing.gmx import extract_dHdl
+from alchemlyb.parsing.gmx import extract_dHdl, extract_u_nk
 from alchemlyb.postprocessors.units import to_kT
+from alchemlyb.preprocessing import (dhdl2series, u_nk2series,
+                                     decorrelate_u_nk, decorrelate_dhdl,
+                                     slicing, statistical_inefficiency,
+                                     equilibrium_detection)
 
 def test_noT():
     '''Test no temperature error'''
@@ -120,3 +124,32 @@ def test_pd_slice():
     df = pd.DataFrame(data=d)
     df.attrs = {1: 1}
     assert df[::2].attrs == {1: 1}
+
+class TestRetainUnit():
+    '''This test tests if the functions that should retain the unit would actually
+     retain the units.'''
+    @staticmethod
+    @pytest.fixture(scope='class')
+    def dhdl():
+        dataset = load_benzene()
+        dhdl = extract_dHdl(dataset['data']['Coulomb'][0], 310)
+        return dhdl
+
+    @staticmethod
+    @pytest.fixture(scope='class')
+    def u_nk():
+        dataset = load_benzene()
+        u_nk = extract_u_nk(dataset['data']['Coulomb'][0], 310)
+        return u_nk
+
+    @pytest.mark.parametrize('func,fixture_in',
+                             [(dhdl2series, 'dhdl'),
+                              (u_nk2series, 'u_nk'),
+                              (decorrelate_u_nk, 'u_nk'),
+                              (decorrelate_dhdl, 'dhdl'),
+                              (slicing, 'dhdl'),
+                              (statistical_inefficiency, 'dhdl'),
+                              (equilibrium_detection, 'dhdl')])
+    def test_function(self, func, fixture_in, request):
+        result = func(request.getfixturevalue(fixture_in))
+        assert result.attrs['energy_unit'] is not None