kernelpack-python/examples/moving_surface_rbc_capstone.py at main · VarShankar/kernelpack-python · GitHub

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
"""Run source-free surface transport on the public IBAMR RBC trajectory."""

import argparse
from pathlib import Path

import numpy as np

from kernelpack import manifold, solvers


def initial_tracer(material_sites):
    center_one = np.array([0.25, 0.80, 0.54])
    center_two = np.array([-0.55, -0.20, 0.81])
    center_one /= np.linalg.norm(center_one)
    center_two /= np.linalg.norm(center_two)
    distance_one = np.arccos(np.clip(material_sites @ center_one, -1.0, 1.0))
    distance_two = np.arccos(np.clip(material_sites @ center_two, -1.0, 1.0))
    return 1.0 + np.exp(-(distance_one / 0.28) ** 2) + 0.65 * np.exp(-(distance_two / 0.22) ** 2)


def main(trajectory_directory=None, *, xi=6, diffusivity=2.0e-3, steps=None):
    root = Path(__file__).resolve().parents[1]
    folder = Path(trajectory_directory or root / "data" / "ibamr_rbc_3d" / "trajectory")
    trajectory = manifold.IBAMRSurfaceTrajectory(folder, xi=xi)
    frame_count = trajectory.times.size if steps is None else min(trajectory.times.size, int(steps) + 1)
    if frame_count < 2:
        raise ValueError("the RBC capstone requires at least two trajectory frames")
    geometry, _ = trajectory.geometry(0)
    neighbors = manifold.build_surface_stencil_graph(
        trajectory.material_sites,
        manifold.surface_stencil_size(manifold.surface_polynomial_degree(xi)),
    )
    operators, cache = manifold.assemble_tangent_plane_operators(
        geometry.points, geometry.normals, neighbors, xi=xi
    )
    concentration = initial_tracer(trajectory.material_sites)
    history = solvers.initialize_moving_surface_history(geometry.points, concentration)
    target_mass = float(geometry.quadrature_weights @ concentration)
    h = np.sqrt(np.sum(geometry.quadrature_weights) / concentration.size)
    calibration = manifold.calibrate_surface_hyperviscosity(
        operators, geometry.points, geometry.normals, h, target_order=xi
    )
    reference_dt = float(trajectory.times[1] - trajectory.times[0])

    for frame_index in range(1, frame_count):
        dt = float(trajectory.times[frame_index] - trajectory.times[frame_index - 1])
        if not np.isclose(dt, reference_dt):
            raise ValueError("the RBC BDF driver requires uniformly spaced trajectory frames")
        geometry, _ = trajectory.geometry(frame_index)
        operators, cache, _ = manifold.update_tangent_plane_operators(
            geometry.points,
            geometry.normals,
            neighbors,
            cache,
            xi=xi,
            tolerance=1.0e-6,
        )
        h = np.sqrt(np.sum(geometry.quadrature_weights) / concentration.size)
        calibration, drift = manifold.predict_surface_hyperviscosity(
            operators, geometry.points, h, calibration
        )
        if drift > 0.15:
            calibration = manifold.calibrate_surface_hyperviscosity(
                operators,
                geometry.points,
                geometry.normals,
                h,
                target_order=xi,
                power=calibration.power,
            )
        history, info = solvers.moving_surface_adr_step(
            history,
            geometry.points,
            operators,
            np.zeros(concentration.size),
            geometry.quadrature_weights,
            target_mass,
            reference_dt,
            diffusivity,
            calibration.gamma,
            min(0.1 * h**xi, 1.0e-7),
            order=min(frame_index, 3),
            hyperviscosity_power=calibration.power,
        )

    concentration = history.concentration[0]
    mass_error = abs(float(geometry.quadrature_weights @ concentration) - target_mass) / abs(target_mass)
    output = root / "artifacts" / "moving_surface_rbc_capstone.npz"
    output.parent.mkdir(parents=True, exist_ok=True)
    np.savez(output, points=geometry.points, concentration=concentration, time=trajectory.times[frame_count - 1])
    print(f"frames={frame_count}, nodes={concentration.size}")
    print(f"relative mass error: {mass_error:.6e}")
    print(f"linear relative residual: {info.relative_residual:.6e}")
    print(output)


if __name__ == "__main__":
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--trajectory-directory")
    parser.add_argument("--xi", type=int, default=6)
    parser.add_argument("--diffusivity", type=float, default=2.0e-3)
    parser.add_argument("--steps", type=int)
    arguments = parser.parse_args()
    main(
        arguments.trajectory_directory,
        xi=arguments.xi,
        diffusivity=arguments.diffusivity,
        steps=arguments.steps,
    )