### Check TensorFlow Version

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Prints the installed TensorFlow version. This is a common check to ensure compatibility with other libraries and code.

```python
# Check Versions
import tensorflow as tf
print(tf.__version__)
```

--------------------------------

### Import Libraries

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample_minimal.ipynb

Imports necessary libraries including NumPy, Matplotlib, Keras, TensorFlow, and OmniFold. This is a standard setup for machine learning tasks.

```python
import numpy as np
from matplotlib import pyplot as plt

from keras.layers import Dense, Input
from keras.models import Model

import omnifold as of
import os
import tensorflow as tf
```

--------------------------------

### Prepare Data for Classification (Step 1)

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Concatenates simulation and data with background for the first classification step.

```python
xvals_1 = np.concatenate((x_sim, x_data_withback))
yvals_1 = np.concatenate((np.zeros(len(x_sim)), np.ones(len(x_data_withback))))
```

--------------------------------

### Initialize Weights

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Initializes weights for pull and push operations to ones.

```python
weights_pull = np.ones(len(x_sim))
weights_push = np.ones(len(x_sim))
```

--------------------------------

### Plotting Initial Distributions

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample_minimal.ipynb

Visualizes the generated synthetic data and the 'natural' data distributions using Matplotlib. This helps in understanding the initial state of the data before correction.

```python
_,_,_=plt.hist(theta0_G,bins=np.linspace(-3,3,20),color='blue',alpha=0.5,label="MC, true")
_,_,_=plt.hist(theta0_S,bins=np.linspace(-3,3,20),histtype="step",color='black',ls=':',label="MC, reco")
_,_,_=plt.hist(theta_unknown_G,bins=np.linspace(-3,3,20),color='orange',alpha=0.5,label="Data, true")
_,_,_=plt.hist(theta_unknown_S,bins=np.linspace(-3,3,20),histtype="step",color='black',label="Data, reco")
plt.xlabel("x")
plt.ylabel("events")
plt.legend(frameon=False)
```

--------------------------------

### Plot Full Data/Simulation with Background

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Visualizes the combined simulated data (signal + background) and real data distributions, including background components. Aids in comparing simulation to observed data.

```python
#Now the full data/sim. including background
binning = np.linspace(-5,5,20)
ns,_,_=plt.hist(x_sim_withback[x_sim_withback!=dummyval],weights=w_synth_withback[x_sim_withback!=dummyval],bins=binning,alpha=0.5,label="Sim. w/ back")
_,_,_=plt.hist(x_background,bins=binning,weights=w_back,alpha=0.5,label="Sim. back.")
dplot=np.histogram(x_data_withback[x_data_withback!=dummyval],bins=binning)
plt.plot(0.5*(dplot[1][1:]+dplot[1][:-1]),dplot[0],label="Data w/ back",ls="",marker='o',color="black")
_,_,_=plt.hist(x_background_real,bins=binning,alpha=0.5,label="Data back.",histtype="step",color="black")
plt.legend(frameon=False,loc="upper left")
plt.ylabel("Number of events")
plt.xlabel("x")
plt.ylim([0,1.8*np.max(ns)])
```

--------------------------------

### Prepare Data for Classification (Step 2)

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Concatenates generated data for the second classification step.

```python
xvals_2 = np.concatenate((x_gen, x_gen))
yvals_2 = np.concatenate((np.zeros(len(x_gen)), np.ones(len(x_gen))))
```

--------------------------------

### Define Gaussian Parameters and Efficiencies

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Sets the mean and standard deviation for the initial Gaussian distribution, along with parameters for 'eff' (true but not reco), 'fake' (reco but not true), and 'back' (background fraction). These parameters are crucial for simulating and unfolding data.

```python
mu0 = 0
sigma0 = 1

eff = 0.1 #fraction of true but not reco
fake = 0.1 #fraction of reco but not true

back = 0.1 #fraction of a background process that we would like to subtract
```

--------------------------------

### Simulate Gaussian Data with Smearing and Selections

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Generates synthetic and natural Gaussian data, applies smearing effects, and simulates selection criteria. Used for setting up simulation studies.

```python
#param = (mu, sigma)
theta0_param = (mu0, sigma0)  # synthetic sample
theta_unknown_param = (0.2, 0.8)  # this is the data (the target)
background_param = (0, 1.2)  # background
epsilon = sigma0 / 2.  # Smearing width

dummyval = -10 #a value for examples that don't pass one of the measured/ideal selections

#Synthetic
x_gen = np.random.normal(theta0_param[0], theta0_param[1],N)
x_sim = np.array([(x + np.random.normal(0, epsilon)) for x in x_gen])
pass_reco = np.random.binomial(1,1.-eff,len(x_gen))
pass_truth = np.random.binomial(1,1.-fake,len(x_gen))
x_gen[pass_truth==0] = dummyval
x_sim[pass_reco==0] = dummyval
w_synth = np.ones(len(x_gen))
w_synth = w_synth*len(w_synth)/np.sum(w_synth)

x_background = np.random.normal(background_param[0],background_param[1], int(N*back))
w_back = reweighting(x_background,background_param[0],background_param[0],background_param[1],background_param[1]*0.8)
w_back = w_back*len(w_back)/np.sum(w_back) #pick weights so that they don't change sum (N.B. this changes statistical power, but I don't care for this example)
x_sim_withback = np.concatenate([x_sim,x_background]) #sim with background added
x_gen_withback= np.concatenate([x_gen,-np.ones(int(N*back))*dummyval]) #gen with background added; since background has no corresponding truth in the fiducial volume, set to dummy value
w_synth_withback = np.concatenate([w_synth,w_back])

x_synth = np.stack([x_gen, x_sim], axis=1)
y_synth = np.zeros(len(x_synth))

#Natural
x_truth = np.random.normal(theta_unknown_param[0],theta_unknown_param[1], N) # Nature, particle-level analog
x_data = np.array([(x + np.random.normal(0, epsilon)) for x in x_truth]) # Measured Data analog
pass_reco = np.random.binomial(1,1.-eff,len(x_truth))
pass_truth = np.random.binomial(1,1.-fake,len(x_truth))
x_truth[pass_truth==0] = dummyval
x_data[pass_reco==0] = dummyval #emulates cuts done in an analysis

x_background_real = np.random.normal(background_param[0],background_param[1]*0.8, int(N*back))
x_data_withback = np.concatenate([x_data,x_background_real]) #Data+Background
x_truth_withback = np.concatenate([x_truth,-np.ones(int(N*back))*dummyval]) #Nature+background

x_natural = np.stack([x_truth, x_data], axis=1)
y_natural = np.ones(len(x_natural))
```

--------------------------------

### Import Libraries

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Imports necessary libraries: numpy for numerical operations, energyflow for data handling, and matplotlib for plotting.

```python
import numpy as np
import energyflow as ef
import matplotlib.pyplot as plt
```

--------------------------------

### Visualize Gaussian Distributions with Reweighting

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Generates and plots histograms for two normal distributions. It also applies reweighting to a third distribution to match a target PDF, visualizing the effect of the reweighting function.

```python
xx = np.linspace(-5,5,20)
plt.hist(np.random.normal(0,1,10000),bins=xx,density=True)
plt.hist(np.random.normal(0,2,10000),bins=xx,alpha=0.5,density=True)
vals = np.random.normal(0,2,10000)
plt.hist(vals,bins=xx,weights=reweighting(vals,0,0,2,1),density=True,histtype="step",color="black")
```

--------------------------------

### Plot Weighted Background Simulation

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Compares the distribution of simulated background events with and without reweighting, against real background events. Useful for validating background modeling.

```python
#Check that I did the background weighting properly.
binning = np.linspace(-5,5,20)
ns,_,_=plt.hist(x_background,bins=binning,alpha=0.5,label="Sim. back., unweighted")
_,_,_=plt.hist(x_background,bins=binning,weights=w_back,alpha=0.5,label="Sim. back., weighted")
_,_,_=plt.hist(x_background_real,bins=binning,histtype="step",color="black",label="Real back.")
plt.legend(frameon=False,loc="upper left")
plt.ylabel("Number of events")
plt.xlabel("x")
plt.ylim([0,1.8*np.max(ns)])
```

--------------------------------

### Prepare Data for Background Subtraction

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Concatenates background and data with background, assigning appropriate weights for training a classification model. This prepares the dataset for distinguishing between signal and background events.

```python
x_data_and_MCback = np.concatenate([x_background[x_background!=dummyval],
                                    x_data_withback[x_data_withback!=dummyval],
                                    x_data_withback[x_data_withback!=dummyval]])
    
y_data_and_MCback = np.concatenate([np.ones(len(x_background[x_background!=dummyval])),
                                    np.ones(len(x_data_withback[x_data_withback!=dummyval])),
                                    np.zeros(len(x_data_withback[x_data_withback!=dummyval]))])
    
W_data_and_MCback = np.concatenate([-1.*w_back[x_background!=dummyval],
                                    np.ones(len(x_data_withback[x_data_withback!=dummyval])),
                                    np.ones(len(x_data_withback[x_data_withback!=dummyval]))])

X_train_1, X_test_1, Y_train_1, Y_test_1, w_train_1, w_test_1 = train_test_split(
    x_data_and_MCback, y_data_and_MCback, W_data_and_MCback)
```

--------------------------------

### GPU Configuration

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample_minimal.ipynb

Lists available GPUs and optionally configures memory limits for TensorFlow. This is useful for managing computational resources.

```python
gpus = tf.config.experimental.list_physical_devices('GPU')
print(gpus)
#tf.config.experimental.set_virtual_device_configuration(
#    gpus[1],[tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)]) #in MB
```

--------------------------------

### Configure GPU Memory Limit

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Lists available GPUs and sets a virtual device configuration to limit the memory usage of the first GPU to 1024 MB. This is useful for managing resources on systems with multiple GPUs or limited memory.

```python
gpus = tf.config.experimental.list_physical_devices('GPU')
print(gpus)
tf.config.experimental.set_virtual_device_configuration(gpus[0], [tf.config.experimental.VirtualDeviceConfiguration(memory_limit=1024)]) #in MB
```

--------------------------------

### Train Neural Network for Background Subtraction

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Trains the compiled neural network model using the prepared training data, incorporating sample weights and early stopping. Validation data is used to monitor performance during training.

```python
model.fit(X_train_1,
          Y_train_1,
          sample_weight=w_train_1,
          epochs=200,
          batch_size=10000,
          validation_data=(X_test_1, Y_test_1, w_test_1),
          callbacks=[earlystopping],
          verbose=0)
```

--------------------------------

### Visualize Background-Subtracted Data

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Generates a histogram comparing simulated data with background, simulated background, and actual data (with and without weights). This visualization helps verify the effectiveness of the background subtraction and reweighting process.

```python
#Now the full data/sim. including background
binning = np.linspace(-5,5,20)
ns,_,_=plt.hist(x_sim_withback[x_sim_withback!=dummyval],weights=w_synth_withback[x_sim_withback!=dummyval],bins=binning,alpha=0.5,label="Sim. w/ back")
_,_,_=plt.hist(x_background,bins=binning,weights=w_back,alpha=0.5,label="Sim. back.")
dplot=np.histogram(x_data_withback[x_data_withback!=dummyval],bins=binning)
plt.plot(0.5*(dplot[1][1:]+dplot[1][:-1]),dplot[0],label="Data w/ back",ls="",marker='o',color="black")

dplot2=np.histogram(x_data[x_data!=dummyval],bins=binning)
plt.plot(0.5*(dplot2[1][1:]+dplot2[1][:-1]),dplot2[0],label="Data w/o back",ls="",marker='o',color="black",markerfacecolor="none")

_,_,_=plt.hist(x_background_real,bins=binning,alpha=0.5,label="Data back.",histtype="step",color="black")
_,_,_=plt.hist(x_data_withback[x_data_withback!=dummyval],weights=w_data[x_data_withback!=dummyval],bins=binning,alpha=0.5,label="Data w/ weights",histtype="step",color="black",ls=":")
plt.legend(frameon=False,loc="upper left",ncols=2)
plt.ylabel("Number of events")
plt.xlabel("x")
plt.ylim([0,1.8*np.max(ns)])
```

--------------------------------

### Initialize Weights for Unfolding

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Initializes an empty numpy array to store weights for unfolding iterations. This array is structured to hold weights for each step and event across multiple unfolding iterations.

```python
iterations = 4

weights = np.empty(shape=(iterations, 2, len(x_gen)))
# shape = (iteration, step, event)
```

--------------------------------

### Load Zjets Data

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Loads Zjets data using the energyflow library. It specifies the data source, number of samples, padding option, and keys to exclude.

```python
zjets = ef.zjets_delphes.load('Pythia21', num_data=1000000, pad=False, exclude_keys=['particles','gen_particles'])
```

--------------------------------

### Plot Generated Jet Width Distribution

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the width distribution of generated jets between the original paper's data and the updated data.

```python
_=plt.hist(zjets['gen_widths'],bins=np.linspace(0,0.5,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['gen_widths'],bins=np.linspace(0,0.5,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Gen. jet width")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Process Custom Data from File

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Parses a text file ('test_Omni.txt') to extract and store jet properties for both generated ('truth') and simulated ('reco') events, filtering for Z pT greater than 200 GeV. Handles potential NaN values in zg.

```python
gen_Zs = []
gen_mults = []
gen_jets = []
gen_widths = []
gen_tau2s = []
gen_zgs = []
gen_sdms = []

sim_mults = []
sim_jets = []
sim_widths = []
sim_tau2s = []
sim_zgs = []
sim_sdms = []
for line in open("test_Omni.txt"):
    if 'truth' in line:
        ZpT = float(line.split()[2])
        if (ZpT > 200):
            gen_Zs+=[[ZpT,0.,0.]]
            gen_mults+=[int(line.split()[14])]
            gen_jets+=[[float(line.split()[3]),float(line.split()[4]),float(line.split()[5]),float(line.split()[6])]]
            gen_widths+=[float(line.split()[7])]
            gen_tau2s+=[float(line.split()[8])]
            zg = float(line.split()[13])
            if np.isnan(zg):
                zg = 0.
                pass
            gen_zgs+=[zg]
            gen_sdms+=[float(line.split()[12])]
            pass
        pass
    if 'reco' in line:
        ZpT = float(line.split()[2])
        if (ZpT > 200):
            sim_mults+=[int(line.split()[14])]
            sim_jets+=[[float(line.split()[3]),float(line.split()[4]),float(line.split()[5]),float(line.split()[6])]]
            sim_widths+=[float(line.split()[7])]
            sim_tau2s+=[float(line.split()[8])]
            zg = float(line.split()[13])
            if np.isnan(zg):
                zg = 0.
                pass
            sim_zgs+=[zg]
            sim_sdms+=[float(line.split()[12])]
            pass
        pass
    pass
zjets_bigger = {}
zjets_bigger['gen_Zs'] = np.array(gen_Zs)
zjets_bigger['gen_mults'] = np.array(gen_mults)
zjets_bigger['gen_jets'] = np.array(gen_jets)
zjets_bigger['gen_widths'] = np.array(gen_widths)
zjets_bigger['gen_tau2s'] = np.array(gen_tau2s)
zjets_bigger['gen_zgs'] = np.array(gen_zgs)
zjets_bigger['gen_sdms'] = np.array(gen_sdms)

zjets_bigger['sim_mults'] = np.array(sim_mults)
zjets_bigger['sim_jets'] = np.array(sim_jets)
zjets_bigger['sim_widths'] = np.array(sim_widths)
zjets_bigger['sim_tau2s'] = np.array(sim_tau2s)
zjets_bigger['sim_zgs'] = np.array(sim_zgs)
zjets_bigger['sim_sdms'] = np.array(sim_sdms)
```

--------------------------------

### Apply Reweighting to Data

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Applies the reweighting function to the background-subtracted data to obtain corrected event weights. This step is crucial for accurate subsequent analysis.

```python
w_data = reweight(x_data_withback)
```

--------------------------------

### Histogram of Simulated Jet Widths

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the distribution of simulated jet widths between the original paper's dataset and an updated dataset. Uses matplotlib for plotting.

```python
_=plt.hist(zjets['sim_widths'],bins=np.linspace(0,0.5,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['sim_widths'],bins=np.linspace(0,0.5,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Sim. jet width")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Plotting Generated, Truth, and Unfolded Data

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

This snippet visualizes the distributions of generated, truth, and unfolded data using matplotlib. It requires numpy and matplotlib to be imported. The data is binned using np.linspace, and weights are applied to the generated data.

```python
#Now the full data/sim. including background
binning = np.linspace(-5,5,20)
ns,_,_=plt.hist(x_gen[x_gen!=dummyval],weights=w_synth[x_gen!=dummyval],bins=binning,alpha=0.5,label="Gen.")
_,_,_=plt.hist(x_truth[x_truth!=dummyval],bins=binning,alpha=0.5,label="Truth")
_,_,_=plt.hist(x_gen[x_gen!=dummyval],weights=w_synth[x_gen!=dummyval]*weights[3, 1:2, :][0][x_gen!=dummyval],bins=binning,alpha=0.5,label="Unfolded",histtype="step",color="black")

plt.legend(frameon=False)
plt.ylabel("Number of events")
plt.xlabel("x")
plt.ylim([0,1.8*np.max(ns)])
```

--------------------------------

### Define and Compile Neural Network Model

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Defines a simple feedforward neural network with three hidden layers for binary classification. The model is compiled using binary crossentropy loss and the Adam optimizer, with early stopping to prevent overfitting.

```python
#I did not optimize this network at all; it should probably be quite good for simple problems.  Could also swap out with a BDT for low-dimensional/tabular data.
inputs = Input((1, ))
hidden_layer_1 = Dense(50, activation='relu')(inputs)
hidden_layer_2 = Dense(50, activation='relu')(hidden_layer_1)
hidden_layer_3 = Dense(50, activation='relu')(hidden_layer_2)
outputs = Dense(1, activation='sigmoid')(hidden_layer_3)

model = Model(inputs=inputs, outputs=outputs)

earlystopping = EarlyStopping(patience=10,
                              verbose=1,
                              restore_best_weights=True)
    
model.compile(loss='binary_crossentropy',
              optimizer='Adam',
              metrics=['accuracy'],
              weighted_metrics=[]) #to silence an annoying warning ... could compute the weighted accuracy
```

--------------------------------

### Histogram of Simulated Jet Groomed Mass

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the distribution of simulated jet groomed masses between the original paper's dataset and an updated dataset. Uses matplotlib for plotting.

```python
_=plt.hist(zjets['sim_sdms'],bins=np.linspace(0,50,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['sim_sdms'],bins=np.linspace(0,50,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Sim. jet groomed mass [GeV]")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Plot Simulated Jet pT Distribution

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the transverse momentum (pT) distribution of simulated jets between the original paper's data and the updated data.

```python
_=plt.hist(zjets['sim_jets'][:,3],bins=np.linspace(0,100,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['sim_jets'][:,3],bins=np.linspace(0,100,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Sim. jet $p_T$ [GeV]")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Iterative Weight Update Loop

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Performs iterative classification and weight updates for simulation and generation.

```python
for i in range(iterations):
    print("\nITERATION: {}\\n".format(i + 1))

    # STEP 1: classify Sim. (which is reweighted by weights_push) to Data
    # weights reweighted Sim. --> Data
    print("STEP 1\n")

    weights_1 = np.concatenate((weights_push*w_synth, w_data))

    X_train_1, X_test_1, Y_train_1, Y_test_1, w_train_1, w_test_1 = train_test_split(
        xvals_1, yvals_1, weights_1)

    model.compile(loss='binary_crossentropy',
                  optimizer='Adam',
                  metrics=['accuracy'],
                  weighted_metrics=[])
    model.fit(X_train_1[X_train_1!=dummyval],
              Y_train_1[X_train_1!=dummyval],
              sample_weight=w_train_1[X_train_1!=dummyval],
              epochs=200,
              batch_size=10000,
              validation_data=(X_test_1[X_test_1!=dummyval], Y_test_1[X_test_1!=dummyval], w_test_1[X_test_1!=dummyval]),
              callbacks=[earlystopping],
              verbose=0)

    weights_pull = weights_push * reweight(x_sim) 
    
    #One option is to take the prior:
    weights_pull[x_sim==dummyval] = 1. 
    
    #Another (more complicated) option is to assign the average weight: <w|x_true>.  To do this, we need to estimate this quantity.
    #xvals_1b = np.concatenate([x_gen[x_sim!=dummyval],x_gen[x_sim!=dummyval]])
    #yvals_1b = np.concatenate([np.ones(len(x_gen[x_sim!=dummyval])),np.zeros(len(x_gen[x_sim!=dummyval]))])
    #weights_1b = np.concatenate([weights_pull[x_sim!=dummyval],np.ones(len(x_gen[x_sim!=dummyval]))])
    
    #X_train_1b, X_test_1b, Y_train_1b, Y_test_1b, w_train_1b, w_test_1b = train_test_split(
    #    xvals_1b, yvals_1b, weights_1b)    
    
    #model.compile(loss='binary_crossentropy',
    #              optimizer='Adam',
    #              metrics=['accuracy'])
    #model.fit(X_train_1b,
    #          Y_train_1b,
    #          sample_weight=w_train_1b,
    #          epochs=200,
    #          batch_size=10000,
    #          validation_data=(X_test_1b, Y_test_1b, w_test_1b),
    #          callbacks=[earlystopping],
    #          verbose=1)
    
    #average_vals = reweight(x_gen[x_sim==dummyval])
    #weights_pull[x_sim==dummyval] = average_vals
    ###
    
    weights[i, :1, :] = weights_pull

    # STEP 2: classify Gen. to reweighted Gen. (which is reweighted by weights_pull)
    # weights Gen. --> reweighted Gen.
    print("\nSTEP 2\n")

    weights_2 = np.concatenate((w_synth, weights_pull*w_synth))

    X_train_2, X_test_2, Y_train_2, Y_test_2, w_train_2, w_test_2 = train_test_split(
        xvals_2, yvals_2, weights_2)

    model.compile(loss='binary_crossentropy',
                  optimizer='Adam',
                  metrics=['accuracy'],
                  weighted_metrics=[])
    model.fit(X_train_2,
              Y_train_2,
              sample_weight=w_train_2,
              epochs=200,
              batch_size=2000,
              validation_data=(X_test_2, Y_test_2, w_test_2),
              callbacks=[earlystopping],
              verbose=0)

    weights_push = reweight(x_gen)
    ###
    #Need to do something with events that don't pass truth    
    
    #One option is to take the prior:
    weights_push[x_gen==dummyval] = 1. 
    
    #Another option is to assign the average weight: <w|x_reco>.  To do this, we need to estimate this quantity.
    #xvals_1b = np.concatenate([x_sim[x_gen!=dummyval],x_sim[x_gen!=dummyval]])
    #yvals_1b = np.concatenate([np.ones(len(x_sim[x_gen!=dummyval])),np.zeros(len(x_sim[x_gen!=dummyval]))])
    #weights_1b = np.concatenate([weights_push[x_gen!=dummyval],np.ones(len(x_sim[x_gen!=dummyval]))])
    
    #X_train_1b, X_test_1b, Y_train_1b, Y_test_1b, w_train_1b, w_test_1b = train_test_split(
    #    xvals_1b, yvals_1b, weights_1b)    
    
    #model.compile(loss='binary_crossentropy',
    #              optimizer='Adam',
    #              metrics=['accuracy'])
    #model.fit(X_train_1b,
    #          Y_train_1b,
    #          sample_weight=w_train_1b,
    #          epochs=200,
    #          batch_size=10000,
    #          validation_data=(X_test_1b, Y_test_1b, w_test_1b),
    #          callbacks=[earlystopping],
    #          verbose=1)
    
    #average_vals = reweight(x_sim[x_gen==dummyval])
    #weights_push[x_gen==dummyval] = average_vals
    ###    
    
    weights[i, 1:2, :] = weights_push
```

--------------------------------

### Histogram of Generated Jet Groomed Mass

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the distribution of generated jet groomed masses between the original paper's dataset and an updated dataset. Uses matplotlib for plotting.

```python
_=plt.hist(zjets['gen_sdms'],bins=np.linspace(0,50,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['gen_sdms'],bins=np.linspace(0,50,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Gen. jet groomed mass [GeV]")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Plot Simulated Jet Multiplicity

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the distribution of the number of constituents in simulated jets between the original paper's data and the updated data.

```python
_=plt.hist(zjets['sim_mults'],bins=np.linspace(0,100,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['sim_mults'],bins=np.linspace(0,100,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Number of constituents in sim. jet")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Import Libraries for OmniFold

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Imports necessary libraries including numpy, matplotlib, scikit-learn, and TensorFlow for data manipulation, visualization, and machine learning tasks.

```python
import numpy as np
from matplotlib import pyplot as plt

from sklearn.model_selection import train_test_split
import tensorflow as tf
from tensorflow.keras.layers import Dense, Input
from tensorflow.keras.models import Model
from tensorflow.keras.callbacks import EarlyStopping

from scipy.stats import norm
```

--------------------------------

### Set Number of Iterations

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Defines the total number of iterations for the process.

```python
iterations = 4
```

--------------------------------

### Plot Generated Jet pT Distribution

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the transverse momentum (pT) distribution of generated jets between the original paper's data and the updated data.

```python
_=plt.hist(zjets['gen_jets'][:,3],bins=np.linspace(0,100,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['gen_jets'][:,3],bins=np.linspace(0,100,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Gen. jet $p_T$ [GeV]")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Histogram of Simulated Jet Zg

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the distribution of simulated jet zg values between the original paper's dataset and an updated dataset. Uses matplotlib for plotting.

```python
_=plt.hist(zjets['sim_zgs'],bins=np.linspace(0,0.5,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['sim_zgs'],bins=np.linspace(0,0.5,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Sim. jet $z_g$")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Plot Generated Jet Multiplicity

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the distribution of the number of constituents in generated jets between the original paper's data and the updated data.

```python
_=plt.hist(zjets['gen_mults'],bins=np.linspace(0,100,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['gen_mults'],bins=np.linspace(0,100,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Number of constituents in gen. jet")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Define Reweighting Function

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Defines a function `reweighting` that calculates the ratio of two Gaussian probability density functions (PDFs). This is used to reweight samples from one distribution to match another, essential for tasks like importance sampling or correcting distributions.

```python
def reweighting(xx,mu0,mu1,sigma0,sigma1):
    starting = norm.pdf((xx-mu0)/sigma0)
    target = norm.pdf((xx-mu1)/sigma1)
    return target/starting
```

--------------------------------

### Histogram of Simulated Jet Tau21

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the distribution of simulated jet tau21 values between the original paper's dataset and an updated dataset. Uses matplotlib for plotting.

```python
_=plt.hist(zjets['sim_tau2s'],bins=np.linspace(0,0.5,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['sim_tau2s'],bins=np.linspace(0,0.5,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Sim. jet $\tau_{21}$")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Set Number of Events

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Sets the total number of events (N) to 100,000. This variable is likely used for generating synthetic data in subsequent steps of the simulation or analysis.

```python
N = 10**5
```

--------------------------------

### Plot Truth Z pT Distribution

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Generates and displays a histogram comparing the truth Z pT distribution from the original paper's data and the updated data.

```python
plt.hist(zjets['gen_Zs'][:,0],bins=np.linspace(200,500,20),density=True,label="Original Paper",alpha=0.5)
plt.hist(zjets_bigger['gen_Zs'][:,0],bins=np.linspace(200,500,20),density=True,histtype="step",color="black",label="Updated")
plt.xlabel("Truth Z $p_T$ [GeV]")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Define Reweighting Function

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample.ipynb

Creates a function to reweight events based on the trained model's predictions. This function calculates the ratio of signal to background probabilities, essential for correcting event weights.

```python
def reweight(events):
    f = model.predict(events, batch_size=10000)
    weights = f / (1. - f) #for binary cross entropy - see refs. above and https://arxiv.org/abs/2101.07263 for more info.
    return np.squeeze(np.nan_to_num(weights))
```

--------------------------------

### Plotting OmniFolded Results

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample_minimal.ipynb

Visualizes the OmniFolded (corrected) data distribution alongside the original MC and data distributions using Matplotlib. This shows the effect of the correction.

```python
_,_,_=plt.hist(theta0_G,bins=np.linspace(-3,3,20),color='blue',alpha=0.5,label="MC, true")
_,_,_=plt.hist(theta_unknown_G,bins=np.linspace(-3,3,20),color='orange',alpha=0.5,label="Data, true")
_,_,_=plt.hist(theta0_G,weights=myweights[-1, 0, :], bins=np.linspace(-3,3,20),color='black',histtype="step",label="OmniFolded",lw="2")
plt.xlabel("x")
plt.ylabel("events")
plt.legend(frameon=False)
```

--------------------------------

### Histogram of Generated Jet Zg

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the distribution of generated jet zg values between the original paper's dataset and an updated dataset. Uses matplotlib for plotting.

```python
_=plt.hist(zjets['gen_zgs'],bins=np.linspace(0,0.5,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['gen_zgs'],bins=np.linspace(0,0.5,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Gen. jet $z_g$")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Generate Synthetic Data

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample_minimal.ipynb

Creates synthetic data for both generator-level (true) and detector-level (reconstructed) distributions using NumPy. This data is used to train and test the OmniFold model.

```python
N = 10**5

#Synthetic
theta0_G = np.random.normal(0.2,0.8,N)  # Generator-level synthetic sample
theta0_S = np.array([(x + np.random.normal(0, 0.5)) for x in theta0_G])  # Detector smearing for synthetic sample

theta0 = np.stack([theta0_G, theta0_S], axis=1)

#Natural
theta_unknown_G = np.random.normal(0,1, N)
theta_unknown_S = np.array([(x + np.random.normal(0, 0.5)) for x in theta_unknown_G]) 
```

--------------------------------

### Histogram of Generated Jet Tau21

Source: https://github.com/hep-lbdl/omnifold/blob/main/Benchmark/CheckOmniFold.ipynb

Compares the distribution of generated jet tau21 values between the original paper's dataset and an updated dataset. Uses matplotlib for plotting.

```python
_=plt.hist(zjets['gen_tau2s'],bins=np.linspace(0,0.5,20),density=True,label="Original Paper",alpha=0.5)
_=plt.hist(zjets_bigger['gen_tau2s'],bins=np.linspace(0,0.5,20),histtype="step",color="black",label="Updated",density=True,)
plt.xlabel("Gen. jet $\tau_{21}$")
plt.ylabel("Normalized to unity")
plt.legend(frameon=False)
```

--------------------------------

### Define Keras Model

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample_minimal.ipynb

Defines a simple feedforward neural network using Keras functional API. This model is used within the OmniFold process for learning the correction function.

```python
inputs = Input((1, ))
hidden_layer_1 = Dense(50, activation='relu')(inputs)
hidden_layer_2 = Dense(50, activation='relu')(hidden_layer_1)
hidden_layer_3 = Dense(50, activation='relu')(hidden_layer_2)
outputs = Dense(1, activation='sigmoid')(hidden_layer_3)
model = Model(inputs=inputs, outputs=outputs)
```

--------------------------------

### Apply OmniFold Correction

Source: https://github.com/hep-lbdl/omnifold/blob/main/GaussianExample_minimal.ipynb

Applies the OmniFold algorithm to correct the 'natural' data distribution using the defined Keras model and synthetic data. This is the core step for data correction.

```python
myweights = of.omnifold(theta0,theta_unknown_S,2,model)
```

=== COMPLETE CONTENT === This response contains all available snippets from this library. No additional content exists. Do not make further requests.