tutorials/01_tutorial.ipynb

{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "2189ac9a-3298-4cd3-885b-f82303149bec",
   "metadata": {},
   "source": [
    "# 01 Introduction\n",
    "\n",
    "Multilevel best linear unbiased estimators (MLBLUEs), orginially developed by Schaden and Ullmann [[paper](https://epubs.siam.org/doi/abs/10.1137/19M1263534)], are optimal multilevel/multifidelity estimators for the expectation of quantities of interest (QoIs)\n",
    "affected by uncertainty. Given a set of models to use in the estimation, MLBLUEs automatically select the optimal selection and combination of models to use as well as the optimal number of samples required.\n",
    "**BLUEST** implements the single- and multi-output MLBLUEs by Croci, Willcox and Wright [[paper](https://arxiv.org/pdf/2301.07831.pdf)], as well as the model selection and sample allocation optimal multilevel and multifidelity Monte Carlo methods (MLMC and MFMC respectively).\n",
    "\n",
    "### Problem definition\n",
    "\n",
    "In this tutorial we approximate $\\mathbb{E}[e^Z]$, the expectation of $e^Z$, where $Z$ is a standard Gaussian random variable.\n",
    "\n",
    "We use the following model set:\n",
    "* We take the high-fidelity model to be $e^Z$.\n",
    "* We define $3$ low-fidelity models by truncating the exponential series after $4$, $3$, and $2$ terms.\n",
    "* We define the lowest-fidelity model to be $\\log(|Z|)$ for no particular reason.\n",
    "\n",
    "### Model set implementation"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ddf9d188-3815-428a-aece-8929da33bb2f",
   "metadata": {},
   "source": [
    "In terms of **BLUEST**, most of the users will only ever need to import the **BLUEST** problem class `BLUEProblem`:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "413f30c2-b673-4fcd-82ab-6b0eb2773783",
   "metadata": {},
   "outputs": [],
   "source": [
    "from bluest import BLUEProblem"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7af2f523-3d03-4bb8-b8ae-2dba6a089cbc",
   "metadata": {},
   "source": [
    "The `BLUEProblem` class must be overloaded to define the estimation problem. It is sufficient to only overload two class member functions: `sampler` and `evaluate`, with syntax:\n",
    "\n",
    "**sampler**\n",
    "\n",
    "The `sampler` function samples the model input random parameters.\n",
    "\n",
    "The `sampler` function takes a list `ls` as input where `ls` is a list containing the indices of the models to be sampled. Note that the **highest-fidelity model must always correspond to the 0 index**.\n",
    "\n",
    "The `sampler` function must return a list `samples` so that `len(samples) == len(ls)` so that `samples[i]` contains the input parameters for the model `ls[i]`. These input parameters must be either scalars or numpy arrays.\n",
    "\n",
    "**IMPORTANT:** for the model samples to be correlated the inputs must also be statistically correlated. Often this simply means that the same outputs should be used for all models.\n",
    "\n",
    "```python\n",
    "def sampler(self, ls):\n",
    "    # compute samples, a list of length len(ls)\n",
    "    return samples\n",
    "\n",
    "```\n",
    "\n",
    "**evaluate**\n",
    "\n",
    "The `evaluate` function takes the model input random parameters sampled with `sampler` and computes the model outputs.\n",
    "\n",
    "The `evaluate` function takes `ls` and `samples` as input, where `ls` and `samples` are exactly the same as in the `sampler` function.\n",
    "\n",
    "The `evaluate` function must return a list of list `out` with the following structure\n",
    "```python\n",
    "    L = len(ls)\n",
    "    out = [[0 for i in range(L)] for n in range(n_outputs)] # n_outputs is an integer containing the number of output quantities of interest.\n",
    "```\n",
    "\n",
    "The syntax of `evaluate` for e.g., $1$ output is thus:\n",
    "\n",
    "```python\n",
    "    n_outputs = 1\n",
    "    def evaluate(self, ls, samples):\n",
    "        L = len(ls)\n",
    "        out = [[0 for i in range(L)] for n in range(n_outputs)]\n",
    "\n",
    "        for i in range(L):\n",
    "            for n in range(n_outputs):\n",
    "                model_number = ls[i]\n",
    "                # recall that samples[ls[i]] contains the input parameters for model ls[i]\n",
    "                # out[n][i] will thus contain the n-th output of the model model_number given the input samples[model_number]\n",
    "                out[n][i] = ... # call model model_number with input samples[model_number]\n",
    "                \n",
    "        return out\n",
    "```\n",
    "\n",
    "We now go back to our problem. In this case the `BLUEProblem` class is overloaded as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "id": "a5d0e154-d8d8-415d-a373-968c2c600198",
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "\n",
    "# 5 models\n",
    "n_models = 5\n",
    "\n",
    "# define a function for the truncated exponential series\n",
    "from scipy.special import gamma\n",
    "def exponential_series(x,i):\n",
    "    ii = np.arange(i+1)\n",
    "    return np.sum(x**ii/gamma(ii+1)) # Euler Gamma function (recall \\Gamma(n+1) = n!\n",
    "\n",
    "# overload BLUEProblem and create a new class, MyProblem.\n",
    "class MyProblem(BLUEProblem):\n",
    "    def sampler(self, ls):\n",
    "        L = len(ls)\n",
    "        Z = np.random.randn()\n",
    "        # IMPORTANT: float(Z) here is used to make copies of the same input, one for each model to be sampled.\n",
    "        # If copies are not made, Python will only pass references to the same object which may lead to subtle bugs later on\n",
    "        # in case the models modify the inputs in-place.\n",
    "        samples = [float(Z) for i in range(L)]\n",
    "        return samples\n",
    "\n",
    "    def evaluate(self, ls, samples):\n",
    "        L = len(ls)\n",
    "        out = [[0 for i in range(L)]] # only one output, output 0\n",
    "\n",
    "        for i in range(L):\n",
    "            if ls[i] == 0:\n",
    "                out[0][i] = np.exp(samples[i]) # high-fidelity model corresponding to model index 0\n",
    "            elif ls[i] < n_models-1:\n",
    "                out[0][i] = exponential_series(samples[i], n_models-ls[i]) # truncated exponential series\n",
    "            else:\n",
    "                out[0][i] = np.log(abs(samples[i])) # log(|Z|)\n",
    "\n",
    "        return out\n"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "0bb72abc-bb1b-4a9d-aee0-aa3ce4c00286",
   "metadata": {},
   "source": [
    "We can now create a `MyProblem` class object to setup the estimation problem:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "id": "946c8696-b462-43d0-bb44-62d46e6e4158",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define costs somewhat arbitrarily. If not provided, they will\n",
    "# be estimated via CPU time, which for this problem makes little sense.\n",
    "costs = np.array([2**(n_models-i) for i in range(n_models)])\n",
    "\n",
    "# Default verbose option is True for debugging, you can see everything that goes on\n",
    "# under the hood if you set it to True. Advised if something breaks!\n",
    "# 32 (or even 20) samples are typically enough for application runs. For debugging and Maths papers, set it to 1000.\n",
    "# These samples won't be re-used. Sample re-use introduces bias and is not implemented here yet.\n",
    "problem = MyProblem(n_models, costs=costs, covariance_estimation_samples=32, verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e64a383a-9d2c-4958-82f9-2344571b443c",
   "metadata": {},
   "source": [
    "The above code will instantiate a `problem` object of class `MyProblem`. **BLUEST** will assume by default that we are solving a single-output problem.\n",
    "\n",
    "Instantiating `problem` will trigger an initial estimation of the covariance between all models using $32$ samples. If `costs` vector hadn't been provided, it would have taken two\n",
    "extra samples to estimate the costs via CPU time."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6dfa1d44-513f-4687-a969-5e7eeedfbebf",
   "metadata": {},
   "source": [
    "### Part 1 - Basic Usage\n",
    "\n",
    "We can ask **BLUEST** for the model covariance and correlation matrices, as well as the cost vector: "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "id": "8c089a99-ce76-4a49-87a8-eac7aa65d9c4",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Covariance matrix:\n",
      "\n",
      "[[2.53500581 2.3583669  2.36312599 1.66331444 0.72897923]\n",
      " [2.3583669  2.27345322 2.07702441 1.67029559 0.80520217]\n",
      " [2.36312599 2.07702441 2.40127866 1.37242096 0.47072499]\n",
      " [1.66331444 1.67029559 1.37242096 1.29027193 0.67275196]\n",
      " [0.72897923 0.80520217 0.47072499 0.67275196 1.56637342]]\n",
      "\n",
      "Correlation matrix:\n",
      "\n",
      "[[1.         0.98237857 0.95780283 0.9196959  0.36582891]\n",
      " [0.98237857 1.         0.88894995 0.97523594 0.42669178]\n",
      " [0.95780283 0.88894995 1.         0.77969623 0.24271595]\n",
      " [0.9196959  0.97523594 0.77969623 1.         0.47322393]\n",
      " [0.36582891 0.42669178 0.24271595 0.47322393 1.        ]]\n",
      "\n",
      "Cost vector:\n",
      "\n",
      "[32 16  8  4  2]\n"
     ]
    }
   ],
   "source": [
    "print(\"Covariance matrix:\\n\")\n",
    "print(problem.get_covariance())\n",
    "print(\"\\nCorrelation matrix:\\n\")\n",
    "print(problem.get_correlation())\n",
    "print(\"\\nCost vector:\\n\")\n",
    "print(problem.get_costs())"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "792f23a9-a510-4ebe-9866-99a303392ad1",
   "metadata": {},
   "source": [
    "Let's start solving the estimatione problem. For this purpose, we need to prescribe a statistical error tolerance $\\varepsilon$ or a computational budget.\n",
    "\n",
    "Let's start by prescribing a tolerance of $0.01$ times the standard deviation of the high-fidelity model:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "id": "b24a1e19-7ae1-466e-842c-c5ab1ae69f58",
   "metadata": {},
   "outputs": [],
   "source": [
    "# define statistical error tolerance\n",
    "eps = 0.01*np.sqrt(problem.get_covariance()[0,0])"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "07e8ad76-61ad-4b53-9164-1f1757fd4cb7",
   "metadata": {},
   "source": [
    "**BLUEST** offers $4$ different estimators: standard Monte Carlo, MLMC, MFMC, and MLBLUE.\n",
    "\n",
    "To use standard Monte Carlo, we can call"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "id": "3bb5853e-d53b-4e66-9d93-49a63d4f45ff",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "Std MC\n",
      "\n",
      "Std MC solution:  [1.641081294514136] \n",
      "Total cost:  320000\n"
     ]
    }
   ],
   "source": [
    "# solve with standard MC\n",
    "sol_MC = problem.solve_mc(eps=eps)\n",
    "print(\"\\n\\nStd MC\\n\")\n",
    "print(\"Std MC solution: \", sol_MC[0], \"\\nTotal cost: \", sol_MC[2]) # sol_MC[1] here contains the statistical error, which in this case should be <= eps"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "097b92f6-f673-42d0-b69d-6cc21a886213",
   "metadata": {},
   "source": [
    "The above will compute the samples and return the estimator. Of course, standard Monte Carlo is the slowest method and **BLUEST** can do better.\n",
    "\n",
    "A feature of **BLUEST** is that we can inquire about the cost of MLMC, MFMC and MLBLUE **without actually running any samples**!"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 7,
   "id": "3c7e9e1f-3ac6-465a-9ef0-951333752127",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "MLMC\n",
      "\n",
      "models :  [0, 1, 3]\n",
      "samples :  [ 1118  2701 14523]\n",
      "errors :  [0.015921651552570866]\n",
      "total_cost :  137084\n",
      "\n",
      "\n",
      "MFMC\n",
      "\n",
      "models :  [0, 3]\n",
      "samples :  [ 2819 18671]\n",
      "errors :  [0.01592076469695972]\n",
      "total_cost :  164892\n",
      "alphas :  [array([1.28911929])]\n",
      "\n",
      "\n",
      "MLBLUE\n",
      "\n",
      "models :  [[2], [3], [0, 2, 3], [1, 2, 3], [0, 1, 2, 3]]\n",
      "samples :  [5938 6797    1  493   43]\n",
      "errors :  [0.01592249]\n",
      "total_cost :  91120\n"
     ]
    }
   ],
   "source": [
    "# MLMC\n",
    "print(\"\\n\\nMLMC\\n\")\n",
    "MLMC_data = problem.setup_mlmc(eps=eps)\n",
    "for key, item in MLMC_data.items(): print(key, \": \", item)\n",
    "\n",
    "# MFMC\n",
    "print(\"\\n\\nMFMC\\n\")\n",
    "MFMC_data = problem.setup_mfmc(eps=eps)\n",
    "for key, item in MFMC_data.items(): print(key, \": \", item)\n",
    "\n",
    "# MLBLUE. K denotes the maximum model group size allowed, see our paper.\n",
    "print(\"\\n\\nMLBLUE\\n\")\n",
    "MLBLUE_data = problem.setup_solver(K=n_models, eps=eps)\n",
    "for key, item in MLBLUE_data.items(): print(key, \": \", item)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8e44f171-9bf0-4a22-a6b6-0af41bc9bb92",
   "metadata": {},
   "source": [
    "Note that here the outputs have slightly different meanings depending on the methods.\n",
    "We recommend the reader to familiarise themselves with the MLBLUE, MLMC and MFMC methods to have a clearer picture.\n",
    "\n",
    "For MLMC and MFMC, the variable models contains a list of the models used, while for MLBLUE it contains a list of the model groupings used. **Note** that **BLUEST** has automatically discarded some models as their use has been found suboptimal. The final model choice is of course method-dependent.\n",
    "\n",
    "For MLMC, `samples[i]` contains the number of samples that MLMC would take on MLMC level i. Recall that MLMC on a given level computes samples of differences between models, so that `samples[0]` contains the number of samples taken of the difference between model number `models[0]` and model number `models[1]`.\n",
    "\n",
    "For MFMC, `samples[i]` simply contains the number of samples taken from model number `models[i]`. The variable `alpha` contains the optimal MFMC control-variate coefficients.\n",
    "\n",
    "For MLBLUE, `samples[i]` simply contains the number of samples taken from all models in the model group `models[i]`.\n",
    "\n",
    "---------------------------------------------------------------------------------------------------------------------------\n",
    "\n",
    "Generally speaking, we recommend that the user should always check the costs of the methods since it is very quick and often very informative. Note that MLBLUE should always be more efficient (or more accurate) than the other methods. However, the setup of MLBLUE requires nonlinear integer optimization and, while unlikely, it may happen that the integer solution found is sub-optimal (see paper).\n",
    "\n",
    "To actually use the methods and start sampling, we can simply call:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 8,
   "id": "5b015709-f052-4b84-8dbc-5ca287f91091",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "MLMC\n",
      "\n",
      "MLMC solution:  [1.6462603578070634]\n",
      "\n",
      "\n",
      "MFMC\n",
      "\n",
      "MFMC solution:  [1.6367389524750158]\n",
      "\n",
      "\n",
      "MLBLUE\n",
      "\n",
      "MLBLUE solution:  [1.6328484153177762]\n",
      "\n",
      "Cost comparison. MLMC: 137084.000000, MFMC: 164892.000000, MLBLUE: 91120.000000\n"
     ]
    }
   ],
   "source": [
    "# Solve with MLMC\n",
    "print(\"\\n\\nMLMC\\n\")\n",
    "sol_MLMC = problem.solve_mlmc(eps=eps)\n",
    "print(\"MLMC solution: \", sol_MLMC[0])\n",
    "\n",
    "# Solve with MFMC\n",
    "print(\"\\n\\nMFMC\\n\")\n",
    "sol_MFMC = problem.solve_mfmc(eps=eps)\n",
    "print(\"MFMC solution: \", sol_MFMC[0])\n",
    "\n",
    "# Solve with MLBLUE. K denotes the maximum group size allowed.\n",
    "print(\"\\n\\nMLBLUE\\n\")\n",
    "sol_MLBLUE = problem.solve(K=n_models, eps=eps)\n",
    "print(\"MLBLUE solution: \", sol_MLBLUE[0])\n",
    "\n",
    "# MLBLUE is more sensitive than the other methods to integer projection,\n",
    "# always good to check all methods. This does not require any sampling.\n",
    "MLMC_data   = problem.setup_mlmc(eps=eps)\n",
    "MFMC_data   = problem.setup_mfmc(eps=eps)\n",
    "MLBLUE_data = problem.setup_solver(K=n_models, eps=eps)\n",
    "print(\"\\nCost comparison. MLMC: %f, MFMC: %f, MLBLUE: %f\" % (MLMC_data[\"total_cost\"], MFMC_data[\"total_cost\"], MLBLUE_data[\"total_cost\"]))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "6d8081dc-ca27-46d9-a166-0cf26aff38cd",
   "metadata": {},
   "source": [
    "Note that calling solve will internally call the methods' setup routines as well. For MLBLUE only, previous setup calls will be saved into the `BLUEProblem` object for speed.\n",
    "This is especially useful when the user wants to try out different MLBLUE setup strategies, e.g., rather than prescribing the maximum model group size, we can just prescribe the allowed\n",
    "model groupings directly (this will of course affect the total estimation cost):"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 9,
   "id": "37ca2b4a-0cbb-420f-873a-bb255f4e6168",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "\n",
      "MLBLUE\n",
      "\n",
      "MLBLUE data:\n",
      "\n",
      "models :  [[1], [0, 3], [0, 1, 2, 3, 4]]\n",
      "samples :  [9882 1794  305]\n",
      "errors :  [0.01592225]\n",
      "total_cost :  241606\n"
     ]
    }
   ],
   "source": [
    "# Alternatively, can specify which groups to use:\n",
    "print(\"\\n\\nMLBLUE\\n\")\n",
    "groups = [[0], [1], [0,3], [4,5], [0,1,2,3,4]]\n",
    "MLBLUE_data = problem.setup_solver(groups=groups, eps=eps)\n",
    "print(\"MLBLUE data:\\n\")\n",
    "for key, item in MLBLUE_data.items(): print(key, \": \", item)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9e99e144-d40d-49b8-ae49-012acf849f0f",
   "metadata": {},
   "source": [
    "It is also possible to prescribe a budget:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 10,
   "id": "afbf05f5-355b-429b-aa9b-f707dc064a8d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can also set a budget rather than a tolerance, e.g.\n",
    "budget = 100*max(costs) # budget corresponding to 100 std MC samples\n",
    "# same syntax for MLMC and MFMC. Never provide both eps and budget at the same time\n",
    "MLBLUE_data = problem.setup_solver(K=n_models, budget=budget)\n",
    "\n",
    "#NOTE: no need for calling setup_solver, can call solve directly, although we do not recommend it\n",
    "#      as it is always better to check first."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "749bb0f2-3b5d-4090-861a-e28f9e14af1c",
   "metadata": {},
   "source": [
    "**IMPORTANT** MLBLUE optimization solver. Most of the time this is not needed, but it may happen that the MLBLUE optimization solver fails due to a lack of feature/bug in CVXPY/CVXOPT for which the optimizer does not recognise it has found a solution, and it keeps iterating until failure. If this happens, you can tweak the optimization solver parameters as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 11,
   "id": "da73d755-517b-4369-861d-8dfd4342cd36",
   "metadata": {},
   "outputs": [],
   "source": [
    "cvxopt_params = {\n",
    "        \"abstol\" : 1.e-7,\n",
    "        \"reltol\" : 1.e-4,\n",
    "        \"maxiters\" : 1000, # called max_iters for \"cvxpy\"\n",
    "        \"feastol\" : 1.0e-6,\n",
    "}\n",
    "MLBLUE_data = problem.setup_solver(K=n_models, budget=budget, solver=\"cvxopt\", optimization_solver_params=cvxopt_params)\n",
    "# Changing feastol is typically enough, sometimes you can increase the other tolerances or reduce maxiters"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "76b536a0-6646-40f4-b6f6-ef706c6bce05",
   "metadata": {},
   "source": [
    "__IMPORTANT:__ We now recommend using CVXPY and the CLARABEL solver. CVXOPT and SDPA also work well.\n",
    "See https://www.cvxpy.org/tutorial/advanced/index.html#solve-method-options\n",
    "To change CVXPY solver, setup the solver parameters as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c8a5c20a-9661-4118-9cc5-2701e84943c4",
   "metadata": {},
   "outputs": [],
   "source": [
    "cvxpy_params = {\n",
    "        \"solver\" : \"CLARABEL\",\n",
    "        \"solver_params\" : {}, # A dictionary with CLARABEL solver parameters. See CVXPY and CLARABEL documentation\n",
    "}\n",
    "MLBLUE_data = problem.setup_solver(K=n_models, budget=budget, solver=\"cvxpy\", optimization_solver_params=cvxpy_params)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7a1686db-75cd-46fb-acb4-2c71af410329",
   "metadata": {},
   "source": [
    "### Part 2 - Parallelization\n",
    "\n",
    "**BLUEST** runs in parallel with MPI. Simply call the script with e.g.:\n",
    "\n",
    "```bash\n",
    "> mpiexec -n NPROCS python3 minimal.py \n",
    "```\n",
    "And all the sampling will be split across `NPROCS` workers and occur in parallel (N.B. number of pilot samples will be\n",
    "rounded up so that it is a multiple of `NPROCS`, but the number of online samples will not be increased).\n",
    "\n",
    "#On computing nodes it is often best to set the `OMP_NUM_THREADS` flag to avoid unwanted multithreading, e.g.,\n",
    "\n",
    "```bash\n",
    "> OMP_NUM_THREADS=1 mpiexec -n NPROCS python3 minimal.py \n",
    "```\n",
    "\n",
    "**NOTE:** BLUEST uses mpi4py for MPI parallelization.\n",
    "\n",
    "**NOTE:** Always make sure that the random number generator you use is thread-safe and that you are using independent generators for each worker by using skipahead functionalities (if implemented) or different random seeds. e.g. in Python:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 12,
   "id": "f24440a1-4b95-4df3-b982-8f75b5aef636",
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "1.764052345967664"
      ]
     },
     "execution_count": 12,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "from mpi4py import MPI\n",
    "from numpy.random import RandomState\n",
    "\n",
    "comm = MPI.COMM_WORLD\n",
    "mpiRank = comm.Get_rank()\n",
    "mpiSize = comm.Get_size()\n",
    "\n",
    "RNG = RandomState(mpiRank) # sets a different seed for each worker's random number generator\n",
    "RNG.randn()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9ec03589-44aa-40b8-95ca-b07840e26e95",
   "metadata": {},
   "source": [
    "BLUEST will handle the rest of the parallelization for you so that you do not have to worry about it apart from the random number generator RNG.\n",
    "\n",
    "**NOTE:** If your models also use MPI, then you need to be careful. Simplest option is to set the MPI communicator of your models to `MPI_COMM_SELF` so that each worker loads the full model and deadlocks are avoided. However, BLUEST does support nested MPI communicators: the `BLUEProblem` takes an mpi4py communicator as input:\n",
    "\n",
    "```python\n",
    "    # here mycomm is an mpi4py communicator\n",
    "    problem = MyProblem(n_models, costs=costs, comm = mycomm, covariance_estimation_samples=32, verbose=False)\n",
    "```\n",
    "\n",
    "In the above, samples will be parallelised across each rank of `mycomm` and it is up to the user to modify the `BLUEProblem` class accordingly so that each rank does the right thing and passes the correct MPI communicator to the models. E.g., if we use $4$ workers split into two pairs, we can create a subcommunicator for each pair and a cross-communicator between the two pairs. In this case, we need to pass the cross-communicator to `BLUEProblem` and use the subcommunicators as the MPI communicators of the models so that samples are parallelised across pairs and each model sample is parallelised using a single pair.\n",
    "\n",
    "For an (alas, contrived) example on the use of nested MPI communicators, see the restrictions_matern example in the example/paper_example folder in the **BLUEST** repository. The developers are available via email for any questions."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5556fb63-b8b3-461a-be00-d007ffce2c3d",
   "metadata": {},
   "source": [
    "### Part 3 - Advanced usage"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 13,
   "id": "6b20c444-a5f5-4a7e-bc1a-5c6577b1cbf6",
   "metadata": {},
   "outputs": [],
   "source": [
    "# First some cleanup\n",
    "# Nothing BLUEST-specific in the next two lines: just creating/cleaning up a temporary directory for this script\n",
    "import shutil; shutil.rmtree(\"/tmp/mlblue/\", ignore_errors=True) # cleaning up from previous runs\n",
    "import os; os.makedirs(\"/tmp/mlblue\",exist_ok=True) # creating temporary directory for the script"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f7dc7bb8-68a9-4b45-9f0b-ab4fe684a3c5",
   "metadata": {},
   "source": [
    "Since the offline covariance estimation can be costly, **BLUEST** allows the user to save offline estimation data and load them later, or to use user-provided covariances and cost vectors. "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 14,
   "id": "98045557-e0e8-40c9-844f-d89ab7e543a5",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "\n",
      "BLUE estimator ready.\n",
      "\n",
      "\n",
      "BLUE estimator ready.\n",
      "\n"
     ]
    },
    {
     "data": {
      "text/plain": [
       "array([[        nan,  3.06620309, 11.97634042,  7.66229986,  6.77110948],\n",
       "       [        nan,         nan,  7.70089361,  6.4406335 ,  2.65369473],\n",
       "       [        nan,         nan,         nan, 12.14401703,  6.38489099],\n",
       "       [        nan,         nan,         nan,         nan,  2.58219465],\n",
       "       [        nan,         nan,         nan,         nan,         nan]])"
      ]
     },
     "execution_count": 14,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "# Can save offline estimation data and load them later.\n",
    "problem.save_graph_data(\"/tmp/mlblue/minimal_data.npz\")\n",
    "problem = MyProblem(n_models, datafile=\"/tmp/mlblue/minimal_data.npz\")\n",
    "\n",
    "# Can also overload cost vector when loading.\n",
    "problem = MyProblem(n_models, costs=costs, datafile=\"/tmp/mlblue/minimal_data.npz\")\n",
    "\n",
    "# Similarly, you can avoid offline sampling if covariance and costs are known.\n",
    "C = np.random.randn(n_models, n_models); C = C.T@C;\n",
    "problem = MyProblem(n_models, C = C, costs=costs, verbose=False)\n",
    "\n",
    "# ONLY FOR MLMC:\n",
    "# For MLMC only, the variance of model differences V[P_i-P_j] must be also provided.\n",
    "# In this case only an upper triangular block is needed,\n",
    "# with dV[i,j] = V[P_i - P_j]. The rest must be set to NaN\n",
    "dV = np.nan*np.ones_like(C)\n",
    "for i in range(n_models):\n",
    "    for j in range(i+1,n_models):\n",
    "        # V[P_i - P_j] = V[P_i] + V[P_j] - 2*C(P_i, P_j)\n",
    "        dV[i,j] = C[i,i] + C[j,j] - 2*C[i,j]\n",
    "\n",
    "problem = MyProblem(n_models, C = C, mlmc_variances=dV, costs=costs, verbose=False)\n",
    "# When estimating the covariance with pilot samples, the MLMC variances will also be estimated\n",
    "# and they will also be saved to file with save_graph_data and loaded automatically afterwards.\n",
    "# you can access them with\n",
    "problem.get_mlmc_variance()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "deb04fb1-794f-46a9-bbe8-388b0503ede7",
   "metadata": {},
   "source": [
    "It is also possible to re-estimate some covariances or to exclude some model couplings:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 15,
   "id": "99c26c63-1642-459d-8129-cffc6237baf9",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Covariance estimation with 32 samples...\n",
      "Sampling of models [0, 1, 2, 3, 4] completed.                                                                                                                                                                      \n",
      "Running Spectral Gradient Descent for Covariance projection...\n",
      "Covariance projected, projection error:  3.5252221702063965e-30\n",
      "\n",
      "BLUE estimator ready.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# Can ask to re-estimate some entries by setting entries of C to NaN. Note that\n",
    "# all models might need to be sampled so you might as well re-estimate everything\n",
    "C[0,0] = np.nan\n",
    "problem = MyProblem(n_models, C = C, costs=costs, covariance_estimation_samples=32, verbose=False)\n",
    "\n",
    "# Can exclude model groups also by setting covariance entries to inf, e.g.\n",
    "C = np.nan*C # setting all entries of C to NaN, they will be re-estimated\n",
    "# The covariance between Model 0 and Model 1 won't be estimated\n",
    "# and the two models will never be sampled together\n",
    "# Model set might be pruned after this in case some models become useless\n",
    "C[0,1] = np.inf; C[1,0] = np.inf\n",
    "problem = MyProblem(n_models, C = C, costs=costs, covariance_estimation_samples=32)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8024a020-dfca-484a-9c60-e0195271d44f",
   "metadata": {},
   "source": [
    "Note that independently from whether the covariance matrices are estimated or user-prescribed, they will be checked for positive definiteness. If they are not positive definite,\n",
    "they will be projected to be positive definite. You can skip this projection at your own risk by setting `skip_projection=True`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 16,
   "id": "c1079875-43df-44cf-b57a-051af569c2be",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Covariance estimation with 32 samples...\n",
      "Sampling of models [0, 1, 2, 3, 4] completed.                                                                                                                                                                      \n",
      "\n",
      "BLUE estimator ready.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# You can skip this projection at your own risk by setting skip_projection=True\n",
    "problem = MyProblem(n_models, C = C, costs=costs, covariance_estimation_samples=32, skip_projection=True)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f24c5fb5-6510-4362-b4f9-d52ed1313193",
   "metadata": {},
   "source": [
    "Note that the MLBLUE method requires the covariance matrix to be positive definite so the best course of action may be to actually remove a model by hand rather than skipping or using the projection.\n",
    "This is one of the cases in which you do want to check whether MLMC of MFMC do better than MLBLUE: MLMC and MFMC are unaffected by singular covariances. See our paper for further details about singular covariances."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "id": "c1986db3-f489-4188-9b42-dd47fd7176e1",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Covariance estimation with 32 samples...\n",
      "Sampling of models [0, 1, 2, 3, 4] completed.                                                                                                                                                                      \n",
      "Running Spectral Gradient Descent for Covariance projection...\n",
      "Covariance projected, projection error:  8.960196473267176e-30\n",
      "\n",
      "BLUE estimator ready.\n",
      "\n"
     ]
    }
   ],
   "source": [
    "# NOT SO IMPORTANT:\n",
    "# You can tweak the parameters for the spd projection, but it is almost never needed\n",
    "spg_params = {\"maxit\" : 10000,\n",
    "              \"maxfc\" : 10000**2,\n",
    "              \"verbose\" : False,\n",
    "              \"spd_threshold\" : 5.0e-14,  # minimum eigenvalue\n",
    "              \"eps\"     : 1.0e-10,        # optimization solver tolerance\n",
    "              \"lmbda_min\" : 10.**-30,\n",
    "              \"lmbda_max\" : 10.**30,\n",
    "              \"linesearch_history_length\" : 10,\n",
    "             }\n",
    "# In some very rare cases with near-singular covariance matrices spd_threshold and eps might require some tweaking. \n",
    "problem = MyProblem(n_models, C = C, costs=costs, covariance_estimation_samples=32, spg_params=spg_params)\n",
    "\n",
    "# The projection uses simple SVD hard thresholding for covariances without excluded model groups. Otherwise BLUEST uses a spectral gradient descent method."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "787442b5-9428-4333-8623-c3f48cb43076",
   "metadata": {},
   "source": [
    "You can ask **BLUEST** to save all sample outputs (e.g. snapshots), and their input parameters. All samples will be saved in different npz files with naming convention `snapshots$MODELS.npz`\n",
    "where `$MODELS` corresponds to which models are sample together. New samples will always be appended."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "id": "ad770214-11b3-4bff-bc5b-05971c706f2d",
   "metadata": {},
   "outputs": [],
   "source": [
    "# You can ask BLUEST to save all sample outputs (e.g. snapshots), and their input parameters\n",
    "C = np.random.randn(n_models, n_models); C = C.T@C;\n",
    "problem = MyProblem(n_models, C = C, costs=costs, samplefile=\"/tmp/mlblue/snapshots.npz\", verbose=False)\n",
    "\n",
    "# all samples will be saved in different npz files with naming convention snapshots$MODELS.npz\n",
    "# where $MODELS corresponds to which models are sample together. New samples will always be appended.\n",
    "sol_MLBLUE = problem.solve(K=n_models, eps=eps)\n",
    "\n",
    "# You can avoid saving the pilot samples by setting a samplefile later\n",
    "problem = MyProblem(n_models, costs=costs, covariance_estimation_samples=32, verbose=False)\n",
    "problem.params[\"samplefile\"] = \"/tmp/mlblue/snapshots.npz\"\n",
    "\n",
    "# You can change samplefile name as you go\n",
    "problem.params[\"samplefile\"] = \"/tmp/mlblue/snapshots_MLMC.npz\" # set specific filename for MLMC\n",
    "sol_MLMC = problem.solve_mlmc(eps=eps)\n",
    "problem.params[\"samplefile\"] = \"/tmp/mlblue/snapshots.npz\" # reset to default"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cee7a82a-ea95-4ae3-97dd-6837136b0fca",
   "metadata": {},
   "source": [
    "Batch sampling (taking multiple samples in one go) is supported, but untested. For batch sampling, it is sufficient to define the `sampler` and `evaluate` functions as follows:\n",
    "\n",
    "```python\n",
    "def sampler(self, ls, Nbatch=1):\n",
    "    ...\n",
    "\n",
    "def evaluate(self, ls, samples, Nbatch=1):\n",
    "    # Here out[0][i][n] is the n_th sample in the batch for model i and output 0\n",
    "    return out\n",
    "```\n",
    "\n",
    "If you are interested in this functionality please let the developers know, especially if you find bugs."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "ae0ee8ae-58b8-4fa1-998a-685708a15345",
   "metadata": {},
   "source": [
    "## Part 4 - Multiple outputs\n",
    "\n",
    "Same problem as before, but now we consider two outputs: $e^Z$ and $e^{2Z}$. We use the same models at before to approximate $e^Z$ (and we square the result to approximate $e^{2Z}$).\n",
    "\n",
    "We define the multi-output `BLUEProblem` as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "id": "47f81d5d-fb5c-47ce-915a-ffd22be5d1e1",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Now consider two outputs: e^Z and (e^Z)^2\n",
    "n_outputs = 2\n",
    "\n",
    "class MyMultiOutputProblem(BLUEProblem):\n",
    "    def sampler(self, ls):\n",
    "        L = len(ls)\n",
    "        Z = RNG.randn()\n",
    "        samples = [float(Z) for i in range(L)]\n",
    "        return samples\n",
    "\n",
    "    def evaluate(self, ls, samples):\n",
    "        L = len(ls)\n",
    "        #NOTE the definition of the output.\n",
    "        out = [[0 for i in range(L)] for n in range(n_outputs)]\n",
    "\n",
    "        pw = [1,2]\n",
    "        for i in range(L):\n",
    "            for n in range(n_outputs):\n",
    "                #NOTE the indexing of out\n",
    "                if ls[i] == 0:\n",
    "                    out[n][i] = np.exp(samples[i])**pw[n]\n",
    "                elif ls[i] < n_models-1:\n",
    "                    out[n][i] = exponential_series(samples[i], n_models-ls[i])**pw[n]\n",
    "                else:\n",
    "                    # let's reuse the same quantity for both outputs, why not?\n",
    "                    out[n][i] = np.log(abs(samples[i]))\n",
    "\n",
    "        return out\n",
    "\n",
    "# NOTE: costs are the same as before, i.e. we assume that each model evaluation costs the same for all QoI that it outputs.\n",
    "#       different costs vectors for each output currently not supported. Raise an issue if needed.\n",
    "costs = np.array([2**(n_models-i) for i in range(n_models)])\n",
    "\n",
    "# Now need to specify how many outputs when instantiating the problem object\n",
    "problem = MyMultiOutputProblem(n_models, n_outputs=n_outputs, costs=costs, covariance_estimation_samples=32, verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b0b1b04c-63e6-45f2-a59c-a9f7bde7e4f2",
   "metadata": {},
   "source": [
    "Note that we now had to specify how many outputs we have to the overloaded `BLUEProblem` class instantiation. The rest is mainly the same as before.\n",
    "\n",
    "We setup the multi-output estimators as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 20,
   "id": "0d70c24b-e005-450e-a687-8c4403df4d0c",
   "metadata": {},
   "outputs": [],
   "source": [
    "# same as before with prescribed budget\n",
    "budget = 1000*max(costs) # budget corresponding to 1000 std MC samples\n",
    "MLBLUE_data = problem.setup_solver(K=n_models, budget=budget)\n",
    "\n",
    "# can prescribe a single statistical error tolerance for all outputs:\n",
    "eps = 0.01*np.sqrt(problem.get_covariance(0)[0,0])\n",
    "MLBLUE_data = problem.setup_solver(K=n_models, eps=eps)\n",
    "# or a different tolerance for different outputs\n",
    "eps = [0.01*np.sqrt(problem.get_covariance(n)[0,0]) for n in range(n_outputs)]\n",
    "MLBLUE_data = problem.setup_solver(K=n_models, eps=eps)\n",
    "\n",
    "# Can prescribe a single group for all\n",
    "groups = [[0], [1], [0,3], [4,5], [0,1,2,3,4]]\n",
    "MLBLUE_data = problem.setup_solver(groups=groups, eps=eps)\n",
    "\n",
    "# Or different groups for each\n",
    "groups_0 = [[0], [1], [0,3], [4,5], [0,1,2,3,4]]\n",
    "groups_1 = [[0], [1], [1,3], [3,5], [0,1,3,4]]\n",
    "multi_groups = [groups_0, groups_1]\n",
    "#NOTE: the follwing is untested, if it breaks please raise an issue\n",
    "MLBLUE_data = problem.setup_solver(multi_groups=multi_groups, eps=eps)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "20f1bc86-3be9-4016-bcf5-e5db161fe5ed",
   "metadata": {},
   "source": [
    "We can still prescribe and/or request the model covariances as follows:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 21,
   "id": "be3778e3-04ab-4ff6-adae-19dbe35aabdb",
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[array([[ 4.65733401,  3.64808022, -0.05067521, -2.72783366, -0.54276508],\n",
      "       [ 3.64808022,  4.96631701,  0.03092927, -2.8008079 , -0.89574066],\n",
      "       [-0.05067521,  0.03092927,  2.52874353,  0.51137754, -0.37679739],\n",
      "       [-2.72783366, -2.8008079 ,  0.51137754,  3.74978472, -1.28059895],\n",
      "       [-0.54276508, -0.89574066, -0.37679739, -1.28059895,  2.19306546]]), array([[ 3.94709133,  1.63920466,  2.40555303,  0.85065015,  1.22081507],\n",
      "       [ 1.63920466,  2.08853369,  0.42852247, -0.02972595, -1.72286796],\n",
      "       [ 2.40555303,  0.42852247,  2.61809801,  0.85195781,  1.466295  ],\n",
      "       [ 0.85065015, -0.02972595,  0.85195781,  0.90699405,  0.37464459],\n",
      "       [ 1.22081507, -1.72286796,  1.466295  ,  0.37464459,  4.86135647]])]\n",
      "[array([[ 1.        ,  0.75853997, -0.01476639, -0.65274871, -0.16983117],\n",
      "       [ 0.75853997,  1.        ,  0.00872771, -0.64902726, -0.27141839],\n",
      "       [-0.01476639,  0.00872771,  1.        ,  0.16606805, -0.16000358],\n",
      "       [-0.65274871, -0.64902726,  0.16606805,  1.        , -0.44656448],\n",
      "       [-0.16983117, -0.27141839, -0.16000358, -0.44656448,  1.        ]]), array([[ 1.        ,  0.57091815,  0.74831318,  0.44958331,  0.27869717],\n",
      "       [ 0.57091815,  1.        ,  0.18325672, -0.02159796, -0.5406952 ],\n",
      "       [ 0.74831318,  0.18325672,  1.        ,  0.55286987,  0.41100743],\n",
      "       [ 0.44958331, -0.02159796,  0.55286987,  1.        ,  0.17841794],\n",
      "       [ 0.27869717, -0.5406952 ,  0.41100743,  0.17841794,  1.        ]])]\n",
      "[array([[        nan,  2.32749058,  7.28742795, 13.86278604,  7.93592964],\n",
      "       [        nan,         nan,  7.43320199, 14.31771753,  8.9508638 ],\n",
      "       [        nan,         nan,         nan,  5.25577316,  5.47540377],\n",
      "       [        nan,         nan,         nan,         nan,  8.50404809],\n",
      "       [        nan,         nan,         nan,         nan,         nan]]), array([[        nan,  2.75721569,  1.75408329,  3.15278508,  6.36681766],\n",
      "       [        nan,         nan,  3.84958676,  3.05497963, 10.39562607],\n",
      "       [        nan,         nan,         nan,  1.82117643,  4.54686448],\n",
      "       [        nan,         nan,         nan,         nan,  5.01906133],\n",
      "       [        nan,         nan,         nan,         nan,         nan]])]\n",
      "[[ 4.65733401  3.64808022 -0.05067521 -2.72783366 -0.54276508]\n",
      " [ 3.64808022  4.96631701  0.03092927 -2.8008079  -0.89574066]\n",
      " [-0.05067521  0.03092927  2.52874353  0.51137754 -0.37679739]\n",
      " [-2.72783366 -2.8008079   0.51137754  3.74978472 -1.28059895]\n",
      " [-0.54276508 -0.89574066 -0.37679739 -1.28059895  2.19306546]]\n",
      "[[ 1.          0.75853997 -0.01476639 -0.65274871 -0.16983117]\n",
      " [ 0.75853997  1.          0.00872771 -0.64902726 -0.27141839]\n",
      " [-0.01476639  0.00872771  1.          0.16606805 -0.16000358]\n",
      " [-0.65274871 -0.64902726  0.16606805  1.         -0.44656448]\n",
      " [-0.16983117 -0.27141839 -0.16000358 -0.44656448  1.        ]]\n",
      "[[        nan  2.32749058  7.28742795 13.86278604  7.93592964]\n",
      " [        nan         nan  7.43320199 14.31771753  8.9508638 ]\n",
      " [        nan         nan         nan  5.25577316  5.47540377]\n",
      " [        nan         nan         nan         nan  8.50404809]\n",
      " [        nan         nan         nan         nan         nan]]\n"
     ]
    }
   ],
   "source": [
    "# now need to prescribe a model covariance (and mlmc_variances) for each QoI:\n",
    "# (can use nan and inf as before)\n",
    "C0 = np.random.randn(n_models, n_models); C0 = C0.T@C0;\n",
    "C1 = np.random.randn(n_models, n_models); C1 = C1.T@C1;\n",
    "C = [C0, C1]\n",
    "dV = [np.nan*np.ones_like(c) for c in C]\n",
    "for n in range(n_outputs):\n",
    "    for i in range(n_models):\n",
    "        for j in range(i+1,n_models):\n",
    "            dV[n][i,j] = C[n][i,i] + C[n][j,j] - 2*C[n][i,j]\n",
    "\n",
    "# each covariance matrix will be projected to be spd\n",
    "problem = MyMultiOutputProblem(n_models, n_outputs=n_outputs, C = C, mlmc_variances=dV, costs=costs, verbose=False)\n",
    "\n",
    "# can get them all\n",
    "C = problem.get_covariances()\n",
    "R = problem.get_correlations()\n",
    "dV = problem.get_mlmc_variances()\n",
    "print(C)\n",
    "print(R)\n",
    "print(dV)\n",
    "\n",
    "# or one at a time (the indices 0 and 1 correspond to outputs 0 and 1 respectively)\n",
    "C0 = problem.get_covariance(0)\n",
    "R0 = problem.get_correlation(0)\n",
    "dV0 = problem.get_mlmc_variance(0)\n",
    "C1 = problem.get_covariance(1)\n",
    "R1 = problem.get_correlation(1)\n",
    "dV1 = problem.get_mlmc_variance(1)\n",
    "\n",
    "print(C0)\n",
    "print(R0)\n",
    "print(dV0)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "234ae1e7-a38c-497d-bdc9-ab1f048e391e",
   "metadata": {},
   "source": [
    "We can still store samples for all outputs (or just for some outputs) to file:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 22,
   "id": "da5fec48-784d-41ac-b44c-588894c206d2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Can still store all the samples for all the outputs to file\n",
    "problem = MyMultiOutputProblem(n_models, n_outputs=n_outputs, C = C, costs=costs, samplefile=\"/tmp/mlblue/snapshots_multi.npz\", verbose=False)\n",
    "# or only some of the outputs. Just put them in a list.\n",
    "outputs_to_save = [1]\n",
    "problem = MyMultiOutputProblem(n_models, n_outputs=n_outputs, C = C, costs=costs, outputs_to_save = outputs_to_save, samplefile=\"/tmp/mlblue/snapshots_multi_again.npz\", verbose=False)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "9a1b292a-6506-4887-9825-643528b7dd12",
   "metadata": {},
   "source": [
    "**NOTE:** everything else should be the same as for the single-output case. Any questions please ask the developers!\n",
    "\n",
    "**NOTE:** full python script available in this folder `01_tutorial.py`."
   ]
  }
 ],
 "metadata": {
  "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.12"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}