gempy3_first_model_generative_full_covariance.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Import all the libraries and packages "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import joblib\n",
    "import seaborn as sns\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "import torch\n",
    "import torch.nn.functional as F\n",
    "\n",
    "import pyro\n",
    "import pyro.distributions as dist\n",
    "from pyro.infer import MCMC, NUTS, Predictive, EmpiricalMarginal\n",
    "from pyro.infer.autoguide import init_to_mean, init_to_median, init_to_value\n",
    "from pyro.infer.inspect import get_dependencies\n",
    "from pyro.infer import SVI, TraceEnum_ELBO, config_enumerate, infer_discrete\n",
    "\n",
    "import gempy as gp\n",
    "import gempy_engine\n",
    "import gempy_viewer as gpv\n",
    "from gempy_engine.core.backend_tensor import BackendTensor\n",
    "import arviz as az\n",
    "from gempy_probability.plot_posterior import default_red, default_blue, PlotPosterior\n",
    "\n",
    "from sklearn.mixture import GaussianMixture\n",
    "from sklearn.mixture import BayesianGaussianMixture\n",
    "from sklearn.cluster import KMeans\n",
    "from scipy.stats import multivariate_normal, norm"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Get the path of all the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "filename_a = './Fw__Hyperspectral_datasets_from_the_KSL_cores/CuSp131.pkl'\n",
    "filename_b = './Fw__Hyperspectral_datasets_from_the_KSL_cores/CuSp133.pkl'\n",
    "filename_c = './Fw__Hyperspectral_datasets_from_the_KSL_cores/CuSp136.pkl'"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Load and get the data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "with open(filename_a, 'rb') as myfile:\n",
    "    a= joblib.load(myfile)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Description about Hyperspectral data\n",
    "## Data is obtained by scannig the core data from different boreholes using hyperspectral sensors. There were around 450 channels for each pixels initially. It was preprocessed and seperated based on 10 different types of rocks. In each of the file we have \"X\", \"Y\", \"Z\" coordinates points corresponding to sensors and corresponding to each rock type we have a transformed RGB correspondence information. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# The variability in \"X\" and \"Y\" are much smaller as compared to \"Z\" direction in case of birehole information. \n",
    "# Therefore, we are trying to build our model considering the \"Z\" direction mostly. \n",
    "# get the z-cordinates of borehole\n",
    "zz = a['XYZ'][:,2]\n",
    "print(zz.shape)\n",
    "# sort the z-cordinates\n",
    "ixx = np.argsort( zz )\n",
    "# mask if values is less than some specified value\n",
    "mask = zz[ixx] < -700\n",
    "ah = a['BR_Anhydrite'][:,0] # correlates to \"anhydrite index\" derived from hyperspectral \n",
    "position_cord , hsi_data = zz[ixx][mask], ah[ixx][mask]/255 # To normalize the hyperspectral spectra, divide it with 255. \n",
    "plt.plot( position_cord , hsi_data )\n",
    "print(position_cord.shape)\n",
    "\n",
    "# define breakpoints \n",
    "# In general it is very difficult to define the breakpoints in the plot. Prepocessing of hyperspectral data is itself a very difficult task\n",
    "# becasue of high correaltion, high dimensional and noisy data. \n",
    "brk1 = -845 \n",
    "brk2 = -825 \n",
    "\n",
    "plt.axvline( brk1, color='r' )\n",
    "plt.axvline( brk2, color='g')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "hsi_data.reshape(1,-1)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Since it is difficult to classify a hyperspectral data in general. We can apply different classical clustering methods to have some starting guess"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "X = np.concatenate(((position_cord.reshape((-1,1))/1000.0), hsi_data.reshape((-1,1))), axis=1)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#gm2 = BayesianGaussianMixture(n_components=3,covariance_type=\"full\", random_state=0).fit(hsi_data.reshape(-1,1))\n",
    "gm2 = BayesianGaussianMixture(n_components=3,covariance_type=\"full\", random_state=0).fit(X)\n",
    "gm2.means_ , gm2.covariances_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loc_mean = torch.tensor(gm2.means_)\n",
    "loc_cov  = torch.tensor(gm2.covariances_)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "correct_order = [0,2,1]\n",
    "loc_mean, loc_cov = loc_mean[correct_order], loc_cov[correct_order]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import multivariate_normal\n",
    "\n",
    "# Create a grid of points\n",
    "x, y = np.meshgrid(np.linspace(-0.9, -0.7, 100), np.linspace(-2, 2, 100))\n",
    "pos = np.dstack((x, y))  # Combine x and y grids into a (100, 100, 2) array\n",
    "for i in range(3):\n",
    "    # Create a multivariate normal distribution\n",
    "    rv = multivariate_normal(loc_mean[i], loc_cov[i])\n",
    "\n",
    "    # Calculate PDF values for each point in the grid\n",
    "    pdf_values = rv.pdf(pos)\n",
    "\n",
    "    # Plot the Gaussian distribution using contour plot\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    plt.contourf(x, y, pdf_values, cmap='viridis')\n",
    "    plt.colorbar(label='Probability Density')\n",
    "    plt.xlabel('spatial')\n",
    "    plt.ylabel('spectral')\n",
    "    plt.title('2D Gaussian Distribution')\n",
    "    plt.scatter(loc_mean[i][0], loc_mean[i][1], color='red', label='Mean')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import multivariate_normal, norm\n",
    "plt.figure(figsize=(8, 6))\n",
    "# Create a grid of points\n",
    "x, y = np.meshgrid(np.linspace(-0.9, -0.7, 100), np.linspace(-2, 2, 100))\n",
    "pos = np.dstack((x, y))  # Combine x and y grids into a (100, 100, 2) array\n",
    "for i in range(3):\n",
    "    # Create a multivariate normal distribution\n",
    "    rv = multivariate_normal(loc_mean[i], loc_cov[i])\n",
    "\n",
    "    # Calculate PDF values for each point in the grid\n",
    "    pdf_values = rv.pdf(pos)\n",
    "\n",
    "    # Plot the Gaussian distribution using contour plot\n",
    "    \n",
    "    plt.contour(x, y, pdf_values, cmap='viridis')\n",
    "    plt.scatter(loc_mean[i][0], loc_mean[i][1], label='Mean'+str(i+1))\n",
    "#plt.colorbar(label='Probability Density')\n",
    "plt.xlabel('spatial')\n",
    "plt.ylabel('spectral')\n",
    "plt.title('2D Gaussian Distribution')\n",
    "\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_gmm_label2 = gm2.predict(X)\n",
    "print(y_gmm_label2)\n",
    "y_gmm_label_arranged2 = np.zeros_like(y_gmm_label2)\n",
    "y_gmm_label_arranged2[y_gmm_label2 == 1] = 3\n",
    "y_gmm_label_arranged2[y_gmm_label2 == 0] = 2\n",
    "y_gmm_label_arranged2[y_gmm_label2 == 2] = 1\n",
    "y_gmm_label_arranged2 = torch.Tensor(y_gmm_label_arranged2)\n",
    "y_gmm_label_arranged2 "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import matplotlib.pyplot as plt\n",
    "# Define colors for each label\n",
    "colors = ['r', 'g', 'b']\n",
    "labels = y_gmm_label_arranged2\n",
    "\n",
    "# Plot the dataset with different colors for each label\n",
    "plt.figure(figsize=(8, 6))\n",
    "for label_val, color in zip([1,2,3], colors):\n",
    "    plt.scatter(position_cord[labels == label_val], hsi_data[labels == label_val], c=color, label=f'Label {label_val}')\n",
    "\n",
    "plt.xlabel('z')\n",
    "plt.ylabel('hsi_data')\n",
    "plt.title('2D Dataset with Label Information')\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# A = np.zeros_like(a['BR_Anhydrite'][:, 0])\n",
    "# B = np.zeros_like(a['BR_Anhydrite'][:, 0])\n",
    "# C = np.zeros_like(a['BR_Anhydrite'][:, 0])\n",
    "\n",
    "A = np.zeros_like(hsi_data)\n",
    "B = np.zeros_like(hsi_data)\n",
    "C = np.zeros_like(hsi_data)\n",
    "\n",
    "# Get indices where the mask is True\n",
    "# indices_A = np.where(mask & (zz[ixx] < brk1))\n",
    "# indices_B = np.where(mask & (zz[ixx] > brk1) & (zz[ixx] < brk2))\n",
    "# indices_C = np.where(mask & (zz[ixx] > brk2))\n",
    "shift =0\n",
    "# Get indices where the mask is True\n",
    "indices_A = np.where(mask & (zz[ixx] < (brk1+shift)))\n",
    "indices_B = np.where(mask & (zz[ixx] > (brk1+shift)) & (zz[ixx] < (brk2+shift)))\n",
    "indices_C = np.where(mask & (zz[ixx] > (brk2+shift)))\n",
    "\n",
    "y_obs_label = torch.ones(234)\n",
    "y_obs_label[indices_A] =3\n",
    "y_obs_label[indices_B] = 2\n",
    "y_obs_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "matched_label2 = y_obs_label == y_gmm_label_arranged2\n",
    "(matched_label2.sum()/len(y_obs_label)) *100"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#y_obs_label = y_gmm_label_arranged\n",
    "y_obs_label = hsi_data.reshape((-1,1))\n",
    "y_obs_label2 = X"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test = gp.create_geomodel(\n",
    "    project_name='Gempy_abc_Test',\n",
    "    extent=[0, 1000, -500, 500, -900, -700],\n",
    "    resolution=[100,100,100],\n",
    "    refinement=3,\n",
    "    structural_frame= gp.data.StructuralFrame.initialize_default_structure()\n",
    "    )\n",
    "p2d = gpv.plot_2d(geo_model_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.grid.active_grids_bool"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gp.add_surface_points(\n",
    "    geo_model=geo_model_test,\n",
    "    x=[100.0, 900.0],\n",
    "    y=[0.0, 0.0],\n",
    "    z=[brk1, brk1],\n",
    "    elements_names=['surface1', 'surface1']\n",
    ")\n",
    "gpv.plot_2d(geo_model_test, cell_number=11)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gp.add_orientations(\n",
    "    geo_model=geo_model_test,\n",
    "    x=[500],\n",
    "    y=[0.0],\n",
    "    z=[brk1],\n",
    "    elements_names=['surface1'],\n",
    "    pole_vector=[[0, 0, 1]]\n",
    ")\n",
    "\n",
    "gpv.plot_2d(geo_model_test, cell_number=5)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.update_transform(gp.data.GlobalAnisotropy.NONE)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gp.compute_model(geo_model_test, engine_config=gp.data.GemPyEngineConfig())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.interpolation_options.kernel_options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpv.plot_2d(geo_model_test, cell_number=[5])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.structural_frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "element2 = gp.data.StructuralElement(\n",
    "    name='surface2',\n",
    "    color=next(geo_model_test.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([100.0, 900.0]),\n",
    "        y=np.array([0.0, 0.0]),\n",
    "        z=np.array([brk2, brk2]),\n",
    "        names='surface2'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test.structural_frame.structural_groups[0].append_element(element2)\n",
    "# Compute and visualize the updated model:\n",
    "gp.compute_model(geo_model_test)\n",
    "gpv.plot_2d(geo_model_test, cell_number=5, legend='force')\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#gpv.plot_3d(geo_model_test, image=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.structural_frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.structural_frame.structural_groups[0].elements[0], geo_model_test.structural_frame.structural_groups[0].elements[1] = \\\n",
    "geo_model_test.structural_frame.structural_groups[0].elements[1], geo_model_test.structural_frame.structural_groups[0].elements[0]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.structural_frame"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gpv.plot_2d(geo_model_test, cell_number=5, legend='force')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Create a custome grid where the observed data information is available "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_loc = 300\n",
    "y_loc = 0\n",
    "z_loc = position_cord\n",
    "xyz_coord = np.array([[x_loc, y_loc, z] for z in z_loc])\n",
    "gp.set_custom_grid(geo_model_test.grid, xyz_coord=xyz_coord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_coords_copy_test = geo_model_test.interpolation_input.surface_points.sp_coords.copy()\n",
    "geo_model_test.transform.apply_inverse(sp_coords_copy_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gp.add_surface_points(\n",
    "    geo_model=geo_model_test,\n",
    "    x=[x_loc, x_loc],\n",
    "    y=[0.0, 0.0],\n",
    "    z=[brk1, brk2],\n",
    "    elements_names=['surface1', 'surface2']\n",
    ")\n",
    "gpv.plot_2d(geo_model_test, cell_number=1)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.surface_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gp.compute_model(geo_model_test)\n",
    "gpv.plot_2d(geo_model_test, cell_number=5, legend='force')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_coords_copy_test = geo_model_test.interpolation_input.surface_points.sp_coords.copy()\n",
    "geo_model_test.transform.apply_inverse(sp_coords_copy_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Change the backend to PyTorch for probabilistic modeling\n",
    "BackendTensor.change_backend_gempy(engine_backend=gp.data.AvailableBackends.PYTORCH)\n",
    "# Set random seed for PyTorch backend\n",
    "torch.manual_seed(42)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.interpolation_options.uni_degree = 0\n",
    "geo_model_test.interpolation_options.mesh_extraction = False\n",
    "geo_model_test.interpolation_options.sigmoid_slope = 1100."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#geo_model_test.solutions.octrees_output[0].last_output_center"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_grid_values = geo_model_test.solutions.octrees_output[0].last_output_center.custom_grid_values\n",
    "custom_grid_values.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_obs_label.shape, y_obs_label2.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.surface_points"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.orientations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.transform.apply_inverse(sp_coords_copy_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "# Set random seed for Pyro\n",
    "pyro.set_rng_seed(42)\n",
    "\n",
    "@config_enumerate\n",
    "def model_test(y_obs_label):\n",
    "    \"\"\"\n",
    "    This Pyro model represents the probabilistic aspects of the geological model.\n",
    "    It defines a prior distribution for the top layer's location and\n",
    "    computes the thickness of the geological layer as an observed variable.\n",
    "    \"\"\"\n",
    "    # Define prior for the top layer's location\n",
    "    prior_mean_surface_1 = sp_coords_copy_test[2, 2]\n",
    "    prior_mean_surface_2 = sp_coords_copy_test[5, 2]\n",
    "    \n",
    "    \n",
    "\n",
    "    \n",
    "    mu_surface_1 = pyro.sample('mu_1', dist.Normal(prior_mean_surface_1, torch.tensor(0.02, dtype=torch.float64)))\n",
    "    mu_surface_2 = pyro.sample('mu_2', dist.Normal(prior_mean_surface_2, torch.tensor(0.02, dtype=torch.float64)))\n",
    "    \n",
    "    # Ensure that mu_surface_1 is greater than mu_surface_2\n",
    "    pyro.sample('condition', dist.Delta(torch.tensor(1.0, dtype=torch.float64)), obs=(mu_surface_1 > mu_surface_2))\n",
    "    # Update the model with the new top layer's location\n",
    "    interpolation_input = geo_model_test.interpolation_input\n",
    "    \n",
    "    \n",
    "    interpolation_input.surface_points.sp_coords = torch.index_put(\n",
    "        interpolation_input.surface_points.sp_coords,\n",
    "        (torch.tensor([2]), torch.tensor([2])),\n",
    "        mu_surface_1\n",
    "    )\n",
    "    interpolation_input.surface_points.sp_coords = torch.index_put(\n",
    "        interpolation_input.surface_points.sp_coords,\n",
    "        (torch.tensor([5]), torch.tensor([2])),\n",
    "        mu_surface_2\n",
    "    )\n",
    "    \n",
    "\n",
    "    \n",
    "    \n",
    "    # # Compute the geological model\n",
    "    geo_model_test.solutions = gempy_engine.compute_model(\n",
    "        interpolation_input=interpolation_input,\n",
    "        options=geo_model_test.interpolation_options,\n",
    "        data_descriptor=geo_model_test.input_data_descriptor,\n",
    "        geophysics_input=geo_model_test.geophysics_input,\n",
    "    )\n",
    "    \n",
    "    # Compute and observe the thickness of the geological layer\n",
    "    \n",
    "    custom_grid_values = geo_model_test.solutions.octrees_output[0].last_output_center.custom_grid_values\n",
    "    \n",
    "    lambda_ = 4\n",
    "    class_label = F.softmax(- lambda_ * (torch.tensor([1,2,3], dtype=torch.float64) - custom_grid_values.reshape(-1,1))**2, dim=1)\n",
    "    \n",
    "    sample =[]\n",
    "    for i in range(loc_mean.shape[0]):\n",
    "        sample_data = pyro.sample(\"sample_data\"+str(i+1), dist.MultivariateNormal(loc=loc_mean[i],covariance_matrix=loc_cov[i]))\n",
    "        sample.append(sample_data)\n",
    "    sample_tesnor = torch.stack(sample, dim=0)\n",
    "    \n",
    "    \n",
    "    with pyro.plate('N='+str(y_obs_label.shape[0]), y_obs_label.shape[0]):\n",
    "        assignment = pyro.sample(\"assignment\", dist.Categorical(class_label))\n",
    "        \n",
    "        obs = pyro.sample(\"obs\", dist.MultivariateNormal(loc=sample_tesnor[assignment],covariance_matrix=loc_cov[assignment]), obs=y_obs_label)\n",
    "        \n",
    "    return obs\n",
    "    \n",
    "dependencies = get_dependencies(model_test, model_args=(torch.tensor(X),))\n",
    "pyro.render_model(model_test, model_args=(torch.tensor(X),),render_distributions=True)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_obs_label =torch.tensor(X)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "\n",
    "model_test(y_obs_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior = Predictive(model_test, num_samples=100)(y_obs_label)\n",
    "\n",
    "# Key to avoid\n",
    "avoid_key = 'condition'\n",
    "\n",
    "# Create sub-dictionary without the avoid_key\n",
    "prior = dict((key, value) for key, value in prior.items() if key != avoid_key)\n",
    "\n",
    "data = az.from_pyro(prior=prior)\n",
    "az.plot_trace(data.prior)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "pyro.primitives.enable_validation(is_validate=True)\n",
    "nuts_kernel = NUTS(model_test, step_size=0.0085, adapt_step_size=True, target_accept_prob=0.9, max_tree_depth=10, init_strategy=init_to_mean)\n",
    "mcmc = MCMC(nuts_kernel, num_samples=200, warmup_steps=50, disable_validation=False)\n",
    "mcmc.run(y_obs_label)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "posterior_samples = mcmc.get_samples()\n",
    "posterior_predictive = Predictive(model_test, posterior_samples)(y_obs_label)\n",
    "data = az.from_pyro(posterior=mcmc, prior=prior, posterior_predictive=posterior_predictive)\n",
    "az.plot_trace(data)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "posterior_samples"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "loc_mean_posterior, loc_cov_posterior =[], []\n",
    "for key, values in posterior_samples.items():\n",
    "    print(key)\n",
    "    if key == \"sample_data1\":\n",
    "        mean = values.mean(dim=0)\n",
    "        cov = np.cov(values.detach().numpy(), rowvar=False)\n",
    "        print(\"mean\\n\",mean)\n",
    "        print(\"cov\\n\", cov)\n",
    "        loc_mean_posterior.append(mean.detach().numpy())\n",
    "        loc_cov_posterior.append(cov)\n",
    "    elif key == \"sample_data2\":\n",
    "        mean = values.mean(dim=0)\n",
    "        cov = np.cov(values.detach().numpy(), rowvar=False)\n",
    "        print(\"mean\\n\",mean)\n",
    "        print(\"cov\\n\", cov)\n",
    "        loc_mean_posterior.append(mean.detach().numpy())\n",
    "        loc_cov_posterior.append(cov)\n",
    "    elif key == \"sample_data3\":\n",
    "        mean = values.mean(dim=0)\n",
    "        cov = np.cov(values.detach().numpy(), rowvar=False)\n",
    "        print(\"mean\\n\",mean)\n",
    "        print(\"cov\\n\", cov)\n",
    "        loc_mean_posterior.append(mean.detach().numpy())\n",
    "        loc_cov_posterior.append(cov)\n",
    "    elif key == \"sigma_data\":\n",
    "        print(\"mean\\n\",values.mean(dim=0), \"\\nstd\\n\", values.std(dim=0))\n",
    "    else:\n",
    "        print(\"mean\\n\",values.mean(), \"\\nstd\\n\", values.std())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create a grid of points\n",
    "x, y = np.meshgrid(np.linspace(-0.9, -0.7, 100), np.linspace(-0.2, 1.2, 100))\n",
    "pos = np.dstack((x, y))  # Combine x and y grids into a (100, 100, 2) array\n",
    "for i in range(3):\n",
    "    # Create a multivariate normal distribution\n",
    "    rv = multivariate_normal(loc_mean_posterior[i], loc_cov_posterior[i])\n",
    "\n",
    "    # Calculate PDF values for each point in the grid\n",
    "    pdf_values = rv.pdf(pos)\n",
    "\n",
    "    # Plot the Gaussian distribution using contour plot\n",
    "    plt.figure(figsize=(8, 6))\n",
    "    plt.contourf(x, y, pdf_values, cmap='viridis')\n",
    "    plt.colorbar(label='Probability Density')\n",
    "    plt.xlabel('spatial')\n",
    "    plt.ylabel('spectral')\n",
    "    plt.title('2D Gaussian Distribution')\n",
    "    plt.scatter(loc_mean_posterior[i][0], loc_mean_posterior[i][1], color='red', label='Mean')\n",
    "    plt.legend()\n",
    "    plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(8, 6))\n",
    "# Create a grid of points\n",
    "x, y = np.meshgrid(np.linspace(-0.9, -0.7, 1000), np.linspace(-2, 2, 1000))\n",
    "pos = np.dstack((x, y))  # Combine x and y grids into a (100, 100, 2) array\n",
    "for i in range(3):\n",
    "    # Calculate distances from the mean for all points in the grid\n",
    "    #distances = np.sqrt((x - loc_mean_posterior[i][0])**2 + (y - loc_mean_posterior[i][1])**2)\n",
    "    # Create a multivariate normal distribution\n",
    "    rv = multivariate_normal(loc_mean_posterior[i], loc_cov_posterior[i])\n",
    "    # Calculate PDF values for each point in the grid\n",
    "    pdf_values = rv.pdf(pos)\n",
    "    \n",
    "   \n",
    "    #pdf_values[distances >0.01]=0\n",
    "    # Plot the Gaussian distribution using contour plot\n",
    "    plt.contour(x, y, pdf_values, extend='min', cmap='viridis')\n",
    "    plt.scatter(loc_mean_posterior[i][0], loc_mean_posterior[i][1], label='Mean'+str(i+1))\n",
    "#plt.colorbar(label='Probability Density')\n",
    "plt.xlabel('spatial')\n",
    "plt.ylabel('spectral')\n",
    "plt.title('2D Gaussian Distribution')\n",
    "\n",
    "plt.legend()\n",
    "plt.grid(True)\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "data.posterior"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_density(\n",
    "    data=[data.posterior, data.prior],\n",
    "    shade=.9,\n",
    "    var_names=['mu_1'],\n",
    "    data_labels=[\"Posterior Predictive\", \"Prior Predictive\"],\n",
    "    colors=[default_red, default_blue],\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "az.plot_density(\n",
    "    data=[data.posterior, data.prior],\n",
    "    shade=.9,\n",
    "    var_names=['mu_2'],\n",
    "    data_labels=[\"Posterior Predictive\", \"Prior Predictive\"],\n",
    "    colors=[default_red, default_blue],\n",
    ")\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Update the model with the new top layer's location\n",
    "interpolation_input = geo_model_test.interpolation_input\n",
    "interpolation_input.surface_points.sp_coords = torch.index_put(\n",
    "    interpolation_input.surface_points.sp_coords,\n",
    "    (torch.tensor([2]), torch.tensor([2])),\n",
    "    posterior_samples[\"mu_1\"].mean()\n",
    ")\n",
    "interpolation_input.surface_points.sp_coords = torch.index_put(\n",
    "    interpolation_input.surface_points.sp_coords,\n",
    "    (torch.tensor([5]), torch.tensor([2])),\n",
    "    posterior_samples[\"mu_2\"].mean()\n",
    ")\n",
    "\n",
    "#print(\"interpolation_input\",interpolation_input.surface_points.sp_coords)\n",
    "\n",
    "# # Compute the geological model\n",
    "geo_model_test.solutions = gempy_engine.compute_model(\n",
    "    interpolation_input=interpolation_input,\n",
    "    options=geo_model_test.interpolation_options,\n",
    "    data_descriptor=geo_model_test.input_data_descriptor,\n",
    "    geophysics_input=geo_model_test.geophysics_input,\n",
    ")\n",
    "\n",
    "gpv.plot_2d(geo_model_test, cell_number=5,legend='force')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_coords_copy_test2 =interpolation_input.surface_points.sp_coords\n",
    "sp_coords_copy_test2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_cord= geo_model_test.transform.apply_inverse(sp_coords_copy_test2.detach().numpy())\n",
    "sp_cord"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#gpv.plot_3d(geo_model_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_grid_values = geo_model_test.solutions.octrees_output[0].last_output_center.custom_grid_values\n",
    "custom_grid_values"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "gempy3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.1.-1"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}