Add POT3D capabilities

psipy/conftest.py

@@ -74,3 +74,8 @@ def pluto_model():
 
 
 fixture_union("model", [mas_model, pluto_model])
+
+
+@fixture(scope="module")
+def pot3d_directory() -> Path:
+    return sample_data.pot3d_sample_data()

psipy/data/sample_data.py

@@ -22,6 +22,14 @@ def _get_url(
     thermo: str = "poly",
     resolution: str = "medium",
 ) -> str:
+    """
+    Get the URL to a MAS dataset with a given
+    - simulation type ('corona' or 'helio')
+    - simulation variable (e.g. magnetic field, density)
+    - carrington rotation
+    - thermodynamic model ('poly' or 'thermo')
+    - simulation resolution
+    """
     if thermo == "poly":
         thermo = "p"
     elif thermo == "thermo":
@@ -63,25 +71,6 @@ def _get_url(
     registry=registry,
 )
 
-# Add some PLUTO data
-pluto_reg: Dict[str, None] = {}
-PLUTO_FILES = [
-    "grid.out",
-    "dbl.out",
-    "rho.0000.dbl",
-    "Bx1.0000.dbl",
-    "Bx2.0000.dbl",
-    "Bx3.0000.dbl",
-]
-for file in PLUTO_FILES:
-    pluto_reg[file] = None
-
-pluto_pooch = pooch.create(
-    path=cache_dir,
-    base_url="doi:10.6084/m9.figshare.19401089.v1/",
-    registry=pluto_reg,
-)
-
 
 def mas_sample_data(sim_type="helio"):
     """
@@ -159,6 +148,37 @@ def mas_high_res_thermo() -> Path:
     return high_res_dir
 
 
+def _create_empty_pooch(files: list[str], base_url: str, cache_subdir: str):
+    """
+    Create an empty Pooch registry containing a list of files
+    available at a given URL.
+    """
+    reg_dict: Dict[str, None] = {}
+    for file in files:
+        reg_dict[file] = None
+
+    pooch_reg = pooch.create(
+        path=cache_dir / cache_subdir,
+        base_url=base_url,
+        registry=reg_dict,
+    )
+    return pooch_reg
+
+
+# PLUTO data
+PLUTO_FILES = [
+    "grid.out",
+    "dbl.out",
+    "rho.0000.dbl",
+    "Bx1.0000.dbl",
+    "Bx2.0000.dbl",
+    "Bx3.0000.dbl",
+]
+pluto_pooch = _create_empty_pooch(
+    PLUTO_FILES, "doi:10.6084/m9.figshare.19401089.v1/", "pluto"
+)
+
+
 def pluto_sample_data() -> Path:
     """
     Get some sample PLUTO data.
@@ -172,3 +192,27 @@ def pluto_sample_data() -> Path:
         path = pluto_pooch.fetch(file, progressbar=True)
 
     return Path(path).parent
+
+
+# POT3D data
+POT3D_FILES = [f"{var}.hdf" for var in ["bp", "br", "bt"]]
+pot3d_pooch = _create_empty_pooch(
+    POT3D_FILES,
+    "https://www.predsci.com/~pete/research/dstansby/2022-06-13-rss25/",
+    "pot3d",
+)
+
+
+def pot3d_sample_data() -> Path:
+    """
+    Get some sample POT3D data.
+
+    Returns
+    -------
+    pathlib.Path
+        Download directory.
+    """
+    for file in POT3D_FILES:
+        path = pot3d_pooch.fetch(file, progressbar=True)
+
+    return Path(path).parent

psipy/io/pot3d.py

@@ -0,0 +1,61 @@
+"""
+Tools for reading POT3D output files.
+"""
+import glob
+from pathlib import Path
+
+import numpy as np
+import xarray as xr
+
+from .util import read_hdf4
+
+__all__ = ["read_pot3d", "get_pot3d_variables"]
+
+
+def read_pot3d(file_path: Path) -> xr.Dataset:
+    """
+    Read a single POT3D file.
+
+    Parameters
+    ----------
+    path :
+        Path to the file.
+    """
+    if not file_path.suffix == ".hdf":
+        raise ValueError("File path must point to a hdf file.")
+
+    data, coords = read_hdf4(file_path.resolve())
+
+    dims = ["phi", "theta", "r"]
+    # Convert from co-latitude to latitude
+    coords[1] = np.pi / 2 - np.array(coords[1])
+    # Get variable name
+    var = file_path.name.split(".")[0]
+    # Add time
+    data = xr.Dataset({var: xr.DataArray(data=data, coords=coords, dims=dims)})
+    return data
+
+
+def get_pot3d_variables(directory):
+    """
+    Return a list of variables present in a given directory.
+
+    Parameters
+    ----------
+    directory :
+        Path to the folder containing the POT3D data files.
+
+    Returns
+    -------
+    var_names : list
+        List of variable names present in the given directory.
+    """
+    files = glob.glob(str(Path(directory) / "*.hdf"))
+    # Get the variable name from the filename
+    var_names = [Path(f).stem for f in files]
+    # Only return unique names
+    var_names = list(set(var_names))
+    if not len(var_names):
+        raise FileNotFoundError(f"No variable files found in {directory}")
+    var_names.sort()
+    return var_names

psipy/io/tests/test_pot3d.py

@@ -0,0 +1,12 @@
+import xarray as xr
+
+from psipy.io.pot3d import get_pot3d_variables, read_pot3d
+
+
+def test_read_pot3d(pot3d_directory):
+    ds = read_pot3d(pot3d_directory / "br.hdf")
+    assert isinstance(ds, xr.Dataset)
+
+
+def test_get_var_names(pot3d_directory):
+    assert get_pot3d_variables(pot3d_directory) == ["bp", "br", "bt"]

psipy/model/__init__.py

@@ -4,4 +4,5 @@
 from .base import *
 from .mas import *
 from .pluto import *
+from .pot3d import *
 from .variable import *

psipy/model/base.py

@@ -115,7 +115,6 @@ def get_runit(self) -> u.Unit:
         Return the units for the radial coordinate.
         """
 
-    @abc.abstractmethod
     def cell_corner_b(self, t_idx: Optional[int] = None) -> xr.DataArray:
         """
         Get the magnetic field vector at the cell corners.

psipy/model/pot3d.py

@@ -0,0 +1,42 @@
+from psipy.io.pot3d import get_pot3d_variables, read_pot3d
+from psipy.model import ModelOutput
+
+__all__ = ["POT3DOutput"]
+
+
+class POT3DOutput(ModelOutput):
+    """
+    The results from a single run of POT3D.
+
+    This is a storage object that contains a number of `Variable` objects. It
+    is designed to be used like::
+
+        mas_output = POT3DOutput('directory')
+        br = mas_output['br']
+
+    Notes
+    -----
+    Variables are loaded on demand. To see the list of available variables
+    use `POT3DOutput.variables`, and to see the list of already loaded variables
+    use `POT3DOutput.loaded_variables`.
+    """
+
+    def get_unit(self, var):
+        # TODO: fill in
+        ...
+
+    def get_runit(self):
+        # TODO: fill in
+        ...
+
+    def get_variables(self):
+        return get_pot3d_variables(self.path)
+
+    def load_file(self, var):
+        return read_pot3d(self.path / (var + ".hdf"))
+
+    def __repr__(self):
+        return f'psipy.model.pot3d.POT3DOutput("{self.path}")'
+
+    def __str__(self):
+        return f"POT3D output in directory {self.path}\n" + super().__str__()

psipy/model/tests/test_pot3d.py

@@ -0,0 +1,28 @@
+from psipy.model import ModelOutput, POT3DOutput
+
+
+def test_pot3d_model(pot3d_directory):
+    model = POT3DOutput(pot3d_directory)
+    # Check that loading a single file works
+    assert isinstance(model, ModelOutput)
+    assert "POT3D" in str(model)
+    assert model.variables == ["bp", "br", "bt"]
+    assert "bp" in str(model)
+
+    '''
+    rho = model["bp"]
+    assert isinstance(rho, base.Variable)
+    assert isinstance(rho.data, xr.DataArray)
+    assert rho.unit == u.dimensionless_unscaled
+    assert rho.n_timesteps == 1
+    assert (
+        str(rho)
+        == """
+Variable
+--------
+Name: rho
+Grid size: (128, 111, 141) (phi, theta, r)
+Timesteps: 1
+"""
+    )
+    '''