5
\$\begingroup\$

I am an astrophysicist working on large simulations as part of the APOSTLE project. The output of the simulations I use are large (TBs) and are stored in tables spread across multiple hdf5 files. Often, I'm interested in studying a particular simulated galaxy, which is represented by a collection of particles of different types and which have different properties.

Prior to writing this code, I would typically need to read tables corresponding to various properties for all particles in the simulation, then construct a selection to extract only those particles belonging to the galaxy of interest, copy those, and discard the large tables.

Given that I often study one or a few galaxies for a while, I thought it would be a good idea to build a little interface which computes selection "masks" for a given galaxy once, the first time it is run, and caches those to disk.

In addition, it caches any additional properties of the galaxy particles that get pulled out of the big tables, and some associated metadata. My solution is a subclass of dict, so once an ApostleObject (i.e. an abstraction of a galaxy) is initialized, call it AO, I can extract e.g. particle properties from the object by AO['T_g']. (In this example 'T_g' refers to the gas temperature. Another class, ApostleFileset, is aware of all available keys and how to get the raw tables for all particles from the disk.) If 'T_g' has been loaded before it will already be in memory and simply be returned, otherwise it will be loaded from the master tables if the key is valid, otherwise a KeyError is raised.

There are two other features of note:

  1. I just implemented __getattr__ as an alternative to __getitem__ since writing AO.T_g is a bit more succinct and readable than AO['T_g'] (though __getitem__ is still useful in many contexts).

  2. This one has to do with how the caching is implemented - the dict subclass itself is actually buried inside a thin wrapper that is just a context manager. This allows me to "lock" the cache file when the object is created (achieved by creating an empty file with the same name but a .lock extension in the same directory), and later remove the lock when the block exits. This is to avoid the scenario where one process reads the cache, then another reads the cache, then the first writes the cache, then the second re-writes the cache, undoing the changes.

    I'm not entirely happy with how this came out, since I'm effectively required to use with ApostleObject(...) as AO:, which is otherwise unnecessary - e.g. handling files is all taken care of within ApostleFileset. I'm considering adding a read-only mode where the cache can still be used, but can't be updated, but am interested in other solutions as well.

Beyond that, I'm also interested in any general comments on style and readability. I have some formal programming training, but all in the context of scientific computing, and I'm mostly self-taught, so I'm curious to know how this all looks to an outsider.

One last note on the imports. numpy and os are of course standard modules. slvars is a thin wrapper around cPickle. ApostleFileset I've described somewhat above. I don't think it's necessary to include it here as long as its purpose is explained, especially since it comes with its own dependencies - reading those big tables quickly is a whole problem in itself! T, which gets read from the same file, is pretty trivial: T = {'g': '0', 'dm': '1', 'b2': '2', 'b3': '3', 's': '4', 'bh': '5'}, these are just suffixes for the different particle types (for the curious: gas, dark matter, boundary2, boundary3, star, black hole).

from apostle_fileset import ApostleFileset, T
from slvars import savevars, loadvars
import numpy as np
import os

class ApostleObject:

    def __enter__(self):

        class _ApostleObject(dict):
            def __init__(self, res=None, vol=None, phys=None, snap=None, fof=None, sub=None, mask_type=None, aperture=0, prefix='./'):

                if (res not in ('L', 'M', 'H')):
                    raise ValueError("ApostleObject: res must be in 'L', 'M', 'H'")
                if (vol not in [str(i) for i in range(1,13)]):
                    raise ValueError("ApostleObject: vol must be in '1', '2', ..., '12'")
                if (phys not in ('fix', 'DMO')):
                    raise ValueError("ApostleObject: phys must be 'fix' or 'DMO'")
                if (type(snap) != int):
                    raise ValueError('ApostleObject: provide snapshot number as integer')
                if (type(fof) != int):
                    raise ValueError('ApostleObject: provide fof number as integer')
                if (type(sub) != int):
                    raise ValueError('ApostleObject: provide subhalo number as integer')
                if (mask_type not in ('fof_sub', 'fof', 'aperture')):
                    raise ValueError("ApostleObject: provide mask_type ('fof_sub', 'fof', 'aperture')")
                if (mask_type == 'aperture') and (aperture == 0):
                    raise ValueError("ApostleObject: provide aperture for mask_type='aperture'")

                self._cache_name = prefix + '/' + 'AOCache_' + phys + '_' + res + '_' + vol + '_' + str(snap) + '_'+str(fof) + '_'+str(sub) + '_' + {'fof_sub': 'fof_sub', 'fof': 'fof', 'aperture':'aperture' + str(int(aperture))}[mask_type]

                if os.path.exists(self._cache_name + '.lock'):
                    raise RuntimeError("ApostleObject '" + self._cache_name + ".pkl' is locked by another instance.")
                else:
                    self._lock()

                if os.path.exists(self._cache_name):
                    D, = loadvars(self._cache_name)
                    self.__dict__.update(D)
                    self._F._init_extractors()

                else:
                    self.res, self.vol, self.phys, self.snap, self.fof, self.sub, self.mask_type, self.aperture = res, vol, phys, snap, fof, sub, mask_type, aperture

                    self._F = ApostleFileset(res=res, vol=vol, phys=phys, snap=snap)
                    self._define_masks()
                    self._cache()

                return

            def __setattr__(self, name, value):
                self.__dict__[name] = value
                return

            def __getattr__(self, name):
                try:
                    return self.__dict__[str(name)]
                except KeyError:
                    self._load_key(name)
                    return self.__dict__[str(name)]

            def __setitem__(self, key, item):
                self.__dict__[key] = item
                return

            def __getitem__(self, key):
                try:
                    return self.__dict__[key]
                except KeyError:
                    self._load_key(key)
                    return self.__dict__[key]

            def __repr__(self):
                return repr(self.__dict__)

            def __len__(self):
                return len(self.__dict__)

            def __delitem__(self, key):
                del self.__dict__[key]
                return

            def clear(self):
                return self.__dict__.clear()

            def copy(self):
                return self.__dict__.copy()

            def has_key(self, key):
                return self.__dict__.has_key(key)

            def pop(self, key, d=None):
                return self.__dict__.pop(key, d)

            def update(self, *args, **kwargs):
                return self.__dict__.update(*args, **kwargs)

            def keys(self):
                return self.__dict__.keys()

            def values(self):
                return self.__dict__.values()

            def items(self):
                return self.__dict__.items()

            def pop(self, *args):
                return self.__dict__.pop(*args)

            def __cmp__(self, d):
                return cmp(self.__dict__, d)

            def __contains__(self, item):
                return item in self.__dict__

            def __iter__(self):
                return iter(self.__dict__)

            def __unicode__(self):
                return unicode(repr(self.__dict__))

            def _define_masks(self):
                loaded_keys = set()
                loaded_keys.update(self._F.load('group', ('gns', 'sgns', 'nfof', 'nID', 'offID')))
                self.gmask = np.logical_and(self._F.gns == self.fof, self._F.sgns == self.sub)
                self.fofmask = np.arange(1, self._F.nfof + 1) == self.fof
                self.idmask = np.s_[self._F.offID[np.logical_and(self._F.gns == self.fof, self._F.sgns == self.sub)][0] : self._F.offID[np.logical_and(self._F.gns == self.fof, self._F.sgns == self.sub)][0] + self._F.nID[np.logical_and(self._F.gns == self.fof, self._F.sgns == self.sub)][0]]
                self.pmasks = {}
                if self.mask_type == 'fof_sub':
                    loaded_keys.update(self._F.load('particle', tuple([field + typesuffix for field in ['ng_', 'nsg_'] for typesuffix in T.keys()])))
                    for typesuffix in T.keys():
                        self.pmasks[typesuffix] = np.logical_and(self._F['ng_' + typesuffix] == self.fof, self._F['nsg_' + typesuffix] == self.sub)
                elif self.mask_type == 'fof':
                    loaded_keys.update(self._F.load('particle', tuple(['ng_' + typesuffix for typesuffix in T.keys()])))
                    for typesuffix in T.keys():
                        self.pmasks[typesuffix] = self._F['ng_' + typesuffix] == self.fof
                elif self.mask_type == 'aperture':
                    loaded_keys.update(self._F.load('group', ('cops', 'vcents')))
                    loaded_keys.update(self._F.load('snapshot', ('xyz_g', 'xyz_dm', 'xyz_b2', 'xyz_b3', 'xyz_s', 'xyz_bh', 'Lbox')))
                    for typesuffix in T.keys():
                        self._F['xyz_' + typesuffix] = self._F['xyz_' + typesuffix] - self._F.cops[self.gmask]
                        self._F['xyz_' + typesuffix][self._F['xyz_' + typesuffix] > self._F.Lbox / 2.] -= self._F.Lbox
                        self._F['xyz_' + typesuffix][self._F['xyz_' + typesuffix] < self._F.Lbox / 2.] += self._F.Lbox
                        cube = (np.abs(self._F['xyz_' + typesuffix]) < self.aperture).all(axis=1)
                        self.pmasks[typesuffix] = np.zeros(self._F['xyz_' + typesuffix].shape[0], dtype=np.bool)
                        self.pmasks[typesuffix][cube] = np.sum(np.power(self._F['xyz_' + typesuffix][cube], 2), axis=1) < np.power(self.aperture, 2)
                for k in loaded_keys:
                    del self._F[k]
                return

            def _load_key(self, key):

                loaded_keys = set()

                if key in self._F.fields('header'):
                    loaded_keys.update(self._F.load('snapshot', (key, )))
                    self[key] = self._F[key]

                elif key in self._F.fields('particle'):
                    ptype = key.split('_')[-1]
                    if self.mask_type == 'aperture':
                        loaded_keys.update(self._F.load('snapshot', (key, )))
                        self[key] = self._F[key][self.pmasks[ptype]]
                    elif (self.mask_type == 'fof') or (self.mask_type == 'fof_sub'):
                        loaded_keys.update(self._F.load('particle', (key, )))
                        self[key] = self._F[key][self.pmasks[ptype]]
                    else:
                        raise ValueError
                    if key.split('_')[0] == 'xyz':
                        loaded_keys.update(self._F.load('snapshot', ('Lbox',)))
                        loaded_keys.update(self._F.load('group', ('cops',)))
                        self[key] -= self._F.cops[self.gmask]
                        self[key][self[key] > self._F.Lbox / 2.] -= self._F.Lbox
                        self[key][self[key] < -self._F.Lbox / 2.] += self._F.Lbox
                    elif key.split('_')[0] == 'vxyz':
                        loaded_keys.update(self._F.load('group', ('vcents',)))
                        self[key] -= self._F.vcents[self.gmask]

                elif key in self._F.fields('group'):
                    loaded_keys.update(self._F.load('group', (key, )))
                    self[key] = self._F[key][self.gmask]

                elif key in self._F.fields('fofgroup'):
                    loaded_keys.update(self._F.load('group', (key, )))
                    self[key] = self._F[key][self.fofmask]

                elif key in self._F.fields('idgroup'):
                    loaded_keys.update(self._F.load('group', (key, )))
                    self[key] = self._F[key][self.idmask]

                else:
                    raise KeyError

                for k in loaded_keys:
                    del self._F[k]

                self._cache()

            def _cache(self):
                del self._F._Extractors
                savevars([self.__dict__], self._cache_name)
                self._F._init_extractors()
                return

            def _lock(self):
                open(self._cache_name + '.lock', 'a').close()
                return

            def _unlock(self):
                os.remove(self._cache_name + '.lock')
                return


        self._AO = _ApostleObject(*self._args, **self._kwargs)
        return self._AO

    def __exit__(self, exc_type, exc_value, traceback):
        self._AO._unlock()

    def __init__(self, *args, **kwargs):
        self._args = args
        self._kwargs = kwargs
\$\endgroup\$
  • \$\begingroup\$ Welcome to Code Review, nice question! As you say the data is large, are performance or memory-optimization a concern? \$\endgroup\$ – Peilonrayz Dec 21 '16 at 0:19
  • \$\begingroup\$ @Peilonrayz they are concerns in the context of the broader project, but not so much for the class at hand. \$\endgroup\$ – Kyle Dec 21 '16 at 4:18
3
\$\begingroup\$

From the look of your code, I'd say ApostleFileset and ApostleObject are both god classes. First _define_masks should be defined in ApostleFileset, but _load_key can almost be moved into it too. The only problem with _load_key is when the key is a particle, and the first value is xyz, you mutate self, rather than just assign to it.

Also using an if statement, and a couple of function calls may lead to a TOCTOU bug. Which was changed in Python 3 by adding the x mode to open.

Since you don't want any more god classes, you should change the input and output to be as simple as possible. You want; the cache file location, and a function to get missing values. This should generate a lock for the file, and open the cache for you. If you want anymore functionality, then you have to subclass it and add it. But you should only really add saving and reading to or from the cache.

This can get you:

import os

class Cache(dict):
    def __init__(self, provider):
        self._provider = provider

    def __missing__(self, key):
        v = self._provider(self, key)
        self[key] = v
        return v


class FileCache(Cache):
    def __init__(self, path, provider):
        super(FileCache, self).__init__(provider)
        self._path = path

    def __enter__(self):
        path = self._path + '.lock'
        if not os.path.isfile(path):
            with open(path, 'a') as f:
                pass
        else:
            raise IOError('file {!r} exists.'.format(path))
        self._file = open(self._path, 'a+')
        return self

    def __exit__(self, exc_type, exc_value, traceback):
        self._file.close()
        os.remove(self._path + '.lock')
        return False

And then you want to subclass FileCache for however you want to change ApostleFileset. As an example, below is an example of how you could extend the class:

class ApostleObject(FileCache):
    def __enter__(self):
        ret = super(ApostleObject, self).__enter__()
        for line in self._file:
            key, value = line.split(' ', 1)
            self[int(key)] = int(value)
        return ret

    def __missing__(self, key):
        value = super(ApostleObject, self).__missing__(key)
        self._file.write('{key} {value}\n'.format(key=key, value=value))
        return value


class ApostleFileset(object):
    def missing(self, cache, key):
        return 2 * key


with ApostleObject('cache', ApostleFileset().missing) as obj:
    print(obj[1])
    print(obj[1])
    print(obj[3])
    # Only to show the content of the file, not for actual use.
    obj._file.seek(0, os.SEEK_SET)
    print(obj._file.read())

with ApostleObject('cache', ApostleFileset().missing) as obj:
    print(obj[2])
    obj._file.seek(0, os.SEEK_SET)
    print(obj._file.read())

Which outputs:

2
2
6
1 2
3 6

4
1 2
3 6
2 4
\$\endgroup\$
  • \$\begingroup\$ Thanks, working through all this. Re the TOCTOU bug, I don't see it, nothing happens in between checking if the lock exists and creating it: if os.path.exists(self._cache_name + '.lock'): raise RuntimeError("ApostleObject '" + self._cache_name + ".pkl' is locked by another instance.") else: self._lock() \$\endgroup\$ – Kyle Dec 21 '16 at 22:49
  • \$\begingroup\$ @Kyle Oh, I don't know how I misread that, yeah I was wrong on that, :) \$\endgroup\$ – Peilonrayz Dec 21 '16 at 23:38
2
\$\begingroup\$

A small improvement:

Right now you duplicate the logic in __getitem__ and __getattr__. In __getattr__ you already know the name will be a string (otherwise it is not a valid attribute name) and you know it exists in __dict__ (or has t be loaded). Therefore you can just delegate it to __getitem__:

def __getattr__(self, name):
    return self.__getitem__(self, name)

or, even easier:

__getattr__ = __getitem__

The same is true for __setitem__ and __setattr__.

Note that your class states that it is a subclass of dict, but actually overwrites all of the methods! You could just replace all self.__dict__ with self and store it in the actual dictionary. This way you don't need to implement __repr__, __len__, __delitem__, pop, update, keys, values, items, pop (you actually defined it twice), __cmp__, __contains__ and __iter__, because dict already has all of them, except for __unicode__!

And going even further, you can use the fact that dict has a __missing__ method, which is normally not implemented, but which is called if a key is missing.

def __missing__(self, key):
    self._load_key(key)
    return dict.__getitem__(self, key)

Note that it then becomes necessary to actually have a different __getattr__ again:

def __getattr__(self, key):
    try:
        return self.__dict__[key]
    except KeyError:
        return self[key]

(This will call __mising__ if self[key] is invoked with a missing key, triggering the loading again.) This is necessary, because otherwise self.update would fail.

This way your class could be:

class _ApostleObject(dict):
    # __init__
    ...

    def __getattr__(self, key):
        try:
            return self.__dict__[key]
        except KeyError:
            return self[key]

    def __missing__(self, key):
        result = self[key] = self._load_key(key) # make _load_key return its result
        return result

    # _lock, _load_key, ...
    ...
\$\endgroup\$
  • \$\begingroup\$ Thanks very much, all useful comments. Any thoughts on the caching functionality? \$\endgroup\$ – Kyle Dec 21 '16 at 4:19

Your Answer

By clicking “Post Your Answer”, you agree to our terms of service, privacy policy and cookie policy

Not the answer you're looking for? Browse other questions tagged or ask your own question.