misc: Cleanup

Author: EdCaunt

devito/operator/operator.py

@@ -653,7 +653,6 @@ def _prepare_arguments(self, autotune=None, estimate_memory=False, **kwargs):
         # the subsequent phases of the arguments processing
         args = kwargs['args'] = ArgumentsMap(args, grid, self)
 
-        # FIXME: Will want to remove this if not using prepare_args to estimate memory
         if estimate_memory:
             # No need to do anything more if only checking the memory
             return args
@@ -1306,7 +1305,6 @@ def nbytes_avail_mapper(self):
         """
         mapper = {}
 
-        # TODO: This doesn't account for the size of the snapshots?
         # The amount of space available on the disk
         usage = shutil.disk_usage(gettempdir())
         mapper[disk_layer] = usage.free
@@ -1324,10 +1322,7 @@ def nbytes_avail_mapper(self):
         mapper[host_layer] = int(ANYCPU.memavail() / nproc)
 
         for layer in (host_layer, device_layer):
-            try:
-                mapper[layer] -= self.nbytes_consumed_operator[layer]
-            except KeyError:
-                continue
+            mapper[layer] -= self.nbytes_consumed_operator.get(layer, 0)
 
         mapper = {k: int(v) for k, v in mapper.items()}
 
@@ -1466,7 +1461,6 @@ def nbytes_consumed_memmapped(self):
 
     @cached_property
     def nbytes_snapshots(self):
-        disk = 0
 
         # Symbols in the operator which may or may not carry data
         op_symbols = FindSymbols().visit(self.op)
@@ -1475,14 +1469,14 @@ def nbytes_snapshots(self):
         op_symbols = [i for i in op_symbols if i.is_AbstractFunction
                       and not i.is_ArrayBasic and not i.alias]
 
+        disk = 0
         for i in op_symbols:
             try:
-                disk += i.size_snapshot
-                print(i.size_snapshot, i._time_size_ideal)
+                disk += i.size_snapshot*i._time_size_ideal*np.dtype(i.dtype).itemsize
             except AttributeError:
                 pass
 
-        return {disk_layer: 0, host_layer: 0, device_layer: 0}
+        return {disk_layer: disk, host_layer: 0, device_layer: 0}
 
 
 def parse_kwargs(**kwargs):

devito/types/sparse.py

@@ -652,7 +652,7 @@ def _arg_defaults(self, alias=None, estimate_memory=False):
             if f is not None:
                 mapper[getattr(self, i)] = f
 
-        if estimate_memory:  # kwargs.get("estimate_memory", False):
+        if estimate_memory:
             # Avoid touching the data in any capacity, and simply return
             # the symbolic objects if merely estimating memory consumption.
             return ReducerMap({v.name: k for k, v in mapper.items()})