Skip to content

Issue 2656 windowed array#2671

Open
fluidnumericsJoe wants to merge 5 commits into
mainfrom
issue-2656-windowed-array
Open

Issue 2656 windowed array#2671
fluidnumericsJoe wants to merge 5 commits into
mainfrom
issue-2656-windowed-array

Conversation

@fluidnumericsJoe

Copy link
Copy Markdown
Contributor

Description

Checklist

AI Disclosure

  • This PR contains AI-generated content.
    • I have tested any AI-generated content in my PR.
    • I take responsibility for any AI-generated content in my PR.
    • Describe how you used it (e.g., by pasting your prompt): Claude Code was used for doc string and test creation; each were hand-edited after creation.

Comment on lines +130 to +157
def to_windowed_arrays(self, *, max_levels: int | None = None):
"""Wrap dask-backed field data in a rolling time-window cache.

Opt-in optimization for forward-marching simulations where all particles
share a single clock. For each dask-backed field ``isel`` then samples
a resident NumPy window instead of re-reading chunks and paying the
dask scheduling overhead on every kernel step.

NumPy-backed fields are left unchanged, so this is safe to call
more than once.

Parameters
----------
max_levels : int, optional
Cap on the number of time levels kept resident per field. ``None``
(default) retains every level that the advancing clock still brackets.

Returns
-------
FieldSet
``self``, to allow chaining.
"""
for field in self.fields.values():
components = (field.U, field.V, field.W) if isinstance(field, VectorField) else (field,)
for component in components:
if component is not None:
component.data = maybe_windowed(component.data, max_levels=max_levels)
return self

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Hmmm, I think that this is actually incompatible with #2646 :'( since that refactors things so that there isn't a field.data object anymore.

Do you think that it would be possible to wrap this WindowArray approach similar to what is done in #2668 ?

"""Bulk, sequential read of one time level into NumPy (the dask->NumPy step)."""
return np.asarray(self._data.isel({self._tdim: int(lvl)}).values)

def _ensure(self, levels: np.ndarray) -> None:

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
def _ensure(self, levels: np.ndarray) -> None:
def _ensure(self, levels: np.ndarray) -> None:
if self._max < len(levels):
raise ValueError("Trying to ensure more levels than we can hold.")

ds["U"].data[:] = 1.0 # steady zonal flow -> in-bounds, deterministic

def run(windowed):
d = ds.chunk({"time": 1}) if chunked else ds

Copy link
Copy Markdown
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

Suggested change
d = ds.chunk({"time": 1}) if chunked else ds
d = ds.chunk({"time": 1}) if windowed else ds

Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

Status: Backlog

Development

Successfully merging this pull request may close these issues.

Out-of-core fields: move unstructured interpolators to .isel, and add a time-windowed array cache (WindowedArray)

3 participants