Add zeros_from_header() utils function

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)

scopesim/utils.py

@@ -340,6 +340,38 @@
     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 = {}