<jupyter_start><jupyter_text>Create SGS model assembliesThis notebook creates and exports wrapper models to be used in OpenFOAM simulations. Each wrapper model hides the single MLP models for diffusive and convective fluxes for one species. The wrapper model also contains constants to scale and rescale model inputs and outputs. DependenciesNotebooks (links work only if the Docker container is running)- [create_sgs_models.ipynb](http://1192.168.127.12:8888/notebooks/create_sgs_models.ipynb)- [reduce_training_data.ipynb](http://127.0.0.1:8888/notebooks/reduce_training_data.ipynb) Model assemblies<jupyter_code>import helper_module as hm
import pandas as pd
import numpy as np
import torch
from sklearn.preprocessing import MinMaxScaler

data_path = "../data/sgs_data/"
output_path_model = "../output/models/"
torch.set_default_tensor_type(torch.DoubleTensor)
reactions = ["phys", "decay", "single", "cons"]
data_files = {}
for r in reactions:
    path = data_path + "{:s}_red.csv".format(r)
    data_files[r] = pd.read_csv(path, header=0)
    
model_dict = {
    "n_inputs" : 2,
    "n_outputs" : 1,
    "n_layers" : 6,
    "n_neurons" : 40,
    "activation" : torch.nn.functional.relu,
    "batch_norm" : False
}
class A_Model(torch.nn.Module):
    def __init__(self, modelGrad_s, modelGrad_f, model_f, X_max, X_min, y_max, y_min):
        super().__init__()
        self.modelGrad_s = modelGrad_s
        self.modelGrad_f = modelGrad_f
        self.model_f = model_f
        self.X_max = torch.tensor(X_max, dtype=torch.float64)
        self.X_min = torch.tensor(X_min, dtype=torch.float64)
        self.grad_s_max = torch.tensor(y_max[0], dtype=torch.float64)
        self.grad_s_min = torch.tensor(y_min[0], dtype=torch.float64)
        self.grad_f_max = torch.tensor(y_max[1], dtype=torch.float64)
        self.grad_f_min = torch.tensor(y_min[1], dtype=torch.float64)
        self.f_max = torch.tensor(y_max[2], dtype=torch.float64)
        self.f_min = torch.tensor(y_min[2], dtype=torch.float64)
        
    def forward(self, Xorg):
        X = Xorg.clone()
        X.requires_grad_(False)
        X = (X - self.X_min) / (self.X_max - self.X_min)
        grad_s = self.modelGrad_s(X)
        grad_f = self.modelGrad_f(X)
        f = self.model_f(X)
        grad_s = grad_s * (self.grad_s_max - self.grad_s_min) + self.grad_s_min
        grad_f = grad_f * (self.grad_f_max - self.grad_f_min) + self.grad_f_min
        f = f * (self.f_max - self.f_min) + self.f_min
        return torch.cat((grad_s, grad_f, f), 1)

    
class BPS_Model(torch.nn.Module):
    def __init__(self, modelGrad_f, model_f, X_max, X_min, y_max, y_min):
        super().__init__()
        self.modelGrad_f = modelGrad_f
        self.model_f = model_f
        self.X_max = torch.tensor(X_max, dtype=torch.float64)
        self.X_min = torch.tensor(X_min, dtype=torch.float64)
        self.grad_f_max = torch.tensor(y_max[0], dtype=torch.float64)
        self.grad_f_min = torch.tensor(y_min[0], dtype=torch.float64)
        self.f_max = torch.tensor(y_max[1], dtype=torch.float64)
        self.f_min = torch.tensor(y_min[1], dtype=torch.float64)
        
    def forward(self, Xorg):
        X = Xorg.clone()
        X.requires_grad_(False)
        X = (X - self.X_min) / (self.X_max - self.X_min)
        grad_f = self.modelGrad_f(X)
        f = self.model_f(X)
        grad_f = grad_f * (self.grad_f_max - self.grad_f_min) + self.grad_f_min
        f = f * (self.f_max - self.f_min) + self.f_min
        return torch.cat((grad_f, f), 1)
    
    
class R_Model(torch.nn.Module):
    def __init__(self, model_r, X_max, X_min, y_max, y_min):
        super().__init__()
        self.model_r = model_r
        self.X_max = torch.tensor(X_max, dtype=torch.float64)
        self.X_min = torch.tensor(X_min, dtype=torch.float64)
        self.r_max = torch.tensor(y_max[0], dtype=torch.float64)
        self.r_min = torch.tensor(y_min[0], dtype=torch.float64)
        
    def forward(self, Xorg):
        X = Xorg.clone()
        X.requires_grad_(False)
        X = (X - self.X_min) / (self.X_max - self.X_min)
        r = self.model_r(X)
        r = r * (self.r_max - self.r_min) + self.r_min
        return r.squeeze()<jupyter_output><empty_output><jupyter_text>Physisorption<jupyter_code>data = data_files["phys"]
scaler_X = MinMaxScaler()
scaler_y = MinMaxScaler()
X01 = scaler_X.fit_transform(data[["dist", "A_av"]].values)
y01 = scaler_y.fit_transform(data[["gradA_s", "gradA_f", "A_f"]].values)
print("feature min: ", scaler_X.data_min_)
print("feature max: ", scaler_X.data_max_)
print("label min: ", scaler_y.data_min_)
print("label max: ", scaler_y.data_max_)
model_dict["n_inputs"] = 2

modelGrad_s = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "phys_gradA_s.pt"
modelGrad_s.load_state_dict(torch.load(model_path))
modelGrad_s.eval()

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "phys_gradA_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "phys_A_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_A_model = A_Model(modelGrad_s, modelGrad_f, model_f,
                       scaler_X.data_max_, scaler_X.data_min_,
                       scaler_y.data_max_, scaler_y.data_min_)

X = data[["dist", "A_av"]].values
model_trace = torch.jit.trace(full_A_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "A_model_phys.pt")
# sanity test
A_model_phys = torch.jit.load(output_path_model + "A_model_phys.pt")
X = data[["dist", "A_av"]].values
y = data[["gradA_s", "gradA_f", "A_f"]].values
ym = A_model_phys(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))<jupyter_output>[  3.21585444e+01   6.45445141e+01   1.17640193e-02]<jupyter_text>Decay reaction<jupyter_code>data = data_files["decay"]
scaler_X_A = MinMaxScaler()
scaler_X_P = MinMaxScaler()
scaler_y_A = MinMaxScaler()
scaler_y_P = MinMaxScaler()
X01_A = scaler_X_A.fit_transform(data[["dist", "A_av"]].values)
y01_A = scaler_y_A.fit_transform(data[["gradA_s", "gradA_f", "A_f"]])
X01_P = scaler_X_P.fit_transform(data[["dist", "A_av", "P_av"]].values)
y01_P = scaler_y_P.fit_transform(data[["gradP_f", "P_f"]])
print("feature A min: ", scaler_X_A.data_min_)
print("feature A max: ", scaler_X_A.data_max_)
print("feature P min: ", scaler_X_P.data_min_)
print("feature P max: ", scaler_X_P.data_max_)
print("label A min: ", scaler_y_A.data_min_)
print("label A max: ", scaler_y_A.data_max_)
print("label P min: ", scaler_y_P.data_min_)
print("label P max: ", scaler_y_P.data_max_)
model_dict["n_inputs"] = 2

modelGrad_s = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "decay_gradA_s.pt"
modelGrad_s.load_state_dict(torch.load(model_path))
modelGrad_s.eval()

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "decay_gradA_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "decay_A_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_A_model = A_Model(modelGrad_s, modelGrad_f, model_f,
                       scaler_X_A.data_max_, scaler_X_A.data_min_,
                       scaler_y_A.data_max_, scaler_y_A.data_min_)

X = data[["dist", "A_av"]].values
model_trace = torch.jit.trace(full_A_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "A_model_decay.pt")
# sanity test
A_model_decay = torch.jit.load(output_path_model + "A_model_decay.pt")
X = data[["dist", "A_av"]].values
y = data[["gradA_s", "gradA_f", "A_f"]].values
ym = A_model_decay(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))
model_dict["n_inputs"] = 3

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "decay_gradP_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "decay_P_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_P_model = BPS_Model(modelGrad_f, model_f,
                         scaler_X_P.data_max_, scaler_X_P.data_min_,
                         scaler_y_P.data_max_, scaler_y_P.data_min_)

X = data[["dist", "A_av", "P_av"]].values
model_trace = torch.jit.trace(full_P_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "P_model_decay.pt")
# sanity test
P_model_decay = torch.jit.load(output_path_model + "P_model_decay.pt")
X = data[["dist", "A_av", "P_av"]].values
y = data[["gradP_f", "P_f"]].values
ym = P_model_decay(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))<jupyter_output>[ 0.6124376   0.00449567]<jupyter_text>Single reaction<jupyter_code>data = data_files["single"]
scaler_X_AB = MinMaxScaler()
scaler_X_P = MinMaxScaler()
scaler_y_A = MinMaxScaler()
scaler_y_B = MinMaxScaler()
scaler_y_P = MinMaxScaler()
X01_AB = scaler_X_AB.fit_transform(data[["dist", "A_av", "B_av"]].values)
y01_A = scaler_y_A.fit_transform(data[["gradA_s", "gradA_f", "A_f"]])
y01_B = scaler_y_B.fit_transform(data[["gradB_f", "B_f"]])
X01_P = scaler_X_P.fit_transform(data[["dist", "A_av", "B_av", "P_av"]].values)
y01_P = scaler_y_P.fit_transform(data[["gradP_f", "P_f"]])
print("feature AB min: ", scaler_X_AB.data_min_)
print("feature AB max: ", scaler_X_AB.data_max_)
print("feature P min: ", scaler_X_P.data_min_)
print("feature P max: ", scaler_X_P.data_max_)
print("label A min: ", scaler_y_A.data_min_)
print("label A max: ", scaler_y_A.data_max_)
print("label B min: ", scaler_y_B.data_min_)
print("label B max: ", scaler_y_B.data_max_)
print("label P min: ", scaler_y_P.data_min_)
print("label P max: ", scaler_y_P.data_max_)
model_dict["n_inputs"] = 3

modelGrad_s = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "single_gradA_s.pt"
modelGrad_s.load_state_dict(torch.load(model_path))
modelGrad_s.eval()

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "single_gradA_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "single_A_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_A_model = A_Model(modelGrad_s, modelGrad_f, model_f,
                       scaler_X_AB.data_max_, scaler_X_AB.data_min_,
                       scaler_y_A.data_max_, scaler_y_A.data_min_)

X = data[["dist", "A_av", "B_av"]].values
model_trace = torch.jit.trace(full_A_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "A_model_single.pt")
# sanity test
A_model_single = torch.jit.load(output_path_model + "A_model_single.pt")
X = data[["dist", "A_av", "B_av"]].values
y = data[["gradA_s", "gradA_f", "A_f"]].values
ym = A_model_single(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))
model_dict["n_inputs"] = 3

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "single_gradB_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "single_B_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_B_model = BPS_Model(modelGrad_f, model_f,
                         scaler_X_AB.data_max_, scaler_X_AB.data_min_,
                         scaler_y_B.data_max_, scaler_y_B.data_min_)

X = data[["dist", "A_av", "B_av"]].values
model_trace = torch.jit.trace(full_B_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "B_model_single.pt")
# sanity test
B_model_single = torch.jit.load(output_path_model + "B_model_single.pt")
X = data[["dist", "A_av", "B_av"]].values
y = data[["gradB_f", "B_f"]].values
ym = B_model_single(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))
model_dict["n_inputs"] = 4

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "single_gradP_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "single_P_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_P_model = BPS_Model(modelGrad_f, model_f,
                         scaler_X_P.data_max_, scaler_X_P.data_min_,
                         scaler_y_P.data_max_, scaler_y_P.data_min_)

X = data[["dist", "A_av", "B_av", "P_av"]].values
model_trace = torch.jit.trace(full_P_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "P_model_single.pt")
# sanity test
P_model_single = torch.jit.load(output_path_model + "P_model_single.pt")
X = data[["dist", "A_av", "B_av", "P_av"]].values
y = data[["gradP_f", "P_f"]].values
ym = P_model_single(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))<jupyter_output>[ 0.61555782  0.00356164]<jupyter_text>Consecutive reaction<jupyter_code>data = data_files["cons"]
scaler_X_ABP = MinMaxScaler()
scaler_X_ABPS = MinMaxScaler()
scaler_y_A = MinMaxScaler()
scaler_y_B = MinMaxScaler()
scaler_y_P = MinMaxScaler()
scaler_y_S = MinMaxScaler()
X01_ABP = scaler_X_ABP.fit_transform(data[["dist", "A_av", "B_av", "P_av"]].values)
y01_A = scaler_y_A.fit_transform(data[["gradA_s", "gradA_f", "A_f"]])
y01_B = scaler_y_B.fit_transform(data[["gradB_f", "B_f"]])
y01_P = scaler_y_P.fit_transform(data[["gradP_f", "P_f"]])
X01_ABPS = scaler_X_ABPS.fit_transform(data[["dist", "A_av", "B_av", "P_av", "S_av"]].values)
y01_S = scaler_y_S.fit_transform(data[["gradS_f", "S_f"]])
print("feature ABP min: ", scaler_X_ABP.data_min_)
print("feature ABP max: ", scaler_X_ABP.data_max_)
print("feature ABPS min: ", scaler_X_ABPS.data_min_)
print("feature ABPS max: ", scaler_X_ABPS.data_max_)
print("label A min: ", scaler_y_A.data_min_)
print("label A max: ", scaler_y_A.data_max_)
print("label B min: ", scaler_y_B.data_min_)
print("label B max: ", scaler_y_B.data_max_)
print("label P min: ", scaler_y_P.data_min_)
print("label P max: ", scaler_y_P.data_max_)
print("label S min: ", scaler_y_S.data_min_)
print("label S max: ", scaler_y_S.data_max_)
model_dict["n_inputs"] = 4

modelGrad_s = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_gradA_s.pt"
modelGrad_s.load_state_dict(torch.load(model_path))
modelGrad_s.eval()

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_gradA_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_A_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_A_model = A_Model(modelGrad_s, modelGrad_f, model_f,
                       scaler_X_ABP.data_max_, scaler_X_ABP.data_min_,
                       scaler_y_A.data_max_, scaler_y_A.data_min_)

X = data[["dist", "A_av", "B_av", "P_av"]].values
model_trace = torch.jit.trace(full_A_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "A_model_cons.pt")
# sanity test
A_model_cons = torch.jit.load(output_path_model + "A_model_cons.pt")
X = data[["dist", "A_av", "B_av", "P_av"]].values
y = data[["gradA_s", "gradA_f", "A_f"]].values
ym = A_model_cons(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))
model_dict["n_inputs"] = 4

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_gradB_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_B_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_B_model = BPS_Model(modelGrad_f, model_f,
                         scaler_X_ABP.data_max_, scaler_X_ABP.data_min_,
                         scaler_y_B.data_max_, scaler_y_B.data_min_)

X = data[["dist", "A_av", "B_av", "P_av"]].values
model_trace = torch.jit.trace(full_B_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "B_model_cons.pt")
# sanity test
B_model_cons = torch.jit.load(output_path_model + "B_model_cons.pt")
X = data[["dist", "A_av", "B_av", "P_av"]].values
y = data[["gradB_f", "B_f"]].values
ym = B_model_cons(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))
model_dict["n_inputs"] = 4

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_gradP_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_P_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_P_model = BPS_Model(modelGrad_f, model_f,
                         scaler_X_ABP.data_max_, scaler_X_ABP.data_min_,
                         scaler_y_P.data_max_, scaler_y_P.data_min_)

X = data[["dist", "A_av", "B_av", "P_av"]].values
model_trace = torch.jit.trace(full_P_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "P_model_cons.pt")
# sanity test
P_model_cons = torch.jit.load(output_path_model + "P_model_cons.pt")
X = data[["dist", "A_av", "B_av", "P_av"]].values
y = data[["gradP_f", "P_f"]].values
ym = P_model_cons(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))
model_dict["n_inputs"] = 5

modelGrad_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_gradS_f.pt"
modelGrad_f.load_state_dict(torch.load(model_path))
modelGrad_f.eval()

model_f = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_S_f.pt"
model_f.load_state_dict(torch.load(model_path))
model_f.eval()

full_S_model = BPS_Model(modelGrad_f, model_f,
                         scaler_X_ABPS.data_max_, scaler_X_ABPS.data_min_,
                         scaler_y_S.data_max_, scaler_y_S.data_min_)

X = data[["dist", "A_av", "B_av", "P_av", "S_av"]].values
model_trace = torch.jit.trace(full_S_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "S_model_cons.pt")
# sanity test
S_model_cons = torch.jit.load(output_path_model + "S_model_cons.pt")
X = data[["dist", "A_av", "B_av", "P_av", "S_av"]].values
y = data[["gradS_f", "S_f"]].values
ym = S_model_cons(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))<jupyter_output>[ 0.1011508   0.00194043]<jupyter_text>Source terms Single reaction<jupyter_code>data = data_files["single"]
scaler_X_rAB = MinMaxScaler()
scaler_y_rAB = MinMaxScaler()
X01_rAB = scaler_X_rAB.fit_transform(data[["dist", "A_av", "B_av"]].values)
y01_rAB = scaler_y_rAB.fit_transform(data[["rAB_av"]].values)
print("X rAB min: ", scaler_X_rAB.data_min_)
print("X rAB max: ", scaler_X_rAB.data_max_)
print("y rAB min: ", scaler_y_rAB.data_min_)
print("y rAB max: ", scaler_y_rAB.data_max_)
model_dict["n_inputs"] = 3

model_r = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "single_rAB_av.pt"
model_r.load_state_dict(torch.load(model_path))
model_r.eval()

full_R_model = R_Model(model_r,
                       scaler_X_rAB.data_max_, scaler_X_rAB.data_min_,
                       scaler_y_rAB.data_max_, scaler_y_rAB.data_min_)

X = data[["dist", "A_av", "B_av"]].values
model_trace = torch.jit.trace(full_R_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "rAB_model_single.pt")
# sanity test
rAB_model_single = torch.jit.load(output_path_model + "rAB_model_single.pt")
X = data[["dist", "A_av", "B_av"]].values
y = data[["rAB_av"]].values[:,0]
ym = rAB_model_single(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))<jupyter_output>0.00648585085187<jupyter_text>Consecutive reaction<jupyter_code>data = data_files["cons"]
scaler_X_rAB = MinMaxScaler()
scaler_y_rAB = MinMaxScaler()
X01_rAB = scaler_X_rAB.fit_transform(data[["dist", "A_av", "B_av"]].values)
y01_rAB = scaler_y_rAB.fit_transform(data[["rAB_av"]].values)
print("X rAB min: ", scaler_X_rAB.data_min_)
print("X rAB max: ", scaler_X_rAB.data_max_)
print("y rAB min: ", scaler_y_rAB.data_min_)
print("y rAB max: ", scaler_y_rAB.data_max_)
model_dict["n_inputs"] = 3

model_r = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_rAB_av.pt"
model_r.load_state_dict(torch.load(model_path))
model_r.eval()

full_R_model = R_Model(model_r,
                       scaler_X_rAB.data_max_, scaler_X_rAB.data_min_,
                       scaler_y_rAB.data_max_, scaler_y_rAB.data_min_)

X = data[["dist", "A_av", "B_av"]].values
model_trace = torch.jit.trace(full_R_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "rAB_model_cons.pt")
# sanity test
rAB_model_cons = torch.jit.load(output_path_model + "rAB_model_cons.pt")
X = data[["dist", "A_av", "B_av"]].values
y = data[["rAB_av"]].values[:,0]
ym = rAB_model_cons(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))
data = data_files["cons"]
scaler_X_rAP = MinMaxScaler()
scaler_y_rAP = MinMaxScaler()
X01_rAP = scaler_X_rAP.fit_transform(data[["dist", "A_av", "P_av"]].values)
y01_rAP = scaler_y_rAP.fit_transform(data[["rAP_av"]].values)
print("X rAP min: ", scaler_X_rAP.data_min_)
print("X rAP max: ", scaler_X_rAP.data_max_)
print("y rAP min: ", scaler_y_rAP.data_min_)
print("y rAP max: ", scaler_y_rAP.data_max_)
model_dict["n_inputs"] = 3

model_r = hm.SimpleMLP(**model_dict)
model_path = output_path_model + "cons_rAP_av.pt"
model_r.load_state_dict(torch.load(model_path))
model_r.eval()

full_R_model = R_Model(model_r,
                       scaler_X_rAP.data_max_, scaler_X_rAP.data_min_,
                       scaler_y_rAP.data_max_, scaler_y_rAP.data_min_)

X = data[["dist", "A_av", "P_av"]].values
model_trace = torch.jit.trace(full_R_model, torch.from_numpy(X[0]).unsqueeze(-1).T)
model_trace.save(output_path_model + "rAP_model_cons.pt")
# sanity test
rAP_model_cons = torch.jit.load(output_path_model + "rAP_model_cons.pt")
X = data[["dist", "A_av", "P_av"]].values
y = data[["rAP_av"]].values[:,0]
ym = rAP_model_cons(torch.from_numpy(X)).detach().numpy()
diff = ym - y
print(np.amax(diff, axis=0))<jupyter_output>0.00176690606511
{'fraction_non_alphanumeric': 0.15195395062283581, 'fraction_numerical': 0.013401088276296262, 'mean_word_length': 5.112699285321606, 'pattern_counts': {'":': 0, '<': 23, '<?xml version=': 0, '>': 23, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 2, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'chain_length': '7', 'dir': 'jupyter-structured-clean-dedup', 'id': '319053', 'license': 'MIT', 'n_tokens_mistral': 9537, 'n_tokens_neox': 9376, 'n_words': 1308, 'path': 'notebooks/create_sgs_model_assembly.ipynb', 'repo_name': 'AndreWeiner/phd_notebooks', 'size': '22231'}