Merge branch 'remove_unit_metadata' of github.com:SWIFTSIM/swiftsimio into remove_unit_metadata

Author: kyleaoman

.github/actions/start-hdfstream/action.yml

@@ -0,0 +1,40 @@
+name: 'Start hdfstream'
+description: 'Start the hdfstream service with a mounted directory'
+inputs:
+  data-dir:
+    description: 'Directory with the files to serve'
+    required: true
+    default: './data/'
+  virtual-prefix:
+    description: 'Virtual directory prefix to use'
+    required: true
+    default: 'Data'
+  image:
+    description: 'Server container image to use'
+    required: true
+    default: ghcr.io/jchelly/hdfstream-api:0.0.6
+runs:
+  using: "composite"
+  steps:
+    - name: Start hdfstream service with tests/data mounted
+      shell: bash
+      run: |
+        docker run -d \
+          --name hdfstream \
+          -p 8080:8080 \
+          -v ${{ inputs.data-dir }}:/opt/hdfstream/data:ro \
+          -e HDFSTREAM_PREFIX=${{ inputs.virtual-prefix }} \
+          ${{ inputs.image }}
+    - name: Wait until hdfstream service is responding
+      shell: bash
+      run: |
+        for i in {1..30}; do
+          status=$(docker inspect --format='{{.State.Health.Status}}' hdfstream)
+          if [ "$status" = "healthy" ]; then
+            echo "Service is ready"
+            exit 0
+          fi
+          sleep 2
+        done
+        echo "Service did not become healthy in time"
+        exit 1

.github/actions/stop-hdfstream/action.yml

@@ -0,0 +1,8 @@
+name: "Stop hdfstream"
+description: "Stop and remove the hdfstream container"
+runs:
+  using: "composite"
+  steps:
+    - name: Stop the hdfstream service
+      shell: bash
+      run: docker rm -f hdfstream

.github/workflows/lint_and_test.yml

@@ -36,7 +36,6 @@ jobs:
     strategy:
       matrix:
         python-version: ["3.10", "3.11", "3.12", "~3.13.0 <= 3.13.3 || ~3.13.5"]  # exclude 3.13.4
-
     steps:
     - name: Checkout swiftsimio
       uses: actions/checkout@v4
@@ -47,9 +46,21 @@ jobs:
     - name: Install dependencies
       run: |
         python -m pip install --upgrade pip
-        pip install -e .[dev]
+        pip install -e ".[dev,hdfstream]"
+    - name: Preload test data files
+      run: |
+        python -c "from tests.conftest import preload_test_data; preload_test_data()"
+    - name: Start hdfstream service
+      uses: ./.github/actions/start-hdfstream
+      with:
+        data-dir: ./test_data
+        virtual-prefix: test_data
+        image: ghcr.io/jchelly/hdfstream-api:0.0.6
     - name: Test with pytest
       run: |
-        pytest
+        pytest --hdfstream-server=http://localhost:8080/hdfstream
       env:
         NUMBA_BOUNDSCHECK: 1
+    - name: Stop hdfstream service
+      uses: ./.github/actions/stop-hdfstream
+      if: always()

docs/source/conf.py

@@ -112,4 +112,6 @@ def setup(app):  # numpydoc ignore=GL08
     swiftgalaxy=("https://swiftgalaxy.readthedocs.io/en/stable/", None),
     velociraptor=("https://velociraptor-python.readthedocs.io/en/latest/", None),
     astropy=("https://docs.astropy.org/en/stable/", None),
+    h5py=("https://docs.h5py.org/en/latest/", None),
+    hdfstream=("https://hdfstream-python.readthedocs.io/en/latest/", None),
 )

docs/source/loading_data/index.rst

@@ -1,3 +1,5 @@
+.. _loading-data:
+
 Loading Data
 ============
 
@@ -269,3 +271,95 @@ in SWIFT will be automatically read.
        "extra_test.hdf5",
    )
 
+
+Reading from an open file
+-------------------------
+
+:mod:`swiftsimio` normally opens and closes the HDF5 snapshot file for
+each operation. This is convenient for interactive use and avoids
+leaving files open for long periods of time, but sometimes it might be
+desirable to minimize the number of file open and close operations.
+
+It is possible to pass an open :obj:`h5py.File` object to
+:mod:`swiftsimio.load` and :mod:`swiftsimio.mask` in place of the
+filename. In this case swiftsimio will do all file access through the
+provided file object. This allows us to read multiple datasets while
+only opening and closing the file once. For example:
+
+.. code-block:: python
+
+   import h5py
+   import swiftsimio as sw
+
+   with h5py.File("cosmo_volume_example.hdf5","r") as snap_file:
+      data = sw.load(snap_file)
+      pos = data.dark_matter.coordinates
+      vel = data.dark_matter.velocities
+      ids = data.dark_matter.particle_ids
+
+This would open the snapshot file, read the dark matter particle
+positions, velocities and IDs, then close the file.
+
+
+Reading from a remote file
+--------------------------
+
+:mod:`swiftsimio` is able to read from snapshots hosted on a remote
+server using the `hdfstream
+<https://hdfstream-python.readthedocs.io/en/latest>`_ python
+module. This is useful if you're interested in accessing a small part
+of a larger snapshot: you can read a small region or a subset of
+particle properties without downloading the whole snapshot.
+
+To open a remote snapshot, you can pass a :obj:`hdfstream.RemoteFile`
+object to :mod:`swiftsimio.load` and :mod:`swiftsimio.mask` in place
+of the filename. For example, you can open one of the SWIFT example
+snapshots with:
+
+.. code-block:: python
+
+   import hdfstream
+   from swiftsimio import load
+
+   snap_file = hdfstream.open("cosma", "Tests/SWIFT/IOExamples/ssio_ci_04_2025/EagleSingle.hdf5")
+   data = load(snap_file)
+
+Here, ``data`` will be a :obj:`swiftsimio.reader.SWIFTDataset`. It
+functions in the same way as described in the :ref:`loading-data`
+section above, except that instead of reading data from a local HDF5
+file, it requests data from the server.
+
+Opening a snapshot like this only downloads a small amount of
+metadata. Accessing particle properties, such as coordinates, will
+trigger another download:
+
+.. code-block:: python
+
+   pos = data.dark_matter.coordinates
+
+This will download the dark matter particle coordinates and return an
+array with units and cosmological factors attached.
+
+To read part of a remote snapshot, we can use swiftsimio's
+:ref:`masking` feature as we would with a local snapshot, but passing
+the remote file to :mod:`swiftsimio.mask` :mod:`swiftsimio.load` in
+place of the filename.
+
+.. code-block:: python
+
+   import hdfstream
+   import swiftsimio as sw
+
+   snap_file = hdfstream.open("cosma", "Tests/SWIFT/IOExamples/ssio_ci_04_2025/EagleSingle.hdf5")
+
+   mask = sw.mask(snap_file)
+   # The full metadata object is available from within the mask
+   boxsize = mask.metadata.boxsize
+   # load_region is a 3x2 list [[left, right], [bottom, top], [front, back]]
+   load_region = [[0.0 * b, 0.5 * b] for b in boxsize]
+
+   # Constrain the mask
+   mask.constrain_spatial(load_region)
+
+   # Now load the snapshot with this mask
+   data = sw.load(snap_file, mask=mask)

docs/source/masking/index.rst

@@ -1,3 +1,5 @@
+.. _masking:
+
 Masking
 =======
 

pyproject.toml

@@ -67,8 +67,11 @@ docs = [
     "recommonmark",
     "sphinx_design",
 ]
+hdfstream = [
+    "hdfstream >= 0.0.23",
+]
 all = [
-    "swiftsimio[dev,docs]"
+    "swiftsimio[dev,docs,hdfstream]"
 ]
 
 [tool.ruff]

swiftsimio/__init__.py

@@ -6,12 +6,12 @@
 """
 
 from pathlib import Path
-import h5py
 
 from .reader import SWIFTDataset
 from .snapshot_writer import SWIFTSnapshotWriter as Writer
 from .masks import SWIFTMask
 from .statistics import SWIFTStatisticsFile
+from ._file_utils import open_path_or_handle
 from .__version__ import __version__
 from .__cite__ import __cite__
 
@@ -82,7 +82,7 @@ def validate_file(filename: str) -> bool:
         If the file is not a SWIFT data file.
     """
     try:
-        with h5py.File(filename, "r") as handle:
+        with open_path_or_handle(filename) as handle:
             if handle["Code"].attrs["Code"] != b"SWIFT":
                 raise KeyError
     except KeyError:
@@ -137,13 +137,11 @@ def mask(
     more expensive, ~bytes per particle instead of ~bytes per cell
     spatial_only=False version).
     """
-    if isinstance(filename, str):
-        filename = Path(filename)
-    with h5py.File(filename, "r") as handle:
-        units = SWIFTUnits(filename, handle=handle)
-        metadata = _metadata_discriminator(filename, units, handle=handle)
+    with open_path_or_handle(filename) as handle:
+        units = SWIFTUnits(handle.filename, handle=handle)
+        metadata = _metadata_discriminator(handle.filename, units, handle=handle)
         mask = SWIFTMask(
-            filename,
+            handle.filename,
             metadata=metadata,
             spatial_only=spatial_only,
             safe_padding=safe_padding,
@@ -169,11 +167,8 @@ def load(filename: str | Path, mask: SWIFTMask | None = None) -> SWIFTDataset:
     SWIFTDataset
         Dataset object providing an interface to the data file.
     """
-    if isinstance(filename, str):
-        filename = Path(filename)
-
-    with h5py.File(filename, "r") as handle:
-        data = SWIFTDataset(filename, mask=mask, handle=handle)
+    with open_path_or_handle(filename) as handle:
+        data = SWIFTDataset(handle.filename, mask=mask, handle=handle)
 
     return data
 

swiftsimio/_file_utils.py

@@ -0,0 +1,102 @@
+"""
+Define functions for handling file and dataset objects.
+
+These are used to handle a few situations where we need to check exactly
+what type of file or dataset object we have.
+"""
+
+from ._handle_provider import HandleProvider
+
+import h5py
+from pathlib import Path
+from swiftsimio.optional_packages import HDFSTREAM_AVAILABLE, hdfstream
+
+
+def is_soft_link(obj: h5py.Group | h5py.Dataset | h5py.SoftLink) -> bool:
+    """
+    Return ``True`` if ``obj`` is a soft link.
+
+    Note that soft links are usually dereferenced automatically, so to check
+    if an object is a soft link a reference to the object must be obtained
+    with::
+
+      obj = group.get(key, getlink=True)
+
+    where ``group`` is the group containing the object and ``key`` is the name
+    of the object.
+
+    Parameters
+    ----------
+    obj : Group, Dataset or SoftLink
+        The object to check.
+
+    Returns
+    -------
+    bool
+        ``True`` if ``obj`` is a soft link.
+    """
+    if HDFSTREAM_AVAILABLE:
+        return isinstance(obj, (h5py.SoftLink, hdfstream.SoftLink))
+    else:
+        return isinstance(obj, h5py.SoftLink)
+
+
+def is_dataset(obj: h5py.Group | h5py.Dataset | h5py.SoftLink) -> bool:
+    """
+    Return ``True`` if ``obj`` is a dataset.
+
+    Parameters
+    ----------
+    obj : Group, Dataset or SoftLink
+        The object to check.
+
+    Returns
+    -------
+    bool
+        ``True`` if ``obj`` is a dataset.
+    """
+    if HDFSTREAM_AVAILABLE:
+        return isinstance(obj, (h5py.Dataset, hdfstream.RemoteDataset))
+    else:
+        return isinstance(obj, h5py.Dataset)
+
+
+def split_path_or_handle(obj: str | Path | h5py.File) -> tuple[Path, h5py.File]:
+    """
+    Given a filename or handle, return a ``(filename, handle)`` tuple.
+
+    Parameters
+    ----------
+    obj : str, Path or h5py.File
+        A path to a file or a file handle object.
+
+    Returns
+    -------
+    tuple[Path, h5py.File]
+        Tuple with the path and file handle.
+    """
+    if isinstance(obj, (str, Path)):
+        filename = Path(obj)
+        handle = None
+    else:
+        filename = Path(obj.filename)
+        handle = obj
+    return filename, handle
+
+
+def open_path_or_handle(obj: str | Path | h5py.File) -> h5py.File:
+    """
+    Context manager to open a file, given a path or handle.
+
+    Parameters
+    ----------
+    obj : str, Path or h5py.File
+        A path to a file or a file handle object.
+
+    Returns
+    -------
+    h5py.File
+        The file handle.
+    """
+    filename, handle = split_path_or_handle(obj)
+    return HandleProvider(filename, handle).open_file()

swiftsimio/_handle_provider.py

@@ -1,5 +1,7 @@
 """Provide a mixin class for managing file handles."""
 
+from contextlib import contextmanager
+from typing import ContextManager
 from pathlib import Path
 import h5py
 
@@ -58,3 +60,23 @@ def _close_handle_if_manager(self) -> None:
         """Close the file handle if this object is the manager of the handle."""
         if self.handle_manager:
             self._handle.close()
+
+    @contextmanager
+    def open_file(self) -> ContextManager[h5py.File]:
+        """
+        Return a context manager that can be used to read the file.
+
+        This will use the existing handle if it is open. If not, we
+        assume that we're reading local HDF5 files using h5py and
+        create a temporary handle.
+
+        Returns
+        -------
+        ContextManager
+            A context manager which can be used to read the file.
+        """
+        if self._handle:
+            yield self._handle
+        else:
+            with h5py.File(self.filename, "r") as handle:
+                yield handle