{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "active-ipynb",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "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": [
    "# LinearSystemComp\n",
    "",
    "\n",
    "The LinearSystemComp solves the linear system Ax = b where A and b are inputs, and x is the output.\n",
    "\n",
    "## LinearSystemComp Options\n",
    ""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input"
    ]
   },
   "outputs": [],
   "source": [
    "import openmdao.api as om\n",
    "om.show_options_table(\"openmdao.components.linear_system_comp.LinearSystemComp\")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## LinearSystemComp Constructor\n",
    "",
    "\n",
    "The call signature for the `LinearSystemComp` constructor is:\n",
    "\n",
    "```{eval-rst}\n",
    "    .. automethod:: openmdao.components.linear_system_comp.LinearSystemComp.__init__\n",
    "        :noindex:\n",
    "```\n",
    "\n",
    "## LinearSystemComp Example\n",
    ""
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import openmdao.api as om\n",
    "\n",
    "model = om.Group()\n",
    "\n",
    "A = np.array([[5.0, -3.0, 2.0], [1.0, 7.0, -4.0], [1.0, 0.0, 8.0]])\n",
    "b = np.array([1.0, 2.0, -3.0])\n",
    "\n",
    "lingrp = model.add_subsystem('lingrp', om.Group(), promotes=['*'])\n",
    "lingrp.add_subsystem('lin', om.LinearSystemComp(size=3))\n",
    "\n",
    "prob = om.Problem(model)\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('lin.A', A)\n",
    "prob.set_val('lin.b', b)\n",
    "\n",
    "lingrp.linear_solver = om.ScipyKrylov()\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('lin.x'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "from openmdao.utils.assert_utils import assert_near_equal\n",
    "\n",
    "assert_near_equal(prob.get_val('lin.x'), np.array([0.36423841, -0.00662252, -0.4205298 ]), .0001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This component can also be vectorized to either solve a single linear system with multiple right hand sides, or to solve multiple independent linear systems.\n",
    "\n",
    "You can solve multiple right hand sides by setting the \"vec_size\" argument, giving it the number of right hand sides. When you do this, the LinearSystemComp creates an input for \"b\" such that each row of b is solved independently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = om.Group()\n",
    "\n",
    "A = np.array([[5.0, -3.0, 2.0], [1.0, 7.0, -4.0], [1.0, 0.0, 8.0]])\n",
    "b = np.array([[2.0, -3.0, 4.0], [1.0, 0.0, -1.0]])\n",
    "\n",
    "lingrp = model.add_subsystem('lingrp', om.Group(), promotes=['*'])\n",
    "lingrp.add_subsystem('lin', om.LinearSystemComp(size=3, vec_size=2))\n",
    "\n",
    "prob = om.Problem(model)\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('lin.A', A)\n",
    "prob.set_val('lin.b', b)\n",
    "\n",
    "lingrp.linear_solver = om.ScipyKrylov()\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('lin.x'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('lin.x'), np.array([[ 0.10596026, -0.16556291,  0.48675497],\n",
    "                                                   [ 0.19205298, -0.11258278, -0.14900662]]), .0001)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "To solve multiple linear systems, you just need to set the \"vectorize_A\" option or argument to True. The A\n",
    "matrix is now a 3-dimensional matrix where the first dimension is the number of linear systems to solve."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "model = om.Group()\n",
    "\n",
    "A = np.array([[[5.0, -3.0, 2.0], [1.0, 7.0, -4.0], [1.0, 0.0, 8.0]],\n",
    "              [[2.0, 3.0, 4.0], [1.0, -1.0, -2.0], [3.0, 2.0, -2.0]]])\n",
    "b = np.array([[-5.0, 2.0, 3.0], [-1.0, 1.0, -3.0]])\n",
    "\n",
    "lingrp = model.add_subsystem('lingrp', om.Group(), promotes=['*'])\n",
    "lingrp.add_subsystem('lin', om.LinearSystemComp(size=3, vec_size=2, vectorize_A=True))\n",
    "\n",
    "prob = om.Problem(model)\n",
    "prob.setup()\n",
    "\n",
    "prob.set_val('lin.A', A)\n",
    "prob.set_val('lin.b', b)\n",
    "\n",
    "lingrp.linear_solver = om.ScipyKrylov()\n",
    "\n",
    "prob.run_model()\n",
    "\n",
    "print(prob.get_val('lin.x'))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {
    "tags": [
     "remove-input",
     "remove-output"
    ]
   },
   "outputs": [],
   "source": [
    "assert_near_equal(prob.get_val('lin.x'), np.array([[-0.78807947,  0.66887417,  0.47350993],\n",
    "                                                [ 0.7       , -1.8       ,  0.75      ]]),\n",
    "                 .0001)"
   ]
  }
 ],
 "metadata": {
  "celltoolbar": "Tags",
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.8.1"
  },
  "orphan": true
 },
 "nbformat": 4,
 "nbformat_minor": 4
}
