{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "0",
   "metadata": {},
   "source": [
    "# Specialized Quantities\n",
    "\n",
    "Beyond `astropy.units.Quantity`, astropy provides further specialized quantities. `astropy-xarray` captures this using the `units.class` attribute."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1",
   "metadata": {},
   "outputs": [],
   "source": [
    "import xarray as xr\n",
    "from astropy import units as u"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "2",
   "metadata": {},
   "source": [
    "## Time and TimeDelta\n",
    "\n",
    "`astropy_xarray` patches `astropy.time.Time` and `astropy.time.TimeDelta` to be recognized xarray, capturing:\n",
    "* format\n",
    "* timescale\n",
    "* precision"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.time import Time, TimeDelta\n",
    "\n",
    "import astropy_xarray  # noqa: F401\n",
    "\n",
    "ds = xr.Dataset(\n",
    "    data_vars={\n",
    "        \"d_time\": (\n",
    "            \"time\",\n",
    "            Time([0, 1, 2], format=\"unix\", scale=\"utc\").copy(\"datetime64\"),\n",
    "        ),\n",
    "        \"s_time\": (\n",
    "            \"time\",\n",
    "            Time([0, 1, 2], format=\"unix\", scale=\"utc\").copy(\"isot\"),\n",
    "        ),\n",
    "        \"interval\": (\"time\", TimeDelta([0.9, 0.9, 0.9], format=\"sec\", scale=\"tai\")),\n",
    "    },\n",
    "    coords=xr.Coordinates(\n",
    "        {\n",
    "            \"f_time\": (\"time\", Time([0, 1, 2], format=\"unix\", scale=\"utc\")),\n",
    "        }\n",
    "    ),\n",
    ")\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4",
   "metadata": {},
   "outputs": [],
   "source": [
    "# dequantified attributes\n",
    "ds.astropy.dequantify()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "5",
   "metadata": {},
   "source": [
    "## SkyCoord, Frame and Angle\n",
    "\n",
    "SkyCoord and Frame fully support conversion to datasets preserving:\n",
    "* Representation\n",
    "* Differential\n",
    "* Frame\n",
    "* Angle, Longitude, Latitude containers\n",
    "\n",
    "Conversion optionally requires coords argument for applying named coordinates to data variable axes.\n",
    "\n",
    "Frame-specific mapped names for data variables coming soon."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6",
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy.coordinates import SkyCoord\n",
    "\n",
    "from astropy_xarray.coordinates.sky_coord import (\n",
    "    skycoord_to_dataset,\n",
    ")\n",
    "\n",
    "sc = SkyCoord(ra=[2, 6, 7, 4] * u.deg, dec=[4, 7, 4, 3] * u.deg, frame=\"icrs\")\n",
    "\n",
    "\n",
    "s = skycoord_to_dataset(sc, coords={\"field\": [\"a\", \"b\", \"c\", \"d\"]})\n",
    "display(s)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "7",
   "metadata": {},
   "source": [
    "## Logarithmic Units and Quantities\n",
    "\n",
    "Astropy provides a generic `LogQuantity` and other various log scale specialized classes and units:\n",
    "\n",
    "$$scalar = A * log_{B}(\\frac{V}{V_0})$$\n",
    "\n",
    "\n",
    "| Name                | Equiv SI Unit | Scale Factor   | Log Base  | 0 Reference                            |\n",
    "| --------------------| ------------- | -------------- | --------- | ---------------------------------------|\n",
    "| Magnitude (absolute)| Unitless      | -2.5           | 10        | Apparent brightness at 10pc away       |\n",
    "| Dex                 | Unitless      | 1              | 10        | 1                                      |\n",
    "| Decibels (power)    | Unitless      | 10             | 10        | Variable                               |\n",
    "| Decibels (amplitude)| Unitless      | 20             | 10        | Variable                               |\n",
    "| Magnitude (apparent)| Unitless      | -2.5           | 10        | Usually Vega or AB from Earth          |\n",
    "| Magnitude (AB)      | Jy            | -2.5           | 10        | 3631 Jy                                |"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "8",
   "metadata": {},
   "source": [
    "### Custom Unit References"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "9",
   "metadata": {},
   "outputs": [],
   "source": [
    "import types\n",
    "\n",
    "import astropy.units\n",
    "import astropy.units as u\n",
    "import numpy as np\n",
    "from astropy.units import Dex, Equivalency, Magnitude, Quantity\n",
    "\n",
    "\n",
    "# extend registry\n",
    "def solar():\n",
    "    solar = types.ModuleType(\"solar\")\n",
    "\n",
    "    # Solar relative metalicity\n",
    "    astropy.units.def_unit(\n",
    "        [\"solH\", \"Xsun\"],\n",
    "        format={\"latex\": \"X_{\\\\odot}\", \"unicode\": \"X☉\"},\n",
    "        namespace=solar.__dict__,\n",
    "    )\n",
    "    astropy.units.def_unit(\n",
    "        [\"solHe\", \"Ysun\"],\n",
    "        format={\"latex\": \"Y_{\\\\odot}\", \"unicode\": \"Y☉\"},\n",
    "        namespace=solar.__dict__,\n",
    "    )\n",
    "    astropy.units.def_unit(\n",
    "        [\"solMetal\", \"Zsun\"],\n",
    "        format={\"latex\": \"Z_{\\\\odot}\", \"unicode\": \"Z☉\"},\n",
    "        namespace=solar.__dict__,\n",
    "    )\n",
    "\n",
    "    return solar\n",
    "\n",
    "\n",
    "if not hasattr(u, \"solar\"):\n",
    "    u.solar = solar()\n",
    "    u.add_enabled_units(u.solar)\n",
    "    u.__dict__.update(u.get_current_unit_registry().registry)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "10",
   "metadata": {},
   "source": [
    "### asplund et. al equivalencies"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "11",
   "metadata": {},
   "outputs": [],
   "source": [
    "def asplund_solar_relative_mass() -> Equivalency:\n",
    "    \"\"\"measurements from https://www.aanda.org/articles/aa/full_html/2021/09/aa40445-21/aa40445-21.html\"\"\"\n",
    "    return Equivalency(\n",
    "        [\n",
    "            (u.Xsun, u.Msun, lambda x: x * 0.7438),\n",
    "            (u.Ysun, u.Msun, lambda x: x * 0.2423),\n",
    "            (u.Zsun, u.Msun, lambda x: x * 0.0139),\n",
    "        ]\n",
    "    )\n",
    "\n",
    "\n",
    "if not hasattr(u, \"asplund_solar_relative_mass\"):\n",
    "    u.asplund_solar_relative_mass = asplund_solar_relative_mass\n",
    "    # uncomment for implicit conversions\n",
    "    # u.add_enabled_equivalencies(u.asplund_solar_relative_mass())"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "12",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds = xr.Dataset(\n",
    "    coords=xr.Coordinates(\n",
    "        {\n",
    "            \"timestamp\": (\"time\", Time([1.7e9], format=\"unix\")),\n",
    "            \"body_label\": (\"body\", [\"sun\", \"proxima centauri\"]),\n",
    "            \"wavelength\": (\"frequency\", [565, 535, 445] * u.nm),\n",
    "        }\n",
    "    ),\n",
    "    data_vars={\n",
    "        \"mass\": ([\"time\", \"body\"], Dex([[0.0, 0.02]], unit=u.dex(u.Msun))),\n",
    "        \"distance\": ([\"time\", \"body\"], [[1.007, 12950]] * u.AU),\n",
    "        # radiant flux\n",
    "        \"luminosity\": (\n",
    "            [\"time\", \"body\"],\n",
    "            Quantity([[1, 0.001]], unit=u.Lsun).to(u.dex(u.Lsun)),\n",
    "        ),\n",
    "        \"metallicity\": ([\"time\", \"body\"], Dex([[0.0, np.nan]], unit=u.dex(u.Zsun))),\n",
    "        # relative to Vega\n",
    "        \"apparent_magnitude\": ([\"time\", \"body\"], Magnitude([[-26.74, 10.67]])),\n",
    "        \"absolute_magnitude\": ([\"time\", \"body\"], Magnitude([[4.83, 15.60]])),\n",
    "    },\n",
    ")\n",
    "ds"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "13",
   "metadata": {},
   "outputs": [],
   "source": [
    "ds.astropy.sel(body=[0]).astropy.to(\n",
    "    {\n",
    "        \"mass\": u.kg,\n",
    "        \"distance\": u.km,\n",
    "        \"illuminance\": u.cd * u.sr / u.m**2,\n",
    "        \"luminosity\": u.kW,\n",
    "        \"metallicity\": u.kg,\n",
    "    },\n",
    "    equivalencies=asplund_solar_relative_mass(),\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "14",
   "metadata": {},
   "outputs": [],
   "source": []
  }
 ],
 "metadata": {
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}