fix load_model

aadya940 · Jun 18, 2024 · da2f8ee · da2f8ee
1 parent fd1bdfe
commit da2f8ee
Show file tree

Hide file tree

Showing 2 changed files with 22 additions and 6 deletions.
diff --git a/chainopy/fileio.py b/chainopy/fileio.py
@@ -151,9 +151,9 @@ def _load_model_markovchain(filepath):
                 _eigenvalues,
                 _eigenvectors,
                 epsilon,
-            ]
+            ]  # 6
 
-        return [transition_matrix, states, eigendecom, epsilon]
+        return [transition_matrix, states, eigendecom, epsilon]  # 4
 
     elif isinstance(data["tpm"], list):
         # Handle Non - Sparse Matrix Cases
@@ -171,9 +171,15 @@ def _load_model_markovchain(filepath):
                 data["eigenvectors-imag"], np.complex_
             )
 
-            return [transition_matrix, states, eigendecom, _eigenvalues, _eigenvectors]
+            return [
+                transition_matrix,
+                states,
+                eigendecom,
+                _eigenvalues,
+                _eigenvectors,
+            ]  # 5
 
-        return [transition_matrix, states, eigendecom]
+        return [transition_matrix, states, eigendecom]  # 3
 
 
 def load_text(path: str):

diff --git a/chainopy/markov_chain.py b/chainopy/markov_chain.py
@@ -544,9 +544,19 @@ def load_model(self, path: str):
         """
 
         result = _load_model_markovchain(path)
-        if len(result) == 4:
+        if len(result) == 3:
             self.tpm, self.states, self.eigendecom, self.epsilon = result
-        else:
+        elif len(result) == 4:
+            self.tpm, self.states, self.eigendecom, self.epsilon = result
+        elif len(result) == 5:
+            (
+                self.tpm,
+                self.states,
+                self.eigendecom,
+                self.eigenvalues,
+                self.eigenvectors,
+            ) = result
+        elif len(result) == 6:
             (
                 self.tpm,
                 self.states,