From 4959f0b736c9be5ba21c99df87d658256c4c97ab Mon Sep 17 00:00:00 2001
From: jac16 <jennifer.clark@nist.gov>
Date: Tue, 12 Nov 2024 14:28:50 -0500
Subject: [PATCH] Bug fix columns, add name attribute

---
 src/alchemlyb/parsing/lammps.py | 36 +++++++++++++++++++++++++--------
 1 file changed, 28 insertions(+), 8 deletions(-)

diff --git a/src/alchemlyb/parsing/lammps.py b/src/alchemlyb/parsing/lammps.py
index fbe4fee1..7cded23b 100644
--- a/src/alchemlyb/parsing/lammps.py
+++ b/src/alchemlyb/parsing/lammps.py
@@ -768,24 +768,42 @@ def extract_u_nk(
                             lambda1, lambda12
                         )
                     )
+                if lambda1 == lambda12 and not np.all(tmp_df2["dU_nk"][0] == 0):
+                    raise ValueError(
+                        f"The difference in dU should be zero when lambda = lambda', {lambda1} = {lambda12},"
+                        " Check that 'column_dU' was defined correctly."
+                    )
+                                        
                 if (
                     u_nk.loc[u_nk[lambda1_col] == lambda1, column_name].shape[0]
                     != tmp_df2["dU_nk"].shape[0]
                 ):
-                    raise ValueError(
-                        "Number of energy values in file, {}, N={}, inconsistent with previous files of length, {}.".format(
+                    old_length = tmp_df2["dU_nk"].shape[0]
+                    start = u_nk.loc[u_nk[lambda1_col] == lambda1, "time"].iloc[0]
+                    stop = u_nk.loc[u_nk[lambda1_col] == lambda1, "time"].iloc[-1]
+                    stepsize = (
+                        u_nk.loc[u_nk[lambda1_col] == lambda1, "time"].iloc[1] 
+                        - u_nk.loc[u_nk[lambda1_col] == lambda1, "time"].iloc[0]
+                    )
+                    # Fill in gaps where 'time' is NaN
+                    nan_index = np.unique(np.where(tmp_df2['time'].isnull())[0])
+                    for index in nan_index:
+                        tmp_df2.loc[index, "time"] = tmp_df2.loc[index-1, "time"] + stepsize
+                    # Add rows of NaN for timesteps that are missing
+                    new_index = pd.Index(np.arange(start, stop+stepsize, stepsize), name="time")
+                    tmp_df2 = tmp_df2.set_index("time").reindex(new_index).reset_index()
+                    warnings.warn(
+                        "Number of energy values in file, {}, N={}, inconsistent with previous".format(
                             file,
-                            tmp_df2["dU_nk"].shape[0],
+                            old_length,
+                        )
+                        + " files of length, {}. Adding NaN to row: {}".format(
                             u_nk.loc[u_nk[lambda1_col] == lambda1, column_name].shape[
                                 0
                             ],
+                            np.unique(np.where(tmp_df2.isna())[0]),
                         )
                     )
-                if lambda1 == lambda12 and not np.all(tmp_df2["dU_nk"][0] == 0):
-                    raise ValueError(
-                        f"The difference in dU should be zero when lambda = lambda', {lambda1} = {lambda12},"
-                        " Check that 'column_dU' was defined correctly."
-                    )
 
                 # calculate reduced potential u_k = dH + pV + U
                 u_nk.loc[u_nk[lambda1_col] == lambda1, column_name] = beta * (
@@ -802,6 +820,8 @@ def extract_u_nk(
         u_nk.set_index(["time", "coul-lambda", "vdw-lambda"], inplace=True)
     u_nk.name = "u_nk"
 
+    u_nk = u_nk.dropna()
+    
     return u_nk