jones.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "id": "5076dee4-776f-427e-b790-ce440f171cbe",
   "metadata": {},
   "outputs": [],
   "source": [
    "!python -m pip install -e . >> /dev/null"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "b75de828-c739-476f-aad5-a6b01e53aef7",
   "metadata": {},
   "outputs": [],
   "source": [
    "import toolviper\n",
    "import xradio\n",
    "import pathlib\n",
    "import numba\n",
    "\n",
    "import numpy as np\n",
    "import pandas as pd\n",
    "import calviper as cv\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from xradio import measurement_set as ms"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6ed9a29d-367e-46ec-ab9d-ac55cfcfad6e",
   "metadata": {},
   "outputs": [],
   "source": [
    "if not pathlib.Path.cwd().joinpath(\"data/gaincaltest2.ps.zarr\").exists():\n",
    "    toolviper.utils.data.download(\"gaincal.test.zarr\", \"data\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "d9493256-8a80-4bee-ae5e-c4145627b157",
   "metadata": {},
   "outputs": [],
   "source": [
    "ps = ms.open_processing_set(\"data/gaincaltest2.ps.zarr\")\n",
    "\n",
    "sub_ps = ps.sel(intents=\"CALIBRATE_DELAY#ON_SOURCE\", scan_number=2)\n",
    "sub_ps.summary()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0613dde0-351b-4c16-a775-206029bb91ae",
   "metadata": {},
   "outputs": [],
   "source": [
    "dataset = sub_ps[\"gaincaltest2_0\"]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "79e116f8-09d8-4151-b30a-ada716b467f3",
   "metadata": {},
   "outputs": [],
   "source": [
    "V = dataset.VISIBILITY.mean(dim=\"time\").data.compute()\n",
    "V.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "de8216b0-d489-4a7c-93f4-c4914fef9e40",
   "metadata": {},
   "outputs": [],
   "source": [
    "V"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dc7448f-20ab-4eca-9e8c-11eba2e971ca",
   "metadata": {},
   "outputs": [],
   "source": [
    "s = dataset.VISIBILITY.shape\n",
    "s"
   ]
  },
  {
   "cell_type": "raw",
   "id": "73177abc-0b35-43d7-a0ed-bcafb9306557",
   "metadata": {},
   "source": [
    "def from_vis(dataset):\n",
    "    shape = dataset.VISIBILITY.shape\n",
    "\n",
    "    # There should be a gain value for each independent antenna. Here we choose antennai names but either\n",
    "    # would work fine.\n",
    "    n_parameters = np.unique(dataset.baseline_antenna1_name).shape[0]\n",
    "    \n",
    "    \n",
    "    identity = np.identity(n_parameters, dtype=np.complex64)\n",
    "\n",
    "    #return np.tile(identity, reps=[*shape, 1, 1])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "0ce2493e-0c7e-4e6f-9711-5b63df1973ef",
   "metadata": {},
   "outputs": [],
   "source": [
    "gain_jones = cv.jones.GainJones()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "e62cd1f3-36ce-44d0-8132-a31cdd4f4fa8",
   "metadata": {},
   "outputs": [],
   "source": [
    "G = gain_jones.from_visibility(dataset=dataset)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "6403894b-a634-45a8-9219-05aed71457f1",
   "metadata": {},
   "outputs": [],
   "source": [
    "G.matrix.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "7ccb777d-9509-4d19-a2c1-77d53c9564f9",
   "metadata": {},
   "outputs": [],
   "source": [
    "G.parameters.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "3ee0983f-5a19-4001-973c-c88216caa061",
   "metadata": {},
   "outputs": [],
   "source": [
    "g = G.matrix.mean(axis=0)\n",
    "g.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "bf6e3883-7923-4a5f-be89-ac0244ef8c17",
   "metadata": {},
   "outputs": [],
   "source": [
    "# This should all be done within the equivalent of VisEquation\n",
    "#v = V[:, 0, 0]\n",
    "v = V\n",
    "\n",
    "index_a, ant = cv.math.tools.encode(dataset.baseline_antenna1_name.to_numpy())\n",
    "index_b, _ = cv.math.tools.encode(dataset.baseline_antenna2_name.to_numpy())\n",
    "\n",
    "v_ = cv.math.tools.build_visibility_matrix(array=v, index_a=index_a, index_b=index_b)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "273a36a3-7f9d-4bfe-93de-3ac5888adab5",
   "metadata": {},
   "outputs": [],
   "source": [
    "v_.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "2ff5cc79-02c8-49a0-9016-c7113a657226",
   "metadata": {},
   "outputs": [],
   "source": [
    "solver = cv.math.solver.least_squares.LeastSquaresSolver()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "99704e1f-e800-40bf-a0e4-5d4715e011b3",
   "metadata": {},
   "outputs": [],
   "source": [
    "gain_solutions = solver.solve(\n",
    "    vis=v_,\n",
    "    iterations=50,\n",
    "    optimizer=cv.math.optimizer.MeanSquaredError(alpha=0.25),\n",
    "    stopping=1e-4\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "4dba3ecf-0693-4462-9956-3a4e9541ba77",
   "metadata": {},
   "outputs": [],
   "source": [
    "t = np.linspace(1, len(solver.losses), len(solver.losses))\n",
    "\n",
    "plt.scatter(solver.losses, t)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "1afafaab-eef7-4960-b6b9-31de7e12579f",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Gain solutions\n",
    "\n",
    "solver.parameter"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "a268274a-56b4-4fd3-b3a7-fb0704abe7c1",
   "metadata": {},
   "outputs": [],
   "source": [
    "solver.parameter.shape"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3 (ipykernel)",
   "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.10.16"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}