Update mx.py

math-econ-code · Dec 31, 2024 · b652595 · b652595
1 parent 52cd8b2
commit b652595
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/mec/mx.py b/mec/mx.py
@@ -3,14 +3,18 @@
 import numpy as np, scipy.sparse as sp, pandas as pd
 
 
-def ivgmm(Y_i,X_i_k,Z_i_l,W_l_l=None):
+def ivgmm(Y_i,X_i_k,Z_i_l,W_l_l=None, require_der = 0):
     I = len(Y_i)
     if W_l_l is None:
         W_l_l = np.linalg.inv( Z_i_l.T @ Z_i_l /  I )
     Pi_i_i =  Z_i_l @ W_l_l @ Z_i_l.T
     beta_k = np.linalg.solve(X_i_k.T @ Pi_i_i @ X_i_k,  X_i_k.T @ Pi_i_i @ Y_i)
     val = (Y_i - X_i_k @ beta_k).T @ Pi_i_i @ (Y_i - X_i_k @ beta_k) / (2* I * I)
-    return beta_k,val
+    if require_der > 0:
+        dval_dY_i = (Pi_i_i @ (Y_i - X_i_k @ beta_k) / (I*I)).flatten()
+        return beta_k,val,dval_dY_i
+    else:
+        return beta_k,val
 
 def efficient_ivgmm(Y_i,X_i_k,Z_i_l, centering = True):
     I=len(Y_i)