From 2c508941fab8528cf96df3cf668eeacc96e5bdce Mon Sep 17 00:00:00 2001
From: Marius <44572437+m-a-r-i-u-s@users.noreply.github.com>
Date: Fri, 6 Sep 2024 17:38:11 +0200
Subject: [PATCH] Update base.py

Fixed minor bug in TDM selection
---
 domhmm/analysis/base.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/domhmm/analysis/base.py b/domhmm/analysis/base.py
index cac2a7e..92c94ec 100644
--- a/domhmm/analysis/base.py
+++ b/domhmm/analysis/base.py
@@ -232,14 +232,14 @@ def __init__(
             for each in tmd_protein_list:
                 for leaflet, query in each.items():
                     if leaflet not in ["0", "1"]:
-                        raise ValueError("TDM Protein list should contain dictionaries in format {'0': ..., '1': ...} "
+                        raise ValueError("Entry for each TDM protein should be a dictionary in the format {'0': ..., '1': ...} "
                                          "where 0 for lower leaflet and 1 for upper leaflet.")
                     if isinstance(query, AtomGroup):
                         # Take center of geometry of three positions
                         cog = np.mean(query.positions, axis=0)
                         self.tmd_protein[leaflet].append(cog)
                     # Character string was provided as input, assume it contains a selection for an MDAnalysis.AtomGroup
-                    elif isinstance(leaflet_select[i], str):
+                    elif isinstance(leaflet_select[int(leaflet)], str):
                         # Try to create a MDAnalysis.AtomGroup, raise a ValueError if not selection group could be
                         # provided
                         try:
@@ -256,7 +256,7 @@ def __init__(
             # An unknown argument is provided for tdm_protein_list
             raise ValueError(
                 "Please provide tdm_protein_list in list format such as [{'0': upper leaflet related 3 atom, "
-                "'1': lower leaflet related 3 atom }, {'0': ..., '1': ...}]")
+                "'1': lower leaflet related 3 atom }, {'0': ..., '1': ...}]. Every dictionary stands for an individual transmembrane protein.")
         # -----------------------------------------------------------------HMMs--------------------------------------- #
 
         # Check for user-specified trained HMM