Fix speed problem with area_weights

[pp-mo]

May actually be a more general problem with Cube.coords(), and certainly relates to aux-factories.
See #2371 (comment)

[pp-mo]

Initial Investigation (unfinished)

Summary of work so far :

• I've had a go with "cProfile" and "kapture"
• I think area_weights may be triggering a problem with repeatedly realising + re-calculating data of aux-factory coords.
• it seems to be very slow just calling "cube.coords()", which is a bit worrying.

cProfile
It seems you can't use the command-line interface with a script that uses unittests.main.
So I chopped out the test class iris.tests.tests_analysis.TestAreaWeightGeneration and added code to make it work with a cProfile.Profiler + spit out a file to analyse with pstats.Stats.
It was apparent that several methods were taking much more cumulative time in dask than with the 'basis' commit.

kapture
I took the same above. To work properly with kapture I had to chop it back further to a single test, and loop over the test body to make it last a good few seconds.
That now basically just looks like :

@tests.skip_data
class TestAreaWeightGeneration(tests.IrisTest):
    def setUp(self):
        self.cube = iris.tests.stock.realistic_4d()

    def test_area_weights_std(self):
        # weights for stock 4d data
        for _n in range(30):
            weights = iris.analysis.cartography.area_weights(self.cube)
            self.assertEqual(weights.shape, self.cube.shape)

From testing this, it seems clear that there are two places in area_weights which call into cube.coords(), and this is the slow bit.
From my notes on the code path that suffers this :

... ABOVE the common Cube.coords() call shown in the 'kapture_get_coords.png",
there are two calls to Cube.coords() from within the
"iris.analysis.cartography.area_weights" code :
    (1) analysis/cartography.py::381
            "lat, lon = _get_lat_lon_coords(cube)"
    (2) analysis/cartography.py::365
            "cs = cube.coord_system("CoordSystem")"

Both of these paths call into Cube.coords(), below which the two columns of
dark green blocks are both showing the same call path,
which is...
   iris.cube::1399 
        in Cube.coords :
        calling  "extract_coord()"
   iris.cube::1390
        in Cube.coords[extract_coord]
        calling  "coord_or_factory.make_coord()"
   iris.aux_factory::451
        in HybridHeightFactory.make_coord 
        calling "calc_bounds()"
   iris.aux_factory::450
        in HybridHeightFactory.make_coord[calc_bounds]
        calling "self._derive()"
   iris.aux_factory::400
        in HybridHeightFactory._derive
        "return delta + sigma * orography"

The implication is that  ***it is doing all this many times***

At present it looks like kapture cannot "see inside" this calculation, and it isn't clear ((to me!)) whether that is occurring in numpy or via _data_manager / dask ??

[DPeterK]

Hi @pp-mo, this analysis seems to make sense: as you have recorded above, cube.coords() calls any aux factory's make_coord method. On master, this would "realise" the LazyArray object that had been constructed to defer calculation of the aux factory. On dask this will again run the calcuation, but return a dask array object so long as any of the input coords are lazy. It's possible that running this is repeatedly reconstructing dask arrays, bringing with it that overhead. It might be quite interesting to look into whether the inputs and outputs are lazy when calling cube.coords() and seeing whether any of them change state as a result of the call.

[pp-mo]

The fix is here : #2604

[pp-mo]

NOTES:

Found that there were a number of problems with the existing code...

Most noticeably, it was no longer deferring calculation of derived coordinate values, so when AuxFactory.make_coord was called it was immediately doing calculations.
This call occurs whenever cube.coords() is invoked, which is potentially quite serious.

This problem does not show very prominently in the testing, as it only shows up with larger coords and it only occurs when the source coords have real data, not when they are lazy.

NOTE:
The actual test doesn't show this very clearly because the data is quite small.
However, messing with iris.tests.stock.realistic_4d() to produce data with an XY region of [1000, 1000] rather than [100, 100] makes this very clear:
Simply fetching this altitude coord was then taking approx 3.0 seconds instead of 0.003
- and the above fix reduces this to about 0.03
That extra *10 above (but only 0.03secs) is the cost of making dask arrays, instead of LazyArray, as creating Dask objects is quite costly.

[pp-mo]

Results comparison on one critical test (as highlighted by test timings analysis):
Test used:

$ IRIS_TEST_TIMINGS=1 python lib/iris/tests/test_analysis.py TestAreaWeightGeneration.test_area_weights_std

Timings compared:

code commit	time
With pre-dask master @2b9c2aa65	0.0056
With existing post-dask-merge @6bd26b779	0.1114
With proposed new code from #2604 @718fad1	0.0367

Note that the new code is still several times slower -- this seems to be due to the overhead of creating dask arrays.

More clearly:
Increasing the XY dims of the test cube from [100, 100] to [1000, 1000] gives the following ...

code commit	time
With pre-dask master @2b9c2aa65	0.0170
With existing post-dask-merge @6bd26b779	9.7063
With proposed new code from #2604 @718fad1	0.0932

[pp-mo]

Further note:
This has highlighted a notable outstanding inefficiency:
Whenever you call cube.coords(), all derived coords are made afresh by the factories. These also make new lazy coord data arrays. All this is costly in time.
It would be much more sensible for either the cube or factory to cache the derived coord.
See : #2605

[bjlittle]

@pp-mo Is this still an open issue?

[pp-mo]

@bjlittle pp-mo Is this still an open issue?

No, fixed by #2604