HYDRAD Interface fixes (#101)

* truncate loop length to nearest Mm

* undo explicit loop length fix

* add functions for cube I/O with xarray

* remove instr classes from top level namespace

* ignore deprecation warning

examples/arcade-rtv.py

@@ -14,7 +14,7 @@
 
 import synthesizAR
 
-from synthesizAR.instruments import InstrumentSDOAIA
+from synthesizAR.instruments.sdo import InstrumentSDOAIA
 from synthesizAR.interfaces import RTVInterface
 from synthesizAR.models import semi_circular_arcade
 

examples/ebtel-nanoflares.py

@@ -17,7 +17,7 @@
 
 import synthesizAR
 
-from synthesizAR.instruments import InstrumentSDOAIA
+from synthesizAR.instruments.sdo import InstrumentSDOAIA
 from synthesizAR.interfaces.ebtel import EbtelInterface
 from synthesizAR.interfaces.ebtel.heating_models import PowerLawNanoflareTrain
 from synthesizAR.models import semi_circular_arcade

examples/loop-bundle-rtv.py

@@ -16,7 +16,7 @@
 
 import synthesizAR
 
-from synthesizAR.instruments import InstrumentSDOAIA
+from synthesizAR.instruments.sdo import InstrumentSDOAIA
 from synthesizAR.interfaces import RTVInterface
 from synthesizAR.models import semi_circular_bundle
 

examples/multi-instrument.py

@@ -23,7 +23,8 @@
 
 import synthesizAR
 
-from synthesizAR.instruments import InstrumentHinodeXRT, InstrumentSDOAIA
+from synthesizAR.instruments.hinode import InstrumentHinodeXRT
+from synthesizAR.instruments.sdo import InstrumentSDOAIA
 from synthesizAR.interfaces import MartensInterface
 from synthesizAR.models import semi_circular_arcade
 

pyproject.toml

@@ -21,6 +21,7 @@ dependencies = [
     "zarr",
     "asdf",
     "ndcube>=2.0",
+    "xarray",
 ]
 
 [project.urls]
@@ -86,6 +87,7 @@ filterwarnings = [
   "error",
   "ignore:The unit 'G' has been deprecated in the VOUnit standard.*:astropy.units.core.UnitsWarning",
   "ignore:'cgi' is deprecated and slated for removal in Python 3.13:DeprecationWarning",
+  "ignore:.*pkg_resources.*:DeprecationWarning"
 ]
 
 [tool.coverage.run]

synthesizAR/instruments/__init__.py

@@ -3,5 +3,3 @@
 """
 from .base import *
 from .physical import *
-from .sdo import *
-from .hinode import *

synthesizAR/instruments/util.py

@@ -1,13 +1,23 @@
 """
 Utility functions for building instrument classes
 """
+import astropy.units as u
+import astropy.wcs
 import copy
+import ndcube
+import numpy as np
+import xarray
 
-import astropy.units as u
 from ndcube.extra_coords.table_coord import QuantityTableCoordinate
 from ndcube.wcs.wrappers import CompoundLowLevelWCS
 
-__all__ = ['get_wave_keys', 'add_wave_keys_to_header', 'extend_celestial_wcs']
+__all__ = [
+    'get_wave_keys',
+    'add_wave_keys_to_header',
+    'extend_celestial_wcs',
+    'read_cube_from_dataset',
+    'write_cube_to_netcdf',
+]
 
 
 @u.quantity_input
@@ -45,3 +55,67 @@ def extend_celestial_wcs(celestial_wcs, array, name, physical_type):
         [celestial_wcs.pixel_n_dim] * temp_table.wcs.pixel_n_dim
     )
     return CompoundLowLevelWCS(celestial_wcs, temp_table.wcs, mapping=mapping)
+
+
+def read_cube_from_dataset(filename, axis_name, physical_type):
+    """
+    Read an `~ndcube.NDCube` from an `xarray` dataset.
+
+    This function reads a data cube from a netCDF file and rebuilds it
+    as an NDCube. The assumption is that the attributes on the stored
+    data array have the keys necessary to reconstitute a celestial FITS
+    WCS and that the axis denoted by `axis_name` is the additional axis
+    along which to extend that celestial WCS. This works only for 3D cubes
+    where two of the axes correspond to spatial, celestial axes.
+
+    Parameters
+    ----------
+    filename: `str`, path-like
+        File to read from, usually a netCDF file
+    axis_name: `str`
+        The addeded coordinate along which to extend the celestial WCS.
+    physical_type: `str`
+        The physical type of `axis_name` as denoted by the IVOA designation.
+    """
+    ds = xarray.load_dataset(filename)
+    meta = ds.attrs
+    data = u.Quantity(ds['data'].data, meta.pop('unit'))
+    mask = ds['mask'].data
+    celestial_wcs = astropy.wcs.WCS(header=meta)
+    axis_array = u.Quantity(ds[axis_name].data, ds[axis_name].attrs.get('unit'))
+    combined_wcs = extend_celestial_wcs(celestial_wcs, axis_array, axis_name, physical_type)
+    return ndcube.NDCube(data, wcs=combined_wcs, meta=meta, mask=mask)
+
+
+def write_cube_to_netcdf(filename, axis_name, cube):
+    """
+    Write a `~ndcube.NDCube` to a netCDF file.
+
+    This function writes an NDCube to a netCDF file by first expressing
+    it as an `xarray.DataArray`. This works only for 3D cubes where two of
+    the axes correspond to spatial, celestial axes.
+
+    Parameters
+    ----------
+    cube: `ndcube.NDCube`
+    axis_name: `str`
+    filename: `str` or path-like
+    """
+    # FIXME: This is not a general solution and is probably really brittle
+    celestial_wcs = cube.wcs.low_level_wcs._wcs[0]
+    wcs_keys = dict(celestial_wcs.to_header())
+    axis_array = cube.axis_world_coords(axis_name)[0]
+    axis_coord = xarray.Variable(axis_name,
+                                 axis_array.value,
+                                 attrs={'unit': axis_array.unit.to_string()})
+    if (mask:=cube.mask) is None:
+        mask = np.full(cube.data.shape, False)
+    cube_xa = xarray.Dataset(
+        {'data': ([axis_name, 'lat', 'lon'], cube.data),
+         'mask': ([axis_name, 'lat', 'lon'], mask)},
+        coords={
+            axis_name: axis_coord,
+        },
+        attrs={**wcs_keys, 'unit': cube.unit.to_string()}
+    )
+    cube_xa.to_netcdf(filename)

synthesizAR/interfaces/hydrad/hydrad.py

@@ -2,8 +2,8 @@
 Model interface for the HYDrodynamics and RADiation (HYDRAD) code
 """
 import astropy.units as u
+import copy
 import numpy as np
-import os
 import pathlib
 import pydrad.parse
 
@@ -69,18 +69,31 @@ def __init__(self,
                  interpolate_to_norm=False):
         self.output_dir = pathlib.Path(output_dir)
         self.base_config = base_config
-        self.hydrad_dir = hydrad_dir if hydrad_dir is None else pathlib.Path(hydrad_dir)
+        self.hydrad_dir = hydrad_dir
         self.heating_model = heating_model
         self.use_gravity = use_gravity
         self.use_magnetic_field = use_magnetic_field
         self.use_initial_conditions = use_initial_conditions
         self.maximum_chromosphere_ratio = maximum_chromosphere_ratio
         self.interpolate_to_norm = interpolate_to_norm
 
-    def configure_input(self, loop):
-        config = self.base_config.copy()
+    @property
+    def hydrad_dir(self):
+        return self._hydrad_dir
+
+    @hydrad_dir.setter
+    def hydrad_dir(self, value):
+        if value is not None:
+            value = pathlib.Path(value)
+        self._hydrad_dir = value
+
+    def _map_strand_to_config_dict(self, loop):
+        # NOTE: This is a separate function for ease of debugging
+        config = copy.deepcopy(self.base_config)
+        # NOTE: Avoids a bug in HYDRAD where seg faults can arise due to inconsistencies between max cell
+        # widths and minimum number of cells
         config['general']['loop_length'] = loop.length
-        config['initial_conditions']['heating_location'] = loop.length / 2.
+        config['initial_conditions']['heating_location'] = loop.length / 2
         if self.maximum_chromosphere_ratio:
             config['general']['footpoint_height'] = min(
                 config['general']['footpoint_height'], self.maximum_chromosphere_ratio * loop.length / 2)
@@ -89,10 +102,17 @@ def configure_input(self, loop):
         if self.use_magnetic_field:
             config['general']['poly_fit_magnetic_field'] = self.configure_magnetic_field_fit(loop)
         config = self.heating_model.calculate_event_properties(config, loop)
-        c = Configure(config)
-        c.setup_simulation(os.path.join(self.output_dir, loop.name),
+        return config
+
+    def configure_input(self, loop, **kwargs):
+        # Import here to avoid circular imports
+        from synthesizAR import log
+        log.debug(f'Configuring HYDRAD for {loop.name}')
+        config_dict = self._map_strand_to_config_dict(loop)
+        c = Configure(config_dict)
+        c.setup_simulation(self.output_dir / loop.name,
                            base_path=self.hydrad_dir,
-                           verbose=False)
+                           **kwargs)
 
     def load_results(self, loop):
         strand = pydrad.parse.Strand(self.output_dir / loop.name)

synthesizAR/visualize/aia.py

@@ -1,14 +1,14 @@
 """
 Plotting functions for easily and quickily visualizing synthesized AIA results
 """
-import os
-
-import numpy as np
+import astropy.units as u
 import h5py
-import matplotlib.pyplot as plt
-import matplotlib.colors
 import matplotlib.animation
-import astropy.units as u
+import matplotlib.colors
+import matplotlib.pyplot as plt
+import numpy as np
+import os
+
 from astropy.coordinates import SkyCoord
 from sunpy.map import Map
 
@@ -22,7 +22,7 @@ def plot_aia_channels(aia, time: u.s, root_dir, corners=None, figsize=None, norm
 
     Parameters
     ----------
-    aia : `synthesizAR.instruments.InstrumentSDOAIA`
+    aia : `synthesizAR.instruments.sdo.InstrumentSDOAIA`
     time : `astropy.Quantity`
     root_dir : `str`
     figsize : `tuple`, optional
@@ -62,7 +62,7 @@ def plot_aia_channels(aia, time: u.s, root_dir, corners=None, figsize=None, norm
         ax.text(0.1*tmp.dimensions.x.value, 0.9*tmp.dimensions.y.value,
                 r'${}$ $\mathrm{{\mathring{{A}}}}$'.format(channel['name']),
                 color='w', fontsize=fontsize)
-    fig.suptitle(r'$t={:.0f}$ {}'.format(time.value, time.unit.to_string()), fontsize=fontsize)
+    fig.suptitle(rf'$t={time.value:.0f}$ {time.unit.to_string()}', fontsize=fontsize)
     if kwargs.get('use_with_animation', False):
         return fig, ims