Sorting an Iris cube by coordinate.

sorted.ipynb

{
 "metadata": {
  "name": ""
 },
 "nbformat": 3,
 "nbformat_minor": 0,
 "worksheets": [
  {
   "cells": [
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Question\n",
      "\n",
      "Is it possible to sort a cube by a coordinate?\n",
      "\n",
      "## Answer\n",
      "\n",
      "There is nothing in Iris to do this currently, but it should be relatively easy to implement."
     ]
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "## Stratergy for implementing `sorted(cube, coord)`\n",
      "\n",
      "Rather than doing anything in-place as ``[].sort()`` does, let's follow the ``sorted(list)`` pattern to produce a new, sorted, cube. As a result we can quickly make use of Iris' indexing to return the sorted indices in order. The secret here is that numpy has a function, ``argsort`` which returns the indices of the sorted array, so all we need to do is index the cube on the dimension of the 1d coordinate to be sorted with the result of passing the coordinate points to ``argsort``."
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import numpy as np\n",
      "\n",
      "def sorted(cube, coord):\n",
      "    coord_to_sort = cube.coord(coord)\n",
      "    assert coord_to_sort.ndim == 1, 'One dim coords only please.'\n",
      "    dim, = cube.coord_dims(coord_to_sort)\n",
      "    index = [slice(None)] * cube.ndim\n",
      "    index[dim] = np.argsort(coord_to_sort.points)\n",
      "    return cube[tuple(index)]"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [],
     "prompt_number": 1
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "Let's test this with a simple 1D cube:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "import iris.coords\n",
      "import iris.cube\n",
      "\n",
      "cube = iris.cube.Cube([0, 1, 2, 3])\n",
      "cube.add_aux_coord(iris.coords.AuxCoord([2, 1, 0, 3], long_name='tbd'), 0)\n",
      "\n",
      "print cube\n",
      "print cube.data\n",
      "print cube.coord('tbd')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "unknown / (unknown)                 (-- : 4)\n",
        "     Auxiliary coordinates:\n",
        "          tbd                           x\n",
        "[0 1 2 3]\n",
        "AuxCoord(array([2, 1, 0, 3]), standard_name=None, units=Unit('1'), long_name='tbd')\n"
       ]
      }
     ],
     "prompt_number": 2
    },
    {
     "cell_type": "markdown",
     "metadata": {},
     "source": [
      "And when we call our ``sorted`` function on the coordinate to sort we see our data is now changed:"
     ]
    },
    {
     "cell_type": "code",
     "collapsed": false,
     "input": [
      "cube = sorted(cube, 'tbd')\n",
      "print cube\n",
      "print cube.data\n",
      "print cube.coord('tbd')"
     ],
     "language": "python",
     "metadata": {},
     "outputs": [
      {
       "output_type": "stream",
       "stream": "stdout",
       "text": [
        "unknown / (unknown)                 (-- : 4)\n",
        "     Auxiliary coordinates:\n",
        "          tbd                           x\n",
        "[2 1 0 3]\n",
        "AuxCoord(array([0, 1, 2, 3]), standard_name=None, units=Unit('1'), long_name='tbd')\n"
       ]
      }
     ],
     "prompt_number": 3
    }
   ],
   "metadata": {}
  }
 ]
}

FWIW to upcast the cube's new coordinate to a DimCoord (such that when saving NetCDF the dimension name comes from the coord's var_name) one would need to do:

import iris.coords

tbd = cube.coord('tbd')
tbd_dims = cube.coord_dims(tbd)

# Create a new coordinate which is a DimCoord.
dim_tbd = iris.coords.DimCoord.from_coord(tbd)

# Remove the AuxCoord and add the DimCoord.
cube.remove_coord(tbd)
cube.add_dim_coord(dim_tbd, tbd_dims)

I like finding answers to questions that I wrote many years ago. That still work today! 🎉

Just wish I'd implemented multiple coordinate sorting, not just one...

I guess it would work just as well to split the cube into single length items on the dimension to sort then to merge them together again...