{
  "cells": [
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "\n# Using Annotated Arrays\n\nThe primary use cases of AnnotatedArrays are to documenting datamodels for\nusers and developers, annotate functions which are expecting specific kinds\nof data, validate data expected to conform to a paritcular model, and to\nfacilitate instantiation datamodels with complex structures.\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "import rmellipse as rme\nimport numpy as np\nfrom pydantic import BaseModel, ConfigDict"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Instantiation\n\nHaving a defined schema makes it easier to initialized empty arrays for\na given datatype by using the AnnotatedArray.zero() method. Coordinate\ndimensions can be supplied, along with metadata you may need. Information\nthat is expected. For example, TimeDomainVoltage2x2 is a rather complex\ndata structure with a lot of requirements, including specific metadata. Using\nthe AnnotatedArray format, you can generate an empty data set fairly\nconcisely.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "class MeasurementMetadata(BaseModel):\n    # Enable extra fields\n    model_config = ConfigDict(extra='allow')\n    operator: str\n    temperature_celcius: float\n\n\nclass TimeAndFrequencySweep(rme.AnnotatedArray):\n    schema = rme.ArraySchema(\n        shape=('N', 'M'),\n        dims=('time', 'frequency'),\n        dtype=float,\n        coords={\n            'time': rme.CoordinateSchema(dtype=float, units='s'),\n            'frequency': rme.CoordinateSchema(dtype=float, units='GHz'),\n        },\n        attrs=MeasurementMetadata,\n    )\n\n\nvoltages = TimeAndFrequencySweep.zeros(\n    time=np.linspace(0, 100, 100),\n    frequency=np.linspace(1, 10, 10),\n    attrs={\n        'operator': 'reader',\n        'temperature_celcius': 23.5,\n    },\n)\n\ncurrent = TimeAndFrequencySweep.zeros(\n    time=np.linspace(0, 100, 100),\n    frequency=np.linspace(1, 10, 10),\n    attrs={\n        'operator': 'reader',\n        'temperature_celcius': 22.5,\n    },\n)\n\n# fill with random values\nvoltages[...] = np.random.rand(*voltages.shape)\ncurrent[...] = np.random.rand(*current.shape)"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Validation\n\nIf you already have an array that you think should perform to a particular\nschema, you can try validating it.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "my_sample = np.zeros((2, 2))\n\ntry:\n    TimeAndFrequencySweep(my_sample).validate()\nexcept rme.ValidationError as e:\n    print(f'array fails validation for : {e}')"
      ]
    },
    {
      "cell_type": "markdown",
      "metadata": {},
      "source": [
        "## Annotating Functions\n\nSince the type is available, you can use it as part of Python's type\nannotation system to help document your code. For example, annotating\nfunctions. Suppose we had a dataset of multiple time sweeps of voltage\nand current measurements at multiple different frequencies. We wanted to\ndetermine what frequency the time average power was the highest, and\nwhat the approximate temperature was. We could write a function to do that,\nand then annotate the inputs and outputs of that function using the\nAnnotatedArray types we defined.\n\n"
      ]
    },
    {
      "cell_type": "code",
      "execution_count": null,
      "metadata": {
        "collapsed": false
      },
      "outputs": [],
      "source": [
        "def max_time_average_power_frequency(\n    voltages: TimeAndFrequencySweep, current: TimeAndFrequencySweep\n) -> tuple[float, float]:\n    \"\"\"\n    Get the frequency corresponding to the maximum time average power.\n\n    Returns the frequency (in GHz) and the approximate temperature based\n    on measurement metadata.\n\n    Parameters\n    ----------\n    voltages : TimeAndFrequencySweep\n        Voltage measurements.\n    current : TimeAndFrequencySweep\n        Current measurements.\n\n    Returns\n    -------\n    frequency : float\n        Frequency where the most time average power was measured\n    temperature : float\n        Approximate temperatue from metadata.\n    \"\"\"\n\n    mean_temp = np.mean(\n        [voltages.attrs['temperature_celcius'], current.attrs['temperature_celcius']]\n    )\n    power = voltages * current\n    time_average = power.mean(dim='time')\n    i_max = np.argmax(time_average.data)\n    return float(time_average.frequency[i_max]), float(mean_temp)\n\n\nfreq, temp = max_time_average_power_frequency(voltages, current)\n\nprint(freq, temp)"
      ]
    }
  ],
  "metadata": {
    "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.12.4"
    }
  },
  "nbformat": 4,
  "nbformat_minor": 0
}