Infer units for derived fields when quickly defining them

Author: cphyc

yt/fields/derived_field.py

@@ -189,6 +189,42 @@ def getDependentFields(self):
     def name(self):
         return f"{self!r}"
 
+    @property
+    def units(self):
+        def helper(term):
+            # Get a Unit object, using the unit registry if possible
+            if hasattr(term, "ds"):
+                registry = term.ds.unit_registry
+            else:
+                registry = None
+
+            unit = getattr(term, "units", "1")
+
+            return Unit(unit, registry=registry)
+
+        match self.op:
+            case (
+                operator.eq
+                | operator.ne
+                | operator.le
+                | operator.lt
+                | operator.ge
+                | operator.gt
+            ):
+                # Boolean operator return a dimensionless quantity
+                return Unit("1")
+            case operator.add | operator.sub:
+                # 'Unit's cannot be added/subtracted direclty but they have to
+                # have the same dimensions
+                units = [helper(_) for _ in self.terms]
+                if units and not all(units[0].same_dimensions_as(u) for u in units):
+                    raise ValueError("Incompatible units")
+                return units[0]
+            case _:
+                # Other operators are mul, truediv and neg
+                units = [helper(_) for _ in self.terms]
+                return reduce(self.op, units)
+
 
 class DerivedField(DerivedFieldBase):
     """

yt/fields/field_type_container.py

@@ -94,7 +94,7 @@ def __getattr__(self, attr):
 
     def __setattr__(self, attr, value):
         if isinstance(value, DerivedFieldCombination):
-            self.ds.add_field((self.field_type, attr), value)
+            self.ds.add_field((self.field_type, attr), value, units=value.units)
         else:
             super().__setattr__(attr, value)