Handle normalised lambda and lambda_grad metadata entries.

Author: lohedges

src/BioSimSpace/FreeEnergy/_relative.py

@@ -1059,15 +1059,20 @@ def _somd2_extract(parquet_file, T=None, estimator="MBAR"):
             raise ValueError("Parquet metadata does not contain 'lambda'.")
         if not is_mbar:
             try:
-                lambda_grad = metadata["lambda_grad"]
+                # Normalise to :.5f strings to match sire energy trajectory column names.
+                lambda_grad = [f"{float(v):.5f}" for v in metadata["lambda_grad"]]
             except:
                 raise ValueError("Parquet metadata does not contain 'lambda grad'")
         else:
             try:
-                lambda_grad = metadata["lambda_grad"]
+                # Normalise to :.5f strings to match sire energy trajectory column names.
+                lambda_grad = [f"{float(v):.5f}" for v in metadata["lambda_grad"]]
             except:
                 lambda_grad = []
 
+        # Key used to index the simulated lambda column in the dataframe.
+        lam_key = f"{lam:.5f}"
+
         # Make sure that the temperature is correct.
         if not T == temperature:
             raise ValueError(
@@ -1083,7 +1088,7 @@ def _somd2_extract(parquet_file, T=None, estimator="MBAR"):
             df = df[[x for x in df.columns if x not in lambda_grad]]
 
             # Subtract the potential at the simulated lambda.
-            df = df.subtract(df[lam], axis=0)
+            df = df.subtract(df[lam_key], axis=0)
 
             # Apply the existing attributes.
             df.attrs = attrs
@@ -1096,19 +1101,19 @@ def _somd2_extract(parquet_file, T=None, estimator="MBAR"):
                 lam_delta = lambda_grad[0]
 
                 # Forward difference.
-                if lam_delta > lam:
-                    incr = lam_delta - lam
-                    grad = (df[lam_delta] - df[lam]) / incr
+                if float(lam_delta) > lam:
+                    incr = float(lam_delta) - lam
+                    grad = (df[lam_delta] - df[lam_key]) / incr
 
                 # Backward difference.
                 else:
-                    incr = lam - lam_delta
-                    grad = (df[lam] - df[lam_delta]) / incr
+                    incr = lam - float(lam_delta)
+                    grad = (df[lam_key] - df[lam_delta]) / incr
 
             # Central difference.
             else:
                 lam_below, lam_above = lambda_grad
-                double_incr = lam_above - lam_below
+                double_incr = float(lam_above) - float(lam_below)
                 grad = (df[lam_above] - df[lam_below]) / double_incr
 
             # Create a DataFrame with the multi-index