gempy3_first_model_generative_SalinasA_generative.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import os\n",
    "import numpy as np\n",
    "import matplotlib.pyplot as plt\n",
    "import torch\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",
    "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"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import scipy.io\n",
    "\n",
    "# Load .mat file\n",
    "SalinasA= np.array(scipy.io.loadmat('./HSI_Salinas/SalinasA.mat')['salinasA'])\n",
    "SalinasA_corrected= np.array(scipy.io.loadmat('./HSI_Salinas/SalinasA_corrected.mat')['salinasA_corrected'])\n",
    "SalinasA_gt= np.array(scipy.io.loadmat('./HSI_Salinas/SalinasA_gt.mat')['salinasA_gt'])\n",
    "\n",
    "i=0\n",
    "label_data = [0,6,1,5,4,3,2]\n",
    "print(np.unique(SalinasA_gt))\n",
    "for ele in np.unique(SalinasA_gt):\n",
    "    #print(ele)\n",
    "    mask = SalinasA_gt==ele\n",
    "    SalinasA_gt[mask] = label_data[i]\n",
    "    #print(SalinasA_gt)\n",
    "    i=i+1\n",
    "    #print(i+1)\n",
    "print(np.unique(SalinasA_gt))\n",
    "\n",
    "print(SalinasA_gt.shape)\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import ListedColormap\n",
    "# Define a colormap with colors for each label\n",
    "colors = ['black', 'red', 'green', 'blue', 'yellow', 'purple', 'orange']\n",
    "cmap = ListedColormap(colors)\n",
    "\n",
    "plt.imshow(SalinasA_gt, cmap=cmap, vmin=0, vmax=len(colors)-1)\n",
    "\n",
    "# Add colorbar to show label-color mapping\n",
    "cbar = plt.colorbar(ticks=np.arange(len(colors)))\n",
    "cbar.set_ticklabels([f'Label {i}' for i in range(len(colors))])\n",
    "# Add a label at a specific coordinate (e.g., x=5, y=5)\n",
    "#label_text = \"A\"\n",
    "#plt.text(20, 10, label_text, color='red', fontsize=4, ha='center', va='center')\n",
    "# Set ticks to align with individual pixels\n",
    "#plt.xticks(np.arange(0.0, 87.0, 1), [])\n",
    "#plt.yticks(np.arange(0.0, 84.0, 1), [])\n",
    "plt.grid(True, which='both', color='black', linewidth=0.5, linestyle=':')\n",
    "plt.show()\n",
    "print(SalinasA_gt)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "SalinasA_gt.shape, SalinasA_corrected.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "column = 20\n",
    "y_obs = torch.tensor(SalinasA_gt[:,column], dtype=torch.float64)\n",
    "mask = y_obs!=0\n",
    "print(mask)\n",
    "y_obs_label = y_obs[mask]\n",
    "print(len(y_obs_label))\n",
    "print(SalinasA_corrected.shape)\n",
    "y_obs_hsi = torch.tensor(SalinasA_corrected[:,column,:], dtype=torch.float64)[mask]\n",
    "y_obs_hsi.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_obs_hsi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from scipy.stats import zscore"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "normalised_hsi = zscore(y_obs_hsi, axis=0)\n",
    "normalised_hsi"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.manifold import TSNE\n",
    "model = TSNE(n_components=2, random_state=0)\n",
    "transformed_hsi = model.fit_transform(normalised_hsi) "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "transformed_hsi.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.figure(figsize=(10,8))\n",
    "plt.scatter(transformed_hsi[:,0],transformed_hsi[:,1],c=y_obs_label, cmap=ListedColormap(colors[1:]))\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from sklearn.mixture import GaussianMixture\n",
    "from sklearn.mixture import BayesianGaussianMixture\n",
    "from sklearn.cluster import KMeans"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gm =gm = BayesianGaussianMixture(n_components=6, random_state=0).fit(normalised_hsi)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gm.means_ , gm.covariances_"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gmm_label = gm.predict(normalised_hsi) +1 \n",
    "gmm_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_obs_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gmm_label2 = torch.zeros_like(y_obs_label)\n",
    "gmm_label2[gmm_label==2]=6\n",
    "gmm_label2[gmm_label==4]=5\n",
    "gmm_label2[gmm_label==1]=4\n",
    "gmm_label2[gmm_label==3]=3\n",
    "gmm_label2[gmm_label==6]=2\n",
    "gmm_label2[gmm_label==5]=1\n",
    "gmm_label2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gm.means_.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "mean_init, cov_init = gm.means_[[4,5,2,0,3,1]], gm.covariances_[[4,5,2,0,3,1]]\n",
    "mean_init.shape, cov_init.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_obs_label - gmm_label2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.scatter(transformed_hsi[:,0],transformed_hsi[:,1],c=gmm_label2,cmap=ListedColormap(colors))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": []
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test = gp.create_geomodel(\n",
    "    project_name='Gempy_abc_Test',\n",
    "    extent=[0, 86, -10, 10, -83, 0],\n",
    "    resolution=[86,20,83],\n",
    "    refinement=3,\n",
    "    structural_frame= gp.data.StructuralFrame.initialize_default_structure()\n",
    "    )\n",
    "\n",
    "gp.add_surface_points(\n",
    "    geo_model=geo_model_test,\n",
    "    x=[70.0, 80.0],\n",
    "    y=[0.0, 0.0],\n",
    "    z=[-77.0, -71.0],\n",
    "    elements_names=['surface1', 'surface1']\n",
    ")\n",
    "\n",
    "gp.add_orientations(\n",
    "    geo_model=geo_model_test,\n",
    "    x=[75],\n",
    "    y=[0.0],\n",
    "    z=[-74],\n",
    "    elements_names=['surface1'],\n",
    "    pole_vector=[[-5/3, 0, 1]]\n",
    ")\n",
    "geo_model_test.update_transform(gp.data.GlobalAnisotropy.NONE)\n",
    "\n",
    "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([200.0, 600.0])/10,\n",
    "        y=np.array([0.0, 0.0]),\n",
    "        z=np.array([-740, -520])/10,\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",
    "\n",
    "element3 = gp.data.StructuralElement(\n",
    "    name='surface3',\n",
    "    color=next(geo_model_test.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([0.0, 300.0, 600])/10,\n",
    "        y=np.array([0.0, 0.0,0.0]),\n",
    "        z=np.array([-720, -830+275, -830+440])/10,\n",
    "        names='surface3'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test.structural_frame.structural_groups[0].append_element(element3)\n",
    "\n",
    "element4 = gp.data.StructuralElement(\n",
    "    name='surface4',\n",
    "    color=next(geo_model_test.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([0.0, 200.0, 600])/10,\n",
    "        y=np.array([0.0, 0.0,0.0]),\n",
    "        z=np.array([-830+220, -830+340, -830+560])/10,\n",
    "        names='surface4'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test.structural_frame.structural_groups[0].append_element(element4)\n",
    "\n",
    "element5 = gp.data.StructuralElement(\n",
    "    name='surface5',\n",
    "    color=next(geo_model_test.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([0.0, 200.0, 400])/10,\n",
    "        y=np.array([0.0, 0.0, 0.0]),\n",
    "        z=np.array([440, 550, 670])/10 -830/10,\n",
    "        names='surface5'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test.structural_frame.structural_groups[0].append_element(element5)\n",
    "\n",
    "element6 = gp.data.StructuralElement(\n",
    "    name='surface6',\n",
    "    color=next(geo_model_test.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([0.0, 200.0,300])/10,\n",
    "        y=np.array([0.0, 0.0, 0.0]),\n",
    "        z=np.array([620, 730, 820])/10 -83,\n",
    "        names='surface6'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test.structural_frame.structural_groups[0].append_element(element6)\n",
    "\n",
    "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[2], geo_model_test.structural_frame.structural_groups[0].elements[3],\\\n",
    "geo_model_test.structural_frame.structural_groups[0].elements[4], geo_model_test.structural_frame.structural_groups[0].elements[5] = \\\n",
    "geo_model_test.structural_frame.structural_groups[0].elements[1], geo_model_test.structural_frame.structural_groups[0].elements[0],\\\n",
    "geo_model_test.structural_frame.structural_groups[0].elements[3], geo_model_test.structural_frame.structural_groups[0].elements[2],\\\n",
    "geo_model_test.structural_frame.structural_groups[0].elements[5], geo_model_test.structural_frame.structural_groups[0].elements[4]  \n",
    "\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": [
    "y_obs_label = 7 -y_obs[mask]\n",
    "y_obs_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_loc = 20\n",
    "y_loc = 0\n",
    "z_loc = np.linspace(0,-82, 83)\n",
    "print(z_loc)\n",
    "xyz_coord = np.array([[x_loc, y_loc, z] for z in z_loc])[mask]\n",
    "print(xyz_coord.shape)\n",
    "gp.set_custom_grid(geo_model_test.grid, xyz_coord=xyz_coord)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geo_model_test.interpolation_options.mesh_extraction = False\n",
    "sol = gp.compute_model(geo_model_test)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "#gpv.plot_2d(geo_model_test, section_names=['section1'])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sp_coords_copy_test = geo_model_test.interpolation_input.surface_points.sp_coords.copy()\n",
    "print(sp_coords_copy_test)\n",
    "geo_model_test.transform.apply_inverse(sp_coords_copy_test)"
   ]
  },
  {
   "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",
    "print(sp_coords_copy_test)\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)"
   ]
  },
  {
   "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 = 50"
   ]
  },
  {
   "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_values1 = geo_model_test.solutions.octrees_output[0].last_output_center.custom_grid_values\n",
    "custom_grid_values1"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    " \n",
    "# # Compute the geological model\n",
    "interpolation_input = geo_model_test.interpolation_input\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_values2 = geo_model_test.solutions.octrees_output[0].last_output_center.custom_grid_values\n",
    "\n",
    "print(custom_grid_values2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(custom_grid_values1) - custom_grid_values2"
   ]
  },
  {
   "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",
    "import torch.nn.functional as F\n",
    "from pyro.infer import SVI, TraceEnum_ELBO, config_enumerate, infer_discrete\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[11, 2]\n",
    "    prior_mean_surface_2 = sp_coords_copy_test[14, 2]\n",
    "\n",
    "    mu_surface_1 = pyro.sample('mu_1', dist.Normal(prior_mean_surface_1, torch.tensor(0.2, dtype=torch.float64)))\n",
    "    mu_surface_2 = pyro.sample('mu_2', dist.Normal(prior_mean_surface_2, torch.tensor(0.2, dtype=torch.float64)))\n",
    "    #print(mu_surface_1, mu_surface_2)\n",
    "    # Ensure that mu_surface_1 is greater than mu_surface_2\n",
    "    pyro.sample('mu_1 < 0', dist.Delta(torch.tensor(1.0, dtype=torch.float64)), obs=(mu_surface_1 < 3.7))\n",
    "    pyro.sample('mu_1 > mu_2', dist.Delta(torch.tensor(1.0, dtype=torch.float64)), obs=(mu_surface_1 > mu_surface_2))\n",
    "    pyro.sample('mu_2 > -49', dist.Delta(torch.tensor(1.0, dtype=torch.float64)), obs=(mu_surface_2 > 1.25))\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([11]), 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([14]), torch.tensor([2])),\n",
    "        mu_surface_2\n",
    "    )\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",
    "    # 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",
    "    lambda_ = 10.0\n",
    "    loc_mean = torch.tensor(mean_init,dtype=torch.float64)\n",
    "    loc_cov =  torch.tensor(cov_init, dtype=torch.float64)\n",
    "    class_label = F.softmax(-lambda_* (torch.tensor([1,2,3,4,5,6], dtype=torch.float64) - custom_grid_values.reshape(-1,1))**2, dim=1)\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",
    "    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",
    "        #print(\"mean\\n\", mean[assignment])\n",
    "        #obs = pyro.sample(\"obs\", dist.Normal(custom_grid_values, 0.5 ), obs=y_obs_label.reshape(-1))\n",
    "        #obs = pyro.sample(\"obs\", dist.Normal(sample[assignment], 0.1 ), obs=y_obs_label.reshape(-1))\n",
    "        #print(sample_tesnor[assignment].shape, loc_cov[assignment].shape, y_obs_label.shape)\n",
    "        obs = pyro.sample(\"obs\", dist.MultivariateNormal(loc=sample_tesnor[assignment],covariance_matrix=loc_cov[assignment]), obs=y_obs_label)\n",
    "        #obs = pyro.sample(\"obs\", dist.Normal(mean[assignment],  sigma[assignment]), obs=y_obs_label.reshape(-1))\n",
    "    #return obs\n",
    "dependencies = get_dependencies(model_test, model_args=(normalised_hsi,))\n",
    "pyro.render_model(model_test, model_args=(normalised_hsi,),render_distributions=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "print(normalised_hsi)\n",
    "model_test(normalised_hsi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "prior = Predictive(model_test, num_samples=100)(normalised_hsi)\n",
    "\n",
    "# Key to avoid\n",
    "avoid_key = 'mu_1 > mu_2'\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",
    "#nuts_kernel = NUTS(model_test, step_size=0.00085, adapt_step_size=True, target_accept_prob=0.9, max_tree_depth=10)\n",
    "#nuts_kernel = NUTS(model_test)\n",
    "#initial_values = {'mu_1': torch.tensor(0.01, dtype=torch.float64),'mu_2': torch.tensor(0.01, dtype=torch.float64) }\n",
    "initial_values = {'mu_1': torch.tensor(0.01, dtype=torch.float64),'mu_2': torch.tensor(0.01, dtype=torch.float64) }\n",
    "#mcmc = MCMC(nuts_kernel, num_samples=200, warmup_steps=50, disable_validation=False, initial_params=initial_values)\n",
    "mcmc = MCMC(nuts_kernel, num_samples=200, warmup_steps=50, disable_validation=False)\n",
    "mcmc.run(normalised_hsi)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "posterior_samples = mcmc.get_samples()\n",
    "posterior_predictive = Predictive(model_test, posterior_samples)(normalised_hsi)\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": [
    "mu_1_mean = posterior_samples[\"mu_1\"].mean()\n",
    "mu_2_mean = posterior_samples[\"mu_2\"].mean()\n"
   ]
  },
  {
   "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_data\":\n",
    "#         mean = values.mean(dim=0)\n",
    "#         cov = values.std(dim=0)\n",
    "#         print(\"mean\\n\",mean)\n",
    "#         print(\"cov\\n\", cov)\n",
    "#         loc_mean_posterior.append(mean.detach().numpy())\n",
    "#         loc_cov_posterior.append(cov.detach().numpy())\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": [
    "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([11]), 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([14]), 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": [
    "geo_model_test_post = gp.create_geomodel(\n",
    "    project_name='Gempy_abc_Test_post',\n",
    "    extent=[0, 86, -10, 10, -83, 0],\n",
    "    resolution=[86,20,83],\n",
    "    refinement=7,\n",
    "    structural_frame= gp.data.StructuralFrame.initialize_default_structure()\n",
    "    )\n",
    "\n",
    "gp.add_surface_points(\n",
    "    geo_model=geo_model_test_post,\n",
    "    x=[70.0, 80.0],\n",
    "    y=[0.0, 0.0],\n",
    "    z=[-77.0, -71.0],\n",
    "    elements_names=['surface1', 'surface1']\n",
    ")\n",
    "\n",
    "gp.add_orientations(\n",
    "    geo_model=geo_model_test_post,\n",
    "    x=[75],\n",
    "    y=[0.0],\n",
    "    z=[-74],\n",
    "    elements_names=['surface1'],\n",
    "    pole_vector=[[-5/3, 0, 1]]\n",
    ")\n",
    "geo_model_test_post.update_transform(gp.data.GlobalAnisotropy.NONE)\n",
    "\n",
    "element2 = gp.data.StructuralElement(\n",
    "    name='surface2',\n",
    "    color=next(geo_model_test_post.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([20.0, 60.0]),\n",
    "        y=np.array([0.0, 0.0]),\n",
    "        z=np.array([-74, -52]),\n",
    "        names='surface2'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test_post.structural_frame.structural_groups[0].append_element(element2)\n",
    "\n",
    "element3 = gp.data.StructuralElement(\n",
    "    name='surface3',\n",
    "    color=next(geo_model_test_post.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([0.0, 30.0, 60]),\n",
    "        y=np.array([0.0, 0.0,0.0]),\n",
    "        z=np.array([-72, -55.5, -39]),\n",
    "        names='surface3'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test_post.structural_frame.structural_groups[0].append_element(element3)\n",
    "\n",
    "element4 = gp.data.StructuralElement(\n",
    "    name='surface4',\n",
    "    color=next(geo_model_test_post.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([0.0, 20.0, 60]),\n",
    "        y=np.array([0.0, 0.0,0.0]),\n",
    "        z=np.array([-61, -49, -27]),\n",
    "        names='surface4'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test_post.structural_frame.structural_groups[0].append_element(element4)\n",
    "\n",
    "element5 = gp.data.StructuralElement(\n",
    "    name='surface5',\n",
    "    color=next(geo_model_test_post.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([0.0, 20, 40]),\n",
    "        y=np.array([0.0, 0.0, 0.0]),\n",
    "        z=np.array([-39, sp_cord[14,2], -16]),\n",
    "        names='surface5'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test_post.structural_frame.structural_groups[0].append_element(element5)\n",
    "\n",
    "element6 = gp.data.StructuralElement(\n",
    "    name='surface6',\n",
    "    color=next(geo_model_test_post.structural_frame.color_generator),\n",
    "    surface_points=gp.data.SurfacePointsTable.from_arrays(\n",
    "        x=np.array([0.0, 20.0,30]),\n",
    "        y=np.array([0.0, 0.0, 0.0]),\n",
    "        z=np.array([-21, sp_cord[11,2], -1]),\n",
    "        names='surface6'\n",
    "    ),\n",
    "    orientations=gp.data.OrientationsTable.initialize_empty()\n",
    ")\n",
    "\n",
    "geo_model_test_post.structural_frame.structural_groups[0].append_element(element6)\n",
    "\n",
    "geo_model_test_post.structural_frame.structural_groups[0].elements[0], geo_model_test_post.structural_frame.structural_groups[0].elements[1],\\\n",
    "geo_model_test_post.structural_frame.structural_groups[0].elements[2], geo_model_test_post.structural_frame.structural_groups[0].elements[3],\\\n",
    "geo_model_test_post.structural_frame.structural_groups[0].elements[4], geo_model_test_post.structural_frame.structural_groups[0].elements[5] = \\\n",
    "geo_model_test_post.structural_frame.structural_groups[0].elements[1], geo_model_test_post.structural_frame.structural_groups[0].elements[0],\\\n",
    "geo_model_test_post.structural_frame.structural_groups[0].elements[3], geo_model_test_post.structural_frame.structural_groups[0].elements[2],\\\n",
    "geo_model_test_post.structural_frame.structural_groups[0].elements[5], geo_model_test_post.structural_frame.structural_groups[0].elements[4]  \n",
    "\n",
    "\n",
    "gp.set_custom_grid(geo_model_test_post.grid, xyz_coord=xyz_coord)\n",
    "gp.compute_model(geo_model_test_post)\n",
    "gpv.plot_2d(geo_model_test_post, cell_number=5, legend='force')\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_grid_values_post = geo_model_test_post.solutions.octrees_output[0].last_output_center.custom_grid_values"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "custom_grid_values_post"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "y_obs_label"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "torch.tensor(custom_grid_values_post, dtype=torch.float64) - y_obs_label"
   ]
  }
 ],
 "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.10.13"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}