Speed-up suggestion - only slice datasets once

[dfulu]

Currently we slice each of our datasets in time and in space using different calls to DataArray.sel() or DataArray.isel(). These .sel() methods are slow and I saw somewhere (though can't find it right now) that xarray recommends combining selections to only do one if possible.

I tried this locally, and messily and got a speed up in the Dataset object of 5-8%. This is a nice speed up, but introduces a little more complexity to our code. Locally I needed to adapt the select_time_slice[_nwp]() functions in order to return dict[str, slice] objects for use with isel() instead of .sel() as we do now. I needed to store these and make the select_spatial_slice() function return dict[str, slice] objects. I then combined these space time slices and applied .isel() once.

This adds some complexity to the code compared to our current way of doing it where we have separate space and time slicing functions which return dataarrays. But the way we do it now comes at a performance cost.

So is it worth to refactor to a single slice, or is the cleaner code worth it to stay using the method we have now?

[Sukh-P]

Thanks for raising this James and nice one getting it working locally and seeing those performance gains! I am in two minds here but I think I am leaning towards not currently implementing this, I think it would make the code harder to follow than what it currently is (although perhaps not loads) for not a huge performance gain and more so it feels like perhaps we need to benchmark where we are at again a bit more in terms of the full training loop performance (how much can we saturate the GPU whilst training with different worker/batch_size configurations, how long does it take to train PVNet now etc)? If you have that info already then that would help make the decision, otherwise this feels like this could be at risk of premature optimisation

[dfulu]

Thanks for your thoughts @Sukh-P, I think I agree. Its feels like it adds tech debt to the code for a modest improvement. It might be possible to do it without it being overly complex, after all we're constructing dicts of slices, but it would be a reasonably big change and might make it harder to test.

There might be an option where we create functions internal to select_time_slice[_nwp]() and select_spatial_slice() but leave those current functions intact. The internal functions would calculate the slices, and the top level functions would run the internal function (something like construct_spatial_islice()) and apply them with an isel().