You can not select more than 25 topics
Topics must start with a letter or number, can include dashes ('-') and can be up to 35 characters long.
215 lines
6.5 KiB
215 lines
6.5 KiB
"""This module defines the Atom object.""" |
|
|
|
import numpy as np |
|
|
|
from ase.data import atomic_numbers, chemical_symbols, atomic_masses |
|
|
|
# Singular, plural, default value: |
|
names = {'position': ('positions', np.zeros(3)), |
|
'number': ('numbers', 0), |
|
'tag': ('tags', 0), |
|
'momentum': ('momenta', np.zeros(3)), |
|
'mass': ('masses', None), |
|
'magmom': ('initial_magmoms', 0.0), |
|
'charge': ('initial_charges', 0.0)} |
|
|
|
|
|
def atomproperty(name, doc): |
|
"""Helper function to easily create Atom attribute property.""" |
|
|
|
def getter(self): |
|
return self.get(name) |
|
|
|
def setter(self, value): |
|
self.set(name, value) |
|
|
|
def deleter(self): |
|
self.delete(name) |
|
|
|
return property(getter, setter, deleter, doc) |
|
|
|
|
|
def abcproperty(index): |
|
"""Helper function to easily create Atom ABC-property.""" |
|
|
|
def getter(self): |
|
return self.scaled_position[index] |
|
|
|
def setter(self, value): |
|
# We can't just do self.scaled_position[i] = value |
|
# because scaled_position is a new buffer, not a view into |
|
# something we can write back to. |
|
# This is a clear bug! |
|
spos = self.scaled_position |
|
spos[index] = value |
|
self.scaled_position = spos |
|
|
|
return property(getter, setter, doc='ABC'[index] + '-coordinate') |
|
|
|
|
|
def xyzproperty(index): |
|
"""Helper function to easily create Atom XYZ-property.""" |
|
|
|
def getter(self): |
|
return self.position[index] |
|
|
|
def setter(self, value): |
|
self.position[index] = value |
|
|
|
return property(getter, setter, doc='XYZ'[index] + '-coordinate') |
|
|
|
|
|
class Atom: |
|
"""Class for representing a single atom. |
|
|
|
Parameters: |
|
|
|
symbol: str or int |
|
Can be a chemical symbol (str) or an atomic number (int). |
|
position: sequence of 3 floats |
|
Atomic position. |
|
tag: int |
|
Special purpose tag. |
|
momentum: sequence of 3 floats |
|
Momentum for atom. |
|
mass: float |
|
Atomic mass in atomic units. |
|
magmom: float or 3 floats |
|
Magnetic moment. |
|
charge: float |
|
Atomic charge. |
|
""" |
|
__slots__ = ['data', 'atoms', 'index'] |
|
|
|
def __init__(self, symbol='X', position=(0, 0, 0), |
|
tag=None, momentum=None, mass=None, |
|
magmom=None, charge=None, |
|
atoms=None, index=None): |
|
|
|
self.data = d = {} |
|
|
|
if atoms is None: |
|
# This atom is not part of any Atoms object: |
|
if isinstance(symbol, str): |
|
d['number'] = atomic_numbers[symbol] |
|
else: |
|
d['number'] = symbol |
|
d['position'] = np.array(position, float) |
|
d['tag'] = tag |
|
if momentum is not None: |
|
momentum = np.array(momentum, float) |
|
d['momentum'] = momentum |
|
d['mass'] = mass |
|
if magmom is not None: |
|
magmom = np.array(magmom, float) |
|
d['magmom'] = magmom |
|
d['charge'] = charge |
|
|
|
self.index = index |
|
self.atoms = atoms |
|
|
|
@property |
|
def scaled_position(self): |
|
pos = self.position |
|
spos = self.atoms.cell.scaled_positions(pos[np.newaxis]) |
|
return spos[0] |
|
|
|
@scaled_position.setter |
|
def scaled_position(self, value): |
|
pos = self.atoms.cell.cartesian_positions(value) |
|
self.position = pos |
|
|
|
def __repr__(self): |
|
s = "Atom('%s', %s" % (self.symbol, list(self.position)) |
|
for name in ['tag', 'momentum', 'mass', 'magmom', 'charge']: |
|
value = self.get_raw(name) |
|
if value is not None: |
|
if isinstance(value, np.ndarray): |
|
value = value.tolist() |
|
s += ', %s=%s' % (name, value) |
|
if self.atoms is None: |
|
s += ')' |
|
else: |
|
s += ', index=%d)' % self.index |
|
return s |
|
|
|
def cut_reference_to_atoms(self): |
|
"""Cut reference to atoms object.""" |
|
for name in names: |
|
self.data[name] = self.get_raw(name) |
|
self.index = None |
|
self.atoms = None |
|
|
|
def get_raw(self, name): |
|
"""Get name attribute, return None if not explicitly set.""" |
|
if name == 'symbol': |
|
return chemical_symbols[self.get_raw('number')] |
|
|
|
if self.atoms is None: |
|
return self.data[name] |
|
|
|
plural = names[name][0] |
|
if plural in self.atoms.arrays: |
|
return self.atoms.arrays[plural][self.index] |
|
else: |
|
return None |
|
|
|
def get(self, name): |
|
"""Get name attribute, return default if not explicitly set.""" |
|
value = self.get_raw(name) |
|
if value is None: |
|
if name == 'mass': |
|
value = atomic_masses[self.number] |
|
else: |
|
value = names[name][1] |
|
return value |
|
|
|
def set(self, name, value): |
|
"""Set name attribute to value.""" |
|
if name == 'symbol': |
|
name = 'number' |
|
value = atomic_numbers[value] |
|
|
|
if self.atoms is None: |
|
assert name in names |
|
self.data[name] = value |
|
else: |
|
plural, default = names[name] |
|
if plural in self.atoms.arrays: |
|
array = self.atoms.arrays[plural] |
|
if name == 'magmom' and array.ndim == 2: |
|
assert len(value) == 3 |
|
array[self.index] = value |
|
else: |
|
if name == 'magmom' and np.asarray(value).ndim == 1: |
|
array = np.zeros((len(self.atoms), 3)) |
|
elif name == 'mass': |
|
array = self.atoms.get_masses() |
|
else: |
|
default = np.asarray(default) |
|
array = np.zeros((len(self.atoms),) + default.shape, |
|
default.dtype) |
|
array[self.index] = value |
|
self.atoms.new_array(plural, array) |
|
|
|
def delete(self, name): |
|
"""Delete name attribute.""" |
|
assert self.atoms is None |
|
assert name not in ['number', 'symbol', 'position'] |
|
self.data[name] = None |
|
|
|
symbol = atomproperty('symbol', 'Chemical symbol') |
|
number = atomproperty('number', 'Atomic number') |
|
position = atomproperty('position', 'XYZ-coordinates') |
|
tag = atomproperty('tag', 'Integer tag') |
|
momentum = atomproperty('momentum', 'XYZ-momentum') |
|
mass = atomproperty('mass', 'Atomic mass') |
|
magmom = atomproperty('magmom', 'Initial magnetic moment') |
|
charge = atomproperty('charge', 'Initial atomic charge') |
|
x = xyzproperty(0) |
|
y = xyzproperty(1) |
|
z = xyzproperty(2) |
|
|
|
a = abcproperty(0) |
|
b = abcproperty(1) |
|
c = abcproperty(2)
|
|
|