A trial to switch back to numpy arrays for performance. We do not even really use any xarray features currently.
For the cascading reservior example, which also does some calculations, this results in approx. 3x speedup.
xarray
to_xarray() (renamed to prepare())xdata argsWhen comparing columns, keep in mind that pushing xarray data basically requires creating the underlying numpy data first. So the most frequent use case is comparing first vs. last column. However, even compared to pushing xarray data we get quite some speedup for larger data.
| benchmark | before (numpy) | before (xarray) | after |
|---|---|---|---|
| run 365, 64x32 | 178 ms | 29 ms | 32 ms |
| run 365, 1024x512 | 289 ms | 134 ms | 127 ms |
| run 365, 2048x1024 | 862 ms | 740 ms | 481 ms |
| run 365, 64x32, copy | 220 ms | 57 ms | 62 ms |
| run 365, 1024x512, copy | 600 ms | 426 ms | 270 ms |
| run 365, 2048x1024, copy | 2395 ms | 2234 ms | 915 ms |
| run 365, 64x32, units | - | 109 ms | 49 ms |
| run 365, 1024x512, units | - | 483 ms | 254 ms |
| run 365, 2048x1024, units | - | 2044 ms | 1138 ms |
| benchmark | before (numpy) | before (xarray) | after |
|---|---|---|---|
| push&pull 512x256 | 460 us | 74 us | 80 us |
| push&pull 1024x512 | 490 us | 74 us | 80 us |
| push&pull 2048x1024 | 475 us | 74 us | 80 us |
| push&pull 512x256, units | 880 us | 501 us | 157 us |
| push&pull 1024x512, units | 1370 us | 1000 us | 297 us |
| push&pull 2048x1024, units | 3940 us | 3570 us | 1870 us |
| benchmark | before (xarray) | after |
|---|---|---|
| get_magnitude | 6 us | 0.4 us |
| check | 13 us | 6 us |
| to_xarray | 18 us |
|
| to_units (no-op) | 8 us | 2.3 us |
| to_units 512x256 | 297 us | 62 us |
| to_units 2048x1024 | 2180 us | 1928 us (?) |
| get_units | 3.5 us | 1.1 us |
| is_quantified | 3.6 us | 0.2 us |
| strip_time | 135 us | 5.7 us |
to_xarray can be further optimized by removing prod from the data dim calculation (see !237 (merged))