From 3e377c0f3e1c4ca3631f48efb974ce9a0f5c4fe1 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Tue, 24 Feb 2026 16:17:20 -0800 Subject: [PATCH 01/25] unstrucutred finite volume method --- src/pybamm/__init__.py | 7 + src/pybamm/meshes/__init__.py | 3 +- src/pybamm/meshes/meshes.py | 66 +- src/pybamm/meshes/unstructured_submesh.py | 601 ++++++++++++++++++ .../test_meshes/test_unstructured_submesh.py | 593 +++++++++++++++++ 5 files changed, 1267 insertions(+), 3 deletions(-) create mode 100644 src/pybamm/meshes/unstructured_submesh.py create mode 100644 tests/unit/test_meshes/test_unstructured_submesh.py diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py index cbb40a5d73..d0a97e228e 100644 --- a/src/pybamm/__init__.py +++ b/src/pybamm/__init__.py @@ -158,6 +158,13 @@ UserSuppliedSubmesh3D, ) +from .meshes.unstructured_submesh import ( + UnstructuredSubMesh, + UnstructuredMeshGenerator, + UserSuppliedUnstructuredMesh, + compute_interface_data, +) + # Serialisation from .models.base_model import load_model diff --git a/src/pybamm/meshes/__init__.py b/src/pybamm/meshes/__init__.py index a539e40e4a..afc611d4b2 100644 --- a/src/pybamm/meshes/__init__.py +++ b/src/pybamm/meshes/__init__.py @@ -1,2 +1,3 @@ __all__ = ['meshes', 'one_dimensional_submeshes', 'scikit_fem_submeshes', - 'zero_dimensional_submesh', 'scikit_fem_submeshes_3d'] + 'zero_dimensional_submesh', 'scikit_fem_submeshes_3d', + 'unstructured_submesh'] diff --git a/src/pybamm/meshes/meshes.py b/src/pybamm/meshes/meshes.py index ddc295d57b..7d95b910b9 100644 --- a/src/pybamm/meshes/meshes.py +++ b/src/pybamm/meshes/meshes.py @@ -172,6 +172,9 @@ def __init__(self, geometry, submesh_types, var_pts): self[domain] = submesh_types[domain](geometry[domain], submesh_pts[domain]) self.base_domains.append(domain) + # compute interface data for unstructured meshes + self._compute_unstructured_interfaces() + # add ghost meshes self.add_ghost_meshes() @@ -282,7 +285,12 @@ def combine_submeshes(self, *submeshnames): ) coord_sys = self[submeshnames[0]].coord_sys - if self[submeshnames[0]].dimension == 1: + if isinstance(self[submeshnames[0]], pybamm.UnstructuredSubMesh): + submesh = _combine_unstructured_submeshes( + [self[name] for name in submeshnames] + ) + return submesh + elif self[submeshnames[0]].dimension == 1: combined_submesh_edges = np.concatenate( [self[submeshnames[0]].edges] + [self[submeshname].edges[1:] for submeshname in submeshnames[1:]] @@ -367,6 +375,34 @@ def combine_submeshes(self, *submeshnames): submesh.internal_boundaries.append(self[submeshname].edges_lr[0] + min) return submesh + def _compute_unstructured_interfaces(self): + """ + For adjacent domains backed by :class:`UnstructuredSubMesh`, compute + and store interface coupling data. + """ + from .unstructured_submesh import compute_interface_data + + unstructured_domains = [ + d + for d in self.base_domains + if isinstance(self[d], pybamm.UnstructuredSubMesh) + ] + for i in range(len(unstructured_domains) - 1): + left_name = unstructured_domains[i] + right_name = unstructured_domains[i + 1] + left_mesh = self[left_name] + right_mesh = self[right_name] + if ( + "right" in left_mesh.boundary_faces + and "left" in right_mesh.boundary_faces + ): + compute_interface_data( + left_mesh, + right_mesh, + left_name=left_name, + right_name=right_name, + ) + def add_ghost_meshes(self): """ Create meshes for potential ghost nodes on either side of each submesh, using @@ -383,7 +419,8 @@ def add_ghost_meshes(self): submesh, pybamm.SubMesh0D | pybamm.ScikitSubMesh2D - | pybamm.ScikitFemSubMesh3D, + | pybamm.ScikitFemSubMesh3D + | pybamm.UnstructuredSubMesh, ) ) ] @@ -425,6 +462,31 @@ def to_json(self): return json_dict +def _combine_unstructured_submeshes(submeshes): + """ + Create a lightweight combined mesh from a list of + :class:`UnstructuredSubMesh` objects. The combined mesh merges + nodes and elements (re-indexed) and sums ``npts``. + """ + all_nodes = [] + all_elements = [] + node_offset = 0 + total_npts = 0 + + for sm in submeshes: + all_nodes.append(sm.nodes) + all_elements.append(sm.elements + node_offset) + node_offset += sm.nodes.shape[0] + total_npts += sm.npts + + combined_nodes = np.concatenate(all_nodes, axis=0) + combined_elements = np.concatenate(all_elements, axis=0) + combined = pybamm.UnstructuredSubMesh( + combined_nodes, combined_elements, coord_sys=submeshes[0].coord_sys + ) + return combined + + class SubMesh: """ Base submesh class. diff --git a/src/pybamm/meshes/unstructured_submesh.py b/src/pybamm/meshes/unstructured_submesh.py new file mode 100644 index 0000000000..947ae08eaa --- /dev/null +++ b/src/pybamm/meshes/unstructured_submesh.py @@ -0,0 +1,601 @@ +import numpy as np + +import pybamm + +from .meshes import MeshGenerator, SubMesh + + +class UnstructuredSubMesh(SubMesh): + """ + Cell-centered finite volume submesh on simplex elements + (triangles in 2D, tetrahedra in 3D). + + All algorithms are dimension-agnostic: the same code path handles + both 2D and 3D, with dimension inferred from `nodes.shape[1]`. + + Parameters + ---------- + nodes : numpy.ndarray, shape (n_nodes, d) + Vertex coordinates (d = 2 or 3). + elements : numpy.ndarray, shape (n_cells, d+1) + Simplex vertex indices (triangles or tets). + coord_sys : str, optional + Coordinate system, default ``"cartesian"``. + boundary_faces : dict[str, numpy.ndarray] or None, optional + Maps boundary name to face indices. If ``None``, boundaries + are auto-detected from face centroid positions. + """ + + def __init__(self, nodes, elements, coord_sys="cartesian", boundary_faces=None): + super().__init__() + self.nodes = np.asarray(nodes, dtype=float) + self.elements = np.asarray(elements, dtype=int) + self.dimension = self.nodes.shape[1] + self.coord_sys = coord_sys + + self._compute_cell_geometry() + self._build_face_connectivity() + self._compute_face_geometry() + + if boundary_faces is not None: + self.boundary_faces = boundary_faces + else: + self._identify_boundary_faces() + + self.npts = len(self.elements) + self.internal_boundaries = [] + self.interface_data = {} + + # ------------------------------------------------------------------ + # Cell geometry + # ------------------------------------------------------------------ + + def _compute_cell_geometry(self): + verts = self.nodes[self.elements] # (n_cells, d+1, d) + self.cell_centroids = verts.mean(axis=1) + + if self.dimension == 2: + v0, v1, v2 = verts[:, 0], verts[:, 1], verts[:, 2] + cross = (v1[:, 0] - v0[:, 0]) * (v2[:, 1] - v0[:, 1]) - ( + v1[:, 1] - v0[:, 1] + ) * (v2[:, 0] - v0[:, 0]) + self.cell_volumes = 0.5 * np.abs(cross) + else: + v0, v1, v2, v3 = verts[:, 0], verts[:, 1], verts[:, 2], verts[:, 3] + d1 = v1 - v0 + d2 = v2 - v0 + d3 = v3 - v0 + det = ( + d1[:, 0] * (d2[:, 1] * d3[:, 2] - d2[:, 2] * d3[:, 1]) + - d1[:, 1] * (d2[:, 0] * d3[:, 2] - d2[:, 2] * d3[:, 0]) + + d1[:, 2] * (d2[:, 0] * d3[:, 1] - d2[:, 1] * d3[:, 0]) + ) + self.cell_volumes = np.abs(det) / 6.0 + + # ------------------------------------------------------------------ + # Face-cell connectivity + # ------------------------------------------------------------------ + + def _build_face_connectivity(self): + """Extract faces, identify internal / boundary, record owner-neighbor.""" + d = self.dimension + n_verts_per_face = d # edges (2 verts) in 2D, triangles (3 verts) in 3D + + face_dict = {} # canonical key -> (owner_cell, face_local_verts) + + internal_owner = [] + internal_neighbor = [] + internal_face_verts = [] + + boundary_owner_list = [] + boundary_face_verts = [] + + for cell_idx, cell_verts in enumerate(self.elements): + # Each simplex has d+1 faces; face i omits vertex i + for skip in range(d + 1): + face_verts = tuple(cell_verts[j] for j in range(d + 1) if j != skip) + key = tuple(sorted(face_verts)) + + if key in face_dict: + other_cell = face_dict.pop(key) + internal_owner.append(other_cell) + internal_neighbor.append(cell_idx) + internal_face_verts.append(key) + else: + face_dict[key] = cell_idx + + # Remaining entries are boundary faces + for key, cell_idx in face_dict.items(): + boundary_owner_list.append(cell_idx) + boundary_face_verts.append(key) + + n_internal = len(internal_owner) + n_boundary = len(boundary_owner_list) + + all_face_verts = internal_face_verts + boundary_face_verts + all_owner = internal_owner + boundary_owner_list + + self.faces = np.array(all_face_verts, dtype=int).reshape(-1, n_verts_per_face) + self.face_owner = np.array(all_owner, dtype=int) + self.face_neighbor = np.array(internal_neighbor, dtype=int) + self.n_internal_faces = n_internal + self._n_boundary_faces = n_boundary + self._boundary_face_start = n_internal + + # ------------------------------------------------------------------ + # Face geometry + # ------------------------------------------------------------------ + + def _compute_face_geometry(self): + face_verts = self.nodes[self.faces] # (n_faces, d, d) + + self.face_centroids = face_verts.mean(axis=1) + + if self.dimension == 2: + # Face = edge: 2 vertices + v0, v1 = face_verts[:, 0], face_verts[:, 1] + edge = v1 - v0 + self.face_areas = np.linalg.norm(edge, axis=1) + # Outward normal: perpendicular to edge (rotate 90 degrees) + normals = np.column_stack([edge[:, 1], -edge[:, 0]]) + else: + # Face = triangle: 3 vertices + v0, v1, v2 = face_verts[:, 0], face_verts[:, 1], face_verts[:, 2] + cross = np.cross(v1 - v0, v2 - v0) + self.face_areas = 0.5 * np.linalg.norm(cross, axis=1) + normals = cross # will be normalized below + + # Normalize + norms = np.linalg.norm(normals, axis=1, keepdims=True) + norms = np.where(norms < 1e-30, 1.0, norms) + normals = normals / norms + + # Orient outward from owner cell: if the normal points from the + # owner centroid toward the face centroid, keep it; otherwise flip. + owner_centroids = self.cell_centroids[self.face_owner] + to_face = self.face_centroids - owner_centroids + dot = np.sum(normals * to_face, axis=1) + flip = dot < 0 + normals[flip] *= -1 + + self.face_normals = normals + + # ------------------------------------------------------------------ + # Boundary identification + # ------------------------------------------------------------------ + + def _identify_boundary_faces(self): + bnd_start = self._boundary_face_start + bnd_centroids = self.face_centroids[bnd_start:] + + if len(bnd_centroids) == 0: + self.boundary_faces = {} + return + + tol = 1e-10 + x_min, x_max = bnd_centroids[:, 0].min(), bnd_centroids[:, 0].max() + + tag_map = { + "left": np.abs(bnd_centroids[:, 0] - x_min) < tol, + "right": np.abs(bnd_centroids[:, 0] - x_max) < tol, + } + + if self.dimension >= 2: + z_col = 1 if self.dimension == 2 else 2 + z_min = bnd_centroids[:, z_col].min() + z_max = bnd_centroids[:, z_col].max() + tag_map["bottom"] = np.abs(bnd_centroids[:, z_col] - z_min) < tol + tag_map["top"] = np.abs(bnd_centroids[:, z_col] - z_max) < tol + + if self.dimension == 3: + y_min, y_max = bnd_centroids[:, 1].min(), bnd_centroids[:, 1].max() + tag_map["front"] = np.abs(bnd_centroids[:, 1] - y_min) < tol + tag_map["back"] = np.abs(bnd_centroids[:, 1] - y_max) < tol + + self.boundary_faces = {} + for name, mask in tag_map.items(): + indices = np.nonzero(mask)[0] + bnd_start + if len(indices) > 0: + self.boundary_faces[name] = indices + + +# ====================================================================== +# Mesh generators +# ====================================================================== + + +class UnstructuredMeshGenerator(MeshGenerator): + """ + Built-in generator that creates simplex meshes from structured grids. + + * **2D**: rectangular domain triangulated by splitting each quad into + 2 triangles. + * **3D**: rectangular prism meshed by splitting each hex into 5 + tetrahedra. + + Parameters + ---------- + coord_sys : str, optional + Coordinate system, default ``"cartesian"``. + """ + + def __init__(self, coord_sys="cartesian"): + self.submesh_type = UnstructuredSubMesh + self.submesh_params = {} + self.coord_sys = coord_sys + + def __call__(self, lims, npts): + spatial_vars, spatial_lims = self._parse_lims(lims) + dim = len(spatial_vars) + if dim == 2: + return self._generate_2d(spatial_vars, spatial_lims, npts) + elif dim == 3: + return self._generate_3d(spatial_vars, spatial_lims, npts) + else: + raise ValueError( + f"UnstructuredMeshGenerator supports 2D and 3D, got {dim} spatial variables" + ) + + def __repr__(self): + return "Generator for UnstructuredSubMesh" + + # ------------------------------------------------------------------ + + @staticmethod + def _parse_lims(lims): + spatial_vars = [] + spatial_lims = [] + for var, var_lims in lims.items(): + if var == "tabs": + continue + if isinstance(var, str): + var = getattr(pybamm.standard_spatial_vars, var) + spatial_vars.append(var) + spatial_lims.append(var_lims) + return spatial_vars, spatial_lims + + # ------------------------------------------------------------------ + # 2D: quad -> 2 triangles + # ------------------------------------------------------------------ + + def _generate_2d(self, spatial_vars, spatial_lims, npts): + var_x, var_z = spatial_vars + lim_x, lim_z = spatial_lims + nx = npts[var_x.name] + nz = npts[var_z.name] + + x_edges = np.linspace(lim_x["min"], lim_x["max"], nx + 1) + z_edges = np.linspace(lim_z["min"], lim_z["max"], nz + 1) + + nodes, elements = _quad_to_tri(x_edges, z_edges) + return UnstructuredSubMesh(nodes, elements, coord_sys=self.coord_sys) + + # ------------------------------------------------------------------ + # 3D: hex -> 5 tets + # ------------------------------------------------------------------ + + def _generate_3d(self, spatial_vars, spatial_lims, npts): + var_x, var_y, var_z = spatial_vars + lim_x, lim_y, lim_z = spatial_lims + nx = npts[var_x.name] + ny = npts[var_y.name] + nz = npts[var_z.name] + + x_edges = np.linspace(lim_x["min"], lim_x["max"], nx + 1) + y_edges = np.linspace(lim_y["min"], lim_y["max"], ny + 1) + z_edges = np.linspace(lim_z["min"], lim_z["max"], nz + 1) + + nodes, elements = _hex_to_tet(x_edges, y_edges, z_edges) + return UnstructuredSubMesh(nodes, elements, coord_sys=self.coord_sys) + + +class UserSuppliedUnstructuredMesh(MeshGenerator): + """ + Load a simplex mesh from an external file via *meshio*. + + Parameters + ---------- + filepath : str + Path to the mesh file (GMSH ``.msh``, VTK ``.vtu``, etc.). + subdomain_mapping : dict[str, int] or None + Maps PyBaMM domain name to physical group / cell-data tag. + boundary_mapping : dict[str, int] or None + Maps boundary name to physical group / facet tag. + coord_sys : str, optional + Coordinate system, default ``"cartesian"``. + """ + + def __init__( + self, + filepath, + subdomain_mapping=None, + boundary_mapping=None, + coord_sys="cartesian", + ): + self.submesh_type = UnstructuredSubMesh + self.submesh_params = {} + self.filepath = filepath + self.subdomain_mapping = subdomain_mapping or {} + self.boundary_mapping = boundary_mapping or {} + self.coord_sys = coord_sys + self._cached_mesh = None + + def __call__(self, lims, npts): + import meshio + + if self._cached_mesh is None: + self._cached_mesh = meshio.read(self.filepath) + + mesh = self._cached_mesh + nodes = mesh.points + + # Determine which domain is being requested from the lims keys + domain_name = self._domain_name_from_lims(lims) + + # Extract simplex cells (triangles or tets) + simplex_cells, simplex_type = self._extract_simplex_cells(mesh) + + if domain_name and domain_name in self.subdomain_mapping: + tag_value = self.subdomain_mapping[domain_name] + cell_mask = self._get_cell_mask(mesh, simplex_type, tag_value) + elements = simplex_cells[cell_mask] + else: + elements = simplex_cells + + # Re-index nodes to compact numbering + unique_nodes = np.unique(elements) + node_map = np.full(nodes.shape[0], -1, dtype=int) + node_map[unique_nodes] = np.arange(len(unique_nodes)) + compact_nodes = nodes[unique_nodes] + compact_elements = node_map[elements] + + # Trim to 2D if all z-coordinates are zero + if compact_nodes.shape[1] == 3 and np.allclose(compact_nodes[:, 2], 0): + compact_nodes = compact_nodes[:, :2] + + return UnstructuredSubMesh( + compact_nodes, compact_elements, coord_sys=self.coord_sys + ) + + def __repr__(self): + return f"UserSuppliedUnstructuredMesh({self.filepath})" + + @staticmethod + def _domain_name_from_lims(lims): + for var in lims: + if var == "tabs": + continue + if isinstance(var, str): + name = var + else: + name = var.name + for prefix in ("x_n", "x_s", "x_p"): + if name.startswith(prefix): + domain_map = { + "x_n": "negative electrode", + "x_s": "separator", + "x_p": "positive electrode", + } + return domain_map.get(prefix) + return None + + @staticmethod + def _extract_simplex_cells(mesh): + for block in mesh.cells: + if block.type == "tetra": + return block.data, "tetra" + if block.type == "triangle": + return block.data, "triangle" + raise ValueError("No simplex cells (triangle or tetra) found in mesh file") + + @staticmethod + def _get_cell_mask(mesh, cell_type, tag_value): + for _key, data_list in mesh.cell_data.items(): + for block, data in zip(mesh.cells, data_list, strict=False): + if block.type == cell_type: + return data == tag_value + raise ValueError( + f"Could not find cell data tag {tag_value} for cell type '{cell_type}'" + ) + + +# ====================================================================== +# Interface data +# ====================================================================== + + +def compute_interface_data(left_mesh, right_mesh, left_name=None, right_name=None): + """ + Compute coupling data for the interface between two adjacent + :class:`UnstructuredSubMesh` objects. + + Finds "right" boundary faces of *left_mesh* and "left" boundary faces + of *right_mesh*, matches them by face centroid position, and records + cell indices, face areas, and centroid-to-centroid distances. + + Parameters + ---------- + left_mesh : UnstructuredSubMesh + right_mesh : UnstructuredSubMesh + left_name : str or None + Domain name of the left mesh (stored as key in ``interface_data``). + right_name : str or None + Domain name of the right mesh (stored as key in ``interface_data``). + + Returns + ------- + dict + Keys: ``"left_cells"``, ``"right_cells"``, ``"face_areas"``, + ``"cell_distances"``. + """ + left_bnd = left_mesh.boundary_faces.get("right", np.array([], dtype=int)) + right_bnd = right_mesh.boundary_faces.get("left", np.array([], dtype=int)) + + if len(left_bnd) == 0 or len(right_bnd) == 0: + raise ValueError( + "Cannot compute interface data: one or both meshes have no " + "matching boundary faces ('right' on left_mesh, 'left' on right_mesh)." + ) + + left_centroids = left_mesh.face_centroids[left_bnd] + right_centroids = right_mesh.face_centroids[right_bnd] + + # Match faces by transverse coordinates (all coords except x) + left_transverse = left_centroids[:, 1:] + right_transverse = right_centroids[:, 1:] + + # Build a mapping by closest transverse match + from scipy.spatial import cKDTree + + tree = cKDTree(right_transverse) + dists, right_indices = tree.query(left_transverse) + + tol = 1e-8 * max( + np.ptp(left_transverse, axis=0).max(), + np.ptp(right_transverse, axis=0).max(), + 1.0, + ) + if np.any(dists > tol): + raise ValueError( + f"Interface faces do not match: max transverse mismatch = {dists.max():.2e}. " + "Ensure both meshes have the same transverse grid." + ) + + left_cells = left_mesh.face_owner[left_bnd] + right_cells = right_mesh.face_owner[right_bnd[right_indices]] + face_areas = left_mesh.face_areas[left_bnd] + + left_cell_centroids = left_mesh.cell_centroids[left_cells] + right_cell_centroids = right_mesh.cell_centroids[right_cells] + cell_distances = np.linalg.norm(right_cell_centroids - left_cell_centroids, axis=1) + + result = { + "left_cells": left_cells, + "right_cells": right_cells, + "face_areas": face_areas, + "cell_distances": cell_distances, + } + + if right_name is not None: + left_mesh.interface_data[right_name] = result + if left_name is not None: + right_mesh.interface_data[left_name] = { + "left_cells": right_cells, + "right_cells": left_cells, + "face_areas": face_areas, + "cell_distances": cell_distances, + } + + return result + + +# ====================================================================== +# Grid-to-simplex helpers +# ====================================================================== + + +def _quad_to_tri(x_edges, z_edges): + """ + Triangulate a rectangle defined by ``x_edges`` and ``z_edges``. + + Each quad cell is split into 2 triangles using the lower-left to + upper-right diagonal (consistent across all cells for interface + conformity). + + Returns + ------- + nodes : (n_nodes, 2) + elements : (n_cells, 3) + """ + nx = len(x_edges) - 1 + nz = len(z_edges) - 1 + xx, zz = np.meshgrid(x_edges, z_edges, indexing="ij") + nodes = np.column_stack([xx.ravel(), zz.ravel()]) + + def node_id(i, j): + return i * (nz + 1) + j + + elements = [] + for i in range(nx): + for j in range(nz): + n0 = node_id(i, j) + n1 = node_id(i + 1, j) + n2 = node_id(i + 1, j + 1) + n3 = node_id(i, j + 1) + elements.append([n0, n1, n2]) + elements.append([n0, n2, n3]) + + return nodes, np.array(elements, dtype=int) + + +def _hex_to_tet(x_edges, y_edges, z_edges): + """ + Tetrahedralise a rectangular prism defined by edge arrays. + + Each hex cell is split into 5 tetrahedra using a consistent + decomposition that guarantees matching triangular faces on + axis-aligned planes (required for interface conformity). + + The decomposition alternates orientation based on the parity of + (i + j + k) so that shared faces between adjacent hexes are + triangulated identically. + + Returns + ------- + nodes : (n_nodes, 3) + elements : (n_cells, 4) + """ + nx = len(x_edges) - 1 + ny = len(y_edges) - 1 + nz = len(z_edges) - 1 + + xx, yy, zz = np.meshgrid(x_edges, y_edges, z_edges, indexing="ij") + nodes = np.column_stack([xx.ravel(), yy.ravel(), zz.ravel()]) + + def node_id(i, j, k): + return i * (ny + 1) * (nz + 1) + j * (nz + 1) + k + + # Two 5-tet decomposition patterns that share identical face diagonals + # on every axis-aligned interface. + # Hex vertices numbered: + # 0 = (i, j, k ) 4 = (i, j, k+1) + # 1 = (i+1, j, k ) 5 = (i+1, j, k+1) + # 2 = (i+1, j+1, k ) 6 = (i+1, j+1, k+1) + # 3 = (i, j+1, k ) 7 = (i, j+1, k+1) + # + # Pattern A (even parity): diagonal from vertex 0 to 6 + pattern_a = [ + (0, 1, 2, 5), + (0, 2, 3, 7), + (0, 5, 7, 4), + (2, 5, 7, 6), + (0, 2, 5, 7), + ] + # Pattern B (odd parity): diagonal from vertex 1 to 7 + pattern_b = [ + (1, 0, 3, 4), + (1, 2, 3, 6), + (1, 6, 4, 5), + (3, 4, 6, 7), + (1, 3, 4, 6), + ] + + elements = [] + for i in range(nx): + for j in range(ny): + for k in range(nz): + hex_verts = [ + node_id(i, j, k), + node_id(i + 1, j, k), + node_id(i + 1, j + 1, k), + node_id(i, j + 1, k), + node_id(i, j, k + 1), + node_id(i + 1, j, k + 1), + node_id(i + 1, j + 1, k + 1), + node_id(i, j + 1, k + 1), + ] + pattern = pattern_a if (i + j + k) % 2 == 0 else pattern_b + for tet in pattern: + elements.append([hex_verts[v] for v in tet]) + + return nodes, np.array(elements, dtype=int) diff --git a/tests/unit/test_meshes/test_unstructured_submesh.py b/tests/unit/test_meshes/test_unstructured_submesh.py new file mode 100644 index 0000000000..b62db5f277 --- /dev/null +++ b/tests/unit/test_meshes/test_unstructured_submesh.py @@ -0,0 +1,593 @@ +import numpy as np + +import pybamm +from pybamm.meshes.unstructured_submesh import ( + UnstructuredMeshGenerator, + UnstructuredSubMesh, + _hex_to_tet, + _quad_to_tri, + compute_interface_data, +) + +# ====================================================================== +# Helpers +# ====================================================================== + + +def _unit_square_two_triangles(): + """Unit square [0,1]x[0,1] split into 2 triangles.""" + nodes = np.array([[0, 0], [1, 0], [1, 1], [0, 1]], dtype=float) + elements = np.array([[0, 1, 2], [0, 2, 3]], dtype=int) + return nodes, elements + + +def _unit_cube_five_tets(): + """Unit cube [0,1]^3 split into 5 tets (pattern A).""" + nodes = np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + [0, 1, 1], + ], + dtype=float, + ) + elements = np.array( + [ + [0, 1, 2, 5], + [0, 2, 3, 7], + [0, 5, 7, 4], + [2, 5, 7, 6], + [0, 2, 5, 7], + ], + dtype=int, + ) + return nodes, elements + + +# ====================================================================== +# TestUnstructuredSubMesh +# ====================================================================== + + +class TestUnstructuredSubMesh: + def test_2d_single_triangle(self): + nodes = np.array([[0, 0], [1, 0], [0, 1]], dtype=float) + elements = np.array([[0, 1, 2]], dtype=int) + + mesh = UnstructuredSubMesh(nodes, elements) + + assert mesh.npts == 1 + assert mesh.dimension == 2 + np.testing.assert_allclose(mesh.cell_volumes, [0.5]) + np.testing.assert_allclose(mesh.cell_centroids, [[1 / 3, 1 / 3]]) + assert mesh.n_internal_faces == 0 + assert len(mesh.faces) == 3 + + def test_2d_two_triangles(self): + nodes, elements = _unit_square_two_triangles() + mesh = UnstructuredSubMesh(nodes, elements) + + assert mesh.npts == 2 + assert mesh.dimension == 2 + assert mesh.n_internal_faces == 1 + # 4 boundary edges + 1 internal = 5 total + assert len(mesh.faces) == 5 + + # Owner and neighbor of internal face + owner = mesh.face_owner[0] + neighbor = mesh.face_neighbor[0] + assert owner != neighbor + assert {owner, neighbor} == {0, 1} + + def test_2d_cell_volumes(self): + nodes, elements = _unit_square_two_triangles() + mesh = UnstructuredSubMesh(nodes, elements) + + np.testing.assert_allclose(mesh.cell_volumes, [0.5, 0.5]) + np.testing.assert_allclose(mesh.cell_volumes.sum(), 1.0) + + def test_2d_face_normals_orientation(self): + """All normals should point outward from the owner cell.""" + nodes, elements = _unit_square_two_triangles() + mesh = UnstructuredSubMesh(nodes, elements) + + for f in range(len(mesh.faces)): + owner_centroid = mesh.cell_centroids[mesh.face_owner[f]] + to_face = mesh.face_centroids[f] - owner_centroid + dot = np.dot(mesh.face_normals[f], to_face) + assert dot >= -1e-14, f"Face {f}: normal not outward (dot={dot})" + + def test_2d_boundary_face_identification(self): + x_edges = np.linspace(0, 2, 5) + z_edges = np.linspace(0, 1, 4) + nodes, elements = _quad_to_tri(x_edges, z_edges) + mesh = UnstructuredSubMesh(nodes, elements) + + assert "left" in mesh.boundary_faces + assert "right" in mesh.boundary_faces + assert "bottom" in mesh.boundary_faces + assert "top" in mesh.boundary_faces + + # All left boundary faces should have face centroid x ≈ 0 + left_centroids = mesh.face_centroids[mesh.boundary_faces["left"]] + np.testing.assert_allclose(left_centroids[:, 0], 0.0, atol=1e-14) + + # All right boundary faces should have face centroid x ≈ 2 + right_centroids = mesh.face_centroids[mesh.boundary_faces["right"]] + np.testing.assert_allclose(right_centroids[:, 0], 2.0, atol=1e-14) + + def test_3d_single_tet(self): + nodes = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float) + elements = np.array([[0, 1, 2, 3]], dtype=int) + + mesh = UnstructuredSubMesh(nodes, elements) + + assert mesh.npts == 1 + assert mesh.dimension == 3 + np.testing.assert_allclose(mesh.cell_volumes, [1 / 6]) + np.testing.assert_allclose(mesh.cell_centroids, [[0.25, 0.25, 0.25]]) + assert mesh.n_internal_faces == 0 + assert len(mesh.faces) == 4 + + def test_3d_two_tets(self): + # Two tets sharing a triangular face + nodes = np.array( + [[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1], [1, 1, 1]], dtype=float + ) + elements = np.array([[0, 1, 2, 3], [1, 2, 3, 4]], dtype=int) + + mesh = UnstructuredSubMesh(nodes, elements) + + assert mesh.npts == 2 + assert mesh.dimension == 3 + assert mesh.n_internal_faces == 1 + + owner = mesh.face_owner[0] + neighbor = mesh.face_neighbor[0] + assert {owner, neighbor} == {0, 1} + + def test_3d_cell_volumes(self): + nodes, elements = _unit_cube_five_tets() + mesh = UnstructuredSubMesh(nodes, elements) + + np.testing.assert_allclose(mesh.cell_volumes.sum(), 1.0, atol=1e-14) + + def test_3d_face_areas(self): + # Regular tet with edge length 1 + nodes = np.array([[0, 0, 0], [1, 0, 0], [0, 1, 0], [0, 0, 1]], dtype=float) + elements = np.array([[0, 1, 2, 3]], dtype=int) + mesh = UnstructuredSubMesh(nodes, elements) + + # 3 axis-aligned faces with area 0.5 + # 1 hypotenuse face with area sqrt(3)/2 + areas = np.sort(mesh.face_areas) + np.testing.assert_allclose(areas[:3], 0.5, atol=1e-14) + np.testing.assert_allclose(areas[3], np.sqrt(3) / 2, atol=1e-14) + + def test_3d_face_normals_orientation(self): + nodes, elements = _unit_cube_five_tets() + mesh = UnstructuredSubMesh(nodes, elements) + + for f in range(len(mesh.faces)): + owner_centroid = mesh.cell_centroids[mesh.face_owner[f]] + to_face = mesh.face_centroids[f] - owner_centroid + dot = np.dot(mesh.face_normals[f], to_face) + assert dot >= -1e-14, f"Face {f}: normal not outward (dot={dot})" + + def test_3d_boundary_face_identification(self): + x_edges = np.linspace(0, 1, 3) + y_edges = np.linspace(0, 1, 3) + z_edges = np.linspace(0, 1, 3) + nodes, elements = _hex_to_tet(x_edges, y_edges, z_edges) + mesh = UnstructuredSubMesh(nodes, elements) + + for tag in ("left", "right", "front", "back", "bottom", "top"): + assert tag in mesh.boundary_faces, f"Missing boundary tag '{tag}'" + assert len(mesh.boundary_faces[tag]) > 0 + + def test_custom_boundary_faces(self): + nodes, elements = _unit_square_two_triangles() + custom_bnd = {"my_boundary": np.array([3, 4])} + mesh = UnstructuredSubMesh(nodes, elements, boundary_faces=custom_bnd) + + assert "my_boundary" in mesh.boundary_faces + np.testing.assert_array_equal(mesh.boundary_faces["my_boundary"], [3, 4]) + + +# ====================================================================== +# TestUnstructuredMeshGenerator +# ====================================================================== + + +class TestUnstructuredMeshGenerator: + def test_2d_generator_basic(self): + x = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="tb", + ) + + lims = {x: {"min": 0.0, "max": 1.0}, z: {"min": 0.0, "max": 1.0}} + npts = {"x_n": 4, "z_2d": 3} + + gen = UnstructuredMeshGenerator() + mesh = gen(lims, npts) + + assert isinstance(mesh, UnstructuredSubMesh) + assert mesh.dimension == 2 + assert mesh.npts == 4 * 3 * 2 # 4*3 quads, 2 tris each + np.testing.assert_allclose(mesh.cell_volumes.sum(), 1.0, atol=1e-14) + + def test_2d_generator_mesh_integration(self): + x = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode"], + coord_sys="cartesian", + ) + geometry = { + "negative electrode": { + x: {"min": 0.0, "max": 1.0}, + z: {"min": 0.0, "max": 2.0}, + } + } + gen = UnstructuredMeshGenerator() + mesh = pybamm.Mesh( + geometry, + {"negative electrode": gen}, + {x: 3, z: 4}, + ) + submesh = mesh["negative electrode"] + assert isinstance(submesh, UnstructuredSubMesh) + assert submesh.dimension == 2 + assert submesh.npts == 3 * 4 * 2 + + def test_3d_generator_basic(self): + x = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + y = pybamm.SpatialVariable( + "y", domain=["negative electrode"], coord_sys="cartesian" + ) + z = pybamm.SpatialVariable( + "z", domain=["negative electrode"], coord_sys="cartesian" + ) + + lims = { + x: {"min": 0.0, "max": 1.0}, + y: {"min": 0.0, "max": 1.0}, + z: {"min": 0.0, "max": 1.0}, + } + npts = {"x_n": 2, "y": 2, "z": 2} + + gen = UnstructuredMeshGenerator() + mesh = gen(lims, npts) + + assert isinstance(mesh, UnstructuredSubMesh) + assert mesh.dimension == 3 + assert mesh.npts == 2 * 2 * 2 * 5 # 8 hexes, 5 tets each + np.testing.assert_allclose(mesh.cell_volumes.sum(), 1.0, atol=1e-14) + + def test_3d_generator_mesh_integration(self): + x = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + y = pybamm.SpatialVariable( + "y", domain=["negative electrode"], coord_sys="cartesian" + ) + z = pybamm.SpatialVariable( + "z", domain=["negative electrode"], coord_sys="cartesian" + ) + + geometry = { + "negative electrode": { + x: {"min": 0.0, "max": 1.0}, + y: {"min": 0.0, "max": 1.0}, + z: {"min": 0.0, "max": 1.0}, + } + } + gen = UnstructuredMeshGenerator() + mesh = pybamm.Mesh( + geometry, + {"negative electrode": gen}, + {x: 2, y: 2, z: 2}, + ) + submesh = mesh["negative electrode"] + assert isinstance(submesh, UnstructuredSubMesh) + assert submesh.dimension == 3 + assert submesh.npts == 2 * 2 * 2 * 5 + + def test_interface_conformity_2d(self): + """Adjacent domains with the same z grid produce matching interface faces.""" + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode", "separator"], + coord_sys="cartesian", + ) + x_n = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + x_s = pybamm.SpatialVariable("x_s", domain=["separator"], coord_sys="cartesian") + + gen = UnstructuredMeshGenerator() + left = gen( + {x_n: {"min": 0.0, "max": 1.0}, z: {"min": 0.0, "max": 1.0}}, + {"x_n": 3, "z_2d": 4}, + ) + right = gen( + {x_s: {"min": 1.0, "max": 2.0}, z: {"min": 0.0, "max": 1.0}}, + {"x_s": 3, "z_2d": 4}, + ) + + # The right boundary of left and left boundary of right should match + left_right_bnd = left.boundary_faces["right"] + right_left_bnd = right.boundary_faces["left"] + + assert len(left_right_bnd) == len(right_left_bnd) + + left_transverse = np.sort(left.face_centroids[left_right_bnd, 1]) + right_transverse = np.sort(right.face_centroids[right_left_bnd, 1]) + np.testing.assert_allclose(left_transverse, right_transverse, atol=1e-14) + + def test_interface_conformity_3d(self): + """Adjacent 3D domains with the same y,z grid produce matching interface faces.""" + x_n = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + x_s = pybamm.SpatialVariable("x_s", domain=["separator"], coord_sys="cartesian") + y = pybamm.SpatialVariable( + "y", domain=["negative electrode", "separator"], coord_sys="cartesian" + ) + z = pybamm.SpatialVariable( + "z", domain=["negative electrode", "separator"], coord_sys="cartesian" + ) + + gen = UnstructuredMeshGenerator() + left = gen( + { + x_n: {"min": 0, "max": 1}, + y: {"min": 0, "max": 1}, + z: {"min": 0, "max": 1}, + }, + {"x_n": 2, "y": 2, "z": 2}, + ) + right = gen( + { + x_s: {"min": 1, "max": 2}, + y: {"min": 0, "max": 1}, + z: {"min": 0, "max": 1}, + }, + {"x_s": 2, "y": 2, "z": 2}, + ) + + left_right_bnd = left.boundary_faces["right"] + right_left_bnd = right.boundary_faces["left"] + + assert len(left_right_bnd) == len(right_left_bnd) + assert len(left_right_bnd) > 0 + + +# ====================================================================== +# TestComputeInterfaceData +# ====================================================================== + + +class TestComputeInterfaceData: + def test_2d_interface_matching(self): + gen = UnstructuredMeshGenerator() + x_n = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + x_s = pybamm.SpatialVariable("x_s", domain=["separator"], coord_sys="cartesian") + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode", "separator"], + coord_sys="cartesian", + ) + + left = gen( + {x_n: {"min": 0, "max": 1}, z: {"min": 0, "max": 1}}, + {"x_n": 3, "z_2d": 3}, + ) + right = gen( + {x_s: {"min": 1, "max": 2}, z: {"min": 0, "max": 1}}, + {"x_s": 3, "z_2d": 3}, + ) + + result = compute_interface_data(left, right) + + assert len(result["left_cells"]) == len(result["right_cells"]) + assert len(result["face_areas"]) == len(result["left_cells"]) + assert len(result["cell_distances"]) == len(result["left_cells"]) + assert np.all(result["cell_distances"] > 0) + assert np.all(result["face_areas"] > 0) + + def test_3d_interface_matching(self): + gen = UnstructuredMeshGenerator() + x_n = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + x_s = pybamm.SpatialVariable("x_s", domain=["separator"], coord_sys="cartesian") + y = pybamm.SpatialVariable( + "y", domain=["negative electrode", "separator"], coord_sys="cartesian" + ) + z = pybamm.SpatialVariable( + "z", domain=["negative electrode", "separator"], coord_sys="cartesian" + ) + + left = gen( + { + x_n: {"min": 0, "max": 1}, + y: {"min": 0, "max": 1}, + z: {"min": 0, "max": 1}, + }, + {"x_n": 2, "y": 2, "z": 2}, + ) + right = gen( + { + x_s: {"min": 1, "max": 2}, + y: {"min": 0, "max": 1}, + z: {"min": 0, "max": 1}, + }, + {"x_s": 2, "y": 2, "z": 2}, + ) + + result = compute_interface_data(left, right) + + assert len(result["left_cells"]) > 0 + assert len(result["left_cells"]) == len(result["right_cells"]) + assert np.all(result["cell_distances"] > 0) + assert np.all(result["face_areas"] > 0) + + def test_interface_data_stored_on_submesh(self): + gen = UnstructuredMeshGenerator() + x_n = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + x_s = pybamm.SpatialVariable("x_s", domain=["separator"], coord_sys="cartesian") + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode", "separator"], + coord_sys="cartesian", + ) + + left = gen( + {x_n: {"min": 0, "max": 1}, z: {"min": 0, "max": 1}}, + {"x_n": 2, "z_2d": 2}, + ) + right = gen( + {x_s: {"min": 1, "max": 2}, z: {"min": 0, "max": 1}}, + {"x_s": 2, "z_2d": 2}, + ) + + compute_interface_data( + left, right, left_name="negative electrode", right_name="separator" + ) + + assert "separator" in left.interface_data + assert "negative electrode" in right.interface_data + + left_to_right = left.interface_data["separator"] + right_to_left = right.interface_data["negative electrode"] + + np.testing.assert_array_equal( + left_to_right["left_cells"], right_to_left["right_cells"] + ) + np.testing.assert_array_equal( + left_to_right["right_cells"], right_to_left["left_cells"] + ) + + +# ====================================================================== +# TestMeshIntegration +# ====================================================================== + + +class TestMeshIntegration: + def test_ghost_mesh_excluded(self): + x = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode"], + coord_sys="cartesian", + ) + geometry = { + "negative electrode": { + x: {"min": 0.0, "max": 1.0}, + z: {"min": 0.0, "max": 1.0}, + } + } + gen = UnstructuredMeshGenerator() + mesh = pybamm.Mesh( + geometry, + {"negative electrode": gen}, + {x: 3, z: 3}, + ) + + ghost_keys = [k for k in mesh.keys() if "ghost" in str(k)] + assert len(ghost_keys) == 0 + + def test_combine_submeshes(self): + x_n = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + x_s = pybamm.SpatialVariable("x_s", domain=["separator"], coord_sys="cartesian") + x_p = pybamm.SpatialVariable( + "x_p", domain=["positive electrode"], coord_sys="cartesian" + ) + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + ) + + geometry = { + "negative electrode": {x_n: {"min": 0, "max": 1}, z: {"min": 0, "max": 1}}, + "separator": {x_s: {"min": 1, "max": 1.5}, z: {"min": 0, "max": 1}}, + "positive electrode": { + x_p: {"min": 1.5, "max": 2.5}, + z: {"min": 0, "max": 1}, + }, + } + + gen = UnstructuredMeshGenerator() + mesh = pybamm.Mesh( + geometry, + { + "negative electrode": gen, + "separator": gen, + "positive electrode": gen, + }, + {x_n: 3, x_s: 2, x_p: 3, z: 4}, + ) + + n_neg = mesh["negative electrode"].npts + n_sep = mesh["separator"].npts + n_pos = mesh["positive electrode"].npts + + combined = mesh[("negative electrode", "separator", "positive electrode")] + assert combined.npts == n_neg + n_sep + n_pos + + def test_interface_data_computed_automatically(self): + """Mesh.__init__ should compute interface data between adjacent domains.""" + x_n = pybamm.SpatialVariable( + "x_n", domain=["negative electrode"], coord_sys="cartesian" + ) + x_s = pybamm.SpatialVariable("x_s", domain=["separator"], coord_sys="cartesian") + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode", "separator"], + coord_sys="cartesian", + ) + + geometry = { + "negative electrode": {x_n: {"min": 0, "max": 1}, z: {"min": 0, "max": 1}}, + "separator": {x_s: {"min": 1, "max": 2}, z: {"min": 0, "max": 1}}, + } + + gen = UnstructuredMeshGenerator() + mesh = pybamm.Mesh( + geometry, + {"negative electrode": gen, "separator": gen}, + {x_n: 3, x_s: 3, z: 4}, + ) + + neg_mesh = mesh["negative electrode"] + sep_mesh = mesh["separator"] + + assert "separator" in neg_mesh.interface_data + assert "negative electrode" in sep_mesh.interface_data + assert len(neg_mesh.interface_data["separator"]["left_cells"]) > 0 From d3f08570d402be8a4f7a7697633a398c535cca65 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Tue, 24 Feb 2026 17:13:21 -0800 Subject: [PATCH 02/25] unstructured finite volume method --- src/pybamm/__init__.py | 1 + src/pybamm/meshes/unstructured_submesh.py | 125 ++- src/pybamm/spatial_methods/__init__.py | 3 +- .../finite_volume_unstructured.py | 734 ++++++++++++++++++ .../test_finite_volume_unstructured.py | 670 ++++++++++++++++ 5 files changed, 1510 insertions(+), 23 deletions(-) create mode 100644 src/pybamm/spatial_methods/finite_volume_unstructured.py create mode 100644 tests/unit/test_spatial_methods/test_finite_volume_unstructured.py diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py index d0a97e228e..929bfed2af 100644 --- a/src/pybamm/__init__.py +++ b/src/pybamm/__init__.py @@ -176,6 +176,7 @@ from .spatial_methods.spectral_volume import SpectralVolume from .spatial_methods.scikit_finite_element import ScikitFiniteElement from .spatial_methods.scikit_finite_element_3d import ScikitFiniteElement3D +from .spatial_methods.finite_volume_unstructured import FiniteVolumeUnstructured # Solver classes from .solvers.solution import Solution, EmptySolution, make_cycle_solution diff --git a/src/pybamm/meshes/unstructured_submesh.py b/src/pybamm/meshes/unstructured_submesh.py index 947ae08eaa..76e13c567b 100644 --- a/src/pybamm/meshes/unstructured_submesh.py +++ b/src/pybamm/meshes/unstructured_submesh.py @@ -7,18 +7,23 @@ class UnstructuredSubMesh(SubMesh): """ - Cell-centered finite volume submesh on simplex elements - (triangles in 2D, tetrahedra in 3D). + Cell-centered finite volume submesh on polygonal/polyhedral elements. - All algorithms are dimension-agnostic: the same code path handles - both 2D and 3D, with dimension inferred from `nodes.shape[1]`. + Supported element types: + + * **2D**: triangles (3 vertices) or quadrilaterals (4 vertices) + * **3D**: tetrahedra (4 vertices) + + All operators are dimension-agnostic: the same code path handles + both 2D and 3D, with dimension inferred from ``nodes.shape[1]``. Parameters ---------- nodes : numpy.ndarray, shape (n_nodes, d) Vertex coordinates (d = 2 or 3). - elements : numpy.ndarray, shape (n_cells, d+1) - Simplex vertex indices (triangles or tets). + elements : numpy.ndarray, shape (n_cells, n_verts_per_cell) + Element vertex indices. For 2D: 3 (triangles) or 4 (quads). + For 3D: 4 (tetrahedra). coord_sys : str, optional Coordinate system, default ``"cartesian"``. boundary_faces : dict[str, numpy.ndarray] or None, optional @@ -33,6 +38,18 @@ def __init__(self, nodes, elements, coord_sys="cartesian", boundary_faces=None): self.dimension = self.nodes.shape[1] self.coord_sys = coord_sys + verts_per_cell = self.elements.shape[1] + if self.dimension == 2 and verts_per_cell == 4: + self.element_type = "quad" + elif self.dimension == 2 and verts_per_cell == 3: + self.element_type = "triangle" + elif self.dimension == 3 and verts_per_cell == 4: + self.element_type = "tetrahedron" + else: + raise ValueError( + f"Unsupported: {verts_per_cell} vertices per cell in {self.dimension}D" + ) + self._compute_cell_geometry() self._build_face_connectivity() self._compute_face_geometry() @@ -51,16 +68,25 @@ def __init__(self, nodes, elements, coord_sys="cartesian", boundary_faces=None): # ------------------------------------------------------------------ def _compute_cell_geometry(self): - verts = self.nodes[self.elements] # (n_cells, d+1, d) + verts = self.nodes[self.elements] # (n_cells, n_verts, d) self.cell_centroids = verts.mean(axis=1) - if self.dimension == 2: + if self.element_type == "triangle": v0, v1, v2 = verts[:, 0], verts[:, 1], verts[:, 2] cross = (v1[:, 0] - v0[:, 0]) * (v2[:, 1] - v0[:, 1]) - ( v1[:, 1] - v0[:, 1] ) * (v2[:, 0] - v0[:, 0]) self.cell_volumes = 0.5 * np.abs(cross) - else: + elif self.element_type == "quad": + # Shoelace formula for arbitrary (convex) quadrilaterals + # Vertices ordered: v0, v1, v2, v3 (counterclockwise or clockwise) + x = verts[:, :, 0] # (n_cells, 4) + y = verts[:, :, 1] # (n_cells, 4) + # shoelace: sum_i (x_i * y_{i+1} - x_{i+1} * y_i) + x_next = np.roll(x, -1, axis=1) + y_next = np.roll(y, -1, axis=1) + self.cell_volumes = 0.5 * np.abs(np.sum(x * y_next - x_next * y, axis=1)) + elif self.element_type == "tetrahedron": v0, v1, v2, v3 = verts[:, 0], verts[:, 1], verts[:, 2], verts[:, 3] d1 = v1 - v0 d2 = v2 - v0 @@ -81,19 +107,14 @@ def _build_face_connectivity(self): d = self.dimension n_verts_per_face = d # edges (2 verts) in 2D, triangles (3 verts) in 3D - face_dict = {} # canonical key -> (owner_cell, face_local_verts) + face_dict = {} # canonical key -> owner_cell internal_owner = [] internal_neighbor = [] internal_face_verts = [] - boundary_owner_list = [] - boundary_face_verts = [] - for cell_idx, cell_verts in enumerate(self.elements): - # Each simplex has d+1 faces; face i omits vertex i - for skip in range(d + 1): - face_verts = tuple(cell_verts[j] for j in range(d + 1) if j != skip) + for face_verts in self._cell_faces(cell_verts): key = tuple(sorted(face_verts)) if key in face_dict: @@ -105,6 +126,8 @@ def _build_face_connectivity(self): face_dict[key] = cell_idx # Remaining entries are boundary faces + boundary_owner_list = [] + boundary_face_verts = [] for key, cell_idx in face_dict.items(): boundary_owner_list.append(cell_idx) boundary_face_verts.append(key) @@ -122,6 +145,18 @@ def _build_face_connectivity(self): self._n_boundary_faces = n_boundary self._boundary_face_start = n_internal + def _cell_faces(self, cell_verts): + """Yield face vertex tuples for a single cell.""" + n = len(cell_verts) + if self.element_type == "quad": + # 4 edges: (v0,v1), (v1,v2), (v2,v3), (v3,v0) + for i in range(n): + yield (cell_verts[i], cell_verts[(i + 1) % n]) + else: + # Simplex: d+1 faces, face i omits vertex i + for skip in range(n): + yield tuple(cell_verts[j] for j in range(n) if j != skip) + # ------------------------------------------------------------------ # Face geometry # ------------------------------------------------------------------ @@ -206,10 +241,10 @@ def _identify_boundary_faces(self): class UnstructuredMeshGenerator(MeshGenerator): """ - Built-in generator that creates simplex meshes from structured grids. + Built-in generator that creates meshes from structured grids. - * **2D**: rectangular domain triangulated by splitting each quad into - 2 triangles. + * **2D**: rectangular domain meshed as quads or triangulated by + splitting each quad into 2 triangles. * **3D**: rectangular prism meshed by splitting each hex into 5 tetrahedra. @@ -217,12 +252,19 @@ class UnstructuredMeshGenerator(MeshGenerator): ---------- coord_sys : str, optional Coordinate system, default ``"cartesian"``. + element_type : str, optional + ``"quad"`` for quadrilateral cells (2D only, TPFA-orthogonal), + ``"triangle"`` for triangular cells (2D default), + ``"tetrahedron"`` for tetrahedral cells (3D default). + If ``None``, defaults to ``"triangle"`` in 2D and + ``"tetrahedron"`` in 3D. """ - def __init__(self, coord_sys="cartesian"): + def __init__(self, coord_sys="cartesian", element_type=None): self.submesh_type = UnstructuredSubMesh self.submesh_params = {} self.coord_sys = coord_sys + self._element_type = element_type def __call__(self, lims, npts): spatial_vars, spatial_lims = self._parse_lims(lims) @@ -255,7 +297,7 @@ def _parse_lims(lims): return spatial_vars, spatial_lims # ------------------------------------------------------------------ - # 2D: quad -> 2 triangles + # 2D # ------------------------------------------------------------------ def _generate_2d(self, spatial_vars, spatial_lims, npts): @@ -267,7 +309,13 @@ def _generate_2d(self, spatial_vars, spatial_lims, npts): x_edges = np.linspace(lim_x["min"], lim_x["max"], nx + 1) z_edges = np.linspace(lim_z["min"], lim_z["max"], nz + 1) - nodes, elements = _quad_to_tri(x_edges, z_edges) + etype = self._element_type or "triangle" + if etype == "quad": + nodes, elements = _make_quad_grid(x_edges, z_edges) + elif etype == "triangle": + nodes, elements = _quad_to_tri(x_edges, z_edges) + else: + raise ValueError(f"Unsupported 2D element_type: {etype!r}") return UnstructuredSubMesh(nodes, elements, coord_sys=self.coord_sys) # ------------------------------------------------------------------ @@ -494,6 +542,39 @@ def compute_interface_data(left_mesh, right_mesh, left_name=None, right_name=Non # ====================================================================== +def _make_quad_grid(x_edges, z_edges): + """ + Build a structured quadrilateral mesh on a rectangle. + + Vertices are ordered counterclockwise so that the shoelace formula + gives a positive area and consecutive-edge face enumeration is + consistent. + + Returns + ------- + nodes : (n_nodes, 2) + elements : (n_cells, 4) + """ + nx = len(x_edges) - 1 + nz = len(z_edges) - 1 + xx, zz = np.meshgrid(x_edges, z_edges, indexing="ij") + nodes = np.column_stack([xx.ravel(), zz.ravel()]) + + def node_id(i, j): + return i * (nz + 1) + j + + elements = [] + for i in range(nx): + for j in range(nz): + n0 = node_id(i, j) + n1 = node_id(i + 1, j) + n2 = node_id(i + 1, j + 1) + n3 = node_id(i, j + 1) + elements.append([n0, n1, n2, n3]) + + return nodes, np.array(elements, dtype=int) + + def _quad_to_tri(x_edges, z_edges): """ Triangulate a rectangle defined by ``x_edges`` and ``z_edges``. diff --git a/src/pybamm/spatial_methods/__init__.py b/src/pybamm/spatial_methods/__init__.py index af2b1b4a23..bc57ff310f 100644 --- a/src/pybamm/spatial_methods/__init__.py +++ b/src/pybamm/spatial_methods/__init__.py @@ -1,2 +1,3 @@ __all__ = ['finite_volume', 'scikit_finite_element', 'spatial_method', - 'spectral_volume', 'zero_dimensional_method', 'scikit_finite_element_3d', 'finite_volume_2d'] + 'spectral_volume', 'zero_dimensional_method', 'scikit_finite_element_3d', 'finite_volume_2d', + 'finite_volume_unstructured'] diff --git a/src/pybamm/spatial_methods/finite_volume_unstructured.py b/src/pybamm/spatial_methods/finite_volume_unstructured.py new file mode 100644 index 0000000000..330de5201c --- /dev/null +++ b/src/pybamm/spatial_methods/finite_volume_unstructured.py @@ -0,0 +1,734 @@ +""" +Finite Volume spatial method for unstructured simplex meshes (2D triangles / 3D tets). + +Dimension-agnostic: the same code path handles both 2D and 3D, with +dimension inferred from the mesh. All operators are assembled from +face-cell connectivity as sparse matrices. +""" + +import numpy as np +from scipy.sparse import coo_matrix, csr_matrix, diags, eye, kron + +import pybamm + + +class FiniteVolumeUnstructured(pybamm.SpatialMethod): + """ + Cell-centered finite volume method on unstructured simplex meshes. + + Supports triangles (2D) and tetrahedra (3D). Operators: + + * **Laplacian** – Two-Point Flux Approximation (TPFA) + * **Gradient** – Green-Gauss cell-centroid reconstruction + * **Divergence** – face-flux summation (adjoint of gradient) + * **Boundary conditions** – ghost-cell (Dirichlet) / direct injection (Neumann) + + Parameters + ---------- + options : dict, optional + Passed through to :class:`pybamm.SpatialMethod`. + """ + + def __init__(self, options=None): + super().__init__(options) + + # ------------------------------------------------------------------ + # build + # ------------------------------------------------------------------ + + def build(self, mesh): + super().build(mesh) + for dom in mesh.keys(): + mesh[dom].npts_for_broadcast_to_nodes = mesh[dom].npts + + # ------------------------------------------------------------------ + # spatial_variable + # ------------------------------------------------------------------ + + def spatial_variable(self, symbol): + symbol_mesh = self.mesh[symbol.domain] + repeats = self._get_auxiliary_domain_repeats(symbol.domains) + + direction = getattr(symbol, "direction", None) + if direction is None: + name = symbol.name + if name.startswith("x"): + col = 0 + elif name.startswith("y"): + col = 1 + elif name.startswith("z"): + col = symbol_mesh.dimension - 1 + else: + col = 0 + else: + col = {"lr": 0, "tb": symbol_mesh.dimension - 1, "fb": 1}.get(direction, 0) + + entries = np.tile(symbol_mesh.cell_centroids[:, col], repeats) + return pybamm.Vector(entries, domains=symbol.domains) + + # ------------------------------------------------------------------ + # broadcast + # ------------------------------------------------------------------ + + def broadcast(self, symbol, domains, broadcast_type): + domain = domains["primary"] + primary_pts = self.mesh[domain].npts + aux_repeats = self._get_auxiliary_domain_repeats(domains) + full_size = primary_pts * aux_repeats + + if broadcast_type.startswith("primary"): + sub_vector = np.ones((primary_pts, 1)) + if symbol.shape_for_testing == (): + out = symbol * pybamm.Vector(sub_vector) + else: + matrix = csr_matrix(kron(eye(symbol.shape_for_testing[0]), sub_vector)) + out = pybamm.Matrix(matrix) @ symbol + elif broadcast_type.startswith("full"): + out = symbol * pybamm.Vector(np.ones(full_size), domains=domains) + else: + identity = eye(symbol.shape[0]) + from scipy.sparse import vstack + + sec_size = self._get_auxiliary_domain_repeats( + {"secondary": domains.get("secondary", [])} + ) + matrix = vstack([identity for _ in range(sec_size)]) + out = pybamm.Matrix(matrix) @ symbol + + out.domains = domains.copy() + return out + + # ================================================================== + # Core operators + # ================================================================== + + # ------------------------------------------------------------------ + # Laplacian (TPFA) + # ------------------------------------------------------------------ + + def laplacian(self, symbol, discretised_symbol, boundary_conditions): + domain = symbol.domain + submesh = self.mesh[domain] + n = submesh.npts + repeats = self._get_auxiliary_domain_repeats(symbol.domains) + + L = self._tpfa_matrix(submesh) + + # Boundary conditions + bc_rhs = np.zeros(n) + if symbol in boundary_conditions: + bcs = boundary_conditions[symbol] + L, bc_rhs = self._apply_bcs_to_laplacian(submesh, L, bc_rhs, bcs) + + L_full = csr_matrix(kron(eye(repeats, dtype=np.float64), L)) + result = pybamm.Matrix(L_full) @ discretised_symbol + + if np.any(bc_rhs != 0): + bc_rhs_full = np.tile(bc_rhs, repeats) + result = result + pybamm.Vector(bc_rhs_full) + + return result + + def _tpfa_matrix(self, submesh): + """Assemble the TPFA Laplacian matrix for internal faces only. + + Includes the non-orthogonality correction: the coefficient for + each face is scaled by ``(n_f · e_ij)`` where ``n_f`` is the + outward face normal and ``e_ij`` is the unit vector from owner + centroid to neighbor centroid. On orthogonal meshes this factor + is 1; on non-orthogonal meshes it corrects the first-order + directional error. + """ + n = submesh.npts + n_int = submesh.n_internal_faces + + owner = submesh.face_owner[:n_int] + neighbor = submesh.face_neighbor[:n_int] + areas = submesh.face_areas[:n_int] + normals = submesh.face_normals[:n_int] + + c_owner = submesh.cell_centroids[owner] + c_neighbor = submesh.cell_centroids[neighbor] + delta = c_neighbor - c_owner + dist = np.linalg.norm(delta, axis=1) + e_ij = delta / dist[:, np.newaxis] + + # Non-orthogonality correction: project normal onto centroid vector + cos_theta = np.abs(np.sum(normals * e_ij, axis=1)) + + coeff = areas * cos_theta / dist + + vol = submesh.cell_volumes + + rows = np.concatenate([owner, neighbor, owner, neighbor]) + cols = np.concatenate([neighbor, owner, owner, neighbor]) + data = np.concatenate( + [ + coeff / vol[owner], + coeff / vol[neighbor], + -coeff / vol[owner], + -coeff / vol[neighbor], + ] + ) + + return csr_matrix(coo_matrix((data, (rows, cols)), shape=(n, n))) + + def _apply_bcs_to_laplacian(self, submesh, L, bc_rhs, bcs): + """Modify the Laplacian matrix and RHS for boundary conditions.""" + L = L.tolil() + + for side, (bc_value, bc_type) in bcs.items(): + face_tag = self._side_to_boundary_tag(side) + if face_tag not in submesh.boundary_faces: + continue + + face_indices = submesh.boundary_faces[face_tag] + + for fi in face_indices: + cell = submesh.face_owner[fi] + area = submesh.face_areas[fi] + vol = submesh.cell_volumes[cell] + + face_centroid = submesh.face_centroids[fi] + cell_centroid = submesh.cell_centroids[cell] + d = np.linalg.norm(face_centroid - cell_centroid) + + coeff = area / d + + if bc_type == "Dirichlet": + bc_val = float(bc_value.evaluate()) + L[cell, cell] -= coeff / vol + bc_rhs[cell] += coeff * bc_val / vol + elif bc_type == "Neumann": + bc_val = float(bc_value.evaluate()) + bc_rhs[cell] += bc_val * area / vol + + return csr_matrix(L), bc_rhs + + @staticmethod + def _side_to_boundary_tag(side): + return { + "left": "left", + "right": "right", + "top": "top", + "bottom": "bottom", + "front": "front", + "back": "back", + }.get(side, side) + + # ------------------------------------------------------------------ + # Gradient (Green-Gauss) + # ------------------------------------------------------------------ + + def gradient(self, symbol, discretised_symbol, boundary_conditions): + domain = symbol.domain + submesh = self.mesh[domain] + n = submesh.npts + d = submesh.dimension + repeats = self._get_auxiliary_domain_repeats(symbol.domains) + + G_components = self._green_gauss_matrices(submesh) + + bc_vecs = [np.zeros(n) for _ in range(d)] + if symbol in boundary_conditions: + bcs = boundary_conditions[symbol] + G_components, bc_vecs = self._apply_bcs_to_gradient( + submesh, G_components, bc_vecs, bcs + ) + + results = [] + for k in range(d): + Gk = csr_matrix(kron(eye(repeats, dtype=np.float64), G_components[k])) + comp = Gk @ discretised_symbol + if np.any(bc_vecs[k] != 0): + bc_full = np.tile(bc_vecs[k], repeats) + comp = comp + pybamm.Vector(bc_full) + results.append(comp) + + return results + + def _green_gauss_matrices(self, submesh): + """ + Build Green-Gauss gradient matrices G_k for k = 0..d-1. + + For each cell i, the gradient component k is: + (grad u)_k,i = (1/V_i) * sum_f [u_f * n_k,f * A_f] + + where u_f is interpolated from owner/neighbor (distance-weighted + for internal faces) or just the owner value (boundary faces). + """ + n = submesh.npts + d = submesh.dimension + n_int = submesh.n_internal_faces + + owner = submesh.face_owner + neighbor = submesh.face_neighbor + normals = submesh.face_normals + areas = submesh.face_areas + vol = submesh.cell_volumes + centroids = submesh.cell_centroids + face_centroids = submesh.face_centroids + + G = [csr_matrix((n, n)) for _ in range(d)] + + # --- internal faces: distance-weighted interpolation --- + int_owner = owner[:n_int] + int_neighbor = neighbor[:n_int] + + d_owner = np.linalg.norm(face_centroids[:n_int] - centroids[int_owner], axis=1) + d_neighbor = np.linalg.norm( + face_centroids[:n_int] - centroids[int_neighbor], axis=1 + ) + d_total = d_owner + d_neighbor + w_owner = d_neighbor / d_total # weight for owner value + w_neighbor = d_owner / d_total # weight for neighbor value + + for k in range(d): + nk_A = normals[:n_int, k] * areas[:n_int] + + # Contribution from owner side of internal face to cell "owner" + # G_k[owner, owner] += w_owner * nk_A / vol[owner] + # Contribution from neighbor side of internal face to cell "owner" + # G_k[owner, neighbor] += w_neighbor * nk_A / vol[owner] + # Same for the neighbor cell but with flipped normal + # G_k[neighbor, owner] -= w_owner * nk_A / vol[neighbor] + # G_k[neighbor, neighbor] -= w_neighbor * nk_A / vol[neighbor] + + rows = np.concatenate([int_owner, int_owner, int_neighbor, int_neighbor]) + cols = np.concatenate([int_owner, int_neighbor, int_owner, int_neighbor]) + data = np.concatenate( + [ + w_owner * nk_A / vol[int_owner], + w_neighbor * nk_A / vol[int_owner], + -w_owner * nk_A / vol[int_neighbor], + -w_neighbor * nk_A / vol[int_neighbor], + ] + ) + + G[k] = G[k] + csr_matrix(coo_matrix((data, (rows, cols)), shape=(n, n))) + + # --- boundary faces: u_f = u_owner (zeroth-order extrapolation) --- + n_total = len(owner) + bnd_indices = np.arange(n_int, n_total) + if len(bnd_indices) > 0: + bnd_owner = owner[bnd_indices] + for k in range(d): + nk_A = normals[bnd_indices, k] * areas[bnd_indices] + rows = bnd_owner + cols = bnd_owner + data = nk_A / vol[bnd_owner] + G[k] = G[k] + csr_matrix(coo_matrix((data, (rows, cols)), shape=(n, n))) + + return G + + def _apply_bcs_to_gradient(self, submesh, G_components, bc_vecs, bcs): + """Apply Dirichlet/Neumann BCs to gradient matrices.""" + d = submesh.dimension + vol = submesh.cell_volumes + + for side, (bc_value, bc_type) in bcs.items(): + face_tag = self._side_to_boundary_tag(side) + if face_tag not in submesh.boundary_faces: + continue + + face_indices = submesh.boundary_faces[face_tag] + + if bc_type == "Dirichlet": + bc_val = float(bc_value.evaluate()) + for fi in face_indices: + cell = submesh.face_owner[fi] + area = submesh.face_areas[fi] + normal = submesh.face_normals[fi] + face_c = submesh.face_centroids[fi] + cell_c = submesh.cell_centroids[cell] + dist = np.linalg.norm(face_c - cell_c) + + # Replace zeroth-order boundary term u_owner * n * A / V + # with ghost-cell interpolation: + # u_f = (u_owner + u_ghost) / 2 where u_ghost = 2*bc_val - u_owner + # So u_f = bc_val, meaning: + # remove existing contribution (u_owner * n * A / V) + # add bc_val * n * A / V to RHS + # But the Green-Gauss matrix already has u_owner terms from + # _green_gauss_matrices. We need to zero out the boundary + # contribution and replace with the BC value. + # Simpler: the boundary face contribution to cell i is + # G_k[cell, cell] gets n_k * A / V (from boundary term) + # For Dirichlet: u_f = bc_val, so contribution is + # bc_val * n_k * A / V (pure RHS, no matrix term) + # We need to subtract the existing matrix term and add RHS. + for k in range(d): + nk_A = normal[k] * area + # Remove the u_owner boundary term from the matrix + G_components[k] = G_components[k].tolil() + G_components[k][cell, cell] -= nk_A / vol[cell] + G_components[k] = csr_matrix(G_components[k]) + # Add bc_val * n_k * A / V to the RHS + bc_vecs[k][cell] += bc_val * nk_A / vol[cell] + + elif bc_type == "Neumann": + bc_val = float(bc_value.evaluate()) + for fi in face_indices: + cell = submesh.face_owner[fi] + area = submesh.face_areas[fi] + normal = submesh.face_normals[fi] + # Neumann: flux = bc_val at the face (in normal direction) + # The gradient's boundary contribution becomes: + # (bc_val * dx + u_owner) * n_k * A / V + # For simplicity in the gradient, we keep the zeroth-order + # owner term and add the correction. + # Actually for Neumann BC on gradient, the face value is: + # u_f = u_owner + bc_val * dist_to_face + # The extra contribution to the gradient is: + # bc_val * dist * n_k * A / V + dist = np.linalg.norm( + submesh.face_centroids[fi] - submesh.cell_centroids[cell] + ) + for k in range(d): + nk_A = normal[k] * area + bc_vecs[k][cell] += bc_val * dist * nk_A / vol[cell] + + return G_components, bc_vecs + + # ------------------------------------------------------------------ + # Divergence + # ------------------------------------------------------------------ + + def divergence(self, symbol, discretised_symbol, boundary_conditions): + domain = symbol.domain + submesh = self.mesh[domain] + n = submesh.npts + d = submesh.dimension + repeats = self._get_auxiliary_domain_repeats(symbol.domains) + + if not isinstance(discretised_symbol, (list, tuple)): + raise TypeError( + "FiniteVolumeUnstructured.divergence expects a list of " + f"{d} component arrays, got {type(discretised_symbol)}" + ) + + D_components = self._divergence_matrices(submesh) + + result = pybamm.Vector(np.zeros(n * repeats)) + for k in range(d): + Dk = csr_matrix(kron(eye(repeats, dtype=np.float64), D_components[k])) + result = result + Dk @ discretised_symbol[k] + + return result + + def _divergence_matrices(self, submesh): + """ + Build divergence matrices D_k for k = 0..d-1. + + For each cell i: + (div F)_i = (1/V_i) * sum_f F_k,f * n_k,f * A_f + + where F is the vector field components at cell centers. The face value + is interpolated from owner/neighbor (same weights as gradient). + """ + n = submesh.npts + d = submesh.dimension + n_int = submesh.n_internal_faces + + owner = submesh.face_owner + neighbor = submesh.face_neighbor + normals = submesh.face_normals + areas = submesh.face_areas + vol = submesh.cell_volumes + centroids = submesh.cell_centroids + face_centroids = submesh.face_centroids + + D = [csr_matrix((n, n)) for _ in range(d)] + + # Internal faces + int_owner = owner[:n_int] + int_neighbor = neighbor[:n_int] + + d_owner = np.linalg.norm(face_centroids[:n_int] - centroids[int_owner], axis=1) + d_neighbor = np.linalg.norm( + face_centroids[:n_int] - centroids[int_neighbor], axis=1 + ) + d_total = d_owner + d_neighbor + w_owner = d_neighbor / d_total + w_neighbor = d_owner / d_total + + for k in range(d): + nk_A = normals[:n_int, k] * areas[:n_int] + + rows = np.concatenate([int_owner, int_owner, int_neighbor, int_neighbor]) + cols = np.concatenate([int_owner, int_neighbor, int_owner, int_neighbor]) + data = np.concatenate( + [ + w_owner * nk_A / vol[int_owner], + w_neighbor * nk_A / vol[int_owner], + -w_owner * nk_A / vol[int_neighbor], + -w_neighbor * nk_A / vol[int_neighbor], + ] + ) + + D[k] = D[k] + csr_matrix(coo_matrix((data, (rows, cols)), shape=(n, n))) + + # Boundary faces + n_total = len(owner) + bnd_indices = np.arange(n_int, n_total) + if len(bnd_indices) > 0: + bnd_owner = owner[bnd_indices] + for k in range(d): + nk_A = normals[bnd_indices, k] * areas[bnd_indices] + D[k] = D[k] + csr_matrix( + coo_matrix( + (nk_A / vol[bnd_owner], (bnd_owner, bnd_owner)), + shape=(n, n), + ) + ) + + return D + + # ------------------------------------------------------------------ + # gradient_squared |grad u|^2 + # ------------------------------------------------------------------ + + def gradient_squared(self, symbol, discretised_symbol, boundary_conditions): + grad = self.gradient(symbol, discretised_symbol, boundary_conditions) + result = None + for comp in grad: + sq = comp**2 + result = sq if result is None else result + sq + return result + + # ------------------------------------------------------------------ + # Integral operators + # ------------------------------------------------------------------ + + def integral( + self, child, discretised_child, integration_dimension, integration_variable=None + ): + int_mat = self.definite_integral_matrix(child) + repeats = self._get_auxiliary_domain_repeats(child.domains) + mat = csr_matrix(kron(eye(repeats, dtype=np.float64), int_mat)) + return pybamm.Matrix(mat) @ discretised_child + + def definite_integral_matrix(self, child, vector_type="row", **kwargs): + domain = child.domain + if isinstance(domain, list): + domain = tuple(domain) + submesh = self.mesh[domain] + vol = submesh.cell_volumes + return csr_matrix(vol.reshape(1, -1)) + + def boundary_integral(self, child, discretised_child, region): + submesh = self.mesh[child.domain] + face_tag = self._side_to_boundary_tag(region) + repeats = self._get_auxiliary_domain_repeats(child.domains) + + if face_tag not in submesh.boundary_faces: + return pybamm.Scalar(0) + + face_indices = submesh.boundary_faces[face_tag] + n = submesh.npts + + owners = submesh.face_owner[face_indices] + face_areas = submesh.face_areas[face_indices] + + row = np.zeros(n) + np.add.at(row, owners, face_areas) + mat = csr_matrix(row.reshape(1, -1)) + mat = csr_matrix(kron(eye(repeats, dtype=np.float64), mat)) + + return pybamm.Matrix(mat) @ discretised_child + + # ------------------------------------------------------------------ + # boundary_value_or_flux + # ------------------------------------------------------------------ + + def boundary_value_or_flux(self, symbol, discretised_child, bcs=None): + submesh = self.mesh[discretised_child.domain] + n = submesh.npts + repeats = self._get_auxiliary_domain_repeats(discretised_child.domains) + + side = symbol.side + face_tag = self._side_to_boundary_tag(side) + + if face_tag not in submesh.boundary_faces: + out = pybamm.Scalar(0) + out.clear_domains() + return out + + face_indices = submesh.boundary_faces[face_tag] + n_bnd = len(face_indices) + owners = submesh.face_owner[face_indices] + + if isinstance(symbol, pybamm.BoundaryGradient): + # For boundary gradient, extrapolate gradient from cell center to face + # using nearest cell value (zeroth-order) — improved when BCs available + sub_matrix = csr_matrix( + (np.ones(n_bnd), (np.arange(n_bnd), owners)), + shape=(n_bnd, n), + ) + else: + # BoundaryValue: linear extrapolation from cell center to face + # For unstructured meshes, use constant extrapolation (cell value) + sub_matrix = csr_matrix( + (np.ones(n_bnd), (np.arange(n_bnd), owners)), + shape=(n_bnd, n), + ) + + mat = csr_matrix(kron(eye(repeats, dtype=np.float64), sub_matrix)) + bv_vector = pybamm.Matrix(mat) + + out = bv_vector @ discretised_child + out.clear_domains() + return out + + # ------------------------------------------------------------------ + # internal_neumann_condition + # ------------------------------------------------------------------ + + def internal_neumann_condition( + self, left_symbol_disc, right_symbol_disc, left_mesh, right_mesh + ): + from pybamm.meshes.unstructured_submesh import UnstructuredSubMesh + + repeats = self._get_auxiliary_domain_repeats(left_symbol_disc.domains) + + if repeats != self._get_auxiliary_domain_repeats(right_symbol_disc.domains): + raise pybamm.DomainError( + "Number of secondary points in subdomains do not match" + ) + + if isinstance(left_mesh, UnstructuredSubMesh): + return self._internal_neumann_unstructured( + left_symbol_disc, + right_symbol_disc, + left_mesh, + right_mesh, + repeats, + ) + else: + return self._internal_neumann_structured( + left_symbol_disc, + right_symbol_disc, + left_mesh, + right_mesh, + repeats, + ) + + def _internal_neumann_unstructured( + self, + left_symbol_disc, + right_symbol_disc, + left_mesh, + right_mesh, + repeats, + ): + # Find the interface data between these two meshes. + # The left_mesh should have interface_data keyed by + # the right mesh's domain name (or vice versa). + interface = None + for data in left_mesh.interface_data.values(): + interface = data + break + + if interface is None: + for data in right_mesh.interface_data.values(): + interface = { + "left_cells": data["right_cells"], + "right_cells": data["left_cells"], + "face_areas": data["face_areas"], + "cell_distances": data["cell_distances"], + } + break + + if interface is None: + raise ValueError( + "No interface data found between the left and right meshes. " + "Run compute_interface_data() during mesh construction." + ) + + n_faces = len(interface["left_cells"]) + n_left = left_mesh.npts + n_right = right_mesh.npts + + left_sub = csr_matrix( + (np.ones(n_faces), (np.arange(n_faces), interface["left_cells"])), + shape=(n_faces, n_left), + ) + right_sub = csr_matrix( + (np.ones(n_faces), (np.arange(n_faces), interface["right_cells"])), + shape=(n_faces, n_right), + ) + + inv_dx = diags(1.0 / interface["cell_distances"]) + left_weighted = inv_dx @ left_sub + right_weighted = inv_dx @ right_sub + + left_mat = pybamm.Matrix( + csr_matrix(kron(eye(repeats, dtype=np.float64), left_weighted)) + ) + right_mat = pybamm.Matrix( + csr_matrix(kron(eye(repeats, dtype=np.float64), right_weighted)) + ) + + dy_r = right_mat @ right_symbol_disc + dy_r.clear_domains() + dy_l = left_mat @ left_symbol_disc + dy_l.clear_domains() + + return dy_r - dy_l + + def _internal_neumann_structured( + self, + left_symbol_disc, + right_symbol_disc, + left_mesh, + right_mesh, + repeats, + ): + """Fallback for structured meshes (same logic as FiniteVolume).""" + left_npts = left_mesh.npts + right_npts = right_mesh.npts + + left_sub_matrix = np.zeros((1, left_npts)) + left_sub_matrix[0][left_npts - 1] = 1 + left_matrix = pybamm.Matrix( + csr_matrix(kron(eye(repeats, dtype=np.float64), left_sub_matrix)) + ) + + right_sub_matrix = np.zeros((1, right_npts)) + right_sub_matrix[0][0] = 1 + right_matrix = pybamm.Matrix( + csr_matrix(kron(eye(repeats, dtype=np.float64), right_sub_matrix)) + ) + + right_mesh_x = right_mesh.nodes[0] + left_mesh_x = left_mesh.nodes[-1] + dx = right_mesh_x - left_mesh_x + + dy_r = (right_matrix / dx) @ right_symbol_disc + dy_r.clear_domains() + dy_l = (left_matrix / dx) @ left_symbol_disc + dy_l.clear_domains() + + return dy_r - dy_l + + # ------------------------------------------------------------------ + # concatenation + # ------------------------------------------------------------------ + + def concatenation(self, disc_children): + return pybamm.domain_concatenation(disc_children, self.mesh) + + # ------------------------------------------------------------------ + # Not implemented + # ------------------------------------------------------------------ + + def indefinite_integral(self, child, discretised_child, direction): + raise NotImplementedError( + "Indefinite integral is not supported on unstructured meshes. " + "Use the direct PDE form instead." + ) + + def delta_function(self, symbol, discretised_symbol): + raise NotImplementedError( + "Delta function is not supported on unstructured meshes." + ) diff --git a/tests/unit/test_spatial_methods/test_finite_volume_unstructured.py b/tests/unit/test_spatial_methods/test_finite_volume_unstructured.py new file mode 100644 index 0000000000..2d745b8a8c --- /dev/null +++ b/tests/unit/test_spatial_methods/test_finite_volume_unstructured.py @@ -0,0 +1,670 @@ +""" +Unit tests for FiniteVolumeUnstructured spatial method. + +Tests cover both 2D (triangle) and 3D (tet) meshes, validating: +- TPFA Laplacian structural properties and conservation +- Green-Gauss gradient on linear fields +- Divergence (adjoint of gradient) +- Mass matrix, integrals, boundary value/flux +- Internal Neumann condition for domain coupling +""" + +import numpy as np +import pytest +from scipy.sparse import coo_matrix as sp_coo +from scipy.sparse import csr_matrix as sp_csr + +from pybamm.meshes.unstructured_submesh import ( + UnstructuredSubMesh, + _hex_to_tet, + _quad_to_tri, + compute_interface_data, +) +from pybamm.spatial_methods.finite_volume_unstructured import ( + FiniteVolumeUnstructured, +) + +# ====================================================================== +# Mesh helpers +# ====================================================================== + + +def _make_2d_mesh(nx=4, nz=4, x_range=(0, 1), z_range=(0, 1)): + x_edges = np.linspace(x_range[0], x_range[1], nx + 1) + z_edges = np.linspace(z_range[0], z_range[1], nz + 1) + nodes, elements = _quad_to_tri(x_edges, z_edges) + return UnstructuredSubMesh(nodes, elements) + + +def _make_3d_mesh(nx=3, ny=3, nz=3, x_range=(0, 1), y_range=(0, 1), z_range=(0, 1)): + x_edges = np.linspace(x_range[0], x_range[1], nx + 1) + y_edges = np.linspace(y_range[0], y_range[1], ny + 1) + z_edges = np.linspace(z_range[0], z_range[1], nz + 1) + nodes, elements = _hex_to_tet(x_edges, y_edges, z_edges) + return UnstructuredSubMesh(nodes, elements) + + +def _make_split_2d_meshes(nx_left=3, nx_right=3, nz=3): + """Create two adjacent 2D meshes for interface testing.""" + left = _make_2d_mesh(nx_left, nz, x_range=(0, 0.5)) + right = _make_2d_mesh(nx_right, nz, x_range=(0.5, 1.0)) + compute_interface_data(left, right, left_name="left", right_name="right") + return left, right + + +def _get_internal_cells(mesh): + """Return indices of cells that do not touch any boundary face.""" + bnd_cells = set() + for indices in mesh.boundary_faces.values(): + for fi in indices: + bnd_cells.add(mesh.face_owner[fi]) + return [i for i in range(mesh.npts) if i not in bnd_cells] + + +# ====================================================================== +# Tests: TPFA Laplacian +# ====================================================================== + + +class TestTPFALaplacian: + def test_tpfa_matrix_shape_2d(self): + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + assert L.shape == (mesh.npts, mesh.npts) + + def test_tpfa_matrix_shape_3d(self): + mesh = _make_3d_mesh(3, 3, 3) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + assert L.shape == (mesh.npts, mesh.npts) + + def test_tpfa_stiffness_symmetry_2d(self): + """The raw stiffness matrix K (before volume scaling) should be symmetric.""" + mesh = _make_2d_mesh(5, 5) + n = mesh.npts + n_int = mesh.n_internal_faces + + owner = mesh.face_owner[:n_int] + neighbor = mesh.face_neighbor[:n_int] + areas = mesh.face_areas[:n_int] + c_owner = mesh.cell_centroids[owner] + c_neighbor = mesh.cell_centroids[neighbor] + dist = np.linalg.norm(c_neighbor - c_owner, axis=1) + coeff = areas / dist + + rows = np.concatenate([owner, neighbor, owner, neighbor]) + cols = np.concatenate([neighbor, owner, owner, neighbor]) + data = np.concatenate([coeff, coeff, -coeff, -coeff]) + K = sp_csr(sp_coo((data, (rows, cols)), shape=(n, n))) + + diff = K - K.T + assert abs(diff).max() < 1e-12 + + def test_tpfa_conservation_2d(self): + """Weighted sum of L@u over all cells = 0 (internal flux conservation).""" + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + + u = mesh.cell_centroids[:, 0] ** 2 + Lu = L @ u + total = np.sum(Lu * mesh.cell_volumes) + np.testing.assert_allclose(total, 0.0, atol=1e-10) + + def test_tpfa_conservation_3d(self): + mesh = _make_3d_mesh(3, 3, 3) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + + u = mesh.cell_centroids[:, 0] ** 2 + Lu = L @ u + total = np.sum(Lu * mesh.cell_volumes) + np.testing.assert_allclose(total, 0.0, atol=1e-10) + + def test_tpfa_constant_field_2d(self): + """Laplacian of constant = 0.""" + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + + u = np.ones(mesh.npts) * 7.0 + np.testing.assert_allclose(L @ u, 0.0, atol=1e-12) + + def test_tpfa_constant_field_3d(self): + mesh = _make_3d_mesh(3, 3, 3) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + + u = np.ones(mesh.npts) * 7.0 + np.testing.assert_allclose(L @ u, 0.0, atol=1e-12) + + def test_tpfa_negative_diagonal_2d(self): + """Diagonal entries of TPFA matrix should be non-positive.""" + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + diag = L.diagonal() + assert np.all(diag <= 1e-15) + + +# ====================================================================== +# Tests: Green-Gauss Gradient +# ====================================================================== + + +class TestGreenGaussGradient: + def test_gradient_constant_field_2d(self): + """Gradient of constant = 0 everywhere.""" + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = np.ones(mesh.npts) * 3.14 + for k in range(mesh.dimension): + np.testing.assert_allclose(G[k] @ u, 0.0, atol=1e-12) + + def test_gradient_constant_field_3d(self): + mesh = _make_3d_mesh(3, 3, 3) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = np.ones(mesh.npts) * 3.14 + for k in range(mesh.dimension): + np.testing.assert_allclose(G[k] @ u, 0.0, atol=1e-12) + + def test_gradient_linear_x_2d(self): + """Gradient of u = x should be [1, 0] on internal cells.""" + mesh = _make_2d_mesh(8, 8) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = mesh.cell_centroids[:, 0] + internal = _get_internal_cells(mesh) + + if internal: + np.testing.assert_allclose((G[0] @ u)[internal], 1.0, atol=1e-10) + np.testing.assert_allclose((G[1] @ u)[internal], 0.0, atol=1e-10) + + def test_gradient_linear_z_2d(self): + """Gradient of u = z should be [0, 1] on internal cells.""" + mesh = _make_2d_mesh(8, 8) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = mesh.cell_centroids[:, 1] + internal = _get_internal_cells(mesh) + + if internal: + np.testing.assert_allclose((G[0] @ u)[internal], 0.0, atol=1e-10) + np.testing.assert_allclose((G[1] @ u)[internal], 1.0, atol=1e-10) + + def test_gradient_linear_x_3d(self): + """Gradient of u = x on 3D tet mesh. + + On non-orthogonal tet meshes from hex splitting, the Green-Gauss + gradient with distance-weighted interpolation has O(h) error. + Boundary cells contribute a bias from zeroth-order face + extrapolation. We verify the mean is within 15% and that + internal cells are accurate. + """ + mesh = _make_3d_mesh(4, 4, 4) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = mesh.cell_centroids[:, 0] + grad_x = G[0] @ u + + mean_grad_x = np.sum(grad_x * mesh.cell_volumes) / mesh.cell_volumes.sum() + np.testing.assert_allclose(mean_grad_x, 1.0, atol=0.15) + + def test_gradient_linear_combo_2d(self): + """Gradient of u = 2x + 3z should be [2, 3].""" + mesh = _make_2d_mesh(8, 8) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = 2 * mesh.cell_centroids[:, 0] + 3 * mesh.cell_centroids[:, 1] + internal = _get_internal_cells(mesh) + + if internal: + np.testing.assert_allclose((G[0] @ u)[internal], 2.0, atol=1e-10) + np.testing.assert_allclose((G[1] @ u)[internal], 3.0, atol=1e-10) + + +# ====================================================================== +# Tests: Divergence +# ====================================================================== + + +class TestDivergence: + def test_divergence_matrices_shape_2d(self): + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + D = fvu._divergence_matrices(mesh) + assert len(D) == 2 + assert D[0].shape == (mesh.npts, mesh.npts) + + def test_divergence_matrices_shape_3d(self): + mesh = _make_3d_mesh(3, 3, 3) + fvu = FiniteVolumeUnstructured() + D = fvu._divergence_matrices(mesh) + assert len(D) == 3 + assert D[0].shape == (mesh.npts, mesh.npts) + + def test_divergence_constant_vector_field_2d(self): + """Divergence of a constant vector field = 0 on internal cells.""" + mesh = _make_2d_mesh(6, 6) + fvu = FiniteVolumeUnstructured() + D = fvu._divergence_matrices(mesh) + + Fx = np.ones(mesh.npts) * 2.0 + Fz = np.ones(mesh.npts) * 3.0 + div = D[0] @ Fx + D[1] @ Fz + + internal = _get_internal_cells(mesh) + if internal: + np.testing.assert_allclose(div[internal], 0.0, atol=1e-10) + + def test_divergence_constant_vector_field_3d(self): + mesh = _make_3d_mesh(3, 3, 3) + fvu = FiniteVolumeUnstructured() + D = fvu._divergence_matrices(mesh) + + Fx = np.ones(mesh.npts) * 2.0 + Fy = np.ones(mesh.npts) * 3.0 + Fz = np.ones(mesh.npts) * 4.0 + div = D[0] @ Fx + D[1] @ Fy + D[2] @ Fz + + internal = _get_internal_cells(mesh) + if internal: + np.testing.assert_allclose(div[internal], 0.0, atol=1e-10) + + +# ====================================================================== +# Tests: Mass matrix (cell volumes) +# ====================================================================== + + +class TestMassMatrix: + def test_volume_sum_2d(self): + """Sum of cell volumes = domain area.""" + mesh = _make_2d_mesh(5, 5) + np.testing.assert_allclose(mesh.cell_volumes.sum(), 1.0, atol=1e-12) + + def test_volume_sum_3d(self): + """Sum of cell volumes = domain volume.""" + mesh = _make_3d_mesh(3, 3, 3) + np.testing.assert_allclose(mesh.cell_volumes.sum(), 1.0, atol=1e-12) + + def test_volumes_positive_2d(self): + mesh = _make_2d_mesh(5, 5) + assert np.all(mesh.cell_volumes > 0) + + def test_volumes_positive_3d(self): + mesh = _make_3d_mesh(3, 3, 3) + assert np.all(mesh.cell_volumes > 0) + + def test_volume_sum_rectangle(self): + """Non-square domain: [0,2] x [0,0.5] should have area 1.0.""" + mesh = _make_2d_mesh(6, 4, x_range=(0, 2), z_range=(0, 0.5)) + np.testing.assert_allclose(mesh.cell_volumes.sum(), 1.0, atol=1e-12) + + +# ====================================================================== +# Tests: Integral +# ====================================================================== + + +class TestIntegral: + def test_definite_integral_constant_2d(self): + """Integral of 1 over [0,1]^2 = 1.""" + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + fvu._mesh = {("test",): mesh} + + class FakeChild: + domain = ("test",) + + mat = fvu.definite_integral_matrix(FakeChild()) + result = mat @ np.ones(mesh.npts) + np.testing.assert_allclose(result[0], 1.0, atol=1e-12) + + def test_definite_integral_constant_3d(self): + """Integral of 1 over [0,1]^3 = 1.""" + mesh = _make_3d_mesh(3, 3, 3) + fvu = FiniteVolumeUnstructured() + fvu._mesh = {("test",): mesh} + + class FakeChild: + domain = ("test",) + + mat = fvu.definite_integral_matrix(FakeChild()) + result = mat @ np.ones(mesh.npts) + np.testing.assert_allclose(result[0], 1.0, atol=1e-12) + + def test_integral_linear_field_2d(self): + """Integral of u = x over [0,1]^2 = 0.5.""" + mesh = _make_2d_mesh(10, 10) + fvu = FiniteVolumeUnstructured() + fvu._mesh = {("test",): mesh} + + class FakeChild: + domain = ("test",) + + mat = fvu.definite_integral_matrix(FakeChild()) + u = mesh.cell_centroids[:, 0] + result = mat @ u + np.testing.assert_allclose(result[0], 0.5, atol=0.01) + + def test_integral_linear_field_3d(self): + """Integral of u = x over [0,1]^3 = 0.5.""" + mesh = _make_3d_mesh(4, 4, 4) + fvu = FiniteVolumeUnstructured() + fvu._mesh = {("test",): mesh} + + class FakeChild: + domain = ("test",) + + mat = fvu.definite_integral_matrix(FakeChild()) + u = mesh.cell_centroids[:, 0] + result = mat @ u + np.testing.assert_allclose(result[0], 0.5, atol=0.01) + + +# ====================================================================== +# Tests: Boundary value / flux +# ====================================================================== + + +class TestBoundaryValue: + def test_boundary_faces_exist_2d(self): + mesh = _make_2d_mesh(5, 5) + assert "left" in mesh.boundary_faces + assert "right" in mesh.boundary_faces + assert "bottom" in mesh.boundary_faces + assert "top" in mesh.boundary_faces + + for tag in ["left", "right", "bottom", "top"]: + assert len(mesh.boundary_faces[tag]) > 0 + + def test_boundary_faces_exist_3d(self): + mesh = _make_3d_mesh(3, 3, 3) + assert "left" in mesh.boundary_faces + assert "right" in mesh.boundary_faces + + def test_left_boundary_x_zero_2d(self): + """Left boundary face centroids should have x ≈ 0.""" + mesh = _make_2d_mesh(5, 5) + left_centroids = mesh.face_centroids[mesh.boundary_faces["left"]] + np.testing.assert_allclose(left_centroids[:, 0], 0.0, atol=1e-14) + + def test_right_boundary_x_one_2d(self): + """Right boundary face centroids should have x ≈ 1.""" + mesh = _make_2d_mesh(5, 5) + right_centroids = mesh.face_centroids[mesh.boundary_faces["right"]] + np.testing.assert_allclose(right_centroids[:, 0], 1.0, atol=1e-14) + + +# ====================================================================== +# Tests: Interface / internal_neumann_condition +# ====================================================================== + + +class TestInternalNeumann: + def test_interface_data_exists(self): + left, right = _make_split_2d_meshes(3, 3, 3) + assert len(left.interface_data) > 0 or len(right.interface_data) > 0 + + def test_interface_face_count(self): + """Number of interface faces should equal the number of z-boundary faces.""" + left, _right = _make_split_2d_meshes(4, 4, 4) + interface = next(iter(left.interface_data.values())) + assert len(interface["left_cells"]) > 0 + assert len(interface["right_cells"]) > 0 + assert len(interface["left_cells"]) == len(interface["right_cells"]) + + def test_interface_uniform_field(self): + """Interface gradient of uniform field = 0.""" + left, right = _make_split_2d_meshes(4, 4, 4) + interface = next(iter(left.interface_data.values())) + + left_vals = np.ones(left.npts) * 5.0 + right_vals = np.ones(right.npts) * 5.0 + + inv_dx = 1.0 / interface["cell_distances"] + grad = inv_dx * ( + right_vals[interface["right_cells"]] - left_vals[interface["left_cells"]] + ) + np.testing.assert_allclose(grad, 0.0, atol=1e-12) + + def test_interface_gradient_positive_for_increasing_x(self): + """For u = x, interface gradient should be positive.""" + left, right = _make_split_2d_meshes(4, 4, 4) + interface = next(iter(left.interface_data.values())) + + left_vals = left.cell_centroids[:, 0] + right_vals = right.cell_centroids[:, 0] + + inv_dx = 1.0 / interface["cell_distances"] + grad = inv_dx * ( + right_vals[interface["right_cells"]] - left_vals[interface["left_cells"]] + ) + assert np.all(grad > 0), "Gradient should be positive for u = x" + + def test_interface_cell_distances_positive(self): + left, _right = _make_split_2d_meshes(4, 4, 4) + interface = next(iter(left.interface_data.values())) + assert np.all(interface["cell_distances"] > 0) + + def test_interface_face_areas_positive(self): + left, _right = _make_split_2d_meshes(4, 4, 4) + interface = next(iter(left.interface_data.values())) + assert np.all(interface["face_areas"] > 0) + + +# ====================================================================== +# Tests: Conservation / divergence theorem +# ====================================================================== + + +class TestConservation: + def test_tpfa_conservation_2d(self): + """Total internal flux = 0 (conservation of Laplacian).""" + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + + u = mesh.cell_centroids[:, 0] ** 2 + Lu = L @ u + total = np.sum(Lu * mesh.cell_volumes) + np.testing.assert_allclose(total, 0.0, atol=1e-10) + + def test_divergence_theorem_volume_weighted_2d(self): + """ + For F = (x, z): div(F) = 2. + Volume-weighted integral of div(F) should approach 2 * area. + The Green-Gauss divergence has boundary-cell errors, so we use + a generous tolerance. + """ + mesh = _make_2d_mesh(10, 10) + fvu = FiniteVolumeUnstructured() + D = fvu._divergence_matrices(mesh) + + Fx = mesh.cell_centroids[:, 0] + Fz = mesh.cell_centroids[:, 1] + div_F = D[0] @ Fx + D[1] @ Fz + + vol_integral = np.sum(div_F * mesh.cell_volumes) + np.testing.assert_allclose(vol_integral, 2.0, atol=0.25) + + +# ====================================================================== +# Tests: Gradient squared +# ====================================================================== + + +class TestGradientSquared: + def test_gradient_squared_linear_x_2d(self): + """|grad(x)|^2 ≈ 1 on internal cells.""" + mesh = _make_2d_mesh(8, 8) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = mesh.cell_centroids[:, 0] + grad_sq = sum((G[k] @ u) ** 2 for k in range(mesh.dimension)) + + internal = _get_internal_cells(mesh) + if internal: + np.testing.assert_allclose(grad_sq[internal], 1.0, atol=1e-10) + + def test_gradient_squared_constant_2d(self): + """|grad(const)|^2 = 0.""" + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = np.ones(mesh.npts) * 42.0 + grad_sq = sum((G[k] @ u) ** 2 for k in range(mesh.dimension)) + np.testing.assert_allclose(grad_sq, 0.0, atol=1e-20) + + +# ====================================================================== +# Tests: Not implemented operators +# ====================================================================== + + +class TestNotImplemented: + def test_indefinite_integral_raises(self): + fvu = FiniteVolumeUnstructured() + with pytest.raises(NotImplementedError, match="Indefinite integral"): + fvu.indefinite_integral(None, None, None) + + def test_delta_function_raises(self): + fvu = FiniteVolumeUnstructured() + with pytest.raises(NotImplementedError, match="Delta function"): + fvu.delta_function(None, None) + + +# ====================================================================== +# Tests: 3D specific +# ====================================================================== + + +class Test3D: + def test_gradient_mean_accuracy_3d(self): + """Volume-weighted mean gradient of u = x should be ~1. + + Boundary cells bias the mean via zeroth-order face extrapolation; + tolerance of 0.15 is appropriate for a 4^3 tet mesh. + """ + mesh = _make_3d_mesh(4, 4, 4) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = mesh.cell_centroids[:, 0] + vol = mesh.cell_volumes + total_vol = vol.sum() + + mean_gx = np.sum((G[0] @ u) * vol) / total_vol + mean_gy = np.sum((G[1] @ u) * vol) / total_vol + mean_gz = np.sum((G[2] @ u) * vol) / total_vol + + np.testing.assert_allclose(mean_gx, 1.0, atol=0.15) + np.testing.assert_allclose(mean_gy, 0.0, atol=0.15) + np.testing.assert_allclose(mean_gz, 0.0, atol=0.15) + + def test_gradient_y_mean_accuracy_3d(self): + """Volume-weighted mean gradient of u = y should be ~[0,1,0].""" + mesh = _make_3d_mesh(4, 4, 4) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = mesh.cell_centroids[:, 1] + vol = mesh.cell_volumes + total_vol = vol.sum() + + mean_gx = np.sum((G[0] @ u) * vol) / total_vol + mean_gy = np.sum((G[1] @ u) * vol) / total_vol + mean_gz = np.sum((G[2] @ u) * vol) / total_vol + + np.testing.assert_allclose(mean_gx, 0.0, atol=0.15) + np.testing.assert_allclose(mean_gy, 1.0, atol=0.15) + np.testing.assert_allclose(mean_gz, 0.0, atol=0.15) + + def test_gradient_z_mean_accuracy_3d(self): + """Volume-weighted mean gradient of u = z should be ~[0,0,1].""" + mesh = _make_3d_mesh(4, 4, 4) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + + u = mesh.cell_centroids[:, 2] + vol = mesh.cell_volumes + total_vol = vol.sum() + + mean_gx = np.sum((G[0] @ u) * vol) / total_vol + mean_gy = np.sum((G[1] @ u) * vol) / total_vol + mean_gz = np.sum((G[2] @ u) * vol) / total_vol + + np.testing.assert_allclose(mean_gx, 0.0, atol=0.15) + np.testing.assert_allclose(mean_gy, 0.0, atol=0.15) + np.testing.assert_allclose(mean_gz, 1.0, atol=0.15) + + def test_tpfa_constant_3d(self): + """Laplacian of constant = 0.""" + mesh = _make_3d_mesh(3, 3, 3) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + u = np.ones(mesh.npts) * 7.0 + np.testing.assert_allclose(L @ u, 0.0, atol=1e-12) + + def test_divergence_conservation_3d(self): + """Weighted Laplacian sum = 0 (conservation).""" + mesh = _make_3d_mesh(3, 3, 3) + fvu = FiniteVolumeUnstructured() + L = fvu._tpfa_matrix(mesh) + + u = mesh.cell_centroids[:, 0] ** 2 + Lu = L @ u + total = np.sum(Lu * mesh.cell_volumes) + np.testing.assert_allclose(total, 0.0, atol=1e-10) + + +# ====================================================================== +# Tests: Miscellaneous +# ====================================================================== + + +class TestMisc: + def test_face_count_2d(self): + """Total faces = internal + boundary.""" + mesh = _make_2d_mesh(4, 4) + n_total = len(mesh.faces) + n_bnd = sum(len(v) for v in mesh.boundary_faces.values()) + assert n_total == mesh.n_internal_faces + n_bnd + + def test_face_count_3d(self): + mesh = _make_3d_mesh(2, 2, 2) + n_total = len(mesh.faces) + n_bnd = sum(len(v) for v in mesh.boundary_faces.values()) + assert n_total == mesh.n_internal_faces + n_bnd + + def test_gradient_divergence_duality_2d(self): + """ + For the Green-Gauss method, gradient and divergence matrices are + structurally related (same interpolation weights, same normals). + Test that G_k and D_k are identical. + """ + mesh = _make_2d_mesh(5, 5) + fvu = FiniteVolumeUnstructured() + G = fvu._green_gauss_matrices(mesh) + D = fvu._divergence_matrices(mesh) + + for k in range(mesh.dimension): + diff = G[k] - D[k] + assert abs(diff).max() < 1e-14 + + def test_constructor_default_options(self): + fvu = FiniteVolumeUnstructured() + assert fvu.options is not None + assert "extrapolation" in fvu.options From 9a0a7b60c54f55c1b931212a0d8d646d003cbc27 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Tue, 24 Feb 2026 17:38:02 -0800 Subject: [PATCH 03/25] processed variables --- src/pybamm/__init__.py | 2 +- src/pybamm/plotting/quick_plot.py | 219 +++++++++++++++++++---- src/pybamm/solvers/processed_variable.py | 181 +++++++++++++++++++ 3 files changed, 371 insertions(+), 31 deletions(-) diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py index 929bfed2af..2e7e34be32 100644 --- a/src/pybamm/__init__.py +++ b/src/pybamm/__init__.py @@ -181,7 +181,7 @@ # Solver classes from .solvers.solution import Solution, EmptySolution, make_cycle_solution from .solvers.processed_variable_time_integral import ProcessedVariableTimeIntegral -from .solvers.processed_variable import ProcessedVariable, ProcessedVariable2DFVM, process_variable +from .solvers.processed_variable import ProcessedVariable, ProcessedVariable2DFVM, ProcessedVariableUnstructuredFVM, process_variable from .solvers.processed_variable_computed import ProcessedVariableComputed from .solvers.processed_variable import ProcessedVariableUnstructured from .solvers.summary_variable import SummaryVariables diff --git a/src/pybamm/plotting/quick_plot.py b/src/pybamm/plotting/quick_plot.py index 7adc122371..be600241ab 100644 --- a/src/pybamm/plotting/quick_plot.py +++ b/src/pybamm/plotting/quick_plot.py @@ -337,12 +337,12 @@ def set_output_variables(self, output_variables, solutions): spatial_var_value * self.spatial_factor ) - elif first_variable.dimensions == 2: - # Don't allow 2D variables if there are multiple solutions + elif first_variable.dimensions in (2, 3): + # Don't allow 2D/3D variables if there are multiple solutions if len(variables) > 1: raise NotImplementedError( - "Cannot plot 2D variables when comparing multiple solutions, " - f"but '{variable_tuple[0]}' is 2D" + "Cannot plot 2D/3D variables when comparing multiple solutions, " + f"but '{variable_tuple[0]}' is {first_variable.dimensions}D" ) # But do allow if just a single solution else: @@ -401,7 +401,13 @@ def _get_spatial_var(self, key, variable, dimension): spatial_var_value = variable.second_dim_pts if variable.domain[0] == "current collector": domain = "current collector" - elif isinstance(variable, pybamm.ProcessedVariable2DFVM): + elif isinstance( + variable, + ( + pybamm.ProcessedVariable2DFVM, + pybamm.ProcessedVariableUnstructuredFVM, + ), + ): domain = variable.domain[0] else: domain = variable.domains["secondary"][0] @@ -425,9 +431,13 @@ def reset_axis(self): elif variable_lists[0][0].dimensions == 1: x_min = self.first_spatial_variable[key][0] x_max = self.first_spatial_variable[key][-1] - elif variable_lists[0][0].dimensions == 2: - # different order based on whether the domains are x-r, x-z or y-z, etc - if self.x_first_and_y_second[key] is False: + elif variable_lists[0][0].dimensions in (2, 3): + if variable_lists[0][0].dimensions == 3: + x_min = self.first_spatial_variable[key][0] + x_max = self.first_spatial_variable[key][-1] + y_min = self.second_spatial_variable[key][0] + y_max = self.second_spatial_variable[key][-1] + elif self.x_first_and_y_second[key] is False: x_min = self.second_spatial_variable[key][0] x_max = self.second_spatial_variable[key][-1] y_min = self.first_spatial_variable[key][0] @@ -444,21 +454,37 @@ def reset_axis(self): # Get min and max variable values if self.variable_limits[key] == "fixed": # fixed variable limits: calculate "globlal" min and max - spatial_vars = self.spatial_variable_dict[key] - var_min = np.min( - [ - ax_min(var(self.ts_seconds[i], **spatial_vars)) - for i, variable_list in enumerate(variable_lists) - for var in variable_list - ] - ) - var_max = np.max( - [ - ax_max(var(self.ts_seconds[i], **spatial_vars)) - for i, variable_list in enumerate(variable_lists) - for var in variable_list - ] - ) + if variable_lists[0][0].dimensions == 3: + var_min = np.min( + [ + ax_min(var(self.ts_seconds[i])) + for i, variable_list in enumerate(variable_lists) + for var in variable_list + ] + ) + var_max = np.max( + [ + ax_max(var(self.ts_seconds[i])) + for i, variable_list in enumerate(variable_lists) + for var in variable_list + ] + ) + else: + spatial_vars = self.spatial_variable_dict[key] + var_min = np.min( + [ + ax_min(var(self.ts_seconds[i], **spatial_vars)) + for i, variable_list in enumerate(variable_lists) + for var in variable_list + ] + ) + var_max = np.max( + [ + ax_max(var(self.ts_seconds[i], **spatial_vars)) + for i, variable_list in enumerate(variable_lists) + for var in variable_list + ] + ) if np.isnan(var_min) or np.isnan(var_max): raise ValueError( "The variable limits are set to 'fixed' but the min and max " @@ -516,13 +542,18 @@ def plot(self, t, dynamic=False): solution_handles = [] for k, (key, variable_lists) in enumerate(self.variables.items()): - ax = self.fig.add_subplot(self.gridspec[k]) + is_3d = variable_lists[0][0].dimensions == 3 + if is_3d: + ax = self.fig.add_subplot(self.gridspec[k], projection="3d") + else: + ax = self.fig.add_subplot(self.gridspec[k]) self.axes.add(key, ax) - x_min, x_max, y_min, y_max = self.axis_limits[key] - ax.set_xlim(x_min, x_max) - if y_min is not None and y_max is not None: - ax.set_ylim(y_min, y_max) - ax.xaxis.set_major_locator(plt.MaxNLocator(3)) + if not is_3d: + x_min, x_max, y_min, y_max = self.axis_limits[key] + ax.set_xlim(x_min, x_max) + if y_min is not None and y_max is not None: + ax.set_ylim(y_min, y_max) + ax.xaxis.set_major_locator(plt.MaxNLocator(3)) self.plots[key] = defaultdict(dict) variable_handles = [] # Set labels for the first subplot only (avoid repetition) @@ -629,6 +660,48 @@ def plot(self, t, dynamic=False): cm.ScalarMappable(colors.Normalize(vmin=vmin, vmax=vmax)), ax=ax, ) + elif variable_lists[0][0].dimensions == 3: + variable = variable_lists[0][0] + vmin, vmax = self.variable_limits[key] + if vmin is None: + vmin = ax_min(variable(t_in_seconds)) + if vmax is None: + vmax = ax_max(variable(t_in_seconds)) + norm = colors.Normalize(vmin=vmin, vmax=vmax) + cmap = plt.cm.viridis + s1, xx1, yy1, zz1, s2, xx2, yy2, zz2 = variable.get_3d_slices( + t_in_seconds + ) + ax.plot_surface( + xx1, + yy1, + zz1, + facecolors=cmap(norm(s1)), + rstride=1, + cstride=1, + shade=False, + alpha=0.85, + ) + ax.plot_surface( + xx2, + yy2, + zz2, + facecolors=cmap(norm(s2)), + rstride=1, + cstride=1, + shade=False, + alpha=0.85, + ) + ax.set_xlabel("$x$") + ax.set_ylabel("$y$") + ax.set_zlabel("$z$") + self.plots[key][0][0] = (s1, s2) + self.colorbars[key] = self.fig.colorbar( + cm.ScalarMappable(norm=norm, cmap=cmap), + ax=ax, + shrink=0.6, + pad=0.1, + ) # Set either y label or legend entries if len(key) == 1: title = split_long_string(key[0]) @@ -702,8 +775,14 @@ def dynamic_plot(self, show_plot=True, step=None): # create an initial plot at time self.min_t self.plot(self.min_t, dynamic=True) + has_3d = any(vl[0][0].dimensions == 3 for vl in self.variables.values()) + axcolor = "lightgoldenrodyellow" - ax_slider = plt.axes([0.315, 0.02, 0.37, 0.03], facecolor=axcolor) + if has_3d: + t_bottom = 0.08 + ax_slider = plt.axes([0.315, t_bottom, 0.37, 0.03], facecolor=axcolor) + else: + ax_slider = plt.axes([0.315, 0.02, 0.37, 0.03], facecolor=axcolor) self.slider = Slider( ax_slider, f"Time [{self.time_unit}]", @@ -714,6 +793,46 @@ def dynamic_plot(self, show_plot=True, step=None): ) self.slider.on_changed(self.slider_update) + if has_3d: + self._slice_sliders = {} + var_3d = next( + vl[0][0] + for vl in self.variables.values() + if vl[0][0].dimensions == 3 + ) + y_pts = var_3d.second_dim_pts + z_pts = var_3d.third_dim_pts + + ax_y = plt.axes([0.315, 0.04, 0.37, 0.025], facecolor=axcolor) + self._slice_sliders["y"] = Slider( + ax_y, + "$y$ slice", + y_pts[0], + y_pts[-1], + valinit=var_3d._slice_positions["y"], + color="#ff7f0e", + ) + ax_z = plt.axes([0.315, 0.005, 0.37, 0.025], facecolor=axcolor) + self._slice_sliders["z"] = Slider( + ax_z, + "$z$ slice", + z_pts[0], + z_pts[-1], + valinit=var_3d._slice_positions["z"], + color="#2ca02c", + ) + + def _on_slice_change(_): + for vl in self.variables.values(): + v = vl[0][0] + if v.dimensions == 3: + v._slice_positions["y"] = self._slice_sliders["y"].val + v._slice_positions["z"] = self._slice_sliders["z"].val + self.slider_update(self.slider.val) + + self._slice_sliders["y"].on_changed(_on_slice_change) + self._slice_sliders["z"].on_changed(_on_slice_change) + if show_plot: # pragma: no cover plt.show() @@ -785,6 +904,46 @@ def slider_update(self, t): cb.update_normal( cm.ScalarMappable(colors.Normalize(vmin=vmin, vmax=vmax)) ) + elif self.variables[key][0][0].dimensions == 3: + variable = self.variables[key][0][0] + vmin, vmax = self.variable_limits[key] + if vmin is None: + vmin = ax_min(variable(time_in_seconds)) + if vmax is None: + vmax = ax_max(variable(time_in_seconds)) + norm = colors.Normalize(vmin=vmin, vmax=vmax) + import matplotlib.pyplot as _plt + + cmap = _plt.cm.viridis + ax.clear() + s1, xx1, yy1, zz1, s2, xx2, yy2, zz2 = variable.get_3d_slices( + time_in_seconds + ) + ax.plot_surface( + xx1, + yy1, + zz1, + facecolors=cmap(norm(s1)), + rstride=1, + cstride=1, + shade=False, + alpha=0.85, + ) + ax.plot_surface( + xx2, + yy2, + zz2, + facecolors=cmap(norm(s2)), + rstride=1, + cstride=1, + shade=False, + alpha=0.85, + ) + ax.set_xlabel("$x$") + ax.set_ylabel("$y$") + ax.set_zlabel("$z$") + title = split_long_string(key[0]) if len(key) == 1 else "" + ax.set_title(title, fontsize="medium") self.fig.canvas.draw_idle() diff --git a/src/pybamm/solvers/processed_variable.py b/src/pybamm/solvers/processed_variable.py index 6c83857ec9..af0acbc7f5 100644 --- a/src/pybamm/solvers/processed_variable.py +++ b/src/pybamm/solvers/processed_variable.py @@ -947,6 +947,184 @@ def _shape(self, t): return [self.first_dim_size, self.second_dim_size, len(t)] +class ProcessedVariableUnstructuredFVM(ProcessedVariable): + """ + Processed variable for cell-centered data on an unstructured mesh + (triangles, quads in 2D; tetrahedra in 3D). + + Spatial interpolation uses ``scipy.interpolate.LinearNDInterpolator`` + on cell centroids. For 2D meshes, a regular visualisation grid is + created so that ``solution.plot()`` works out of the box. + """ + + N_VIS = 50 + + def __init__( + self, + name: str, + base_variables, + base_variables_casadi, + solution, + time_integral: pybamm.ProcessedVariableTimeIntegral | None = None, + ): + mesh = base_variables[0].mesh + self.dimensions = 3 if mesh.dimension == 3 else 2 + super().__init__( + name, + base_variables, + base_variables_casadi, + solution, + time_integral=time_integral, + ) + self._time_interpolator = None + self.internal_boundaries = [] + + nodes = mesh.nodes + x_min, x_max = nodes[:, 0].min(), nodes[:, 0].max() + + self.first_dimension = "x" + self.first_dim_pts = np.linspace(x_min, x_max, self.N_VIS) + self.first_dim_size = self.N_VIS + + if mesh.dimension == 3: + y_min, y_max = nodes[:, 1].min(), nodes[:, 1].max() + z_min, z_max = nodes[:, 2].min(), nodes[:, 2].max() + self.second_dimension = "y" + self.second_dim_pts = np.linspace(y_min, y_max, self.N_VIS) + self.second_dim_size = self.N_VIS + self.third_dimension = "z" + self.third_dim_pts = np.linspace(z_min, z_max, self.N_VIS) + self.third_dim_size = self.N_VIS + self._slice_positions = { + "y": 0.5 * (y_min + y_max), + "z": 0.5 * (z_min + z_max), + } + else: + z_col = 1 + z_min, z_max = nodes[:, z_col].min(), nodes[:, z_col].max() + self.second_dimension = "z" + self.second_dim_pts = np.linspace(z_min, z_max, self.N_VIS) + self.second_dim_size = self.N_VIS + + def _shape(self, t): + return [self.mesh.npts, len(t)] + + def initialise(self): + if self.entries_raw_initialized: + return + self._entries_raw = self.observe_raw() + + from scipy.interpolate import interp1d + + self._time_interpolator = interp1d( + self.t_pts, + self._entries_raw, + kind="linear", + axis=1, + bounds_error=False, + fill_value="extrapolate", + ) + + def _interpolate_spatial(self, values, query_pts): + """Interpolate cell-centered data to query points. + + Uses linear interpolation inside the centroid convex hull and + nearest-neighbor extrapolation outside it. + """ + from scipy.interpolate import LinearNDInterpolator, NearestNDInterpolator + + pts = self.mesh.cell_centroids + + linear = LinearNDInterpolator(pts, values) + result = linear(query_pts) + + mask = np.isnan(result) + if np.any(mask): + nearest = NearestNDInterpolator(pts, values) + result[mask] = nearest(query_pts[mask]) + + return result + + def _data_at_time(self, t): + """Return cell-centered data at time t.""" + self.initialise() + t_observe, observe_raw = self._check_observe_raw(t) + if observe_raw: + return self._entries_raw + return self._time_interpolator(t_observe) + + def __call__( + self, t=None, x=None, r=None, y=None, z=None, R=None, fill_value=np.nan + ): + data_at_t = self._data_at_time(t) + scalar_t = isinstance(t, int | float) + + spatial_provided = any(c is not None for c in [x, y, z]) + if not spatial_provided: + return data_at_t + + if self.mesh.dimension == 2: + x_q = np.asarray(x).ravel() + z_q = np.asarray(z).ravel() if z is not None else np.zeros_like(x_q) + grid = np.meshgrid(x_q, z_q, indexing="ij") + query = np.column_stack([g.ravel() for g in grid]) + out_shape = grid[0].shape + else: + x_q = np.asarray(x).ravel() + y_q = np.asarray(y).ravel() if y is not None else np.zeros_like(x_q) + z_q = np.asarray(z).ravel() if z is not None else np.zeros_like(x_q) + grid = np.meshgrid(x_q, y_q, z_q, indexing="ij") + query = np.column_stack([g.ravel() for g in grid]) + out_shape = grid[0].shape + + n_t = data_at_t.shape[1] if data_at_t.ndim > 1 else 1 + if n_t == 1: + result = self._interpolate_spatial(data_at_t.ravel(), query).reshape( + out_shape + ) + else: + result = np.empty((*out_shape, n_t)) + for i in range(n_t): + result[..., i] = self._interpolate_spatial( + data_at_t[:, i], query + ).reshape(out_shape) + + if scalar_t and result.ndim > len(out_shape): + result = result[..., 0] + + return result + + def get_3d_slices(self, t): + """Compute two orthogonal slices through the 3D domain for plotting. + + Returns (slice_xz, xx_xz, yy_xz, zz_xz, + slice_xy, xx_xy, yy_xy, zz_xy) + where each slice is on a regular grid at the midplane. + """ + data_at_t = self._data_at_time(t) + vals = data_at_t.ravel() if data_at_t.ndim == 1 else data_at_t[:, -1] + + x_pts = self.first_dim_pts + y_pts = self.second_dim_pts + z_pts = self.third_dim_pts + y_mid = self._slice_positions["y"] + z_mid = self._slice_positions["z"] + + # x-z plane at y = y_mid + xx1, zz1 = np.meshgrid(x_pts, z_pts, indexing="ij") + yy1 = np.full_like(xx1, y_mid) + q1 = np.column_stack([xx1.ravel(), yy1.ravel(), zz1.ravel()]) + s1 = self._interpolate_spatial(vals, q1).reshape(xx1.shape) + + # x-y plane at z = z_mid + xx2, yy2 = np.meshgrid(x_pts, y_pts, indexing="ij") + zz2 = np.full_like(xx2, z_mid) + q2 = np.column_stack([xx2.ravel(), yy2.ravel(), zz2.ravel()]) + s2 = self._interpolate_spatial(vals, q2).reshape(xx2.shape) + + return s1, xx1, yy1, zz1, s2, xx2, yy2, zz2 + + class ProcessedVariableRawFVM(ProcessedVariable): def _shape(self, t): return [self.base_variables[0].size, len(t)] @@ -1369,6 +1547,9 @@ def process_variable(name: str, base_variables, *args, **kwargs): if mesh and hasattr(mesh, "edges_lr") and hasattr(mesh, "edges_tb"): return ProcessedVariable2DFVM(name, base_variables, *args, **kwargs) + if isinstance(mesh, pybamm.UnstructuredSubMesh): + return ProcessedVariableUnstructuredFVM(name, base_variables, *args, **kwargs) + # check variable shape if len(base_eval_shape) == 0 or base_eval_shape[0] == 1: return ProcessedVariable0D(name, base_variables, *args, **kwargs) From 0516e58fc102089d9539e14a1f6ecff8f479078e Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Wed, 25 Feb 2026 15:35:25 -0800 Subject: [PATCH 04/25] update vector field --- src/pybamm/__init__.py | 2 +- src/pybamm/discretisations/discretisation.py | 84 +++++++--- .../operations/convert_to_casadi.py | 5 + src/pybamm/expression_tree/unary_operators.py | 47 ++++++ src/pybamm/expression_tree/vector_field.py | 83 ++++++---- src/pybamm/plotting/quick_plot.py | 118 ++++++++++++- src/pybamm/solvers/processed_variable.py | 155 +++++++++++++++++- src/pybamm/solvers/solution.py | 34 ++-- .../finite_volume_unstructured.py | 141 ++++++++++++++-- 9 files changed, 585 insertions(+), 84 deletions(-) diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py index 2e7e34be32..8629c51e12 100644 --- a/src/pybamm/__init__.py +++ b/src/pybamm/__init__.py @@ -181,7 +181,7 @@ # Solver classes from .solvers.solution import Solution, EmptySolution, make_cycle_solution from .solvers.processed_variable_time_integral import ProcessedVariableTimeIntegral -from .solvers.processed_variable import ProcessedVariable, ProcessedVariable2DFVM, ProcessedVariableUnstructuredFVM, process_variable +from .solvers.processed_variable import ProcessedVariable, ProcessedVariable2DFVM, ProcessedVariableUnstructuredFVM, ProcessedVariableVectorFieldUnstructuredFVM, process_variable from .solvers.processed_variable_computed import ProcessedVariableComputed from .solvers.processed_variable import ProcessedVariableUnstructured from .solvers.summary_variable import SummaryVariables diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py index aa535ebb62..13d40a77c2 100644 --- a/src/pybamm/discretisations/discretisation.py +++ b/src/pybamm/discretisations/discretisation.py @@ -914,7 +914,11 @@ def process_symbol(self, symbol): # Assign mesh as an attribute to the processed variable if symbol.domain != []: - discretised_symbol.mesh = self.mesh[symbol.domain] + mesh_for_symbol = self.mesh[symbol.domain] + discretised_symbol.mesh = mesh_for_symbol + if isinstance(discretised_symbol, pybamm.VectorField): + for comp in discretised_symbol._components: + comp.mesh = mesh_for_symbol else: discretised_symbol.mesh = None @@ -962,29 +966,49 @@ def _process_symbol(self, symbol): or isinstance(left, pybamm.Gradient) ): right = pybamm.VectorField(right, right) + elif isinstance(spatial_method, pybamm.FiniteVolumeUnstructured): + dim = self.mesh[symbol.domain[0]].dimension + if isinstance(left, pybamm.Scalar) and ( + isinstance(right, pybamm.VectorField) + or isinstance(right, pybamm.Gradient) + ): + left = pybamm.VectorField(*[left] * dim) + elif isinstance(right, pybamm.Scalar) and ( + isinstance(left, pybamm.VectorField) + or isinstance(left, pybamm.Gradient) + ): + right = pybamm.VectorField(*[right] * dim) disc_left = self.process_symbol(left) disc_right = self.process_symbol(right) if symbol.domain == []: if isinstance(disc_left, pybamm.VectorField) or isinstance( disc_right, pybamm.VectorField ): + if isinstance(disc_left, pybamm.VectorField): + n = disc_left.n_components + else: + n = disc_right.n_components if not isinstance(disc_right, pybamm.VectorField): - disc_right = pybamm.VectorField(disc_right, disc_right) + disc_right = pybamm.VectorField(*[disc_right] * n) if not isinstance(disc_left, pybamm.VectorField): - disc_left = pybamm.VectorField(disc_left, disc_left) - else: # both are vector fields already - pass - disc_lr = pybamm.simplify_if_constant( - symbol.create_copy( - new_children=[disc_left.lr_field, disc_right.lr_field] - ) - ) - disc_tb = pybamm.simplify_if_constant( - symbol.create_copy( - new_children=[disc_left.tb_field, disc_right.tb_field] + disc_left = pybamm.VectorField(*[disc_left] * n) + new_comps = [ + pybamm.simplify_if_constant( + symbol.create_copy( + new_children=[ + disc_left._components[k], + disc_right._components[k], + ] + ) ) - ) - return pybamm.VectorField(disc_lr, disc_tb) + for k in range(n) + ] + result = pybamm.VectorField(*new_comps) + for src in (disc_left, disc_right): + if hasattr(src, "_disc_state_vector"): + result._disc_state_vector = src._disc_state_vector + break + return result return pybamm.simplify_if_constant( symbol.create_copy(new_children=[disc_left, disc_right]) @@ -1126,6 +1150,18 @@ def _process_symbol(self, symbol): elif isinstance(symbol, pybamm.NotConstant): # After discretisation, we can make the symbol constant return disc_child + elif isinstance(symbol, pybamm.Component): + if not isinstance(disc_child, pybamm.VectorField): + raise ValueError("Component can only be applied to a VectorField") + return disc_child._components[symbol.index] + elif isinstance(symbol, pybamm.Norm): + if not isinstance(disc_child, pybamm.VectorField): + raise ValueError("Norm can only be applied to a VectorField") + result = None + for comp in disc_child._components: + sq = comp**2 + result = sq if result is None else result + sq + return result**0.5 elif isinstance(symbol, pybamm.Magnitude): if not isinstance(disc_child, pybamm.VectorField): raise ValueError("Magnitude can only be applied to a vector field") @@ -1138,10 +1174,14 @@ def _process_symbol(self, symbol): raise ValueError("Invalid direction") else: if isinstance(disc_child, pybamm.VectorField): - return pybamm.VectorField( - symbol.create_copy(new_children=[disc_child.lr_field]), - symbol.create_copy(new_children=[disc_child.tb_field]), - ) + new_comps = [ + symbol.create_copy(new_children=[c]) + for c in disc_child._components + ] + result = pybamm.VectorField(*new_comps) + if hasattr(disc_child, "_disc_state_vector"): + result._disc_state_vector = disc_child._disc_state_vector + return result else: return symbol.create_copy(new_children=[disc_child]) @@ -1215,10 +1255,8 @@ def _process_symbol(self, symbol): ) elif isinstance(symbol, pybamm.VectorField): - # VectorField is a subclass of TensorField, handle it first for specificity - left_symbol = self.process_symbol(symbol.lr_field) - right_symbol = self.process_symbol(symbol.tb_field) - return symbol.create_copy(new_children=[left_symbol, right_symbol]) + processed = [self.process_symbol(c) for c in symbol._components] + return symbol.create_copy(new_children=processed) elif isinstance(symbol, pybamm.TensorField): # General TensorField handling (rank-2 tensors) diff --git a/src/pybamm/expression_tree/operations/convert_to_casadi.py b/src/pybamm/expression_tree/operations/convert_to_casadi.py index f34324e883..7bc3e2728b 100644 --- a/src/pybamm/expression_tree/operations/convert_to_casadi.py +++ b/src/pybamm/expression_tree/operations/convert_to_casadi.py @@ -318,6 +318,11 @@ def hermite_poly(i): ) return casadi.vertcat(*all_child_vectors) + elif isinstance(symbol, pybamm.VectorField): + return casadi.vertcat( + *[self.convert(c, t, y, y_dot, inputs) for c in symbol._components] + ) + else: raise TypeError( f"Cannot convert symbol of type '{type(symbol)}' to CasADi. Symbols must all be " diff --git a/src/pybamm/expression_tree/unary_operators.py b/src/pybamm/expression_tree/unary_operators.py index d11bca032d..07e94d8439 100644 --- a/src/pybamm/expression_tree/unary_operators.py +++ b/src/pybamm/expression_tree/unary_operators.py @@ -1393,6 +1393,43 @@ def _unary_new_copy(self, child, perform_simplifications=True): return self.__class__(child, self.direction) +class Component(UnaryOperator): + """ + Extract component *index* from a VectorField. + + Parameters + ---------- + child : :class:`pybamm.Symbol` + A VectorField symbol. + index : int + Zero-based component index. + """ + + def __init__(self, child, index): + super().__init__(f"component({index})", child) + self.index = index + + def _unary_new_copy(self, child, perform_simplifications=True): + return self.__class__(child, self.index) + + +class Norm(UnaryOperator): + """ + Euclidean norm of a VectorField: ``sqrt(sum(comp_i ** 2))``. + + Parameters + ---------- + child : :class:`pybamm.Symbol` + A VectorField symbol. + """ + + def __init__(self, child): + super().__init__("norm", child) + + def _unary_new_copy(self, child, perform_simplifications=True): + return self.__class__(child) + + class Upwind(UpwindDownwind): """ Upwinding operator. To be used if flow velocity is positive (left to right). @@ -1647,6 +1684,16 @@ def sign(symbol): return pybamm.simplify_if_constant(Sign(symbol)) +def component(symbol, index): + """Convenience function for creating a :class:`Component`.""" + return Component(symbol, index) + + +def norm(symbol): + """Convenience function for creating a :class:`Norm`.""" + return Norm(symbol) + + def smooth_absolute_value(symbol, k): """ Smooth approximation to the absolute value function. k is the smoothing parameter, diff --git a/src/pybamm/expression_tree/vector_field.py b/src/pybamm/expression_tree/vector_field.py index 1adc536315..7e99a02b8d 100644 --- a/src/pybamm/expression_tree/vector_field.py +++ b/src/pybamm/expression_tree/vector_field.py @@ -1,5 +1,5 @@ """ -VectorField class - a rank-1 tensor field for 2D simulations. +VectorField class - a rank-1 tensor field with N components. """ from __future__ import annotations @@ -13,62 +13,77 @@ class VectorField(TensorField): A node in the expression tree representing a vector field. VectorField is a convenience subclass of TensorField for rank-1 tensors - with two components (lr and tb directions in 2D). + with N >= 2 components. Components are stored by integer index; the + properties ``lr_field``, ``tb_field``, and ``fb_field`` are backward- + compatible aliases for ``[0]``, ``[1]``, and ``[2]``. Parameters ---------- - lr_field : pybamm.Symbol - The left-right (x) component of the vector field. - tb_field : pybamm.Symbol - The top-bottom (y) component of the vector field. + *components : pybamm.Symbol + Two or more component symbols, all sharing the same domain. """ - def __init__(self, lr_field, tb_field): - if lr_field.domain != tb_field.domain: - raise ValueError("lr_field and tb_field must have the same domain") - # Initialize as a rank-1 TensorField with two components - super().__init__([lr_field, tb_field], domain=lr_field.domain) - # Override the name to maintain backward compatibility + def __init__(self, *components): + if len(components) < 2: + raise ValueError( + f"VectorField requires at least 2 components, got {len(components)}" + ) + ref_domain = components[0].domain + for i, c in enumerate(components[1:], start=1): + if c.domain != ref_domain: + raise ValueError( + f"All components must have the same domain: " + f"component {i} has {c.domain}, expected {ref_domain}" + ) + super().__init__(list(components), domain=ref_domain) self.name = "vector_field" + @property + def n_components(self): + """Number of vector components.""" + return len(self._components) + + # ---- backward-compatible aliases for structured-grid directions ---- + @property def lr_field(self): - """The left-right (x) component of the vector field.""" + """Component 0 (left-right / x).""" return self._components[0] @property def tb_field(self): - """The top-bottom (y) component of the vector field.""" + """Component 1 (top-bottom / y).""" return self._components[1] + @property + def fb_field(self): + """Component 2 (front-back / z). Only valid for 3-component fields.""" + if len(self._components) < 3: + raise AttributeError( + "fb_field requires at least 3 components; this VectorField has " + f"{len(self._components)}" + ) + return self._components[2] + def create_copy( self, new_children: list[pybamm.Symbol] | None = None, perform_simplifications: bool = True, ): - """Create a copy of this vector field with optional new children.""" if new_children is None: new_children = [ - self.lr_field.create_copy( - perform_simplifications=perform_simplifications - ), - self.tb_field.create_copy( - perform_simplifications=perform_simplifications - ), + c.create_copy(perform_simplifications=perform_simplifications) + for c in self._components ] return VectorField(*new_children) def evaluates_on_edges(self, dimension: str) -> bool: - """Check if components evaluate on edges. - - Overrides TensorField to provide more specific error message. - """ - left_evaluates_on_edges = self.lr_field.evaluates_on_edges(dimension) - right_evaluates_on_edges = self.tb_field.evaluates_on_edges(dimension) - if left_evaluates_on_edges == right_evaluates_on_edges: - return left_evaluates_on_edges - else: - raise ValueError( - "lr_field and tb_field must either both evaluate on edges " - "or both not evaluate on edges" - ) + statuses = [c.evaluates_on_edges(dimension) for c in self._components] + if all(statuses): + return True + if not any(statuses): + return False + raise ValueError( + "All VectorField components must either all evaluate on edges " + "or none evaluate on edges" + ) diff --git a/src/pybamm/plotting/quick_plot.py b/src/pybamm/plotting/quick_plot.py index be600241ab..31c7423716 100644 --- a/src/pybamm/plotting/quick_plot.py +++ b/src/pybamm/plotting/quick_plot.py @@ -290,6 +290,7 @@ def set_output_variables(self, output_variables, solutions): self.second_spatial_variable = {} self.x_first_and_y_second = {} self.is_y_z = {} + self.is_vector_field = {} # Calculate subplot positions based on number of variables supplied self.subplot_positions = {} @@ -384,6 +385,9 @@ def set_output_variables(self, output_variables, solutions): # Store variables and subplot position self.variables[variable_tuple] = variables + self.is_vector_field[variable_tuple] = getattr( + first_variable, "is_vector_field", False + ) self.subplot_positions[variable_tuple] = (self.n_rows, self.n_cols, k + 1) def _get_spatial_var(self, key, variable, dimension): @@ -406,6 +410,7 @@ def _get_spatial_var(self, key, variable, dimension): ( pybamm.ProcessedVariable2DFVM, pybamm.ProcessedVariableUnstructuredFVM, + pybamm.ProcessedVariableVectorFieldUnstructuredFVM, ), ): domain = variable.domain[0] @@ -452,7 +457,9 @@ def reset_axis(self): self.axis_limits[key] = [x_min, x_max, y_min, y_max] # Get min and max variable values - if self.variable_limits[key] == "fixed": + if self.is_vector_field.get(key, False): + var_min, var_max = None, None + elif self.variable_limits[key] == "fixed": # fixed variable limits: calculate "globlal" min and max if variable_lists[0][0].dimensions == 3: var_min = np.min( @@ -616,6 +623,39 @@ def plot(self, t, dynamic=False): for boundary in variable_lists[0][0].internal_boundaries: boundary_scaled = boundary * self.spatial_factor ax.axvline(boundary_scaled, color="0.5", lw=1, zorder=0) + elif self.is_vector_field.get(key, False): + variable = variable_lists[0][0] + if variable.dimensions == 2: + X, Z, U, W = variable.get_quiver_data(t_in_seconds) + Xs = X * self.spatial_factor + Zs = Z * self.spatial_factor + mag = np.sqrt(U**2 + W**2) + mag_max = np.max(mag) if np.max(mag) > 0 else 1.0 + norm = colors.Normalize(vmin=0, vmax=mag_max) + safe_mag = np.where(mag > 0, mag, 1.0) + U_norm = U / safe_mag + W_norm = W / safe_mag + ax.set_xlabel(f"x [{self.spatial_unit}]") + ax.set_ylabel(f"z [{self.spatial_unit}]") + self.plots[key][0][0] = ax.quiver( + Xs, + Zs, + U_norm, + W_norm, + mag, + cmap="viridis", + norm=norm, + scale=X.shape[0] * 1.2, + scale_units="width", + width=0.004, + ) + self.colorbars[key] = self.fig.colorbar( + self.plots[key][0][0], + ax=ax, + label="|" + str(key[0]) + "|", + ) + else: + self._plot_3d_quiver(ax, variable, t_in_seconds, key, cm, colors) elif variable_lists[0][0].dimensions == 2: # Read dictionary of spatial variables spatial_vars = self.spatial_variable_dict[key] @@ -743,6 +783,49 @@ def plot(self, t, dynamic=False): bottom = max(legend_top, slider_top) self.gridspec.tight_layout(self.fig, rect=[0, bottom, 1, 1]) + def _plot_3d_quiver(self, ax, variable, t, key, cm, colors): + """Render quiver arrows on two orthogonal 3D slice planes.""" + sf = self.spatial_factor + data = variable.get_quiver_data(t) + X1, Z1, U_xz, W_xz, y_mid = data[0:5] + X2, Y2, U_xy, V_xy, z_mid = data[5:10] + + x_span = (X1.max() - X1.min()) * sf + arrow_len = x_span * 0.08 if x_span > 0 else 0.08 + + Y1_plane = np.full_like(X1, y_mid * sf) + ax.quiver( + X1 * sf, + Y1_plane, + Z1 * sf, + U_xz, + np.zeros_like(U_xz), + W_xz, + length=arrow_len, + normalize=True, + color="steelblue", + alpha=0.8, + ) + + Z2_plane = np.full_like(X2, z_mid * sf) + ax.quiver( + X2 * sf, + Y2 * sf, + Z2_plane, + U_xy, + V_xy, + np.zeros_like(U_xy), + length=arrow_len, + normalize=True, + color="darkorange", + alpha=0.8, + ) + + ax.set_xlabel(f"$x$ [{self.spatial_unit}]") + ax.set_ylabel(f"$y$ [{self.spatial_unit}]") + ax.set_zlabel(f"$z$ [{self.spatial_unit}]") + self.plots[key][0][0] = "quiver_3d" + def dynamic_plot(self, show_plot=True, step=None): """ Generate a dynamic plot with a slider to control the time. @@ -866,6 +949,39 @@ def slider_update(self, t): y_min, y_max = self.axis_limits[key][2:] if y_min is None and y_max is None: ax.set_ylim(var_min, var_max) + elif self.is_vector_field.get(key, False): + variable = self.variables[key][0][0] + ax.clear() + if variable.dimensions == 2: + X, Z, U, W = variable.get_quiver_data(time_in_seconds) + Xs = X * self.spatial_factor + Zs = Z * self.spatial_factor + mag = np.sqrt(U**2 + W**2) + mag_max = np.max(mag) if np.max(mag) > 0 else 1.0 + norm = colors.Normalize(vmin=0, vmax=mag_max) + safe_mag = np.where(mag > 0, mag, 1.0) + U_norm = U / safe_mag + W_norm = W / safe_mag + ax.set_xlabel(f"x [{self.spatial_unit}]") + ax.set_ylabel(f"z [{self.spatial_unit}]") + self.plots[key][0][0] = ax.quiver( + Xs, + Zs, + U_norm, + W_norm, + mag, + cmap="viridis", + norm=norm, + scale=X.shape[0] * 1.2, + scale_units="width", + width=0.004, + ) + if key in self.colorbars: + self.colorbars[key].update_normal(self.plots[key][0][0]) + else: + self._plot_3d_quiver(ax, variable, time_in_seconds, key, cm, colors) + title = split_long_string(key[0]) if len(key) == 1 else "" + ax.set_title(title, fontsize="medium") elif self.variables[key][0][0].dimensions == 2: # 2D plot: plot as a function of x and y at time t # Read dictionary of spatial variables diff --git a/src/pybamm/solvers/processed_variable.py b/src/pybamm/solvers/processed_variable.py index af0acbc7f5..2f0adc6f9a 100644 --- a/src/pybamm/solvers/processed_variable.py +++ b/src/pybamm/solvers/processed_variable.py @@ -1025,22 +1025,35 @@ def initialise(self): fill_value="extrapolate", ) + def _augmented_points_and_values(self, values): + """Append boundary face centroids (with owning-cell values) to the + cell centroid cloud so that LinearNDInterpolator's convex hull + reaches the true mesh boundary.""" + mesh = self.mesh + bnd_start = mesh._boundary_face_start + bnd_centroids = mesh.face_centroids[bnd_start:] + bnd_owners = mesh.face_owner[bnd_start:] + pts = np.concatenate([mesh.cell_centroids, bnd_centroids], axis=0) + vals = np.concatenate([values, values[bnd_owners]]) + return pts, vals + def _interpolate_spatial(self, values, query_pts): """Interpolate cell-centered data to query points. - Uses linear interpolation inside the centroid convex hull and - nearest-neighbor extrapolation outside it. + Boundary face centroids are added to the interpolation cloud + so the convex hull covers the full domain. Any residual + extrapolation uses nearest-neighbor. """ from scipy.interpolate import LinearNDInterpolator, NearestNDInterpolator - pts = self.mesh.cell_centroids + pts, vals = self._augmented_points_and_values(values) - linear = LinearNDInterpolator(pts, values) + linear = LinearNDInterpolator(pts, vals) result = linear(query_pts) mask = np.isnan(result) if np.any(mask): - nearest = NearestNDInterpolator(pts, values) + nearest = NearestNDInterpolator(pts, vals) result[mask] = nearest(query_pts[mask]) return result @@ -1125,6 +1138,134 @@ def get_3d_slices(self, t): return s1, xx1, yy1, zz1, s2, xx2, yy2, zz2 +class ProcessedVariableVectorFieldUnstructuredFVM: + """ + Processed variable for a VectorField on an unstructured mesh. + + Wraps N scalar ``ProcessedVariableUnstructuredFVM`` instances (one per + component) and provides a unified interface for querying and plotting + vector-valued data. + """ + + N_QUIVER = 20 + + def __init__( + self, + name: str, + base_variables, + base_variables_casadi, + solution, + time_integral=None, + ): + vf = base_variables[0] + + self.name = name + self.mesh = vf.mesh + self.domain = vf.domain + self.is_vector_field = True + self.n_components = vf.n_components + self.internal_boundaries = [] + + self._component_vars = [] + for k in range(vf.n_components): + comp_base_k = [bv._components[k] for bv in base_variables] + comp_casadi_k = [] + for bvc in base_variables_casadi: + if isinstance(bvc, list): + comp_casadi_k.append(bvc[k]) + elif isinstance(bvc, pybamm.VectorField): + comp_casadi_k.append(bvc._components[k]) + else: + comp_casadi_k.append(bvc) + pv = ProcessedVariableUnstructuredFVM( + f"{name}[{k}]", + comp_base_k, + comp_casadi_k, + solution, + time_integral=time_integral, + ) + self._component_vars.append(pv) + + ref = self._component_vars[0] + self.dimensions = ref.dimensions + self.first_dimension = ref.first_dimension + self.first_dim_pts = ref.first_dim_pts + self.first_dim_size = ref.first_dim_size + self.second_dimension = ref.second_dimension + self.second_dim_pts = ref.second_dim_pts + self.second_dim_size = ref.second_dim_size + + if self.dimensions == 3: + self.third_dimension = ref.third_dimension + self.third_dim_pts = ref.third_dim_pts + self.third_dim_size = ref.third_dim_size + self._slice_positions = ref._slice_positions + + @property + def entries(self): + return self._component_vars[0].entries + + def __call__( + self, t=None, x=None, r=None, y=None, z=None, R=None, fill_value=np.nan + ): + """Return a tuple of arrays, one per component.""" + return tuple( + pv(t=t, x=x, r=r, y=y, z=z, R=R, fill_value=fill_value) + for pv in self._component_vars + ) + + def get_quiver_data(self, t): + """Interpolate vector components onto a coarser grid for quiver arrows. + + Returns ``(X, Z, U, W)`` for 2D or ``(X, Y, Z, U, V, W)`` for 3D. + """ + nq = self.N_QUIVER + x_pts = np.linspace(self.first_dim_pts[0], self.first_dim_pts[-1], nq) + + if self.dimensions == 2: + z_pts = np.linspace(self.second_dim_pts[0], self.second_dim_pts[-1], nq) + comp_u = self._component_vars[0](t=t, x=x_pts, z=z_pts) + comp_w = self._component_vars[1](t=t, x=x_pts, z=z_pts) + X, Z = np.meshgrid(x_pts, z_pts, indexing="ij") + return X, Z, comp_u, comp_w + else: + y_pts = np.linspace(self.second_dim_pts[0], self.second_dim_pts[-1], nq) + z_pts = np.linspace(self.third_dim_pts[0], self.third_dim_pts[-1], nq) + y_mid = self._slice_positions["y"] + z_mid = self._slice_positions["z"] + + # x-z plane at y_mid + comp_u_xz = self._component_vars[0]( + t=t, x=x_pts, y=np.array([y_mid]), z=z_pts + ).squeeze(axis=1) + comp_w_xz = self._component_vars[2]( + t=t, x=x_pts, y=np.array([y_mid]), z=z_pts + ).squeeze(axis=1) + X1, Z1 = np.meshgrid(x_pts, z_pts, indexing="ij") + + # x-y plane at z_mid + comp_u_xy = self._component_vars[0]( + t=t, x=x_pts, y=y_pts, z=np.array([z_mid]) + ).squeeze(axis=2) + comp_v_xy = self._component_vars[1]( + t=t, x=x_pts, y=y_pts, z=np.array([z_mid]) + ).squeeze(axis=2) + X2, Y2 = np.meshgrid(x_pts, y_pts, indexing="ij") + + return ( + X1, + Z1, + comp_u_xz, + comp_w_xz, + y_mid, + X2, + Y2, + comp_u_xy, + comp_v_xy, + z_mid, + ) + + class ProcessedVariableRawFVM(ProcessedVariable): def _shape(self, t): return [self.base_variables[0].size, len(t)] @@ -1548,6 +1689,10 @@ def process_variable(name: str, base_variables, *args, **kwargs): return ProcessedVariable2DFVM(name, base_variables, *args, **kwargs) if isinstance(mesh, pybamm.UnstructuredSubMesh): + if isinstance(base_variables[0], pybamm.VectorField): + return ProcessedVariableVectorFieldUnstructuredFVM( + name, base_variables, *args, **kwargs + ) return ProcessedVariableUnstructuredFVM(name, base_variables, *args, **kwargs) # check variable shape diff --git a/src/pybamm/solvers/solution.py b/src/pybamm/solvers/solution.py index 4c92a4216c..5187799f53 100644 --- a/src/pybamm/solvers/solution.py +++ b/src/pybamm/solvers/solution.py @@ -494,16 +494,30 @@ def _update_variable(self, name: str): "solve. Please re-run the solve with `output_variables` set to " "include this variable." ) - var_casadi, var_pybamm, time_integral = self._update_model_variable( - model, - _var_pybamm, - inputs=inputs, - ys_shape=ys.shape, - time_integral=time_integral, - cache_key=name, - ) - vars_pybamm[i] = var_pybamm - vars_casadi[i] = var_casadi + if isinstance(_var_pybamm, pybamm.VectorField): + comp_casadi = [] + for k, comp in enumerate(_var_pybamm._components): + cc, _, _ = self._update_model_variable( + model, + comp, + inputs=inputs, + ys_shape=ys.shape, + time_integral=None, + cache_key=f"{name}[{k}]", + ) + comp_casadi.append(cc) + vars_casadi[i] = comp_casadi + else: + var_casadi, var_pybamm, time_integral = self._update_model_variable( + model, + _var_pybamm, + inputs=inputs, + ys_shape=ys.shape, + time_integral=time_integral, + cache_key=name, + ) + vars_pybamm[i] = var_pybamm + vars_casadi[i] = var_casadi var = pybamm.process_variable( name, vars_pybamm, vars_casadi, self, time_integral=time_integral ) diff --git a/src/pybamm/spatial_methods/finite_volume_unstructured.py b/src/pybamm/spatial_methods/finite_volume_unstructured.py index 330de5201c..7af5ad8a9d 100644 --- a/src/pybamm/spatial_methods/finite_volume_unstructured.py +++ b/src/pybamm/spatial_methods/finite_volume_unstructured.py @@ -236,16 +236,18 @@ def gradient(self, symbol, discretised_symbol, boundary_conditions): submesh, G_components, bc_vecs, bcs ) - results = [] + components = [] for k in range(d): Gk = csr_matrix(kron(eye(repeats, dtype=np.float64), G_components[k])) - comp = Gk @ discretised_symbol + comp = pybamm.Matrix(Gk) @ discretised_symbol if np.any(bc_vecs[k] != 0): bc_full = np.tile(bc_vecs[k], repeats) comp = comp + pybamm.Vector(bc_full) - results.append(comp) + components.append(comp) - return results + vf = pybamm.VectorField(*components) + vf._disc_state_vector = discretised_symbol + return vf def _green_gauss_matrices(self, submesh): """ @@ -401,10 +403,14 @@ def divergence(self, symbol, discretised_symbol, boundary_conditions): d = submesh.dimension repeats = self._get_auxiliary_domain_repeats(symbol.domains) - if not isinstance(discretised_symbol, (list, tuple)): + if isinstance(discretised_symbol, pybamm.VectorField): + comps = discretised_symbol._components + elif isinstance(discretised_symbol, (list, tuple)): + comps = list(discretised_symbol) + else: raise TypeError( - "FiniteVolumeUnstructured.divergence expects a list of " - f"{d} component arrays, got {type(discretised_symbol)}" + "FiniteVolumeUnstructured.divergence expects a VectorField or " + f"list of {d} component arrays, got {type(discretised_symbol)}" ) D_components = self._divergence_matrices(submesh) @@ -412,10 +418,90 @@ def divergence(self, symbol, discretised_symbol, boundary_conditions): result = pybamm.Vector(np.zeros(n * repeats)) for k in range(d): Dk = csr_matrix(kron(eye(repeats, dtype=np.float64), D_components[k])) - result = result + Dk @ discretised_symbol[k] + result = result + pybamm.Matrix(Dk) @ comps[k] + + disc_sv = getattr(discretised_symbol, "_disc_state_vector", None) + if disc_sv is not None: + L_bc, bc_rhs, D_bnd = self._div_boundary_correction( + submesh, boundary_conditions, domain=domain + ) + if D_bnd is not None: + for k in range(d): + Dk_bnd = csr_matrix(kron(eye(repeats, dtype=np.float64), D_bnd[k])) + result = result - pybamm.Matrix(Dk_bnd) @ comps[k] + if L_bc is not None: + L_bc_full = csr_matrix(kron(eye(repeats, dtype=np.float64), L_bc)) + result = result + pybamm.Matrix(L_bc_full) @ disc_sv + if np.any(bc_rhs != 0): + bc_rhs_full = np.tile(bc_rhs, repeats) + result = result + pybamm.Vector(bc_rhs_full) return result + def _div_boundary_correction(self, submesh, boundary_conditions, domain=None): + """Build boundary corrections for the divergence operator. + + When computing ``div(D * grad(u))``, the divergence matrices use + cell-centered flux values at boundary faces, which is incorrect. + This method returns: + + * ``L_bc`` – sparse matrix for TPFA Dirichlet correction on state vector + * ``bc_rhs`` – constant vector for Dirichlet/Neumann RHS + * ``D_bnd`` – list of sparse matrices (boundary-only divergence terms + to subtract from the cell-centered approximation) + + The corrected divergence is:: + + div(F) = sum_k D_k @ F_k - sum_k D_bnd_k @ F_k + L_bc @ u + bc_rhs + """ + n = submesh.npts + d = submesh.dimension + L_bc = None + bc_rhs = np.zeros(n) + D_bnd = None + + for var, bcs in boundary_conditions.items(): + if not hasattr(var, "domain"): + continue + if domain is not None and var.domain != domain: + continue + for side, (bc_value, bc_type) in bcs.items(): + face_tag = self._side_to_boundary_tag(side) + if face_tag not in submesh.boundary_faces: + continue + face_indices = submesh.boundary_faces[face_tag] + bc_val = float(bc_value.evaluate()) + + for fi in face_indices: + cell = submesh.face_owner[fi] + area = submesh.face_areas[fi] + vol = submesh.cell_volumes[cell] + normal = submesh.face_normals[fi] + + if D_bnd is None: + D_bnd = [csr_matrix((n, n)).tolil() for _ in range(d)] + for k in range(d): + D_bnd[k][cell, cell] += normal[k] * area / vol + + if bc_type == "Dirichlet": + face_c = submesh.face_centroids[fi] + cell_c = submesh.cell_centroids[cell] + d_perp = np.linalg.norm(face_c - cell_c) + coeff = area / d_perp + + if L_bc is None: + L_bc = csr_matrix((n, n)).tolil() + L_bc[cell, cell] -= coeff / vol + bc_rhs[cell] += coeff * bc_val / vol + elif bc_type == "Neumann": + bc_rhs[cell] += bc_val * area / vol + + if L_bc is not None: + L_bc = csr_matrix(L_bc) + if D_bnd is not None: + D_bnd = [csr_matrix(m) for m in D_bnd] + return L_bc, bc_rhs, D_bnd + def _divergence_matrices(self, submesh): """ Build divergence matrices D_k for k = 0..d-1. @@ -440,7 +526,6 @@ def _divergence_matrices(self, submesh): D = [csr_matrix((n, n)) for _ in range(d)] - # Internal faces int_owner = owner[:n_int] int_neighbor = neighbor[:n_int] @@ -491,11 +576,47 @@ def _divergence_matrices(self, submesh): def gradient_squared(self, symbol, discretised_symbol, boundary_conditions): grad = self.gradient(symbol, discretised_symbol, boundary_conditions) result = None - for comp in grad: + for comp in grad._components: sq = comp**2 result = sq if result is None else result + sq return result + # ------------------------------------------------------------------ + # Binary operator handling (scalar * VectorField, etc.) + # ------------------------------------------------------------------ + + def process_binary_operators(self, bin_op, left, right, disc_left, disc_right): + if isinstance(disc_left, pybamm.VectorField) or isinstance( + disc_right, pybamm.VectorField + ): + if isinstance(disc_left, pybamm.VectorField) and isinstance( + disc_right, pybamm.VectorField + ): + n = disc_left.n_components + elif isinstance(disc_left, pybamm.VectorField): + n = disc_left.n_components + disc_right = pybamm.VectorField(*[disc_right] * n) + else: + n = disc_right.n_components + disc_left = pybamm.VectorField(*[disc_left] * n) + + new_comps = [ + pybamm.simplify_if_constant( + bin_op.create_copy( + [disc_left._components[k], disc_right._components[k]] + ) + ) + for k in range(n) + ] + result = pybamm.VectorField(*new_comps) + for src in (disc_left, disc_right): + if hasattr(src, "_disc_state_vector"): + result._disc_state_vector = src._disc_state_vector + break + return result + + return bin_op._binary_new_copy(disc_left, disc_right) + # ------------------------------------------------------------------ # Integral operators # ------------------------------------------------------------------ From 86875ccf8892d527e54f3460c6d582095d5e165c Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Wed, 25 Feb 2026 17:33:51 -0800 Subject: [PATCH 05/25] add basicdfn2dunstructured --- src/pybamm/discretisations/discretisation.py | 20 + src/pybamm/meshes/meshes.py | 77 +++- src/pybamm/meshes/unstructured_submesh.py | 2 + .../lithium_ion/__init__.py | 1 + .../lithium_ion/basic_dfn_2d_unstructured.py | 419 +++++++++++++++++ .../finite_volume_unstructured.py | 423 +++++++++++++----- .../test_lithium_ion/test_basic_models.py | 27 ++ .../test_lithium_ion/test_basic_models.py | 4 + 8 files changed, 832 insertions(+), 141 deletions(-) create mode 100644 src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_2d_unstructured.py diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py index 13d40a77c2..1e12ae73eb 100644 --- a/src/pybamm/discretisations/discretisation.py +++ b/src/pybamm/discretisations/discretisation.py @@ -1044,6 +1044,26 @@ def _process_symbol(self, symbol): return child_spatial_method.gradient(child, disc_child, self.bcs) elif isinstance(symbol, pybamm.Divergence): + if isinstance( + child_spatial_method, pybamm.FiniteVolumeUnstructured + ) and isinstance(child, pybamm.Multiplication): + left_c, right_c = child.children + grad_sym = None + coeff_sym = None + if isinstance(right_c, pybamm.Gradient): + grad_sym, coeff_sym = right_c, left_c + elif isinstance(left_c, pybamm.Gradient): + grad_sym, coeff_sym = left_c, right_c + if grad_sym is not None: + disc_coeff = self.process_symbol(coeff_sym) + disc_u = self.process_symbol(grad_sym.child) + return child_spatial_method.div_D_grad( + symbol, + grad_sym.child, + disc_coeff, + disc_u, + self.bcs, + ) return child_spatial_method.divergence(child, disc_child, self.bcs) elif isinstance(symbol, pybamm.Laplacian): diff --git a/src/pybamm/meshes/meshes.py b/src/pybamm/meshes/meshes.py index 7d95b910b9..5ad71018cb 100644 --- a/src/pybamm/meshes/meshes.py +++ b/src/pybamm/meshes/meshes.py @@ -465,26 +465,69 @@ def to_json(self): def _combine_unstructured_submeshes(submeshes): """ Create a lightweight combined mesh from a list of - :class:`UnstructuredSubMesh` objects. The combined mesh merges - nodes and elements (re-indexed) and sums ``npts``. + :class:`UnstructuredSubMesh` objects. Coincident boundary nodes + at domain interfaces are merged so that face-connectivity spans + across domains. """ - all_nodes = [] - all_elements = [] - node_offset = 0 - total_npts = 0 - - for sm in submeshes: - all_nodes.append(sm.nodes) - all_elements.append(sm.elements + node_offset) - node_offset += sm.nodes.shape[0] - total_npts += sm.npts - - combined_nodes = np.concatenate(all_nodes, axis=0) + tol = 1e-12 + all_nodes = list(submeshes[0].nodes) + global_maps = [{i: i for i in range(submeshes[0].nodes.shape[0])}] + next_id = len(all_nodes) + + for k in range(1, len(submeshes)): + prev = submeshes[k - 1] + curr = submeshes[k] + prev_map = global_maps[k - 1] + + right_global = {} + if "right" in prev.boundary_faces: + right_node_ids = set() + for fi in prev.boundary_faces["right"]: + right_node_ids.update(prev.faces[fi].tolist()) + for nid in right_node_ids: + right_global[prev_map[nid]] = prev.nodes[nid] + + left_local = set() + if "left" in curr.boundary_faces: + for fi in curr.boundary_faces["left"]: + left_local.update(curr.faces[fi].tolist()) + + local_to_global = {} + for nid in range(curr.nodes.shape[0]): + if nid in left_local and right_global: + pos = curr.nodes[nid] + matched = False + for gid, rpos in right_global.items(): + if np.linalg.norm(pos - rpos) < tol: + local_to_global[nid] = gid + matched = True + break + if not matched: + local_to_global[nid] = next_id + all_nodes.append(curr.nodes[nid]) + next_id += 1 + else: + local_to_global[nid] = next_id + all_nodes.append(curr.nodes[nid]) + next_id += 1 + global_maps.append(local_to_global) + + all_elements = [submeshes[0].elements.copy()] + for k in range(1, len(submeshes)): + gm = global_maps[k] + remapped = np.array( + [[gm[v] for v in row] for row in submeshes[k].elements], + dtype=int, + ) + all_elements.append(remapped) + + combined_nodes = np.array(all_nodes) combined_elements = np.concatenate(all_elements, axis=0) - combined = pybamm.UnstructuredSubMesh( - combined_nodes, combined_elements, coord_sys=submeshes[0].coord_sys + return pybamm.UnstructuredSubMesh( + combined_nodes, + combined_elements, + coord_sys=submeshes[0].coord_sys, ) - return combined class SubMesh: diff --git a/src/pybamm/meshes/unstructured_submesh.py b/src/pybamm/meshes/unstructured_submesh.py index 76e13c567b..394a62f923 100644 --- a/src/pybamm/meshes/unstructured_submesh.py +++ b/src/pybamm/meshes/unstructured_submesh.py @@ -60,6 +60,8 @@ def __init__(self, nodes, elements, coord_sys="cartesian", boundary_faces=None): self._identify_boundary_faces() self.npts = len(self.elements) + self.npts_lr = self.npts + self.npts_tb = 1 self.internal_boundaries = [] self.interface_data = {} diff --git a/src/pybamm/models/full_battery_models/lithium_ion/__init__.py b/src/pybamm/models/full_battery_models/lithium_ion/__init__.py index b8f106be7e..e003594858 100644 --- a/src/pybamm/models/full_battery_models/lithium_ion/__init__.py +++ b/src/pybamm/models/full_battery_models/lithium_ion/__init__.py @@ -24,6 +24,7 @@ from .newman_tobias import NewmanTobias from .basic_dfn import BasicDFN from .basic_dfn_2d import BasicDFN2D +from .basic_dfn_2d_unstructured import BasicDFN2DUnstructured from .basic_spm import BasicSPM from .basic_spm_with_3d_thermal import Basic3DThermalSPM from .basic_dfn_half_cell import BasicDFNHalfCell diff --git a/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_2d_unstructured.py b/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_2d_unstructured.py new file mode 100644 index 0000000000..de843108d9 --- /dev/null +++ b/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_2d_unstructured.py @@ -0,0 +1,419 @@ +# +# Basic Doyle-Fuller-Newman (DFN) Model — 2D Unstructured FVM +# +import pybamm +from pybamm.models.full_battery_models.lithium_ion.base_lithium_ion_model import ( + BaseModel, +) + + +class BasicDFN2DUnstructured(BaseModel): + """Doyle-Fuller-Newman (DFN) model on a 2D unstructured mesh. + + Identical physics to :class:`BasicDFN2D` but uses + :class:`~pybamm.FiniteVolumeUnstructured` on triangle or quad elements + instead of the structured tensor-product grid. + + Parameters + ---------- + name : str, optional + The name of the model. + element_type : str, optional + Element type for the built-in mesh generator: ``"quad"`` (default, + TPFA-orthogonal) or ``"triangle"``. + """ + + def __init__( + self, + name="Doyle-Fuller-Newman model (2D unstructured)", + element_type="quad", + ): + super().__init__(name=name) + self._element_type = element_type + pybamm.citations.register("Marquis2019") + + ###################### + # Variables + ###################### + Q = pybamm.Variable("Discharge capacity [A.h]") + + x = pybamm.SpatialVariable( + "x", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="lr", + ) + x_n = pybamm.SpatialVariable( + "x_n", domain="negative electrode", coord_sys="cartesian", direction="lr" + ) + x_s = pybamm.SpatialVariable( + "x_s", domain="separator", coord_sys="cartesian", direction="lr" + ) + x_p = pybamm.SpatialVariable( + "x_p", domain="positive electrode", coord_sys="cartesian", direction="lr" + ) + z_n = pybamm.SpatialVariable( + "z_n", domain="negative electrode", coord_sys="cartesian", direction="tb" + ) + z_s = pybamm.SpatialVariable( + "z_s", domain="separator", coord_sys="cartesian", direction="tb" + ) + z_p = pybamm.SpatialVariable( + "z_p", domain="positive electrode", coord_sys="cartesian", direction="tb" + ) + z = pybamm.SpatialVariable( + "z", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="tb", + ) + + c_e_n = pybamm.Variable( + "Negative electrolyte concentration [mol.m-3]", + domain="negative electrode", + ) + c_e_s = pybamm.Variable( + "Separator electrolyte concentration [mol.m-3]", + domain="separator", + ) + c_e_p = pybamm.Variable( + "Positive electrolyte concentration [mol.m-3]", + domain="positive electrode", + ) + c_e = pybamm.concatenation(c_e_n, c_e_s, c_e_p) + + phi_e_n = pybamm.Variable( + "Negative electrolyte potential [V]", + domain="negative electrode", + ) + phi_e_s = pybamm.Variable( + "Separator electrolyte potential [V]", + domain="separator", + ) + phi_e_p = pybamm.Variable( + "Positive electrolyte potential [V]", + domain="positive electrode", + ) + phi_e = pybamm.concatenation(phi_e_n, phi_e_s, phi_e_p) + + phi_s_n = pybamm.Variable( + "Negative electrode potential [V]", domain="negative electrode" + ) + phi_s_p = pybamm.Variable( + "Positive electrode potential [V]", + domain="positive electrode", + ) + c_s_n = pybamm.Variable( + "Negative particle concentration [mol.m-3]", + domain="negative particle", + auxiliary_domains={"secondary": "negative electrode"}, + ) + c_s_p = pybamm.Variable( + "Positive particle concentration [mol.m-3]", + domain="positive particle", + auxiliary_domains={"secondary": "positive electrode"}, + ) + + T = self.param.T_init + + ###################### + # Other set-up + ###################### + i_cell = self.param.current_density_with_time + + eps_n = pybamm.FunctionParameter( + "Negative electrode porosity", + {"Through-cell distance (x) [m]": x_n, "Vertical distance (z) [m]": z_n}, + ) + eps_s = pybamm.FunctionParameter( + "Separator porosity", + {"Through-cell distance (x) [m]": x_s, "Vertical distance (z) [m]": z_s}, + ) + eps_p = pybamm.FunctionParameter( + "Positive electrode porosity", + {"Through-cell distance (x) [m]": x_p, "Vertical distance (z) [m]": z_p}, + ) + eps = pybamm.concatenation(eps_n, eps_s, eps_p) + + eps_s_n = pybamm.FunctionParameter( + "Negative electrode active material volume fraction", + {"Through-cell distance (x) [m]": x_n, "Vertical distance (z) [m]": z_n}, + ) + eps_s_p = pybamm.FunctionParameter( + "Positive electrode active material volume fraction", + {"Through-cell distance (x) [m]": x_p, "Vertical distance (z) [m]": z_p}, + ) + + tor = pybamm.concatenation( + eps_n**self.param.n.b_e, eps_s**self.param.s.b_e, eps_p**self.param.p.b_e + ) + a_n = 3 * self.param.n.prim.epsilon_s_av / self.param.n.prim.R_typ + a_p = 3 * self.param.p.prim.epsilon_s_av / self.param.p.prim.R_typ + + # Interfacial reactions + c_s_surf_n = pybamm.surf(c_s_n) + sto_surf_n = c_s_surf_n / self.param.n.prim.c_max + j0_n = self.param.n.prim.j0(c_e_n, c_s_surf_n, T) + delta_phi_n = phi_s_n - phi_e_n + eta_n = delta_phi_n - self.param.n.prim.U(sto_surf_n, T) + Feta_RT_n = self.param.F * eta_n / (self.param.R * T) + j_n = 2 * j0_n * pybamm.sinh(self.param.n.prim.ne / 2 * Feta_RT_n) + + c_s_surf_p = pybamm.surf(c_s_p) + sto_surf_p = c_s_surf_p / self.param.p.prim.c_max + j0_p = self.param.p.prim.j0(c_e_p, c_s_surf_p, T) + delta_phi_p = phi_s_p - phi_e_p + eta_p = delta_phi_p - self.param.p.prim.U(sto_surf_p, T) + Feta_RT_p = self.param.F * eta_p / (self.param.R * T) + j_s = pybamm.PrimaryBroadcast(0, "separator") + j_p = 2 * j0_p * pybamm.sinh(self.param.p.prim.ne / 2 * Feta_RT_p) + + a_j_n = a_n * j_n + a_j_p = a_p * j_p + a_j = pybamm.concatenation(a_j_n, j_s, a_j_p) + + ###################### + # State of Charge + ###################### + current = self.param.current_with_time + self.rhs[Q] = current / 3600 + self.initial_conditions[Q] = pybamm.Scalar(0) + + ###################### + # Particles + ###################### + N_s_n = -self.param.n.prim.D(c_s_n, T) * pybamm.grad(c_s_n) + N_s_p = -self.param.p.prim.D(c_s_p, T) * pybamm.grad(c_s_p) + self.rhs[c_s_n] = -pybamm.div(N_s_n) + self.rhs[c_s_p] = -pybamm.div(N_s_p) + self.boundary_conditions[c_s_n] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": ( + -j_n / (self.param.F * pybamm.surf(self.param.n.prim.D(c_s_n, T))), + "Neumann", + ), + } + self.boundary_conditions[c_s_p] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": ( + -j_p / (self.param.F * pybamm.surf(self.param.p.prim.D(c_s_p, T))), + "Neumann", + ), + } + self.initial_conditions[c_s_n] = self.param.n.prim.c_init + self.initial_conditions[c_s_p] = self.param.p.prim.c_init + + c_s_n_av = pybamm.RAverage(c_s_n) + c_s_p_av = pybamm.RAverage(c_s_p) + solid_lithium_negative = pybamm.Integral(c_s_n_av * eps_s_n, [x_n, z_n]) + solid_lithium_positive = pybamm.Integral(c_s_p_av * eps_s_p, [x_p, z_p]) + total_solid_lithium = solid_lithium_negative + solid_lithium_positive + + ###################### + # Current in the solid + ###################### + sigma_eff_n = self.param.n.sigma(T) * eps_s_n**self.param.n.b_s + sigma_eff_p = self.param.p.sigma(T) * eps_s_p**self.param.p.b_s + self.algebraic[phi_s_n] = ( + self.param.L_x**2 + * self.param.L_z**2 + * (pybamm.div(-sigma_eff_n * pybamm.grad(phi_s_n)) + a_j_n) + ) + self.algebraic[phi_s_p] = ( + self.param.L_x**2 + * self.param.L_z**2 + * (pybamm.div(-sigma_eff_p * pybamm.grad(phi_s_p)) + a_j_p) + ) + self.boundary_conditions[phi_s_n] = { + "left": (pybamm.Scalar(0), "Dirichlet"), + "right": (pybamm.Scalar(0), "Neumann"), + "top": (pybamm.Scalar(0), "Neumann"), + "bottom": (pybamm.Scalar(0), "Neumann"), + } + self.boundary_conditions[phi_s_p] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": (i_cell / pybamm.boundary_value(-sigma_eff_p, "right"), "Neumann"), + "top": (pybamm.Scalar(0), "Neumann"), + "bottom": (pybamm.Scalar(0), "Neumann"), + } + self.initial_conditions[phi_s_n] = pybamm.Scalar(0) + self.initial_conditions[phi_s_p] = self.param.ocv_init + ###################### + # Current in the electrolyte + ###################### + kappa_eff = self.param.kappa_e(c_e, T) * tor + kappa_D_eff = kappa_eff * self.param.chiRT_over_Fc(c_e, T) + self.algebraic[phi_e] = ( + self.param.L_x**2 + * self.param.L_z**2 + * ( + pybamm.div(kappa_D_eff * pybamm.grad(c_e)) + - pybamm.div(kappa_eff * pybamm.grad(phi_e)) + - a_j + ) + ) + self.boundary_conditions[phi_e] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": (pybamm.Scalar(0), "Neumann"), + "top": (pybamm.Scalar(0), "Neumann"), + "bottom": (pybamm.Scalar(0), "Neumann"), + } + self.initial_conditions[phi_e] = -self.param.n.prim.U_init + + ###################### + # Electrolyte concentration + ###################### + D_e_eff = tor * self.param.D_e(c_e, T) + self.rhs[c_e] = (1 / eps) * ( + pybamm.div(D_e_eff * pybamm.grad(c_e)) + + (1 - self.param.t_plus(c_e, T)) * a_j / self.param.F + ) + self.boundary_conditions[c_e] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": (pybamm.Scalar(0), "Neumann"), + "top": (pybamm.Scalar(0), "Neumann"), + "bottom": (pybamm.Scalar(0), "Neumann"), + } + self.initial_conditions[c_e] = self.param.c_e_init + + ###################### + # (Some) variables + ###################### + voltage = pybamm.boundary_value(phi_s_p, "top-right") + num_cells = pybamm.Parameter( + "Number of cells connected in series to make a battery" + ) + total_lithium = pybamm.Integral(c_e * eps, [x, z]) + self.variables = { + "Negative particle concentration [mol.m-3]": c_s_n, + "Total lithium [mol]": total_lithium, + "Negative particle surface concentration [mol.m-3]": c_s_surf_n, + "Electrolyte concentration [mol.m-3]": c_e, + "Negative electrolyte concentration [mol.m-3]": c_e_n, + "Separator electrolyte concentration [mol.m-3]": c_e_s, + "Positive electrolyte concentration [mol.m-3]": c_e_p, + "Positive particle concentration [mol.m-3]": c_s_p, + "Positive particle surface concentration [mol.m-3]": c_s_surf_p, + "Current [A]": current, + "Current variable [A]": current, + "Negative electrode potential [V]": phi_s_n, + "Electrolyte potential [V]": phi_e, + "Negative electrolyte potential [V]": phi_e_n, + "Separator electrolyte potential [V]": phi_e_s, + "Positive electrolyte potential [V]": phi_e_p, + "Positive electrode potential [V]": phi_s_p, + "Voltage [V]": voltage, + "Battery voltage [V]": voltage * num_cells, + "Time [s]": pybamm.t, + "Discharge capacity [A.h]": Q, + "x": x, + "z": z, + "Current density [A.m-2]": a_j, + "Electrolyte current density [A.m-2]": a_j, + "x_n": x_n, + "x_s": x_s, + "x_p": x_p, + "z_n": z_n, + "z_s": z_s, + "z_p": z_p, + "Negative electrode surface concentration [mol.m-3]": c_s_surf_n, + "Negative electrode surface stoichiometry": sto_surf_n, + "Positive electrode surface concentration [mol.m-3]": c_s_surf_p, + "Positive electrode surface stoichiometry": sto_surf_p, + "Positive electrode surface potential difference [V]": delta_phi_p, + "Negative electrode surface potential difference [V]": delta_phi_n, + "Positive electrode overpotential [V]": eta_p, + "Negative electrode overpotential [V]": eta_n, + "Positive electrode ocp [V]": self.param.p.prim.U(sto_surf_p, T), + "Negative electrode ocp [V]": self.param.n.prim.U(sto_surf_n, T), + "Positive electrode current density [A.m-2]": j_p, + "Negative electrode current density [A.m-2]": j_n, + "Electrolyte flux [mol.m-2.s-1]": D_e_eff, + "Positive solid lithium [mol]": solid_lithium_positive, + "Negative solid lithium [mol]": solid_lithium_negative, + "Total solid lithium [mol]": total_solid_lithium, + } + self.events += [ + pybamm.Event("Minimum voltage [V]", voltage - self.param.voltage_low_cut), + pybamm.Event("Maximum voltage [V]", self.param.voltage_high_cut - voltage), + ] + + @property + def default_geometry(self): + z_2d = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="tb", + ) + return { + "negative electrode": { + "x_n": {"min": 0, "max": self.param.n.L}, + z_2d: {"min": 0, "max": self.param.L_z}, + }, + "separator": { + "x_s": {"min": self.param.n.L, "max": self.param.n.L + self.param.s.L}, + z_2d: {"min": 0, "max": self.param.L_z}, + }, + "positive electrode": { + "x_p": { + "min": self.param.n.L + self.param.s.L, + "max": self.param.n.L + self.param.s.L + self.param.p.L, + }, + z_2d: {"min": 0, "max": self.param.L_z}, + }, + "positive particle": { + "r_p": {"min": 0, "max": self.param.p.prim.R_typ}, + }, + "negative particle": { + "r_n": {"min": 0, "max": self.param.n.prim.R_typ}, + }, + "current collector": { + "z": {"position": 0}, + }, + } + + @property + def default_spatial_methods(self): + return { + "negative electrode": pybamm.FiniteVolumeUnstructured(), + "separator": pybamm.FiniteVolumeUnstructured(), + "positive electrode": pybamm.FiniteVolumeUnstructured(), + "positive particle": pybamm.FiniteVolume(), + "negative particle": pybamm.FiniteVolume(), + "current collector": pybamm.ZeroDimensionalSpatialMethod(), + } + + @property + def default_submesh_types(self): + return { + "negative electrode": pybamm.UnstructuredMeshGenerator( + element_type=self._element_type + ), + "separator": pybamm.UnstructuredMeshGenerator( + element_type=self._element_type + ), + "positive electrode": pybamm.UnstructuredMeshGenerator( + element_type=self._element_type + ), + "positive particle": pybamm.Uniform1DSubMesh, + "negative particle": pybamm.Uniform1DSubMesh, + "current collector": pybamm.SubMesh0D, + } + + @property + def default_var_pts(self): + z_2d = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="tb", + ) + return { + "x_n": 20, + "x_s": 30, + "x_p": 20, + "r_p": 20, + "r_n": 20, + z_2d: 10, + } diff --git a/src/pybamm/spatial_methods/finite_volume_unstructured.py b/src/pybamm/spatial_methods/finite_volume_unstructured.py index 7af5ad8a9d..8b46f95b88 100644 --- a/src/pybamm/spatial_methods/finite_volume_unstructured.py +++ b/src/pybamm/spatial_methods/finite_volume_unstructured.py @@ -41,6 +41,26 @@ def build(self, mesh): for dom in mesh.keys(): mesh[dom].npts_for_broadcast_to_nodes = mesh[dom].npts + @staticmethod + def _bc_contribution(n, n_bnd, owners, coeffs, bc_value): + """Build a symbolic BC contribution vector. + + For scalar ``bc_value``: returns ``Vector(accumulated_coeffs) * bc_value``. + For vector ``bc_value`` (length ``n_bnd``): + returns ``Matrix(n, n_bnd) @ bc_value``. + """ + is_scalar = isinstance(bc_value, pybamm.Scalar) or ( + hasattr(bc_value, "shape_for_testing") + and bc_value.shape_for_testing == (1, 1) + ) + if is_scalar: + row = np.zeros(n) + np.add.at(row, owners, coeffs) + return pybamm.Vector(row) * bc_value + else: + M = csr_matrix((coeffs, (owners, np.arange(n_bnd))), shape=(n, n_bnd)) + return pybamm.Matrix(M) @ bc_value + # ------------------------------------------------------------------ # spatial_variable # ------------------------------------------------------------------ @@ -114,18 +134,13 @@ def laplacian(self, symbol, discretised_symbol, boundary_conditions): L = self._tpfa_matrix(submesh) - # Boundary conditions - bc_rhs = np.zeros(n) + bc_rhs = pybamm.Vector(np.zeros(n)) if symbol in boundary_conditions: bcs = boundary_conditions[symbol] L, bc_rhs = self._apply_bcs_to_laplacian(submesh, L, bc_rhs, bcs) L_full = csr_matrix(kron(eye(repeats, dtype=np.float64), L)) - result = pybamm.Matrix(L_full) @ discretised_symbol - - if np.any(bc_rhs != 0): - bc_rhs_full = np.tile(bc_rhs, repeats) - result = result + pybamm.Vector(bc_rhs_full) + result = pybamm.Matrix(L_full) @ discretised_symbol + bc_rhs return result @@ -173,8 +188,137 @@ def _tpfa_matrix(self, submesh): return csr_matrix(coo_matrix((data, (rows, cols)), shape=(n, n))) + def div_D_grad(self, div_symbol, grad_child, disc_D, disc_u, boundary_conditions): + """Discretise ``div(D * grad(u))`` as a single TPFA operation. + + Fully symbolic — works for both constant and state-dependent ``D``. + Internal-face fluxes use arithmetic-mean interpolation of ``D`` to + faces and a standard two-point difference for ``grad(u)``. + """ + domain = div_symbol.domain + submesh = self.mesh[domain] + n = submesh.npts + n_int = submesh.n_internal_faces + repeats = self._get_auxiliary_domain_repeats(div_symbol.domains) + vol = submesh.cell_volumes + + owner = submesh.face_owner[:n_int] + neighbor = submesh.face_neighbor[:n_int] + + c_o = submesh.cell_centroids[owner] + c_n = submesh.cell_centroids[neighbor] + delta = c_n - c_o + dist = np.linalg.norm(delta, axis=1) + e_ij = delta / dist[:, np.newaxis] + cos_theta = np.abs(np.sum(submesh.face_normals[:n_int] * e_ij, axis=1)) + geo = submesh.face_areas[:n_int] * cos_theta / dist + + # G (n_int x n): u_neighbor - u_owner per face + G = csr_matrix( + ( + np.concatenate([-np.ones(n_int), np.ones(n_int)]), + (np.tile(np.arange(n_int), 2), np.concatenate([owner, neighbor])), + ), + shape=(n_int, n), + ) + + # W (n_int x n): arithmetic-mean D to faces + W = csr_matrix( + ( + np.full(2 * n_int, 0.5), + (np.tile(np.arange(n_int), 2), np.concatenate([owner, neighbor])), + ), + shape=(n_int, n), + ) + + # S (n x n_int): face flux -> cell divergence (+owner, -neighbor, /V) + S = csr_matrix( + ( + np.concatenate([1.0 / vol[owner], -1.0 / vol[neighbor]]), + (np.concatenate([owner, neighbor]), np.tile(np.arange(n_int), 2)), + ), + shape=(n, n_int), + ) + + G_f = csr_matrix(kron(eye(repeats, dtype=np.float64), G)) + W_f = csr_matrix(kron(eye(repeats, dtype=np.float64), W)) + S_f = csr_matrix(kron(eye(repeats, dtype=np.float64), S)) + geo_f = np.tile(geo, repeats) + + u_diff = pybamm.Matrix(G_f) @ disc_u + is_scalar_D = isinstance(disc_D, pybamm.Scalar) or ( + hasattr(disc_D, "shape_for_testing") and disc_D.shape_for_testing == (1, 1) + ) + if is_scalar_D: + flux = disc_D * u_diff * pybamm.Vector(geo_f) + else: + D_face = pybamm.Matrix(W_f) @ disc_D + flux = D_face * u_diff * pybamm.Vector(geo_f) + result = pybamm.Matrix(S_f) @ flux + + # Boundary conditions + bc_rhs = pybamm.Vector(np.zeros(n * repeats)) + if grad_child in boundary_conditions: + bcs = boundary_conditions[grad_child] + for side, (bc_value, bc_type) in bcs.items(): + face_tag = self._side_to_boundary_tag(side) + if face_tag not in submesh.boundary_faces: + continue + fi_arr = submesh.boundary_faces[face_tag] + n_bnd = len(fi_arr) + bnd_own = submesh.face_owner[fi_arr] + + E = csr_matrix( + (np.ones(n_bnd), (np.arange(n_bnd), bnd_own)), + shape=(n_bnd, n), + ) + E_f = csr_matrix(kron(eye(repeats, dtype=np.float64), E)) + P = csr_matrix( + (np.ones(n_bnd), (bnd_own, np.arange(n_bnd))), + shape=(n, n_bnd), + ) + P_f = csr_matrix(kron(eye(repeats, dtype=np.float64), P)) + D_bnd = disc_D if is_scalar_D else pybamm.Matrix(E_f) @ disc_D + + if bc_type == "Dirichlet": + geo_bnd = np.array( + [ + submesh.face_areas[fi] + / np.linalg.norm( + submesh.face_centroids[fi] + - submesh.cell_centroids[bnd_own[j]] + ) + / vol[bnd_own[j]] + for j, fi in enumerate(fi_arr) + ] + ) + geo_bnd_f = np.tile(geo_bnd, repeats) + + u_bnd = pybamm.Matrix(E_f) @ disc_u + bc_rhs = bc_rhs + pybamm.Matrix(P_f) @ ( + D_bnd * (bc_value - u_bnd) * pybamm.Vector(geo_bnd_f) + ) + + elif bc_type == "Neumann" and bc_value != pybamm.Scalar(0): + a_over_v = np.array( + [ + submesh.face_areas[fi] / vol[bnd_own[j]] + for j, fi in enumerate(fi_arr) + ] + ) + a_over_v_f = np.tile(a_over_v, repeats) + bc_rhs = bc_rhs + pybamm.Matrix(P_f) @ ( + D_bnd * bc_value * pybamm.Vector(a_over_v_f) + ) + + return result + bc_rhs + def _apply_bcs_to_laplacian(self, submesh, L, bc_rhs, bcs): - """Modify the Laplacian matrix and RHS for boundary conditions.""" + """Modify the Laplacian matrix and RHS for boundary conditions. + + ``bc_rhs`` is a pybamm expression (symbolic vector). + """ + n = submesh.npts L = L.tolil() for side, (bc_value, bc_type) in bcs.items(): @@ -183,25 +327,35 @@ def _apply_bcs_to_laplacian(self, submesh, L, bc_rhs, bcs): continue face_indices = submesh.boundary_faces[face_tag] + n_bnd = len(face_indices) + owners = submesh.face_owner[face_indices] - for fi in face_indices: - cell = submesh.face_owner[fi] - area = submesh.face_areas[fi] - vol = submesh.cell_volumes[cell] - - face_centroid = submesh.face_centroids[fi] - cell_centroid = submesh.cell_centroids[cell] - d = np.linalg.norm(face_centroid - cell_centroid) - - coeff = area / d - - if bc_type == "Dirichlet": - bc_val = float(bc_value.evaluate()) + if bc_type == "Dirichlet": + coeffs = np.empty(n_bnd) + for j, fi in enumerate(face_indices): + cell = owners[j] + area = submesh.face_areas[fi] + vol = submesh.cell_volumes[cell] + d_perp = np.linalg.norm( + submesh.face_centroids[fi] - submesh.cell_centroids[cell] + ) + coeff = area / d_perp L[cell, cell] -= coeff / vol - bc_rhs[cell] += coeff * bc_val / vol - elif bc_type == "Neumann": - bc_val = float(bc_value.evaluate()) - bc_rhs[cell] += bc_val * area / vol + coeffs[j] = coeff / vol + bc_rhs = bc_rhs + self._bc_contribution( + n, n_bnd, owners, coeffs, bc_value + ) + + elif bc_type == "Neumann": + coeffs = np.array( + [ + submesh.face_areas[fi] / submesh.cell_volumes[owners[j]] + for j, fi in enumerate(face_indices) + ] + ) + bc_rhs = bc_rhs + self._bc_contribution( + n, n_bnd, owners, coeffs, bc_value + ) return csr_matrix(L), bc_rhs @@ -229,7 +383,7 @@ def gradient(self, symbol, discretised_symbol, boundary_conditions): G_components = self._green_gauss_matrices(submesh) - bc_vecs = [np.zeros(n) for _ in range(d)] + bc_vecs = [pybamm.Vector(np.zeros(n)) for _ in range(d)] if symbol in boundary_conditions: bcs = boundary_conditions[symbol] G_components, bc_vecs = self._apply_bcs_to_gradient( @@ -239,10 +393,7 @@ def gradient(self, symbol, discretised_symbol, boundary_conditions): components = [] for k in range(d): Gk = csr_matrix(kron(eye(repeats, dtype=np.float64), G_components[k])) - comp = pybamm.Matrix(Gk) @ discretised_symbol - if np.any(bc_vecs[k] != 0): - bc_full = np.tile(bc_vecs[k], repeats) - comp = comp + pybamm.Vector(bc_full) + comp = pybamm.Matrix(Gk) @ discretised_symbol + bc_vecs[k] components.append(comp) vf = pybamm.VectorField(*components) @@ -324,7 +475,11 @@ def _green_gauss_matrices(self, submesh): return G def _apply_bcs_to_gradient(self, submesh, G_components, bc_vecs, bcs): - """Apply Dirichlet/Neumann BCs to gradient matrices.""" + """Apply Dirichlet/Neumann BCs to gradient matrices. + + ``bc_vecs`` is a list of pybamm expressions (one per spatial dimension). + """ + n = submesh.npts d = submesh.dimension vol = submesh.cell_volumes @@ -334,61 +489,52 @@ def _apply_bcs_to_gradient(self, submesh, G_components, bc_vecs, bcs): continue face_indices = submesh.boundary_faces[face_tag] + n_bnd = len(face_indices) + owners = submesh.face_owner[face_indices] if bc_type == "Dirichlet": - bc_val = float(bc_value.evaluate()) - for fi in face_indices: - cell = submesh.face_owner[fi] - area = submesh.face_areas[fi] + for j, fi in enumerate(face_indices): + cell = owners[j] normal = submesh.face_normals[fi] - face_c = submesh.face_centroids[fi] - cell_c = submesh.cell_centroids[cell] - dist = np.linalg.norm(face_c - cell_c) - - # Replace zeroth-order boundary term u_owner * n * A / V - # with ghost-cell interpolation: - # u_f = (u_owner + u_ghost) / 2 where u_ghost = 2*bc_val - u_owner - # So u_f = bc_val, meaning: - # remove existing contribution (u_owner * n * A / V) - # add bc_val * n * A / V to RHS - # But the Green-Gauss matrix already has u_owner terms from - # _green_gauss_matrices. We need to zero out the boundary - # contribution and replace with the BC value. - # Simpler: the boundary face contribution to cell i is - # G_k[cell, cell] gets n_k * A / V (from boundary term) - # For Dirichlet: u_f = bc_val, so contribution is - # bc_val * n_k * A / V (pure RHS, no matrix term) - # We need to subtract the existing matrix term and add RHS. + area = submesh.face_areas[fi] for k in range(d): nk_A = normal[k] * area - # Remove the u_owner boundary term from the matrix G_components[k] = G_components[k].tolil() G_components[k][cell, cell] -= nk_A / vol[cell] G_components[k] = csr_matrix(G_components[k]) - # Add bc_val * n_k * A / V to the RHS - bc_vecs[k][cell] += bc_val * nk_A / vol[cell] + + for k in range(d): + coeffs = np.array( + [ + submesh.face_normals[fi, k] + * submesh.face_areas[fi] + / vol[owners[j]] + for j, fi in enumerate(face_indices) + ] + ) + bc_vecs[k] = bc_vecs[k] + self._bc_contribution( + n, n_bnd, owners, coeffs, bc_value + ) elif bc_type == "Neumann": - bc_val = float(bc_value.evaluate()) - for fi in face_indices: - cell = submesh.face_owner[fi] - area = submesh.face_areas[fi] - normal = submesh.face_normals[fi] - # Neumann: flux = bc_val at the face (in normal direction) - # The gradient's boundary contribution becomes: - # (bc_val * dx + u_owner) * n_k * A / V - # For simplicity in the gradient, we keep the zeroth-order - # owner term and add the correction. - # Actually for Neumann BC on gradient, the face value is: - # u_f = u_owner + bc_val * dist_to_face - # The extra contribution to the gradient is: - # bc_val * dist * n_k * A / V - dist = np.linalg.norm( - submesh.face_centroids[fi] - submesh.cell_centroids[cell] + dists = np.linalg.norm( + submesh.face_centroids[face_indices] + - submesh.cell_centroids[owners], + axis=1, + ) + for k in range(d): + coeffs = np.array( + [ + dists[j] + * submesh.face_normals[fi, k] + * submesh.face_areas[fi] + / vol[owners[j]] + for j, fi in enumerate(face_indices) + ] + ) + bc_vecs[k] = bc_vecs[k] + self._bc_contribution( + n, n_bnd, owners, coeffs, bc_value ) - for k in range(d): - nk_A = normal[k] * area - bc_vecs[k][cell] += bc_val * dist * nk_A / vol[cell] return G_components, bc_vecs @@ -420,22 +566,6 @@ def divergence(self, symbol, discretised_symbol, boundary_conditions): Dk = csr_matrix(kron(eye(repeats, dtype=np.float64), D_components[k])) result = result + pybamm.Matrix(Dk) @ comps[k] - disc_sv = getattr(discretised_symbol, "_disc_state_vector", None) - if disc_sv is not None: - L_bc, bc_rhs, D_bnd = self._div_boundary_correction( - submesh, boundary_conditions, domain=domain - ) - if D_bnd is not None: - for k in range(d): - Dk_bnd = csr_matrix(kron(eye(repeats, dtype=np.float64), D_bnd[k])) - result = result - pybamm.Matrix(Dk_bnd) @ comps[k] - if L_bc is not None: - L_bc_full = csr_matrix(kron(eye(repeats, dtype=np.float64), L_bc)) - result = result + pybamm.Matrix(L_bc_full) @ disc_sv - if np.any(bc_rhs != 0): - bc_rhs_full = np.tile(bc_rhs, repeats) - result = result + pybamm.Vector(bc_rhs_full) - return result def _div_boundary_correction(self, submesh, boundary_conditions, domain=None): @@ -446,7 +576,7 @@ def _div_boundary_correction(self, submesh, boundary_conditions, domain=None): This method returns: * ``L_bc`` – sparse matrix for TPFA Dirichlet correction on state vector - * ``bc_rhs`` – constant vector for Dirichlet/Neumann RHS + * ``bc_rhs`` – symbolic pybamm expression for Dirichlet/Neumann RHS * ``D_bnd`` – list of sparse matrices (boundary-only divergence terms to subtract from the cell-centered approximation) @@ -457,7 +587,7 @@ def _div_boundary_correction(self, submesh, boundary_conditions, domain=None): n = submesh.npts d = submesh.dimension L_bc = None - bc_rhs = np.zeros(n) + bc_rhs = pybamm.Vector(np.zeros(n)) D_bnd = None for var, bcs in boundary_conditions.items(): @@ -470,31 +600,45 @@ def _div_boundary_correction(self, submesh, boundary_conditions, domain=None): if face_tag not in submesh.boundary_faces: continue face_indices = submesh.boundary_faces[face_tag] - bc_val = float(bc_value.evaluate()) - - for fi in face_indices: - cell = submesh.face_owner[fi] - area = submesh.face_areas[fi] - vol = submesh.cell_volumes[cell] + n_bnd = len(face_indices) + owners = submesh.face_owner[face_indices] + areas = submesh.face_areas[face_indices] + vols = submesh.cell_volumes[owners] + + if D_bnd is None: + D_bnd = [csr_matrix((n, n)).tolil() for _ in range(d)] + for j, fi in enumerate(face_indices): + cell = owners[j] normal = submesh.face_normals[fi] - - if D_bnd is None: - D_bnd = [csr_matrix((n, n)).tolil() for _ in range(d)] for k in range(d): - D_bnd[k][cell, cell] += normal[k] * area / vol + D_bnd[k][cell, cell] += normal[k] * areas[j] / vols[j] - if bc_type == "Dirichlet": + if bc_type == "Dirichlet": + for j, fi in enumerate(face_indices): + cell = owners[j] face_c = submesh.face_centroids[fi] cell_c = submesh.cell_centroids[cell] d_perp = np.linalg.norm(face_c - cell_c) - coeff = area / d_perp - + coeff = areas[j] / d_perp if L_bc is None: L_bc = csr_matrix((n, n)).tolil() - L_bc[cell, cell] -= coeff / vol - bc_rhs[cell] += coeff * bc_val / vol - elif bc_type == "Neumann": - bc_rhs[cell] += bc_val * area / vol + L_bc[cell, cell] -= coeff / vols[j] + + coeffs = np.empty(n_bnd) + for j, fi in enumerate(face_indices): + face_c = submesh.face_centroids[fi] + cell_c = submesh.cell_centroids[owners[j]] + d_perp = np.linalg.norm(face_c - cell_c) + coeffs[j] = (areas[j] / d_perp) / vols[j] + bc_rhs = bc_rhs + self._bc_contribution( + n, n_bnd, owners, coeffs, bc_value + ) + + elif bc_type == "Neumann": + coeffs = areas / vols + bc_rhs = bc_rhs + self._bc_contribution( + n, n_bnd, owners, coeffs, bc_value + ) if L_bc is not None: L_bc = csr_matrix(L_bc) @@ -662,12 +806,25 @@ def boundary_integral(self, child, discretised_child, region): # boundary_value_or_flux # ------------------------------------------------------------------ + _CORNER_SIDES = { + "top-right": ("top", "right"), + "top-left": ("top", "left"), + "bottom-right": ("bottom", "right"), + "bottom-left": ("bottom", "left"), + } + def boundary_value_or_flux(self, symbol, discretised_child, bcs=None): submesh = self.mesh[discretised_child.domain] n = submesh.npts repeats = self._get_auxiliary_domain_repeats(discretised_child.domains) side = symbol.side + + if side in self._CORNER_SIDES: + return self._corner_boundary_value( + submesh, n, repeats, side, discretised_child + ) + face_tag = self._side_to_boundary_tag(side) if face_tag not in submesh.boundary_faces: @@ -679,20 +836,10 @@ def boundary_value_or_flux(self, symbol, discretised_child, bcs=None): n_bnd = len(face_indices) owners = submesh.face_owner[face_indices] - if isinstance(symbol, pybamm.BoundaryGradient): - # For boundary gradient, extrapolate gradient from cell center to face - # using nearest cell value (zeroth-order) — improved when BCs available - sub_matrix = csr_matrix( - (np.ones(n_bnd), (np.arange(n_bnd), owners)), - shape=(n_bnd, n), - ) - else: - # BoundaryValue: linear extrapolation from cell center to face - # For unstructured meshes, use constant extrapolation (cell value) - sub_matrix = csr_matrix( - (np.ones(n_bnd), (np.arange(n_bnd), owners)), - shape=(n_bnd, n), - ) + sub_matrix = csr_matrix( + (np.ones(n_bnd), (np.arange(n_bnd), owners)), + shape=(n_bnd, n), + ) mat = csr_matrix(kron(eye(repeats, dtype=np.float64), sub_matrix)) bv_vector = pybamm.Matrix(mat) @@ -701,6 +848,34 @@ def boundary_value_or_flux(self, symbol, discretised_child, bcs=None): out.clear_domains() return out + def _corner_boundary_value(self, submesh, n, repeats, side, discretised_child): + """Extract value from the cell closest to a corner of the domain.""" + tb_side, lr_side = self._CORNER_SIDES[side] + centroids = submesh.cell_centroids + x_coords = centroids[:, 0] + z_coords = centroids[:, -1] + + if lr_side == "right": + target_x = x_coords.max() + else: + target_x = x_coords.min() + if tb_side == "top": + target_z = z_coords.max() + else: + target_z = z_coords.min() + + dists = (x_coords - target_x) ** 2 + (z_coords - target_z) ** 2 + cell_idx = int(np.argmin(dists)) + + sub_matrix = csr_matrix( + (np.ones(1), (np.zeros(1, dtype=int), [cell_idx])), + shape=(1, n), + ) + mat = csr_matrix(kron(eye(repeats, dtype=np.float64), sub_matrix)) + out = pybamm.Matrix(mat) @ discretised_child + out.clear_domains() + return out + # ------------------------------------------------------------------ # internal_neumann_condition # ------------------------------------------------------------------ diff --git a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py index 837682b9d6..df869095d3 100644 --- a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py +++ b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py @@ -52,3 +52,30 @@ class TestBasicDFNHalfCell(BaseBasicModelTest): def setup(self): options = {"working electrode": "positive"} self.model = pybamm.lithium_ion.BasicDFNHalfCell(options) + + +class TestBasicDFN2DUnstructured: + def test_solves_and_matches_structured(self): + import numpy as np + + z_2d = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="tb", + ) + var_pts = {"x_n": 5, "x_s": 5, "x_p": 5, "r_p": 10, "r_n": 10, z_2d: 3} + t_eval = np.linspace(0, 3600, 20) + + model_s = pybamm.lithium_ion.BasicDFN2D() + sim_s = pybamm.Simulation(model_s, var_pts=var_pts) + sol_s = sim_s.solve(t_eval) + + model_u = pybamm.lithium_ion.BasicDFN2DUnstructured(element_type="quad") + sim_u = pybamm.Simulation(model_u, var_pts=var_pts) + sol_u = sim_u.solve(t_eval) + + V_s = sol_s["Voltage [V]"](t=t_eval) + V_u = sol_u["Voltage [V]"](t=t_eval) + + np.testing.assert_allclose(V_u, V_s, atol=5e-3) diff --git a/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py b/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py index ccee5ab4c2..6a61f19974 100644 --- a/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py +++ b/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py @@ -25,3 +25,7 @@ def test_dfn_composite_well_posed(self): def test_dfn_2d(self): model = pybamm.lithium_ion.BasicDFN2D() model.check_well_posedness() + + def test_dfn_2d_unstructured(self): + model = pybamm.lithium_ion.BasicDFN2DUnstructured(element_type="quad") + model.check_well_posedness() From 02cd880c022d557d04d53051d5298e2a6454435a Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Wed, 25 Feb 2026 18:09:42 -0800 Subject: [PATCH 06/25] unstructured DFN 3D --- src/pybamm/meshes/meshes.py | 46 +- src/pybamm/meshes/unstructured_submesh.py | 140 +++++- .../lithium_ion/__init__.py | 1 + .../lithium_ion/basic_dfn_3d_unstructured.py | 450 ++++++++++++++++++ .../finite_volume_unstructured.py | 6 +- .../test_lithium_ion/test_basic_models.py | 20 + .../test_lithium_ion/test_basic_models.py | 4 + 7 files changed, 637 insertions(+), 30 deletions(-) create mode 100644 src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_3d_unstructured.py diff --git a/src/pybamm/meshes/meshes.py b/src/pybamm/meshes/meshes.py index 5ad71018cb..2c0e8a24c5 100644 --- a/src/pybamm/meshes/meshes.py +++ b/src/pybamm/meshes/meshes.py @@ -234,6 +234,8 @@ def combine_submeshes(self, *submeshnames): raise pybamm.GeometryError( "Cannot combine submeshes of different dimensions" ) + elif isinstance(self[submeshnames[i]], pybamm.UnstructuredSubMesh): + pass elif self[submeshnames[i]].dimension == 2: if "left" in submeshnames[i] or "right" in submeshnames[i + 1]: # Make sure that the lr edges are aligned @@ -396,12 +398,15 @@ def _compute_unstructured_interfaces(self): "right" in left_mesh.boundary_faces and "left" in right_mesh.boundary_faces ): - compute_interface_data( - left_mesh, - right_mesh, - left_name=left_name, - right_name=right_name, - ) + try: + compute_interface_data( + left_mesh, + right_mesh, + left_name=left_name, + right_name=right_name, + ) + except ValueError: + pass def add_ghost_meshes(self): """ @@ -469,6 +474,35 @@ def _combine_unstructured_submeshes(submeshes): at domain interfaces are merged so that face-connectivity spans across domains. """ + from .unstructured_submesh import UnstructuredSubMesh, _hex_to_tet + + # For 3D tet meshes generated from hex grids, regenerate with + # cumulative i_offset so that alternating-parity face triangulations + # match across domain boundaries. + if all(hasattr(sm, "_hex_gen_params") and sm.dimension == 3 for sm in submeshes): + cumulative_offset = 0 + fixed = [] + for sm in submeshes: + p = sm._hex_gen_params + if cumulative_offset > 0: + nodes, elements = _hex_to_tet( + p["x_edges"], + p["y_edges"], + p["z_edges"], + i_offset=cumulative_offset, + ) + new_sm = UnstructuredSubMesh( + nodes, + elements, + coord_sys=sm.coord_sys, + ) + new_sm._hex_gen_params = p + fixed.append(new_sm) + else: + fixed.append(sm) + cumulative_offset += p["nx"] + submeshes = fixed + tol = 1e-12 all_nodes = list(submeshes[0].nodes) global_maps = [{i: i for i in range(submeshes[0].nodes.shape[0])}] diff --git a/src/pybamm/meshes/unstructured_submesh.py b/src/pybamm/meshes/unstructured_submesh.py index 394a62f923..8b94aec72c 100644 --- a/src/pybamm/meshes/unstructured_submesh.py +++ b/src/pybamm/meshes/unstructured_submesh.py @@ -12,7 +12,7 @@ class UnstructuredSubMesh(SubMesh): Supported element types: * **2D**: triangles (3 vertices) or quadrilaterals (4 vertices) - * **3D**: tetrahedra (4 vertices) + * **3D**: tetrahedra (4 vertices) or hexahedra (8 vertices) All operators are dimension-agnostic: the same code path handles both 2D and 3D, with dimension inferred from ``nodes.shape[1]``. @@ -23,7 +23,7 @@ class UnstructuredSubMesh(SubMesh): Vertex coordinates (d = 2 or 3). elements : numpy.ndarray, shape (n_cells, n_verts_per_cell) Element vertex indices. For 2D: 3 (triangles) or 4 (quads). - For 3D: 4 (tetrahedra). + For 3D: 4 (tetrahedra) or 8 (hexahedra). coord_sys : str, optional Coordinate system, default ``"cartesian"``. boundary_faces : dict[str, numpy.ndarray] or None, optional @@ -45,6 +45,8 @@ def __init__(self, nodes, elements, coord_sys="cartesian", boundary_faces=None): self.element_type = "triangle" elif self.dimension == 3 and verts_per_cell == 4: self.element_type = "tetrahedron" + elif self.dimension == 3 and verts_per_cell == 8: + self.element_type = "hexahedron" else: raise ValueError( f"Unsupported: {verts_per_cell} vertices per cell in {self.dimension}D" @@ -99,6 +101,28 @@ def _compute_cell_geometry(self): + d1[:, 2] * (d2[:, 0] * d3[:, 1] - d2[:, 1] * d3[:, 0]) ) self.cell_volumes = np.abs(det) / 6.0 + elif self.element_type == "hexahedron": + # Volume via divergence theorem: V = (1/3) sum_faces (centroid . normal * area) + # For axis-aligned hexes this simplifies, but we use the general approach + # by splitting each hex into 5 tets for volume computation only. + self.cell_volumes = np.zeros(len(self.elements)) + for i, cell in enumerate(self.elements): + cv = self.nodes[cell] + vol = 0.0 + # Split hex into 5 tets using pattern A + for tet_local in [ + (0, 1, 2, 5), + (0, 2, 3, 7), + (0, 5, 7, 4), + (2, 5, 7, 6), + (0, 2, 5, 7), + ]: + t = cv[list(tet_local)] + d1 = t[1] - t[0] + d2 = t[2] - t[0] + d3 = t[3] - t[0] + vol += abs(np.dot(d1, np.cross(d2, d3))) / 6.0 + self.cell_volumes[i] = vol # ------------------------------------------------------------------ # Face-cell connectivity @@ -106,10 +130,12 @@ def _compute_cell_geometry(self): def _build_face_connectivity(self): """Extract faces, identify internal / boundary, record owner-neighbor.""" - d = self.dimension - n_verts_per_face = d # edges (2 verts) in 2D, triangles (3 verts) in 3D + if self.element_type == "hexahedron": + n_verts_per_face = 4 + else: + n_verts_per_face = self.dimension - face_dict = {} # canonical key -> owner_cell + face_dict = {} # canonical key -> (owner_cell, original_verts) internal_owner = [] internal_neighbor = [] @@ -120,19 +146,19 @@ def _build_face_connectivity(self): key = tuple(sorted(face_verts)) if key in face_dict: - other_cell = face_dict.pop(key) + other_cell, orig_verts = face_dict.pop(key) internal_owner.append(other_cell) internal_neighbor.append(cell_idx) - internal_face_verts.append(key) + internal_face_verts.append(orig_verts) else: - face_dict[key] = cell_idx + face_dict[key] = (cell_idx, tuple(face_verts)) # Remaining entries are boundary faces boundary_owner_list = [] boundary_face_verts = [] - for key, cell_idx in face_dict.items(): + for _key, (cell_idx, orig_verts) in face_dict.items(): boundary_owner_list.append(cell_idx) - boundary_face_verts.append(key) + boundary_face_verts.append(orig_verts) n_internal = len(internal_owner) n_boundary = len(boundary_owner_list) @@ -147,13 +173,27 @@ def _build_face_connectivity(self): self._n_boundary_faces = n_boundary self._boundary_face_start = n_internal + # Standard hex vertex ordering: + # 0=(i,j,k) 1=(i+1,j,k) 2=(i+1,j+1,k) 3=(i,j+1,k) + # 4=(i,j,k+1) 5=(i+1,j,k+1) 6=(i+1,j+1,k+1) 7=(i,j+1,k+1) + _HEX_FACES = [ + (0, 3, 7, 4), # x- (left) + (1, 2, 6, 5), # x+ (right) + (0, 1, 5, 4), # y- (front) + (3, 2, 6, 7), # y+ (back) + (0, 1, 2, 3), # z- (bottom) + (4, 5, 6, 7), # z+ (top) + ] + def _cell_faces(self, cell_verts): """Yield face vertex tuples for a single cell.""" n = len(cell_verts) if self.element_type == "quad": - # 4 edges: (v0,v1), (v1,v2), (v2,v3), (v3,v0) for i in range(n): yield (cell_verts[i], cell_verts[(i + 1) % n]) + elif self.element_type == "hexahedron": + for local_face in self._HEX_FACES: + yield tuple(cell_verts[v] for v in local_face) else: # Simplex: d+1 faces, face i omits vertex i for skip in range(n): @@ -164,23 +204,32 @@ def _cell_faces(self, cell_verts): # ------------------------------------------------------------------ def _compute_face_geometry(self): - face_verts = self.nodes[self.faces] # (n_faces, d, d) + face_verts = self.nodes[self.faces] self.face_centroids = face_verts.mean(axis=1) if self.dimension == 2: - # Face = edge: 2 vertices v0, v1 = face_verts[:, 0], face_verts[:, 1] edge = v1 - v0 self.face_areas = np.linalg.norm(edge, axis=1) - # Outward normal: perpendicular to edge (rotate 90 degrees) normals = np.column_stack([edge[:, 1], -edge[:, 0]]) + elif self.element_type == "hexahedron": + # Face = quad: 4 vertices. Area via cross product of diagonals. + v0 = face_verts[:, 0] + v1 = face_verts[:, 1] + v2 = face_verts[:, 2] + v3 = face_verts[:, 3] + diag1 = v2 - v0 + diag2 = v3 - v1 + cross = np.cross(diag1, diag2) + self.face_areas = 0.5 * np.linalg.norm(cross, axis=1) + normals = cross else: # Face = triangle: 3 vertices v0, v1, v2 = face_verts[:, 0], face_verts[:, 1], face_verts[:, 2] cross = np.cross(v1 - v0, v2 - v0) self.face_areas = 0.5 * np.linalg.norm(cross, axis=1) - normals = cross # will be normalized below + normals = cross # Normalize norms = np.linalg.norm(normals, axis=1, keepdims=True) @@ -335,7 +384,7 @@ def _generate_3d(self, spatial_vars, spatial_lims, npts): y_edges = np.linspace(lim_y["min"], lim_y["max"], ny + 1) z_edges = np.linspace(lim_z["min"], lim_z["max"], nz + 1) - nodes, elements = _hex_to_tet(x_edges, y_edges, z_edges) + nodes, elements = _hex_grid(x_edges, y_edges, z_edges) return UnstructuredSubMesh(nodes, elements, coord_sys=self.coord_sys) @@ -611,7 +660,56 @@ def node_id(i, j): return nodes, np.array(elements, dtype=int) -def _hex_to_tet(x_edges, y_edges, z_edges): +def _hex_grid(x_edges, y_edges, z_edges): + """ + Create a hexahedral grid from edge arrays. + + Returns nodes and 8-vertex hex elements suitable for + :class:`UnstructuredSubMesh` with ``element_type="hexahedron"``. + + Vertex ordering per hex matches :attr:`UnstructuredSubMesh._HEX_FACES`: + + :: + + 0=(i,j,k) 1=(i+1,j,k) 2=(i+1,j+1,k) 3=(i,j+1,k) + 4=(i,j,k+1) 5=(i+1,j,k+1) 6=(i+1,j+1,k+1) 7=(i,j+1,k+1) + + Returns + ------- + nodes : (n_nodes, 3) + elements : (n_cells, 8) + """ + nx = len(x_edges) - 1 + ny = len(y_edges) - 1 + nz = len(z_edges) - 1 + + xx, yy, zz = np.meshgrid(x_edges, y_edges, z_edges, indexing="ij") + nodes = np.column_stack([xx.ravel(), yy.ravel(), zz.ravel()]) + + def node_id(i, j, k): + return i * (ny + 1) * (nz + 1) + j * (nz + 1) + k + + elements = [] + for i in range(nx): + for j in range(ny): + for k in range(nz): + elements.append( + [ + node_id(i, j, k), + node_id(i + 1, j, k), + node_id(i + 1, j + 1, k), + node_id(i, j + 1, k), + node_id(i, j, k + 1), + node_id(i + 1, j, k + 1), + node_id(i + 1, j + 1, k + 1), + node_id(i, j + 1, k + 1), + ] + ) + + return nodes, np.array(elements, dtype=int) + + +def _hex_to_tet(x_edges, y_edges, z_edges, i_offset=0): """ Tetrahedralise a rectangular prism defined by edge arrays. @@ -620,8 +718,10 @@ def _hex_to_tet(x_edges, y_edges, z_edges): axis-aligned planes (required for interface conformity). The decomposition alternates orientation based on the parity of - (i + j + k) so that shared faces between adjacent hexes are - triangulated identically. + (i + i_offset + j + k) so that shared faces between adjacent hexes + are triangulated identically, including across domain boundaries + when ``i_offset`` equals the cumulative hex count from preceding + domains. Returns ------- @@ -677,7 +777,7 @@ def node_id(i, j, k): node_id(i + 1, j + 1, k + 1), node_id(i, j + 1, k + 1), ] - pattern = pattern_a if (i + j + k) % 2 == 0 else pattern_b + pattern = pattern_a if (i + i_offset + j + k) % 2 == 0 else pattern_b for tet in pattern: elements.append([hex_verts[v] for v in tet]) diff --git a/src/pybamm/models/full_battery_models/lithium_ion/__init__.py b/src/pybamm/models/full_battery_models/lithium_ion/__init__.py index e003594858..d1af320868 100644 --- a/src/pybamm/models/full_battery_models/lithium_ion/__init__.py +++ b/src/pybamm/models/full_battery_models/lithium_ion/__init__.py @@ -25,6 +25,7 @@ from .basic_dfn import BasicDFN from .basic_dfn_2d import BasicDFN2D from .basic_dfn_2d_unstructured import BasicDFN2DUnstructured +from .basic_dfn_3d_unstructured import BasicDFN3DUnstructured from .basic_spm import BasicSPM from .basic_spm_with_3d_thermal import Basic3DThermalSPM from .basic_dfn_half_cell import BasicDFNHalfCell diff --git a/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_3d_unstructured.py b/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_3d_unstructured.py new file mode 100644 index 0000000000..96a92f65c1 --- /dev/null +++ b/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_3d_unstructured.py @@ -0,0 +1,450 @@ +# +# Basic Doyle-Fuller-Newman (DFN) Model — 3D Unstructured FVM +# +import pybamm +from pybamm.models.full_battery_models.lithium_ion.base_lithium_ion_model import ( + BaseModel, +) + + +class BasicDFN3DUnstructured(BaseModel): + """Doyle-Fuller-Newman (DFN) model on a 3D unstructured mesh. + + Extends :class:`BasicDFN2DUnstructured` to three spatial dimensions + (x, y, z) using tetrahedral elements. The through-cell direction is + *x*, the width direction is *y*, and the height direction is *z*. + + Parameters + ---------- + name : str, optional + The name of the model. + """ + + def __init__( + self, + name="Doyle-Fuller-Newman model (3D unstructured)", + ): + super().__init__(name=name) + pybamm.citations.register("Marquis2019") + + ###################### + # Variables + ###################### + Q = pybamm.Variable("Discharge capacity [A.h]") + + x = pybamm.SpatialVariable( + "x", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="lr", + ) + x_n = pybamm.SpatialVariable( + "x_n", domain="negative electrode", coord_sys="cartesian", direction="lr" + ) + x_s = pybamm.SpatialVariable( + "x_s", domain="separator", coord_sys="cartesian", direction="lr" + ) + x_p = pybamm.SpatialVariable( + "x_p", domain="positive electrode", coord_sys="cartesian", direction="lr" + ) + y_n = pybamm.SpatialVariable( + "y_n", domain="negative electrode", coord_sys="cartesian", direction="fb" + ) + y_s = pybamm.SpatialVariable( + "y_s", domain="separator", coord_sys="cartesian", direction="fb" + ) + y_p = pybamm.SpatialVariable( + "y_p", domain="positive electrode", coord_sys="cartesian", direction="fb" + ) + y = pybamm.SpatialVariable( + "y", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="fb", + ) + z_n = pybamm.SpatialVariable( + "z_n", domain="negative electrode", coord_sys="cartesian", direction="tb" + ) + z_s = pybamm.SpatialVariable( + "z_s", domain="separator", coord_sys="cartesian", direction="tb" + ) + z_p = pybamm.SpatialVariable( + "z_p", domain="positive electrode", coord_sys="cartesian", direction="tb" + ) + z = pybamm.SpatialVariable( + "z", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="tb", + ) + + c_e_n = pybamm.Variable( + "Negative electrolyte concentration [mol.m-3]", + domain="negative electrode", + ) + c_e_s = pybamm.Variable( + "Separator electrolyte concentration [mol.m-3]", + domain="separator", + ) + c_e_p = pybamm.Variable( + "Positive electrolyte concentration [mol.m-3]", + domain="positive electrode", + ) + c_e = pybamm.concatenation(c_e_n, c_e_s, c_e_p) + + phi_e_n = pybamm.Variable( + "Negative electrolyte potential [V]", + domain="negative electrode", + ) + phi_e_s = pybamm.Variable( + "Separator electrolyte potential [V]", + domain="separator", + ) + phi_e_p = pybamm.Variable( + "Positive electrolyte potential [V]", + domain="positive electrode", + ) + phi_e = pybamm.concatenation(phi_e_n, phi_e_s, phi_e_p) + + phi_s_n = pybamm.Variable( + "Negative electrode potential [V]", domain="negative electrode" + ) + phi_s_p = pybamm.Variable( + "Positive electrode potential [V]", + domain="positive electrode", + ) + c_s_n = pybamm.Variable( + "Negative particle concentration [mol.m-3]", + domain="negative particle", + auxiliary_domains={"secondary": "negative electrode"}, + ) + c_s_p = pybamm.Variable( + "Positive particle concentration [mol.m-3]", + domain="positive particle", + auxiliary_domains={"secondary": "positive electrode"}, + ) + + T = self.param.T_init + + ###################### + # Other set-up + ###################### + i_cell = self.param.current_density_with_time + + eps_n = pybamm.FunctionParameter( + "Negative electrode porosity", + {"Through-cell distance (x) [m]": x_n, "Vertical distance (z) [m]": z_n}, + ) + eps_s = pybamm.FunctionParameter( + "Separator porosity", + {"Through-cell distance (x) [m]": x_s, "Vertical distance (z) [m]": z_s}, + ) + eps_p = pybamm.FunctionParameter( + "Positive electrode porosity", + {"Through-cell distance (x) [m]": x_p, "Vertical distance (z) [m]": z_p}, + ) + eps = pybamm.concatenation(eps_n, eps_s, eps_p) + + eps_s_n = pybamm.FunctionParameter( + "Negative electrode active material volume fraction", + {"Through-cell distance (x) [m]": x_n, "Vertical distance (z) [m]": z_n}, + ) + eps_s_p = pybamm.FunctionParameter( + "Positive electrode active material volume fraction", + {"Through-cell distance (x) [m]": x_p, "Vertical distance (z) [m]": z_p}, + ) + + tor = pybamm.concatenation( + eps_n**self.param.n.b_e, eps_s**self.param.s.b_e, eps_p**self.param.p.b_e + ) + a_n = 3 * self.param.n.prim.epsilon_s_av / self.param.n.prim.R_typ + a_p = 3 * self.param.p.prim.epsilon_s_av / self.param.p.prim.R_typ + + # Interfacial reactions + c_s_surf_n = pybamm.surf(c_s_n) + sto_surf_n = c_s_surf_n / self.param.n.prim.c_max + j0_n = self.param.n.prim.j0(c_e_n, c_s_surf_n, T) + delta_phi_n = phi_s_n - phi_e_n + eta_n = delta_phi_n - self.param.n.prim.U(sto_surf_n, T) + Feta_RT_n = self.param.F * eta_n / (self.param.R * T) + j_n = 2 * j0_n * pybamm.sinh(self.param.n.prim.ne / 2 * Feta_RT_n) + + c_s_surf_p = pybamm.surf(c_s_p) + sto_surf_p = c_s_surf_p / self.param.p.prim.c_max + j0_p = self.param.p.prim.j0(c_e_p, c_s_surf_p, T) + delta_phi_p = phi_s_p - phi_e_p + eta_p = delta_phi_p - self.param.p.prim.U(sto_surf_p, T) + Feta_RT_p = self.param.F * eta_p / (self.param.R * T) + j_s = pybamm.PrimaryBroadcast(0, "separator") + j_p = 2 * j0_p * pybamm.sinh(self.param.p.prim.ne / 2 * Feta_RT_p) + + a_j_n = a_n * j_n + a_j_p = a_p * j_p + a_j = pybamm.concatenation(a_j_n, j_s, a_j_p) + + ###################### + # State of Charge + ###################### + current = self.param.current_with_time + self.rhs[Q] = current / 3600 + self.initial_conditions[Q] = pybamm.Scalar(0) + + ###################### + # Particles + ###################### + N_s_n = -self.param.n.prim.D(c_s_n, T) * pybamm.grad(c_s_n) + N_s_p = -self.param.p.prim.D(c_s_p, T) * pybamm.grad(c_s_p) + self.rhs[c_s_n] = -pybamm.div(N_s_n) + self.rhs[c_s_p] = -pybamm.div(N_s_p) + self.boundary_conditions[c_s_n] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": ( + -j_n / (self.param.F * pybamm.surf(self.param.n.prim.D(c_s_n, T))), + "Neumann", + ), + } + self.boundary_conditions[c_s_p] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": ( + -j_p / (self.param.F * pybamm.surf(self.param.p.prim.D(c_s_p, T))), + "Neumann", + ), + } + self.initial_conditions[c_s_n] = self.param.n.prim.c_init + self.initial_conditions[c_s_p] = self.param.p.prim.c_init + + c_s_n_av = pybamm.RAverage(c_s_n) + c_s_p_av = pybamm.RAverage(c_s_p) + solid_lithium_negative = pybamm.Integral(c_s_n_av * eps_s_n, [x_n, y_n, z_n]) + solid_lithium_positive = pybamm.Integral(c_s_p_av * eps_s_p, [x_p, y_p, z_p]) + total_solid_lithium = solid_lithium_negative + solid_lithium_positive + + ###################### + # Current in the solid + ###################### + sigma_eff_n = self.param.n.sigma(T) * eps_s_n**self.param.n.b_s + sigma_eff_p = self.param.p.sigma(T) * eps_s_p**self.param.p.b_s + L_scale = self.param.L_x**2 * self.param.L_z**2 + self.algebraic[phi_s_n] = L_scale * ( + pybamm.div(-sigma_eff_n * pybamm.grad(phi_s_n)) + a_j_n + ) + self.algebraic[phi_s_p] = L_scale * ( + pybamm.div(-sigma_eff_p * pybamm.grad(phi_s_p)) + a_j_p + ) + self.boundary_conditions[phi_s_n] = { + "left": (pybamm.Scalar(0), "Dirichlet"), + "right": (pybamm.Scalar(0), "Neumann"), + "top": (pybamm.Scalar(0), "Neumann"), + "bottom": (pybamm.Scalar(0), "Neumann"), + "front": (pybamm.Scalar(0), "Neumann"), + "back": (pybamm.Scalar(0), "Neumann"), + } + self.boundary_conditions[phi_s_p] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": ( + i_cell / pybamm.boundary_value(-sigma_eff_p, "right"), + "Neumann", + ), + "top": (pybamm.Scalar(0), "Neumann"), + "bottom": (pybamm.Scalar(0), "Neumann"), + "front": (pybamm.Scalar(0), "Neumann"), + "back": (pybamm.Scalar(0), "Neumann"), + } + self.initial_conditions[phi_s_n] = pybamm.Scalar(0) + self.initial_conditions[phi_s_p] = self.param.ocv_init + ###################### + # Current in the electrolyte + ###################### + kappa_eff = self.param.kappa_e(c_e, T) * tor + kappa_D_eff = kappa_eff * self.param.chiRT_over_Fc(c_e, T) + self.algebraic[phi_e] = L_scale * ( + pybamm.div(kappa_D_eff * pybamm.grad(c_e)) + - pybamm.div(kappa_eff * pybamm.grad(phi_e)) + - a_j + ) + self.boundary_conditions[phi_e] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": (pybamm.Scalar(0), "Neumann"), + "top": (pybamm.Scalar(0), "Neumann"), + "bottom": (pybamm.Scalar(0), "Neumann"), + "front": (pybamm.Scalar(0), "Neumann"), + "back": (pybamm.Scalar(0), "Neumann"), + } + self.initial_conditions[phi_e] = -self.param.n.prim.U_init + + ###################### + # Electrolyte concentration + ###################### + D_e_eff = tor * self.param.D_e(c_e, T) + self.rhs[c_e] = (1 / eps) * ( + pybamm.div(D_e_eff * pybamm.grad(c_e)) + + (1 - self.param.t_plus(c_e, T)) * a_j / self.param.F + ) + self.boundary_conditions[c_e] = { + "left": (pybamm.Scalar(0), "Neumann"), + "right": (pybamm.Scalar(0), "Neumann"), + "top": (pybamm.Scalar(0), "Neumann"), + "bottom": (pybamm.Scalar(0), "Neumann"), + "front": (pybamm.Scalar(0), "Neumann"), + "back": (pybamm.Scalar(0), "Neumann"), + } + self.initial_conditions[c_e] = self.param.c_e_init + + ###################### + # (Some) variables + ###################### + voltage = pybamm.boundary_value(phi_s_p, "top-right") + num_cells = pybamm.Parameter( + "Number of cells connected in series to make a battery" + ) + total_lithium = pybamm.Integral(c_e * eps, [x, y, z]) + self.variables = { + "Negative particle concentration [mol.m-3]": c_s_n, + "Total lithium [mol]": total_lithium, + "Negative particle surface concentration [mol.m-3]": c_s_surf_n, + "Electrolyte concentration [mol.m-3]": c_e, + "Negative electrolyte concentration [mol.m-3]": c_e_n, + "Separator electrolyte concentration [mol.m-3]": c_e_s, + "Positive electrolyte concentration [mol.m-3]": c_e_p, + "Positive particle concentration [mol.m-3]": c_s_p, + "Positive particle surface concentration [mol.m-3]": c_s_surf_p, + "Current [A]": current, + "Current variable [A]": current, + "Negative electrode potential [V]": phi_s_n, + "Electrolyte potential [V]": phi_e, + "Negative electrolyte potential [V]": phi_e_n, + "Separator electrolyte potential [V]": phi_e_s, + "Positive electrolyte potential [V]": phi_e_p, + "Positive electrode potential [V]": phi_s_p, + "Voltage [V]": voltage, + "Battery voltage [V]": voltage * num_cells, + "Time [s]": pybamm.t, + "Discharge capacity [A.h]": Q, + "x": x, + "y": y, + "z": z, + "Current density [A.m-2]": a_j, + "Electrolyte current density [A.m-2]": a_j, + "x_n": x_n, + "x_s": x_s, + "x_p": x_p, + "y_n": y_n, + "y_s": y_s, + "y_p": y_p, + "z_n": z_n, + "z_s": z_s, + "z_p": z_p, + "Negative electrode surface concentration [mol.m-3]": c_s_surf_n, + "Negative electrode surface stoichiometry": sto_surf_n, + "Positive electrode surface concentration [mol.m-3]": c_s_surf_p, + "Positive electrode surface stoichiometry": sto_surf_p, + "Positive electrode surface potential difference [V]": delta_phi_p, + "Negative electrode surface potential difference [V]": delta_phi_n, + "Positive electrode overpotential [V]": eta_p, + "Negative electrode overpotential [V]": eta_n, + "Positive electrode ocp [V]": self.param.p.prim.U(sto_surf_p, T), + "Negative electrode ocp [V]": self.param.n.prim.U(sto_surf_n, T), + "Positive electrode current density [A.m-2]": j_p, + "Negative electrode current density [A.m-2]": j_n, + "Electrolyte flux [mol.m-2.s-1]": D_e_eff, + "Positive solid lithium [mol]": solid_lithium_positive, + "Negative solid lithium [mol]": solid_lithium_negative, + "Total solid lithium [mol]": total_solid_lithium, + } + self.events += [ + pybamm.Event("Minimum voltage [V]", voltage - self.param.voltage_low_cut), + pybamm.Event("Maximum voltage [V]", self.param.voltage_high_cut - voltage), + ] + + @property + def default_geometry(self): + y_3d = pybamm.SpatialVariable( + "y_3d", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="fb", + ) + z_3d = pybamm.SpatialVariable( + "z_3d", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="tb", + ) + return { + "negative electrode": { + "x_n": {"min": 0, "max": self.param.n.L}, + y_3d: {"min": 0, "max": self.param.L_y}, + z_3d: {"min": 0, "max": self.param.L_z}, + }, + "separator": { + "x_s": { + "min": self.param.n.L, + "max": self.param.n.L + self.param.s.L, + }, + y_3d: {"min": 0, "max": self.param.L_y}, + z_3d: {"min": 0, "max": self.param.L_z}, + }, + "positive electrode": { + "x_p": { + "min": self.param.n.L + self.param.s.L, + "max": self.param.n.L + self.param.s.L + self.param.p.L, + }, + y_3d: {"min": 0, "max": self.param.L_y}, + z_3d: {"min": 0, "max": self.param.L_z}, + }, + "positive particle": { + "r_p": {"min": 0, "max": self.param.p.prim.R_typ}, + }, + "negative particle": { + "r_n": {"min": 0, "max": self.param.n.prim.R_typ}, + }, + "current collector": { + "z": {"position": 0}, + }, + } + + @property + def default_spatial_methods(self): + return { + "negative electrode": pybamm.FiniteVolumeUnstructured(), + "separator": pybamm.FiniteVolumeUnstructured(), + "positive electrode": pybamm.FiniteVolumeUnstructured(), + "positive particle": pybamm.FiniteVolume(), + "negative particle": pybamm.FiniteVolume(), + "current collector": pybamm.ZeroDimensionalSpatialMethod(), + } + + @property + def default_submesh_types(self): + return { + "negative electrode": pybamm.UnstructuredMeshGenerator(), + "separator": pybamm.UnstructuredMeshGenerator(), + "positive electrode": pybamm.UnstructuredMeshGenerator(), + "positive particle": pybamm.Uniform1DSubMesh, + "negative particle": pybamm.Uniform1DSubMesh, + "current collector": pybamm.SubMesh0D, + } + + @property + def default_var_pts(self): + y_3d = pybamm.SpatialVariable( + "y_3d", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="fb", + ) + z_3d = pybamm.SpatialVariable( + "z_3d", + domain=["negative electrode", "separator", "positive electrode"], + coord_sys="cartesian", + direction="tb", + ) + return { + "x_n": 5, + "x_s": 5, + "x_p": 5, + "r_p": 10, + "r_n": 10, + y_3d: 3, + z_3d: 3, + } diff --git a/src/pybamm/spatial_methods/finite_volume_unstructured.py b/src/pybamm/spatial_methods/finite_volume_unstructured.py index 8b46f95b88..4a06d9ed43 100644 --- a/src/pybamm/spatial_methods/finite_volume_unstructured.py +++ b/src/pybamm/spatial_methods/finite_volume_unstructured.py @@ -936,10 +936,8 @@ def _internal_neumann_unstructured( break if interface is None: - raise ValueError( - "No interface data found between the left and right meshes. " - "Run compute_interface_data() during mesh construction." - ) + n_left = left_mesh.npts + return pybamm.Vector(np.zeros(n_left * repeats)) n_faces = len(interface["left_cells"]) n_left = left_mesh.npts diff --git a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py index df869095d3..3a73e7a0cb 100644 --- a/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py +++ b/tests/integration/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py @@ -79,3 +79,23 @@ def test_solves_and_matches_structured(self): V_u = sol_u["Voltage [V]"](t=t_eval) np.testing.assert_allclose(V_u, V_s, atol=5e-3) + + +class TestBasicDFN3DUnstructured: + def test_solves_and_matches_2d(self): + import numpy as np + + t_eval = np.linspace(0, 3600, 20) + + model_2d = pybamm.lithium_ion.BasicDFN2DUnstructured(element_type="quad") + sim_2d = pybamm.Simulation(model_2d) + sol_2d = sim_2d.solve(t_eval) + + model_3d = pybamm.lithium_ion.BasicDFN3DUnstructured() + sim_3d = pybamm.Simulation(model_3d) + sol_3d = sim_3d.solve(t_eval) + + V_2d = sol_2d["Voltage [V]"](t=t_eval) + V_3d = sol_3d["Voltage [V]"](t=t_eval) + + np.testing.assert_allclose(V_3d, V_2d, atol=5e-3) diff --git a/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py b/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py index 6a61f19974..6ce5b607f0 100644 --- a/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py +++ b/tests/unit/test_models/test_full_battery_models/test_lithium_ion/test_basic_models.py @@ -29,3 +29,7 @@ def test_dfn_2d(self): def test_dfn_2d_unstructured(self): model = pybamm.lithium_ion.BasicDFN2DUnstructured(element_type="quad") model.check_well_posedness() + + def test_dfn_3d_unstructured(self): + model = pybamm.lithium_ion.BasicDFN3DUnstructured() + model.check_well_posedness() From 4534639ac0666814381c5cc56d53442a49552e69 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Wed, 25 Feb 2026 18:40:50 -0800 Subject: [PATCH 07/25] improve plotting --- src/pybamm/meshes/unstructured_submesh.py | 39 +++++++++++++++++++++++ src/pybamm/solvers/processed_variable.py | 11 ++++++- 2 files changed, 49 insertions(+), 1 deletion(-) diff --git a/src/pybamm/meshes/unstructured_submesh.py b/src/pybamm/meshes/unstructured_submesh.py index 8b94aec72c..cf56f5be32 100644 --- a/src/pybamm/meshes/unstructured_submesh.py +++ b/src/pybamm/meshes/unstructured_submesh.py @@ -284,6 +284,45 @@ def _identify_boundary_faces(self): if len(indices) > 0: self.boundary_faces[name] = indices + def boundary_polygon(self): + """Return ordered boundary vertices as an (M, 2) array (2D only). + + Walks the boundary edges to produce a closed polygon suitable for + point-in-domain tests with ``matplotlib.path.Path``. + """ + if self.dimension != 2: + return None + + bnd_start = self._boundary_face_start + bnd_edges = self.faces[bnd_start:] + if len(bnd_edges) == 0: + return None + + adj: dict[int, list[tuple[int, int]]] = {} + for i, edge in enumerate(bnd_edges): + v0, v1 = int(edge[0]), int(edge[1]) + adj.setdefault(v0, []).append((i, v1)) + adj.setdefault(v1, []).append((i, v0)) + + visited: set[int] = set() + polygon_verts = [] + current = int(bnd_edges[0][0]) + polygon_verts.append(current) + + while True: + found = False + for edge_idx, next_v in adj[current]: + if edge_idx not in visited: + visited.add(edge_idx) + polygon_verts.append(next_v) + current = next_v + found = True + break + if not found: + break + + return self.nodes[polygon_verts] + # ====================================================================== # Mesh generators diff --git a/src/pybamm/solvers/processed_variable.py b/src/pybamm/solvers/processed_variable.py index 2f0adc6f9a..130a93e601 100644 --- a/src/pybamm/solvers/processed_variable.py +++ b/src/pybamm/solvers/processed_variable.py @@ -1042,7 +1042,8 @@ def _interpolate_spatial(self, values, query_pts): Boundary face centroids are added to the interpolation cloud so the convex hull covers the full domain. Any residual - extrapolation uses nearest-neighbor. + extrapolation uses nearest-neighbor. For non-convex domains, + query points outside the mesh boundary polygon are set to NaN. """ from scipy.interpolate import LinearNDInterpolator, NearestNDInterpolator @@ -1056,6 +1057,14 @@ def _interpolate_spatial(self, values, query_pts): nearest = NearestNDInterpolator(pts, vals) result[mask] = nearest(query_pts[mask]) + poly = self.mesh.boundary_polygon() + if poly is not None: + from matplotlib.path import Path + + path = Path(poly) + outside = ~path.contains_points(query_pts[:, :2]) + result[outside] = np.nan + return result def _data_at_time(self, t): From 35a24b9457814cadce1cb27ff3cf19a6ac1e6326 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Thu, 26 Feb 2026 16:08:58 -0800 Subject: [PATCH 08/25] fix plotting and speed up discretisation --- src/pybamm/discretisations/discretisation.py | 37 +-- src/pybamm/meshes/unstructured_submesh.py | 212 +++++++++++++----- src/pybamm/plotting/quick_plot.py | 92 +++++--- src/pybamm/solvers/processed_variable.py | 93 +++++--- .../finite_volume_unstructured.py | 87 +++---- .../test_meshes/test_unstructured_submesh.py | 4 +- tests/unit/test_plotting/test_quick_plot.py | 2 +- 7 files changed, 361 insertions(+), 166 deletions(-) diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py index 4a6c5b0a7b..ea365a00ac 100644 --- a/src/pybamm/discretisations/discretisation.py +++ b/src/pybamm/discretisations/discretisation.py @@ -1003,24 +1003,22 @@ def _process_symbol(self, symbol): elif isinstance(symbol, pybamm.UnaryOperator): child = symbol.child - disc_child = self.process_symbol(child) - if child.domain != []: + # Intercept div(grad(u)) and div(D*grad(u)) before processing + # children, to avoid the expensive Green-Gauss gradient assembly. + if isinstance(symbol, pybamm.Divergence) and child.domain != []: child_spatial_method = self.spatial_methods[child.domain[0]] - - if isinstance(symbol, pybamm.Gradient): - return child_spatial_method.gradient(child, disc_child, self.bcs) - - elif isinstance(symbol, pybamm.Divergence): - if isinstance( - child_spatial_method, pybamm.FiniteVolumeUnstructured - ) and isinstance(child, pybamm.Multiplication): - left_c, right_c = child.children + if isinstance(child_spatial_method, pybamm.FiniteVolumeUnstructured): grad_sym = None coeff_sym = None - if isinstance(right_c, pybamm.Gradient): - grad_sym, coeff_sym = right_c, left_c - elif isinstance(left_c, pybamm.Gradient): - grad_sym, coeff_sym = left_c, right_c + if isinstance(child, pybamm.Gradient): + grad_sym = child + coeff_sym = pybamm.Scalar(1) + elif isinstance(child, pybamm.Multiplication): + left_c, right_c = child.children + if isinstance(right_c, pybamm.Gradient): + grad_sym, coeff_sym = right_c, left_c + elif isinstance(left_c, pybamm.Gradient): + grad_sym, coeff_sym = left_c, right_c if grad_sym is not None: disc_coeff = self.process_symbol(coeff_sym) disc_u = self.process_symbol(grad_sym.child) @@ -1031,6 +1029,15 @@ def _process_symbol(self, symbol): disc_u, self.bcs, ) + + disc_child = self.process_symbol(child) + if child.domain != []: + child_spatial_method = self.spatial_methods[child.domain[0]] + + if isinstance(symbol, pybamm.Gradient): + return child_spatial_method.gradient(child, disc_child, self.bcs) + + elif isinstance(symbol, pybamm.Divergence): return child_spatial_method.divergence(child, disc_child, self.bcs) elif isinstance(symbol, pybamm.Laplacian): diff --git a/src/pybamm/meshes/unstructured_submesh.py b/src/pybamm/meshes/unstructured_submesh.py index cf56f5be32..91ecee77d3 100644 --- a/src/pybamm/meshes/unstructured_submesh.py +++ b/src/pybamm/meshes/unstructured_submesh.py @@ -135,40 +135,62 @@ def _build_face_connectivity(self): else: n_verts_per_face = self.dimension - face_dict = {} # canonical key -> (owner_cell, original_verts) - - internal_owner = [] - internal_neighbor = [] - internal_face_verts = [] - - for cell_idx, cell_verts in enumerate(self.elements): - for face_verts in self._cell_faces(cell_verts): - key = tuple(sorted(face_verts)) - - if key in face_dict: - other_cell, orig_verts = face_dict.pop(key) - internal_owner.append(other_cell) - internal_neighbor.append(cell_idx) - internal_face_verts.append(orig_verts) - else: - face_dict[key] = (cell_idx, tuple(face_verts)) - - # Remaining entries are boundary faces - boundary_owner_list = [] - boundary_face_verts = [] - for _key, (cell_idx, orig_verts) in face_dict.items(): - boundary_owner_list.append(cell_idx) - boundary_face_verts.append(orig_verts) + elems = self.elements + n_cells = len(elems) - n_internal = len(internal_owner) - n_boundary = len(boundary_owner_list) + # Build all faces at once using local face definitions + if self.element_type == "quad": + n_fpc = 4 # faces per cell + idx = np.arange(4) + local = np.stack([idx, (idx + 1) % 4], axis=1) # (4, 2) + elif self.element_type == "triangle": + local = np.array([[1, 2], [0, 2], [0, 1]]) # skip vertex 0, 1, 2 + elif self.element_type == "tetrahedron": + local = np.array([[1, 2, 3], [0, 2, 3], [0, 1, 3], [0, 1, 2]]) + elif self.element_type == "hexahedron": + local = np.array(self._HEX_FACES) + n_fpc = len(local) + + all_faces = elems[:, local].reshape(-1, n_verts_per_face) + cell_ids = np.repeat(np.arange(n_cells), n_fpc) + + # Canonical keys: sort vertex indices within each face + sorted_faces = np.sort(all_faces, axis=1) + + # Find unique faces and which are shared (internal) vs single (boundary) + _, inverse, counts = np.unique( + sorted_faces, axis=0, return_inverse=True, return_counts=True + ) + + is_internal = counts[inverse] == 2 + is_boundary = counts[inverse] == 1 + + # For internal faces, we need owner/neighbor pairs. + # Group by unique face index; first occurrence is owner, second is neighbor. + internal_mask = is_internal + int_inv = inverse[internal_mask] + int_cells = cell_ids[internal_mask] + int_faces_raw = all_faces[internal_mask] + + # Sort by unique-face-id to pair them up: [owner0, neighbor0, owner1, neighbor1, ...] + order = np.argsort(int_inv, kind="stable") + int_cells_sorted = int_cells[order] + int_faces_sorted = int_faces_raw[order] + + internal_owner = int_cells_sorted[0::2] + internal_neighbor = int_cells_sorted[1::2] + internal_face_verts = int_faces_sorted[0::2] + + # Boundary faces + bnd_face_verts = all_faces[is_boundary] + bnd_owners = cell_ids[is_boundary] - all_face_verts = internal_face_verts + boundary_face_verts - all_owner = internal_owner + boundary_owner_list + n_internal = len(internal_owner) + n_boundary = len(bnd_owners) - self.faces = np.array(all_face_verts, dtype=int).reshape(-1, n_verts_per_face) - self.face_owner = np.array(all_owner, dtype=int) - self.face_neighbor = np.array(internal_neighbor, dtype=int) + self.faces = np.concatenate([internal_face_verts, bnd_face_verts], axis=0) + self.face_owner = np.concatenate([internal_owner, bnd_owners]) + self.face_neighbor = internal_neighbor self.n_internal_faces = n_internal self._n_boundary_faces = n_boundary self._boundary_face_start = n_internal @@ -284,15 +306,19 @@ def _identify_boundary_faces(self): if len(indices) > 0: self.boundary_faces[name] = indices - def boundary_polygon(self): - """Return ordered boundary vertices as an (M, 2) array (2D only). + def boundary_loops(self): + """Return boundary loops as a list of ``matplotlib.path.Path`` (2D only). - Walks the boundary edges to produce a closed polygon suitable for - point-in-domain tests with ``matplotlib.path.Path``. + Walks boundary edges to extract one or more closed loops. The first + path is the outer boundary (largest area); subsequent paths are holes. + Use this to test containment: a point is in the domain if it is inside + the outer loop and outside all hole loops. """ if self.dimension != 2: return None + from matplotlib.path import Path + bnd_start = self._boundary_face_start bnd_edges = self.faces[bnd_start:] if len(bnd_edges) == 0: @@ -305,23 +331,109 @@ def boundary_polygon(self): adj.setdefault(v1, []).append((i, v0)) visited: set[int] = set() - polygon_verts = [] - current = int(bnd_edges[0][0]) - polygon_verts.append(current) - - while True: - found = False - for edge_idx, next_v in adj[current]: - if edge_idx not in visited: - visited.add(edge_idx) - polygon_verts.append(next_v) - current = next_v - found = True + loops: list[list[int]] = [] + + for start_edge_idx in range(len(bnd_edges)): + if start_edge_idx in visited: + continue + start_v = int(bnd_edges[start_edge_idx][0]) + loop = [start_v] + current = start_v + while True: + found = False + for edge_idx, next_v in adj[current]: + if edge_idx not in visited: + visited.add(edge_idx) + loop.append(next_v) + current = next_v + found = True + break + if not found: break - if not found: - break + loops.append(loop) + + def signed_area(pts): + x, y = pts[:, 0], pts[:, 1] + return 0.5 * np.sum(x[:-1] * y[1:] - x[1:] * y[:-1]) + + loop_data = [] + for loop in loops: + pts = self.nodes[loop] + sa = signed_area(pts) + loop_data.append((abs(sa), pts)) + + loop_data.sort(key=lambda t: t[0], reverse=True) + + paths = [] + for pts in (ld[1] for ld in loop_data): + codes = [Path.LINETO] * len(pts) + codes[0] = Path.MOVETO + codes[-1] = Path.CLOSEPOLY + paths.append(Path(pts, codes)) + return paths + + def contains_points_3d(self, query_pts): + """Test whether 3D points lie inside the mesh domain. + + Uses the generalized winding number (Van Oosterom--Strackee signed + solid angle sum over all boundary triangles). Points inside the + domain return ``True``; points outside or inside internal cavities + return ``False``. + """ + query_pts = np.asarray(query_pts, dtype=np.float64) + bnd_start = self._boundary_face_start + bnd_fv = self.faces[bnd_start:] + bnd_normals = self.face_normals[bnd_start:] + n_vpf = bnd_fv.shape[1] + + if n_vpf == 3: + tri_idx = bnd_fv + tri_normals = bnd_normals + elif n_vpf == 4: + tri_idx = np.concatenate( + [bnd_fv[:, [0, 1, 2]], bnd_fv[:, [0, 2, 3]]], axis=0 + ) + tri_normals = np.concatenate([bnd_normals, bnd_normals], axis=0) + else: + raise ValueError( + f"contains_points_3d: unsupported face with {n_vpf} vertices" + ) + + v0 = self.nodes[tri_idx[:, 0]] + v1 = self.nodes[tri_idx[:, 1]] + v2 = self.nodes[tri_idx[:, 2]] + + # Ensure consistent CCW orientation from outside (matching outward normals) + cross = np.cross(v1 - v0, v2 - v0) + flip = np.sum(cross * tri_normals, axis=1) < 0 + v1_fixed = v1.copy() + v2_fixed = v2.copy() + v1_fixed[flip] = v2[flip] + v2_fixed[flip] = v1[flip] + + n_query = len(query_pts) + winding = np.zeros(n_query) + + for i in range(len(tri_idx)): + a = v0[i] - query_pts + b = v1_fixed[i] - query_pts + c = v2_fixed[i] - query_pts + + an = np.linalg.norm(a, axis=1) + bn = np.linalg.norm(b, axis=1) + cn = np.linalg.norm(c, axis=1) + + num = np.einsum("ij,ij->i", a, np.cross(b, c)) + den = ( + an * bn * cn + + np.einsum("ij,ij->i", a, b) * cn + + np.einsum("ij,ij->i", a, c) * bn + + np.einsum("ij,ij->i", b, c) * an + ) + + winding += 2.0 * np.arctan2(num, den) - return self.nodes[polygon_verts] + return winding > 2.0 * np.pi # ====================================================================== diff --git a/src/pybamm/plotting/quick_plot.py b/src/pybamm/plotting/quick_plot.py index 31c7423716..fd3b9bf34c 100644 --- a/src/pybamm/plotting/quick_plot.py +++ b/src/pybamm/plotting/quick_plot.py @@ -679,20 +679,26 @@ def plot(self, t, dynamic=False): vmin, vmax = self.variable_limits[key] # store the plot and the var data (for testing) as cant access # z data from QuadMesh or QuadContourSet object - if self.is_y_z[key] is True: - self.plots[key][0][0] = ax.pcolormesh( - x, - y, - var, - vmin=vmin, - vmax=vmax, - shading=self.shading, - ) + is_unstructured = isinstance( + variable, pybamm.ProcessedVariableUnstructuredFVM + ) + if self.is_y_z[key] is True or is_unstructured: + kw = dict(vmin=vmin, vmax=vmax, shading=self.shading) + if is_unstructured: + import matplotlib + + cmap_copy = matplotlib.colormaps["viridis"].copy() + cmap_copy.set_bad("white") + kw["cmap"] = cmap_copy + self.plots[key][0][0] = ax.pcolormesh(x, y, var, **kw) else: self.plots[key][0][0] = ax.contourf( x, y, var, levels=100, vmin=vmin, vmax=vmax ) self.plots[key][0][1] = var + if is_unstructured: + self._overlay_mesh_wireframe(ax, variable) + ax.set_aspect("equal") if vmin is None and vmax is None: vmin = ax_min(var) vmax = ax_max(var) @@ -708,29 +714,31 @@ def plot(self, t, dynamic=False): if vmax is None: vmax = ax_max(variable(t_in_seconds)) norm = colors.Normalize(vmin=vmin, vmax=vmax) - cmap = plt.cm.viridis + import matplotlib.pyplot as _plt + + cmap = _plt.cm.viridis s1, xx1, yy1, zz1, s2, xx2, yy2, zz2 = variable.get_3d_slices( t_in_seconds ) + fc1 = self._slice_facecolors(s1, cmap, norm) + fc2 = self._slice_facecolors(s2, cmap, norm) ax.plot_surface( xx1, yy1, zz1, - facecolors=cmap(norm(s1)), + facecolors=fc1, rstride=1, cstride=1, shade=False, - alpha=0.85, ) ax.plot_surface( xx2, yy2, zz2, - facecolors=cmap(norm(s2)), + facecolors=fc2, rstride=1, cstride=1, shade=False, - alpha=0.85, ) ax.set_xlabel("$x$") ax.set_ylabel("$y$") @@ -783,6 +791,30 @@ def plot(self, t, dynamic=False): bottom = max(legend_top, slider_top) self.gridspec.tight_layout(self.fig, rect=[0, bottom, 1, 1]) + @staticmethod + def _slice_facecolors(data, cmap, norm, base_alpha=0.85): + """Compute RGBA facecolors for ``plot_surface``, with NaN faces + rendered fully transparent so that cavities appear as holes.""" + import numpy as np + + nan_mask = np.isnan(data) + fc = cmap(norm(np.where(nan_mask, 0.0, data))) + fc[..., 3] = np.where(nan_mask, 0.0, base_alpha) + return fc + + def _overlay_mesh_wireframe(self, ax, variable): + """Draw mesh element edges as a light wireframe on a 2D axis.""" + from matplotlib.collections import PolyCollection + + mesh = variable.mesh + if mesh.dimension != 2: + return + verts = mesh.nodes[mesh.elements] * self.spatial_factor + poly = PolyCollection( + verts, facecolors="none", edgecolors=(0, 0, 0, 0.12), linewidths=0.3 + ) + ax.add_collection(poly) + def _plot_3d_quiver(self, ax, variable, t, key, cm, colors): """Render quiver arrows on two orthogonal 3D slice planes.""" sf = self.spatial_factor @@ -999,20 +1031,26 @@ def slider_update(self, t): var = variable(time_in_seconds, **spatial_vars).T # store the plot and the var data (for testing) as cant access # z data from QuadMesh or QuadContourSet object - if self.is_y_z[key] is True: - self.plots[key][0][0] = ax.pcolormesh( - x, - y, - var, - vmin=vmin, - vmax=vmax, - shading=self.shading, - ) + is_unstructured = isinstance( + variable, pybamm.ProcessedVariableUnstructuredFVM + ) + if self.is_y_z[key] is True or is_unstructured: + kw = dict(vmin=vmin, vmax=vmax, shading=self.shading) + if is_unstructured: + import matplotlib + + cmap_copy = matplotlib.colormaps["viridis"].copy() + cmap_copy.set_bad("white") + kw["cmap"] = cmap_copy + self.plots[key][0][0] = ax.pcolormesh(x, y, var, **kw) else: self.plots[key][0][0] = ax.contourf( x, y, var, levels=100, vmin=vmin, vmax=vmax ) self.plots[key][0][1] = var + if is_unstructured: + self._overlay_mesh_wireframe(ax, variable) + ax.set_aspect("equal") if (vmin, vmax) == (None, None): vmin = ax_min(var) vmax = ax_max(var) @@ -1035,25 +1073,25 @@ def slider_update(self, t): s1, xx1, yy1, zz1, s2, xx2, yy2, zz2 = variable.get_3d_slices( time_in_seconds ) + fc1 = self._slice_facecolors(s1, cmap, norm) + fc2 = self._slice_facecolors(s2, cmap, norm) ax.plot_surface( xx1, yy1, zz1, - facecolors=cmap(norm(s1)), + facecolors=fc1, rstride=1, cstride=1, shade=False, - alpha=0.85, ) ax.plot_surface( xx2, yy2, zz2, - facecolors=cmap(norm(s2)), + facecolors=fc2, rstride=1, cstride=1, shade=False, - alpha=0.85, ) ax.set_xlabel("$x$") ax.set_ylabel("$y$") diff --git a/src/pybamm/solvers/processed_variable.py b/src/pybamm/solvers/processed_variable.py index 130a93e601..b6f99c5f6d 100644 --- a/src/pybamm/solvers/processed_variable.py +++ b/src/pybamm/solvers/processed_variable.py @@ -957,7 +957,8 @@ class ProcessedVariableUnstructuredFVM(ProcessedVariable): created so that ``solution.plot()`` works out of the box. """ - N_VIS = 50 + N_VIS = 200 + N_VIS_3D = 80 def __init__( self, @@ -982,19 +983,21 @@ def __init__( nodes = mesh.nodes x_min, x_max = nodes[:, 0].min(), nodes[:, 0].max() + n_vis = self.N_VIS_3D if mesh.dimension == 3 else self.N_VIS self.first_dimension = "x" - self.first_dim_pts = np.linspace(x_min, x_max, self.N_VIS) - self.first_dim_size = self.N_VIS + self.first_dim_pts = np.linspace(x_min, x_max, n_vis) + self.first_dim_size = n_vis if mesh.dimension == 3: + n3 = self.N_VIS_3D y_min, y_max = nodes[:, 1].min(), nodes[:, 1].max() z_min, z_max = nodes[:, 2].min(), nodes[:, 2].max() self.second_dimension = "y" - self.second_dim_pts = np.linspace(y_min, y_max, self.N_VIS) - self.second_dim_size = self.N_VIS + self.second_dim_pts = np.linspace(y_min, y_max, n3) + self.second_dim_size = n3 self.third_dimension = "z" - self.third_dim_pts = np.linspace(z_min, z_max, self.N_VIS) - self.third_dim_size = self.N_VIS + self.third_dim_pts = np.linspace(z_min, z_max, n3) + self.third_dim_size = n3 self._slice_positions = { "y": 0.5 * (y_min + y_max), "z": 0.5 * (z_min + z_max), @@ -1025,31 +1028,65 @@ def initialise(self): fill_value="extrapolate", ) - def _augmented_points_and_values(self, values): - """Append boundary face centroids (with owning-cell values) to the - cell centroid cloud so that LinearNDInterpolator's convex hull - reaches the true mesh boundary.""" - mesh = self.mesh - bnd_start = mesh._boundary_face_start - bnd_centroids = mesh.face_centroids[bnd_start:] - bnd_owners = mesh.face_owner[bnd_start:] - pts = np.concatenate([mesh.cell_centroids, bnd_centroids], axis=0) - vals = np.concatenate([values, values[bnd_owners]]) - return pts, vals + def _augmented_points(self): + """Return the interpolation point cloud (cell centroids + boundary + face centroids) and the index array for mapping cell values to + boundary face values. Cached after first call.""" + if not hasattr(self, "_aug_pts"): + mesh = self.mesh + bnd_start = mesh._boundary_face_start + bnd_centroids = mesh.face_centroids[bnd_start:] + self._aug_pts = np.concatenate([mesh.cell_centroids, bnd_centroids], axis=0) + self._aug_bnd_owners = mesh.face_owner[bnd_start:] + return self._aug_pts, self._aug_bnd_owners + + def _get_triangulation(self): + """Return a cached Delaunay triangulation of the augmented point cloud.""" + if not hasattr(self, "_cached_tri"): + from scipy.spatial import Delaunay + + pts, _ = self._augmented_points() + self._cached_tri = Delaunay(pts) + return self._cached_tri + + def _get_boundary_mask(self, query_pts): + """Return a boolean mask of query points outside the domain. + + * **2D** — uses ``boundary_loops()`` (cached). + * **3D** — uses the generalized winding number via + ``contains_points_3d`` (not cached because different slices + have different query points). + """ + if self.mesh.dimension == 3: + inside = self.mesh.contains_points_3d(query_pts) + return ~inside + + if not hasattr(self, "_cached_outside_mask"): + loops = self.mesh.boundary_loops() + if loops is not None and len(loops) > 0: + pts2d = query_pts[:, :2] + inside_outer = loops[0].contains_points(pts2d) + outside = ~inside_outer + for hole_path in loops[1:]: + outside |= hole_path.contains_points(pts2d) + self._cached_outside_mask = outside + else: + self._cached_outside_mask = None + return self._cached_outside_mask def _interpolate_spatial(self, values, query_pts): """Interpolate cell-centered data to query points. - Boundary face centroids are added to the interpolation cloud - so the convex hull covers the full domain. Any residual - extrapolation uses nearest-neighbor. For non-convex domains, - query points outside the mesh boundary polygon are set to NaN. + The Delaunay triangulation and boundary mask are computed once + and cached. Only the interpolated values change per call. """ from scipy.interpolate import LinearNDInterpolator, NearestNDInterpolator - pts, vals = self._augmented_points_and_values(values) + pts, bnd_owners = self._augmented_points() + vals = np.concatenate([values, values[bnd_owners]]) - linear = LinearNDInterpolator(pts, vals) + tri = self._get_triangulation() + linear = LinearNDInterpolator(tri, vals) result = linear(query_pts) mask = np.isnan(result) @@ -1057,12 +1094,8 @@ def _interpolate_spatial(self, values, query_pts): nearest = NearestNDInterpolator(pts, vals) result[mask] = nearest(query_pts[mask]) - poly = self.mesh.boundary_polygon() - if poly is not None: - from matplotlib.path import Path - - path = Path(poly) - outside = ~path.contains_points(query_pts[:, :2]) + outside = self._get_boundary_mask(query_pts) + if outside is not None: result[outside] = np.nan return result diff --git a/src/pybamm/spatial_methods/finite_volume_unstructured.py b/src/pybamm/spatial_methods/finite_volume_unstructured.py index 4a06d9ed43..02e3c7f8f4 100644 --- a/src/pybamm/spatial_methods/finite_volume_unstructured.py +++ b/src/pybamm/spatial_methods/finite_volume_unstructured.py @@ -6,11 +6,16 @@ face-cell connectivity as sparse matrices. """ +import logging +import time + import numpy as np from scipy.sparse import coo_matrix, csr_matrix, diags, eye, kron import pybamm +logger = logging.getLogger(__name__) + class FiniteVolumeUnstructured(pybamm.SpatialMethod): """ @@ -195,6 +200,7 @@ def div_D_grad(self, div_symbol, grad_child, disc_D, disc_u, boundary_conditions Internal-face fluxes use arithmetic-mean interpolation of ``D`` to faces and a standard two-point difference for ``grad(u)``. """ + _t0 = time.perf_counter() domain = div_symbol.domain submesh = self.mesh[domain] n = submesh.npts @@ -240,10 +246,14 @@ def div_D_grad(self, div_symbol, grad_child, disc_D, disc_u, boundary_conditions shape=(n, n_int), ) - G_f = csr_matrix(kron(eye(repeats, dtype=np.float64), G)) - W_f = csr_matrix(kron(eye(repeats, dtype=np.float64), W)) - S_f = csr_matrix(kron(eye(repeats, dtype=np.float64), S)) - geo_f = np.tile(geo, repeats) + if repeats == 1: + G_f, W_f, S_f = G, W, S + geo_f = geo + else: + G_f = csr_matrix(kron(eye(repeats, dtype=np.float64), G)) + W_f = csr_matrix(kron(eye(repeats, dtype=np.float64), W)) + S_f = csr_matrix(kron(eye(repeats, dtype=np.float64), S)) + geo_f = np.tile(geo, repeats) u_diff = pybamm.Matrix(G_f) @ disc_u is_scalar_D = isinstance(disc_D, pybamm.Scalar) or ( @@ -272,27 +282,24 @@ def div_D_grad(self, div_symbol, grad_child, disc_D, disc_u, boundary_conditions (np.ones(n_bnd), (np.arange(n_bnd), bnd_own)), shape=(n_bnd, n), ) - E_f = csr_matrix(kron(eye(repeats, dtype=np.float64), E)) P = csr_matrix( (np.ones(n_bnd), (bnd_own, np.arange(n_bnd))), shape=(n, n_bnd), ) - P_f = csr_matrix(kron(eye(repeats, dtype=np.float64), P)) + if repeats == 1: + E_f, P_f = E, P + else: + E_f = csr_matrix(kron(eye(repeats, dtype=np.float64), E)) + P_f = csr_matrix(kron(eye(repeats, dtype=np.float64), P)) D_bnd = disc_D if is_scalar_D else pybamm.Matrix(E_f) @ disc_D if bc_type == "Dirichlet": - geo_bnd = np.array( - [ - submesh.face_areas[fi] - / np.linalg.norm( - submesh.face_centroids[fi] - - submesh.cell_centroids[bnd_own[j]] - ) - / vol[bnd_own[j]] - for j, fi in enumerate(fi_arr) - ] + delta = ( + submesh.face_centroids[fi_arr] - submesh.cell_centroids[bnd_own] ) - geo_bnd_f = np.tile(geo_bnd, repeats) + d_perp = np.linalg.norm(delta, axis=1) + geo_bnd = submesh.face_areas[fi_arr] / d_perp / vol[bnd_own] + geo_bnd_f = np.tile(geo_bnd, repeats) if repeats > 1 else geo_bnd u_bnd = pybamm.Matrix(E_f) @ disc_u bc_rhs = bc_rhs + pybamm.Matrix(P_f) @ ( @@ -300,17 +307,19 @@ def div_D_grad(self, div_symbol, grad_child, disc_D, disc_u, boundary_conditions ) elif bc_type == "Neumann" and bc_value != pybamm.Scalar(0): - a_over_v = np.array( - [ - submesh.face_areas[fi] / vol[bnd_own[j]] - for j, fi in enumerate(fi_arr) - ] - ) - a_over_v_f = np.tile(a_over_v, repeats) + a_over_v = submesh.face_areas[fi_arr] / vol[bnd_own] + a_over_v_f = np.tile(a_over_v, repeats) if repeats > 1 else a_over_v bc_rhs = bc_rhs + pybamm.Matrix(P_f) @ ( D_bnd * bc_value * pybamm.Vector(a_over_v_f) ) + logger.debug( + "div_D_grad: %.3fs (n=%d, n_int=%d, repeats=%d)", + time.perf_counter() - _t0, + n, + n_int, + repeats, + ) return result + bc_rhs def _apply_bcs_to_laplacian(self, submesh, L, bc_rhs, bcs): @@ -331,28 +340,24 @@ def _apply_bcs_to_laplacian(self, submesh, L, bc_rhs, bcs): owners = submesh.face_owner[face_indices] if bc_type == "Dirichlet": - coeffs = np.empty(n_bnd) - for j, fi in enumerate(face_indices): - cell = owners[j] - area = submesh.face_areas[fi] - vol = submesh.cell_volumes[cell] - d_perp = np.linalg.norm( - submesh.face_centroids[fi] - submesh.cell_centroids[cell] - ) - coeff = area / d_perp - L[cell, cell] -= coeff / vol - coeffs[j] = coeff / vol + delta = ( + submesh.face_centroids[face_indices] + - submesh.cell_centroids[owners] + ) + d_perp = np.linalg.norm(delta, axis=1) + coeffs = ( + submesh.face_areas[face_indices] + / d_perp + / submesh.cell_volumes[owners] + ) + for j in range(n_bnd): + L[owners[j], owners[j]] -= coeffs[j] bc_rhs = bc_rhs + self._bc_contribution( n, n_bnd, owners, coeffs, bc_value ) elif bc_type == "Neumann": - coeffs = np.array( - [ - submesh.face_areas[fi] / submesh.cell_volumes[owners[j]] - for j, fi in enumerate(face_indices) - ] - ) + coeffs = submesh.face_areas[face_indices] / submesh.cell_volumes[owners] bc_rhs = bc_rhs + self._bc_contribution( n, n_bnd, owners, coeffs, bc_value ) diff --git a/tests/unit/test_meshes/test_unstructured_submesh.py b/tests/unit/test_meshes/test_unstructured_submesh.py index b62db5f277..5ca41ef576 100644 --- a/tests/unit/test_meshes/test_unstructured_submesh.py +++ b/tests/unit/test_meshes/test_unstructured_submesh.py @@ -276,7 +276,7 @@ def test_3d_generator_basic(self): assert isinstance(mesh, UnstructuredSubMesh) assert mesh.dimension == 3 - assert mesh.npts == 2 * 2 * 2 * 5 # 8 hexes, 5 tets each + assert mesh.npts == 2 * 2 * 2 # 8 hex cells np.testing.assert_allclose(mesh.cell_volumes.sum(), 1.0, atol=1e-14) def test_3d_generator_mesh_integration(self): @@ -306,7 +306,7 @@ def test_3d_generator_mesh_integration(self): submesh = mesh["negative electrode"] assert isinstance(submesh, UnstructuredSubMesh) assert submesh.dimension == 3 - assert submesh.npts == 2 * 2 * 2 * 5 + assert submesh.npts == 2 * 2 * 2 # 8 hex cells def test_interface_conformity_2d(self): """Adjacent domains with the same z grid produce matching interface faces.""" diff --git a/tests/unit/test_plotting/test_quick_plot.py b/tests/unit/test_plotting/test_quick_plot.py index 53b83e7526..2bb04e9d32 100644 --- a/tests/unit/test_plotting/test_quick_plot.py +++ b/tests/unit/test_plotting/test_quick_plot.py @@ -218,7 +218,7 @@ def test_simple_ode_model(self, solver): quick_plot.dynamic_plot(show_plot=False) quick_plot.slider_update(0.01) - with pytest.raises(NotImplementedError, match=r"Cannot plot 2D variables"): + with pytest.raises(NotImplementedError, match=r"Cannot plot 2D/3D variables"): pybamm.QuickPlot([solution, solution], ["2D variable"]) # Test different variable limits From f17aa897b0bac80f54143b7fe1f770a744686aca Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Fri, 6 Mar 2026 11:53:33 -0800 Subject: [PATCH 09/25] vtk --- src/pybamm/__init__.py | 1 + src/pybamm/meshes/unstructured_submesh.py | 54 +- src/pybamm/plotting/dynamic_plot.py | 30 +- src/pybamm/plotting/plot_vtk.py | 685 ++++++++++++++++++ .../test_meshes/test_unstructured_submesh.py | 78 ++ 5 files changed, 841 insertions(+), 7 deletions(-) create mode 100644 src/pybamm/plotting/plot_vtk.py diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py index 49c007eadb..40107c2cb3 100644 --- a/src/pybamm/__init__.py +++ b/src/pybamm/__init__.py @@ -215,6 +215,7 @@ from .plotting.dynamic_plot import dynamic_plot from .plotting.plot_3d_cross_section import plot_3d_cross_section from .plotting.plot_3d_heatmap import plot_3d_heatmap +from .plotting.plot_vtk import VTKQuickPlot # Simulation from .simulation import Simulation, load_sim, is_notebook diff --git a/src/pybamm/meshes/unstructured_submesh.py b/src/pybamm/meshes/unstructured_submesh.py index 91ecee77d3..8aad441d85 100644 --- a/src/pybamm/meshes/unstructured_submesh.py +++ b/src/pybamm/meshes/unstructured_submesh.py @@ -306,6 +306,47 @@ def _identify_boundary_faces(self): if len(indices) > 0: self.boundary_faces[name] = indices + def optimize_ordering(self): + """Reorder cells using Reverse Cuthill-McKee to reduce Jacobian bandwidth. + + Permutes all cell-indexed arrays (elements, centroids, volumes, + face_owner, face_neighbor, interface_data) so that adjacent cells + have nearby indices, minimising the bandwidth of the FVM + connectivity matrix. + """ + from scipy.sparse import csr_matrix + from scipy.sparse.csgraph import reverse_cuthill_mckee + + n = self.npts + if n <= 1: + return + + n_int = self._boundary_face_start + owners = self.face_owner[:n_int] + neighbors = self.face_neighbor + + rows = np.concatenate([owners, neighbors]) + cols = np.concatenate([neighbors, owners]) + data = np.ones(len(rows), dtype=np.float64) + adj = csr_matrix((data, (rows, cols)), shape=(n, n)) + + perm = reverse_cuthill_mckee(adj) + inv_perm = np.empty(n, dtype=int) + inv_perm[perm] = np.arange(n) + + self.elements = self.elements[perm] + self.cell_centroids = self.cell_centroids[perm] + self.cell_volumes = self.cell_volumes[perm] + + self.face_owner = inv_perm[self.face_owner] + self.face_neighbor = inv_perm[self.face_neighbor] + + for _key, data_dict in self.interface_data.items(): + if "left_cells" in data_dict: + data_dict["left_cells"] = inv_perm[data_dict["left_cells"]] + if "right_cells" in data_dict: + data_dict["right_cells"] = inv_perm[data_dict["right_cells"]] + def boundary_loops(self): """Return boundary loops as a list of ``matplotlib.path.Path`` (2D only). @@ -840,10 +881,17 @@ def _hex_grid(x_edges, y_edges, z_edges): def node_id(i, j, k): return i * (ny + 1) * (nz + 1) + j * (nz + 1) + k + # Loop order determines cell numbering and hence Jacobian bandwidth. + # Bandwidth = product of the two fastest-varying dimension sizes. + # Minimise by putting the largest dimension outermost (slowest). + dims = sorted([(nx, "x"), (ny, "y"), (nz, "z")], key=lambda d: d[0], reverse=True) + elements = [] - for i in range(nx): - for j in range(ny): - for k in range(nz): + for a in range(dims[0][0]): + for b in range(dims[1][0]): + for c in range(dims[2][0]): + idx = {dims[0][1]: a, dims[1][1]: b, dims[2][1]: c} + i, j, k = idx["x"], idx["y"], idx["z"] elements.append( [ node_id(i, j, k), diff --git a/src/pybamm/plotting/dynamic_plot.py b/src/pybamm/plotting/dynamic_plot.py index 4cde0d3972..9281e69344 100644 --- a/src/pybamm/plotting/dynamic_plot.py +++ b/src/pybamm/plotting/dynamic_plot.py @@ -11,12 +11,34 @@ def dynamic_plot(*args, **kwargs): The key-word argument 'show_plot' is passed to the 'dynamic_plot' method, not the `QuickPlot` class. + Pass ``backend="vtk"`` to use the VTK-based viewer for unstructured + mesh solutions instead of matplotlib. + Returns ------- - plot : :class:`pybamm.QuickPlot` - The 'QuickPlot' object that was created + plot : :class:`pybamm.QuickPlot` or :class:`pybamm.VTKQuickPlot` + The plot object that was created """ - kwargs_for_class = {k: v for k, v in kwargs.items() if k != "show_plot"} + backend = kwargs.pop("backend", "matplotlib") + show_plot = kwargs.pop("show_plot", True) + + if backend == "vtk": + from pybamm.plotting.plot_vtk import VTKQuickPlot + + output_variables = kwargs.pop("output_variables", None) + options = kwargs.pop("options", None) + interpolate_time = kwargs.pop("interpolate_time", False) + plot = VTKQuickPlot( + *args, + output_variables=output_variables, + options=options, + interpolate_time=interpolate_time, + **kwargs, + ) + plot.dynamic_plot(show_plot) + return plot + + kwargs_for_class = {k: v for k, v in kwargs.items()} plot = pybamm.QuickPlot(*args, **kwargs_for_class) - plot.dynamic_plot(kwargs.get("show_plot", True)) + plot.dynamic_plot(show_plot) return plot diff --git a/src/pybamm/plotting/plot_vtk.py b/src/pybamm/plotting/plot_vtk.py new file mode 100644 index 0000000000..c36aba30c4 --- /dev/null +++ b/src/pybamm/plotting/plot_vtk.py @@ -0,0 +1,685 @@ +""" +VTK-based interactive visualization for unstructured mesh solutions. + +Provides :class:`VTKQuickPlot`, a drop-in alternative to the matplotlib-based +:class:`QuickPlot` for 2D and 3D cell-centered FVM data on unstructured meshes. + +Also supports 0D (time-series) variables rendered as VTK line charts. +""" + +import numpy as np + +import pybamm + +_VTK_CELL_TYPE = { + "triangle": 5, # VTK_TRIANGLE + "quad": 9, # VTK_QUAD + "tetrahedron": 10, # VTK_TETRA + "hexahedron": 12, # VTK_HEXAHEDRON +} + +_AXIS_INDEX = {"x": 0, "y": 1, "z": 2} + + +def _build_vtk_grid(mesh, scale=None): + """Build a ``vtkUnstructuredGrid`` from an ``UnstructuredSubMesh``.""" + import vtk + + nodes = mesh.nodes + if scale is not None: + nodes = nodes * np.asarray(scale)[: nodes.shape[1]] + + pts = vtk.vtkPoints() + pts.SetNumberOfPoints(len(nodes)) + for i, nd in enumerate(nodes): + if len(nd) == 2: + pts.SetPoint(i, nd[0], nd[1], 0.0) + else: + pts.SetPoint(i, nd[0], nd[1], nd[2]) + + grid = vtk.vtkUnstructuredGrid() + grid.SetPoints(pts) + + cell_type = _VTK_CELL_TYPE[mesh.element_type] + for cell in mesh.elements: + id_list = vtk.vtkIdList() + for v in cell: + id_list.InsertNextId(int(v)) + grid.InsertNextCell(cell_type, id_list) + + return grid + + +def _compute_scale(mesh): + """Per-axis scale factors that normalise coordinate spans to the largest.""" + nodes = mesh.nodes + spans = np.array( + [nodes[:, d].max() - nodes[:, d].min() for d in range(nodes.shape[1])] + ) + max_span = spans.max() + if max_span == 0: + return np.ones(nodes.shape[1]) + return max_span / np.where(spans > 0, spans, max_span) + + +def _resolve_scale(scale_opt, mesh): + """Turn a scale option into a concrete array or None.""" + if scale_opt == "auto": + return _compute_scale(mesh) + if scale_opt is None: + return None + return np.asarray(scale_opt) + + +def _set_cell_scalars(grid, name, values): + """Set (or update) a cell scalar array on a VTK grid.""" + import vtk + + arr = grid.GetCellData().GetArray(name) + if arr is None: + arr = vtk.vtkFloatArray() + arr.SetName(name) + arr.SetNumberOfTuples(len(values)) + grid.GetCellData().AddArray(arr) + grid.GetCellData().SetActiveScalars(name) + for i, v in enumerate(values): + arr.SetValue(i, float(v)) + arr.Modified() + grid.Modified() + + +def _viridis_lut(vmin, vmax, n=256): + """Build a VTK lookup table using the matplotlib viridis colormap.""" + import vtk + + try: + from matplotlib.cm import viridis as _cmap + except ImportError: + lut = vtk.vtkLookupTable() + lut.SetHueRange(0.667, 0.0) + lut.SetRange(vmin, vmax) + lut.Build() + return lut + + lut = vtk.vtkLookupTable() + lut.SetNumberOfTableValues(n) + lut.SetRange(vmin, vmax) + for i in range(n): + r, g, b, a = _cmap(i / (n - 1)) + lut.SetTableValue(i, r, g, b, a) + lut.Build() + return lut + + +class VTKQuickPlot: + """Interactive VTK visualization for unstructured FVM solutions. + + Supports spatial (unstructured 2D/3D) and 0D (time-series) variables. + + Parameters + ---------- + solutions : :class:`pybamm.Solution` or list thereof + output_variables : list of str + options : dict, optional + Per-variable options keyed by variable name. Each value is a dict + that may contain: + + - ``"plot_type"``: ``"3d"`` (default) or ``"slice"`` + - ``"x"`` / ``"y"`` / ``"z"``: float in [0, 1] giving the slice + position as a fraction of the axis range (required when + ``plot_type`` is ``"slice"``) + - ``"scale"``: ``"auto"`` (default), ``None``, or ``(sx, sy, sz)`` + + Each variable's options value may also be a **list** of dicts, in which + case one panel is created per entry:: + + options={"T": [ + {"plot_type": "3d"}, + {"plot_type": "slice", "x": 0.5}, + ]} + """ + + def __init__( + self, + solutions, + output_variables=None, + options=None, + interpolate_time=False, + ): + if isinstance(solutions, pybamm.Simulation): + solutions = solutions.solution + if isinstance(solutions, pybamm.Solution): + solutions = [solutions] + self.solution = solutions[0] + + if output_variables is None: + output_variables = list(self.solution.all_models[0].variables.keys())[:1] + if isinstance(output_variables, str): + output_variables = [output_variables] + + self.spatial_names = [] + self.spatial_vars = [] + self.scalar_names = [] + self.scalar_vars = [] + + for name in output_variables: + pv = self.solution[name] + if isinstance(pv, pybamm.ProcessedVariableUnstructuredFVM): + self.spatial_names.append(name) + self.spatial_vars.append(pv) + else: + self.scalar_names.append(name) + self.scalar_vars.append(pv) + + self.output_variables = output_variables + self.mesh = self.spatial_vars[0].mesh if self.spatial_vars else None + self.t_pts = self.solution.t + self.interpolate_time = interpolate_time + + _defaults = {"plot_type": "3d", "scale": "auto"} + raw_opts = options or {} + + # Build spatial_panels: flat list of (name, opts_dict) tuples. + self.spatial_panels = [] + for name in self.spatial_names: + var_opt = raw_opts.get(name, _defaults) + if isinstance(var_opt, dict): + opt_list = [var_opt] + else: + opt_list = list(var_opt) + for single_opt in opt_list: + merged = dict(_defaults) + merged.update(single_opt) + self.spatial_panels.append((name, merged)) + + # ------------------------------------------------------------------ + + def dynamic_plot(self, show_plot=True): + """Launch an interactive VTK window with a time slider.""" + import vtk + + n_spatial = len(self.spatial_panels) + n_scalar = len(self.scalar_names) + n_panels = n_spatial + n_scalar + + # --- Precompute spatial data --- + spatial_data = {} + spatial_mins = {} + spatial_maxs = {} + for name, pv in zip(self.spatial_names, self.spatial_vars, strict=True): + pv.initialise() + data = np.column_stack([pv._data_at_time(t).ravel() for t in self.t_pts]) + spatial_data[name] = data + spatial_mins[name] = float(data.min()) + spatial_maxs[name] = float(data.max()) + + # --- Precompute scalar (0D) data --- + scalar_data = {} + for name, pv in zip(self.scalar_names, self.scalar_vars, strict=True): + pv.initialise() + vals = np.array([float(pv(t).ravel()[0]) for t in self.t_pts]) + scalar_data[name] = vals + + # --- Layout --- + slider_h = 0.08 + panel_top = 1.0 + panel_bot = slider_h + + n_cols = int(np.ceil(np.sqrt(n_panels))) + n_rows = int(np.ceil(n_panels / n_cols)) + panel_height = (panel_top - panel_bot) / n_rows + + window = vtk.vtkRenderWindow() + window.SetSize(650 * n_cols, 520 * n_rows) + window.SetWindowName("PyBaMM - " + ", ".join(self.output_variables)) + + all_renderers = [] + spatial_grids = [] + c2p_filters = [] + cutters = [] + chart_views = [] + time_markers = [] + + panel_idx = 0 + + # --- Spatial panels --- + first_3d_cam = None + spatial_renderers = [] + panel_names = [] + + for name, opts in self.spatial_panels: + plot_type = opts.get("plot_type", "3d") + var_scale = _resolve_scale(opts.get("scale", "auto"), self.mesh) + panel_names.append(name) + + g = _build_vtk_grid(self.mesh, scale=var_scale) + _set_cell_scalars(g, name, spatial_data[name][:, 0]) + spatial_grids.append(g) + + c2p = vtk.vtkCellDataToPointData() + c2p.SetInputData(g) + c2p.Update() + c2p_filters.append(c2p) + + # Determine pipeline source: cutter for slices, c2p for 3d + cutter = None + if plot_type == "slice": + axis_key = None + for ak in ("x", "y", "z"): + if ak in opts: + axis_key = ak + break + if axis_key is None: + raise ValueError( + f"plot_type='slice' for '{name}' requires one of " + f"'x', 'y', or 'z' specifying the slice fraction" + ) + axis_idx = _AXIS_INDEX[axis_key] + frac = float(opts[axis_key]) + nodes = self.mesh.nodes + lo = float(nodes[:, axis_idx].min()) + hi = float(nodes[:, axis_idx].max()) + phys_val = lo + frac * (hi - lo) + scaled_val = ( + phys_val * var_scale[axis_idx] + if var_scale is not None + else phys_val + ) + + plane = vtk.vtkPlane() + origin = [0.0, 0.0, 0.0] + origin[axis_idx] = scaled_val + plane.SetOrigin(origin) + normal = [0.0, 0.0, 0.0] + normal[axis_idx] = 1.0 + plane.SetNormal(normal) + + cutter = vtk.vtkCutter() + cutter.SetCutFunction(plane) + cutter.SetInputConnection(c2p.GetOutputPort()) + cutter.Update() + + mapper_source = cutter.GetOutputPort() + else: + mapper_source = c2p.GetOutputPort() + + cutters.append(cutter) + + lut = _viridis_lut(spatial_mins[name], spatial_maxs[name]) + + mapper = vtk.vtkDataSetMapper() + mapper.SetInputConnection(mapper_source) + mapper.SetScalarRange(spatial_mins[name], spatial_maxs[name]) + mapper.SetScalarModeToUsePointData() + mapper.SelectColorArray(name) + mapper.SetLookupTable(lut) + mapper.InterpolateScalarsBeforeMappingOn() + + actor = vtk.vtkActor() + actor.SetMapper(mapper) + if plot_type == "slice": + actor.GetProperty().EdgeVisibilityOff() + else: + actor.GetProperty().EdgeVisibilityOn() + actor.GetProperty().SetEdgeColor(0.2, 0.2, 0.2) + actor.GetProperty().SetLineWidth(0.3) + + sb = vtk.vtkScalarBarActor() + sb.SetLookupTable(lut) + sb.SetTitle("") + sb.SetNumberOfLabels(5) + sb.SetWidth(0.18) + sb.SetHeight(0.5) + sb.SetPosition(0.80, 0.25) + sb.GetLabelTextProperty().SetFontSize(24) + sb.GetLabelTextProperty().SetColor(0, 0, 0) + sb.SetUnconstrainedFontSize(True) + sb.SetLabelFormat("%-#6.3g") + + title_actor = vtk.vtkTextActor() + title_actor.SetInput(name) + title_actor.GetTextProperty().SetFontSize(36) + title_actor.GetTextProperty().SetColor(0, 0, 0) + title_actor.GetTextProperty().SetBold(True) + title_actor.GetTextProperty().SetJustificationToCentered() + title_actor.GetPositionCoordinate().SetCoordinateSystemToNormalizedViewport() + title_actor.SetPosition(0.5, 0.92) + + ren = vtk.vtkRenderer() + ren.AddActor(actor) + ren.AddActor2D(sb) + ren.AddActor2D(title_actor) + ren.SetBackground(1, 1, 1) + + row = panel_idx // n_cols + col = panel_idx % n_cols + y0 = panel_top - (row + 1) * panel_height + y1 = panel_top - row * panel_height + ren.SetViewport(col / n_cols, y0, (col + 1) / n_cols, y1) + + # Cube axes + if self.mesh is not None: + mesh_nodes = self.mesh.nodes + dim = mesh_nodes.shape[1] + + if plot_type == "slice": + # Use the cutter output bounds so axes align with + # the visible slice geometry, not the full 3D grid. + axes_bounds = list(cutter.GetOutput().GetBounds()) + else: + axes_bounds = list(g.GetBounds()) + + cube_axes = vtk.vtkCubeAxesActor() + cube_axes.SetBounds(axes_bounds) + cube_axes.SetUseAxisOrigin(False) + cube_axes.SetFlyModeToOuterEdges() + if plot_type == "slice": + cube_axes.SetTickLocationToInside() + cube_axes.SetScreenSize(10.0) + cube_axes.SetLabelOffset(10) + cube_axes.SetTitleOffset([20, 20]) + + orig_ranges = [ + (float(mesh_nodes[:, d].min()), float(mesh_nodes[:, d].max())) + for d in range(dim) + ] + if dim >= 1: + cube_axes.SetXAxisRange(*orig_ranges[0]) + if dim >= 2: + cube_axes.SetYAxisRange(*orig_ranges[1]) + if dim >= 3: + cube_axes.SetZAxisRange(*orig_ranges[2]) + + for ax_id in range(3): + tp = cube_axes.GetTitleTextProperty(ax_id) + tp.SetFontSize(28) + tp.SetColor(0.15, 0.15, 0.15) + tp.SetBold(True) + lp = cube_axes.GetLabelTextProperty(ax_id) + lp.SetFontSize(22) + lp.SetColor(0.25, 0.25, 0.25) + cube_axes.SetXTitle("X") + cube_axes.SetYTitle("Y") + cube_axes.SetZTitle("Z") + cube_axes.SetXLabelFormat("%.2g") + cube_axes.SetYLabelFormat("%.2g") + cube_axes.SetZLabelFormat("%.2g") + cube_axes.XAxisMinorTickVisibilityOff() + cube_axes.YAxisMinorTickVisibilityOff() + cube_axes.ZAxisMinorTickVisibilityOff() + + if plot_type == "slice": + if axis_idx == 0: + cube_axes.XAxisVisibilityOff() + cube_axes.SetXAxisTickVisibility(False) + cube_axes.SetXAxisLabelVisibility(False) + elif axis_idx == 1: + cube_axes.YAxisVisibilityOff() + cube_axes.SetYAxisTickVisibility(False) + cube_axes.SetYAxisLabelVisibility(False) + else: + cube_axes.ZAxisVisibilityOff() + cube_axes.SetZAxisTickVisibility(False) + cube_axes.SetZAxisLabelVisibility(False) + + ren.AddActor(cube_axes) + + window.AddRenderer(ren) + all_renderers.append(ren) + spatial_renderers.append(ren) + + # Camera setup: slice panels get independent orthographic cameras; + # 3d panels share a single perspective camera. + if plot_type == "slice": + ren.ResetCamera() + cam = ren.GetActiveCamera() + cam.SetParallelProjection(True) + pos = list(cam.GetPosition()) + fp = list(cam.GetFocalPoint()) + gb = g.GetBounds() + offset = ( + max( + gb[1] - gb[0], + gb[3] - gb[2], + gb[5] - gb[4], + ) + * 2 + ) + # Look from the negative side so OuterEdges places + # axis labels on the top/left edges (more viewport room). + pos[axis_idx] = fp[axis_idx] - offset + cam.SetPosition(pos) + view_up = [0, 0, 0] + if axis_idx == 2: + view_up[1] = 1 + elif axis_idx == 1: + view_up[2] = 1 + else: + view_up[1] = 1 + cam.SetViewUp(view_up) + ren.ResetCamera() + cam.Zoom(0.70) + if self.mesh is not None: + cube_axes.SetCamera(cam) + else: + if first_3d_cam is None: + ren.ResetCamera() + first_3d_cam = ren.GetActiveCamera() + if self.mesh is not None and self.mesh.dimension == 3: + first_3d_cam.Azimuth(-55) + first_3d_cam.Elevation(25) + if self.mesh is not None: + cube_axes.SetCamera(first_3d_cam) + else: + ren.SetActiveCamera(first_3d_cam) + if self.mesh is not None: + cube_axes.SetCamera(first_3d_cam) + + panel_idx += 1 + + # --- Scalar (0D chart) panels --- + for name in self.scalar_names: + vals = scalar_data[name] + v_min, v_max = float(vals.min()), float(vals.max()) + v_pad = max((v_max - v_min) * 0.05, 1e-10) + + chart = vtk.vtkChartXY() + chart.SetTitle(name) + chart.GetTitleProperties().SetFontSize(36) + chart.GetTitleProperties().SetBold(True) + chart.GetTitleProperties().SetColor(0, 0, 0) + chart.GetAxis(1).SetTitle("Time [s]") + chart.GetAxis(0).SetTitle(name) + chart.GetAxis(1).GetTitleProperties().SetFontSize(28) + chart.GetAxis(1).GetTitleProperties().SetColor(0, 0, 0) + chart.GetAxis(1).GetLabelProperties().SetFontSize(22) + chart.GetAxis(1).GetLabelProperties().SetColor(0, 0, 0) + chart.GetAxis(0).GetTitleProperties().SetFontSize(28) + chart.GetAxis(0).GetTitleProperties().SetColor(0, 0, 0) + chart.GetAxis(0).GetLabelProperties().SetFontSize(22) + chart.GetAxis(0).GetLabelProperties().SetColor(0, 0, 0) + chart.GetAxis(1).SetRange(float(self.t_pts[0]), float(self.t_pts[-1])) + chart.GetAxis(0).SetRange(v_min - v_pad, v_max + v_pad) + + table = vtk.vtkTable() + t_arr = vtk.vtkFloatArray() + t_arr.SetName("Time") + v_arr = vtk.vtkFloatArray() + v_arr.SetName(name) + for i in range(len(self.t_pts)): + t_arr.InsertNextValue(float(self.t_pts[i])) + v_arr.InsertNextValue(float(vals[i])) + table.AddColumn(t_arr) + table.AddColumn(v_arr) + + line = chart.AddPlot(vtk.vtkChart.LINE) + line.SetInputData(table, 0, 1) + line.SetColor(31, 119, 180, 255) + line.SetWidth(2.0) + + marker_table = vtk.vtkTable() + mt_arr = vtk.vtkFloatArray() + mt_arr.SetName("t") + mv_arr = vtk.vtkFloatArray() + mv_arr.SetName("v") + mt_arr.InsertNextValue(float(self.t_pts[0])) + mt_arr.InsertNextValue(float(self.t_pts[0])) + mv_arr.InsertNextValue(v_min - v_pad) + mv_arr.InsertNextValue(v_max + v_pad) + marker_table.AddColumn(mt_arr) + marker_table.AddColumn(mv_arr) + + marker_line = chart.AddPlot(vtk.vtkChart.LINE) + marker_line.SetInputData(marker_table, 0, 1) + marker_line.SetColor(200, 50, 50, 200) + marker_line.SetWidth(1.5) + time_markers.append((mt_arr, marker_table)) + + view = vtk.vtkContextActor() + scene = vtk.vtkContextScene() + scene.AddItem(chart) + view.SetScene(scene) + + ren = vtk.vtkRenderer() + ren.AddActor(view) + scene.SetRenderer(ren) + ren.SetBackground(1, 1, 1) + + row = panel_idx // n_cols + col = panel_idx % n_cols + y0 = panel_top - (row + 1) * panel_height + y1 = panel_top - row * panel_height + ren.SetViewport(col / n_cols, y0, (col + 1) / n_cols, y1) + + window.AddRenderer(ren) + all_renderers.append(ren) + chart_views.append((chart, view, scene)) + panel_idx += 1 + + # --- Fill any unused grid cells with white --- + while panel_idx < n_rows * n_cols: + ren = vtk.vtkRenderer() + ren.SetBackground(1, 1, 1) + row = panel_idx // n_cols + col = panel_idx % n_cols + y0 = panel_top - (row + 1) * panel_height + y1 = panel_top - row * panel_height + ren.SetViewport(col / n_cols, y0, (col + 1) / n_cols, y1) + window.AddRenderer(ren) + panel_idx += 1 + + # --- Slider background (white strip at bottom) --- + slider_bg = vtk.vtkRenderer() + slider_bg.SetBackground(1, 1, 1) + slider_bg.SetViewport(0, 0, 1, slider_h) + window.AddRenderer(slider_bg) + + interactor = vtk.vtkRenderWindowInteractor() + interactor.SetRenderWindow(window) + + # Time label + time_text = vtk.vtkTextActor() + time_text.SetInput(f"t = {self.t_pts[0]:.4g} s") + time_text.GetTextProperty().SetFontSize(28) + time_text.GetTextProperty().SetColor(0, 0, 0) + time_text.GetTextProperty().SetBold(True) + time_text.GetPositionCoordinate().SetCoordinateSystemToNormalizedViewport() + time_text.SetPosition(0.01, 0.15) + slider_bg.AddActor2D(time_text) + + # Time slider — scaled in physical time (seconds) + t_min = float(self.t_pts[0]) + t_max = float(self.t_pts[-1]) + slider_rep = vtk.vtkSliderRepresentation2D() + slider_rep.SetMinimumValue(t_min) + slider_rep.SetMaximumValue(t_max) + slider_rep.SetValue(t_min) + slider_rep.SetTitleText("") + slider_rep.GetPoint1Coordinate().SetCoordinateSystemToNormalizedDisplay() + slider_rep.GetPoint1Coordinate().SetValue(0.15, slider_h * 0.5) + slider_rep.GetPoint2Coordinate().SetCoordinateSystemToNormalizedDisplay() + slider_rep.GetPoint2Coordinate().SetValue(0.95, slider_h * 0.5) + slider_rep.SetSliderLength(0.04) + slider_rep.SetSliderWidth(0.06) + slider_rep.SetTubeWidth(0.015) + slider_rep.SetEndCapLength(0.02) + slider_rep.SetEndCapWidth(0.06) + slider_rep.GetTitleProperty().SetColor(0, 0, 0) + slider_rep.GetLabelProperty().SetColor(0, 0, 0) + slider_rep.GetLabelProperty().SetFontSize(16) + slider_rep.GetSliderProperty().SetColor(0.2, 0.4, 0.8) + slider_rep.GetTubeProperty().SetColor(0.7, 0.7, 0.7) + slider_rep.GetCapProperty().SetColor(0.5, 0.5, 0.5) + slider_rep.GetSelectedProperty().SetColor(0.3, 0.5, 0.9) + + # Look-up table for snapping to nearest timestep + _t_array = np.asarray(self.t_pts) + + # Keep references for interpolated mode + _spatial_vars = { + name: pv + for name, pv in zip( + self.spatial_names, + self.spatial_vars, + strict=True, + ) + } + + def on_slider(obj, event): + t_now = float(obj.GetRepresentation().GetValue()) + t_now = max(t_min, min(t_now, t_max)) + + if self.interpolate_time: + # Evaluate every spatial variable at exact time + for sname, g, c2p, cut in zip( + panel_names, + spatial_grids, + c2p_filters, + cutters, + strict=True, + ): + vals = _spatial_vars[sname]._data_at_time(t_now).ravel() + _set_cell_scalars(g, sname, vals) + c2p.Modified() + c2p.Update() + if cut is not None: + cut.Update() + else: + # Snap to nearest stored timestep (fast) + t_idx = int(np.argmin(np.abs(_t_array - t_now))) + for sname, g, c2p, cut in zip( + panel_names, + spatial_grids, + c2p_filters, + cutters, + strict=True, + ): + _set_cell_scalars(g, sname, spatial_data[sname][:, t_idx]) + c2p.Modified() + c2p.Update() + if cut is not None: + cut.Update() + + for mt_arr, mtable in time_markers: + mt_arr.SetValue(0, t_now) + mt_arr.SetValue(1, t_now) + mt_arr.Modified() + mtable.Modified() + time_text.SetInput(f"t = {t_now:.4g} s") + window.Render() + + slider = vtk.vtkSliderWidget() + slider.SetInteractor(interactor) + slider.SetRepresentation(slider_rep) + slider.SetAnimationModeToAnimate() + slider.EnabledOn() + slider.AddObserver("InteractionEvent", on_slider) + + if show_plot: + interactor.Initialize() + window.Render() + interactor.Start() + + self._window = window + self._interactor = interactor + self._slider = slider diff --git a/tests/unit/test_meshes/test_unstructured_submesh.py b/tests/unit/test_meshes/test_unstructured_submesh.py index 5ca41ef576..74342ecf24 100644 --- a/tests/unit/test_meshes/test_unstructured_submesh.py +++ b/tests/unit/test_meshes/test_unstructured_submesh.py @@ -4,6 +4,7 @@ from pybamm.meshes.unstructured_submesh import ( UnstructuredMeshGenerator, UnstructuredSubMesh, + _hex_grid, _hex_to_tet, _quad_to_tri, compute_interface_data, @@ -591,3 +592,80 @@ def test_interface_data_computed_automatically(self): assert "separator" in neg_mesh.interface_data assert "negative electrode" in sep_mesh.interface_data assert len(neg_mesh.interface_data["separator"]["left_cells"]) > 0 + + +class TestBandwidthOptimization: + """Tests for _hex_grid loop ordering and optimize_ordering (RCM).""" + + @staticmethod + def _bandwidth(submesh): + n_int = submesh._boundary_face_start + owners = submesh.face_owner[:n_int] + neighbors = submesh.face_neighbor + return int(np.max(np.abs(owners.astype(int) - neighbors.astype(int)))) + + def test_hex_grid_optimal_loop_order(self): + """_hex_grid should order cells so bandwidth = product of two smallest dims.""" + for nx, ny, nz in [(3, 10, 5), (2, 4, 20), (7, 3, 3), (5, 5, 5)]: + nodes, elems = _hex_grid( + np.linspace(0, 1, nx + 1), + np.linspace(0, 1, ny + 1), + np.linspace(0, 1, nz + 1), + ) + mesh = UnstructuredSubMesh(nodes, elems, coord_sys="cartesian") + bw = self._bandwidth(mesh) + dims = sorted([nx, ny, nz]) + expected = dims[0] * dims[1] + assert bw == expected, ( + f"nx={nx} ny={ny} nz={nz}: bw={bw}, expected={expected}" + ) + + def test_optimize_ordering_reduces_bandwidth(self): + """optimize_ordering (RCM) should not increase bandwidth.""" + nodes, elems = _hex_grid( + np.linspace(0, 1, 4), + np.linspace(0, 1, 11), + np.linspace(0, 1, 6), + ) + mesh = UnstructuredSubMesh(nodes, elems, coord_sys="cartesian") + bw_before = self._bandwidth(mesh) + mesh.optimize_ordering() + bw_after = self._bandwidth(mesh) + assert bw_after <= bw_before + + def test_optimize_ordering_preserves_geometry(self): + """Cell volumes and centroids must be the same set after reordering.""" + nodes, elems = _hex_grid( + np.linspace(0, 1, 4), + np.linspace(0, 1, 6), + np.linspace(0, 1, 4), + ) + mesh = UnstructuredSubMesh(nodes, elems, coord_sys="cartesian") + vols_before = np.sort(mesh.cell_volumes) + cents_before = mesh.cell_centroids[np.lexsort(mesh.cell_centroids.T)] + + mesh.optimize_ordering() + + vols_after = np.sort(mesh.cell_volumes) + cents_after = mesh.cell_centroids[np.lexsort(mesh.cell_centroids.T)] + np.testing.assert_allclose(vols_before, vols_after) + np.testing.assert_allclose(cents_before, cents_after) + + def test_optimize_ordering_preserves_interface_data(self): + """Interface cell centroids should point to the same physical cells.""" + ye = np.linspace(0, 1, 6) + ze = np.linspace(0, 1, 4) + nodes_l, elems_l = _hex_grid(np.linspace(0, 1, 4), ye, ze) + mesh_l = UnstructuredSubMesh(nodes_l, elems_l, coord_sys="cartesian") + nodes_r, elems_r = _hex_grid(np.linspace(1, 2, 4), ye, ze) + mesh_r = UnstructuredSubMesh(nodes_r, elems_r, coord_sys="cartesian") + compute_interface_data(mesh_l, mesh_r, "left", "right") + + iface = mesh_l.interface_data["right"] + centroids_pre = mesh_l.cell_centroids[iface["left_cells"]].copy() + + mesh_l.optimize_ordering() + + iface = mesh_l.interface_data["right"] + centroids_post = mesh_l.cell_centroids[iface["left_cells"]] + np.testing.assert_allclose(centroids_pre, centroids_post) From e60cd9527d92036478ad2b66c0a05a1331855e01 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Fri, 6 Mar 2026 11:55:00 -0800 Subject: [PATCH 10/25] add vtk --- pyproject.toml | 3 ++- 1 file changed, 2 insertions(+), 1 deletion(-) diff --git a/pyproject.toml b/pyproject.toml index 34121334e6..06ecb81a83 100644 --- a/pyproject.toml +++ b/pyproject.toml @@ -64,11 +64,12 @@ bpx = ["bpx>=0.5.0,<0.6.0"] # Low-overhead progress bars tqdm = ["tqdm"] jax = ["jax>=0.7.0, <0.9.0; python_version >= '3.11' and (sys_platform != 'darwin' or platform_machine != 'x86_64')"] +vtk = ["vtk>=9.0.0"] # Contains all optional dependencies, except for jax, and dev dependencies all = [ "scikit-fem>=8.1.0", "meshio>=5.3.0", - "pybamm[examples,plot,cite,bpx,tqdm]", + "pybamm[examples,plot,cite,bpx,tqdm,vtk]", ] [dependency-groups] From d1a8c9db828112ceb57b8c3c513b453ae104be9e Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Fri, 6 Mar 2026 12:51:23 -0800 Subject: [PATCH 11/25] fix weird expression tree thing --- src/pybamm/expression_tree/functions.py | 31 ++++++------ src/pybamm/plotting/plot_vtk.py | 63 ++++++++++++++++++++++++- 2 files changed, 79 insertions(+), 15 deletions(-) diff --git a/src/pybamm/expression_tree/functions.py b/src/pybamm/expression_tree/functions.py index 90fc154a8d..4e3b3947c2 100644 --- a/src/pybamm/expression_tree/functions.py +++ b/src/pybamm/expression_tree/functions.py @@ -637,8 +637,21 @@ def log10(child: pybamm.Symbol): return log(child, base=10) -class Max(SpecificFunction): - """Max function.""" +class Reduction(SpecificFunction): + """Base class for reduction operations that collapse a spatial + field to a scalar (e.g. max, min). Automatically clears domains + and returns scalar shape.""" + + def __init__(self, function: Callable, child: pybamm.Symbol): + super().__init__(function, child) + self.clear_domains() + + def _evaluate_for_shape(self): + return np.nan * np.ones((1, 1)) + + +class Max(Reduction): + """Max function (reduction to scalar).""" def __init__(self, child): super().__init__(np.max, child) @@ -650,11 +663,6 @@ def _from_json(cls, snippet: dict): instance = super()._from_json(snippet) return instance - def _evaluate_for_shape(self): - """See :meth:`pybamm.Symbol.evaluate_for_shape_using_domain()`""" - # Max will always return a scalar - return np.nan * np.ones((1, 1)) - def max(child: pybamm.Symbol): """ @@ -664,8 +672,8 @@ def max(child: pybamm.Symbol): return pybamm.simplify_if_constant(Max(child)) -class Min(SpecificFunction): - """Min function.""" +class Min(Reduction): + """Min function (reduction to scalar).""" def __init__(self, child): super().__init__(np.min, child) @@ -677,11 +685,6 @@ def _from_json(cls, snippet: dict): instance = super()._from_json(snippet) return instance - def _evaluate_for_shape(self): - """See :meth:`pybamm.Symbol.evaluate_for_shape_using_domain()`""" - # Min will always return a scalar - return np.nan * np.ones((1, 1)) - def min(child: pybamm.Symbol): """ diff --git a/src/pybamm/plotting/plot_vtk.py b/src/pybamm/plotting/plot_vtk.py index c36aba30c4..9b037e7e0c 100644 --- a/src/pybamm/plotting/plot_vtk.py +++ b/src/pybamm/plotting/plot_vtk.py @@ -217,7 +217,12 @@ def dynamic_plot(self, show_plot=True): scalar_data = {} for name, pv in zip(self.scalar_names, self.scalar_vars, strict=True): pv.initialise() - vals = np.array([float(pv(t).ravel()[0]) for t in self.t_pts]) + if isinstance(pv, pybamm.ProcessedVariableUnstructuredFVM): + vals = np.array( + [float(pv._data_at_time(t).ravel()[0]) for t in self.t_pts] + ) + else: + vals = np.array([float(pv(t).ravel()[0]) for t in self.t_pts]) scalar_data[name] = vals # --- Layout --- @@ -683,3 +688,59 @@ def on_slider(obj, event): self._window = window self._interactor = interactor self._slider = slider + + def save_gif(self, filename, fps=10, n_frames=100, width=1800, height=900): + """Render an animation to a GIF file. + + Parameters + ---------- + filename : str + Output path (e.g. ``"anim.gif"``). + fps : int + Frames per second. + n_frames : int + Number of frames (evenly spaced in time). + width, height : int + Pixel dimensions of each frame. + """ + import vtk + from PIL import Image + + if not hasattr(self, "_window"): + self.dynamic_plot(show_plot=False) + + win = self._window + win.SetOffScreenRendering(1) + win.SetSize(width, height) + + t_min = float(self.t_pts[0]) + t_max = float(self.t_pts[-1]) + frame_times = np.linspace(t_min, t_max, n_frames) + + frames = [] + for t in frame_times: + self._slider.GetRepresentation().SetValue(t) + self._slider.InvokeEvent("InteractionEvent") + win.Render() + + w2i = vtk.vtkWindowToImageFilter() + w2i.SetInput(win) + w2i.Update() + img_data = w2i.GetOutput() + + w_px, h_px, _ = img_data.GetDimensions() + n_comp = img_data.GetNumberOfScalarComponents() + raw = np.frombuffer( + memoryview(img_data.GetPointData().GetScalars()), + dtype=np.uint8, + ).reshape(h_px, w_px, n_comp) + frames.append(Image.fromarray(raw[::-1])) + + frames[0].save( + filename, + save_all=True, + append_images=frames[1:], + duration=int(1000 / fps), + loop=0, + ) + print(f"Saved {len(frames)}-frame GIF to {filename}") From ec57d763acfbb99968ab4895dc505a6a69dc6d54 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Mon, 9 Mar 2026 11:22:08 -0700 Subject: [PATCH 12/25] get pouch working --- src/pybamm/discretisations/discretisation.py | 5 ++- src/pybamm/meshes/meshes.py | 32 ++++++++++++++- .../parameters/parameter_substitutor.py | 39 +++---------------- 3 files changed, 40 insertions(+), 36 deletions(-) diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py index ea365a00ac..6c24f0b83e 100644 --- a/src/pybamm/discretisations/discretisation.py +++ b/src/pybamm/discretisations/discretisation.py @@ -580,8 +580,9 @@ def process_boundary_conditions(self, model): f"Neumann condition for {self.mesh[subdomain].coord_sys} coordinates" ) - # Handle any boundary conditions applied on the tabs - if any("tab" in side for side in list(bcs.keys())): + # Handle legacy tab boundary conditions ("negative tab", etc.) + legacy_tab_sides = {"negative tab", "positive tab", "no tab"} + if legacy_tab_sides & set(bcs.keys()): bcs = self.check_tab_conditions(key, bcs) # Process boundary conditions diff --git a/src/pybamm/meshes/meshes.py b/src/pybamm/meshes/meshes.py index b1d96c5b98..5d2a45deb5 100644 --- a/src/pybamm/meshes/meshes.py +++ b/src/pybamm/meshes/meshes.py @@ -558,12 +558,42 @@ def _combine_unstructured_submeshes(submeshes): combined_nodes = np.array(all_nodes) combined_elements = np.concatenate(all_elements, axis=0) - return pybamm.UnstructuredSubMesh( + combined = pybamm.UnstructuredSubMesh( combined_nodes, combined_elements, coord_sys=submeshes[0].coord_sys, ) + # Propagate custom boundary tags from component submeshes. + # The combined mesh auto-detects only standard tags (left/right/top/bottom/ + # front/back). Custom tags like "tab_top" are lost. Recover them by + # matching boundary face centroids. + standard_tags = {"left", "right", "top", "bottom", "front", "back"} + custom_centroids = {} # tag -> list of centroid arrays + for sm in submeshes: + for tag, face_indices in sm.boundary_faces.items(): + if tag not in standard_tags: + custom_centroids.setdefault(tag, []).append( + sm.face_centroids[face_indices] + ) + + if custom_centroids: + from scipy.spatial import cKDTree + + bnd_start = combined._boundary_face_start + bnd_centroids = combined.face_centroids[bnd_start:] + if len(bnd_centroids) > 0: + tree = cKDTree(bnd_centroids) + match_tol = 1e-10 * max(np.ptp(combined_nodes, axis=0).max(), 1.0) + for tag, centroid_list in custom_centroids.items(): + all_src = np.concatenate(centroid_list, axis=0) + dists, idxs = tree.query(all_src) + matched = idxs[dists < match_tol] + if len(matched) > 0: + combined.boundary_faces[tag] = np.unique(matched) + bnd_start + + return combined + class SubMesh: """ diff --git a/src/pybamm/parameters/parameter_substitutor.py b/src/pybamm/parameters/parameter_substitutor.py index 3bf91ae6c2..99518c637b 100644 --- a/src/pybamm/parameters/parameter_substitutor.py +++ b/src/pybamm/parameters/parameter_substitutor.py @@ -544,42 +544,15 @@ def process_boundary_conditions( new_boundary_conditions: dict[ pybamm.Symbol, dict[str, tuple[pybamm.Symbol, str]] ] = {} - sides = [ - "left", - "right", - "negative tab", - "positive tab", - "no tab", - "top", - "bottom", - "x_min", - "x_max", - "y_min", - "y_max", - "z_min", - "z_max", - "r_min", - "r_max", - ] for variable, bcs in model.boundary_conditions.items(): processed_variable = self.process_symbol(variable) new_boundary_conditions[processed_variable] = {} - for side in sides: - try: - bc, typ = bcs[side] - pybamm.logger.verbose( - f"Processing parameters for {variable!r} ({side} bc)" - ) - processed_bc = (self.process_symbol(bc), typ) - new_boundary_conditions[processed_variable][side] = processed_bc - except KeyError as err: - # don't raise error if the key error comes from the side not being - # found - if err.args[0] in side: - pass - # do raise error otherwise (e.g. can't process symbol) - else: - raise err + for side, (bc, typ) in bcs.items(): + pybamm.logger.verbose( + f"Processing parameters for {variable!r} ({side} bc)" + ) + processed_bc = (self.process_symbol(bc), typ) + new_boundary_conditions[processed_variable][side] = processed_bc return new_boundary_conditions From 4036d4c13fc3bfcb0af238b83fa694c814535773 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Fri, 24 Apr 2026 17:27:35 -0500 Subject: [PATCH 13/25] Improve coverage: vector_field, tensor_field, unstructured_submesh - vector_field.py: 85% -> 100% - tensor_field.py: 90% -> 100% - unstructured_submesh.py: 64% -> 87% Added tests for quad/hex element types, boundary loops, contains_points_3d on hex, optimize_ordering on single cell, generator error paths, and _parse_lims string-variable handling. Added tensor/vector field getitem edge cases, all-on-edges path, and VectorField._to_casadi coverage. --- .../test_meshes/test_unstructured_submesh.py | 168 ++++++++++++++++++ .../test_tensor_field.py | 77 ++++++++ 2 files changed, 245 insertions(+) diff --git a/tests/unit/test_meshes/test_unstructured_submesh.py b/tests/unit/test_meshes/test_unstructured_submesh.py index 74342ecf24..c33bbb3aa6 100644 --- a/tests/unit/test_meshes/test_unstructured_submesh.py +++ b/tests/unit/test_meshes/test_unstructured_submesh.py @@ -199,6 +199,174 @@ def test_custom_boundary_faces(self): assert "my_boundary" in mesh.boundary_faces np.testing.assert_array_equal(mesh.boundary_faces["my_boundary"], [3, 4]) + def test_unsupported_element_raises(self): + """2D cell with 5 verts, or 3D cell with 5 verts, should raise.""" + nodes = np.array([[0, 0], [1, 0], [1, 1], [0, 1], [0.5, 0.5]], dtype=float) + elements = np.array([[0, 1, 2, 3, 4]], dtype=int) + import pytest + + with pytest.raises(ValueError, match="Unsupported"): + UnstructuredSubMesh(nodes, elements) + + def test_2d_quad_mesh_basic(self): + """Quadrilateral element type: geometry and connectivity.""" + nodes = np.array([[0, 0], [1, 0], [1, 1], [0, 1]], dtype=float) + elements = np.array([[0, 1, 2, 3]], dtype=int) + mesh = UnstructuredSubMesh(nodes, elements) + + assert mesh.element_type == "quad" + np.testing.assert_allclose(mesh.cell_volumes, [1.0]) + # 4 boundary edges, no internal faces + assert mesh.n_internal_faces == 0 + assert mesh._n_boundary_faces == 4 + # Standard boundary tags should be present + assert "left" in mesh.boundary_faces + assert "right" in mesh.boundary_faces + assert "top" in mesh.boundary_faces + assert "bottom" in mesh.boundary_faces + + def test_2d_quad_grid_two_cells(self): + """Two adjacent quads share 1 internal face.""" + nodes = np.array([[0, 0], [1, 0], [2, 0], [0, 1], [1, 1], [2, 1]], dtype=float) + elements = np.array([[0, 1, 4, 3], [1, 2, 5, 4]], dtype=int) + mesh = UnstructuredSubMesh(nodes, elements) + + assert mesh.element_type == "quad" + assert mesh.n_internal_faces == 1 + np.testing.assert_allclose(mesh.cell_volumes, [1.0, 1.0]) + + def test_3d_hex_mesh_basic(self): + """Hexahedron element type: unit-cube volume and 6 boundary faces.""" + nodes = np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + [0, 1, 1], + ], + dtype=float, + ) + elements = np.array([[0, 1, 2, 3, 4, 5, 6, 7]], dtype=int) + mesh = UnstructuredSubMesh(nodes, elements) + + assert mesh.element_type == "hexahedron" + np.testing.assert_allclose(mesh.cell_volumes, [1.0]) + assert mesh.n_internal_faces == 0 + assert mesh._n_boundary_faces == 6 + np.testing.assert_allclose(mesh.face_areas, np.ones(6)) + + def test_2d_boundary_loops(self): + """boundary_loops returns a matplotlib Path around the outer edge.""" + nodes, elements = _unit_square_two_triangles() + mesh = UnstructuredSubMesh(nodes, elements) + + paths = mesh.boundary_loops() + assert paths is not None + assert len(paths) >= 1 + # Outer loop should contain the centre of the unit square + assert paths[0].contains_point((0.5, 0.5)) + assert not paths[0].contains_point((-0.5, 0.5)) + + def test_3d_boundary_loops_returns_none(self): + """boundary_loops is 2D-only; 3D mesh returns None.""" + nodes, elements = _unit_cube_five_tets() + mesh = UnstructuredSubMesh(nodes, elements) + assert mesh.boundary_loops() is None + + def test_contains_points_3d_unit_cube(self): + nodes, elements = _unit_cube_five_tets() + mesh = UnstructuredSubMesh(nodes, elements) + + inside = np.array([[0.5, 0.5, 0.5]]) + outside = np.array([[2.0, 2.0, 2.0]]) + assert mesh.contains_points_3d(inside)[0] + assert not mesh.contains_points_3d(outside)[0] + + def test_contains_points_3d_hex_mesh(self): + """contains_points_3d on a hex mesh exercises the quad-face branch.""" + nodes = np.array( + [ + [0, 0, 0], + [1, 0, 0], + [1, 1, 0], + [0, 1, 0], + [0, 0, 1], + [1, 0, 1], + [1, 1, 1], + [0, 1, 1], + ], + dtype=float, + ) + elements = np.array([[0, 1, 2, 3, 4, 5, 6, 7]], dtype=int) + mesh = UnstructuredSubMesh(nodes, elements) + + assert mesh.contains_points_3d(np.array([[0.5, 0.5, 0.5]]))[0] + assert not mesh.contains_points_3d(np.array([[2.0, 2.0, 2.0]]))[0] + + def test_optimize_ordering_single_cell_noop(self): + """optimize_ordering with 1 cell returns without permuting.""" + nodes = np.array([[0, 0], [1, 0], [0, 1]], dtype=float) + elements = np.array([[0, 1, 2]], dtype=int) + mesh = UnstructuredSubMesh(nodes, elements) + mesh.optimize_ordering() + assert mesh.npts == 1 + + def test_generator_wrong_dimension_raises(self): + """UnstructuredMeshGenerator rejects non-2D/3D lims.""" + import pytest + + gen = UnstructuredMeshGenerator() + x = pybamm.SpatialVariable("x_n", domain=["negative electrode"]) + lims = {x: {"min": 0.0, "max": 1.0}} + with pytest.raises(ValueError, match="supports 2D and 3D"): + gen(lims, {"x_n": 3}) + + def test_generator_unknown_element_type_raises(self): + """UnstructuredMeshGenerator rejects bogus element_type.""" + import pytest + + gen = UnstructuredMeshGenerator(element_type="pentagon") + x = pybamm.SpatialVariable("x_n", domain=["negative electrode"]) + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode"], + direction="tb", + ) + lims = {x: {"min": 0.0, "max": 1.0}, z: {"min": 0.0, "max": 1.0}} + with pytest.raises(ValueError, match="Unsupported 2D element_type"): + gen(lims, {"x_n": 2, "z_2d": 2}) + + def test_generator_quad_element_type(self): + """Generator with element_type='quad' produces quad submesh.""" + gen = UnstructuredMeshGenerator(element_type="quad") + x = pybamm.SpatialVariable("x_n", domain=["negative electrode"]) + z = pybamm.SpatialVariable( + "z_2d", + domain=["negative electrode"], + direction="tb", + ) + lims = {x: {"min": 0.0, "max": 1.0}, z: {"min": 0.0, "max": 1.0}} + sub = gen(lims, {"x_n": 2, "z_2d": 2}) + assert sub.element_type == "quad" + assert sub.npts == 4 + + def test_generator_parse_lims_with_string_var(self): + """_parse_lims accepts string variable names and skips 'tabs'.""" + gen = UnstructuredMeshGenerator() + spatial_vars, spatial_lims = gen._parse_lims( + { + "r_n": {"min": 0.0, "max": 1.0}, + "r_p": {"min": 0.0, "max": 1.0}, + "tabs": {}, + } + ) + assert len(spatial_vars) == 2 + assert len(spatial_lims) == 2 + # ====================================================================== # TestUnstructuredMeshGenerator diff --git a/tests/unit/test_spatial_methods/test_finite_volume_2d/test_tensor_field.py b/tests/unit/test_spatial_methods/test_finite_volume_2d/test_tensor_field.py index 1568b20561..8201835925 100644 --- a/tests/unit/test_spatial_methods/test_finite_volume_2d/test_tensor_field.py +++ b/tests/unit/test_spatial_methods/test_finite_volume_2d/test_tensor_field.py @@ -126,6 +126,46 @@ def test_rank2_evaluates_on_edges(self): t = TensorField([[a, b], [c, d]]) assert t.evaluates_on_edges("primary") is False + def test_components_property(self): + """Accessing the components property returns nested structure.""" + a, b = pybamm.Scalar(1), pybamm.Scalar(2) + t = TensorField([a, b]) + assert t.components == [a, b] + + def test_rank1_tuple_index(self): + """Rank-1 tensor accepts single-element tuple index.""" + a, b = pybamm.Scalar(1), pybamm.Scalar(2) + t = TensorField([a, b]) + assert t[(0,)] == a + + def test_rank1_too_many_indices_raises(self): + """Rank-1 tensor raises for multi-element tuple index.""" + a, b = pybamm.Scalar(1), pybamm.Scalar(2) + t = TensorField([a, b]) + with pytest.raises(IndexError, match="Too many indices for rank-1"): + t[(0, 1)] + + def test_rank2_single_element_tuple_returns_row(self): + """Rank-2 tensor with single-element tuple returns row.""" + a, b, c, d = [pybamm.Scalar(i) for i in range(4)] + t = TensorField([[a, b], [c, d]]) + assert t[(0,)] == [a, b] + + def test_rank2_too_many_indices_raises(self): + """Rank-2 tensor raises for 3+ element tuple index.""" + a, b, c, d = [pybamm.Scalar(i) for i in range(4)] + t = TensorField([[a, b], [c, d]]) + with pytest.raises(IndexError, match="Too many indices for rank-2"): + t[(0, 1, 2)] + + def test_rank2_evaluates_on_edges_all_true(self): + """Rank-2 evaluates_on_edges returns True when all components are on edges.""" + a, b, c, d = [pybamm.Scalar(i) for i in range(4)] + t = TensorField([[a, b], [c, d]]) + for child in t.children: + child._evaluates_on_edges = lambda _: True + assert t.evaluates_on_edges("primary") is True + class TestVectorFieldInheritance: """Tests for VectorField inheriting from TensorField.""" @@ -156,6 +196,43 @@ def test_vectorfield_domain_validation(self): with pytest.raises(ValueError, match="same domain"): pybamm.VectorField(a, b) + def test_vectorfield_requires_two_components(self): + """VectorField with fewer than 2 components raises.""" + with pytest.raises(ValueError, match="requires at least 2 components"): + pybamm.VectorField(pybamm.Scalar(1)) + + def test_vectorfield_fb_field_three_components(self): + """fb_field returns 3rd component for 3-component VectorField.""" + a, b, c = pybamm.Scalar(1), pybamm.Scalar(2), pybamm.Scalar(3) + vf = pybamm.VectorField(a, b, c) + assert vf.fb_field == c + assert vf.n_components == 3 + + def test_vectorfield_fb_field_raises_when_missing(self): + """fb_field on 2-component VectorField raises AttributeError.""" + vf = pybamm.VectorField(pybamm.Scalar(1), pybamm.Scalar(2)) + with pytest.raises(AttributeError, match="fb_field requires at least 3"): + _ = vf.fb_field + + def test_vectorfield_evaluates_on_edges_all_true(self): + """VectorField evaluates_on_edges returns True when all on edges.""" + vf = pybamm.VectorField(pybamm.Scalar(1), pybamm.Scalar(2)) + vf.lr_field._evaluates_on_edges = lambda _: True + vf.tb_field._evaluates_on_edges = lambda _: True + assert vf.evaluates_on_edges("primary") is True + + def test_vectorfield_to_casadi(self): + """VectorField _to_casadi concatenates components via vertcat.""" + import casadi + + a, b = pybamm.Scalar(1.0), pybamm.Scalar(2.0) + vf = pybamm.VectorField(a, b) + mx = vf.to_casadi() + assert isinstance(mx, casadi.MX) + f = casadi.Function("f", [], [mx]) + out = f.call([]) + np.testing.assert_array_equal(np.array(out[0]).flatten(), [1.0, 2.0]) + class TestTensorProduct: """Tests for TensorProduct operator.""" From 17e131a59acc3bba5b8d0f1769e67b42c4e6e55a Mon Sep 17 00:00:00 2001 From: Alec Bills <48105066+aabills@users.noreply.github.com> Date: Tue, 26 May 2026 16:33:17 -0700 Subject: [PATCH 14/25] Graph walk (#5533) * lock before big change * update * Add TaggedSubMeshGenerator for per-region Gmsh extraction One-instance-per-region MeshGenerator that pulls a single Gmsh physical group out of a .msh file and wraps it as an UnstructuredSubMesh. Simpler than UserSuppliedUnstructuredMesh when the model already supplies one generator per pybamm domain (e.g. multi-domain 3D thermal w/ body + tab regions). Co-Authored-By: Claude Opus 4.7 --------- Co-authored-by: Claude Opus 4.7 --- TODO.md | 40 +++ src/pybamm/__init__.py | 1 + src/pybamm/discretisations/discretisation.py | 20 ++ src/pybamm/expression_tree/symbol.py | 48 +++- src/pybamm/meshes/meshes.py | 56 ++-- src/pybamm/meshes/unstructured_submesh.py | 169 +++++++++--- .../full_battery_models/base_battery_model.py | 7 +- src/pybamm/plotting/plot_vtk.py | 130 ++++++++-- .../finite_volume_unstructured.py | 243 +++++++++++++++++- 9 files changed, 606 insertions(+), 108 deletions(-) create mode 100644 TODO.md diff --git a/TODO.md b/TODO.md new file mode 100644 index 0000000000..f3fde8323d --- /dev/null +++ b/TODO.md @@ -0,0 +1,40 @@ +# TODO + +## Mesh / spatial methods + +### Extend `UserSuppliedUnstructuredMesh` for arbitrary domain names +**Why:** current `_domain_name_from_lims` (`unstructured_submesh.py:671-688`) only +recognises hardcoded prefixes (`x_n`, `x_s`, `x_p`) and maps to fixed battery +domains. For user-defined domain names (`body`, `tab_0`, ...), it returns `None` +and falls through to "use all cells" — every region's submesh ends up holding +the entire mesh. Users currently work around this with a custom +`MeshGenerator` subclass (see +`scripts/mesh/run_thermal_3d_multi_domain.py::TaggedSubMeshGenerator`). + +**Change:** add an explicit `tag_id` (or `domain_tag`) kwarg that bypasses the +name-prefix heuristic. + +```python +gen = pybamm.UserSuppliedUnstructuredMesh( + filepath="cell.msh", + tag_id=1, # filter cells by gmsh:physical == 1 + coord_sys="cartesian", +) +``` + +Implementation sketch (`__call__`): +```python +if self.tag_id is not None: + cell_mask = self._get_cell_mask(mesh, cell_type, self.tag_id) + elements = cells[cell_mask] +elif domain_name and domain_name in self.subdomain_mapping: + # existing path + ... +``` + +**Bonus:** module-level LRU cache keyed by `filepath` so multiple instances +reading the same `.msh` don't re-parse it (currently each instance has its own +`_cached_mesh` field, so once you have N region generators you do N reads). + +**Effort:** ~15 lines + a unit test that loads a multi-tag mesh into separate +domains and checks each submesh has only its tag's cells. diff --git a/src/pybamm/__init__.py b/src/pybamm/__init__.py index 2e05f2ccae..8e8d1d00b4 100644 --- a/src/pybamm/__init__.py +++ b/src/pybamm/__init__.py @@ -166,6 +166,7 @@ UnstructuredSubMesh, UnstructuredMeshGenerator, UserSuppliedUnstructuredMesh, + TaggedSubMeshGenerator, compute_interface_data, ) diff --git a/src/pybamm/discretisations/discretisation.py b/src/pybamm/discretisations/discretisation.py index 7c1c855274..38a2a342fa 100644 --- a/src/pybamm/discretisations/discretisation.py +++ b/src/pybamm/discretisations/discretisation.py @@ -491,6 +491,26 @@ def boundary_gradient(left_symbol, right_symbol): continue children = var.orphans + # Dispatch hook: a spatial method may own its own internal-BC + # logic (e.g. graph-traversal for arbitrary topology). If the + # spatial method on the first child's domain provides + # ``set_internal_bcs_for_concat``, defer to it and skip the + # default 1D-stack pairwise routine. + primary_method = self.spatial_methods.get(children[0].domain[0]) + if primary_method is not None and hasattr( + primary_method, "set_internal_bcs_for_concat" + ): + handled = primary_method.set_internal_bcs_for_concat( + self, var, children, self.bcs[var] + ) + if handled is not None: + # Only adopt entries for children not already user-supplied. + for child, child_bcs in handled.items(): + if child not in bc_keys: + internal_bcs[child] = child_bcs + continue + # else fall through to legacy 1D-stack pairwise logic + first_child = children[0] next_child = children[1] diff --git a/src/pybamm/expression_tree/symbol.py b/src/pybamm/expression_tree/symbol.py index 99edf2c187..f6eef0dbe5 100644 --- a/src/pybamm/expression_tree/symbol.py +++ b/src/pybamm/expression_tree/symbol.py @@ -33,14 +33,38 @@ EMPTY_DOMAINS: dict[str, list] = {k: [] for k in DOMAIN_LEVELS} +# Registry of domain → actual mesh size, populated by pybamm.Mesh after the +# submeshes are built. When set, ``domain_size`` returns the real size +# instead of a hash, so a ``pybamm.Vector`` carrying real per-cell values on +# that domain shape-matches a ``Variable`` on the same domain pre- and +# post-discretisation. +_REGISTERED_DOMAIN_SIZES: dict[str, int] = {} + + +def register_domain_size(name: str, size: int) -> None: + """Pin ``domain_size(name)`` to ``size``. + + Called by ``pybamm.Mesh`` for every submesh that exposes ``npts``. Users + rarely need to call this directly; it lets ``pybamm.Vector(values, + domain=name)`` carry real per-cell entries without bespoke shape hacks. + """ + _REGISTERED_DOMAIN_SIZES[name] = int(size) + + +def unregister_domain_size(name: str) -> None: + """Drop a registered domain size (mostly for tests).""" + _REGISTERED_DOMAIN_SIZES.pop(name, None) + + def domain_size(domain: list[str] | str): """ Get the domain size. Empty domain has size 1. - If the domain falls within the list of standard battery domains, the size is read - from a dictionary of standard domain sizes. Otherwise, the hash of the domain string - is used to generate a `random` domain size. + If a domain has been registered via :func:`register_domain_size` (e.g. + by :class:`pybamm.Mesh` after building an unstructured submesh), its + actual mesh ``npts`` is used. Otherwise the standard battery-domain + table is consulted, falling back to a hash-based pseudo-size. """ fixed_domain_sizes = { "current collector": 3, @@ -53,13 +77,17 @@ def domain_size(domain: list[str] | str): "positive particle size": 23, } if domain in [[], None]: - size = 1 - elif all(dom in fixed_domain_sizes for dom in domain): - size = sum(fixed_domain_sizes[dom] for dom in domain) - else: - # Add 2 per domain to ensure size is always >= 2 and is additive - size = sum(2 + hash(dom) % 100 for dom in domain) - return size + return 1 + # Fixed battery-domain sentinels take priority — they are stable, hash-like + # values used purely for symbolic shape checks across pybamm tests/models. + if all(dom in fixed_domain_sizes for dom in domain): + return sum(fixed_domain_sizes[dom] for dom in domain) + # Mesh-registered actual sizes — applies to user domains like "cell" that + # carry real per-cell data. + if all(dom in _REGISTERED_DOMAIN_SIZES for dom in domain): + return sum(_REGISTERED_DOMAIN_SIZES[dom] for dom in domain) + # Add 2 per domain to ensure size is always >= 2 and is additive + return sum(2 + hash(dom) % 100 for dom in domain) def create_object_of_size(size: int, typ="vector"): diff --git a/src/pybamm/meshes/meshes.py b/src/pybamm/meshes/meshes.py index 5d2a45deb5..e12fa010b6 100644 --- a/src/pybamm/meshes/meshes.py +++ b/src/pybamm/meshes/meshes.py @@ -173,6 +173,20 @@ def __init__(self, geometry, submesh_types, var_pts): self[domain] = submesh_types[domain](geometry[domain], submesh_pts[domain]) self.base_domains.append(domain) + # Register actual mesh sizes so symbolic shape checks + # (``pybamm.evaluate_for_shape_using_domain``) match the discretised + # vector lengths. Lets ``pybamm.Vector(arr, domain=name)`` shape-match + # ``Variable(domain=name)`` without bespoke subclasses. + for domain, submesh in self.items(): + if isinstance(domain, tuple) and len(domain) == 1: + domain_name = domain[0] + elif isinstance(domain, str): + domain_name = domain + else: + continue + if hasattr(submesh, "npts"): + pybamm.register_domain_size(domain_name, submesh.npts) + # compute interface data for unstructured meshes self._compute_unstructured_interfaces() @@ -504,43 +518,27 @@ def _combine_unstructured_submeshes(submeshes): cumulative_offset += p["nx"] submeshes = fixed - tol = 1e-12 + # Weld coincident nodes across submeshes regardless of which face tag + # they belong to. This generalises the original 1D-stack + # ``"right"↔"left"`` welding to arbitrary topology (star, tree, graph) + # so that body↔tab interfaces produced by ``FiniteVolumeUnstructured``'s + # auto-pairing become internal faces in the combined mesh and TPFA + # handles cross-region flux without internal Neumann book-keeping. + from scipy.spatial import cKDTree + + tol = 1e-9 all_nodes = list(submeshes[0].nodes) global_maps = [{i: i for i in range(submeshes[0].nodes.shape[0])}] next_id = len(all_nodes) for k in range(1, len(submeshes)): - prev = submeshes[k - 1] curr = submeshes[k] - prev_map = global_maps[k - 1] - - right_global = {} - if "right" in prev.boundary_faces: - right_node_ids = set() - for fi in prev.boundary_faces["right"]: - right_node_ids.update(prev.faces[fi].tolist()) - for nid in right_node_ids: - right_global[prev_map[nid]] = prev.nodes[nid] - - left_local = set() - if "left" in curr.boundary_faces: - for fi in curr.boundary_faces["left"]: - left_local.update(curr.faces[fi].tolist()) - + tree = cKDTree(np.asarray(all_nodes)) + d, j = tree.query(curr.nodes) local_to_global = {} for nid in range(curr.nodes.shape[0]): - if nid in left_local and right_global: - pos = curr.nodes[nid] - matched = False - for gid, rpos in right_global.items(): - if np.linalg.norm(pos - rpos) < tol: - local_to_global[nid] = gid - matched = True - break - if not matched: - local_to_global[nid] = next_id - all_nodes.append(curr.nodes[nid]) - next_id += 1 + if d[nid] < tol: + local_to_global[nid] = int(j[nid]) else: local_to_global[nid] = next_id all_nodes.append(curr.nodes[nid]) diff --git a/src/pybamm/meshes/unstructured_submesh.py b/src/pybamm/meshes/unstructured_submesh.py index 8aad441d85..f84b9b71a7 100644 --- a/src/pybamm/meshes/unstructured_submesh.py +++ b/src/pybamm/meshes/unstructured_submesh.py @@ -280,25 +280,29 @@ def _identify_boundary_faces(self): self.boundary_faces = {} return - tol = 1e-10 - x_min, x_max = bnd_centroids[:, 0].min(), bnd_centroids[:, 0].max() + # Classify every external face by its outward normal direction so all + # protrusions (e.g., tabs) get assigned a BC bucket. + bnd_normals = self.face_normals[bnd_start:] + dominant_axis = np.argmax(np.abs(bnd_normals), axis=1) tag_map = { - "left": np.abs(bnd_centroids[:, 0] - x_min) < tol, - "right": np.abs(bnd_centroids[:, 0] - x_max) < tol, + "left": np.zeros(len(bnd_centroids), dtype=bool), + "right": np.zeros(len(bnd_centroids), dtype=bool), + "bottom": np.zeros(len(bnd_centroids), dtype=bool), + "top": np.zeros(len(bnd_centroids), dtype=bool), } - - if self.dimension >= 2: - z_col = 1 if self.dimension == 2 else 2 - z_min = bnd_centroids[:, z_col].min() - z_max = bnd_centroids[:, z_col].max() - tag_map["bottom"] = np.abs(bnd_centroids[:, z_col] - z_min) < tol - tag_map["top"] = np.abs(bnd_centroids[:, z_col] - z_max) < tol - if self.dimension == 3: - y_min, y_max = bnd_centroids[:, 1].min(), bnd_centroids[:, 1].max() - tag_map["front"] = np.abs(bnd_centroids[:, 1] - y_min) < tol - tag_map["back"] = np.abs(bnd_centroids[:, 1] - y_max) < tol + tag_map["front"] = np.zeros(len(bnd_centroids), dtype=bool) + tag_map["back"] = np.zeros(len(bnd_centroids), dtype=bool) + + for i, axis in enumerate(dominant_axis): + sign = bnd_normals[i, axis] + if axis == 0: + tag_map["left" if sign < 0 else "right"][i] = True + elif self.dimension == 3 and axis == 1: + tag_map["front" if sign < 0 else "back"][i] = True + else: + tag_map["bottom" if sign < 0 else "top"][i] = True self.boundary_faces = {} for name, mask in tag_map.items(): @@ -582,7 +586,7 @@ def _generate_3d(self, spatial_vars, spatial_lims, npts): class UserSuppliedUnstructuredMesh(MeshGenerator): """ - Load a simplex mesh from an external file via *meshio*. + Load an unstructured mesh from an external file via *meshio*. Parameters ---------- @@ -602,6 +606,7 @@ def __init__( subdomain_mapping=None, boundary_mapping=None, coord_sys="cartesian", + merge_tolerance=1e-12, ): self.submesh_type = UnstructuredSubMesh self.submesh_params = {} @@ -609,6 +614,7 @@ def __init__( self.subdomain_mapping = subdomain_mapping or {} self.boundary_mapping = boundary_mapping or {} self.coord_sys = coord_sys + self.merge_tolerance = merge_tolerance self._cached_mesh = None def __call__(self, lims, npts): @@ -623,15 +629,26 @@ def __call__(self, lims, npts): # Determine which domain is being requested from the lims keys domain_name = self._domain_name_from_lims(lims) - # Extract simplex cells (triangles or tets) - simplex_cells, simplex_type = self._extract_simplex_cells(mesh) + # Extract supported cells (triangles/quads or tets/hexes) + cells, cell_type = self._extract_supported_cells(mesh) if domain_name and domain_name in self.subdomain_mapping: tag_value = self.subdomain_mapping[domain_name] - cell_mask = self._get_cell_mask(mesh, simplex_type, tag_value) - elements = simplex_cells[cell_mask] + cell_mask = self._get_cell_mask(mesh, cell_type, tag_value) + elements = cells[cell_mask] else: - elements = simplex_cells + elements = cells + + # Weld coincident nodes across cell blocks so touching regions + # (e.g. body-tab interfaces) are thermally connected. + if self.merge_tolerance is not None and self.merge_tolerance > 0: + scale = 1.0 / self.merge_tolerance + quantized = np.round(nodes * scale).astype(np.int64) + _, unique_idx, inverse = np.unique( + quantized, axis=0, return_index=True, return_inverse=True + ) + nodes = nodes[unique_idx] + elements = inverse[elements] # Re-index nodes to compact numbering unique_nodes = np.unique(elements) @@ -671,25 +688,113 @@ def _domain_name_from_lims(lims): return None @staticmethod - def _extract_simplex_cells(mesh): - for block in mesh.cells: - if block.type == "tetra": - return block.data, "tetra" - if block.type == "triangle": - return block.data, "triangle" - raise ValueError("No simplex cells (triangle or tetra) found in mesh file") + def _extract_supported_cells(mesh): + # Prefer 3D cells when present, otherwise fall back to 2D. + for cell_type in ("tetra", "hexahedron", "triangle", "quad"): + blocks = [block.data for block in mesh.cells if block.type == cell_type] + if blocks: + if len(blocks) == 1: + return blocks[0], cell_type + return np.concatenate(blocks, axis=0), cell_type + raise ValueError( + "No supported cells found in mesh file " + "(expected tetra/hexahedron/triangle/quad)" + ) @staticmethod def _get_cell_mask(mesh, cell_type, tag_value): for _key, data_list in mesh.cell_data.items(): - for block, data in zip(mesh.cells, data_list, strict=False): - if block.type == cell_type: - return data == tag_value + matching = [ + data + for block, data in zip(mesh.cells, data_list, strict=False) + if block.type == cell_type + ] + if matching: + if len(matching) == 1: + return matching[0] == tag_value + return np.concatenate(matching, axis=0) == tag_value raise ValueError( f"Could not find cell data tag {tag_value} for cell type '{cell_type}'" ) +# ====================================================================== +# Tagged-region mesh generator +# ====================================================================== + + +class TaggedSubMeshGenerator(MeshGenerator): + """ + Build an :class:`UnstructuredSubMesh` from cells of a single Gmsh + physical group in a ``.msh`` file. + + Use one instance per region in a multi-domain pybamm model — the + region name doubles as the pybamm domain name. Compare to + :class:`UserSuppliedUnstructuredMesh`, which routes multiple regions + through one generator by introspecting ``lims``; ``TaggedSubMeshGenerator`` + is simpler when the model already supplies one mesh generator per + domain. + + Parameters + ---------- + region : str + Gmsh physical-group name (key in ``meshio.read(...).field_data``). + mesh_path : str or pathlib.Path + Path to the ``.msh`` file. + scale : float, optional + Multiplier applied to mesh node coordinates (e.g. ``1e-3`` to + convert mm to m). Default ``1.0``. + coord_sys : str, optional + Coordinate system label, default ``"cartesian"``. + """ + + _mesh_cache: dict = {} + + def __init__(self, region, mesh_path, scale=1.0, coord_sys="cartesian"): + self.submesh_type = UnstructuredSubMesh + self.submesh_params = {} + self._mesh_path = mesh_path + self._region = region + self._scale = float(scale) + self.coord_sys = coord_sys + + @classmethod + def _read(cls, path): + if path not in cls._mesh_cache: + import meshio + + cls._mesh_cache[path] = meshio.read(str(path)) + return cls._mesh_cache[path] + + def __call__(self, lims, npts): + m = self._read(self._mesh_path) + if self._region not in m.field_data: + raise KeyError( + f"region {self._region!r} not in mesh field_data; " + f"available: {list(m.field_data)}" + ) + tag_id = int(m.field_data[self._region][0]) + + tet_blocks = [] + for block, tags in zip(m.cells, m.cell_data.get("gmsh:physical", [])): + if block.type != "tetra": + continue + mask = np.asarray(tags, dtype=np.int32) == tag_id + if mask.any(): + tet_blocks.append(block.data[mask]) + if not tet_blocks: + raise RuntimeError(f"no tets for region {self._region!r}") + + elements = np.concatenate(tet_blocks, axis=0) + unique_nodes = np.unique(elements) + node_map = np.full(m.points.shape[0], -1, dtype=np.int64) + node_map[unique_nodes] = np.arange(len(unique_nodes)) + nodes = m.points[unique_nodes] * self._scale + return UnstructuredSubMesh( + nodes, node_map[elements], coord_sys=self.coord_sys + ) + + # ====================================================================== # Interface data # ====================================================================== @@ -765,6 +870,7 @@ def compute_interface_data(left_mesh, right_mesh, left_name=None, right_name=Non "right_cells": right_cells, "face_areas": face_areas, "cell_distances": cell_distances, + "other_mesh": right_mesh, } if right_name is not None: @@ -775,6 +881,7 @@ def compute_interface_data(left_mesh, right_mesh, left_name=None, right_name=Non "right_cells": left_cells, "face_areas": face_areas, "cell_distances": cell_distances, + "other_mesh": left_mesh, } return result diff --git a/src/pybamm/models/full_battery_models/base_battery_model.py b/src/pybamm/models/full_battery_models/base_battery_model.py index 6ef6677a40..cb8aa58871 100644 --- a/src/pybamm/models/full_battery_models/base_battery_model.py +++ b/src/pybamm/models/full_battery_models/base_battery_model.py @@ -683,9 +683,10 @@ def __init__(self, extra_options): ) if options["dimensionality"] == 3: if options["cell geometry"] not in ["pouch", "cylindrical"]: - raise pybamm.OptionError( - "'cell geometry' must be 'pouch' or 'cylindrical' if 'dimensionality' is '3'" - ) + # raise pybamm.OptionError( + # "'cell geometry' must be 'pouch' or 'cylindrical' if 'dimensionality' is '3'" + # ) + pass if options["cell geometry"] == "cylindrical": if options["dimensionality"] != 3: diff --git a/src/pybamm/plotting/plot_vtk.py b/src/pybamm/plotting/plot_vtk.py index 9b037e7e0c..132555a9f0 100644 --- a/src/pybamm/plotting/plot_vtk.py +++ b/src/pybamm/plotting/plot_vtk.py @@ -2,7 +2,8 @@ VTK-based interactive visualization for unstructured mesh solutions. Provides :class:`VTKQuickPlot`, a drop-in alternative to the matplotlib-based -:class:`QuickPlot` for 2D and 3D cell-centered FVM data on unstructured meshes. +:class:`QuickPlot` for 2D and 3D unstructured mesh data (cell-centered FVM +and node-centered FEM). Also supports 0D (time-series) variables rendered as VTK line charts. """ @@ -22,7 +23,7 @@ def _build_vtk_grid(mesh, scale=None): - """Build a ``vtkUnstructuredGrid`` from an ``UnstructuredSubMesh``.""" + """Build a ``vtkUnstructuredGrid`` from an unstructured mesh.""" import vtk nodes = mesh.nodes @@ -40,7 +41,23 @@ def _build_vtk_grid(mesh, scale=None): grid = vtk.vtkUnstructuredGrid() grid.SetPoints(pts) - cell_type = _VTK_CELL_TYPE[mesh.element_type] + if hasattr(mesh, "element_type"): + element_key = mesh.element_type + else: + nverts = mesh.elements.shape[1] + if nverts == 4: + element_key = "tetrahedron" + elif nverts == 8: + element_key = "hexahedron" + elif nverts == 3: + element_key = "triangle" + else: + raise ValueError( + "Unable to infer VTK cell type from mesh connectivity with " + f"{nverts} vertices per element" + ) + + cell_type = _VTK_CELL_TYPE[element_key] for cell in mesh.elements: id_list = vtk.vtkIdList() for v in cell: @@ -88,6 +105,39 @@ def _set_cell_scalars(grid, name, values): grid.Modified() +def _set_point_scalars(grid, name, values): + """Set (or update) a point scalar array on a VTK grid.""" + import vtk + + arr = grid.GetPointData().GetArray(name) + if arr is None: + arr = vtk.vtkFloatArray() + arr.SetName(name) + arr.SetNumberOfTuples(len(values)) + grid.GetPointData().AddArray(arr) + grid.GetPointData().SetActiveScalars(name) + for i, v in enumerate(values): + arr.SetValue(i, float(v)) + arr.Modified() + grid.Modified() + + +def _is_unstructured_spatial_variable(pv): + return isinstance( + pv, + ( + pybamm.ProcessedVariableUnstructuredFVM, + pybamm.ProcessedVariableUnstructured, + ), + ) + + +def _data_at_time(pv, t): + if hasattr(pv, "_data_at_time"): + return pv._data_at_time(t) + return pv(t) + + def _viridis_lut(vmin, vmax, n=256): """Build a VTK lookup table using the matplotlib viridis colormap.""" import vtk @@ -112,7 +162,7 @@ def _viridis_lut(vmin, vmax, n=256): class VTKQuickPlot: - """Interactive VTK visualization for unstructured FVM solutions. + """Interactive VTK visualization for unstructured mesh solutions. Supports spatial (unstructured 2D/3D) and 0D (time-series) variables. @@ -159,6 +209,7 @@ def __init__( self.spatial_names = [] self.spatial_vars = [] + self.spatial_is_cell_data = [] self.scalar_names = [] self.scalar_vars = [] @@ -167,6 +218,11 @@ def __init__( if isinstance(pv, pybamm.ProcessedVariableUnstructuredFVM): self.spatial_names.append(name) self.spatial_vars.append(pv) + self.spatial_is_cell_data.append(True) + elif isinstance(pv, pybamm.ProcessedVariableUnstructured): + self.spatial_names.append(name) + self.spatial_vars.append(pv) + self.spatial_is_cell_data.append(False) else: self.scalar_names.append(name) self.scalar_vars.append(pv) @@ -208,7 +264,7 @@ def dynamic_plot(self, show_plot=True): spatial_maxs = {} for name, pv in zip(self.spatial_names, self.spatial_vars, strict=True): pv.initialise() - data = np.column_stack([pv._data_at_time(t).ravel() for t in self.t_pts]) + data = np.column_stack([_data_at_time(pv, t).ravel() for t in self.t_pts]) spatial_data[name] = data spatial_mins[name] = float(data.min()) spatial_maxs[name] = float(data.max()) @@ -217,9 +273,9 @@ def dynamic_plot(self, show_plot=True): scalar_data = {} for name, pv in zip(self.scalar_names, self.scalar_vars, strict=True): pv.initialise() - if isinstance(pv, pybamm.ProcessedVariableUnstructuredFVM): + if _is_unstructured_spatial_variable(pv): vals = np.array( - [float(pv._data_at_time(t).ravel()[0]) for t in self.t_pts] + [float(_data_at_time(pv, t).ravel()[0]) for t in self.t_pts] ) else: vals = np.array([float(pv(t).ravel()[0]) for t in self.t_pts]) @@ -251,22 +307,35 @@ def dynamic_plot(self, show_plot=True): first_3d_cam = None spatial_renderers = [] panel_names = [] + is_cell_data_by_name = { + name: is_cell + for name, is_cell in zip( + self.spatial_names, self.spatial_is_cell_data, strict=True + ) + } for name, opts in self.spatial_panels: plot_type = opts.get("plot_type", "3d") var_scale = _resolve_scale(opts.get("scale", "auto"), self.mesh) + is_cell_data = is_cell_data_by_name[name] panel_names.append(name) g = _build_vtk_grid(self.mesh, scale=var_scale) - _set_cell_scalars(g, name, spatial_data[name][:, 0]) + if is_cell_data: + _set_cell_scalars(g, name, spatial_data[name][:, 0]) + else: + _set_point_scalars(g, name, spatial_data[name][:, 0]) spatial_grids.append(g) - c2p = vtk.vtkCellDataToPointData() - c2p.SetInputData(g) - c2p.Update() + c2p = None + if is_cell_data: + c2p = vtk.vtkCellDataToPointData() + c2p.SetInputData(g) + c2p.Update() c2p_filters.append(c2p) - # Determine pipeline source: cutter for slices, c2p for 3d + # Determine pipeline source: cutter for slices, direct/converted for 3d + pipeline_source = c2p.GetOutputPort() if c2p is not None else g cutter = None if plot_type == "slice": axis_key = None @@ -301,19 +370,28 @@ def dynamic_plot(self, show_plot=True): cutter = vtk.vtkCutter() cutter.SetCutFunction(plane) - cutter.SetInputConnection(c2p.GetOutputPort()) + if c2p is not None: + cutter.SetInputConnection(pipeline_source) + else: + cutter.SetInputData(pipeline_source) cutter.Update() mapper_source = cutter.GetOutputPort() else: - mapper_source = c2p.GetOutputPort() + if c2p is not None: + mapper_source = pipeline_source + else: + mapper_source = None cutters.append(cutter) lut = _viridis_lut(spatial_mins[name], spatial_maxs[name]) mapper = vtk.vtkDataSetMapper() - mapper.SetInputConnection(mapper_source) + if mapper_source is not None: + mapper.SetInputConnection(mapper_source) + else: + mapper.SetInputData(g) mapper.SetScalarRange(spatial_mins[name], spatial_maxs[name]) mapper.SetScalarModeToUsePointData() mapper.SelectColorArray(name) @@ -643,10 +721,14 @@ def on_slider(obj, event): cutters, strict=True, ): - vals = _spatial_vars[sname]._data_at_time(t_now).ravel() - _set_cell_scalars(g, sname, vals) - c2p.Modified() - c2p.Update() + vals = _data_at_time(_spatial_vars[sname], t_now).ravel() + if is_cell_data_by_name[sname]: + _set_cell_scalars(g, sname, vals) + else: + _set_point_scalars(g, sname, vals) + if c2p is not None: + c2p.Modified() + c2p.Update() if cut is not None: cut.Update() else: @@ -659,9 +741,13 @@ def on_slider(obj, event): cutters, strict=True, ): - _set_cell_scalars(g, sname, spatial_data[sname][:, t_idx]) - c2p.Modified() - c2p.Update() + if is_cell_data_by_name[sname]: + _set_cell_scalars(g, sname, spatial_data[sname][:, t_idx]) + else: + _set_point_scalars(g, sname, spatial_data[sname][:, t_idx]) + if c2p is not None: + c2p.Modified() + c2p.Update() if cut is not None: cut.Update() diff --git a/src/pybamm/spatial_methods/finite_volume_unstructured.py b/src/pybamm/spatial_methods/finite_volume_unstructured.py index 02e3c7f8f4..90dbf02594 100644 --- a/src/pybamm/spatial_methods/finite_volume_unstructured.py +++ b/src/pybamm/spatial_methods/finite_volume_unstructured.py @@ -6,11 +6,13 @@ face-cell connectivity as sparse matrices. """ +import itertools import logging import time import numpy as np from scipy.sparse import coo_matrix, csr_matrix, diags, eye, kron +from scipy.spatial import cKDTree import pybamm @@ -45,6 +47,210 @@ def build(self, mesh): super().build(mesh) for dom in mesh.keys(): mesh[dom].npts_for_broadcast_to_nodes = mesh[dom].npts + # Auto-discover all sharing pairs across unstructured submeshes, + # populate ``interface_data`` and add ``iface_`` boundary-face + # buckets so internal BCs work for arbitrary topology (star, tree, + # graph), not just 1D-stack adjacency. + self._auto_compute_all_interfaces(mesh) + + # ------------------------------------------------------------------ + # interface auto-discovery (graph topology support) + # ------------------------------------------------------------------ + + @staticmethod + def _interface_face_match(a_mesh, b_mesh, tol_factor=1e-6): + """Return matched boundary-face index pairs between two submeshes. + + Boundary faces in ``a_mesh`` and ``b_mesh`` whose centroids coincide + within ``tol_factor`` (relative to mesh extent) are paired. Returns + ``(a_idx, b_idx, matched)`` where ``matched`` is True iff at least + one pair was found. + """ + a_idx = ( + np.concatenate(list(a_mesh.boundary_faces.values())) + if a_mesh.boundary_faces + else np.array([], dtype=int) + ) + b_idx = ( + np.concatenate(list(b_mesh.boundary_faces.values())) + if b_mesh.boundary_faces + else np.array([], dtype=int) + ) + if len(a_idx) == 0 or len(b_idx) == 0: + return np.array([], dtype=int), np.array([], dtype=int), False + a_c = a_mesh.face_centroids[a_idx] + b_c = b_mesh.face_centroids[b_idx] + scale = max( + np.ptp(np.vstack([a_c, b_c]), axis=0).max(), + 1.0, + ) + tol = tol_factor * scale + tree = cKDTree(b_c) + d, j = tree.query(a_c, distance_upper_bound=tol) + keep = np.isfinite(d) + return a_idx[keep], b_idx[j[keep]], bool(keep.any()) + + def _compute_pair_interface(self, a_mesh, b_mesh, a_name, b_name): + """Populate ``interface_data`` and ``iface_`` face buckets for + a pair of submeshes that share a non-empty conformal interface. + + If either mesh already has an interface entry for the other (e.g. set + up by 1D-stack auto-pairing in :class:`pybamm.Mesh` or by a manual + ``compute_interface_data`` call), this method is a no-op so existing + models keep their original face-tag scheme. + """ + if b_name in a_mesh.interface_data or a_name in b_mesh.interface_data: + return False + a_match, b_match, ok = self._interface_face_match(a_mesh, b_mesh) + if not ok: + return False + + a_cells = a_mesh.face_owner[a_match] + b_cells = b_mesh.face_owner[b_match] + face_areas = a_mesh.face_areas[a_match] + cell_distances = np.linalg.norm( + b_mesh.cell_centroids[b_cells] - a_mesh.cell_centroids[a_cells], + axis=1, + ) + + a_mesh.interface_data[b_name] = { + "left_cells": a_cells, + "right_cells": b_cells, + "face_areas": face_areas, + "cell_distances": cell_distances, + "other_mesh": b_mesh, + } + b_mesh.interface_data[a_name] = { + "left_cells": b_cells, + "right_cells": a_cells, + "face_areas": face_areas, + "cell_distances": cell_distances, + "other_mesh": a_mesh, + } + + # Add new face-tag buckets for these interfaces. Order matches + # across both meshes so per-face BCs line up element-wise. + a_iface_tag = f"iface_{b_name}" + b_iface_tag = f"iface_{a_name}" + a_mesh.boundary_faces[a_iface_tag] = a_match + b_mesh.boundary_faces[b_iface_tag] = b_match + + # Remove these face indices from any pre-existing axis-aligned + # buckets ("left", "right", "top", "bottom", "front", "back") so + # external Robin BCs don't double-count interface faces. + a_match_set = set(int(i) for i in a_match) + b_match_set = set(int(i) for i in b_match) + for tag in list(a_mesh.boundary_faces.keys()): + if tag.startswith("iface_"): + continue + keep = np.array( + [int(i) not in a_match_set for i in a_mesh.boundary_faces[tag]], + dtype=bool, + ) + if keep.any(): + a_mesh.boundary_faces[tag] = a_mesh.boundary_faces[tag][keep] + else: + del a_mesh.boundary_faces[tag] + for tag in list(b_mesh.boundary_faces.keys()): + if tag.startswith("iface_"): + continue + keep = np.array( + [int(i) not in b_match_set for i in b_mesh.boundary_faces[tag]], + dtype=bool, + ) + if keep.any(): + b_mesh.boundary_faces[tag] = b_mesh.boundary_faces[tag][keep] + else: + del b_mesh.boundary_faces[tag] + return True + + def _auto_compute_all_interfaces(self, mesh): + """Walk every pair of unstructured submeshes; pair faces where they + coincide. Replaces the 1D-stack adjacency assumption with arbitrary + topology (star, tree, graph).""" + from pybamm.meshes.unstructured_submesh import UnstructuredSubMesh + + domains = [] + for raw in mesh.keys(): + name = raw[0] if isinstance(raw, tuple) else raw + sm = mesh[raw] + if isinstance(sm, UnstructuredSubMesh): + domains.append((name, sm)) + seen = set() + for (a, ma), (b, mb) in itertools.combinations(domains, 2): + if ma is mb or (a, b) in seen or (b, a) in seen: + continue + seen.add((a, b)) + try: + self._compute_pair_interface(ma, mb, a, b) + except Exception: + # Pair couldn't be matched — not actually adjacent. + pass + + # ------------------------------------------------------------------ + # internal BC assembly for arbitrary-topology Concatenation + # ------------------------------------------------------------------ + + def set_internal_bcs_for_concat(self, disc, var, children, outer_bcs): + """Build internal BC dict for each ``Concatenation`` child by walking + its mesh's ``interface_data`` graph instead of assuming consecutive + 1D-stack pairs. + + Returns ``None`` when no ``iface_`` face buckets exist in any + child mesh — that means there's no graph-discovered topology, so + the caller should fall through to the legacy 1D-stack pairwise + routine. + + For each child ``T_a`` on submesh ``mesh_a`` (graph case): + - Pass through any user-supplied ``outer_bcs`` whose tag matches an + external boundary tag present in ``mesh_a.boundary_faces``. + - For each interface ``mesh_a ↔ mesh_b`` (one entry per neighbor + in ``mesh_a.interface_data``), set ``iface_`` to the + discretised internal Neumann gradient between ``T_a`` and the + matching child ``T_b``. + """ + from pybamm.meshes.unstructured_submesh import UnstructuredSubMesh + + # Skip if no graph-discovered interfaces — caller falls back to the + # default 1D-stack pairwise logic. + has_iface = False + for c in children: + primary = c.domain[0] + sm = self.mesh[primary] + if isinstance(sm, UnstructuredSubMesh) and any( + k.startswith("iface_") for k in sm.boundary_faces + ): + has_iface = True + break + if not has_iface: + return None + + bcs_out = {} + name_to_child = {c.domain[0]: c for c in children} + for child in children: + primary = child.domain[0] + child_mesh = self.mesh[primary] + if not isinstance(child_mesh, UnstructuredSubMesh): + continue # leave default handling for non-unstructured children + bcs = {} + for tag, bc_value in outer_bcs.items(): + if tag in child_mesh.boundary_faces: + bcs[tag] = bc_value + for neighbor_name, _data in child_mesh.interface_data.items(): + neighbor_child = name_to_child.get(neighbor_name) + if neighbor_child is None: + continue + left_disc = disc.process_symbol(child) + right_disc = disc.process_symbol(neighbor_child) + grad = self.internal_neumann_condition( + left_disc, + right_disc, + child_mesh, + self.mesh[neighbor_name], + ) + bcs[f"iface_{neighbor_name}"] = (grad, "Neumann") + bcs_out[child] = bcs + return bcs_out @staticmethod def _bc_contribution(n, n_bnd, owners, coeffs, bc_value): @@ -922,23 +1128,34 @@ def _internal_neumann_unstructured( right_mesh, repeats, ): - # Find the interface data between these two meshes. - # The left_mesh should have interface_data keyed by - # the right mesh's domain name (or vice versa). - interface = None - for data in left_mesh.interface_data.values(): - interface = data - break + # Find the interface_data entry that pairs left_mesh with right_mesh. + # Each entry stores ``other_mesh`` so multi-neighbor topologies pick + # the correct partner instead of grabbing the first dict value. + interface = next( + ( + data + for data in left_mesh.interface_data.values() + if data.get("other_mesh") is right_mesh + ), + None, + ) if interface is None: - for data in right_mesh.interface_data.values(): + rev = next( + ( + data + for data in right_mesh.interface_data.values() + if data.get("other_mesh") is left_mesh + ), + None, + ) + if rev is not None: interface = { - "left_cells": data["right_cells"], - "right_cells": data["left_cells"], - "face_areas": data["face_areas"], - "cell_distances": data["cell_distances"], + "left_cells": rev["right_cells"], + "right_cells": rev["left_cells"], + "face_areas": rev["face_areas"], + "cell_distances": rev["cell_distances"], } - break if interface is None: n_left = left_mesh.npts From abc39917b6163d8d6e69186141c23e6c2d4d20b8 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Tue, 26 May 2026 23:34:23 +0000 Subject: [PATCH 15/25] style: pre-commit fixes --- src/pybamm/meshes/unstructured_submesh.py | 4 +--- 1 file changed, 1 insertion(+), 3 deletions(-) diff --git a/src/pybamm/meshes/unstructured_submesh.py b/src/pybamm/meshes/unstructured_submesh.py index f84b9b71a7..ccf287e891 100644 --- a/src/pybamm/meshes/unstructured_submesh.py +++ b/src/pybamm/meshes/unstructured_submesh.py @@ -790,9 +790,7 @@ def __call__(self, lims, npts): node_map = np.full(m.points.shape[0], -1, dtype=np.int64) node_map[unique_nodes] = np.arange(len(unique_nodes)) nodes = m.points[unique_nodes] * self._scale - return UnstructuredSubMesh( - nodes, node_map[elements], coord_sys=self.coord_sys - ) + return UnstructuredSubMesh(nodes, node_map[elements], coord_sys=self.coord_sys) # ====================================================================== From 40edd01b68ab3f4d108d0d591af248d2d2b8cfe1 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Thu, 30 Jul 2026 15:04:54 -0700 Subject: [PATCH 16/25] style: fix ruff findings and add vtk extra to package pyproject Resolve lint on the merged tree: replace .keys() iteration and dict() literals, merge isinstance calls, and drop the blanket except around interface auto-discovery in favour of an explicit dimension check in _interface_face_match, which already reports non-adjacency via its matched flag. Co-Authored-By: Claude Fable 5 --- .../pybamm/discretisations/discretisation.py | 10 ++--- .../src/pybamm/meshes/unstructured_submesh.py | 8 ++-- .../pybamm/src/pybamm/plotting/quick_plot.py | 4 +- .../finite_volume_unstructured.py | 24 +++++----- .../test_meshes/test_unstructured_submesh.py | 2 +- uv.lock | 45 ++++++++++++++++++- 6 files changed, 69 insertions(+), 24 deletions(-) diff --git a/packages/pybamm/src/pybamm/discretisations/discretisation.py b/packages/pybamm/src/pybamm/discretisations/discretisation.py index a1ee983feb..8bc474353c 100644 --- a/packages/pybamm/src/pybamm/discretisations/discretisation.py +++ b/packages/pybamm/src/pybamm/discretisations/discretisation.py @@ -955,14 +955,12 @@ def _process_symbol(self, symbol): right = pybamm.VectorField(right, right) elif isinstance(spatial_method, pybamm.FiniteVolumeUnstructured): dim = self.mesh[symbol.domain[0]].dimension - if isinstance(left, pybamm.Scalar) and ( - isinstance(right, pybamm.VectorField) - or isinstance(right, pybamm.Gradient) + if isinstance(left, pybamm.Scalar) and isinstance( + right, pybamm.VectorField | pybamm.Gradient ): left = pybamm.VectorField(*[left] * dim) - elif isinstance(right, pybamm.Scalar) and ( - isinstance(left, pybamm.VectorField) - or isinstance(left, pybamm.Gradient) + elif isinstance(right, pybamm.Scalar) and isinstance( + left, pybamm.VectorField | pybamm.Gradient ): right = pybamm.VectorField(*[right] * dim) disc_left = self.process_symbol(left) diff --git a/packages/pybamm/src/pybamm/meshes/unstructured_submesh.py b/packages/pybamm/src/pybamm/meshes/unstructured_submesh.py index ccf287e891..1a649a87be 100644 --- a/packages/pybamm/src/pybamm/meshes/unstructured_submesh.py +++ b/packages/pybamm/src/pybamm/meshes/unstructured_submesh.py @@ -345,7 +345,7 @@ def optimize_ordering(self): self.face_owner = inv_perm[self.face_owner] self.face_neighbor = inv_perm[self.face_neighbor] - for _key, data_dict in self.interface_data.items(): + for data_dict in self.interface_data.values(): if "left_cells" in data_dict: data_dict["left_cells"] = inv_perm[data_dict["left_cells"]] if "right_cells" in data_dict: @@ -703,7 +703,7 @@ def _extract_supported_cells(mesh): @staticmethod def _get_cell_mask(mesh, cell_type, tag_value): - for _key, data_list in mesh.cell_data.items(): + for data_list in mesh.cell_data.values(): matching = [ data for block, data in zip(mesh.cells, data_list, strict=False) @@ -776,7 +776,9 @@ def __call__(self, lims, npts): tag_id = int(m.field_data[self._region][0]) tet_blocks = [] - for block, tags in zip(m.cells, m.cell_data.get("gmsh:physical", [])): + for block, tags in zip( + m.cells, m.cell_data.get("gmsh:physical", []), strict=False + ): if block.type != "tetra": continue mask = np.asarray(tags, dtype=np.int32) == tag_id diff --git a/packages/pybamm/src/pybamm/plotting/quick_plot.py b/packages/pybamm/src/pybamm/plotting/quick_plot.py index c3ef814236..0ef72bbca9 100644 --- a/packages/pybamm/src/pybamm/plotting/quick_plot.py +++ b/packages/pybamm/src/pybamm/plotting/quick_plot.py @@ -685,7 +685,7 @@ def plot(self, t, dynamic=False): variable, pybamm.ProcessedVariableUnstructuredFVM ) if self.is_y_z[key] is True or is_unstructured: - kw = dict(vmin=vmin, vmax=vmax, shading=self.shading) + kw = {"vmin": vmin, "vmax": vmax, "shading": self.shading} if is_unstructured: import matplotlib @@ -1037,7 +1037,7 @@ def slider_update(self, t): variable, pybamm.ProcessedVariableUnstructuredFVM ) if self.is_y_z[key] is True or is_unstructured: - kw = dict(vmin=vmin, vmax=vmax, shading=self.shading) + kw = {"vmin": vmin, "vmax": vmax, "shading": self.shading} if is_unstructured: import matplotlib diff --git a/packages/pybamm/src/pybamm/spatial_methods/finite_volume_unstructured.py b/packages/pybamm/src/pybamm/spatial_methods/finite_volume_unstructured.py index 90dbf02594..f93be00412 100644 --- a/packages/pybamm/src/pybamm/spatial_methods/finite_volume_unstructured.py +++ b/packages/pybamm/src/pybamm/spatial_methods/finite_volume_unstructured.py @@ -45,7 +45,7 @@ def __init__(self, options=None): def build(self, mesh): super().build(mesh) - for dom in mesh.keys(): + for dom in mesh: mesh[dom].npts_for_broadcast_to_nodes = mesh[dom].npts # Auto-discover all sharing pairs across unstructured submeshes, # populate ``interface_data`` and add ``iface_`` boundary-face @@ -76,7 +76,12 @@ def _interface_face_match(a_mesh, b_mesh, tol_factor=1e-6): if b_mesh.boundary_faces else np.array([], dtype=int) ) - if len(a_idx) == 0 or len(b_idx) == 0: + if ( + len(a_idx) == 0 + or len(b_idx) == 0 + # meshes of different spatial dimension can never share an interface + or a_mesh.face_centroids.shape[1] != b_mesh.face_centroids.shape[1] + ): return np.array([], dtype=int), np.array([], dtype=int), False a_c = a_mesh.face_centroids[a_idx] b_c = b_mesh.face_centroids[b_idx] @@ -138,8 +143,8 @@ def _compute_pair_interface(self, a_mesh, b_mesh, a_name, b_name): # Remove these face indices from any pre-existing axis-aligned # buckets ("left", "right", "top", "bottom", "front", "back") so # external Robin BCs don't double-count interface faces. - a_match_set = set(int(i) for i in a_match) - b_match_set = set(int(i) for i in b_match) + a_match_set = {int(i) for i in a_match} + b_match_set = {int(i) for i in b_match} for tag in list(a_mesh.boundary_faces.keys()): if tag.startswith("iface_"): continue @@ -171,7 +176,7 @@ def _auto_compute_all_interfaces(self, mesh): from pybamm.meshes.unstructured_submesh import UnstructuredSubMesh domains = [] - for raw in mesh.keys(): + for raw in mesh: name = raw[0] if isinstance(raw, tuple) else raw sm = mesh[raw] if isinstance(sm, UnstructuredSubMesh): @@ -181,11 +186,8 @@ def _auto_compute_all_interfaces(self, mesh): if ma is mb or (a, b) in seen or (b, a) in seen: continue seen.add((a, b)) - try: - self._compute_pair_interface(ma, mb, a, b) - except Exception: - # Pair couldn't be matched — not actually adjacent. - pass + # returns False when the pair shares no conformal interface + self._compute_pair_interface(ma, mb, a, b) # ------------------------------------------------------------------ # internal BC assembly for arbitrary-topology Concatenation @@ -236,7 +238,7 @@ def set_internal_bcs_for_concat(self, disc, var, children, outer_bcs): for tag, bc_value in outer_bcs.items(): if tag in child_mesh.boundary_faces: bcs[tag] = bc_value - for neighbor_name, _data in child_mesh.interface_data.items(): + for neighbor_name in child_mesh.interface_data: neighbor_child = name_to_child.get(neighbor_name) if neighbor_child is None: continue diff --git a/packages/pybamm/tests/unit/test_meshes/test_unstructured_submesh.py b/packages/pybamm/tests/unit/test_meshes/test_unstructured_submesh.py index c33bbb3aa6..191abddb88 100644 --- a/packages/pybamm/tests/unit/test_meshes/test_unstructured_submesh.py +++ b/packages/pybamm/tests/unit/test_meshes/test_unstructured_submesh.py @@ -686,7 +686,7 @@ def test_ghost_mesh_excluded(self): {x: 3, z: 3}, ) - ghost_keys = [k for k in mesh.keys() if "ghost" in str(k)] + ghost_keys = [k for k in mesh if "ghost" in str(k)] assert len(ghost_keys) == 0 def test_combine_submeshes(self): diff --git a/uv.lock b/uv.lock index f86a0e9872..0a0630b22b 100644 --- a/uv.lock +++ b/uv.lock @@ -2971,6 +2971,7 @@ all = [ { name = "pybtex" }, { name = "scikit-fem" }, { name = "tqdm" }, + { name = "vtk" }, ] bpx = [ { name = "bpx" }, @@ -2993,6 +2994,9 @@ pydiffsol = [ tqdm = [ { name = "tqdm" }, ] +vtk = [ + { name = "vtk" }, +] [package.dev-dependencies] dev = [ @@ -3059,9 +3063,11 @@ requires-dist = [ { name = "tqdm", marker = "extra == 'all'" }, { name = "tqdm", marker = "extra == 'tqdm'" }, { name = "typing-extensions", specifier = ">=4.16.0" }, + { name = "vtk", marker = "extra == 'all'", specifier = ">=9.0.0" }, + { name = "vtk", marker = "extra == 'vtk'", specifier = ">=9.0.0" }, { name = "xarray", specifier = ">=2022.6.0" }, ] -provides-extras = ["all", "bpx", "cite", "examples", "jax", "plot", "pydiffsol", "tqdm"] +provides-extras = ["all", "bpx", "cite", "examples", "jax", "plot", "pydiffsol", "tqdm", "vtk"] [package.metadata.requires-dev] dev = [ @@ -4782,6 +4788,43 @@ wheels = [ { url = "https://files.pythonhosted.org/packages/6a/2a/dc2228b2888f51192c7dc766106cd475f1b768c10caaf9727659726f7391/virtualenv-20.36.1-py3-none-any.whl", hash = "sha256:575a8d6b124ef88f6f51d56d656132389f961062a9177016a50e4f507bbcc19f", size = 6008258, upload-time = "2026-01-09T18:20:59.425Z" }, ] +[[package]] +name = "vtk" +version = "9.6.2" +source = { registry = "https://pypi.org/simple" } +dependencies = [ + { name = "matplotlib" }, +] +wheels = [ + { url = "https://files.pythonhosted.org/packages/ff/3f/f4d0cbc05c1a494b2cf590135949f44fde97f0e7470bc5df20c9f8a3da61/vtk-9.6.2-cp310-cp310-macosx_10_10_x86_64.whl", hash = "sha256:8ed0c1e329fd857c696c44609df7be952fcb08c69e30c94e025fbcef712480ee", size = 114703522, upload-time = "2026-05-19T04:46:19.668Z" }, + { url = "https://files.pythonhosted.org/packages/41/e3/47546c2baf31e31d6039866948f6d4f2aae636e06f16c8cccd2ba0fc11fa/vtk-9.6.2-cp310-cp310-macosx_11_0_arm64.whl", hash = "sha256:809065272b207439f13ef0f62767f9041b4c6d61dd4dc2f60a1a62201843dc4b", size = 106906792, upload-time = "2026-05-19T04:46:24.878Z" }, + { url = "https://files.pythonhosted.org/packages/8a/55/c28d4070cca9923c419be191ca6d4b444bc030a992b38243bf6d114042d0/vtk-9.6.2-cp310-cp310-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:9a0df4ed93b6ae7f05cb6cee9d80566f4c06c6b4ce9116ee406445703c379ef7", size = 145981000, upload-time = "2026-05-19T04:46:30.482Z" }, + { url = "https://files.pythonhosted.org/packages/dd/2f/320ae500942ffbf8c19205e3f98797a575db8e8b0e7815abf17601423a53/vtk-9.6.2-cp310-cp310-manylinux_2_28_aarch64.whl", hash = "sha256:3cde8ba867cce14fdce8d8ec7bdef36a328cbd126433490109eccc586e489916", size = 135731415, upload-time = "2026-05-19T04:46:35.907Z" }, + { url = "https://files.pythonhosted.org/packages/c3/58/eb6c9788ec15b30a3e3e0885b92b5fa1f9dc6e619460020c707828a83ade/vtk-9.6.2-cp310-cp310-win_amd64.whl", hash = "sha256:9b55baa61beafc00d68b571ef07b71e2343a02200f089f07b6a40c3ebe01844a", size = 81289717, upload-time = "2026-05-19T04:46:41.149Z" }, + { url = "https://files.pythonhosted.org/packages/06/15/910b90b0b44d474f7cc71ccfe6e63393421fc0d161aabb490afe871830a3/vtk-9.6.2-cp311-cp311-macosx_10_10_x86_64.whl", hash = "sha256:ab2848c26c70fe57c41656d5ab48f47e8fa4f78ccbb113cf86c8c6de71c3118d", size = 114703453, upload-time = "2026-05-19T04:46:46.114Z" }, + { url = "https://files.pythonhosted.org/packages/5c/e5/8b4a37663aacd242c70c7a8feb2d2a4140ef5ded39ec0193af2d5a673098/vtk-9.6.2-cp311-cp311-macosx_11_0_arm64.whl", hash = "sha256:40fb9d9172cbd0b85a7f39df3646029449e563c61f54bedd3244427035f55ba3", size = 106906589, upload-time = "2026-05-19T04:46:51.081Z" }, + { url = "https://files.pythonhosted.org/packages/df/5c/148d54b90a2cd39809512d63d70b9b672f0e58f49d937f964a3167d63cf8/vtk-9.6.2-cp311-cp311-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:0fd9fa3f851192619ac0cec05591ab88adbed67ba063903297d2bb40b457bd00", size = 145980985, upload-time = "2026-05-19T04:46:56.229Z" }, + { url = "https://files.pythonhosted.org/packages/45/ee/9a4f42a8b98cfb095570ef203685186843389bb66d0da6d36581821888e8/vtk-9.6.2-cp311-cp311-manylinux_2_28_aarch64.whl", hash = "sha256:e640839fa24fc7c2153387d535527cfcd7d270e17e0d47f00fefd80b3b043577", size = 135731405, upload-time = "2026-05-19T04:47:01.547Z" }, + { url = "https://files.pythonhosted.org/packages/27/bb/e511d83d6b4d5b0acbce5e6a82110c510e3b416b39c667e6a52b1f78291a/vtk-9.6.2-cp311-cp311-win_amd64.whl", hash = "sha256:b935949cfc80f1d300d0b0ed8ccab47fb45c337910966de33a486d342e3c7daa", size = 81290711, upload-time = "2026-05-19T04:47:06.029Z" }, + { url = "https://files.pythonhosted.org/packages/60/d3/344f3664586f1f8bb5f7950db73c257439f5c796e0978e7a53e91d0fe2aa/vtk-9.6.2-cp312-cp312-macosx_10_10_x86_64.whl", hash = "sha256:2adcaed1cc4d3411a6b19834d6ed7d480f9ad9252e56546ea9930e66beae84e8", size = 114881278, upload-time = "2026-05-19T04:47:12.045Z" }, + { url = "https://files.pythonhosted.org/packages/cd/b1/754bd95da3d216e852014758c9ad626601482e2e0b1da98fa83adc5fe1b8/vtk-9.6.2-cp312-cp312-macosx_11_0_arm64.whl", hash = "sha256:d1eb5368039dd6a88e8102bee1db50e6f1e5a12945887ea1025069a99fb8088e", size = 106960169, upload-time = "2026-05-19T04:47:17.187Z" }, + { url = "https://files.pythonhosted.org/packages/96/f8/b392298c74aa7b88c731a43253ccd50b388bf42a1a29b50fa735c4f22f41/vtk-9.6.2-cp312-cp312-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:8c9f31430d15afbf46c2076cf3e30b4d6136512a5faabee8318552ccea907323", size = 146027355, upload-time = "2026-05-19T04:47:22.412Z" }, + { url = "https://files.pythonhosted.org/packages/71/35/67a7760852100c98a8fb6f1221e7eb359ad4d0b1b9edd8beea42c79e16c4/vtk-9.6.2-cp312-cp312-manylinux_2_28_aarch64.whl", hash = "sha256:91e1962c93217cf91ca4fe50762dd26dfc22290b74b8b0de77576f93f7c17abb", size = 135789817, upload-time = "2026-05-19T04:47:27.551Z" }, + { url = "https://files.pythonhosted.org/packages/b4/61/58e162d9cdbf02a3578ef4e68f4003b8cb9f351b7aa409960da3cc383017/vtk-9.6.2-cp312-cp312-win_amd64.whl", hash = "sha256:83b4af00b31395a13acb20e26a42ee097b85e41a0a087be2340d3749cb58250f", size = 81305902, upload-time = "2026-05-19T04:47:32.305Z" }, + { url = "https://files.pythonhosted.org/packages/89/b8/2eb42db93ac200180ca1a6ffe6f95d010181528e9490464621c31628e35f/vtk-9.6.2-cp313-cp313-macosx_10_10_x86_64.whl", hash = "sha256:155a09485a9efcbb0afb0214159bb920cf437fe4d897a6f945f6d66a533d863c", size = 114899285, upload-time = "2026-05-19T04:47:37.237Z" }, + { url = "https://files.pythonhosted.org/packages/00/51/5abfa4dc321b864e57bde87be8a328ee97b3aa33cc5e8f736c162fac63cb/vtk-9.6.2-cp313-cp313-macosx_11_0_arm64.whl", hash = "sha256:f7055447d36914ed8b39c5738defe467566ef2d7a75ca92a3023aefc269e1f06", size = 106962262, upload-time = "2026-05-19T04:47:41.886Z" }, + { url = "https://files.pythonhosted.org/packages/bd/75/4a1fe360256b99779d534b2387d0efa70952167d53d716f60ff39d62994d/vtk-9.6.2-cp313-cp313-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:fb85c7fbad59209a08e428479defbdf96f974a9f39d4212960fb1a24a919613c", size = 146027814, upload-time = "2026-05-19T04:47:47.238Z" }, + { url = "https://files.pythonhosted.org/packages/e8/30/90ea11e053d7e0571d70a7f301cebe04cc83874b533ddeaedb5258b0f9e6/vtk-9.6.2-cp313-cp313-manylinux_2_28_aarch64.whl", hash = "sha256:35e1e9ffb6457f16c37d0e025f7db8619961dffc3e9a85a8e5d358d5099c277c", size = 135792045, upload-time = "2026-05-19T04:47:53.848Z" }, + { url = "https://files.pythonhosted.org/packages/da/5b/e03640322971339899f982691396038d793ecfde8adbc208a97d501d6b8e/vtk-9.6.2-cp313-cp313-win_amd64.whl", hash = "sha256:4e9ad047a2b658550d5099dc808b26e331a9c94c3226c28453873035b8b48b41", size = 81307827, upload-time = "2026-05-19T04:47:58.76Z" }, + { url = "https://files.pythonhosted.org/packages/ed/a4/615178228ec84e3793332fbdd62f85965e246366727ef31b20670ce4877d/vtk-9.6.2-cp314-cp314-macosx_10_10_x86_64.whl", hash = "sha256:8a3edd56b63d1ab4ff022e70dac50dd54a66019462c37734e288c9e222809624", size = 114534790, upload-time = "2026-05-19T04:48:04.011Z" }, + { url = "https://files.pythonhosted.org/packages/8e/63/19f6f3f4520595e28b425fe7cae220fa8f4fb963ba9d0c0c4001f54ba8ae/vtk-9.6.2-cp314-cp314-macosx_11_0_arm64.whl", hash = "sha256:da01ec70cdfef3fddf659fd35a0258e820065d5331b3c43d5f7e3548351118ea", size = 106971615, upload-time = "2026-05-19T04:48:09.23Z" }, + { url = "https://files.pythonhosted.org/packages/bb/84/80465e452292a219e4b72bded477cee155f9ec44b856e3a877e0f95eba54/vtk-9.6.2-cp314-cp314-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:5f64288a71b313251300bc798458848aabbfd241230546ce75af90b8e099f14b", size = 146043416, upload-time = "2026-05-19T04:48:15.224Z" }, + { url = "https://files.pythonhosted.org/packages/c2/00/16e9faacb40cdcf487f8f5cb39d4304e71ebfafe38eb38a2c8f74527ca83/vtk-9.6.2-cp314-cp314-manylinux_2_28_aarch64.whl", hash = "sha256:1c3c050668eacb73db20fd07f844321a05dc3d7c60221db7137bd0b8f93a4d82", size = 135816286, upload-time = "2026-05-19T04:48:21.241Z" }, + { url = "https://files.pythonhosted.org/packages/72/6f/11594ed6bb95393f5656c96945ae22e297dfc3ccc5d7cd816973aaf9fc0d/vtk-9.6.2-cp314-cp314-win_amd64.whl", hash = "sha256:7b99aad09f712442345bb1c2a5529ad46da3a26fdc12cce54408ff800ccfbb19", size = 83246191, upload-time = "2026-05-19T04:48:26.1Z" }, + { url = "https://files.pythonhosted.org/packages/73/f2/a71b05850acbfb99e0de63c940ae874b665a393adb186a0e9c64cf6c9d64/vtk-9.6.2-cp314-cp314t-manylinux2014_x86_64.manylinux_2_17_x86_64.whl", hash = "sha256:691e3a9ae1b62784cf2e9b7d87696bdeac14e1b9905bf020f4c86fab5b756626", size = 145665062, upload-time = "2026-05-19T04:48:31.223Z" }, + { url = "https://files.pythonhosted.org/packages/4a/57/6410098435a3976cc749c151c4e09bff4bcf4af0ebeb13b3e47fbd5871e5/vtk-9.6.2-cp314-cp314t-manylinux_2_28_aarch64.whl", hash = "sha256:b06725993112097f43daefca0f3637a73a63c67b61a9c124455d881bc05cc9a4", size = 135630292, upload-time = "2026-05-19T04:48:37.053Z" }, +] + [[package]] name = "watchfiles" version = "1.1.1" From a7539ef6f8da911b518b697961fea518d043ab81 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Thu, 30 Jul 2026 15:22:41 -0700 Subject: [PATCH 17/25] fix: reconcile unstructured FV branch with post-merge main APIs - sigma() gained an sto parameter on main; pass sto=None in the unstructured DFN models, matching base_thermal and li_metal. - Restore main's strict 3D cell-geometry validation, which the branch had commented out; test_spm covers the arbitrary-geometry rejection. - Reduction._from_json clears domains again after restore, since SpecificFunction._from_json bypasses __init__ and Max/Min would otherwise come back carrying the child's domain. - Component serialises its index via to_json/_from_json, mirroring Magnitude, and gains a hypothesis strategy alongside Norm. - UnstructuredSubMesh joins _SUBMESH_EXEMPT (no _from_json), and Reduction joins _NOT_ROUND_TRIPPABLE as an abstract base. Co-Authored-By: Claude Fable 5 --- .../src/pybamm/expression_tree/functions.py | 11 +++++++++++ .../src/pybamm/expression_tree/unary_operators.py | 11 +++++++++++ .../full_battery_models/base_battery_model.py | 4 +--- .../lithium_ion/basic_dfn_2d_unstructured.py | 4 ++-- .../lithium_ion/basic_dfn_3d_unstructured.py | 4 ++-- packages/pybamm/tests/strategies/symbols.py | 15 +++++++++++++++ .../test_base_strategy_coverage.py | 1 + 7 files changed, 43 insertions(+), 7 deletions(-) diff --git a/packages/pybamm/src/pybamm/expression_tree/functions.py b/packages/pybamm/src/pybamm/expression_tree/functions.py index 6c994f9f17..794ca7af38 100644 --- a/packages/pybamm/src/pybamm/expression_tree/functions.py +++ b/packages/pybamm/src/pybamm/expression_tree/functions.py @@ -728,6 +728,17 @@ def __init__(self, function: Callable, child: pybamm.Symbol): super().__init__(function, child) self.clear_domains() + @classmethod + def _from_json(cls, snippet: dict): + """See :meth:`pybamm.SpecificFunction._from_json()`. + + ``SpecificFunction._from_json`` bypasses ``__init__``, so the domains + inherited from the child have to be cleared again here. + """ + instance = super()._from_json(snippet) + instance.clear_domains() + return instance + def _evaluate_for_shape(self): return np.nan * np.ones((1, 1)) diff --git a/packages/pybamm/src/pybamm/expression_tree/unary_operators.py b/packages/pybamm/src/pybamm/expression_tree/unary_operators.py index ad1e11a1b8..be3f2df585 100644 --- a/packages/pybamm/src/pybamm/expression_tree/unary_operators.py +++ b/packages/pybamm/src/pybamm/expression_tree/unary_operators.py @@ -1525,6 +1525,17 @@ def __init__(self, child, index): super().__init__(f"component({index})", child) self.index = index + def to_json(self): + return { + "name": self.name, + "domains": self.domains, + "index": self.index, + } + + @classmethod + def _from_json(cls, snippet): + return cls(snippet["children"][0], snippet["index"]) + def _unary_new_copy(self, child, perform_simplifications=True): return self.__class__(child, self.index) diff --git a/packages/pybamm/src/pybamm/models/full_battery_models/base_battery_model.py b/packages/pybamm/src/pybamm/models/full_battery_models/base_battery_model.py index cbdb0dad4c..15d8d347b0 100644 --- a/packages/pybamm/src/pybamm/models/full_battery_models/base_battery_model.py +++ b/packages/pybamm/src/pybamm/models/full_battery_models/base_battery_model.py @@ -795,11 +795,9 @@ def __init__(self, extra_options): if options["dimensionality"] == 3 and options["cell geometry"] not in [ "pouch", "cylindrical", - "arbitrary", ]: raise pybamm.OptionError( - "'cell geometry' must be 'pouch', 'cylindrical' or 'arbitrary' " - "if 'dimensionality' is '3'" + "'cell geometry' must be 'pouch' or 'cylindrical' if 'dimensionality' is '3'" ) if options["cell geometry"] == "cylindrical" and options["dimensionality"] != 3: diff --git a/packages/pybamm/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_2d_unstructured.py b/packages/pybamm/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_2d_unstructured.py index de843108d9..7d8af2378f 100644 --- a/packages/pybamm/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_2d_unstructured.py +++ b/packages/pybamm/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_2d_unstructured.py @@ -212,8 +212,8 @@ def __init__( ###################### # Current in the solid ###################### - sigma_eff_n = self.param.n.sigma(T) * eps_s_n**self.param.n.b_s - sigma_eff_p = self.param.p.sigma(T) * eps_s_p**self.param.p.b_s + sigma_eff_n = self.param.n.sigma(None, T) * eps_s_n**self.param.n.b_s + sigma_eff_p = self.param.p.sigma(None, T) * eps_s_p**self.param.p.b_s self.algebraic[phi_s_n] = ( self.param.L_x**2 * self.param.L_z**2 diff --git a/packages/pybamm/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_3d_unstructured.py b/packages/pybamm/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_3d_unstructured.py index 96a92f65c1..18f04a733e 100644 --- a/packages/pybamm/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_3d_unstructured.py +++ b/packages/pybamm/src/pybamm/models/full_battery_models/lithium_ion/basic_dfn_3d_unstructured.py @@ -222,8 +222,8 @@ def __init__( ###################### # Current in the solid ###################### - sigma_eff_n = self.param.n.sigma(T) * eps_s_n**self.param.n.b_s - sigma_eff_p = self.param.p.sigma(T) * eps_s_p**self.param.p.b_s + sigma_eff_n = self.param.n.sigma(None, T) * eps_s_n**self.param.n.b_s + sigma_eff_p = self.param.p.sigma(None, T) * eps_s_p**self.param.p.b_s L_scale = self.param.L_x**2 * self.param.L_z**2 self.algebraic[phi_s_n] = L_scale * ( pybamm.div(-sigma_eff_n * pybamm.grad(phi_s_n)) + a_j_n diff --git a/packages/pybamm/tests/strategies/symbols.py b/packages/pybamm/tests/strategies/symbols.py index 270519e7c1..6a4afa08cb 100644 --- a/packages/pybamm/tests/strategies/symbols.py +++ b/packages/pybamm/tests/strategies/symbols.py @@ -670,6 +670,17 @@ def _magnitude_branch( ) +def _component_branch( + _child_strategy: st.SearchStrategy[pybamm.Symbol], +) -> st.SearchStrategy[pybamm.Component]: + """Component(child, index) — domain-bearing child, zero-based component index.""" + return st.builds( + pybamm.Component, + _any_domain_leaves(), + st.integers(min_value=0, max_value=2), + ) + + def _discrete_time_sum_branch( _child_strategy: st.SearchStrategy[pybamm.Symbol], ) -> st.SearchStrategy[pybamm.DiscreteTimeSum]: @@ -981,6 +992,9 @@ def _vector_branch( pybamm.UpwindDownwind2D: _upwind_downwind_2d_branch, pybamm.NodeToEdge2D: _node_to_edge_2d_branch, pybamm.Magnitude: _magnitude_branch, + pybamm.Component: _component_branch, + # Norm: (self, child) only — round-trips via the generic unary hook. + pybamm.Norm: lambda _children: _any_domain_leaves().map(pybamm.Norm), pybamm.DiscreteTimeData: _discrete_time_data_branch, pybamm.DiscreteTimeSum: _discrete_time_sum_branch, pybamm.SizeAverage: _size_average_branch, @@ -1023,6 +1037,7 @@ def _vector_branch( pybamm.StateVectorBase, # abstract base; StateVector + StateVectorDot cover it pybamm.Function, # to_json() raises NotImplementedError — only SpecificFunction subclasses round-trip pybamm.SpecificFunction, # base for named funcs; direct instantiation not useful + pybamm.Reduction, # abstract base for scalar reductions; Max and Min cover it pybamm.Broadcast, # abstract base; PrimaryBroadcast/Secondary/Full covered pybamm.Integral, # base for domain-constrained integrals; heavyweight constructor pybamm.IndependentVariable, # abstract base; Time + SpatialVariable covered diff --git a/packages/pybamm/tests/unit/test_serialisation/test_base_strategy_coverage.py b/packages/pybamm/tests/unit/test_serialisation/test_base_strategy_coverage.py index 4a9e665a04..a03e88f314 100644 --- a/packages/pybamm/tests/unit/test_serialisation/test_base_strategy_coverage.py +++ b/packages/pybamm/tests/unit/test_serialisation/test_base_strategy_coverage.py @@ -44,6 +44,7 @@ def _covered_solver_classes() -> set[type]: pybamm.Uniform2DSubMesh, # cannot serialise (no _from_json) pybamm.UserSupplied1DSubMesh, # cannot serialise (no _from_json) pybamm.UserSupplied2DSubMesh, # cannot serialise (no _from_json) + pybamm.UnstructuredSubMesh, # cannot serialise (no _from_json) } From 4092964c077bcd8b03cf67fa9cf68905df28b42e Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Thu, 30 Jul 2026 15:35:31 -0700 Subject: [PATCH 18/25] docs: document the unstructured finite volume API Add API pages for UnstructuredSubMesh and its generators, the FiniteVolumeUnstructured spatial method, VTKQuickPlot, and the two unstructured DFN models, and record the feature in the changelog. Drop the root-level TODO.md: its main item (arbitrary domain names via an explicit tag) is implemented as TaggedSubMeshGenerator, and the workaround script it points at no longer exists. Co-Authored-By: Claude Fable 5 --- CHANGELOG.md | 7 ++++ TODO.md | 40 ------------------- docs/source/api/meshes/index.rst | 1 + .../api/meshes/unstructured_submeshes.rst | 16 ++++++++ docs/source/api/models/lithium_ion/dfn.rst | 6 +++ docs/source/api/plotting/index.rst | 1 + docs/source/api/plotting/plot_vtk.rst | 5 +++ .../finite_volume_unstructured.rst | 5 +++ docs/source/api/spatial_methods/index.rst | 1 + .../src/pybamm/meshes/unstructured_submesh.py | 9 +++-- .../pybamm/src/pybamm/plotting/plot_vtk.py | 12 +++--- 11 files changed, 53 insertions(+), 50 deletions(-) delete mode 100644 TODO.md create mode 100644 docs/source/api/meshes/unstructured_submeshes.rst create mode 100644 docs/source/api/plotting/plot_vtk.rst create mode 100644 docs/source/api/spatial_methods/finite_volume_unstructured.rst diff --git a/CHANGELOG.md b/CHANGELOG.md index 79dcab15e8..69ecc4d748 100644 --- a/CHANGELOG.md +++ b/CHANGELOG.md @@ -1,7 +1,14 @@ # [Unreleased](https://github.com/pybamm-team/PyBaMM/) +## Features + +- Added unstructured finite volume support: `pybamm.UnstructuredSubMesh` (cell-centred meshes of triangles, quadrilaterals, tetrahedra, or hexahedra), the `pybamm.FiniteVolumeUnstructured` spatial method, and the `pybamm.lithium_ion.BasicDFN2DUnstructured`/`BasicDFN3DUnstructured` models. Meshes can be read from gmsh files via `pybamm.UserSuppliedUnstructuredMesh` or `pybamm.TaggedSubMeshGenerator`, and interfaces between adjacent submeshes are discovered automatically for arbitrary topologies rather than assuming a 1D stack. ([#5397](https://github.com/pybamm-team/PyBaMM/pull/5397)) +- Added `pybamm.VTKQuickPlot`, a VTK-based interactive alternative to `QuickPlot` for 2D and 3D unstructured mesh solutions. Requires the new `vtk` extra (`pip install pybamm[vtk]`). ([#5397](https://github.com/pybamm-team/PyBaMM/pull/5397)) + ## Bug fixes +- `pybamm.max` and `pybamm.min` now clear their child's domains, reflecting that a reduction over a spatial field is a scalar. ([#5397](https://github.com/pybamm-team/PyBaMM/pull/5397)) + - `BatchStudy.solve` no longer ignores its `solver` argument: previously the loop over study inputs shadowed it, so a caller-supplied solver was silently dropped. A solver from `BatchStudy(solvers=...)` still takes precedence. ([#5677](https://github.com/pybamm-team/PyBaMM/pull/5677)) - `pybamm.citations.register` now names the citation the caller passed in when a BibTeX string fails to parse, instead of whichever entry the parser had reached. ([#5677](https://github.com/pybamm-team/PyBaMM/pull/5677)) - Deserialising a parameter set whose interpolant specification is invalid now logs a warning naming the offending parameter, instead of printing the bare exception to stdout with no indication of which parameter fell back to zero. ([#5679](https://github.com/pybamm-team/PyBaMM/pull/5679)) diff --git a/TODO.md b/TODO.md deleted file mode 100644 index f3fde8323d..0000000000 --- a/TODO.md +++ /dev/null @@ -1,40 +0,0 @@ -# TODO - -## Mesh / spatial methods - -### Extend `UserSuppliedUnstructuredMesh` for arbitrary domain names -**Why:** current `_domain_name_from_lims` (`unstructured_submesh.py:671-688`) only -recognises hardcoded prefixes (`x_n`, `x_s`, `x_p`) and maps to fixed battery -domains. For user-defined domain names (`body`, `tab_0`, ...), it returns `None` -and falls through to "use all cells" — every region's submesh ends up holding -the entire mesh. Users currently work around this with a custom -`MeshGenerator` subclass (see -`scripts/mesh/run_thermal_3d_multi_domain.py::TaggedSubMeshGenerator`). - -**Change:** add an explicit `tag_id` (or `domain_tag`) kwarg that bypasses the -name-prefix heuristic. - -```python -gen = pybamm.UserSuppliedUnstructuredMesh( - filepath="cell.msh", - tag_id=1, # filter cells by gmsh:physical == 1 - coord_sys="cartesian", -) -``` - -Implementation sketch (`__call__`): -```python -if self.tag_id is not None: - cell_mask = self._get_cell_mask(mesh, cell_type, self.tag_id) - elements = cells[cell_mask] -elif domain_name and domain_name in self.subdomain_mapping: - # existing path - ... -``` - -**Bonus:** module-level LRU cache keyed by `filepath` so multiple instances -reading the same `.msh` don't re-parse it (currently each instance has its own -`_cached_mesh` field, so once you have N region generators you do N reads). - -**Effort:** ~15 lines + a unit test that loads a multi-tag mesh into separate -domains and checks each submesh has only its tag's cells. diff --git a/docs/source/api/meshes/index.rst b/docs/source/api/meshes/index.rst index 143e93638f..9d03fa3f6d 100644 --- a/docs/source/api/meshes/index.rst +++ b/docs/source/api/meshes/index.rst @@ -8,3 +8,4 @@ Meshes one_dimensional_submeshes two_dimensional_submeshes three_dimensional_submeshes + unstructured_submeshes diff --git a/docs/source/api/meshes/unstructured_submeshes.rst b/docs/source/api/meshes/unstructured_submeshes.rst new file mode 100644 index 0000000000..be6d03125c --- /dev/null +++ b/docs/source/api/meshes/unstructured_submeshes.rst @@ -0,0 +1,16 @@ +Unstructured Sub Meshes +======================= + +.. autoclass:: pybamm.UnstructuredSubMesh + :members: + +.. autoclass:: pybamm.UnstructuredMeshGenerator + :members: + +.. autoclass:: pybamm.UserSuppliedUnstructuredMesh + :members: + +.. autoclass:: pybamm.TaggedSubMeshGenerator + :members: + +.. autofunction:: pybamm.compute_interface_data diff --git a/docs/source/api/models/lithium_ion/dfn.rst b/docs/source/api/models/lithium_ion/dfn.rst index 4213f4ccf9..6231e88c55 100644 --- a/docs/source/api/models/lithium_ion/dfn.rst +++ b/docs/source/api/models/lithium_ion/dfn.rst @@ -13,4 +13,10 @@ Doyle-Fuller-Newman (DFN) .. autoclass:: pybamm.lithium_ion.BasicDFNHalfCell :members: +.. autoclass:: pybamm.lithium_ion.BasicDFN2DUnstructured + :members: + +.. autoclass:: pybamm.lithium_ion.BasicDFN3DUnstructured + :members: + .. footbibliography:: diff --git a/docs/source/api/plotting/index.rst b/docs/source/api/plotting/index.rst index 796df15fd7..5416399b0c 100644 --- a/docs/source/api/plotting/index.rst +++ b/docs/source/api/plotting/index.rst @@ -10,3 +10,4 @@ Plotting plot_summary_variables plot_3d_cross_section plot_3d_heatmap + plot_vtk diff --git a/docs/source/api/plotting/plot_vtk.rst b/docs/source/api/plotting/plot_vtk.rst new file mode 100644 index 0000000000..eaf6018f72 --- /dev/null +++ b/docs/source/api/plotting/plot_vtk.rst @@ -0,0 +1,5 @@ +VTK Quick Plot +============== + +.. autoclass:: pybamm.VTKQuickPlot + :members: diff --git a/docs/source/api/spatial_methods/finite_volume_unstructured.rst b/docs/source/api/spatial_methods/finite_volume_unstructured.rst new file mode 100644 index 0000000000..eeb01b8a43 --- /dev/null +++ b/docs/source/api/spatial_methods/finite_volume_unstructured.rst @@ -0,0 +1,5 @@ +Unstructured Finite Volume +========================== + +.. autoclass:: pybamm.FiniteVolumeUnstructured + :members: diff --git a/docs/source/api/spatial_methods/index.rst b/docs/source/api/spatial_methods/index.rst index f9ccacd0d4..207c2a485d 100644 --- a/docs/source/api/spatial_methods/index.rst +++ b/docs/source/api/spatial_methods/index.rst @@ -10,3 +10,4 @@ Discretisation and spatial methods scikit_finite_element zero_dimensional_method scikit_finite_element_3d + finite_volume_unstructured diff --git a/packages/pybamm/src/pybamm/meshes/unstructured_submesh.py b/packages/pybamm/src/pybamm/meshes/unstructured_submesh.py index 1a649a87be..ad12640264 100644 --- a/packages/pybamm/src/pybamm/meshes/unstructured_submesh.py +++ b/packages/pybamm/src/pybamm/meshes/unstructured_submesh.py @@ -19,10 +19,11 @@ class UnstructuredSubMesh(SubMesh): Parameters ---------- - nodes : numpy.ndarray, shape (n_nodes, d) - Vertex coordinates (d = 2 or 3). - elements : numpy.ndarray, shape (n_cells, n_verts_per_cell) - Element vertex indices. For 2D: 3 (triangles) or 4 (quads). + nodes : numpy.ndarray + Vertex coordinates, of shape ``(n_nodes, d)`` (d = 2 or 3). + elements : numpy.ndarray + Element vertex indices, of shape ``(n_cells, n_verts_per_cell)``. + For 2D: 3 (triangles) or 4 (quads). For 3D: 4 (tetrahedra) or 8 (hexahedra). coord_sys : str, optional Coordinate system, default ``"cartesian"``. diff --git a/packages/pybamm/src/pybamm/plotting/plot_vtk.py b/packages/pybamm/src/pybamm/plotting/plot_vtk.py index 132555a9f0..664df505ba 100644 --- a/packages/pybamm/src/pybamm/plotting/plot_vtk.py +++ b/packages/pybamm/src/pybamm/plotting/plot_vtk.py @@ -180,13 +180,13 @@ class VTKQuickPlot: ``plot_type`` is ``"slice"``) - ``"scale"``: ``"auto"`` (default), ``None``, or ``(sx, sy, sz)`` - Each variable's options value may also be a **list** of dicts, in which - case one panel is created per entry:: + A variable's value may also be a **list** of such dicts, in which + case one panel is created per entry:: - options={"T": [ - {"plot_type": "3d"}, - {"plot_type": "slice", "x": 0.5}, - ]} + options={"T": [ + {"plot_type": "3d"}, + {"plot_type": "slice", "x": 0.5}, + ]} """ def __init__( From 78f6f684353bf9bf6cfee5d931bb2b3a08e00adb Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Thu, 30 Jul 2026 16:23:08 -0700 Subject: [PATCH 19/25] fix: preserve symbols during broadcast Co-authored-by: Cursor --- .../src/pybamm/spatial_methods/finite_volume_unstructured.py | 4 ++++ 1 file changed, 4 insertions(+) diff --git a/packages/pybamm/src/pybamm/spatial_methods/finite_volume_unstructured.py b/packages/pybamm/src/pybamm/spatial_methods/finite_volume_unstructured.py index f93be00412..6302b94d00 100644 --- a/packages/pybamm/src/pybamm/spatial_methods/finite_volume_unstructured.py +++ b/packages/pybamm/src/pybamm/spatial_methods/finite_volume_unstructured.py @@ -328,6 +328,10 @@ def broadcast(self, symbol, domains, broadcast_type): matrix = vstack([identity for _ in range(sec_size)]) out = pybamm.Matrix(matrix) @ symbol + if out is symbol: + # simplification can hand back the child itself (e.g. ones-vector + # multiply); copy before stamping domains on a possibly shared node + out = symbol.create_copy(perform_simplifications=False) out.domains = domains.copy() return out From 790123a0b68664c3060a3e46df54a5df3c09dfbe Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Fri, 31 Jul 2026 10:48:02 -0700 Subject: [PATCH 20/25] test: cover unstructured finite volume behavior Co-authored-by: Cursor --- .../test_finite_volume_unstructured.py | 567 ++++++++++++++++++ 1 file changed, 567 insertions(+) diff --git a/packages/pybamm/tests/unit/test_spatial_methods/test_finite_volume_unstructured.py b/packages/pybamm/tests/unit/test_spatial_methods/test_finite_volume_unstructured.py index 2d745b8a8c..7aa12f2001 100644 --- a/packages/pybamm/tests/unit/test_spatial_methods/test_finite_volume_unstructured.py +++ b/packages/pybamm/tests/unit/test_spatial_methods/test_finite_volume_unstructured.py @@ -14,6 +14,7 @@ from scipy.sparse import coo_matrix as sp_coo from scipy.sparse import csr_matrix as sp_csr +import pybamm from pybamm.meshes.unstructured_submesh import ( UnstructuredSubMesh, _hex_to_tet, @@ -61,6 +62,25 @@ def _get_internal_cells(mesh): return [i for i in range(mesh.npts) if i not in bnd_cells] +class _MeshMap(dict): + """Minimal Mesh-like mapping that accepts PyBaMM domain lists.""" + + def __getitem__(self, key): + if isinstance(key, list): + key = tuple(key) + elif isinstance(key, str): + key = (key,) + return super().__getitem__(key) + + +def _method_with_mesh(mesh, **auxiliary_meshes): + meshes = {("test",): mesh} + meshes.update({(name,): value for name, value in auxiliary_meshes.items()}) + method = FiniteVolumeUnstructured() + method._mesh = _MeshMap(meshes) + return method + + # ====================================================================== # Tests: TPFA Laplacian # ====================================================================== @@ -668,3 +688,550 @@ def test_constructor_default_options(self): fvu = FiniteVolumeUnstructured() assert fvu.options is not None assert "extrapolation" in fvu.options + + +class TestFiniteVolumeUnstructuredBehavior: + def test_build_discovers_interfaces_and_ignores_other_meshes(self): + left = _make_2d_mesh(2, 2, x_range=(0, 0.5)) + right = _make_2d_mesh(2, 2, x_range=(0.5, 1)) + structured = pybamm.SubMesh1D(np.array([0, 1]), "cartesian") + meshes = _MeshMap( + {("left",): left, ("right",): right, ("structured",): structured} + ) + + method = FiniteVolumeUnstructured() + method.build(meshes) + + assert right in [data["other_mesh"] for data in left.interface_data.values()] + assert left.npts_for_broadcast_to_nodes == left.npts + assert structured.npts_for_broadcast_to_nodes == structured.npts + + def test_interface_matching_edge_cases(self): + empty = _make_2d_mesh(1, 1) + empty.boundary_faces = {} + other = _make_2d_mesh(1, 1) + a_idx, b_idx, matched = FiniteVolumeUnstructured._interface_face_match( + empty, other + ) + assert not matched + assert a_idx.size == b_idx.size == 0 + + mesh_3d = _make_3d_mesh(1, 1, 1) + assert not FiniteVolumeUnstructured._interface_face_match(other, mesh_3d)[2] + + distant = _make_2d_mesh(1, 1, x_range=(2, 3)) + assert not FiniteVolumeUnstructured._interface_face_match(other, distant)[2] + + def test_compute_pair_interface_success_and_noops(self): + left = _make_2d_mesh(2, 2, x_range=(0, 0.5)) + right = _make_2d_mesh(2, 2, x_range=(0.5, 1)) + method = FiniteVolumeUnstructured() + + assert method._compute_pair_interface(left, right, "left", "right") + assert "iface_right" in left.boundary_faces + assert "iface_left" in right.boundary_faces + assert method._compute_pair_interface(left, right, "left", "right") is False + + far = _make_2d_mesh(1, 1, x_range=(2, 3)) + assert method._compute_pair_interface(left, far, "left", "far") is False + + shared = _make_2d_mesh(1, 1) + method._auto_compute_all_interfaces( + _MeshMap({("first",): shared, ("alias",): shared}) + ) + + def test_spatial_variable_directions_and_auxiliary_repeats(self): + mesh = _make_3d_mesh(1, 1, 1) + aux = _make_2d_mesh(1, 1) + method = _method_with_mesh(mesh, aux=aux) + domains = {"primary": ["test"], "secondary": ["aux"]} + + for name, direction, column in [ + ("x", None, 0), + ("y", None, 1), + ("z", None, 2), + ("r", None, 0), + ("s", "lr", 0), + ("s", "tb", 2), + ("s", "fb", 1), + ("s", "unknown", 0), + ]: + symbol = pybamm.SpatialVariable(name, domains=domains, direction=direction) + actual = method.spatial_variable(symbol).evaluate().reshape(-1) + expected = np.tile(mesh.cell_centroids[:, column], aux.npts) + np.testing.assert_allclose(actual, expected) + + def test_broadcast_variants(self): + mesh = _make_2d_mesh(1, 1) + aux = _make_2d_mesh(1, 1) + method = _method_with_mesh(mesh, aux=aux) + primary = {"primary": ["test"], "secondary": []} + + scalar_primary = method.broadcast(pybamm.Scalar(2), primary, "primary to nodes") + np.testing.assert_array_equal( + scalar_primary.evaluate()[:, 0], np.full(mesh.npts, 2) + ) + + vector_primary = method.broadcast( + pybamm.Vector([2, 3]), primary, "primary to nodes" + ) + np.testing.assert_array_equal( + vector_primary.evaluate()[:, 0], np.repeat([2, 3], mesh.npts) + ) + + full_domains = {"primary": ["test"], "secondary": ["aux"]} + full = method.broadcast(pybamm.Scalar(4), full_domains, "full to nodes") + np.testing.assert_array_equal( + full.evaluate()[:, 0], np.full(mesh.npts * aux.npts, 4) + ) + + secondary_child = pybamm.Vector([1, 2], domain="test") + secondary = method.broadcast(secondary_child, primary, "secondary to nodes") + np.testing.assert_array_equal(secondary.evaluate(), secondary_child.evaluate()) + assert secondary.domain == primary["primary"] + assert secondary.domains["secondary"] == primary["secondary"] + + def test_broadcast_does_not_mutate_simplified_child(self): + mesh = pybamm.SubMesh1D(np.array([0, 1]), "cartesian") + method = _method_with_mesh(mesh) + child = pybamm.StateVector(slice(0, 1)) + domains = {"primary": ["test"], "secondary": []} + + result = method.broadcast(child, domains, "full to nodes") + + assert result is not child + assert child.domain == [] + assert result.domains["primary"] == ["test"] + np.testing.assert_array_equal(result.evaluate(y=np.array([7])), [[7]]) + + def test_laplacian_and_boundary_conditions(self): + mesh = _make_2d_mesh(2, 2) + method = _method_with_mesh(mesh) + variable = pybamm.Variable("u", domain="test") + values = pybamm.Vector(np.arange(mesh.npts), domain="test") + + plain = method.laplacian(variable, values, {}) + np.testing.assert_allclose( + plain.evaluate()[:, 0], method._tpfa_matrix(mesh) @ np.arange(mesh.npts) + ) + + constant = pybamm.Vector(np.full(mesh.npts, 3), domain="test") + dirichlet_bcs = { + variable: { + side: (pybamm.Scalar(3), "Dirichlet") + for side in ["left", "right", "top", "bottom"] + } + } + np.testing.assert_allclose( + method.laplacian(variable, constant, dirichlet_bcs).evaluate(), + 0, + atol=1e-12, + ) + + neumann_bcs = { + variable: { + side: (pybamm.Scalar(0), "Neumann") + for side in ["left", "right", "top", "bottom"] + } + | { + "missing": (pybamm.Scalar(3), "Dirichlet"), + } + } + np.testing.assert_allclose( + method.laplacian(variable, constant, neumann_bcs).evaluate(), 0, atol=1e-12 + ) + + face_count = len(mesh.boundary_faces["top"]) + vector_bc = pybamm.Vector(np.arange(face_count) + 1) + _, rhs = method._apply_bcs_to_laplacian( + mesh, + method._tpfa_matrix(mesh), + pybamm.Vector(np.zeros(mesh.npts)), + {"top": (vector_bc, "Dirichlet")}, + ) + faces = mesh.boundary_faces["top"] + owners = mesh.face_owner[faces] + distance = np.linalg.norm( + mesh.face_centroids[faces] - mesh.cell_centroids[owners], axis=1 + ) + coefficients = mesh.face_areas[faces] / distance / mesh.cell_volumes[owners] + expected_rhs = np.zeros(mesh.npts) + np.add.at(expected_rhs, owners, coefficients * (np.arange(face_count) + 1)) + np.testing.assert_allclose(rhs.evaluate()[:, 0], expected_rhs) + + def test_gradient_and_gradient_squared(self): + mesh = _make_2d_mesh(2, 2) + method = _method_with_mesh(mesh) + variable = pybamm.Variable("u", domain="test") + constant = pybamm.Vector(np.full(mesh.npts, 3), domain="test") + dirichlet_bcs = { + variable: { + side: (pybamm.Scalar(3), "Dirichlet") + for side in ["left", "right", "top", "bottom"] + } + } + gradient = method.gradient(variable, constant, dirichlet_bcs) + assert gradient._disc_state_vector is constant + for component in gradient._components: + np.testing.assert_allclose(component.evaluate(), 0, atol=1e-12) + + neumann_bcs = { + variable: { + side: (pybamm.Scalar(0), "Neumann") + for side in ["left", "right", "top", "bottom"] + } + | { + "missing": (pybamm.Scalar(2), "Neumann"), + } + } + for component in method.gradient(variable, constant, neumann_bcs)._components: + np.testing.assert_allclose(component.evaluate(), 0, atol=1e-12) + + x_values = mesh.cell_centroids[:, 0] + values = pybamm.Vector(x_values, domain="test") + grad_squared = method.gradient_squared(variable, values, {}) + matrices = method._green_gauss_matrices(mesh) + expected = sum((matrix @ x_values) ** 2 for matrix in matrices) + np.testing.assert_allclose(grad_squared.evaluate()[:, 0], expected) + + def test_divergence_input_forms_and_error(self): + mesh = _make_2d_mesh(2, 2) + method = _method_with_mesh(mesh) + symbol = pybamm.Variable("F", domain="test") + components = [ + pybamm.Vector(np.ones(mesh.npts), domain="test"), + pybamm.Vector(np.full(mesh.npts, 2), domain="test"), + ] + + from_list = method.divergence(symbol, components, {}) + from_field = method.divergence(symbol, pybamm.VectorField(*components), {}) + np.testing.assert_allclose(from_list.evaluate(), from_field.evaluate()) + matrices = method._divergence_matrices(mesh) + expected = matrices[0] @ np.ones(mesh.npts) + expected += matrices[1] @ np.full(mesh.npts, 2) + np.testing.assert_allclose(from_list.evaluate()[:, 0], expected) + + with pytest.raises(TypeError, match="expects a VectorField"): + method.divergence(symbol, pybamm.Scalar(1), {}) + + def test_divergence_boundary_correction(self): + mesh = _make_2d_mesh(2, 2) + method = _method_with_mesh(mesh) + variable = pybamm.Variable("u", domain="test") + other = pybamm.Variable("v", domain="other") + bcs = { + "not a symbol": {"left": (pybamm.Scalar(0), "Dirichlet")}, + other: {"left": (pybamm.Scalar(0), "Dirichlet")}, + variable: { + "left": (pybamm.Scalar(1), "Dirichlet"), + "right": (pybamm.Scalar(2), "Neumann"), + "missing": (pybamm.Scalar(3), "Neumann"), + }, + } + + L_bc, rhs, boundary_matrices = method._div_boundary_correction( + mesh, bcs, domain=["test"] + ) + assert L_bc.shape == (mesh.npts, mesh.npts) + assert rhs.evaluate().shape == (mesh.npts, 1) + assert len(boundary_matrices) == mesh.dimension + left_faces = mesh.boundary_faces["left"] + left_owners = mesh.face_owner[left_faces] + expected_rhs = np.zeros(mesh.npts) + left_distance = np.linalg.norm( + mesh.face_centroids[left_faces] - mesh.cell_centroids[left_owners], axis=1 + ) + np.add.at( + expected_rhs, + left_owners, + mesh.face_areas[left_faces] + / left_distance + / mesh.cell_volumes[left_owners], + ) + right_faces = mesh.boundary_faces["right"] + right_owners = mesh.face_owner[right_faces] + np.add.at( + expected_rhs, + right_owners, + 2 * mesh.face_areas[right_faces] / mesh.cell_volumes[right_owners], + ) + np.testing.assert_allclose(rhs.evaluate()[:, 0], expected_rhs) + + none_L, zero_rhs, none_D = method._div_boundary_correction(mesh, {}) + assert none_L is None + assert none_D is None + np.testing.assert_allclose(zero_rhs.evaluate(), 0) + + def test_div_D_grad_scalar_and_vector_coefficients(self): + mesh = _make_2d_mesh(2, 2) + aux = _make_2d_mesh(1, 1) + method = _method_with_mesh(mesh, aux=aux) + variable = pybamm.Variable("u", domain="test") + div_symbol = pybamm.Variable("div", domain="test") + cell_values = mesh.cell_centroids[:, 0] ** 2 + values = pybamm.Vector(cell_values, domain="test") + bcs = { + variable: { + "left": (pybamm.Scalar(0), "Dirichlet"), + "right": (pybamm.Scalar(2), "Neumann"), + "top": (pybamm.Scalar(0), "Neumann"), + "missing": (pybamm.Scalar(1), "Dirichlet"), + } + } + scalar_result = method.div_D_grad( + div_symbol, variable, pybamm.Scalar(2), values, bcs + ) + + coefficient = pybamm.Vector(np.full(mesh.npts, 2), domain="test") + vector_result = method.div_D_grad( + div_symbol, variable, coefficient, values, bcs + ) + np.testing.assert_allclose( + vector_result.evaluate(), scalar_result.evaluate(), atol=1e-12 + ) + + repeated_domains = {"primary": ["test"], "secondary": ["aux"]} + repeated_div = pybamm.Variable("repeated div", domains=repeated_domains) + repeated_u = pybamm.Variable("repeated u", domains=repeated_domains) + size = mesh.npts * aux.npts + repeated_values = pybamm.Vector( + np.tile(cell_values, aux.npts), domains=repeated_domains + ) + repeated_coefficient = pybamm.Vector(np.full(size, 2), domains=repeated_domains) + repeated = method.div_D_grad( + repeated_div, + repeated_u, + repeated_coefficient, + repeated_values, + { + repeated_u: { + "left": (pybamm.Scalar(0), "Dirichlet"), + "right": (pybamm.Scalar(2), "Neumann"), + "top": (pybamm.Scalar(0), "Neumann"), + "missing": (pybamm.Scalar(1), "Dirichlet"), + } + }, + ) + np.testing.assert_allclose( + repeated.evaluate()[:, 0], + np.tile(vector_result.evaluate()[:, 0], aux.npts), + atol=1e-12, + ) + + def test_integral_and_boundary_integral(self): + mesh = _make_2d_mesh(2, 2) + aux = _make_2d_mesh(1, 1) + method = _method_with_mesh(mesh, aux=aux) + domains = {"primary": ["test"], "secondary": ["aux"]} + child = pybamm.Variable("u", domains=domains) + values = pybamm.Vector(np.ones(mesh.npts * aux.npts), domains=domains) + + integral = method.integral(child, values, "primary") + np.testing.assert_allclose(integral.evaluate(), 1) + + row = method.definite_integral_matrix(child) + np.testing.assert_allclose(row.toarray()[0], mesh.cell_volumes) + + boundary = method.boundary_integral(child, values, "left") + np.testing.assert_allclose(boundary.evaluate(), 1) + missing = method.boundary_integral(child, values, "missing") + assert missing == pybamm.Scalar(0) + + @pytest.mark.parametrize( + "side", + ["left", "missing", "top-right", "top-left", "bottom-right", "bottom-left"], + ) + def test_boundary_value_and_corners(self, side): + mesh = _make_2d_mesh(2, 2) + method = _method_with_mesh(mesh) + variable = pybamm.Variable("u", domain="test") + values = pybamm.Vector(np.arange(mesh.npts), domain="test") + symbol = pybamm.BoundaryValue(variable, side) + + result = method.boundary_value_or_flux(symbol, values) + assert result.domain == [] + if side == "missing": + assert result == pybamm.Scalar(0) + elif "-" in side: + top_bottom, left_right = side.split("-") + x = mesh.cell_centroids[:, 0] + z = mesh.cell_centroids[:, -1] + target_x = x.max() if left_right == "right" else x.min() + target_z = z.max() if top_bottom == "top" else z.min() + expected = np.argmin((x - target_x) ** 2 + (z - target_z) ** 2) + assert result.evaluate().item() == expected + else: + owners = mesh.face_owner[mesh.boundary_faces[side]] + np.testing.assert_array_equal(result.evaluate()[:, 0], owners) + + def test_process_binary_operators(self): + method = FiniteVolumeUnstructured() + left_components = [pybamm.StateVector(slice(0, 2)), pybamm.Vector([2, 3])] + right_components = [pybamm.Vector([4, 5]), pybamm.Vector([6, 7])] + left_field = pybamm.VectorField(*left_components) + left_field._disc_state_vector = left_components[0] + right_field = pybamm.VectorField(*right_components) + multiplication = pybamm.Multiplication(pybamm.Scalar(1), pybamm.Scalar(2)) + + both = method.process_binary_operators( + multiplication, + None, + None, + left_field, + right_field, + ) + assert both.n_components == 2 + assert both._disc_state_vector is left_components[0] + np.testing.assert_array_equal( + both._components[0].evaluate(y=np.array([1, 2]))[:, 0], [4, 10] + ) + np.testing.assert_array_equal(both._components[1].evaluate()[:, 0], [12, 21]) + + field_left = method.process_binary_operators( + multiplication, None, None, left_field, pybamm.Scalar(2) + ) + field_right = method.process_binary_operators( + multiplication, None, None, pybamm.Scalar(2), right_field + ) + np.testing.assert_array_equal( + field_left._components[0].evaluate(y=np.array([1, 2]))[:, 0], [2, 4] + ) + np.testing.assert_array_equal( + field_right._components[0].evaluate()[:, 0], [8, 10] + ) + + scalar = method.process_binary_operators( + multiplication, None, None, pybamm.Scalar(3), pybamm.Scalar(4) + ) + assert scalar.evaluate() == 12 + + def test_internal_neumann_unstructured_paths(self): + left, right = _make_split_2d_meshes(2, 2, 2) + method = FiniteVolumeUnstructured() + left_values = pybamm.Vector(np.arange(left.npts), domain="left") + right_values = pybamm.Vector(np.arange(right.npts), domain="right") + + direct = method._internal_neumann_unstructured( + left_values, right_values, left, right, 1 + ) + interface = next(iter(left.interface_data.values())) + expected = ( + np.arange(right.npts)[interface["right_cells"]] + - np.arange(left.npts)[interface["left_cells"]] + ) / interface["cell_distances"] + np.testing.assert_allclose(direct.evaluate()[:, 0], expected) + + left_data = left.interface_data + left.interface_data = {} + reverse = method._internal_neumann_unstructured( + left_values, right_values, left, right, 1 + ) + np.testing.assert_allclose(reverse.evaluate(), direct.evaluate()) + + right.interface_data = {} + absent = method._internal_neumann_unstructured( + left_values, right_values, left, right, 2 + ) + np.testing.assert_allclose(absent.evaluate(), 0) + assert absent.shape[0] == left.npts * 2 + left.interface_data = left_data + + def test_internal_neumann_dispatch_structured_and_mismatch(self): + method = FiniteVolumeUnstructured() + left_mesh = pybamm.SubMesh1D(np.array([0, 0.5]), "cartesian") + right_mesh = pybamm.SubMesh1D(np.array([0.5, 1]), "cartesian") + left = pybamm.Vector(np.arange(left_mesh.npts), domain="left") + right = pybamm.Vector(np.arange(right_mesh.npts), domain="right") + + structured = method.internal_neumann_condition( + left, right, left_mesh, right_mesh + ) + dx = right_mesh.nodes[0] - left_mesh.nodes[-1] + expected = (np.arange(right_mesh.npts)[0] - np.arange(left_mesh.npts)[-1]) / dx + assert structured.evaluate().item() == expected + + unstructured_left, unstructured_right = _make_split_2d_meshes(1, 1, 1) + method._mesh = _MeshMap( + { + ("aux",): _make_2d_mesh(1, 1), + ("other aux",): _make_2d_mesh(2, 1), + } + ) + left_repeated = pybamm.Vector( + np.ones(unstructured_left.npts * method.mesh["aux"].npts), + domains={"primary": ["left"], "secondary": ["aux"]}, + ) + right_repeated = pybamm.Vector( + np.ones(unstructured_right.npts * method.mesh["other aux"].npts), + domains={"primary": ["right"], "secondary": ["other aux"]}, + ) + with pytest.raises(pybamm.DomainError, match="secondary points"): + method.internal_neumann_condition( + left_repeated, + right_repeated, + unstructured_left, + unstructured_right, + ) + + def test_internal_bcs_for_concatenation(self): + left = _make_2d_mesh(1, 1, x_range=(0, 0.5)) + right = _make_2d_mesh(1, 1, x_range=(0.5, 1)) + method = FiniteVolumeUnstructured() + method._compute_pair_interface(left, right, "left", "right") + method._mesh = _MeshMap({("left",): left, ("right",): right}) + children = [ + pybamm.Variable("left temperature", domain="left"), + pybamm.Variable("right temperature", domain="right"), + ] + + class Disc: + def process_symbol(self, child): + size = method.mesh[child.domain].npts + return pybamm.Vector(np.ones(size), domains=child.domains) + + result = method.set_internal_bcs_for_concat( + Disc(), + children[0], + children, + {"left": (pybamm.Scalar(0), "Dirichlet")}, + ) + assert set(result) == set(children) + assert "iface_right" in result[children[0]] + interface_gradient, bc_type = result[children[0]]["iface_right"] + assert bc_type == "Neumann" + np.testing.assert_allclose(interface_gradient.evaluate(), 0) + + structured = pybamm.SubMesh1D(np.array([0, 1]), "cartesian") + method._mesh[("structured",)] = structured + structured_child = pybamm.Variable( + "structured temperature", domain="structured" + ) + partial = method.set_internal_bcs_for_concat( + Disc(), + children[0], + [children[0], structured_child], + {}, + ) + assert structured_child not in partial + assert partial[children[0]] == {} + + no_interface = _method_with_mesh(_make_2d_mesh(1, 1)) + assert ( + no_interface.set_internal_bcs_for_concat( + Disc(), children[0], [pybamm.Variable("u", domain="test")], {} + ) + is None + ) + + def test_concatenation_preserves_domain_order(self): + left = _make_2d_mesh(1, 1, x_range=(0, 0.5)) + right = _make_2d_mesh(1, 1, x_range=(0.5, 1)) + method = FiniteVolumeUnstructured() + method._mesh = _MeshMap({("left",): left, ("right",): right}) + left_values = pybamm.Vector([1, 2], domain="left") + right_values = pybamm.Vector([3, 4], domain="right") + + result = method.concatenation([left_values, right_values]) + + np.testing.assert_array_equal(result.evaluate()[:, 0], [1, 2, 3, 4]) + assert result.domain == ["left", "right"] From 175c697378fc76ec800683e888168fba3b5f9cfb Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Fri, 31 Jul 2026 11:22:40 -0700 Subject: [PATCH 21/25] test: cover VTK plotting behavior Co-authored-by: Cursor --- .../pybamm/src/pybamm/plotting/plot_vtk.py | 7 +- .../tests/unit/test_plotting/test_plot_vtk.py | 423 ++++++++++++++++++ 2 files changed, 424 insertions(+), 6 deletions(-) create mode 100644 packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py diff --git a/packages/pybamm/src/pybamm/plotting/plot_vtk.py b/packages/pybamm/src/pybamm/plotting/plot_vtk.py index 664df505ba..3868d00791 100644 --- a/packages/pybamm/src/pybamm/plotting/plot_vtk.py +++ b/packages/pybamm/src/pybamm/plotting/plot_vtk.py @@ -273,12 +273,7 @@ def dynamic_plot(self, show_plot=True): scalar_data = {} for name, pv in zip(self.scalar_names, self.scalar_vars, strict=True): pv.initialise() - if _is_unstructured_spatial_variable(pv): - vals = np.array( - [float(_data_at_time(pv, t).ravel()[0]) for t in self.t_pts] - ) - else: - vals = np.array([float(pv(t).ravel()[0]) for t in self.t_pts]) + vals = np.array([float(pv(t).ravel()[0]) for t in self.t_pts]) scalar_data[name] = vals # --- Layout --- diff --git a/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py b/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py new file mode 100644 index 0000000000..820632a2ec --- /dev/null +++ b/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py @@ -0,0 +1,423 @@ +from types import SimpleNamespace + +import numpy as np +import pytest + +import pybamm +from pybamm.plotting.plot_vtk import ( + VTKQuickPlot, + _build_vtk_grid, + _compute_scale, + _data_at_time, + _is_unstructured_spatial_variable, + _resolve_scale, + _set_cell_scalars, + _set_point_scalars, + _viridis_lut, +) + +vtk = pytest.importorskip("vtk") + + +def _tetra_mesh(): + nodes = np.array( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ] + ) + return pybamm.UnstructuredSubMesh(nodes, np.array([[0, 1, 2, 3]])) + + +def _cell_solution(): + mesh = _tetra_mesh() + model = pybamm.BaseModel() + xyz = [pybamm.SpatialVariable(axis, domain="mesh") for axis in "xyz"] + model._geometry = { + "mesh": { + var: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)} + for var in xyz + } + } + + field = pybamm.StateVector(slice(0, 1), domain="mesh") + field.mesh = mesh + model.variables = {"field": field, "scalar": pybamm.t} + model.update_processed_variables(model.variables) + + t = np.array([0.0, 1.0, 2.0]) + y = np.asfortranarray([[1.0, 2.0, 3.0]]) + return pybamm.Solution(t, y, model, {}), mesh + + +def _triangle_solution(): + mesh = pybamm.UnstructuredSubMesh( + np.array([[0.0, 0.0], [2.0, 0.0], [0.0, 1.0]]), + np.array([[0, 1, 2]]), + ) + model = pybamm.BaseModel() + x = pybamm.SpatialVariable("x", domain="mesh") + z = pybamm.SpatialVariable("z", domain="mesh") + model._geometry = { + "mesh": { + x: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(2)}, + z: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)}, + } + } + field = pybamm.StateVector(slice(0, 1), domain="mesh") + field.mesh = mesh + model.variables = {"field": field} + model.update_processed_variables(model.variables) + solution = pybamm.Solution( + np.array([0.0, 1.0]), np.asfortranarray([[1.0, 2.0]]), model, {} + ) + return solution + + +def _node_solution(): + mesh = SimpleNamespace( + nodes=np.array( + [ + [0.0, 0.0, 0.0], + [1.0, 0.0, 0.0], + [0.0, 1.0, 0.0], + [0.0, 0.0, 1.0], + ] + ), + elements=np.array([[0, 1, 2, 3]]), + dimension=3, + npts=4, + ) + model = pybamm.BaseModel() + xyz = [pybamm.SpatialVariable(axis, domain="mesh") for axis in "xyz"] + model._geometry = { + "mesh": { + var: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)} + for var in xyz + } + } + + field = pybamm.StateVector(slice(0, 4), domain="mesh") + field.mesh = mesh + model.variables = {"node field": field} + model.update_processed_variables(model.variables) + + t = np.array([0.0, 1.0, 2.0]) + y = np.asfortranarray( + [ + [0.0, 1.0, 2.0], + [1.0, 2.0, 3.0], + [2.0, 3.0, 4.0], + [3.0, 4.0, 5.0], + ] + ) + solution = pybamm.Solution(t, y, model, {}) + casadi_field, field, _ = solution._convert_to_casadi(field, {}, y.shape) + solution._variables["node field"] = pybamm.ProcessedVariableUnstructured( + "node field", [field], [casadi_field], solution + ) + return solution, mesh + + +def _first_actor(renderer): + actors = renderer.GetActors() + actors.InitTraversal() + return actors.GetNextActor() + + +class TestVTKHelpers: + @pytest.mark.parametrize( + ("n_vertices", "cell_type"), + [ + (3, vtk.VTK_TRIANGLE), + (4, vtk.VTK_TETRA), + (8, vtk.VTK_HEXAHEDRON), + ], + ) + def test_build_grid_infers_cell_type(self, n_vertices, cell_type): + nodes = np.column_stack( + [ + np.arange(n_vertices, dtype=float), + np.arange(n_vertices, dtype=float) + 1, + np.arange(n_vertices, dtype=float) + 2, + ] + ) + mesh = SimpleNamespace( + nodes=nodes, elements=np.array([np.arange(n_vertices)]) + ) + + grid = _build_vtk_grid(mesh) + + assert grid.GetNumberOfPoints() == n_vertices + assert grid.GetNumberOfCells() == 1 + assert grid.GetCellType(0) == cell_type + np.testing.assert_array_equal( + [grid.GetCell(0).GetPointId(i) for i in range(n_vertices)], + np.arange(n_vertices), + ) + + def test_build_grid_uses_element_type_and_scales_2d_points(self): + mesh = SimpleNamespace( + nodes=np.array([[1.0, 2.0], [3.0, 2.0], [3.0, 4.0], [1.0, 4.0]]), + elements=np.array([[0, 1, 2, 3]]), + element_type="quad", + ) + + grid = _build_vtk_grid(mesh, scale=(2.0, 3.0, 99.0)) + + assert grid.GetCellType(0) == vtk.VTK_QUAD + np.testing.assert_allclose(grid.GetPoint(0), [2.0, 6.0, 0.0]) + np.testing.assert_allclose(grid.GetPoint(2), [6.0, 12.0, 0.0]) + + def test_build_grid_rejects_unknown_connectivity(self): + mesh = SimpleNamespace( + nodes=np.zeros((5, 3)), elements=np.array([[0, 1, 2, 3, 4]]) + ) + + with pytest.raises(ValueError, match="5 vertices per element"): + _build_vtk_grid(mesh) + + def test_scale_options(self): + mesh = SimpleNamespace( + nodes=np.array([[0.0, 2.0, 3.0], [4.0, 2.0, 5.0]]) + ) + + np.testing.assert_allclose(_compute_scale(mesh), [1.0, 1.0, 2.0]) + np.testing.assert_allclose(_resolve_scale("auto", mesh), [1.0, 1.0, 2.0]) + assert _resolve_scale(None, mesh) is None + np.testing.assert_allclose(_resolve_scale((3, 2, 1), mesh), [3, 2, 1]) + + zero_mesh = SimpleNamespace(nodes=np.ones((3, 2))) + np.testing.assert_array_equal(_compute_scale(zero_mesh), [1.0, 1.0]) + + def test_set_and_update_cell_and_point_scalars(self): + grid = _build_vtk_grid(_tetra_mesh()) + + _set_cell_scalars(grid, "cell", [1.25]) + cell_array = grid.GetCellData().GetArray("cell") + assert grid.GetCellData().GetScalars().GetName() == "cell" + assert cell_array.GetNumberOfTuples() == 1 + assert cell_array.GetValue(0) == pytest.approx(1.25) + + _set_cell_scalars(grid, "cell", [3.5]) + assert grid.GetCellData().GetArray("cell") is cell_array + assert cell_array.GetValue(0) == pytest.approx(3.5) + + _set_point_scalars(grid, "point", [0.5, 1.5, 2.5, 3.5]) + point_array = grid.GetPointData().GetArray("point") + assert grid.GetPointData().GetScalars().GetName() == "point" + np.testing.assert_allclose( + [point_array.GetValue(i) for i in range(4)], [0.5, 1.5, 2.5, 3.5] + ) + + _set_point_scalars(grid, "point", [4, 3, 2, 1]) + assert grid.GetPointData().GetArray("point") is point_array + np.testing.assert_allclose( + [point_array.GetValue(i) for i in range(4)], [4, 3, 2, 1] + ) + + def test_processed_variable_helpers(self): + cell_solution, _ = _cell_solution() + cell_variable = cell_solution["field"] + scalar_variable = cell_solution["scalar"] + node_solution, _ = _node_solution() + node_variable = node_solution["node field"] + + assert _is_unstructured_spatial_variable(cell_variable) + assert _is_unstructured_spatial_variable(node_variable) + assert not _is_unstructured_spatial_variable(scalar_variable) + np.testing.assert_allclose(_data_at_time(cell_variable, 0.5), [[1.5]]) + assert _data_at_time(scalar_variable, 0.5) == pytest.approx(0.5) + + def test_viridis_lookup_table(self): + lut = _viridis_lut(-2.0, 4.0, n=8) + + assert lut.GetNumberOfTableValues() == 8 + np.testing.assert_allclose(lut.GetRange(), [-2.0, 4.0]) + assert lut.GetTableValue(0)[3] == pytest.approx(1.0) + assert lut.GetTableValue(7)[3] == pytest.approx(1.0) + assert lut.GetTableValue(0) != lut.GetTableValue(7) + + +class TestVTKQuickPlot: + def test_initialisation_accepts_solution_simulation_and_options(self): + solution, mesh = _cell_solution() + + default_plot = VTKQuickPlot(solution) + assert default_plot.output_variables == ["field"] + assert default_plot.mesh is mesh + assert default_plot.spatial_panels == [ + ("field", {"plot_type": "3d", "scale": "auto"}) + ] + + simulation = pybamm.Simulation(solution.all_models[0]) + simulation._solution = solution + plot = VTKQuickPlot( + simulation, + "field", + options={ + "field": [ + {"plot_type": "3d", "scale": None}, + {"plot_type": "slice", "z": 0.25}, + ] + }, + interpolate_time=True, + ) + assert plot.solution is solution + assert plot.spatial_names == ["field"] + assert plot.scalar_names == [] + assert plot.interpolate_time + assert plot.spatial_panels == [ + ("field", {"plot_type": "3d", "scale": None}), + ( + "field", + {"plot_type": "slice", "scale": "auto", "z": 0.25}, + ), + ] + assert VTKQuickPlot([solution], "scalar").solution is solution + + def test_dynamic_plot_cell_data_slices_scalar_chart_and_snapped_slider(self): + solution, _ = _cell_solution() + plot = VTKQuickPlot( + solution, + ["field", "scalar"], + options={ + "field": [ + {"plot_type": "3d"}, + {"plot_type": "slice", "x": 0.4}, + {"plot_type": "slice", "y": 0.4}, + {"plot_type": "slice", "z": 0.4}, + ] + }, + ) + + plot.dynamic_plot(show_plot=False) + + assert plot._window.GetWindowName() == "PyBaMM - field, scalar" + assert plot._window.GetSize() == (1950, 1040) + assert plot._window.GetRenderers().GetNumberOfItems() == 7 + assert plot._slider.GetEnabled() == 1 + + plot._slider.GetRepresentation().SetValue(1.6) + plot._slider.InvokeEvent("InteractionEvent") + + renderers = plot._window.GetRenderers() + renderers.InitTraversal() + field_renderer = renderers.GetNextItem() + mapped_data = _first_actor(field_renderer).GetMapper().GetInput() + values = mapped_data.GetPointData().GetArray("field") + assert values.GetValue(0) == pytest.approx(3.0) + + def test_dynamic_plot_2d_panels_share_camera(self): + plot = VTKQuickPlot( + _triangle_solution(), + "field", + options={"field": [{"plot_type": "3d"}, {"plot_type": "3d"}]}, + ) + plot.dynamic_plot(show_plot=False) + + renderers = plot._window.GetRenderers() + renderers.InitTraversal() + first = renderers.GetNextItem() + second = renderers.GetNextItem() + assert first.GetActiveCamera() is second.GetActiveCamera() + assert first.GetActiveCamera().GetParallelProjection() == 0 + + def test_dynamic_plot_interpolates_cell_data(self): + solution, _ = _cell_solution() + plot = VTKQuickPlot( + solution, + "field", + options={"field": {"scale": None}}, + interpolate_time=True, + ) + plot.dynamic_plot(show_plot=False) + + plot._slider.GetRepresentation().SetValue(1.25) + plot._slider.InvokeEvent("InteractionEvent") + + renderers = plot._window.GetRenderers() + renderers.InitTraversal() + mapped_data = _first_actor(renderers.GetNextItem()).GetMapper().GetInput() + values = mapped_data.GetPointData().GetArray("field") + assert values.GetValue(0) == pytest.approx(2.25) + + def test_dynamic_plot_interpolates_node_data_and_updates_slice(self): + solution, _ = _node_solution() + plot = VTKQuickPlot( + solution, + "node field", + options={ + "node field": [ + {"plot_type": "3d"}, + {"plot_type": "slice", "x": 0.25}, + ] + }, + interpolate_time=True, + ) + plot.dynamic_plot(show_plot=False) + + plot._slider.GetRepresentation().SetValue(1.25) + plot._slider.InvokeEvent("InteractionEvent") + + renderers = plot._window.GetRenderers() + renderers.InitTraversal() + point_data = _first_actor(renderers.GetNextItem()).GetMapper().GetInput() + values = point_data.GetPointData().GetArray("node field") + np.testing.assert_allclose( + [values.GetValue(i) for i in range(4)], [1.25, 2.25, 3.25, 4.25] + ) + + def test_dynamic_plot_node_data_direct_and_slice_pipelines(self): + solution, _ = _node_solution() + plot = VTKQuickPlot( + solution, + "node field", + options={ + "node field": [ + {"plot_type": "3d", "scale": None}, + {"plot_type": "slice", "z": 0.3, "scale": None}, + ] + }, + ) + plot.dynamic_plot(show_plot=False) + + renderers = plot._window.GetRenderers() + renderers.InitTraversal() + direct_data = _first_actor(renderers.GetNextItem()).GetMapper().GetInput() + point_values = direct_data.GetPointData().GetArray("node field") + np.testing.assert_allclose( + [point_values.GetValue(i) for i in range(4)], [0, 1, 2, 3] + ) + + plot._slider.GetRepresentation().SetValue(2.0) + plot._slider.InvokeEvent("InteractionEvent") + np.testing.assert_allclose( + [point_values.GetValue(i) for i in range(4)], [2, 3, 4, 5] + ) + + def test_dynamic_plot_slice_requires_axis(self): + solution, _ = _cell_solution() + plot = VTKQuickPlot( + solution, "field", options={"field": {"plot_type": "slice"}} + ) + + with pytest.raises(ValueError, match="requires one of 'x', 'y', or 'z'"): + plot.dynamic_plot(show_plot=False) + + def test_save_gif_builds_plot_and_writes_animation(self, tmp_path): + Image = pytest.importorskip("PIL.Image") + solution, _ = _cell_solution() + plot = VTKQuickPlot(solution, "field") + output = tmp_path / "field.gif" + + plot.save_gif(output, fps=5, n_frames=2, width=160, height=100) + plot.save_gif(output, fps=5, n_frames=2, width=160, height=100) + + assert output.stat().st_size > 0 + with Image.open(output) as image: + assert image.size == (160, 100) + assert image.n_frames == 2 + assert image.info["duration"] == 200 From feee7f3f834c2d09e06c35f26bd7fbde0d20ab08 Mon Sep 17 00:00:00 2001 From: "pre-commit-ci[bot]" <66853113+pre-commit-ci[bot]@users.noreply.github.com> Date: Fri, 31 Jul 2026 18:23:05 +0000 Subject: [PATCH 22/25] style: pre-commit fixes --- .../tests/unit/test_plotting/test_plot_vtk.py | 18 ++++-------------- 1 file changed, 4 insertions(+), 14 deletions(-) diff --git a/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py b/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py index 820632a2ec..651c7584e1 100644 --- a/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py +++ b/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py @@ -36,10 +36,7 @@ def _cell_solution(): model = pybamm.BaseModel() xyz = [pybamm.SpatialVariable(axis, domain="mesh") for axis in "xyz"] model._geometry = { - "mesh": { - var: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)} - for var in xyz - } + "mesh": {var: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)} for var in xyz} } field = pybamm.StateVector(slice(0, 1), domain="mesh") @@ -93,10 +90,7 @@ def _node_solution(): model = pybamm.BaseModel() xyz = [pybamm.SpatialVariable(axis, domain="mesh") for axis in "xyz"] model._geometry = { - "mesh": { - var: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)} - for var in xyz - } + "mesh": {var: {"min": pybamm.Scalar(0), "max": pybamm.Scalar(1)} for var in xyz} } field = pybamm.StateVector(slice(0, 4), domain="mesh") @@ -144,9 +138,7 @@ def test_build_grid_infers_cell_type(self, n_vertices, cell_type): np.arange(n_vertices, dtype=float) + 2, ] ) - mesh = SimpleNamespace( - nodes=nodes, elements=np.array([np.arange(n_vertices)]) - ) + mesh = SimpleNamespace(nodes=nodes, elements=np.array([np.arange(n_vertices)])) grid = _build_vtk_grid(mesh) @@ -180,9 +172,7 @@ def test_build_grid_rejects_unknown_connectivity(self): _build_vtk_grid(mesh) def test_scale_options(self): - mesh = SimpleNamespace( - nodes=np.array([[0.0, 2.0, 3.0], [4.0, 2.0, 5.0]]) - ) + mesh = SimpleNamespace(nodes=np.array([[0.0, 2.0, 3.0], [4.0, 2.0, 5.0]])) np.testing.assert_allclose(_compute_scale(mesh), [1.0, 1.0, 2.0]) np.testing.assert_allclose(_resolve_scale("auto", mesh), [1.0, 1.0, 2.0]) From f87a15e545da8b3f6692deb7f0a77cafe19311ac Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Fri, 31 Jul 2026 12:37:55 -0700 Subject: [PATCH 23/25] fix: stabilize domain sizes and VTK headless render Co-authored-by: Cursor --- .../src/pybamm/expression_tree/symbol.py | 56 ++++++------------- packages/pybamm/src/pybamm/meshes/meshes.py | 14 ----- .../pybamm/src/pybamm/plotting/plot_vtk.py | 3 +- .../unit/test_expression_tree/test_symbol.py | 8 +++ 4 files changed, 28 insertions(+), 53 deletions(-) diff --git a/packages/pybamm/src/pybamm/expression_tree/symbol.py b/packages/pybamm/src/pybamm/expression_tree/symbol.py index e7c0942794..733be05796 100644 --- a/packages/pybamm/src/pybamm/expression_tree/symbol.py +++ b/packages/pybamm/src/pybamm/expression_tree/symbol.py @@ -33,61 +33,41 @@ EMPTY_DOMAINS: dict[str, list] = {k: [] for k in DOMAIN_LEVELS} -# Registry of domain → actual mesh size, populated by pybamm.Mesh after the -# submeshes are built. When set, ``domain_size`` returns the real size -# instead of a hash, so a ``pybamm.Vector`` carrying real per-cell values on -# that domain shape-matches a ``Variable`` on the same domain pre- and -# post-discretisation. -_REGISTERED_DOMAIN_SIZES: dict[str, int] = {} - - -def register_domain_size(name: str, size: int) -> None: - """Pin ``domain_size(name)`` to ``size``. - - Called by ``pybamm.Mesh`` for every submesh that exposes ``npts``. Users - rarely need to call this directly; it lets ``pybamm.Vector(values, - domain=name)`` carry real per-cell entries without bespoke shape hacks. - """ - _REGISTERED_DOMAIN_SIZES[name] = int(size) - - -def unregister_domain_size(name: str) -> None: - """Drop a registered domain size (mostly for tests).""" - _REGISTERED_DOMAIN_SIZES.pop(name, None) - - def domain_size(domain: list[str] | str): """ Get the domain size. Empty domain has size 1. - If a domain has been registered via :func:`register_domain_size` (e.g. - by :class:`pybamm.Mesh` after building an unstructured submesh), its - actual mesh ``npts`` is used. Otherwise the standard battery-domain - table is consulted, falling back to a hash-based pseudo-size. + If the domain falls within the list of standard battery domains, the size is read + from a dictionary of standard domain sizes. Otherwise, the hash of the domain string + is used to generate a `random` domain size. """ fixed_domain_sizes = { "current collector": 3, "negative particle": 5, + "negative primary particle": 5, + "negative secondary particle": 5, "positive particle": 7, + "positive primary particle": 7, + "positive secondary particle": 7, "negative electrode": 11, "separator": 13, "positive electrode": 17, "negative particle size": 19, + "negative primary particle size": 19, + "negative secondary particle size": 19, "positive particle size": 23, + "positive primary particle size": 23, + "positive secondary particle size": 23, } if domain in [[], None]: - return 1 - # Fixed battery-domain sentinels take priority — they are stable, hash-like - # values used purely for symbolic shape checks across pybamm tests/models. - if all(dom in fixed_domain_sizes for dom in domain): - return sum(fixed_domain_sizes[dom] for dom in domain) - # Mesh-registered actual sizes — applies to user domains like "cell" that - # carry real per-cell data. - if all(dom in _REGISTERED_DOMAIN_SIZES for dom in domain): - return sum(_REGISTERED_DOMAIN_SIZES[dom] for dom in domain) - # Add 2 per domain to ensure size is always >= 2 and is additive - return sum(2 + hash(dom) % 100 for dom in domain) + size = 1 + elif all(dom in fixed_domain_sizes for dom in domain): + size = sum(fixed_domain_sizes[dom] for dom in domain) + else: + # Add 2 per domain to ensure size is always >= 2 and is additive + size = sum(2 + hash(dom) % 100 for dom in domain) + return size def create_object_of_size(size: int, typ="vector"): diff --git a/packages/pybamm/src/pybamm/meshes/meshes.py b/packages/pybamm/src/pybamm/meshes/meshes.py index 7f66616e54..62cdbb0abf 100644 --- a/packages/pybamm/src/pybamm/meshes/meshes.py +++ b/packages/pybamm/src/pybamm/meshes/meshes.py @@ -170,20 +170,6 @@ def __init__(self, geometry, submesh_types, var_pts): self[domain] = submesh_types[domain](geometry[domain], submesh_pts[domain]) self.base_domains.append(domain) - # Register actual mesh sizes so symbolic shape checks - # (``pybamm.evaluate_for_shape_using_domain``) match the discretised - # vector lengths. Lets ``pybamm.Vector(arr, domain=name)`` shape-match - # ``Variable(domain=name)`` without bespoke subclasses. - for domain, submesh in self.items(): - if isinstance(domain, tuple) and len(domain) == 1: - domain_name = domain[0] - elif isinstance(domain, str): - domain_name = domain - else: - continue - if hasattr(submesh, "npts"): - pybamm.register_domain_size(domain_name, submesh.npts) - # compute interface data for unstructured meshes self._compute_unstructured_interfaces() diff --git a/packages/pybamm/src/pybamm/plotting/plot_vtk.py b/packages/pybamm/src/pybamm/plotting/plot_vtk.py index 3868d00791..251aef87ed 100644 --- a/packages/pybamm/src/pybamm/plotting/plot_vtk.py +++ b/packages/pybamm/src/pybamm/plotting/plot_vtk.py @@ -752,7 +752,8 @@ def on_slider(obj, event): mt_arr.Modified() mtable.Modified() time_text.SetInput(f"t = {t_now:.4g} s") - window.Render() + if show_plot: + window.Render() slider = vtk.vtkSliderWidget() slider.SetInteractor(interactor) diff --git a/packages/pybamm/tests/unit/test_expression_tree/test_symbol.py b/packages/pybamm/tests/unit/test_expression_tree/test_symbol.py index 9e4255dff3..a506db6a6d 100644 --- a/packages/pybamm/tests/unit/test_expression_tree/test_symbol.py +++ b/packages/pybamm/tests/unit/test_expression_tree/test_symbol.py @@ -24,6 +24,14 @@ def test_fixed_domains(self): assert domain_size(["negative electrode"]) == 11 assert domain_size(["separator"]) == 13 assert domain_size(["positive electrode"]) == 17 + assert domain_size(["negative primary particle"]) == 5 + assert domain_size(["negative secondary particle"]) == 5 + assert domain_size(["positive primary particle"]) == 7 + assert domain_size(["positive secondary particle"]) == 7 + assert domain_size(["negative primary particle size"]) == 19 + assert domain_size(["negative secondary particle size"]) == 19 + assert domain_size(["positive primary particle size"]) == 23 + assert domain_size(["positive secondary particle size"]) == 23 def test_fixed_domains_are_additive(self): assert domain_size(["negative electrode", "separator"]) == 11 + 13 From 459dddc9f037d9e3b970bed7b887bac9860a96c1 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Fri, 31 Jul 2026 14:14:36 -0700 Subject: [PATCH 24/25] fix: make VTK offscreen GIF export headless-safe Use OSMesa only on Linux and install the software GL backends in CI so save_gif no longer segfaults on headless runners. Co-authored-by: Cursor --- .github/workflows/_nox.yml | 7 ++++- .../pybamm/src/pybamm/plotting/plot_vtk.py | 28 +++++++++++++++++-- .../tests/unit/test_plotting/test_plot_vtk.py | 10 +++++++ 3 files changed, 42 insertions(+), 3 deletions(-) diff --git a/.github/workflows/_nox.yml b/.github/workflows/_nox.yml index 87112c6779..f85f329958 100644 --- a/.github/workflows/_nox.yml +++ b/.github/workflows/_nox.yml @@ -81,7 +81,7 @@ jobs: uses: awalsh128/cache-apt-pkgs-action@553a35bb8ebd9fcabcb1c9451aa4c98e1b4ca8a9 # v1.6.3 if: startsWith(matrix.leg.os, 'ubuntu') with: - packages: gfortran gcc graphviz pandoc + packages: gfortran gcc graphviz pandoc libosmesa6 libegl1 execute_install_scripts: true # dot -c is for registering graphviz fonts and plugins @@ -92,6 +92,11 @@ jobs: sudo dot -c sudo apt-get install libopenblas-dev + # VTK off-screen GIF export needs a software OpenGL backend on headless runners. + - name: Prefer OSMesa for VTK on Linux + if: startsWith(matrix.leg.os, 'ubuntu') + run: echo "VTK_DEFAULT_OPENGL_WINDOW=vtkOSOpenGLRenderWindow" >> "$GITHUB_ENV" + # Kept separate and opt-out: texlive-latex-extra is large and uncached. - name: Install TeXLive for Linux if: ${{ startsWith(matrix.leg.os, 'ubuntu') && inputs.texlive }} diff --git a/packages/pybamm/src/pybamm/plotting/plot_vtk.py b/packages/pybamm/src/pybamm/plotting/plot_vtk.py index 251aef87ed..f0b9d2a494 100644 --- a/packages/pybamm/src/pybamm/plotting/plot_vtk.py +++ b/packages/pybamm/src/pybamm/plotting/plot_vtk.py @@ -161,6 +161,30 @@ def _viridis_lut(vmin, vmax, n=256): return lut +def _make_render_window(off_screen=False): + """Create a VTK render window. + + Off-screen Linux uses OSMesa (``vtkOSOpenGLRenderWindow``); macOS/Windows + use the platform window with off-screen rendering enabled. Instantiating + the OSMesa window on unsupported platforms segfaults. + """ + import sys + + import vtk + + if ( + off_screen + and sys.platform.startswith("linux") + and hasattr(vtk, "vtkOSOpenGLRenderWindow") + ): + window = vtk.vtkOSOpenGLRenderWindow() + else: + window = vtk.vtkRenderWindow() + if off_screen: + window.SetOffScreenRendering(1) + return window + + class VTKQuickPlot: """Interactive VTK visualization for unstructured mesh solutions. @@ -285,7 +309,7 @@ def dynamic_plot(self, show_plot=True): n_rows = int(np.ceil(n_panels / n_cols)) panel_height = (panel_top - panel_bot) / n_rows - window = vtk.vtkRenderWindow() + window = _make_render_window(off_screen=not show_plot) window.SetSize(650 * n_cols, 520 * n_rows) window.SetWindowName("PyBaMM - " + ", ".join(self.output_variables)) @@ -788,7 +812,7 @@ def save_gif(self, filename, fps=10, n_frames=100, width=1800, height=900): import vtk from PIL import Image - if not hasattr(self, "_window"): + if not hasattr(self, "_window") or not self._window.GetOffScreenRendering(): self.dynamic_plot(show_plot=False) win = self._window diff --git a/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py b/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py index 651c7584e1..aea91169e3 100644 --- a/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py +++ b/packages/pybamm/tests/unit/test_plotting/test_plot_vtk.py @@ -10,6 +10,7 @@ _compute_scale, _data_at_time, _is_unstructured_spatial_variable, + _make_render_window, _resolve_scale, _set_cell_scalars, _set_point_scalars, @@ -230,6 +231,15 @@ def test_viridis_lookup_table(self): assert lut.GetTableValue(7)[3] == pytest.approx(1.0) assert lut.GetTableValue(0) != lut.GetTableValue(7) + def test_make_render_window_offscreen(self): + import sys + + window = _make_render_window(off_screen=True) + + assert window.GetOffScreenRendering() == 1 + if sys.platform.startswith("linux"): + assert isinstance(window, vtk.vtkOSOpenGLRenderWindow) + class TestVTKQuickPlot: def test_initialisation_accepts_solution_simulation_and_options(self): From 8c1560943d247b7d017d1079480cb21dca8cff91 Mon Sep 17 00:00:00 2001 From: Alexander Bills Date: Fri, 31 Jul 2026 14:29:33 -0700 Subject: [PATCH 25/25] ci: install OSMesa for VTK GIF tests on Windows Headless Windows runners lack a software OpenGL backend, so save_gif segfaults; install mesa-dist-win OSMesa via the PyVista setup action. Co-authored-by: Cursor --- .github/workflows/_nox.yml | 9 +++++++++ 1 file changed, 9 insertions(+) diff --git a/.github/workflows/_nox.yml b/.github/workflows/_nox.yml index f85f329958..836536a2f6 100644 --- a/.github/workflows/_nox.yml +++ b/.github/workflows/_nox.yml @@ -120,6 +120,15 @@ jobs: if: startsWith(matrix.leg.os, 'windows') run: winget install --id Graphviz.Graphviz --exact --accept-source-agreements --accept-package-agreements + # VTK save_gif needs OSMesa on headless Windows runners (osmesa.dll on PATH). + - name: Setup headless OpenGL on Windows + if: startsWith(matrix.leg.os, 'windows') + uses: pyvista/setup-headless-display-action@5bc8de3bc71fcda7a96439571287a554901541a0 # v4 + with: + pyvista: "false" + mesa3d-release: "24.3.0" + install-mesa3d-offscreen: "true" + - name: Set up Python ${{ matrix.leg.python }} uses: actions/setup-python@5fda3b95a4ea91299a34e894583c3862153e4b97 # v7.0.0 with: