Files
firedrake-conda/verify.py
2026-02-08 21:04:59 -06:00

277 lines
9.4 KiB
Python
Executable File

#!/usr/bin/env python3
"""Verify that the Firedrake conda environment is correctly configured.
Checks:
1. Environment variables (PETSC_DIR, CC, etc.)
2. Conda-provided libraries (PETSc, MPI, HDF5)
3. Python imports (petsc4py, mpi4py, firedrake, icepack)
4. JIT compilation (the most likely failure point)
5. A minimal PDE solve
Run with:
conda activate firedrake
python verify.py
"""
import os
import sys
import importlib
import ctypes
from pathlib import Path
PASS = "\033[92m✓\033[0m"
FAIL = "\033[91m✗\033[0m"
WARN = "\033[93m!\033[0m"
passed = 0
failed = 0
warned = 0
def check(description, condition, detail=""):
global passed, failed
if condition:
print(f" {PASS} {description}")
passed += 1
return True
else:
msg = f" {FAIL} {description}"
if detail:
msg += f" ({detail})"
print(msg)
failed += 1
return False
def warn_check(description, detail=""):
global warned
msg = f" {WARN} {description}"
if detail:
msg += f" ({detail})"
print(msg)
warned += 1
# ═══════════════════════════════════════════════════════════════════════════
print("\n1. Environment variables")
print("" * 40)
conda_prefix = os.environ.get("CONDA_PREFIX", "")
check("CONDA_PREFIX is set", bool(conda_prefix), conda_prefix or "not set")
petsc_dir = os.environ.get("PETSC_DIR", "")
check(
"PETSC_DIR points to conda prefix",
petsc_dir == conda_prefix,
f"PETSC_DIR={petsc_dir!r}, expected {conda_prefix!r}"
)
petsc_arch = os.environ.get("PETSC_ARCH", "")
check(
"PETSC_ARCH is empty (prefix install)",
petsc_arch == "",
f"PETSC_ARCH={petsc_arch!r}"
)
cc = os.environ.get("CC", "")
check("CC is set to mpicc", cc == "mpicc", f"CC={cc!r}")
hdf5_mpi = os.environ.get("HDF5_MPI", "")
check("HDF5_MPI=ON", hdf5_mpi == "ON", f"HDF5_MPI={hdf5_mpi!r}")
pyop2_cflags = os.environ.get("PYOP2_CFLAGS", "")
check("PYOP2_CFLAGS has -fPIC", "-fPIC" in pyop2_cflags, f"PYOP2_CFLAGS={pyop2_cflags!r}")
pyop2_ldflags = os.environ.get("PYOP2_LDFLAGS", "")
check("PYOP2_LDFLAGS has -shared", "-shared" in pyop2_ldflags, f"PYOP2_LDFLAGS={pyop2_ldflags!r}")
# ═══════════════════════════════════════════════════════════════════════════
print("\n2. Shared libraries")
print("" * 40)
if conda_prefix:
lib_dir = Path(conda_prefix) / "lib"
ext = "dylib" if sys.platform == "darwin" else "so"
for lib_name in ["libpetsc", "libmpi", "libhdf5"]:
lib_path = lib_dir / f"{lib_name}.{ext}"
if not lib_path.exists():
# Try finding any matching file
matches = list(lib_dir.glob(f"{lib_name}*"))
found = len(matches) > 0
detail = str(matches[0]) if found else f"no {lib_name}* in {lib_dir}"
else:
found = True
detail = str(lib_path)
check(f"{lib_name} found", found, detail)
else:
warn_check("Skipping library checks (CONDA_PREFIX not set)")
# ═══════════════════════════════════════════════════════════════════════════
print("\n3. Python imports")
print("" * 40)
modules = [
("numpy", "numpy"),
("mpi4py", "mpi4py"),
("petsc4py", "petsc4py"),
("h5py", "h5py"),
("ufl", "ufl"),
("FIAT", "FIAT"),
("finat", "finat"),
("loopy", "loopy"),
("tsfc", "tsfc"),
("pyop2", "pyop2"),
("firedrake", "firedrake"),
("icepack", "icepack"),
]
imported = {}
for display_name, module_name in modules:
try:
mod = importlib.import_module(module_name)
version = getattr(mod, "__version__", "?")
check(f"import {display_name}", True, f"v{version}")
imported[module_name] = mod
except Exception as e:
check(f"import {display_name}", False, str(e))
# ═══════════════════════════════════════════════════════════════════════════
print("\n4. PETSc configuration")
print("" * 40)
if "petsc4py" in imported:
from petsc4py import PETSc as petsc
# Check that PETSc has key external packages
has_mumps = petsc.Sys.hasExternalPackage("mumps")
check("PETSc has MUMPS", has_mumps)
has_hypre = petsc.Sys.hasExternalPackage("hypre")
check("PETSc has hypre", has_hypre)
has_superlu = petsc.Sys.hasExternalPackage("superlu_dist")
check("PETSc has SuperLU_dist", has_superlu)
# Scalar type
import numpy as np
scalar = petsc.ScalarType
is_real = scalar in (float, np.float64, np.float32)
check("PETSc scalar type is real", is_real, f"scalar={scalar}")
else:
warn_check("Skipping PETSc config checks (import failed)")
# ═══════════════════════════════════════════════════════════════════════════
print("\n5. JIT compilation test")
print("" * 40)
if "firedrake" in imported:
try:
import firedrake
# This triggers the full JIT pipeline:
# UFL form → TSFC → loopy → C code → compile → dlopen
mesh = firedrake.UnitSquareMesh(4, 4)
V = firedrake.FunctionSpace(mesh, "CG", 1)
u = firedrake.TrialFunction(V)
v = firedrake.TestFunction(V)
a = firedrake.inner(firedrake.grad(u), firedrake.grad(v)) * firedrake.dx
A = firedrake.assemble(a)
check("JIT compilation works", True, "assembled a Laplacian matrix")
except Exception as e:
check("JIT compilation works", False, str(e))
# Try to find and display the actual compiler error
import glob
err_files = sorted(
glob.glob("/tmp/pyop2-tempcache-*/**/*.err", recursive=True),
key=lambda f: Path(f).stat().st_mtime,
reverse=True,
)
log_files = sorted(
glob.glob("/tmp/pyop2-tempcache-*/**/*.log", recursive=True),
key=lambda f: Path(f).stat().st_mtime,
reverse=True,
)
for label, files in [("COMPILE ERRORS", err_files), ("COMPILE LOG", log_files)]:
if files:
content = Path(files[0]).read_text().strip()
if content:
print(f"\n ── {label}: {files[0]} ──")
for line in content.splitlines()[-20:]:
print(f" {line}")
else:
warn_check("Skipping JIT test (firedrake import failed)")
# ═══════════════════════════════════════════════════════════════════════════
print("\n6. Minimal PDE solve")
print("" * 40)
if "firedrake" in imported:
try:
import firedrake
from firedrake import *
mesh = UnitSquareMesh(8, 8)
V = FunctionSpace(mesh, "CG", 1)
u = TrialFunction(V)
v = TestFunction(V)
x, y = SpatialCoordinate(mesh)
f = sin(pi * x) * sin(pi * y)
a = inner(grad(u), grad(v)) * dx
L = f * v * dx
bc = DirichletBC(V, 0, "on_boundary")
u_sol = Function(V)
solve(a == L, u_sol, bcs=bc)
# Check the solution is reasonable
import numpy as np
u_data = u_sol.dat.data_ro
check(
"Poisson solve produces valid output",
np.all(np.isfinite(u_data)) and np.max(np.abs(u_data)) > 0,
f"max|u| = {np.max(np.abs(u_data)):.6f}",
)
except Exception as e:
check("Poisson solve", False, str(e))
else:
warn_check("Skipping PDE solve (firedrake import failed)")
# ═══════════════════════════════════════════════════════════════════════════
print("\n7. icepack smoke test")
print("" * 40)
if "icepack" in imported and "firedrake" in imported:
try:
import icepack
import firedrake
# Check that core model classes exist and instantiate
model = icepack.models.IceShelf()
check("icepack.models.IceShelf()", True)
model = icepack.models.IceStream()
check("icepack.models.IceStream()", True)
# Check rate factor function
A = icepack.rate_factor(260.0)
check(
"icepack.rate_factor(260 K)",
A > 0,
f"A = {A:.4e}",
)
except Exception as e:
check("icepack smoke test", False, str(e))
else:
warn_check("Skipping icepack test (import failed)")
# ═══════════════════════════════════════════════════════════════════════════
print("\n" + "" * 50)
print(f"Results: {PASS} {passed} passed {FAIL} {failed} failed {WARN} {warned} warnings")
print("" * 50)
sys.exit(0 if failed == 0 else 1)