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

[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]

hey @chongchonghe , i'm drumming up support for a yt feature release (version 4.5.0) in January, it'd be great to get this merged so it's included in the release! Do you mind fixing the merge conflict and checking out my last batch of comments? they're pretty minor and should be quick.

[chongchonghe]

hey @chongchonghe , i'm drumming up support for a yt feature release (version 4.5.0) in January, it'd be great to get this merged so it's included in the release! Do you mind fixing the merge conflict and checking out my last batch of comments? they're pretty minor and should be quick.

Sure. I'll do it today.

[chongchonghe]

pre-commit.ci autofix

[chongchonghe]

hey @chongchonghe , i'm drumming up support for a yt feature release (version 4.5.0) in January, it'd be great to get this merged so it's included in the release! Do you mind fixing the merge conflict and checking out my last batch of comments? they're pretty minor and should be quick.

hi @chrishavlin . Sorry for the late response. I have addressed the comments and resolved the merge conflicts. Can you review and approve this PR?

[chrishavlin]

great, thanks @chongchonghe ! planning to take a look in the next few days.