Add support for focused sampling with SOMD2

python/BioSimSpace/FreeEnergy/_relative.py

@@ -1064,15 +1064,16 @@ def _somd2_extract(parquet_file, T=None, estimator="MBAR"):
             lam = float(metadata["lambda"])
         except:
             raise ValueError("Parquet metadata does not contain 'lambda'.")
-        try:
-            lambda_array = metadata["lambda_array"]
-        except:
-            raise ValueError("Parquet metadata does not contain 'lambda array'")
         if not is_mbar:
             try:
                 lambda_grad = metadata["lambda_grad"]
             except:
                 raise ValueError("Parquet metadata does not contain 'lambda grad'")
+        else:
+            try:
+                lambda_grad = metadata["lambda_grad"]
+            except:
+                lambda_grad = []
 
         # Make sure that the temperature is correct.
         if not T == temperature:
@@ -1085,8 +1086,8 @@ def _somd2_extract(parquet_file, T=None, estimator="MBAR"):
         df = table.to_pandas()
 
         if is_mbar:
-            # Extract the columns correspodning to the lambda array.
-            df = df[[x for x in lambda_array]]
+            # Extract all columns other than those used for the gradient.
+            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)
@@ -1311,10 +1312,21 @@ def _analyse_internal(files, temperatures, lambdas, engine, estimator, **kwargs)
 
         # Preprocess the data.
         try:
-            processed_data = Relative._preprocess_data(data, estimator, **kwargs)
-            processed_data = _alchemlyb.concat(processed_data)
-        except:
-            _warnings.warn("Could not preprocess the data!")
+            preprocess = kwargs.pop("preprocess", True)
+        except KeyError:
+            preprocess = True
+
+        if not isinstance(preprocess, bool):
+            raise TypeError("'preprocess' must be of type 'bool'.")
+
+        if preprocess:
+            try:
+                processed_data = Relative._preprocess_data(data, estimator, **kwargs)
+                processed_data = _alchemlyb.concat(processed_data)
+            except:
+                _warnings.warn("Could not preprocess the data!")
+                processed_data = _alchemlyb.concat(data)
+        else:
             processed_data = _alchemlyb.concat(data)
 
         mbar_method = None