{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "50da7606-a0e2-4546-92fe-1ab65c30c3d2",
   "metadata": {},
   "source": [
    "# Multicompartment in the mouse brain\n",
    "\n",
    "\n",
    "This is the script for the 7 compartments Inulin test case\n",
    "This scriipt cqn qlso be used to simulate the 4-compartment case (just put to zero the fluid transfer between blood and PVS compartments)\n",
    "\n",
    "\n",
    "The equations are\n",
    "1) equations for the pore pressures\n",
    "2) advection-diffusion equations for the tracer concentration within the pores\n",
    "\n",
    "\n",
    "For this script we consider 4 compartments:\n",
    " - interstitial space (0)\n",
    " - arteries (1)\n",
    " - veins (2)\n",
    " - capillaries (3)\n",
    " - PVS arteries (4)\n",
    " - PVS veins (5)\n",
    " - PVS capillaries (6)\n",
    "\n",
    "Units are: mm for space and second for time"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "87e45910-1f5d-45f5-ad7d-f29fe85a7fec",
   "metadata": {},
   "source": [
    "## Clone the multicompartment solute transport GitHub repo"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "eb547002-b602-4c99-8621-17495d24f985",
   "metadata": {},
   "outputs": [],
   "source": [
    "!git clone -q https://github.com/annefou/multicompartment-solute-transport.git multicompartment-solute-transport"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "ef346836-9c2a-4d2a-ae98-6b9b41ded04d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# system\n",
    "import sys\n",
    "sys.path.insert(0, 'multicompartment-solute-transport')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4cf3c911-2f0a-452b-876a-0c533263bd06",
   "metadata": {
    "tags": []
   },
   "outputs": [],
   "source": [
    "from dolfin import *\n",
    "from IPython import embed\n",
    "import numpy as np\n",
    "from ts_storage import TimeSeriesStorage\n",
    "from pathlib import Path\n",
    "import matplotlib.pyplot as plt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0e957822-cd7c-4b26-8b11-a51b1f157894",
   "metadata": {},
   "source": [
    "## Choose Boundary type"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2557a26c-c711-45b7-84ea-7d567e9c7ba0",
   "metadata": {},
   "outputs": [],
   "source": [
    "#BC_type = \"Homogeneous\"\n",
    "#BC_type = \"Conservation\"  # You have three choices: Conservation, Homogeneous, Decay\n",
    "BC_type = \"Decay\""
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6447812b-ef63-4669-be40-ab4d942c0aed",
   "metadata": {},
   "source": [
    "### Temporal parameters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "65cb3ebd-e15d-415f-83c6-457b2dc69bbc",
   "metadata": {},
   "outputs": [],
   "source": [
    "dt = 600.\n",
    "T = 3600.*6\n",
    "decay = 0.01/60."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "755ae9a8-dc48-49d9-94f1-5718a2745366",
   "metadata": {},
   "source": [
    "## Load mesh"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5bf3f02b-d7c3-43bb-ab36-efd85eecbadb",
   "metadata": {},
   "outputs": [],
   "source": [
    "res = 16\n",
    "meshfile = Path(f\"mesh/mesh{res}.h5\")\n",
    "\n",
    "# Load mesh\n",
    "mesh = Mesh()\n",
    "hdf = HDF5File(mesh.mpi_comm(), str(meshfile), \"r\")\n",
    "hdf.read(mesh, \"/mesh\", False)\n",
    "SD = MeshFunction(\"size_t\", mesh,mesh.topology().dim())\n",
    "hdf.read(SD, \"/subdomains\")\n",
    "bnd = MeshFunction(\"size_t\", mesh,mesh.topology().dim()-1)\n",
    "hdf.read(bnd, \"/boundaries\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ecf324c2-0635-44d0-b57b-1d49b06400ab",
   "metadata": {},
   "source": [
    "## Define compartments"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "fa59fe8a-4c41-4a18-ad28-3f701f3bb365",
   "metadata": {},
   "outputs": [],
   "source": [
    "comp = ['IS', 'arteries', 'veins', 'capillaries', 'PVS arteries', 'PVS veins', 'PVS capillaries']\n",
    "ncomp = len(comp)\n",
    "geo = mesh.ufl_cell()\n",
    "h = mesh.hmin()\n",
    "print(h)\n",
    "\n",
    "results_path = Path(f\"results/results-{BC_type}-mesh{res}-dt{int(dt)}-inulin-{len(comp)}comps\")\n",
    "results_path.mkdir(parents=True, exist_ok=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "710a70e6-04d4-4fdd-bb17-6be28eec0520",
   "metadata": {},
   "source": [
    "## Finite element functions (P1 for everything)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "73ac6066-05cf-4c5b-a069-ee3629d2d33e",
   "metadata": {},
   "outputs": [],
   "source": [
    "P1 = FiniteElement('CG',geo,1)\n",
    "P2 = FiniteElement('Lagrange',geo,2)\n",
    "ME = MixedElement(ncomp*[P2])\n",
    "Q = FunctionSpace(mesh,ME)\n",
    "V = VectorFunctionSpace(mesh, 'Lagrange', 2)\n",
    "VV = FunctionSpace(mesh,P2)\n",
    "print(Q.dim())\n",
    "p = TrialFunctions(Q)\n",
    "q = TestFunctions(Q)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "442ea409-9984-47cb-b911-b8204c9e7b14",
   "metadata": {},
   "source": [
    "## Load measure data"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "30f14471-f3b3-4dc8-85e2-b8528c77d55d",
   "metadata": {},
   "outputs": [],
   "source": [
    "dx = Measure(\"dx\", domain=mesh, subdomain_data=SD) # Volume\n",
    "ds = Measure('ds')() # surface"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8a8e3262-fa9e-4bee-b785-701c9ee025d7",
   "metadata": {},
   "source": [
    "## Compute surface area of the mesh and its volume"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "77cd8857-f7ba-4888-872f-fc552f3d13c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "surface_area = assemble(1.0*ds(domain=mesh))  # in mm^2\n",
    "print('Surface area: ', surface_area, ' mm^2')\n",
    "brain_volume = assemble(1*dx(domain=mesh))\n",
    "print('brain volume: ', brain_volume, ' mm^3')\n",
    "Vcsf = 0.1 * brain_volume\n",
    "print('brain volume: ', Vcsf, ' mm^3')\n",
    "n = FacetNormal(mesh)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "43b63276-0a07-449e-b379-4c777b4955be",
   "metadata": {},
   "source": [
    "###  PHYSICAL PARAMETERS"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7dcb740-47c0-4940-bd4c-dbb99df45560",
   "metadata": {},
   "source": [
    "#### porosity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "38942f90-cda5-42cd-96fe-ac2277acc8b8",
   "metadata": {},
   "outputs": [],
   "source": [
    "phi0 = ncomp*[Constant(5e-8)]  # Porosity V_i/V_Total\n",
    "phi0[0] = 0.14\n",
    "phi0[1] = 0.00658\n",
    "phi0[2] = 0.02303\n",
    "phi0[3] = 0.00329\n",
    "phi0[4] = 6.e-4\n",
    "phi0[5] = 2.1e-3\n",
    "phi0[6] = 3.e-4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "931b845e-b7b6-4e71-a5bf-93390de9122c",
   "metadata": {},
   "source": [
    "#### viscosity"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3af0347b-2177-48cd-a3bb-a1a389311e5c",
   "metadata": {},
   "outputs": [],
   "source": [
    "nu = np.array([0.0]*ncomp)\n",
    "nu[0] = 7.0e-4\n",
    "nu[1] = 2.67e-3 # blood viscosity\n",
    "nu[2] = 2.67e-3\n",
    "nu[3] = 2.67e-3\n",
    "nu[4] = 7.0e-4\n",
    "nu[5] = 7.0e-4\n",
    "nu[6] = 7.0e-4"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0797245e-999e-44b7-9550-7e603a3620d1",
   "metadata": {},
   "source": [
    "#### Permeability of fluid"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3e3434b4-2f23-45a4-9f50-5144afb9f183",
   "metadata": {},
   "outputs": [],
   "source": [
    "kappa_f = np.array([0.0]*ncomp)\n",
    "kappa_f[0] = 2.0e-11\n",
    "kappa_f[1] = 3.29478e-06\n",
    "kappa_f[2] = 6.58956e-06\n",
    "kappa_f[3] = 1.14276e-09\n",
    "kappa_f[4] = 1.0e-11\n",
    "kappa_f[5] = 6.514285714285714e-09\n",
    "kappa_f[6] = 3.5359801488833745e-13"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "69e0c580-4f52-4973-9850-faf06107bd25",
   "metadata": {},
   "source": [
    "#### transfer coefficients"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "678bd683-123a-4d66-8147-f1a908d93339",
   "metadata": {},
   "outputs": [],
   "source": [
    "w_apa = 5.89e-09\n",
    "w_vpv = 1.26e-10\n",
    "w_cpc = 2.98e-09\n",
    "# WARNING: Comment the floowing 3 lines to have 7 compartments\n",
    "# w_apa = 0\n",
    "# w_vpv = 0\n",
    "# w_cpc = 0\n",
    "\n",
    "w_pae = 2.1932017763179826e-07\n",
    "w_pve = 1.9533203320332029e-07\n",
    "w_pce = 9.976258977745193e-10\n",
    "\n",
    "w_ac = 1.05010501050105e-07\n",
    "w_cv = 5.250525052505252e-07\n",
    "w_papc = 2.5002500250025003e-08\n",
    "w_pcpv = 1.0001000100010001e-07"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0626fd42-9566-4bf5-a94f-c315a25e1206",
   "metadata": {},
   "source": [
    "#### Fluid pressure exchange"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c6e7a0ec-63de-49c0-984d-9cf25c5fc5d1",
   "metadata": {},
   "outputs": [],
   "source": [
    "gamma = np.array([[0,0,0,0,w_pae,w_pve, w_pce],\n",
    "                  [0,0,0,w_ac,w_apa,0,0],\n",
    "                  [0,0,0,w_cv,0,w_vpv,0],\n",
    "                  [0,w_ac,w_cv,0,0,0,w_cpc],\n",
    "                  [w_pae,w_apa,0,0,0,0,w_papc],\n",
    "                  [w_pve,0,w_vpv,0,0,0,w_pcpv],\n",
    "                  [w_pce,0,0,w_cpc,w_papc,w_pcpv,0]])\n",
    "print(gamma)\n",
    "\n",
    "osmo_cap = 20.0*133.33\n",
    "#osmo_e = 0.*osmo_cap  # If you consider the blood compartments, change this line to:\n",
    "osmo_e = 0.2*osmo_cap\n",
    "osmo = np.array([osmo_e,osmo_cap,osmo_cap,osmo_cap,osmo_e,osmo_e,osmo_e])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "99d0caca-5ac9-4abd-adea-001cf765a3b3",
   "metadata": {},
   "source": [
    "### INITIAL CONDITIONS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "286ef2d4-5719-492f-a5b8-3219411ed2dc",
   "metadata": {},
   "outputs": [],
   "source": [
    "F=0\n",
    "# Solving steady state pressure\n",
    "for i in range(ncomp):\n",
    "\tF += kappa_f[i]/nu[i]*inner(grad(p[i]),grad(q[i]))*dx\n",
    "\n",
    "\tfor j in range(ncomp):\n",
    "\t\tif i != j:\n",
    "                        F -=  Constant(gamma[i][j])*inner(p[j]-p[i]-(osmo[j]-osmo[i]),q[i])*dx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0176aa4e-496c-451d-beda-5b182e697a5d",
   "metadata": {},
   "source": [
    "#### Boundary conditions"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ebc5799-7747-41c8-b3a8-58d76c220bc9",
   "metadata": {},
   "outputs": [],
   "source": [
    "F -= 1/(0.01*133.32*60.0/1000.0*1.0e6)*(120.0*133.33*q[1]*ds - p[1]*q[1]*ds)\n",
    "F -= (3.13e-7)*(3.26*133.33*q[0]*ds - p[0]*q[0]*ds)\n",
    "\n",
    "p_CSF = 4.74*133.333\n",
    "F -= (1.25e-6)*(p_CSF*q[4]*ds - p[4]*q[4]*ds)\n",
    "\n",
    "bcs = [DirichletBC(Q.sub(2), Constant(7.0*133.33), 'on_boundary'),DirichletBC(Q.sub(5), Constant(3.26*133.33), 'on_boundary')]"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0b392ac5-00d6-4065-8a55-8da8b6d534c0",
   "metadata": {},
   "source": [
    "#### Solving pressure equations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "81dc2168-7f80-4631-b0c9-58cfad57b847",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Solving pressure system...\", sep=\"\")\n",
    "A = assemble(lhs(F))\n",
    "b = assemble(rhs(F))\n",
    "[bc.apply(A,b) for bc in bcs]\n",
    "p_ = Function(Q)\n",
    "solve(A, p_.vector(), b, 'gmres', 'ilu')\n",
    "print(\"Done.\")\n",
    "# Some manips to get the pressure fields for the pressure equations\n",
    "p_new = p_.split(True)\n",
    "p1,p2,p3,p4,p5,p6,p7 =  p_.split(True)\n",
    "p_0 = Function(VV)\n",
    "assign(p_0,p_new[0])\n",
    "p_1 = Function(VV)\n",
    "assign(p_1,p_new[4])\n",
    "p_2 = Function(VV)\n",
    "assign(p_2,p_new[5])\n",
    "p_3 = Function(VV)\n",
    "assign(p_3,p_new[6])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "40f47ec0-94b1-4d02-9b77-3b3daf390f5a",
   "metadata": {},
   "source": [
    "## PART 2: the diffusion-convection equations\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8df13374-0e1a-4a2c-9d5f-7db0b9214e6b",
   "metadata": {},
   "outputs": [],
   "source": [
    "## Go to 4 compartments because we do not have transport in blood\n",
    "comp = ['IS', 'PVS arteries', 'PVS veins', 'PVS capillaries']\n",
    "ncomp = len(comp)\n",
    "\n",
    "# Diffusion Coefficient\n",
    "D_free = 2.98e-4\n",
    "D_eff =  1.03e-4\n",
    "# porosity\n",
    "phi0 = ncomp*[Constant(5e-8)]  # Porosity V_i/V_Total\n",
    "phi0[0] = 0.14\n",
    "phi0[1] = 6.e-4\n",
    "phi0[3] = 3.e-4\n",
    "phi0[2] = 2.1e-3\n",
    "# viscosity\n",
    "nu = np.array([0.0]*ncomp)\n",
    "nu[0] = 7.0e-4\n",
    "nu[1] = 7.0e-4\n",
    "nu[2] = 7.0e-4\n",
    "nu[3] = 7.0e-4\n",
    "# Permeability of fluid\n",
    "kappa_f = np.array([0.0]*ncomp)\n",
    "kappa_f[0] = 2.0e-11\n",
    "kappa_f[1] = 1.0e-11\n",
    "kappa_f[2] = 6.514285714285714e-09\n",
    "kappa_f[3] = 3.5359801488833745e-13\n",
    "# INULIN exchange\n",
    "sigma_reflect_AEF = 0.2\n",
    "g_ae = w_pae*(1-sigma_reflect_AEF)\n",
    "g_ce = w_pce*(1-sigma_reflect_AEF)\n",
    "g_ve = w_pve*(1-sigma_reflect_AEF)\n",
    "g_ac = w_papc\n",
    "g_cv = w_pcpv\n",
    "gamma_tilde = np.array([[0,g_ae,g_ve,g_ce],[g_ae,0,0,g_ac],[g_ve,0,0,g_cv],[g_ce,g_ac,g_cv,0]])\n",
    "# From diffusion\n",
    "l_ae = 3.744331208456129e-05\n",
    "l_ce = 1.7198523872626686e-05\n",
    "l_ve = 3.744331208456129e-05\n",
    "l_ac = 0\n",
    "l_cv = 0\n",
    "lmbd = np.array([[0,l_ae,l_ve,l_ce],[l_ae,0,0,l_ac],[l_ve,0,0,l_cv],[l_ce,l_ac,l_cv,0]])\n",
    "\n",
    "\n",
    "ME = MixedElement(ncomp*[P2])\n",
    "Q = FunctionSpace(mesh,ME)\n",
    "p_new = Function(Q)\n",
    "assign(p_new, [p_0, p_1,p_2,p_3])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ee1f0c54-b458-4ee9-b555-810baf5060f8",
   "metadata": {},
   "source": [
    "### CONCENTRATIONS EQUATIONS"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c13a62b7-251b-468f-93d0-5cabc65aad06",
   "metadata": {},
   "outputs": [],
   "source": [
    "P1 = FiniteElement('CG',geo,1)\n",
    "P2 = FiniteElement('Lagrange',geo,2)\n",
    "ME = MixedElement(ncomp*[P1])\n",
    "Q = FunctionSpace(mesh,ME)\n",
    "V = VectorFunctionSpace(mesh, 'Lagrange', 1)\n",
    "VV = FunctionSpace(mesh,P1)\n",
    "print(Q.dim())\n",
    "p = TrialFunctions(Q)\n",
    "q = TestFunctions(Q)\n",
    "# Gaussian initial condition.\n",
    "print(\"Computing initial conditions\")\n",
    "center =  (4., 2., 3.)\n",
    "spread = 1.0\n",
    "u0 = Expression(\n",
    "    \"exp(- (pow(x[0]-s[0], 2) + pow(x[1]-s[1], 2) + pow(x[2]-s[2], 2)) / (b * b))\",\n",
    "   degree=1, b=Constant(spread) , s=Constant(center)\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "29114167-2dd1-4cc9-8ff3-8d7b5bd8fd01",
   "metadata": {},
   "source": [
    "#### Project u0 to have a homogeneous Dirichlet boundary (might exists a lot better approaches)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2c2f4dfd-a95d-49da-8537-b90a6905ca75",
   "metadata": {},
   "outputs": [],
   "source": [
    "V = FunctionSpace(mesh, \"Lagrange\", 1)\n",
    "u = TrialFunction(V)\n",
    "v = TestFunction(V)\n",
    "a0 = u * v * dx\n",
    "L0 = u0 * v * dx\n",
    "init_c_ecs = Function(V)\n",
    "solve(a0 == L0, init_c_ecs, bcs=[DirichletBC(V, Constant(0.), \"on_boundary\")])\n",
    "init_c_ecs = project(0.14*init_c_ecs/(phi0[0]),V)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b36e9e4f-cf20-4288-9f34-1b40e89f1ca5",
   "metadata": {},
   "outputs": [],
   "source": [
    "G = 0 # init variational formsplit\n",
    "cn = Function(Q) # function to save concentrations at the previous time step"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "70428007-39df-4fd5-bda7-f9744914bbce",
   "metadata": {},
   "source": [
    "### Variational form for the tracer concentration equations"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "544ee6b7-6aec-456a-8df1-df198955e639",
   "metadata": {},
   "outputs": [],
   "source": [
    "def Max(a, b): return (a+b+abs(a-b))/Constant(2)\n",
    "def Min(a, b): return (a+b-abs(a-b))/Constant(2)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "534f61d5-014d-418a-a8b3-37c68827f61b",
   "metadata": {},
   "outputs": [],
   "source": [
    "for i in range(ncomp):\n",
    "    if i == 0: # If in ISF grey and white matter diffusion\n",
    "        G += Constant(D_eff)*inner(grad(p[i]),grad(q[i]))*dx\n",
    "    else: # Otherwise free diffusion\n",
    "        G += Constant(D_eff)*inner(grad(p[i]),grad(q[i]))*dx\n",
    "\n",
    "    G += 1/dt*inner(p[i]-cn[i],q[i])*dx\n",
    "\n",
    "    G += Constant(kappa_f[i]/(phi0[i]*nu[i]))*inner(p[i],inner(grad(p_new[i]),grad(q[i])))*dx\n",
    "    # mass exchange\n",
    "    for j in range(ncomp):\n",
    "        if i != j:\n",
    "            G += Constant(lmbd[i][j])*inner(p[i]-p[j],q[i])*1./phi0[i]*dx\n",
    "            #G += Constant(gamma_tilde[i][j])*Max(p_new[i]-p_new[j],0)*p[i]*q[i]*1./phi0[i]*dx\n",
    "            #G += Constant(gamma_tilde[i][j])*Min(p_new[i]-p_new[j],0)*p[j]*q[i]*1./phi0[i]*dx\n",
    "            G += Constant(gamma_tilde[i][j])*(p_new[i]-p_new[j])*(p[j]+p[i])*0.5*q[i]*1./phi0[i]*dx"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "225d7288-0dfc-4574-814d-528aade28a19",
   "metadata": {},
   "source": [
    "### Boundary conditions for tracer concentration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e398b012-a158-4a67-a281-82488a5d0fa9",
   "metadata": {},
   "outputs": [],
   "source": [
    "if  BC_type == \"Homogeneous\":\n",
    "    A_in = Expression('C', C=0, degree=1)\n",
    "    bcs_conc = [DirichletBC(Q.sub(1), A_in, 'on_boundary'),DirichletBC(Q.sub(2), A_in, 'on_boundary'),DirichletBC(Q.sub(0), A_in, 'on_boundary')]\n",
    "elif BC_type == \"Conservation\":\n",
    "    g0 = 0.0 # Zero concentration at the beginning\n",
    "    g = Constant(g0)\n",
    "    g1 = Constant(g)\n",
    "    g2 = Constant(g)\n",
    "    g00 = Constant(g)\n",
    "    bcs_conc = [DirichletBC(Q.sub(1), g1, 'on_boundary'),DirichletBC(Q.sub(2), g2, 'on_boundary'),DirichletBC(Q.sub(0), g00, 'on_boundary')]\n",
    "elif BC_type == \"Decay\":\n",
    "    g0 = 0.0 # Zero concentration at the beginning\n",
    "    g = Constant(g0)\n",
    "    g1 = Constant(g)\n",
    "    g2 = Constant(g)\n",
    "    g00 = Constant(g)\n",
    "    bcs_conc = [DirichletBC(Q.sub(1), g1, 'on_boundary'),DirichletBC(Q.sub(2), g2, 'on_boundary'),DirichletBC(Q.sub(0), g00, 'on_boundary')]\n",
    "elif BC_type == \"zeroNeum\":\n",
    "    bcs_conc = []\n",
    "else:\n",
    "    print('Wrong Boundary conditions')\n",
    "    exit(0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "48d490c4-8272-43bb-b128-52e57b038214",
   "metadata": {},
   "source": [
    "### ASSEMBLING"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3faeb1a2-fef7-4ced-9dac-f9068f7d7c94",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(\"Assembling diffusion problem\")\n",
    "a_conc = lhs(G)\n",
    "L_conc = rhs(G)\n",
    "A_conc = assemble(a_conc)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1662c04b-b9a9-4d6b-b123-ed090739c57d",
   "metadata": {},
   "outputs": [],
   "source": [
    "czero = Expression('0.0',degree=1)\n",
    "#init_c_ecs = interpolate(u0, Q.sub(0).collapse())\n",
    "init_other = interpolate(czero,V)\n",
    "c_ = Function(Q) # Function to save solution\n",
    "assign(c_, [init_c_ecs, init_other,init_other,init_other])\n",
    "[bc.apply(c_.vector()) for bc in bcs_conc]\n",
    "c1 = c_.split(True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "dc41c1e1-160d-409c-94b5-a758826a66fc",
   "metadata": {},
   "outputs": [],
   "source": [
    "storage_cecs = TimeSeriesStorage(\"w\", results_path, mesh=mesh, V=VV, name=\"ecs\")\n",
    "storage_carteries = TimeSeriesStorage(\"w\", results_path, mesh=mesh, V=VV, name=\"arteries\")\n",
    "storage_cveins = TimeSeriesStorage(\"w\", results_path, mesh=mesh, V=VV, name=\"veins\")\n",
    "storage_ccap = TimeSeriesStorage(\"w\", results_path, mesh=mesh, V=VV, name=\"capillaries\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de39c663-858d-401f-a9f8-c91256997632",
   "metadata": {},
   "outputs": [],
   "source": [
    "init_c_ecs.rename('concentration', '')\n",
    "storage_cecs.write(c1[0], 0.)\n",
    "storage_carteries.write(c1[1], 0.)\n",
    "storage_cveins.write(c1[2], 0.)\n",
    "storage_ccap.write(c1[3], 0.)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "8540846c-07f9-4ef3-83db-c3dd61fa00d9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Store some vector values for total inulin amount in brain\n",
    "N0 = Constant(phi0[0] * assemble(init_c_ecs * dx))\n",
    "amount = np.zeros(int(T / dt + 1) )\n",
    "amount[0] = 1.0"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "225d2b21-7d3f-4cf9-9173-5035298cbc3b",
   "metadata": {},
   "source": [
    "### Storage for point concentration"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "716975db-a678-4e91-a708-586533506471",
   "metadata": {},
   "outputs": [],
   "source": [
    "concentration_p = np.zeros_like(amount)\n",
    "concentration_p[0] = init_c_ecs(center)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2dfc7634-d470-4305-963c-f6eb2450629e",
   "metadata": {},
   "source": [
    "### Total tracer amount in system"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0c85a66-2e4b-4689-ad2c-43f85151105b",
   "metadata": {},
   "outputs": [],
   "source": [
    "N = np.zeros_like(amount)\n",
    "N[0] = N0\n",
    "print(\"mass_init =\" +str(N[0]))\n",
    "times = np.zeros(len(amount))\n",
    "t=0.0\n",
    "t+=dt\n",
    "it =1\n",
    "transfer_arteries =  np.zeros_like(amount)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "db07faf1-8398-4653-9e98-08f2b606f57d",
   "metadata": {},
   "source": [
    "### Time stepping"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a3cae57f-1f89-4225-8524-a69e0748784c",
   "metadata": {},
   "outputs": [],
   "source": [
    "while t< T + dt/2: # Added dt/2 to ensure final time included.\n",
    "    print('t = ', t)\n",
    "    cn.assign(c_)\n",
    "\n",
    "    b_conc = assemble(L_conc)\n",
    "    [bc.apply(A_conc,b_conc) for bc in bcs_conc]\n",
    "\n",
    "    # Solve\n",
    "    print(\"solving diffusion equations\")\n",
    "\n",
    "    solve(A_conc, c_.vector(), b_conc, 'gmres', 'ilu')\n",
    "    # c1 = c_.split(True)\n",
    "    # storage_cecs.write(c1[0], t) # store the ISF concentration\n",
    "    # storage_carteries.write(c1[1], t) # store the arterial concentration\n",
    "    # storage_cveins.write(c1[2], t) # store the venous concentration\n",
    "\n",
    "\n",
    "\n",
    "    mass_in = 0\n",
    "    for j in range(ncomp):\n",
    "            mass_in += assemble(phi0[j]*c1[j]*dx)\n",
    "    # storage_ccap.write(c1[3], t) # store the capillary concentration\n",
    "\n",
    "    amount[it] = assemble((phi0[0]*c1[0]+phi0[1]*c1[1]+phi0[2]*c1[2]+phi0[3]*c1[3]) * dx)/float(N0)\n",
    "    times[it] = t/60\n",
    "    concentration_p[it] = c1[0](center)\n",
    "\n",
    "\n",
    "\n",
    "         # Prepare boundary conditions\n",
    "    if  BC_type == \"Conservation\":\n",
    "        mass_out = 0\n",
    "        mass_in = 0\n",
    "        mmass = 0\n",
    "        for j in range(ncomp):\n",
    "                mass_in += assemble(phi0[j]*c_[j]*dx)\n",
    "                mmass += assemble(phi0[j]*(c_[j] - cn[j])/dt *dx)\n",
    "        transfer =0.0\n",
    "        for i in range(ncomp):\n",
    "                for j in range(ncomp):\n",
    "                        if i != j:\n",
    "                                transfer += assemble(lmbd[i][j]*(c_[i]-c_[j])*dx)\n",
    "                                transfer += assemble(gamma_tilde[i][j]*(p_new[i]-p_new[j])*(c_[j]+c_[i])*0.5*dx)\n",
    "\n",
    "        print(\"real mass total = \" + str(amount[0]))\n",
    "        print(\"mass inside = \" + str(mass_in))\n",
    "        print(\"g = \" + str(float(g)))\n",
    "        print(\"Total transfer = \" + str(transfer))\n",
    "        g.assign(g - dt*mmass/Vcsf)\n",
    "\n",
    "        g1.assign(g)\n",
    "        g2.assign(g)\n",
    "        g00.assign(g)\n",
    "\n",
    "    elif  BC_type == \"Decay\":\n",
    "        #mass_out = assemble(phi0[0]*D_eff/Vcsf* grad(c1[0])*n * ds + phi0[0]*kappa_f[0]/(Vcsf*nu[0])*c1[0]*grad(p_new[0])*n *ds )\n",
    "        mass_in = 0\n",
    "        mmass = 0\n",
    "        for j in range(ncomp):\n",
    "                mass_in += assemble(phi0[j]*c_[j]*dx)\n",
    "                mmass += assemble(phi0[j]*(c_[j] - cn[j])/dt *dx)\n",
    "        g.assign((1/(1+ decay*dt))*(g - dt/Vcsf *(mmass  )))\n",
    "        g1.assign(g)\n",
    "        g2.assign(g)\n",
    "        g00.assign(g)\n",
    "\n",
    "    it += 1\n",
    "\n",
    "    t+=dt"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "68ebd02d-e132-476b-93a2-1b874b1e3e05",
   "metadata": {},
   "source": [
    "## Storage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d6fb4dee-4224-48a0-b0d8-e7c5c1318407",
   "metadata": {},
   "outputs": [],
   "source": [
    "storage_cecs.write(c1[0], t) # store the ISF concentration\n",
    "storage_carteries.write(c1[1], t) # store the arterial concentration\n",
    "storage_cveins.write(c1[2], t) # store the venous concentration\n",
    "storage_ccap.write(c1[3], t) # store the capillary concentration\n",
    "\n",
    "storage_cecs.store_info()\n",
    "storage_carteries.store_info()\n",
    "storage_cveins.store_info()\n",
    "storage_ccap.store_info()\n",
    "\n",
    "storage_cecs.close()\n",
    "storage_carteries.close()\n",
    "storage_cveins.close()\n",
    "storage_ccap.close()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "49dd2f23-9884-4bb8-85ab-d4ac0a16ff2f",
   "metadata": {},
   "source": [
    "## Compare clearence with diffusion"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b7cb791e-f578-45a9-9d20-f397b39e12b2",
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(times, amount, label=\"Multicompartment\")\n",
    "plt.legend()\n",
    "# plt.ylim(0, None, auto=True)\n",
    "plt.savefig(results_path / f\"inulin-multicomp-{BC_type}-dt{dt}-res{res}-{len(comp)}comps.png\", bbox_inches='tight')\n",
    "plt.show()\n",
    "# save the clearance\n",
    "np.savetxt(results_path / f\"amount-multicomp-{BC_type}-dt{dt}-res{res}-{len(comp)}comps.csv\", amount, delimiter=\",\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "5a89bd37-5e97-4b90-bf61-96440331343e",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "multicompartment",
   "language": "python",
   "name": "multicompartment"
  },
  "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": 5
}