Use cleaner implementation

Author: cphyc

yt/data_objects/selection_objects/cut_region.py

@@ -8,7 +8,7 @@
     YTSelectionContainer3D,
 )
 from yt.data_objects.static_output import Dataset
-from yt.fields.derived_field import DerivedField
+from yt.fields.derived_field import DerivedField, DerivedFieldCombination
 from yt.funcs import iter_fields, validate_object
 from yt.geometry.selection_routines import points_in_cells
 from yt.utilities.exceptions import YTIllDefinedCutRegion
@@ -57,7 +57,8 @@ def __init__(
         validate_object(data_source, YTSelectionContainer)
         conditionals = list(always_iterable(conditionals))
         for condition in conditionals:
-            validate_object(condition, (str, DerivedField))
+            validate_object(condition, (str, DerivedFieldCombination))
+
         validate_object(ds, Dataset)
         validate_object(field_parameters, dict)
         validate_object(base_object, YTSelectionContainer)
@@ -83,8 +84,8 @@ def __init__(
     def _check_filter_fields(self):
         fields = []
         for cond in self.conditionals:
-            if isinstance(cond, DerivedField):
-                fields.append(cond.name)
+            if isinstance(cond, DerivedFieldCombination):
+                fields.extend(cond.getDependentFields())
                 continue
 
             for field in re.findall(r"\[([A-Za-z0-9_,.'\"\(\)]+)\]", cond):
@@ -152,8 +153,8 @@ def _cond_ind(self):
         locals["obj"] = obj
         with obj._field_parameter_state(self.field_parameters):
             for cond in self.conditionals:
-                if isinstance(cond, DerivedField):
-                    res = cond(obj)
+                if isinstance(cond, DerivedFieldCombination):
+                    res = cond(None, obj)
                 else:
                     res = eval(cond, locals)
                 if ind is None:

yt/data_objects/static_output.py

@@ -35,7 +35,11 @@
 from yt.data_objects.particle_filters import ParticleFilter, filter_registry
 from yt.data_objects.region_expression import RegionExpression
 from yt.data_objects.unions import ParticleUnion
-from yt.fields.derived_field import DerivedField, ValidateSpatial
+from yt.fields.derived_field import (
+    DerivedField,
+    DerivedFieldCombination,
+    ValidateSpatial,
+)
 from yt.fields.field_type_container import FieldTypeContainer
 from yt.fields.fluid_fields import setup_gradient_fields
 from yt.funcs import iter_fields, mylog, set_intersection, setdefaultattr
@@ -883,6 +887,9 @@ def add_particle_filter(self, filter):
             used = self._setup_filtered_type(f)
             if used:
                 filter = f
+        elif isinstance(filter, DerivedFieldCombination):
+            filter_registry[filter.name] = filter
+            used = self._setup_filtered_type(filter)
         else:
             used = self._setup_filtered_type(filter)
         if not used:
@@ -1757,15 +1764,12 @@ def add_field(
         """
         from yt.fields.field_functions import validate_field_function
 
-        if not isinstance(function, DerivedField):
+        if not isinstance(function, DerivedFieldCombination):
             if sampling_type is None:
                 raise ValueError("You must specify a sampling_type for the field.")
             validate_field_function(function)
         else:
             sampling_type = function.sampling_type
-            kwargs.setdefault("units", function.units)
-
-            function = function._function
 
         self.index
         if force_override and name in self.index.field_list:

yt/fields/derived_field.py

@@ -1,9 +1,11 @@
+import abc
 import contextlib
 import inspect
 import operator
 import re
 from collections.abc import Callable, Iterable
-from typing import Optional, Union
+from functools import reduce
+from typing import Optional
 
 from more_itertools import always_iterable
 
@@ -59,7 +61,125 @@ def _DeprecatedFieldFunc(field, data):
     return _DeprecatedFieldFunc
 
 
-class DerivedField:
+class DerivedFieldBase(abc.ABC):
+    @abc.abstractmethod
+    def __call__(self, field, data):
+        pass
+
+    @abc.abstractmethod
+    def __repr__(self) -> str:
+        pass
+
+    # Multiplication (left and right side)
+    def __mul__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.mul)
+
+    def __rmul__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.mul)
+
+    # Division (left side)
+    def __truediv__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.truediv)
+
+    def __rtruediv__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([other, self], op=operator.truediv)
+
+    # Addition (left and right side)
+    def __add__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.add)
+
+    def __radd__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.add)
+
+    # Subtraction (left and right side)
+    def __sub__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.sub)
+
+    def __rsub__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([other, self], op=operator.sub)
+
+    # Unary minus
+    def __neg__(self) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self], op=operator.neg)
+
+    # Comparison operators
+    def __leq__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.le)
+
+    def __lt__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.lt)
+
+    def __geq__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.ge)
+
+    def __gt__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.gt)
+
+    # def __eq__(self, other) -> "DerivedFieldCombination":
+    #     return DerivedFieldCombination([self, other], op=operator.eq)
+
+    def __ne__(self, other) -> "DerivedFieldCombination":
+        return DerivedFieldCombination([self, other], op=operator.ne)
+
+
+class DerivedFieldCombination(DerivedFieldBase):
+    sampling_type: str | None
+    terms: list
+    op: Callable | None
+
+    def __init__(self, terms: list, op=None):
+        if not terms:
+            raise ValueError("DerivedFieldCombination requires at least one term.")
+
+        # Make sure all terms have the same sampling type
+        sampling_types = set()
+        for term in terms:
+            if isinstance(term, DerivedField):
+                sampling_types.add(term.sampling_type)
+
+        if len(sampling_types) > 1:
+            raise ValueError(
+                "All terms in a DerivedFieldCombination must "
+                "have the same sampling type."
+            )
+        self.sampling_type = sampling_types.pop() if sampling_types else None
+        self.terms = terms
+        self.op = op
+
+    def __call__(self, field, data):
+        """
+        Return the value of the field in a given data object.
+        """
+        qties = []
+        for term in self.terms:
+            if isinstance(term, DerivedField):
+                qties.append(data[term.name])
+            elif isinstance(term, DerivedFieldCombination):
+                qties.append(term(field, data))
+            else:
+                qties.append(term)
+
+        if len(qties) == 1:
+            return self.op(qties[0])
+        else:
+            return reduce(self.op, qties)
+
+    def __repr__(self):
+        return f"DerivedFieldCombination(terms={self.terms!r}, op={self.op!r})"
+
+    def getDependentFields(self):
+        fields = []
+        for term in self.terms:
+            if isinstance(term, DerivedField):
+                fields.append(term.name)
+            elif isinstance(term, DerivedFieldCombination):
+                fields.extend(term.getDependentFields())
+            else:
+                continue
+        return fields
+
+
+class DerivedField(DerivedFieldBase):
     """
     This is the base class used to describe a cell-by-cell derived field.
 
@@ -277,11 +397,11 @@ def __call__(self, data):
         """Return the value of the field in a given *data* object."""
         self.check_available(data)
         original_fields = data.keys()  # Copy
-        if self._function is NullFunc:
-            raise RuntimeError(
-                "Something has gone terribly wrong, _function is NullFunc "
-                + f"for {self.name}"
-            )
+        # if self._function is NullFunc:
+        #     raise RuntimeError(
+        #         "Something has gone terribly wrong, _function is NullFunc "
+        #         + f"for {self.name}"
+        #     )
         with self.unit_registry(data):
             dd = self._function(self, data)
         for field_name in data.keys():
@@ -499,128 +619,6 @@ def __copy__(self):
             nodal_flag=self.nodal_flag,
         )
 
-    def _operator(
-        self, other: Union["DerivedField", float], op: Callable
-    ) -> "DerivedField":
-        my_units = self.ds.get_unit_from_registry(self.units)
-        if isinstance(other, DerivedField):
-            if self.sampling_type != other.sampling_type:
-                raise TypeError(
-                    f"Cannot {op} fields with different sampling types: "
-                    f"{self.sampling_type} and {other.sampling_type}"
-                )
-
-            def wrapped(field, data):
-                return op(self(data), other(data))
-
-            other_name = other.name[1]
-            other_units = self.ds.get_unit_from_registry(other.units)
-
-        else:
-            # Special case when passing (value, "unit") tuple
-            if isinstance(other, tuple) and len(other) == 2:
-                other = self.ds.quan(*other)
-
-            def wrapped(field, data):
-                return op(self(data), other)
-
-            other_name = str(other)
-            other_units = getattr(other, "units", self.ds.get_unit_from_registry("1"))
-
-        if op in (operator.add, operator.sub, operator.eq):
-            assert my_units.same_dimensions_as(other_units)
-            new_units = my_units
-        elif op in (operator.mul, operator.truediv):
-            new_units = op(my_units, other_units)
-        elif op in (operator.le, operator.lt, operator.ge, operator.gt, operator.ne):
-            # Comparison yield unitless fields
-            new_units = Unit("1")
-        else:
-            raise TypeError(f"Unsupported operator {op} for DerivedField")
-
-        return DerivedField(
-            name=(self.name[0], f"{self.name[1]}_{op.__name__}_{other_name}"),
-            sampling_type=self.sampling_type,
-            function=wrapped,
-            units=new_units,
-            ds=self.ds,
-        )
-
-    # Multiplication (left and right side)
-    def __mul__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.mul)
-
-    def __rmul__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.mul)
-
-    # Division (left side)
-    def __truediv__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.truediv)
-
-    # Addition (left and right side)
-    def __add__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.add)
-
-    def __radd__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.add)
-
-    # Subtraction (left and right side)
-    def __sub__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.sub)
-
-    def __rsub__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(-other, op=operator.add)
-
-    # Unary minus
-    def __neg__(self) -> "DerivedField":
-        def wrapped(field, data):
-            return -self(data)
-
-        return DerivedField(
-            name=(self.name[0], f"neg_{self.name[1]}"),
-            sampling_type=self.sampling_type,
-            function=wrapped,
-            units=self.units,
-            ds=self.ds,
-        )
-
-    # Division (right side, a bit more complex)
-    def __rtruediv__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        units = self.ds.get_unit_from_registry(self.units)
-
-        def wrapped(field, data):
-            return 1 / self(data)
-
-        inverse_self = DerivedField(
-            name=(self.name[0], f"inverse_{self.name[1]}"),
-            sampling_type=self.sampling_type,
-            function=wrapped,
-            units=units**-1,
-            ds=self.ds,
-        )
-
-        return inverse_self * other
-
-    # Comparison operators
-    def __leq__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.le)
-
-    def __lt__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.lt)
-
-    def __geq__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.ge)
-
-    def __gt__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.gt)
-
-    # Somehow, makes yt not work?
-    # def __eq__(self, other: Union["DerivedField", float]) -> "DerivedField":
-    #     return self._operator(other, op=operator.eq)
-
-    def __ne__(self, other: Union["DerivedField", float]) -> "DerivedField":
-        return self._operator(other, op=operator.ne)
-
 
 class FieldValidator:
     """

yt/fields/field_type_container.py

@@ -8,7 +8,7 @@
 from functools import cached_property
 
 from yt._maintenance.ipython_compat import IPYWIDGETS_ENABLED
-from yt.fields.derived_field import DerivedField
+from yt.fields.derived_field import DerivedField, DerivedFieldCombination
 
 
 def _fill_values(values):
@@ -93,7 +93,7 @@ def __getattr__(self, attr):
         return ds.field_info[ft, attr]
 
     def __setattr__(self, attr, value):
-        if isinstance(value, DerivedField):
+        if isinstance(value, DerivedFieldCombination):
             self.ds.add_field((self.field_type, attr), value)
         else:
             super().__setattr__(attr, value)