LF2.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "34906bd7",
   "metadata": {},
   "source": [
    "# Model LF2"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8538ce6",
   "metadata": {},
   "source": [
    "This is the script used to run experiment LF2 in the publication:\n",
    "\n",
    "Youseph Ibrahim, Patrice F. Rey; The Role of Isostasy in the Evolution and Architecture of Fold and Thrust Belts. Lithosphere 2023;; 2023 (1): 8503619. doi: https://doi-org.ezproxy.library.sydney.edu.au/10.2113/2023/8503619\n",
    " \n",
    "By Youseph Ibrahim \n",
    "\n",
    "This was run using UWGeodynamics V2.9"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "56811214",
   "metadata": {},
   "source": [
    "### Import Modules"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d2a2536c",
   "metadata": {},
   "outputs": [],
   "source": [
    "import underworld.function as fn\n",
    "import UWGeodynamics as GEO\n",
    "import numpy as np\n",
    "import gflex"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5365b062",
   "metadata": {},
   "source": [
    "### Scaling"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d7a4a897",
   "metadata": {},
   "outputs": [],
   "source": [
    "u = GEO.UnitRegistry\n",
    "\n",
    "half_rate = 10 * u.millimeter / u.year\n",
    "model_length = 44.8e3 * u.meter\n",
    "surfaceTemp = 273.15 * u.degK\n",
    "baseModelTemp = 1603.15 * u.degK\n",
    "bodyforce = 2700 * u.kilogram / u.metre**3 * 9.81 * u.meter / u.second**2\n",
    "rigidbasedensity = 2800. * u.kilogram / u.metre**3\n",
    "\n",
    "KL = model_length\n",
    "Kt = KL / half_rate\n",
    "KM = bodyforce * KL**2 * Kt**2\n",
    "KT = (baseModelTemp - surfaceTemp)\n",
    "\n",
    "GEO.scaling_coefficients[\"[length]\"] = KL\n",
    "GEO.scaling_coefficients[\"[time]\"] = Kt\n",
    "GEO.scaling_coefficients[\"[mass]\"] = KM\n",
    "GEO.scaling_coefficients[\"[temperature]\"] = KT"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "63640e06",
   "metadata": {},
   "source": [
    "### Model Domain and Resolution"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "569857d5",
   "metadata": {},
   "outputs": [],
   "source": [
    "Model = GEO.Model(elementRes=(800, 200),\n",
    "                  minCoord=(0. * u.kilometer, -12. * u.kilometer),\n",
    "                  maxCoord=(64. * u.kilometer, 4. * u.kilometer),\n",
    "                  gravity=(0.0, -9.81 * u.meter / u.second**2))\n",
    "\n",
    "Model.outputDir = \"1cm_20Te_v7_O\" #output directory\n",
    "Model.diffusivity = 9e-7 * u.metre**2 / u.second\n",
    "Model.capacity = 1000. * u.joule / (u.kelvin * u.kilogram)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1d7b63c5",
   "metadata": {},
   "source": [
    "### Defining Materials"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4f8212e",
   "metadata": {},
   "outputs": [],
   "source": [
    "air = Model.add_material(name=\"Air\", shape=GEO.shapes.Layer(\n",
    "    top=Model.top, bottom=0 * u.kilometer))\n",
    "air.density = 1. * u.kilogram / u.metre**3\n",
    "air.diffusivity = 1e-6 * u.metre**2 / u.second\n",
    "air.capacity = 1000. * u.joule / (u.kelvin * u.kilogram)\n",
    "\n",
    "Loose_Sediment = Model.add_material(name=\"Loose_Sediment\", shape=GEO.shapes.Layer2D(\n",
    "    top=0. * u.kilometer, bottom=-0.2 * u.kilometer))\n",
    "Loose_Sediment.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Loose_Sediment.density = 2000. * u.kilogram / u.metre**3\n",
    "\n",
    "Strong_1 = Model.add_material(name=\"Strong_1\", shape=GEO.shapes.Layer2D(\n",
    "    top=-0.2 * u.kilometer, bottom=-0.7 * u.kilometer))\n",
    "Strong_1.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Strong_1.density = 2600. * u.kilogram / u.metre**3\n",
    "\n",
    "Weak_1 = Model.add_material(name=\"Weak_1\", shape=GEO.shapes.Layer2D(\n",
    "    top=Strong_1.bottom, bottom=-1.2 * u.kilometer))\n",
    "Weak_1.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Weak_1.density = 2200. * u.kilogram / u.metre**3\n",
    "\n",
    "Strong_2 = Model.add_material(name=\"Strong_2\", shape=GEO.shapes.Layer2D(\n",
    "    top=Weak_1.bottom, bottom=-1.7 * u.kilometer))\n",
    "Strong_2.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Strong_2.density = 2600. * u.kilogram / u.metre**3\n",
    "\n",
    "Weak_2 = Model.add_material(name=\"Weak_2\", shape=GEO.shapes.Layer2D(\n",
    "    top=Strong_2.bottom, bottom=-2.2 * u.kilometer))\n",
    "Weak_2.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Weak_2.density = 2200. * u.kilogram / u.metre**3\n",
    "\n",
    "Strong_3 = Model.add_material(name=\"Strong_3\", shape=GEO.shapes.Layer2D(\n",
    "    top=Weak_2.bottom, bottom=-2.7 * u.kilometer))\n",
    "Strong_3.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Strong_3.density = 2600. * u.kilogram / u.metre**3\n",
    "\n",
    "Weak_3 = Model.add_material(name=\"Weak_3\", shape=GEO.shapes.Layer2D(\n",
    "    top=Strong_3.bottom, bottom=-3.2 * u.kilometer))\n",
    "Weak_3.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Weak_3.density = 2300. * u.kilogram / u.metre**3\n",
    "\n",
    "Strong_4 = Model.add_material(name=\"Strong_3\", shape=GEO.shapes.Layer2D(\n",
    "    top=Weak_3.bottom, bottom=-3.7 * u.kilometer))\n",
    "Strong_4.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Strong_4.density = 2600. * u.kilogram / u.metre**3\n",
    "\n",
    "Weak_4 = Model.add_material(name=\"Weak_3\", shape=GEO.shapes.Layer2D(\n",
    "    top=Strong_4.bottom, bottom=-4.2 * u.kilometer))\n",
    "Weak_4.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Weak_4.density = 2300. * u.kilogram / u.metre**3\n",
    "\n",
    "Basement = Model.add_material(name=\"Continental Crust\", shape=GEO.shapes.Layer2D(top=Weak_4.bottom, bottom=-7.5 * u.kilometer))\n",
    "Basement.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Basement.density  = 2720. * u.kilogram / u.metre**3\n",
    "\n",
    "Beam = Model.add_material(name=\"Beam\", shape=GEO.shapes.Layer2D(top=Basement.bottom, bottom=Model.bottom))\n",
    "Beam.radiogenicHeatProd = 7.67e-7 * u.watt / u.meter**3\n",
    "Beam.density  = 2720. * u.kilogram / u.metre**3"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7d7fcd7b",
   "metadata": {},
   "source": [
    "### Material Viscosity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9e8acb66",
   "metadata": {},
   "outputs": [],
   "source": [
    "rh=GEO.ViscousCreepRegistry()\n",
    "\n",
    "Model.minViscosity=5e18 * u.pascal * u.second\n",
    "Model.maxViscosity=5e23 * u.pascal * u.second\n",
    "\n",
    "air.viscosity=5e18 * u.pascal * u.second\n",
    "\n",
    "Loose_Sediment.viscosity=1e20 * u.pascal * u.second\n",
    "\n",
    "Strong_1.viscosity=1e22 * u.pascal * u.second\n",
    "\n",
    "Strong_2.viscosity=1e22 * u.pascal * u.second\n",
    "\n",
    "Strong_3.viscosity=1e22 * u.pascal * u.second\n",
    "\n",
    "Strong_4.viscosity=1e22 * u.pascal * u.second\n",
    "\n",
    "\n",
    "Weak_1.viscosity=5e20 * u.pascal * u.second\n",
    "\n",
    "Weak_2.viscosity=5e20 * u.pascal * u.second\n",
    "\n",
    "Weak_3.viscosity=5e20 * u.pascal * u.second\n",
    "\n",
    "Weak_4.viscosity=5e20 * u.pascal * u.second\n",
    "\n",
    "Basement.viscosity=1e23 * u.pascal * u.second\n",
    "\n",
    "Beam.viscosity=1e23 * u.pascal * u.second"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "eab8b529",
   "metadata": {},
   "source": [
    "### Material Plasticity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ed3f41c",
   "metadata": {},
   "outputs": [],
   "source": [
    "Loose_Sediment.plasticity=GEO.DruckerPrager(name = \"Strong_1\",\n",
    "                                              cohesion = 0. * u.megapascal,\n",
    "                                              cohesionAfterSoftening = 0.0 * u.megapascal,\n",
    "                                              frictionCoefficient = 0.01,\n",
    "                                              frictionAfterSoftening = 0.001,\n",
    "                                              epsilon1 = 0., epsilon2 = 0.25)\n",
    "\n",
    "Strong_1.plasticity=GEO.DruckerPrager(name = \"Strong_1\",\n",
    "                                        cohesion=5. * u.megapascal,\n",
    "                                        cohesionAfterSoftening=0.5 * u.megapascal,\n",
    "                                        frictionCoefficient=0.1,\n",
    "                                        frictionAfterSoftening=0.01,\n",
    "                                        epsilon1=0., epsilon2=0.25)\n",
    "Weak_1.plasticity = GEO.DruckerPrager(name=\"Weak_1\",\n",
    "                                      cohesion=5. * u.megapascal,\n",
    "                                      cohesionAfterSoftening=0.5 * u.megapascal,\n",
    "                                      frictionCoefficient=0.1,\n",
    "                                      frictionAfterSoftening=0.01,\n",
    "                                      epsilon1=0., epsilon2=0.25)\n",
    "Strong_2.plasticity = GEO.DruckerPrager(name=\"Strong_2\",\n",
    "                                        cohesion=5. * u.megapascal,\n",
    "                                        cohesionAfterSoftening=0.5 * u.megapascal,\n",
    "                                        frictionCoefficient=0.1,\n",
    "                                        frictionAfterSoftening=0.01,\n",
    "                                        epsilon1=0., epsilon2=0.25)\n",
    "Weak_2.plasticity = GEO.DruckerPrager(name=\"Weak_2\",\n",
    "                                      cohesion=5. * u.megapascal,\n",
    "                                      cohesionAfterSoftening=0.5 * u.megapascal,\n",
    "                                      frictionCoefficient=0.1,\n",
    "                                      frictionAfterSoftening=0.01,\n",
    "                                      epsilon1=0., epsilon2=0.25)\n",
    "Strong_3.plasticity = GEO.DruckerPrager(name=\"Strong_3\",\n",
    "                                        cohesion=5. * u.megapascal,\n",
    "                                        cohesionAfterSoftening=0.5 * u.megapascal,\n",
    "                                        frictionCoefficient=0.1,\n",
    "                                        frictionAfterSoftening=0.01,\n",
    "                                        epsilon1=0., epsilon2=0.25)\n",
    "Weak_3.plasticity = GEO.DruckerPrager(name=\"Weak_3\",\n",
    "                                      cohesion=5. * u.megapascal,\n",
    "                                      cohesionAfterSoftening=0.5 * u.megapascal,\n",
    "                                      frictionCoefficient=0.1,\n",
    "                                      frictionAfterSoftening=0.01,\n",
    "                                      epsilon1=0., epsilon2=0.25)\n",
    "Strong_4.plasticity = GEO.DruckerPrager(name=\"Strong_3\",\n",
    "                                        cohesion=5. * u.megapascal,\n",
    "                                        cohesionAfterSoftening=0.5 * u.megapascal,\n",
    "                                        frictionCoefficient=0.1,\n",
    "                                        frictionAfterSoftening=0.01,\n",
    "                                        epsilon1=0., epsilon2=0.25)\n",
    "Weak_4.plasticity = GEO.DruckerPrager(name=\"Weak_3\",\n",
    "                                      cohesion=5. * u.megapascal,\n",
    "                                      cohesionAfterSoftening=0.5 * u.megapascal,\n",
    "                                      frictionCoefficient=0.1,\n",
    "                                      frictionAfterSoftening=0.01,\n",
    "                                      epsilon1=0., epsilon2=0.25)\n",
    "Basement.plasticity = GEO.DruckerPrager(name=\"Basement\",\n",
    "                                        cohesion=40. * u.megapascal,\n",
    "                                        cohesionAfterSoftening=4. * u.megapascal,\n",
    "                                        frictionCoefficient=0.6,\n",
    "                                        frictionAfterSoftening=0.06,\n",
    "                                        epsilon1=0.1, epsilon2=0.25)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "52e2e6f0",
   "metadata": {},
   "source": [
    "### Material Elasticity"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "36c728f4",
   "metadata": {},
   "source": [
    "#### Calculating elastic properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef5a7f38",
   "metadata": {},
   "outputs": [],
   "source": [
    "eta = 1e23 * u.pascal * u.second  # Viscosity\n",
    "mu = 2e9 * u.pascal  # Shear Modulus\n",
    "\n",
    "alpha = eta / mu         # Maxwell relaxation time\n",
    "dt_e = 20e3 * u.year  # Load relaxation time\n",
    "eta_eff = (eta * dt_e) / (alpha + dt_e)  # effective viscosity\n",
    "\n",
    "yieldStrength = 12e6 * (u.kilogram * u.meter**-1 * u.second**-2)\n",
    "\n",
    "minVisc = 1e19 * (u.kilogram * u.meter**-1 * u.second**-2) * u.second\n",
    "maxVisc = 1e24 * (u.kilogram * u.meter**-1 * u.second**-2) * u.second\n",
    "\n",
    "density = 2700 * u.kilogram / u.metre**3\n",
    "gravity = 9.81 * u.metre / u.second**2\n",
    "\n",
    "shearVelocity = 0.5 * u.centimetre / u.year\n",
    "\n",
    "print('Maxwell relaxation time = ', alpha.to(u.years))\n",
    "print(\"Observation time        = \", dt_e.to(u.year), dt_e)\n",
    "print(\"effective viscosity     = \", eta_eff.to(u.pascal * u.second))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "25fbae72",
   "metadata": {},
   "source": [
    "#### Defining elastic material properties"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "57d5a7b0",
   "metadata": {
    "lines_to_next_cell": 2
   },
   "outputs": [],
   "source": [
    "Strong_1.elasticity = GEO.Elasticity(shear_modulus=2e9 * u.pascal,\n",
    "                                     observation_time=20000 * u.year)\n",
    "Weak_1.elasticity = GEO.Elasticity(shear_modulus=2e9 * u.pascal,\n",
    "                                   observation_time=20000 * u.year)\n",
    "Strong_2.elasticity = GEO.Elasticity(shear_modulus=2e9 * u.pascal,\n",
    "                                     observation_time=20000 * u.year)\n",
    "Weak_2.elasticity = GEO.Elasticity(shear_modulus=2e9 * u.pascal,\n",
    "                                   observation_time=20000 * u.year)\n",
    "Strong_3.elasticity = GEO.Elasticity(shear_modulus=2e9 * u.pascal,\n",
    "                                     observation_time=20000 * u.year)\n",
    "Weak_3.elasticity = GEO.Elasticity(shear_modulus=2e9 * u.pascal,\n",
    "                                   observation_time=20000 * u.year)\n",
    "Strong_4.elasticity = GEO.Elasticity(shear_modulus=2e9 * u.pascal,\n",
    "                                     observation_time=20000 * u.year)\n",
    "Weak_4.elasticity = GEO.Elasticity(shear_modulus=2e9 * u.pascal,\n",
    "                                   observation_time=20000 * u.year)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "1be3b70c",
   "metadata": {},
   "source": [
    "### Temperature Boundary Conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "158e8f79",
   "metadata": {},
   "outputs": [],
   "source": [
    "Model.set_temperatureBCs(top=293.15 * u.degK, materials=[(air, 293.15*u.degK)])\n",
    "\n",
    "Model.set_heatFlowBCs(bottom=(-0.044 * u.watt / u.metre**2, Basement))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8f87bfe8",
   "metadata": {},
   "source": [
    "### Velocity Boundary Conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fceca016",
   "metadata": {},
   "outputs": [],
   "source": [
    "velocity = 1. * u.centimeter / u.year\n",
    "\n",
    "conditions = [(Model.y <= GEO.nd(Beam.top), GEO.nd(-velocity)),\n",
    "              (Model.y >  GEO.nd(Beam.top),\n",
    "               GEO.nd(0. * u.centimeter / u.year)),\n",
    "              (True, GEO.nd(0. * u.centimeter / u.year))]\n",
    "\n",
    "fn_condition = fn.branching.conditional(conditions)\n",
    "\n",
    "Model.set_velocityBCs(left=[fn_condition, None],\n",
    "                      right=[-velocity, 0.],\n",
    "                      top=[None, None],\n",
    "                      bottom=[-velocity, 0.])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f624b171",
   "metadata": {},
   "source": [
    "### Initialise Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d048dfdf",
   "metadata": {},
   "outputs": [],
   "source": [
    "Model.init_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "77c33ce3",
   "metadata": {},
   "source": [
    "### Tracers"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9af0f986",
   "metadata": {},
   "outputs": [],
   "source": [
    "x = np.linspace(GEO.nd(Model.minCoord[0]), GEO.nd(Model.maxCoord[0]), 1000)\n",
    "y = 0.\n",
    "\n",
    "surface_tracers = Model.add_passive_tracers(name=\"Surface\", vertices=[x, y])\n",
    "moho_tracers = Model.add_passive_tracers(\n",
    "    name=\"Moho\", vertices=[x, y-GEO.nd(24.*u.kilometer)])\n",
    "\n",
    "npoints = int(Model.maxCoord[0].to(u.kilometer).magnitude)\n",
    "x_surface = np.linspace(\n",
    "    GEO.nd(Model.minCoord[0]), GEO.nd(Model.maxCoord[0]), npoints)\n",
    "y_surface = -0.01 * u.kilometer\n",
    "\n",
    "surface_tracers_no_erosion = Model.add_passive_tracers(\n",
    "    name=\"Surface-NoErosion\", vertices=[x_surface, y_surface], zOnly=True)\n",
    "surface_tracers_erosion = Model.add_passive_tracers(\n",
    "    name=\"Surface-Erosion\", vertices=[x_surface, y_surface], zOnly=True)\n",
    "\n",
    "x_c, y_c = GEO.circles_grid(radius=0.1*u.kilometer,\n",
    "                            minCoord=[Model.minCoord[0], Basement.bottom],\n",
    "                            maxCoord=[Model.maxCoord[0], 0.*u.kilometer])\n",
    "\n",
    "FSE_Crust = Model.add_passive_tracers(name=\"FSE_Crust\", vertices=[x_c, y_c])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ce2d8aba",
   "metadata": {},
   "source": [
    "### Surface Processes"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e1dc919",
   "metadata": {},
   "source": [
    "#### Hillslope function"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "06751b60",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Hillslope_diffusion_basic():\n",
    "    from scipy.interpolate import interp1d\n",
    "    from scipy.interpolate import InterpolatedUnivariateSpline\n",
    "\n",
    "    x = GEO.dimensionalise(\n",
    "        surface_tracers_erosion.data[:, 0], u.meter).magnitude\n",
    "    z = GEO.dimensionalise(\n",
    "        surface_tracers_erosion.data[:, 1], u.meter).magnitude\n",
    "\n",
    "    dx = (Model.maxCoord[0].to(u.meter).magnitude)/npoints\n",
    "\n",
    "    total_time = (GEO.dimensionalise(Model._dt, u.year)).magnitude\n",
    "\n",
    "    D = ((1.0e3 * u.meter**2 / u.year).to(u.meter**2 / u.year)).magnitude\n",
    "\n",
    "    dt = min((0.2 * dx * dx / D), total_time)\n",
    "\n",
    "    nts = int(round(total_time/dt))\n",
    "\n",
    "    print('total time:', total_time, 'timestep:', dt, 'No. of its:', nts)\n",
    "\n",
    "    z_orig = z.copy()\n",
    "\n",
    "    for i in range(nts):\n",
    "        qs = -D * np.diff(z)/dx\n",
    "        dzdt = -np.diff(qs)/dx\n",
    "        z[1:-1] += dzdt*dt\n",
    "\n",
    "    x_nd = GEO.nd(x*u.meter)\n",
    "\n",
    "    z_nd = GEO.nd(z*u.meter)\n",
    "\n",
    "    if x_nd.shape[0] > 0.:\n",
    "        f1 = interp1d(x_nd, z_nd, fill_value='extrapolate', kind='nearest')\n",
    "        y_eroded_surface = f1(x_nd)\n",
    "        y_eroded_surface[x_nd < GEO.nd(Update_material_LHS_Length *\n",
    "                                       u.kilometer)] = 0.\n",
    "\n",
    "        surface_tracers_erosion.data[:, 1] = y_eroded_surface\n",
    "\n",
    "        Model.materialField.data[(Model.swarm.data[:, 1] > f1(Model.swarm.data[:, 0])) & (\n",
    "            Model.materialField.data[:, 0] != air.index)] = air.index\n",
    "        Model.materialField.data[(Model.swarm.data[:, 1] < f1(Model.swarm.data[:, 0])) & (\n",
    "            Model.materialField.data[:, 0] == air.index)] = Sediment.index"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ba07b1b4",
   "metadata": {},
   "source": [
    "### GFlex Setup"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "89bab35f",
   "metadata": {},
   "outputs": [],
   "source": [
    "sample_depth = Model.bottom\n",
    "gFlex_interval = 5000. * u.years\n",
    "ps_points = np.zeros((Model.elementRes[0], Model.mesh.dim))\n",
    "ps_points[:, 0] = np.linspace(GEO.nd(Model.minCoord[0]), GEO.nd(\n",
    "    Model.maxCoord[0]), Model.elementRes[0])\n",
    "ps_points[:, -1] = GEO.nd(sample_depth)\n",
    "flex = gflex.F1D()\n",
    "\n",
    "def gFlexSetup():\n",
    "    ''' gFlex setup function\n",
    "\n",
    "    Include general parameters for gFlex algorithm. \n",
    "    As gFlex is only being run on the root proc we can assign this to only proc 0'''\n",
    "\n",
    "    if GEO.comm.rank == 0:\n",
    "\n",
    "        # gFlex mode\n",
    "        flex.Quiet = False\n",
    "\n",
    "        flex.Method = 'FD' \n",
    "\n",
    "        flex.PlateSolutionType = 'vWC1994'  \n",
    "     \n",
    "        flex.Solver = 'direct'  \n",
    "        flex.g = 9.8 \n",
    "        flex.E = 7.e10  \n",
    "        flex.nu = 0.25  \n",
    "        flex.rho_m = 3300. \n",
    "        flex.rho_fill = 0.  \n",
    "        flex.Te = 20.e3  \n",
    "\n",
    "        flex.dx = GEO.dimensionalise(\n",
    "            (ps_points[1, 0]-ps_points[0, 0]), u.km).m_as('m')\n",
    "\n",
    "    GEO.comm.Barrier()\n",
    "gFlexSetup()\n",
    "\n",
    "elSize = [(Model.mesh.maxCoord[i] - Model.mesh.minCoord[i]) /\n",
    "          Model.elementRes[i] for i in range(Model.mesh.dim)]\n",
    "\n",
    "\n",
    "def run_gFlex():\n",
    "    ''' An algorithm to incorporate gFlex elastic compensation into a running UWGeo model:\n",
    "\n",
    "    Collect pressure values on the root proc, solve gFlex and broadcast deformation information to all procs.\n",
    "\n",
    "    All procs then deform Model.swarm. Population control and an iterative deformation method strategy \n",
    "    ensure particles remain in cells after the entire deformation occurs.\n",
    "    '''\n",
    "    \n",
    "    rank = GEO.comm.rank\n",
    "\n",
    "    if rank == 0:\n",
    "        print(\"UW: Starting gFlex callback\")\n",
    "\n",
    "    current_pressure = Model.pressureField.evaluate_global(ps_points)\n",
    "\n",
    "    from scipy import interpolate\n",
    "    uw_deflection = []  \n",
    "\n",
    "    if rank == 0:\n",
    "        global previous_pressure\n",
    "        current_pressure = GEO.dimensionalise(current_pressure, u.Pa).magnitude\n",
    "        np.savetxt(\n",
    "            fname='Outputs/previous_pressure{}.txt'.format(Model.step), X=previous_pressure)\n",
    "        np.savetxt(\n",
    "            fname='Outputs/current_pressure{}.txt'.format(Model.step), X=current_pressure)\n",
    "\n",
    "        sloads = current_pressure - previous_pressure\n",
    "\n",
    "        np.savetxt(fname='Outputs/sloads{}.txt'.format(Model.step), X=sloads)\n",
    "        flex.qs = sloads.reshape(-1)  # surface load stresses\n",
    "\n",
    "        flex.BC_W = '0Moment0Shear'  # west boundary condition\n",
    "        flex.BC_E = '0Displacement0Slope'  # east boundary condition\n",
    "\n",
    "        flex.initialize()\n",
    "        flex.run()  # flex.w is the gFlex output\n",
    "        flex.finalize()\n",
    "\n",
    "        print(\"UW: Finished gFlex\")\n",
    "\n",
    "        np.savetxt(fname='Outputs/raw_gFlex{}'.format(Model.step),\n",
    "                   X=flex.w)\n",
    "\n",
    "        uw_deflection = GEO.nd(flex.w*u.m)\n",
    "\n",
    "    uw_deflection = GEO.comm.bcast(uw_deflection, root=0)\n",
    "    max_d = 0.25*elSize[1]\n",
    "    i = 1  \n",
    "    while np.any(np.abs(uw_deflection) > 0.):\n",
    "\n",
    "        if rank == 0:\n",
    "            print(f\"gFlex deformation iteration {i}\")\n",
    "\n",
    "        it_d = np.clip(uw_deflection, -max_d, max_d)\n",
    "\n",
    "        tck = interpolate.splrep(ps_points[:, 0], it_d)\n",
    "\n",
    "        with Model.swarm.deform_swarm():\n",
    "            dy = interpolate.splev(\n",
    "                Model.swarm.particleCoordinates.data[:, 0], tck)\n",
    "            Model.swarm.particleCoordinates.data[:, 1] += dy\n",
    "\n",
    "        Model.population_control.repopulate()\n",
    "\n",
    "        uw_deflection[:] = uw_deflection[:] - it_d[:]\n",
    "\n",
    "        i = i + 1\n",
    "\n",
    "    if rank == 0:\n",
    "        print(\"UW: Finished updating swarm\")\n",
    "\n",
    "    previous_pressure = Model.pressureField.evaluate_global(ps_points)\n",
    "\n",
    "    if type(previous_pressure) == np.ndarray and len(previous_pressure) > 0:\n",
    "\n",
    "        previous_pressure = GEO.dimensionalise(\n",
    "            previous_pressure, u.Pa).magnitude\n",
    "\n",
    "    if rank == 0:\n",
    "        print(\"updated previous pressure with new swarm\")\n",
    "\n",
    "def callback():\n",
    "    if Model.step == 1:  \n",
    "        global previous_pressure\n",
    "\n",
    "        print('running first pressure evauluate at step 1')\n",
    "        previous_pressure = Model.pressureField.evaluate_global(ps_points)\n",
    "\n",
    "        if type(previous_pressure) == np.ndarray and len(previous_pressure) > 0:\n",
    "            previous_pressure = GEO.dimensionalise(\n",
    "                previous_pressure, u.Pa).magnitude\n",
    "\n",
    "            print('previous_pressure is defined')\n",
    "            np.savetxt(fname='Outputs/previous_pressure_1.txt',\n",
    "                       X=previous_pressure)\n",
    "\n",
    "    if not hasattr(callback, '_t0'):\n",
    "        callback._t0 = GEO.nd(gFlex_interval)\n",
    "\n",
    "    until = callback._t0 - Model._ndtime\n",
    "    if until <= 0:\n",
    "        run_gFlex()\n",
    "        callback._t0 = Model._ndtime + GEO.nd(gFlex_interval)\n",
    "\n",
    "Model.post_solve_functions[\"gFlex_integration\"] = callback"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e3a5b9b1",
   "metadata": {},
   "source": [
    "### Reinitilalise Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cb5fd06c",
   "metadata": {},
   "outputs": [],
   "source": [
    "Model.init_model()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d7d04722",
   "metadata": {},
   "source": [
    "### Run Model"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "70ee21df",
   "metadata": {},
   "outputs": [],
   "source": [
    "GEO.rcParams[\"default.outputs\"].append(\"projStrainTensor\")\n",
    "GEO.rcParams[\"default.outputs\"].append(\"projStressTensor\")\n",
    "Model.solver.set_inner_method(\"mumps\")\n",
    "Model.solver.set_penalty(1e6)\n",
    "Model.run_for(2200000. * u.year, restartStep=-1,\n",
    "              restartDir=\"1cm_20Te_v7_O\",  checkpoint_interval=5000. * u.year)"
   ]
  }
 ],
 "metadata": {
  "jupytext": {
   "cell_metadata_filter": "-all",
   "main_language": "python",
   "notebook_metadata_filter": "-all"
  },
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.6"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}