Source code for opentidalfarm.memoize

import os
import signal
import pickle
from dolfin import log, INFO, WARNING
from .helpers import cpu0only

def to_tuple(obj):
    if hasattr(obj, '__iter__'):
        return tuple([to_tuple(o) for o in obj])
    else:
        return obj


[docs]class MemoizeMutable: ''' Implements a memoization function to avoid duplicated functional (derivative) evaluations ''' def get_key(self, args, kwds): h1 = to_tuple(args) h2 = to_tuple(list(kwds.items())) h = tuple([h1, h2]) # Often useful to have a explicit # turbine parameter -> functional value mapping, # i.e. no hashing on the key if self.hash_keys: h = hash(h) return h def __init__(self, fn, hash_keys=False): ''' sigint_save: Create a checkpoint file in case a sigint signal is received. ''' self.fn = fn self.memo = {} self.hash_keys = hash_keys def __call__(self, *args, **kwds): h = self.get_key(args, kwds) if h not in self.memo: self.memo[h] = self.fn(*args, **kwds) else: log(INFO, "Use checkpoint value.") return self.memo[h] def has_cache(self, *args, **kwds): h = self.get_key(args, kwds) return h in self.memo # Insert a function value into the cache manually. def __add__(self, value, *args, **kwds): h = self.get_key(args, kwds) self.memo[h] = value @cpu0only def save_checkpoint(self, filename): def sig_save(sig, stack): print("Received signal %i. Writing final checkpoint to disk before exiting..." % sig) pickle.dump(self.memo, open(filename, "wb")) print("Checkpoint writing finished. Bye.") os._exit(sig) # Make sure we save successfully, even if the user sends a signal print("Save checkpoint.") old_handler = signal.signal(signal.SIGINT, sig_save) pickle.dump(self.memo, open(filename, "wb")) signal.signal(signal.SIGINT, old_handler) def load_checkpoint(self, filename): try: self.memo = pickle.load(open(filename, "rb")) except IOError: log(WARNING, "Warning: Checkpoint file '%s' not found." % filename) except ValueError: log(WARNING, "Error: Checkpoint file '%s' is invalid." % filename)