Update geoshapley.py

Incorporate the suggestion from #10

geoshapley/geoshapley.py

@@ -70,8 +70,9 @@ def _kernel_geoshap_single(self, x, reference):
         y = self.predict_f(V).reshape(-1)
 
         #Solve WLS
-        ZTw = np.dot(Z.T, np.diag(weights))
-
+        #ZTw = np.dot(Z.T, np.diag(weights))
+        ZTw = Z.T * weights
+
         phi = np.linalg.solve(np.dot(ZTw, Z), np.dot(ZTw, y))
 
         return phi
@@ -383,4 +384,4 @@ def check_additivity(self,atol=1e-5):
         total = np.sum(self.primary,axis=1) + self.geo + np.sum(self.geo_intera,axis=1)
 
         print("Components add up to model prediction: ", 
-              np.allclose(total+self.base_value, self.predict_f(self.X_geo).reshape(-1), atol=atol))
+              np.allclose(total+self.base_value, self.predict_f(self.X_geo).reshape(-1), atol=atol))