diff --git a/python/BioSimSpace/FreeEnergy/_relative.py b/python/BioSimSpace/FreeEnergy/_relative.py index 1ba69e79..218803cf 100644 --- a/python/BioSimSpace/FreeEnergy/_relative.py +++ b/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