cell_properties.py

# Frobulator takes an Atoms object as an argument and implements `frobulate` which does stuff to the diagonal of an orthorhombic unit cell.
class Frobulator(object):
    def __init__(self, atoms):
        self.atoms = atoms
    
    # This would likely either be in the get_cell() routine or in ase.utils, depending on the direction we go
    def is_orthorhombic(self, cell):
        cr = cell.ravel()
        if np.all(np.abs(cr[[0, 4, 8]]) > 0) and np.all(np.abs(cr[[1, 2, 3, 5, 6, 7]]) < 1e-8):
            return True
        else:
            return False
        
    # Option 1: the status quo
    def frobulate1(self):
        cell = self.atoms.get_cell()
        # Check to see if cell has its default value of np.eye(3):
        if np.all(np.abs(cell - np.eye(3)) < 1e-8):
            raise NotImplementedError("Cell is undefined!")
        if not self.is_orthorhombic(cell):
            raise NotImplementedError("Frobulator only works for orthorhombic unit cells!")
        self.do_stuff(cell.diagonal())

    # Option 2: 3x3 array or None
    def frobulate2(self):
        cell = self.atoms.get_cell()
        if cell is None:
            raise NotImplementedError("Unit cell is not defined!")
        if not self.is_orthorhombic(cell):
            raise NotImplementedError("Frobulator only works for orthorhombic unit cells!")
        self.do_stuff(cell.diagonal())
        
    # Option 3: get_cell() with additional logic
    def frobulate3(self):
        # All of the magic happens in atoms.get_cell(). Exact kwarg syntax TBD.
        cell = self.atoms.get_cell(orthorhombic=True, diagonal=True)
        self.do_stuff(cell)
    
    # Option 4: cell is an ndarray subclass:
    def frobulate4(self):
        cell = self.atoms.get_cell()
        # cell.is_orthorhombic() only returns True if the cell is defined *and* orthorhombic.
        if not cell.is_orthorhombic():
            raise NotImplementedError("Frobulator only works with orthorhombic unit cells!")
        self.do_stuff(cell.diagonal())
        
    # Option 5: additional getters for specific kinds of cells:
    def frobulate5(self):
        cell = self.atoms.get_orthorhombic_cell_diagonal()
        self.do_stuff(cell)