diff --git a/ase/atoms.py b/ase/atoms.py index fa8c655..3e94454 100644 --- a/ase/atoms.py +++ b/ase/atoms.py @@ -526,7 +526,7 @@ class Atoms(object): the masses argument is not given or for those elements of the masses list that are None, standard values are set.""" - if masses == 'defaults': + if isinstance(masses, str) and masses == 'defaults': masses = atomic_masses[self.arrays['numbers']] elif isinstance(masses, (list, tuple)): newmasses = [] diff --git a/ase/constraints.py b/ase/constraints.py index 232223d..4d53229 100644 --- a/ase/constraints.py +++ b/ase/constraints.py @@ -103,17 +103,18 @@ class FixAtoms(FixConstraint): raise ValueError('Use only one of "indices" and "mask".') if mask is not None: - self.index = np.asarray(mask, bool) + indices = np.arange(len(mask))[np.asarray(mask, bool)] else: - # Check for duplicates + # Check for duplicates: srt = np.sort(indices) + assert (srt == indices).all() for i in range(len(indices) - 1): if srt[i] == srt[i + 1]: raise ValueError( 'FixAtoms: The indices array contained duplicates. ' 'Perhaps you wanted to specify a mask instead, but ' 'forgot the mask= keyword.') - self.index = np.asarray(indices, int) + self.index = np.asarray(indices, int) if self.index.ndim != 1: raise ValueError('Wrong argument to FixAtoms class!') @@ -126,29 +127,20 @@ class FixAtoms(FixConstraint): def index_shuffle(self, atoms, ind): # See docstring of superclass - if self.index.dtype == bool: - self.index = self.index[ind] - else: - index = [] - for new, old in slice2enlist(ind, len(atoms)): - if old in self.index: - index.append(new) - if len(index) == 0: - raise IndexError('All indices in FixAtoms not part of slice') - self.index = np.asarray(index, int) + index = [] + for new, old in slice2enlist(ind, len(atoms)): + if old in self.index: + index.append(new) + if len(index) == 0: + raise IndexError('All indices in FixAtoms not part of slice') + self.index = np.asarray(index, int) def __repr__(self): - if self.index.dtype == bool: - return 'FixAtoms(mask=%s)' % ints2string(self.index.astype(int)) return 'FixAtoms(indices=%s)' % ints2string(self.index) def todict(self): - dct = {'name': 'FixAtoms'} - if self.index.dtype == bool: - dct['kwargs'] = {'mask': self.index} - else: - dct['kwargs'] = {'indices': self.index} - return dct + return {'name': 'FixAtoms', + 'kwargs': {'indices': self.index}} def repeat(self, m, n): i0 = 0 @@ -160,31 +152,22 @@ class FixAtoms(FixConstraint): for m1 in range(m[1]): for m0 in range(m[0]): i1 = i0 + n - if self.index.dtype == bool: - index_new.extend(self.index) - else: - index_new += [i + natoms for i in self.index] + index_new += [i + natoms for i in self.index] i0 = i1 natoms += n - if self.index.dtype == bool: - self.index = np.asarray(index_new, bool) - else: - self.index = np.asarray(index_new, int) + self.index = np.asarray(index_new, int) return self def delete_atom(self, ind): """ Removes atom number ind from the index array, if present. Required for removing atoms with existing FixAtoms constraints. """ - if self.index.dtype == bool: - self.index = np.delete(self.index, ind) - else: - if ind in self.index: - i = list(self.index).index(ind) - self.index = np.delete(self.index, i) - for i in range(len(self.index)): - if self.index[i] >= ind: - self.index[i] -= 1 + if ind in self.index: + i = list(self.index).index(ind) + self.index = np.delete(self.index, i) + for i in range(len(self.index)): + if self.index[i] >= ind: + self.index[i] -= 1 def ints2string(x, threshold=10): diff --git a/ase/test/build.py b/ase/test/build.py index 690c042..041a9f8 100644 --- a/ase/test/build.py +++ b/ase/test/build.py @@ -4,15 +4,15 @@ from ase import Atoms, Atom a = Atoms([Atom('Cu')]) a.positions[:] += 1.0 print(a.get_positions(), a.positions) -a=a+a -a+=a +a = a + a +a += a a.append(Atom('C')) a += Atoms([]) a += Atom('H', magmom=1) print(a.get_initial_magnetic_moments()) print(a[0].number) -print(a[[0,1]].get_atomic_numbers()) -print(a[np.array([1,1,0,0,1], bool)].get_atomic_numbers()) +print(a[[0, 1]].get_atomic_numbers()) +print(a[np.array([1, 1, 0, 0, 1, 0], bool)].get_atomic_numbers()) print(a[::2].get_atomic_numbers()) print(a.get_chemical_symbols()) del a[2] diff --git a/ase/test/repeat_FixAtoms.py b/ase/test/repeat_FixAtoms.py index 865b39e..c8bd227 100644 --- a/ase/test/repeat_FixAtoms.py +++ b/ase/test/repeat_FixAtoms.py @@ -4,36 +4,31 @@ from ase.constraints import FixAtoms N = 2 atoms = molecule('CO2') -atoms.set_cell((15.,15.,15.)) +atoms.set_cell((15, 15, 15)) -print('indices method') +# Indices method: atomsi = atoms.copy() -atomsi.set_constraint(FixAtoms(indices=[0,])) -atomsi = atomsi.repeat((N,1,1)) +atomsi.set_constraint(FixAtoms(indices=[0])) +atomsi = atomsi.repeat((N, 1, 1)) -atomsiref = atoms.copy().repeat((N,1,1)) -atomsiref.set_constraint(FixAtoms(indices=[0, N + 1])) +atomsiref = atoms.copy().repeat((N, 1, 1)) +atomsiref.set_constraint(FixAtoms(indices=list(range(0, 3 * N, 3)))) lcatomsi = list(atomsi.constraints[0].index) lcatomsiref = list(atomsiref.constraints[0].index) assert lcatomsi == lcatomsiref -print('mask method') +# Mask method: atomsm = atoms.copy() atomsm.set_constraint(FixAtoms(mask=[True, False, False])) -atomsm = atomsm.repeat((N,1,1)) +atomsm = atomsm.repeat((N, 1, 1)) -atomsmref = atoms.copy().repeat((N,1,1)) +atomsmref = atoms.copy().repeat((N, 1, 1)) atomsmref.set_constraint(FixAtoms(mask=[True, False, False] * N)) lcatomsm = list(atomsm.constraints[0].index) lcatomsmref = list(atomsmref.constraints[0].index) assert lcatomsm == lcatomsmref - -# http://stackoverflow.com/questions/3873361/finding-multiple-occurrences-of-a-string-within-a-string-in-python - -lcatomsm2i = [n for (n, e) in enumerate(lcatomsm) if e == True] - -assert lcatomsm2i == lcatomsi +assert lcatomsm == lcatomsi