Add a frontend for Quokka dataset

[chongchonghe]

PR Summary

This PR add a frontend for loading Quokka dataset. This frontend fully support uniform-grid or AMR datasets in cartesian coordinates and particles. QUOKKA is based on the AMReX framework, so we define a new QuokkaDataset class that inherits from AMReXDataset.

This PR is fully tested on my computer. We have also added Answer Tests in frontends/amrex/test_outputs.py to validate new features with Quokka data that we provided.

We prepared a Jupyter Notebook with examples and test dataset: https://github.com/Rongjun-ANU/README-of-yt-frontend-for-QUOKKA/blob/main/README.ipynb

Authors: @Rongjun-ANU @chongchonghe

Quokka is a two-moment AMR radiation hydrodynamics (with self-gravity, particles, and chemistry) on CPUs/GPUs for astrophysics.

PR Checklist

• New features are documented, with docstrings and narrative docs
• Adds a test for any bugs fixed. Adds tests for new features.

[chongchonghe]

@BenWibking @Rongjun-ANU

[chongchonghe]

I couldn't understand the remaining pre-commit.ci failsures. I would appreciate some explanation or help from the reviewers.

[chongchonghe]

pre-commit.ci autofix

[chongchonghe]

pre-commit.ci autofix

[chrishavlin]

@chongchonghe I should be able to take a detailed look late next week.

one quick question: do you have a sample dataset you can share and upload? I think it'd be particularly helpful to have one that has the face-centered data as well. The way you've set up this frontend to have a nested dataset if the face-centered data exists is intriguing. But it may introduce some confusing bugs -- the yt dataset-related classes rely on many class attributes and loading multiple datasets of the same class at once can lead to errors (e.g., like the one that PR #5159 fixed). So it'd be good if we can thoroughly test that case.

Usually, we try to get sample datasets uploaded to the yt website so that new frontend tests can utilize the data (usually done via a pull request to https://github.com/yt-project/website like this one or this one, but let me check back with you before you do so. The server that the full test suite runs on is having some issues, so I'm not sure we'll follow the usual protocol right now... at the very least it'd be useful if you could share some datasets so that I can download and run the new test(s) locally.

[chongchonghe]

@chrishavlin

This repo contains face-centered Quokka dataset (sample/HydroWave) and a script (test.py) to test that:
https://github.com/Rongjun-ANU/README-of-yt-frontend-for-QUOKKA . Running test.py using the YT from this PR is successful on my computer.

The dataset HydroWave and RadiatingParticles in the repo above is used in the Answer Test for this frontend. However, I don't know how to upload these data to the server for CI tests.

[chongchonghe]

Unfortunately, I couldn't get pytest to work on my computer. I always get this error:

E   ModuleNotFoundError: No module named 'imp'

It seems like the answer_testing in yt is using nose, which relies on imp, which is deprecated. However, I run the new tests I added in this PR on my computer as a python script and it worked, so I'm sure the unit tests will pass once we upload the dataset to the server.

[chrishavlin]

thanks for the dataset! hopefully will have time to try things out later this week.

[chongchonghe]

Hi @chrishavlin , did you have time to try out this PR?

[chrishavlin]

Hi @chongchonghe thanks for the ping, but unfortunately not yet. Had to bump it to this week's to-do list .

[chrishavlin]

So I'm still working through the code, but wanted to bring up a high-level question.

Normally, when frontends introduce dependencies, those dependencies should be added only to the frontend that requires them. But right now you're making PyYaml a dependency of all the amrex subclass frontends. It's a pretty small dependency, so I'm inclined to let it go in this instance, especially since it's already a dependency for a full yt install (via pyaml)... @neutrinoceros do you have an opinion on this?

For reference, here's a diff that would make PyYaml a dependency only for Quokka (adjusting pyproject.toml, moving the imports to the functions where they're used):

Diff is here

diff --git a/pyproject.toml b/pyproject.toml
index 03300c836..16a5f9b41 100644
--- a/pyproject.toml
+++ b/pyproject.toml
@@ -89,7 +89,7 @@ Fortran = ["f90nml>=1.1"]
 # We also normalize all target names to lower case for consistency.
 adaptahop = []
 ahf = []
-amrex = ["PyYAML>=6.0.1"]
+amrex = []
 amrvac = ["yt[Fortran]"]
 art = []
 arepo = ["yt[HDF5]"]
@@ -121,7 +121,7 @@ open-pmd = ["yt[HDF5]"]
 owls = ["yt[HDF5]"]
 owls-subfind = ["yt[HDF5]"]
 parthenon = ["yt[HDF5]"]
-quokka = ["PyYAML>=6.0.1"]
+quokka = ["yt[amrex]", "PyYAML>=6.0.1"]
 ramses = ["yt[Fortran]", "scipy"]
 rockstar = []
 sdf = ["requests>=2.20.0"]
diff --git a/yt/frontends/amrex/data_structures.py b/yt/frontends/amrex/data_structures.py
index 7a27f4efb..4e584de6b 100644
--- a/yt/frontends/amrex/data_structures.py
+++ b/yt/frontends/amrex/data_structures.py
@@ -6,7 +6,6 @@ from functools import cached_property
 from stat import ST_CTIME
 
 import numpy as np
-import yaml
 
 from yt.data_objects.index_subobjects.grid_patch import AMRGridPatch
 from yt.data_objects.static_output import Dataset
@@ -1373,6 +1372,7 @@ class QuokkaDataset(AMReXDataset):
     _field_info_class = QuokkaFieldInfo
     _subtype_keyword = ""
     _default_cparam_filename = "metadata.yaml"
+    _load_requirements = ["yaml"]
 
     def __init__(
         self,
@@ -1465,6 +1465,8 @@ class QuokkaDataset(AMReXDataset):
                 mylog.debug(f"No face-centered {direction} dataset found at {fc_dir}")
 
     def _parse_parameter_file(self):
+        import yaml
+
         # Call parent method to initialize core setup by yt
         super()._parse_parameter_file()
 
@@ -1686,6 +1688,8 @@ class QuokkaDataset(AMReXDataset):
         mylog.debug("Parsed header parameters: %s", self.parameters)
 
     def _parse_metadata_file(self):
+        import yaml
+
         # Construct the full path to the metadata file
         metadata_filename = os.path.join(self.output_dir, self.cparam_filename)
         try:

[neutrinoceros]

@neutrinoceros do you have an opinion on this?

Well, I think I'd have been more lenient with letting this one through... if you hadn't already done the work to separate dependencies more clearly. Now I'm in favor of committing your patch :)

[chrishavlin]

Now I'm in favor of committing your patch :)

Hah, ok, then let's go with that. @chongchonghe , could you update the code so that pyyaml is only needed for quokka (check the diff in my previous comment)?

[chongchonghe]

@chrishavlin I've implemented the changes as you suggested. Thanks a lot!

I have also been trying to optimize how we deal with face-centering variables. I came to the conclusion that the current approach is the optimal approach. QUOKKA does not currently support writing FC variables as 'raw' data as WarpX does. So, I did a hack by loading the face-centered quantities as if they are cell-centered. Raw and "nodal" fields will be supported in the future.

Sorry for my delayed response. I've been overwhelmed with a proposal recently. Please let me know if any revisions are required, or if you need any additional information from me before we can proceed with merging this PR.

[chrishavlin]

Great, thanks! I'll do another read through and will likely have some code suggestions, but probably won't be anything major.

[neutrinoceros]

I am not currently available to review a large PR as this one, I apologize for the inconvenience.

[chongchonghe]

@chrishavlin Just want to pin you for a review on this PR.

[chrishavlin]

Thanks for your patience @chongchonghe ! I have some minor suggestions here to clean up some of the new code. I'll also see if anyone else is available for a quick review since we like to have at least 2 reviews for new frontends: but I recognize this has been around for quite a while and it'd be nice to get it merged ASAP.

And I did download the test data and confirmed that the new tests run successfully. I'll make a companion PR to the yt website repo with the data in a convenient form so it ends up where it needs to eventually, but note that the tests won't actually run for now while yt's continuous integration server is down.

[chrishavlin]

Given how you have https://github.com/Rongjun-ANU/README-of-yt-frontend-for-QUOKKA/tree/main set up, it will be easier to just use paths here:

Suggested change

	quokka = "quokka_RadiatingParticles_plt026"
	quokka = "quokka/RadiatingParticles/plt026"

Then the sample folder from https://github.com/Rongjun-ANU/README-of-yt-frontend-for-QUOKKA/tree/main can be copied to a quokka directory within yt's test data directory (rather than renaming, e.g., sample/RadiatingParticles/plt026 to sample_RadiatingParticles_plt026).

[chrishavlin]

see previous comment.

Suggested change

	quokka_face_centered = "quokka_HydroWave_plt00004"
	quokka_face_centered = "quokka/HydroWave/plt00004"

[chrishavlin]

Both of the test files you provided contain a quokka_version entry in the metadata yaml. If this is expected to be true of all quokka files, you should change this line:

Suggested change

	_subtype_keyword = ""
	_subtype_keyword = "quokka"

I confirmed that your test datasets still load.

[chrishavlin]

after getting through the rest of the code, I saw that there may be cases where quokka_version is not defined... so you can ignore this suggestion. It'd be preferable if there were a way to identify quokka files beyond just the presence of a Header and metadata.yaml file, but if that's not possible it's OK.

[chrishavlin]

More conventional to just use _ here:

Suggested change

	for _i in range(num_fields):
	for _ in range(num_fields):

[chrishavlin]

Was there a particular reason for using update here? It's a tad roundabout to create a dict just to use update when the following will work just as well:

Suggested change

	self.parameters.update(
	{
	"particles": len(detected_particle_types),
	"particle_types": tuple(detected_particle_types),
	"particle_info": particle_info,
	}
	)
	self.parameters["particles"] = len(detected_particle_types)
	self.parameters["particle_types"] = tuple(detected_particle_types),
	self.parameters["particle_info"] = particle_info

[chrishavlin]

yt has packaging as a minimal dependency, you can use it here to simplify this logic:

Suggested change

	quokka_version = metadata.get("quokka_version", "0.0")
	# Define version comparison function
	def is_newer_or_equal(ver, compare_to):
	return tuple(map(int, str(ver).split("."))) >= tuple(
	map(int, str(compare_to).split("."))
	)
	# Compare with version 25.03 (year 2025, month 3)
	if not is_newer_or_equal(quokka_version, "25.03"):
	from packaging.version import Version
	quokka_version = Version(metadata.get("quokka_version", "0.0"))
	if quokka_version < Version("25.03"):

[chrishavlin]

(you can move that from packaging.version import Version import to the top of this file: I put it in the suggestion here for clarity)

[chrishavlin]

Might as well put the actual version in here

Suggested change

	mylog.debug("Metadata loaded successfully (version 25.03 or newer)")
	mylog.debug(f"Metadata loaded successfully (version 25.03 or newer): {quokka_version}")