File: test_precon_assembly.py

package info (click to toggle)
python-ase 3.22.1-3
  • links: PTS, VCS
  • area: main
  • in suites: bookworm
  • size: 14,344 kB
  • sloc: python: 126,379; xml: 946; makefile: 111; javascript: 47
file content (112 lines) | stat: -rw-r--r-- 3,266 bytes parent folder | download
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
import numpy as np
import pytest

from ase.build import bulk
from ase.constraints import FixAtoms, UnitCellFilter
from ase.calculators.emt import EMT
from ase.optimize.precon import make_precon, Precon
from ase.neighborlist import neighbor_list
from ase.utils.ff import Bond


@pytest.fixture
def ref_atoms():
    atoms = bulk('Al', a=3.994) * 2
    atoms.calc = EMT()
    i, j, r = neighbor_list('ijd', atoms, cutoff=3.0)
    bonds = [Bond(I, J, k=1.0, b0=R) for (I, J, R) in zip(i, j, r)]
    return atoms, bonds


@pytest.fixture
def atoms(ref_atoms):
    atoms, bonds = ref_atoms
    atoms.rattle(stdev=0.1, seed=7)
    return atoms, bonds


@pytest.fixture
def var_cell(atoms):
    atoms, bonds = atoms
    return UnitCellFilter(atoms), bonds


@pytest.fixture
def fixed_atoms(atoms):
    atoms, bonds = atoms
    atoms.set_constraint(FixAtoms(indices=[0]))
    return atoms, bonds


def check_assembly(precon, system):
    atoms, bonds = system
    kwargs = {}
    if precon == 'FF' or precon == 'Exp_FF':
        kwargs['bonds'] = bonds
    precon = make_precon(precon, atoms, **kwargs)
    assert isinstance(precon, Precon)
    # check its a symmetric positive definite matrix of expected size
    N = 3 * len(atoms)
    P = precon.asarray()
    assert P.shape == (N, N)
    assert np.abs(P - P.T).max() < 1e-6
    assert np.all(np.linalg.eigvalsh(P)) > 0


precons = [None, 'C1', 'Exp', 'Pfrommer', 'FF', 'Exp_FF']


@pytest.mark.parametrize('precon', precons)
def test_assembly_ref_atoms(precon, ref_atoms):
    check_assembly(precon, ref_atoms)


@pytest.mark.parametrize('precon', precons)
def test_assembly_atoms(precon, atoms):
    check_assembly(precon, atoms)

    
@pytest.mark.parametrize('precon', precons)
def test_assembly_var_cell(precon, var_cell):
    check_assembly(precon, var_cell)
    

@pytest.mark.parametrize('precon', precons)
def test_assembly_fixed_atoms(precon, fixed_atoms):
    check_assembly(precon, fixed_atoms)

    
def check_apply(precon, system):
    atoms, bonds = system
    kwargs = {}
    if precon == 'FF' or precon == 'Exp_FF':
        kwargs['bonds'] = bonds
    precon = make_precon(precon, atoms, **kwargs)
    forces = atoms.get_forces().reshape(-1)
    precon_forces, residual = precon.apply(forces, atoms)
    residual_P = np.linalg.norm(precon_forces, np.inf)
    print(f'|F| = {residual:.3f} '
          f'|F|_P = {np.linalg.norm(precon_forces, np.inf):.3f}')
    
    # force norm should not get much bigger when we precondition:
    # in this case all norms get smaller, but this will not be true in general
    assert residual_P <= residual
    
    # forces on any fixed atoms be zero, and remain zero after applying precon
    fixed_atoms = []
    for constraint in atoms.constraints:
        if isinstance(constraint, FixAtoms):
            fixed_atoms.extend(list(constraint.index))
    if len(fixed_atoms) != 0:
        assert np.linalg.norm(forces[fixed_atoms], np.inf) < 1e-8
        assert np.linalg.norm(precon_forces[fixed_atoms], np.inf) < 1e-8

    
@pytest.mark.parametrize('precon', precons)
def test_apply_atoms(precon, atoms):
    check_apply(precon, atoms)


@pytest.mark.parametrize('precon', precons)
def test_apply_fixed_atoms(precon, fixed_atoms):
    check_apply(precon, fixed_atoms)