{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "# SkyCoord, Frame and Angles as Dataset\n",
    "\n",
    "SkyCoord and Frame fully support conversion to datasets preserving:\n",
    "* Representation\n",
    "* Differential\n",
    "* Frame\n",
    "* Angle (Longitude and Latitude)\n",
    "\n",
    "Conversion to dataset supports optionally providing axis (xarray coordinate) labels."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import astropy.units as u\n",
    "import numpy as np\n",
    "import xarray as xr\n",
    "from astropy.coordinates import ICRS, SkyCoord\n",
    "from astropy.time import Time\n",
    "\n",
    "from astropy_xarray.coordinates.sky_coord import (\n",
    "    skycoord_to_dataset,\n",
    ")\n",
    "\n",
    "sky_direction = SkyCoord(\n",
    "    ra=[[2, 6, 7, 4]] * u.deg,\n",
    "    dec=[[4, 7, 4, 3]] * u.deg,\n",
    "    pm_ra_cosdec=[[1, 1, 1, 1]] * u.mas / u.yr,\n",
    "    pm_dec=[[1, 1, 1, 1]] * u.mas / u.yr,\n",
    "    frame=ICRS(\n",
    "        representation_type=\"unitspherical\",\n",
    "        differential_type=\"unitsphericalcoslat\",\n",
    "    ),\n",
    ")\n",
    "display(sky_direction)\n",
    "skycoord_to_dataset(\n",
    "    sky_direction,\n",
    "    coords={\n",
    "        \"timestamp\": (\"time\", Time([1.7e9], format=\"unix\")),\n",
    "        \"field_label\": (\"field\", [\"a\", \"b\", \"c\", \"d\"]),\n",
    "    },\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "SkyCoordinate representation and differential types for both display and data are detected. In other words, radial_velocity when **implicit** is not recorded as a data_var.\n",
    "\n",
    "* *display* differential_type is `'sphericalcoslat'`, but\n",
    "* *data* differential_type is `'unitsphericalcoslat'`"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sky_position = SkyCoord(\n",
    "    ra=[2, 6, 7, 4] * u.deg,\n",
    "    dec=[4, 7, 4, 3] * u.deg,\n",
    "    distance=[4, 3, 6, 4] * u.parsec,\n",
    "    pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    frame=ICRS(\n",
    "        representation_type=\"spherical\",\n",
    "        differential_type=\"unitsphericalcoslat\",\n",
    "    ),\n",
    ")\n",
    "display(sky_position)\n",
    "display(\n",
    "    skycoord_to_dataset(\n",
    "        sky_position, coords={\"field_label\": (\"field\", [\"a\", \"b\", \"c\", \"d\"])}\n",
    "    )\n",
    ")"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Convert Dataset to SkyCoord\n",
    "\n",
    "SkyCoord datasets use a frame attribute for converting to a SkyCoordinate."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from astropy_xarray.coordinates import dataset_to_skycoord, load_frame\n",
    "\n",
    "ds = skycoord_to_dataset(\n",
    "    sky_position, coords={\"field_label\": (\"field\", [\"a\", \"b\", \"c\", \"d\"])}\n",
    ")\n",
    "\n",
    "# frame specific info stored as dataset attribute\n",
    "assert load_frame(ds.attrs[\"frame\"]) == sky_position.replicate_without_data()\n",
    "\n",
    "# dataset\n",
    "result = dataset_to_skycoord(ds)\n",
    "display(result)\n",
    "np.testing.assert_array_equal(result.ra, sky_position.ra)\n",
    "np.testing.assert_array_equal(result.dec, sky_position.dec)\n",
    "np.testing.assert_array_equal(result.distance, sky_position.distance)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# accessor available\n",
    "ds.astropy.to_skycoord()"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## Serialization and Deserialization\n",
    "\n",
    "Using dequantifing, all type metadata can be converted entirely to json-like attributes."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# serializable layout\n",
    "dds = skycoord_to_dataset(\n",
    "    sky_direction,\n",
    "    coords={\n",
    "        \"timestamp\": (\"time\", Time([1.7e9], format=\"unix\")),\n",
    "        \"field_label\": (\"field\", [0, 1, 2, 3]),\n",
    "    },\n",
    ").astropy.dequantify()\n",
    "dds"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "This allows for compatibility with simple serialization implementations (no custom type handling required)."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "import msgpack_numpy\n",
    "\n",
    "\n",
    "def to_msgpack_numpy(ds: xr.Dataset):\n",
    "    return msgpack_numpy.packb(ds.to_dict(data=\"array\"))\n",
    "\n",
    "\n",
    "def from_msgpack_numpy(buf: memoryview):\n",
    "    return xr.Dataset.from_dict(msgpack_numpy.unpackb(buf))\n",
    "\n",
    "\n",
    "sc = SkyCoord(ra=np.random.random(1000) * u.rad, dec=np.random.random(1000) * u.rad)\n",
    "ds = skycoord_to_dataset(sc)\n",
    "display(ds)\n",
    "\n",
    "raw = to_msgpack_numpy(ds.astropy.dequantify())\n",
    "display(raw)\n",
    "\n",
    "result = from_msgpack_numpy(raw).astropy.quantify()\n",
    "display(result)"
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "## SkyCoordinate Quirks"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# SkyCoord frame doesn't strongly infer directional reprentation and differential types\n",
    "import astropy.units as u\n",
    "from astropy.coordinates import SkyCoord\n",
    "\n",
    "sky_direction = SkyCoord(\n",
    "    ra=[2, 6, 7, 4] * u.deg,\n",
    "    dec=[4, 7, 4, 3] * u.deg,\n",
    "    pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    frame=\"icrs\",\n",
    ")\n",
    "assert (\n",
    "    sky_direction.representation_type.name == \"spherical\"\n",
    ")  #  may have expected \"unitspherical\"\n",
    "assert (\n",
    "    sky_direction.differential_type.name == \"sphericalcoslat\"\n",
    ")  # may have expected \"unitsphericalcoslat\"\n",
    "\n",
    "display(sky_direction)\n",
    "\n",
    "# additional members are therefore implicitly available\n",
    "display(sky_direction.distance)\n",
    "display(sky_direction.radial_velocity)\n",
    "\n",
    "display(sky_direction)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "sky_position = SkyCoord(\n",
    "    ra=[2, 6, 7, 4] * u.deg,\n",
    "    dec=[4, 7, 4, 3] * u.deg,\n",
    "    distance=[4, 3, 6, 4] * u.parsec,\n",
    "    pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    frame=\"icrs\",\n",
    ")\n",
    "display(sky_position)\n",
    "# using positional representation with direction differentials causes accessor of\n",
    "# `radial_velocity` to mutate data of the skycoordinate\n",
    "display(sky_position.radial_velocity)\n",
    "# skycoord_to_dataset will then subsequently include these radial velocity values when this occurs\n",
    "display(sky_position)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# `represent_as` can also update cached data\n",
    "sky_direction = SkyCoord(\n",
    "    ra=[2, 6, 7, 4] * u.deg,\n",
    "    dec=[4, 7, 4, 3] * u.deg,\n",
    "    distance=[1, 1, 1, 1] * u.dimensionless_unscaled,\n",
    "    pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    frame=\"icrs\",\n",
    ")\n",
    "assert sky_direction.representation_type.name == \"spherical\"\n",
    "assert sky_direction.differential_type.name == \"sphericalcoslat\"\n",
    "sky_direction_old_str = str(sky_direction)\n",
    "\n",
    "sky_direction.represent_as(\"spherical\", \"sphericalcoslat\")  # cache updated\n",
    "assert sky_direction.representation_type.name == \"spherical\"\n",
    "assert sky_direction.differential_type.name == \"sphericalcoslat\"\n",
    "assert str(sky_direction) != sky_direction_old_str  # string method changed"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# To avoid mutable cache quirks, use explicit repr and diff types for unit sphericals\n",
    "sky_position = SkyCoord(\n",
    "    ra=[2, 6, 7, 4] * u.deg,\n",
    "    dec=[4, 7, 4, 3] * u.deg,\n",
    "    distance=[1, 1, 1, 1] * u.dimensionless_unscaled,\n",
    "    pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    frame=ICRS(\n",
    "        representation_type=\"spherical\",\n",
    "        differential_type=\"unitsphericalcoslat\",\n",
    "    ),\n",
    ")\n",
    "sky_position_old_str = str(sky_position)\n",
    "sky_position.represent_as(\"spherical\", \"sphericalcoslat\")\n",
    "assert str(sky_position) == sky_position_old_str"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# astropy skycoord frame type casting between directional and positional\n",
    "# representation may alias data\n",
    "sky_direction = SkyCoord(\n",
    "    ra=[2, 6, 7, 4] * u.deg,\n",
    "    dec=[4, 7, 4, 3] * u.deg,\n",
    "    pm_ra_cosdec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    pm_dec=[1, 1, 1, 1] * u.mas / u.yr,\n",
    "    frame=\"icrs\",\n",
    ")\n",
    "print(\"SkyCoord:\", sky_direction.ra.data[1])\n",
    "\n",
    "# astropy-xarray dataset prevents unit conversion loss\n",
    "ds = skycoord_to_dataset(sky_direction, coords={\"field\": [\"a\", \"b\", \"c\", \"d\"]})\n",
    "print(\"Dataset:\", ds.ra.data[1].value)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "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": 2
}