DeepIV Value Overflow

The following is a slightly modification to the DeepIV notebook, but it no longer works. I guess it is because of the fact that I increased the max value of the X and T, which caused the overflow somewhere:

from econml.deepiv import DeepIVEstimator
from econml.bootstrap import BootstrapEstimator
import keras
import numpy as np
import matplotlib.pyplot as plt

get_ipython().run_line_magic('matplotlib', 'inline')


# ## Synthetic data
# 
# To demonstrate the Deep IV approach, we'll construct a synthetic dataset obeying the requirements set out above.  In this case, we'll take `X`, `Z`, `T` to come from the following distribution: 

# In[ ]:


# Initialize exogenous variables; normal errors, uniformly distributed covariates and instruments
e = np.random.normal(size=(445,1))
x = np.random.uniform(low=0.0, high=1000, size=(445,1))
z = np.random.uniform(low=0.0, high=1000, size=(445,1))

# Initialize treatment variable
t = np.sqrt((x+2) * z) + e

# Show the marginal distribution of t
plt.hist(t)
plt.xlabel("t")
plt.show()

plt.scatter(z[x < 1], t[x < 1], label='low X')
plt.scatter(z[(x > 4.5) * (x < 5.5)], t[(x > 4.5) * (x < 5.5)], label='moderate X')
plt.scatter(z[x > 9], t[x > 9], label='high X')
plt.legend()
plt.xlabel("z")
plt.ylabel("t")
plt.show()


# Here, we'll imagine that `Z` and `X` are causally affecting `T`; as you can see in the plot above, low or high values of `Z` drive moderate values of `T` and moderate values of `Z` cause `T` to have a bi-modal distribution when `X` is high, but a unimodal distribution centered on 0 when `X` is low.  The instrument is positively correlated with the treatment and treatments tend to be bigger at high values of x. The instrument has higher power at higher values of x 

# In[ ]:


# Outcome equation 
y = t*t / 10 - x*t / 10 + e

# The endogeneity problem is clear, the latent error enters both treatment and outcome equally
plt.scatter(t,z, label ='raw data')
tticks = np.arange(-2,12)
yticks2 = tticks*tticks/10 - 0.2 * tticks
yticks5 = tticks*tticks/10 - 0.5 * tticks
yticks8 = tticks*tticks/10 - 0.8 * tticks
plt.plot(tticks,yticks2, 'r--', label = 'truth, x=2')
plt.plot(tticks,yticks5, 'g--', label = 'truth, x=5')
plt.plot(tticks,yticks8, 'y--', label = 'truth, x=8')
plt.xlabel("t")
plt.ylabel("y")
plt.legend()
plt.show()


# `Y` is a non-linear function of `T` and `X` with no direct dependence on `Z` plus additive noise (as required).  We want to estimate the effect of particular `T` and `X` values on `Y`.
# 
# The plot makes it clear that looking at the raw data is highly misleading as to the treatment effect. Moreover the treatment effects are both non-linear and heterogeneous in x, so this is a hard problem!
# 
# ## Defining the neural network models
# 
# Now we'll define simple treatment and response models using the Keras `Sequential` model built up of a series of layers.  Each model will have an `input_shape` of 2 (to match the sums of the dimensions of `X` plus `Z` in the treatment case and `T` plus `X` in the response case).

# In[ ]:


treatment_model = keras.Sequential([keras.layers.Dense(128, activation='relu', input_shape=(2,)),
                                    keras.layers.Dropout(0.17),
                                    keras.layers.Dense(64, activation='relu'),
                                    keras.layers.Dropout(0.17),
                                    keras.layers.Dense(32, activation='relu'),
                                    keras.layers.Dropout(0.17)])

response_model = keras.Sequential([keras.layers.Dense(128, activation='relu', input_shape=(2,)),
                                   keras.layers.Dropout(0.17),
                                   keras.layers.Dense(64, activation='relu'),
                                   keras.layers.Dropout(0.17),
                                   keras.layers.Dense(32, activation='relu'),
                                   keras.layers.Dropout(0.17),
                                   keras.layers.Dense(1)])


# Now we'll instantiate the `DeepIVEstimator` class using these models.  Defining the response model *outside* of the lambda passed into constructor is important, because (depending on the settings for the loss) it can be used multiple times in the second stage and we want the same weights to be used every time.

# In[ ]:


keras_fit_options = { "epochs": 30,
                      "validation_split": 0.1,
                      "callbacks": [keras.callbacks.EarlyStopping(patience=2, restore_best_weights=True)]}

deepIvEst = DeepIVEstimator(n_components = 10, # number of gaussians in our mixture density network
                            m = lambda z, x : treatment_model(keras.layers.concatenate([z,x])), # treatment model
                            h = lambda t, x : response_model(keras.layers.concatenate([t,x])),  # response model
                            n_samples = 1, # number of samples to use to estimate the response
                            use_upper_bound_loss = False, # whether to use an approximation to the true loss
                            n_gradient_samples = 1, # number of samples to use in second estimate of the response (to make loss estimate unbiased)
                            optimizer='adam', # Keras optimizer to use for training - see https://keras.io/optimizers/ 
                            first_stage_options = keras_fit_options, # options for training treatment model
                            second_stage_options = keras_fit_options) # options for training response model


# ## Fitting and predicting using the model
# Now we can fit our model to the data:

# In[ ]:


boot_est = BootstrapEstimator(deepIvEst, n_bootstrap_samples=2, n_jobs=1)
boot_est.fit(Y=y,T=t,X=x,Z=z)


# And now we can create a new set of data and see whether our predicted effect matches the true effect `T*T-X*X`:

# In[ ]:


n_test = 500
for i, x_enum in enumerate([2, 5, 8]):
    t = np.linspace(0,10,num = 100)
    y_true = t*t / 10 - x_enum*t/10
    y_pred = boot_est.predict(t, np.full_like(t, x_enum))
    plt.plot(t, y_true, label='true y, x={0}'.format(x_enum),color='C'+str(i))
    plt.plot(t, y_pred, label='pred y, x={0}'.format(x_enum),color='C'+str(i),ls='--')
plt.xlabel('t')
plt.ylabel('y')
plt.legend()
plt.show()


# You can see that despite the fact that the response surface varies with x, our model was able to fit the data reasonably well. Where is does worst is where the instrument has the least power, which is in the low x case.  There it fits a straight line rather than a quadratic, which suggests that the regularization at least is perfoming well.  

# ## Estimating the Confidence Interval of the Model

# In[ ]:


treatment_effect_interval = boot_est.marginal_effect_interval(X=x, T=t, lower=1, upper=99)

Could you help and investigate why?

Start by running the provided DeepIV notebook reproduction with X and Z ranges set to 0–1000, then capture the exact overflow and the stage where it occurs. Trace that failure through DeepIVEstimator and its treatment and response models; done means the cause is isolated and a concrete correction or documented limitation is validated on the reproduction.