Cut regions

Author: cphyc

yt/data_objects/selection_objects/cut_region.py

@@ -8,7 +8,8 @@
     YTSelectionContainer3D,
 )
 from yt.data_objects.static_output import Dataset
-from yt.funcs import iter_fields, validate_object, validate_sequence
+from yt.fields.derived_field import DerivedField
+from yt.funcs import iter_fields, validate_object
 from yt.geometry.selection_routines import points_in_cells
 from yt.utilities.exceptions import YTIllDefinedCutRegion
 from yt.utilities.on_demand_imports import _scipy
@@ -54,14 +55,14 @@ def __init__(
         if locals is None:
             locals = {}
         validate_object(data_source, YTSelectionContainer)
-        validate_sequence(conditionals)
+        conditionals = list(always_iterable(conditionals))
         for condition in conditionals:
-            validate_object(condition, str)
+            validate_object(condition, (str, DerivedField))
         validate_object(ds, Dataset)
         validate_object(field_parameters, dict)
         validate_object(base_object, YTSelectionContainer)
 
-        self.conditionals = list(always_iterable(conditionals))
+        self.conditionals = conditionals
         if isinstance(data_source, YTCutRegion):
             # If the source is also a cut region, add its conditionals
             # and set the source to be its source.
@@ -82,6 +83,10 @@ def __init__(
     def _check_filter_fields(self):
         fields = []
         for cond in self.conditionals:
+            if isinstance(cond, DerivedField):
+                fields.append(cond.name)
+                continue
+
             for field in re.findall(r"\[([A-Za-z0-9_,.'\"\(\)]+)\]", cond):
                 fd = field.replace('"', "").replace("'", "")
                 if "," in fd:
@@ -125,8 +130,11 @@ def blocks(self):
             m = m.copy()
             with obj._field_parameter_state(self.field_parameters):
                 for cond in self.conditionals:
-                    ss = eval(cond)
-                    m = np.logical_and(m, ss, m)
+                    if isinstance(cond, DerivedField):
+                        ss = cond(obj)
+                    else:
+                        ss = eval(cond)
+                    m &= ss
             if not np.any(m):
                 continue
             yield obj, m
@@ -144,12 +152,15 @@ def _cond_ind(self):
         locals["obj"] = obj
         with obj._field_parameter_state(self.field_parameters):
             for cond in self.conditionals:
-                res = eval(cond, locals)
+                if isinstance(cond, DerivedField):
+                    res = cond(obj)
+                else:
+                    res = eval(cond, locals)
                 if ind is None:
                     ind = res
                 if ind.shape != res.shape:
                     raise YTIllDefinedCutRegion(self.conditionals)
-                np.logical_and(res, ind, ind)
+                ind &= res
         return ind
 
     def _part_ind_KDTree(self, ptype):

yt/fields/derived_field.py

@@ -524,11 +524,16 @@ def wrapped(field, data):
             other_name = str(other)
             other_units = self.ds.get_unit_from_registry(getattr(other, "units", "1"))
 
-        if op in (operator.add, operator.sub):
+        if op in (operator.add, operator.sub, operator.eq):
             assert my_units.same_dimensions_as(other_units)
             new_units = my_units
-        else:
+        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}"),
@@ -593,6 +598,26 @@ def wrapped(field, data):
 
         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:
     """