Repository: sunpy/sunpy

Issue:
`make_fitswcs_header` should have an option to set the unit
### Describe the feature

`make_fitswcs_header` provides kwargs for setting several non-WCS keywords such as wavelength, exposure time, instrument, etc. We should also provide the ability to set the unit by setting the `"BUNIT"` key. 

### Proposed solution

Add a `unit=...` kwarg that takes in either an `astropy.Unit` or a `str` that is then used to set the `"BUNIT"` key in the new header.

Discussion:

Patch:
diff --git a/changelog/6499.feature.1.rst b/changelog/6499.feature.1.rst
new file mode 100644
index 000000000..36de226af
--- /dev/null
+++ b/changelog/6499.feature.1.rst
@@ -0,0 +1,3 @@
+`~sunpy.map.make_fitswcs_header` now includes the keyword argument ``unit`` for setting the
+``BUNIT`` FITS keyword in the resulting header.
+This will take precedence over any unit information attached to ``data``.
diff --git a/changelog/6499.feature.2.rst b/changelog/6499.feature.2.rst
new file mode 100644
index 000000000..8857bb859
--- /dev/null
+++ b/changelog/6499.feature.2.rst
@@ -0,0 +1,2 @@
+If the ``data`` argument to `~sunpy.map.make_fitswcs_header` is an `~astropy.units.Quantity`,
+the associated unit will be used to set the ``BUNIT`` FITS keyword in the resulting header.
diff --git a/sunpy/map/header_helper.py b/sunpy/map/header_helper.py
index 364e8959c..74b997474 100644
--- a/sunpy/map/header_helper.py
+++ b/sunpy/map/header_helper.py
@@ -4,6 +4,7 @@ import astropy.units as u
 import astropy.wcs
 from astropy.coordinates import SkyCoord
 
+from sunpy import log
 from sunpy.coordinates import frames, sun
 from sunpy.util import MetaDict
 
@@ -37,7 +38,8 @@ def make_fitswcs_header(data, coordinate,
                         observatory=None,
                         wavelength: u.angstrom = None,
                         exposure: u.s = None,
-                        projection_code="TAN"):
+                        projection_code="TAN",
+                        unit=None):
     """
     Function to create a FITS-WCS header from a coordinate object
     (`~astropy.coordinates.SkyCoord`) that is required to
@@ -45,7 +47,7 @@ def make_fitswcs_header(data, coordinate,
 
     Parameters
     ----------
-    data : `~numpy.ndarray` or `tuple`
+    data : `~numpy.ndarray`, `~astropy.units.Quantity`, or `tuple`
         Array data of Map for which a header is required, or the shape of the
         data array (in numpy order, i.e. ``(y_size, x_size)``).
     coordinate : `~astropy.coordinates.SkyCoord` or `~astropy.coordinates.BaseCoordinateFrame`
@@ -81,6 +83,10 @@ def make_fitswcs_header(data, coordinate,
         Exposure time of the observation
     projection_code : `str`, optional
         The FITS standard projection code for the new header.
+    unit : `~astropy.units.Unit`, optional
+        Units of the array data of the Map. This will populate the the ``'bunit'`` meta keyword.
+        If ``data`` is a `~astropy.units.Quantity`, the unit specified here will take precedence
+        over the unit information attached to ``data``.
 
     Returns
     -------
@@ -113,9 +119,15 @@ def make_fitswcs_header(data, coordinate,
     else:
         shape = data
 
+    if hasattr(data, "unit"):
+        if unit is None:
+            unit = data.unit
+        else:
+            log.info("Overwriting data's current unit with specified unit.")
+
     meta_wcs = _get_wcs_meta(coordinate, projection_code)
 
-    meta_wcs = _set_instrument_meta(meta_wcs, instrument, telescope, observatory, wavelength, exposure)
+    meta_wcs = _set_instrument_meta(meta_wcs, instrument, telescope, observatory, wavelength, exposure, unit)
     meta_wcs = _set_transform_params(meta_wcs, coordinate, reference_pixel, scale, shape)
     meta_wcs = _set_rotation_params(meta_wcs, rotation_angle, rotation_matrix)
 
@@ -262,7 +274,7 @@ def get_observer_meta(observer, rsun: (u.Mm, None) = None):
     return coord_meta
 
 
-def _set_instrument_meta(meta_wcs, instrument, telescope, observatory, wavelength, exposure):
+def _set_instrument_meta(meta_wcs, instrument, telescope, observatory, wavelength, exposure, unit):
     """
     Function to correctly name keywords from keyword arguments
     """
@@ -277,6 +289,8 @@ def _set_instrument_meta(meta_wcs, instrument, telescope, observatory, wavelengt
         meta_wcs['waveunit'] = wavelength.unit.to_string("fits")
     if exposure is not None:
         meta_wcs['exptime'] = exposure.to_value(u.s)
+    if unit is not None:
+        meta_wcs['bunit'] = unit.to_string("fits")
 
     return meta_wcs
 


Test patch:
diff --git a/sunpy/map/tests/test_header_helper.py b/sunpy/map/tests/test_header_helper.py
index 4bc2dd989..b990c81d6 100644
--- a/sunpy/map/tests/test_header_helper.py
+++ b/sunpy/map/tests/test_header_helper.py
@@ -164,14 +164,41 @@ def test_hgs_header(hgs_header, hgs_coord):
 
 
 def test_instrument_keyword(map_data, hpc_coord):
-    header = make_fitswcs_header(map_data, hpc_coord, instrument='test name')
-    assert header['instrume'] == 'test name'
+    instrument_kwargs = {
+        'instrument': 'test name',
+        'observatory': 'test observatory',
+        'telescope': 'test telescope',
+        'wavelength': 171 * u.Angstrom,
+        'exposure': 2 * u.s,
+        'unit': u.Unit('ct s-1'),
+    }
+    header = make_fitswcs_header(map_data, hpc_coord, **instrument_kwargs)
+    assert header['instrume'] == instrument_kwargs['instrument']
+    assert header['obsrvtry'] == instrument_kwargs['observatory']
+    assert header['telescop'] == instrument_kwargs['telescope']
+    assert header['wavelnth'] == instrument_kwargs['wavelength'].to_value()
+    assert header['waveunit'] == instrument_kwargs['wavelength'].unit.to_string("fits")
+    assert header['exptime'] == instrument_kwargs['exposure'].to_value('s')
+    assert header['bunit'] == instrument_kwargs['unit'].to_string("fits")
 
     # Check returned MetaDict will make a `sunpy.map.Map`
     map_test = sunpy.map.Map(map_data, header)
     assert isinstance(map_test, sunpy.map.mapbase.GenericMap)
 
 
+def test_quantity_input(map_data, hpc_coord):
+    # Test that unit information is extracted correctly when data array is a quantity
+    map_unit = u.Unit('ct / s')
+    map_quantity = u.Quantity(map_data, map_unit)
+    header = make_fitswcs_header(map_quantity, hpc_coord)
+    assert header['bunit'] == map_unit.to_string('fits')
+    # In cases where unit is specified and data is a quantity, specified unit will take
+    # precedence
+    override_unit = u.Unit('erg cm-2 s-1')
+    header = make_fitswcs_header(map_quantity, hpc_coord, unit=override_unit)
+    assert header['bunit'] == override_unit.to_string('fits')
+
+
 def test_invalid_inputs(map_data, hcc_coord, hpc_coord_notime, hpc_coord):
     # Raise the HCC error
     with pytest.raises(ValueError):