Implement .to_chunk_cached_arrays() with ChunkCachedArray

[VeckoTheGecko]

Description

Provides an opt-in optimization that wraps dask-backed field data in
chunk-level LRU caches.
vectorized .isel() calls hit an in-memory cache instead of recomputing
dask task graphs, giving large speedups for particle simulations.

This PR adds ChunkCachedArray according to the approach detailed in #2854 , providing an opt-in optimization wrapping dash-backed field data with (dask)chunk-level LRU caches on the array level. This PR:

• Adds a subpackage parcels._chunk_cached_array with the implementation of the Array and the needed LRU cache. Note the cache is on the array level with configurable memory limits per array.
• Exposes a function parcels._chunk_cached_array.wrap_dataset providing the main entrypoint for Parcels for wrapping Xarray dataset objects
• Adds .to_chunk_cached_arrays to FieldSet and ModelData classes.
• Adds equivalence testing (test_backends.py) on a NEMO dataset, asserting that the same results are received regardless of backend.

Still needed:

• (separate PR) Upstream testing
• (separate PR?) Backend equivalency testing on Unstructured grids?
• Benchmarking with realworld simulations
• (pre-merge) Remove temporary testing files (data-generation.py,benchmark_chunk_cache.py,benchmark_chunk_cache.png)

Future work:

• max_cache_bytes on FieldSet.to_chunk_cached_arrays operates per variable (i.e., per Field). This requires the user to think about how many Fields are in their simulation in order to tune their performance. Down the line, we could tune this to allocate a total cache of (e.g.,) 50% of the system memory (which would allow for other memory overhead of coordinates (not stored in the cache), the Python process, or other user processes). This would require significant refactoring not only on how the chunk cache is handled, but also backends in general.

(cc @erikvansebille, @wyatt-fluidnumerics )

Performance

I have some some preliminary profiling with a 20Gb idealised dataset, and found the following performance profile which looks like a promising first step.

I've included data-generation.py and benchmark_chunk_cache.py here for testing in case it helps.

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Checklist

• Contributes to New Parcels backend - ChunkCachedArray #2854
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

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): I guided an LLM to explore the Xarray codebase, and to help work on the implementation of the Array class. I also used an LLM to generate and iterate on the profiling script, and write docstrings.

[erikvansebille]

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Thanks for this exciting PR, @VeckoTheGecko! I plan to do some real-world performance testing (also including/comparing #2846 and v3) later this week. Will report back here when I know more!

[VeckoTheGecko]

Description

Provides an opt-in optimization that wraps dask-backed field data in
chunk-level LRU caches.
vectorized .isel() calls hit an in-memory cache instead of recomputing
dask task graphs, giving large speedups for particle simulations.

This PR adds ChunkCachedArray according to the approach detailed in #2854 , providing an opt-in optimization wrapping dash-backed field data with (dask)chunk-level LRU caches on the array level. This PR:

• Adds a subpackage parcels._chunk_cached_array with the implementation of the Array and the needed LRU cache. Note the cache is on the array level with configurable memory limits per array.
• Exposes a function parcels._chunk_cached_array.wrap_dataset providing the main entrypoint for Parcels for wrapping Xarray dataset objects
• Adds .to_chunk_cached_arrays to FieldSet and ModelData classes.
• Adds equivalence testing (test_backends.py) on a NEMO dataset, asserting that the same results are received regardless of backend.

Still needed:

• (separate PR) Upstream testing
• (separate PR?) Backend equivalency testing on Unstructured grids?
• Benchmarking with realworld simulations
• (pre-merge) Remove temporary testing files (data-generation.py,benchmark_chunk_cache.py,benchmark_chunk_cache.png)

Future work:

• max_cache_bytes on FieldSet.to_chunk_cached_arrays operates per variable (i.e., per Field). This requires the user to think about how many Fields are in their simulation in order to tune their performance. Down the line, we could tune this to allocate a total cache of (e.g.,) 50% of the system memory (which would allow for other memory overhead of coordinates (not stored in the cache), the Python process, or other user processes). This would require significant refactoring not only on how the chunk cache is handled, but also backends in general.

(cc @erikvansebille, @wyatt-fluidnumerics )

Performance

I have some some preliminary profiling with a 20Gb idealised dataset, and found the following performance profile which looks like a promising first step.

I've included data-generation.py and benchmark_chunk_cache.py here for testing in case it helps.

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Checklist

• Contributes to New Parcels backend - ChunkCachedArray #2854
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

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): I guided an LLM to explore the Xarray codebase, and to help work on the implementation of the Array class. I also used an LLM to generate and iterate on the profiling script, and write docstrings.

[erikvansebille]

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Thanks for this exciting PR, @VeckoTheGecko! I plan to do some real-world performance testing (also including/comparing #2846 and v3) later this week. Will report back here when I know more!

[VeckoTheGecko]

Description

Provides an opt-in optimization that wraps dask-backed field data in
chunk-level LRU caches.
vectorized .isel() calls hit an in-memory cache instead of recomputing
dask task graphs, giving large speedups for particle simulations.

This PR adds ChunkCachedArray according to the approach detailed in #2854 , providing an opt-in optimization wrapping dash-backed field data with (dask)chunk-level LRU caches on the array level. This PR:

• Adds a subpackage parcels._chunk_cached_array with the implementation of the Array and the needed LRU cache. Note the cache is on the array level with configurable memory limits per array.
• Exposes a function parcels._chunk_cached_array.wrap_dataset providing the main entrypoint for Parcels for wrapping Xarray dataset objects
• Adds .to_chunk_cached_arrays to FieldSet and ModelData classes.
• Adds equivalence testing (test_backends.py) on a NEMO dataset, asserting that the same results are received regardless of backend.

Still needed:

• (separate PR) Upstream testing
• (separate PR?) Backend equivalency testing on Unstructured grids?
• Benchmarking with realworld simulations
• (pre-merge) Remove temporary testing files (data-generation.py,benchmark_chunk_cache.py,benchmark_chunk_cache.png)

Future work:

• max_cache_bytes on FieldSet.to_chunk_cached_arrays operates per variable (i.e., per Field). This requires the user to think about how many Fields are in their simulation in order to tune their performance. Down the line, we could tune this to allocate a total cache of (e.g.,) 50% of the system memory (which would allow for other memory overhead of coordinates (not stored in the cache), the Python process, or other user processes). This would require significant refactoring not only on how the chunk cache is handled, but also backends in general.

(cc @erikvansebille, @wyatt-fluidnumerics )

Performance

I have some some preliminary profiling with a 20Gb idealised dataset, and found the following performance profile which looks like a promising first step.

I've included data-generation.py and benchmark_chunk_cache.py here for testing in case it helps.

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Checklist

• Contributes to New Parcels backend - ChunkCachedArray #2854
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

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): I guided an LLM to explore the Xarray codebase, and to help work on the implementation of the Array class. I also used an LLM to generate and iterate on the profiling script, and write docstrings.

[erikvansebille]

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Thanks for this exciting PR, @VeckoTheGecko! I plan to do some real-world performance testing (also including/comparing #2846 and v3) later this week. Will report back here when I know more!

[VeckoTheGecko]

Description

Provides an opt-in optimization that wraps dask-backed field data in
chunk-level LRU caches.
vectorized .isel() calls hit an in-memory cache instead of recomputing
dask task graphs, giving large speedups for particle simulations.

This PR adds ChunkCachedArray according to the approach detailed in #2854 , providing an opt-in optimization wrapping dash-backed field data with (dask)chunk-level LRU caches on the array level. This PR:

• Adds a subpackage parcels._chunk_cached_array with the implementation of the Array and the needed LRU cache. Note the cache is on the array level with configurable memory limits per array.
• Exposes a function parcels._chunk_cached_array.wrap_dataset providing the main entrypoint for Parcels for wrapping Xarray dataset objects
• Adds .to_chunk_cached_arrays to FieldSet and ModelData classes.
• Adds equivalence testing (test_backends.py) on a NEMO dataset, asserting that the same results are received regardless of backend.

Still needed:

• (separate PR) Upstream testing
• (separate PR?) Backend equivalency testing on Unstructured grids?
• Benchmarking with realworld simulations
• (pre-merge) Remove temporary testing files (data-generation.py,benchmark_chunk_cache.py,benchmark_chunk_cache.png)

Future work:

• max_cache_bytes on FieldSet.to_chunk_cached_arrays operates per variable (i.e., per Field). This requires the user to think about how many Fields are in their simulation in order to tune their performance. Down the line, we could tune this to allocate a total cache of (e.g.,) 50% of the system memory (which would allow for other memory overhead of coordinates (not stored in the cache), the Python process, or other user processes). This would require significant refactoring not only on how the chunk cache is handled, but also backends in general.

(cc @erikvansebille, @wyatt-fluidnumerics )

Performance

I have some some preliminary profiling with a 20Gb idealised dataset, and found the following performance profile which looks like a promising first step.

I've included data-generation.py and benchmark_chunk_cache.py here for testing in case it helps.

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Checklist

• Contributes to New Parcels backend - ChunkCachedArray #2854
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

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): I guided an LLM to explore the Xarray codebase, and to help work on the implementation of the Array class. I also used an LLM to generate and iterate on the profiling script, and write docstrings.

[erikvansebille]

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Thanks for this exciting PR, @VeckoTheGecko! I plan to do some real-world performance testing (also including/comparing #2846 and v3) later this week. Will report back here when I know more!

[VeckoTheGecko]

Description

Provides an opt-in optimization that wraps dask-backed field data in
chunk-level LRU caches.
vectorized .isel() calls hit an in-memory cache instead of recomputing
dask task graphs, giving large speedups for particle simulations.

This PR adds ChunkCachedArray according to the approach detailed in #2854 , providing an opt-in optimization wrapping dash-backed field data with (dask)chunk-level LRU caches on the array level. This PR:

• Adds a subpackage parcels._chunk_cached_array with the implementation of the Array and the needed LRU cache. Note the cache is on the array level with configurable memory limits per array.
• Exposes a function parcels._chunk_cached_array.wrap_dataset providing the main entrypoint for Parcels for wrapping Xarray dataset objects
• Adds .to_chunk_cached_arrays to FieldSet and ModelData classes.
• Adds equivalence testing (test_backends.py) on a NEMO dataset, asserting that the same results are received regardless of backend.

Still needed:

• (separate PR) Upstream testing
• (separate PR?) Backend equivalency testing on Unstructured grids?
• Benchmarking with realworld simulations
• (pre-merge) Remove temporary testing files (data-generation.py,benchmark_chunk_cache.py,benchmark_chunk_cache.png)

Future work:

• max_cache_bytes on FieldSet.to_chunk_cached_arrays operates per variable (i.e., per Field). This requires the user to think about how many Fields are in their simulation in order to tune their performance. Down the line, we could tune this to allocate a total cache of (e.g.,) 50% of the system memory (which would allow for other memory overhead of coordinates (not stored in the cache), the Python process, or other user processes). This would require significant refactoring not only on how the chunk cache is handled, but also backends in general.

(cc @erikvansebille, @wyatt-fluidnumerics )

Performance

I have some some preliminary profiling with a 20Gb idealised dataset, and found the following performance profile which looks like a promising first step.

I've included data-generation.py and benchmark_chunk_cache.py here for testing in case it helps.

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Checklist

• Contributes to New Parcels backend - ChunkCachedArray #2854
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

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): I guided an LLM to explore the Xarray codebase, and to help work on the implementation of the Array class. I also used an LLM to generate and iterate on the profiling script, and write docstrings.

[erikvansebille]

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Thanks for this exciting PR, @VeckoTheGecko! I plan to do some real-world performance testing (also including/comparing #2846 and v3) later this week. Will report back here when I know more!

[VeckoTheGecko]

Description

Provides an opt-in optimization that wraps dask-backed field data in
chunk-level LRU caches.
vectorized .isel() calls hit an in-memory cache instead of recomputing
dask task graphs, giving large speedups for particle simulations.

This PR adds ChunkCachedArray according to the approach detailed in #2854 , providing an opt-in optimization wrapping dash-backed field data with (dask)chunk-level LRU caches on the array level. This PR:

• Adds a subpackage parcels._chunk_cached_array with the implementation of the Array and the needed LRU cache. Note the cache is on the array level with configurable memory limits per array.
• Exposes a function parcels._chunk_cached_array.wrap_dataset providing the main entrypoint for Parcels for wrapping Xarray dataset objects
• Adds .to_chunk_cached_arrays to FieldSet and ModelData classes.
• Adds equivalence testing (test_backends.py) on a NEMO dataset, asserting that the same results are received regardless of backend.

Still needed:

• (separate PR) Upstream testing
• (separate PR?) Backend equivalency testing on Unstructured grids?
• Benchmarking with realworld simulations
• (pre-merge) Remove temporary testing files (data-generation.py,benchmark_chunk_cache.py,benchmark_chunk_cache.png)

Future work:

• max_cache_bytes on FieldSet.to_chunk_cached_arrays operates per variable (i.e., per Field). This requires the user to think about how many Fields are in their simulation in order to tune their performance. Down the line, we could tune this to allocate a total cache of (e.g.,) 50% of the system memory (which would allow for other memory overhead of coordinates (not stored in the cache), the Python process, or other user processes). This would require significant refactoring not only on how the chunk cache is handled, but also backends in general.

(cc @erikvansebille, @wyatt-fluidnumerics )

Performance

I have some some preliminary profiling with a 20Gb idealised dataset, and found the following performance profile which looks like a promising first step.

I've included data-generation.py and benchmark_chunk_cache.py here for testing in case it helps.

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Checklist

• Contributes to New Parcels backend - ChunkCachedArray #2854
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

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): I guided an LLM to explore the Xarray codebase, and to help work on the implementation of the Array class. I also used an LLM to generate and iterate on the profiling script, and write docstrings.

[erikvansebille]

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Thanks for this exciting PR, @VeckoTheGecko! I plan to do some real-world performance testing (also including/comparing #2846 and v3) later this week. Will report back here when I know more!

[VeckoTheGecko]

Description

Provides an opt-in optimization that wraps dask-backed field data in
chunk-level LRU caches.
vectorized .isel() calls hit an in-memory cache instead of recomputing
dask task graphs, giving large speedups for particle simulations.

This PR adds ChunkCachedArray according to the approach detailed in #2854 , providing an opt-in optimization wrapping dash-backed field data with (dask)chunk-level LRU caches on the array level. This PR:

• Adds a subpackage parcels._chunk_cached_array with the implementation of the Array and the needed LRU cache. Note the cache is on the array level with configurable memory limits per array.
• Exposes a function parcels._chunk_cached_array.wrap_dataset providing the main entrypoint for Parcels for wrapping Xarray dataset objects
• Adds .to_chunk_cached_arrays to FieldSet and ModelData classes.
• Adds equivalence testing (test_backends.py) on a NEMO dataset, asserting that the same results are received regardless of backend.

Still needed:

• (separate PR) Upstream testing
• (separate PR?) Backend equivalency testing on Unstructured grids?
• Benchmarking with realworld simulations
• (pre-merge) Remove temporary testing files (data-generation.py,benchmark_chunk_cache.py,benchmark_chunk_cache.png)

Future work:

• max_cache_bytes on FieldSet.to_chunk_cached_arrays operates per variable (i.e., per Field). This requires the user to think about how many Fields are in their simulation in order to tune their performance. Down the line, we could tune this to allocate a total cache of (e.g.,) 50% of the system memory (which would allow for other memory overhead of coordinates (not stored in the cache), the Python process, or other user processes). This would require significant refactoring not only on how the chunk cache is handled, but also backends in general.

(cc @erikvansebille, @wyatt-fluidnumerics )

Performance

I have some some preliminary profiling with a 20Gb idealised dataset, and found the following performance profile which looks like a promising first step.

I've included data-generation.py and benchmark_chunk_cache.py here for testing in case it helps.

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Checklist

• Contributes to New Parcels backend - ChunkCachedArray #2854
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

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): I guided an LLM to explore the Xarray codebase, and to help work on the implementation of the Array class. I also used an LLM to generate and iterate on the profiling script, and write docstrings.

[erikvansebille]

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Thanks for this exciting PR, @VeckoTheGecko! I plan to do some real-world performance testing (also including/comparing #2846 and v3) later this week. Will report back here when I know more!

[VeckoTheGecko]

Description

Provides an opt-in optimization that wraps dask-backed field data in
chunk-level LRU caches.
vectorized .isel() calls hit an in-memory cache instead of recomputing
dask task graphs, giving large speedups for particle simulations.

This PR adds ChunkCachedArray according to the approach detailed in #2854 , providing an opt-in optimization wrapping dash-backed field data with (dask)chunk-level LRU caches on the array level. This PR:

• Adds a subpackage parcels._chunk_cached_array with the implementation of the Array and the needed LRU cache. Note the cache is on the array level with configurable memory limits per array.
• Exposes a function parcels._chunk_cached_array.wrap_dataset providing the main entrypoint for Parcels for wrapping Xarray dataset objects
• Adds .to_chunk_cached_arrays to FieldSet and ModelData classes.
• Adds equivalence testing (test_backends.py) on a NEMO dataset, asserting that the same results are received regardless of backend.

Still needed:

• (separate PR) Upstream testing
• (separate PR?) Backend equivalency testing on Unstructured grids?
• Benchmarking with realworld simulations
• (pre-merge) Remove temporary testing files (data-generation.py,benchmark_chunk_cache.py,benchmark_chunk_cache.png)

Future work:

• max_cache_bytes on FieldSet.to_chunk_cached_arrays operates per variable (i.e., per Field). This requires the user to think about how many Fields are in their simulation in order to tune their performance. Down the line, we could tune this to allocate a total cache of (e.g.,) 50% of the system memory (which would allow for other memory overhead of coordinates (not stored in the cache), the Python process, or other user processes). This would require significant refactoring not only on how the chunk cache is handled, but also backends in general.

(cc @erikvansebille, @wyatt-fluidnumerics )

Performance

I have some some preliminary profiling with a 20Gb idealised dataset, and found the following performance profile which looks like a promising first step.

I've included data-generation.py and benchmark_chunk_cache.py here for testing in case it helps.

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Checklist

• Contributes to New Parcels backend - ChunkCachedArray #2854
• Tests added
• This PR targets the correct branch (main for normal development, v3-support for v3 support)

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): I guided an LLM to explore the Xarray codebase, and to help work on the implementation of the Array class. I also used an LLM to generate and iterate on the profiling script, and write docstrings.

[erikvansebille]

Its not clear to me the performance with real-world simulations and how that scales (I only have access to my laptop, and found working/debugging on Lorenz quite frustrating with my high latency from travels). Let me know if there's anything I can do to help here @erikvansebille .

Thanks for this exciting PR, @VeckoTheGecko! I plan to do some real-world performance testing (also including/comparing #2846 and v3) later this week. Will report back here when I know more!