Repository: yt-project/yt

Issue:
BUG/ENH: create_firefly_object only allow fields that are common to all particle types
<!--To help us understand and resolve your issue, please fill out the form to
the best of your ability.-->
<!--You can feel free to delete the sections that do not apply.-->

### Bug report

**Bug summary**

<!--A short 1-2 sentences that succinctly describes the bug-->
The documentation of `fields_to_include` in the `data_containers.create_firefly_object()` method does not specify that only fields that are common to all particle types are allowed. This should either be mentioned in the documentation (Option 1) or the behavior changed to allow fields that are not common to all particle types (Option 2, preferred). (Open to other possibilities though)

**Code for reproduction**

<!--A minimum code snippet required to reproduce the bug, also minimizing the
number of dependencies required.-->

<!-- If you need to use a data file to trigger the issue you're having, consider
using one of the datasets from the yt data hub (http://yt-project.org/data). If
your issue cannot be triggered using a public dataset, you can use the yt
curldrop (https://docs.hub.yt/services.html#curldrop) to share data
files. Please include a link to the dataset in the issue if you use the
curldrop.-->

```python
# This uses the Gizmo gizmo_mhd_mwdisk sample dataset instead of the Ramses one in the
# "Exporting an Example Dataset to Firefly", but is otherwise mostly equivalent to that example.
# I've changed the names (ramses -> gizmo) and included the bounding box/unit base to 
# ensure the data is loaded correctly. This dataset includes six particle types: PartType0, which
# has the field "Temperature" and PartType1-5, which do not.
bbox = [[-600.0, 600.0], [-600.0, 600.0], [-600.0, 600.0]]
unit_base = {
    "length": (1.0, "kpc"),
    "velocity": (1.0, "km/s"),
    "mass": (1.0, "Msun"),
}
ge = yt.load("sample_data/gizmo_mhd_mwdisk/gizmo_mhd_mwdisk.hdf5",bounding_box=bbox,unit_base=unit_base)
region = ge.sphere(ge.domain_center, (1000, "kpc"))
reader = region.create_firefly_object(
       "GizmoExample",
       fields_to_include=["Temperature"],  # Temperature is only defined for PartType0
       fields_units=["code_temperature"],  
)

## adjust some of the options
reader.settings["color"]["io"] = [1, 1, 0, 1]  ## set default color
reader.particleGroups[0].decimation_factor = 100  ## increase the decimation factor

reader.writeToDisk()
```

**Actual outcome**

<!--The output produced by the above code, which may be a screenshot, console
output, etc.-->
This produces the error:

```
YTFieldNotFound: Could not find field ('PartType1', 'Temperature') in gizmo_mhd_mwdisk.
Did you mean:
	('PartType0', 'Temperature')
	('PartType0', 'temperature')
	('gas', 'temperature')
```
And nothing is written to the disk.

**Expected outcome**

<!--A description of the expected outcome from the code snippet-->
<!--If this used to work in an earlier version of yt, please note the
version it used to work on-->
The expected behavior depends a bit on what option is chosen. 
- For the documentation change (Option 1), I guess I would expect a more useful error, perhaps: 
```
YTFieldNotFound: Could not find field ('PartType1', 'Temperature') in gizmo_mhd_mwdisk.
Please include only fields that are common to all particle types.
Did you mean:
	('PartType0', 'Temperature')
	('PartType0', 'temperature')
	('gas', 'temperature')
```
with nothing written to the disk.
 - For Option 2, skipping the field if not present for that particle type, I would expect something like:
```
datadir: /XXXX/GizmoExample -- is not a sub-directory of firefly/static/data. 
This may produce confusing or inoperable results. As such, we will create a symlink for you when you writeToDisk.
Make sure each field_array (1) has a field_radius_flag (0), assuming False.
Make sure each field_array (1) has a field_radius_flag (0), assuming False.
Make sure each field_array (1) has a field_radius_flag (0), assuming False.
Make sure each field_array (1) has a field_radius_flag (0), assuming False.
Make sure each field_array (1) has a field_radius_flag (0), assuming False.
Make sure each field_array (1) has a field_radius_flag (0), assuming False.
Warning: Not all particle types contained field(s): Temperature                            <--- New
PartType0 - 4091/409013 particles - 3 tracked fields
PartType1 - 4375/437500 particles - 2 tracked fields
PartType2 - 2500/250000 particles - 2 tracked fields
PartType3 - 1250/125000 particles - 2 tracked fields
PartType4 - 766/76546 particles - 2 tracked fields
PartType5 - 1/1 particles - 2 tracked fields
```
And in Firefly, the ability to set Temperature as a selectable field.

**Version Information**
<!--Please specify your platform and versions of the relevant libraries you are
using:-->
  * yt version: 4.1.1
  * Other Libraries (if applicable): firefly v 3.2.0

<!--Please tell us how you installed yt and python e.g., from source,
pip, conda. If you installed from conda, please specify which channel you used
if not the default-->
yt installed from source 
python from anaconda
firefly from PyPI


Discussion:
Hi, and welcome to yt!  Thanks for opening your first issue.  We have an issue template that helps us to gather relevant information to help diagnosing and fixing the issue.

I'm happy to work on this, I just don't know which way would be preferred aka Option 1 vs 2 (vs 3?).
I think fields that aren't available in every particle types should be allowed, as long as there's a clear warning about it.

Resolving patch:
diff --git a/yt/data_objects/data_containers.py b/yt/data_objects/data_containers.py
--- a/yt/data_objects/data_containers.py
+++ b/yt/data_objects/data_containers.py
@@ -714,6 +714,7 @@ def create_firefly_object(
         show_unused_fields=0,
         *,
         JSONdir=None,
+        match_any_particle_types=True,
         **kwargs,
     ):
         r"""This function links a region of data stored in a yt dataset
@@ -727,7 +728,7 @@ def create_firefly_object(
             Path to where any `.json` files should be saved. If a relative
             path will assume relative to `${HOME}`. A value of `None` will default to `${HOME}/Data`.
 
-        fields_to_include : array_like of strings
+        fields_to_include : array_like of strings or field tuples
             A list of fields that you want to include in your
             Firefly visualization for on-the-fly filtering and
             colormapping.
@@ -754,6 +755,12 @@ def create_firefly_object(
             A flag to optionally print the fields that are available, in the
             dataset but were not explicitly requested to be tracked.
 
+        match_any_particle_types : boolean
+            If True, when any of the fields_to_include match multiple particle
+            groups then the field will be added for all matching particle
+            groups. If False, an error is raised when encountering an ambiguous
+            field. Default is True.
+
         Any additional keyword arguments are passed to
         firefly.data_reader.Reader.__init__
 
@@ -813,6 +820,45 @@ def create_firefly_object(
             datadir=datadir, clean_datadir=True, **kwargs
         )
 
+        ## Ensure at least one field type contains every field requested
+        if match_any_particle_types:
+            # Need to keep previous behavior: single string field names that
+            # are ambiguous should bring in any matching ParticleGroups instead
+            # of raising an error
+            # This can be expanded/changed in the future to include field
+            # tuples containing some sort of special "any" ParticleGroup
+            unambiguous_fields_to_include = []
+            unambiguous_fields_units = []
+            for field, field_unit in zip(fields_to_include, fields_units):
+                if isinstance(field, tuple):
+                    # skip tuples, they'll be checked with _determine_fields
+                    unambiguous_fields_to_include.append(field)
+                    unambiguous_fields_units.append(field_unit)
+                    continue
+                _, candidates = self.ds._get_field_info_helper(field)
+                if len(candidates) == 1:
+                    # Field is unambiguous, add in tuple form
+                    # This should be equivalent to _tupleize_field
+                    unambiguous_fields_to_include.append(candidates[0])
+                    unambiguous_fields_units.append(field_unit)
+                else:
+                    # Field has multiple candidates, add all of them instead
+                    # of original field. Note this may bring in aliases and
+                    # equivalent particle fields
+                    for c in candidates:
+                        unambiguous_fields_to_include.append(c)
+                        unambiguous_fields_units.append(field_unit)
+            fields_to_include = unambiguous_fields_to_include
+            fields_units = unambiguous_fields_units
+        # error if any requested field is unknown or (still) ambiguous
+        # This is also sufficient if match_any_particle_types=False
+        fields_to_include = self._determine_fields(fields_to_include)
+        ## Also generate equivalent of particle_fields_by_type including
+        ## derived fields
+        kysd = defaultdict(list)
+        for k, v in self.ds.derived_field_list:
+            kysd[k].append(v)
+
         ## create a ParticleGroup object that contains *every* field
         for ptype in sorted(self.ds.particle_types_raw):
             ## skip this particle type if it has no particles in this dataset
@@ -836,17 +882,23 @@ def create_firefly_object(
 
             ## explicitly go after the fields we want
             for field, units in zip(fields_to_include, fields_units):
+                ## Only interested in fields with the current particle type,
+                ## whether that means general fields or field tuples
+                ftype, fname = field
+                if ftype not in (ptype, "all"):
+                    continue
+
                 ## determine if you want to take the log of the field for Firefly
                 log_flag = "log(" in units
 
                 ## read the field array from the dataset
-                this_field_array = self[ptype, field]
+                this_field_array = self[ptype, fname]
 
                 ## fix the units string and prepend 'log' to the field for
                 ##  the UI name
                 if log_flag:
                     units = units[len("log(") : -1]
-                    field = f"log{field}"
+                    fname = f"log{fname}"
 
                 ## perform the unit conversion and take the log if
                 ##  necessary.
@@ -856,7 +908,7 @@ def create_firefly_object(
 
                 ## add this array to the tracked arrays
                 field_arrays += [this_field_array]
-                field_names = np.append(field_names, [field], axis=0)
+                field_names = np.append(field_names, [fname], axis=0)
 
             ## flag whether we want to filter and/or color by these fields
             ##  we'll assume yes for both cases, this can be changed after
@@ -864,15 +916,32 @@ def create_firefly_object(
             field_filter_flags = np.ones(len(field_names))
             field_colormap_flags = np.ones(len(field_names))
 
+            ## field_* needs to be explicitly set None if empty
+            ## so that Firefly will correctly compute the binary
+            ## headers
+            if len(field_arrays) == 0:
+                if len(fields_to_include) > 0:
+                    mylog.warning("No additional fields specified for %s", ptype)
+                field_arrays = None
+                field_names = None
+                field_filter_flags = None
+                field_colormap_flags = None
+
+            ## Check if particles have velocities
+            if "relative_particle_velocity" in kysd[ptype]:
+                velocities = self[ptype, "relative_particle_velocity"].in_units(
+                    velocity_units
+                )
+            else:
+                velocities = None
+
             ## create a firefly ParticleGroup for this particle type
             pg = firefly.data_reader.ParticleGroup(
                 UIname=ptype,
                 coordinates=self[ptype, "relative_particle_position"].in_units(
                     coordinate_units
                 ),
-                velocities=self[ptype, "relative_particle_velocity"].in_units(
-                    velocity_units
-                ),
+                velocities=velocities,
                 field_arrays=field_arrays,
                 field_names=field_names,
                 field_filter_flags=field_filter_flags,


Test patch:
diff --git a/tests/tests.yaml b/tests/tests.yaml
--- a/tests/tests.yaml
+++ b/tests/tests.yaml
@@ -228,6 +228,7 @@ other_tests:
      - "--ignore-file=test_add_field\\.py"
      - "--ignore-file=test_glue\\.py"
      - "--ignore-file=test_geometries\\.py"
+     - "--ignore-file=test_firefly\\.py"
      - "--ignore-file=test_callable_grids\\.py"
      - "--exclude-test=yt.frontends.gdf.tests.test_outputs.TestGDF"
      - "--exclude-test=yt.frontends.adaptahop.tests.test_outputs"
diff --git a/yt/data_objects/tests/test_firefly.py b/yt/data_objects/tests/test_firefly.py
--- a/yt/data_objects/tests/test_firefly.py
+++ b/yt/data_objects/tests/test_firefly.py
@@ -1,9 +1,12 @@
-import tempfile
+import numpy as np
+import pytest
+from numpy.testing import assert_array_equal
 
 from yt.testing import fake_particle_ds, requires_module
+from yt.utilities.exceptions import YTFieldNotFound
 
 
-@requires_module("Firefly")
+@requires_module("firefly")
 def test_firefly_JSON_string():
 
     ds = fake_particle_ds()
@@ -14,7 +17,7 @@ def test_firefly_JSON_string():
         coordinate_units="cm",
     )
 
-    reader.dumpToJSON(write_jsons_to_disk=False)
+    reader.writeToDisk(write_to_disk=False, file_extension=".json")
 
     ## reader.JSON was not output to string correctly
     ##  either Firefly is damaged or needs a hotfix-- try reinstalling.
@@ -23,9 +26,9 @@ def test_firefly_JSON_string():
     assert len(reader.JSON) > 0
 
 
-@requires_module("Firefly")
-def test_firefly_write_to_disk():
-    tmpdir = tempfile.mkdtemp()
+@requires_module("firefly")
+def test_firefly_write_to_disk(tmp_path):
+    tmpdir = str(tmp_path)  # create_firefly_object needs a str, not PosixPath
 
     ds = fake_particle_ds()
     ad = ds.all_data()
@@ -33,6 +36,246 @@ def test_firefly_write_to_disk():
         tmpdir,
         velocity_units="cm/s",
         coordinate_units="cm",
+        match_any_particle_types=True,  # Explicitly specifying to avoid deprecation warning
     )
 
-    reader.dumpToJSON()
+    reader.writeToDisk()
+
+
+@pytest.fixture
+def firefly_test_dataset():
+    # create dataset
+    ds_fields = [
+        # Assumed present
+        ("pt1", "particle_position_x"),
+        ("pt1", "particle_position_y"),
+        ("pt1", "particle_position_z"),
+        ("pt2", "particle_position_x"),
+        ("pt2", "particle_position_y"),
+        ("pt2", "particle_position_z"),
+        # User input
+        ("pt1", "common_field"),
+        ("pt2", "common_field"),
+        ("pt2", "pt2only_field"),
+    ]
+    ds_field_units = ["code_length"] * 9
+    ds_negative = [0] * 9
+    ds = fake_particle_ds(
+        fields=ds_fields,
+        units=ds_field_units,
+        negative=ds_negative,
+    )
+    return ds
+
+
+@requires_module("firefly")
+@pytest.mark.parametrize(
+    "fields_to_include,fields_units",
+    [
+        (None, None),  # Test default values
+        ([], []),  # Test empty fields
+    ],
+)
+def test_field_empty_specification(
+    firefly_test_dataset, fields_to_include, fields_units
+):
+    dd = firefly_test_dataset.all_data()
+    reader = dd.create_firefly_object(
+        fields_to_include=fields_to_include,
+        fields_units=fields_units,
+        coordinate_units="code_length",
+    )
+    assert_array_equal(
+        dd[("pt1", "relative_particle_position")].d,
+        reader.particleGroups[0].coordinates,
+    )
+    assert_array_equal(
+        dd[("pt2", "relative_particle_position")].d,
+        reader.particleGroups[1].coordinates,
+    )
+
+
+@requires_module("firefly")
+def test_field_unique_string_specification(firefly_test_dataset):
+    # Test unique field (pt2only_field)
+    dd = firefly_test_dataset.all_data()
+    # Unique field string will fallback to "all" field type and fail
+    # as nonexistent ("all", "pt2only_field") unless we set
+    # match_any_particle_types=True
+    reader = dd.create_firefly_object(
+        fields_to_include=["pt2only_field"],
+        fields_units=["code_length"],
+        coordinate_units="code_length",
+        match_any_particle_types=True,
+    )
+
+    pt1 = reader.particleGroups[0]
+    pt2 = reader.particleGroups[1]
+    assert_array_equal(
+        dd[("pt1", "relative_particle_position")].d,
+        pt1.coordinates,
+    )
+    assert_array_equal(
+        dd[("pt2", "relative_particle_position")].d,
+        pt2.coordinates,
+    )
+    assert "pt2only_field" not in pt1.field_names
+    assert "pt2only_field" in pt2.field_names
+    arrind = np.flatnonzero(pt2.field_names == "pt2only_field")[0]
+    assert_array_equal(dd[("pt2", "pt2only_field")].d, pt2.field_arrays[arrind])
+
+
+@requires_module("firefly")
+def test_field_common_string_specification(firefly_test_dataset):
+    # Test common field (common_field)
+    dd = firefly_test_dataset.all_data()
+    # Common field string will be ambiguous and fail
+    # unless we set match_any_particle_types=True
+    reader = dd.create_firefly_object(
+        fields_to_include=["common_field"],
+        fields_units=["code_length"],
+        coordinate_units="code_length",
+        match_any_particle_types=True,
+    )
+
+    pt1 = reader.particleGroups[0]
+    pt2 = reader.particleGroups[1]
+    assert_array_equal(
+        dd[("pt1", "relative_particle_position")].d,
+        pt1.coordinates,
+    )
+    assert_array_equal(
+        dd[("pt2", "relative_particle_position")].d,
+        pt2.coordinates,
+    )
+    assert "common_field" in pt1.field_names
+    assert "common_field" in pt2.field_names
+    arrind = np.flatnonzero(pt1.field_names == "common_field")[0]
+    assert_array_equal(dd[("pt1", "common_field")].d, pt1.field_arrays[arrind])
+    arrind = np.flatnonzero(pt2.field_names == "common_field")[0]
+    assert_array_equal(dd[("pt2", "common_field")].d, pt2.field_arrays[arrind])
+
+
+@requires_module("firefly")
+@pytest.mark.parametrize(
+    "fields_to_include,fields_units",
+    [
+        (
+            [("pt2", "pt2only_field")],
+            ["code_length"],
+        ),  # Test existing field tuple (pt2, pt2only_field)
+        (
+            [("pt1", "common_field")],
+            ["code_length"],
+        ),  # Test that tuples only bring in referenced particleGroup
+        (
+            [("all", "common_field")],
+            ["code_length"],
+        ),  # Test that "all" brings in all particleGroups
+    ],
+)
+def test_field_tuple_specification(
+    firefly_test_dataset,
+    fields_to_include,
+    fields_units,
+):
+    dd = firefly_test_dataset.all_data()
+    reader = dd.create_firefly_object(
+        fields_to_include=fields_to_include,
+        fields_units=fields_units,
+        coordinate_units="code_length",
+    )
+    assert_array_equal(
+        dd[("pt1", "relative_particle_position")].d,
+        reader.particleGroups[0].coordinates,
+    )
+    assert_array_equal(
+        dd[("pt2", "relative_particle_position")].d,
+        reader.particleGroups[1].coordinates,
+    )
+    all_pgs = reader.particleGroups
+    all_pgs_names = ["pt1", "pt2"]
+    for field in fields_to_include:
+        ftype, fname = field
+        for pgi in range(2):
+            pg = all_pgs[pgi]
+            if ftype == all_pgs_names[pgi]:
+                assert fname in pg.field_names
+                arrind = np.flatnonzero(pg.field_names == fname)[0]
+                assert_array_equal(dd[field].d, pg.field_arrays[arrind])
+            elif ftype == "all":
+                assert fname in pg.field_names
+                this_pg_name = all_pgs_names[pgi]
+                arrind = np.flatnonzero(pg.field_names == fname)[0]
+                assert_array_equal(dd[this_pg_name, fname].d, pg.field_arrays[arrind])
+            else:
+                assert fname not in pg.field_names
+
+
+@requires_module("firefly")
+@pytest.mark.parametrize(
+    "fields_to_include,fields_units,ErrorType",
+    [
+        (
+            ["dinos"],
+            ["code_length"],
+            YTFieldNotFound,
+        ),  # Test nonexistent field (dinos)
+        (
+            ["pt2only_field"],
+            ["code_length"],
+            YTFieldNotFound,
+        ),  # Test unique field (match_any_particle_types=False)
+        (
+            ["common_field"],
+            ["code_length"],
+            ValueError,
+        ),  # Test ambiguous field (match_any_particle_types=False)
+        (
+            [("pt1", "pt2only_field")],
+            ["code_length"],
+            YTFieldNotFound,
+        ),  # Test nonexistent field tuple (pt1, pt2only_field)
+        (
+            ["pt2only_field", ("pt1", "common_field")],
+            ["code_length", "code_length"],
+            YTFieldNotFound,
+        ),  # Test mixed field spec (match_any_particle_types=False)
+    ],
+)
+def test_field_invalid_specification(
+    firefly_test_dataset, fields_to_include, fields_units, ErrorType
+):
+
+    dd = firefly_test_dataset.all_data()
+    # Note that we have specified match_any_particle_types as False since
+    # that is the behavior expected in the future
+    with pytest.raises(ErrorType):
+        dd.create_firefly_object(
+            fields_to_include=fields_to_include,
+            fields_units=fields_units,
+            coordinate_units="code_length",
+            match_any_particle_types=False,
+        )
+
+
+@requires_module("firefly")
+def test_field_mixed_specification(firefly_test_dataset):
+    dd = firefly_test_dataset.all_data()
+
+    reader = dd.create_firefly_object(
+        fields_to_include=["pt2only_field", ("pt1", "common_field")],
+        fields_units=["code_length", "code_length"],
+    )
+
+    pt1 = reader.particleGroups[0]
+    pt2 = reader.particleGroups[1]
+    assert "common_field" in pt1.field_names
+    assert "common_field" not in pt2.field_names
+    arrind = np.flatnonzero(pt1.field_names == "common_field")[0]
+    assert_array_equal(dd[("pt1", "common_field")].d, pt1.field_arrays[arrind])
+
+    assert "pt2only_field" not in pt1.field_names
+    assert "pt2only_field" in pt2.field_names
+    arrind = np.flatnonzero(pt2.field_names == "pt2only_field")[0]
+    assert_array_equal(dd[("pt2", "pt2only_field")].d, pt2.field_arrays[arrind])