Example notebook to convert from actions to positions and velocities (galpy.orbit.Orbit) using optimization. Candidate actions are calculated assuming an action-angle object (galpy.actionAngle.actionAngle) and compared with target actions using a simple likelihood function. Positions and velocities are determined using scipy.optimize.minimize. M…

action-inverter-optimization.ipynb

{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [],
   "source": [
    "import numpy as np\n",
    "import time\n",
    "import warnings\n",
    "from galpy import orbit, potential, actionAngle as aA\n",
    "from scipy import optimize"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Function to convert actions to orbit using optimization\n",
    "def accs_to_orbit_optimize(accs,aA_obj,fun,xvo=[1.,0.1,0.1,0.1],\n",
    "    minimize_kwargs={},return_opt=False):\n",
    "    '''accs_to_orbit_optimize:\n",
    "\n",
    "    Convert actions to an orbit by optimizing a simple log-likelihood function \n",
    "    which calculates candidate actions and compares to target actions.\n",
    "\n",
    "    Args:\n",
    "        accs (array): Array of target actions (JR,Lz,Jz), shape (3)\n",
    "        aA_obj (galpy.actionAngle.actionAngle): actionAngle instance. __call__ \n",
    "            method must take (R,vR,vT,z,vz) and return (jr,jphi,jz)\n",
    "        fun (callable) - Objective function to minimize. Must take \n",
    "            (xv,accs,aA_obj) as arguments and return a scalar.\n",
    "        xvo (array): Initial guess for physical coordinates (R,vR,z,vz), \n",
    "            shape (4) [default [1.,0.1,0.1,0.1]]\n",
    "        minimize_kwargs (dict): Keyword arguments for scipy.optimize.minimize\n",
    "            [default {}]\n",
    "        return_opt (bool): If True, return the scipy.optimize.OptimizeResult\n",
    "            [default False]\n",
    "    \n",
    "    Returns:\n",
    "        orb (Orbit) - orbit corresponding to action diamond coordinates\n",
    "    '''\n",
    "    opt = optimize.minimize(fun,xvo,args=(accs,aA_obj),**minimize_kwargs)\n",
    "    if not opt.success:\n",
    "        warnings.warn('Optimization failed for target actions {0}, message: {1}'.format(accs,opt.message))\n",
    "    R,vR,z,vz = opt.x\n",
    "    orb = orbit.Orbit([R,vR,accs[1]/R,z,vz,0.])\n",
    "    if return_opt:\n",
    "        return orb,opt\n",
    "    else:\n",
    "        return orb\n",
    "\n",
    "# Define a function to act as the 'likelihood' for the action inverter\n",
    "def xv_accs_log_likelihood(xv,accs,aA_obj):\n",
    "    '''xv_accs_log_likelihood:\n",
    "\n",
    "    Simple log-likelihood function for actions. The phi action will be used \n",
    "    to set vT, and the phi coordinate will be held at zero.\n",
    "\n",
    "    Args:\n",
    "        xv (array): Array of coordinates (R,vR,z,vz), shape (4)\n",
    "        accs (array): Array of target actions (JR,Lz,Jz), shape (3)\n",
    "        aA_obj (galpy.actionAngle.actionAngle): actionAngle instance. __call__ \n",
    "            method must take (R,vR,vT,z,vz) and return (jr,jphi,jz)\n",
    "    \n",
    "    Returns:\n",
    "        log_likelihood (float) - log-likelihood of the actions for a given orbit\n",
    "    '''\n",
    "    # Unpack\n",
    "    R,vR,z,vz = xv\n",
    "    jr,jphi,jz = accs\n",
    "    # Calculate the actions\n",
    "    jr_calc,_,jz_calc = aA_obj(R,vR,jphi/R,z,vz)\n",
    "    # Calculate the log-likelihood\n",
    "    return (jr-jr_calc)**2 + (jz-jz_calc)**2"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 3,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Potential and actionAngle object, use Staeckel\n",
    "pot = potential.MWPotential2014\n",
    "delta = 0.4\n",
    "aAS = aA.actionAngleStaeckel(pot=pot, delta=delta, c=True)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 4,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "[0.5, 1.2, 0.0]\n",
      "[0.49999977] [1.2] [1.58472372e-05]\n"
     ]
    }
   ],
   "source": [
    "# Candidate actions (JR,Lz,Jz)\n",
    "accs = [0.5,1.2,0.]\n",
    "# Convert to orbit\n",
    "o,opt = accs_to_orbit_optimize(accs,aAS,xv_accs_log_likelihood,\n",
    "    minimize_kwargs={'method':'BFGS'},return_opt=True)\n",
    "# Recalculate actions\n",
    "jr,jphi,jz = aAS(o)\n",
    "# Print the actions\n",
    "print(accs)\n",
    "print(jr,jphi,jz)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "/var/folders/y2/0ghj8gl95972kwy339x5wc140000gn/T/ipykernel_65719/3503498068.py:27: UserWarning: Optimization failed for target actions [2.965390323496496, -1.636332394128209, 0.5811965616972211], message: Desired error not necessarily achieved due to precision loss.\n",
      "  warnings.warn('Optimization failed for target actions {0}, message: {1}'.format(accs,opt.message))\n",
      "\n",
      "/var/folders/y2/0ghj8gl95972kwy339x5wc140000gn/T/ipykernel_65719/3503498068.py:27: UserWarning: Optimization failed for target actions [2.8879816082637433, 1.2379670559218887, 0.3400638236633117], message: Desired error not necessarily achieved due to precision loss.\n",
      "  warnings.warn('Optimization failed for target actions {0}, message: {1}'.format(accs,opt.message))\n",
      "\n",
      "total time:  9.474062442779541\n",
      "median time:  0.09055972099304199\n",
      "max time:  0.29350781440734863\n"
     ]
    }
   ],
   "source": [
    "# Do a test for multiple sets of actions and benchmark the speed\n",
    "n = 100\n",
    "jr,jz = np.random.uniform(0.1,3.,(n,2)).T\n",
    "jphi = np.random.uniform(-2.,2.,n)\n",
    "minimize_kwargs = {'method':'BFGS'}\n",
    "\n",
    "ts = []\n",
    "for i in range(n):\n",
    "    # Time the conversion\n",
    "    t1 = time.time()\n",
    "    accs = [jr[i],jphi[i],jz[i]]\n",
    "    o = accs_to_orbit_optimize(accs,aAS,xv_accs_log_likelihood,\n",
    "        minimize_kwargs=minimize_kwargs)\n",
    "    t2 = time.time()\n",
    "    ts.append(t2-t1)\n",
    "    # Check that the actions are recovered\n",
    "    _jr,_jphi,_jz = aAS(o)\n",
    "    _accs = [_jr[0],_jphi[0],_jz[0]]\n",
    "    assert np.allclose(accs,_accs,atol=1e-3),\\\n",
    "        'Target actions were {0}, recovered actions were {1}'.format(accs,_accs)\n",
    "\n",
    "print('total time: ',np.sum(ts))\n",
    "print('median time: ',np.median(ts))\n",
    "print('max time: ',np.max(ts))\n"
   ]
  },
  {
   "attachments": {},
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "Other solvers will recover all target actions without failure, but at a slower rate. Even though BFGS (default) doesn't reach default accuracy for some sets of actions, the actions that are recovered are within reasonable tolerance."
   ]
  },
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": []
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "james",
   "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.8"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "39281e2d1d208f5d5f10b92e49c40383b657448f443f22dc78686b7e0f8179a6"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}