File: test_snapshot_checkpoint.py

package info (click to toggle)
io4dolfinx 1.1.2-4
  • links: PTS, VCS
  • area: main
  • in suites: sid
  • size: 832 kB
  • sloc: python: 8,419; sh: 34; makefile: 3
file content (122 lines) | stat: -rw-r--r-- 4,255 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
113
114
115
116
117
118
119
120
121
122
from pathlib import Path

from mpi4py import MPI

import basix.ufl
import dolfinx
import numpy as np
import pytest

from io4dolfinx import FileMode, snapshot_checkpoint


def suffix(backend: str) -> str:
    if backend == "adios2":
        return ".bp"
    elif backend == "h5py":
        return ".h5"
    else:
        raise NotImplementedError(f"Unsupported backend {backend}")


triangle = dolfinx.mesh.CellType.triangle
quad = dolfinx.mesh.CellType.quadrilateral
tetra = dolfinx.mesh.CellType.tetrahedron
hex = dolfinx.mesh.CellType.hexahedron


@pytest.mark.parametrize(
    "cell_type, family", [(triangle, "N1curl"), (triangle, "RT"), (quad, "RTCF")]
)
@pytest.mark.parametrize("degree", [1, 4])
def test_read_write_2D(family, degree, cell_type, tmp_path, backend):
    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 10, 10, cell_type=cell_type)
    el = basix.ufl.element(family, mesh.basix_cell(), degree)

    def f(x):
        return (np.full(x.shape[1], np.pi) + x[0], x[1])

    V = dolfinx.fem.functionspace(mesh, el)
    u = dolfinx.fem.Function(V)
    u.interpolate(f)

    fname = MPI.COMM_WORLD.bcast(tmp_path, root=0)
    file = fname / Path("snapshot_2D_vs").with_suffix(suffix(backend))
    snapshot_checkpoint(u, file, FileMode.write, backend=backend)

    v = dolfinx.fem.Function(V)
    snapshot_checkpoint(v, file, FileMode.read, backend=backend)
    assert np.allclose(u.x.array, v.x.array)


@pytest.mark.parametrize("cell_type, family", [(tetra, "N1curl"), (tetra, "RT"), (hex, "NCF")])
@pytest.mark.parametrize("degree", [1, 4])
def test_read_write_3D(family, degree, cell_type, tmp_path, backend):
    mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 3, 3, 3, cell_type=cell_type)
    el = basix.ufl.element(family, mesh.basix_cell(), degree)

    def f(x):
        return (np.full(x.shape[1], np.pi) + x[0], x[1], x[1] * x[2])

    V = dolfinx.fem.functionspace(mesh, el)
    u = dolfinx.fem.Function(V)
    u.interpolate(f)

    fname = MPI.COMM_WORLD.bcast(tmp_path, root=0)
    file = fname / Path("snapshot_3D_vs").with_suffix(suffix(backend))
    snapshot_checkpoint(u, file, FileMode.write, backend=backend)

    v = dolfinx.fem.Function(V)
    snapshot_checkpoint(v, file, FileMode.read, backend=backend)
    assert np.allclose(u.x.array, v.x.array)


@pytest.mark.parametrize(
    "cell_type", [dolfinx.mesh.CellType.triangle, dolfinx.mesh.CellType.quadrilateral]
)
@pytest.mark.parametrize("family", ["Lagrange", "DG"])
@pytest.mark.parametrize("degree", [1, 4])
def test_read_write_P_2D(family, degree, cell_type, tmp_path, backend):
    mesh = dolfinx.mesh.create_unit_square(MPI.COMM_WORLD, 5, 5, cell_type=cell_type)
    el = basix.ufl.element(family, mesh.basix_cell(), degree, shape=(mesh.geometry.dim,))

    def f(x):
        return (np.full(x.shape[1], np.pi) + x[0], x[1])

    V = dolfinx.fem.functionspace(mesh, el)
    u = dolfinx.fem.Function(V)
    u.interpolate(f)

    fname = MPI.COMM_WORLD.bcast(tmp_path, root=0)
    file = fname / Path("snapshot_2D_p").with_suffix(suffix(backend))
    snapshot_checkpoint(u, file, FileMode.write, backend=backend)

    v = dolfinx.fem.Function(V)
    snapshot_checkpoint(v, file, FileMode.read, backend=backend)
    assert np.allclose(u.x.array, v.x.array)


@pytest.mark.parametrize(
    "cell_type", [dolfinx.mesh.CellType.tetrahedron, dolfinx.mesh.CellType.hexahedron]
)
@pytest.mark.parametrize("family", ["Lagrange", "DG"])
@pytest.mark.parametrize("degree", [1, 4])
def test_read_write_P_3D(family, degree, cell_type, tmp_path, backend):
    mesh = dolfinx.mesh.create_unit_cube(MPI.COMM_WORLD, 5, 5, 5, cell_type=cell_type)
    el = basix.ufl.element(family, mesh.basix_cell(), degree, shape=(mesh.geometry.dim,))

    def f(x):
        return (np.full(x.shape[1], np.pi) + x[0], x[1] + 2 * x[0], np.cos(x[2]))

    V = dolfinx.fem.functionspace(mesh, el)
    u = dolfinx.fem.Function(V)
    u.interpolate(f)

    fname = MPI.COMM_WORLD.bcast(tmp_path, root=0)
    file = fname / Path("snapshot_3D_p").with_suffix(suffix(backend))

    snapshot_checkpoint(u, file, FileMode.write, backend=backend)

    v = dolfinx.fem.Function(V)
    snapshot_checkpoint(v, file, FileMode.read, backend=backend)
    assert np.allclose(u.x.array, v.x.array)