{
  "cells": [
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "%matplotlib inline"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Coupled MetaSWAP - Modflow6 model\n\nThis example illustrates how to setup a simple MetaSWAP model coupled to a\nModflow 6 model model using the ``imod`` package and associated packages.\n\nOverview of steps made:\n\n    * Create Modflow 6 model\n    * Create MetaSWAP model\n    * Write coupled models\n"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We'll start with the following imports:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import numpy as np\nimport pandas as pd\nimport xarray as xr\n\nimport imod\nfrom imod import couplers, mf6, msw"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Modflow 6 model\n\nNext, we initiate the Modflow 6 groundwater model:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "gwf_model = mf6.GroundwaterFlowModel()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Create grid\n\nWe'll then define the Modflow 6 grid.\nIt consists of 3 layers of 9 by 9 cells rasters.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "shape = nlay, nrow, ncol = 3, 9, 9\n\ndx = 10.0\ndy = -10.0\ndz = np.array([1.0, 2.0, 10.0])\nxmin = 0.0\nxmax = dx * ncol\nymin = 0.0\nymax = abs(dy) * nrow\ndims = (\"layer\", \"y\", \"x\")\n\nlayer = np.arange(1, nlay + 1)\ny = np.arange(ymax, ymin, dy) + 0.5 * dy\nx = np.arange(xmin, xmax, dx) + 0.5 * dx\ncoords = {\"layer\": layer, \"y\": y, \"x\": x}\n\nidomain = xr.DataArray(np.ones(shape, dtype=int), coords=coords, dims=dims)\n\ntop = 0.0\nbottom = top - xr.DataArray(\n    np.cumsum(layer * dz), coords={\"layer\": layer}, dims=\"layer\"\n)"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "gwf_model[\"dis\"] = mf6.StructuredDiscretization(idomain=idomain, top=top, bottom=bottom)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Hydrogeology\n\n#### Hydraulic conductivity\n\nAssign the node property flow package, which specifies the hydraulic\nconductivities. The middle layer is an aquitard.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "k = xr.DataArray([10.0, 0.1, 10.0], {\"layer\": layer}, (\"layer\",))\nk33 = xr.DataArray([1.0, 0.01, 1.0], {\"layer\": layer}, (\"layer\",))\ngwf_model[\"npf\"] = mf6.NodePropertyFlow(\n    icelltype=0,\n    k=k,\n    k33=k33,\n    save_flows=True,\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Storage\n\nCells are set to non-convertible (convertible = 0). This is a requirement for\nMetaSWAP, because, once coupled, MetaSWAP is responsible for computing the\nstorage coefficient instead of Modflow.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "gwf_model[\"sto\"] = mf6.SpecificStorage(\n    specific_storage=1e-3, specific_yield=0.0, transient=True, convertible=0\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Initial conditions\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "gwf_model[\"ic\"] = mf6.InitialConditions(head=0.5)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Output Control\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "gwf_model[\"oc\"] = mf6.OutputControl(save_head=\"last\", save_budget=\"last\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Boundary conditions\n\n#### Constant head\nWe'll create constant head cells at the most left and right columns of the grid,\nrepresenting two ditches.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "head = xr.full_like(idomain, np.nan, dtype=float)\nhead[0, :, 0] = -1.0\nhead[0, :, -1] = -1.0\n\ngwf_model[\"chd\"] = mf6.ConstantHead(\n    head, print_input=True, print_flows=True, save_flows=True\n)\n\nhead.isel(layer=0).plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "#### Dummy boundary conditions\n\nThe iMOD Coupler requires a dummy recharge package, and well package if\nMetaSWAP's sprinkling is enabled. This to let Modflow 6 allocate the\nappropriate matrices needed in the exchange of states during model\ncomputation.\n\n##### Recharge\n\nWe'll start off with the recharge package, which has no recharge cells\nat the location of our ditches.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "recharge = xr.zeros_like(idomain.sel(layer=1), dtype=float)\nrecharge[:, 0] = np.nan\nrecharge[:, -1] = np.nan\n\ngwf_model[\"rch_msw\"] = mf6.Recharge(recharge)\n\nrecharge.plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "##### Wells\n\nWe'll create a dummy well package as well. imod.mf6.WellDisStructured needs\nits input data provided as long tables instead of grids to so therefore we'll\ncreate 1d arrays by calling ``np.tile`` on the column indices,\nand ``np.repeat`` on the row indices.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "wel_layer = 3\n\nix = np.tile(np.arange(ncol) + 1, nrow)\niy = np.repeat(np.arange(nrow) + 1, ncol)\nrate = np.zeros(ix.shape)\nlayer = np.full_like(ix, wel_layer)\n\ngwf_model[\"wells_msw\"] = mf6.WellDisStructured(\n    layer=layer, row=iy, column=ix, rate=rate\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Initiate a Modflow 6 simulatation and attach the groundwater model to it.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "simulation = mf6.Modflow6Simulation(\"test\")\nsimulation[\"GWF_1\"] = gwf_model\n\n# Define solver settings, we'll use a preset that is sufficient for this example.\n\nsimulation[\"solver\"] = mf6.SolutionPresetSimple(\n    modelnames=[\"GWF_1\"], print_option=\"summary\", csv_output=False, no_ptc=True\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Create time discretization, we'll model 2 days.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "freq = \"D\"\ntimes = pd.date_range(start=\"1/1/1971\", end=\"1/3/1971\", freq=freq)\n\nsimulation.create_time_discretization(additional_times=times)\n\ntimes"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## MetaSWAP model\n\nThe next step is initating a ``MetaSwapModel``. Critical is setting the right\npath to MetaSWAP's soil physical database, which contains the lookup table\nwith the soil physical relationships. Without access to this database MetaSWAP\ncannot function. [The full database can be downloaded here.](http://www.nhi.nu/nl/index.php/data/nhi-lhm/invoer/onverzadigde-zone/)\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "msw_model = msw.MetaSwapModel(unsaturated_database=\"./path/to/unsaturated/database\")\n\n# Create grid\n# ```````````\n#\n# We'll start off specifying the grids required for MetaSWAP. The x,y values\n# of this grid should be identical as the Modflow6 model, but it should\n# not have a layer dimension.\n\nmsw_grid = idomain.sel(layer=1, drop=True).astype(float)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We do not want MetaSWAP cells in the cells where the ditches are located in\nModflow 6. We can specify where MetaSWAP cells are active with the \"active\"\ngrid, which is a grid of booleans (i.e. True/False).\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "active = msw_grid.astype(bool)\nactive[..., 0] = False\nactive[..., -1] = False\n\nactive"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Another crucial grid is the \"area\" grid. The area grid denotes the area in\neach cell, for each \"subunit\". A subunit represent a separate landuse in the\ngrid. We'll create a grid with two separate land uses.\n\nEach grid which specifies parameters related to landuse (e.g. landuse,\nrootzone_depth, ponding depth) requires a subunit dimension. In contrast,\ngrids specifying parameters not induced by landuse (e.g. soil type, elevation,\nprecipitation) cannot contain a subunit dimension.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "subunit = [0, 1]\n\ntotal_cell_area = abs(dx * dy)\nequal_area_per_subunit = total_cell_area / len(subunit)\n\ntotal_cell_area\n\n# Create a full grid equal to the msw_grid. And expand_dims() to broadcast this\n# grid along a new dimension, named \"subunit\"\narea = (\n    xr.full_like(msw_grid, equal_area_per_subunit, dtype=float)\n    .expand_dims(subunit=subunit)\n    .copy()  # expand_dims creates a view, so copy it to allow setting values.\n)\n\n# To the left we only have subunit 0\narea[0, :, :3] = total_cell_area\narea[1, :, :3] = np.nan\n\n# To the right we only have subunit 1\narea[0, :, -3:] = np.nan\narea[1, :, -3:] = total_cell_area\n\narea"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Landuse\n\nDefine a grid with landuse classes.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "landuse = xr.full_like(area, 1, dtype=np.int16)\nlanduse[1, :, :] = 2\n\nlanduse"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Soil types\n\nDefine soil type classes. These will be looked up in MetaSWAP's giant lookup\ntable for the national Staring series describing Dutch soils. `The full\ndatabase can be downloaded here.\n<http://www.nhi.nu/nl/index.php/data/nhi-lhm/invoer/onverzadigde-zone/>` In\nprevious examples we set values in our DataArray using numpy indexing. But we\ncan also use xarray's ``where()`` method to set values.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "slt = xr.full_like(msw_grid, 1, dtype=np.int16)\n# Set all cells on the right half to 2.\nslt = slt.where((slt.x < (xmax / 2)), 2)\n\nslt"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Finishing the grid\nTo finish specifying the landuse grid data, we'll require a rootzone depth\nfor each subunit, and a grid with surface elevations\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "rootzone_depth = xr.full_like(area, 0.5)\nsurface_elevation = xr.full_like(msw_grid, 2.0)\n\nmsw_model[\"grid\"] = msw.GridData(\n    area=area,\n    landuse=landuse,\n    rootzone_depth=rootzone_depth,\n    surface_elevation=surface_elevation,\n    soil_physical_unit=slt,\n    active=active,\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Initial Condition\n\nThere are four options to specify initial conditions,\nsee this for page for an explanation ---link-here---.\nIn this case we opt for an inital pF value of 2.2.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "msw_model[\"ic\"] = msw.InitialConditionsRootzonePressureHead(initial_pF=2.2)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Meteorology\n\nMeteorological information should be provided as grids with a ``time``\ndimension.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "precipitation = msw_grid.expand_dims(time=times[:-1])\nevapotranspiration = precipitation * 1.5\n\nmsw_model[\"meteo_grid\"] = msw.MeteoGrid(precipitation, evapotranspiration)\nmsw_model[\"mapping_prec\"] = msw.PrecipitationMapping(precipitation)\nmsw_model[\"mapping_evt\"] = msw.EvapotranspirationMapping(evapotranspiration)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Ponding\n\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "msw_model[\"ponding\"] = msw.Ponding(\n    ponding_depth=xr.full_like(area, 0.0),\n    runon_resistance=xr.full_like(area, 1.0),\n    runoff_resistance=xr.full_like(area, 1.0),\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Scaling Factors\n\nScaling factors can be defined to adapt some parameters in the soil physical\ndatabase. With this you can investigate the sensitivity of parameters in soil\nphysical database. Furthermore, with this package you can specify the depth of\nthe perched water table.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "msw_model[\"scaling\"] = msw.ScalingFactors(\n    scale_soil_moisture=xr.full_like(area, 1.0),\n    scale_hydraulic_conductivity=xr.full_like(area, 1.0),\n    scale_pressure_head=xr.full_like(area, 1.0),\n    depth_perched_water_table=xr.full_like(msw_grid, 1.0),\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Infiltration Factors\n\nSet the infiltration parameters. We set the resistances to -9999.0, which\nmakes MetaSWAP ignore them.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "msw_model[\"infiltration\"] = msw.Infiltration(\n    infiltration_capacity=xr.full_like(area, 1.0),\n    downward_resistance=xr.full_like(msw_grid, -9999.0),\n    upward_resistance=xr.full_like(msw_grid, -9999.0),\n    bottom_resistance=xr.full_like(msw_grid, -9999.0),\n    extra_storage_coefficient=xr.full_like(msw_grid, 0.1),\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Landuse options\n\nThe landuse option class constructs a lookup table which is used to map\nlanduse indices to a set of parameters. In this example, 3 stands for\npotatoes. This means that for every cell in the ``landuse`` grid with a 3, the\nparameters for a crop with ``vegetation_index == 3`` are associate, which in this\ncase are potatoes.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "vegetation_index = [1, 2, 3]\nnames = [\"grassland\", \"maize\", \"potatoes\"]\n\nlanduse_index = [1, 2, 3]\ncoords = {\"landuse_index\": landuse_index}\n\nlanduse_names = xr.DataArray(data=names, coords=coords, dims=(\"landuse_index\",))\nvegetation_index_da = xr.DataArray(\n    data=vegetation_index, coords=coords, dims=(\"landuse_index\",)\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Because there are a lot of parameters to define, we'll create a DataArray of\nones (``lu``) to more easily broadcast all the different parameters.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "lu = xr.ones_like(vegetation_index_da, dtype=float)\n\nmsw_model[\"landuse_options\"] = msw.LanduseOptions(\n    landuse_name=landuse_names,\n    vegetation_index=vegetation_index_da,\n    jarvis_o2_stress=xr.ones_like(lu),\n    jarvis_drought_stress=xr.ones_like(lu),\n    feddes_p1=xr.full_like(lu, 99.0),\n    feddes_p2=xr.full_like(lu, 99.0),\n    feddes_p3h=lu * [-2.0, -4.0, -3.0],\n    feddes_p3l=lu * [-8.0, -5.0, -5.0],\n    feddes_p4=lu * [-80.0, -100.0, -100.0],\n    feddes_t3h=xr.full_like(lu, 5.0),\n    feddes_t3l=xr.full_like(lu, 1.0),\n    threshold_sprinkling=lu * [-8.0, -5.0, -5.0],\n    fraction_evaporated_sprinkling=xr.full_like(lu, 0.05),\n    gift=xr.full_like(lu, 20.0),\n    gift_duration=xr.full_like(lu, 0.25),\n    rotational_period=lu * [10, 7, 7],\n    start_sprinkling_season=lu * [120, 180, 150],\n    end_sprinkling_season=lu * [230, 230, 240],\n    interception_option=xr.ones_like(lu, dtype=int),\n    interception_capacity_per_LAI=xr.zeros_like(lu),\n    interception_intercept=xr.ones_like(lu),\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Crop Growth\n\nCrop growth tables are specified as a two-dimensional array, with the day of\nyear as one dimension, and the vegetation index on the other. In the vegetation\nfactors, we'll show how to bring some distinction between different crops.\n\nWe'll start off specifiying the coordinates:\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "day_of_year = np.arange(1, 367)\nvegetation_index = np.arange(1, 4)\n\ncoords = {\"day_of_year\": day_of_year, \"vegetation_index\": vegetation_index}"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We can use the coordinates to specify the soil cover of each plant.\nWe'll start with a grid of zeros\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "soil_cover = xr.DataArray(\n    data=np.zeros(day_of_year.shape + vegetation_index.shape),\n    coords=coords,\n    dims=(\"day_of_year\", \"vegetation_index\"),\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The simplest soil cover specification is a step function. In this case soil\ncover equals 1.0 for days 133 to 255 (mind Python's 0-based index here), and\nfor the rest of the days it equals zero.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "soil_cover[132:254, :] = 1.0\n\nsoil_cover.sel(vegetation_index=1).plot()"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We'll simply triple the soil cover to get a leaf area index\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "leaf_area_index = soil_cover * 3"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "Vegetation factors are used to convert the Makkink reference\nevapotranspiration to a potential evapotranspiration for a certain vegetation\ntype. We'll specify some simple crop schemes for the three crops as vegetation\nfactors. Mind that the vegetation factor array has two dimensions:\n``day_of_year`` and ``vegetation_index``\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "vegetation_names = [\"grass\", \"maize\", \"potatoes\"]\n\nvegetation_factor = xr.zeros_like(soil_cover)\n\nvegetation_factor[120:132, :] = [1.0, 0.5, 0.0]\nvegetation_factor[132:142, :] = [1.0, 0.7, 0.7]\nvegetation_factor[142:152, :] = [1.0, 0.8, 0.9]\nvegetation_factor[152:162, :] = [1.0, 0.9, 1.0]\nvegetation_factor[162:172, :] = [1.0, 1.0, 1.2]\nvegetation_factor[172:182, :] = [1.0, 1.2, 1.2]\nvegetation_factor[182:192, :] = [1.0, 1.3, 1.2]\nvegetation_factor[192:244, :] = [1.0, 1.2, 1.1]\nvegetation_factor[244:254, :] = [1.0, 1.2, 0.7]\nvegetation_factor[254:283, :] = [1.0, 1.2, 0.0]\n\n# Since grass is the reference crop, force all grass to 1.0\nvegetation_factor[:, 0] = 1.0\n\n\n# Assign vegetation names for the plot\nvegetation_factor.assign_coords(\n    vegetation_names=(\"vegetation_index\", vegetation_names)\n).plot.line(x=\"day_of_year\", hue=\"vegetation_names\")"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "We'll leave the interception capacity at zero, and the other factors at\none, and assign these to the AnnualCropFactors package.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "msw_model[\"crop_factors\"] = msw.AnnualCropFactors(\n    soil_cover=soil_cover,\n    leaf_area_index=leaf_area_index,\n    interception_capacity=xr.zeros_like(soil_cover),\n    vegetation_factor=vegetation_factor,\n    interception_factor=xr.ones_like(soil_cover),\n    bare_soil_factor=xr.ones_like(soil_cover),\n    ponding_factor=xr.ones_like(soil_cover),\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### Output Control\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "msw_model[\"oc_idf\"] = msw.IdfMapping(area, -9999.0)\nmsw_model[\"oc_var\"] = msw.VariableOutputControl()\nmsw_model[\"oc_time\"] = msw.TimeOutputControl(time=times)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "### MetaSWAP Mappings\n\nMetaSWAP requires its own mapping of SVAT to MODFLOW cells, for internal use.\nWe therefore provide the mf6.StructuredDiscretization and mf6.Well package to\nmf6.CouplerMapping.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "msw_model[\"mod2svat\"] = msw.CouplerMapping(\n    modflow_dis=gwf_model[\"dis\"], well=gwf_model[\"wells_msw\"]\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "The sprinkling package also requires the Modflow6 wells.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "msw_model[\"sprinkling\"] = msw.Sprinkling(\n    max_abstraction_groundwater=xr.full_like(msw_grid, 100.0),\n    max_abstraction_surfacewater=xr.full_like(msw_grid, 100.0),\n    well=gwf_model[\"wells_msw\"],\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Coupler mapping\n\nThe MetaSWAP model and Modflow 6 simulation are provided to the MetaMod class,\nwhich takes care of connecting (= \"mapping\") the two models. Make sure to\nprovide the keys of the dummy Modflow 6 boundary conditions where MetaSWAP is\ncoupled to, so iMOD Python knows where to look: It is technically possible to\ndefine multiple WEL and RCH packages in Modflow 6.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "metamod = couplers.MetaMod(\n    msw_model=msw_model,\n    mf6_simulation=simulation,\n    mf6_rch_pkgkey=\"rch_msw\",\n    mf6_wel_pkgkey=\"wells_msw\",\n)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "By providing a few necessary paths to the modflow and metaswap libraries\nfor iMOD Coupler, we can write the coupled models. You can download the\nmodflow and metaswap libraries as part of the [the last iMOD5 release](https://oss.deltares.nl/web/imod/download-imod5)\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "metamod_dir = imod.util.temporary_directory()\nmf6_dll = \"./path/to/mf6.dll\"\nmetaswap_dll = \"./path/to/metaswap.dll\"\nmetaswap_dll_dependency = \"./path/to/metaswap/dll/dependency\"\n\nmetamod.write(metamod_dir, mf6_dll, metaswap_dll, metaswap_dll_dependency)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Running the models\n\nIn order to run the models, make sure you install ``imod_coupler``. [You can\nfind the installation instructions here.](https://github.com/Deltares/imod_coupler)\n\n"
      ]
    }
  ],
  "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.10.8"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}