Selecting objects

Question

I have created an Atom object as follows:

class Atom(object):
    def __init__(self, symbol, x, y, z)
        self.symbol = symbol
        self.position = (x, y, z)

and a Selection class that contains atoms selected by some criteria:

class Selection(object):
    def __init__(self, a_system, atom_list=[]):
        for atom in a_system:
            atom_list.append(atom)
        self.atom_list = atom_list

    def by_symbol(self, symbol):
        r_list = []
        for atom in self.atom_list:
            if atom.symbol is symbol:
                r_list.append(atom)
        self.atom_list = r_list

    def by_zrange(self, zmin, zmax):
        r_list = []
        for atom in self.atom_list:
            pos = atom.position[2]
            if pos > zmin and pos < zmax:
                r_list.append(atom)
        self.atom_list = r_list

So as you can see I can say for example:

# my_system is a list of atoms objects
group = Selection(my_system)

and then say:

group.by_symbol('H')

and I'll have in object group all hydrogen atoms. Then if I do:

group.by_zrange(1, 2)

and I'll have in object group all hydrogen atoms which have z-coordinate between 1 and 2.

I have other selections criteria but in general they have all the same structure, to know:

r_list = []
for atom in self.atom_list:
    # Some criteria here
        r_list.append(atom)
self.atom_list = r_list

So the question is: is there something I can do in order to avoid writing the above structure for each selection criteria?

If you know there is a simpler way to accomplish my purpose I'll be glad to hear it.

*WARNING*: you have supplied a mutable default argument to `Selection.__init__(atom_list)`, This is likely to make your life difficult. — SingleNegationElimination, Nov 02 '13 at 03:43
Thanks, but why is that? Any info i can use to know about this? — ezitoc, Nov 02 '13 at 11:58
yes, see this question: http://stackoverflow.com/questions/1132941/least-astonishment-in-python-the-mutable-default-argument — SingleNegationElimination, Nov 02 '13 at 12:13

MadeOfAir · Answer 1 · 2013-11-02T08:57:06.080

You might use the built-in filter() function, it does the looping automatically for you and is arguably more elegant:

def by_symbol(self, symbol):
    res = filter(lambda atom: atom.symbol == symbol, self.atom_list)
    self.atom_list.extend(res)

If you need more complex filtering, you might need to write a nested function and pass that instead of lambda. It should be a one-argument function and returns True on the correct results.

score 1 · Accepted Answer · answered Nov 02 '13 at 03:25

Here's a working example of how you could make use of built-in filter() function.

The code below also includes a few other enhancements to your classes as well as some embellishments on the idea. In particular note that theby_symbol()andby_zrange()methods end withreturn selfwhich makes it easier to print results as well as chain them together as illustrated in the example usage.

from collections import namedtuple

Point = namedtuple('Point', 'x, y, z')

class Atom(object):
    def __init__(self, symbol, x, y, z):
        self.symbol = symbol
        self.position = Point(x, y, z)

    def __repr__(self):
        return '{name}({sym!r}, {pos.x}, {pos.y}, {pos.z})'.format(
            name=self.__class__.__name__, sym=self.symbol, pos=self.position)

class Selection(object):
    def __init__(self, a_system, atom_list=None):
        if atom_list is None:
            atom_list = []
        for atom in a_system:
            atom_list.append(atom)
        self.atom_list = atom_list

    def __repr__(self):
        return '{name}({atoms})'.format(
            name=self.__class__.__name__, atoms=self.atom_list)

    def _filter(self, func):
        return filter(func, self.atom_list)

    def by_symbol(self, symbol):
        self.atom_list = self._filter(lambda a: a.symbol == symbol)
        return self

    def by_zrange(self, zmin, zmax):
        def zrange(a):
            return zmin <= a.position.z <= zmax
        self.atom_list = self._filter(zrange)
        return self

Examples of usage:

my_system = [Atom('H', 0, 1, 2),
             Atom('N', 3, 4, 5),
             Atom('C', 6, 7, 8),
             Atom('H', 9, 10, 11),]

group = Selection(my_system)
print group
print group.by_symbol('H')
print group.by_zrange(1, 2)
print
group = Selection(my_system)
print group.by_symbol('H').by_zrange(1, 2)

Output:

Selection([Atom('H', 0, 1, 2), Atom('N', 3, 4, 5), Atom('C', 6, 7, 8),
           Atom('H', 9, 10, 11)])
Selection([Atom('H', 0, 1, 2), Atom('H', 9, 10, 11)])
Selection([Atom('H', 0, 1, 2)])

Selection([Atom('H', 0, 1, 2)])

Selecting objects

2 Answers2