{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "a1b2c3d4",
   "metadata": {},
   "source": [
    "# Array Variables\n",
    "\n",
    "This notebook demonstrates how to define and use `ArrayVariable` for multi-dimensional inputs and outputs in the Standard Evaluator library. You will learn how `ArrayVariable` extends `FloatVariable` with shape information, how array data is represented in both rolled (NumPy array in a DataFrame cell) and unrolled (individual scalar columns) forms, and how array evaluators integrate with OpenMDAO."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b2c3d4e5",
   "metadata": {},
   "source": [
    "## Overview: ArrayVariable Fields and Concepts\n",
    "\n",
    "`ArrayVariable` is a subclass of `FloatVariable` that represents multi-dimensional NumPy arrays. It adds the following fields:\n",
    "\n",
    "| Field | Type | Description |\n",
    "| ----- | ---- | ----------- |\n",
    "| `shape` | `tuple[int, ...]` | Shape of the array (e.g., `(3,)` for 1-D, `(2, 3)` for 2-D) |\n",
    "| `bounds` | `tuple[array, array]` | Lower and upper bounds — can be scalar (broadcast to shape) or full arrays |\n",
    "| `shift` | `float \\| array` | Shift value(s) for the variable (broadcast to shape if scalar) |\n",
    "| `scale` | `float \\| array` | Scale value(s) for the variable (broadcast to shape if scalar, cannot be zero) |\n",
    "| `default` | `ndarray` | Default values as a NumPy array matching the shape |\n",
    "| `units` | `str` | Units of the variable |\n",
    "\n",
    "### Relationship to FloatVariable\n",
    "\n",
    "`ArrayVariable` inherits from `FloatVariable` and reuses its name validation and general structure, but replaces scalar bounds/shift/scale with array-aware versions. While a `FloatVariable` represents a single scalar value, an `ArrayVariable` represents an entire NumPy array as a single logical variable.\n",
    "\n",
    "### Rolled vs Unrolled Representations\n",
    "\n",
    "- **Rolled**: The array is stored as a single NumPy array inside one DataFrame cell. This is what the `__call__` method expects and returns.\n",
    "- **Unrolled**: Each element of the array becomes its own scalar column (e.g., `arr[0]`, `arr[1]`, `arr[2]` for a shape `(3,)` array). This is how `eval_np` and `eval_list` receive and return data."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c3d4e5f6",
   "metadata": {},
   "source": [
    "## Setup\n",
    "\n",
    "We import the Standard Evaluator library and supporting packages."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d4e5f6a7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import standard_evaluator as se\n",
    "from standard_evaluator.evaluators.numpy_evaluator import NumpyEvaluator\n",
    "from standard_evaluator.components import EvaluatorOpenMdaoComponent\n",
    "import openmdao.api as om\n",
    "import numpy as np\n",
    "import pandas as pd"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "e5f6a7b8",
   "metadata": {},
   "source": [
    "## Creating ArrayVariables\n",
    "\n",
    "We create a 1-D `ArrayVariable` representing a 3-element vector and a 2-D `ArrayVariable` representing a 2×3 matrix. Both have explicit bounds set."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "f6a7b8c9",
   "metadata": {},
   "outputs": [],
   "source": [
    "# 1-D ArrayVariable: a vector of length 3\n",
    "vec_var = se.ArrayVariable(\n",
    "    name=\"vec\",\n",
    "    shape=(3,),\n",
    "    bounds=(-10.0, 10.0),\n",
    "    default=np.array([1.0, 2.0, 3.0]),\n",
    ")\n",
    "\n",
    "print(\"1-D ArrayVariable:\")\n",
    "print(f\"  Name: {vec_var.name}\")\n",
    "print(f\"  Shape: {vec_var.shape}\")\n",
    "print(f\"  Bounds lower: {vec_var.bounds[0]}\")\n",
    "print(f\"  Bounds upper: {vec_var.bounds[1]}\")\n",
    "print(f\"  Default: {vec_var.default}\")\n",
    "print()\n",
    "\n",
    "# 2-D ArrayVariable: a 2x3 matrix\n",
    "mat_var = se.ArrayVariable(\n",
    "    name=\"matrix\",\n",
    "    shape=(2, 3),\n",
    "    bounds=(0.0, 100.0),\n",
    "    default=np.array([[1.0, 2.0, 3.0], [4.0, 5.0, 6.0]]),\n",
    ")\n",
    "\n",
    "print(\"2-D ArrayVariable:\")\n",
    "print(f\"  Name: {mat_var.name}\")\n",
    "print(f\"  Shape: {mat_var.shape}\")\n",
    "print(f\"  Bounds lower:\\n{mat_var.bounds[0]}\")\n",
    "print(f\"  Bounds upper:\\n{mat_var.bounds[1]}\")\n",
    "print(f\"  Default:\\n{mat_var.default}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a7b8c9d0",
   "metadata": {},
   "source": [
    "The scalar bounds `(-10.0, 10.0)` and `(0.0, 100.0)` are automatically broadcast to match the array shape. You can also pass full NumPy arrays for per-element bounds."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "b8c9d0e1",
   "metadata": {},
   "source": [
    "## Using ArrayVariables in a NumpyEvaluator\n",
    "\n",
    "We define an evaluator that takes a 3-element vector input and produces a 3-element vector output by squaring each element. This demonstrates how `__call__` (rolled) and `eval_np` (unrolled) handle array data differently."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "c9d0e1f2",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Define the optimization problem with array input and output\n",
    "opt_problem = se.OptProblem(\n",
    "    name=\"array_demo\",\n",
    "    variables=[\n",
    "        se.ArrayVariable(\n",
    "            name=\"x\",\n",
    "            shape=(3,),\n",
    "            bounds=(-5.0, 5.0),\n",
    "            default=np.array([1.0, 2.0, 3.0]),\n",
    "        ),\n",
    "    ],\n",
    "    responses=[\n",
    "        se.ArrayVariable(\n",
    "            name=\"y\",\n",
    "            shape=(3,),\n",
    "            bounds=(-np.inf, np.inf),\n",
    "        ),\n",
    "    ],\n",
    "    objectives=[],\n",
    "    constraints=[],\n",
    ")\n",
    "\n",
    "\n",
    "# Implement the evaluator: y = x^2 (element-wise)\n",
    "class SquareArray(NumpyEvaluator):\n",
    "    \"\"\"Computes y = x^2 element-wise for a 3-element vector.\"\"\"\n",
    "\n",
    "    def eval_np(self, sites, **kwargs):\n",
    "        # sites is unrolled: shape (n_sites, 3) where columns are x[0], x[1], x[2]\n",
    "        return sites ** 2\n",
    "\n",
    "\n",
    "evaluator = SquareArray(opt_problem=opt_problem)\n",
    "print(\"Inputs:\", evaluator.inputs)\n",
    "print(\"Outputs:\", evaluator.outputs)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d0e1f2a3",
   "metadata": {},
   "source": [
    "### Rolled Invocation via `__call__`\n",
    "\n",
    "With `__call__`, we pass a DataFrame where array columns contain full NumPy arrays in each cell. The evaluator returns results in the same rolled format."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e1f2a3b4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create rolled input DataFrame: each cell in the 'x' column holds a NumPy array\n",
    "sites_rolled = pd.DataFrame({\"x\": [np.array([1.0, 2.0, 3.0]), np.array([4.0, -1.0, 0.5])]})\n",
    "\n",
    "# Call the evaluator (modifies DataFrame in place)\n",
    "evaluator(sites_rolled)\n",
    "\n",
    "print(\"Rolled results (arrays in DataFrame cells):\")\n",
    "print(sites_rolled)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f2a3b4c5",
   "metadata": {},
   "source": [
    "Each cell in the `y` column contains a NumPy array with the squared values. The first row computes `[1, 4, 9]` and the second computes `[16, 1, 0.25]`."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a3b4c5d6",
   "metadata": {},
   "source": [
    "### Unrolled Invocation via `eval_np`\n",
    "\n",
    "With `eval_np`, input and output arrays are unrolled into individual scalar columns. A shape `(3,)` array becomes 3 separate columns in the input and output NumPy arrays."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b4c5d6e7",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Unrolled input: each row has 3 scalar values (x[0], x[1], x[2])\n",
    "sites_unrolled = np.array([\n",
    "    [1.0, 2.0, 3.0],\n",
    "    [4.0, -1.0, 0.5],\n",
    "])\n",
    "\n",
    "# eval_np receives and returns unrolled numpy arrays\n",
    "results_unrolled = evaluator.eval_np(sites_unrolled)\n",
    "\n",
    "print(\"Unrolled input (shape\", sites_unrolled.shape, \"):\")\n",
    "print(sites_unrolled)\n",
    "print()\n",
    "print(\"Unrolled output (shape\", results_unrolled.shape, \"):\")\n",
    "print(results_unrolled)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c5d6e7f8",
   "metadata": {},
   "source": [
    "The unrolled output is a 2D NumPy array where each column corresponds to `y[0]`, `y[1]`, `y[2]`. Both approaches produce the same mathematical result, just in different data representations."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "d6e7f8a9",
   "metadata": {},
   "source": [
    "## OpenMDAO Integration with Array Variables\n",
    "\n",
    "Array evaluators integrate seamlessly with OpenMDAO via `EvaluatorOpenMdaoComponent`. Array variables are registered as OpenMDAO inputs/outputs with their full shape."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e7f8a9b0",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Wrap the array evaluator as an OpenMDAO component\n",
    "prob = om.Problem()\n",
    "prob.model.add_subsystem(\n",
    "    \"comp\", EvaluatorOpenMdaoComponent(evaluator), promotes=[\"*\"]\n",
    ")\n",
    "prob.setup()\n",
    "\n",
    "# Set array input and run\n",
    "prob.set_val(\"x\", np.array([2.0, 3.0, 4.0]))\n",
    "prob.run_model()\n",
    "\n",
    "print(f\"Input x = {prob.get_val('x')}\")\n",
    "print(f\"Output y = x^2 = {prob.get_val('y')}\")\n",
    "print()\n",
    "\n",
    "# Verify shapes are registered correctly\n",
    "comp = prob.model.comp\n",
    "input_meta = comp.get_io_metadata(iotypes=\"input\", metadata_keys=[\"shape\"])\n",
    "output_meta = comp.get_io_metadata(iotypes=\"output\", metadata_keys=[\"shape\"])\n",
    "for name, meta in input_meta.items():\n",
    "    print(f\"Input '{meta['prom_name']}': shape = {meta['shape']}\")\n",
    "for name, meta in output_meta.items():\n",
    "    print(f\"Output '{meta['prom_name']}': shape = {meta['shape']}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "f8a9b0c1",
   "metadata": {},
   "source": [
    "The component correctly computes `y = [4, 9, 16]` for input `x = [2, 3, 4]`. OpenMDAO registers both input and output with shape `(3,)`, matching the `ArrayVariable` definition."
   ]
  },
  {
   "cell_type": "markdown",
   "id": "a9b0c1d2",
   "metadata": {},
   "source": [
    "## Unrolled Name Generation\n",
    "\n",
    "When array variables are unrolled (for `eval_np`, `eval_list`, or internal processing), each element gets a unique name following NumPy indexing syntax:\n",
    "\n",
    "### 1-D Arrays\n",
    "\n",
    "For a variable named `input` with shape `(3,)`, the unrolled names are:\n",
    "\n",
    "- `input[0]`\n",
    "- `input[1]`\n",
    "- `input[2]`\n",
    "\n",
    "### Multi-dimensional Arrays\n",
    "\n",
    "For a variable named `matrix` with shape `(2, 3)`, the unrolled names follow row-major (C-order) iteration:\n",
    "\n",
    "- `matrix[0,0]`, `matrix[0,1]`, `matrix[0,2]`\n",
    "- `matrix[1,0]`, `matrix[1,1]`, `matrix[1,2]`\n",
    "\n",
    "The `generate_names` function and the `OptProblem.unroll_names` method produce these names."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b0c1d2e3",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Demonstrate unrolled name generation\n",
    "from standard_evaluator.problem import generate_names\n",
    "\n",
    "# 1-D names\n",
    "names_1d = generate_names(\"input\", (3,))\n",
    "print(\"1-D unrolled names:\", names_1d)\n",
    "\n",
    "# 2-D names\n",
    "names_2d = generate_names(\"matrix\", (2, 3))\n",
    "print(\"2-D unrolled names:\", names_2d)\n",
    "\n",
    "# Using OptProblem.unroll_names\n",
    "print()\n",
    "print(\"Unrolled variable names from opt_problem:\")\n",
    "print(opt_problem.unroll_names(opt_problem.variables))\n",
    "print(\"Unrolled response names from opt_problem:\")\n",
    "print(opt_problem.unroll_names(opt_problem.responses))"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "c1d2e3f4",
   "metadata": {},
   "source": [
    "The unrolled names use comma-separated indices in brackets, following NumPy's multi-dimensional indexing convention. This naming scheme is used internally when converting between rolled and unrolled representations."
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "name": "python",
   "version": "3.11"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}