• Home
  • Features
  • Pricing
  • Docs
  • Announcements
  • Sign In

maurergroup / dfttoolkit / 15065810745

16 May 2025 09:56AM UTC coverage: 30.583% (+8.8%) from 21.747%
15065810745

Pull #59

github

b2c890
web-flow
Merge f52b00038 into e895278a4
Pull Request #59: Vibrations refactor

1164 of 3806 relevant lines covered (30.58%)

0.31 hits per line

Source File
Press 'n' to go to next uncovered line, 'b' for previous

21.03
dfttoolkit/geometry.py
1
import ast
1✔
2
import colorsys
1✔
3
import copy
1✔
4
import itertools
1✔
5
import re
1✔
6
import warnings
1✔
7
from collections import defaultdict
1✔
8
from collections.abc import Iterable
1✔
9
from copy import deepcopy
1✔
10
from fractions import Fraction
1✔
11
from typing import List, Tuple, Union
1✔
12

13
import ase
1✔
14
import matplotlib as mpl
1✔
15
import matplotlib.cm as cmx
1✔
16
import matplotlib.colors
1✔
17
import matplotlib.pyplot as plt
1✔
18
import networkx as nx
1✔
19
import numpy as np
1✔
20
import numpy.typing as npt
1✔
21
import scipy.linalg
1✔
22
import scipy.spatial
1✔
23
import scipy.spatial.distance
1✔
24
import spglib
1✔
25
from ase.constraints import FixAtoms
1✔
26
from scipy.spatial import distance_matrix
1✔
27

28
import dfttoolkit.utils.math_utils as utils
1✔
29
from dfttoolkit.utils import units
1✔
30
from dfttoolkit.utils.periodic_table import PeriodicTable
1✔
31

32

33
class Geometry:
1✔
34
    """
35
    This class represents a geometry file for (in principle) any DFT code.
36
    In practice it has only been fully implemented for FHI-aims geometry.in
37
    files.
38

39
    """
40

41
    def __init__(self, filename: Union[str, None] = None):
1✔
42
        """
43
        Filename : str
44
            Path to text file of geometry.
45

46
        center: dict
47
            atom indices and linear combination of them to define the center of
48
            a molecule. Used only for mapping to first unit cell.
49
            Example: if the center should be the middle between the first three
50
            atoms, center would be {1:1/3,2:1/3,3:1/3}
51
        """
52
        self.species = []
1✔
53
        self.lattice_vectors = np.zeros([3, 3])
1✔
54
        self.comment_lines = []
1✔
55
        self.constrain_relax = np.zeros([0, 3], bool)
1✔
56
        self.external_force = np.zeros([0, 3], np.float64)
1✔
57
        self.calculate_friction = np.zeros([0], np.float64)
1✔
58
        self.initial_moment = []
1✔
59
        self.initial_charge = []
1✔
60
        self.name = filename
1✔
61
        self.center = None
1✔
62
        self.energy = None
1✔
63
        self.forces = None
1✔
64
        self.hessian = None
1✔
65
        self.geometry_parts = []  # list of lists: indices of each geometry part
1✔
66
        self.geometry_part_descriptions = []  # list of strings:  name of each geometry part
1✔
67
        self.symmetry_axes = None
1✔
68
        self.inversion_index = None
1✔
69
        self.vacuum_level = None
1✔
70
        self.multipoles = []
1✔
71
        self.homogeneous_field = None
1✔
72
        self.read_as_fractional_coords = False
1✔
73
        self.symmetry_params = None
1✔
74
        self.n_symmetry_params = None
1✔
75
        self.symmetry_LVs = None
1✔
76
        self.symmetry_frac_coords = None  # symmetry_frac_coords should have str values, not float, to include the parameters
1✔
77
        self.original_lattice_vectors = None  # Save the original lattice vectors if the geometry is periodically replicated
1✔
78

79
        if filename is None:
1✔
80
            self.n_atoms = 0
1✔
81
            self.coords = np.zeros([self.n_atoms, 3])
1✔
82
        else:
83
            self.read_from_file(filename)
1✔
84

85
        self.periodic_table = PeriodicTable
1✔
86

87
    def __eq__(self, other):
1✔
88
        if len(self) != len(other):
1✔
89
            equal = False
×
90
        else:
91
            equal = np.allclose(self.coords, other.coords)
1✔
92
            equal = equal and np.allclose(self.lattice_vectors, other.lattice_vectors)
1✔
93
            equal = equal and self.species == other.species
1✔
94
        return equal
1✔
95

96
    def __len__(self):
1✔
97
        return self.n_atoms
1✔
98

99
    def __add__(self, other_geometry):
1✔
100
        geom = copy.deepcopy(self)
×
101
        geom += other_geometry
×
102
        return geom
×
103

104
    def __iadd__(self, other_geometry):
1✔
105
        self.add_geometry(other_geometry)
×
106
        return self
×
107

108
    def get_instance_of_other_type(self, geometry_type):
1✔
109
        if geometry_type == "aims":
1✔
110
            new_geometry = AimsGeometry()
1✔
111
        elif geometry_type == "vasp":
×
112
            new_geometry = VaspGeometry()
×
113
        elif geometry_type == "xyz":
×
114
            new_geometry = XYZGeometry()
×
115
        elif geometry_type == "xsf":
×
116
            new_geometry = XSFGeometry()
×
117
        else:
118
            raise ValueError(f'Type "{geometry_type}" is not availlable.')
×
119

120
        new_geometry.__dict__ = self.__dict__
1✔
121
        return new_geometry
1✔
122

123
    ###########################################################################
124
    #                             INPUT PARSER                                #
125
    ###########################################################################
126
    def read_from_file(self, filename: str) -> None:
1✔
127
        """
128
        Parses a geometry file
129

130
        Parameters
131
        ----------
132
        filename : str
133
            Path to input file.
134

135
        Returns
136
        -------
137
        None.
138

139
        """
140
        with open(filename) as f:
1✔
141
            text = f.read()
1✔
142

143
        self.parse(text)
1✔
144

145
    def parse(self, text):
1✔
146
        raise NotImplementedError
×
147

148
    def add_top_comment(self, comment_string: str) -> None:
1✔
149
        """
150
        Add comments that are saved at the top of the geometry file.
151

152
        Parameters
153
        ----------
154
        comment_string : str
155
            Comment string.
156

157
        Returns
158
        -------
159
        None.
160
        """
161
        lines = comment_string.split("\n")
×
162
        for line in lines:
×
163
            commented_line = "# " + line if not line.startswith("#") else line
×
164
            self.comment_lines.append(commented_line)
×
165

166
    ###########################################################################
167
    #                             OUTPUT PARSER                               #
168
    ###########################################################################
169
    def save_to_file(self, filename, **kwargs):
1✔
170
        geometry_type = get_file_format_from_ending(filename)
1✔
171

172
        new_geometry = self.get_instance_of_other_type(geometry_type)
1✔
173

174
        text = new_geometry.get_text(**kwargs)
1✔
175

176
        # Enforce linux file ending, even if running on windows machine by
177
        # using binary mode
178
        with open(filename, "w", newline="\n") as f:
1✔
179
            f.write(text)
1✔
180

181
    def get_text(self, **kwargs):
1✔
182
        raise NotImplementedError
×
183

184
    ###########################################################################
185
    #                          Data exchange with ASE                         #
186
    ###########################################################################
187
    def get_from_ase_atoms_object(self, atoms) -> None:
1✔
188
        """
189
        Reads an ASE.Atoms object. Taken from ase.io.aims and adapted. Only basic features are implemented.
190

191
        Parameters
192
        ----------
193
        atoms : ASE atoms object
194
            Atoms object from ASE that should be converted into geometry.
195

196
        Raises
197
        ------
198
        RuntimeError
199
            If atoms object is erroneous.
200

201
        Returns
202
        -------
203
        None.
204

205
        """
206
        if isinstance(atoms, (list, tuple)):
×
207
            if len(atoms) > 1:
×
208
                raise RuntimeError(
×
209
                    "Don't know how to save more than one image to FHI-aims input"
210
                )
211
            atoms = atoms[0]  # pyright:ignore
×
212

213
        if atoms.get_pbc().any():
×
214
            self.lattice_vectors = np.array(atoms.get_cell())
×
215

216
        fix_cart = np.zeros([len(atoms), 3])
×
217
        if atoms.constraints:
×
218
            for constr in atoms.constraints:
×
219
                if isinstance(constr, FixAtoms):
×
220
                    fix_cart[constr.index] = [True, True, True]
×
221
        constrain_relax = fix_cart
×
222

223
        coords = []
×
224
        species_list = []
×
225
        for i, atom in enumerate(atoms):
×
226
            species = atom.symbol
×
227
            if isinstance(species, int):
×
228
                species = self.periodic_table.get_chemical_symbol(species)
×
229
            species_list.append(species)
×
230
            coords.append(atom.position)
×
231
        coords = np.array(coords)
×
232
        self.add_atoms(coords, species_list, constrain_relax)
×
233

234
    def get_as_ase(self) -> None:
1✔
235
        """
236
        Convert geometry file to ASE object.
237

238
        Returns
239
        -------
240
        None.
241

242
        """
243
        # atoms_string = ""
244
        atom_coords = []
×
245
        atom_numbers = []
×
246
        atom_constraints = []
×
247
        for i in range(self.n_atoms):
×
248
            # Do not export 'emptium" atoms
249
            if self.species[i] != "Em":
×
250
                atom_coords.append(self.coords[i, :])
×
251
                atom_numbers.append(
×
252
                    self.periodic_table.get_atomic_number(self.species[i])
253
                )
254

255
                if np.any(self.constrain_relax[i]):
×
256
                    atom_constraints.append(i)
×
257

258
        ase_system = ase.Atoms(numbers=atom_numbers, positions=atom_coords)
×
259
        ase_system.cell = self.lattice_vectors
×
260

261
        c = FixAtoms(indices=atom_constraints)
×
262
        ase_system.set_constraint(c)
×
263

264
        if not np.sum(self.lattice_vectors) == 0.0:
×
265
            ase_system.pbc = [1, 1, 1]
×
266

267
        if self.energy is not None:
×
268
            ase_system.info["energy"] = self.energy
×
269

270
        if self.forces is not None:
×
271
            ase_system.arrays["forces"] = self.forces
×
272

273
        return ase_system
×
274

275
    ###########################################################################
276
    #                    Adding and removing atoms (in place)                 #
277
    ###########################################################################
278
    def add_atoms(
1✔
279
        self,
280
        cartesian_coords,
281
        species,
282
        constrain_relax=None,
283
        initial_moment=None,
284
        initial_charge=None,
285
        external_force=None,
286
        calculate_friction=None,
287
    ) -> None:
288
        """
289
        Add additional atoms to the current geometry file.
290

291
        Parameters
292
        ----------
293
        cartesion_coords : List of numpy arrays of shape [nx3]
294
            coordinates of new atoms
295
        species : list of strings
296
            element symbol for each atom
297
        constrain_relax : list of lists of bools (optional)
298
            [bool,bool,bool] (for [x,y,z] axis) for all atoms that should be
299
            constrained during a geometry relaxation
300

301
        Retruns
302
        -------
303
        None.
304

305
        """
306
        if constrain_relax is None or len(constrain_relax) == 0:
1✔
307
            constrain_relax = np.zeros([len(species), 3], bool)
×
308
        if external_force is None or len(external_force) == 0:
1✔
309
            external_force = np.zeros([len(species), 3], np.float64)
1✔
310
        if calculate_friction is None:
1✔
311
            calculate_friction = np.array([False] * len(species))
1✔
312
        if initial_moment is None:
1✔
313
            initial_moment = [0.0] * len(species)
1✔
314
        if initial_charge is None:
1✔
315
            initial_charge = [0.0] * len(species)
1✔
316
        # TODO: this should not be necessary as self.coords should always be a np.array
317
        if not hasattr(self, "coords") or self.coords is None:
1✔
318
            assert isinstance(cartesian_coords, np.ndarray)
×
319
            self.coords = cartesian_coords
×
320
        else:
321
            # make sure that coords are 2D
322
            cartesian_coords = np.atleast_2d(cartesian_coords)
1✔
323
            self.coords = np.concatenate((self.coords, cartesian_coords), axis=0)
1✔
324
        self.species += species
1✔
325
        self.n_atoms = self.coords.shape[0]
1✔
326

327
        self.constrain_relax = np.concatenate(
1✔
328
            (self.constrain_relax, constrain_relax), axis=0
329
        )
330
        self.external_force = np.concatenate(
1✔
331
            (self.external_force, external_force), axis=0
332
        )
333
        self.calculate_friction = np.concatenate(
1✔
334
            (self.calculate_friction, calculate_friction)
335
        )
336
        self.initial_moment += initial_moment
1✔
337
        self.initial_charge += initial_charge
1✔
338

339
    def add_geometry(self, geometry) -> None:
1✔
340
        """
341
        Adds full geometry to initial geometry.
342

343
        Parameters
344
        ----------
345
        geometry : Instance of geometry
346
            New geometry to be added to current geometry.
347

348
        Returns
349
        -------
350
        None.
351

352
        """
353
        # check parts: (needs to be done before adding atoms to self)
354
        if hasattr(self, "geometry_parts") and hasattr(geometry, "geometry_parts"):
×
355
            for part, name in zip(
×
356
                geometry.geometry_parts, geometry.geometry_part_descriptions
357
            ):
358
                if len(part) > 0:
×
359
                    self.geometry_parts.append([i + self.n_atoms for i in part])
×
360
                    self.geometry_part_descriptions.append(name)
×
361

362
        # some lines of code in order to preserve backwards compatibility
363
        if not hasattr(geometry, "external_force"):
×
364
            geometry.external_force = np.zeros([0, 3], np.float64)
×
365
        self.add_atoms(
×
366
            geometry.coords,
367
            geometry.species,
368
            constrain_relax=geometry.constrain_relax,
369
            initial_moment=geometry.initial_moment,
370
            initial_charge=geometry.initial_charge,
371
            external_force=geometry.external_force,
372
            calculate_friction=geometry.calculate_friction,
373
        )
374

375
        # check lattice vectors:
376
        # g has lattice and self not:
377
        if not np.any(self.lattice_vectors) and np.any(geometry.lattice_vectors):
×
378
            self.lattice_vectors = np.copy(geometry.lattice_vectors)
×
379

380
        # both have lattice vectors:
381
        elif np.any(self.lattice_vectors) and np.any(geometry.lattice_vectors):
×
382
            warnings.warn(
×
383
                "Caution: The lattice vectors of the first file will be used!"
384
            )
385

386
        # add multipoles
387
        self.add_multipoles(geometry.multipoles)
×
388

389
        # check center:
390
        # g has center and self not:
391
        if hasattr(self, "center") and hasattr(geometry, "center"):
×
392
            if self.center is None and geometry.center is not None:
×
393
                self.center = geometry.center.copy()
×
394
            # both have a center:
395
            elif self.center is not None and geometry.center is not None:
×
396
                warnings.warn("Caution: The center of the first file will be used!")
×
397

398
    def add_multipoles(self, multipoles: Union[List[float], List[list]]) -> None:
1✔
399
        """
400
        Adds multipoles to the the geometry.
401

402
        Parameters
403
        ----------
404
        multipoles: Union[List[float], List[list]]
405
            Each multipole is defined as a list -> [x, y, z, order, charge]
406
            With x,y,z cartesian coordinates
407
            order: 0 for monopoles, 1 for dipoles
408
            charge: charge
409

410
        Returns
411
        -------
412
        None
413

414
        """
415
        # if multiple multipoles are given: indented lists
416
        if len(multipoles) == 0:
×
417
            return
×
418
        if isinstance(multipoles[0], list):
×
419
            for x in multipoles:
×
420
                self.multipoles.append(x)
×
421
        # else: one list
422
        else:
423
            self.multipoles.append(multipoles)
×
424

425
    def remove_atoms(self, atom_inds: npt.NDArray[np.int64]) -> None:
1✔
426
        """
427
        Remove atoms with indices atom_inds. If no indices are specified, all
428
        atoms are removed.
429

430
        Parameters
431
        ----------
432
        atom_inds : np.array
433
            Indices of atoms to be removed.
434

435
        Returns
436
        -------
437
        None
438

439
        """
440
        if hasattr(self, "geometry_parts") and len(self.geometry_parts) > 0:
×
441
            warnings.warn(
×
442
                "CAUTION: geometry_parts indices are not updated after atom"
443
                "deletion!!\n You are welcome to implement this!!"
444
            )
445
        mask = np.ones(len(self.species), dtype=bool)
×
446
        mask[atom_inds] = False
×
447

448
        self.species = list(np.array(self.species)[mask])
×
449
        self.constrain_relax = self.constrain_relax[mask, :]
×
450
        self.external_force = self.external_force[mask, :]
×
451
        self.calculate_friction = self.calculate_friction[mask]
×
452
        self.coords = self.coords[mask, :]
×
453
        self.n_atoms = len(self.constrain_relax)
×
454

455
        if hasattr(self, "hessian") and self.hessian is not None:
×
456
            flat_mask = np.kron(mask, np.ones(3, dtype=bool))
×
457
            new_dim = np.sum(flat_mask)
×
458
            a, b = np.meshgrid(flat_mask, flat_mask)
×
459
            hess_mask = np.logical_and(a, b)
×
460
            new_hessian = self.hessian[hess_mask].reshape(new_dim, new_dim)
×
461
            self.hessian = new_hessian
×
462

463
    def remove_atoms_by_species(self, species: str) -> None:
1✔
464
        """
465
        Removes all atoms of a given species.
466

467
        Parameters
468
        ----------
469
        species : str
470
            Atom species to be removed.
471

472
        Returns
473
        -------
474
        None
475

476
        """
477
        L = np.array(self.species) == species
×
478
        atom_inds = np.where(L)[0]
×
479
        self.remove_atoms(atom_inds)
×
480

481
    def remove_constrained_atoms(self) -> None:
1✔
482
        """
483
        Remove all atoms where all coordinates are constrained.
484

485
        Returns
486
        -------
487
        None
488

489
        """
490
        remove_inds = self.get_constrained_atoms()
×
491
        self.remove_atoms(remove_inds)
×
492

493
    def remove_unconstrained_atoms(self):
1✔
494
        """
495
        Remove all atoms where all coordinates are constrained.
496

497
        Returns
498
        -------
499
        None
500

501
        """
502
        remove_inds = self.get_unconstrained_atoms()
×
503
        self.remove_atoms(remove_inds)
×
504

505
    def truncate(self, n_atoms: int) -> None:
1✔
506
        """
507
        Keep only the first n_atoms atoms
508

509
        Parameters
510
        ----------
511
        n_atoms : int
512
            Number of atoms to be kept.
513

514
        Returns
515
        -------
516
        None
517

518
        """
519
        self.species = self.species[:n_atoms]
×
520
        self.constrain_relax = self.constrain_relax[:n_atoms]
×
521
        self.external_force = self.external_force[:n_atoms]
×
522
        self.calculate_friction = self.calculate_friction[:n_atoms]
×
523
        self.coords = self.coords[:n_atoms, :]
×
524
        self.n_atoms = n_atoms
×
525

526
    def crop(
1✔
527
        self,
528
        xlim=(-np.inf, np.inf),
529
        ylim=(-np.inf, np.inf),
530
        zlim=(-np.inf, np.inf),
531
        auto_margin: bool = False,
532
    ) -> None:
533
        """
534
        Removes all atoms that are outside specified bounds.
535

536
        Parameters
537
        ----------
538
        xlim : tuple, optional
539
            Limit in x-direction. The default is (-np.inf, np.inf).
540
        ylim : tuple, optional
541
            Limit in y-direction. The default is (-np.inf, np.inf).
542
        zlim : tuple, optional
543
            Limit in z-direction. The default is (-np.inf, np.inf).
544
        auto_margin : TYPE, optional
545
            If auto_margin == True then an additional margin of the maximum
546
            covalent radius is added to all borders. The default is False.
547

548
        Returns
549
        -------
550
        None.
551

552
        """
553
        indices_to_remove = self.get_cropping_indices(xlim, ylim, zlim, auto_margin)
×
554
        self.remove_atoms(indices_to_remove)
×
555

556
    def crop_inverse(
1✔
557
        self,
558
        xlim=(-np.inf, np.inf),
559
        ylim=(-np.inf, np.inf),
560
        zlim=(-np.inf, np.inf),
561
        auto_margin=False,
562
    ) -> None:
563
        """
564
        Removes all atoms that are inside specified bounds.
565

566
        Parameters
567
        ----------
568
        xlim : tuple, optional
569
            Limit in x-direction. The default is (-np.inf, np.inf).
570
        ylim : tuple, optional
571
            Limit in y-direction. The default is (-np.inf, np.inf).
572
        zlim : tuple, optional
573
            Limit in z-direction. The default is (-np.inf, np.inf).
574
        auto_margin : TYPE, optional
575
            If auto_margin == True then an additional margin of the maximum
576
            covalent radius is added to all borders. The default is False.
577

578
        Returns
579
        -------
580
        None.
581

582
        """
583
        indices_to_keep = self.get_cropping_indices(xlim, ylim, zlim, auto_margin)
×
584
        indices_to_remove = np.array(
×
585
            [i for i in range(self.n_atoms) if i not in indices_to_keep]
586
        )
587
        self.remove_atoms(indices_to_remove)
×
588

589
    def crop_to_unit_cell(  # noqa: PLR0912
1✔
590
        self, lattice=None, frac_coord_factors: list[int] | None = None
591
    ) -> None:
592
        """
593
        Remove all atoms that are outside the given unit cell.
594

595
        Similar to `self.crop()` but allows for arbitrary unit cells
596

597
        Parameters
598
        ----------
599
        lattice : array like, optional
600
            Lattice vectors. The default is None.
601
        frac_coord_factors : list, optional
602
            The default is [0, 1].
603

604
        Returns
605
        -------
606
        None.
607

608
        """
609
        # Atoms that have fractional coordinates outside the defined frac_coord
610
        # factors are removed. Per default frac_coord_factors=[0,1]
611

612
        if frac_coord_factors is None:
×
613
            frac_coord_factors = [0, 1]
×
614
        if lattice is None:
×
615
            lattice = self.lattice_vectors
×
616
        frac_coords = utils.get_fractional_coords(self.coords, lattice)
×
617

618
        remove_inds = []
×
619
        remove_inds += list(np.where(frac_coords[:, 0] >= frac_coord_factors[1])[0])
×
620
        remove_inds += list(np.where(frac_coords[:, 1] >= frac_coord_factors[1])[0])
×
621
        remove_inds += list(np.where(frac_coords[:, 0] < frac_coord_factors[0])[0])
×
622
        remove_inds += list(np.where(frac_coords[:, 1] < frac_coord_factors[0])[0])
×
623
        remove_inds += list(np.where(frac_coords[:, 2] > frac_coord_factors[1])[0])
×
624
        remove_inds += list(np.where(frac_coords[:, 2] < frac_coord_factors[0])[0])
×
625

626
        remove_inds = np.array(set(remove_inds))
×
627

628
        self.remove_atoms(remove_inds)
×
629
        self.lattice_vectors = lattice
×
630

631
        # In the following all redundant atoms, i.e. atoms that are multiplied
632
        # at the same position when the unitcell is repeated periodically, are
633
        # removed from the new unit cell
634
        epsilon = 0.1
×
635
        # Distance in Angstrom for which two atoms are assumed to be in the
636
        # same position
637
        init_geom = self
×
638
        allcoords = init_geom.coords
×
639

640
        # generate all possible translation vectors that could map an atom of
641
        # the unit cell into itsself
642
        prim_lat_vec = [
×
643
            [init_geom.lattice_vectors[i], -init_geom.lattice_vectors[i]]
644
            for i in range(3)
645
        ]
646
        self_mapping_translation_vectors = []
×
647

648
        self_mapping_translation_vectors.extend(
×
649
            i[sign] for i in prim_lat_vec for sign in range(2)
650
        )
651

652
        for i in range(3):
×
653
            for sign0 in range(2):
×
654
                for j in range(3):
×
655
                    for sign1 in range(2):
×
656
                        if i != j:
×
657
                            single_addition_vector = (
×
658
                                prim_lat_vec[i][sign0] + prim_lat_vec[j][sign1]
659
                            )
660
                            self_mapping_translation_vectors.append(
×
661
                                single_addition_vector
662
                            )
663

664
        for i in range(3):
×
665
            for sign0 in range(2):
×
666
                for j in range(3):
×
667
                    for sign1 in range(2):
×
668
                        for k in range(3):
×
669
                            for sign2 in range(2):
×
670
                                if i not in (j, k) and j != k:
×
671
                                    single_addition_vector = (
×
672
                                        prim_lat_vec[i][sign0]
673
                                        + prim_lat_vec[j][sign1]
674
                                        + prim_lat_vec[k][sign2]
675
                                    )
676
                                    self_mapping_translation_vectors.append(
×
677
                                        single_addition_vector
678
                                    )
679

680
        # Find the indices of those atoms that are equivalent, i.e. atoms that
681
        # are doubled when the unit cell is repeated periodically
682

683
        doubleindices = []  # list of pairs of atom indices that are equivalent
×
684
        for i, coords_i in enumerate(allcoords):
×
685
            for trans_l in self_mapping_translation_vectors:
×
686
                coords_i_shift_l = copy.deepcopy(coords_i)
×
687
                coords_i_shift_l += trans_l
×
688
                for j, coords_j in enumerate(allcoords):
×
689
                    if j != i:
×
690
                        distance_i_shift_l_j = np.linalg.norm(
×
691
                            coords_i_shift_l - coords_j
692
                        )
693
                        if distance_i_shift_l_j < epsilon:
×
694
                            doubleindices.append([i, j])
×
695

696
        for i in range(len(doubleindices)):
×
697
            doubleindices[i].sort()
×
698

699
        ###################################################################
700
        ##Create a list of redundant atoms according to the atoms that are equivalent
701
        # according to all the pairs in doubleindices
702

703
        liste = doubleindices
×
704
        to_be_killed = []  # List of all atom indicess that are redundant
×
705
        for liste_i in liste:
×
706
            replacer = liste_i[0]
×
707
            to_be_replaced = liste_i[1]
×
708
            to_be_killed.append(to_be_replaced)
×
709
            for j, liste_j in enumerate(liste):
×
710
                for k in range(2):
×
711
                    if liste_j[k] == to_be_replaced:
×
712
                        liste[j][k] = replacer
×
713
        remainers = [j[0] for j in liste]
×
714
        for r in remainers:
×
715
            for k in to_be_killed:
×
716
                if k == r:
×
717
                    to_be_killed.remove(k)
×
718

719
        self.remove_atoms(np.array(to_be_killed))
×
720

721
    def remove_metal_atoms(self) -> None:
1✔
722
        """
723
        Removes all atoms with atomic number > 18 and atomic numbers
724
        3,4,11,12,13,14
725

726
        Returns
727
        -------
728
        None.
729

730
        """
731
        metal_atoms = self.get_indices_of_metal()
×
732

733
        print(type(metal_atoms))
×
734

735
        self.remove_atoms(metal_atoms)
×
736

737
    def remove_non_metallic_atoms(self) -> None:
1✔
738
        """
739
        Removes all atoms that are not metal
740

741
        Returns
742
        -------
743
        None.
744

745
        """
746
        mol_inds = self.get_indices_of_molecules()
×
747
        self.remove_atoms(mol_inds)
×
748

749
    def remove_substrate(self, primitive_substrate) -> None:
1✔
750
        """
751
        Removes all substrate atoms given the primitive substrate by
752
        identifying species and height
753

754
        Parameters
755
        ----------
756
        primitive_substrate: Geometry
757
            primitive substrate file of system
758

759
        dimension: int
760
            dimension to use as z-axis
761

762
        threshold: float
763
            height threshold in A
764
        """
765
        substrate_indices = self.get_substrate_indices(
×
766
            primitive_substrate=primitive_substrate
767
        )
768
        self.remove_atoms(substrate_indices)
×
769

770
    def remove_adsorbates(self, primitive_substrate=None) -> None:
1✔
771
        """
772
        Removes all atoms that are not part of the substrate
773

774
        Returns
775
        -------
776
        None.
777

778
        """
779
        adsorbate_indices = self.get_adsorbate_indices(
×
780
            primitive_substrate=primitive_substrate
781
        )
782
        self.remove_atoms(adsorbate_indices)
×
783

784
    def remove_collisions(self, keep_latest: Union[bool, slice] = True) -> None:
1✔
785
        """
786
        Removes all atoms that are in a collision group as given by
787
        GeometryFile.getCollidingGroups.
788

789
        Parameters
790
        ----------
791
        keep_latest : Union[bool, slice], optional
792
            Whether to keep the earliest or latest added. If a slice object is
793
            given, the selection is used to determine which atoms to keep.
794
            Defaults to True.
795

796
        Returns
797
        -------
798
        None.
799

800
        """
801
        indices = []
×
802
        if isinstance(keep_latest, bool):
×
803
            if keep_latest:
×
804
                selection = slice(None, -1, None)
×
805
            else:
806
                selection = slice(1, None, None)
×
807
        elif isinstance(keep_latest, slice):
×
808
            selection = keep_latest
×
809
        collisions = self.get_colliding_groups()
×
810
        for group in collisions:
×
811
            indices += group[selection]
×
812
        self.remove_atoms(np.array(indices))
×
813

814
    ###########################################################################
815
    #                         Transformations (in place)                      #
816
    ###########################################################################
817
    def map_to_first_unit_cell(
1✔
818
        self, lattice_vectors=None, dimensions=np.array(range(3))
819
    ) -> None:
820
        """
821
        Aps the coordinate of a geometry in multiples of the substrate lattice
822
        vectors to a point that is closest to the origin
823

824
        Parameters
825
        ----------
826
        lattice_vectors : float-array, optional
827
            Lattice vectors of the substrate. The default is None.
828
        dimensions : float-array, optional
829
            Dimensions (x, y, z) where the mapping should be done. The default
830
            is np.array(range(3)).
831

832
        Returns
833
        -------
834
        None.
835

836
        """
837
        if lattice_vectors is None:
×
838
            lattice_vectors = self.lattice_vectors
×
839

840
        assert not np.allclose(lattice_vectors, np.zeros([3, 3])), (
×
841
            "Lattice vector must be defined in Geometry or given as function parameter"
842
        )
843

844
        frac_coords = utils.get_fractional_coords(self.coords, lattice_vectors)
×
845
        # modulo 1 maps all coordinates to first unit cell
846
        frac_coords[:, dimensions] = frac_coords[:, dimensions] % 1
×
847
        new_coords = utils.get_cartesian_coords(frac_coords, lattice_vectors)
×
848
        self.coords = new_coords
×
849

850
    def map_center_of_atoms_to_first_unit_cell(self, lattice_vectors=None):
1✔
851
        if lattice_vectors is None:
×
852
            lattice_vectors = self.lattice_vectors
×
853

854
        assert not np.allclose(lattice_vectors, np.zeros([3, 3])), (
×
855
            "Lattice vector must be defined in Geometry or given as function parameter"
856
        )
857

858
        offset = self.get_geometric_center()
×
859
        frac_offset = utils.get_fractional_coords(offset, lattice_vectors)
×
860
        frac_offset = np.floor(frac_offset)
×
861
        self.move_all_atoms_by_fractional_coords(-frac_offset, lattice_vectors)
×
862

863
    def center_coordinates(
1✔
864
        self,
865
        ignore_center_attribute: bool = False,
866
        dimensions=np.array(range(3)),
867
    ):
868
        """
869
        Shift the coordinates of a geometry such that the "center of mass" or
870
        specified center lies at (0,0,0)
871

872
        Parameters
873
        ----------
874
        ignore_center_attribute : bool
875
            Switch usage of *center* attribute off/on. The default is False.
876

877
        dimensions: np.array
878
            Dimensions that should be cnetered. The default is False [0, 1, 2].
879

880
        Returns
881
        -------
882
        None.
883

884
        """
885
        offset = self.get_geometric_center(
×
886
            ignore_center_attribute=ignore_center_attribute
887
        )[dimensions]
888
        self.coords[:, dimensions] -= offset
×
889
        return offset
×
890

891
    def move_all_atoms(self, shift):
1✔
892
        """
893
        Translates the whole geometry by vector 'shift'
894

895
        """
896
        self.coords += shift
×
897

898
    def move_all_atoms_by_fractional_coords(self, frac_shift, lattice_vectors=None):
1✔
899
        if lattice_vectors is None:
×
900
            lattice_vectors = self.lattice_vectors
×
901

902
        self.coords += utils.get_cartesian_coords(frac_shift, lattice_vectors)
×
903

904
    def move_adsorbates(self, shift, primitive_substrate=None):
1✔
905
        """
906
        Shifts the adsorbates in Cartesian coordinates
907
        """
908
        adsorbates = self.get_adsorbates(primitive_substrate=primitive_substrate)
×
909
        adsorbates.coords += shift
×
910

911
        self.remove_adsorbates(primitive_substrate=primitive_substrate)
×
912
        self += adsorbates
×
913

914
    def rotate_lattice_around_axis(
1✔
915
        self,
916
        angle_in_degree: float,
917
        axis: npt.NDArray[np.float64] = np.array([0.0, 0.0, 1.0]),
918
    ) -> None:
919
        """
920
        Rotates lattice around a given axis.
921

922
        Parameters
923
        ----------
924
        angle_in_degree : float
925
            angle of rotation
926

927
        axis: np.array
928
            Axis around which to rotate. The default is
929

930
        Returns
931
        -------
932
        None.
933

934
        """
935
        R = utils.get_rotation_matrix_around_axis(axis, angle_in_degree * np.pi / 180)
×
936
        self.lattice_vectors = np.dot(self.lattice_vectors, R)
×
937

938
    def rotate_coords_around_axis(
1✔
939
        self,
940
        angle_in_degree: float,
941
        axis: npt.NDArray[np.float64] = np.array([0.0, 0.0, 1.0]),
942
        center=None,
943
        indices=None,
944
    ) -> None:
945
        """
946
        Rotates structure COUNTERCLOCKWISE around a point defined by <center>.
947

948
        Parameters
949
        ----------
950
        angle_in_degree : float
951
            Angle of rotation.
952
        axis : npt.NDArray[np.float64], optional
953
            Axis of rotation. The default is np.array([0.0, 0.0, 1.0]).
954
        center : npt.NDArray[np.float64], optional
955
            If center == None, the geometric center of the structure (as
956
            defined by self.get_geometric_center()). The default is None.
957
        indices : list, optional
958
            Indices of atoms to be manipulated. If indices == None, all atoms
959
            will be used. The default is None.
960

961
        Returns
962
        -------
963
        None.
964

965
        """
966
        if indices is None:
×
967
            indices = np.arange(self.n_atoms)
×
968
        if center is None:
×
969
            center = self.get_geometric_center(indices=indices)
×
970

971
        R = utils.get_rotation_matrix_around_axis(axis, angle_in_degree * np.pi / 180)
×
972
        temp_coords = copy.deepcopy(self.coords[indices])
×
973
        temp_coords -= center
×
974
        temp_coords = np.dot(temp_coords, R)
×
975
        temp_coords += center
×
976
        self.coords[indices] = temp_coords
×
977

978
    def mirror_through_plane(self, normal_vector: npt.NDArray[np.float64]) -> None:
1✔
979
        """
980
        Mirrors the geometry through the plane defined by the normal vector.
981

982
        Parameters
983
        ----------
984
        normal_vector : npt.NDArray[np.float64]
985
            Normal vector of mirror plane.
986

987
        Returns
988
        -------
989
        None.
990

991
        """
992
        mirror_matrix = utils.get_mirror_matrix(normal_vector=normal_vector)
×
993
        self.transform(mirror_matrix)
×
994

995
    def align_into_xy_plane(self, atom_indices):
1✔
996
        """
997
        Rotates a planar molecule (defined by 3 atom indices) into the XY
998
        plane. Double check results, use with caution
999

1000
        Parameters
1001
        ----------
1002
        atom_indices
1003
            Indices of atoms that should be aligned.
1004

1005
        Returns
1006
        -------
1007
        None.
1008

1009
        """
1010
        p1 = self.coords[atom_indices[0]]
×
1011
        p2 = self.coords[atom_indices[1]]
×
1012
        p3 = self.coords[atom_indices[2]]
×
1013

1014
        X = np.zeros([3, 3])
×
1015
        X[0, :] = p2 - p1
×
1016
        X[1, :] = p3 - p1
×
1017
        X[2, :] = np.cross(X[0], X[1])
×
1018
        for i in range(3):
×
1019
            X[i] /= np.linalg.norm(X[i])
×
1020
        X[1, :] = np.cross(X[2], X[0])
×
1021

1022
        U = np.linalg.inv(X)
×
1023
        self.transform(U)
×
1024

1025
    def align_with_z_vector(self, new_z_vec: npt.NDArray[np.float64]) -> None:
1✔
1026
        """
1027
        Transforms the coordinate system of the geometry file to a new z-vector
1028
        calculates rotation martrix for coordinate transformation to new z-vector
1029
        and uses it to transform coordinates of geometry object
1030

1031
        Parameters
1032
        ----------
1033
        new_z_vec : npt.NDArray[np.float64]
1034
            The vector to align with the z-axis.
1035

1036
        Returns
1037
        -------
1038
        None
1039

1040
        """
1041
        # get old_positions
1042
        old_positions = self.coords
×
1043

1044
        # normalize new_z_vec
1045
        new_z_vec = new_z_vec / np.linalg.norm(new_z_vec)
×
1046

1047
        # Check if the desired vector is antiparallel to the z-axis
1048
        if np.allclose(new_z_vec, -np.array([0, 0, 1])):
×
1049
            rotation_matrix = np.diag([-1, -1, 1])  # Antiparallel case
×
1050
        else:
1051
            # Calculate the rotation matrix
1052
            z_axis = np.array([0, 0, 1])
×
1053
            cross_product = np.cross(new_z_vec, z_axis)
×
1054
            dot_product = np.dot(new_z_vec, z_axis)
×
1055
            skew_symmetric_matrix = np.array(
×
1056
                [
1057
                    [0, -cross_product[2], cross_product[1]],
1058
                    [cross_product[2], 0, -cross_product[0]],
1059
                    [-cross_product[1], cross_product[0], 0],
1060
                ]
1061
            )
1062
            rotation_matrix = (
×
1063
                np.eye(3)
1064
                + skew_symmetric_matrix
1065
                + np.dot(skew_symmetric_matrix, skew_symmetric_matrix)
1066
                * (1 - dot_product)
1067
                / (np.linalg.norm(cross_product) ** 2)
1068
            )
1069

1070
        # Apply the rotation to all positions
1071
        rotated_positions = np.dot(old_positions, rotation_matrix.T)
×
1072

1073
        self.coords = rotated_positions
×
1074

1075
    def align_vector_to_vector(
1✔
1076
        self,
1077
        vector: npt.NDArray[np.float64],
1078
        vector_to_align: npt.NDArray[np.float64],
1079
    ):
1080
        """
1081
        Aligns a vector and the atomic coordiantes to a given vector.
1082

1083
        Parameters
1084
        ----------
1085
        vector : npt.NDArray[np.float64]
1086
            vector for alignment
1087

1088
        vector_to_align : npt.NDArray[np.float64]
1089
            index of the lattice vector that should be aligned
1090

1091
        Returns
1092
        -------
1093
        None
1094

1095
        """
1096
        vector_to_align_normed = vector_to_align / np.linalg.norm(vector_to_align)
×
1097

1098
        vector_normed = vector / np.linalg.norm(vector)
×
1099

1100
        R = utils.get_rotation_matrix(vector_normed, vector_to_align_normed)
×
1101

1102
        self.lattice_vectors = np.dot(self.lattice_vectors, R)
×
1103
        self.coords = np.dot(self.coords, R)
×
1104

1105
    def align_lattice_vector_to_vector(self, vector, lattice_vector_index):
1✔
1106
        """
1107
        Aligns a lattice vector and the atomic coordiantes to a given axis.
1108

1109
        Parameters
1110
        ----------
1111
        vector : array
1112
            vector for alignment
1113

1114
        lattice_vector_index : int
1115
            index of the lattice vector that should be aligned
1116

1117
        Returns
1118
        -------
1119
        None
1120

1121
        """
1122
        lattice_vector_normed = self.lattice_vectors[
×
1123
            lattice_vector_index
1124
        ] / np.linalg.norm(self.lattice_vectors[lattice_vector_index])
1125

1126
        self.align_vector_to_vector(vector, lattice_vector_normed)
×
1127

1128
    def align_cartiesian_axis_to_vector(self, vector, axis_index):
1✔
1129
        """
1130
        Aligns a lattice vector and the atomic coordiantes to a given axis.
1131

1132
        Parameters
1133
        ----------
1134
        vector : array
1135
            vector for alignment
1136

1137
        axis_index : int
1138
            index of the axis that should be aligned
1139

1140
        """
1141
        axis = np.zeros(3, dtype=np.float64)
×
1142
        axis[axis_index] = 1.0
×
1143

1144
        self.align_vector_to_vector(vector, axis)
×
1145

1146
    def align_with_view_direction(
1✔
1147
        self, view_direction: npt.NDArray[np.float64]
1148
    ) -> None:
1149
        view_direction /= np.linalg.norm(view_direction)
×
1150

1151
        vec_z = np.array([0.0, 0.0, -1.0])
×
1152

1153
        view_direction_y = np.cross(vec_z, view_direction)
×
1154
        norm_y = np.linalg.norm(view_direction_y)
×
1155

1156
        if norm_y == 0.0:
×
1157
            sign = np.dot(vec_z, view_direction)
×
1158
            self.lattice_vectors[2] *= sign
×
1159
            self.coords[:, 2] *= sign
×
1160

1161
        else:
1162
            # Orient z-axis in view direction
1163
            self.align_cartiesian_axis_to_vector(-view_direction, 2)
×
1164

1165
    def align_main_axis_along_xyz(self) -> None:
1✔
1166
        """
1167
        Align coordinates of rodlike molecules along specified axis
1168

1169
        """
1170
        vals, vecs = self.get_main_axes()
×
1171
        R = np.linalg.inv(vecs.T)
×
1172
        self.coords = np.dot(self.coords, R)
×
1173

1174
    def transform(
1✔
1175
        self,
1176
        R: npt.NDArray[np.float64],
1177
        t: npt.NDArray[np.float64] = np.array([0, 0, 0]),
1178
        rotation_center: Union[npt.NDArray[np.float64], None] = None,
1179
        atom_indices: Union[npt.NDArray[np.int64], None] = None,
1180
    ) -> None:
1181
        """
1182
        Performs a symmetry transformation of the coordinates by rotation and
1183
        translation. The transformation is applied as
1184
        x_new[3x1] = x_old[3x1] x R[3x3] + t[3x1]
1185

1186
        Parameters
1187
        ----------
1188
        R : np.array
1189
            Rotation matrix in Catesian coordinates.
1190
        t : np.array, optional
1191
            Translation vector in Catesian coordinates. The default is np.array([0,0,0]).
1192
        rotation_center : np.array | None, optional
1193
            Centre of rotation. The default is None.
1194
        atom_indices : np.array | None, optional
1195
            List of indexes of atoms that should be transformed. The default is
1196
            None.
1197

1198
        Returns
1199
        -------
1200
        None
1201

1202
        """
1203
        if atom_indices is None:
×
1204
            atom_indices = np.arange(self.n_atoms)
×
1205
        if rotation_center is None:
×
1206
            temp_coords = np.dot(self.coords[atom_indices, :], R) + t
×
1207
            self.coords[atom_indices, :] = temp_coords
×
1208
        else:
1209
            temp_coords = copy.deepcopy(self.coords[atom_indices, :])
×
1210
            temp_coords -= rotation_center
×
1211
            temp_coords = np.dot(temp_coords, R) + t
×
1212
            temp_coords += rotation_center
×
1213
            self.coords[atom_indices, :] = temp_coords
×
1214

1215
    def transform_fractional(
1✔
1216
        self,
1217
        R: npt.NDArray[np.float64],
1218
        t: npt.NDArray[np.float64],
1219
        lattice=None,
1220
    ):
1221
        """
1222
        Transforms the coordinates by rotation and translation, where R,t are
1223
        given in fractional coordinates
1224
        The transformation is applied as c_new[3x1] = R[3x3] * c_old[3x1] + t[3x1]
1225

1226
        """
1227
        if lattice is None:
×
1228
            lattice = self.lattice_vectors
×
1229
        coords_frac = utils.get_fractional_coords(self.coords, lattice)
×
1230
        coords_frac = np.dot(coords_frac, R.T) + t.reshape([1, 3])
×
1231
        self.coords = utils.get_cartesian_coords(coords_frac, lattice)
×
1232

1233
    def transform_lattice(
1✔
1234
        self,
1235
        R: npt.NDArray[np.float64],
1236
        t: npt.NDArray[np.float64] = np.array([0, 0, 0]),
1237
    ) -> None:
1238
        """
1239
        Transforms the lattice vectors by rotation and translation.
1240
        The transformation is applied as x_new[3x1] = x_old[3x1] x R[3x3] + t[3x1]
1241
        Notice that this works in cartesian coordinates.
1242
        Use transform_fractional if you got your R and t from get_symmetries
1243

1244
        Parameters
1245
        ----------
1246
        R : np.array
1247
            Rotation matrix in Catesian coordinates.
1248
        t : np.array, optional
1249
            Translation vector in Catesian coordinates. The default is np.array([0,0,0]).
1250

1251
        Returns
1252
        -------
1253
        None
1254

1255
        """
1256
        new_lattice_vectors = np.dot(self.lattice_vectors, R) + t
×
1257
        self.lattice_vectors = new_lattice_vectors
×
1258

1259
    def transform_lattice_fractional(self, R, t, lattice):
1✔
1260
        """Transforms the lattice vectors by rotation and translation.
1261
        The transformation is applied as x_new[3x1] = x_old[3x1] x R[3x3] + t[3x1]
1262
        """
1263
        coords_frac = utils.get_fractional_coords(self.lattice_vectors, lattice)
×
1264
        coords_frac = np.dot(coords_frac, R.T) + t.reshape([1, 3])
×
1265
        self.lattice_vectors = utils.get_cartesian_coords(coords_frac, lattice)
×
1266

1267
    def swap_lattice_vectors(self, axis_1=0, axis_2=1):
1✔
1268
        """
1269
        Can be used to interchange two lattice vectors
1270
        Attention! Other values - for instance k_grid - will stay unchanged!!
1271
        :param axis_1 integer [0,1,2]
1272
        :param axis_2 integer [0,1,2]     axis_1 !=axis_2
1273
        :return:
1274

1275
        """
1276
        self.lattice_vectors[[axis_1, axis_2], :] = self.lattice_vectors[
×
1277
            [axis_2, axis_1], :
1278
        ]
1279
        self.coords[[axis_1, axis_2], :] = self.coords[[axis_2, axis_1], :]
×
1280

1281
    def symmetrize(self, symmetry_operations, center=None):
1✔
1282
        """
1283
        Symmetrizes Geometry with given list of symmetry operation matrices
1284
        after transferring it to the origin.
1285
        Do not include the unity matrix for symmetrizing, as it is already the
1286
        first geometry!
1287
        ATTENTION: use symmetrize_periodic to reliably symmetrize periodic
1288
        structures
1289

1290
        """
1291
        if center is not None:
×
1292
            offset = center
×
1293
            self.coords -= center
×
1294
        else:
1295
            offset = np.mean(self.coords, axis=0)
×
1296
            self.center_coordinates()
×
1297
        temp_coords = copy.deepcopy(
×
1298
            self.coords
1299
        )  # this corresponds to the unity matrix symmetry operation
1300
        for R in symmetry_operations:
×
1301
            new_geom = copy.deepcopy(self)
×
1302
            new_geom.transform(R)
×
1303
            new_geom.reorder_atoms(self.get_transformation_indices(new_geom))
×
1304
            temp_coords += new_geom.coords
×
1305
        self.coords = temp_coords / (len(symmetry_operations) + 1) + offset
×
1306

1307
    def average_with(self, other_geometries) -> None:
1✔
1308
        """
1309
        Average self.coords with those of other_geometries and apply on self
1310
        ATTENTION: this can change bond lengths etc.!Ok n
1311

1312
        Parameters
1313
        ----------
1314
        other_geometries: List of Geometrys ... might be nice to accept list of
1315
        coords too
1316

1317
        Returns
1318
        -------
1319
        None
1320

1321
        """
1322
        if len(other_geometries) > 0:
×
1323
            offset = (
×
1324
                self.get_geometric_center()
1325
            )  # Attribute center should be used if it exists
1326
            self.coords -= offset
×
1327

1328
            for other_geom in other_geometries:
×
1329
                geom = copy.deepcopy(other_geom)
×
1330
                # center all other geometries to remove offset
1331
                geom.center_coordinates()
×
1332
                # all geometries have to be ordered like first geometry in
1333
                # order to sum them
1334
                geom.reorder_atoms(self.get_transformation_indices(geom))
×
1335
                self.coords += geom.coords
×
1336
            self.coords /= len(other_geometries) + 1  # +1 for this geometry itself
×
1337
            self.coords += offset
×
1338

1339
    def reorder_atoms(self, inds: npt.NDArray[np.int64]) -> None:
1✔
1340
        """
1341
        Eorders Atoms with index list
1342

1343
        Parameters
1344
        ----------
1345
        inds : npt.NDArray[np.int64]
1346
            Array of indices with new order of atoms.
1347

1348
        Returns
1349
        -------
1350
        None
1351

1352
        """
1353
        self.coords = self.coords[inds, :]
×
1354
        self.species = [self.species[i] for i in inds]
×
1355
        self.constrain_relax = self.constrain_relax[inds, :]
×
1356
        self.initial_charge = [self.initial_charge[i] for i in inds]
×
1357
        self.initial_moment = [self.initial_moment[i] for i in inds]
×
1358

1359
        inds_hessian = []
×
1360
        for ind in inds:
×
1361
            inds_new = np.array([0, 1, 2]) + 3 * ind
×
1362
            inds_hessian.append(inds_new)
×
1363

1364
        inds_hessian = np.array(inds_hessian).flatten()
×
1365

1366
        if self.hessian is not None:
×
1367
            self.hessian = self.hessian[inds_hessian, :]
×
1368
            self.hessian = self.hessian[:, inds_hessian]
×
1369

1370
    def shift_to_bottom(self):
1✔
1371
        """
1372
        Shifts the coordiantes such that the coordinate with the samllest
1373
        z-value ists at (x, y, 0).
1374

1375
        Returns
1376
        -------
1377
        None.
1378

1379
        """
1380
        min_z = np.min(self.coords[:, -1])
×
1381
        self.coords[:, -1] -= min_z
×
1382

1383
    def displace_atoms(
1✔
1384
        self,
1385
        displacement_strength: float,
1386
        displacement_indices: npt.NDArray[np.int64],
1387
    ) -> None:
1388
        """
1389
        Displaces atoms randomly.
1390

1391
        Parameters
1392
        ----------
1393
        displacement_strength : float
1394
            Scaling factor for the strenght of the dispacements.
1395
        displacement_indices : npt.NDArray[np.int64]
1396
            Indices where atoms should be dispaced.
1397

1398
        Returns
1399
        -------
1400
        None
1401

1402
        """
1403
        displacements = np.random.rand(len(displacement_indices), 3) - 0.5
1✔
1404
        displacements *= displacement_strength
1✔
1405

1406
        self.coords[displacement_indices, :] += displacements
1✔
1407

1408
    def move_multipoles(self, shift: npt.NDArray[np.float64]) -> None:
1✔
1409
        """
1410
        Moves all the multipoles by a shift vector
1411
        :param shift: list or array, len==3
1412
        :return:
1413
        """
1414
        assert len(shift) == 3
×
1415
        for m in self.multipoles:
×
1416
            m[0] += shift[0]
×
1417
            m[1] += shift[1]
×
1418
            m[2] += shift[2]
×
1419

1420
    ###########################################################################
1421
    #                      Set Properties of the Geometry                     #
1422
    ###########################################################################
1423
    def set_vacuum_height(self, vac_height, bool_shift_to_bottom=False) -> None:
1✔
1424
        if bool_shift_to_bottom:
×
1425
            self.shift_to_bottom()
×
1426
        min_z = np.min(self.coords[:, -1])
×
1427
        max_z = np.max(self.coords[:, -1])
×
1428
        self.lattice_vectors[-1, -1] = max_z + vac_height - min_z
×
1429

1430
        if vac_height < min_z:
×
1431
            raise Exception(
×
1432
                """set_vacuum_height: the defined vacuum height is smaller than
1433
                height of the lowest atom. Shift unit cell either manually or
1434
                by the keyword bool_shift_to_bottom towards the bottom
1435
                of the unit cell."""
1436
            )
1437
        self.lattice_vectors[-1, -1] = max_z + vac_height - min_z
×
1438

1439
    def set_vacuum_level(self, vacuum_level: float) -> None:
1✔
1440
        """
1441
        Sets vacuum level of geometry calculation
1442

1443
        Parameters
1444
        ----------
1445
        vacuum_level : float
1446
            Height of the vacuum level.
1447

1448
        Returns
1449
        -------
1450
        None
1451

1452
        """
1453
        self.vacuum_level = vacuum_level
×
1454

1455
    def set_multipoles_charge(self, charge: npt.NDArray[np.float64]) -> None:
1✔
1456
        """
1457
        Sets the charge of all multipoles
1458

1459
        Parameters
1460
        ----------
1461
        charge : list or float or int
1462

1463
        Returns
1464
        -------
1465
        None
1466

1467
        """
1468
        if isinstance(charge, list):
×
1469
            assert len(charge) == len(self.multipoles)
×
1470
            for i, m in enumerate(self.multipoles):
×
1471
                m[4] = charge[i]
×
1472
        else:
1473
            for i, m in enumerate(self.multipoles):
×
1474
                m[4] = charge
×
1475

1476
    def set_constraints(
1✔
1477
        self,
1478
        indices_of_atoms_to_constrain: list,
1479
        constrain_dim_flags=None,
1480
    ):
1481
        """
1482
        Sets a constraint for a few atoms in the system (identified by
1483
        'indices_of_atoms_to_constrain') for a geometry relaxation.
1484
        Since the relaxation constraint can be in any and/or all dimensions
1485
        the second parameter, 'constraint_dim_flags', makes it possible to
1486
        set which dimension(s) should be constrained for which molecule.
1487
        By default all dimensions are to be constrained for all atoms are
1488
        constrained. If the dimension to constrain should be set individually
1489
        for different atoms, you need to provide a list of booleans of the
1490
        shap len(indices_of_atoms_to_constrain) x 3, which contains the
1491
        constrain flags for each dimension for each atom.
1492

1493
        Parameters
1494
        ----------
1495
        indices_of_atoms_to_constrain : list
1496
            List of atoms to constrain.
1497
        constrain_dim_flags : list[boolean]
1498
            The default is: [True, True, True].
1499

1500
        Returns
1501
        -------
1502
        None
1503

1504
        """
1505
        if constrain_dim_flags is None:
×
1506
            constrain_dim_flags = [True, True, True]
×
1507

1508
        self.constrain_relax[indices_of_atoms_to_constrain, :] = constrain_dim_flags
×
1509

1510
    def set_constraints_based_on_space(
1✔
1511
        self,
1512
        xlim=(-np.inf, np.inf),
1513
        ylim=(-np.inf, np.inf),
1514
        zlim=(-np.inf, np.inf),
1515
        constrain_dim_flags=None,
1516
    ):
1517
        """
1518
        Constrain all atoms that are within a cuboid (defined by
1519
        limits in all dimensions: xlim, etc.) for a geometry relaxation.
1520

1521
        It is possible to define which dimension will be constrained, but since
1522
        the number of atoms in the cuboid is only calculated at runtime
1523
        the dimensions may only be set for all atoms at once. If you need to
1524
        set them individually please use set_constraints.
1525

1526
        Parameters
1527
        ----------
1528
        zlim
1529
        xlim
1530
        ylim
1531
        constrain_dim_flags : list[boolean]
1532
            The default is: [True, True, True].
1533

1534
        Returns
1535
        -------
1536
        None
1537

1538
        """
1539
        if constrain_dim_flags is None:
×
1540
            constrain_dim_flags = [True, True, True]
×
1541

1542
        # --- get indices of all atoms outside the required interval ---
1543
        indices_outside = self.get_cropping_indices(
×
1544
            xlim=xlim, ylim=ylim, zlim=zlim, auto_margin=False
1545
        )
1546
        # ---
1547

1548
        # --- Filter all that are outside ---
1549
        # The indices of the atoms of relevance to us are all that are NOT
1550
        # outside of the cuboid
1551
        indices_inside = [i for i in range(len(self)) if i not in indices_outside]
×
1552
        # ---
1553

1554
        self.set_constraints(indices_inside, constrain_dim_flags)
×
1555

1556
    def free_all_constraints(self) -> None:
1✔
1557
        """
1558
        Frees all constraints.
1559

1560
        Returns
1561
        -------
1562
        None
1563

1564
        """
1565
        self.constrain_relax = np.zeros([len(self.species), 3], bool)
×
1566

1567
    def set_calculate_friction(
1✔
1568
        self, indices_of_atoms: list, calculate_friction: bool = True
1569
    ) -> None:
1570
        """
1571
        Sets to calculate electronic friction for atoms specified by the given
1572
        list of indices.
1573

1574
        Parameters
1575
        ----------
1576
        indices_of_atoms : list
1577
            List of atoms from which electronic friction should be calculated.
1578
        calculate_friction : bool, optional
1579
            Calculate friction for these atims. The default is True.
1580

1581
        Returns
1582
        -------
1583
        None
1584

1585
        """
1586
        self.calculate_friction[indices_of_atoms] = calculate_friction
×
1587

1588
    def free_all_calculate_friction(self) -> None:
1✔
1589
        """
1590
        Set calculate electronic friction to false on all atoms.
1591

1592
        Returns
1593
        -------
1594
        None
1595

1596
        """
1597
        self.calculate_friction = np.array([False] * len(self))
×
1598

1599
    def set_external_forces(
1✔
1600
        self,
1601
        indices_of_atoms: Union[int, npt.NDArray[np.int64]],
1602
        external_force: npt.NDArray[np.float64],
1603
    ) -> None:
1604
        """
1605
        Sets a constraint for a few atoms in the system (identified by
1606
        'indices_of_atoms_to_constrain') for a geometry relaxation.
1607
        Since the relaxation constraint can be in any and/or all dimensions
1608
        the second parameter, 'constraint_dim_flags', makes it possible to
1609
        set which dimension(s) should be constrained for which molecule.
1610
        By default all dimensions are to be constrained for all atoms are
1611
        constrained. If the dimension to constrain should be set individually
1612
        for different atoms, you need to provide a list of booleans of the
1613
        shape len(indices_of_atoms_to_constrain) x 3, which contains the
1614
        constrain flags for each dimension for each atom.
1615

1616
        Parameters
1617
        ----------
1618
        indices_of_atoms_to_constrain : int or npt.NDArray[np.int64]
1619
            Index of atoms to which atoms should should be applied.
1620
        constrain_dim_flags : npt.NDArray[np.float64]
1621
            Force that should act on a given atom.
1622
        """
1623
        self.external_force[indices_of_atoms, :] = external_force
×
1624

1625
    def free_all_external_forces(self) -> None:
1✔
1626
        """
1627
        Set calculate electronic friction to false on all atoms.
1628

1629
        Returns
1630
        -------
1631
        None
1632

1633
        """
1634
        self.external_force = np.zeros((len(self), 3))
×
1635

1636
    def remove_periodicity(self):
1✔
1637
        """
1638
        Makes geometry non-periodic by setting its lattice vectors to zero
1639

1640
        """
1641
        self.lattice_vectors = np.zeros((3, 3), dtype=float)
×
1642

1643
    def set_homogeneous_field(self, E):
1✔
1644
        """Field should be a numpy array (Ex, Ey, Ez) with the Field in V/A"""
1645
        assert len(E) == 3, "Expected E-field components [Ex, Ey, Ez], but got " + str(
×
1646
            E
1647
        )
1648
        self.homogeneous_field = np.asarray(E)
×
1649

1650
    def free_homogeneous_field(self):
1✔
1651
        self.homogeneous_field = np.array([0.0, 0.0, 0.0])
×
1652

1653
    def set_initial_magnetic_moments(self, moments: List[float]) -> None:
1✔
1654
        self.initial_moment = moments
×
1655

1656
    ###############################################################################
1657
    #                      Get Properties of the Geometry                         #
1658
    ###############################################################################
1659
    def get_is_periodic(self) -> bool:
1✔
1660
        """
1661
        Checks if the geometry is periodic.
1662

1663
        Returns
1664
        -------
1665
        bool
1666
            Ture if geometry is periodic.
1667

1668
        """
1669
        return not np.allclose(self.lattice_vectors, np.zeros([3, 3]))
1✔
1670

1671
    def get_reassembled_molecule(self, threshold: float = 2.0):
1✔
1672
        geom_replica = self.get_periodic_replica(
×
1673
            (1, 1, 1),
1674
            explicit_replications=[[-1, 0, 1], [-1, 0, 1], [-1, 0, 1]],
1675
        )
1676

1677
        tree = scipy.spatial.KDTree(geom_replica.coords)
×
1678
        pairs = tree.query_pairs(threshold)
×
1679

1680
        new_cluster = True
×
1681

1682
        while new_cluster:
×
1683
            clusters = []
×
1684
            new_cluster = False
×
1685

1686
            for pair in pairs:
×
1687
                in_culster = False
×
1688
                for ind, indices in enumerate(clusters):
×
1689
                    for p in pair:
×
1690
                        if p in indices:
×
1691
                            clusters[ind] = set(list(indices) + list(pair))
×
1692
                            new_cluster = True
×
1693
                            in_culster = True
×
1694
                            break
×
1695

1696
                if not in_culster:
×
1697
                    clusters.append(set(pair))
×
1698

1699
            pairs = copy.deepcopy(clusters)
×
1700

1701
        for index_array in pairs:
×
1702
            if len(index_array) == len(self):
×
1703
                final_geom = geom_replica.get_atoms_by_indices(
×
1704
                    np.sort(np.array(list(index_array), dtype=np.int32))
1705
                )
1706
                final_geom.lattice_vectors = self.lattice_vectors
×
1707
                final_geom.map_center_of_atoms_to_first_unit_cell()
×
1708

1709
                return final_geom
×
1710

1711
        warnings.warn(
×
1712
            "Geometry.getReassembledMolecule could not reassemble \
1713
                      molecule. Returning original Geometry."
1714
        )
1715
        return self
×
1716

1717
    def get_scaled_copy(self, scaling_factor: Union[float, List]) -> object:
1✔
1718
        """
1719
        Returns a copy of the geometry, scaled by scaling_factor.
1720

1721
        Both the coordinates of the atoms and the length of the lattice vectors are
1722
        affected
1723

1724
        Parameters
1725
        ----------
1726
        scaling_factor : Union[float, Iterator]
1727
            Scaling factor for the geometry. If float, the volume of the geometry
1728
            will be scaled accordingly. If iterable, the length of the lattice vectors
1729
            will be scaled accordingly.
1730

1731
        Returns
1732
        -------
1733
        scaled_geometry : Geometry
1734
            Geometry object with scaled coordinates and lattice vectors
1735
        """
1736
        assert hasattr(self, "lattice_vectors"), (
×
1737
            "This function only works for geometries with a Unit Cell"
1738
        )
1739

1740
        if isinstance(scaling_factor, float):
×
1741
            scaling_factors = [
×
1742
                scaling_factor ** (1 / 3),
1743
            ] * 3
1744

1745
        else:
1746
            assert len(scaling_factor) == 3
×
1747
            scaling_factors = scaling_factor
×
1748

1749
        scaled_geom = deepcopy(self)
×
1750
        lattice_vectors = deepcopy(self.lattice_vectors)
×
1751
        lattice_vectors[0] *= scaling_factors[0]
×
1752
        lattice_vectors[1] *= scaling_factors[1]
×
1753
        lattice_vectors[2] *= scaling_factors[2]
×
1754

1755
        new_coords = utils.get_cartesian_coords(
×
1756
            self.get_fractional_coords(), lattice_vectors
1757
        )
1758
        scaled_geom.lattice_vectors = lattice_vectors
×
1759
        scaled_geom.coords = new_coords
×
1760

1761
        return scaled_geom
×
1762

1763
    def get_displaced_atoms(
1✔
1764
        self,
1765
        displacement_strength: float,
1766
        displace_only_unconstrained: bool = True,
1767
    ):
1768
        """
1769
        Returns a copy of the geometry, where the atoms have been displaced
1770
        randomly.
1771

1772
        Parameters
1773
        ----------
1774
        displacement_strength : float
1775
            Scaling factor for the strenght of the dispacements.
1776
        displace_only_unconstrained : bool
1777
            Indices where atoms should be dispaced. The default is True.
1778

1779
        Returns
1780
        -------
1781
        geometry.
1782

1783
        """
1784
        if displace_only_unconstrained:
1✔
1785
            displacement_indices = self.get_unconstrained_atoms()
1✔
1786
        else:
1787
            displacement_indices = np.array(range(len(self)))
×
1788

1789
        new_geometry = deepcopy(self)
1✔
1790
        new_geometry.displace_atoms(displacement_strength, displacement_indices)
1✔
1791

1792
        return new_geometry
1✔
1793

1794
    def get_fractional_coords(self, lattice_vectors=None):
1✔
1795
        if lattice_vectors is None:
×
1796
            lattice_vectors = self.lattice_vectors
×
1797

1798
        assert not np.allclose(lattice_vectors, np.zeros([3, 3])), (
×
1799
            "Lattice vector must be defined in Geometry or given as function parameter"
1800
        )
1801

1802
        fractional_coords = np.linalg.solve(lattice_vectors.T, self.coords.T)
×
1803
        return fractional_coords.T
×
1804

1805
    def get_fractional_lattice_vectors(self, lattice_vectors=None):
1✔
1806
        """
1807
        Calculate the fractional representation of lattice vectors of the
1808
        geometry file in another basis.
1809
        Useful to calculate epitaxy matrices
1810

1811
        """
1812
        fractional_coords = np.linalg.solve(lattice_vectors.T, self.lattice_vectors.T)
×
1813
        return fractional_coords.T
×
1814

1815
    def get_reciprocal_lattice_vectors(self) -> npt.NDArray[np.float64]:
1✔
1816
        """
1817
        Calculate the reciprocal lattice of the Geometry lattice_vectors in standard form
1818
        For convention see en.wikipedia.org/wiki/Reciprocal_lattice
1819

1820
        Returns
1821
        -------
1822
        recip_lattice : npt.NDArray[np.float64]
1823
            Row-wise reciprocal lattice vectors (3x3)
1824

1825
        """
1826
        a1 = self.lattice_vectors[0]
×
1827
        a2 = self.lattice_vectors[1]
×
1828
        a3 = self.lattice_vectors[2]
×
1829

1830
        volume = np.cross(a1, a2).dot(a3)
×
1831

1832
        b1 = np.cross(a2, a3)
×
1833
        b2 = np.cross(a3, a1)
×
1834
        b3 = np.cross(a1, a2)
×
1835

1836
        recip_lattice = np.array([b1, b2, b3]) * 2 * np.pi / volume
×
1837
        return recip_lattice
×
1838

1839
    def get_main_axes(
1✔
1840
        self, weights: Union[str, npt.NDArray[np.float64]] = "unity"
1841
    ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.float64]]:
1842
        """
1843
        Get main axes and eigenvalues of a molecule
1844
        https://de.wikipedia.org/wiki/Tr%C3%A4gheitstensor
1845

1846
        Parameters
1847
        ----------
1848
        weights : Union[str, npt.NDArray[np.float64]] = "unity"
1849
            Specifies how the atoms are weighted.
1850
            "unity": all same weight
1851
            "Z": weighted by atomic number.
1852
            The default is 'unity'.
1853

1854
        Returns
1855
        -------
1856
        vals : npt.NDArray[np.float64]
1857
            TODO DESCRIPTION.
1858
        vecs : npt.NDArray[np.float64]
1859
            TODO DESCRIPTION.
1860

1861
        """
1862
        if weights == "unity":
×
1863
            weights = np.ones(self.n_atoms)
×
1864
        elif weights == "Z":
×
1865
            weights = np.array(
×
1866
                [self.periodic_table.get_atomic_mass(s) for s in self.species]
1867
            )
1868

1869
        coords = self.coords - np.mean(self.coords, axis=0)
×
1870
        diag_entry = np.sum(np.sum(coords**2, axis=1) * weights)
×
1871
        ident = np.eye(3) * diag_entry
×
1872
        for i in range(self.n_atoms):
×
1873
            ident -= weights[i] * np.outer(coords[i, :], coords[i, :])
×
1874

1875
        vals, vecs = scipy.linalg.eigh(ident)
×
1876
        sort_ind = np.argsort(vals)
×
1877

1878
        return vals[sort_ind], vecs[:, sort_ind]
×
1879

1880
    def get_distance_between_all_atoms(self) -> npt.NDArray[np.float64]:
1✔
1881
        """
1882
        Get the distance between all atoms in the current Geometry
1883
        object. Gives an symmetric array where distances between atom i and j
1884
        are denoted in the array elements (ij) and (ji).
1885

1886
        """
1887
        distances = scipy.spatial.distance.cdist(self.coords, self.coords)
×
1888
        return distances
×
1889

1890
    def get_closest_atoms(self, indices, species=None, n_closest=1):
1✔
1891
        """
1892
        Get the indices of the closest atom(s) for the given index or list of
1893
        indices
1894

1895
        Parameters
1896
        ----------
1897
        index: int or iterable
1898
            atoms for which the closest indices are  to be found
1899
        species: None or list of species identifiers
1900
            species to consider for closest atoms. This allows to get only the
1901
            closest atoms of the same or another species
1902
        n_closest: int
1903
            number of closest atoms to return
1904

1905
        Returns
1906
        -------
1907
        closest_atoms_list: list or list of lists
1908
            closest atoms for each entry in index
1909
        """
1910
        all_distances = self.get_distance_between_all_atoms()
×
1911

1912
        if species is None:
×
1913
            species_to_consider = list(set(self.species))
×
1914
        else:
1915
            assert isinstance(species, list), (
×
1916
                "species must be a list of species identifiers or None if all atoms should be probed"
1917
            )
1918
            species_to_consider = species
×
1919

1920
        return_single_list = False
×
1921
        if not isinstance(indices, Iterable):
×
1922
            return_single_list = True
×
1923
            indices = [indices]
×
1924

1925
        indices_to_consider = []
×
1926
        for i, s in enumerate(self.species):
×
1927
            if (s in species_to_consider) and (i not in indices):
×
1928
                indices_to_consider.append(i)
×
1929
        indices_to_consider = np.array(indices_to_consider)
×
1930

1931
        closest_atoms_list = []
×
1932
        for index in indices:
×
1933
            distances = all_distances[index, indices_to_consider]
×
1934
            distance_indices = np.argsort(distances)
×
1935
            closest_atoms = indices_to_consider[distance_indices]
×
1936
            if len(closest_atoms) > n_closest:
×
1937
                closest_atoms = closest_atoms[:n_closest]
×
1938

1939
            closest_atoms_list.append(closest_atoms.tolist())
×
1940

1941
        if return_single_list:  # can only be true if only a single index was specified
×
1942
            return closest_atoms_list[0]
×
1943
        return closest_atoms_list
×
1944

1945
    def get_distance_between_two_atoms(self, atom_indices: list) -> np.float64:
1✔
1946
        """Get the distance between two atoms in the current Geometry
1947
        object.
1948
        """
1949
        atom1 = self.coords[atom_indices[0], :]
×
1950
        atom2 = self.coords[atom_indices[1], :]
×
1951
        vec = atom2 - atom1
×
1952
        dist = np.linalg.norm(vec)
×
1953

1954
        return dist
×
1955

1956
    def get_volume_of_unit_cell(self) -> np.float64:
1✔
1957
        """
1958
        Calcualtes the volume of the unit cell.
1959

1960
        Returns
1961
        -------
1962
        volume : float
1963
            Volume of the unit cell.
1964

1965
        """
1966
        a1 = self.lattice_vectors[0]
×
1967
        a2 = self.lattice_vectors[1]
×
1968
        a3 = self.lattice_vectors[2]
×
1969
        volume = np.cross(a1, a2).dot(a3)
×
1970
        return volume
×
1971

1972
    def get_geometric_center(
1✔
1973
        self, ignore_center_attribute=False, indices=None
1974
    ) -> npt.NDArray[np.float64]:
1975
        """
1976
        Returns the center of the geometry. If the attribute *center* is set,
1977
        it is used as the definition for the center of the geometry.
1978

1979
        Parameters
1980
        ----------
1981
        ignore_center_attribute : Bool
1982
            If True, the attribute self.center is used
1983
            Otherwise, the function returns the geometric center of the structure,
1984
            i.e. the average over the position of all atoms.
1985

1986
        indices: iterable of indices or None
1987
            indices of all atoms to consider when calculating the center. Useful to calculate centers of adsorbates only
1988
            if None, all atoms are used
1989

1990
        Returns
1991
        -------
1992
        center : npt.NDArray[np.float64]
1993
            Center of the geometry
1994
        """
1995
        if (
×
1996
            not hasattr(self, "center")
1997
            or self.center is None
1998
            or ignore_center_attribute
1999
            or indices is not None
2000
        ):
2001
            if indices is None:
×
2002
                indices = np.arange(self.n_atoms)
×
2003
            center = np.mean(self.coords[indices], axis=0)
×
2004
        else:
2005
            center = np.zeros([3])
×
2006
            for i, weight in self.center.items():
×
2007
                center += self.coords[i, :] * weight
×
2008
        return center
×
2009

2010
    def get_center_of_mass(self) -> npt.NDArray[np.float64]:
1✔
2011
        """
2012
        Mind the difference to self.get_geometric_center
2013

2014
        Returns
2015
        -------
2016
        center_of_mass: npt.NDArray[np.float64]
2017
            The 3D-coordinate of the center of mass
2018

2019
        """
2020
        species_helper = []
×
2021
        for si in self.species:
×
2022
            si_new = si.split("_")[0]
×
2023
            species_helper.append(si_new)
×
2024

2025
        masses_np = np.array(
×
2026
            [self.periodic_table.get_atomic_mass(s) for s in species_helper],
2027
            dtype=np.float64,
2028
        )
2029
        center_of_mass = self.coords.T.dot(masses_np) / masses_np.sum()
×
2030
        return center_of_mass
×
2031

2032
    def get_symmetries(
1✔
2033
        self,
2034
        symmetry_precision: float = 1e-05,
2035
        remove_refelction_in_z: bool = False,
2036
    ):
2037
        """
2038
        Returns symmetries (rotation and translation matrices) from spglig.
2039
        works only for unitcell and supercell geometries (lattice vecotrs must
2040
        not be 0)
2041

2042
        Beware: The returned symmetry matrices are given with respect to
2043
        fractional coordinates, not Cartesian ones!
2044

2045
        See https://atztogo.github.io/spglib/python-spglib.html#get-symmetry
2046
        for details
2047

2048
        Parameters
2049
        ----------
2050
        save_directory : str
2051
            save directory in string format, file will be name symmetry.pickle
2052
            (default = None --> symmetry is not saved)
2053

2054
        Raises
2055
        ------
2056
        ValueError: If lattice vectors are 0
2057
        """
2058
        if np.count_nonzero(self.lattice_vectors) == 0:
1✔
2059
            print(
×
2060
                "Lattice vectors must not be 0! getSymmetry requires a unitcell-like"
2061
                "geometry file!"
2062
            )
2063
            raise ValueError(self.lattice_vectors)
×
2064

2065
        unit_cell = self.get_spglib_cell()
1✔
2066
        symmetry = spglib.get_symmetry(unit_cell, symprec=symmetry_precision)
1✔
2067

2068
        rotations = symmetry["rotations"]
1✔
2069
        translations = symmetry["translations"]
1✔
2070

2071
        if remove_refelction_in_z:
1✔
2072
            upside = rotations[:, 2, 2] == 1
×
2073
            rotations = rotations[upside, :, :]
×
2074
            translations = translations[upside, :]
×
2075

2076
        return rotations, translations
1✔
2077

2078
    def get_atomic_numbers_of_atoms(self) -> npt.NDArray[np.float64]:
1✔
2079
        """Get the atomic numbers of all atoms in the geometry file"""
2080
        species = [self.periodic_table.get_atomic_number(s) for s in self.species]
×
2081
        return np.array(species)
×
2082

2083
    def get_number_of_electrons(self) -> float:
1✔
2084
        """
2085
        Determines the number of electrons.
2086

2087
        Returns
2088
        -------
2089
        float
2090
            Number of electrons.
2091

2092
        """
2093
        electrons = []
1✔
2094
        for s in self.species:
1✔
2095
            try:
1✔
2096
                if "_" in s:
1✔
2097
                    curr_species = s.split("_")[0]
×
2098
                else:
2099
                    curr_species = s
1✔
2100
                electrons.append(self.periodic_table.get_atomic_number(curr_species))
1✔
2101

2102
            except KeyError:
×
2103
                KeyError(f"Species {s} is not known")
×
2104
        return np.sum(electrons)
1✔
2105

2106
    def get_atomic_masses(self) -> npt.NDArray[np.float64]:
1✔
2107
        """
2108
        Determines the atomic mass for all atoms.
2109

2110
        Returns
2111
        -------
2112
        masses : np.array
2113
            List of atomic masses for all atoms in the same order as
2114
            the atoms.
2115

2116
        """
2117
        masses = []
×
2118
        for s in self.species:
×
2119
            try:
×
2120
                if "_" in s:
×
2121
                    curr_species = s.split("_")[0]
×
2122
                else:
2123
                    curr_species = s
×
2124

2125
                mass = self.periodic_table.get_atomic_mass(curr_species)
×
2126
                masses.append(mass)
×
2127

2128
            except KeyError:
×
2129
                KeyError(f"Atomic mass for species {s} is not known")
×
2130
        return np.array(masses)
×
2131

2132
    def get_total_mass(self) -> float:
1✔
2133
        """
2134
        Determines the atomic mass of the entrie geometry.
2135

2136
        Returns
2137
        -------
2138
        atomic_mass : float
2139
            Atomic mass of the entrie geometry.
2140

2141
        """
2142
        atomic_mass = 0
×
2143

2144
        for s in self.species:
×
2145
            atomic_mass += self.periodic_table.get_atomic_mass(s)
×
2146

2147
        return atomic_mass
×
2148

2149
    def get_largest_atom_distance(self, dims_to_consider=(0, 1, 2)) -> float:
1✔
2150
        """
2151
        Find largest distance between atoms in geometry.
2152
        #search tags; molecule length, maximum size
2153

2154
        Parameters
2155
        ----------
2156
        dims_to_consider : list
2157
            Dimensions along which largest distance should be calculated.
2158

2159
        Returns
2160
        -------
2161
        geometry_size : float
2162
            Largest distance between two atoms in the unit cell.
2163

2164
        """
2165
        mask = np.array([i in dims_to_consider for i in range(3)], dtype=bool)
×
2166
        geometry_size = 0.0
×
2167
        for ind1 in range(self.n_atoms):
×
2168
            for ind2 in range(ind1, self.n_atoms):
×
2169
                geometry_size_test = np.linalg.norm(
×
2170
                    self.coords[ind1][mask] - self.coords[ind2][mask]
2171
                )
2172

2173
                if geometry_size_test > geometry_size:
×
2174
                    geometry_size = float(geometry_size_test)
×
2175

2176
        return geometry_size
×
2177

2178
    def get_area(self) -> np.float64:
1✔
2179
        """
2180
        Returns the area of the surface described by lattice_vectors 0 and 1 of
2181
        the geometry, assuming that the lattice_vector 2 is orthogonal to both.
2182

2183
        Returns
2184
        -------
2185
        area : float
2186
            Area of the unit cell.
2187

2188
        """
2189
        a = deepcopy(self.lattice_vectors[0, :])
×
2190
        b = deepcopy(self.lattice_vectors[1, :])
×
2191
        area = np.abs(np.cross(a, b)[-1])
×
2192

2193
        return area
×
2194

2195
    def get_area_in_atom_numbers(
1✔
2196
        self,
2197
        substrate_indices: Union[npt.NDArray[np.float64], None] = None,
2198
        substrate=None,
2199
    ) -> float:
2200
        """
2201
        Returns the area of the unit cell in terms of substrate atoms in the
2202
        topmost substrate layer. The substrate is determined by
2203
        default by the function self.getSubstrate(). For avoiding a faulty
2204
        substrate determination it can be either given through indices or
2205
        through the substrate geometry itself
2206

2207
        Parameters
2208
        ----------
2209
        substrate_indices : Union[np.array, None], optional
2210
            List of indices of substrate atoms. The default is None.
2211
        substrate : TYPE, optional
2212
            Geometry of substrate. The default is None.
2213

2214
        Returns
2215
        -------
2216
        float
2217
            Area of the unit cell in units of the area of the substrate.
2218

2219
        """
2220
        topmost_sub_layer = self.get_substrate_layers(
×
2221
            layer_indices=[0],
2222
            substrate_indices=substrate_indices,
2223
            primitive_substrate=substrate,
2224
        )
2225
        return topmost_sub_layer.n_atoms
×
2226

2227
    def get_bond_lengths(self, bond_factor: float = 1.5) -> npt.NDArray[np.float64]:
1✔
2228
        """
2229
        Parameters
2230
        ----------
2231
        Parameter for bond detection based on atom-distance : float, optional
2232
            DESCRIPTION. The default is 1.5.
2233

2234
        Returns
2235
        -------
2236
        bond_lengths : npt.NDArray[np.float64]
2237
            List of bond lengths for neighbouring atoms.
2238

2239
        """
2240
        raise NotImplementedError
×
2241

2242
        # TODO write the below function
2243
        neighbouring_atoms = self.get_all_neighbouring_atoms(bond_factor=bond_factor)
2244

2245
        bond_lengths = []
2246
        for v in neighbouring_atoms.values():
2247
            bond_lengths.append(v[1])
2248

2249
        bond_lengths = np.array(bond_lengths)
2250
        return bond_lengths
2251

2252
    def get_number_of_atom_layers(self, threshold: float = 1e-2) -> Tuple[dict, float]:
1✔
2253
        """
2254
        Return the number of atom layers.
2255

2256
        Parameters
2257
        ----------
2258
        threshold : float, optional
2259
            Threshold to determine the number of layers. The default is 1e-2.
2260

2261
        Returns
2262
        -------
2263
        dict
2264
            Number of layers per atom species.
2265
        float
2266
            Total number of layers.
2267

2268
        """
2269
        layers = self.get_atom_layers_indices(threshold=threshold)
×
2270

2271
        total_number_layers = 0
×
2272
        for atom_species in layers:
×
2273
            layers[atom_species] = len(layers[atom_species])
×
2274
            total_number_layers += layers[atom_species]
×
2275

2276
        return layers, total_number_layers
×
2277

2278
    def get_unit_cell_parameters(
1✔
2279
        self,
2280
    ) -> Tuple[np.float64, np.float64, np.float64, float, float, float]:
2281
        """
2282
        Determines unit cell parameters.
2283

2284
        Returns
2285
        -------
2286
        a : float
2287
            Length of lattice vector 1.
2288
        b : float
2289
            Length of lattice vector 2.
2290
        c : float
2291
            Length of lattice vector 3.
2292
        alpha : float
2293
            Angle between lattice vectors 1 and 2.
2294
        beta : float
2295
            Angle between lattice vectors 2 and 3.
2296
        gamma : float
2297
            Angle between lattice vectors 1 and 3.
2298

2299
        """
2300
        cell = self.lattice_vectors
×
2301

2302
        a = np.linalg.norm(cell[0])
×
2303
        b = np.linalg.norm(cell[1])
×
2304
        c = np.linalg.norm(cell[2])
×
2305

2306
        alpha = np.arccos(np.dot(cell[1], cell[2]) / (c * b))
×
2307
        gamma = np.arccos(np.dot(cell[1], cell[0]) / (a * b))
×
2308
        beta = np.arccos(np.dot(cell[2], cell[0]) / (a * c))
×
2309

2310
        return a, b, c, alpha, beta, gamma
×
2311

2312
    def get_spglib_cell(
1✔
2313
        self,
2314
    ) -> Tuple[npt.NDArray[np.float64], npt.NDArray[np.int64], list]:
2315
        """
2316
        Returns the unit cell in a format that can be used in spglib (to find symmetries)
2317

2318
        return : tuple
2319
            (lattice vectors, frac coordinates of atoms, atomic numbers)
2320
        """
2321
        coordinates = utils.get_fractional_coords(self.coords, self.lattice_vectors)
1✔
2322

2323
        atom_number = []
1✔
2324
        for atom_name in self.species:
1✔
2325
            atom_number.append(self.periodic_table.get_atomic_number(atom_name))
1✔
2326

2327
        return (self.lattice_vectors, coordinates, atom_number)
1✔
2328

2329
    def get_orientation_of_main_axis(self) -> float:
1✔
2330
        """
2331
        Get the orientation of the main axis relative to the x axis
2332

2333
        This is transformed such that it always points in the upper half of cartesian
2334
        space
2335

2336
        Returns
2337
        -------
2338
        float
2339
            angle between main axis and x axis in degree
2340
        """
2341
        main_ax = self.get_main_axes()[1][:, 0]
×
2342

2343
        if main_ax[1] < 0:
×
2344
            main_ax *= -1
×
2345

2346
        return np.arctan2(main_ax[1], main_ax[0]) * 180 / np.pi
×
2347

2348
    def get_constrained_atoms(self) -> npt.NDArray[np.int64]:
1✔
2349
        """
2350
        Returns indice of constrained atoms.
2351

2352
        Returns
2353
        -------
2354
        inds : npt.NDArray[np.int64]
2355
            Indice of constrained atoms.
2356

2357
        """
2358
        constrain = np.any(self.constrain_relax, axis=1)
1✔
2359
        inds = [i for i, c in enumerate(constrain) if c]
1✔
2360
        inds = np.array(inds)
1✔
2361
        return inds
1✔
2362

2363
    def get_unconstrained_atoms(self) -> npt.NDArray[np.int64]:
1✔
2364
        """
2365
        Returns indice of unconstrained atoms.
2366

2367
        Returns
2368
        -------
2369
        inds : npt.NDArray[np.int64]
2370
            Indice of unconstrained atoms.
2371

2372
        """
2373
        all_inds = list(range(len(self)))
1✔
2374
        keep_inds = self.get_constrained_atoms()
1✔
2375
        inds = np.array(list(set(all_inds) - set(keep_inds)))
1✔
2376
        return inds
1✔
2377

2378
    def get_cropping_indices(
1✔
2379
        self,
2380
        xlim=(-np.inf, np.inf),
2381
        ylim=(-np.inf, np.inf),
2382
        zlim=(-np.inf, np.inf),
2383
        auto_margin=False,
2384
        inverse=False,
2385
    ) -> npt.NDArray[np.int64]:
2386
        """
2387
        Gets indices of all atoms that are outside specified bounds.
2388
        If auto_margin == True then an additional margin of the maximum covalent radius
2389
        is added to all borders
2390
        :param inverse : if True, gets indices of all atoms INSIDE specified bounds
2391

2392
        """
2393
        assert len(xlim) == 2, "xlim must have a lower and an upper bound"
×
2394
        assert len(ylim) == 2, "ylim must have a lower and an upper bound"
×
2395
        assert len(zlim) == 2, "zlim must have a lower and an upper bound"
×
2396
        if auto_margin:
×
2397
            margin = max(
×
2398
                [self.periodic_table.get_covalent_radius(s) for s in self.species]
2399
            )
2400
            xlim = xlim[0] - margin, xlim[1] + margin
×
2401
            ylim = ylim[0] - margin, ylim[1] + margin
×
2402
            zlim = zlim[0] - margin, zlim[1] + margin
×
2403

2404
        remove = np.zeros(len(self), bool)
×
2405
        remove = remove | (self.coords[:, 0] < xlim[0])
×
2406
        remove = remove | (self.coords[:, 0] > xlim[1])
×
2407
        remove = remove | (self.coords[:, 1] < ylim[0])
×
2408
        remove = remove | (self.coords[:, 1] > ylim[1])
×
2409
        remove = remove | (self.coords[:, 2] < zlim[0])
×
2410
        remove = remove | (self.coords[:, 2] > zlim[1])
×
2411
        indices_to_remove = np.arange(len(self))[remove]
×
2412
        if inverse:
×
2413
            indices_to_remove = [
×
2414
                i for i in range(self.n_atoms) if i not in indices_to_remove
2415
            ]
2416

2417
        indices_to_remove = np.array(indices_to_remove)
×
2418

2419
        return indices_to_remove
×
2420

2421
    def get_substrate_indices_from_parts(
1✔
2422
        self, do_warn=True
2423
    ) -> Union[None, npt.NDArray[np.int64]]:
2424
        """
2425
        Returns indices of those atoms that a part of the substrate. The
2426
        definition of the substrate does NOT rely of it being a metal or the
2427
        height or something like that. Instead a geometry part named
2428
        'substrate' must be defined.
2429

2430
        Parameters
2431
        ----------
2432
        do_warn : boolean
2433
            Can be set to False to suppress warnings
2434

2435
        Returns
2436
        -------
2437
        substrate_indices : npt.NDArray[np.int64]
2438

2439
        """
2440
        substrate_key = "substrate"  # should probably moved to a class variable, but first finished folder-read-write
×
2441

2442
        if not hasattr(self, "geometry_parts") or not hasattr(
×
2443
            self, "geometry_part_descriptions"
2444
        ):
2445
            if do_warn:
×
2446
                print("getSubstrate: geometry parts not defined")
×
2447
            return None
×
2448

2449
        if substrate_key not in self.geometry_part_descriptions:
×
2450
            if do_warn:
×
2451
                print(
×
2452
                    "getSubstrate: geometry parts are defined, "
2453
                    f"but part '{substrate_key}' not found"
2454
                )
2455
            return None
×
2456

2457
        index_of_geometry_parts_substrate = self.geometry_part_descriptions.index(
×
2458
            substrate_key
2459
        )
2460
        substrate_indices = self.geometry_parts[index_of_geometry_parts_substrate]
×
2461
        substrate_indices = np.array(substrate_indices)
×
2462
        return substrate_indices
×
2463

2464
    def get_indices_of_metal(self) -> npt.NDArray[np.int64]:
1✔
2465
        """
2466
        Gets indices of all atoms with atomic number > 18 and atomic numbers
2467
        3,4,11,12,13,14
2468

2469
        Returns
2470
        -------
2471
        metal_atoms : npt.NDArray[np.int64]
2472
            Inidices of metallic atoms.
2473

2474
        """
2475
        atom_inds = [self.periodic_table.get_atomic_number(s) for s in self.species]
×
2476
        metal_atoms = []
×
2477
        for i, ind in enumerate(atom_inds):
×
2478
            if (ind > 18) or (ind in [3, 4, 11, 12, 13, 14]):
×
2479
                metal_atoms.append(i)
×
2480

2481
        metal_atoms = np.array(metal_atoms, dtype=np.int64)
×
2482
        return metal_atoms
×
2483

2484
    def get_indices_of_molecules(self, substrate_species=None) -> npt.NDArray[np.int64]:
1✔
2485
        """
2486
        Fetches the indices of the substrate atoms, but it defaults to
2487
        just returning all non-metal atom's indices!
2488

2489
        If substrate_species is given indices of all atoms that are not of the
2490
        substrate species are returned.
2491

2492
        """
2493
        if substrate_species:
×
2494
            substrate_indices = self.get_indices_of_species(substrate_species)
×
2495
        else:
2496
            substrate_indices = self.get_indices_of_metal()
×
2497

2498
        molecules = [i for i in range(self.n_atoms) if i not in substrate_indices]
×
2499
        molecules = np.array(molecules)
×
2500
        return molecules
×
2501

2502
    def get_indices_of_species(self, species) -> npt.NDArray[np.int64]:
1✔
2503
        """
2504
        Returns all indices of atoms the are of species defined in the input.
2505
        species can be a string of a list
2506

2507
        """
2508
        # make sure species is a list
2509
        if isinstance(species, str):
×
2510
            species = [species]
×
2511

2512
        species_indices = [i for i in range(self.n_atoms) if self.species[i] in species]
×
2513
        species_indices = np.array(species_indices)
×
2514
        return species_indices
×
2515

2516
    def get_substrate_indices(
1✔
2517
        self, primitive_substrate=None, dimension=2, threshold=0.3
2518
    ) -> npt.NDArray[np.int64]:
2519
        """
2520
        This method returns the indices of all atoms that are part of the substrate.
2521
        Often these are simply all metal atoms of the geometry.
2522
        But the substrate can also be organic in which case it can't be picked by being a metal.
2523
        And there might be multiple adsorbate layers in which case only the molecules of the highest layer
2524
        shall be counted as adsorbates and the others are part of the substrate.
2525

2526
        dimension: int      only used if primitive_substrate is not None
2527
        threshold: float    only used if primitive_substrate is not None
2528

2529
        Returns
2530
        -------
2531
        indices of all substrate atoms: npt.NDArray[np.int64]
2532

2533
        """
2534
        # case 1: if a primitive_substrate was passed, use that one for the decision
2535
        # (copied from self.removeSubstrate)
2536

2537
        substrate_indices = []
×
2538

2539
        # case 2: if no substrate was passed but a geometry_parts "substrate" is defined in geometry_parts, use that one
2540
        substrate_indices_from_parts = self.get_substrate_indices_from_parts(
×
2541
            do_warn=False
2542
        )
2543

2544
        if primitive_substrate is not None:
×
2545
            substrate_species = set(primitive_substrate.species)
×
2546
            substrate_heights = primitive_substrate.coords[:, dimension]
×
2547
            substrate_candidate_indices = [
×
2548
                i for i, s in enumerate(self.species) if s in substrate_species
2549
            ]
2550
            for c in substrate_candidate_indices:
×
2551
                if np.any(
×
2552
                    np.absolute(substrate_heights - self.coords[c, dimension])
2553
                    < threshold
2554
                ):
2555
                    substrate_indices.append(c)
×
2556

2557
        elif substrate_indices_from_parts is not None:
×
2558
            substrate_indices = substrate_indices_from_parts
×
2559

2560
        else:
2561
            # case 3: if neither a substrate was passed, nor a geometry_parts "substrate" is defined,
2562
            #            use a fallback solution: assume that the substrate (and nothing else) is a metal
2563
            warnings.warn(
×
2564
                "Geometry.getIndicesOfAdsorbates: Substrate is not explicitly defined. "
2565
                "Using fallback solution of counting all metal atoms as substrate."
2566
            )
2567
            substrate_indices = self.get_indices_of_metal()
×
2568

2569
        return np.array(substrate_indices)
×
2570

2571
    def get_adsorbate_indices(self, primitive_substrate=None) -> npt.NDArray[np.int64]:
1✔
2572
        """
2573
        This method returns the indices of all atoms that are NOT part of the substrate.
2574
        In a classical organic monolayer on a metal substrate these are simply all molecules.
2575
        But the substrate can also be organic in which case it can't be picked by being a metal.
2576
        And there might be multiple adsorbate layers in which case only the molecules of the highest layer
2577
        shall be counted as adsorbates.
2578

2579
        Returns
2580
        -------
2581
        indices of all adsorbate atoms: npt.NDArray[np.int64]
2582
        """
2583
        substrate_indices = self.get_substrate_indices(
×
2584
            primitive_substrate=primitive_substrate
2585
        )
2586
        # invert:
2587
        return np.array(
×
2588
            [i for i in self.get_indices_of_all_atoms() if i not in substrate_indices]
2589
        )
2590

2591
    def get_indices_of_all_atoms(self, species=None):
1✔
2592
        if species is None:
×
2593
            return [i for i in range(self.n_atoms)]
×
2594
        return [i for i in range(self.n_atoms) if self.species[i] == species]
×
2595

2596
    def get_atom_layers_indices(self, threshold: float = 1e-2) -> dict:
1✔
2597
        """
2598
        Returns a dict of the following form.
2599

2600
        Parameters
2601
        ----------
2602
        threshold : float, optional
2603
            Treshold within which atoms are considered to be in the same layer.
2604
            The default is 1e-2.
2605

2606
        Returns
2607
        -------
2608
        layers : dict
2609
            {<Element symbol>: {height: [indices of atoms of element at height]}}.
2610

2611
        """
2612
        layers = {}
×
2613

2614
        for ind, atom_coord in enumerate(self.coords):
×
2615
            atom_species = self.species[ind]
×
2616
            if atom_species not in layers:
×
2617
                layers[atom_species] = {}
×
2618

2619
            add_new_z_coord = True
×
2620
            for z_coord in layers[atom_species].keys():
×
2621
                if abs(atom_coord[2] - z_coord) < threshold:
×
2622
                    layers[atom_species][z_coord].append(ind)
×
2623
                    add_new_z_coord = False
×
2624

2625
            if add_new_z_coord:
×
2626
                layers[atom_species][atom_coord[2]] = [ind]
×
2627

2628
        return layers
×
2629

2630
    def get_atom_layers_indices_by_height(self, threshold: float = 1e-2) -> dict:
1✔
2631
        """
2632
        Similarly to get_atom_layers_indices this function returns a dict
2633
        continaing info about height and the indices of atoms at that height.
2634

2635
        Parameters
2636
        ----------
2637
        threshold : float, optional
2638
            reshold within which atoms are considered to be in the same layer.
2639
            The default is 1e-2.
2640

2641
        Returns
2642
        -------
2643
        layers_by_height : dict
2644
            {height: [indices of atoms at that height]}.
2645

2646
        """
2647
        layers_by_species = self.get_atom_layers_indices(threshold=threshold)
×
2648

2649
        layers_by_height = defaultdict(list)
×
2650

2651
        # --- merge height-indices dicts ---
2652
        for data in layers_by_species.values():
×
2653
            for height, indices in data.items():
×
2654
                new = True
×
2655
                for new_height in layers_by_height.keys():
×
2656
                    if abs(height - new_height) < threshold:
×
2657
                        layers_by_height[new_height] += indices
×
2658
                        new = False
×
2659
                if new:
×
2660
                    layers_by_height[height] += indices
×
2661

2662
        # sort dictionary by descending height
2663
        layers_by_height = dict(sorted(layers_by_height.items(), reverse=True))
×
2664
        return layers_by_height
×
2665

2666
    def get_list_of_neighbouring_atoms(self, bond_factor: float = 1.5) -> dict:
1✔
2667
        """
2668
        Get a dictionary of neighbouring atoms.
2669

2670
        Parameters
2671
        ----------
2672
        bond_factor : float, optional
2673
            Multiply for covelent radius. If the distance between two atoms is
2674
            smaller thand (r_0+r_1)*bond_factor, the atoms are considered
2675
            neighbours. The default is 1.5.
2676

2677
        Returns
2678
        -------
2679
        neighbouring_atoms : dict
2680
            {(index_0, index_1) : [(element_0, element_1), atom distance]}.
2681

2682
        """
2683
        coords = self.coords
×
2684
        species = self.species
×
2685

2686
        all_species = set(species)
×
2687
        all_species_pairs = itertools.product(all_species, repeat=2)
×
2688

2689
        bond_thresholds = {}
×
2690
        for pair in all_species_pairs:
×
2691
            bond_thresholds[pair] = (
×
2692
                self.periodic_table.get_covalent_radius(pair[0])
2693
                + self.periodic_table.get_covalent_radius(pair[1])
2694
            ) * bond_factor
2695

2696
        neighbouring_atoms = {}
×
2697

2698
        for i, coord in enumerate(coords):
×
2699
            for j, coord_test in enumerate(coords):
×
2700
                if i >= j:
×
2701
                    continue
×
2702

2703
                pair_index = (i, j)
×
2704

2705
                if pair_index not in neighbouring_atoms:
×
2706
                    dist = np.linalg.norm(coord - coord_test)
×
2707

2708
                    pair_species = (species[i], species[j])
×
2709

2710
                    if dist < bond_thresholds[pair_species]:
×
2711
                        neighbouring_atoms[pair_index] = [pair_species, dist]
×
2712

2713
        return neighbouring_atoms
×
2714

2715
    def get_principal_moments_of_inertia(self) -> npt.NDArray[np.float64]:
1✔
2716
        """
2717
        Calculates the eigenvalues of the moments of inertia matrix
2718

2719
        Returns
2720
        -------
2721
        moments : np.array, shape=(3,), dtype=np.float64
2722
            principal moments of inertia in kg * m**2
2723

2724
        """
2725
        masses_kg = [
×
2726
            units.ATOMIC_MASS_IN_KG * self.periodic_table.get_atomic_mass(s)
2727
            for s in self.species
2728
        ]
2729

2730
        center_of_mass = self.get_center_of_mass()
×
2731
        r_to_center_in_m = units.ANGSTROM_IN_METER * (self.coords - center_of_mass)
×
2732

2733
        ###########
2734
        # begin: code based on ase/atoms.py: get_moments_of_inertia
2735
        # (GNU Lesser General Public License)
2736
        # Initialize elements of the inertial tensor
2737
        I11 = I22 = I33 = I12 = I13 = I23 = 0.0
×
2738
        for i in range(len(self)):
×
2739
            x, y, z = r_to_center_in_m[i]
×
2740
            m = masses_kg[i]
×
2741

2742
            I11 += m * (y**2 + z**2)
×
2743
            I22 += m * (x**2 + z**2)
×
2744
            I33 += m * (x**2 + y**2)
×
2745
            I12 += -m * x * y
×
2746
            I13 += -m * x * z
×
2747
            I23 += -m * y * z
×
2748

2749
        I = np.array(
×
2750
            [[I11, I12, I13], [I12, I22, I23], [I13, I23, I33]],
2751
            dtype=np.float64,
2752
        )
2753

2754
        moments, _ = np.linalg.eigh(I)
×
2755
        return moments
×
2756

2757
    def get_homogeneous_field(self):
1✔
2758
        """Field is a numpy array (Ex, Ey, Ez) with the Field in V/A"""
2759
        if not hasattr(self, "_homogeneous_field"):
×
2760
            self.homogeneous_field = None
×
2761

2762
        return self.homogeneous_field
×
2763

2764
    ###############################################################################
2765
    #               Get Properties in comparison to other Geometry                #
2766
    ###############################################################################
2767
    def get_distance_to_equivalent_atoms(self, other_geometry) -> np.float64:
1✔
2768
        """
2769
        Calculate the maximum distance that atoms of geom would have to be moved,
2770
        to coincide with the atoms of self.
2771

2772
        """
2773
        _, dist = self.get_transformation_indices(other_geometry, get_distances=True)
×
2774
        return np.max(dist)
×
2775

2776
    def get_transformation_indices(
1✔
2777
        self,
2778
        other_geometry,
2779
        get_distances=False,
2780
        periodic_2D=False,
2781
    ):
2782
        """
2783
        Associates every atom in self to the closest atom of the same specie in
2784
        other_geometry.
2785

2786
        If self should be orderd like other_geometry then this is done in the
2787
        following way:
2788
        >>> transformation_indices = other_geometry.get_transformation_indices(self)
2789
        >>> self.reorder_atoms(transformation_indices)
2790

2791
        Parameters
2792
        ----------
2793
        other_geometry: geometry
2794
        norm_threshold : float
2795
        get_distances : bool
2796
            The default is False.
2797
        periodic_2D : bool
2798
            The default is False.
2799

2800
        Returns
2801
        -------
2802
        transformation_indices : np.array.
2803
            The positions on the array correspond to the atoms in self;
2804
            the values of the array correspond to the atoms in other_geometry
2805

2806
        """
2807
        assert len(self) == len(other_geometry), (
×
2808
            f"Geometries have different number of atoms {len(self)} != {len(other_geometry)}"
2809
        )
2810

2811
        # Replicate other geometry to also search in neighbouring cells
2812
        if periodic_2D:
×
2813
            other_geometry = other_geometry.get_periodic_replica((3, 3, 1))
×
2814
            other_geometry.move_all_atoms_by_fractional_coords([-1 / 3.0, -1 / 3.0, 0])
×
2815

2816
        # Get the atomic numbers of each geometry file: Later only compare matching atom types
2817
        Z_values_1 = np.array(
×
2818
            [self.periodic_table.get_atomic_number(s) for s in self.species],
2819
            np.int64,
2820
        )
2821
        Z_values_2 = np.array(
×
2822
            [self.periodic_table.get_atomic_number(s) for s in other_geometry.species],
2823
            np.int64,
2824
        )
2825
        unique_Z = set(Z_values_1)
×
2826

2827
        # Allocate output arrays
2828
        transformation_indices = np.zeros(len(self), np.int64)
×
2829
        distances = np.zeros(len(self))
×
2830
        atom2_indices = np.arange(len(other_geometry)) % len(self)
×
2831

2832
        # Loop over all types of atoms
2833
        for Z in unique_Z:
×
2834
            # Select all the coordinates that belong to that species
2835
            select1 = Z_values_1 == Z
×
2836
            select2 = Z_values_2 == Z
×
2837
            # Calculate all distances between the geometries
2838
            dist_matrix = scipy.spatial.distance_matrix(
×
2839
                self.coords[select1, :], other_geometry.coords[select2, :]
2840
            )
2841
            # For each row (= atom in self with species Z) find the index of the other_geometry (with species Z) that is closest
2842
            index_of_smallest_mismatch = np.argmin(dist_matrix, axis=1)
×
2843
            transformation_indices[select1] = atom2_indices[select2][
×
2844
                index_of_smallest_mismatch
2845
            ]
2846
            if get_distances:
×
2847
                distances[select1] = [
×
2848
                    dist_matrix[i, index_of_smallest_mismatch[i]]
2849
                    for i in range(len(dist_matrix))
2850
                ]
2851

2852
        if get_distances:
×
2853
            return transformation_indices, distances
×
2854
        return transformation_indices
×
2855

2856
    def is_equivalent(
1✔
2857
        self, geom, tolerance=0.01, check_neightbouring_cells=False
2858
    ) -> bool:
2859
        """
2860
        Check if this geometry is equivalent to another given geometry.
2861
        The function checks that the same atoms sit on the same positions
2862
        (but possibly in some permutation)
2863

2864
        :param check_neightbouring_cells: for periodic structures, recognizes two structures as equivalent, even if one\
2865
         of them has its atoms distributed in different unit cells compared to the other. More complete, but slower.
2866
        """
2867
        # Check that both geometries have same number of atoms
2868
        # If not, they cannot be equivalent
2869
        if geom.n_atoms != self.n_atoms:
×
2870
            return False
×
2871

2872
        # check in neighbouring cells, to account for geometries 'broken' around the cell border
2873
        if check_neightbouring_cells:
×
2874
            if self.lattice_vectors is not None:
×
2875
                geom = geom.get_periodic_replica(
×
2876
                    (3, 3), explicit_replications=[[-1, 0, 1], [-1, 0, 1]]
2877
                )
2878
            else:
2879
                print("Non periodic structure. Ignoring check_neighbouring_cells")
×
2880

2881
        n_atoms = self.n_atoms
×
2882
        n_atoms_geom = geom.n_atoms
×
2883
        # Check for each atom in coords1 that is has a matching atom in coords2
2884
        for n1 in range(n_atoms):
×
2885
            is_ok = False
×
2886
            for n2 in range(n_atoms_geom):
×
2887
                if self.species[n1] == geom.species[n2]:
×
2888
                    d = np.linalg.norm(self.coords[n1, :] - geom.coords[n2, :])
×
2889
                    if d < tolerance:
×
2890
                        # Same atom and same position
2891
                        is_ok = True
×
2892
                        break
×
2893

2894
            if not is_ok:
×
2895
                return False
×
2896
        return True
×
2897

2898
    def is_equivalent_up_to_translation(
1✔
2899
        self,
2900
        geom,
2901
        get_translation=False,
2902
        tolerance=0.01,
2903
        check_neighbouring_cells=False,
2904
    ):
2905
        """
2906
            Returns True if self can be transformed into geom by a translation
2907
                (= without changing the geometry itself).
2908

2909
        Parameters
2910
        ----------
2911
        geom : geometry
2912
            Geometry to compare to.
2913
        get_translation : bool
2914
            Additionally return the found translation
2915
        tolerance : float
2916
            Tolerance threshold for larget mismatch between two atoms, below
2917
            which they are still considered to be at the sameposition.
2918
        check_neightbouring_cells : bool
2919
            For periodic structures, recognizes two structures as equivalent,
2920
            even if one of them has its atoms distributed in different
2921
            unit cells compared to the other. More complete, but slower.
2922

2923
        Returns
2924
        -------
2925
            is_equivalent : bool
2926
                True if equivalent.
2927
            translation : np.array
2928
                Translation vetror between equivalent geometries
2929
        """
2930
        # shift both geometries to origin, get their relative translation.
2931
        # Ignore center attribute (GeometryFile.center), if defined
2932
        meanA = self.get_geometric_center(ignore_center_attribute=True)
×
2933
        meanB = geom.get_geometric_center(ignore_center_attribute=True)
×
2934
        translation = meanA - meanB
×
2935
        self.center_coordinates(ignore_center_attribute=True)
×
2936
        geom.center_coordinates(ignore_center_attribute=True)
×
2937

2938
        # check if they are equivalent (up to permutation)
2939
        is_equivalent = self.is_equivalent(
×
2940
            geom, tolerance, check_neightbouring_cells=check_neighbouring_cells
2941
        )
2942

2943
        # undo shifting to origin
2944
        self.coords += meanA
×
2945
        geom.coords += meanB
×
2946

2947
        if get_translation:
×
2948
            return is_equivalent, translation
×
2949
        return is_equivalent
×
2950

2951
    ###########################################################################
2952
    #                          Get Part of a Geometry                         #
2953
    ###########################################################################
2954
    def get_atoms_by_indices(self, atom_indices: npt.NDArray[np.int64]) -> "Geometry":
1✔
2955
        """
2956
        Return a geometry instance with the atoms listed in atom_indices
2957

2958
        Parameters
2959
        ----------
2960
        atom_indices : Union[int, np.array]
2961
            List of integers, indices of those atoms which should be copied to
2962
            new geometry
2963

2964
        Returns
2965
        -------
2966
        new_geom : Geometry
2967

2968
        """
2969
        new_geom = self.__class__()
×
2970
        new_geom.add_atoms(
×
2971
            self.coords[atom_indices, :],
2972
            [self.species[i] for i in atom_indices],
2973
            constrain_relax=self.constrain_relax[atom_indices],
2974
            initial_moment=[self.initial_moment[i] for i in atom_indices],
2975
            initial_charge=[self.initial_charge[i] for i in atom_indices],
2976
        )
2977
        new_geom.lattice_vectors = self.lattice_vectors
×
2978
        return new_geom
×
2979

2980
    def get_atoms_by_species(self, species) -> "Geometry":
1✔
2981
        """
2982
        Get new geometry file with specific atom species
2983
        """
2984
        L = np.array(self.species) == species
×
2985
        atom_indices = np.where(L)[0]
×
2986
        return self.get_atoms_by_indices(atom_indices)
×
2987

2988
    def get_primitive_slab(self, surface, threshold=1e-6):
1✔
2989
        """
2990
        Generates a primitive slab unit cell with the z-direction perpendicular
2991
        to the surface.
2992

2993
        Arguments:
2994
        ----------
2995
        surface : array_like
2996
            miller indices, eg. (1,1,1)
2997

2998
        threshold : float
2999
            numerical threshold for symmetry operations
3000

3001
        Returns
3002
        -------
3003
        primitive_slab : Geometry
3004
        """
3005
        lattice, scaled_positions, atomic_numbers = spglib.standardize_cell(
×
3006
            self.get_spglib_cell()
3007
        )
3008

3009
        surface_vector = (
×
3010
            surface[0] * lattice[0, :]
3011
            + surface[1] * lattice[1, :]
3012
            + surface[2] * lattice[2, :]
3013
        )
3014

3015
        # TODO: this way of building lattice vectors parallel to the surface
3016
        # is not ideal for certain surfaces
3017
        dot_0 = surface_vector.dot(surface_vector)
×
3018
        dot_1 = surface_vector.dot(lattice[0, :])
×
3019
        dot_2 = surface_vector.dot(lattice[1, :])
×
3020
        dot_3 = surface_vector.dot(lattice[2, :])
×
3021

3022
        if abs(dot_1) > threshold:
×
3023
            frac = Fraction(dot_0 / dot_1).limit_denominator(1000)
×
3024
            n, m = frac.numerator, frac.denominator
×
3025
            v1 = m * surface_vector - n * lattice[0, :]
×
3026
        else:
3027
            v1 = lattice[0, :]
×
3028

3029
        if abs(dot_2) > threshold:
×
3030
            frac = Fraction(dot_0 / dot_2).limit_denominator(1000)
×
3031
            n, m = frac.numerator, frac.denominator
×
3032
            v2 = m * surface_vector - n * lattice[1, :]
×
3033
        else:
3034
            v2 = lattice[1, :]
×
3035

3036
        if abs(dot_3) > threshold:
×
3037
            frac = Fraction(dot_0 / dot_3).limit_denominator(1000)
×
3038
            n, m = frac.numerator, frac.denominator
×
3039
            v3 = m * surface_vector - n * lattice[2, :]
×
3040
        else:
3041
            v3 = lattice[2, :]
×
3042

3043
        surface_lattice = np.zeros((3, 3))
×
3044
        surface_lattice[0, :] = surface_vector
×
3045

3046
        ind = 1
×
3047
        for v in [v1, v2, v3]:
×
3048
            if not np.linalg.norm(v) == 0:
×
3049
                surface_lattice[ind, :] = v
×
3050
                rank = np.linalg.matrix_rank(surface_lattice)
×
3051

3052
                if rank == ind + 1:
×
3053
                    ind += 1
×
3054
                    if ind == 3:
×
3055
                        break
×
3056

3057
        # flip surface lattice such that surface normal becomes the z-axis
3058
        surface_lattice = np.flip(surface_lattice, 0)
×
3059

3060
        slab = self.__class__()
×
3061
        slab.lattice_vectors = surface_lattice
×
3062

3063
        # shellsize 100 such that the code does not run infinitely
3064
        shellsize = 100
×
3065
        for shell in range(shellsize):
×
3066
            add_next_shell = False
×
3067
            for h in range(-shell, shell + 1):
×
3068
                for k in range(-shell, shell + 1):
×
3069
                    for l in range(-shell, shell + 1):
×
3070
                        if (abs(h) < shell) and (abs(k) < shell) and (abs(l) < shell):
×
3071
                            continue
×
3072

3073
                        for new_species, coord in zip(atomic_numbers, scaled_positions):
×
3074
                            new_coord = coord.dot(lattice) + np.array([h, k, l]).dot(
×
3075
                                lattice
3076
                            )
3077
                            frac_new_coord = utils.get_fractional_coords(
×
3078
                                new_coord, surface_lattice
3079
                            )
3080

3081
                            L1 = np.sum(frac_new_coord >= 1 - threshold)
×
3082
                            L2 = np.sum(frac_new_coord < -threshold)
×
3083

3084
                            if not L1 and not L2:
×
3085
                                slab.add_atoms(
×
3086
                                    [new_coord],
3087
                                    [
3088
                                        self.periodic_table.get_chemical_symbol(
3089
                                            new_species
3090
                                        )
3091
                                    ],
3092
                                )
3093
                                add_next_shell = True
×
3094

3095
            if not shell == 0 and not add_next_shell:
×
3096
                break
×
3097

3098
            if shell == 100:
×
3099
                warnings.warn(
×
3100
                    "<Geometry.get_primitive_slab> could not build a correct slab."
3101
                )
3102

3103
        slab.align_lattice_vector_to_vector(np.array([0, 0, 1]), 2)
×
3104
        slab.align_lattice_vector_to_vector(np.array([1, 0, 0]), 0)
×
3105

3106
        scaled_slab_lattice = np.array(slab.lattice_vectors)
×
3107
        # break symmetry in z-direction
3108
        scaled_slab_lattice[2, :] *= 2
×
3109
        frac_coords = utils.get_fractional_coords(slab.coords, scaled_slab_lattice)
×
3110
        species = [self.periodic_table.get_atomic_number(s) for s in slab.species]
×
3111

3112
        (
×
3113
            primitive_slab_lattice,
3114
            primitive_slab_scaled_positions,
3115
            primitive_slab_atomic_numbers,
3116
        ) = spglib.find_primitive(
3117
            (scaled_slab_lattice, frac_coords, species), symprec=1e-5
3118
        )
3119

3120
        primitive_slab_species = [
×
3121
            self.periodic_table.get_chemical_symbol(s)
3122
            for s in primitive_slab_atomic_numbers
3123
        ]
3124
        primitive_slab_coords = primitive_slab_scaled_positions.dot(
×
3125
            primitive_slab_lattice
3126
        )
3127
        # replace lattice vector in z-direction
3128
        primitive_slab_lattice[2, :] = slab.lattice_vectors[2, :]
×
3129

3130
        primitive_slab = self.__class__()
×
3131
        primitive_slab.lattice_vectors = primitive_slab_lattice
×
3132
        primitive_slab.add_atoms(primitive_slab_coords, primitive_slab_species)
×
3133
        primitive_slab.map_to_first_unit_cell()
×
3134

3135
        # Sanity check: primitive_slab must be reducable to the standard unit cell
3136
        check_lattice, _, _ = spglib.standardize_cell(primitive_slab.get_spglib_cell())
×
3137

3138
        assert np.allclose(check_lattice, lattice), (
×
3139
            "<Geometry.get_primitive_slab> the slab that was constructed \
3140
        could not be reduced to the original bulk unit cell. Something \
3141
        must have gone wrong."
3142
        )
3143

3144
        return primitive_slab
×
3145

3146
    def get_slab(
1✔
3147
        self,
3148
        layers: int,
3149
        surface: Union[npt.NDArray[np.int64], None] = None,
3150
        threshold: float = 1e-6,
3151
        surface_replica: Tuple[int, int] = (1, 1),
3152
        vacuum_height: Union[float, None] = None,
3153
        bool_shift_slab_to_bottom: bool = False,
3154
    ) -> "Geometry":
3155
        """
3156
        Generates a slab.
3157

3158
        Parameters
3159
        ----------
3160
        layers : int
3161
            Number of layers of the slab.
3162
        surface : npt.NDArray[np.int64] | None, optional
3163
            miller indices, eg. (1,1,1)
3164
        threshold : float, optional
3165
            numerical threshold for symmetry operations
3166
        surface_replica : Tuple[int, int], optional
3167
            Replications of surface. The default is (1,1).
3168
        vacuum_height : float | None, optional
3169
            DESCRIPTION. The default is None.
3170
        bool_shift_slab_to_bottom : bool, optional
3171
            DESCRIPTION. The default is False.
3172

3173
        Returns
3174
        -------
3175
        slab_new : Geometry
3176
            New Geometry
3177

3178
        """
3179
        if surface is not None:
×
3180
            primitive_slab = self.get_primitive_slab(surface, threshold=threshold)
×
3181
        else:
3182
            primitive_slab = self
×
3183

3184
        slab_layers = primitive_slab.get_number_of_atom_layers()[1]
×
3185

3186
        replica = np.array([1, 1, int(np.ceil(layers / slab_layers))], dtype=np.int32)
×
3187
        replica[:2] = surface_replica
×
3188
        slab_new = primitive_slab.get_periodic_replica(tuple(replica))
×
3189

3190
        slab_new_layers = slab_new.get_atom_layers_indices()
×
3191

3192
        for atom_species in slab_new_layers:
×
3193
            z_coords = list(slab_new_layers[atom_species])
×
3194
            z_coords = sorted(z_coords)
×
3195

3196
            n_layers_to_remove = len(z_coords) - layers
×
3197

3198
            atom_indices_to_remove = []
×
3199
            for ind in range(n_layers_to_remove):
×
3200
                atom_indices_to_remove += slab_new_layers[atom_species][z_coords[ind]]
×
3201

3202
            slab_new.remove_atoms(np.array(atom_indices_to_remove, dtype=np.int32))
×
3203

3204
            if vacuum_height is not None:
×
3205
                slab_new.set_vacuum_height(
×
3206
                    vac_height=vacuum_height,
3207
                    bool_shift_to_bottom=bool_shift_slab_to_bottom,
3208
                )
3209
            elif bool_shift_slab_to_bottom:
×
3210
                self.shift_to_bottom()
×
3211

3212
        return slab_new
×
3213

3214
    def get_colliding_groups(self, distance_threshold=1e-2, check_3D=False):
1✔
3215
        """
3216
        Remove atoms that are too close too each other from the geometry file.
3217
        This approach is useful if one maps back atoms into a different cell
3218
        and then needs to get rid of overlapping atoms
3219

3220
        Parameters
3221
        ----------
3222
        distance_threshold: float
3223
            maximum distance between atoms below which they are counted as
3224
            duplicates
3225

3226
        Returns
3227
        -------
3228
        """
3229
        # get all distances between all atoms
3230

3231
        z_period = [-1, 0, 1] if check_3D else [0]
×
3232
        index_tuples = []
×
3233
        for i in [-1, 0, 1]:
×
3234
            for j in [-1, 0, 1]:
×
3235
                for k in z_period:
×
3236
                    curr_shift = (
×
3237
                        i * self.lattice_vectors[0, :]
3238
                        + j * self.lattice_vectors[1, :]
3239
                        + k * self.lattice_vectors[2, :]
3240
                    )
3241

3242
                    atom_distances = scipy.spatial.distance.cdist(
×
3243
                        self.coords, self.coords + curr_shift
3244
                    )
3245
                    index_tuples += self._get_collision_indices(
×
3246
                        atom_distances, distance_threshold
3247
                    )
3248
        if len(index_tuples) > 0:
×
3249
            G = nx.Graph()
×
3250
            G = nx.from_edgelist(
×
3251
                itertools.chain.from_iterable(
3252
                    itertools.pairwise(e) for e in index_tuples
3253
                )
3254
            )
3255
            G.add_nodes_from(set.union(*map(set, index_tuples)))  # adding single items
×
3256
            atoms_to_remove = list(nx.connected_components(G))
×
3257
            return [sorted(list(s)) for s in atoms_to_remove]
×
3258
        return []
×
3259

3260
    def get_cropped_geometry(
1✔
3261
        self,
3262
        xlim=(-np.inf, np.inf),
3263
        ylim=(-np.inf, np.inf),
3264
        zlim=(-np.inf, np.inf),
3265
        auto_margin=False,
3266
    ):
3267
        """
3268
        Returns a copy of the object to which self.crop has been applied
3269

3270
        """
3271
        newgeom = deepcopy(self)
×
3272
        newgeom.crop(xlim=xlim, ylim=ylim, zlim=zlim, auto_margin=auto_margin)
×
3273
        return newgeom
×
3274

3275
    def get_substrate(self, primitive_substrate=None):
1✔
3276
        substrate_indices = self.get_substrate_indices(
×
3277
            primitive_substrate=primitive_substrate
3278
        )
3279
        return self.get_atoms_by_indices(substrate_indices)
×
3280

3281
    def get_adsorbates(self, primitive_substrate=None):
1✔
3282
        adsorbate_indices = self.get_adsorbate_indices(
×
3283
            primitive_substrate=primitive_substrate
3284
        )
3285
        return self.get_atoms_by_indices(adsorbate_indices)
×
3286

3287
    def get_periodic_replica(
1✔
3288
        self,
3289
        replications: tuple,
3290
        lattice: Union[npt.NDArray[np.float64], None] = None,
3291
        explicit_replications: Union[list, None] = None,
3292
    ):
3293
        """
3294
        Return a new geometry file that is a periodic replica of the original
3295
        file. Repeats the geometry N-1 times in all given directions:
3296
        (1,1,1) returns the original file
3297

3298
        Parameters
3299
        ----------
3300
        TODO Fix types
3301
        replications : tuple or list
3302
            number of replications for each dimension
3303
        lattice : numpy array of shape [3x3]
3304
            super-lattice vectors to use for periodic replication
3305
            if lattice is None (default) the lattice vectors from the current
3306
            geometry file are used.
3307
        explicit_replications : iterable of iterables
3308
             a way to explicitly define which replicas should be made.
3309
             example: [[-1, 0, 1], [0, 1, 2, 3], [0]] will repeat 3 times in x
3310
             (centered) and 4 times in y (not centered)
3311

3312
        Returns
3313
        -------
3314
        New geometry
3315
        """
3316
        # TODO implement geometry_parts the right way (whatever this is)
3317
        if lattice is None:
×
3318
            lattice = np.array(self.lattice_vectors)
×
3319

3320
        if explicit_replications:
×
3321
            rep = explicit_replications
×
3322
            lattice_multipliers = [np.max(t) - np.min(t) for t in explicit_replications]
×
3323
        else:
3324
            rep = [list(range(r)) for r in replications]
×
3325
            lattice_multipliers = replications
×
3326

3327
        # old: n_replicas = np.abs(np.prod(replications))
3328
        n_replicas = np.prod([len(i) for i in rep])
×
3329
        n_atoms_new = n_replicas * self.n_atoms
×
3330
        new_coords = np.zeros([n_atoms_new, 3])
×
3331
        new_species = list(self.species) * n_replicas
×
3332
        new_constrain = list(self.constrain_relax) * n_replicas
×
3333

3334
        insert_pos = 0
×
3335
        # itertools.product = nested for loop
3336
        for frac_offset in itertools.product(*rep):
×
3337
            frac_shift = np.zeros([1, 3])
×
3338
            frac_shift[0, : len(frac_offset)] = frac_offset
×
3339
            offset = utils.get_cartesian_coords(frac_shift, lattice)
×
3340
            new_coords[insert_pos : insert_pos + self.n_atoms, :] = self.coords + offset
×
3341
            insert_pos += self.n_atoms
×
3342

3343
        new_geom = self.__class__()
×
3344

3345
        new_geom.add_atoms(new_coords, new_species, new_constrain)
×
3346
        new_geom.lattice_vectors = lattice
×
3347

3348
        # save original lattice vector for visualization
3349
        if hasattr(self, "original_lattice_vectors"):
×
3350
            new_geom.original_lattice_vectors = copy.deepcopy(
×
3351
                self.original_lattice_vectors
3352
            )
3353
        else:
3354
            new_geom.original_lattice_vectors = copy.deepcopy(self.lattice_vectors)
×
3355

3356
        for i, r in enumerate(lattice_multipliers):
×
3357
            new_geom.lattice_vectors[i, :] *= r
×
3358
        return new_geom
×
3359

3360
    def get_split_into_molecules(self, threshold) -> list:
1✔
3361
        """
3362
        Splits a structure into individual molecules. Two distinct molecules
3363
        A and B are defined as two sets of atoms, such that no atom in A is
3364
        closer than the selected thresold to any atom of B
3365

3366
        """
3367
        coords = deepcopy(self.coords)
×
3368
        distances = distance_matrix(coords, coords)
×
3369
        distances[distances <= threshold] = 1
×
3370
        distances[distances > threshold] = 0
×
3371

3372
        def scan_line(line_index, matrix, already_scanned_lines_indices):
×
3373
            already_scanned_lines_indices.append(line_index)
×
3374
            line = matrix[line_index]
×
3375
            links = np.nonzero(line)[0]
×
3376
            links = [l for l in links if l not in already_scanned_lines_indices]
×
3377
            return links, already_scanned_lines_indices
×
3378

3379
        molecules_indices_sets = []
×
3380
        scanned_lines_indices = []
×
3381
        indices_set = []
×
3382
        # scan lines one by one, but skips those that have already been examined
3383
        for i in range(len(distances)):
×
3384
            if i in scanned_lines_indices:
×
3385
                continue
×
3386
            # add line to the present set
3387
            indices_set.append(i)
×
3388
            # get indices of the lines connected to the examined one
3389
            links, scanned_lines_indices = scan_line(
×
3390
                i, distances, scanned_lines_indices
3391
            )
3392
            indices_set += links
×
3393
            # as long as new links are found, adds the new lines to the present set
3394
            while len(links) > 0:
×
3395
                new_links = []
×
3396
                for l in links:
×
3397
                    if l not in scanned_lines_indices:
×
3398
                        new_links_part, scanned_lines_indices = scan_line(
×
3399
                            l, distances, scanned_lines_indices
3400
                        )
3401
                        new_links += new_links_part
×
3402
                links = set(new_links)
×
3403
                indices_set += links
×
3404
            # once no more links are found, stores the present set and starts a new one
3405
            molecules_indices_sets.append(indices_set)
×
3406
            indices_set = []
×
3407

3408
        molecules = []
×
3409
        for molecule_indices in molecules_indices_sets:
×
3410
            complementary_indices = [
×
3411
                x for x in self.get_indices_of_all_atoms() if x not in molecule_indices
3412
            ]
3413
            g = deepcopy(self)
×
3414
            g.remove_atoms(np.array(complementary_indices))
×
3415
            molecules.append(g)
×
3416

3417
        return molecules
×
3418

3419
    def get_layers(self, layer_indices: list, threshold: float = 1e-2):
1✔
3420
        """
3421
        Get substrate layer by indices. The substrate is determined by
3422
        default by the function self.get_substrate. For avoiding a faulty
3423
        substrate determination it can be either given through indices or
3424
        through the substrate geometry itself.
3425

3426
        Parameters
3427
        ----------
3428
        layer_indices : list
3429
            List of indices of layers that shall be returned.
3430
        substrate_indices : Union[None, list], optional
3431
            List of indices of substrate atoms. The default is None.
3432
        substrate : Union[None, Geometry()], optional
3433
            Geometry of substrate. The default is None.
3434
        threshold : float, optional
3435
            DESCRIPTION. The default is 1e-2.
3436
        primitive_substrate : Geometry, optional
3437
            DESCRIPTION. The default is None.
3438

3439
        Returns
3440
        -------
3441
        geometry_of_layer : Geometry
3442
            Geometry of layer.
3443

3444
        """
3445
        layers = self.get_atom_layers_indices_by_height(threshold=threshold)
×
3446

3447
        heights = list(layers.keys())
×
3448
        heights = np.sort(heights)
×
3449
        heights = heights[::-1]
×
3450

3451
        # if layer_indices is an empty list, keeps the substrate as a whole
3452
        if not layer_indices:
×
3453
            geometry_of_layer = self
×
3454
        else:
3455
            geometry_of_layer = self.__class__()
×
3456
            geometry_of_layer.lattice_vectors = self.lattice_vectors
×
3457
            for layer_ind in layer_indices:
×
3458
                geometry_of_layer += self.get_atoms_by_indices(
×
3459
                    layers[heights[layer_ind]]
3460
                )
3461

3462
        return geometry_of_layer
×
3463

3464
    def get_substrate_layers(
1✔
3465
        self,
3466
        layer_indices: list,
3467
        threshold: float = 1e-2,
3468
        primitive_substrate=None,
3469
        substrate_indices: Union[None, npt.NDArray[np.int64]] = None,
3470
    ):
3471
        """
3472
        Get substrate layer by indices. The substrate is determined by
3473
        default by the function self.get_substrate. For avoiding a faulty
3474
        substrate determination it can be either given through indices or
3475
        through the substrate geometry itself.
3476

3477
        Parameters
3478
        ----------
3479
        layer_indices : list
3480
            List of indices of layers that shall be returned.
3481
        threshold : float, optional
3482
            DESCRIPTION. The default is 1e-2.
3483
        primitive_substrate : TYPE, optional
3484
            DESCRIPTION. The default is None.
3485
        substrate_indices : Union[None, list], optional
3486
            List of indices of substrate atoms. The default is None.
3487

3488
        Returns
3489
        -------
3490
        geometry_of_layer : TYPE
3491
            Geometry of substrate layer.
3492

3493
        """
3494
        if substrate_indices is not None:
×
3495
            sub = self.get_atoms_by_indices(substrate_indices)
×
3496
        else:
3497
            sub = self.get_substrate(primitive_substrate=primitive_substrate)
×
3498

3499
        return sub.get_layers(layer_indices=layer_indices, threshold=threshold)
×
3500

3501
    ###########################################################################
3502
    #                           Evaluation Functions                          #
3503
    ###########################################################################
3504
    def check_symmetry(
1✔
3505
        self,
3506
        tolerance: float,
3507
        R: npt.NDArray[np.float64],
3508
        t: npt.NDArray[np.float64] = np.array([0.0, 0.0, 0.0]),
3509
        return_symmetrical: bool = False,
3510
    ):
3511
        """
3512
        Returns True if the geometry is symmetric with respect to the
3513
        transformation, and False if it is not. If the geometry is periodic,
3514
        transformation can be tuple (rotation, translation) or np.array
3515
        (only rotation), otherwise it can only be np.array
3516

3517
        Parameters
3518
        ----------
3519
        tolerance ; float
3520
            Tolerance for checking symmetry.
3521

3522
        R : npt.NDArray[np.float64]
3523
            Symmetry transformation agaist which geometry should be checked.
3524

3525
        t: npt.NDArray[np.float64]
3526
            Translation vector
3527

3528
        return_symmetrical : bool
3529
            Return the corresponding transformed geometry together with the
3530
            result.
3531

3532
        Returns
3533
        -------
3534
        is_symmetric : bool
3535
        symm_geometry : Geometry
3536

3537
        or
3538

3539
        is_symmetric : bool
3540

3541
        """
3542
        # original structure with all atoms in the first unit cell
3543
        self_1UC = copy.deepcopy(self)
×
3544
        self_1UC.map_to_first_unit_cell()
×
3545

3546
        # original structure centered for reordering
3547
        centered_geometry = copy.deepcopy(self)
×
3548
        centered_geometry.center_coordinates()
×
3549

3550
        # apply transformation
3551
        symm_geometry = copy.deepcopy(self)
×
3552
        symm_geometry.transform_fractional(R, np.array([0, 0, 0]), self.lattice_vectors)
×
3553
        symm_geometry.move_all_atoms_by_fractional_coords(t)
×
3554

3555
        # prevent problems if the center is very close to the edge
3556
        center = utils.get_fractional_coords(
×
3557
            symm_geometry.get_geometric_center(), symm_geometry.lattice_vectors
3558
        )
3559
        center[:2] %= 1.0
×
3560
        if 1 - center[0] < 0.001:
×
3561
            adjust = -(center[0] - 0.0001)
×
3562
            symm_geometry.move_all_atoms_by_fractional_coords([adjust, 0, 0])
×
3563
        if 1 - center[1] < 0.001:
×
3564
            adjust = -(center[1] - 0.0001)
×
3565
            symm_geometry.move_all_atoms_by_fractional_coords([0, adjust, 0])
×
3566

3567
        symm_geometry.map_center_of_atoms_to_first_unit_cell(
×
3568
            lattice_vectors=self.lattice_vectors
3569
        )
3570

3571
        # reorder atoms
3572
        offset_symm = np.mean(symm_geometry.coords, axis=0)
×
3573
        symm_geometry.center_coordinates()
×
3574
        indices = centered_geometry.get_transformation_indices(symm_geometry)
×
3575
        symm_geometry.reorder_atoms(np.array(indices))
×
3576
        symm_geometry.move_all_atoms(offset_symm)
×
3577

3578
        # compare in first unit cell
3579
        symm_geometry_1UC = copy.deepcopy(symm_geometry)
×
3580
        symm_geometry_1UC.map_to_first_unit_cell()
×
3581
        is_symmetric = symm_geometry_1UC.is_equivalent(
×
3582
            self_1UC, tolerance=tolerance, check_neightbouring_cells=True
3583
        )
3584

3585
        if return_symmetrical:
×
3586
            return is_symmetric, symm_geometry
×
3587
        return is_symmetric
×
3588

3589
    ###############################################################################
3590
    #                                 Visualisation                               #
3591
    ###############################################################################
3592
    def visualise(
1✔
3593
        self,
3594
        axes=[0, 1],
3595
        min_zorder=0,
3596
        value_list=None,
3597
        maxvalue=None,
3598
        minvalue=None,
3599
        cbar_label="",
3600
        hide_axes=False,
3601
        axis_labels=True,
3602
        auto_limits=True,
3603
        crop_ratio=None,
3604
        brightness_modifier=None,
3605
        print_lattice_vectors=False,
3606
        alpha=1.0,
3607
        linewidth=1,
3608
        lattice_linewidth=None,
3609
        lattice_color="k",
3610
        atom_scale=1,
3611
        highlight_inds=[],
3612
        highlight_color="C2",
3613
        color_list=None,
3614
        cmap=None,
3615
        ax=None,
3616
        xlim=None,
3617
        ylim=None,
3618
        zlim=None,
3619
        plot_method="circles",
3620
        invert_colormap=False,
3621
        edge_color=None,
3622
        show_colorbar=True,
3623
        reverse_sort_inds=False,
3624
        axis_labels_format="/",
3625
    ) -> None:
3626
        """
3627
        Generates at plt-plot of the current geometry. This function has a
3628
        large number of options. In most cases the following examples will
3629
        work well:
3630

3631
        - Visualise the geometry:
3632
        geometry.visualise()
3633

3634
        - Turn aff axis:
3635
        geometry.visualise(hide_axes=True)
3636

3637
        - Turn off axis and set limits:
3638
        geometry.visualise(hide_axes=True, xlim=(-10, 10))
3639

3640
        - If you want to look at the geoemtry in the xz-plane:
3641
        geometry.visualise(axes=[0,2], hide_axes=True, xlim=(-10, 10))
3642

3643
        Visualise is one of the most useful things about geometry. Reading
3644
        through this code you may think that it is very ugly and on to of that
3645
        if has it's own imports. Still it is a great function and it must be
3646
        part of geometry. If you think otherwise you are wrong.
3647

3648
        Note from Dylan: I completely agree if you think it's ugly and should be killed
3649
        with fire. Do not listen to Lukas. He is wrong. Fortunately I managed to
3650
        convince him to get rid of the function-scoped imports so you're welcome.
3651

3652
        Parameter:
3653
        ----------
3654
        axes : list of 2 int elements
3655
            axis that should be visualized, x=0, y=1, z=2
3656
            By default, we look at the geometry from:
3657
            the "top" (our viewpoint is at z = +infinity) when we visualize the xy plane;
3658
            the "right" (our viewpoint is at x = +infinity) when we visualize the yz plane;
3659
            the "front" (our viewpoint is at y = -infinity) when we visualize the xz plane.
3660
            In order to visualize the geometry from the opposite viewpoints, one needs
3661
            to use the reverse_sort_inds flag, and invert the axis when necessary
3662
            (= set axis limits so that the first value is larger than the second value)
3663

3664
        min_zorder : int
3665
            plotting layer
3666

3667
        value_list : None or list of length nr. atoms
3668

3669
        maxvalue : None
3670

3671
        cbar_label : str
3672

3673
        hide_axes : bool
3674
            hide axis
3675

3676
        axis_labels : bool
3677
            generates automatic axis labels
3678

3679
        auto_limits : bool
3680
            set xlim, ylim automatically
3681

3682
        crop_ratio: float
3683
            defines the ratio between xlim and ylim if auto_limits is enabled
3684

3685
        brightness_modifier : float or list/array with length equal to the number of atoms
3686
            modifies the brightness of selected atoms. If brightness_modifier is a
3687
            list/array, then brightness_modifier[i] sets the brightness for atom i,
3688
            otherwise all atoms are set to the same brightness value. This is done by
3689
            tweaking the 'lightness' value of said atoms' color in the HSL
3690
            (hue-saturation-lightness) colorspace. Effect of brightness_modifier in
3691
            detail:
3692
            -1.0 <= brightness_modifier < 0.0  : darker color
3693
            brightness_modifier == 0.0 or None : original color
3694
            0.0 < brightness_modifier <= 1.0   :  brighter color
3695

3696
        print_lattice_vectors : bool
3697
            display lattice vectors
3698

3699
        print_unit_cell : bool
3700
            display original unit cell
3701

3702
        alpha : float between 0 and 1
3703

3704
        color_list : list or string
3705
            choose colors for visualizing each atom. If only one color is passed, all
3706
            atoms will have that color.
3707

3708
        plot_method: str
3709
            circles: show filled circles for each atom
3710
            wireframe: show molecular wireframe, standard settings: don't show H,
3711

3712
        reverse_sort_inds: bool
3713
            if set to True, inverts the order at which atoms are visualized, allowing
3714
            to visualise the geometry from the "bottom", from the "left" or from the
3715
            "back".
3716
            Example: if one wants to visualise the geometry from the "left"
3717
            (= viewpoint at x=-infinity), atoms at lower x values should be
3718
            visualised after atoms at high x values, and hide them. This is the
3719
            opposite of the default behavior of this function, and can be achieved with
3720
            reverse_sort_inds=True
3721
            NOTE: in order to correctly visualize the structure from these non-default
3722
            points of view, setting this flag to True is not sufficient: one must also
3723
            invert the XY axes of the plot where needed.
3724
            Example: when visualising from the "left", atoms with negative y values
3725
            should appear on the right side of the plot, and atoms with positive y
3726
            values should appear on the left side of the plot. But if one simply sets
3727
            reverse_sort_inds=True, atoms with negative y values will appear on the left
3728
            side of the plot (because the x axis of the plot, the horizontal axis, goes
3729
            from left to right!) and vice-versa. This is equivalent to visualising a
3730
            mirrored image of the structure. To visualise the structure correctly, one
3731
            should then set the x_limits of the plot with a first value smaller than the
3732
            second value, so the x axis is inverted, and shows y-negative values on the
3733
            left and viceversa.
3734

3735
        Returns
3736
        -------
3737
        None
3738

3739
        """
3740
        # default for lattice_linewidth (which is used to draw the lattice)
3741
        if lattice_linewidth is None:
×
3742
            lattice_linewidth = 2 * linewidth
×
3743

3744
        orig_inds = np.arange(self.n_atoms)
×
3745
        remove_inds = []
×
3746
        if xlim is not None:
×
3747
            remove_x = self.get_cropping_indices(xlim=xlim, auto_margin=True)
×
3748
            remove_inds += list(remove_x)
×
3749
        if ylim is not None:
×
3750
            remove_y = self.get_cropping_indices(ylim=ylim, auto_margin=True)
×
3751
            remove_inds += list(remove_y)
×
3752
        if zlim is not None:
×
3753
            remove_z = self.get_cropping_indices(zlim=zlim, auto_margin=True)
×
3754
            remove_inds += list(remove_z)
×
3755

3756
        crop_inds = list(set(remove_inds))
×
3757

3758
        if len(crop_inds) > 0:
×
3759
            orig_inds = [orig_inds[i] for i in orig_inds if i not in crop_inds]
×
3760
            cropped_geom = copy.deepcopy(self)
×
3761
            cropped_geom.remove_atoms(np.array(crop_inds))
×
3762
        else:
3763
            cropped_geom = self
×
3764

3765
        if ax is None:
×
3766
            ax = plt.gca()
×
3767

3768
        axnames = ["x", "y", "z"]
×
3769
        orig_coords = cropped_geom.coords
×
3770
        orig_species = cropped_geom.species
×
3771
        #        orig_constrain = cropped_geom.constrain_relax
3772

3773
        # sorting along projecting dimension.
3774
        # If sort_ind == 1, which means that we look at XZ, along the Y axis, in order to enforce our default behaviour
3775
        # of looking at the XZ from "under" (== from the negative side of the Y axis), we need to flip the order
3776
        # at which we see atoms, so we reverse the order of sort inds.
3777
        # If the flat reverse_sort_inds is set to True, the order will be flipped again, to bring us out of our default.
3778
        for i in range(3):
×
3779
            if i not in axes:
×
3780
                sort_ind = i
×
3781

3782
        inds = np.argsort(orig_coords[:, sort_ind])
×
3783

3784
        if sort_ind == 1:
×
3785
            inds = inds[::-1]
×
3786
        if reverse_sort_inds:
×
3787
            inds = inds[::-1]
×
3788

3789
        orig_inds = [orig_inds[i] for i in inds]
×
3790
        coords = orig_coords[inds]
×
3791
        #        constrain = orig_constrain[inds]
3792
        species = [orig_species[i] for i in inds]
×
3793
        n_atoms = len(species)
×
3794
        circlesize = [
×
3795
            self.periodic_table.get_covalent_radius(s) * atom_scale for s in species
3796
        ]
3797

3798
        # Specify atom colors by value list or default atom colors
3799
        if value_list is None and color_list is None:
×
3800
            colors = [self.periodic_table.get_species_colours(s) for s in species]
×
3801
            colors = np.array(colors)
×
3802
        elif color_list is not None:
×
3803
            if len(color_list) == 1:
×
3804
                colors = list(color_list) * len(self.species)
×
3805
                colors = [mpl.colors.to_rgb(colors[i]) for i in inds]
×
3806
            else:
3807
                assert len(species) == len(color_list), (
×
3808
                    "Color must be specified for all atoms or none!"
3809
                    + f" Expected {len(species)}, but got {len(color_list)} values"
3810
                )
3811
                colors = [
×
3812
                    mpl.colors.to_rgb(color_list[i]) for i in inds
3813
                ]  # converting all types of color inputs to rgba here
3814
            colors = np.array(colors)
×
3815
        elif value_list is not None:
×
3816
            assert len(value_list) == self.n_atoms, (
×
3817
                "Number of Values does not match number of atoms in geometry"
3818
            )
3819
            values = [value_list[i] for i in orig_inds]
×
3820

3821
            if minvalue is not None:
×
3822
                assert maxvalue is not None, (
×
3823
                    "Error! If minvalue is defined also maxvalue must be defined"
3824
                )
3825

3826
            if maxvalue is None and minvalue is None:
×
3827
                maxvalue = np.max(np.abs(value_list))
×
3828
                minvalue = -maxvalue
×
3829

3830
                if maxvalue < 1e-5:
×
3831
                    maxvalue = 1e-5
×
3832
                    print(
×
3833
                        "Maxvalue for colormap not specified and smaller 1E-5, \nsetting it automatically to: ",
3834
                        maxvalue,
3835
                    )
3836
                else:
3837
                    print(
×
3838
                        "Maxvalue for colormap not specified, \nsetting it automatically to: ",
3839
                        maxvalue,
3840
                    )
3841

3842
            if maxvalue is not None and minvalue is None:
×
3843
                minvalue = -maxvalue
×
3844

3845
            if cmap is None:
×
3846
                if invert_colormap:
×
3847
                    cw = plt.get_cmap("coolwarm_r")
×
3848
                else:
3849
                    cw = plt.get_cmap("coolwarm")
×
3850
            else:
3851
                cw = plt.get_cmap(cmap)
×
3852

3853
            cNorm = matplotlib.colors.Normalize(vmin=minvalue, vmax=maxvalue)
×
3854
            scalarMap = cmx.ScalarMappable(norm=cNorm, cmap=cw)
×
3855

3856
            a = np.array([[minvalue, maxvalue]])
×
3857
            img = plt.imshow(a, cmap=cw)
×
3858
            img.set_visible(False)
×
3859

3860
            colors = []
×
3861
            for v in values:
×
3862
                colors.append(scalarMap.to_rgba(v)[:3])  # reomve alpha channel
×
3863
            colors = np.array(colors)
×
3864

3865
        # make specified atoms brighter by adding color_offset to all rgb values
3866

3867
        if brightness_modifier is not None:
×
3868
            # Check if brightness modifier is flat (i.e. a single value) or per atom (list of length n_atoms)
3869
            if isinstance(brightness_modifier, float) or isinstance(
×
3870
                brightness_modifier, int
3871
            ):
3872
                brightness_modifier = brightness_modifier * np.ones(n_atoms)
×
3873

3874
            else:
3875
                # Sort list according to orig_inds (which is already cropped if necessary!)
3876
                assert len(brightness_modifier) == self.n_atoms, (
×
3877
                    "Argument 'brightness_modifier' must either be a "
3878
                    "scalar (float or int) or a list with length equal "
3879
                    "to the number of atoms"
3880
                )
3881
                brightness_modifier = [brightness_modifier[i] for i in orig_inds]
×
3882

3883
            assert len(brightness_modifier) == n_atoms, (
×
3884
                "Something went wrong while reformatting brightness_modifier!"
3885
            )
3886
            for i in range(n_atoms):
×
3887
                # TODO fix the pyright errors
3888
                hls_color = np.array(
×
3889
                    colorsys.rgb_to_hls(*colors[i, :])  # pyright:ignore
3890
                )
3891
                hls_color[1] += brightness_modifier[i] * (1 - hls_color[1])
×
3892
                hls_color = np.clip(hls_color, 0, 1)
×
3893
                colors[i, :] = colorsys.hls_to_rgb(*hls_color)  # pyright:ignore
×
3894
        else:
3895
            brightness_modifier = np.zeros(n_atoms)
×
3896

3897
        zorder = min_zorder
×
3898

3899
        if plot_method == "circles":
×
3900
            for i in range(len(species)):
×
3901
                if plot_method == "circles":
×
3902
                    x1 = coords[i, axes[0]]
×
3903
                    x2 = coords[i, axes[1]]
×
3904
                    if orig_inds[i] not in highlight_inds:
×
3905
                        if edge_color is None:
×
3906
                            curr_edge_color = (
×
3907
                                np.zeros(3) + brightness_modifier[i]
3908
                                if brightness_modifier[i] > 0
3909
                                else np.zeros(3)
3910
                            )
3911
                        else:
3912
                            curr_edge_color = edge_color
×
3913

3914
                        ax.add_artist(
×
3915
                            mpl.patches.Circle(
3916
                                (x1, x2),
3917
                                circlesize[i],
3918
                                color=colors[i],
3919
                                zorder=zorder,
3920
                                linewidth=linewidth,
3921
                                alpha=alpha,
3922
                                ec=curr_edge_color,
3923
                            )
3924
                        )
3925
                    else:
3926
                        if edge_color is None:
×
3927
                            curr_edge_color = highlight_color
×
3928
                        else:
3929
                            curr_edge_color = edge_color
×
3930
                        ax.add_artist(
×
3931
                            mpl.patches.Circle(
3932
                                [x1, x2],
3933
                                circlesize[i],
3934
                                color=colors[i],
3935
                                zorder=zorder,
3936
                                linewidth=linewidth,
3937
                                alpha=alpha,
3938
                                ec=curr_edge_color,
3939
                            )
3940
                        )
3941
                    zorder += 2
×
3942

3943
        elif plot_method == "wireframe":
×
3944
            raise NotImplementedError("self.visualize_wireframe is not implemented")
×
3945
            # self.visualizeWireframe(coords=coords, species=species,
3946
            #                         linewidth=linewidth, min_zorder=min_zorder,
3947
            #                         axes=axes, alpha=alpha, **kwargs)
3948

3949
        if print_lattice_vectors:
×
3950
            ax.add_artist(
×
3951
                plt.arrow(
3952
                    0,
3953
                    0,
3954
                    *cropped_geom.lattice_vectors[0, axes],
3955
                    zorder=zorder,
3956
                    fc=lattice_color,
3957
                    ec=lattice_color,
3958
                    head_width=0.5,
3959
                    head_length=1,
3960
                )
3961
            )
3962
            ax.add_artist(
×
3963
                plt.arrow(
3964
                    0,
3965
                    0,
3966
                    *cropped_geom.lattice_vectors[1, axes],
3967
                    zorder=zorder,
3968
                    fc=lattice_color,
3969
                    ec=lattice_color,
3970
                    head_width=0.5,
3971
                    head_length=1,
3972
                )
3973
            )
3974

3975
        # scale:
3976
        xmax = np.max(coords[:, axes[0]]) + 2
×
3977
        xmin = np.min(coords[:, axes[0]]) - 2
×
3978
        ymax = np.max(coords[:, axes[1]]) + 2
×
3979
        ymin = np.min(coords[:, axes[1]]) - 2
×
3980

3981
        if auto_limits:
×
3982
            if print_lattice_vectors:
×
3983
                xmin_lattice = np.min(cropped_geom.lattice_vectors[:, axes[0]]) - 1
×
3984
                xmax_lattice = np.max(cropped_geom.lattice_vectors[:, axes[0]]) + 1
×
3985
                ymin_lattice = np.min(cropped_geom.lattice_vectors[:, axes[1]]) - 1
×
3986
                ymax_lattice = np.max(cropped_geom.lattice_vectors[:, axes[1]]) + 1
×
3987

3988
                ax_xmin = min(xmin, xmin_lattice)
×
3989
                ax_xmax = max(xmax, xmax_lattice)
×
3990
                ax_ymin = min(ymin, ymin_lattice)
×
3991
                ax_ymax = max(ymax, ymax_lattice)
×
3992

3993
            else:
3994
                ax_xmin, ax_xmax, ax_ymin, ax_ymax = xmin, xmax, ymin, ymax
×
3995
                # allow for a fixed ratio when defining the limits
3996
                # For this calculate the lengths and make the smaller limit longer so that the ratio fits
3997

3998
            if crop_ratio is not None:
×
3999
                len_xlim = ax_xmax - ax_xmin
×
4000
                len_ylim = ax_ymax - ax_ymin
×
4001
                curr_crop_ratio = len_xlim / len_ylim
×
4002

4003
                if curr_crop_ratio > crop_ratio:
×
4004
                    # make y limits larger
4005
                    y_padding_fac = len_xlim / (crop_ratio * len_ylim)
×
4006
                    y_padding = len_ylim * (y_padding_fac - 1)
×
4007
                    ax_ymin -= y_padding / 2
×
4008
                    ax_ymax += y_padding / 2
×
4009

4010
                else:
4011
                    # make x limits larger
4012
                    x_padding_fac = (crop_ratio * len_ylim) / len_xlim
×
4013
                    x_padding = len_xlim * (x_padding_fac - 1)
×
4014
                    ax_xmin -= x_padding / 2
×
4015
                    ax_xmax += x_padding / 2
×
4016

4017
            # TODO: fix pyright linting errors
4018
            ax.set_xlim([ax_xmin, ax_xmax])  # pyright:ignore
×
4019
            ax.set_ylim([ax_ymin, ax_ymax])  # pyright:ignore
×
4020

4021
        # If limits are given, set them
4022
        limits = [xlim, ylim, zlim]
×
4023
        x1lim = limits[axes[0]]
×
4024
        x2lim = limits[axes[1]]
×
4025
        if x1lim is not None:
×
4026
            ax.set_xlim(x1lim)
×
4027
        if x2lim is not None:
×
4028
            ax.set_ylim(x2lim)
×
4029

4030
        if axis_labels:
×
4031
            if axis_labels_format == "/":
×
4032
                ax.set_xlabel(rf"{axnames[axes[0]]} / $\AA$")
×
4033
                ax.set_ylabel(rf"{axnames[axes[1]]} / $\AA$")
×
4034
            elif axis_labels_format == "[]":
×
4035
                ax.set_xlabel(rf"{axnames[axes[0]]} [$\AA$]")
×
4036
                ax.set_ylabel(rf"{axnames[axes[1]]} [$\AA$]")
×
4037

4038
        if show_colorbar and (value_list is not None):
×
4039
            cbar = plt.colorbar(ax=ax)
×
4040
            cbar.ax.set_ylabel(cbar_label)
×
4041

4042
        ax.set_aspect("equal")
×
4043
        plt.grid(False)
×
4044
        if hide_axes:
×
4045
            ax.set_axis_off()
×
4046
            ax.get_xaxis().set_visible(False)
×
4047
            ax.get_yaxis().set_visible(False)
×
4048

4049
    ###############################################################################
4050
    #                            Protected Functions                              #
4051
    ###############################################################################
4052
    def _get_collision_indices(self, atom_distances, distance_threshold=1e-2):
1✔
4053
        """Helper function for removeDuplicateAtoms
4054

4055
        Parameters
4056
        ----------
4057
        distance_threshold: float
4058
            maximum distance between atoms below which they are counted as duplicates
4059

4060
        Returns
4061
        -------
4062
        atoms_to_remove : list
4063
            indices of atoms that can be removed due to collision
4064
        """
4065
        # get all distances between all atoms
4066
        is_collision = atom_distances < distance_threshold
×
4067

4068
        colliding_atoms_dict = {}
×
4069
        colliding_atoms_list = []
×
4070

4071
        # loop over all atoms
4072
        for i in range(self.n_atoms):
×
4073
            # evaluate only if atom is not already on the black list
4074
            if i not in colliding_atoms_list:
×
4075
                colliding_atoms_dict[i] = []
×
4076
                # loop over all distances to other atoms, neglecting the diagonal (thus i+1)
4077
                for j in range(i + 1, self.n_atoms):
×
4078
                    if is_collision[i, j]:
×
4079
                        colliding_atoms_dict[i].append(j)
×
4080
                        colliding_atoms_list.append(j)
×
4081

4082
        return [
×
4083
            (k, ind) for k, value in colliding_atoms_dict.items() for ind in list(value)
4084
        ]
4085

4086

4087
class AimsGeometry(Geometry):
1✔
4088
    def parse(self, text):
1✔
4089
        """
4090
        Parses text from AIMS geometry file and sets all necessary parameters
4091
        in AimsGeometry.
4092

4093
        Parameters
4094
        ----------
4095
        text : str
4096
            line wise text file of AIMS geometry.
4097
        """
4098
        atom_lines = []
1✔
4099
        is_fractional = False
1✔
4100
        is_own_hessian = False
1✔
4101
        self.trust_radius = False
1✔
4102
        self.vacuum_level = None
1✔
4103
        self.constrain_relax = []
1✔
4104
        self.external_force = []
1✔
4105
        self.calculate_friction = []
1✔
4106
        self.multipoles = []
1✔
4107
        self.homogeneous_field = None
1✔
4108
        self.symmetry_params = None
1✔
4109
        self.n_symmetry_params = None
1✔
4110
        self.symmetry_LVs = None  # symmetry_LVs should have str values, not float, to allow for the inclusion of the parameters
1✔
4111
        symmetry_LVs_lines = []
1✔
4112
        self.symmetry_frac_coords = None  # symmetry_frac_coords should have str values, not float, to allow for the inclusion of the parameters
1✔
4113
        symmetry_frac_lines = []
1✔
4114
        lattice_vector_lines = []
1✔
4115
        atom_line_ind = []
1✔
4116
        hessian_lines = []
1✔
4117
        text_lines = text.split("\n")
1✔
4118

4119
        for ind_line, line in enumerate(text_lines):
1✔
4120
            line = line.strip()  # Remove leading and trailing space in line
1✔
4121
            # Comment in input file
4122
            if line.startswith("#"):
1✔
4123
                if "DFT_ENERGY " in line:
×
4124
                    self.DFT_energy = float(line.split()[2])
×
4125
                elif "ADSORPTION_ENERGY " in line:
×
4126
                    self.E_ads = float(line.split()[2])
×
4127
                elif "ADSORPTION_ENERGY_UNRELAXED " in line:
×
4128
                    self.E_ads_sp = float(line.split()[2])
×
4129
                elif "CENTER" in line:
×
4130
                    self.center = ast.literal_eval(" ".join(line.split()[2:]))
×
4131
                # check if it is an own Hessian and not from a geometry optimization
4132
                elif "own_hessian" in line:
×
4133
                    is_own_hessian = True
×
4134

4135
                # PARTS defines parts of the geometry that can later on be treated separately.
4136
                # intended for distinction between different molecules and substrate
4137
                elif "PARTS" in line:
×
4138
                    part_definition = ast.literal_eval(" ".join(line.split()[2:]))
×
4139
                    if isinstance(part_definition, dict):
×
4140
                        for k, v in part_definition.items():
×
4141
                            self.geometry_part_descriptions.append(k)
×
4142
                            self.geometry_parts.append(v)
×
4143
                    elif isinstance(part_definition, list):
×
4144
                        if isinstance(part_definition[0], list):
×
4145
                            for part in part_definition:
×
4146
                                self.geometry_part_descriptions.append("")
×
4147
                                self.geometry_parts.append(part)
×
4148
                        else:
4149
                            self.geometry_parts.append(part)
×
4150
                            self.geometry_part_descriptions.append("")
×
4151

4152
                else:
4153
                    # Remove '#' at beginning of line, then remove any leading whitespace
4154
                    line_comment = line[1:].lstrip()
×
4155
                    # Finally add line comment to self.comment_lines
4156
                    self.comment_lines.append(line_comment)
×
4157

4158
            else:
4159
                # Extract all lines that define atoms, lattice vectors, multipoles or the Hessian matrix
4160
                if "atom" in line:
1✔
4161
                    atom_lines.append(line)
1✔
4162
                    atom_line_ind.append(ind_line)
1✔
4163
                if "lattice_vector" in line:
1✔
4164
                    lattice_vector_lines.append(line)
1✔
4165
                # c Check for fractional coordinates
4166
                if "_frac" in line:
1✔
4167
                    is_fractional = True
1✔
4168
                if "hessian_block" in line:
1✔
4169
                    hessian_lines.append(line)
×
4170
                if "trust_radius" in line:
1✔
4171
                    self.trust_radius = float(line.split()[-1])
×
4172
                if "set_vacuum_level" in line:
1✔
4173
                    self.vacuum_level = float(line.split()[1])
×
4174
                if "multipole" in line:
1✔
4175
                    multipole = [float(x) for x in list(line.split())[1:]]
×
4176
                    assert len(multipole) == 5
×
4177
                    self.multipoles.append(multipole)
×
4178
                # extract lines concerning symmetry params
4179
                if "symmetry_n_params" in line:
1✔
4180
                    self.n_symmetry_params = [int(x) for x in list(line.split())[1:]]
×
4181
                if "symmetry_params" in line:
1✔
4182
                    self.symmetry_params = list(line.split())[1:]
×
4183
                if "symmetry_lv" in line:
1✔
4184
                    symmetry_LVs_lines.append(line)
×
4185
                if "symmetry_frac" in line:
1✔
4186
                    symmetry_frac_lines.append(line)
×
4187
                if "homogeneous_field" in line:
1✔
4188
                    self.homogeneous_field = np.asarray(
×
4189
                        list(map(float, line.split()[1:4]))
4190
                    )
4191

4192
        # c Read all constraints/ moments and spins
4193
        for i, l in enumerate(atom_line_ind):
1✔
4194
            constraints = [False, False, False]
1✔
4195
            external_force = np.zeros(3)
1✔
4196
            calculate_friction = False
1✔
4197
            charge = 0.0
1✔
4198
            moment = 0.0
1✔
4199
            if i < len(atom_line_ind) - 1:
1✔
4200
                last_line = atom_line_ind[i + 1]
1✔
4201
            else:
4202
                last_line = len(text_lines)
1✔
4203
            for j in range(l, last_line):
1✔
4204
                line = text_lines[j]
1✔
4205
                if not line.startswith("#"):
1✔
4206
                    if "initial_moment" in line:
1✔
4207
                        moment = float(line.split()[1])
×
4208
                    elif "initial_charge" in line:
1✔
4209
                        charge = float(line.split()[1])
×
4210
                    elif "constrain_relaxation" in line:
1✔
4211
                        directions = line.split("constrain_relaxation")[1].lower()
1✔
4212
                        if ".true." in directions:
1✔
4213
                            constraints = [True, True, True]
1✔
4214
                        if "x" in directions:
1✔
4215
                            constraints[0] = True
×
4216
                        if "y" in directions:
1✔
4217
                            constraints[1] = True
×
4218
                        if "z" in directions:
1✔
4219
                            constraints[2] = True
×
4220
                    elif "external_force" in line:
1✔
4221
                        external_force[0] = float(line.split()[1])
×
4222
                        external_force[1] = float(line.split()[2])
×
4223
                        external_force[2] = float(line.split()[3])
×
4224
                    elif "calculate_friction" in line:
1✔
4225
                        if ".true." in line:
×
4226
                            calculate_friction = True
×
4227

4228
            self.constrain_relax.append(constraints)
1✔
4229
            self.external_force.append(external_force)
1✔
4230
            self.calculate_friction.append(calculate_friction)
1✔
4231
            self.initial_charge.append(charge)
1✔
4232
            self.initial_moment.append(moment)
1✔
4233

4234
        # read the atom species and coordinates
4235
        self.n_atoms = len(atom_lines)
1✔
4236
        self.coords = np.zeros([self.n_atoms, 3])
1✔
4237
        for i, l in enumerate(atom_lines):
1✔
4238
            tokens = l.split()
1✔
4239
            # self.species.append(tokens[-1])
4240
            self.species.append(tokens[4])
1✔
4241
            self.coords[i, :] = [float(x) for x in tokens[1:4]]
1✔
4242

4243
        # store symmetry_lv and symmetry_frac
4244
        if len(symmetry_LVs_lines) != 0:
1✔
4245
            self.symmetry_LVs = []
×
4246
            if len(symmetry_LVs_lines) != 3:
×
4247
                print(
×
4248
                    "Warning: Number of symmetry_LVs is: "
4249
                    + str(len(symmetry_LVs_lines))
4250
                )
4251
            for i, l in enumerate(symmetry_LVs_lines):
×
4252
                l = l[11:]
×
4253
                terms = [t.strip() for t in l.split(",")]
×
4254
                self.symmetry_LVs.append(terms)
×
4255
        if len(symmetry_frac_lines) != 0:
1✔
4256
            self.symmetry_frac_coords = []
×
4257
            for i, l in enumerate(symmetry_frac_lines):
×
4258
                l = l[13:]
×
4259
                terms = [t.strip() for t in l.split(",")]
×
4260
                # self.species.append(tokens[-1])
4261
                self.symmetry_frac_coords.append(terms)
×
4262

4263
        # read the hessian matrix if it is an own Hessian
4264
        if is_own_hessian:
1✔
4265
            # hessian has three coordinates for every atom
4266
            self.hessian = np.zeros([self.n_atoms * 3, self.n_atoms * 3])
×
4267
            for i, l in enumerate(hessian_lines):
×
4268
                tokens = l.split()
×
4269
                ind_1 = int(tokens[1])
×
4270
                ind_2 = int(tokens[2])
×
4271
                value_line = np.array([float(x) for x in tokens[3:12]])
×
4272
                self.hessian[
×
4273
                    (ind_1 - 1) * 3 : ind_1 * 3, (ind_2 - 1) * 3 : ind_2 * 3
4274
                ] = value_line.reshape((3, 3))
4275
            # self.hessian += np.tril(self.hessian.T, -1) # make symmetric hessian matrix
4276

4277
        if len(lattice_vector_lines) != 3 and len(lattice_vector_lines) != 0:
1✔
4278
            print(
×
4279
                "Warning: Number of lattice vectors is: "
4280
                + str(len(lattice_vector_lines))
4281
            )
4282
        for i, l in enumerate(lattice_vector_lines):
1✔
4283
            tokens = l.split()
1✔
4284
            self.lattice_vectors[i, :] = [float(x) for x in tokens[1:4]]
1✔
4285

4286
        # convert to cartesian coordinates
4287
        if is_fractional:
1✔
4288
            self.coords = utils.get_cartesian_coords(self.coords, self.lattice_vectors)
1✔
4289
            self.read_as_fractional_coords = True
1✔
4290

4291
        self.constrain_relax = np.array(self.constrain_relax)
1✔
4292
        self.external_force = np.array(self.external_force)
1✔
4293
        self.calculate_friction = np.array(self.calculate_friction)
1✔
4294

4295
        # update Part list and add all atoms that are not yet in the list
4296
        if len(self.geometry_parts) > 0:
1✔
4297
            already_indexed = list(itertools.chain.from_iterable(self.geometry_parts))
×
4298
            if len(already_indexed) < self.n_atoms:
×
4299
                additional_indices = [
×
4300
                    i for i in range(self.n_atoms) if i not in already_indexed
4301
                ]
4302
                self.geometry_parts.append(additional_indices)
×
4303
                self.geometry_part_descriptions.append("rest")
×
4304

4305
    def get_text(self, is_fractional=None):
1✔
4306
        """
4307
        If symmetry_params are to be used, the coordinates need to be fractional.
4308
        So, if symmetry_params are found, is_fractional is overridden to true.
4309
        """
4310
        if is_fractional is None:
1✔
4311
            if hasattr(self, "symmetry_params") and self.symmetry_params is not None:
1✔
4312
                is_fractional = True
×
4313
            else:
4314
                is_fractional = False
1✔
4315
        elif is_fractional is False:
×
4316
            if hasattr(self, "symmetry_params") and self.symmetry_params is not None:
×
4317
                warnings.warn(
×
4318
                    "The symmetry parameters of your geometry will be lost. "
4319
                    "To keep them set is_fractional to True"
4320
                )
4321

4322
        text = ""
1✔
4323
        for l in self.comment_lines:
1✔
4324
            if l.startswith("#"):
×
4325
                text += l + "\n"
×
4326
            else:
4327
                text += (
×
4328
                    "# " + l.lstrip() + "\n"
4329
                )  # str.lstrip() removes leading whitespace in comment line 'l'
4330

4331
        # If set, write 'center' dict ( see docstring of Geometry.__init__ ) to file
4332
        if hasattr(self, "center") and isinstance(self.center, dict):
1✔
4333
            center_string = "# CENTER " + str(self.center)
×
4334
            text += center_string + "\n"
×
4335

4336
        if hasattr(self, "geometry_parts") and (len(self.geometry_parts) > 0):
1✔
4337
            part_string = "# PARTS "
×
4338
            part_dict = {}
×
4339
            for part, name in zip(self.geometry_parts, self.geometry_part_descriptions):
×
4340
                if not name == "rest":
×
4341
                    if name not in part_dict:
×
4342
                        part_dict[name] = part
×
4343
                    else:
4344
                        warnings.warn(
×
4345
                            "Multiple equally named parts in file, renaming automatically!"
4346
                        )
4347
                        part_dict[name + "_1"] = part
×
4348
            part_string += str(part_dict) + "\n"
×
4349
            text += part_string
×
4350

4351
        if hasattr(self, "vacuum_level") and (self.vacuum_level is not None):
1✔
4352
            text += f"set_vacuum_level {self.vacuum_level: 15.10f}" + "\n"
×
4353

4354
        # Lattice vector relaxation constraints
4355
        constrain_vectors = np.zeros([3, 3], dtype=bool)
1✔
4356
        # if is_2D:
4357
        #    constrain_vectors[0, 2], constrain_vectors[1, 2], constrain_vectors[2] = True, True, 3*[True]
4358

4359
        # TODO: Some sort of custom lattice vector relaxation constraints parser
4360

4361
        if (self.lattice_vectors != 0).any():
1✔
4362
            for i in range(3):
×
4363
                line = "lattice_vector"
×
4364
                for j in range(3):
×
4365
                    line += f"     {self.lattice_vectors[i, j]:.8f}"
×
4366
                text += line + "\n"
×
4367
                cr = "\tconstrain_relaxation "
×
4368
                if constrain_vectors.any():
×
4369
                    if constrain_vectors[i].all():
×
4370
                        text += f"{cr}.true.\n"
×
4371
                    else:
4372
                        if constrain_vectors[i, 0]:
×
4373
                            text += f"{cr}x\n"
×
4374
                        if constrain_vectors[i, 1]:
×
4375
                            text += f"{cr}y\n"
×
4376
                        if constrain_vectors[i, 2]:
×
4377
                            text += f"{cr}z\n"
×
4378

4379
        # write down the homogeneous field if any is present
4380
        if self.homogeneous_field is not None:
1✔
4381
            text += "homogeneous_field {} {} {}\n".format(*self.homogeneous_field)
×
4382

4383
        if is_fractional:
1✔
4384
            coords = utils.get_fractional_coords(self.coords, self.lattice_vectors)
×
4385
            line_start = "atom_frac"
×
4386
        else:
4387
            coords = self.coords
1✔
4388
            line_start = "atom"
1✔
4389

4390
        for n in range(self.n_atoms):
1✔
4391
            if self.species[n] == "Em":  # do not save "Emptium" atoms
1✔
4392
                warnings.warn("Emptium atom was removed!!")
×
4393
                continue
×
4394
            line = line_start
1✔
4395
            for j in range(3):
1✔
4396
                line += f"     {coords[n, j]:.8f}"
1✔
4397
            line += " " + self.species[n]
1✔
4398
            text += line + "\n"
1✔
4399
            # backwards compatibilty for old-style constrain_relax
4400
            if isinstance(self.constrain_relax[0], bool):
1✔
4401
                if self.constrain_relax[n]:
×
4402
                    text += "constrain_relaxation .true.\n"
×
4403
            elif all(self.constrain_relax[n]):
1✔
4404
                text += "constrain_relaxation .true.\n"
1✔
4405
            else:
4406
                if self.constrain_relax[n][0]:
1✔
4407
                    text += "constrain_relaxation x\n"
×
4408
                if self.constrain_relax[n][1]:
1✔
4409
                    text += "constrain_relaxation y\n"
×
4410
                if self.constrain_relax[n][2]:
1✔
4411
                    text += "constrain_relaxation z\n"
×
4412
            if not len(self.initial_charge) == 0 and self.initial_charge[n] != 0.0:
1✔
4413
                text += f"initial_charge {self.initial_charge[n]: .6f}\n"
×
4414
            if not len(self.initial_moment) == 0 and self.initial_moment[n] != 0.0:
1✔
4415
                text += f"initial_moment {self.initial_moment[n]: .6f}\n"
×
4416
            if (
1✔
4417
                hasattr(self, "external_force")
4418
                and np.linalg.norm(self.external_force[n]) != 0.0
4419
            ):
4420
                text += f"external_force {self.external_force[n][0]: .6f} {self.external_force[n][1]: .6f} {self.external_force[n][2]: .6f}\n"
×
4421
            if hasattr(self, "calculate_friction") and self.calculate_friction[n]:
1✔
4422
                text += "calculate_friction .true.\n"
×
4423

4424
        if hasattr(self, "hessian") and self.hessian is not None:
1✔
4425
            text += "# own_hessian\n# This is a self calculated Hessian, not from a geometry optimization!\n"
×
4426
            for i in range(self.n_atoms):
×
4427
                for j in range(self.n_atoms):
×
4428
                    s = f"hessian_block  {i + 1} {j + 1}"
×
4429
                    H_block = self.hessian[3 * i : 3 * (i + 1), 3 * j : 3 * (j + 1)]
×
4430
                    # max_diff = np.max(np.abs(H_block-H_block.T))
4431
                    # print("Max diff in H: {:.3f}".format(max_diff))
4432
                    for h in H_block.flatten():
×
4433
                        s += f"  {h:.6f}"
×
4434
                    text += s + "\n"
×
4435

4436
        # write down symmetry_params and related data
4437
        if is_fractional:
1✔
4438
            if self.symmetry_params is not None:
×
4439
                l = "symmetry_params "
×
4440
                for p in self.symmetry_params:
×
4441
                    l += f"{p} "
×
4442
                l += "\n"
×
4443
                text += "\n" + l
×
4444
            if self.n_symmetry_params is not None:
×
4445
                l = "symmetry_n_params "
×
4446
                for n in self.n_symmetry_params:
×
4447
                    l += f"{n} "
×
4448
                text += l + "\n"
×
4449
                text += "\n"
×
4450
            if self.symmetry_LVs is not None:
×
4451
                for i in range(3):
×
4452
                    line = "symmetry_lv     {}  ,  {}  ,  {}".format(
×
4453
                        *self.symmetry_LVs[i]
4454
                    )
4455
                    text += line + "\n"
×
4456
                text += "\n"
×
4457
            if self.symmetry_frac_coords is not None:
×
4458
                for c in self.symmetry_frac_coords:
×
4459
                    line = "symmetry_frac     {}  ,  {}  ,  {}".format(*c)
×
4460
                    text += line + "\n"
×
4461
                text += "\n"
×
4462

4463
        # write down multipoles
4464
        for m in self.multipoles:
1✔
4465
            text += "multipole {}   {}   {}   {}   {}\n".format(*m)
×
4466
        return text
1✔
4467

4468

4469
class VaspGeometry(Geometry):
1✔
4470
    def parse(self, text):
1✔
4471
        """
4472
        Read the VASP structure definition in the typical POSCAR format
4473
        (also used by CONTCAR files, for example) from the file with the given filename
4474

4475
        Returns
4476
        -------
4477
        dic:
4478
            The dictionary holds the following keys:
4479
            systemname:
4480
            The name of the system as given in the first line of the POSCAR file.
4481
            vecs:
4482
            The unit cell vector as a 3x3 numpy.array. vecs[0,:] is the first unit
4483
            cell vector, vecs[:,0] are the x-coordinates of the three unit cell cevtors.
4484
            scaling:
4485
            The scaling factor of the POSCAR as given in the second line. However, this
4486
            information is not processed, it is up to the user to use this information
4487
            to scale whatever needs to be scaled.
4488
            coordinates:
4489
            The coordinates of all the atoms. Q[k,:] are the coordinates of the k-th atom
4490
            (the index starts with 0, as usual). Q[:,0] are the x-coordinates of all the atoms. These coordinates are always given in Cartesian coordinates.
4491
            elementtypes:
4492
            A list of as many entries as there are atoms. Gives the type specification
4493
            for every atom (typically the atom name). elementtypes[k] is the species of
4494
            the k-th atom.
4495
            typenames:
4496
            The names of all the species. This list contains as many elements as there are species.
4497
            numberofelements:
4498
            Gives the number of atoms per species. This list contains as many elements as there are species.
4499
            elementid:
4500
            Gives the index (from 0 to the number of atoms-1) of the first atom of a
4501
            certain species. This list contains as many elements as there are species.
4502
            cartesian:
4503
            A logical value whether the coordinates were given in Cartesian form (True)
4504
            or as direct coordinates (False).
4505
            originalcoordinates:
4506
            The original coordinates as read from the POSCAR file. It has the same
4507
            format as coordinates. For Cartesian coordinates (cartesian == True) this
4508
            is identical to coordinates, for direct coordinates (cartesian == False)
4509
            this contains the direct coordinates.
4510
            selective:
4511
            True or False: whether selective dynamics is on.
4512
            selectivevals:
4513
            Consists of as many rows as there are atoms, three colums: True if
4514
            selective dynamics is on for this coordinate for the atom, else False.
4515
            Only if selective is True.
4516
        """
4517
        lino = 0
×
4518
        vecs = []
×
4519
        scaling = 1.0
×
4520
        typenames = []
×
4521
        nelements = []
×
4522
        cartesian = False
×
4523
        selective = False
×
4524
        selectivevals = []
×
4525
        P = []
×
4526
        fi = text.split("\n")
×
4527

4528
        for line in fi:
×
4529
            lino += 1
×
4530
            line = line.strip()
×
4531

4532
            if lino == 1:
×
4533
                self.add_top_comment(line)
×
4534
            if lino == 2:
×
4535
                scaling = float(line)
×
4536
                # RB: now the scaling should be taken account for below when the lattice vectors and coordinates
4537
                # if scaling != 1.0:
4538
                #    print("WARNING (readin_struct): universal scaling factor is not one. This is ignored.")
4539

4540
            if lino in (3, 4, 5):
×
4541
                vecs.append(list(map(float, line.split())))
×
4542
            if lino == 6:
×
4543
                if line[0] in [
×
4544
                    "0",
4545
                    "1",
4546
                    "2",
4547
                    "3",
4548
                    "4",
4549
                    "5",
4550
                    "6",
4551
                    "7",
4552
                    "8",
4553
                    "9",
4554
                ]:
4555
                    lino += 1
×
4556
                else:
4557
                    typenames = line.split()
×
4558
            if lino == 7:
×
4559
                splitline = line.split()
×
4560
                nelements = list(map(int, splitline))
×
4561
                elementid = np.cumsum(np.array(nelements))
×
4562
                self.n_atoms = elementid[-1]
×
4563
            if lino == 8:
×
4564
                if line[0] in ("S", "s"):
×
4565
                    selective = True
×
4566
                else:
4567
                    lino += 1
×
4568
            if lino == 9:
×
4569
                if line[0] in ("K", "k", "C", "c"):  # cartesian coordinates
×
4570
                    cartesian = True
×
4571
            if lino >= 10:
×
4572
                if lino >= 10 + self.n_atoms:
×
4573
                    break
×
4574
                P.append(list(map(float, line.split()[0:3])))
×
4575

4576
                if selective:
×
4577
                    # TODO: experimental...
4578
                    constraints = list(
×
4579
                        map(lambda x: x in ("F", "f"), line.split()[3:6])
4580
                    )
4581
                    if len(constraints) != 3:
×
4582
                        self.constrain_relax.append([False, False, False])
×
4583
                    else:
4584
                        self.constrain_relax.append(constraints)
×
4585
                    selectivevals.append(constraints)
×
4586
                else:
4587
                    self.constrain_relax.append([False, False, False])
×
4588
                    # TODO: write true value
4589
                    self.initial_charge.append(0)
×
4590
                    self.initial_moment.append(0)
×
4591

4592
                self.external_force = np.append(
×
4593
                    self.external_force, np.atleast_2d(np.zeros(3)), axis=0
4594
                )
4595
                self.calculate_friction = np.append(
×
4596
                    self.calculate_friction, np.array([False])
4597
                )
4598

4599
        vecs = np.array(vecs)
×
4600
        P = np.array(P)
×
4601
        if not cartesian:
×
4602
            Q = np.dot(P, vecs)
×
4603
        else:
4604
            Q = P
×
4605
        if len(typenames) > 0:
×
4606
            for k in range(Q.shape[0]):
×
4607
                self.species.append(typenames[np.min(np.where(elementid > k)[0])])
×
4608

4609
        self.lattice_vectors = vecs
×
4610
        self.coords = Q
×
4611
        self.constrain_relax = np.array(self.constrain_relax)
×
4612

4613
        # RB: include the scaling. should work for both direct and cartesian settings
4614
        self.lattice_vectors = vecs * scaling
×
4615
        self.coords = Q * scaling
×
4616

4617
    def get_text(self, comment="POSCAR file written by Geometry.py"):
1✔
4618
        comment = comment.replace("\n", " ")
×
4619
        text = comment + "\n"
×
4620
        text += "1\n"
×
4621
        if (self.lattice_vectors != 0).any():
×
4622
            for i in range(3):
×
4623
                line = ""
×
4624
                for j in range(3):
×
4625
                    line += f"     {self.lattice_vectors[i, j]:-4.8f}"
×
4626
                text += line.strip() + "\n"
×
4627

4628
        all_species = sorted(
×
4629
            list(set(self.species))
4630
        )  # get unique species and sort alphabetically
4631
        text += " ".join(all_species) + "\n"
×
4632
        species_coords = {}
×
4633
        n_of_species = {}
×
4634
        # R.B. relax constraints
4635
        relax_constraints = {}
×
4636
        ## R.B. relax constraints end
4637

4638
        for species in all_species:
×
4639
            is_right_species = np.array(
×
4640
                [s == species for s in self.species], dtype=bool
4641
            )
4642
            curr_species_coords = self.coords[is_right_species, :]
×
4643
            species_coords[species] = curr_species_coords
×
4644
            n_of_species[species] = curr_species_coords.shape[0]
×
4645

4646
            # R.B. relax constraints
4647
            curr_species_constrain_relax = self.constrain_relax[is_right_species, :]
×
4648
            relax_constraints[species] = curr_species_constrain_relax
×
4649
            ## R.B. relax constraints end
4650

4651
        # add number of atoms per species
4652
        text += " ".join([str(n_of_species[s]) for s in all_species]) + "\n"
×
4653

4654
        # R.B. Write out selective dynamics so that the relaxation constraints are read
4655
        text += "Selective dynamics" + "\n"
×
4656

4657
        text += "Cartesian" + "\n"
×
4658

4659
        for species in all_species:
×
4660
            curr_coords = species_coords[species]
×
4661
            n_atoms = n_of_species[species]
×
4662

4663
            ## R.B. relax constraints
4664
            curr_relax_constr = relax_constraints[species]
×
4665
            ## R.B. relax constraints end
4666

4667
            for n in range(n_atoms):
×
4668
                line = ""
×
4669
                for j in range(3):
×
4670
                    if j > 0:
×
4671
                        line += "    "
×
4672
                    line += f"{curr_coords[n, j]: 2.8f}"
×
4673

4674
                ## R.B. relax constraints
4675
                for j in range(3):
×
4676
                    if curr_relax_constr[n, j] is True:
×
4677
                        line += "  " + "F"
×
4678
                    elif curr_relax_constr[n, j] is False:
×
4679
                        line += "  " + "T"
×
4680
                ## R.B. relax constraints end
4681

4682
                text += line + "\n"
×
4683

4684
        return text
×
4685

4686

4687
class XYZGeometry(Geometry):
1✔
4688
    def parse(self, text):
1✔
4689
        """
4690
        Reads a .xyz file. Designed to work with .xyz files produced by Avogadro
4691

4692
        """
4693
        # to use add_atoms we need to initialize coords the same as for Geometry
4694
        self.n_atoms = 0
×
4695
        self.coords = np.zeros([self.n_atoms, 3])
×
4696

4697
        read_natoms = None
×
4698
        count_natoms = 0
×
4699
        coords = []
×
4700
        forces = []
×
4701
        species = []
×
4702
        fi = text.split("\n")
×
4703

4704
        # parse will assume first few lines are comments
4705
        started_parsing_atoms = False
×
4706

4707
        for ind, line in enumerate(fi):
×
4708
            if ind == 0:
×
4709
                if len(line.split()) == 1:
×
4710
                    read_natoms = int(line.split()[0])
×
4711
                    continue
×
4712

4713
            # look for lattice vectors
4714
            if "Lattice" in line:
×
4715
                # split_line = line.split('Lattice')[1]
4716
                split_line = line.split('"')[1]
×
4717

4718
                lattice_parameters = re.findall(r"\d+\.\d+", split_line)
×
4719

4720
                if len(lattice_parameters) == 9:
×
4721
                    lattice_parameters = np.array(lattice_parameters, dtype=np.float64)
×
4722
                    self.lattice_vectors = np.reshape(lattice_parameters, (3, 3))
×
4723

4724
            if "energy" in line:
×
4725
                split_line = line.split("energy")[1]
×
4726

4727
                energy = re.findall(r"-?[\d.]+(?:e-?\d+)?", split_line)
×
4728

4729
                if len(energy) > 0:
×
4730
                    self.energy = np.float64(energy[0])
×
4731

4732
            # words = re.findall('[a-zA-Z]+', line)
4733
            # numbers = re.findall('-?[\d.]+(?:e-?\d+)?', line)
4734

4735
            # n_words = len( words )
4736
            # n_floats = len( numbers )
4737

4738
            split_line = line.split()
×
4739

4740
            n_words = 0
×
4741
            n_floats = 0
×
4742

4743
            for l in split_line:
×
4744
                n_words_new = len(re.findall("[a-zA-Z]+", l))
×
4745
                n_floats_new = len(re.findall(r"-?[\d.]+(?:e-?\d+)?", l))
×
4746

4747
                if n_words_new == 1 and n_floats_new == 1:
×
4748
                    n_floats += 1
×
4749
                else:
4750
                    n_words += n_words_new
×
4751
                    n_floats += n_floats_new
×
4752

4753
            # first few lines may be comments or properties
4754
            if not started_parsing_atoms:
×
4755
                if n_words == 1 and (n_floats == 3 or n_floats == 6):
×
4756
                    continue
×
4757
                started_parsing_atoms = True
×
4758

4759
            else:
4760
                if split_line == []:
×
4761
                    break
×
4762
                assert n_words == 1 and (n_floats == 3 or n_floats == 6), (
×
4763
                    "Bad atoms specification: "
4764
                    + str(split_line)
4765
                    + f"{n_words} {n_floats}"
4766
                )
4767

4768
                # write atoms
4769
                species.append(str(split_line[0]))
×
4770
                coords.append(np.array(split_line[1:4], dtype=np.float64))
×
4771

4772
                if n_floats == 6:
×
4773
                    forces.append(np.array(split_line[4:], dtype=np.float64))
×
4774

4775
                count_natoms += 1
×
4776

4777
        if not started_parsing_atoms:
×
4778
            raise RuntimeError("Not atoms found in xyz file!")
×
4779

4780
        if read_natoms is not None:
×
4781
            assert read_natoms == count_natoms, "Not all atoms found!"
×
4782

4783
        coords = np.asarray(coords)
×
4784
        self.add_atoms(coords, species)
×
4785

4786
        if forces:
×
4787
            forces = np.asarray(forces)
×
4788
            self.forces = forces
×
4789

4790
    def get_text(self, comment="XYZ file written by Geometry.py"):
1✔
4791
        text = str(self.n_atoms) + "\n"
×
4792
        comment = comment.replace("\n", " ")
×
4793
        text += comment + "\n"
×
4794
        for index in range(self.n_atoms):
×
4795
            element = self.species[index]
×
4796
            x, y, z = self.coords[index]
×
4797
            text += f"{element}    {x:-4.8f}    {y:-4.8f}    {z:-4.8f}" + "\n"
×
4798
        return text
×
4799

4800

4801
class XSFGeometry(Geometry):
1✔
4802
    def get_text(self):
1✔
4803
        text = ""
×
4804
        text += "CRYSTAL\n"
×
4805
        text += "PRIMVEC\n"
×
4806
        for i in range(3):
×
4807
            line = ""
×
4808
            for j in range(3):
×
4809
                line += f"    {self.lattice_vectors[i, j]:.8f}"
×
4810
            text += line + "\n"
×
4811
        text += "PRIMCOORD\n"
×
4812
        # the 1 is mysterious but is needed for primcoord according to XSF docu
4813
        text += str(self.n_atoms) + " 1\n"
×
4814
        for i in range(self.n_atoms):
×
4815
            if self.constrain_relax[i]:
×
4816
                raise NotImplementedError(
×
4817
                    "Constrained relaxation not supported for XSF output file"
4818
                )
4819
            line = str(self.periodic_table.get_atomic_number(self.species[i]))
×
4820
            for j in range(3):
×
4821
                line += f"    {self.coords[i, j]:.8f}"
×
4822
            text += line + "\n"
×
4823
        return text
×
4824

4825

4826
###############################################################################
4827
#                            Auxiliary Functions                              #
4828
###############################################################################
4829
def get_file_format_from_ending(filename):
1✔
4830
    if filename.endswith(".in") or filename.endswith(".next_step"):
1✔
4831
        return "aims"
1✔
4832
    if filename.endswith(".xsf"):
×
4833
        return "xsf"
×
4834
    if filename.endswith(".molden"):
×
4835
        return "molden"
×
4836
    if filename.endswith("POSCAR") or filename.endswith("CONTCAR"):
×
4837
        return "vasp"
×
4838
    if filename.endswith(".xyz"):
×
4839
        return "xyz"
×
4840
    return None
×
STATUS · Troubleshooting · Open an Issue · Sales · Support · CAREERS · ENTERPRISE · START FREE · SCHEDULE DEMO
ANNOUNCEMENTS · TWITTER · TOS & SLA · Supported CI Services · What's a CI service? · Automated Testing

© 2026 Coveralls, Inc