Add register_fields method to register or override on-disk field properties.

Author: brittonsmith

yt/data_objects/static_output.py

@@ -180,6 +180,7 @@ class Dataset(abc.ABC):
     _ionization_label_format = "roman_numeral"
     _determined_fields: dict[str, list[FieldKey]] | None = None
     fields_detected = False
+    _registered_fields = None
 
     # these are set in self._parse_parameter_file()
     domain_left_edge = MutableAttribute(True)
@@ -1716,6 +1717,44 @@ def quan(self):
         self._quan = functools.partial(YTQuantity, registry=self.unit_registry)
         return self._quan
 
+    def register_field(
+        self, name, units=None, aliases=None, display_name=None
+    ):
+        """
+        Register properties for an on-disk field.
+
+        Register or override units, aliases, or the display name for an
+        on-disk field.
+
+        Note, this must be called immediately after yt.load (i.e., before
+        the list of fields is generated).
+
+        Parameters
+        ----------
+
+        name : str
+           name of the field
+        units : str
+           units of the field
+        aliases : list
+           a list of alias names
+        display_name: str
+           the name used in plots
+
+        """
+        if self._registered_fields is None:
+            self._registered_fields = {}
+
+        entry = {}
+        if units is not None:
+            entry["units"] = units
+        if aliases is not None:
+            entry["alias"] = alias
+        if display_name is not None:
+            entry["display_name"] = display_name
+
+        self._registered_fields[name] = entry
+
     def add_field(
         self, name, function, sampling_type, *, force_override=False, **kwargs
     ):

yt/fields/field_info_container.py

@@ -150,7 +150,12 @@ def setup_particle_fields(self, ptype, ftype="gas", num_neighbors=64):
                 raise RuntimeError
             if field[0] not in self.ds.particle_types:
                 continue
+
             units = self.ds.field_units.get(field, None)
+            rfields = getattr(self.ds, "_registered_fields", {})
+            if entry := rfields.get(field):
+                units = entry.get("units", units)
+
             if units is None:
                 try:
                     units = ytcfg.get("fields", *field, "units")
@@ -256,6 +261,7 @@ def setup_fluid_aliases(self, ftype: FieldType = "gas") -> None:
                 raise RuntimeError
             if field[0] in self.ds.particle_types:
                 continue
+
             args = known_other_fields.get(field[1], None)
             if args is not None:
                 units, aliases, display_name = args
@@ -273,6 +279,14 @@ def setup_fluid_aliases(self, ftype: FieldType = "gas") -> None:
             # field *name* is in there, then the field *tuple*.
             units = self.ds.field_units.get(field[1], units)
             units = self.ds.field_units.get(field, units)
+
+            # allow user to override with call to ds.register_fields
+            rfields = getattr(self.ds, "_registered_fields", {})
+            if entry := rfields.get(field):
+                units = entry.get("units", units)
+                aliases = entry.get("aliases", aliases)
+                display_name = entry.get("display_name", display_name)
+
             self.add_output_field(
                 field, sampling_type="cell", units=units, display_name=display_name
             )