Coverage for /builds/ase/ase/ase/atom.py: 94.44%
126 statements
« prev ^ index » next coverage.py v7.5.3, created at 2025-08-02 00:12 +0000
1# fmt: off
3"""This module defines the Atom object."""
5import numpy as np
7from ase.data import atomic_masses, atomic_numbers, chemical_symbols
9# Singular, plural, default value:
10names = {'position': ('positions', np.zeros(3)),
11 'number': ('numbers', 0),
12 'tag': ('tags', 0),
13 'momentum': ('momenta', np.zeros(3)),
14 'mass': ('masses', None),
15 'magmom': ('initial_magmoms', 0.0),
16 'charge': ('initial_charges', 0.0)}
19def atomproperty(name, doc):
20 """Helper function to easily create Atom attribute property."""
22 def getter(self):
23 return self.get(name)
25 def setter(self, value):
26 self.set(name, value)
28 def deleter(self):
29 self.delete(name)
31 return property(getter, setter, deleter, doc)
34def abcproperty(index):
35 """Helper function to easily create Atom ABC-property."""
37 def getter(self):
38 return self.scaled_position[index]
40 def setter(self, value):
41 # We can't just do self.scaled_position[i] = value
42 # because scaled_position is a new buffer, not a view into
43 # something we can write back to.
44 # This is a clear bug!
45 spos = self.scaled_position
46 spos[index] = value
47 self.scaled_position = spos
49 return property(getter, setter, doc='ABC'[index] + '-coordinate')
52def xyzproperty(index):
53 """Helper function to easily create Atom XYZ-property."""
55 def getter(self):
56 return self.position[index]
58 def setter(self, value):
59 self.position[index] = value
61 return property(getter, setter, doc='XYZ'[index] + '-coordinate')
64class Atom:
65 """Class for representing a single atom.
67 Parameters:
69 symbol: str or int
70 Can be a chemical symbol (str) or an atomic number (int).
71 position: sequence of 3 floats
72 Atomic position.
73 tag: int
74 Special purpose tag.
75 momentum: sequence of 3 floats
76 Momentum for atom.
77 mass: float
78 Atomic mass in atomic units.
79 magmom: float or 3 floats
80 Magnetic moment.
81 charge: float
82 Atomic charge.
83 """
84 __slots__ = ['data', 'atoms', 'index']
86 def __init__(self, symbol='X', position=(0, 0, 0),
87 tag=None, momentum=None, mass=None,
88 magmom=None, charge=None,
89 atoms=None, index=None):
91 self.data = d = {}
93 if atoms is None:
94 # This atom is not part of any Atoms object:
95 if isinstance(symbol, str):
96 d['number'] = atomic_numbers[symbol]
97 else:
98 d['number'] = symbol
99 d['position'] = np.array(position, float)
100 d['tag'] = tag
101 if momentum is not None:
102 momentum = np.array(momentum, float)
103 d['momentum'] = momentum
104 d['mass'] = mass
105 if magmom is not None:
106 magmom = np.array(magmom, float)
107 d['magmom'] = magmom
108 d['charge'] = charge
110 self.index = index
111 self.atoms = atoms
113 @property
114 def scaled_position(self):
115 pos = self.position
116 spos = self.atoms.cell.scaled_positions(pos[np.newaxis])
117 return spos[0]
119 @scaled_position.setter
120 def scaled_position(self, value):
121 pos = self.atoms.cell.cartesian_positions(value)
122 self.position = pos
124 def __repr__(self):
125 s = f"Atom('{self.symbol}', {list(self.position)}"
126 for name in ['tag', 'momentum', 'mass', 'magmom', 'charge']:
127 value = self.get_raw(name)
128 if value is not None:
129 if isinstance(value, np.ndarray):
130 value = value.tolist()
131 s += f', {name}={value}'
132 if self.atoms is None:
133 s += ')'
134 else:
135 s += ', index=%d)' % self.index
136 return s
138 def cut_reference_to_atoms(self):
139 """Cut reference to atoms object."""
140 for name in names:
141 self.data[name] = self.get_raw(name)
142 self.index = None
143 self.atoms = None
145 def get_raw(self, name):
146 """Get name attribute, return None if not explicitly set."""
147 if name == 'symbol':
148 return chemical_symbols[self.get_raw('number')]
150 if self.atoms is None:
151 return self.data[name]
153 plural = names[name][0]
154 if plural in self.atoms.arrays:
155 return self.atoms.arrays[plural][self.index]
156 else:
157 return None
159 def get(self, name):
160 """Get name attribute, return default if not explicitly set."""
161 value = self.get_raw(name)
162 if value is None:
163 if name == 'mass':
164 value = atomic_masses[self.number]
165 else:
166 value = names[name][1]
167 return value
169 def set(self, name, value):
170 """Set name attribute to value."""
171 if name == 'symbol':
172 name = 'number'
173 value = atomic_numbers[value]
175 if self.atoms is None:
176 assert name in names
177 self.data[name] = value
178 else:
179 plural, default = names[name]
180 if plural in self.atoms.arrays:
181 array = self.atoms.arrays[plural]
182 if name == 'magmom' and array.ndim == 2:
183 assert len(value) == 3
184 array[self.index] = value
185 else:
186 if name == 'magmom' and np.asarray(value).ndim == 1:
187 array = np.zeros((len(self.atoms), 3))
188 elif name == 'mass':
189 array = self.atoms.get_masses()
190 else:
191 default = np.asarray(default)
192 array = np.zeros((len(self.atoms),) + default.shape,
193 default.dtype)
194 array[self.index] = value
195 self.atoms.new_array(plural, array)
197 def delete(self, name):
198 """Delete name attribute."""
199 assert self.atoms is None
200 assert name not in ['number', 'symbol', 'position']
201 self.data[name] = None
203 symbol = atomproperty('symbol', 'Chemical symbol')
204 number = atomproperty('number', 'Atomic number')
205 position = atomproperty('position', 'XYZ-coordinates')
206 tag = atomproperty('tag', 'Integer tag')
207 momentum = atomproperty('momentum', 'XYZ-momentum')
208 mass = atomproperty('mass', 'Atomic mass')
209 magmom = atomproperty('magmom', 'Initial magnetic moment')
210 charge = atomproperty('charge', 'Initial atomic charge')
211 x = xyzproperty(0)
212 y = xyzproperty(1)
213 z = xyzproperty(2)
215 a = abcproperty(0)
216 b = abcproperty(1)
217 c = abcproperty(2)