Merge pull request #5001 from jzuhone/particle_vlos

ENH: Support line-of-sight fields for particle types

Author: neutrinoceros

doc/source/analyzing/fields.rst

@@ -702,16 +702,20 @@ you want to examine a line-of-sight vector within a 3-D data object, set the
     # Set to a three-vector for an off-axis component
     dd.set_field_parameter("axis", [0.3, 0.4, -0.7])
     print(dd["gas", "velocity_los"])
+    # particle fields are supported too!
+    print(dd["all", "particle_velocity_los"])
 
 .. warning::
 
     If you need to change the axis of the line of sight on the *same* data container
-    (sphere, box, cylinder, or whatever), you will need to delete the field using
+    (sphere, box, cylinder, or whatever), you will need to delete the field using (e.g.)
     ``del dd['velocity_los']`` and re-generate it.
 
 At this time, this functionality is enabled for the velocity and magnetic vector
-fields, ``('gas', 'velocity_los')`` and ``('gas', 'magnetic_field_los')``. The
-following fields built into yt make use of these line-of-sight fields:
+fields, ``('gas', 'velocity_los')`` and ``('gas', 'magnetic_field_los')`` for
+the ``"gas"`` field type, as well as every particle type with
+a velocity field, e.g. ``("all", "particle_velocity_los")``. The following fields
+built into yt make use of these line-of-sight fields:
 
 * ``('gas', 'sz_kinetic')`` uses ``('gas', 'velocity_los')``
 * ``('gas', 'rotation_measure')`` uses ``('gas', 'magnetic_field_los')``

yt/data_objects/tests/test_data_containers.py

@@ -160,6 +160,7 @@ def test_derived_field(self):
         # their parent field to be created
         ds = fake_particle_ds()
         dd = ds.all_data()
+        dd.set_field_parameter("axis", 0)
 
         @particle_filter(requires=["particle_mass"], filtered_type="io")
         def massive(pfilter, data):

yt/fields/particle_fields.py

@@ -22,7 +22,7 @@
 )
 
 from .field_functions import get_radius
-from .vector_operations import create_magnitude_field
+from .vector_operations import create_los_field, create_magnitude_field
 
 sph_whitelist_fields = (
     "density",
@@ -329,6 +329,14 @@ def _particle_velocity_magnitude(field, data):
         units=unit_system["velocity"],
     )
 
+    create_los_field(
+        registry,
+        "particle_velocity",
+        unit_system["velocity"],
+        ftype=ptype,
+        sampling_type="particle",
+    )
+
     def _particle_specific_angular_momentum(field, data):
         """Calculate the angular of a particle velocity.
 

yt/fields/tests/test_particle_fields.py

@@ -1,3 +1,5 @@
+import numpy as np
+
 from yt.testing import assert_allclose_units, requires_file, requires_module
 from yt.utilities.answer_testing.framework import data_dir_load
 
@@ -19,3 +21,28 @@ def test_relative_particle_fields():
     assert_allclose_units(
         sp["all", "relative_particle_velocity"], sp["all", "particle_velocity"] - bv
     )
+
+
+@requires_module("h5py")
+@requires_file(g30)
+def test_los_particle_fields():
+    ds = data_dir_load(g30)
+    offset = ds.arr([0.1, -0.2, 0.3], "code_length")
+    c = ds.domain_center + offset
+    sp = ds.sphere(c, (10, "kpc"))
+    bv = ds.arr([1.0, 2.0, 3.0], "code_velocity")
+    sp.set_field_parameter("bulk_velocity", bv)
+    ax = [0.1, 0.2, -0.3]
+    sp.set_field_parameter("axis", ax)
+    ax /= np.linalg.norm(ax)
+    vlos = (
+        sp["all", "relative_particle_velocity_x"] * ax[0]
+        + sp["all", "relative_particle_velocity_y"] * ax[1]
+        + sp["all", "relative_particle_velocity_z"] * ax[2]
+    )
+    assert_allclose_units(sp["all", "particle_velocity_los"], vlos)
+    sp.clear_data()
+    ax = 2
+    sp.set_field_parameter("axis", ax)
+    vlos = sp["all", "relative_particle_velocity_z"]
+    assert_allclose_units(sp["all", "particle_velocity_los"], vlos)

yt/fields/vector_operations.py

@@ -104,18 +104,34 @@ def _relative_vector(field, data):
         )
 
 
-def create_los_field(registry, basename, field_units, ftype="gas", slice_info=None):
+def create_los_field(
+    registry,
+    basename,
+    field_units,
+    ftype="gas",
+    slice_info=None,
+    *,
+    sampling_type="local",
+):
     axis_order = registry.ds.coordinates.axis_order
 
+    # Here we need to check if we are a particle field, so that we can
+    # correctly identify the "bulk" field parameter corresponding to
+    # this vector field.
+    if sampling_type == "particle":
+        basenm = basename.removeprefix("particle_")
+    else:
+        basenm = basename
+
     validators = [
-        ValidateParameter(f"bulk_{basename}"),
+        ValidateParameter(f"bulk_{basenm}"),
         ValidateParameter("axis", {"axis": [0, 1, 2]}),
     ]
 
     field_comps = [(ftype, f"{basename}_{ax}") for ax in axis_order]
 
     def _los_field(field, data):
-        if data.has_field_parameter(f"bulk_{basename}"):
+        if data.has_field_parameter(f"bulk_{basenm}"):
             fns = [(fc[0], f"relative_{fc[1]}") for fc in field_comps]
         else:
             fns = field_comps
@@ -132,7 +148,7 @@ def _los_field(field, data):
 
     registry.add_field(
         (ftype, f"{basename}_los"),
-        sampling_type="local",
+        sampling_type=sampling_type,
         function=_los_field,
         units=field_units,
         validators=validators,