Save weak references to all fields built on a distributor, to help track fields and memory usage

Author: kburns

dedalus/core/distributor.py

@@ -9,6 +9,7 @@
 from collections import OrderedDict
 from math import prod
 import numbers
+from weakref import WeakSet
 
 from .coords import CoordinateSystem, DirectProduct
 from ..tools.array import reshape_vector
@@ -112,6 +113,8 @@ def __init__(self, coordsystems, comm=None, mesh=None, dtype=None):
         self.comm_coords = np.array(self.comm_cart.coords, dtype=int)
         # Build layout objects
         self._build_layouts()
+        # Keep set of weak field references
+        self.fields = WeakSet()
 
     @CachedAttribute
     def cs_by_axis(self):

dedalus/core/field.py

@@ -572,6 +572,8 @@ def __init__(self, dist, bases=None, name=None, tensorsig=None, dtype=None):
         self.layout = self.dist.get_layout_object('c')
         # Change scales to build buffer and data
         self.preset_scales((1,) * self.dist.dim)
+        # Add weak reference to distributor
+        dist.fields.add(self)
 
     def __getitem__(self, layout):
         """Return data viewed in specified layout."""