diff --git a/scopesim/tests/test_utils_functions.py b/scopesim/tests/test_utils_functions.py index 7fd38c24..c7ebd929 100644 --- a/scopesim/tests/test_utils_functions.py +++ b/scopesim/tests/test_utils_functions.py @@ -275,3 +275,42 @@ def test_setting_instpkgspath(): def test_unit_includes_per_physical_type(): unit = u.Unit("photlam") / u.arcsec**2 assert utils.unit_includes_per_physical_type(unit, "solid angle") + + +class TestZerosFromHeader: + def test_1d(self): + hdr = fits.Header({ + "NAXIS": 1, + "NAXIS1": 42, + }) + data = utils.zeros_from_header(hdr) + assert data.shape == (42,) + + def test_2d(self): + hdr = fits.Header({ + "NAXIS": 2, + "NAXIS1": 3, + "NAXIS2": 5, + }) + data = utils.zeros_from_header(hdr) + assert data.shape == (5, 3) + + def test_3d(self): + hdr = fits.Header({ + "NAXIS": 3, + "NAXIS1": 3, + "NAXIS2": 5, + "NAXIS3": 7, + }) + data = utils.zeros_from_header(hdr) + assert data.shape == (7, 5, 3) # founding date of Rome + + def test_2d_from_3d(self): + hdr = fits.Header({ + "NAXIS": 3, + "NAXIS1": 3, + "NAXIS2": 5, + "NAXIS3": 7, + }) + data = utils.zeros_from_header(hdr, ndims=2) + assert data.shape == (5, 3) diff --git a/scopesim/utils.py b/scopesim/utils.py index 7421676f..0921a90e 100644 --- a/scopesim/utils.py +++ b/scopesim/utils.py @@ -340,6 +340,38 @@ def zenith_dist_to_airmass(zenith_dist): return 1. / np.cos(np.deg2rad(zenith_dist)) +def zeros_from_header( + header: fits.Header, + dtype: type | np.dtype = float, + ndims: int | None = None, +) -> np.ndarray: + """ + Create all-zero array of the shape given by NAXISn keywords in `header`. + + .. versionadded:: PLACEHOLDER_NEXT_RELEASE_VERSION + + Parameters + ---------- + header : fits.Header + FITS header containing relevant keywords. + dtype : type | np.dtype, optional + Valid type or numpy dtype for output array. The default is float. + ndims : int | None, optional + Override number of dimensions. If None (the default), the number of + dimensions is taken from the "NAXIS" header keyword. This argument is + useful for e.g. creating a 2D image array from a 3D header. + + Returns + ------- + arr : np.ndarray + All-zero array of desired shape and dtype. + + """ + ndims = ndims or header["NAXIS"] + shape = tuple(header[f"NAXIS{i+1}"] for i in reversed(range(ndims))) + return np.zeros(shape, dtype=dtype) + + def convert_table_comments_to_dict(tbl): comments_dict = {}