{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output",
     "active-ipynb"
    ]
   },
   "outputs": [],
   "source": [
    "%matplotlib inline\n",
    "from ipyparallel import Client, error  # noqa: F401\n",
    "cluster=Client(profile=\"mpi\")\n",
    "view=cluster[:]\n",
    "view.block=True\n",
    "\n",
    "try:\n",
    "    from openmdao.utils.notebook_utils import notebook_mode  # noqa: F401\n",
    "except ImportError:\n",
    "    !python -m pip install openmdao[notebooks]"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# AnalysisDriver\n",
    "\n",
    "AnalysisDriver serves to run a number of case _samples_ without optimization.\n",
    "The intent is to provide data about the response of a model's outputs when various inputs are changed.\n",
    "\n",
    "## AnalysisDriver Options"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "om.show_options_table(\"openmdao.drivers.analysis_driver.AnalysisDriver\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifying Samples for Execution\n",
    "\n",
    "Upon initialization, samples for AnalysisDriver can be provided in two ways.\n",
    "\n",
    "First, if a list or tuple is provided as the `samples` argument, it contains one or more dictionaries.\n",
    "Each dictionary in the sequence maps the keys (the promoted variable names to be set) to the associated metadata.\n",
    "Key `val` is required, while keys `units` and `indices` are optional.\n",
    "\n",
    "In the following example, nine samples of values for `x` and `y` are provided."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "from openmdao.test_suite.components.paraboloid import Paraboloid\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "prob.model.add_subsystem('comp', Paraboloid(), promotes=['*'])\n",
    "\n",
    "samples_3x3 = [\n",
    "    {'x': {'val': 0.}, 'y': {'val': 0.}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 0.}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 0.}},\n",
    "\n",
    "    {'x': {'val': 0.}, 'y': {'val': 0.5}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 0.5}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 0.5}},\n",
    "\n",
    "    {'x': {'val': 0.}, 'y': {'val': 1.}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 1.}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 1.}},\n",
    "]\n",
    "\n",
    "prob.driver = om.AnalysisDriver(samples=samples_3x3)\n",
    "prob.driver.add_recorder(om.SqliteRecorder('cases.sql'))\n",
    "\n",
    "prob.driver.add_response('f_xy', units=None, indices=[0])\n",
    "\n",
    "prob.setup()\n",
    "prob.run_driver()\n",
    "prob.cleanup()\n",
    "\n",
    "cr = om.CaseReader(str(prob.get_outputs_dir() / \"cases.sql\"))\n",
    "cases = cr.get_cases(source='driver')\n",
    "\n",
    "print(f'Recorded {len(cases)} cases.')\n",
    "\n",
    "print(f'{\"x\":^6}     {\"y\":^6}     {\"f_xy\":^6}')\n",
    "for case in cases:\n",
    "    x = case.get_val('x')[0]\n",
    "    y = case.get_val('y')[0]\n",
    "    f_xy = case.get_val('f_xy')[0]\n",
    "    print(f'{x:6.3f}     {y:6.3f}     {f_xy:6.3f}')"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Specifying Responses for Derivative Recording\n",
    "\n",
    "Since AnalysisDriver does not rely upon the constraints and objectives defined within a model, if the recording\n",
    "of derivatives is requested when executing the AnalysisDriver, the user needs to specify those outputs (responses)\n",
    "whose derivatives are required.\n",
    "\n",
    "AnalysisDriver provides an `add_response` method used to indicate those outputs that should be\n",
    "recorded. When recording derivatives, derivatives will be recorded for all responses.\n",
    "\n",
    "```{eval-rst}\n",
    ".. automethod:: openmdao.drivers.analysis_driver.AnalysisDriver.add_response\n",
    "    :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "from openmdao.test_suite.components.paraboloid import Paraboloid\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "prob.model.add_subsystem('comp', Paraboloid(), promotes=['*'])\n",
    "\n",
    "samples_3x3 = [\n",
    "    {'x': {'val': 0.}, 'y': {'val': 0.}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 0.}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 0.}},\n",
    "\n",
    "    {'x': {'val': 0.}, 'y': {'val': 0.5}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 0.5}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 0.5}},\n",
    "\n",
    "    {'x': {'val': 0.}, 'y': {'val': 1.}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 1.}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 1.}},\n",
    "]\n",
    "\n",
    "prob.driver = om.AnalysisDriver(samples=samples_3x3)\n",
    "prob.driver.add_recorder(om.SqliteRecorder('cases.sql'))\n",
    "prob.driver.recording_options['record_derivatives'] = True\n",
    "\n",
    "prob.driver.add_response('f_xy', units=None, indices=[0])\n",
    "\n",
    "prob.setup()\n",
    "prob.run_driver()\n",
    "prob.cleanup()\n",
    "\n",
    "cr = om.CaseReader(str(prob.get_outputs_dir() / \"cases.sql\"))\n",
    "cases = cr.get_cases(source='driver')\n",
    "\n",
    "\n",
    "print(f'{\"x\":^6}     {\"y\":^6}     {\"f_xy\":^6}     {\"df/dx\":^6}     {\"df/dy\":^6}')\n",
    "for case in cases:\n",
    "    x = case.get_val('x')[0]\n",
    "    y = case.get_val('y')[0]\n",
    "    f_xy = case.get_val('f_xy')[0]\n",
    "    df_dx = case.derivatives['f_xy', 'x']\n",
    "    df_dy = case.derivatives['f_xy', 'y']\n",
    "    print(f'{x:6.3f}     {y:6.3f}     {f_xy:6.3f}     {df_dx[0][0]}     {df_dy[0][0]}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "To quickly record multiple outputs without having an `add_response` call for each one, the user\n",
    "can use the `add_responses` method. This method can accept a sequence of strings, to quickly add\n",
    "outputs to be recorded. Alternatively, providing responses as a dictionary allows the units,\n",
    "indices, and other metadata to be added in a manner similar to samples.\n",
    "\n",
    "```{eval-rst}\n",
    ".. automethod:: openmdao.drivers.analysis_driver.AnalysisDriver.add_responses\n",
    "    :noindex:\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "from openmdao.test_suite.components.paraboloid import Paraboloid\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "prob.model.add_subsystem('comp', Paraboloid(), promotes=['*'])\n",
    "prob.model.add_subsystem('execcomp', om.ExecComp('z = x + y'), promotes=['*'])\n",
    "prob.model.set_input_defaults('x', val=1.0)\n",
    "prob.model.set_input_defaults('y', val=1.0)\n",
    "\n",
    "samples = [\n",
    "    {'x': {'val': 0.}, 'y': {'val': 0.}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 0.25}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 0.125}},\n",
    "]\n",
    "\n",
    "prob.driver = om.AnalysisDriver(samples=samples)\n",
    "prob.driver.add_recorder(om.SqliteRecorder('cases.sql'))\n",
    "\n",
    "prob.driver.add_responses(('f_xy', 'z'))\n",
    "\n",
    "prob.setup()\n",
    "prob.run_driver()\n",
    "prob.cleanup()\n",
    "\n",
    "cr = om.CaseReader(str(prob.get_outputs_dir() / \"cases.sql\"))\n",
    "cases = cr.get_cases(source='driver')\n",
    "\n",
    "print(f'{\"x\":^6}     {\"y\":^6}     {\"f_xy\":^6}     {\"z\":^6}')\n",
    "for case in cases:\n",
    "    x = case.get_val('x')[0]\n",
    "    y = case.get_val('y')[0]\n",
    "    f_xy = case.get_val('f_xy')[0]\n",
    "    z = case.get_val('z')[0]\n",
    "    print(f'{x:6.3f}     {y:6.3f}     {f_xy:6.3f}     {z:6.3f}')\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "### How does AnalysisDriver compare to DOEDriver?\n",
    "\n",
    "AnalysisDriver is intended by be a generalization of DOEDriver that is more capable.\n",
    "\n",
    "DOEDriver was written for a more optimization-centric mindset.  By default, it would only perturb variables marked as design variables of the Problem.\n",
    "\n",
    "AnalysisDriver takes sample values for any input in the model, as well as implicit outputs.\n",
    "Setting the value of an explicit output in the model would have no effect, as the value would be overwritten\n",
    "upon the execution of the model.\n",
    "\n",
    "Unlike DOEDriver, it also allows units and indices to be specified for the variables in each sample.\n",
    "Essentially, any argument to `set_val` can be provided as part of the sample data.\n",
    "\n",
    "DOEDriver also realizes all of the design points to be analyized, while AnalysisDriver's generators do this in a lazy way, which could potentially save some memory when running large data sets.\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Generating Cases\n",
    "\n",
    "Creating a sequence of more than a handful of sample cases can be daunting as the number of model inputs increases.\n",
    "In addition, exceedingly large numbers of cases for large models may begin to tax memory resources if they are all kept in memory.\n",
    "To alleviate these issues, samples for AnalysisDriver can be provided by a lazy pythonic generator descended from an AnalysisGenerator.\n",
    "\n",
    "There are currently three types of AnalysisGenerator included with OpenMDAO\n",
    "\n",
    "- ZipGenerator\n",
    "- ProductGenerator\n",
    "- CSVGenerator"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## ZipGenerator \n",
    "\n",
    "ZipGenerator takes a dictionary that is similar to those given as a sample, except it provids all values to be assumed for each variable.\n",
    "It then uses python's `zip` function to create a sample for each value specified.\n",
    "\n",
    "For instance, the following ZipGenerator will generate three sample cases:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = om.ZipGenerator({'x': {'val': [0.0, 1.0, 2.0], 'units': None},\n",
    "                       'y': {'val': [4.0, 5.0, 6.0], 'units': None}})\n",
    "\n",
    "for i, sample in enumerate(gen):\n",
    "    print(f'{i}:', sample)\n"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Note that the number of values given for each variable to be sampled must be the same.\n",
    "\n",
    "## ProductGenerator \n",
    "\n",
    "ProductGenerator, like ZipGenerator, takes a dictionary that provides values and optionally units and indices for\n",
    "each variable to be sampled.\n",
    "It uses Python's `itertools.product` to produce every possible combination of the values of sampled variables.\n",
    "\n",
    "For instance, the following example will generate 12 sample cases:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "gen = om.ProductGenerator({'x': {'val': [0.0, 1.0, 2.0], 'units': None},\n",
    "                           'y': {'val': [4.0, 5.0, 6.0, 7.0], 'units': None}})\n",
    "\n",
    "for i, sample in enumerate(gen):\n",
    "    print(f'{i}:', sample)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "\n",
    "Unlike ZipGenerator, each variable can have a different number of sample values.\n",
    "\n",
    "## CSVGenerator\n",
    "\n",
    "CSVGenerator provides the ability to read cases from a comma-separated-values file.\n",
    "That is, this allows the user to create a table of cases to be run in a spreadsheet and then execute those cases via the AnalysisDriver.\n",
    "\n",
    "The first row of the CSV file should provide the promoted names of the variables to be sampled.\n",
    "The second row may optionally provide units for each variable to be sampled.\n",
    "The next row may optionally provide the indices to be set for each variable to be sampled.\n",
    "Following that, the remainder of rows provide values for each sampled variable.\n",
    "\n",
    "That is, the following csv file would provide the same nine cases as the ProductGenerator above.\n",
    "The second line contains the units of `x` and `y`.\n",
    "The third line provides the indices being set.\n",
    "These metadata values are unnecessary in this case but provided for illustration.\n",
    "\n",
    "```\n",
    "x, y\n",
    "None, None\n",
    "[0], [0]\n",
    "0.0, 4.0\n",
    "0.0, 5.0\n",
    "0.0, 6.0\n",
    "1.0, 4.0\n",
    "1.0, 5.0\n",
    "1.0, 6.0\n",
    "2.0, 4.0\n",
    "2.0, 5.0\n",
    "2.0, 6.0\n",
    "```"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running AnalysisDriver in Parallel\n",
    "\n",
    "To speed execution, AnalysisDriver can execute samples in parallel when multiple processors are available. This is done by setting the `run_parallel` option to True as shown in the following example and running your script using MPI."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "```{note}\n",
    "This feature requires MPI, and may not be able to be run on Colab or Binder.\n",
    "```"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "import openmdao.api as om\n",
    "from openmdao.test_suite.components.paraboloid import Paraboloid\n",
    "from openmdao.utils.mpi import MPI\n",
    "\n",
    "samples_3x3 = [\n",
    "    {'x': {'val': 0.}, 'y': {'val': 0.}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 0.}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 0.}},\n",
    "\n",
    "    {'x': {'val': 0.}, 'y': {'val': 0.5}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 0.5}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 0.5}},\n",
    "\n",
    "    {'x': {'val': 0.}, 'y': {'val': 1.}},\n",
    "    {'x': {'val': .5}, 'y': {'val': 1.}},\n",
    "    {'x': {'val': 1.}, 'y': {'val': 1.}},\n",
    "]\n",
    "\n",
    "rank = MPI.COMM_WORLD.rank\n",
    "\n",
    "prob = om.Problem()\n",
    "\n",
    "prob.model.add_subsystem('comp', Paraboloid(), promotes=['*'])\n",
    "\n",
    "prob.driver = om.AnalysisDriver(samples=samples_3x3, run_parallel=True)\n",
    "prob.driver.add_recorder(om.SqliteRecorder('cases.sql'))\n",
    "\n",
    "prob.driver.add_response('f_xy', units=None, indices=[0])\n",
    "\n",
    "prob.setup()\n",
    "prob.run_driver()\n",
    "prob.cleanup()\n",
    "\n",
    "cr = om.CaseReader(str(prob.get_outputs_dir() / f'cases.sql_{rank}'))\n",
    "cases = cr.get_cases(source='driver')\n",
    "\n",
    "print(f'Recorded {len(cases)} cases on proc {rank}.')\n",
    "\n",
    "print(f'{\"x\":^6}     {\"y\":^6}     {\"f_xy\":^6}')\n",
    "for case in cases:\n",
    "    x = case.get_val('x')[0]\n",
    "    y = case.get_val('y')[0]\n",
    "    f_xy = case.get_val('f_xy')[0]\n",
    "    print(f'{x:6.3f}     {y:6.3f}     {f_xy:6.3f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "%%px\n",
    "\n",
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "import numpy as np\n",
    "\n",
    "assert(len(cases) == 3 if rank == 0 else 2)\n",
    "# arrays in expected are [x, y, f_xy]\n",
    "if rank == 0:\n",
    "    expected = [[0.0, 0.0, 22.0], [0.5, 0.5, 23.75], [1.0, 1.0, 27.00]]\n",
    "elif rank == 1:\n",
    "    expected = [[0.5, 0.0, 19.25], [1.0, 0.5, 21.75]]\n",
    "elif rank == 2:\n",
    "    expected = [[1.0, 0.0, 17.00], [0.0, 1.0, 31.00]]\n",
    "else:\n",
    "    expected = [[0.0, 0.5, 26.25], [0.5, 1.0, 28.75]]\n",
    "\n",
    "for i, case in enumerate(cases):\n",
    "    x = case.get_val('x')[0]\n",
    "    y = case.get_val('y')[0]\n",
    "    f_xy = case.get_val('f_xy')[0]\n",
    "    assert_near_equal([x, y, f_xy], expected[i])"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Running MPI-enabled models in parallel with AnalysisDriver.\n",
    "\n",
    "If the model being executed supports parallelization, the user should also specify `procs_per_model` to dictate how many processor cores each model instance uses.\n",
    "\n",
    "The following example runs a model that can run two systems simultaneously `c1` and `c2` within a `ParallelGroup`.  In this case, each model can utilize two processors (`procs_per_model = 2`).\n",
    "\n",
    "When providing the number of processors via the `mpirun -n {num_procs} python myscript.py`, the total number of processors must be a multiple of `procs_per_model`."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": "%%px\n\nimport numpy as np\nimport openmdao.api as om\n\nclass FanInGrouped(om.Group):\n    \"\"\"\n    Topology where two components in a Group feed a single component\n    outside of that Group.\n    \"\"\"\n\n    def __init__(self):\n        super().__init__()\n\n        self.set_input_defaults('x1', 1.0)\n        self.set_input_defaults('x2', 1.0)\n\n        self.sub = self.add_subsystem('sub', om.ParallelGroup(),\n                                      promotes_inputs=['x1', 'x2'])\n\n        self.sub.add_subsystem('c1', om.ExecComp(['y=-2.0*x']),\n                               promotes_inputs=[('x', 'x1')])\n        self.sub.add_subsystem('c2', om.ExecComp(['y=5.0*x']),\n                               promotes_inputs=[('x', 'x2')])\n\n        self.add_subsystem('c3', om.ExecComp(['y=3.0*x1+7.0*x2']))\n\n        self.connect(\"sub.c1.y\", \"c3.x1\")\n        self.connect(\"sub.c2.y\", \"c3.x2\")\n\nprob = om.Problem(FanInGrouped())\n\n# Note the absense of adding design varaibles here, compared to DOEGenerator\n\nprob.driver = om.AnalysisDriver(om.ProductGenerator({'x1': {'val': np.linspace(0.0, 1.0, 10)},\n                                                     'x2': {'val': np.linspace(0.0, 1.0, 10)}}))\n\nprob.driver.add_recorder(om.SqliteRecorder(\"cases.sql\"))\n\nprob.driver.add_response('c3.y')\n\n# # the FanInGrouped model uses 2 processes, so we can run\n# # two instances of the model at a time, each using 2 of our 4 procs\nprob.driver.options['run_parallel'] = True\nprob.driver.options['procs_per_model'] = procs_per_model = 2\n\nprob.setup()\nprob.run_driver()\nprob.cleanup()\n\n# a separate case file will be written by rank 0 of each parallel model\n# (the top two global ranks)\nrank = prob.comm.rank\n\nnum_models = prob.comm.size // procs_per_model\n\nif rank < num_models:\n    filename = f'cases.sql_{rank}'\n\n    cr = om.CaseReader(prob.get_outputs_dir() / filename)\n    cases = cr.list_cases('driver', out_stream=None)\n\n    values = []\n    for case in cases:\n        outputs = cr.get_case(case).outputs\n        values.append((outputs['x1'].item(), outputs['x2'].item(), outputs['c3.y'].item()))\n\n    print(\"\\n\"+\"\\n\".join([\"x1: %5.2f, x2: %5.2f, c3.y: %6.2f\" % (x1, x2, y) for x1, x2, y in values]))"
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "py312",
   "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.12.8"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}