From 50e0a188bcbeaf29bfac9bf3bcd89b4f84a3ead1 Mon Sep 17 00:00:00 2001 From: Konstantin Malanchev Date: Sat, 18 Apr 2026 16:11:30 +0000 Subject: [PATCH] HEALPixFITSQuery: use memmap, defer byte-swap to query time --- dustmaps/csfd.py | 4 +- dustmaps/healpix_map.py | 26 ++++++----- dustmaps/tests/__init__.py | 13 ++++++ dustmaps/tests/test_csfd.py | 80 +++++++++++++++++++++++++++++++++ dustmaps/tests/test_lenz2017.py | 79 ++++++++++++++++++++++++++++++++ dustmaps/tests/test_planck.py | 7 +++ 6 files changed, 197 insertions(+), 12 deletions(-) create mode 100644 dustmaps/tests/test_csfd.py create mode 100644 dustmaps/tests/test_lenz2017.py diff --git a/dustmaps/csfd.py b/dustmaps/csfd.py index 9011ed3..0cb4f9f 100644 --- a/dustmaps/csfd.py +++ b/dustmaps/csfd.py @@ -55,9 +55,9 @@ def __init__(self, map_fname=None, mask_fname=None): try: with fits.open(map_fname) as hdulist: - ebv_data = hdulist['xtension'].data[:]['T'].flatten() + ebv_data = hdulist['xtension'].data['T'].ravel() with fits.open(mask_fname) as hdulist: - mask_data = hdulist['xtension'].data[:]['T'].flatten() + mask_data = hdulist['xtension'].data['T'].ravel() except IOError as error: print(dustexceptions.data_missing_message('csfd', 'CSFD (Chiang 2023)')) diff --git a/dustmaps/healpix_map.py b/dustmaps/healpix_map.py index 5eac6a4..44bacb6 100644 --- a/dustmaps/healpix_map.py +++ b/dustmaps/healpix_map.py @@ -109,6 +109,8 @@ def __init__(self, fname, coord_frame, hdu=0, field=None, scale (Optional[:obj:`float`]): Scale factor to be multiplied into the data. """ + self._out_dtype = np.dtype(dtype) + self._scale = scale close_file = False if isinstance(fname, six.string_types): @@ -126,17 +128,9 @@ def __init__(self, fname, coord_frame, hdu=0, field=None, '`BinTableHDU`.') if field is None: - pix_val = np.array(hdu.data[:].ravel().astype(dtype)) + pix_val = hdu.data[:].ravel() else: - pix_val = np.array(hdu.data[field][:].ravel().astype(dtype)) - - if scale is not None: - names = pix_val.dtype.names - if names is None: - pix_val *= scale - else: - for n in names: - pix_val[n] *= scale + pix_val = hdu.data[field] nest = hdu.header.get('ORDERING', 'NESTED').strip() == 'NESTED' @@ -144,3 +138,15 @@ def __init__(self, fname, coord_frame, hdu=0, field=None, hdulist.close() super(HEALPixFITSQuery, self).__init__(pix_val, nest, coord_frame) + + def query(self, coords, return_flags=False): + result = super(HEALPixFITSQuery, self).query(coords, + return_flags=return_flags) + if return_flags: + sel_pix, flags = result + else: + sel_pix = result + sel_pix = sel_pix.astype(self._out_dtype) + if self._scale is not None: + sel_pix *= self._scale + return (sel_pix, flags) if return_flags else sel_pix diff --git a/dustmaps/tests/__init__.py b/dustmaps/tests/__init__.py index c8ae193..ac76cd1 100644 --- a/dustmaps/tests/__init__.py +++ b/dustmaps/tests/__init__.py @@ -15,3 +15,16 @@ # You should have received copies of the GNU General Public License # and the BSD License along with this program. # + +from mmap import mmap +import numpy as np + + +def is_mmap(arr): + """Return True if arr is backed by a memory-mapped file (np.memmap or mmap).""" + obj = arr + while obj is not None: + if isinstance(obj, (np.memmap, mmap)): + return True + obj = getattr(obj, 'base', None) + return False diff --git a/dustmaps/tests/test_csfd.py b/dustmaps/tests/test_csfd.py new file mode 100644 index 0000000..2cda1f7 --- /dev/null +++ b/dustmaps/tests/test_csfd.py @@ -0,0 +1,80 @@ +#!/usr/bin/env python +# +# test_csfd.py +# Test query code for the Corrected SFD (CSFD) dust map of Chiang (2023). +# +# Copyright (C) 2026 Gregory M. Green +# +# dustmaps is free software: you can redistribute it and/or modify +# it under the terms of either: +# +# - The GNU General Public License as published by the Free Software Foundation, +# either version 2 of the License, or (at your option) any later version, or +# - The 2-Clause BSD License (also known as the Simplified BSD License). +# +# You should have received copies of the GNU General Public License +# and the BSD License along with this program. +# + +from __future__ import print_function, division + +import unittest +import time + +import numpy as np +import astropy.coordinates as coords +import astropy.units as units + +from .. import csfd +from . import is_mmap + + +class TestCSFD(unittest.TestCase): + + @classmethod + def setUpClass(self): + print('Loading CSFD dust map ...') + t0 = time.time() + self._csfd = csfd.CSFDQuery() + t1 = time.time() + print('Loaded CSFD test data in {:.5f} s.'.format(t1-t0)) + + def test_shape(self): + """ + Test that the output shapes are as expected with input coordinate arrays + of different shapes. + """ + for reps in range(5): + n_dim = np.random.randint(1, 4) + shape = np.random.randint(1, 7, size=(n_dim,)) + ra = np.random.uniform(-180., 180., size=shape) * units.deg + dec = np.random.uniform(-90., 90., size=shape) * units.deg + c = coords.SkyCoord(ra, dec, frame='icrs') + E = self._csfd(c) + np.testing.assert_equal(E.shape, shape) + + def test_frame(self): + """ + Test that the results are independent of the coordinate frame. + """ + frames = ('icrs', 'galactic', 'fk5', 'fk4', 'barycentrictrueecliptic') + shape = (100,) + ra = np.random.uniform(-180., 180., size=shape) * units.deg + dec = np.random.uniform(-90., 90., size=shape) * units.deg + c = coords.SkyCoord(ra, dec, frame='icrs') + E0 = self._csfd(c) + for fr in frames: + cc = c.transform_to(fr) + E = self._csfd(cc) + np.testing.assert_equal(E, E0) + + def test_mmap(self): + """ + Test that the underlying pixel data is memory-mapped (not copied into RAM). + """ + self.assertTrue(is_mmap(self._csfd._pix_val)) + self.assertTrue(is_mmap(self._csfd._flags)) + + +if __name__ == '__main__': + unittest.main() diff --git a/dustmaps/tests/test_lenz2017.py b/dustmaps/tests/test_lenz2017.py new file mode 100644 index 0000000..29d3bc9 --- /dev/null +++ b/dustmaps/tests/test_lenz2017.py @@ -0,0 +1,79 @@ +#!/usr/bin/env python +# +# test_lenz2017.py +# Test query code for the Lenz, Hensley & Doré (2017) dust map. +# +# Copyright (C) 2026 Gregory M. Green +# +# dustmaps is free software: you can redistribute it and/or modify +# it under the terms of either: +# +# - The GNU General Public License as published by the Free Software Foundation, +# either version 2 of the License, or (at your option) any later version, or +# - The 2-Clause BSD License (also known as the Simplified BSD License). +# +# You should have received copies of the GNU General Public License +# and the BSD License along with this program. +# + +from __future__ import print_function, division + +import unittest +import time + +import numpy as np +import astropy.coordinates as coords +import astropy.units as units + +from .. import lenz2017 +from . import is_mmap + + +class TestLenz2017(unittest.TestCase): + + @classmethod + def setUpClass(self): + print('Loading Lenz2017 dust map ...') + t0 = time.time() + self._lenz = lenz2017.Lenz2017Query() + t1 = time.time() + print('Loaded Lenz2017 test data in {:.5f} s.'.format(t1-t0)) + + def test_shape(self): + """ + Test that the output shapes are as expected with input coordinate arrays + of different shapes. + """ + for reps in range(5): + n_dim = np.random.randint(1, 4) + shape = np.random.randint(1, 7, size=(n_dim,)) + ra = np.random.uniform(-180., 180., size=shape) * units.deg + dec = np.random.uniform(-90., 90., size=shape) * units.deg + c = coords.SkyCoord(ra, dec, frame='icrs') + E = self._lenz(c) + np.testing.assert_equal(E.shape, shape) + + def test_frame(self): + """ + Test that the results are independent of the coordinate frame. + """ + frames = ('icrs', 'galactic', 'fk5', 'fk4', 'barycentrictrueecliptic') + shape = (100,) + ra = np.random.uniform(-180., 180., size=shape) * units.deg + dec = np.random.uniform(-90., 90., size=shape) * units.deg + c = coords.SkyCoord(ra, dec, frame='icrs') + E0 = self._lenz(c) + for fr in frames: + cc = c.transform_to(fr) + E = self._lenz(cc) + np.testing.assert_equal(E, E0) + + def test_mmap(self): + """ + Test that the underlying pixel data is memory-mapped (not copied into RAM). + """ + self.assertTrue(is_mmap(self._lenz._pix_val)) + + +if __name__ == '__main__': + unittest.main() diff --git a/dustmaps/tests/test_planck.py b/dustmaps/tests/test_planck.py index 656128e..716cd43 100644 --- a/dustmaps/tests/test_planck.py +++ b/dustmaps/tests/test_planck.py @@ -29,6 +29,7 @@ from .. import planck from ..std_paths import * +from . import is_mmap class TestPlanck(unittest.TestCase): @@ -68,6 +69,12 @@ def test_shape(self): np.testing.assert_equal(E.shape, shape) + def test_mmap(self): + """ + Test that the underlying pixel data is memory-mapped (not copied into RAM). + """ + self.assertTrue(is_mmap(self._planck._pix_val)) + def test_frame(self): """ Test that the results are independent of the coordinate frame.