Выбор объектов

Question

Я создал объект Atom следующим образом:

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

и Selection класс, который содержит атомы, выбранные по некоторым критериям:

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 

Так как вы можете видеть, что я могу сказать, например, :

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

, а затем сказать:

group.by_symbol('H') 

и у меня в объекте group все атомы водорода. Тогда, если я:

group.by_zrange(1, 2) 

и я буду иметь в объекте group всех атомы водорода, которые имеют Z-координату между 1 и 2.

У меня есть другие критерии отбора, но в целом они все равно структура, чтобы узнать:

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

Итак, вопрос: есть ли что-то я могу сделать для того, чтобы избежать написания выше структуру для каждого критерия отбора?

Если вы знаете, что есть более простой способ выполнить мою задачу, я буду рад услышать это.

Answer 1

Вот пример того, как вы можете использовать встроенную функцию filter().

Код ниже также включает в себя несколько других улучшений для ваших классов, а также некоторые приукрашивания идеи. В частности, обратите внимание, что методы и by_zrange() заканчиваются return self, что упрощает печать результатов, а также связывает их вместе, как показано в примере использования.

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 

Примеры использования:

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) 

Выход:

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)]) 

Answer 2

Вы можете использовать встроенный в filter() функции, он выполняет цикл автоматически для вас и, возможно, более элегантна:

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

Если вам нужна более сложная фильтрация, вам, возможно, придется написать вложенную функцию и пройти , что вместо lambda. Он должен быть функцией с одним аргументом и возвращает True с правильными результатами.