porting_pcnn_experiment.py

import matplotlib.pyplot as plt
import numpy as np
import pandas as pd

import lightning as L
from pathlib import Path
import pickle
import torch

from bems.helpers.data import load_simulation_results
from error_compensation.models import Features as F
from bems.config import EMSConfig
from bems.helpers.data import prepare_pcnn_training_data

from model_fitting.pcnn import S_PCNN_MLP, train_pcnn, setup_pcnn_residual_estimator
from model_fitting.interfaces import Topology
from model_fitting.utility import build_state_space_model, get_parameters_from_bems_config

#%%
# results_path = Path(r"C:\Users\re902174\workspaces\Hierarchical MPC\robustification\bems-cosimos\analysis\journal_publication_study\final_results\year_monolithic_dt_240412_163803")
results_path = Path(r"C:\Users\re902174\workspaces\Hierarchical MPC\robustification\bems-cosimos\simulations\digital_twin\results\year_simple_rbc_dt_240830_130239")
dump_directory = Path(r"C:\Users\re902174\workspaces\Hierarchical MPC\robustification\bems-cosimos\data\fitted_models\simple_rbc")
compensator_directory = Path(r"C:\Users\re902174\workspaces\Hierarchical MPC\robustification\bems-cosimos\data\compensators\fitted_models\simple_rbc")

# linear_parameter_source = r"C:\Users\re902174\workspaces\Hierarchical MPC\robustification\bems-cosimos\data\fitted_models\fitted_parameters_exact_other_2021.pkl"
linear_parameter_source = Path(
    r"C:\Users\re902174\workspaces\Hierarchical MPC\robustification\bems-cosimos\data\fitted_models\simple_rbc\fitted_parameters_simple_rbc_other__2021.pkl"
)
train_on_exact_linear = True

experiment_name = "pcnn_simple_rbc"

data = load_simulation_results(results_path)

config = EMSConfig()

features, labels = prepare_pcnn_training_data(data)

topology = Topology(
    rooms=list(range(0, 9)),
    outside=list(range(0, 9)),
    couplings=[(1, 8), (2, 3), (4, 5), (4, 7), (5, 7)]
)

# topology = Topology(
#     rooms=list(range(0, 7)),
#     outside=list(range(0, 7)),
#     neighbors=[(2, 3), (4, 5)]
# )

# define system and topology
Q_total_cols = [f"Q_total_{i}" for i in topology.room_numbers]
nn_feature_columns = F.base_features + F.cyclic_tod_features + F.cyclic_dow_features + F.server_load

if not train_on_exact_linear:
    with open(linear_parameter_source, "rb") as fh:
        initial_fitted_parameters = pickle.load(fh)
else:
    initial_fitted_parameters = get_parameters_from_bems_config(config)

#%%
pcnn, training_features, scaler_x, scaler_y = train_pcnn(
    topology=topology,
    features=features,
    labels=labels,
    temperature_columns=F.interior_temp,
    power_columns=Q_total_cols,
    ambient_temperature_column=F.ambient_temp[0],
    nn_feature_columns=nn_feature_columns,
    initial_parameter_values=initial_fitted_parameters
)

fitted_state_space_model = build_state_space_model(topology, pcnn.linear_parameters)
fitted_parameters = pcnn.linear_parameters
pcnn_residual_compensator = setup_pcnn_residual_estimator(pcnn, scaler_x, scaler_y)

dump = False

if dump:
    dump_directory.mkdir(parents=True, exist_ok=True)
    compensator_directory.mkdir(parents=True, exist_ok=True)

    parametrization_string = f"{experiment_name}_"

    if not train_on_exact_linear:
        parametrization_string += (
            linear_parameter_source.name
            .replace("fitted_parameters_", "")
            .replace(".pkl", "")
        )
    else:
        parametrization_string = "exact"

    import pickle
    file_name = f"fitted_state_space_{parametrization_string}.pkl"
    with open(dump_directory.joinpath(file_name), "wb") as fh:
        pickle.dump(fitted_state_space_model, fh)

    file_name = f"fitted_parameters_{parametrization_string}.pkl"
    with open(dump_directory.joinpath(file_name), "wb") as fh:
        pickle.dump(fitted_parameters, fh)

    pcnn_residual_compensator.dump(dump_directory.joinpath(f"pcnn_res_{parametrization_string}.pkl"))

#%%
# TODO: plot predictions
pcnn.eval()
theta_cols = [f"theta_{i}" for i in topology.room_numbers]
y_hat = pcnn(torch.tensor(scaler_x.transform(training_features), dtype=torch.float32))
y_hat_pd = pd.DataFrame(data=y_hat.detach().numpy(), columns=theta_cols)
# for i, col in enumerate([f"theta_{i}" for i in range(1, 10)]):
#     y_hat_pd[col] = pcnn._inverse_normalize(y_hat_pd[col].values, "room", i)

fig, axs = plt.subplots(3, 3, figsize=(10, 10))
for i, ax in enumerate(axs.flatten()):
    if i == topology.n_rooms:
        break

    #labels.reset_index()[theta_cols[i]].plot(ax=ax, label="$y_\mathrm{real}$", alpha=0.6)
    ax.plot(scaler_y.transform(labels)[:, i], label="$y_\mathrm{real}$", alpha=0.6)
    y_hat_pd[theta_cols[i]].plot(ax=ax, label="$y_\mathrm{pred}$", alpha=0.6)

    ax.grid(True)
    ax.legend()
    ax.set_ylabel(rf"$\theta_{{{i+1}, \mathrm{{norm}}}}$")

fig.tight_layout()
plt.show()

#%%
fig, axs = plt.subplots(3, 3, figsize=(10, 10))
for i, ax in enumerate(axs.flatten()):
    if i == topology.n_rooms:
        break

    ax.hist(scaler_y.transform(labels)[:, i], label="$y_\mathrm{real}$", bins=100, alpha=0.6)
    y_hat_pd[theta_cols[i]].hist(ax=ax, label="$y_\mathrm{pred}$", bins=100, alpha=0.6)

    ax.grid(True)
    ax.legend()
    ax.set_ylabel(rf"$\theta_{{{i+1}, \mathrm{{norm}}}}$")

fig.tight_layout()
plt.show()

#%%
theta_cols = [f"theta_{i}" for i in topology.room_numbers]
y_hat_lin, y_hat_nn = pcnn.forward_split(torch.tensor(scaler_x.transform(training_features), dtype=torch.float32))
y_hat_lin_pd = pd.DataFrame(data=y_hat_lin.detach().numpy(), columns=theta_cols)
y_hat_nn_pd = pd.DataFrame(data=y_hat_nn.detach().numpy(), columns=theta_cols)

fig, axs = plt.subplots(3, 3, figsize=(10, 10))
for i, ax in enumerate(axs.flatten()):
    if i == topology.n_rooms:
        break

    #labels.reset_index()[theta_cols[i]].plot(ax=ax, label="$y_\mathrm{real}$", alpha=0.6)
    ax.plot(scaler_y.transform(labels)[:, i], label="$y_\mathrm{real}$", alpha=0.6)
    y_hat_lin_pd[theta_cols[i]].plot(ax=ax, label="$y_\mathrm{pred,lin}$", alpha=0.6)
    y_hat_nn_pd[theta_cols[i]].plot(ax=ax, label="$y_\mathrm{pred,NN}$", alpha=0.6)

    ax.grid(True)
    ax.legend()
    ax.set_ylabel(rf"$\theta_{{{i+1}, \mathrm{{norm}}}}$")

fig.tight_layout()
plt.show()

#%%
y_pred_res = pd.DataFrame(data=pcnn_residual_compensator.predict(features), columns=theta_cols)

fig, axs = plt.subplots(3, 3, figsize=(10, 10))
y_hat_nn_rev = scaler_y.inverse_transform(y_hat_nn.detach().numpy())
y_hat_nn_rev_pd = pd.DataFrame(data=y_hat_nn_rev, columns=theta_cols)
y_hat_rev = scaler_y.inverse_transform(y_hat.detach().numpy())
y_hat_rev_pd = pd.DataFrame(data=y_hat_rev, columns=theta_cols)
for i, ax in enumerate(axs.flatten()):
    if i == topology.n_rooms:
        break

    y_pred_res[theta_cols[i]].plot(ax=ax, label="dumped comp", alpha=0.6)
    y_hat_nn_rev_pd[theta_cols[i]].plot(ax=ax, label="res nn", linestyle=":", color="k", alpha=0.6)
    # y_hat_rev_pd[theta_cols[i]].plot(ax=ax, label="full", alpha=0.6)
    # labels.reset_index()[labels.columns[i]].plot(ax=ax, label="true", alpha=0.6)

    ax.grid(True)
    ax.legend()
    ax.set_ylabel(rf"$\theta_{{{i+1}, \mathrm{{norm}}}}$")

fig.tight_layout()
plt.show()
#%%
from error_compensation.models import PCNNCompensator


#res = test.predict(features)

#%%
#neg_eps_c_m = pcnn.neg_eps_c_m.cpu().detach().numpy()[0]
# eps_c_m = -0.02 * np.exp(float(pcnn.neg_eps_c_m._parameters['log_weight']))
# eps_c_c = pcnn.eps_c_c.cpu().detach().numpy()[0]
# eps_c_0 = pcnn.eps_c_0.cpu().detach().numpy()[0]
#
# eps_c_dyn = lambda theta_air: np.minimum(eps_c_m * theta_air + eps_c_c, eps_c_0)
# fig, ax = plt.subplots()
# eps_c_dyn_eval = eps_c_dyn(features["theta_air_0"])
# # Q = eps * P
# ax.scatter(features["P_cool_b"], features["Q_cool_b"], label="true")
# ax.scatter(features["P_cool_b"], eps_c_dyn_eval * features["P_cool_b"], label="dyn")
# ax.scatter(features["P_cool_b"], config.eps_c * features["P_cool_b"], label="const")
# ax.set_xlabel("P_cool_b")
# ax.set_ylabel("P_cool_b")
# ax.grid(True)
# ax.legend()
# fig.tight_layout()
# plt.show()
#
# #%%
# fig, ax = plt.subplots()
#
# ax.scatter(features["theta_air_0"], eps_c_dyn_eval)
# ax.set_xlabel("theta_air")
# ax.set_ylabel("eps_c")
# ax.grid(True)
# ax.legend()
# fig.tight_layout()
# plt.show()