xref: /libCEED/examples/python/ex_common.py (revision f8fffaefe2b0bf336e4f45347da405b61e8c02d8)
1# Copyright (c) 2017-2025, Lawrence Livermore National Security, LLC and other CEED contributors
2# All Rights Reserved. See the top-level LICENSE and NOTICE files for details.
3#
4# SPDX-License-Identifier: BSD-2-Clause
5#
6# This file is part of CEED:  http://github.com/ceed
7
8import sys
9import os
10from sysconfig import get_config_var
11import argparse
12import math
13import numpy as np
14import libceed
15import ctypes
16
17
18def parse_arguments():
19    """Parse command line arguments for surface area computation
20
21    Returns:
22        Namespace: Parsed arguments with fields:
23            ceed: CEED resource specifier
24            dim: Problem dimension (1-3)
25            mesh_degree: Mesh polynomial degree
26            solution_degree: Solution polynomial degree
27            num_qpts: Number of quadrature points
28            problem_size: Approximate problem size
29            test: Test mode flag
30            quiet: Suppress output flag
31            gallery: Use gallery QFunctions flag
32    """
33    parser = argparse.ArgumentParser(description="libCEED surface area example")
34    parser.add_argument("-c", "--ceed", default="/cpu/self",
35                        help="libCEED resource specifier (default: /cpu/self)")
36    parser.add_argument("-d", "--dim", type=int, default=3,
37                        help="Problem dimension (1-3) (default: 3)")
38    parser.add_argument("-m", "--mesh-degree", type=int, default=4,
39                        help="Mesh polynomial degree (default: 4)")
40    parser.add_argument("-p", "--solution-degree", type=int, default=4,
41                        help="Solution polynomial degree (default: 4)")
42    parser.add_argument("-q", "--quadrature-points", type=int, default=6,
43                        help="Number of quadrature points (default: 6)")
44    parser.add_argument("-s", "--problem-size", type=int, default=-1,
45                        help="Approximate problem size (default: ~256k)")
46    parser.add_argument("-t", "--test", action="store_true",
47                        help="Test mode (reduced problem size)")
48    parser.add_argument("--quiet", action="store_true",
49                        help="Suppress output")
50    parser.add_argument("-g", "--gallery", action="store_true",
51                        help="Use gallery QFunctions")
52
53    args = parser.parse_args()
54    if args.dim not in [1, 2, 3]:
55        parser.error("Dimension must be 1, 2, or 3")
56    return args
57
58
59def get_cartesian_mesh_size(dim, degree, prob_size):
60    """Determine Cartesian mesh size for given problem size
61
62    Args:
63        dim: Spatial dimension (1-3)
64        degree: Polynomial degree
65        prob_size: Target problem size
66
67    Returns:
68        list: Number of elements in each dimension
69    """
70    # Calculate number of elements needed
71    num_elem = prob_size // (degree ** dim)
72
73    # Find smallest power of 2 >= num_elem
74    s = 0
75    while num_elem > 1:
76        num_elem = num_elem / 2
77        s += 1
78
79    # Distribute across dimensions
80    r = s % dim
81    num_xyz = []
82    for d in range(dim):
83        sd = s // dim
84        if r > 0:
85            sd += 1
86            r -= 1
87        num_xyz.append(1 << sd)
88    return num_xyz
89
90
91def build_cartesian_restriction(ceed, dim, num_xyz, degree, num_comp, num_q_comp, num_qpts, create_qdata=False):
92    """Build element restriction for Cartesian grid
93
94    Args:
95        ceed: libCEED context
96        dim: Spatial dimension
97        num_xyz: Elements per dimension
98        degree: Polynomial degree
99        num_comp: Number of components
100        num_q_comp: Number of quadrature data components
101        num_qpts: Quadrature points per dimension
102        build_qdata: Flag to build restriction for quadrature data
103
104    Returns:
105        tuple: (elem_restriction, size, q_data_restriction, num_elem, elem_qpts)
106    """
107    p = degree + 1  # Nodes per element per dimension
108    num_nodes = p ** dim
109    elem_qpts = num_qpts ** dim
110
111    # Calculate grid parameters
112    nd = []
113    num_elem = 1
114    scalar_size = 1
115    for d in range(dim):
116        num_elem *= num_xyz[d]
117        nd.append(num_xyz[d] * (p - 1) + 1)  # Nodes per dimension
118        scalar_size *= nd[d]
119
120    size = scalar_size * num_comp
121
122    # Create element connectivity
123    elem_nodes = np.zeros(num_elem * num_nodes, dtype=np.int32)
124    for e in range(num_elem):
125        # Get element coordinates
126        e_xyz = [0] * dim
127        re = e
128        for d in range(dim):
129            e_xyz[d] = re % num_xyz[d]
130            re //= num_xyz[d]
131
132        # Calculate global node numbers
133        for n in range(num_nodes):
134            g_node = 0
135            g_stride = 1
136            r_node = n
137            for d in range(dim):
138                g_node += (e_xyz[d] * (p - 1) + r_node % p) * g_stride
139                g_stride *= nd[d]
140                r_node //= p
141            elem_nodes[e * num_nodes + n] = g_node
142
143    # Create restrictions
144    elem_restriction = ceed.ElemRestriction(
145        num_elem, num_nodes, num_comp, scalar_size, size, elem_nodes)
146
147    q_data_restriction = None
148    if create_qdata:
149        strides = np.array([1, elem_qpts, elem_qpts * num_q_comp], dtype=np.int32)
150        q_data_restriction = ceed.StridedElemRestriction(
151            num_elem, elem_qpts, num_q_comp, num_elem * elem_qpts * num_q_comp, strides)
152
153    return elem_restriction, size, q_data_restriction, num_elem, elem_qpts
154
155
156def set_cartesian_mesh_coords(ceed, dim, num_xyz, mesh_degree, mesh_coords):
157    """Create Cartesian mesh coordinates
158
159    Args:
160        ceed: libCEED context
161        dim: Spatial dimension
162        num_xyz: Elements per dimension
163        mesh_degree: Mesh polynomial degree
164        mesh_coords: CeedVector to hold mesh coordinates
165
166    Returns:
167        Vector: Mesh coordinates
168    """
169    p = mesh_degree + 1
170    nd = []
171    scalar_size = 1
172    for d in range(dim):
173        nd.append(num_xyz[d] * (p - 1) + 1)
174        scalar_size *= nd[d]
175
176    # Get Lobatto nodes (quadrature points)
177    nodes, _ = ceed.lobatto_quadrature(p)
178    nodes = 0.5 + 0.5 * nodes  # Map from [-1,1] to [0,1]
179
180    # Create coordinates
181    coords = np.zeros(scalar_size * dim)
182    for gs_node in range(scalar_size):
183        r_node = gs_node
184        for d in range(dim):
185            d_1d = r_node % nd[d]
186            elem_id = d_1d // (p - 1)
187            node_id = d_1d % (p - 1)
188            coords[gs_node + scalar_size * d] = (elem_id + nodes[node_id]) / num_xyz[d]
189            r_node //= nd[d]
190
191    mesh_coords.set_array(coords, cmode=libceed.COPY_VALUES)
192    return mesh_coords
193
194
195def transform_mesh_coords(dim, mesh_size, mesh_coords, use_sin=True):
196    """Transform mesh coordinates and return exact surface area
197
198    Args:
199        dim: Spatial dimension
200        mesh_size: Total mesh vector size
201        mesh_coords: Mesh coordinates vector
202        use_sin: Use sinusoidal transformation
203
204    Returns:
205        float: Tuple with exact volume and surface area for transformed mesh
206    """
207    exact_volume = {1: 1.0, 2: 3. / 4. * np.pi, 3: 3. / 4. * np.pi}[dim]
208    exact_area = {1: 2.0, 2: 4.0, 3: 6.0}[dim]
209
210    # Apply transformation to coordinates
211    num_nodes = mesh_size // dim
212    with mesh_coords.array_write() as coords:
213        if dim == 1:
214            for i in range(num_nodes):
215                x = coords[i] - 0.5
216                coords[i] = 0.5 + (1.0 / np.sqrt(3.0)) * np.sin((2.0 / 3.0) * np.pi * x)
217        else:
218            if use_sin:
219                for i in range(num_nodes):
220                    u = 1. + coords[i]
221                    v = np.pi / 2. * coords[i + num_nodes]
222                    coords[i] = u * np.cos(v)
223                    coords[i + num_nodes] = u * np.sin(v)
224            else:
225                for i in range(num_nodes):
226                    x = coords[i] - 0.5
227                    coords[i] = 0.5 + (1.0 / np.sqrt(3.0)) * np.sin((2.0 / 3.0) * np.pi * x)
228
229    return (exact_volume, exact_area)
230
231
232def find_qfs_so(name, path):
233    """Find the QFunctions shared library.
234    Returns:
235        Filepath to shared library object
236    """
237    for root, dirs, files in os.walk(path):
238        if name in files:
239            return os.path.join(root, name)
240
241
242def load_qfs_so():
243    """Load the QFunctions shared library.
244    Returns:
245        Loaded shared library object
246    """
247    file_dir = os.path.join(
248        os.path.dirname(os.path.abspath(__file__)),
249        "build")
250    qfs_so = find_qfs_so(
251        "libceed_c_qfunctions" + get_config_var("EXT_SUFFIX"),
252        file_dir)
253
254    # Load library
255    return ctypes.cdll.LoadLibrary(qfs_so)
256