add optional verbosity to _ft

CamDavidsonPilon · Apr 21, 2015 · a0bbb81 · a0bbb81
1 parent 7826a01
commit a0bbb81
Show file tree

Hide file tree

Showing 2 changed files with 11 additions and 6 deletions.
diff --git a/lifetimes/estimation.py b/lifetimes/estimation.py
@@ -38,7 +38,7 @@ class ParetoNBDFitter(BaseFitter):
     def __init__(self, penalizer_coef=0.):
         self.penalizer_coef = penalizer_coef
 
-    def fit(self, frequency, recency, T, iterative_fitting=1, initial_params=None):
+    def fit(self, frequency, recency, T, iterative_fitting=1, initial_params=None, verbose=False):
         """
         This methods fits the data to the Pareto/NBD model.
 
@@ -49,6 +49,8 @@ def fit(self, frequency, recency, T, iterative_fitting=1, initial_params=None):
             iterative_fitting: perform `iterative_fitting` additional fits to find the best
                 parameters for the model. Setting to 0 will improve peformance but possibly
                 hurt estimates.
+            initial_params: set intial params for the iterative fitter.
+            verbose: set to true to print out convergence diagnostics.
 
         Returns:
             self, with additional properties and methods like params_ and plot
@@ -58,7 +60,7 @@ def fit(self, frequency, recency, T, iterative_fitting=1, initial_params=None):
         recency = np.asarray(recency)
         T = np.asarray(T)
 
-        params, self._negative_log_likelihood_ = _fit(self._negative_log_likelihood, frequency, recency, T, iterative_fitting, self.penalizer_coef, initial_params)
+        params, self._negative_log_likelihood_ = _fit(self._negative_log_likelihood, frequency, recency, T, iterative_fitting, self.penalizer_coef, initial_params, verbose)
 
         self.params_ = OrderedDict(zip(['r', 'alpha', 's', 'beta'], params))
         self.data = pd.DataFrame(np.c_[frequency, recency, T], columns=['frequency', 'recency', 'T'])
@@ -196,7 +198,7 @@ class BetaGeoFitter(BaseFitter):
     def __init__(self, penalizer_coef=0.):
         self.penalizer_coef = penalizer_coef
 
-    def fit(self, frequency, recency, T, iterative_fitting=1, initial_params=None):
+    def fit(self, frequency, recency, T, iterative_fitting=1, initial_params=None, verbose=False):
         """
         This methods fits the data to the BG/NBD model.
 
@@ -207,6 +209,9 @@ def fit(self, frequency, recency, T, iterative_fitting=1, initial_params=None):
             iterative_fitting: perform `iterative_fitting` additional fits to find the best
                 parameters for the model. Setting to 0 will improve peformance but possibly
                 hurt estimates.
+            initial_params: set the initial parameters for the fitter.
+            verbose: set to true to print out convergence diagnostics.
+
 
         Returns:
             self, with additional properties and methods like params_ and plot
@@ -216,7 +221,7 @@ def fit(self, frequency, recency, T, iterative_fitting=1, initial_params=None):
         recency = np.asarray(recency)
         T = np.asarray(T)
 
-        params, self._negative_log_likelihood_ = _fit(self._negative_log_likelihood, frequency, recency, T, iterative_fitting, self.penalizer_coef, initial_params)
+        params, self._negative_log_likelihood_ = _fit(self._negative_log_likelihood, frequency, recency, T, iterative_fitting, self.penalizer_coef, initial_params, verbose)
 
         self.params_ = OrderedDict(zip(['r', 'alpha', 'a', 'b'], params))
         self.data = pd.DataFrame(np.c_[frequency, recency, T], columns=['frequency', 'recency', 'T'])

diff --git a/lifetimes/utils.py b/lifetimes/utils.py
@@ -142,7 +142,7 @@ def calculate_alive_path(model, transactions, datetime_col, t, freq='D'):
     return customer_history.apply(lambda row: model.conditional_probability_alive(row['frequency'], row['recency'], row['T']), axis=1)
 
 
-def _fit(minimizing_function, frequency, recency, T, iterative_fitting, penalizer_coef, initial_params):
+def _fit(minimizing_function, frequency, recency, T, iterative_fitting, penalizer_coef, initial_params, disp):
     ll = []
     sols = []
     methods = ['Powell', 'Nelder-Mead']
@@ -151,7 +151,7 @@ def _fit(minimizing_function, frequency, recency, T, iterative_fitting, penalize
         fit_method = methods[i % len(methods)]
         params_init = np.random.exponential(0.5, size=4) if initial_params is None else initial_params
         output = minimize(minimizing_function, method=fit_method, tol=1e-6,
-                          x0=params_init, args=(frequency, recency, T, penalizer_coef), options={'disp': False})
+                          x0=params_init, args=(frequency, recency, T, penalizer_coef), options={'disp': disp})
         ll.append(output.fun)
         sols.append(output.x)
     minimizing_params = sols[np.argmin(ll)]