<jupyter_start><jupyter_text>Generate Data<jupyter_code>def model(theta, x):
    """
    Creates the noise free observation at x
    """
    x_powers = x ** jnp.array([0, 1, 2])
    return jnp.dot(theta, x_powers)

def generate_data(
    theta,
    num_data,
    noise_mean,
    noise_var,
    x_seed=0,
    noise_seed=1
):
    """
    Generate data with IID gaussian noise with the quadratic model
    defined above.
    
    NOTE: the bounds of the x-value data generate are hard-coded.
    
    Parameters:
    -----------
        theta      (np arr) : true parameter value
        num_data   (int)    : number of data draws
        noise_mean (float)  : use 0 for unbiased noise
        noise_var  (float)  : noise variance
        x_seed     (int)    : random seed for generating x values
        noise_seed (int)    : random seed for noise generation
    
    Returns:
    --------
        x_s (np arr) : array of sampled x values
        y   (np arr) : array of outputs (func + noise)
    
    """
    # random x-values
    np.random.seed(x_seed)
    x_s = stats.uniform(loc=0, scale=5).rvs(num_data)
    
    # generate noise
    np.random.seed(noise_seed)
    noise = stats.norm(loc=noise_mean, scale=np.sqrt(noise_var)).rvs(num_data)
    
    # map points according to model and add noise
    y = np.zeros_like(x_s)
    for i in range(num_data):
        y[i] = model(x=x_s[i], theta=theta) + noise[i]
        
    return x_s, y
# generate data -- using values from overleaf
THETA = np.array([1, 0.5, 1.0])
NOISE_MEAN = 0
SIGMA2 = 0.1
NUM_DATA = 100

# generate data
x_s, y = generate_data(
    theta=THETA,
    num_data=NUM_DATA,
    noise_mean=NOISE_MEAN,
    noise_var=SIGMA2
)
plt.figure(figsize=(12.5,6))
plt.scatter(x_s, y)
plt.xlabel('x')
plt.ylabel('y')
plt.show()<jupyter_output><empty_output><jupyter_text>Optimize with identity map as the QoI<jupyter_code># define the QoI to simply be the identity map
id_map = lambda theta: theta
# set optimization parameters
EPSILON = 0.0001                   # stopping criterion for odad_min_e_ball algo
CONF_LEV = 0.95                   # 1 - beta_alpha - i.e., prob not violating
THETA_INIT = "multisample"   # initialization startgey
THETA_DIMS = 3               # dimensionality of theta
MU = 1000                          # strength of penalty
MAX_EPOCH = 100000
MIN_IMPROVMENT = 0.01
PATINCE = 1000
LR = 0.001      
forward_model = lambda theta: jax.vmap(model, in_axes=[None, 0])(theta, x_s)
# create objective function and optimizer objects
objective_obj = DifferentaibleMeritFunc(
    forward_model=forward_model, 
    mu=MU, 
    data=y,
    qoi_func=id_map
)
optimizer = GdOpt(objective=objective_obj)

mle_theta, M_alpha, S, center, radius_0 = optimizer.optimize_min_e_ball(
    sigma_2=SIGMA2,
    data=y,
    initial_theta=THETA_INIT,
    theta_dims=THETA_DIMS,
    raduis_eps=EPSILON,
    conf_level=CONF_LEV,
    max_epoch=MAX_EPOCH,
    lr=LR  
)

print("----- Center and Radius -----")
print(S)
print(center)
print(radius_0)

# perform optimization to find dirac weights
p_opt = optimizer.weight_optimization(S)
print("----- Dirac Weights -----")
print(p_opt)<jupyter_output>----- Center and Radius -----
[[0.84085654 0.68584845 0.96372252]
 [1.31100809 0.2392041  1.04450683]
 [0.84698032 0.69250884 0.96099987]
 [0.84578468 0.69170232 0.96138787]]
[1.07839638 0.46545321 1.00294735]
0.3271453130068976
----- Dirac Weights -----
[6.19632069e-23 5.00000000e-01 7.57937668e-23 5.00000000e-01]<jupyter_text>Optimize with a 1d Projection<jupyter_code># define a projection for the QoI
proj_func = jax.jit(lambda theta: theta[0])
# set optimization parameters
EPSILON = 0.0001                   # stopping criterion for odad_min_e_ball algo
CONF_LEV = 0.95                   # 1 - beta_alpha - i.e., prob not violating
THETA_INIT = "multisample"   # initialization startgey
THETA_DIMS = 3               # dimensionality of theta
MU = 1000                          # strength of penalty
MAX_EPOCH = 100000
MIN_IMPROVMENT = 0.01
PATINCE = 1000
LR = 0.001      
# create objective function and optimizer objects
objective_obj = DifferentaibleMeritFunc(
    forward_model=forward_model, 
    mu=MU, 
    data=y,
    qoi_func=proj_func
)
optimizer = GdOpt(objective=objective_obj)

mle_theta, M_alpha, S, center, radius_0 = optimizer.optimize_min_e_ball(
    sigma_2=SIGMA2,
    data=y,
    initial_theta=THETA_INIT,
    theta_dims=THETA_DIMS,
    raduis_eps=EPSILON,
    conf_level=CONF_LEV,
    max_epoch=MAX_EPOCH,
    lr=LR,     
)

print("----- Center and Radius -----")
print(S)
print(center)
print(radius_0)

# perform optimization to find dirac weights
p_opt = optimizer.weight_optimization(S=S.astype(np.float64))
print("----- Dirac Weights -----")
print(p_opt)<jupyter_output>----- Center and Radius -----
[0.83506882 1.32197781]
[1.07852332]
0.24345449476849773
----- Dirac Weights -----
[0.5 0.5]<jupyter_text>Optimize with $y$ value as the QoI<jupyter_code>x_test = 10
y_true = model(THETA, x_test)
(x_test, y_true)
predict = lambda theta: model(theta, x_test)
# create objective function and optimizer objects
objective_obj = DifferentaibleMeritFunc(
    forward_model=forward_model, 
    mu=MU, 
    data=y,
    qoi_func=predict
)
optimizer = GdOpt(objective=objective_obj)

mle_theta, M_alpha, S, center, radius_0 = optimizer.optimize_min_e_ball(
    sigma_2=SIGMA2,
    data=y,
    initial_theta=THETA_INIT,
    theta_dims=THETA_DIMS,
    raduis_eps=EPSILON,
    conf_level=CONF_LEV,
    max_epoch=MAX_EPOCH,
    lr=LR    
)

print("----- Center and Radius -----")
print(S)
print(center)
print(radius_0)

# perform optimization to find dirac weights
p_opt = optimizer.weight_optimization(S)
print("----- Dirac Weights -----")
print(p_opt)<jupyter_output>----- Center and Radius -----
[108.64814788 103.36511579]
[106.00663184]
2.641516047316081
----- Dirac Weights -----
[0.5 0.5]
{'fraction_non_alphanumeric': 0.11713554987212277, 'fraction_numerical': 0.06871270247229326, 'mean_word_length': 3.0399449035812673, 'pattern_counts': {'":': 0, '<': 14, '<?xml version=': 0, '>': 14, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 2, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '4', 'dir': 'jupyter-structured-clean-dedup', 'id': '441983', 'license': 'MIT', 'n_tokens_mistral': 2467, 'n_tokens_neox': 2146, 'n_words': 494, 'path': 'examples/demo_quadratic_model_gradient.ipynb', 'repo_name': 'JPLMLIA/UQ4K', 'size': '5865'}