fix #4 type hints added

Author: ulf1

CHANGES.md

@@ -1,3 +1,6 @@
+# 0.1.2 / 2019-09-26
+
+  * add type hints
 
 # 0.1.1 / 2019-03-23
 

panelctmc/panel_to_datalist.py

@@ -1,9 +1,12 @@
 
 from yearfrac import yearfrac
 import numpy as np
+import datetime
 
 
-def panel_to_datalist(data, lastdate=None):
+def panel_to_datalist(data: np.ndarray,
+                      lastdate: datetime.datetime = None
+                      ) -> list:
     """Transforms array/list to ctmc's internal data format
 
     Parameters:

panelctmc/panelctmc_class.py

@@ -2,13 +2,18 @@
 from sklearn.base import BaseEstimator
 from .panelctmc_func import panelctmc
 from ctmc import simulate
+import datetime
+import numpy as np
 
 
 class PanelCtmc(BaseEstimator):
     """Continous Time Markov Chain for Panel Data, sklearn API class"""
 
-    def __init__(self, mapping=None, lastdate=None,
-                 transintv=1.0, toltime=1e-8, debug=False):
+    def __init__(self, mapping: list = None,
+                 lastdate: datetime.datetime = None,
+                 transintv: float = 1.0,
+                 toltime: float = 1e-8,
+                 debug: bool = False):
         self.mapping = mapping
         self.lastdate = lastdate
         self.transintv = transintv
@@ -20,7 +25,7 @@ def __init__(self, mapping=None, lastdate=None,
         self.statetime = None
         self.datalist = None
 
-    def fit(self, X, y=None):
+    def fit(self, X: np.ndarray, y=None):
         (
             self.transmat,
             self.genmat,
@@ -35,5 +40,5 @@ def fit(self, X, y=None):
             debug=self.debug)
         return self
 
-    def predict(self, X, steps=1):
+    def predict(self, X: np.ndarray, steps: int = 1) -> np.ndarray:
         return simulate(X, self.transmat, steps)

panelctmc/panelctmc_func.py

@@ -6,8 +6,13 @@
 from ctmc import ctmc, datacorrection
 
 
-def panelctmc(paneldata, mapping, lastdate=None,
-              transintv=1.0, toltime=1e-8, debug=True):
+def panelctmc(paneldata: np.ndarray,
+              mapping: list,
+              lastdate: datetime = None,
+              transintv: float = 1.0,
+              toltime: float = 1e-8,
+              debug: bool = True) -> (np.ndarray, np.ndarray, np.ndarray,
+                                      np.ndarray, list):
     # check if numpy array
     if not isinstance(paneldata, np.ndarray):
         raise Exception("'paneldata' is not a numpy array")