model clarity improvement

Model.py

@@ -2,10 +2,10 @@ import numpy as np
 import matplotlib.pyplot as plt
 from scipy.integrate import ode
 
-R = 8.314  # ideal gas constant (J/(mol*K))
-F = 96.485  # coulomb_per_millimole, Faraday constant
-T = 298  # room temperature (K)
-RT = R * T
+R: float = 8.314  # ideal gas constant (J/(mol*K))
+F: float = 96.485  # coulomb_per_millimole, Faraday constant
+T: float = 298  # room temperature (K)
+RT: float = R * T  # J/mol
 
 
 class Model:
@@ -22,13 +22,13 @@ class Model:
         self.Nai: float = 11000  # uM, Myoplasmic Na+ concentration
         self.Nao: float = 150000  # uM, Extracellular Na+ concentration
 
-        self.eta: float = 0.35  # Controls voltage dependance of Na/Ca2+ exchange
+        self.eta: float = 0.35  # Controls voltage dependance of Na/Ca2+ excng
 
         self.km_Na: float = (
             87500  # uM, Na+ half-saturation constant for Na+/Ca2+ exchange
         )
         self.k_sat: float = (
-            0.1  # Na+/Ca2+ exchange saturation factor at very negative potentials
+            0.1  # Na+/Ca2+ exchange saturation factor at very neg potentials
         )
         self.km_Ca: float = (
             1380  # uM, Ca2+ half-saturation constant for Na+/Ca2+ exchange
@@ -132,7 +132,7 @@ class Model:
         )
         self.V_JSR: float = 0.12e-6  # ul, Junctional SR volume
 
-    def ode_system(self, t, states):
+    def ode_system(self, t, states: list[float]) -> list[float]:
 
         (
             Cass,
@@ -331,7 +331,7 @@ class Model:
             dFCadt,
         ]
 
-    def mem_potential(self, t):
+    def mem_potential(self, t: np.ndarray) -> np.ndarray:
         t = np.asarray(t)
         v1 = 0
         t0 = 100
@@ -339,7 +339,7 @@ class Model:
         n = (t // self.period) * self.period
         return np.where((t0 + n <= t) & (t < t1 + n), v1, self.V_mem_rest)
 
-    def get_initial_values(self):
+    def get_initial_values(self) -> list[float]:
         Cai_0: float = 0.11712  # Cai
         FCa_0: float = (self.k_on * self.F_tot * Cai_0) / (
             self.k_on * Cai_0 + self.k_off
@@ -365,11 +365,13 @@ class Model:
             FCa_0,  # FCa
         ]
 
-    def calculated_current(self, states, V):
-        ICaL = self.gCaL * states[1, :] * (V - self.ECaL)
-        return ICaL
+    def calculated_current(self, states: list[float], V: np.ndarray) -> np.ndarray:
+        Os = states[1]
 
-    def solve(self, initial_values, tspan, dt, times):
+        I_CaL = self.gCaL * Os * (V - self.ECaL)
+        return I_CaL
+
+    def solve(self, initial_values: list[float], tspan, dt, times):
         times = np.arange(*tspan, dt)
 
         r = ode(self.ode_system)