skewed_plotting_test.py

# TODO generic version of the dataset creation
# TODO need to be able to sample parameters in a reasonable way
# TODO load in a file
# TODO create some automatic plotting as well
import torch

import util
import plotting
import numpy as np
from matplotlib.patches import Ellipse, Rectangle
from pathlib import Path
import os
import math
from scipy.stats import skewnorm

import scipy
from scipy.stats import multivariate_normal, norm

from distributions import (
    MultivariateGaussian,
    MultivariateSkewNormal,
    MultivariateExtendedSkewNormal,
)


def gaussian_pdf(data, mean, cov):
    # gaussian = multivariate_normal(mean, cov)
    # z_og = gaussian.pdf(data)# * 2.0
    # print(f"given z: {z_og[0, 0]}")

    flat_data = data.reshape(data.shape[0] ** 2, data.shape[-1])

    diff = (flat_data - mean)[:, :, None]

    z = (
        1
        / math.sqrt(math.pow(2.0 * np.pi, mean.shape[-1]) * np.linalg.det(cov))
        * np.exp(-0.5 * np.transpose(diff, axes=[0, 2, 1]) @ np.linalg.inv(cov) @ diff)
    )
    z = z.reshape(data.shape[0], data.shape[1])

    return z
    # print(f"got z:\n{z}")
    # print(f"max: {np.max(z)}")


def gaussian_sigma_points(mean, cov, axs, label=None, color="red"):
    L = np.linalg.cholesky(cov)
    print(f"gaussian standard L:\n{L}")
    sigma_points = np.zeros((4, 2))
    for i in range(2):
        sigma_points[i, :] = -L[..., i]
        sigma_points[i + 2, ...] = L[..., i]
    sigma_points_xy = sigma_points + mean
    if label is None:
        label = "gaussian sigma"
    axs.scatter(
        sigma_points_xy[:, 0],
        sigma_points_xy[:, 1],
        marker="+",
        color=color,
        label=label,
        s=400,
        zorder=0,
    )
    return sigma_points


def cov2cor(covariance):
    v = np.sqrt(np.diag(covariance))
    outer_v = np.outer(v, v)
    correlation = covariance / outer_v
    correlation[covariance == 0] = 0
    return correlation


def sun_pdf(data, ksi, Omega, Delta, tau, Gamma):
    if len(data.shape) > 1:
        flat_data = data.reshape(data.shape[0] ** 2, data.shape[-1])
    else:
        flat_data = data

    omega = np.sqrt(Omega.diagonal())
    Omega_bar = cov2cor(Omega)
    inv_Omega = np.linalg.inv(Omega_bar)

    bottom_cdf = multivariate_normal(np.zeros_like(tau), Gamma)
    bottom_cdf_value = bottom_cdf.cdf(tau)

    gaussian = multivariate_normal(ksi, Omega)
    gaussian_pdf = gaussian.pdf(data)  # should be offset

    top_cdf = multivariate_normal(
        np.zeros_like(tau), Gamma - Delta.T @ inv_Omega @ Delta
    )
    # print(f"wat?: {(flat_data-ksi).shape}")
    diff = (flat_data - ksi) / omega
    if len(data.shape) > 1:
        diff = diff[..., None]
        tau = tau[..., None]
    input = (tau + Delta.T @ inv_Omega @ diff).squeeze()
    top_cdf_values = top_cdf.cdf(input).reshape(gaussian_pdf.shape)

    return gaussian_pdf * top_cdf_values / bottom_cdf_value


# TODO turn into methods for clarity
# TODO try a different skewed distribution to see if it gives the right results


def one_d_example(pdf_file, levels):
    fig, axs = plotting.createPlot("skewed gaussian 1D", [""], horizontal=True)

    mean = 0
    sigma = 1.0
    alpha = 2.0
    tau = 0
    Gamma = 1.0

    omega = 1.0 / sigma  # TODO something is off here

    delta = alpha / math.sqrt(1 + alpha**2)
    print(f"got old delta: {delta}")
    # these are equivalent
    delta = (1 + alpha * sigma * alpha) ** -0.5 * sigma * alpha
    print(f"new delta: {delta}")

    new_alpha = (1 - delta * sigma ** (-1) * delta) ** (-0.5) * sigma ** (-1) * delta
    print(f"new alpha: {new_alpha}")

    # axs.axvline(mean, label="gaussian mean", color="#1f77b4")

    x_vals = np.arange(-3, 3, 0.05)
    pdf_calc = norm(loc=mean, scale=sigma)
    axs.plot(x_vals, pdf_calc.pdf(x_vals), alpha=1.0, label="gaussian pdf")

    true_sn_pdf = skewnorm.pdf(x_vals, alpha, loc=mean, scale=omega)
    # axs.plot(x_vals, true_sn_pdf, alpha=1.0, label="true skewed pdf")
    unit_cdf = norm(loc=0, scale=1.0)

    skewed_pdf = (2 / omega) * pdf_calc.pdf(x_vals) * unit_cdf.cdf(alpha * x_vals)
    np.testing.assert_allclose(skewed_pdf, true_sn_pdf)
    axs.plot(x_vals, skewed_pdf, alpha=1.0, label="skewed pdf")

    print(
        f"top cdf val for ext: {unit_cdf.cdf(tau * math.sqrt(1 + alpha**2) + alpha * x_vals)}"
    )
    extended_sn_pdf = (
        pdf_calc.pdf(x_vals)
        * unit_cdf.cdf(tau * math.sqrt(1 + alpha**2) + alpha * x_vals)
        / unit_cdf.cdf(tau)
    )
    axs.plot(x_vals, extended_sn_pdf, alpha=1.0, label="extended pdf")

    inv_Omega = 1 / sigma
    top_pdf_scale = Gamma - delta * inv_Omega * delta
    print(f"top cdf scale {top_pdf_scale}")
    top_cdf_calc = norm(loc=0, scale=top_pdf_scale)
    bottom_cdf_calc = norm(loc=0, scale=Gamma)
    print(
        f"top cdf val for sun: {top_cdf_calc.cdf(tau + delta * inv_Omega * (1.0 / omega) * (x_vals - mean))}"
    )
    sun_pdf = (
        pdf_calc.pdf(x_vals)
        * top_cdf_calc.cdf(tau + delta * inv_Omega * (1.0 / omega) * (x_vals - mean))
        / bottom_cdf_calc.cdf(tau)
    )
    axs.plot(x_vals, sun_pdf, alpha=1.0, label="SUN pdf")

    plotting.savePlotToPDF(pdf_file, fig, axs, legend_loc="upper left")

    x_vals = np.arange(-4, 4, 0.05)

    fig, axs = plotting.createPlot("skewed gaussian 1D", [""], horizontal=True)

    gaussian_sigma_points = np.zeros((5, 2))
    gaussian_sigma_points[0, 0] = mean
    gaussian_sigma_points[1, 0] = -sigma + mean
    gaussian_sigma_points[2, 0] = sigma + mean
    gaussian_sigma_points[3, 0] = -2 * sigma + mean
    gaussian_sigma_points[4, 0] = 2 * sigma + mean
    gaussian_sigma_points[:, 1] = 0.5 * (
        1
        + scipy.special.erf(
            (gaussian_sigma_points[:, 0] - mean) / (sigma * math.sqrt(2))
        )
    )
    axs.scatter(
        gaussian_sigma_points[:, 0],
        gaussian_sigma_points[:, 1],
        marker="D",
        color="black",
        label="gaussian sigma points",
        s=100,
    )
    print(f"got sigma points:\n{gaussian_sigma_points}")

    true_sn_pdf = skewnorm.pdf(x_vals, alpha, loc=mean, scale=omega)
    axs.plot(x_vals, true_sn_pdf, alpha=1.0, label="true skewed pdf")

    # np.testing.assert_allclose(
    #     true_sn_pdf, skewed_gaussian_pdf, atol=1.0e-5
    # )

    true_sn_mean, true_sn_var, true_sn_skew, true_sn_kurt = skewnorm.stats(
        alpha, moments="mvsk", loc=mean, scale=omega
    )

    print(
        f"true statistics: {true_sn_mean} {true_sn_var} {true_sn_skew} {true_sn_kurt}"
    )

    true_sn_cdf = skewnorm.cdf(x_vals, alpha, loc=mean, scale=omega)
    axs.plot(x_vals, true_sn_cdf, alpha=1.0, label="true skewed cdf")

    skewed_gaussian_cdf = 0.5 * (
        1 + scipy.special.erf((x_vals - mean) / (sigma * math.sqrt(2)))
    ) - 2 * scipy.special.owens_t(x_vals - mean, np.ones_like(x_vals) * alpha)
    # axs.plot(x_vals, skewed_gaussian_cdf, alpha=1.0, label="skewed cdf")

    # np.testing.assert_allclose(
    #     true_sn_cdf, skewed_gaussian_cdf, atol=1.0e-5
    # )

    mu_z = math.sqrt(2 / math.pi) * delta
    skewed_mean = mean + omega * mu_z
    np.testing.assert_allclose(true_sn_mean, skewed_mean)
    print(f"skewed mean: {skewed_mean}")
    axs.axvline(
        skewed_mean,
        label="skewed mean",
        color="#ff7f0e",
    )

    sigma_z = math.sqrt(1 - mu_z**2)
    skewed_var = (omega * sigma_z) ** 2
    np.testing.assert_allclose(true_sn_var, skewed_var)
    print(f"got skewed var: {skewed_var}")

    skewed_skew = 0.5 * (4 - math.pi) * (mu_z) ** 3 / sigma_z**3
    print(f"skewed skew: {skewed_skew}")
    np.testing.assert_allclose(true_sn_skew, skewed_skew)

    # print(f"my way for skew: {2 * mu_z**3 * omega**3 - mu_z * omega}")

    skewed_kurtosis = 3 + 2 * (math.pi - 3) * (mu_z) ** 4 / sigma_z**4
    print(f"got skewed kurtosis: {skewed_kurtosis}")
    np.testing.assert_allclose(true_sn_kurt + 3, skewed_kurtosis)

    # P = skewed_var
    # S = skewed_skew * (math.sqrt(P) ** 3)
    # K = skewed_kurtosis * (P**2)
    # print(f"got S: {S} and K: {K} P: {P}")
    # u = 0.5 * (
    #     -S * math.pow(math.sqrt(P), -3)
    #     + math.sqrt(4 * K * math.pow(P, -2) - 3 * S**2 * math.pow(P, -3))
    # )
    # v = u + S * math.pow(math.sqrt(P), -3)
    C = math.sqrt(skewed_var)
    print(f"\ngot C: {C}")

    print(f"kurt definition {skewed_kurtosis} >= {skewed_skew**2 + 1}")
    assert skewed_kurtosis >= skewed_skew**2 + 1
    # assert K > (S**2) / P
    u = 0.5 * (
        -(skewed_skew / C**3)
        + math.sqrt(4 * skewed_kurtosis / C**4 - 3 * (skewed_skew / C**3) ** 2)
    )
    # u = math.sqrt(skewed_var)
    v = u + skewed_skew / (C**3)
    print(f"got u {u} v {v}")

    assert u > 0
    assert v > 0

    w_2 = 1 / (v * (u + v))
    w_1 = (v / u) * w_2
    w_0 = 1 - w_1 - w_2
    print(f"got -u: {-u} and v: {v}")
    print(f"\ngot weights {w_0} {w_1} {w_2}")
    print(f"weights sum {w_0 + w_1 +w_2} = 1")
    print(f"-w_1 u + w_2 v = 0 = {-w_1 * u + w_2 * v}")
    print(f"w_1 u**2 + w_2 v**2 = 1 = {w_1 * u**2 + w_2 * v**2}")
    print(f"-w_1 u**3 + w_2 v**3 = {-w_1 * u **3 +w_2 * v**3} = {skewed_skew}")
    print(f"w_1 u**4 + w_2 v**4 = {w_1 * u **4 +w_2 * v**4} = {skewed_kurtosis}")

    skewed_sigma_points = np.zeros((3, 2))
    sigma_point_1 = skewed_mean - C * u
    sigma_point_2 = skewed_mean + C * v
    skewed_sigma_points[0, 0] = skewed_mean
    skewed_sigma_points[1, 0] = sigma_point_1
    skewed_sigma_points[2, 0] = sigma_point_2
    weights = np.array([w_0, w_1, w_2])
    skewed_sigma_points[:, 1] = skewnorm.cdf(
        skewed_sigma_points[:, 0], alpha, loc=mean, scale=omega
    )
    print(f"skewed sigma points\n{skewed_sigma_points}")
    print(
        f"sample mean: {np.sum(skewed_sigma_points[:, 0] * weights)} vs real mean {skewed_mean}"
    )
    print(
        f"sample var: {np.sum(weights * (skewed_sigma_points[:, 0] - skewed_mean)**2)} vs real var {skewed_var}"
    )

    # print(f"distance to center of sigma points: {skewed_sigma_points[0,1] - skewed_sigma_points[1,1]} vs "
    #       f"{skewed_sigma_points[2,1] - skewed_sigma_points[0, 1]}")
    print(
        f"distance to center of sigma points: {0.5 - skewed_sigma_points[1,1]} vs "
        f"{skewed_sigma_points[2,1] - 0.5}"
    )

    # TODO other method

    axs.scatter(
        skewed_sigma_points[:, 0],
        skewed_sigma_points[:, 1],
        marker="D",
        color="red",
        label="skewed sigma points",
        s=100,
    )

    axs.axvline(sigma_point_1, label="skewed u", color="green")
    axs.axvline(sigma_point_2, label="skewed v", color="purple")

    # TODO is there a way to use the augmented form?

    plotting.savePlotToPDF(pdf_file, fig, axs, legend_loc="upper left")


def two_d_example(pdf_file, levels):
    fig, axs = plotting.createPlot("standard gaussian 2D", [""], horizontal=True)

    # cov = np.array([
    #     [1, 0.88], [0.88, 1]
    # ])
    # cov = np.array([
    #     [1, 0.7], [0.7, 1]
    # ])
    cov = np.array([[2, -0.7], [-0.7, 2]])
    # cov = np.array([[2, 2], [2, 4]])

    mean = np.array([0.0, 0.0])
    # alpha = np.array([3, -5])
    alpha = np.array([5.0, -3.0])
    tau = np.array(
        [1.0]
    )  # negative moves density in the direction of the arrow, positive away from it

    axs.scatter(
        mean[0],
        mean[1],
        marker="D",
        color="red",
        label="gaussian mean",
        s=50,
        zorder=1,
    )

    x, y = np.mgrid[mean[0] - 4 : mean[0] + 4 : 0.05, mean[1] - 4 : mean[1] + 4 : 0.05]
    data = np.dstack((x, y))

    flat_data = data.reshape(data.shape[0] ** 2, data.shape[-1])

    mvn = MultivariateGaussian()

    gaussian_pdf_vals = mvn.pdf(
        torch.from_numpy(flat_data), torch.from_numpy(mean), torch.from_numpy(cov)
    ).reshape(data.shape[0], data.shape[1])
    CS = axs.contour(x, y, gaussian_pdf_vals, levels=levels)
    axs.clabel(CS, fontsize=10)
    gaussian_sp = gaussian_sigma_points(mean, cov, axs)
    # print(f"gaussian_sp:\n{gaussian_sp}")

    # this is extremely useless for plotting
    # gaussian_cdf = multivariate_normal(mean, cov).cdf(data)
    # print(f"cdf: {gaussian_cdf}")
    # CS = axs.contour(x, y, gaussian_cdf)
    # axs.clabel(CS, fontsize=10)

    def md(data, mean, cov):
        inv_cov = np.linalg.inv(cov).reshape(1, 2, 2).repeat(data.shape[0] ** 2, axis=0)
        flat_data = data.reshape(data.shape[0] ** 2, data.shape[-1])
        diff = (flat_data - mean).reshape(flat_data.shape[0], 2, 1)
        print(f"diff {diff.shape} vs {inv_cov.shape}")
        dist = np.sqrt(np.transpose(diff, [0, 2, 1]) @ inv_cov @ diff)
        return dist.reshape(data.shape[0], data.shape[1])

    # mahalanobis_dist = md(data, mean, cov)
    # CS = axs.contour(x, y, mahalanobis_dist, levels=[0.5, 1.0, 2.0, 3.0])
    # axs.clabel(CS, fontsize=10)

    axs.axis("equal")
    plotting.savePlotToPDF(pdf_file, fig, axs, legend_loc="upper left")
    fig, axs = plotting.createPlot("skewed gaussian 2D", [""], horizontal=True)

    CS = axs.contour(x, y, gaussian_pdf_vals, levels=levels, alpha=0.5, linestyles="--")
    axs.clabel(CS, fontsize=10)

    # plot the skew direction given by alpha
    axs.arrow(
        mean[0],
        mean[1],
        alpha[0] / np.linalg.norm(alpha),
        alpha[1] / np.linalg.norm(alpha),
        color="#ff7f0e",
        width=0.01,
        zorder=100,
    )

    skewed_obj = MultivariateSkewNormal()
    xi = torch.from_numpy(mean.copy())
    Omega = torch.from_numpy(cov.copy())

    skewed_mean = skewed_obj.mean(xi, Omega, torch.from_numpy(alpha)).numpy()
    print(f"skewed mean: {skewed_mean}")
    axs.scatter(
        skewed_mean[0],
        skewed_mean[1],
        marker="D",
        color="#1f77b4",
        label="skewed mean",
        s=50,
        zorder=1,
    )

    skewed_var = skewed_obj.variance(xi, Omega, torch.from_numpy(alpha)).numpy()
    print(f"got skewed var:\n{skewed_var}")

    gaussian_sigma_points(
        skewed_mean.squeeze(), skewed_var, axs, label="skewed gaussian chol"
    )

    gamma1 = skewed_obj.gamma1(xi, Omega, torch.from_numpy(alpha)).numpy()
    print(f"got gamma1: {gamma1.squeeze()}")

    skewed_pdf_vals = skewed_obj.pdf(
        torch.from_numpy(flat_data),
        torch.from_numpy(mean),
        torch.from_numpy(cov),
        torch.from_numpy(alpha),
    ).reshape(data.shape[0], data.shape[1])
    CS = axs.contour(x, y, skewed_pdf_vals, levels=levels)
    axs.clabel(CS, fontsize=10)

    C = np.linalg.cholesky(skewed_var)
    print(f"C:\n{C}")
    # print(f"C**3:\n{np.pow(C, 3)}")
    C_cubed_inv = (np.linalg.inv(np.pow(C, 3)) @ gamma1).squeeze()
    print(f"C_cubed_inv:\n{C_cubed_inv}")

    kurt_approx = np.pow(C, 4) @ np.pow(C_cubed_inv, 2)
    kurt_approx = kurt_approx * 1.1
    # kurt_approx = np.ones(2) * 0.5885733 / 2 + 3
    print(f"kurt approx: {kurt_approx}")

    C_four_inv = (np.linalg.inv(np.pow(C, 4)) @ kurt_approx).squeeze()
    # print(f"C_four_inv:\n{C_four_inv}")

    u = 0.5 * (-C_cubed_inv + np.sqrt(4 * C_four_inv - 3 * np.pow(C_cubed_inv, 2)))
    print(f"got u: {u}")
    v = u + C_cubed_inv
    print(f"got v: {v}")

    skewed_sigma_points = np.zeros((5, 2))
    print(f"wat? {skewed_sigma_points[1, 0]} {u[0]}")

    print(f"diag u:\n{np.diag(u)}")
    skewed_sigma_points[1:3, :] = (-C @ np.diag(u)).T
    skewed_sigma_points[3:, :] = (C @ np.diag(v)).T
    print(f"got skewed sigma:\n{skewed_sigma_points}")
    skewed_sigma_points = skewed_sigma_points + skewed_mean.squeeze()
    print(f"sigma points with mean:\n{skewed_sigma_points}")
    axs.scatter(
        skewed_sigma_points[:, 0],
        skewed_sigma_points[:, 1],
        marker="+",
        color="purple",
        label="skewed sigma points",
        s=400,
        zorder=0,
    )

    # TODO skewed mode value

    # correct_skewed_sigma_points = np.array([[0.7570, -0.6930],
    #                                [0.6084, -0.6629],
    #                                [0.9065, -1.5089],
    #                                [2.5278, -1.0524],
    #                                [0.7445, -0.6246]])
    # print(f"{correct_skewed_sigma_points.shape}")
    # print(f"corrected skewed sigma points:\n{correct_skewed_sigma_points}")
    # axs.scatter(
    #     correct_skewed_sigma_points[:, 0],
    #     correct_skewed_sigma_points[:, 1],
    #     marker="+",
    #     color="g",
    #     label="corrected skewed sigma points",
    #     s=400,
    #     zorder=0,
    # )

    axs.axis("equal")
    plotting.savePlotToPDF(pdf_file, fig, axs, legend_loc="upper left")
    fig, axs = plotting.createPlot("extended skewed gaussian 2D", [""], horizontal=True)

    CS = axs.contour(x, y, skewed_pdf_vals, levels=levels, alpha=0.5, linestyles="--")
    axs.clabel(CS, fontsize=10)

    extended_skewed_obj = MultivariateExtendedSkewNormal()

    extended_skewed_pdf_vals = extended_skewed_obj.pdf(
        torch.from_numpy(flat_data),
        torch.from_numpy(mean),
        torch.from_numpy(cov),
        torch.from_numpy(alpha),
        torch.from_numpy(tau),
    ).reshape(data.shape[0], data.shape[1])

    CS = axs.contour(x, y, extended_skewed_pdf_vals, levels=levels)
    axs.clabel(CS, fontsize=10)

    axs.scatter(
        skewed_mean[0],
        skewed_mean[1],
        marker="D",
        color="purple",
        label="skewed mean",
        s=50,
        zorder=1,
    )

    extended_skewed_mean = extended_skewed_obj.mean(
        xi, Omega, torch.from_numpy(alpha), torch.from_numpy(tau)
    ).numpy()
    print(f"extended skewed mean: {skewed_mean}")
    axs.scatter(
        extended_skewed_mean[0],
        extended_skewed_mean[1],
        marker="D",
        color="#1f77b4",
        label="extended skewed mean",
        s=50,
        zorder=1,
    )

    extended_skewed_var = extended_skewed_obj.variance(
        xi, Omega, torch.from_numpy(alpha), torch.from_numpy(tau)
    ).numpy()
    print(f"got extended skewed var:\n{extended_skewed_var}")

    # mahalanobis_dist = md(data, skewed_mean.squeeze(), skewed_var)
    # CS = axs.contour(x, y, mahalanobis_dist, levels=[0.5, 1.0, 2.0, 3.0])
    # axs.clabel(CS, fontsize=10)

    # axs.set_xlim((-1.0, 5.0))
    axs.axis("equal")
    plotting.savePlotToPDF(pdf_file, fig, axs, legend_loc="upper left")


def two_d_example_SUN(pdf_file, levels):
    fig, axs = plotting.createPlot("SUN 2D", [""], horizontal=True)

    mean = np.array([0, 0])
    cov = np.array([[1, 0.88], [0.88, 1]])
    # a3
    tau = np.array([0, 0])
    Omega = np.copy(cov)
    Gamma = np.array([[1, 0.96], [0.96, 1]])
    Delta = np.array([[0.61, 0.76], [0.88, 0.97]])

    # a1
    # tau = np.array([0])
    # Omega = np.copy(cov)
    # Gamma = np.array([1])
    # Delta = np.array([0.88, 0.61]).reshape(2, 1)

    axs.scatter(
        mean[0],
        mean[1],
        marker="D",
        color="red",
        label="gaussian mean",
        s=50,
        zorder=1,
    )

    x, y = np.mgrid[mean[0] - 4 : mean[0] + 4 : 0.05, mean[1] - 4 : mean[1] + 4 : 0.05]
    data = np.dstack((x, y))

    z = gaussian_pdf(data, mean, cov)
    CS = axs.contour(x, y, z, levels=levels, linestyles="dashed", alpha=0.5)
    axs.clabel(CS, fontsize=10)
    gaussian_sigma_points(mean, cov, axs)

    Omega_bar = cov2cor(Omega)
    inv_Omega = np.linalg.inv(Omega_bar)
    alpha = (inv_Omega @ Delta).squeeze()
    print(f"alpha: {alpha}")
    # TODO can be multivariate
    if np.linalg.norm(alpha) > 0:
        if len(alpha.shape) > 1:
            for i in range(alpha.shape[0]):
                axs.arrow(
                    mean[0],
                    mean[1],
                    alpha[i, 0] / np.linalg.norm(alpha),
                    alpha[i, 1] / np.linalg.norm(alpha),
                    color="#ff7f0e",
                    width=0.01,
                    zorder=100,
                )
        else:
            axs.arrow(
                mean[0],
                mean[1],
                alpha[0] / np.linalg.norm(alpha),
                alpha[1] / np.linalg.norm(alpha),
                color="#ff7f0e",
                width=0.01,
                zorder=100,
            )

    print(f"sun pdf: {sun_pdf(np.array([0, 0]), mean, Omega, Delta, tau, Gamma)}")

    z = sun_pdf(data, mean, Omega, Delta, tau, Gamma)
    CS = axs.contour(x, y, z, levels=levels)
    axs.clabel(CS, fontsize=10)

    # TODO in the R code
    # take the ESN and convert to SUN
    # compute the third moment and make sure that is a tensor and is correct
    # make sure we can recreate that without using the SUN density

    # TODO compute the mean, var, skewness values for matrix
    # TODO use that for sigma point calculation with janky setup for u somehow so we do not need the kurtosis

    # moments for the density
    omega = np.sqrt(Omega.diagonal())
    Omega_bar = cov2cor(Omega)
    inv_Omega = np.linalg.inv(Omega_bar)
    inv_Gamma = np.linalg.inv(Gamma)
    Lambda = (omega * Delta) @ inv_Gamma
    print(f"Lambda:\n{Lambda}")

    axs.axis("equal")
    plotting.savePlotToPDF(pdf_file, fig, axs, legend_loc="upper left")


if __name__ == "__main__":
    full_path = Path(os.path.abspath(__file__)).parent / "skew_test"
    if full_path.exists():
        full_path.unlink()
    pdf_file = plotting.createPDF(full_path)

    # levels = [0.05, 0.1, 0.15, 0.2, 0.3]
    levels = [0.1, 0.2, 0.3, 0.4, 0.5]
    levels = None

    # one_d_example(pdf_file, levels)
    two_d_example(pdf_file, levels)
    # print(f"\n\n SUN example")
    # two_d_example_SUN(pdf_file, levels)

    pdf_file.close()