Second round of changes

AUTHORS

@@ -35,4 +35,8 @@ Chronological list of authors
   - Hyungro Lee (@lee212)
   - Mohammad S. Barhaghi (@msoroush)
 2020
-  - Zhiyi Wu (@xiki-tempula)
+  - Zhiyi Wu (@xiki-tempula)
+
+2021
+  - Jérôme Hénin (@jhenin)
+  - Thomas T. Joseph (@ttjoseph)

src/alchemlyb/parsing/namd.py

@@ -3,13 +3,17 @@
 """
 import pandas as pd
 import numpy as np
+import logging
 from .util import anyopen
 from . import _init_attrs
 from ..postprocessors.units import R_kJmol, kJ2kcal
 
+logger = logging.getLogger("alchemlyb.parsers.NAMD")
+
 k_b = R_kJmol * kJ2kcal
 
-def get_lambdas(fep_files):
+
+def _get_lambdas(fep_files):
     """Retrieves all lambda values included in the FEP files provided.
 
     We have to do this in order to tolerate truncated and restarted fepout files.
@@ -46,16 +50,16 @@ def get_lambdas(fep_files):
 
                 # Make sure the lambda2 values are consistent
                 if lambda1 in lambda_fwd_map and lambda_fwd_map[lambda1] != lambda2:
-                    print(f'namd.py: get_lambdas: Error: fwd: lambda1 {lambda1} has lambda2 {lambda_fwd_map[lambda1]} but it should be {lambda2}')
-                    return None
+                    logger.error(f'fwd: lambda1 {lambda1} has lambda2 {lambda_fwd_map[lambda1]} but it should be {lambda2}')
+                    raise ValueError('Inconsistent lambda values')
 
                 lambda_fwd_map[lambda1] = lambda2
 
                 # Make sure the lambda_idws values are consistent
                 if lambda_idws is not None:
                     if lambda1 in lambda_bwd_map and lambda_bwd_map[lambda1] != lambda_idws:
-                        print(f'namd.py: get_lambdas: Error: bwd: lambda1 {lambda1} has lambda_idws {lambda_bwd_map[lambda1]} but it should be {lambda_idws}')
-                        return None
+                        logger.error(f'bwd: lambda1 {lambda1} has lambda_idws {lambda_bwd_map[lambda1]} but it should be {lambda_idws}')
+                        raise ValueError('Inconsistent lambda values')
                     lambda_bwd_map[lambda1] = lambda_idws
 
     all_lambdas = set()
@@ -103,7 +107,7 @@ def extract_u_nk(fep_files, T):
 
     time = 0
     # Extract the lambda values only from the fepouts
-    all_lambdas = get_lambdas(fep_files)
+    all_lambdas = _get_lambdas(fep_files)
     # open and get data from fep file.
     # We sort the list of fep files in case some of them represent restarted windows.
     # The assumption is that they make sense in lexicographic order.
@@ -142,14 +146,13 @@ def extract_u_nk(fep_files, T):
                     else:
                         lambda_idws = None
 
-                    # If the lambdas are not what we thought they would be, return None, ensuring the calculation
+                    # If the lambdas are not what we thought they would be, raise an exception to ensure the calculation
                     # fails. This can happen if fepouts where one window spans multiple fepouts are processed out of order
                     if lambda1_at_start is not None \
                         and (lambda1, lambda2, lambda_idws) != (lambda1_at_start, lambda2_at_start, lambda_idws_at_start):
-                        print("namd.py: extract_u_nk: Error: Lambdas changed unexpectedly while processing", fep_file)
-                        print(f"namd.py: extract_u_nk: Error: l1, l2, lidws: {lambda1_at_start}, {lambda2_at_start}, {lambda_idws_at_start} changed to {lambda1}, {lambda2}, {lambda_idws}")
-                        print(f"namd.py: extract_u_nk: Error: fep_file = {fep_file}; has_idws = {has_idws}")
-                        return None
+                        logger.error("Lambdas changed unexpectedly while processing", fep_file)
+                        logger.error(f"l1, l2, lidws: {lambda1_at_start}, {lambda2_at_start}, {lambda_idws_at_start} changed to {lambda1}, {lambda2}, {lambda_idws}")
+                        raise ValueError("Inconsistent lambda values")
 
                     # As we are at the end of a window, convert last window's work and times values to np arrays
                     # (with energy unit kT since they were kcal/mol in the fepouts)
@@ -203,7 +206,7 @@ def extract_u_nk(fep_files, T):
                     parsing = True
 
     if len(win_de) != 0 or len(win_de_back) != 0:
-        print('Warning: trailing data without footer line (\"#Free energy...\"). Interrupted run?')
+        logger.warning('Trailing data without footer line (\"#Free energy...\"). Interrupted run?')
 
     if lambda2 in (0.0, 1.0):
         # this excludes the IDWS case where a dataframe already exists for both endpoints

src/alchemlyb/tests/parsing/test_namd.py

@@ -6,6 +6,7 @@
 from alchemlyb.parsing.namd import extract_u_nk
 from alchemtest.namd import load_tyr2ala
 from alchemtest.namd import load_idws
+from alchemtest.namd import load_restarted
 
 
 @pytest.fixture(scope="module")
@@ -34,3 +35,14 @@ def test_u_nk_idws():
 
     assert u_nk.index.names == ['time', 'fep-lambda']
     assert u_nk.shape == (29252, 11)
+
+def test_u_nk_restarted():
+    """Test that u_nk has the correct form when extracted from an IDWS
+    FEP run that includes a termination and restart.
+    """
+
+    filenames = load_restarted()['data']['both']
+    u_nk = extract_u_nk(filenames, T=300)
+
+    assert u_nk.index.names == ['time', 'fep-lambda']
+    assert u_nk.shape == (30061, 11)