diff --git a/baybe/utils/chemistry.py b/baybe/utils/chemistry.py index 86e5c80b5..7e7e93d03 100644 --- a/baybe/utils/chemistry.py +++ b/baybe/utils/chemistry.py @@ -149,15 +149,14 @@ def smiles_to_fingerprint_features( ) name = f"{encoding.name}_" prefix = prefix + "_" if prefix else "" - if all("fingerprint" in f for f in fingerprint_encoder.get_feature_names_out()): - col_names = [ - prefix + name + f.split("fingerprint")[1] - for f in fingerprint_encoder.get_feature_names_out() - ] - else: - col_names = [ - prefix + name + f for f in fingerprint_encoder.get_feature_names_out() - ] + no_descriptor_names = all( + "fingerprint" in f for f in fingerprint_encoder.get_feature_names_out() + ) + suffixes = [ + f.split("fingerprint")[1] if no_descriptor_names else f + for f in fingerprint_encoder.get_feature_names_out() + ] + col_names = [prefix + name + suffix for suffix in suffixes] df = pd.DataFrame(features, columns=col_names, dtype=DTypeFloatNumpy) return df