Visualizing Geometry

%matplotlib notebook

from os import listdir

import pyvista as pv
from pyvista import themes

from qlbm.infra import AmplitudeResult
from qlbm.lattice import MSLattice
from qlbm.tools.utils import create_directory_and_parents

pv.set_plot_theme(themes.ParaViewTheme())

from qlbm.lattice.lattices.ab_lattice import ABLattice


lattice_2d = MSLattice(
    {
        "lattice": {
            "dim": {"x": 32, "y": 32},
            "velocities": {"x": 4, "y": 4},
        },
        "geometry": [
            {"shape": "cuboid", "x": [18, 19], "y": [7, 14], "boundary": "specular"},
            {"shape": "cuboid", "x": [18, 19], "y": [19, 26], "boundary": "specular"},
            {"shape": "cuboid", "x": [26, 27], "y": [19, 26], "boundary": "specular"},
            {"shape": "cuboid", "x": [26, 27], "y": [7, 14], "boundary": "specular"},
            {"shape": "cuboid", "x": [20, 25], "y": [4, 5], "boundary": "specular"},
            {"shape": "cuboid", "x": [20, 25], "y": [16, 17], "boundary": "specular"},
            {"shape": "cuboid", "x": [20, 25], "y": [28, 29], "boundary": "specular"},
        ],
    }
)
lattice_3d = MSLattice(
    {
        "lattice": {
            "dim": {"x": 16, "y": 128, "z": 16},
            "velocities": {"x": 4, "y": 4, "z": 4},
        },
        "geometry": [
            {
                "shape": "cuboid",
                "x": [10, 13],
                "y": [6, 88],
                "z": [10, 13],
                "boundary": "specular",
            }
        ],
    }
)

ab_lattice_2d = ABLattice(
    {
        "lattice": {
            "dim": {"x": 64, "y": 256},
            "velocities": "d2q9",
        },
        "geometry": [
            {
                "shape": "ymonomial",
                "comparator": "<=",
                "exponent": 2,
                "boundary": "bounceback",
            },
        ],
    }
)

root_directory_2d = "qlbm-output/visualization-components-2d"
root_directory_2d_monomial = "qlbm-output/visualization-components-2d-monomial"
root_directory_3d = "qlbm-output/visualization-components-3d"
create_directory_and_parents(root_directory_2d)
create_directory_and_parents(root_directory_2d_monomial)
create_directory_and_parents(root_directory_3d)

# Will output seven 2D stl files under `qlbm-output/visualization-components-2d/paraview`
AmplitudeResult(lattice_2d, root_directory_2d).visualize_geometry()

# Will output seven 2D stl files under `qlbm-output/visualization-components-2d-monomial/paraview`
AmplitudeResult(ab_lattice_2d, root_directory_2d_monomial).visualize_geometry()

# Will output one 3D stl files under `qlbm-output/visualization-components-3d/paraview`
AmplitudeResult(lattice_3d, root_directory_3d).visualize_geometry()

pv.read(
    [
        f"{root_directory_2d}/paraview/{fname}"
        for fname in listdir(f"{root_directory_2d}/paraview")
    ]
).plot(cpos="xy", show_scalar_bar=True, jupyter_backend="static")

2026-05-12 23:00:32.523 (   2.693s) [    7F4390F64B80]vtkXOpenGLRenderWindow.:1460  WARN| bad X server connection. DISPLAY=

pl = pv.Plotter()
mesh = pv.read(
    [
        f"{root_directory_2d_monomial}/paraview/{fname}"
        for fname in listdir(f"{root_directory_2d_monomial}/paraview")
    ]
).plot(show_scalar_bar=True, jupyter_backend="interactive")

# pl.add_mesh(mesh)
# pl.camera_position = "xy"
# pl.save_graphic("ymonomial.pdf")

pv.read(
    [
        f"{root_directory_3d}/paraview/{fname}"
        for fname in listdir(f"{root_directory_3d}/paraview")
    ]
).plot(show_scalar_bar=True, jupyter_backend="static")