sdf-xarray 0.9.0__tar.gz → 0.10.0__tar.gz

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: sdf-xarray
-Version: 0.9.0
+Version: 0.10.0
 Summary: Provides a backend for xarray to read SDF files as created by the EPOCH plasma PIC code.
 Author-Email: Peter Hill <peter.hill@york.ac.uk>, Joel Adams <joel.adams@york.ac.uk>, Shaun Doherty <shaun.doherty@york.ac.uk>, Chris Herdman <chris.herdman@york.ac.uk>, Liam Pattinson <liam.pattinson@york.ac.uk>, Sviatoslav Shekhanov <sviatoslav.shekhanov@york.ac.uk>
 License-Expression: BSD-3-Clause

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/docs/animation.md

@@ -12,7 +12,6 @@ mimic <inv:#xarray.DataArray.plot>.
 
 ```{code-cell} ipython3
 import sdf_xarray as sdfxr
-import xarray as xr
 import matplotlib.pyplot as plt
 from matplotlib.animation import FuncAnimation
 from IPython.display import HTML
@@ -139,31 +138,23 @@ EPOCH allows for simulations that have a moving simulation window
 `move_window` and change the x-axis limits accordingly.
 
 ```{warning}
-<project:#sdf_xarray.open_mfdataset> does not currently function with moving window data.
-You must use <inv:#xarray.open_mfdataset> and specify arguments in the following way.
+Importing some datasets with moving windows can cause vertical banding
+in the <inv:#xarray.Dataset>, which will affect the animation. Each output stores
+its own x grid, and the same position can differ by floating point rounding between
+outputs. Both values are then kept on the combined x axis, so each time step has
+nan values in some columns. Setting `join = "override"` in
+<inv:#xarray.open_mfdataset> does not fix this for moving window data, as it
+relabels every time step with the first output's x grid.
 ```
 
 ```{code-cell} ipython3
-ds = xr.open_mfdataset(
-   "tutorial_dataset_2d_moving_window/*.sdf",
-   preprocess = sdfxr.SDFPreprocess(),
-   combine = "nested",
-   join = "outer",
-   compat="no_conflicts",
-   concat_dim="time",
-   )
+ds = sdfxr.open_mfdataset("tutorial_dataset_2d_moving_window/*.sdf")
 
 da = ds["Derived_Number_Density_Beam_Electrons"]
 anim = da.epoch.animate(move_window = True, fps = 5)
 anim.show()
 ```
 
-```{warning}
-Importing some datasets with moving windows can cause vertical banding
-in the <inv:#xarray.Dataset>, which will affect the animation. The cause for
-this is unknown but can be circumvented by setting `join = "override"`.
-```
-
 ## Customisation
 
 The animation can be customised in much the same way as <inv:#xarray.DataArray.plot>,

sdf_xarray-0.10.0/docs/known_issues.md

@@ -0,0 +1,8 @@
+# Known Issues
+
+These are the known 'quirks' in `sdf_xarray`:
+
+- [Issue #112](https://github.com/epochpic/sdf-xarray/issues/112) Due to the 
+way [cylindrical EPOCH](https://github.com/epochpic/cylindrical_epoch) sets up 
+grids we can't currently load it's SDF files. If there is enough interest in this
+we will attempt to find a solution otherwise it shall remain low priority.

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/docs/loading_data.md

@@ -85,7 +85,11 @@ is by using the <project:#sdf_xarray.open_mfdataset>.
 ```{tip}
 If your simulation includes multiple `output` blocks that specify different variables
 for output at various time steps, variables not present at a specific step will default
-to a nan value. To remove these nan values we suggest using the <inv:#xarray.DataArray.dropna>
+to a nan value. This includes distribution functions and particle data that are only
+output at some time steps. Grids that change between outputs (such as a moving window,
+or a distribution function with an automatic range) are combined onto the union of
+every output's grid points, with nan values outside the part of the grid each output
+covers. To remove these nan values we suggest using the <inv:#xarray.DataArray.dropna>
 function or following our implmentation in [](#loading-sparse-data).
 ```
 
@@ -98,13 +102,20 @@ sdfxr.open_mfdataset("tutorial_dataset_1d/*.sdf")
 Alternatively, files can be loaded using <inv:#xarray.open_mfdataset>. However, when loading in
 all the files we have do some processing of the data so that we can correctly align it along
 the time dimension; This is done via the `preprocess` parameter utilising the
-<project:#sdf_xarray.SDFPreprocess> function.
+<project:#sdf_xarray.SDFPreprocess> function. Passing `combine="nested"` and
+`concat_dim="time"` makes sure the files are only joined along time. Otherwise
+<inv:#xarray> treats any coordinate that differs between files (such as a moving
+window's x grid) as another dimension to concatenate along, which can fail or request
+far more memory than needed. Unlike <project:#sdf_xarray.open_mfdataset>, this does
+not handle particle positions that are missing from some files.
 
 ```{code-cell} ipython3
 import xarray as xr
 
 xr.open_mfdataset(
    "tutorial_dataset_1d/*.sdf",
+   combine="nested",
+   concat_dim="time",
    join="outer",
    compat="no_conflicts",
    preprocess=sdfxr.SDFPreprocess()
@@ -123,6 +134,9 @@ sdfxr.open_mfdatatree("tutorial_dataset_1d/*.sdf")
 
 ## Loading sparse data
 
+By default, <project:#sdf_xarray.open_mfdataset> puts every variable on a single
+time dimension, and a variable missing from an output (including distribution
+functions and particle data) is filled with nan values at that time step.
 When dealing with sparse data (where different variables are saved at different,
 non-overlapping time steps) you can optimize memory usage by loading the data with
 <project:#sdf_xarray.open_mfdataset> using the parameter `separate_times=True`. This
@@ -140,12 +154,17 @@ sdfxr.open_mfdataset("tutorial_dataset_1d/*.sdf", separate_times=True)
 ## Loading particle data
 
 ```{warning}
-It is **not recommended** to use <inv:#xarray.open_mfdataset> or
-<project:#sdf_xarray.open_mfdataset> to load particle data from multiple
-SDF outputs. The number of particles often varies between outputs,
-which can lead to inconsistent array shapes that these functions
-cannot handle. Instead, consider loading each file individually and
-then concatenating them manually.
+<project:#sdf_xarray.open_mfdataset> can load particle data from multiple
+SDF outputs, and outputs that don't include particles are filled with nan
+values. However, every output that does include a species must contain the
+same number of particles of that species. The number of particles often
+varies between outputs (for example with particle injection, losses or probes),
+in which case loading fails with a conflicting dimension sizes error, and you
+should load each file individually instead.
+
+The `ID_<species>` dimension is the position in each output's particle array,
+not a particle identity, so the same index at two time steps is not
+necessarily the same particle.
 ```
 
 ```{note}
@@ -157,7 +176,7 @@ Failing to do so will result in dimension name conflicts.
 
 By default, particle data isn't kept as it takes up a lot of space.
 Pass `keep_particles=True` as a keyword argument to
-<inv:#xarray.open_dataset> (for single files) or <inv:#xarray.open_mfdataset> (for
+<project:#sdf_xarray.open_dataset> (for single files) or <project:#sdf_xarray.open_mfdataset> (for
 multiple files).
 
 ```{code-cell} ipython3
@@ -168,14 +187,10 @@ sdfxr.open_dataset("tutorial_dataset_1d/0010.sdf", keep_particles=True)
 
 ## Loading specific variables
 
-When loading datasets containing several (`>10`) coordinates/dimensions
-using <project:#sdf_xarray.open_mfdataset>, <inv:#xarray> may struggle to locate
-the necessary RAM to concatenate all of the data (as seen in
-[Issue #57](https://github.com/epochpic/sdf-xarray/issues/57)).
-In this instance, you can optimize memory usage by loading only the data
-you need using the keyword argument `data_vars` and passing one or more
-variables. This creates a dataset consisting only of the given variable(s)
-and the relevant coordinates/dimensions, significantly reducing memory
+If you only need some of the variables, you can load just those using the
+keyword argument `data_vars` and passing one or more variables to
+<project:#sdf_xarray.open_mfdataset>. This creates a dataset consisting only of
+the given variable(s) and the relevant coordinates/dimensions, reducing memory
 consumption.
 
 ```{code-cell} ipython3

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/src/sdf_xarray/__init__.py

@@ -239,30 +239,42 @@ def combine_datasets(
 ) -> xr.Dataset:
     """
     Combine all datasets using a single time dimension, optionally extract
-    data from only the listed data_vars
+    data from only the listed data_vars.
+
+    Coordinates with the "point_data" attribute come from SDF point meshes,
+    i.e. particle positions such as ``X_Particles_proton``. They have one
+    value per particle along an ``ID_<species>`` dimension rather than a
+    grid, and SDFPreprocess gives them a time dimension because particles
+    move. They are usually only in some files (with ``keep_particles=True``,
+    on dumps that output particles). When concatenating along time, xarray
+    fills a data variable missing from a file with ``NaN``'s, but raises if a
+    coordinate is missing. So each file's point coordinates are temporarily
+    turned into data variables and turned back into coordinates after
+    combining. Their names are recorded here because particle data variables
+    (e.g. ``Particles_Px_proton``) carry "point_data" too, so they can't be
+    told apart by the attribute later.
     """
+    point_coords: set[str] = set()
+    preprocess = SDFPreprocess(data_vars=data_vars)
 
-    if data_vars is not None:
-        ds = xr.open_mfdataset(
-            path_glob,
-            join="outer",
-            coords="different",
-            compat="no_conflicts",
-            combine="nested",
-            concat_dim="time",
-            preprocess=SDFPreprocess(data_vars=data_vars),
-            **kwargs,
-        )
-    else:
-        ds = xr.open_mfdataset(
-            path_glob,
-            data_vars="all",
-            coords="different",
-            compat="no_conflicts",
-            join="outer",
-            preprocess=SDFPreprocess(),
-            **kwargs,
-        )
+    def _preprocess(ds: xr.Dataset) -> xr.Dataset:
+        ds = preprocess(ds)
+        names = [k for k, v in ds.coords.items() if v.attrs.get("point_data", False)]
+        point_coords.update(names)
+        return ds.reset_coords(names)
+
+    ds = xr.open_mfdataset(
+        path_glob,
+        join="outer",
+        coords="minimal",
+        compat="no_conflicts",
+        combine="nested",
+        concat_dim="time",
+        data_vars="all",
+        preprocess=_preprocess,
+        **kwargs,
+    )
+    ds = ds.set_coords(sorted(point_coords))
 
     ds.attrs["deck"] = _load_deck(ds.attrs["filename"], deck_path)
 
@@ -626,33 +638,91 @@ def make_time_dims(path_glob):
 
 
 class SDFBackendArray(BackendArray):
-    """Adapater class required for lazy loading"""
+    """Adapter class required for lazy loading.
+
+    Reads go through :py:meth:`SDFFile.read`, which only reads the range of the
+    last axis that is actually requested. SDF arrays are Fortran ordered, so
+    that range is one contiguous block on disk (see
+    `sdf_xarray.SDFBackendArray.raw_indexing_method`).
+
+    ``component`` is only used for particle positions. SDF stores the positions
+    of a particle species as one point mesh block, holding a 1D array of every
+    particle's position for each spatial direction. Each of these directions is
+    a component, and becomes its own coordinate, e.g. ``X_Particles_Electron``
+    and ``Y_Particles_Electron``. ``component`` is the index of the direction
+    this array holds, in the order of the mesh's labels (e.g. 0 for ``X``, 1
+    for ``Y``), and ``variable_name`` is then the name of the point mesh. For
+    every other variable ``component`` is ``None``.
+    """
 
-    __slots__ = ("datastore", "dtype", "shape", "variable_name")
+    __slots__ = ("component", "datastore", "dtype", "shape", "variable_name")
 
-    def __init__(self, variable_name, datastore, shape, dtype):
+    def __init__(self, variable_name, datastore, shape, dtype, component=None):
         self.datastore = datastore
         self.variable_name = variable_name
         self.shape = shape
         self.dtype = dtype
+        self.component = component
 
-    def get_array(self, needs_lock=True):
-        with self.datastore.acquire_context(needs_lock) as ds:
-            return ds.variables[self.variable_name]
-
-    def __getitem__(self, key: indexing.ExplicitIndexer) -> np.typing.ArrayLike:
+    def __getitem__(self, indexer: indexing.ExplicitIndexer) -> np.typing.ArrayLike:
         return indexing.explicit_indexing_adapter(
-            key,
+            indexer,
             self.shape,
-            indexing.IndexingSupport.OUTER,
-            self._raw_indexing_method,
+            indexing.IndexingSupport.BASIC,
+            self.raw_indexing_method,
         )
 
-    def _raw_indexing_method(self, key: tuple) -> np.typing.ArrayLike:
-        # thread safe method that access to data on disk
-        with self.datastore.acquire_context():
-            original_array = self.get_array(needs_lock=False)
-            return original_array.data[key]
+    def raw_indexing_method(self, indexer: tuple) -> np.typing.ArrayLike:
+        """Read the data selected from disk.
+
+        xarray only calls this once the data is needed (e.g. by ``.values``,
+        ``load()``, ``compute()`` or dask), with one int or slice per
+        dimension, and expects back the same array numpy would return for
+        ``full_array[indexer]``. Any negative ints or steps have already
+        been made positive by xarray.
+
+        SDF stores arrays in Fortran order, where the first index changes
+        fastest. A 2D array of shape (nx, ny) is stored as all nx values at
+        y=0, then all nx values at y=1, and so on, so for nx=3 and ny=4 each
+        element's position in the file is::
+
+                  y=0  y=1  y=2  y=3
+            x=0    0    3    6    9
+            x=1    1    4    7   10
+            x=2    2    5    8   11
+
+        Each column (one y, every x) is contiguous, and the columns are stored
+        one after another. So a range of the *last* index, with every other
+        dimension in full, is one contiguous block where ``[:, 1:3]`` is positions
+        3 to 8. A range of any other index is scattered through the file where
+        ``[1:3, :]`` is positions 1, 2, 4, 5, 7, 8, 10 and 11. So we only read
+        the range of the last dimension that is needed from disk, then index the
+        other dimensions in memory.
+        """
+        *other_indices, last_index = indexer
+
+        if isinstance(last_index, slice):
+            start, stop, step = last_index.indices(self.shape[-1])
+            # An empty slice such as 5:1 has stop < start, so read nothing
+            stop = max(start, stop)
+            data = self.read_last_dim(start, stop)
+            # SDF can't skip elements, so a step such as ::2 is applied here
+            return data[*other_indices, ::step]
+
+        data = self.read_last_dim(last_index, last_index + 1)
+        # Index with 0, not 0:1, so the dimension is dropped as numpy would
+        return data[*other_indices, 0]
+
+    def read_last_dim(self, start: int, stop: int) -> np.ndarray:
+        """Read the range [start, stop) of the last dimension from disk."""
+        with self.datastore.acquire_context() as sdf_file:
+            if self.component is None:
+                variable = sdf_file.variables[self.variable_name]
+                return sdf_file.read(variable, start, stop)
+
+            # Particle positions: one component (e.g. x) of a point mesh
+            mesh = sdf_file.grids[self.variable_name]
+            return sdf_file.read_mesh_component(mesh, self.component, start, stop)
 
 
 class SDFDataStore(AbstractDataStore):
@@ -765,7 +835,25 @@ class SDFDataStore(AbstractDataStore):
 
             base_name = _process_grid_name(value.name, _norm_grid_name)
 
-            for label, coord, unit in zip(value.labels, value.data, value.units):
+            if value.is_point_data:
+                # Particle positions are not dimension coordinates, so they
+                # don't need to be in memory to build an index; keep them lazy
+                coord_data = [
+                    indexing.LazilyIndexedArray(
+                        SDFBackendArray(
+                            key,
+                            self,
+                            shape=value.shape[:1],
+                            dtype=value.dtype,
+                            component=component,
+                        )
+                    )
+                    for component in range(len(value.labels))
+                ]
+            else:
+                coord_data = value.data
+
+            for label, coord, unit in zip(value.labels, coord_data, value.units):
                 full_name = f"{label}_{base_name}"
                 dim_name = (
                     f"ID_{_process_grid_name(key, _grid_species_name)}"
@@ -801,7 +889,14 @@ class SDFDataStore(AbstractDataStore):
                 # scalar, or because it's an array over something
                 # else. We have no more information, so just make up
                 # some (hopefully) unique dimension names
-                shape = getattr(value.data, "shape", ())
+                if isinstance(value, Constant):
+                    shape = ()
+                    data = value.data
+                else:
+                    shape = value.shape
+                    data = indexing.LazilyIndexedArray(
+                        SDFBackendArray(key, self, shape=shape, dtype=value.dtype)
+                    )
                 dims = [f"dim_{key}_{n}" for n, _ in enumerate(shape)]
                 base_name = _rename_with_underscore(key)
 
@@ -811,7 +906,7 @@ class SDFDataStore(AbstractDataStore):
                 if value.units is not None:
                     data_attrs["units"] = value.units
 
-                var = Variable(dims, value.data, attrs=data_attrs)
+                var = Variable(dims, data, attrs=data_attrs)
 
                 # Provide preferred_chunks for constants so dask aligns to natural shapes
                 var.encoding["preferred_chunks"] = dict(zip(dims, shape))
@@ -884,14 +979,14 @@ class SDFDataStore(AbstractDataStore):
                 "long_name": long_name,
             }
             lazy_data = indexing.LazilyIndexedArray(
-                SDFBackendArray(key, self, shape=value.shape, dtype=value.data.dtype)
+                SDFBackendArray(key, self, shape=value.shape, dtype=value.dtype)
             )
             var = Variable(var_coords, lazy_data, data_attrs)
             # Set preferred chunks to match on-disk layout
             # For point data (1D): full dimension
             # For grid data (N-D): individual grid chunk sizes
             if value.is_point_data:
-                var.encoding["preferred_chunks"] = {var_coords[0]: len(value.data)}
+                var.encoding["preferred_chunks"] = {var_coords[0]: value.shape[0]}
             else:
                 # Align with on-disk grid structure
                 chunk_dict = {}

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/src/sdf_xarray/_version.py

@@ -18,7 +18,7 @@ version_tuple: tuple[int | str, ...]
 commit_id: str | None
 __commit_id__: str | None
 
-__version__ = version = '0.9.0'
-__version_tuple__ = version_tuple = (0, 9, 0)
+__version__ = version = '0.10.0'
+__version_tuple__ = version_tuple = (0, 10, 0)
 
-__commit_id__ = commit_id = 'g65bb865da'
+__commit_id__ = commit_id = 'g878b38c62'

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/src/sdf_xarray/csdf.pxd

@@ -61,6 +61,7 @@ cdef extern from "sdf.h":
         int64_t dims[SDF_MAXDIMS]
         int64_t local_dims[SDF_MAXDIMS]
         int64_t data_length
+        int64_t data_location
         int32_t ndims, geometry, datatype, blocktype
         int32_t stagger, datatype_out
         char const_value[16]
@@ -125,3 +126,4 @@ cdef extern from "sdf_helper.h":
 cdef extern from "stack_allocator.h":
     void sdf_stack_destroy(sdf_file_t *h)
     void sdf_stack_init(sdf_file_t *h)
+    void sdf_stack_free_block(sdf_file_t *h, sdf_block_t *b)

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/src/sdf_xarray/download.py

@@ -58,18 +58,25 @@ def fetch_dataset(
     import pooch  # noqa: PLC0415
 
     logger = pooch.get_logger()
-    datasets = pooch.create(
-        path=pooch.os_cache("sdf_datasets"),
-        base_url="doi:10.5281/zenodo.17618509",
-        registry=None,
-        retry_if_failed=10,
-    )
-    datasets.load_registry_from_doi()
 
-    datasets.fetch(
-        f"{dataset_name}.zip", processor=pooch.Unzip(extract_dir="."), progressbar=True
-    )
-    cache_path = Path(datasets.path) / dataset_name
+    cache_dir = Path(pooch.os_cache("sdf_datasets"))
+    cache_path = cache_dir / dataset_name
+
+    # Only download dataset if a local copy is not found
+    if not cache_path.is_dir():
+        datasets = pooch.create(
+            path=cache_dir,
+            base_url="doi:10.5281/zenodo.17618509",
+            registry=None,
+            retry_if_failed=10,
+        )
+        datasets.load_registry_from_doi()
+
+        datasets.fetch(
+            f"{dataset_name}.zip",
+            processor=pooch.Unzip(extract_dir="."),
+            progressbar=True,
+        )
 
     if save_path is not None:
         save_path = Path(save_path)

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/src/sdf_xarray/sdf_interface.pyx

@@ -167,6 +167,7 @@ cdef class SDFFile:
     cdef public dict[str, Mesh] grids
 
     def __cinit__(self, filename: str):
+        self.filename = filename
         self._c_sdf_file = csdf.sdf_open(
             filename.encode("UTF-8"), 0, csdf.SDF_READ, False
         )
@@ -254,8 +255,14 @@ cdef class SDFFile:
 
             block = block.next
 
-    cpdef read(self, var: Block):
+    cpdef read(self, var: Block, start=None, stop=None):
         """Read a variable from the file, returning numpy array
+
+        For variables, ``start`` and ``stop`` select a range of the last index,
+        i.e. ``var[..., start:stop]``, so that only that part of the variable is
+        read from disk. As SDF arrays are Fortran ordered (the first index
+        changes fastest), this range is one contiguous block on disk, so it is
+        a single read.
         """
 
         if self._c_sdf_file is NULL:
@@ -276,23 +283,112 @@ cdef class SDFFile:
         if block is NULL:
             raise RuntimeError(f"Could not read variable '{var.name}'")
 
+        shape = var.shape
+        cdef bint is_section = False
+        cdef csdf.int64_t starts[csdf.SDF_MAXDIMS]
+        cdef csdf.int64_t ends[csdf.SDF_MAXDIMS]
+        if not is_mesh and shape and (start is not None or stop is not None):
+            start = 0 if start is None else start
+            stop = shape[-1] if stop is None else stop
+            if not 0 <= start <= stop <= shape[-1]:
+                raise IndexError(
+                    f"Invalid range [{start}, {stop}) for '{var.name}' "
+                    f"with shape {shape}"
+                )
+            if start == stop:
+                # Nothing to read, and the library can't read an empty section
+                return np.empty((*shape[:-1], 0), dtype=var.dtype)
+            is_section = (start, stop) != (0, shape[-1])
+
+        if is_section and block.blocktype not in (
+            csdf.SDF_BLOCKTYPE_PLAIN_VARIABLE, csdf.SDF_BLOCKTYPE_POINT_VARIABLE
+        ):
+            # Only these readers support array sections (`sdf_read_array`, for
+            # example, ignores them), so read everything and slice it here
+            return self.read(var)[..., start:stop].copy()
+
+        if is_section:
+            for i in range(block.ndims):
+                # The array section is clamped to the local (per-rank) dims,
+                # which for point data are only set during the read itself.
+                # For serial reads these are always the global dims
+                block.local_dims[i] = block.dims[i]
+                starts[i] = 0
+                ends[i] = block.dims[i]
+            starts[block.ndims - 1] = start
+            ends[block.ndims - 1] = stop
+            csdf.sdf_block_set_array_section(block, block.ndims, starts, ends, NULL)
+            shape = (*shape[:-1], stop - start)
+
         self._c_sdf_file.current_block = block
-        csdf.sdf_helper_read_data(self._c_sdf_file, block)
+        try:
+            csdf.sdf_helper_read_data(self._c_sdf_file, block)
+
+            if is_mesh:
+                # Meshes store the data for separate dimensions in block.grids
+                if not (block.grids is not NULL and block.grids[0] is not NULL):
+                    raise RuntimeError(f"Could not read variable '{var.name}'")
+
+                data = []
+                for i, dim in enumerate(shape):
+                    data.append(
+                        self._make_array((dim,), var.dtype, block.grids[i])
+                    )
+                return tuple(data)
+
+            # Normal variables
+            return self._make_array(shape, var.dtype, block.data)
+        finally:
+            # The library keeps every block it reads cached until the file is
+            # closed, but we've just copied it into numpy, so release it now.
+            # Otherwise we hold two copies of everything that's been read
+            csdf.sdf_stack_free_block(self._c_sdf_file, block)
+            if is_section:
+                csdf.sdf_block_set_array_section(block, block.ndims, NULL, NULL, NULL)
+
+    cpdef read_mesh_component(self, mesh: Mesh, int component, start=None, stop=None):
+        """Read one component (e.g. just ``X``) of a mesh, returning numpy array
+
+        For point meshes, ``start`` and ``stop`` select a range of points so
+        that only those are read from disk.
+        """
+
+        if self._c_sdf_file is NULL:
+            raise RuntimeError(
+                f"Can't read '{mesh.name}', file '{self.filename}' is closed"
+            )
 
-        if is_mesh:
-            # Meshes store the data for separate dimensions in block.grids
-            if not (block.grids is not NULL and block.grids[0] is not NULL):
-                raise RuntimeError(f"Could not read variable '{var.name}'")
+        if not mesh.is_point_data or mesh.parent:
+            return self.read(mesh)[component][start:stop]
 
-            data = []
-            for i, dim in enumerate(var.shape):
-                data.append(
-                    self._make_array((dim,), var.dtype, block.grids[i])
-                )
-            return tuple(data)
+        cdef csdf.sdf_block_t* block = csdf.sdf_find_block_by_name(
+            self._c_sdf_file, mesh.name.encode("utf-8")
+        )
 
-        # Normal variables
-        return self._make_array(var.shape, var.dtype, block.data)
+        if block is NULL:
+            raise RuntimeError(f"Could not read variable '{mesh.name}'")
+
+        if block.datatype != block.datatype_out:
+            # The library is converting the type on read, so let it
+            return self.read(mesh)[component][start:stop]
+
+        # The library always reads every component of a point mesh, and its
+        # array sections don't work for them. The components are just stored
+        # one after the other though (see `sdf_read_point_mesh`), so read the
+        # part we want directly
+        npoints = block.dims[0]
+        start = 0 if start is None else start
+        stop = npoints if stop is None else stop
+        if not 0 <= component < block.ndims or not 0 <= start <= stop <= npoints:
+            raise IndexError(
+                f"Invalid component {component} or range [{start}, {stop}) "
+                f"for '{mesh.name}' with {block.ndims} components of {npoints} points"
+            )
+
+        offset = block.data_location + (component * npoints + start) * mesh.dtype.itemsize
+        return np.fromfile(
+            self.filename, dtype=mesh.dtype, count=stop - start, offset=offset
+        )
 
     cdef _read_mid_grid(self, mesh: Mesh):
         """Read a midpoint grid"""

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/tests/test_cython.py

@@ -138,3 +138,49 @@ if __name__ == "__main__":
         pprint(f.variables["Electric Field/Ex"].data)
         print("grid")
         pprint(f.grids["grid"].data)
+
+
+def test_read_section_matches_full_read():
+    with SDFFile(str(TEST_FILES_DIR / "0010.sdf")) as f:
+        dropped_key_affixes = ["cpu", "output file", "boundary"]
+        variables = [
+            v
+            for v in f.variables.values()
+            if getattr(v, "shape", ())
+            and all(key not in v.name.lower() for key in dropped_key_affixes)
+        ]
+        assert any(v.is_point_data for v in variables)
+        for var in variables:
+            full = f.read(var)
+            size = var.shape[-1]
+            ranges = [(0, 0), (0, 1), (1, size), (size // 3, size // 2 + 1)]
+            for start, stop in ranges:
+                npt.assert_array_equal(
+                    f.read(var, start, stop), full[..., start:stop], err_msg=var.name
+                )
+            # Reading a section must not affect later full reads
+            npt.assert_array_equal(f.read(var), full, err_msg=var.name)
+
+
+def test_read_section_out_of_range():
+    with SDFFile(str(TEST_FILES_DIR / "0010.sdf")) as f:
+        var = f.variables["Electric Field/Ex"]
+        with pytest.raises(IndexError):
+            f.read(var, 0, var.shape[-1] + 1)
+
+
+def test_read_mesh_component_matches_full_read():
+    with SDFFile(str(TEST_FILES_DIR / "0010.sdf")) as f:
+        meshes = [g for g in f.grids.values() if g.is_point_data]
+        assert meshes
+        for mesh in meshes:
+            full = f.read(mesh)
+            size = mesh.shape[0]
+            for component in range(len(mesh.labels)):
+                npt.assert_array_equal(
+                    f.read_mesh_component(mesh, component), full[component]
+                )
+                npt.assert_array_equal(
+                    f.read_mesh_component(mesh, component, 1, size // 2),
+                    full[component][1 : size // 2],
+                )

{sdf_xarray-0.9.0 → sdf_xarray-0.10.0}/tests/test_dataset.py

@@ -7,6 +7,7 @@ import xarray as xr
 
 import sdf_xarray as sdfxr
 from sdf_xarray import (
+    SDFFile,
     SDFPreprocess,
     _load_deck,
     _process_latex_name,
@@ -19,6 +20,7 @@ TEST_MISMATCHED_FILES_DIR = download.fetch_dataset("test_mismatched_files")
 TEST_ARRAYS_DIR = download.fetch_dataset("test_array_no_grids")
 TEST_3D_DIST_FN = download.fetch_dataset("test_dist_fn")
 TEST_2D_PARTICLE_DATA = download.fetch_dataset("test_two_probes_2D")
+TEST_2D_MOVING_WINDOW = download.fetch_dataset("test_files_2D_moving_window")
 
 
 # Type hinting support
@@ -206,6 +208,33 @@ def test_multiple_files_multiple_time_dims():
         assert df["Absorption_Total_Laser_Energy_Injected"].shape == (11,)
 
 
+def test_open_mfdataset_varying_grid():
+    # A grid that differs between files must not become a concat dimension
+    paths = sorted(TEST_2D_MOVING_WINDOW.glob("*.sdf"))
+    with sdfxr.open_mfdataset(paths) as df:
+        assert df.sizes["time"] == len(paths)
+        assert df["Wall_time"].dims == ("time",)
+        density = df["Derived_Number_Density_electron"]
+        assert density.dims == ("time", "X_Grid_mid", "Y_Grid_mid")
+        assert df.indexes["X_Grid_mid"].is_monotonic_increasing
+
+        for i, path in enumerate(paths):
+            with xr.open_dataset(path) as single:
+                expected = single["Derived_Number_Density_electron"]
+                actual = density.isel(time=i).sel(X_Grid_mid=expected["X_Grid_mid"])
+                npt.assert_array_equal(actual, expected)
+
+
+def test_open_mfdataset_sparse_dist_fn():
+    # The first file's dist_fn and its grid are missing from the last file
+    with sdfxr.open_mfdataset(TEST_3D_DIST_FN.glob("*.sdf")) as df:
+        dist_fn = df["dist_fn_x_px_py_Electron"]
+        assert df.sizes["time"] == 3
+        assert dist_fn.shape == (3, 16, 20, 20)
+        assert dist_fn.isel(time=slice(0, 2)).notnull().all()
+        assert dist_fn.isel(time=2).isnull().all()
+
+
 def test_resolve_paths_from_string_pattern():
     pattern = str(TEST_FILES_DIR / "*.sdf")
     result = resolve_paths(pattern)
@@ -594,6 +623,30 @@ def test_open_mfdataset_data_vars_sparse_time():
         npt.assert_allclose(time_values, time.values, rtol=1e-6)
 
 
+def test_open_mfdataset_data_vars_sparse_particles():
+    with sdfxr.open_mfdataset(
+        TEST_FILES_DIR.glob("*.sdf"),
+        keep_particles=True,
+        data_vars=["Particles_Px_proton"],
+    ) as df:
+        px_protons = df["Particles_Px_proton"]
+        assert sorted(px_protons.coords) == sorted(("X_Particles_proton", "time"))
+        assert px_protons.shape == (11, 1920)
+        has_data = px_protons.notnull().any("ID_proton").values
+        assert has_data.tolist() == [False] * 10 + [True]
+
+
+def test_open_mfdataset_data_vars_sparse_dist_fn():
+    with sdfxr.open_mfdataset(
+        TEST_3D_DIST_FN.glob("*.sdf"),
+        data_vars=["dist_fn_x_px_py_Electron"],
+    ) as df:
+        dist_fn = df["dist_fn_x_px_py_Electron"]
+        assert dist_fn.shape == (3, 16, 20, 20)
+        assert dist_fn.isel(time=slice(0, 2)).notnull().all()
+        assert dist_fn.isel(time=2).isnull().all()
+
+
 def test_open_mfdataset_data_vars_separate_times_single():
     with sdfxr.open_mfdataset(
         TEST_FILES_DIR.glob("*.sdf"),
@@ -765,3 +818,37 @@ def test_epydeck_parsing_failing_deck(tmp_path, capsys):
     _load_deck("", deck)
     out, _ = capsys.readouterr()
     assert "error occurred while trying to load the input deck" in out
+
+
+def test_lazy_particle_coords():
+    with sdfxr.open_dataset(TEST_FILES_DIR / "0010.sdf", keep_particles=True) as df:
+        x_coord = df["X_Particles_proton"]
+        assert not x_coord.variable._in_memory
+        with SDFFile(str(TEST_FILES_DIR / "0010.sdf")) as f:
+            mesh = next(
+                g for g in f.grids.values() if g.name == x_coord.attrs["full_name"]
+            )
+            expected = f.read(mesh)[mesh.labels.index("X")]
+        npt.assert_array_equal(x_coord[2:-3].values, expected[2:-3])
+        npt.assert_array_equal(x_coord.values, expected)
+
+
+@pytest.mark.parametrize(
+    "indexers",
+    [
+        {"dim0": [0, 2, 5], "dim1": [1, 3]},
+        {"dim0": [5, 0], "dim2": slice(10, 2, -2)},
+        {"dim1": -1, "dim2": [0, -1]},
+    ],
+)
+def test_lazy_indexing_matches_eager(indexers):
+    path = TEST_3D_DIST_FN / "0000.sdf"
+    name = "dist_fn_x_px_py_Electron"
+    with xr.open_dataset(path) as eager:
+        expected = eager[name].load()
+    with xr.open_dataset(path) as lazy:
+        dims = lazy[name].dims
+        indexers = {dims[int(k[-1])]: v for k, v in indexers.items()}
+        npt.assert_array_equal(
+            lazy[name].isel(indexers).values, expected.isel(indexers).values
+        )

sdf_xarray-0.9.0/docs/known_issues.md

@@ -1,15 +0,0 @@
-# Known Issues
-
-There are a couple of known 'quirks' in `sdf_xarray`:
-
-- [Issue #57](https://github.com/epochpic/sdf-xarray/issues/57) Loading multiple
-SDF files with <project:#sdf_xarray.open_mfdataset> can lead to out-of-memory
-errors. The issue is believed to stem from how the underlying <inv:#xarray> library
-handles coordinates, causing it to infer an excessively large array shape that
-requests far more memory than is needed. Due to the significant architectural
-changes required for a fix, the maintainers do not plan to resolve this. The
-recommended solution is to load the files individually or in smaller batches.
-- [Issue #112](https://github.com/epochpic/sdf-xarray/issues/112) Due to the 
-way [cylindrical EPOCH](https://github.com/epochpic/cylindrical_epoch) sets up 
-grids we can't currently load it's SDF files. If there is enough interest in this
-we will attempt to find a solution otherwise it shall remain low priority.