orbits.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "from jax import config\n",
    "\n",
    "config.update(\"jax_enable_x64\", True)\n",
    "\n",
    "import jax.numpy as jnp\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "from jax_gw.detector.orbits import (\n",
    "    EARTH_TILT,\n",
    "    axial_tilt,\n",
    "    create_cartwheel_arm_lengths,\n",
    "    create_cartwheel_orbit,\n",
    "    create_circular_orbit_xy,\n",
    "    ecliptic_timeshift,\n",
    "    equatorial_timeshift,\n",
    "    flatten_pairs,\n",
    "    get_arm_lengths,\n",
    "    get_receiver_positions,\n",
    "    get_separations,\n",
    "    lat_lon_to_cartesian,\n",
    "    earthbound_ifo_pipeline,\n",
    ")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "AU_per_billion_meters = 149.597871\n",
    "L_target = 2.5\n",
    "R_target = 1.0\n",
    "ecc = L_target / (AU_per_billion_meters * 2 * jnp.sqrt(3) * R_target)\n",
    "N_LISA = 3"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "times = jnp.linspace(0, 1, 1000)\n",
    "orbits = create_cartwheel_orbit(ecc, R_target, N_LISA, times)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orbits.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111, projection=\"3d\")\n",
    "for i in range(N_LISA):\n",
    "    ax.plot(\n",
    "        orbits[i, 0, :],\n",
    "        orbits[i, 1, :],\n",
    "        orbits[i, 2, :],\n",
    "        label=f\"Spacecraft {i + 1}\",\n",
    "    )\n",
    "# equal aspect ratio for 3d plot\n",
    "# plt.gca().set_aspect(\"equal\", adjustable=\"box\")\n",
    "plt.xlim(-1.2,1.2)\n",
    "plt.ylim(-1.2,1.2)\n",
    "ax.set_zlim([-1.2, 1.2])\n",
    "ax.legend()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r = get_separations(orbits)\n",
    "print(r.shape)\n",
    "L = get_arm_lengths(r) \n",
    "print(L.shape)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L_12 = L[:, 0, 1]\n",
    "L_13 = L[:, 0, 2]\n",
    "L_23 = L[:, 1, 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "jnp.linalg.norm(orbits[0,:,0]- orbits[1,:,0]) * AU_per_billion_meters"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "d_analytic = create_cartwheel_arm_lengths(ecc, R_target, N_LISA, times)\n",
    "d_analytic.shape\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L12_analytic = d_analytic[:, 0, 1]\n",
    "L13_analytic = d_analytic[:, 0, 2]\n",
    "L23_analytic = d_analytic[:, 1, 2]"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# plot separations and compare to analytic\n",
    "fig = plt.figure()\n",
    "ax = fig.add_subplot(111)\n",
    "for dist in [L_12, L_13, L_23]:\n",
    "    ax.plot(times, dist*AU_per_billion_meters, linestyle='-', label='numerical', linewidth=3)\n",
    "for dist in [L12_analytic, L13_analytic, L23_analytic]:\n",
    "    ax.plot(times, dist*AU_per_billion_meters, linestyle='--', label='analytic')\n",
    "ax.set_xlabel(\"Time\")\n",
    "ax.set_ylabel(\"Separation\")\n",
    "ax.legend()\n",
    "plt.tight_layout()\n",
    "plt.show()\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "receiver_orbits = get_receiver_positions(orbits)\n",
    "print(receiver_orbits.shape)\n",
    "receiver_positions = flatten_pairs(receiver_orbits)\n",
    "receiver_positions.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "separations = get_separations(orbits)\n",
    "flatten_pairs(separations).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "flatten_pairs(L).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "%timeit flatten_pairs(L)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "FREQ_CENTER_ORBIT = 1  # in 1/year\n",
    "FREQ_ROTATION = 365.25  # in 1/year\n",
    "times = jnp.linspace(0, 1 / FREQ_ROTATION, 1000)\n",
    "r = 1.0 # in AU\n",
    "r_orbital = create_circular_orbit_xy(r, FREQ_CENTER_ORBIT, times)\n",
    "\n",
    "# calculate x, y, z coordinates of detector with respect to the guiding center\n",
    "# at time t=0\n",
    "# detector_lat = EARTH_Z_LAT\n",
    "# detector_lon = EARTH_Z_LON\n",
    "detector_lat, detector_lon = (\n",
    "    46.455140209119214 * jnp.pi / 180,\n",
    "    -119.40746331631823 * jnp.pi / 180,\n",
    ")\n",
    "r_detector_initial_equatorial = lat_lon_to_cartesian(detector_lat, detector_lon)\n",
    "print(r_detector_initial_equatorial)\n",
    "hour_angle = 2.0 * jnp.pi * FREQ_ROTATION * times\n",
    "r_detector = equatorial_timeshift(r_detector_initial_equatorial, hour_angle)\n",
    "r_detector.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "r_detector = axial_tilt(r_detector, EARTH_TILT)\n",
    "print(r_detector[:,0])\n",
    "r_detector.shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(times, r_detector[0, :], label='x')\n",
    "plt.plot(times, r_detector[1, :], label='y')\n",
    "plt.plot(times, r_detector[2, :], label='z')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "L_arm = 4 # in km\n",
    "r_earth_in_km = 6371.0\n",
    "\n",
    "psi_H = (90+36) * jnp.pi / 180\n",
    "psi_L = None\n",
    "psi = psi_H\n",
    "beta_arm = jnp.pi / 2.0\n",
    "\n",
    "# x_arm_ecliptic_initial = jnp.array([L_arm / r_earth_in_km, 0.0, 0.0])\n",
    "# y_arm_ecliptic_initial = jnp.array([0.0, L_arm / r_earth_in_km, 0.0])\n",
    "north_pole_equatorial = jnp.array([0.0, 0.0, 1.0])\n",
    "local_east = jnp.cross(north_pole_equatorial, r_detector_initial_equatorial)\n",
    "local_east = local_east / jnp.linalg.norm(local_east)\n",
    "# rotate the arms by psi with respect to r_detector_initial_equatorial\n",
    "# by applying the matrix form of Rodrigues' rotation formula\n",
    "K_matrix = jnp.array(\n",
    "    [\n",
    "        [0.0, -r_detector_initial_equatorial[2], r_detector_initial_equatorial[1]],\n",
    "        [r_detector_initial_equatorial[2], 0.0, -r_detector_initial_equatorial[0]],\n",
    "        [-r_detector_initial_equatorial[1], r_detector_initial_equatorial[0], 0.0],\n",
    "    ]\n",
    ")\n",
    "rotation_matrix_psi = jnp.eye(3) + jnp.sin(psi) * K_matrix + (1-jnp.cos(psi)) * K_matrix @ K_matrix\n",
    "rotation_matrix_beta = jnp.eye(3) + jnp.sin(beta_arm) * K_matrix + (1-jnp.cos(beta_arm)) * K_matrix @ K_matrix\n",
    "x_arm_direction = rotation_matrix_psi @ local_east\n",
    "y_arm_direction = rotation_matrix_beta @ x_arm_direction\n",
    "\n",
    "print(x_arm_direction)\n",
    "print(y_arm_direction)\n",
    "\n",
    "arm_length = L_arm / r_earth_in_km\n",
    "x_arm_local_equatorial_initial = arm_length * x_arm_direction\n",
    "y_arm_local_equatorial_initial = arm_length * y_arm_direction\n",
    "\n",
    "# convert to ecliptic coordinates\n",
    "x_arm_ecliptic_initial = axial_tilt(x_arm_local_equatorial_initial, +EARTH_TILT)\n",
    "print(x_arm_ecliptic_initial)\n",
    "y_arm_ecliptic_initial = axial_tilt(y_arm_local_equatorial_initial, +EARTH_TILT)\n",
    "print(y_arm_ecliptic_initial)\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# show the equatorial location of the detector in a lat-lon 2d plot\n",
    "plt.figure()\n",
    "ax = plt.gca()\n",
    "\n",
    "plt.scatter(\n",
    "    detector_lon * 180 / jnp.pi,\n",
    "    detector_lat * 180 / jnp.pi,\n",
    "    marker=\"o\",\n",
    "    color=\"red\",\n",
    "    label=\"detector\",\n",
    "    s=1,\n",
    ")\n",
    "# plot the arms\n",
    "# x_arm_lat_lon_dir_x = jnp.cos(psi)\n",
    "# x_arm_lat_lon_dir_y = jnp.sin(psi)\n",
    "\n",
    "# y_arm_lat_lon_dir_x = jnp.cos(psi+beta_arm)\n",
    "# y_arm_lat_lon_dir_y = jnp.sin(psi+beta_arm)\n",
    "\n",
    "# verify the direction of the arms in lat-lon coordinates\n",
    "# angle with respect to East\n",
    "abs_phi_x_arm = jnp.arccos(jnp.dot(x_arm_direction, local_east))\n",
    "abs_phi_y_arm = jnp.arccos(jnp.dot(y_arm_direction, local_east))\n",
    "phi_x_arm = jnp.sign(x_arm_direction[2]) * abs_phi_x_arm\n",
    "phi_y_arm = jnp.sign(y_arm_direction[2]) * abs_phi_y_arm\n",
    "print(phi_x_arm * 180 / jnp.pi)\n",
    "print(phi_y_arm * 180 / jnp.pi)\n",
    "x_arm_lat_lon_dir_x = jnp.cos(phi_x_arm)\n",
    "x_arm_lat_lon_dir_y = jnp.sin(phi_x_arm)\n",
    "\n",
    "y_arm_lat_lon_dir_x = jnp.cos(phi_y_arm)\n",
    "y_arm_lat_lon_dir_y = jnp.sin(phi_y_arm)\n",
    "\n",
    "plt.quiver(\n",
    "    detector_lon * 180 / jnp.pi,\n",
    "    detector_lat * 180 / jnp.pi,\n",
    "    x_arm_lat_lon_dir_x,\n",
    "    x_arm_lat_lon_dir_y,\n",
    "    color=\"red\",\n",
    "    width=0.005,\n",
    ")\n",
    "plt.quiver(\n",
    "    detector_lon * 180 / jnp.pi,\n",
    "    detector_lat * 180 / jnp.pi,\n",
    "    -x_arm_lat_lon_dir_x,\n",
    "    -x_arm_lat_lon_dir_y,\n",
    "    color=\"red\",\n",
    "    width=0.005,\n",
    ")\n",
    "\n",
    "plt.quiver(\n",
    "    detector_lon * 180 / jnp.pi,\n",
    "    detector_lat * 180 / jnp.pi,\n",
    "    y_arm_lat_lon_dir_x,\n",
    "    y_arm_lat_lon_dir_y,\n",
    "    pivot=\"tip\",\n",
    "    color=\"red\",\n",
    "    width=0.005,\n",
    ")\n",
    "\n",
    "plt.quiver(\n",
    "    detector_lon * 180 / jnp.pi,\n",
    "    detector_lat * 180 / jnp.pi,\n",
    "    -y_arm_lat_lon_dir_x,\n",
    "    -y_arm_lat_lon_dir_y,\n",
    "    pivot=\"tip\",\n",
    "    color=\"red\",\n",
    "    width=0.005,\n",
    ")\n",
    "\n",
    "plt.xlim(-180, 180)\n",
    "plt.ylim(-90, 90)\n",
    "\n",
    "plt.xlabel(\"Longitude (degrees)\")\n",
    "plt.ylabel(\"Latitude (degrees)\")\n",
    "plt.legend()\n",
    "plt.grid(\n",
    "    linestyle=\"--\",\n",
    ")\n",
    "plt.show()\n",
    "\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_arm = ecliptic_timeshift(x_arm_ecliptic_initial, hour_angle, EARTH_TILT)\n",
    "y_arm = ecliptic_timeshift(y_arm_ecliptic_initial, hour_angle, EARTH_TILT)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(times, x_arm[0, :], label='x_arm (x)')\n",
    "plt.plot(times, x_arm[1, :], label='x_arm (y)')\n",
    "plt.plot(times, x_arm[2, :], label='x_arm (z)')\n",
    "plt.plot(times, y_arm[0, :], label='y_arm (x)')\n",
    "plt.plot(times, y_arm[1, :], label='y_arm (y)')\n",
    "plt.plot(times, y_arm[2, :], label='y_arm (z)')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# add a rotation around this guiding center, assuming a solid body like the Earth\n",
    "earth_radius_per_km = 6371.0\n",
    "AU_per_earth_radius = (AU_per_billion_meters * 1e9) / (earth_radius_per_km * 1e3)\n",
    "print(AU_per_earth_radius)\n",
    "\n",
    "r_beam_splitter = r_orbital + r_detector / AU_per_earth_radius\n",
    "# convert x_arm, y_arm to double precision, divide by AU_per_earth_radius, and add to r_beam_splitter\n",
    "r_beam_splitter = jnp.array(r_beam_splitter, dtype=jnp.float64)\n",
    "x_arm = jnp.array(x_arm, dtype=jnp.float64) / AU_per_earth_radius\n",
    "y_arm = jnp.array(y_arm, dtype=jnp.float64) / AU_per_earth_radius\n",
    "\n",
    "x_arm = r_beam_splitter + x_arm\n",
    "y_arm = r_beam_splitter + y_arm\n",
    "\n",
    "# stack the beam splitter, the x Arm, and the y Arm into a 3x3xlen(times) array\n",
    "orbits = jnp.stack([r_beam_splitter, x_arm, y_arm], axis=0)\n",
    "orbits.dtype"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_arm"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "rot_orbits = orbits - r_beam_splitter\n",
    "# plt.plot(rot_orbits[0, 0, :], rot_orbits[0, 1, :], label='beam splitter')\n",
    "plt.plot(rot_orbits[1, 0, :], rot_orbits[1, 1, :], label='x arm')\n",
    "plt.plot(rot_orbits[2, 0, :], rot_orbits[2, 1, :], label='y arm')\n",
    "# plt.axis('equal')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "separations = get_separations(orbits)\n",
    "flatten_pairs(separations).shape"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "arm_lengths = flatten_pairs(get_arm_lengths(separations))"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "plt.plot(times, arm_lengths[0, :], label='x arm length', linewidth=3)\n",
    "plt.plot(times, arm_lengths[2, :], label='y arm length')\n",
    "plt.plot(times, arm_lengths[4, :], label='unphysical arm', linewidth=5, alpha=0.5)\n",
    "plt.plot(times, jnp.sqrt(2) * arm_lengths[0, :], label='x arm * sqrt2', linestyle='--')\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "orbits = earthbound_ifo_pipeline(\n",
    "    lat=detector_lat,\n",
    "    lon=detector_lon,\n",
    "    times=times,\n",
    "    r=r,\n",
    "    L_arm=L_arm,\n",
    "    psi=psi,\n",
    "    beta_arm=beta_arm,\n",
    ")\n",
    "orbits.shape\n",
    "    "
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "x_arm = (orbits[1]-orbits[0])[:,0]\n",
    "y_arm = (orbits[2]-orbits[0])[:,0]\n",
    "\n",
    "angle_beta = jnp.arccos(jnp.dot(x_arm, y_arm) / (jnp.linalg.norm(x_arm) * jnp.linalg.norm(y_arm)))\n",
    "print(angle_beta * 180 / jnp.pi)"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.16 ('jax-gwb')",
   "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.9.16"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "9015ef353a66a28b1cf10b4768999d3c6c5ace8f5466ddae6f166734f75ffc03"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}