fast_inversion_mwe.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Unmixing river sediments for the elemental geochemistry of their source regions \n",
    "\n",
    "This notebook provides a minimum working example demonstrating how to invert river sediment geochemistry for the composition of their source regions. \n",
    "\n",
    "Unlike the previously published method [(Lipp et al. 2021)](https://agupubs.onlinelibrary.wiley.com/doi/full/10.1029/2021GC009838) this is significantly faster. This speed-up is achieved through using Powell's algorithm instead of the Nelder-Mead to optimise the objective function and by parallelising the flow accumulation. We also initialise the solution closer to the solution than if just a constant initial composition is used but this has just a marginal impact on the runtime. \n",
    "\n",
    "### Set-up \n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import netCDF4\n",
    "import time\n",
    "import scipy as sp\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import matplotlib.pyplot as plt\n",
    "from matplotlib.colors import LogNorm, TwoSlopeNorm\n",
    "from landlab import RasterModelGrid\n",
    "from landlab.components import FlowAccumulator,SinkFillerBarnes\n",
    "from landlab.components.flow_accum.flow_accum_bw import find_drainage_area_and_discharge\n",
    "from landlab.utils import get_watershed_mask\n",
    "\n",
    "import accum # Wrapper to interface with cython parallelised accumulation functions   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Data preprocessing \n",
    "\n",
    "### Topography \n",
    "\n",
    "Topography is loaded in and used to set up a `LandLab` model grid. This topographic grid has depressions filled prior to drainage-routing. The resulting topographic grid is shown below. Topographic displayed here is `SRTM1s` data downsampled to a 200x200 grid."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "z = np.loadtxt(\"data/topo.dat\")\n",
    "mg = RasterModelGrid(z.shape,xy_spacing=(200,200))\n",
    "zr = mg.add_zeros('node', 'topographic__elevation')\n",
    "zr += np.reshape(z,z.size)\n",
    "dx = 200\n",
    "\n",
    "# Set up the boundary conditions on the square grid\n",
    "mg.set_fixed_value_boundaries_at_grid_edges(True,True,True,True)\n",
    "\n",
    "flat_shape = zr.shape # a tuple to flatten arrays [number of nodes long]\n",
    "full_shape = mg.shape # the full shape of the grid [rows, columns]\n",
    "\n",
    "sfb = SinkFillerBarnes(mg, method='D8',fill_flat=False) # Fill pits\n",
    "sfb.run_one_step()\n",
    "\n",
    "# Display topography\n",
    "plt.figure()\n",
    "plt.title(\"Topographic Elevation\")\n",
    "plt.imshow(zr.reshape(full_shape),cmap='gist_earth',origin='lower',norm=TwoSlopeNorm(vcenter=10))\n",
    "plt.colorbar()\n",
    "plt.xlabel('Horizontal grid cells')\n",
    "plt.ylabel('Vertical grid cells')\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Drainage \n",
    "\n",
    "Drainage is extracted from the topography using the `D8` algorithm. This modelled drainage network is then used to calculate the upstream area for every point in the grid. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "frr = FlowAccumulator(\n",
    "    mg,\n",
    "    'topographic__elevation',\n",
    "    flow_director = 'FlowDirectorD8')\n",
    "frr.run_one_step()  # flow routing\n",
    "\n",
    "a, tot_sed_flux = find_drainage_area_and_discharge(mg.at_node['flow__upstream_node_order'], mg.at_node['flow__receiver_node']) # a is number of nodes\n",
    "\n",
    "area = mg.at_node['drainage_area']\n",
    "area_threshold = 25 \n",
    "is_drainage = area > (area_threshold*1000000) #km2 to m2\n",
    "channel_xy = np.flip(np.transpose(np.where(is_drainage.reshape(mg.shape))),axis=1)*200 # xy coordinates of channels"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now we extract the ordered nodes, and identify the locations of sink nodes. This allows us to parallelise.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "upst_node_order = mg.at_node['flow__upstream_node_order']\n",
    "receivers =  mg.at_node['flow__receiver_node']\n",
    "\n",
    "sinks = np.where(mg.at_node[\"flow__sink_flag\"])[0] # IDs of sink nodes\n",
    "starts = np.zeros(sinks.size,dtype=int)\n",
    "for i in np.arange(starts.size):\n",
    "    starts[i] = np.where(upst_node_order==sinks[i])[0]   "
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Sample Localities \n",
    "\n",
    "The location of samples on channels is fitted the drainage network automatically, but then manually checked and adjusted where necessary. This manual check makes sure the locality's location on the _model_ drainage network is the correction location. These localities are displayed below. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fitted_locs = np.loadtxt(\"data/fitted_samp_locs.dat\") # Fitted drainage locations\n",
    "sample_data = np.loadtxt(\"data/samples.dat\",dtype=str) # [sample #, locality #, x, y] # x,y, in lat long coordinates\n",
    "\n",
    "\n",
    "loc_indxs = np.transpose(np.flip((fitted_locs).astype(int),axis=1))\n",
    "loc_nodes = np.ravel_multi_index(loc_indxs,dims=full_shape) # model grid node IDs\n",
    "\n",
    "# Load in sample localities \n",
    "plt.imshow(zr.reshape(full_shape),cmap='gist_earth',origin='lower',norm=TwoSlopeNorm(vcenter=10))\n",
    "plt.scatter(x=channel_xy[:,0]/mg.dx, y=channel_xy[:,1]/mg.dx,c='red', s=0.1)\n",
    "plt.scatter(x=fitted_locs[:,0], y=fitted_locs[:,1], marker=\"+\",c='b', s=40)\n",
    "plt.xlabel('Horizontal grid cells')\n",
    "plt.ylabel('Vertical grid cells')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Geochemical Data \n",
    "\n",
    "The geochemical data gathered at each of the above localities is loaded into memory and reformatted in a way that is easy indexable (using pandas). "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "geochem_raw = np.loadtxt('data/geochem.dat',dtype=str) # Read in data\n",
    "geochem_raw = np.delete(geochem_raw,[7,53],1) # Delete columns for S and Bi (too many NAs)\n",
    "elems = geochem_raw[0,1:54] # List of element strings\n",
    "obs_data = pd.DataFrame(geochem_raw[1:,],columns=geochem_raw[0,:]) # Cast to DataFrame for quick access\n",
    "obs_data[elems]=obs_data[elems].astype(float) # Cast numeric data to float"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Setting up an inversion grid \n",
    "\n",
    "We perform an inversion on a grid generally much coarser than the base topographic grid shown above. Generally we refer to this coarse grid as consisting of 'blocks'. Hence, we define an algorithm that upsamples, or 'expands' a coarser inversion grid onto the same resolution as the base topography. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model_width = mg.shape[1] # number of x cells in topographic model grid\n",
    "model_height = mg.shape[0] # number of y cells in topographic model grid\n",
    "\n",
    "def expand(block_grid,block_x,block_y):\n",
    "    \"\"\"Expands low res array of block heights into \n",
    "    model grid array that can be fed into topographic\n",
    "    model. Note that blocks at the upper and eastern \n",
    "    perimeter are clipped if number of blocks doesn't \n",
    "    divide number of model cells. \n",
    "    \n",
    "    block_x and block_y are the number of model cells \n",
    "    in each block in x and y dir respectively\"\"\"\n",
    "    return(block_grid.repeat(block_y, axis=0).repeat(block_x, axis=1)[:model_height,:model_width])\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Dividing our model grid up into a suite of 'blocks' may mean that some 'blocks' lie outside of our drainage basins. We don't want to consider these blocks as our samples will not contain any useful information about this part of the grid, as they lie outside of their upstream area. Hence, now we determine which blocks, for a given coarse grid dimensions, like outside of the 'active' study area. The active study area is defined as the region lying upstream of the most downstream samples on each drainage basin. \n",
    "\n",
    "An example of the coverage area, and the corresponding 'active blocks' for the example of 10x10 inversion grid is displayed."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "lowest_tay_sample = loc_nodes[sample_data[:,1]=='55'] # Locality 55\n",
    "tay_catchment = get_watershed_mask(mg,lowest_tay_sample) # extract upstream area of most downstream tay sample \n",
    "\n",
    "lowest_dee_sample = loc_nodes[sample_data[:,1]=='7'] # Locality 7. \n",
    "dee_catchment = get_watershed_mask(mg,lowest_dee_sample)\n",
    "\n",
    "lowest_don_sample = loc_nodes[sample_data[:,1]=='34'] # Locality 34. \n",
    "don_catchment = get_watershed_mask(mg,lowest_don_sample)\n",
    "\n",
    "lowest_spey_sample = loc_nodes[sample_data[:,1]=='28'] # Locality 28. \n",
    "spey_catchment = get_watershed_mask(mg,lowest_spey_sample)\n",
    "\n",
    "lowest_deveron_sample = loc_nodes[sample_data[:,1]=='40'] # Locality 28. \n",
    "deveron_catchment = get_watershed_mask(mg,lowest_deveron_sample)\n",
    "\n",
    "active_area = tay_catchment | dee_catchment | don_catchment | spey_catchment | deveron_catchment\n",
    "\n",
    "def get_active_blocks(nx,ny):\n",
    "    \"\"\"For a given number of blocks in the x \n",
    "    and y direction (nx & ny), returns a (ny,nx) \n",
    "    bool array saying if cell overlaps with active \n",
    "    area or not. \"\"\"\n",
    "    \n",
    "    block_width = np.ceil(model_width/nx) \n",
    "    block_height = np.ceil(model_height/ny)\n",
    "\n",
    "    model_grid_block_indices = np.zeros((model_height,model_width,2))    \n",
    "    for i in np.arange(model_height):\n",
    "        for j in np.arange(model_width):\n",
    "            model_grid_block_indices[i,j,0] = i//block_height\n",
    "            model_grid_block_indices[i,j,1] = j//block_width\n",
    "    model_grid_block_indices = model_grid_block_indices.astype(int)        \n",
    "    # 3D array that contains index of block that model cell corresponds to \n",
    "    # ([:,:,0] = y coordinate; [:,:,1] = x coordinate)\n",
    "    \n",
    "    out = np.zeros((ny,nx)).astype(bool)\n",
    "    for i in np.arange(ny):\n",
    "        for j in np.arange(nx):\n",
    "            # Boolean array of model cells that correspond to block indeix (i,j)\n",
    "            cells_in_block = np.logical_and(model_grid_block_indices[:,:,0] == i, model_grid_block_indices[:,:,1] == j)\n",
    "            # Returns if block overlap with active area:\n",
    "            out[i,j] = np.any(np.logical_and(cells_in_block,active_area.reshape(full_shape)))\n",
    "    return(out)\n",
    "\n",
    "active_blocks = get_active_blocks(10,10)\n",
    "\n",
    "plt.subplot(1,2,1)\n",
    "plt.imshow(active_area.reshape(full_shape),origin='lower')\n",
    "plt.title(\"Active Area\")\n",
    "plt.xlabel(\"Horizontal Model Cells\")\n",
    "plt.ylabel(\"Vertical Model Cells\")\n",
    "\n",
    "plt.subplot(1,2,2)\n",
    "plt.imshow(active_blocks,origin='lower')\n",
    "plt.title(\"Active Blocks\")\n",
    "plt.xlabel(\"Horizontal Inversion Blocks\")\n",
    "plt.ylabel(\"Vertical Inversion Blocks\")\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initial conditions\n",
    "\n",
    "As part of the initial conditions we need to know the weighted average of the most downstream samples. This is calculated now.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "obs_elems = obs_data[elems]\n",
    "spey_mth_comp = np.asarray(obs_elems[obs_data['locality'] == '28'])\n",
    "dee_mth_comp = np.asarray(obs_elems[obs_data['locality'] == '7'])\n",
    "dev_mth_comp = np.asarray(obs_elems[obs_data['locality'] == '40'])\n",
    "tay_mth_comp = np.asarray(obs_elems[obs_data['locality'] == '55'])\n",
    "don_mth_comp = np.asarray(obs_elems[obs_data['locality'] == '34'])\n",
    "\n",
    "prior_wtd_avg = pd.DataFrame((spey_mth_comp*3012.24 + dee_mth_comp*2009.64 + dev_mth_comp*1407.12 + \n",
    "           tay_mth_comp*5096.36 + don_mth_comp*1330.4)/(3012.24+2009.64+1407.12+5096.36+1330.4))\n",
    "prior_wtd_avg.columns = elems\n",
    "prior_wtd_avg = np.mean(prior_wtd_avg,axis=0)\n",
    "prior_wtd_avg_log = np.log(prior_wtd_avg)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Initiating inversion\n",
    "\n",
    "We start the optimisation close to the optimal solution by setting each unique nested catchment to the composition of that sample. Then, for each block in the inversion we then take the average of the upstream geochemistry. This initial condition is a better starting solution than a simple flat initial condition. \n",
    "\n",
    "First we identify the full upstream area of each sample site (note this is a little slow but only needs to be done once)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "unique_locs = np.unique(sample_data[:,1])\n",
    "loc_areas = []\n",
    "for loc_num in unique_locs:\n",
    "    sample_node_num = loc_nodes[sample_data[:,1]==loc_num]\n",
    "    if(sample_node_num.size==2): # Catch duplicate sample exception\n",
    "        sample_node_num = sample_node_num[0]\n",
    "    upst_area = get_watershed_mask(mg,sample_node_num)\n",
    "    loc_areas = loc_areas + [upst_area]\n",
    "loc_areas = np.array(loc_areas) # The full (not unique) upstream area for each sample site."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, we define functions that fetch the unique upstream area for each sample site.  "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def loc_to_area(loc):\n",
    "    \"\"\"Returns the catchment mask for a given locality\"\"\"\n",
    "    return(loc_areas[np.where(unique_locs == loc)])\n",
    "\n",
    "def b_is_nested_in_a(a,b):\n",
    "    \"\"\"Is catchment 'b' a nested subcatchment of catchment 'a'?\"\"\"\n",
    "    return(not(np.any(np.invert(np.logical_or(a,np.invert(b))))))\n",
    "\n",
    "def which_locs_are_contained_upstream_of(loc_a):\n",
    "    \"\"\"Which localities define subcatchments of area defined by loc 'a'\"\"\"\n",
    "    loc_area = loc_to_area(loc_a)\n",
    "    out_locs = []\n",
    "    for loc_num in unique_locs:\n",
    "        if(not(loc_num==loc_a)):\n",
    "            upst_area = loc_to_area(loc_num)\n",
    "            if(b_is_nested_in_a(loc_area,upst_area)):\n",
    "                out_locs = out_locs + [loc_num]\n",
    "    return(out_locs)\n",
    "\n",
    "def find_unique_seg(loc):\n",
    "    \"\"\" Returns the unique upstream area for a sample, that exists between two samples \n",
    "    (or up to a drainage divide if it is most upstream sample)\"\"\"\n",
    "    locs_upstream_of_ = which_locs_are_contained_upstream_of(loc)\n",
    "    downstream_area = loc_to_area(loc)\n",
    "    out = np.zeros(active_area.size).astype(bool)\n",
    "    for upstream_loc in locs_upstream_of_:\n",
    "        out = np.logical_or(out,loc_to_area(upstream_loc))\n",
    "    unique_seg = np.logical_and(downstream_area,np.invert(out))\n",
    "    return(unique_seg.reshape(active_area.shape))"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Now, with these functions defined we can define a 'smart' initiation that sets the nodes to a composition close to the nearest downstream location."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def initiate_blocky_inversion_smart(nx,ny,elem):\n",
    "    \"\"\"Initiates an inversion grid for given number of \n",
    "    cells and element. \"\"\"\n",
    "    \n",
    "    # Define full-res starting solution\n",
    "    full_init = np.zeros(active_area.shape) + prior_wtd_avg_log[elem]  \n",
    "    for loc_num in unique_locs:\n",
    "        values = np.asarray(obs_elems[elem][obs_data['locality'] == loc_num])\n",
    "        if(values.size==2): # Catch duplicate sample exception\n",
    "            full_init[find_unique_seg(loc_num)] = np.mean(np.log(values))\n",
    "        else:\n",
    "            full_init[find_unique_seg(loc_num)] = np.log(values)\n",
    "\n",
    "    # Define inversion nodes\n",
    "    blox = np.zeros((ny,nx))\n",
    "    active_blox = get_active_blocks(nx,ny) # Active cells\n",
    "    \n",
    "    block_x_step = np.ceil(model_width/nx) # Block width\n",
    "    block_y_step = np.ceil(model_height/ny) # Block height\n",
    "    \n",
    "    # Downsample initial guess\n",
    "    model_grid_block_indices = np.zeros((model_height,model_width,2))\n",
    "    for i in np.arange(model_height):\n",
    "        for j in np.arange(model_width):\n",
    "            model_grid_block_indices[i,j,0] = i//block_y_step\n",
    "            model_grid_block_indices[i,j,1] = j//block_x_step\n",
    "    model_grid_block_indices = model_grid_block_indices.astype(int)\n",
    "    # 3D array that contains index of block that model cell corresponds to\n",
    "    # ([:,:,0] = y coordinate; [:,:,1] = x coozzzrdinate)\n",
    "     \n",
    "    for i in np.arange(ny):\n",
    "        for j in np.arange(nx):\n",
    "            # Boolean array of model cells that correspond to block indeix (i,j)\n",
    "            cells_in_block = np.logical_and(model_grid_block_indices[:,:,0] == i, model_grid_block_indices[:,:,1] == j)\n",
    "            # Returns if block overlap with active area:\n",
    "            blox[i,j] = np.mean(full_init.reshape(full_shape)[cells_in_block])\n",
    "    blox[np.invert(active_blox)] =  prior_wtd_avg_log[elem]     \n",
    "    return blox,active_blox,block_x_step,block_y_step"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Cost functions\n",
    "\n",
    "Here we define the cost-functions we use to define an objective function and hence invert our downstream samples. \n",
    "\n",
    "### Data misfit\n",
    "\n",
    "This function calculates the data-misfit between the predicted geochemical composition at sample sites and the true observations. This works by accumulating the tracer along the drainage network. As we do this many times we have sped this up by parallelisation, accumulating basin by basin. Even faster accumulation can be achieved by only accumulating drainage basins that are covered by our samples (`active_area`) but we don't consider that here. (See `accum.get_active_flux_basins_pll`)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def data_misfit(bdrck_arr,elem):\n",
    "    \"\"\"Returns L2norm data misfit for a given bedrock input array (mgkg), calculating predictions for given element assuming homogenous incision\"\"\"\n",
    "    sed_comp = accum.get_flux_basins_pll(mg.at_node['flow__upstream_node_order'],\n",
    "                                                mg.at_node['flow__receiver_node'], starts,\n",
    "                                                tracer_rate = bdrck_arr)  # composition but homogeneous erosion\n",
    "    sed_comp_norm = sed_comp/tot_sed_flux \n",
    "    l2norm = np.linalg.norm(np.log10(obs_data[elem]) - np.log10(sed_comp_norm[loc_nodes]))\n",
    "    return(l2norm)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Roughness\n",
    "\n",
    "This function calculates the x and y components of roughness for a given inversion array"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def cost_roughness(blox,active_blox,block_x,block_y):\n",
    "    \"\"\"Returns l2 norm of roughness in both directions. \n",
    "    Assumes Von Neumann BCs dC/dx = dC/dy = 0\"\"\"\n",
    "    copy = np.copy(blox)\n",
    "    # Set von neumann BCs\n",
    "    copy[np.invert(active_blox)] = 'nan' # set inactive nodes to 'nan'\n",
    "    padded = np.pad(copy,pad_width=1,mode='constant',constant_values='nan') # pad with 'nan' too\n",
    "    x_diffs = np.diff(padded,axis=1)/block_x # dC/dx\n",
    "    y_diffs = np.diff(padded,axis=0)/block_y # dC/dy\n",
    "    x_rough = np.sqrt(np.nansum(x_diffs**2)) # sqrt(SUM((dC/dx)^2)), NB 'nans' are treated as zeros using nansum.\n",
    "    y_rough = np.sqrt(np.nansum(y_diffs**2)) # sqrt(SUM((dC/dy)^2))\n",
    "    return(x_rough,y_rough) # return tuple of roughness along both axes\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### Objective Function \n",
    "\n",
    "These two separate cost functions are hence combined to create an objective function that penalises both data-misfit and roughness. The smoothing coefficient `lamda` must be specified. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "def smoothed_objective(param_arr,blox,active_blox,block_xstep,block_ystep,elem,lamda_):\n",
    "    \"\"\"Tests a given parameter array `param_arr` for the given inversion setup. \n",
    "    Returns the least squares damped cost. Each iteration is ~25 ms\"\"\"\n",
    "    blox[active_blox] = param_arr # Update model grid with new parameters; 1.25 us\n",
    "    bedrock = expand(np.exp(blox),block_xstep,block_ystep) # Convert log blocks into continuous grid in mg/kg; 3.5 ms\n",
    "    data_sq = data_misfit(bedrock.reshape(flat_shape),elem)**2 # Calculate data misfit; 19.4 ms\n",
    "    rough_x,rough_y = cost_roughness(blox,active_blox,block_xstep,block_ystep) # Calculate roughness; 68 us\n",
    "    roughness_sq = (lamda_**2)*(rough_x**2 + rough_y**2) # Roughness squared; 0.6 us\n",
    "    return(data_sq+roughness_sq)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# Inverting\n",
    "\n",
    "The following chunk represents a minimum working example inversion using the code defined above, but optimised using [Powell's algorithm](https://en.wikipedia.org/wiki/Powell%27s_method). We invert for the best fitting upstream geochemistry of Magnesium using an inversion resolution of ~20x20 km (inversion grid of 10x10) and a smoothing coefficient of 2 ~ $10^{0.3}$. This takes ~2 minutes to converge on a standard desktop with Intel i7 processor. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "nx=10 #  # <<<<<<<<<<<<<<<<   Change number of x blocks in inversion grid\n",
    "ny=10 #  # <<<<<<<<<<<<<<<<   Change number of y blocks in inversion grid\n",
    "lamda=2 # <<<<<<<<<<<<<<<<   Change smoothing coefficient \n",
    "elem='Mg' # <<<<<<<<<<<<<<<<  Change element being inverted\n",
    "\n",
    "### Initiating inversion ####\n",
    "\n",
    "blocks,active_blocks,block_width,block_height = initiate_blocky_inversion_smart(nx,ny,elem)\n",
    "parameters = np.copy(blocks[active_blocks]) # The array we will vary. \n",
    "\n",
    "#### Perform inversion ####\n",
    "\n",
    "start = time.time()\n",
    "res_nm = sp.optimize.minimize(fun=smoothed_objective,args=(blocks,active_blocks,block_width,block_height,elem,lamda),x0=parameters,method='Powell',\n",
    "                                  options={'disp':True,'xtol':1e-3,'ftol':1e-3})\n",
    "end = time.time()\n",
    "\n",
    "#### Finish ####\n",
    "\n",
    "print(\"############ results ############\")\n",
    "print(\"Runtime = \",end-start,\"s\")\n",
    "print(res_nm.success) \n",
    "print(res_nm.status)\n",
    "print(res_nm.message)\n",
    "print(res_nm.nit)\n",
    "\n",
    "expanded = expand(np.exp(blocks),block_width,block_height)\n",
    "expanded[np.invert(active_area.reshape(full_shape))] = 'nan'\n",
    "plt.imshow(expanded,origin='lower',norm=LogNorm())\n",
    "plt.colorbar()\n",
    "plt.title(\"Output Mg Concentrations\")\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "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.8.5"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 4
}