Disallow indexing of non-variables in state

1 files changed, 4 insertions, 0 deletions

diff --git a/mdpoly/state.py b/mdpoly/state.py
index 001d9b6..f22d502 100644
--- a/mdpoly/state.py
+++ b/mdpoly/state.py
@@ -22,6 +22,10 @@ class State:
 
     def index(self, var: Var) -> Index:
         """ Get the index for a variable. """
+        from .leaves import Var
+        if not isinstance(var, Var):
+            raise IndexError(f"Only variables (type {Var}) can be indexed.")
+        
         if var not in self.variables.keys():
             new_index = self._make_index()
             self.variables[var] = new_index