Reduce memory consuption in axis_world_coords (again) #798
Conversation
|
This PR needs tests which test at least with some DKIST or IRIS WCSes that you can get low-memory axis world coords for uncorrelated axes, and as a counterpoint that correlated axes still work properly. For example, I have tested that |
|
I do question why I bother to do any programming if its always wrong. This solution works: Command line: /home/nabil/.mamba/envs/iris-dev/bin/memray run -f -o output.bin examples/umbral_flashes.py The biggest allocations are from loading the data:
|
|
This is great news! |
|
@nabobalis I can't add a visp WCS without a |
I am the same |
|
Could we try Wills WCS from his issue? |
nabobalis
force-pushed
the
maybe_this_will_work
branch
from
deee62c to
72d8859
Compare
ndcube/conftest.py
Outdated
| @@ -519,6 +560,25 @@ def ndcube_3d_rotated(wcs_3d_ln_lt_t_rotated, simple_extra_coords_3d): | |||
| return cube | |||
|
|
|||
|
|
|||
|
|
|||
| @pytest.fixture | |||
| def ndcube_3d_coupled(wcs_3d_ln_lt_l_coupled): | |||
There was a problem hiding this comment.
Not sure if this is good.
| # WCS for a 3D data cube with two celestial axes and one wavelength axis. | ||
| # The latitudinal dimension is coupled to the third pixel dimension through | ||
| # a single off diagonal element in the PCij matrix | ||
| header = { |
There was a problem hiding this comment.
Taken from Wills original issue.
nabobalis
force-pushed
the
maybe_this_will_work
branch
2 times, most recently
from
ae8bac4 to
e15bc13
Compare
Cadair
force-pushed
the
maybe_this_will_work
branch
2 times, most recently
from
17ec1ec to
ff9db90
Compare
If the pixel axes are not correlated with a requested world axis then don't generate the coords in the mesh
Cadair
force-pushed
the
maybe_this_will_work
branch
from
ff9db90 to
c440645
Compare
| @@ -230,6 +230,13 @@ def test_axis_world_coords_wave_ec(ndcube_3d_l_ln_lt_ectime): | |||
| assert coords[0].shape == (5,) | |||
|
|
|||
|
|
|||
| @pytest.mark.limit_memory("12 MB") | |||
There was a problem hiding this comment.
I grafted this test into the commit history before @nabobalis original #780 and it fails (using 3GB of RAM).
This test doesn't actually fail before this PR though, as this change fixes the case where you want to extract a single world coordinate correlated with two pixel axes where those pixel axes are also correlated with other world axes, i.e. extracting time from this VISP WCS:
<dkist.dataset.dataset.Dataset object at 0x7f0300d66ba0>
This VISP Dataset ADDMM consists of 10444 frames.
Files are stored in /data/dkist/prod/pid_2_114/ADDMM
This Dataset has 4 pixel and 4 world dimensions.
The data are represented by a <class 'dask.array.core.Array'> object:
dask.array<reshape, shape=(4, 2611, 996, 2545), dtype=float64, chunksize=(1, 1, 996, 2545), chunktype=numpy.ndarray>
Array Dim Axis Name Data size Bounds
0 raster map repeat number 4 None
1 raster scan step number 2611 None
2 dispersion axis 996 None
3 spatial along slit 2545 None
World Dim Axis Name Physical Type Units
3 time time s
2 helioprojective latitude custom:pos.helioprojective.lat arcsec
1 wavelength em.wl nm
0 helioprojective longitude custom:pos.helioprojective.lon arcsec
Correlation between pixel and world axes:
| PIXEL DIMENSIONS
| spatial | dispersion | raster | raster map
| along slit | axis | scan step | repeat
WORLD DIMENSIONS | | | number | number
------------------------- | ---------- | ---------- | ---------- | ----------
helioprojective longitude | x | | x | x
wavelength | | x | |
helioprojective latitude | x | | x | x
time | | | x | x
|
I would be nice to get a test which tests the 2D usecase. |
|
How does one do that? |
This PR works, as an alternative to #780 by not generating mesh coords for unneeded pixel axes.
i.e. if the pixel axes are not correlated with a requested world axis then don't generate the coords in the mesh.