test.py

import atexit
import subprocess
import sys
import torch
import util
import yaml
import numpy as np
from tqdm import trange
import pathlib
from system import DynamicsModelFullDataset
import plotting
from torch.utils.data import DataLoader, SequentialSampler, BatchSampler
from model import load_models
from model.loss import load_loss_function

QUANTILES = [0.1, 0.25, 0.5, 0.667, 0.75, 0.9, 0.95]

# TODO change the plots to try and plot things by checking if it is in the dictionary or not
# TODO simplify the plotting to use more class methods, just ignore on the compilation
# TODO always plot the full box plots, not the zoomed in versions


def _terminate_server(proc):
    if proc.poll() is None:
        proc.terminate()
        try:
            proc.wait(timeout=5)
        except subprocess.TimeoutExpired:
            proc.kill()
            proc.wait()


def createMarkdownEntry(name, T, error, var):
    local_dict = {}
    local_dict["Name"] = name
    local_dict["T"] = T
    local_dict["Error"] = error
    local_dict["Var"] = var
    return local_dict


def createQuantileMarkdownEntry(name, T, quantile, range, num_bins, value):
    return {
        "Name": name,
        "T": T,
        "Quantile": quantile,
        "Quantile Bin Size": (range[1] - range[0]) / num_bins,
        "Value": value,
    }


@torch.no_grad()
def run_model(
    dataloader,
    model,
    loss_fn,
    model_results_path=None,
    random_state_plot_num=5,
    max_batches=None,
    num_bins=1000,
    T_delta=0.5,
):
    # things we need for all values
    # batch_size = 1000
    # num_batches = len(dataloader)
    # old_state_error = np.zeros(
    #     (num_batches * batch_size, model.system.traj_length, model.system.error_dim)
    # )
    util.clearGPUCache()
    system = model.system

    if max_batches is None:
        max_batches = len(dataloader)
    else:
        max_batches = min(max_batches, len(dataloader))

    plot_bad_traj = False
    if plot_bad_traj:
        random_state_plot_num = max_batches

    # things we only need to randomly sampled values for plotting
    size = random_state_plot_num + 2
    all_state_error = np.zeros((size, model.system.traj_length, model.system.error_dim))
    all_states = np.zeros(
        (
            size,
            model.system.init_length + model.system.traj_length + 1,
            model.system.getJacobianStateDim(),
        )
    )
    all_predicted_states = None
    all_inputs = np.zeros(
        (
            size,
            model.system.traj_length + model.system.init_length,
            model.system.input_dim,
        )
    )
    total_forward = None
    all_model_outputs = np.zeros(
        (size, model.system.traj_length, model.model_output_dim)
    )
    all_loss = np.zeros((size, model.system.traj_length))

    # create randomly sampled indexes +2 for min/max
    # (batch number, index in batch)
    np.random.seed(0)
    batch_size = None
    if hasattr(dataloader, "batch_sampler"):
        batch_size = dataloader.batch_sampler.batch_size
    else:
        batch_size = len(dataloader)

    indexes = np.random.randint(0, batch_size, size=(size, 2))  # batch, index
    # ensure that batch index is bounded
    indexes[:, 0] = indexes[:, 0] % max_batches
    # min batch set to -1
    indexes[0, 0] = -1
    # max batch set to -1
    indexes[1, 0] = -1
    min_loss = None
    max_loss = None

    if plot_bad_traj:
        for i in range(random_state_plot_num):
            indexes[i + 2, 0] = i

    num_points = 0

    prior_mean = np.zeros((model.system.traj_length, model.system.error_dim))
    prior_var = np.zeros((model.system.traj_length, model.system.error_dim))
    prior_abs_mean = None
    prior_abs_var = None

    prior_loss_mean = np.zeros((model.system.traj_length))
    prior_loss_var = np.zeros((model.system.traj_length))
    extra = {}
    extra["test_name"] = [None] * size
    extra["check_valid_states"] = np.zeros(
        (
            size,
            model.system.traj_length + model.system.init_length + 1,
            model.system.check_valid_dim,
        )
    )
    extra["val_loss"] = 0

    # num_bins = 5000
    # T_delta = 0.25  # how often to sample the error for histogram
    extra["T_delta"] = T_delta
    extra["error_bins"] = []
    extra["loss_bins"] = []
    for t in range(1 + int(system.T / T_delta)):
        extra["error_bins"].append(
            util.createEmptyHistogram(system.getErrorHistRanges(), num_bins)
        )
        extra["loss_bins"].append(
            util.createEmptyHistogram(system.getLossHistRange(), num_bins)
        )
    # TODO add in single step for histogram maybe? for initial noise or something
    extra["total_loss_bins"] = util.createEmptyHistogram(
        system.getTotalLossHistRange(), num_bins
    )

    # Calibration accumulators for uncertainty models (full-dataset z-score and
    # Mahalanobis histograms, one per T_delta time bin).
    # calib_z_bins[t][k]   : 1D histogram of signed z-scores for cov dim k at bin t
    # calib_mahal_bins[t]  : 1D histogram of squared Mahalanobis distances at bin t
    _calib_cov_dim = getattr(system, "covariance_dim", None)
    if _calib_cov_dim is not None:
        _num_calib_time = 1 + int(system.T / T_delta)
        calib_z_bins = [
            [
                util.createEmptyHistogram([-10.0, 10.0], num_bins)
                for _ in range(_calib_cov_dim)
            ]
            for _ in range(_num_calib_time)
        ]
        calib_mahal_bins = [
            util.createEmptyHistogram([0.0, float(_calib_cov_dim) * 25.0], num_bins)
            for _ in range(_num_calib_time)
        ]
    else:
        calib_z_bins = None
        calib_mahal_bins = None

    old_model_outputs = model.model_outputs
    if model.model_output_dim > 0:
        model.setModelOutputs(True)

    it = iter(dataloader)
    for cur_index in trange(max_batches, desc="evaluating", unit="batch"):
        # computes all the values
        input, state, check_valid_states, meta_data = next(it)
        input = input[
            :, : (model.system.traj_length + model.system.init_length), :
        ].contiguous()
        state = state[
            :, : (model.system.traj_length + 1 + model.system.init_length), :
        ].contiguous()
        check_valid_states = check_valid_states[
            :, : (model.system.traj_length + 1 + model.system.init_length), :
        ]
        assert state.shape[1] == model.system.traj_length + 1 + model.system.init_length

        input = input.to(model.device, non_blocking=True).contiguous()
        state = state.to(model.device, non_blocking=True).contiguous()
        output, state_prediction, memory, result_model_output = model.predictTrajectory(
            input, state
        )
        # extra_loss = model.extraLoss(memory, output, state, input).cpu().detach().numpy()

        loss_individual = loss_fn.getLoss(
            state_prediction[:, 1:], state[:, model.system.init_length + 1 :, :]
        )
        loss = torch.sum(loss_individual, dim=1).detach().cpu().numpy()
        extra["val_loss"] += (
            loss_fn.reduce(loss_individual).detach().cpu().numpy() * input.shape[0]
        )
        loss_individual = loss_individual.cpu().detach().numpy()

        state_error = model.system.computeError(
            state,
            state_prediction,
        )

        if plot_bad_traj and cur_index < random_state_plot_num:
            indexes[cur_index, 0] = np.argmax(loss)

        state_error = state_error.cpu().detach().numpy()

        for t in range(1 + int(system.T / T_delta)):
            t_index = 0 if t == 0 else int((t * T_delta) / system.dt) - 1
            extra["error_bins"][t] = util.computeHistogramBatchExpand(
                state_error[:, t_index, :], extra["error_bins"][t]
            )
            # gets the loss
            extra["loss_bins"][t] = util.computeHistogramBatchExpand(
                loss_individual[:, t_index], extra["loss_bins"][t]
            )

        # gets the loss
        extra["total_loss_bins"] = util.computeHistogramBatchExpand(
            loss, extra["total_loss_bins"]
        )

        # Accumulate calibration statistics for uncertainty models
        if calib_z_bins is not None:
            _pred_cov = system.getCovarianceFromState(state_prediction)
            if _pred_cov is not None:
                _ssd = system.subsystem_state_dim
                _cov_idx = system.cov_indexes
                _pred_cov_np = _pred_cov[:, 1:].cpu().numpy()
                # e = distance(x_true, x_pred)  — respects angle wrapping (B, T, ssd)
                _err_full = (
                    system.distance(
                        state[:, system.init_length + 1 :, :_ssd],
                        state_prediction[:, 1:, :_ssd],
                    )
                    .cpu()
                    .numpy()
                )
                for _t in range(1 + int(system.T / T_delta)):
                    _t_time = system.dt if _t == 0 else _t * T_delta
                    _t_step = min(int(_t_time / system.dt) - 1, system.traj_length - 1)
                    # e_k : error on covariance dimensions at this time step  (B, cov_dim)
                    _err = _err_full[:, _t_step, _cov_idx]
                    # σ_k = sqrt(Σ_kk)  — marginal standard deviation from diagonal of Σ
                    _std = np.sqrt(
                        np.maximum(
                            np.diagonal(_pred_cov_np[:, _t_step], axis1=1, axis2=2),
                            1e-12,
                        )
                    )
                    # z_k = e_k / σ_k  — marginal z-score
                    _z = _err / _std
                    for _k in range(_calib_cov_dim):
                        calib_z_bins[_t][_k] = util.computeHistogramBatchExpand(
                            _z[:, _k], calib_z_bins[_t][_k]
                        )
                    _cov_t = _pred_cov_np[:, _t_step]  # (B, cov_dim, cov_dim)
                    _invertible = np.abs(np.linalg.det(_cov_t)) > 0
                    _n_singular = int((~_invertible).sum())
                    if _n_singular > 0:
                        print(
                            f"WARNING: {_n_singular} singular covariance matrices at t={_t_time:.3f}s"
                        )
                    _mahal = np.full(len(_err), np.nan)
                    if _invertible.any():
                        _inv_cov = np.linalg.inv(_cov_t[_invertible])
                        # d² = eᵀ Σ⁻¹ e  — squared Mahalanobis distance (Σ is PD by construction)
                        _mahal[_invertible] = np.einsum(
                            "...i,...ij,...j->...",
                            _err[_invertible],
                            _inv_cov,
                            _err[_invertible],
                        )
                    _valid = np.isfinite(_mahal) & (_mahal >= 0)
                    calib_mahal_bins[_t] = util.computeHistogramBatchExpand(
                        _mahal[_valid], calib_mahal_bins[_t]
                    )

        prior_abs_mean, prior_abs_var = util.computeMeanVarBatch(
            np.abs(state_error),
            prior_abs_mean,
            prior_abs_var,
            num_points,
            add_points=False,
        )
        prior_loss_mean, prior_loss_var = util.computeMeanVarBatch(
            loss_individual,
            prior_loss_mean,
            prior_loss_var,
            num_points,
            add_points=False,
        )
        prior_mean, prior_var, num_points = util.computeMeanVarBatch(
            state_error, prior_mean, prior_var, num_points
        )

        # if first batch
        if min_loss is None:
            min_loss = (np.min(loss), np.argmin(loss))
            max_loss = (np.max(loss), np.argmax(loss))

        if all_predicted_states is None:
            all_predicted_states = np.zeros(
                (size, model.system.traj_length + 1, state_prediction.shape[-1])
            )
        if total_forward is None:
            total_forward = np.zeros((size, model.system.traj_length, output.shape[-1]))

        # if we have a new min loss
        if np.min(loss) <= min_loss[0]:
            min_index = np.argmin(loss)
            min_loss = (np.min(loss), min_index * cur_index * 1000)
            all_states[0, :, :] = state[min_index, :, :].cpu()
            all_predicted_states[0, :, :] = state_prediction[min_index, :, :].cpu()
            total_forward[0, :, :] = output[min_index, :, :].cpu()
            all_inputs[0, :, :] = input[min_index, :, :].cpu()
            all_model_outputs[0, :, :] = result_model_output[min_index, :, :].cpu()
            all_state_error[0, :, :] = state_error[min_index, :, :]
            all_loss[0, :] = loss_individual[min_index, :]
            extra["test_name"][0] = "".join([chr(x) for x in meta_data[min_index]])
            extra["check_valid_states"][0, :, :] = check_valid_states[
                min_index, :, :
            ].cpu()

        if np.max(loss) >= max_loss[0]:
            max_index = np.argmax(loss)
            max_loss = (np.max(loss), max_index * cur_index * 1000)
            all_states[1, :, :] = state[max_index, :, :].cpu()
            all_predicted_states[1, :, :] = state_prediction[max_index, :, :].cpu()
            total_forward[1, :, :] = output[max_index, :, :].cpu()
            all_inputs[1, :, :] = input[max_index, :, :].cpu()
            all_model_outputs[1, :, :] = result_model_output[max_index, :, :].cpu()
            all_state_error[1, :, :] = state_error[max_index, :, :]
            all_loss[1, :] = loss_individual[max_index, :]
            extra["test_name"][1] = "".join([chr(x) for x in meta_data[max_index]])
            extra["check_valid_states"][1, :, :] = check_valid_states[
                max_index, :, :
            ].cpu()

        index_mask = (indexes[:, 0] == cur_index).nonzero()
        for index in index_mask[0]:
            batch_index = indexes[index, 1]
            if indexes[index, 1] >= state.shape[0]:
                batch_index = (
                    round(
                        indexes[index, 1]
                        / dataloader.batch_sampler.batch_size
                        * state.shape[0]
                    )
                    - 1
                )
            all_states[index, :, :] = state[batch_index, :, :].cpu()
            all_predicted_states[index, :, :] = state_prediction[
                batch_index, :, :
            ].cpu()
            total_forward[index, :, :] = output[batch_index, :, :].cpu()
            all_inputs[index, :, :] = input[batch_index, :, :].cpu()
            all_model_outputs[index, :, :] = result_model_output[
                batch_index, :, :
            ].cpu()
            all_state_error[index, :, :] = state_error[batch_index, :, :]
            all_loss[index, :] = loss_individual[batch_index, :]
            extra["test_name"][index] = "".join(
                [chr(x) for x in meta_data[batch_index]]
            )
            extra["check_valid_states"][index, :, :] = check_valid_states[
                batch_index, :, :
            ].cpu()

        del input, state, output, state_prediction, state_error

    if model.model_output_dim > 0:
        model.setModelOutputs(old_model_outputs)

    del it
    util.clearGPUCache()
    extra["val_loss"] /= num_points
    # extra["num_bins"] = num_bins
    # extra["error_bin_ranges"] = error_hist_ranges
    # extra["loss_hist_range"] = loss_hist_range
    # extra["total_loss_hist_range"] = total_loss_hist_range
    extra["num_state_plots"] = random_state_plot_num

    # quantiles = quantile, time in traj, index
    extra["quantiles"] = QUANTILES
    extra["dt"] = system.dt
    extra["error_quantiles"] = np.zeros(
        (len(QUANTILES), 1 + int(system.T / T_delta), system.error_dim)
    )
    extra["loss_quantiles"] = np.zeros((len(QUANTILES), 1 + int(system.T / T_delta)))

    extra["total_loss_quantiles"] = util.computeQuantiles(extra["total_loss_bins"])

    for t in range(1 + int(system.T / T_delta)):
        extra["error_quantiles"][:, t, :] = util.computeQuantiles(
            extra["error_bins"][t]
        )
        extra["loss_quantiles"][:, t] = util.computeQuantiles(extra["loss_bins"][t])

    extra["prior_mean"] = prior_mean
    extra["prior_var"] = prior_var
    extra["prior_abs_mean"] = prior_abs_mean
    extra["prior_abs_var"] = prior_abs_var
    extra["prior_loss_mean"] = prior_loss_mean
    extra["prior_loss_var"] = prior_loss_var
    extra["all_state_error"] = all_state_error
    extra["all_predicted_states"] = all_predicted_states
    extra["all_states"] = all_states
    extra["total_forward"] = total_forward
    extra["all_inputs"] = all_inputs
    extra["all_model_output"] = all_model_outputs
    extra["model_output_labels"] = [e.name for e in model.ModelOutputs]
    extra["all_loss"] = all_loss
    extra["calib_z_bins"] = calib_z_bins
    extra["calib_mahal_bins"] = calib_mahal_bins
    return extra


@torch.inference_mode()
def test_main():
    args = util.loadSharedArgumentsModelLearning("test")
    uid = args["uid"]

    models = load_models(args["model"], args)

    if args.get("compile"):
        torch._dynamo.config.recompile_limit = len(models) * 4

    total_error = {}

    colors = {}
    first = True

    combined_results_path = None

    # this is a specific config using the hostname of the computer
    # computer_config = util.getComputerConfig()
    computer_config = {
        "max_batch_size": 10000,
        "test": {"num_workers": 10, "prefetch_factor": 5},
    }
    model_config_file = yaml.load(open(args["model"], "r"), Loader=yaml.FullLoader)

    # TODO what am I doing here??
    all_same_system = True
    if len(models.values()) == 0:
        raise RuntimeError("No models found")
    prev_system_type = list(models.values())[0].system.system_type
    system_types = [prev_system_type]
    for model in models.values():
        if model.system.system_type != prev_system_type:
            all_same_system = False
            system_types.append(model.system.system_type)

    if not all_same_system:
        print(
            f"WARNING: Got multiple systems in config, will put them in top level folder {system_types}"
        )

    for model_name, model in models.items():
        if first:
            first = False
            if all_same_system:
                combined_results_path = util.setup_results_dir(
                    model.system.name, uid, "", "test"
                )
            else:
                combined_results_path = util.setup_results_dir("", uid, "", "test")

            if "folder_name" in args and args["folder_name"] is not None:
                print(f"creating symbolic link with name ", args["folder_name"])
                sym_path = combined_results_path.parents[0]
                util.createSymLink(
                    combined_results_path,
                    sym_path
                    / str(args["folder_name"] + "_" + combined_results_path.name),
                    relative=True,
                )

            print(f"using combined results path {combined_results_path}")
            plotting.write_server_script(combined_results_path)
            server_proc = subprocess.Popen(
                [
                    sys.executable,
                    str(combined_results_path / "run.py"),
                    str(args["port"]),
                ],
                stdout=subprocess.DEVNULL,
            )
            atexit.register(_terminate_server, server_proc)
            plotting.write_test_landing_page(
                combined_results_path / "index.html",
                list(models.keys()),
                completed_names=[],
                individual_plots=args["individual_plots"],
            )

        print(f"\n***************************************************")
        print(
            f"model name: {model_name},"
            # f" size total params: {model.getParamsCount()},"
            # f" recurring params: {model.getRecurringParamsCount()}"
        )
        # util.clean_memory()
        # util.print_memory_usage(True)
        saved_model_path = pathlib.Path(model_config_file[model_name]["path"])
        if saved_model_path.exists():
            print(f"loading the model saved at {saved_model_path}")
            model.loadModel(pathlib.Path(model_config_file[model_name]["path"]))
        else:
            print(
                f"WARNING: not loading a model for {model_name} at {saved_model_path}"
            )

        if args["gpu"]:
            model.to("cuda")
        if args["compile"]:
            model.compile()

        loss_fn = load_loss_function(
            model.system, model_config_file[model_name]["loss_function"]
        )

        system = model.system

        assert pathlib.Path(args["dataset"]).suffix == ".hdf5"
        dataset_path = pathlib.Path(args["dataset"])

        if system.system_type not in str(dataset_path):
            raise RuntimeError(f"{system.system_type} should be in {dataset_path}")

        validation_dataset = DynamicsModelFullDataset(
            dataset_path, "validation", system.init_length, system.traj_length
        )
        pin_memory = False
        if "cuda" in str(torch.get_default_device()):
            pin_memory = True

        batch_size = computer_config["max_batch_size"]
        if (
            "training" in model_config_file[model_name].keys()
            and "batch" in model_config_file[model_name]["training"].keys()
        ):
            batch_size = min(
                batch_size, model_config_file[model_name]["training"]["batch"]
            )
        num_workers = computer_config["test"]["num_workers"]
        if (
            "training" in model_config_file[model_name].keys()
            and "num_workers" in model_config_file[model_name]["training"].keys()
        ):
            num_workers = min(
                num_workers, model_config_file[model_name]["training"]["num_workers"]
            )
        prefetch_factor = computer_config["test"]["prefetch_factor"]
        if (
            "training" in model_config_file[model_name].keys()
            and "prefetch_factor" in model_config_file[model_name]["training"].keys()
        ):
            prefetch_factor = min(
                prefetch_factor,
                model_config_file[model_name]["training"]["prefetch_factor"],
            )
        validation_dataloader = DataLoader(
            validation_dataset,
            num_workers=num_workers,
            prefetch_factor=prefetch_factor,
            batch_sampler=BatchSampler(
                SequentialSampler(validation_dataset),
                batch_size=batch_size,
                drop_last=False,
            ),
            pin_memory=pin_memory,
            multiprocessing_context="forkserver" if num_workers > 0 else None,
        )

        # TODO if no color generate a random one
        # colors[model_name] = model.color

        model_results_path = combined_results_path / model_name

        print(f"Running on total dataset size {len(validation_dataset)} {dataset_path}")
        print(f"***************************************************")

        error = run_model(
            validation_dataloader,
            model,
            loss_fn,
            model_results_path,
            random_state_plot_num=args["random_state_plot_num"],
            max_batches=args["max_batches"],
            num_bins=5000,
            T_delta=0.25,
        )

        total_error[model_name] = {}
        total_error[model_name] = error
        if args["individual_plots"]:
            total_error[model_name]["individual_plots_results"] = (
                plotting.plotly_plot_single_model_error(
                    model_results_path, error, model_name, model
                )
            )
        total_error[model_name]["dt"] = system.dt

        plotting.write_test_landing_page(
            combined_results_path / "index.html",
            list(models.keys()),
            completed_names=list(total_error.keys()),
            individual_plots=args["individual_plots"],
        )

        # model.GPUClean()

    print(f"\n***************************************************")
    print(f"combined model plots")
    print(f"***************************************************")
    plotting.plotly_plot_multiple_model_errors(
        combined_results_path, system, total_error, args, models=models
    )
    plotting.write_test_landing_page(
        combined_results_path / "index.html",
        list(models.keys()),
        completed_names=list(total_error.keys()),
        individual_plots=args["individual_plots"],
    )

    if not args["no_pause"]:
        print(
            f"\nResults: http://localhost:{args['port']}  (serving {combined_results_path})"
        )
        print("Press Ctrl+C to stop the server.")
        try:
            server_proc.wait()
        except KeyboardInterrupt:
            pass  # atexit handler will terminate the server on exit


if __name__ == "__main__":
    test_main()